From 32f02a59cfb2ff6fd97af6c672d09693f65ff472 Mon Sep 17 00:00:00 2001 From: Lmy271828 Date: Tue, 18 Aug 2026 15:26:00 +0800 Subject: [PATCH 1/2] tools: Omega-QVLA dit_svdquant_v1 pack -> E0M3/UE4M3 converter + format doc + fixture --- .gitignore | 1 + docs/omega_pack_e0m3.md | 386 +++++++++++++++++++++++++++++++ tools/check_omega_e0m3_layer.py | 259 +++++++++++++++++++++ tools/convert_omega_pack_e0m3.py | 178 ++++++++++++++ tools/gen_omega_pack_fixture.py | 117 ++++++++++ 5 files changed, 941 insertions(+) create mode 100644 docs/omega_pack_e0m3.md create mode 100644 tools/check_omega_e0m3_layer.py create mode 100644 tools/convert_omega_pack_e0m3.py create mode 100644 tools/gen_omega_pack_fixture.py diff --git a/.gitignore b/.gitignore index 3b475834..b9b19e66 100644 --- a/.gitignore +++ b/.gitignore @@ -44,3 +44,4 @@ internal-tests/ internal-docs/ notes/ third_party/cutlass +*.pt diff --git a/docs/omega_pack_e0m3.md b/docs/omega_pack_e0m3.md new file mode 100644 index 00000000..0ced2998 --- /dev/null +++ b/docs/omega_pack_e0m3.md @@ -0,0 +1,386 @@ +# Omega-QVLA pack format and the E0M3 consumption contract + +## 0. Background concepts (90 seconds) + +- **4-bit quantization**: store `round(x / s)` (a small integer) plus the + "ruler" `s` (the scale) instead of `x`. Compute happens as + `integer × s`. Fewer bits = less memory bandwidth, more rounding error. +- **fake-quant**: quantize then *immediately dequantize*, staying in float. + Simulates quantization error without needing integer hardware — Omega's + whole runtime is fake-quant emulation on plain PyTorch matmuls. +- **scale granularity**: how many elements share one ruler. *Per-channel* = + each of the K input channels gets its own (Omega's choice). *Per-16 + block* = 16 adjacent elements share one (the hardware format's choice). + Coarser granularity = fewer scales to store, but elements of different + magnitudes get crushed under a shared ruler. +- **static vs. dynamic scale**: *static* = measured offline on calibration + data, stored in the pack (Omega's `act_scale_table`). *Dynamic* = + computed per token at runtime from the actual data (`amax / 7`). Dynamic + is fresher but constrains the layout to what hardware computes cheaply. +- **E0M3**: the 4-bit element format here — sign + 3 mantissa-ish bits + decoding the uniform integer grid −7..+7. "Uniform" = evenly spaced + levels, unlike E2M1 (NVFP4) whose levels bunch near zero. +- **UE4M3**: an unsigned 4-exponent/3-mantissa mini-float used *only for + scales* (the ruler itself is quantized too). Per-16 scales on both + operands are UE4M3. +- **packed + SFA/SFB**: 4-bit elements are stored two per byte ("packed"); + the per-16 scales live in a separate buffer in CUTLASS's tile-interleaved + layout (SFA for the activation operand, SFB for the weight operand). +- **DuQuant rotation / permutation**: a learned orthogonal transform + (64×64 blocks + a channel shuffle) applied to activations before + quantization. Its job: even out channel magnitudes so no single outlier + channel dominates a shared scale. Orthogonal = length- and + angle-preserving, so it is mathematically free. +- **tcgen05 MMA**: the SM100/SM110 tensor-core instruction that consumes + packed 4-bit operands + UE4M3 scales directly in hardware. This is the + payoff: Omega's math runs as emulation today; this instruction makes it + native. +- **cosine similarity**: the fidelity metric. 1.0 = identical direction; + per-token cos 0.98 means the quantized output vector points in nearly + the same direction as the reference, with ~2% orthogonal noise. + +With those nine, every section below should read top to bottom without +external references. + +Status: recon complete, converter/harness in `tools/` (Milestone 1). +Scope: `packs_hf/pi05_long/quantized.pt` (4.8 GB, pi0.5 LIBERO-10 recipe +`paligemma=svdh+gptq, expert=svdh+rtn+perstep`). Other Omega packs share the +`dit_svdquant_v1` record format but were not inspected. + +## 1. Container + +Plain `torch.save` dict, loadable with `weights_only=True` (no custom +classes). 253 top-level keys: + +- 126 expert records: + `paligemma_with_expert.gemma_expert.model.layers.{0..17}.{self_attn.{q,k,v,o}_proj,mlp.{gate,up,down}_proj}` +- 126 PaliGemma records: + `paligemma_with_expert.paligemma.model.language_model.layers.{0..17}.` +- `__meta__`: `{"recipe": str, "suite": "10", "fresh": bool}` + +Small projections (state_proj, action_in/out_proj, time_mlp) are +deliberately absent — they break under A4 and stay BF16 at runtime. + +## 2. Record schema (`format == "dit_svdquant_v1"`) + +| field | shape / dtype | meaning | +|---|---|---| +| `weight_res_q` | `(out, in)` fp16 | fake-quantized-then-dequantized weight, **already in the rotated + permuted domain** | +| `lowrank_A` / `lowrank_B` | `(out, 0)` / `(in, 0)` fp16 | SVDQuant low-rank branch; rank = 0 in this pack (INT4-only path) | +| `act_scale_table` | `(num_steps, in)` fp32 | per-denoise-step, per-channel activation scales. Expert: `num_steps = 10`; PaliGemma: `1` | +| `duquant_rotation_blocks` | `(in/64, 64, 64)` fp16 | block-diagonal input rotation R_in | +| `duquant_rotation_perm` | `(in,)` int64 | input-channel permutation (applied before R_in) | +| `duquant_rotation_out_blocks` | `(out/64, 64, 64)` fp16 | block-diagonal output rotation (restore) | +| `weight_bits` / `a_bits` | int | 4 / 4 (both sides). The *runtime* `DuQuantLinear` path defaults activations to A8 (`GR00T_DUQUANT_ABITS`), which is where the "W4A8 PaliGemma" label comes from; pack records consumed through `GptqLinear` use their own `a_bits` (4) | +| `in_features` / `out_features` | int | redundant with tensor shapes | +| `n_calib_*`, `act_percentile`, `gptq_damp_percent` | scalars | calibration provenance | + +Notably absent (vs. a classic GPTQ pack): no packed int4 bitstream, no +`qweight`/`qzeros`/group scales — the weight survives only as dequantized +fp16 on the 4-bit grid (~8k unique values per tensor). No `smooth_scale`. + +## 3. Consumer math (Omega `gr00t/quantization/gptq_layers.py`, verified) + +``` +x2 = bmm(x[..., perm].view(N, in/64, 64), R_in_blocks) # input rotation, runtime +x_q = clamp(round(x2 / s_t), -8, 7) * s_t # s_t = act_scale_table[step] +y' = x_q @ W_res_q^T # bf16-promoted accumulate +y = bmm(y'.view(N, out/64, 64), R_out_blocks) + bias # output rotation restore +``` + +PaliGemma (`duquant_layers.py`) is identical in structure with A8 +activations and a single-row scale table. + +Consequences for a FlashRT consumer: + +- The input rotation **cannot** be folded into `weight_res_q`: fake-quant + sits between rotation and GEMM. It must run on activations (torch bmm, or + a prologue kernel). Same for the output restore. +- The rotation is an exact orthonormal transform, so it does not by itself + affect GEMM fidelity; fidelity questions live entirely in the quantizers. +- `weight_res_q` being plain fp16 means the converter re-quantizes from + fp16 — no GPTQ bitstream decoding needed. + +## 4. Mapping to the FlashRT E0M3 contract + +FlashRT SM110 path (`csrc/gemm/fp4/cutlass_fp4_gemm_e0m3w_sm100.cuh`, +bindings in `csrc/fp4_bindings.cpp`): + +- Weights: fp16 `[N, K]` → `quantize_e0m3_dynamic_sfa_fp16(..., is_sfb=True)` + → packed E0M3 `[N, K/2]` + SFB tile-interleaved UE4M3 (per-16, amax/7). +- Activations: same kernel with `is_sfb=False` → packed + SFA. +- GEMM: `cutlass_fp4_gemm_e0m3w(A, SFA, B, SFB, D, M, N, K, α, β, stream, + a_format)` with `a_format=0` for E0M3 activations (1 = E2M1). +- Buffer sizing: `flash_rt_fp4.sfa_size_bytes(N, K, is_sfb)`; scale buffers + must be zero-initialized (tile-interleave pads K to 64-element atoms; + garbage padding decodes as UE4M3 NaN). + +Grid differences vs. Omega fake-quant: + +| | Omega A4 | FlashRT E0M3 | +|---|---|---| +| element grid | int `[-8, 7]` (asymmetric clamp) | sign-magnitude uniform `[-7, 7]` | +| scale | static calibrated, **per-channel** fp32 | dynamic amax/7, **per-16** UE4M3 | +| weight grid | int4 per-channel-group (already dequantized) | per-16 UE4M3 | + +The scale-granularity mismatch (per-channel static table vs. per-16 dynamic) +is the one real fidelity risk. Two candidate strategies, both implemented in +`tools/check_omega_e0m3_layer.py`: + +- **S0 (drop the table)**: `A = e0m3(x2)`, `B = e0m3(W)`. Loses all + calibration information. +- **S1 (fold step-mean table into W, per-step residual into A)**: + `A = e0m3(x2 / s_t)` per step, `B = e0m3(W · diag(s̄))` once, where + `s̄ = mean_t(s_t)`. Exact for the mean step; residual error scales with + the table's step-to-step spread (measured: std/mean ≈ 10% on expert + layer-0 q_proj). + + Mathematically S1 relies on `Σ_k q_k s_k W_nk = Σ_k q_k (s_k W_nk)`: + a per-K-column scale commutes into the weight. A true per-step fold would + need 10 weight copies (unacceptable), hence the mean fold. + +RHT (per-16 Hadamard, `use_rht=1` variants) is orthogonal to the DuQuant +rotation — `(x2·H)(W·H)^T = x2·W^T` — and can be ablated on top of either +strategy if per-block distributions remain problematic. + +### Measured + +Emulation mode (torch, synthetic activations calibrated to q999 = 7·s_t, +M = 256 tokens), per-token cosine vs. the unquantized-activation reference: + +| layer | omega vs fp | S0 vs fp | S1 vs fp | +|---|---|---|---| +| expert L0 q_proj (K=1024) | 0.9927 | 0.9834 | 0.9052 | +| expert L0 down_proj (K=4096) | 0.9928 | 0.9825 | 0.9799 | +| expert L11 o_proj (K=2048) | 0.9929 | 0.9825 | 0.9615 | +| paligemma L0 gate_proj (K=2048) | 0.9928 | 0.9810 | 0.9775 | + +Kernel mode (real tcgen05 GEMM, Thor SM110, same seed, per-token mean +cosine vs. the unquantized-activation reference): + +| layer | omega vs fp | S0 vs fp | S1 vs fp | +|---|---|---|---| +| expert L0 q_proj (K=1024) | 0.99247 | **0.99322** | 0.158 | +| expert L0 down_proj (K=4096) | 0.99274 | **0.99257** | 0.98999 | +| expert L11 o_proj (K=2048) | 0.99275 | **0.99303** | 0.85557 | +| paligemma L0 gate_proj (K=2048) | 0.99206 | **0.99239** | 0.98204 | + +Three findings: + +1. **S0 is lossless on real hardware on every layer tested** — within + ±0.001 of Omega's own fake-quant everywhere (the tiny edges come from + dynamic per-token amax beating a static table on data calibrated only + at the q999 point). Error independence holds wherever checked: + cos(S0, fp)·cos(omega, fp) ≈ measured cos(S0, omega), i.e. S0's + residual is fresh rounding noise, not a systematic shift. +2. **S1's collapse on real hardware is scale-magnitude-dependent.** + Mechanism: `W · diag(s̄)` shrinks weights by the mean table value, + pushing per-16 block scales toward the UE4M3 subnormal floor (2⁻⁹), + where scale mantissas disintegrate and whole blocks quantize to + garbage. Layers with small s̄ die hard (q_proj 0.16, o_proj 0.86); + layers whose table happens to be larger merely degrade (down_proj + 0.99 — still worse than S0). The emulator's lenient subnormal + handling masked the severe cases. +3. The pure-torch references reproduce across machines to 5 decimal + places (0.992700 Thor vs. 0.992707 x86), cross-validating the harness. + +**S0 wins; S1 is dead.** The table's per-channel scale spread (~4×) +distorts weights when folded, while S0's per-token dynamic per-16 amax is +a *better* quantizer than Omega's static per-channel table — the DuQuant +rotation+perm has already whitened per-channel magnitudes, so the table +is only a second-order correction. + +Decision: **the converter emits S0 (`--fold none`) as the production +format**; `--fold mean` is kept for ablation only. This also shrinks the +runtime story — no per-step scale dispatch is needed on the E0M3 path. + +**Follow-up: `actnorm` (floor-safe S1) — also dead (2026-08-18).** A +reviewer-natural fix for S1's floor problem is to normalize before +folding: decompose `s̄ = c·r̄` with `c = geomean(s̄)`, fold only `r̄` +(O(1), geomean 1) into the weights, divide activations by `s̄` at +runtime, and absorb `c` into the GEMM alpha. This is exactly +`(x/s̄) @ (W·r̄)^T · c = x @ W^T`, and it does fix the floor (0% of +block scales below 2⁻⁹ vs 100% for raw S1 on q_proj). But measured on +Thor (consumer-level, fp16 reference, real pack): + +| layer | S0 vs fp16 | actnorm vs fp16 | +|---|---|---| +| q_proj | 0.9935 | 0.9886 | +| down_proj | 0.9929 | 0.9869 | +| o_proj | 0.9936 | 0.9885 | + +actnorm is *worse* than S0 everywhere. Mechanism: the fold is a zero-sum +redistribution — dividing activations by `s̄` whitens the activation +blocks, but multiplying weights by `r̄` (range 0.43–2.68 on q_proj) +re-opens intra-block magnitude spread on the weight side, where per-16 +single-scale 4-bit pays for it. DuQuant's rotation had already whitened +both operands; any per-channel re-scaling of either side undoes that. +**Per-channel calibration tables are fundamentally incompatible with +per-16 block quantization — the information has to live on one side and +always de-whitens it.** Dynamic per-16 amax is the optimum at this +granularity; S0 is the endpoint, not a compromise. (`--fold actnorm` + +consumer support remain in the tree, `OMEGA_E0M3_ACT_TABLE=0`/artifact +driven, as the documented ablation.) The residual end-to-end gap vs. +the full Omega recipe (90.4% vs 93.2%, concentrated in task9) is not +recoverable by table injection; remaining options are mixed precision +for sensitive layers or acceptance. + +Remaining caveats: synthetic activations (lognormal + outlier channels, +calibrated only at the q999 point) — real activation tails differ. Next: +captured real activations, then LIBERO paired SR (Milestone 2). + +## 5. Roadmap and Milestone-1 deliverables + +**Milestone 1 — offline converter + format doc + single-layer gates +(done).** Deliverables below; acceptance: 252/252 records converted, S0 +per-token cosine ≥ Omega fake-quant on real hardware (4/4 layers, §4). + +**Milestone 2 — runtime consumption (next).** Wire the converted pack +into the pi0.5 Thor pipeline: load `packed`/`sfb` as decoder GEMM +operands, run the DuQuant input rotation (perm + 64×64 block bmm) and +output restore around each replaced Linear, keep the small projections +(state/action/time) BF16 from the checkpoint. Acceptance: end-to-end +action cosine vs. the Omega fake-quant server, then LIBERO-10 ×500 +paired SR vs. the BF16 baseline (target: no measurable loss, matching +the pack's own 93.2% vs 91.6%). + +**Milestone 3 — upstream PRs (after M2 evidence).** Split per +`CONTRIBUTING.fork.md` §6, each with LIBERO paired SR + action cosine + +p50/p95 latency: ① converter + this format doc (pure additive, easiest); +② runtime wiring as a flag-gated `weight_format` branch (the S0 result +shrunk this from the originally-planned per-step scale path); ③ SVDQuant +low-rank epilogue (deferred — rank = 0 in this pack). + +**Milestone 2 status (done, incl. 2c/2d landed after the original +write-up):** + +- **M2a/b — consumer + serving (done).** `tools/omega_e0m3_linear.py` + (`OmegaE0M3Linear`, drop-in for gr00t's `GptqLinear` via + `tools/serve_omega_e0m3.py` monkeypatch) + `tools/check_omega_e0m3_consumer.py` + gate: per-layer cosine 0.978–0.982 vs. GptqLinear, 1.2× layer latency. + Server smoke 10/10 ≙ arm D; **LIBERO-10 ×500 paired: 90.4%** vs. BF16 + 91.6% (McNemar p = 0.53) and vs. fake-quant arm D 93.2% (p = 0.070) — + no significant loss; 58 s/episode vs. 148 s fake-quant (2.6×). Eager + mode with `torch.compile` disabled: the pybind kernels graph-break and + the HF KV cache recompiles per step (~25 min/episode stall) — see the + env flags in `serve_omega_e0m3.py`. +- **M2d — hand-rolled CUDA graph over the whole denoise loop (done).** + `tools/omega_e0m3_graph.py` captures all 10 flow-matching steps + (unrolled, pi05_thor style) into one `torch.cuda.CUDAGraph`: static KV + slabs behind a `DynamicCache` shell, static mask/position buffers + filled by `copy_` per inference, adaRMS conditioning precomputed for + the deterministic time grid. The eager blockers removed are documented + in the module docstring (device-scalar `while`, per-step H2D mask + upload, per-call KV allocation). Enable with `OMEGA_E0M3_CUDA_GRAPH=1` + (see `tools/start_e0m3_server.sh`); falls back to eager permanently on + any capture failure. Thor validation: capture succeeds + (`prefix_len=968, layers=18, steps=10`), smoke 10/10, ~43–50 + s/episode vs. ~58 s eager. +- **M2e — PaliGemma E0M3 (route A: official all-W4A4 pack recipe).** The + converter already emits all 252 records, PaliGemma included; serving + with `OMEGA_E0M3_PATCH_DUQUANT=1` (the default in + `tools/start_e0m3_server.sh`) substitutes the runtime `DuQuantLinear` + wraps with `OmegaE0M3Linear` consumers built from the pack's PaliGemma + records — GPTQ W4A4 weights instead of runtime RTN, single-row scale + table dropped per the S0 decision. With `omega_e0m3_graph.py`'s prefix + graph (default on), the prefix prefill is captured too, so the E0M3 + pybind kernels run inside a CUDA graph on this path as well — the + capture smoke (`check_omega_e0m3_graph_smoke.py`) covers their + capturability. + Validation ladder on Thor: artifact coverage check (252 records) → + per-layer consumer gate on PaliGemma layers → 10-episode smoke → + LIBERO-10 ×500 paired SR + per-episode latency. + +Deliverables: + +- `tools/convert_omega_pack_e0m3.py` — offline pack → E0M3 converter + (S0 weight emission + aux tensors: perm, R_in/R_out blocks, + act_scale_table). Runs where `flash_rt_fp4` is built (Thor). +- `tools/check_omega_e0m3_layer.py` — single-layer cosine harness: + Omega fake-quant reference vs. FlashRT E0M3 GEMM (S0/S1), plus a pure + torch emulation mode that runs without the extension for pre-checks. + Emulation results and the S0 decision are in §4. +- `tools/omega_e0m3_linear.py` — `OmegaE0M3Linear` consumer (M2a). +- `tools/check_omega_e0m3_consumer.py` — consumer-vs-GptqLinear gate + (cosine + layer latency). +- `tools/serve_omega_e0m3.py` — openpi serving entry (monkeypatches + gr00t's wrap classes; env-gated compile kill switch). +- `tools/omega_e0m3_graph.py` — CUDA-graph capture of the 10-step + denoise loop (M2d), `OMEGA_E0M3_CUDA_GRAPH=1`. +- `tools/check_omega_e0m3_graph_smoke.py` — P0 capture gate for a single + consumer layer (pybind capturability check). +- `tools/start_e0m3_server.sh` — Thor server launcher (repo-relative + paths, env overrides). + +Deferred: SVDQuant low-rank epilogue (rank = 0 everywhere in this pack), +per-step weight tables (10× memory; also refuted by the S0 result), +per-step activation scale dispatch (refuted by the S0 result). + +### Reproducing + +**Self-contained fixture round-trip (no Omega-QVLA pack or checkout +needed — the PR review path):** + +```bash +cd third_party/flashrt +# 1. Synthetic miniature pack: schema-identical records, random +# orthogonal rotations, outlier-channel weights (pure CPU, seconds) +python tools/gen_omega_pack_fixture.py --out /tmp/fixture_pack.pt +# 2. Convert (Thor) +python tools/convert_omega_pack_e0m3.py --pack /tmp/fixture_pack.pt \ + --out /tmp/fixture_e0m3.pt --fold none +# 3. Consumer vs fp16 reference, gr00t-free (Thor) +PYTHONPATH=$PWD/tools python tools/check_omega_e0m3_consumer.py \ + --reference fp16 --pack /tmp/fixture_pack.pt \ + --artifact /tmp/fixture_e0m3.pt +``` + +**Full pack (development path):** + +```bash +# Point at an Omega pack (any machine for emulate, Thor for kernel/convert) +export OMEGA_PACK=/path/to/Omega-QVLA/packs_hf/pi05_long/quantized.pt +cd third_party/flashrt + +# 1. Local pre-check, no extension needed (pure torch, CPU is fine) +python tools/check_omega_e0m3_layer.py --pack "$OMEGA_PACK" --mode emulate + +# 2. Hardware check — real tcgen05 GEMM (Thor, flash_rt_fp4 built) +python tools/check_omega_e0m3_layer.py --pack "$OMEGA_PACK" --mode kernel +``` + +```bash +# 2. Hardware check output +layer: paligemma_with_expert.gemma_expert.model.layers.0.self_attn.q_proj N(out)=2048 K(in)=1024 table=(10,1024) step=0 + +references:... +``` +```bash +python tools/check_omega_e0m3_layer.py --pack "$OMEGA_PACK" --mode kernel \ + --layer paligemma_with_expert.gemma_expert.model.layers.0.mlp.down_proj + # --layer paligemma_with_expert.paligemma.model.language_model.layers.0.mlp.gate_proj + # --layer paligemma_with_expert.gemma_expert.model.layers.11.self_attn.o_proj +# 3. Full conversion (252 layers, ~1.3 GB output) +python tools/convert_omega_pack_e0m3.py \ + --pack "$OMEGA_PACK" --out pi05_long_e0m3.pt --fold none +``` + +Gate for accepting the conversion: per-token cosine of S0 vs. fp ≥ +Omega's own fake-quant (per-layer, same seed). Currently met on every +layer tested (see §4). + +## 6. Accuracy context (pi0.5 LIBERO-10, 500 episodes) + +The 93.2% figure was measured on the hybrid deployment: expert records from +this pack (W4A4, `GptqLinear`) + PaliGemma via the *runtime* DuQuant path +(W4 weights, A8 activations by the `GR00T_DUQUANT_ABITS` default) — vs. BF16 +baseline 91.6% (McNemar p = 0.32, no significant difference) on the Omega +PyTorch fake-quant path. The pack itself is the official Omega recipe, +which is W4A4 on both sides (PaliGemma records carry `a_bits=4` and a +single-row `act_scale_table`); consuming the PaliGemma records through the +E0M3 consumer (`OMEGA_E0M3_PATCH_DUQUANT=1`) therefore *is* the official +recipe, and additionally replaces runtime RTN weights with the pack's GPTQ +weights. The E0M3 migration target +is therefore "no measurable SR loss against an already lossless baseline" — +the single-layer cosine gates are the leading indicator, LIBERO the final +one. diff --git a/tools/check_omega_e0m3_layer.py b/tools/check_omega_e0m3_layer.py new file mode 100644 index 00000000..94177cd2 --- /dev/null +++ b/tools/check_omega_e0m3_layer.py @@ -0,0 +1,259 @@ +#!/usr/bin/env python3 +"""Single-layer cosine harness: Omega fake-quant vs. FlashRT E0M3 GEMM. + +Quantifies the fidelity cost of migrating one Omega-QVLA dit_svdquant_v1 +record (docs/omega_pack_e0m3.md) onto the FlashRT SM110 E0M3 path. + +References (torch, fp32 accumulate): + y_fp = x2 @ W^T — no activation quant (ceiling; + W is already fake-quant dequant) + y_omega = fakequant(x2, s_t) @ W^T — exact Omega consumer semantics + (per-channel table, int [-8,7]) + +Variants: + S0: A = e0m3(x2), B = e0m3(W) — drops the scale table + S1: A = e0m3(x2 / s_t), B = e0m3(W * diag(s_mean)) — mean-fold strategy + +x2 is the rotated activation (perm + 64x64 block rotation), computed in +fp16 exactly like the Omega runtime. The output rotation is an exact +orthonormal transform applied to BOTH compared vectors, so it is skipped +(cosines are invariant to it). + +Modes: + --mode emulate : pure-torch E0M3 emulation (per-16 amax/7, UE4M3-rounded + scales, int [-7,7]). Runs anywhere; approximates the + tcgen05 result up to accumulation order and UE4M3 + rounding corner cases. Use for local pre-checks. + --mode kernel : real kernels via flash_rt_fp4 (quantize + GEMM, + a_format=0). Requires CUDA + built extension (Thor). + +Usage: + python tools/check_omega_e0m3_layer.py \ + --pack packs_hf/pi05_long/quantized.pt --mode emulate + python tools/check_omega_e0m3_layer.py \ + --pack packs_hf/pi05_long/quantized.pt --mode kernel \ + --layer paligemma_with_expert.gemma_expert.model.layers.0.mlp.down_proj +""" + +from __future__ import annotations + +import argparse +import sys + +import torch + +DEFAULT_LAYER = ("paligemma_with_expert.gemma_expert.model.layers.0." + "self_attn.q_proj") + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + p.add_argument("--pack", required=True) + p.add_argument("--layer", default=DEFAULT_LAYER) + p.add_argument("--mode", choices=("emulate", "kernel"), default="emulate") + p.add_argument("--tokens", type=int, default=256, help="M (rows)") + p.add_argument("--step", type=int, default=0, + help="denoise step index into act_scale_table") + p.add_argument("--seed", type=int, default=0) + p.add_argument("--device", default="", + help="torch device for reference/emulation math " + "(default: cuda in kernel mode, cpu in emulate mode)") + return p.parse_args() + + +# ──────────────────────────────────────────────────────────────────── +# Omega consumer semantics (mirror gr00t/quantization/gptq_layers.py) +# ──────────────────────────────────────────────────────────────────── +def apply_input_rotation(x: torch.Tensor, perm: torch.Tensor, + blocks: torch.Tensor) -> torch.Tensor: + """x[N, in] fp16 -> bmm(x[:, perm].view(N, nb, B), blocks). fp16 in/out.""" + nb, b, _ = blocks.shape + dev = x.device + x2 = x.index_select(dim=-1, index=perm.to(dev)) + x2 = x2.reshape(-1, nb, b) + x2 = torch.bmm(x2.transpose(0, 1).contiguous(), + blocks.to(device=dev, dtype=x.dtype)) + return x2.transpose(0, 1).contiguous().reshape(x.shape[0], nb * b) + + +def fake_quant_omega(x: torch.Tensor, scale: torch.Tensor, + bits: int = 4) -> torch.Tensor: + """clamp(round(x/s), -2^(b-1), 2^(b-1)-1) * s — Omega's asymmetric grid.""" + qmax = 2 ** (bits - 1) - 1 + return (torch.clamp(torch.round(x / scale), -qmax - 1, qmax) + * scale).float() + + +# ──────────────────────────────────────────────────────────────────── +# E0M3 emulation (approximates quantize_e0m3_dynamic_sfa_fp16) +# ──────────────────────────────────────────────────────────────────── +def ue4m3_round(x: torch.Tensor) -> torch.Tensor: + """Round positive tensor to the nearest UE4M3 value (E4M3 without sign, + exp bias 7, 3 mantissa bits, subnormals below 2^-6). Approximation for + emulation mode; the kernel's exact rounding may differ at bin edges.""" + x = x.clamp_min(1e-12) + e = torch.floor(torch.log2(x)) + # normals: value = 2^E * (1 + M/8), M in 0..7 + base = torch.pow(2.0, e) + m = torch.round(x / base - 1.0).clamp(0, 8) + overflow = m == 8 + e = e + overflow.float() + m = m * (~overflow).float() + normal = torch.pow(2.0, e) * (1.0 + m / 8.0) + # subnormals: step 2^-9 + sub = torch.round(x / 2 ** -9) * 2 ** -9 + return torch.where(x < 2 ** -6, sub, normal).clamp(max=480.0) + + +def e0m3_emulate(t: torch.Tensor) -> torch.Tensor: + """Per-16 dynamic E0M3 fake-quant along the last dim. Returns fp32 + dequantized tensor with the same shape.""" + *lead, d = t.shape + assert d % 16 == 0 + v = t.float().reshape(-1, d // 16, 16) + scale = ue4m3_round(v.abs().amax(dim=-1, keepdim=True) / 7.0) + # all-zero blocks round to scale 0; clamp to the smallest UE4M3 + # subnormal so emulation never divides by zero (kernel writes a real + # scale here, exact corner behavior is hardware-specific) + scale = scale.clamp_min(2 ** -9) + q = torch.clamp(torch.round(v / scale), -7, 7) + return (q * scale).reshape(*lead, d) + + +# ──────────────────────────────────────────────────────────────────── +# Kernel path (requires flash_rt_fp4, i.e. Thor) +# ──────────────────────────────────────────────────────────────────── +def e0m3_kernel_gemm(a_fp16: torch.Tensor, b_fp16: torch.Tensor, + fvk_fp4) -> torch.Tensor: + """Quantize A[M,K] and B[N,K] with the real E0M3 kernels and run the + tcgen05 block-scaled GEMM (a_format=0). Returns fp16 D[M, N].""" + a_fp16 = a_fp16.contiguous() + b_fp16 = b_fp16.contiguous() + m, k = a_fp16.shape + n, kb = b_fp16.shape + assert k == kb and k % 16 == 0 + + a_packed = torch.empty(m, k // 2, dtype=torch.uint8, device="cuda") + a_sfa = torch.zeros(fvk_fp4.sfa_size_bytes(m, k, False), + dtype=torch.uint8, device="cuda") + rc = fvk_fp4.quantize_e0m3_dynamic_sfa_fp16( + a_fp16.data_ptr(), a_packed.data_ptr(), a_sfa.data_ptr(), + m, k, False, 0) + if rc != 0: + raise RuntimeError(f"A quantize failed rc={rc}") + + b_packed = torch.empty(n, k // 2, dtype=torch.uint8, device="cuda") + b_sfb = torch.zeros(fvk_fp4.sfa_size_bytes(n, k, True), + dtype=torch.uint8, device="cuda") + rc = fvk_fp4.quantize_e0m3_dynamic_sfa_fp16( + b_fp16.data_ptr(), b_packed.data_ptr(), b_sfb.data_ptr(), + n, k, True, 0) + if rc != 0: + raise RuntimeError(f"B quantize failed rc={rc}") + + d = torch.empty(m, n, dtype=torch.float16, device="cuda") + rc = fvk_fp4.cutlass_fp4_gemm_e0m3w( + a_packed.data_ptr(), a_sfa.data_ptr(), + b_packed.data_ptr(), b_sfb.data_ptr(), d.data_ptr(), + m, n, k, 1.0, 0.0, 0, 0) + if rc != 0: + raise RuntimeError(f"cutlass_fp4_gemm_e0m3w failed rc={rc:#x}") + torch.cuda.synchronize() + return d + + +# ──────────────────────────────────────────────────────────────────── +def cosine_stats(a: torch.Tensor, b: torch.Tensor) -> tuple: + """(global cos, per-row cos mean, per-row cos min), fp32 inputs.""" + a = a.float() + b = b.float() + glob = torch.dot(a.flatten(), b.flatten()) / ( + a.norm() * b.norm()).item() + per = torch.nn.functional.cosine_similarity(a, b, dim=-1) + return glob, per.mean().item(), per.min().item() + + +def report(tag: str, a: torch.Tensor, b: torch.Tensor) -> None: + g, mean, mn = cosine_stats(a, b) + print(f" {tag:<28} global {g:.6f} per-token mean {mean:.6f} " + f"min {mn:.6f}") + + +def main() -> int: + args = parse_args() + pack = torch.load(args.pack, map_location="cpu", weights_only=True) + if args.layer not in pack: + print(f"error: layer '{args.layer}' not in pack", file=sys.stderr) + return 2 + rec = pack[args.layer] + table = rec["act_scale_table"].float() + if not 0 <= args.step < table.shape[0]: + print(f"error: --step {args.step} out of range " + f"(table has {table.shape[0]} steps)", file=sys.stderr) + return 2 + s_t = table[args.step] + s_mean = table.mean(dim=0) + + w = rec["weight_res_q"].float() # (out, in), rotated+permuted domain + out_f, in_f = w.shape + print(f"layer: {args.layer} N(out)={out_f} K(in)={in_f} " + f"table=({table.shape[0]},{table.shape[1]}) step={args.step}") + + # Synthetic activations with realistic per-channel heterogeneity: + # lognormal gains plus a few strong outlier channels (DuQuant's target). + g = torch.Generator().manual_seed(args.seed) + gains = torch.exp(torch.randn(in_f, generator=g)) + outlier_idx = torch.randperm(in_f, generator=g)[: in_f // 128 + 1] + gains[outlier_idx] *= 10.0 + x = (torch.randn(args.tokens, in_f, generator=g) * gains).half() + + dev = args.device or ("cuda" if args.mode == "kernel" else "cpu") + x = x.to(dev) + x2 = apply_input_rotation( + x, rec["duquant_rotation_perm"], + rec["duquant_rotation_blocks"].to(dev)) + w = w.to(dev) + s_t = s_t.to(dev) + s_mean = s_mean.to(dev) + + # Calibrate synthetic activations to the pack's scale table: the table + # is q99.9(|x2|)/7 on real (rotated) activations, so rescale synthetic + # x2 per channel to match. Without this the fake-quant clips almost + # everything and the comparison measures clipping, not format migration. + q999 = torch.quantile(x2.abs().float(), 0.999, dim=0).clamp_min(1e-8) + x2 = (x2.float() * (7.0 * s_t / q999)).half() + + # References (fp32 accumulate). + y_fp = x2.float() @ w.t() + y_omega = fake_quant_omega(x2.float(), s_t) @ w.t() + print("\nreferences:") + report("omega vs fp (own quant cost)", y_omega, y_fp) + + if args.mode == "emulate": + y_s0 = e0m3_emulate(x2) @ e0m3_emulate(w).t() + y_s1 = (e0m3_emulate((x2.float() / s_t).half().float()) + @ e0m3_emulate(w * s_mean).t()) + else: + if not torch.cuda.is_available(): + print("error: --mode kernel requires CUDA", file=sys.stderr) + return 2 + try: + import flash_rt.flash_rt_fp4 as fvk_fp4 + except ImportError: + print("error: flash_rt_fp4 not importable — run on Thor", + file=sys.stderr) + return 2 + y_s0 = e0m3_kernel_gemm(x2, w.half(), fvk_fp4).float() + y_s1 = e0m3_kernel_gemm((x2.float() / s_t).half(), + (w * s_mean).half(), fvk_fp4).float() + + print(f"\nvariants vs references (mode={args.mode}):") + report("S0 (drop table) vs omega", y_s0, y_omega) + report("S1 (mean fold) vs omega", y_s1, y_omega) + report("S0 vs fp", y_s0, y_fp) + report("S1 vs fp", y_s1, y_fp) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/convert_omega_pack_e0m3.py b/tools/convert_omega_pack_e0m3.py new file mode 100644 index 00000000..ef6ec42e --- /dev/null +++ b/tools/convert_omega_pack_e0m3.py @@ -0,0 +1,178 @@ +#!/usr/bin/env python3 +r"""Offline converter: Omega-QVLA dit_svdquant_v1 pack -> FlashRT E0M3 weights. + +Reads an Omega-QVLA quantized pack (see docs/omega_pack_e0m3.md for the +record schema) and re-quantizes every `weight_res_q` tensor into the +FlashRT SM110 E0M3 operand format: packed 4-bit elements [N, K/2] plus +tile-interleaved UE4M3 SFB scales (per-16, amax/7), via the +`quantize_e0m3_dynamic_sfa_fp16` kernel. No GPTQ bitstream decoding is +needed — the pack stores weights as dequantized fp16. + +Scale-fold strategies (--fold): + none : B = e0m3(W). The activation-side per-channel calibration table is + not represented anywhere (strategy S0 in the doc). + mean : B = e0m3(W * diag(s_mean)), s_mean = act_scale_table.mean(dim=0). + The runtime must then divide activations by s_t per step before + quantization (strategy S1). Exact for the mean step; residual is + the table's step-to-step spread. BROKEN on hardware: raw s_mean + (~1e-2) shrinks weight columns, pressing per-16 block scales + below the UE4M3 subnormal floor (2^-9). Kept for reference. + actnorm : floor-safe S1. Decompose s_mean = c * r with c = geomean + (per-layer scalar) and r = s_mean / c (geomean 1, O(1) entries): + weights fold r (magnitudes preserved, no floor issue), + activations are divided by s_mean at runtime (static — no + per-step dispatch), and c is absorbed into the GEMM alpha. + Identity: (x/s_mean) @ (W*r)^T * c == x @ W^T. + +Auxiliary tensors needed by a runtime consumer (input/output rotations, +permutation, scale tables) are copied through unchanged into the output. + +Requires: CUDA + the compiled flash_rt_fp4 extension (i.e. run on Thor; +the quantize kernels are plain CUDA but the GEMM they feed is SM110). + +Usage: + python tools/convert_omega_pack_e0m3.py \ + --pack /path/to/Omega-QVLA/packs_hf/pi05_long/quantized.pt \ + --out pi05_long_e0m3.pt --fold none + # subset for bring-up: + python tools/convert_omega_pack_e0m3.py --pack ... --out /tmp/one.pt \ + --fold none --layer-regex 'layers\.0\.self_attn\.q_proj' +""" + +from __future__ import annotations + +import argparse +import re +import sys +import time + +import torch + +OUTPUT_FORMAT = "omega_e0m3_v1" + +# Record fields copied verbatim into the output's per-layer aux entry. +AUX_TENSORS = ( + "duquant_rotation_blocks", + "duquant_rotation_perm", + "duquant_rotation_out_blocks", + "act_scale_table", +) +AUX_SCALARS = ("weight_bits", "a_bits", "in_features", "out_features", "rank") + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + p.add_argument("--pack", required=True, help="input Omega quantized.pt") + p.add_argument("--out", required=True, help="output .pt path") + p.add_argument("--fold", choices=("none", "mean", "actnorm"), + default="none", + help="scale-table fold strategy (default: none = S0; " + "mean/actnorm are ablation-only, see docstring)") + p.add_argument("--layer-regex", default="", + help="only convert layers matching this regex") + p.add_argument("--keep-fp16", action="store_true", + help="also store the (possibly folded) fp16 weight, " + "for offline reference checks") + return p.parse_args() + + +def main() -> int: + args = parse_args() + + if not torch.cuda.is_available(): + print("error: CUDA is required (quantize kernels run on GPU)", + file=sys.stderr) + return 2 + try: + import flash_rt.flash_rt_fp4 as fvk_fp4 + except ImportError: + print("error: flash_rt_fp4 extension not importable — run this on a " + "machine with FlashRT built (Thor)", file=sys.stderr) + return 2 + + pack = torch.load(args.pack, map_location="cpu", weights_only=True) + meta = pack.get("__meta__", {}) + names = sorted(k for k in pack if k != "__meta__") + if args.layer_regex: + rx = re.compile(args.layer_regex) + names = [n for n in names if rx.search(n)] + if not names: + print("error: no layers matched", file=sys.stderr) + return 2 + + device = torch.device("cuda") + weights: dict = {} + aux: dict = {} + t0 = time.time() + for i, name in enumerate(names): + rec = pack[name] + if rec.get("format") != "dit_svdquant_v1": + print(f"skip {name}: format={rec.get('format')!r}") + continue + w = rec["weight_res_q"].to(device=device, dtype=torch.float16, + non_blocking=False).contiguous() + table = rec["act_scale_table"].float() + act_out_scale = None + if args.fold == "mean": + s_mean = table.mean(dim=0) # (in,) + w = (w * s_mean.to(device=device, dtype=torch.float16) + .unsqueeze(0)).contiguous() + elif args.fold == "actnorm": + s_mean = table.mean(dim=0).clamp_min(1e-12) # (in,) + c = float(torch.exp(torch.log(s_mean).mean())) + r = s_mean / c # geomean 1, O(1) entries + w = (w * r.to(device=device, dtype=torch.float16) + .unsqueeze(0)).contiguous() + act_out_scale = c + n, k = w.shape + if k % 16 != 0: + print(f"skip {name}: K={k} not divisible by 16") + continue + + packed = torch.empty(n, k // 2, dtype=torch.uint8, device=device) + # Zero-init: tile-interleaved SFB pads K to 64-element atoms and the + # kernel never writes padding entries (see fp4_utils.py). + sfb = torch.zeros(fvk_fp4.sfa_size_bytes(n, k, True), + dtype=torch.uint8, device=device) + rc = fvk_fp4.quantize_e0m3_dynamic_sfa_fp16( + w.data_ptr(), packed.data_ptr(), sfb.data_ptr(), n, k, True, 0) + if rc != 0: + raise RuntimeError(f"quantize_e0m3_dynamic_sfa_fp16 failed on " + f"{name}: rc={rc}") + + entry = {"packed": packed.cpu(), "sfb": sfb.cpu(), "N": n, "K": k} + if args.keep_fp16: + entry["weight_fp16_folded"] = w.cpu() + weights[name] = entry + + aux_entry = {f: rec[f].clone() for f in AUX_TENSORS if f in rec} + aux_entry.update({f: rec[f] for f in AUX_SCALARS if f in rec}) + aux_entry["fold"] = args.fold + if args.fold == "actnorm": + # Consumer contract: divide activations by act_scale_static + # (post-rotation, pre-quantize) and pass act_out_scale as the + # GEMM alpha. See the --fold actnorm note in the docstring. + aux_entry["act_scale_static"] = s_mean.clone() + aux_entry["act_out_scale"] = act_out_scale + aux[name] = aux_entry + + if (i + 1) % 21 == 0 or i + 1 == len(names): + print(f"[{i + 1}/{len(names)}] {name} N={n} K={k} " + f"({time.time() - t0:.1f}s)") + + torch.cuda.synchronize() + out = { + "format": OUTPUT_FORMAT, + "source_pack_meta": meta, + "fold": args.fold, + "weights": weights, + "aux": aux, + } + torch.save(out, args.out) + print(f"wrote {args.out}: {len(weights)} layers, fold={args.fold}, " + f"{time.time() - t0:.1f}s total") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/gen_omega_pack_fixture.py b/tools/gen_omega_pack_fixture.py new file mode 100644 index 00000000..74c6d940 --- /dev/null +++ b/tools/gen_omega_pack_fixture.py @@ -0,0 +1,117 @@ +#!/usr/bin/env python3 +"""Generate a synthetic miniature dit_svdquant_v1 pack (fixture). + +Lets anyone exercise the converter + consumer round-trip WITHOUT the +real 4.8GB Omega-QVLA pack: + + python tools/gen_omega_pack_fixture.py --out /tmp/fixture_pack.pt + python tools/convert_omega_pack_e0m3.py --pack /tmp/fixture_pack.pt \ + --out /tmp/fixture_e0m3.pt --fold none --keep-fp16 # Thor + python tools/check_omega_e0m3_consumer.py --reference fp16 \ + --artifact /tmp/fixture_e0m3.pt + +Design notes (what makes the fixture a real test and not a toy): + +- Schema is field-for-field identical to the real pack (checked against + packs_hf/pi05_long/quantized.pt): format tag, fp16 `weight_res_q`, + empty rank-0 lowrank tensors, fp16 64x64 rotation blocks, int64 perm, + float32 act_scale_table (expert style: T=10 rows; paligemma style: + T=1), plus the calibration metadata scalars. +- Rotations are RANDOM ORTHOGONAL (QR of gaussian), not identity — + otherwise the consumer's bmm rotate path would pass trivially. +- Weights carry per-input-channel lognormal scale outliers, the regime + that stresses E0M3 per-16 quantization and the UE4M3 subnormal floor + (the S1 failure mode in docs/omega_pack_e0m3.md). +- Shapes mix one full-size realistic layer (16384x2048) with small ones + (256x1024) to cover kernel edge sizes; all N,K are multiples of 64 as + the 64x64 rotation blocks require. + +Pure CPU, deterministic under --seed. Output ~85MB with defaults. +""" + +from __future__ import annotations + +import argparse + +import torch + +# (name, N=out, K=in, table_rows) — names/shapes mirror the real pack. +LAYERS = [ + ("paligemma_with_expert.gemma_expert.model.layers.0.self_attn.q_proj", + 2048, 1024, 10), + ("paligemma_with_expert.gemma_expert.model.layers.0.self_attn.v_proj", + 256, 1024, 10), + ("paligemma_with_expert.gemma_expert.model.layers.0.mlp.down_proj", + 1024, 4096, 10), + ("paligemma_with_expert.paligemma.model.language_model.layers.0." + "mlp.gate_proj", 16384, 2048, 1), +] + + +def _orthogonal_blocks(n_blocks: int, g: torch.Generator) -> torch.Tensor: + """(nb, 64, 64) random orthogonal blocks, fp16 (as the real pack).""" + a = torch.randn(n_blocks, 64, 64, generator=g) + q, r = torch.linalg.qr(a) + # QR sign convention: make diag(R) positive so Q is uniform-ish. + sign = torch.sign(torch.diagonal(r, dim1=1, dim2=2)) + q = q * sign[:, None, :] + return q.to(torch.float16) + + +def make_record(name: str, n: int, k: int, t_rows: int, + g: torch.Generator) -> dict: + assert k % 64 == 0 and n % 64 == 0, "64x64 rotation blocks need it" + # Weight: gaussian body x per-input-channel lognormal outliers. + chan_scale = torch.exp(torch.randn(k, generator=g) * 0.8) + w = (torch.randn(n, k, generator=g) + * chan_scale[None, :] * 0.02).to(torch.float16) + + # Act scale table: positive, ~1.0, mild per-step + per-channel jitter. + table = torch.exp(torch.randn(t_rows, k, generator=g) * 0.10 + + torch.randn(k, generator=g)[None, :] * 0.05) + + return { + "format": "dit_svdquant_v1", + "weight_res_q": w, + "lowrank_A": torch.empty(n, 0, dtype=torch.float16), + "lowrank_B": torch.empty(k, 0, dtype=torch.float16), + "act_scale_table": table.float(), + "duquant_rotation_blocks": _orthogonal_blocks(k // 64, g), + "duquant_rotation_perm": torch.randperm(k, generator=g), + "duquant_rotation_out_blocks": _orthogonal_blocks(n // 64, g), + "weight_bits": 4, + "a_bits": 4, + "rank": 0, + "in_features": k, + "out_features": n, + "n_calib_total": 100 * t_rows, + "n_calib_per_step": [100] * t_rows, + "act_percentile": 99.9, + "gptq_damp_percent": 0.05, + } + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + ap.add_argument("--out", required=True, help="output .pt path") + ap.add_argument("--seed", type=int, default=0) + args = ap.parse_args() + + g = torch.Generator().manual_seed(args.seed) + pack = {"__meta__": { + "recipe": "synthetic fixture (tools/gen_omega_pack_fixture.py)", + "suite": "fixture", + "fresh": True, + "seed": args.seed, + }} + for name, n, k, t in LAYERS: + pack[name] = make_record(name, n, k, t, g) + print(f" {name} N={n} K={k} table=({t},{k})") + + torch.save(pack, args.out) + print(f"wrote {args.out}: {len(LAYERS)} layers, seed={args.seed}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) From 7355c3c6be473b4b4bc95a392ae156349bf20c7c Mon Sep 17 00:00:00 2001 From: LiangSu8899 <7thuniversels@gmail.com> Date: Sun, 23 Aug 2026 13:41:49 -0400 Subject: [PATCH 2/2] fix(tools): harden Omega E0M3 conversion --- .gitignore | 2 +- docs/omega_pack_e0m3.md | 512 +++++++++---------------------- tests/test_omega_e0m3_tools.py | 203 ++++++++++++ tools/check_omega_e0m3_layer.py | 121 +++++++- tools/convert_omega_pack_e0m3.py | 173 +++++++++-- tools/gen_omega_pack_fixture.py | 4 +- 6 files changed, 623 insertions(+), 392 deletions(-) create mode 100644 tests/test_omega_e0m3_tools.py diff --git a/.gitignore b/.gitignore index b9b19e66..e244a554 100644 --- a/.gitignore +++ b/.gitignore @@ -44,4 +44,4 @@ internal-tests/ internal-docs/ notes/ third_party/cutlass -*.pt +/artifacts/omega_e0m3/*.pt diff --git a/docs/omega_pack_e0m3.md b/docs/omega_pack_e0m3.md index 0ced2998..41e3a546 100644 --- a/docs/omega_pack_e0m3.md +++ b/docs/omega_pack_e0m3.md @@ -1,386 +1,160 @@ -# Omega-QVLA pack format and the E0M3 consumption contract - -## 0. Background concepts (90 seconds) - -- **4-bit quantization**: store `round(x / s)` (a small integer) plus the - "ruler" `s` (the scale) instead of `x`. Compute happens as - `integer × s`. Fewer bits = less memory bandwidth, more rounding error. -- **fake-quant**: quantize then *immediately dequantize*, staying in float. - Simulates quantization error without needing integer hardware — Omega's - whole runtime is fake-quant emulation on plain PyTorch matmuls. -- **scale granularity**: how many elements share one ruler. *Per-channel* = - each of the K input channels gets its own (Omega's choice). *Per-16 - block* = 16 adjacent elements share one (the hardware format's choice). - Coarser granularity = fewer scales to store, but elements of different - magnitudes get crushed under a shared ruler. -- **static vs. dynamic scale**: *static* = measured offline on calibration - data, stored in the pack (Omega's `act_scale_table`). *Dynamic* = - computed per token at runtime from the actual data (`amax / 7`). Dynamic - is fresher but constrains the layout to what hardware computes cheaply. -- **E0M3**: the 4-bit element format here — sign + 3 mantissa-ish bits - decoding the uniform integer grid −7..+7. "Uniform" = evenly spaced - levels, unlike E2M1 (NVFP4) whose levels bunch near zero. -- **UE4M3**: an unsigned 4-exponent/3-mantissa mini-float used *only for - scales* (the ruler itself is quantized too). Per-16 scales on both - operands are UE4M3. -- **packed + SFA/SFB**: 4-bit elements are stored two per byte ("packed"); - the per-16 scales live in a separate buffer in CUTLASS's tile-interleaved - layout (SFA for the activation operand, SFB for the weight operand). -- **DuQuant rotation / permutation**: a learned orthogonal transform - (64×64 blocks + a channel shuffle) applied to activations before - quantization. Its job: even out channel magnitudes so no single outlier - channel dominates a shared scale. Orthogonal = length- and - angle-preserving, so it is mathematically free. -- **tcgen05 MMA**: the SM100/SM110 tensor-core instruction that consumes - packed 4-bit operands + UE4M3 scales directly in hardware. This is the - payoff: Omega's math runs as emulation today; this instruction makes it - native. -- **cosine similarity**: the fidelity metric. 1.0 = identical direction; - per-token cos 0.98 means the quantized output vector points in nearly - the same direction as the reference, with ~2% orthogonal noise. - -With those nine, every section below should read top to bottom without -external references. - -Status: recon complete, converter/harness in `tools/` (Milestone 1). -Scope: `packs_hf/pi05_long/quantized.pt` (4.8 GB, pi0.5 LIBERO-10 recipe -`paligemma=svdh+gptq, expert=svdh+rtn+perstep`). Other Omega packs share the -`dit_svdquant_v1` record format but were not inspected. - -## 1. Container - -Plain `torch.save` dict, loadable with `weights_only=True` (no custom -classes). 253 top-level keys: - -- 126 expert records: - `paligemma_with_expert.gemma_expert.model.layers.{0..17}.{self_attn.{q,k,v,o}_proj,mlp.{gate,up,down}_proj}` -- 126 PaliGemma records: - `paligemma_with_expert.paligemma.model.language_model.layers.{0..17}.` -- `__meta__`: `{"recipe": str, "suite": "10", "fresh": bool}` - -Small projections (state_proj, action_in/out_proj, time_mlp) are -deliberately absent — they break under A4 and stay BF16 at runtime. - -## 2. Record schema (`format == "dit_svdquant_v1"`) - -| field | shape / dtype | meaning | -|---|---|---| -| `weight_res_q` | `(out, in)` fp16 | fake-quantized-then-dequantized weight, **already in the rotated + permuted domain** | -| `lowrank_A` / `lowrank_B` | `(out, 0)` / `(in, 0)` fp16 | SVDQuant low-rank branch; rank = 0 in this pack (INT4-only path) | -| `act_scale_table` | `(num_steps, in)` fp32 | per-denoise-step, per-channel activation scales. Expert: `num_steps = 10`; PaliGemma: `1` | -| `duquant_rotation_blocks` | `(in/64, 64, 64)` fp16 | block-diagonal input rotation R_in | -| `duquant_rotation_perm` | `(in,)` int64 | input-channel permutation (applied before R_in) | -| `duquant_rotation_out_blocks` | `(out/64, 64, 64)` fp16 | block-diagonal output rotation (restore) | -| `weight_bits` / `a_bits` | int | 4 / 4 (both sides). The *runtime* `DuQuantLinear` path defaults activations to A8 (`GR00T_DUQUANT_ABITS`), which is where the "W4A8 PaliGemma" label comes from; pack records consumed through `GptqLinear` use their own `a_bits` (4) | -| `in_features` / `out_features` | int | redundant with tensor shapes | -| `n_calib_*`, `act_percentile`, `gptq_damp_percent` | scalars | calibration provenance | - -Notably absent (vs. a classic GPTQ pack): no packed int4 bitstream, no -`qweight`/`qzeros`/group scales — the weight survives only as dequantized -fp16 on the 4-bit grid (~8k unique values per tensor). No `smooth_scale`. - -## 3. Consumer math (Omega `gr00t/quantization/gptq_layers.py`, verified) - +# Omega-QVLA pack to FlashRT E0M3 format + +This document describes the **Milestone 1 offline toolchain** shipped in this +repository. It converts rank-0 Omega-QVLA `dit_svdquant_v1` records into the +packed E0M3/UE4M3 weight format consumed by FlashRT's existing SM110 GEMM. + +This change does not connect the artifact to a frontend or pipeline. It adds no +runtime route, server, CUDA graph, CMake source, binding, or public Python API. +Runtime-consumer results mentioned in development discussions were produced on +an external experimental branch and are not capabilities of this repository. + +## Included tools + +- `tools/convert_omega_pack_e0m3.py`: validates and converts a pack. +- `tools/check_omega_e0m3_layer.py`: CPU emulation diagnostics and a Thor + kernel round-trip that directly consumes converted `packed` and `sfb` data. +- `tools/gen_omega_pack_fixture.py`: creates a four-record synthetic fixture. + +## Input container + +The input is a plain `torch.save` mapping loadable with +`torch.load(..., weights_only=True)`. A full inspected pi0.5 pack contains 252 +records plus `__meta__`: 126 expert records and 126 PaliGemma records. + +Each selected record must have `format == "dit_svdquant_v1"` and satisfy: + +| Field | Required contract | +|---|---| +| `weight_res_q` | finite floating tensor `[out_features, in_features]` | +| `rank` | exactly `0` | +| `lowrank_A`, `lowrank_B` | `[out_features, 0]`, `[in_features, 0]` | +| `weight_bits`, `a_bits` | both `4` | +| `act_scale_table` | finite positive floating tensor `[steps, in_features]` | +| `duquant_rotation_blocks` | finite `[in_features / 64, 64, 64]` | +| `duquant_rotation_perm` | int64 permutation of `[0, in_features)` | +| `duquant_rotation_out_blocks` | finite `[out_features / 64, 64, 64]` | +| `in_features`, `out_features` | exact match for the weight tensor | + +Both feature dimensions must be positive multiples of 64. The converter +validates every selected record before importing the CUDA extension or writing +an artifact. Unsupported records are errors; they are never skipped. + +For an unfiltered non-fixture pack, the default coverage gate is 252 records. +`--layer-regex` and fixture packs expect every selected record. Use +`--expected-records` when intentionally converting a different complete pack. + +## Output container + +The output format is `omega_e0m3_v1`, schema version 1: + +```text +{ + "format": "omega_e0m3_v1", + "schema_version": 1, + "source_pack_meta": {...}, + "source_record_count": int, + "selected_record_count": int, + "selected_layers": [str, ...], + "fold": "none" | "mean" | "actnorm", + "weights": { + layer: { + "packed": uint8[N, K / 2], + "sfb": uint8[sfa_size_bytes(N, K, True)], + "N": int, + "K": int, + } + }, + "aux": {layer: {...}} +} ``` -x2 = bmm(x[..., perm].view(N, in/64, 64), R_in_blocks) # input rotation, runtime -x_q = clamp(round(x2 / s_t), -8, 7) * s_t # s_t = act_scale_table[step] -y' = x_q @ W_res_q^T # bf16-promoted accumulate -y = bmm(y'.view(N, out/64, 64), R_out_blocks) + bias # output rotation restore + +The converter writes through a sibling temporary file and atomically replaces +the destination only after all selected records convert successfully. The +`weights`, `aux`, and `selected_layers` sets must be identical. + +## Conversion math + +Omega stores `weight_res_q` as a dequantized fp16 tensor already in the +rotated and permuted domain. No GPTQ bitstream decoding is required. + +FlashRT converts each `[N, K]` weight with: + +```text +quantize_e0m3_dynamic_sfa_fp16(weight, packed, sfb, N, K, is_sfb=True) ``` -PaliGemma (`duquant_layers.py`) is identical in structure with A8 -activations and a single-row scale table. - -Consequences for a FlashRT consumer: - -- The input rotation **cannot** be folded into `weight_res_q`: fake-quant - sits between rotation and GEMM. It must run on activations (torch bmm, or - a prologue kernel). Same for the output restore. -- The rotation is an exact orthonormal transform, so it does not by itself - affect GEMM fidelity; fidelity questions live entirely in the quantizers. -- `weight_res_q` being plain fp16 means the converter re-quantizes from - fp16 — no GPTQ bitstream decoding needed. - -## 4. Mapping to the FlashRT E0M3 contract - -FlashRT SM110 path (`csrc/gemm/fp4/cutlass_fp4_gemm_e0m3w_sm100.cuh`, -bindings in `csrc/fp4_bindings.cpp`): - -- Weights: fp16 `[N, K]` → `quantize_e0m3_dynamic_sfa_fp16(..., is_sfb=True)` - → packed E0M3 `[N, K/2]` + SFB tile-interleaved UE4M3 (per-16, amax/7). -- Activations: same kernel with `is_sfb=False` → packed + SFA. -- GEMM: `cutlass_fp4_gemm_e0m3w(A, SFA, B, SFB, D, M, N, K, α, β, stream, - a_format)` with `a_format=0` for E0M3 activations (1 = E2M1). -- Buffer sizing: `flash_rt_fp4.sfa_size_bytes(N, K, is_sfb)`; scale buffers - must be zero-initialized (tile-interleave pads K to 64-element atoms; - garbage padding decodes as UE4M3 NaN). - -Grid differences vs. Omega fake-quant: - -| | Omega A4 | FlashRT E0M3 | -|---|---|---| -| element grid | int `[-8, 7]` (asymmetric clamp) | sign-magnitude uniform `[-7, 7]` | -| scale | static calibrated, **per-channel** fp32 | dynamic amax/7, **per-16** UE4M3 | -| weight grid | int4 per-channel-group (already dequantized) | per-16 UE4M3 | - -The scale-granularity mismatch (per-channel static table vs. per-16 dynamic) -is the one real fidelity risk. Two candidate strategies, both implemented in -`tools/check_omega_e0m3_layer.py`: - -- **S0 (drop the table)**: `A = e0m3(x2)`, `B = e0m3(W)`. Loses all - calibration information. -- **S1 (fold step-mean table into W, per-step residual into A)**: - `A = e0m3(x2 / s_t)` per step, `B = e0m3(W · diag(s̄))` once, where - `s̄ = mean_t(s_t)`. Exact for the mean step; residual error scales with - the table's step-to-step spread (measured: std/mean ≈ 10% on expert - layer-0 q_proj). - - Mathematically S1 relies on `Σ_k q_k s_k W_nk = Σ_k q_k (s_k W_nk)`: - a per-K-column scale commutes into the weight. A true per-step fold would - need 10 weight copies (unacceptable), hence the mean fold. - -RHT (per-16 Hadamard, `use_rht=1` variants) is orthogonal to the DuQuant -rotation — `(x2·H)(W·H)^T = x2·W^T` — and can be ablated on top of either -strategy if per-block distributions remain problematic. - -### Measured - -Emulation mode (torch, synthetic activations calibrated to q999 = 7·s_t, -M = 256 tokens), per-token cosine vs. the unquantized-activation reference: - -| layer | omega vs fp | S0 vs fp | S1 vs fp | -|---|---|---|---| -| expert L0 q_proj (K=1024) | 0.9927 | 0.9834 | 0.9052 | -| expert L0 down_proj (K=4096) | 0.9928 | 0.9825 | 0.9799 | -| expert L11 o_proj (K=2048) | 0.9929 | 0.9825 | 0.9615 | -| paligemma L0 gate_proj (K=2048) | 0.9928 | 0.9810 | 0.9775 | - -Kernel mode (real tcgen05 GEMM, Thor SM110, same seed, per-token mean -cosine vs. the unquantized-activation reference): - -| layer | omega vs fp | S0 vs fp | S1 vs fp | -|---|---|---|---| -| expert L0 q_proj (K=1024) | 0.99247 | **0.99322** | 0.158 | -| expert L0 down_proj (K=4096) | 0.99274 | **0.99257** | 0.98999 | -| expert L11 o_proj (K=2048) | 0.99275 | **0.99303** | 0.85557 | -| paligemma L0 gate_proj (K=2048) | 0.99206 | **0.99239** | 0.98204 | - -Three findings: - -1. **S0 is lossless on real hardware on every layer tested** — within - ±0.001 of Omega's own fake-quant everywhere (the tiny edges come from - dynamic per-token amax beating a static table on data calibrated only - at the q999 point). Error independence holds wherever checked: - cos(S0, fp)·cos(omega, fp) ≈ measured cos(S0, omega), i.e. S0's - residual is fresh rounding noise, not a systematic shift. -2. **S1's collapse on real hardware is scale-magnitude-dependent.** - Mechanism: `W · diag(s̄)` shrinks weights by the mean table value, - pushing per-16 block scales toward the UE4M3 subnormal floor (2⁻⁹), - where scale mantissas disintegrate and whole blocks quantize to - garbage. Layers with small s̄ die hard (q_proj 0.16, o_proj 0.86); - layers whose table happens to be larger merely degrade (down_proj - 0.99 — still worse than S0). The emulator's lenient subnormal - handling masked the severe cases. -3. The pure-torch references reproduce across machines to 5 decimal - places (0.992700 Thor vs. 0.992707 x86), cross-validating the harness. - -**S0 wins; S1 is dead.** The table's per-channel scale spread (~4×) -distorts weights when folded, while S0's per-token dynamic per-16 amax is -a *better* quantizer than Omega's static per-channel table — the DuQuant -rotation+perm has already whitened per-channel magnitudes, so the table -is only a second-order correction. - -Decision: **the converter emits S0 (`--fold none`) as the production -format**; `--fold mean` is kept for ablation only. This also shrinks the -runtime story — no per-step scale dispatch is needed on the E0M3 path. - -**Follow-up: `actnorm` (floor-safe S1) — also dead (2026-08-18).** A -reviewer-natural fix for S1's floor problem is to normalize before -folding: decompose `s̄ = c·r̄` with `c = geomean(s̄)`, fold only `r̄` -(O(1), geomean 1) into the weights, divide activations by `s̄` at -runtime, and absorb `c` into the GEMM alpha. This is exactly -`(x/s̄) @ (W·r̄)^T · c = x @ W^T`, and it does fix the floor (0% of -block scales below 2⁻⁹ vs 100% for raw S1 on q_proj). But measured on -Thor (consumer-level, fp16 reference, real pack): - -| layer | S0 vs fp16 | actnorm vs fp16 | -|---|---|---| -| q_proj | 0.9935 | 0.9886 | -| down_proj | 0.9929 | 0.9869 | -| o_proj | 0.9936 | 0.9885 | - -actnorm is *worse* than S0 everywhere. Mechanism: the fold is a zero-sum -redistribution — dividing activations by `s̄` whitens the activation -blocks, but multiplying weights by `r̄` (range 0.43–2.68 on q_proj) -re-opens intra-block magnitude spread on the weight side, where per-16 -single-scale 4-bit pays for it. DuQuant's rotation had already whitened -both operands; any per-channel re-scaling of either side undoes that. -**Per-channel calibration tables are fundamentally incompatible with -per-16 block quantization — the information has to live on one side and -always de-whitens it.** Dynamic per-16 amax is the optimum at this -granularity; S0 is the endpoint, not a compromise. (`--fold actnorm` + -consumer support remain in the tree, `OMEGA_E0M3_ACT_TABLE=0`/artifact -driven, as the documented ablation.) The residual end-to-end gap vs. -the full Omega recipe (90.4% vs 93.2%, concentrated in task9) is not -recoverable by table injection; remaining options are mixed precision -for sensitive layers or acceptance. - -Remaining caveats: synthetic activations (lognormal + outlier channels, -calibrated only at the q999 point) — real activation tails differ. Next: -captured real activations, then LIBERO paired SR (Milestone 2). - -## 5. Roadmap and Milestone-1 deliverables - -**Milestone 1 — offline converter + format doc + single-layer gates -(done).** Deliverables below; acceptance: 252/252 records converted, S0 -per-token cosine ≥ Omega fake-quant on real hardware (4/4 layers, §4). - -**Milestone 2 — runtime consumption (next).** Wire the converted pack -into the pi0.5 Thor pipeline: load `packed`/`sfb` as decoder GEMM -operands, run the DuQuant input rotation (perm + 64×64 block bmm) and -output restore around each replaced Linear, keep the small projections -(state/action/time) BF16 from the checkpoint. Acceptance: end-to-end -action cosine vs. the Omega fake-quant server, then LIBERO-10 ×500 -paired SR vs. the BF16 baseline (target: no measurable loss, matching -the pack's own 93.2% vs 91.6%). - -**Milestone 3 — upstream PRs (after M2 evidence).** Split per -`CONTRIBUTING.fork.md` §6, each with LIBERO paired SR + action cosine + -p50/p95 latency: ① converter + this format doc (pure additive, easiest); -② runtime wiring as a flag-gated `weight_format` branch (the S0 result -shrunk this from the originally-planned per-step scale path); ③ SVDQuant -low-rank epilogue (deferred — rank = 0 in this pack). - -**Milestone 2 status (done, incl. 2c/2d landed after the original -write-up):** - -- **M2a/b — consumer + serving (done).** `tools/omega_e0m3_linear.py` - (`OmegaE0M3Linear`, drop-in for gr00t's `GptqLinear` via - `tools/serve_omega_e0m3.py` monkeypatch) + `tools/check_omega_e0m3_consumer.py` - gate: per-layer cosine 0.978–0.982 vs. GptqLinear, 1.2× layer latency. - Server smoke 10/10 ≙ arm D; **LIBERO-10 ×500 paired: 90.4%** vs. BF16 - 91.6% (McNemar p = 0.53) and vs. fake-quant arm D 93.2% (p = 0.070) — - no significant loss; 58 s/episode vs. 148 s fake-quant (2.6×). Eager - mode with `torch.compile` disabled: the pybind kernels graph-break and - the HF KV cache recompiles per step (~25 min/episode stall) — see the - env flags in `serve_omega_e0m3.py`. -- **M2d — hand-rolled CUDA graph over the whole denoise loop (done).** - `tools/omega_e0m3_graph.py` captures all 10 flow-matching steps - (unrolled, pi05_thor style) into one `torch.cuda.CUDAGraph`: static KV - slabs behind a `DynamicCache` shell, static mask/position buffers - filled by `copy_` per inference, adaRMS conditioning precomputed for - the deterministic time grid. The eager blockers removed are documented - in the module docstring (device-scalar `while`, per-step H2D mask - upload, per-call KV allocation). Enable with `OMEGA_E0M3_CUDA_GRAPH=1` - (see `tools/start_e0m3_server.sh`); falls back to eager permanently on - any capture failure. Thor validation: capture succeeds - (`prefix_len=968, layers=18, steps=10`), smoke 10/10, ~43–50 - s/episode vs. ~58 s eager. -- **M2e — PaliGemma E0M3 (route A: official all-W4A4 pack recipe).** The - converter already emits all 252 records, PaliGemma included; serving - with `OMEGA_E0M3_PATCH_DUQUANT=1` (the default in - `tools/start_e0m3_server.sh`) substitutes the runtime `DuQuantLinear` - wraps with `OmegaE0M3Linear` consumers built from the pack's PaliGemma - records — GPTQ W4A4 weights instead of runtime RTN, single-row scale - table dropped per the S0 decision. With `omega_e0m3_graph.py`'s prefix - graph (default on), the prefix prefill is captured too, so the E0M3 - pybind kernels run inside a CUDA graph on this path as well — the - capture smoke (`check_omega_e0m3_graph_smoke.py`) covers their - capturability. - Validation ladder on Thor: artifact coverage check (252 records) → - per-layer consumer gate on PaliGemma layers → 10-episode smoke → - LIBERO-10 ×500 paired SR + per-episode latency. - -Deliverables: - -- `tools/convert_omega_pack_e0m3.py` — offline pack → E0M3 converter - (S0 weight emission + aux tensors: perm, R_in/R_out blocks, - act_scale_table). Runs where `flash_rt_fp4` is built (Thor). -- `tools/check_omega_e0m3_layer.py` — single-layer cosine harness: - Omega fake-quant reference vs. FlashRT E0M3 GEMM (S0/S1), plus a pure - torch emulation mode that runs without the extension for pre-checks. - Emulation results and the S0 decision are in §4. -- `tools/omega_e0m3_linear.py` — `OmegaE0M3Linear` consumer (M2a). -- `tools/check_omega_e0m3_consumer.py` — consumer-vs-GptqLinear gate - (cosine + layer latency). -- `tools/serve_omega_e0m3.py` — openpi serving entry (monkeypatches - gr00t's wrap classes; env-gated compile kill switch). -- `tools/omega_e0m3_graph.py` — CUDA-graph capture of the 10-step - denoise loop (M2d), `OMEGA_E0M3_CUDA_GRAPH=1`. -- `tools/check_omega_e0m3_graph_smoke.py` — P0 capture gate for a single - consumer layer (pybind capturability check). -- `tools/start_e0m3_server.sh` — Thor server launcher (repo-relative - paths, env overrides). - -Deferred: SVDQuant low-rank epilogue (rank = 0 everywhere in this pack), -per-step weight tables (10× memory; also refuted by the S0 result), -per-step activation scale dispatch (refuted by the S0 result). - -### Reproducing - -**Self-contained fixture round-trip (no Omega-QVLA pack or checkout -needed — the PR review path):** +This produces packed 4-bit E0M3 elements and tile-interleaved UE4M3 per-16 +scales. SFB storage is allocated with +`flash_rt_fp4.sfa_size_bytes(N, K, True)` and zero-initialized because the +tile-interleaved layout contains padding. + +The default `--fold none` strategy intentionally does not fold the +per-channel activation table into the weight. Thor experiments on four real +layers measured E0M3-vs-fp16 cosine around 0.9924-0.9932, comparable to the +source fake-quant path. `mean` and `actnorm` remain ablation modes; they are not +the recommended artifact format. + +## Reproduction + +### CPU fixture and emulation ```bash -cd third_party/flashrt -# 1. Synthetic miniature pack: schema-identical records, random -# orthogonal rotations, outlier-channel weights (pure CPU, seconds) python tools/gen_omega_pack_fixture.py --out /tmp/fixture_pack.pt -# 2. Convert (Thor) -python tools/convert_omega_pack_e0m3.py --pack /tmp/fixture_pack.pt \ - --out /tmp/fixture_e0m3.pt --fold none -# 3. Consumer vs fp16 reference, gr00t-free (Thor) -PYTHONPATH=$PWD/tools python tools/check_omega_e0m3_consumer.py \ - --reference fp16 --pack /tmp/fixture_pack.pt \ - --artifact /tmp/fixture_e0m3.pt + +python tools/convert_omega_pack_e0m3.py \ + --pack /tmp/fixture_pack.pt \ + --out /tmp/unused.pt \ + --validate-only + +python tools/check_omega_e0m3_layer.py \ + --pack /tmp/fixture_pack.pt \ + --mode emulate ``` -**Full pack (development path):** +The validation command performs no CUDA work and does not create its output. + +### Thor artifact round-trip ```bash -# Point at an Omega pack (any machine for emulate, Thor for kernel/convert) -export OMEGA_PACK=/path/to/Omega-QVLA/packs_hf/pi05_long/quantized.pt -cd third_party/flashrt - -# 1. Local pre-check, no extension needed (pure torch, CPU is fine) -python tools/check_omega_e0m3_layer.py --pack "$OMEGA_PACK" --mode emulate - -# 2. Hardware check — real tcgen05 GEMM (Thor, flash_rt_fp4 built) -python tools/check_omega_e0m3_layer.py --pack "$OMEGA_PACK" --mode kernel +python tools/convert_omega_pack_e0m3.py \ + --pack /tmp/fixture_pack.pt \ + --out /tmp/fixture_e0m3.pt \ + --fold none + +python tools/check_omega_e0m3_layer.py \ + --pack /tmp/fixture_pack.pt \ + --artifact /tmp/fixture_e0m3.pt \ + --mode kernel \ + --min-artifact-cos 0.98 ``` -```bash -# 2. Hardware check output -layer: paligemma_with_expert.gemma_expert.model.layers.0.self_attn.q_proj N(out)=2048 K(in)=1024 table=(10,1024) step=0 +With `--artifact`, the checker quantizes only the activation. It loads the +weight's `packed` and `sfb` tensors directly from the converted artifact, +validates their shape and byte count, runs the GEMM, and fails if cosine is +below the requested threshold. This is the converter round-trip gate. + +### Full pack -references:... -``` ```bash -python tools/check_omega_e0m3_layer.py --pack "$OMEGA_PACK" --mode kernel \ - --layer paligemma_with_expert.gemma_expert.model.layers.0.mlp.down_proj - # --layer paligemma_with_expert.paligemma.model.language_model.layers.0.mlp.gate_proj - # --layer paligemma_with_expert.gemma_expert.model.layers.11.self_attn.o_proj -# 3. Full conversion (252 layers, ~1.3 GB output) python tools/convert_omega_pack_e0m3.py \ - --pack "$OMEGA_PACK" --out pi05_long_e0m3.pt --fold none + --pack /path/to/quantized.pt \ + --out /path/to/pi05_e0m3.pt \ + --fold none ``` -Gate for accepting the conversion: per-token cosine of S0 vs. fp ≥ -Omega's own fake-quant (per-layer, same seed). Currently met on every -layer tested (see §4). - -## 6. Accuracy context (pi0.5 LIBERO-10, 500 episodes) - -The 93.2% figure was measured on the hybrid deployment: expert records from -this pack (W4A4, `GptqLinear`) + PaliGemma via the *runtime* DuQuant path -(W4 weights, A8 activations by the `GR00T_DUQUANT_ABITS` default) — vs. BF16 -baseline 91.6% (McNemar p = 0.32, no significant difference) on the Omega -PyTorch fake-quant path. The pack itself is the official Omega recipe, -which is W4A4 on both sides (PaliGemma records carry `a_bits=4` and a -single-row `act_scale_table`); consuming the PaliGemma records through the -E0M3 consumer (`OMEGA_E0M3_PATCH_DUQUANT=1`) therefore *is* the official -recipe, and additionally replaces runtime RTN weights with the pack's GPTQ -weights. The E0M3 migration target -is therefore "no measurable SR loss against an already lossless baseline" — -the single-layer cosine gates are the leading indicator, LIBERO the final -one. +Without a layer filter, this command fails unless all 252 records validate and +convert. A representative artifact layer should then be checked with the same +`--artifact --mode kernel` command above. Release evidence must include both +the direct packed/SFB GEMM result and `252/252` coverage. + +## Scope and roadmap + +Milestone 1 is limited to conversion, schema validation, a synthetic fixture, +CPU emulation, and direct artifact GEMM verification. + +A future runtime PR may load these artifacts in a pi0.5 Thor frontend, apply +the DuQuant input/output rotations, and establish end-to-end accuracy and +latency gates. Such a PR must independently add and test its frontend, +pipeline, graph-lifecycle, and serving contracts. None of that runtime surface +is provided here. diff --git a/tests/test_omega_e0m3_tools.py b/tests/test_omega_e0m3_tools.py new file mode 100644 index 00000000..87132fc8 --- /dev/null +++ b/tests/test_omega_e0m3_tools.py @@ -0,0 +1,203 @@ +import copy +import importlib.util +import subprocess +import sys +from pathlib import Path + +import pytest +import torch + + +ROOT = Path(__file__).resolve().parents[1] + + +def _load_tool(name: str): + path = ROOT / "tools" / f"{name}.py" + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +converter = _load_tool("convert_omega_pack_e0m3") +checker = _load_tool("check_omega_e0m3_layer") + + +def _record(n: int = 64, k: int = 64) -> dict: + return { + "format": "dit_svdquant_v1", + "weight_res_q": torch.randn(n, k, dtype=torch.float16), + "lowrank_A": torch.empty(n, 0, dtype=torch.float16), + "lowrank_B": torch.empty(k, 0, dtype=torch.float16), + "act_scale_table": torch.ones(2, k), + "duquant_rotation_blocks": torch.eye(64).half().view(1, 64, 64), + "duquant_rotation_perm": torch.arange(k, dtype=torch.int64), + "duquant_rotation_out_blocks": torch.eye(64).half().view(1, 64, 64), + "weight_bits": 4, + "a_bits": 4, + "rank": 0, + "in_features": k, + "out_features": n, + } + + +def test_converter_accepts_complete_rank_zero_record(): + assert converter.validate_record("layer", _record()) == (64, 64) + + +@pytest.mark.parametrize( + ("mutation", "message"), + [ + (lambda r: r.update(format="other"), "unsupported format"), + (lambda r: r.update(rank=1), "only rank=0"), + (lambda r: r.update(in_features=128), "metadata shape mismatch"), + (lambda r: r.update(act_scale_table=torch.ones(2, 32)), + "act_scale_table"), + (lambda r: r.update(duquant_rotation_blocks=torch.eye(64).half()), + "input rotation shape"), + (lambda r: r.update(duquant_rotation_perm=torch.zeros(64, dtype=torch.int64)), + "not a permutation"), + ], +) +def test_converter_rejects_incomplete_or_inconsistent_records(mutation, message): + record = copy.deepcopy(_record()) + mutation(record) + with pytest.raises(ValueError, match=message): + converter.validate_record("layer", record) + + +def test_full_pack_count_defaults_to_252_but_fixture_and_subset_are_exact(): + assert converter.expected_record_count( + {}, layer_regex="", selected_count=251, explicit=None) == 252 + assert converter.expected_record_count( + {"recipe": "synthetic fixture (test)"}, layer_regex="", selected_count=4, + explicit=None) == 4 + assert converter.expected_record_count( + {}, layer_regex="layer.0", selected_count=1, explicit=None) == 1 + + +def test_validate_only_runs_without_cuda_and_writes_no_artifact(tmp_path): + pack_path = tmp_path / "fixture_pack.pt" + out_path = tmp_path / "must_not_exist.pt" + torch.save({ + "__meta__": {"recipe": "synthetic fixture (test)"}, + "layer": _record(), + }, pack_path) + + result = subprocess.run( + [ + sys.executable, + str(ROOT / "tools" / "convert_omega_pack_e0m3.py"), + "--pack", str(pack_path), + "--out", str(out_path), + "--validate-only", + ], + text=True, + capture_output=True, + check=False, + ) + assert result.returncode == 0, result.stderr + assert "validated 1/1 records" in result.stdout + assert not out_path.exists() + + +def test_validate_only_fails_before_cuda_for_bad_records(tmp_path): + pack_path = tmp_path / "bad_pack.pt" + out_path = tmp_path / "must_not_exist.pt" + record = _record() + record["rank"] = 1 + torch.save({ + "__meta__": {"recipe": "synthetic fixture (test)"}, + "layer": record, + }, pack_path) + + result = subprocess.run( + [ + sys.executable, + str(ROOT / "tools" / "convert_omega_pack_e0m3.py"), + "--pack", str(pack_path), + "--out", str(out_path), + "--validate-only", + ], + text=True, + capture_output=True, + check=False, + ) + assert result.returncode == 2 + assert "only rank=0 is supported" in result.stderr + assert not out_path.exists() + + +def test_converter_refuses_to_overwrite_source_pack(tmp_path): + pack_path = tmp_path / "fixture_pack.pt" + torch.save({ + "__meta__": {"recipe": "synthetic fixture (test)"}, + "layer": _record(), + }, pack_path) + result = subprocess.run( + [ + sys.executable, + str(ROOT / "tools" / "convert_omega_pack_e0m3.py"), + "--pack", str(pack_path), + "--out", str(pack_path), + ], + text=True, + capture_output=True, + check=False, + ) + assert result.returncode == 2 + assert "must not overwrite" in result.stderr + loaded = torch.load(pack_path, map_location="cpu", weights_only=True) + assert "layer" in loaded + + +def test_artifact_schema_requires_complete_direct_packed_sfb_coverage(): + class FakeFp4: + @staticmethod + def sfa_size_bytes(n, k, is_sfb): + assert is_sfb + return 32 + + artifact = { + "format": "omega_e0m3_v1", + "schema_version": 1, + "fold": "none", + "selected_record_count": 1, + "selected_layers": ["layer"], + "weights": { + "layer": { + "packed": torch.zeros(64, 32, dtype=torch.uint8), + "sfb": torch.zeros(32, dtype=torch.uint8), + "N": 64, + "K": 64, + } + }, + "aux": {"layer": {"fold": "none"}}, + } + entry, aux = checker.validate_artifact( + artifact, "layer", 64, 64, FakeFp4()) + assert entry["packed"].shape == (64, 32) + assert aux["fold"] == "none" + + broken = copy.deepcopy(artifact) + broken["aux"] = {} + with pytest.raises(ValueError, match="do not cover"): + checker.validate_artifact(broken, "layer", 64, 64, FakeFp4()) + + +def test_docs_and_gitignore_only_describe_milestone_one_surface(): + docs = (ROOT / "docs" / "omega_pack_e0m3.md").read_text() + gitignore = (ROOT / ".gitignore").read_text().splitlines() + + for nonexistent in ( + "omega_e0m3_linear.py", + "check_omega_e0m3_consumer.py", + "serve_omega_e0m3.py", + "omega_e0m3_graph.py", + "check_omega_e0m3_graph_smoke.py", + "start_e0m3_server.sh", + ): + assert nonexistent not in docs + assert "*.pt" not in gitignore + assert "/artifacts/omega_e0m3/*.pt" in gitignore diff --git a/tools/check_omega_e0m3_layer.py b/tools/check_omega_e0m3_layer.py index 94177cd2..1ab9e27c 100644 --- a/tools/check_omega_e0m3_layer.py +++ b/tools/check_omega_e0m3_layer.py @@ -33,12 +33,16 @@ python tools/check_omega_e0m3_layer.py \ --pack packs_hf/pi05_long/quantized.pt --mode kernel \ --layer paligemma_with_expert.gemma_expert.model.layers.0.mlp.down_proj + python tools/check_omega_e0m3_layer.py \ + --pack /tmp/fixture_pack.pt --artifact /tmp/fixture_e0m3.pt \ + --mode kernel --min-artifact-cos 0.98 """ from __future__ import annotations import argparse import sys +from collections.abc import Mapping import torch @@ -58,6 +62,11 @@ def parse_args() -> argparse.Namespace: p.add_argument("--device", default="", help="torch device for reference/emulation math " "(default: cuda in kernel mode, cpu in emulate mode)") + p.add_argument("--artifact", + help="converted omega_e0m3_v1 artifact; in kernel mode, " + "use its packed/SFB weight directly instead of re-quantizing") + p.add_argument("--min-artifact-cos", type=float, default=0.98, + help="minimum artifact-vs-fp16 global cosine (default: 0.98)") return p.parse_args() @@ -162,13 +171,84 @@ def e0m3_kernel_gemm(a_fp16: torch.Tensor, b_fp16: torch.Tensor, return d +def validate_artifact(artifact: Mapping, layer: str, n: int, k: int, + fvk_fp4) -> tuple[Mapping, Mapping]: + if not isinstance(artifact, Mapping) or artifact.get("format") != "omega_e0m3_v1": + raise ValueError("artifact format must be 'omega_e0m3_v1'") + if artifact.get("schema_version") != 1: + raise ValueError( + f"artifact schema_version must be 1, got {artifact.get('schema_version')!r}") + weights = artifact.get("weights") + aux = artifact.get("aux") + if not isinstance(weights, Mapping) or not isinstance(aux, Mapping): + raise ValueError("artifact weights and aux must be mappings") + selected = artifact.get("selected_layers") + count = artifact.get("selected_record_count") + if not isinstance(selected, list) or count != len(selected): + raise ValueError("artifact selected layer metadata is inconsistent") + if set(selected) != set(weights) or set(selected) != set(aux): + raise ValueError("artifact weights/aux do not cover every selected layer") + if layer not in weights or layer not in aux: + raise ValueError(f"artifact does not contain layer {layer!r}") + entry = weights[layer] + aux_entry = aux[layer] + if not isinstance(entry, Mapping) or not isinstance(aux_entry, Mapping): + raise ValueError(f"artifact layer {layer!r} entries must be mappings") + if entry.get("N") != n or entry.get("K") != k: + raise ValueError( + f"artifact layer shape mismatch: expected N={n}, K={k}, got " + f"N={entry.get('N')!r}, K={entry.get('K')!r}") + packed = entry.get("packed") + sfb = entry.get("sfb") + if not isinstance(packed, torch.Tensor) or packed.dtype != torch.uint8 \ + or tuple(packed.shape) != (n, k // 2): + raise ValueError( + f"artifact packed must be uint8 shape ({n},{k // 2})") + expected_sfb = fvk_fp4.sfa_size_bytes(n, k, True) + if not isinstance(sfb, torch.Tensor) or sfb.dtype != torch.uint8 \ + or sfb.numel() != expected_sfb: + raise ValueError( + f"artifact sfb must be uint8 with {expected_sfb} bytes") + fold = artifact.get("fold") + if fold not in {"none", "mean", "actnorm"} or aux_entry.get("fold") != fold: + raise ValueError("artifact fold metadata is missing or inconsistent") + return entry, aux_entry + + +def e0m3_artifact_gemm(a_fp16: torch.Tensor, entry: Mapping, + fvk_fp4, *, alpha: float = 1.0) -> torch.Tensor: + """Quantize A, then consume artifact packed/SFB without re-quantizing B.""" + a_fp16 = a_fp16.contiguous() + m, k = a_fp16.shape + n = int(entry["N"]) + a_packed = torch.empty(m, k // 2, dtype=torch.uint8, device="cuda") + a_sfa = torch.zeros(fvk_fp4.sfa_size_bytes(m, k, False), + dtype=torch.uint8, device="cuda") + rc = fvk_fp4.quantize_e0m3_dynamic_sfa_fp16( + a_fp16.data_ptr(), a_packed.data_ptr(), a_sfa.data_ptr(), + m, k, False, 0) + if rc != 0: + raise RuntimeError(f"A quantize failed rc={rc}") + b_packed = entry["packed"].to(device="cuda", non_blocking=False).contiguous() + b_sfb = entry["sfb"].to(device="cuda", non_blocking=False).contiguous() + d = torch.empty(m, n, dtype=torch.float16, device="cuda") + rc = fvk_fp4.cutlass_fp4_gemm_e0m3w( + a_packed.data_ptr(), a_sfa.data_ptr(), + b_packed.data_ptr(), b_sfb.data_ptr(), d.data_ptr(), + m, n, k, alpha, 0.0, 0, 0) + if rc != 0: + raise RuntimeError(f"artifact cutlass_fp4_gemm_e0m3w failed rc={rc:#x}") + torch.cuda.synchronize() + return d + + # ──────────────────────────────────────────────────────────────────── def cosine_stats(a: torch.Tensor, b: torch.Tensor) -> tuple: """(global cos, per-row cos mean, per-row cos min), fp32 inputs.""" a = a.float() b = b.float() - glob = torch.dot(a.flatten(), b.flatten()) / ( - a.norm() * b.norm()).item() + glob = (torch.dot(a.flatten(), b.flatten()) / ( + a.norm() * b.norm())).item() per = torch.nn.functional.cosine_similarity(a, b, dim=-1) return glob, per.mean().item(), per.min().item() @@ -181,6 +261,9 @@ def report(tag: str, a: torch.Tensor, b: torch.Tensor) -> None: def main() -> int: args = parse_args() + if args.artifact and args.mode != "kernel": + print("error: --artifact requires --mode kernel", file=sys.stderr) + return 2 pack = torch.load(args.pack, map_location="cpu", weights_only=True) if args.layer not in pack: print(f"error: layer '{args.layer}' not in pack", file=sys.stderr) @@ -243,6 +326,40 @@ def main() -> int: print("error: flash_rt_fp4 not importable — run on Thor", file=sys.stderr) return 2 + if args.artifact: + artifact = torch.load( + args.artifact, map_location="cpu", weights_only=True) + try: + entry, aux_entry = validate_artifact( + artifact, args.layer, out_f, in_f, fvk_fp4) + except ValueError as exc: + print(f"error: {exc}", file=sys.stderr) + return 2 + fold = artifact["fold"] + if fold == "none": + artifact_input, alpha = x2, 1.0 + elif fold == "mean": + artifact_input, alpha = (x2.float() / s_t).half(), 1.0 + else: + static = aux_entry.get("act_scale_static") + alpha = aux_entry.get("act_out_scale") + if not isinstance(static, torch.Tensor) \ + or tuple(static.shape) != (in_f,) or not isinstance(alpha, float): + print("error: invalid actnorm metadata", file=sys.stderr) + return 2 + artifact_input = (x2.float() / static.to(x2.device)).half() + y_artifact = e0m3_artifact_gemm( + artifact_input, entry, fvk_fp4, alpha=alpha).float() + artifact_cos = cosine_stats(y_artifact, y_fp)[0] + print("\nconverted artifact round-trip:") + report(f"artifact ({fold}) vs omega", y_artifact, y_omega) + report(f"artifact ({fold}) vs fp", y_artifact, y_fp) + if artifact_cos < args.min_artifact_cos: + print( + f"error: artifact cosine {artifact_cos:.6f} is below " + f"{args.min_artifact_cos:.6f}", file=sys.stderr) + return 1 + return 0 y_s0 = e0m3_kernel_gemm(x2, w.half(), fvk_fp4).float() y_s1 = e0m3_kernel_gemm((x2.float() / s_t).half(), (w * s_mean).half(), fvk_fp4).float() diff --git a/tools/convert_omega_pack_e0m3.py b/tools/convert_omega_pack_e0m3.py index ef6ec42e..d5a6f720 100644 --- a/tools/convert_omega_pack_e0m3.py +++ b/tools/convert_omega_pack_e0m3.py @@ -42,9 +42,13 @@ from __future__ import annotations import argparse +import os import re import sys +import tempfile import time +from collections.abc import Mapping +from pathlib import Path import torch @@ -58,6 +62,98 @@ "act_scale_table", ) AUX_SCALARS = ("weight_bits", "a_bits", "in_features", "out_features", "rank") +EXPECTED_FULL_PACK_RECORDS = 252 + + +def _tensor(rec: Mapping, name: str, field: str) -> torch.Tensor: + value = rec.get(field) + if not isinstance(value, torch.Tensor): + raise ValueError(f"{name}: {field} must be a tensor") + return value + + +def validate_record(name: str, rec: Mapping) -> tuple[int, int]: + """Validate the complete rank-0 ``dit_svdquant_v1`` record contract.""" + if not isinstance(rec, Mapping): + raise ValueError(f"{name}: record must be a mapping") + if rec.get("format") != "dit_svdquant_v1": + raise ValueError( + f"{name}: unsupported format {rec.get('format')!r}; " + "expected 'dit_svdquant_v1'") + if rec.get("rank") != 0: + raise ValueError(f"{name}: only rank=0 is supported, got {rec.get('rank')!r}") + if rec.get("weight_bits") != 4 or rec.get("a_bits") != 4: + raise ValueError( + f"{name}: expected weight_bits=a_bits=4, got " + f"{rec.get('weight_bits')!r}/{rec.get('a_bits')!r}") + + weight = _tensor(rec, name, "weight_res_q") + if weight.ndim != 2 or not weight.is_floating_point(): + raise ValueError(f"{name}: weight_res_q must be a floating [N,K] tensor") + if not bool(torch.isfinite(weight).all()): + raise ValueError(f"{name}: weight_res_q contains NaN or Inf") + n, k = weight.shape + if n <= 0 or k <= 0 or n % 64 or k % 64: + raise ValueError( + f"{name}: N and K must be positive multiples of 64, got N={n}, K={k}") + if rec.get("out_features") != n or rec.get("in_features") != k: + raise ValueError( + f"{name}: metadata shape mismatch: tensor=({n},{k}), " + f"out_features/in_features={rec.get('out_features')!r}/" + f"{rec.get('in_features')!r}") + + lowrank_a = _tensor(rec, name, "lowrank_A") + lowrank_b = _tensor(rec, name, "lowrank_B") + if tuple(lowrank_a.shape) != (n, 0) or tuple(lowrank_b.shape) != (k, 0): + raise ValueError( + f"{name}: rank-0 lowrank shapes must be ({n},0)/({k},0), got " + f"{tuple(lowrank_a.shape)}/{tuple(lowrank_b.shape)}") + + table = _tensor(rec, name, "act_scale_table") + if table.ndim != 2 or table.shape[0] < 1 or table.shape[1] != k: + raise ValueError( + f"{name}: act_scale_table must have shape [steps,{k}], got " + f"{tuple(table.shape)}") + if not table.is_floating_point() or not bool(torch.isfinite(table).all()) \ + or not bool((table > 0).all()): + raise ValueError(f"{name}: act_scale_table must contain finite positive values") + + rotation_in = _tensor(rec, name, "duquant_rotation_blocks") + rotation_out = _tensor(rec, name, "duquant_rotation_out_blocks") + if tuple(rotation_in.shape) != (k // 64, 64, 64): + raise ValueError( + f"{name}: input rotation shape must be ({k // 64},64,64), got " + f"{tuple(rotation_in.shape)}") + if tuple(rotation_out.shape) != (n // 64, 64, 64): + raise ValueError( + f"{name}: output rotation shape must be ({n // 64},64,64), got " + f"{tuple(rotation_out.shape)}") + if not bool(torch.isfinite(rotation_in).all()) \ + or not bool(torch.isfinite(rotation_out).all()): + raise ValueError(f"{name}: rotation tensors contain NaN or Inf") + + perm = _tensor(rec, name, "duquant_rotation_perm") + if perm.dtype != torch.int64 or tuple(perm.shape) != (k,): + raise ValueError( + f"{name}: permutation must be int64 shape ({k},), got " + f"{perm.dtype} {tuple(perm.shape)}") + if not torch.equal(torch.sort(perm).values, torch.arange(k, dtype=torch.int64)): + raise ValueError(f"{name}: duquant_rotation_perm is not a permutation of [0,{k})") + return n, k + + +def expected_record_count(meta: Mapping, *, layer_regex: str, + selected_count: int, + explicit: int | None) -> int: + if explicit is not None: + if explicit < 1: + raise ValueError("--expected-records must be positive") + return explicit + recipe = meta.get("recipe", "") + if layer_regex or (isinstance(recipe, str) + and recipe.startswith("synthetic fixture")): + return selected_count + return EXPECTED_FULL_PACK_RECORDS def parse_args() -> argparse.Namespace: @@ -73,25 +169,29 @@ def parse_args() -> argparse.Namespace: p.add_argument("--keep-fp16", action="store_true", help="also store the (possibly folded) fp16 weight, " "for offline reference checks") + p.add_argument("--expected-records", type=int, + help="required converted record count; defaults to 252 for " + "a full pack and to the selected count for fixtures/subsets") + p.add_argument("--validate-only", action="store_true", + help="validate schema and record coverage without CUDA") return p.parse_args() def main() -> int: args = parse_args() - - if not torch.cuda.is_available(): - print("error: CUDA is required (quantize kernels run on GPU)", - file=sys.stderr) - return 2 - try: - import flash_rt.flash_rt_fp4 as fvk_fp4 - except ImportError: - print("error: flash_rt_fp4 extension not importable — run this on a " - "machine with FlashRT built (Thor)", file=sys.stderr) + if not args.validate_only \ + and Path(args.pack).resolve() == Path(args.out).resolve(): + print("error: --out must not overwrite the source pack", file=sys.stderr) return 2 pack = torch.load(args.pack, map_location="cpu", weights_only=True) + if not isinstance(pack, Mapping): + print("error: pack must be a mapping", file=sys.stderr) + return 2 meta = pack.get("__meta__", {}) + if not isinstance(meta, Mapping): + print("error: __meta__ must be a mapping", file=sys.stderr) + return 2 names = sorted(k for k in pack if k != "__meta__") if args.layer_regex: rx = re.compile(args.layer_regex) @@ -100,15 +200,38 @@ def main() -> int: print("error: no layers matched", file=sys.stderr) return 2 + try: + expected = expected_record_count( + meta, layer_regex=args.layer_regex, selected_count=len(names), + explicit=args.expected_records) + if len(names) != expected: + raise ValueError( + f"record coverage mismatch: selected {len(names)}, expected {expected}") + shapes = {name: validate_record(name, pack[name]) for name in names} + except (KeyError, TypeError, ValueError) as exc: + print(f"error: {exc}", file=sys.stderr) + return 2 + print(f"validated {len(names)}/{expected} records") + if args.validate_only: + return 0 + + if not torch.cuda.is_available(): + print("error: CUDA is required (quantize kernels run on GPU)", + file=sys.stderr) + return 2 + try: + import flash_rt.flash_rt_fp4 as fvk_fp4 + except ImportError: + print("error: flash_rt_fp4 extension not importable — run this on a " + "machine with FlashRT built (Thor)", file=sys.stderr) + return 2 + device = torch.device("cuda") weights: dict = {} aux: dict = {} t0 = time.time() for i, name in enumerate(names): rec = pack[name] - if rec.get("format") != "dit_svdquant_v1": - print(f"skip {name}: format={rec.get('format')!r}") - continue w = rec["weight_res_q"].to(device=device, dtype=torch.float16, non_blocking=False).contiguous() table = rec["act_scale_table"].float() @@ -124,10 +247,7 @@ def main() -> int: w = (w * r.to(device=device, dtype=torch.float16) .unsqueeze(0)).contiguous() act_out_scale = c - n, k = w.shape - if k % 16 != 0: - print(f"skip {name}: K={k} not divisible by 16") - continue + n, k = shapes[name] packed = torch.empty(n, k // 2, dtype=torch.uint8, device=device) # Zero-init: tile-interleaved SFB pads K to 64-element atoms and the @@ -161,14 +281,31 @@ def main() -> int: f"({time.time() - t0:.1f}s)") torch.cuda.synchronize() + if set(weights) != set(names) or set(aux) != set(names): + raise RuntimeError( + "internal coverage error: not all validated records were converted") out = { "format": OUTPUT_FORMAT, + "schema_version": 1, "source_pack_meta": meta, + "source_record_count": len(pack) - int("__meta__" in pack), + "selected_record_count": len(names), + "selected_layers": names, "fold": args.fold, "weights": weights, "aux": aux, } - torch.save(out, args.out) + out_path = Path(args.out) + out_path.parent.mkdir(parents=True, exist_ok=True) + with tempfile.NamedTemporaryFile( + prefix=out_path.name + ".", suffix=".tmp", + dir=out_path.parent, delete=False) as tmp: + tmp_path = Path(tmp.name) + try: + torch.save(out, tmp_path) + os.replace(tmp_path, out_path) + finally: + tmp_path.unlink(missing_ok=True) print(f"wrote {args.out}: {len(weights)} layers, fold={args.fold}, " f"{time.time() - t0:.1f}s total") return 0 diff --git a/tools/gen_omega_pack_fixture.py b/tools/gen_omega_pack_fixture.py index 74c6d940..0deeefc4 100644 --- a/tools/gen_omega_pack_fixture.py +++ b/tools/gen_omega_pack_fixture.py @@ -7,8 +7,8 @@ python tools/gen_omega_pack_fixture.py --out /tmp/fixture_pack.pt python tools/convert_omega_pack_e0m3.py --pack /tmp/fixture_pack.pt \ --out /tmp/fixture_e0m3.pt --fold none --keep-fp16 # Thor - python tools/check_omega_e0m3_consumer.py --reference fp16 \ - --artifact /tmp/fixture_e0m3.pt + python tools/check_omega_e0m3_layer.py --pack /tmp/fixture_pack.pt \ + --artifact /tmp/fixture_e0m3.pt --mode kernel # Thor Design notes (what makes the fixture a real test and not a toy):