NVFP4 support - #1
Open
sirish-gambhira wants to merge 14 commits into
Open
Conversation
LazyTurtle copied checkpoint tensors into shell params with a plain dtype cast. For an mxfp8 checkpoint the stored tensor is only the F8_E4M3 mantissa payload; the magnitude lives in a sibling weight_scale_inv holding E8M0 exponent bytes, one per 32-element block. Casting to bf16 silently dropped that scale, so every quantized weight was wrong by 2^(e-127). The scale was never even requested: t_params comes from the shell module's named_parameters(), and a plain nn.Linear has no weight_scale_inv attribute, so _resolve_checkpoint_tensor_source was never asked for it. Measured on Minimax-M3-0602, layers.3.self_attn.q_proj: the materialized tensor was bit-identical to the raw fp8 payload (absmax 448, std 148.6) instead of the scaled weight (absmax 0.2812, std 0.0381). Downstream, MoE routing collapsed onto 6.8 of 128 experts per layer with one expert taking every token in 55 of 57 layers; with scales applied it reaches 127.9 of 128 with no expert above a 17.5% share. dequantize_fp8() could not be reused: _expand_scale treats scale_inv as a float multiplier and broadcasts by rank, so E8M0 bytes would multiply by 120.0 rather than 2^-7, and a [8192,192] scale will not broadcast against a [8192,6144] weight. Adds dequantize_block_fp8 / block_scale_multiplier, which decode E8M0 and expand each dim by target_dim // scale_dim. Verified bit-for-bit (maxdiff 0) against modelopt's MXFP8QTensor.dequantize on attention, MoE expert, and shared-expert weights. Routes every materialization read through _read_checkpoint_tensor, covering both _copy_checkpoint_tensors_into_submodule and the parallel _materialize_direct_meta_tensors path (param/buffer, single and concat). Layers in ignored_layers carry no scale tensor, so the lookup misses and they keep the plain-cast path -- which is why the F32 routers loaded correctly even while the bug was live. Not covered: _load_checkpoint_tensors_for_module_path still returns raw payloads. It has 8 callers in models/base.py and returns the _scale_inv entries alongside the weights, so dequantizing there could double-apply; those call sites need review separately. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
GPTQ minimizes the layer output error. For weight-only quantization inference computes Ŵ·X, so H = 2·X·Xᵀ is the right curvature. Under w4a4 the kernel computes Ŵ·Q(X), and a Hessian built from X solves for inputs that never occur. At 4-bit activations that gap is large; at 16-bit it does not exist, which is why this has not mattered until now. Activations reach the solve through exactly one path -- process_batch -> compute_hessian_xtx -> self.H -- so quantizing there also propagates to dead-column detection, act-order, act-group-aware grouping, and the damped Cholesky inverse. Adds GPTQConfig.act_format / act_group_size. act_format=None keeps today's weight-only behavior exactly: no activation quantizer is constructed and process_batch is unchanged. The quantizer is applied explicitly rather than captured by a hook. This is deliberate: a forward_pre_hook fires before a module's forward body and so observes activations *before* any input quantizer, which would silently accumulate a clean Hessian while every config flag looked correct. The regression test asserts H actually differs (~4.8% relative on random input) precisely to catch that failure mode. Ordering in process_batch matters twice: the quantizer runs before the TP zero-padding, so pad columns never enter a 16-wide block's amax, and after any input-scale transform already baked into `inp`, since that is what the kernel quantizes. fake_quantize_nvfp4 wraps torchao's nvfp4_quantize + NVFP4Tensor.dequantize rather than reimplementing the grid, so it cannot drift from the weight path. Following that path, per_tensor_scale is omitted -- block scales only. Scope: this is tier A of docs/w4a4_hessian.md, min ‖(W−Ŵ)·Q(X)‖², matching Model-Optimizer's GPTQHelper. It is an approximation: the exact objective min ‖W·X − Ŵ·Q(X)‖² carries a W·(X−Q(X)) bias no Ŵ can cancel, which needs a second accumulator (tier B, cf. Quark GPTAQ's ∆(XXᵀ)). Known gap: NVFP4 activations use a static per-tensor global scale that is calibrated, plus dynamic per-block scales. Nothing populates the global scale yet -- Model-Optimizer gets it from a separate max_calibrate stage that GPTQModel has no equivalent of. Until that lands, the calibration grid is close to but not identical to a static-global-scale serving kernel. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
NVFP4 activations are scaled twice: a static per-tensor global scale frozen at
calibration, and dynamic per-16-block E4M3 scales. Only the block scales can be
derived while accumulating the Hessian, so w4a4 calibration was quantizing
against a block-only grid -- close to, but not the same grid a static-global
kernel consumes.
The global scale cannot be computed here by construction: it is a property of
the whole calibration set, and we are inside the pass that would produce it.
Both reference pipelines resolve this by ordering rather than by a two-pass
GPTQ -- Model-Optimizer freezes it in an earlier max_calibrate stage, and
tore-quant emits it from Phase 1 (`ptq.py --emit-amax-sidecar`) for Phase 2 to
consume. This follows the latter: read the sidecar the pipeline already
produces instead of adding a calibration stage.
Scale formula mirrors tore-quant's amax_to_scale exactly:
global_scale = amax / 6.0 / 448.0 # amax / F4_E2M1_MAX / F8_E4M3_MAX
The amax is a plain max in that pipeline, not percentile- or MSE-clipped.
act_amax_key_rules exists because the scale is keyed by *input site*, not by
module. w1 and w3 consume the same MoE input and therefore share one entry
(`layer{N}.moe_input`) while w2 sees the post-SwiGLU activation and gets its
own; q/k/v share their attention input by the same logic. A per-module lookup
would hand siblings different scales than the kernel applies -- the exact
mismatch the sidecar exists to prevent -- so a test pins the sharing.
Resolution is lazy rather than done in __init__: `self.name` is not final at
construction time (the looper names modules afterwards), so an early lookup
reads the wrong key. This was caught by the new tests, and in production would
have surfaced as a wrong scale rather than a crash.
A missing key raises instead of falling back to block-only scales; a silent
fallback would quantize against a different grid than the kernel with no
signal. Non-finite and non-positive amax are rejected.
The sidecar loader is lru_cached so thousands of modules read the file once,
and accepts both a flat {key: amax} mapping and the pipeline's
{"amax": {...}} shape.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Broly Security ScanNote Summary 1 actionable finding(s) in this PR
All actionable items are in the table below. 2 finding(s) below the No finding is at or above
Dismiss false positivesTick a box to dismiss the finding; untick it to bring the finding back. That is the same as replying
Note Re-scan this PR anytime with
|
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
…rics
GPTQModel's NVFP4 export stored only per-block E4M3 scales. ModelOpt, Quark
and tore-quant all use two-level scaling: a per-tensor weight_scale_2 =
amax/(6*448) plus block scales. Without it the block scales carry absolute
magnitudes, and for realistic weight tensors block_amax/6 sits below E4M3's
smallest normal (2^-6) -- measured 100% of blocks clamped at the floor on an
N(0, 0.02) weight -- so every block quantized on a compressed grid. Two-level
round-trip error is 1.36x lower on that tensor (0.096 vs 0.131 relative).
RTN path: quantize_nvfp4_weight computes the canonical greedy global scale
(verified equal to ModelOpt's global_amax/(E2M1_MAX*m_fp8) and tore-quant's
amax_to_scale) and returns a 3-tuple.
GPTQ path: the loop has already fixed each block scale s_i and the packed
codes are encoded against them, so the split must satisfy
block_e4m3_i * ws2 == s_i exactly. A greedy ws2 cannot (the ratio of two E4M3
mantissas is generally unrepresentable); a power-of-two ws2 =
2^ceil(log2(max(s)/448)) only shifts exponents, leaving mantissas intact.
Verified bit-exact recovery; an equality check raises rather than silently
emitting codes that no longer match their scales.
Both read paths apply it: dequantize_weight folds ws2 into the block scales
(one pass, exact in fp32), _native_weight passes it as NVFP4Tensor's
per_tensor_scale. TorchFP4Linear -- which executes directly from checkpoint
weights -- now accepts weight_scale_2 too; it previously ignored the tensor
that ModelOpt-layout checkpoints (e.g. MiniMax-M3-NVFP4) carry on every
quantized module, dequantizing every block off by a constant factor.
Also switch x/s to x * s.reciprocal() in NVFP4Quantizer.quantize and
_pack_with_gptq_scales, matching torchao's nvfp4_quantize ("multiply by
reciprocal instead of dividing to match MSLK triton kernel numerics"). The
1-ULP difference can flip an E2M1 code on rounding boundaries, which breaks
bit-identical packing claims.
An unpacked module defaults weight_scale_2 to 1.0 (identity), so one-level
checkpoints load unchanged. test_direct_pack_matches_torchao updated to
assert the two-level layout.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ze driver, run.sh
Three-stage pipeline for quantizing Minimax-M3-0602 routed experts to NVFP4
with w4a4 calibration:
emit_act_amax.py (step 0) -- emits the per-layer activation amax sidecar the
static NVFP4 global scale needs. Two keys per MoE layer, keyed by input site
because modules sharing an input share a scale: layer{N}.moe_input
(gate/up_proj) and layer{N}.w2_input (down_proj, post-SwiGLU). Plain max over
|x| reduced across a layer's experts, matching tore-quant's save_amax_sidecar
peer-max semantics. Captures via down_proj pre-hooks when the loader exposes
per-expert Linears (GPTQModel does: it splits the checkpoint's fused 3D expert
tensors), falling back to wrapping the fused module's _apply_gate -- whose
return value IS the down_proj input -- and refusing to emit a partial sidecar
if neither works. Streams layer windows through one GPU with per-window
resume, the p1b_expert_coverage harness pattern.
quantize_w4a4.py -- GPTQ->NVFP4 driver. Quantizes only routed experts (57
layers x 128 experts x 3 projections); attention, shared experts, dense MLP
0-2, routers, embeddings and the vision tower pass through, matching the
tore-quant recipe format map. RTN-fallback thresholds are per-projection at
each weight's in_features (gate/up 6144, down 3072): H = XtX has
rank <= min(n_tokens, cols), so a module below its own input dim would be
GPTQ-solved on a singular Hessian propped up by damping. The library's 0.5%
default sits below both dims and is scale-invariant besides (RTN iff routing
share < pct/(100*top_k) -- more calibration data cannot lower it). With
ACT_AMAX set, the Hessian is accumulated on Q(X) with the sidecar-derived
static global scale; without it the run is weight-only and says so. Post-check
asserts the saved layout (uint8 weight + e4m3 weight_scale + fp32
weight_scale_2) instead of trusting the log.
run.sh -- preflight (reads file CONTENT, not stat: the /scratch checkpoint
copy passes ls and fails on read), disk/GPU checks, predicted per-projection
RTN rates from the measured coverage matrix, then the driver. All knobs are
env vars; PREFLIGHT_ONLY=1 and LAYERS=N smoke-test modes included.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…nt repair Config (from_quant_config path): accept the `fp4` format alias when the method resolves to NVFP4, matching the constructor path's _normalize_nvfp4_kwargs, and route NVFP4Config targets through the same kwargs normalization. Config (GPTQ validation): format=nvfp4 now rejects what cannot land on the NVFP4 grid -- non-RTN fallback strategies (asymmetric integer grids whose zero-points the packer cannot represent), fallback smoothing (per-block factors make the stored scale non-E4M3-exact, so packing stops being lossless), and mock_quantization (integer rounding, one scale per 128-column block). GPTQ NaN-loss path: under format=nvfp4 a NaN loss now goes to the RTN fallback (which quantizes through NVFP4Quantizer, staying on-grid) instead of the mock-quantization retry, which is off-grid for the same reasons as above. Also repairs tests/test_nvfp4_export.py from commit 6466766, where zero-context partial staging misplaced two lines -- the two-level docstring landed inside a constructor argument list (SyntaxError: positional argument follows keyword argument) and a weight_scale_2 assert landed in the wrong test. The worktree version was always correct (37 tests green); only the committed copy was broken. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
- Point MODEL at /data/huggingface/Minimax-M3-0602: the /scratch copy returns EIO on config.json and every safetensors shard (btrfs read errors) while passing `ls`, so the failure only shows up on read. - WINDOW=5 (was 1): the co-tenant pipeline that forced single-layer windows is gone; 5-layer windows amortize materialize/free and the peak (~100 GiB transient sdpa masks) fits an idle B200. - Checkpoint after every window (CKPT_EVERY=1) so a kill resumes at the last 5-layer boundary; two earlier runs were lost to co-tenant OOM kills. - OOM retry helper for materialization under memory pressure. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
… dequant The GPTQ quantize flow does not read weights through the LazyTurtle materialization path fixed in 3e01e1c. The looper materializes quant targets with role="quant_source" (module_looper.py:1099) and forward modules with role="forward" (:1067); both pull raw tensors via checkpoint_tensors_for_submodule and decode fp8 through _build_decoder_quant_source_module -> dequantize_fp8. That function treated scale_inv as a float multiplier: mxfp8's weight_scale_inv is uint8 UE8M0 exponent bytes (~110-130), which exceed 1, so _fast_scale_arg flipped into divide mode -- dividing by the raw byte value instead of multiplying by 2^(b-127). Measured on Minimax-M3-0602 layers.3 q_proj: absmax 3.94 (divide by ~120) vs correct 0.2812. Every weight GPTQ would have solved against was wrong by a block-dependent factor, silently. A uint8 scale can only be E8M0 exponent bytes -- a real float scale is never serialized as uint8 -- so dequantize_f8_e4m3 now routes uint8 scale/scale_inv straight through dequantize_block_fp8 (verified bit-identical to modelopt's MXFP8QTensor.dequantize). Decoding the bytes in place and continuing was not enough: the fast CPU path also mishandles 2D block-grid scales on small tensors (the reference path, taken for large tensors, is fine), which is why the earlier large-tensor verification passed while a small synthetic failed. Float scales keep their existing paths untouched. Fixes the quant_source and forward decoder roles, TorchFP8Linear, and every other dequantize_fp8 caller in one place. Verified bit-identical to the reference decode on real M3 attention and expert tensors; regression tests pin the multiplicative semantics (4.0 * 2^-3 = 0.5, not 4.0/124), uint8-vs- reference equality, float-path stability, and the real-checkpoint magnitude. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The looper names modules layer-relative -- GPTQ's `self.name` for an M3 expert is `mlp.experts.1.gate_proj`, with no layer index for the sidecar key rules to capture. Resolution now uses NamedModule.full_name (the full dotted path, `...layers.3.mlp.experts.1.gate_proj`), falling back to `self.name` when no NamedModule wraps the module (bare-module tests, HF-optimum path). Found by the LAYERS=3 smoke run: every rule missed, the lookup fell through to the raw name, and the fail-loud KeyError stopped the run before any wrong scale was applied -- the failure mode the guard exists for. The regression test reproduces the exact shape: layer-relative name + full_name on the wrapper, asserting the rule captures the layer index from the latter. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…w4a4 pipeline
splice_source_passthrough.py (new, run.sh phase 3, SPLICE=1 default): GPTQModel's
save materializes every module, so passthrough tensors come out dequantized-BF16
under runtime naming -- numerically right but not the serving format map, and not
loadable by a source-arch stack (both checkpoints declare
MiniMaxM3SparseForConditionalGeneration, but gptqmodel writes
model.language_model...mlp.experts.N.gate_proj where the source and serving
loaders expect language_model.model...block_sparse_moe.experts.N.w1). The splice
walks the SOURCE tensor list and swaps in the quantized experts under source
naming: passthrough stays byte-identical MXFP8/BF16/F32, quantized experts gain
{weight U8, weight_scale E4M3, weight_scale_2 F32, input_scale F32}, with
input_scale = amax/6/448 from the same sidecar the Hessian was calibrated on --
so the serving activation scale and the calibration scale agree by construction.
Verified on the LAYERS=3 smoke checkpoint: index keyset exact (46,606 keys),
17/17 sampled passthrough tensors byte-identical, 18/18 sampled expert quads
bit-equal with exact f32 input_scale.
quantize_w4a4.py: post-check now probes inside the quantized layer set -- the
layer-blind probe hit passthrough layer 10 first and failed a good checkpoint.
gptq.py: RTN-fallback WARN logs full_name -- self.name is layer-relative and
repeats across all 57 layers, which would make multi-layer fallback rosters
collide silently.
run.sh: COLUMNS=360 so headless logbar stops truncating the per-module loss
column to `...`; emits <OUT>/quant_roster.json mapping every RTN module with
observed tokens + threshold (GPTQ-solved = the complement); phase 3 wiring.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…al scale The weight_scale_2 fix (6466766) covered the pack stage and the no-scales RTN path, but pack deliberately preserves the loop's scales -- and the loop's NVFP4Quantizer.find_params still derived them single-level: clamp(block_amax/6, 2^-6, 448). For realistic weight magnitudes those absolute values sit below E4M3's smallest normal, so blocks pin at the floor and quantize on a compressed grid. Measured on the LAYERS=3 smoke checkpoint against MiniMax-M3-NVFP4 (ModelOpt RTN of the same weights): 96.8% of block scales at the floor, reconstruction error 1.4-1.8x the donor's (0.12-0.15 vs 0.082 relative) -- on our RTN-fallback modules too, which should have matched the donor's RTN nearly exactly. cos(ours, donor) was 0.993+, so every layout and losslessness check passed while quality was silently off. find_params only ever sees one 16-wide group, so the per-tensor scale must be fixed from the full weight before the loop: prime_global_scale(W), called in the main quantize path and in _fallback_quantize. Primed, find_params derives block scales as E4M3 ratios against it (torchao's two-level path, verified bit-equal); unprimed behaviour is unchanged single-level. The global scale is snapped to a power of two, 2^ceil(log2(amax/(6*448))) (tore-quant's lossless-conversion snap), so the loop's effective scales split exactly at pack time with no plumbing changes: _pack_with_gptq_scales re-derives gs -- or gs/2 when the max block ratio E4M3-rounds down to exactly 224; either decomposition reproduces the identical product, and the in-pack equality guard holds. Cost vs the donor's greedy per-tensor scale is at most one bit of E4M3 range headroom. Tests: primed find_params bit-equal to nvfp4_quantize(per_tensor_scale=gs); floor-clamp fraction must drop below 50% on N(0, 0.02) weights (was ~100%); pack split exact with a power-of-two recovered scale in {gs, gs/2}. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…() baseline base.py passes quantize_config.fallback into module_looper.loop(); it threads through run_layer_stage -> preprocess -> clone_gptq_config_for_module, which applied it AFTER the dynamic-override block -- so every per-module dynamic `fallback` was silently clobbered by the global value in the standard quantize() path. Found via the LAYERS=3 smoke roster: all three expert projections fell back at threshold=6144 even though the per-projection rank thresholds set down_proj to 3072 (its in_features), needlessly RTN'ing 13 down_proj experts whose Hessians were full-rank. The caller-supplied fallback is now the baseline, applied before the dynamic block, so a dynamic override merges onto it and wins. Regression test pins both directions: an overridden module resolves to its dynamic threshold, a non-overridden module keeps the baseline. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…arams raises The unprimed branch of NVFP4Quantizer.find_params derived absolute amax/6 block scales, which floor-clamp at E4M3's 2^-6 normal minimum on realistic weight magnitudes (measured 96.8% of blocks on Minimax-M3 experts, 1.4x reconstruction error). Keeping it as a silent fallback meant any caller that forgot prime_global_scale would reproduce that bug undetectably -- it cost a full smoke-run cycle plus a donor-checkpoint comparison to find the first time. An unprimed caller is a bug, not a mode: find_params now raises with the priming instruction. Audit of every path that can hold an NVFP4Quantizer, as the precondition for making this fatal: - GPTQ.quantize: primed before the first full-weight find_params; every other site (static_groups precompute, act-order branches, CPU-fallback re-runs, per-group loop) is downstream in the same call. - GPTQ._fallback_quantize: primed before the RTN per-block loop. - GPTAQ / FOEM: subclass GPTQ and inherit the NVFP4 quantizer but run their own quantize() with full-weight find_params -- these were UNPRIMED and would have produced floor-clamped NVFP4 silently; now primed. - rtn.py constructs a plain Quantizer (never NVFP4); the weight-only NVFP4 lifecycle packs via quantize_nvfp4_weight with no quantizer object. The old unprimed-vs-torchao-single-level parity test is repurposed as the raise regression; the lossless-pack test primes as production does; primed two-level parity was already covered. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
|
||
| if [[ -n "${ACT_AMAX}" ]]; then | ||
| [[ -f "${ACT_AMAX}" ]] || die "ACT_AMAX sidecar not found: ${ACT_AMAX}" | ||
| "${PY}" -c "import json,sys; d=json.load(open('${ACT_AMAX}')); d=d.get('amax',d); print(f' act amax ${ACT_AMAX} ({len(d)} entries)')" \ |
| the NVFP4 conversion pipeline emits. | ||
| """ | ||
|
|
||
| with open(path, "r") as handle: |
|
|
||
| start_layer = 0 | ||
| if os.path.exists(CKPT): | ||
| state = torch.load(CKPT, map_location="cpu", weights_only=False) |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Describe the bug, fix, or feature clearly and briefly.
What Changed
Tests
Every working PR must include at least one new simple, fast, targeted unit test when the change affects behavior, a bug fix, or a regression path.
Paste the exact test commands and results here:
Review Requirements
AI-assisted code is welcome.
Every changed file must still be properly reviewed by a human before the PR is opened as ready for review.
We will not accept PRs that are effectively unreviewed AI output. Non-human-reviewed changes often introduce obscure structure, mismatched APIs, project-inconsistent code patterns, or unnecessary monkeypatching instead of a correct fix or clean feature expansion.
Notes
Add any migration notes, risks, compatibility concerns, or reviewer guidance here.