diff --git a/docs/w4a4_hessian.md b/docs/w4a4_hessian.md new file mode 100644 index 000000000..692e30562 --- /dev/null +++ b/docs/w4a4_hessian.md @@ -0,0 +1,286 @@ +# w4a4: accumulating the GPTQ Hessian on quantized activations + +## The problem + +GPTQ minimizes the layer output error. For weight-only quantization inference +computes `Ŵ·X`, so the objective and its Hessian are: + +``` +min_Ŵ ‖ W·X − Ŵ·X ‖² = min_Ŵ ‖ (W − Ŵ)·X ‖² -> H = 2·X·Xᵀ +``` + +Under w4a4 inference computes `Ŵ·Q(X)`. If calibration accumulates `H` on `X` +while the kernel feeds `Q(X)`, the solved weights are optimal for inputs the +model never sees. At 4-bit activations that gap is large — unlike w4a16, where +`Q(X) = X` and the question does not arise. + +In GPTQModel, activations enter in exactly one place: `GPTQ.process_batch` +(`gptqmodel/quantization/gptq.py`) reshapes the layer input and hands it to +`compute_hessian_xtx`; `add_batch` accumulates into `self.H`. After that `X` is +never referenced again. `self.H` feeds dead-column detection, act-order +permutation, act-group-aware grouping, and the damped Cholesky inverse — so +substituting at accumulation time propagates to all of them. + +## Two tiers + +| tier | objective | statistics | reference implementation | +|---|---|---|---| +| **A** | `min ‖(W−Ŵ)·Q(X)‖²` | `H` only | Model-Optimizer `GPTQHelper` | +| **B** | `min ‖W·X − Ŵ·Q(X)‖²` | `H` + `∆(XXᵀ)` | Quark GPTAQ / Qronos | + +Tier A substitutes the input and treats activation-quantization error as an +irreducible constant. It is an approximation: the true residual `W·X − Ŵ·Q(X)` +is **not** of the form `(W − Ŵ)·M` for any `M`, because the bias `W·(X − Q(X))` +cannot be cancelled by any choice of `Ŵ`. It is nonetheless strictly better than +the status quo. + +Tier B is the exact objective and absorbs that bias, at the cost of a second +accumulator and changes to the `fasterquant` solve. + +Tier A is a prerequisite for B. Recommendation: land A, measure against a w4a4 +perplexity baseline, then decide on B. + +--- + +## What the reference implementations do + +### Model-Optimizer — does exactly tier A + +`modelopt/torch/quantization/utils/calib_utils.py`, `GPTQHelper.setup()`: + +```python +def hessian_forward(self, input, *args, **kwargs): + inp = input.to_local() if hasattr(input, "to_local") else input + if self.input_quantizer is not None and self.input_quantizer.is_enabled: + hessian_input = self.input_quantizer(inp) # Q(X) + else: + hessian_input = inp # X + gptq_helper.hessian, gptq_helper.n_samples = update_hessian( + hessian_input, gptq_helper.hessian, gptq_helper.n_samples + ) + return self._forward_no_gptq_hessian(input, *args, **kwargs) + +bind_forward_method(self.module, hessian_forward, self.CACHE_NAME) +``` + +Three things to take from this: + +1. **One branch, no separate w4a4 path.** The Hessian goes on `Q(X)` whenever the + input quantizer is enabled. +2. **It reuses the same `input_quantizer` object the inference path uses** + (`conversion.py` does `input = self.input_quantizer(input)`), so the + calibration scheme cannot drift from the serving scheme. +3. **It patches `forward`, it does not use a hook.** This is essential. A + `forward_pre_hook` fires before the forward body and therefore sees the input + *before* the internal `input_quantizer` call. **The quantized activation + cannot be captured with a hook — the quantizer must be applied explicitly.** + A port that uses hooks will silently accumulate clean activations. + +`update_hessian` itself is standard (`H += (2/n)·XXᵀ` with incremental +rescaling), including a zero-token guard for MoE experts that receive no tokens. + +#### How the static activation scale is resolved + +NVFP4 activations use a **static per-tensor global scale** plus **dynamic +per-16-block E4M3 scales**. The global scale is a property of the whole +calibration set, which appears to create a chicken-and-egg with the calibration +pass itself. + +Model-Optimizer resolves it by **ordering, not by a two-pass GPTQ**: the input +quantizer's amax is calibrated in an earlier, separate stage (`max_calibrate`), +so by the time the GPTQ stage runs the quantizer is already calibrated and +enabled. GPTQ stays single-pass. This works because the scale belongs to the +quantizer, not to GPTQ. + +#### Paths in Model-Optimizer that deliberately do *not* do this + +- `local_hessian_calibrate` (`model_calib.py`) feeds **clean** activations. Its + hooks are `register_forward_pre_hook`s capturing `args[0]` — pre-quantizer, per + the point above. Correct there, because it is a different object: a per-cin-block + Hessian used to weight a *weight*-error metric + (`local_hessian_error(x, xq)` computes `dw = x − xq` over `cout`). +- **AWQ-lite / AWQ-clip** *do* use quantized activations, gated on + `is_input_quantized = module.input_quantizer.is_enabled`, then searching + against `self.input_quantizer(input)`. Note AWQ-clip calls `max_calibrate` on + the current batch and uses the result immediately — a running max. Acceptable + for a scale search; not advisable under a matrix that gets Cholesky-inverted. + +### Quark — a different axis, plus the tier-B machinery + +Quark's GPTQ (`quark/torch/algorithm/gptq/gptq.py`) exposes +`add_batch_quantized` / `add_batch_nonquantized`, which resembles this +distinction but is not it. In Quark, "quantized" means **the input produced by +running the block with already-quantized preceding weights** — sequential +weight-error propagation. Evidence: + +- The hook is a `register_forward_hook` capturing `input[0]`, i.e. pre-quantizer. +- The two passes differ by `layer_inputs` vs `orig_layer_inputs`, and the clean + pass runs under `RestoreOriginalWeights(layer)`. Restoring *weights* is what + produces the clean input, so the axis is weights, not activations. +- Plain GPTQ's `add_batch_nonquantized` raises + `ValueError("We don't need it for GPTQ")`. + +So Quark is not a precedent for activation-quantized Hessians. It does, however, +implement the asymmetric objective needed for tier B, in +`quark/torch/algorithm/gptaq/gptaq.py` (GPTAQ, arXiv 2504.02692) and +`qronos/qronos.py`: + +```python +# add_batch_quantized: H = X̃ X̃ᵀ +self.H += inp.matmul(inp.t()) + +# add_batch_nonquantized: ∆(XXᵀ) = (X − X̃) X̃ᵀ +delta_input = original_input - self.q_input +self.delta_X_Xt += delta_input.matmul(self.q_input.t()) +``` + +That is `min_Ŵ ‖W·X − Ŵ·X̃‖²`, structurally identical to the w4a4 objective with +`X̃ = Q(X)`; `∆(XXᵀ)` is the cross term tier A drops. GPTAQ also runs two +`block_forward`s per batch, so a second pass is accepted practice. + +--- + +## Proposed change to GPTQModel (tier A) + +### Config + +`grep` for `act_bits|activation_bits|act_quant|activation_scheme` in +`gptqmodel/quantization/config.py` returns nothing — GPTQModel has no notion of +activation quantization today. This adds a capability, not a flag. Add to +`GPTQConfig`: + +```python +act_format: Optional[FORMAT] = None # None = weight-only (today's behavior) +act_group_size: int = 16 # NVFP4 block +``` + +Validation: `act_format=nvfp4` should require `format=nvfp4`, and should reject +`mock_quantization` for the same reason the weight path does. + +### `gptq.py` + +Insert in `process_batch`, after `reshaped_inp = reshaped_inp.contiguous()` and +**before** the `_tp_pad_cols` zero-padding: + +```python +reshaped_inp = reshaped_inp.contiguous() + +# w4a4: the kernel sees Q(X), so the Hessian must describe Q(X), not X. +# Quantize before TP zero-padding so pad columns do not enter block statistics. +if self._act_quantizer is not None: + reshaped_inp = self._act_quantizer(reshaped_inp) + +if self._tp_pad_cols: + ... +``` + +Ordering matters twice: + +- **Before TP padding** — `_tp_pad_cols` appends zero columns; quantizing after + would let those zeros enter a 16-wide block's amax and shift its scale. +- **After any input-scale / smoothing transform** — `Q` must see the + post-transform activation, because that is what the kernel quantizes. + +### The activation quantizer + +`NVFP4Quantizer` (`gptqmodel/quantization/quantizer.py`) is **not** reusable: +`find_params` raises `NotImplementedError` for `weight=False`, and it derives one +scale per row of a group-sliced weight. Activations need per-token × per-16-block +along the last dim of a `(tokens, cin)` tensor. + +Use torchao's `nvfp4_quantize(x, block_size=16)` — already imported in +`gptqmodel/quantization/dtype.py` and used by the weight path — then dequantize +back to the staging dtype so `compute_hessian_xtx` is untouched. + +Following Model-Optimizer, the quantizer should be an object shared with (or +provably identical to) whatever the serving kernel applies, rather than a +reimplementation kept in sync by hand. + +## Consequences to watch + +- **Dead columns.** `dead = diag(H) == 0` currently sets `H[dead, dead] = 1`, + silently skipping correction for that column. Small-magnitude channels can + flush to zero under FP4, creating dead columns that did not exist in the clean + Hessian. Log a count when activation quantization is enabled. +- **act-order / act-group-aware.** Both sort on `diag(H)`. With `Q(X)` the + ordering is by quantized-activation importance — arguably more correct for + w4a4, but group assignments will differ from a weight-only run, so results are + not comparable across the flag. +- **Cost.** One quantize+dequantize per batch, negligible against the `XᵀX` GEMM, + but it adds a full-size fp32 temporary — relevant given the existing OOM + fallback in `process_batch`. +- **MoE.** Each expert sees only its routed tokens. With few tokens per expert a + dynamic scale is noisier; the static global scale mitigates this, which is + another argument for matching the serving scheme exactly. + +## The static global scale: how `tore-quant` does it + +`tore-quant` runs the production NVFP4 pipeline and answers the scale questions +concretely. + +**Formula** (`src/tore_quant/convert_mxfp4_to_nvfp4.py`): + +```python +def amax_to_scale(amax: torch.Tensor) -> torch.Tensor: + """NVFP4 second-level scale formula.""" + return amax.float() / 6.0 / 448.0 +``` + +i.e. `global_scale = amax / (F4_E2M1_MAX × F8_E4M3_MAX)`, serialized per module +as a fp32 scalar `.input_scale`. + +**The amax is a plain max, not clipped.** It comes from modelopt max-calibration +buffers, distilled by `save_amax_sidecar` (`src/tore_quant/ptq.py:525`). + +**Two phases, not one pass.** Phase 1 (`ptq.py --emit-amax-sidecar`) writes +`amax_per_layer.json`; Phase 2 (`convert_mxfp4_to_nvfp4`) consumes it. Same shape +as Model-Optimizer: the global scale is frozen by an earlier stage, so the +quantization step never has to solve the chicken-and-egg. + +**The scale is shared by input site, not per module** — the part that matters +most for GPTQModel: + +``` +layer{L}.moe_input <- max input amax over routed-expert w1/w3 +layer{L}.w2_input <- max input amax over routed-expert w2 +``` + +`w1` and `w3` consume the same MoE input, so they share one scale; `w2` sees the +post-SwiGLU activation and gets its own. A peer-max sync +(`fixup_moe_expert_amax`) then makes the value identical across every expert in a +layer and across ranks, so a per-layer scalar loses nothing. + +By the same logic `q/k/v` would share their attention input. + +**MoE calibration caveat, from `--calib-all-experts` (default False):** forcing +every token through every expert "distorts per-expert input amax". Native top-k +routing plus the peer-max sync is the supported path. + +### Consequence for this implementation + +Tier A currently passes `per_tensor_scale=None`, so it uses block scales only. +Closing that gap needs more than a number: `GPTQ` is per-module and has no way to +express "these modules share an activation scale". A faithful implementation +needs the scale keyed by **input site** and supplied from a sidecar, either + +- loaded from an `amax_per_layer.json`-style file (matches the existing + pipeline, no new calibration stage), or +- collected by a GPTQModel-side amax pass with sibling-sharing logic. + +The first is much cheaper and reuses a file the pipeline already produces. + +## Open questions + +1. Whether the input-scale (`inscale1`) transform applies to activations at + inference. If so `Q` sits after it, and any amax must be measured + post-transform. +2. Whether `X̃` should compose both error sources — upstream weight quantization + *and* activation quantization, i.e. `X̃ = Q(X_weightpath)`. That is what the + deployed model sees, and neither reference implementation does both at once. + +## Verification plan + +- A layer where `H_quant` and `H_clean` differ measurably (otherwise the change + is a no-op and something is misconfigured — e.g. quantizing via a hook). +- End-to-end: quantizing with `H_quant` should beat `H_clean` on w4a4 perplexity. + Without this the change is unfalsifiable. diff --git a/gptqmodel/looper/gptq_processor.py b/gptqmodel/looper/gptq_processor.py index a1dc0e7a4..2c08ce1e1 100644 --- a/gptqmodel/looper/gptq_processor.py +++ b/gptqmodel/looper/gptq_processor.py @@ -57,6 +57,12 @@ def clone_gptq_config_for_module( qcfg_clone = copy.deepcopy(qcfg) + # The caller-supplied fallback (base.py passes quantize_config.fallback into + # loop()) is the BASELINE, applied before dynamic so per-module `fallback` + # overrides win. Applying it after (the previous order) silently clobbered + # every dynamic fallback override in the standard quantize() path. + qcfg_clone.fallback = normalize_fallback(fallback, qcfg_clone.fallback) + # dynamic overrides if qcfg.dynamic is not None: qcfg_clone.bits = qcfg.dynamic_get(module_full_name, "bits", qcfg_clone.bits) @@ -102,7 +108,6 @@ def clone_gptq_config_for_module( qcfg_clone._resolve_activation_ordering(desc_act_override, act_group_aware_override) - qcfg_clone.fallback = normalize_fallback(fallback, qcfg_clone.fallback) return qcfg_clone class GPTQProcessor(LoopProcessor): diff --git a/gptqmodel/looper/weight_only_looper.py b/gptqmodel/looper/weight_only_looper.py index 87b9938ac..77e698e2e 100644 --- a/gptqmodel/looper/weight_only_looper.py +++ b/gptqmodel/looper/weight_only_looper.py @@ -30,7 +30,7 @@ from ..models import BaseQModel from ..models._const import CPU, SUPPORTS_MODULE_TYPES from ..nn_modules.converter import MODULE_CONVERTER_MAP -from ..quantization.config import BitsAndBytesConfig, FP8Config, GGUFConfig, RTNConfig, VramStrategy +from ..quantization.config import BitsAndBytesConfig, FP8Config, GGUFConfig, NVFP4Config, RTNConfig, VramStrategy from ..utils import has_gil_disabled from ..utils.device import get_device from ..utils.device_telemetry import emit_device_telemetry @@ -414,7 +414,7 @@ def _quantize_named_module( self, named: NamedModule, target_device: torch.device, - ) -> Tuple[NamedModule, Optional[RTNConfig | GGUFConfig | FP8Config | BitsAndBytesConfig]]: + ) -> Tuple[NamedModule, Optional[RTNConfig | GGUFConfig | FP8Config | NVFP4Config | BitsAndBytesConfig]]: """Run one module's weight-only quantization on its assigned device.""" self._prepare_named_module_for_quantization(named, target_device) @@ -427,7 +427,7 @@ def _quantize_named_module( def _finalize_quantized_module( self, named: NamedModule, - active_qcfg: RTNConfig | GGUFConfig | FP8Config | BitsAndBytesConfig, + active_qcfg: RTNConfig | GGUFConfig | FP8Config | NVFP4Config | BitsAndBytesConfig, ) -> str: """Move a quantized module back to CPU, pack it, and optionally offload it.""" @@ -477,7 +477,7 @@ def _finalize_quantized_module( def _finalize_target_device( self, - active_qcfg: RTNConfig | GGUFConfig | FP8Config | BitsAndBytesConfig, + active_qcfg: RTNConfig | GGUFConfig | FP8Config | NVFP4Config | BitsAndBytesConfig, ) -> torch.device: """Resolve the worker device for one finalize task.""" @@ -488,7 +488,7 @@ def _finalize_target_device( def _finalize_subset_modules( self, - quantized_modules: List[Tuple[NamedModule, RTNConfig | GGUFConfig | FP8Config | BitsAndBytesConfig]], + quantized_modules: List[Tuple[NamedModule, RTNConfig | GGUFConfig | FP8Config | NVFP4Config | BitsAndBytesConfig]], ) -> None: """Finalize one subset, using the device pool when multiple finalize targets exist.""" @@ -576,10 +576,10 @@ def _advance_finalize_progress(named: NamedModule, module_label: str) -> None: def _quantize_subset_modules( self, named_modules: List[NamedModule], - ) -> List[Tuple[NamedModule, RTNConfig | GGUFConfig | FP8Config | BitsAndBytesConfig]]: + ) -> List[Tuple[NamedModule, RTNConfig | GGUFConfig | FP8Config | NVFP4Config | BitsAndBytesConfig]]: """Quantize one subset, using multiple devices when available.""" - results_by_name: Dict[str, Tuple[NamedModule, RTNConfig | GGUFConfig | FP8Config | BitsAndBytesConfig]] = {} + results_by_name: Dict[str, Tuple[NamedModule, RTNConfig | GGUFConfig | FP8Config | NVFP4Config | BitsAndBytesConfig]] = {} task_specs = [ (named, self._assign_quant_device_for_module(named, CPU)) for named in named_modules @@ -669,10 +669,10 @@ def _offload_quantized_module(self, module: NamedModule) -> None: def loop(self, **kwargs): """Quantize layers directly from weights without calibration forwards.""" quant_config = self.gptq_model.quantize_config - if not isinstance(quant_config, (RTNConfig, GGUFConfig, FP8Config, BitsAndBytesConfig)): + if not isinstance(quant_config, (RTNConfig, GGUFConfig, FP8Config, NVFP4Config, BitsAndBytesConfig)): raise NotImplementedError( "Weight-only looper only supports `RTNConfig`, `GGUFConfig`, " - "`FP8Config`, and `BitsAndBytesConfig` today." + "`FP8Config`, `NVFP4Config`, and `BitsAndBytesConfig` today." ) if quant_config.lm_head: diff --git a/gptqmodel/looper/weight_only_processor.py b/gptqmodel/looper/weight_only_processor.py index 31b08abe7..1e4615aef 100644 --- a/gptqmodel/looper/weight_only_processor.py +++ b/gptqmodel/looper/weight_only_processor.py @@ -29,6 +29,7 @@ FP8Config, GGUFConfig, METHOD, + NVFP4Config, RTNConfig, clone_weight_only_config_for_module, resolve_quant_format, @@ -48,9 +49,9 @@ class WeightOnlyProcessor(LoopProcessor): def __init__( self, tokenizer, - qcfg: RTNConfig | GGUFConfig | FP8Config | BitsAndBytesConfig, + qcfg: RTNConfig | GGUFConfig | FP8Config | NVFP4Config | BitsAndBytesConfig, ): - """Initializes a weight-only processor for RTN, GGUF, FP8, or BitsAndBytes.""" + """Initializes a weight-only processor for RTN, GGUF, FP8, NVFP4, or BitsAndBytes.""" super().__init__( tokenizer=tokenizer, @@ -74,10 +75,10 @@ def is_skipped(self, module: NamedModule) -> bool: return self.qcfg.dynamic_get(layer_name=module.full_name) is False @staticmethod - def _uses_direct_pack(qcfg: RTNConfig | GGUFConfig | FP8Config | BitsAndBytesConfig) -> bool: + def _uses_direct_pack(qcfg: RTNConfig | GGUFConfig | FP8Config | NVFP4Config | BitsAndBytesConfig) -> bool: """Returns whether the method packs directly from the original dense weights.""" - return qcfg.method in {METHOD.GGUF, METHOD.FP8, METHOD.BITSANDBYTES} + return qcfg.method in {METHOD.GGUF, METHOD.FP8, METHOD.NVFP4, METHOD.BITSANDBYTES} def _update_logged_loss(self, module: NamedModule, avg_loss: str) -> None: """Backfills the logged loss field after late dequant-error measurement.""" @@ -93,7 +94,7 @@ def quantize_module( module: NamedModule, *, device: Optional[torch.device] = None, - ) -> Optional[RTNConfig | GGUFConfig | FP8Config | BitsAndBytesConfig]: + ) -> Optional[RTNConfig | GGUFConfig | FP8Config | NVFP4Config | BitsAndBytesConfig]: """Clones per-module config, quantizes weights, and logs the result.""" qcfg_clone = clone_weight_only_config_for_module(self.qcfg, module.full_name) @@ -149,7 +150,7 @@ def submodule_finalize( module: NamedModule, model: BaseQModel, *, - qcfg: Optional[RTNConfig | GGUFConfig | FP8Config | BitsAndBytesConfig] = None, + qcfg: Optional[RTNConfig | GGUFConfig | FP8Config | NVFP4Config | BitsAndBytesConfig] = None, **kwargs, ): """Creates and packs the final quantized module into the model graph.""" @@ -268,6 +269,8 @@ def name(self) -> str: return "weight_only_gguf" if self.qcfg.method == METHOD.FP8: return "weight_only_fp8" + if self.qcfg.method == METHOD.NVFP4: + return "weight_only_nvfp4" if self.qcfg.method == METHOD.BITSANDBYTES: return "weight_only_bitsandbytes" return "weight_only_rtn" diff --git a/gptqmodel/models/auto.py b/gptqmodel/models/auto.py index 478afc2f9..ed2699749 100644 --- a/gptqmodel/models/auto.py +++ b/gptqmodel/models/auto.py @@ -730,6 +730,8 @@ def export(model_id_or_path: str, target_path: str, format: str, trust_remote_co backend = BACKEND.PAROQUANT_CUDA elif normalized_method == METHOD.FP8.value: backend = BACKEND.FP8_TORCH + elif normalized_method == METHOD.NVFP4.value: + backend = BACKEND.NVFP4_TORCH elif normalized_method == METHOD.EXL3.value: backend = BACKEND.EXL3_TORCH else: diff --git a/gptqmodel/models/base.py b/gptqmodel/models/base.py index e045547a3..aab73a22e 100644 --- a/gptqmodel/models/base.py +++ b/gptqmodel/models/base.py @@ -896,6 +896,9 @@ def quantize( preferred_backend = BACKEND.AUTO elif self.quantize_config.method == METHOD.FP8: preferred_backend = BACKEND.FP8_TORCH + elif export_quant_method == METHOD.NVFP4 or format_code == FORMAT.NVFP4: + # Covers both weight-only NVFP4 and GPTQ-on-NVFP4-grid exports. + preferred_backend = BACKEND.NVFP4_TORCH elif self.quantize_config.method == METHOD.BITSANDBYTES: preferred_backend = BACKEND.BITSANDBYTES else: diff --git a/gptqmodel/nn_modules/qlinear/fp4.py b/gptqmodel/nn_modules/qlinear/fp4.py index 0eb6a0df8..1a3ab5756 100644 --- a/gptqmodel/nn_modules/qlinear/fp4.py +++ b/gptqmodel/nn_modules/qlinear/fp4.py @@ -9,11 +9,25 @@ import torch.nn as nn import torch.nn.functional as F +from ...adapter.adapter import Adapter, Lora +from ...models._const import DEVICE, PLATFORM +from ...quantization import FORMAT, METHOD +from ...quantization.config import NVFP4_WEIGHT_BLOCK_SIZE +from ...quantization.dtype import dequantize_f4_e2m1, device_supports_native_fp4 +from ...utils.backend import BACKEND +from . import WeightOnlyQuantLinear +from .fp8 import _weight_to_matrix +from .gguf import _apply_optional_smoother + try: - from torchao.prototype.mx_formats.nvfp4_tensor import NVFP4Tensor + from torchao.prototype.mx_formats.kernels import f32_to_f4_unpacked, pack_uint4 + from torchao.prototype.mx_formats.nvfp4_tensor import NVFP4Tensor, nvfp4_quantize except Exception: NVFP4Tensor = None + nvfp4_quantize = None + f32_to_f4_unpacked = None + pack_uint4 = None class TorchFP4Linear(nn.Module): @@ -29,6 +43,7 @@ def __init__( weight_block_size: int, orig_dtype: torch.dtype, bias: Optional[torch.Tensor] = None, + weight_scale_2: Optional[torch.Tensor] = None, ) -> None: super().__init__() self.in_features = in_features @@ -37,6 +52,12 @@ def __init__( self.orig_dtype = orig_dtype self.register_buffer("weight", weight) self.register_buffer("weight_scale", weight_scale) + # ModelOpt/Quark NVFP4 checkpoints carry a second-level per-tensor scale. + # Ignoring it would dequantize every block by a constant factor off. + if isinstance(weight_scale_2, torch.Tensor): + self.register_buffer("weight_scale_2", weight_scale_2.to(torch.float32).reshape(())) + else: + self.weight_scale_2 = None if isinstance(bias, torch.Tensor): self.register_buffer("bias", bias) else: @@ -52,6 +73,7 @@ def _native_weight(self) -> "NVFP4Tensor": self.weight_scale, block_size=self.weight_block_size, orig_dtype=self.orig_dtype, + per_tensor_scale=self.weight_scale_2, ) def forward(self, x: torch.Tensor) -> torch.Tensor: @@ -61,3 +83,408 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: if isinstance(bias, torch.Tensor) and bias.dtype != x.dtype: bias = bias.to(dtype=x.dtype) return F.linear(x, self._native_weight(), bias) + + +# NVFP4's second-level (per-tensor) scale. ModelOpt, Quark and tore-quant all +# derive it the same way -- amax / (E2M1_MAX * F8E4M3_MAX) -- and serialize it as +# `weight_scale_2`. Without it the E4M3 block scales carry absolute magnitudes: +# for typical weights `block_amax/6` lands below E4M3's smallest normal (2^-6), +# so `clamp(min=E4M3_EPS)` floors them and the block quantizes on a compressed +# grid. With it, block scales become 448*block_amax/global_amax and use E4M3's +# full range. +_F4_E2M1_MAX = 6.0 +_F8_E4M3_MAX = 448.0 + + +def compute_nvfp4_global_scale(weight: torch.Tensor) -> torch.Tensor: + """NVFP4 `weight_scale_2` for one tensor: amax / (6 * 448).""" + + amax = weight.detach().abs().max().to(torch.float32) + scale = amax / (_F4_E2M1_MAX * _F8_E4M3_MAX) + # An all-zero tensor would give 0 and make every block scale non-finite. + return torch.where(scale > 0, scale, torch.ones_like(scale)).reshape(()) + + +def quantize_nvfp4_weight( + weight: torch.Tensor, + *, + weight_block_size: int = NVFP4_WEIGHT_BLOCK_SIZE, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Quantize one dense 2D weight to NVFP4: packed bytes, E4M3 block scales, global scale.""" + + if nvfp4_quantize is None: + raise RuntimeError("NVFP4 quantization requires torchao with NVFP4 support.") + if weight.ndim != 2: + raise ValueError(f"NVFP4 quantization expects a 2D weight matrix, got shape {tuple(weight.shape)}.") + if int(weight_block_size) != NVFP4_WEIGHT_BLOCK_SIZE: + raise ValueError(f"NVFP4 quantization requires `weight_block_size={NVFP4_WEIGHT_BLOCK_SIZE}`.") + if weight.shape[1] % NVFP4_WEIGHT_BLOCK_SIZE != 0: + raise ValueError( + f"NVFP4 quantization expects in_features divisible by {NVFP4_WEIGHT_BLOCK_SIZE}, " + f"got shape {tuple(weight.shape)}." + ) + + weight = weight.to(device="cpu", dtype=torch.float32).contiguous() + global_scale = compute_nvfp4_global_scale(weight) + scales, packed = nvfp4_quantize( + weight, block_size=NVFP4_WEIGHT_BLOCK_SIZE, per_tensor_scale=global_scale + ) + # Store packed nibbles as uint8 so safetensors serialization matches the + # ModelOpt/torchao NVFP4 checkpoint layout regardless of torch fp4 dtype support. + if packed.dtype is not torch.uint8: + packed = packed.view(torch.uint8) + return packed.contiguous(), scales.contiguous(), global_scale.contiguous() + + +class TorchNVFP4Linear(WeightOnlyQuantLinear): + SUPPORTS_BACKENDS = [BACKEND.NVFP4_TORCH] + # METHOD.GPTQ is supported with FORMAT.NVFP4: the GPTQ loop quantizes on the + # NVFP4 grid and this kernel stores/executes the packed result. + SUPPORTS_METHODS = [METHOD.NVFP4, METHOD.GPTQ] + SUPPORTS_FORMATS = {FORMAT.NVFP4: 15} + SUPPORTS_BITS = [4] + SUPPORTS_GROUP_SIZE = [-1, NVFP4_WEIGHT_BLOCK_SIZE] + SUPPORTS_SYM = [True] + SUPPORTS_DESC_ACT = [False] + SUPPORTS_SHARDS = True + SUPPORTS_TRAINING = False + SUPPORTS_AUTO_PADDING = False + SUPPORTS_IN_FEATURES_DIVISIBLE_BY = [NVFP4_WEIGHT_BLOCK_SIZE] + SUPPORTS_OUT_FEATURES_DIVISIBLE_BY = [1] + SUPPORTS_DEVICES = [DEVICE.CPU, DEVICE.CUDA, DEVICE.ROCM, DEVICE.XPU, DEVICE.MPS] + SUPPORTS_PLATFORM = [PLATFORM.ALL] + SUPPORTS_PACK_DTYPES = [torch.int8, torch.int16, torch.int32, torch.int64] + SUPPORTS_ADAPTERS = [Lora] + SUPPORTS_DTYPES = [torch.float16, torch.bfloat16, torch.float32] + + QUANT_TYPE = "nvfp4" + + def __init__( + self, + bits: int, + group_size: int, + sym: bool, + desc_act: bool, + in_features: int, + out_features: int, + bias: bool = False, + pack_dtype: torch.dtype = torch.int32, + adapter: Adapter = None, + register_buffers: bool = True, + weight_block_size: int = NVFP4_WEIGHT_BLOCK_SIZE, + **kwargs, + ): + if int(weight_block_size) != NVFP4_WEIGHT_BLOCK_SIZE: + raise ValueError( + f"TorchNVFP4Linear requires `weight_block_size={NVFP4_WEIGHT_BLOCK_SIZE}`." + ) + self.weight_block_size = NVFP4_WEIGHT_BLOCK_SIZE + self._native_linear_hard_disabled = False + + super().__init__( + bits=bits, + in_features=in_features, + out_features=out_features, + bias=bias, + backend=kwargs.pop("backend", BACKEND.NVFP4_TORCH), + adapter=adapter, + register_buffers=False, + pack_dtype=pack_dtype, + **kwargs, + ) + + if register_buffers: + self._allocate_buffers(bias=bias) + + @classmethod + def validate_once(cls): + if NVFP4Tensor is None or nvfp4_quantize is None: + return False, RuntimeError("TorchNVFP4Linear requires torchao with NVFP4 support.") + return True, None + + def smooth_block_size(self) -> int: + return self.weight_block_size + + def _allocate_buffers(self, *, bias: bool) -> None: + weight = torch.zeros((self.out_features, self.in_features // 2), dtype=torch.uint8) + scale = torch.zeros( + (self.out_features, self.in_features // self.weight_block_size), + dtype=torch.float8_e4m3fn, + ) + + if "weight" in self._buffers: + self.weight = weight + else: + self.register_buffer("weight", weight) + + if "weight_scale" in self._buffers: + self.weight_scale = scale + else: + self.register_buffer("weight_scale", scale) + + # NVFP4's second-level per-tensor scale. 1.0 is the identity, so an + # un-packed module still dequantizes correctly. + scale_2 = torch.ones((), dtype=torch.float32) + if "weight_scale_2" in self._buffers: + self.weight_scale_2 = scale_2 + else: + self.register_buffer("weight_scale_2", scale_2) + + if bias: + bias_tensor = torch.zeros(self.out_features, dtype=torch.float16) + if "bias" in self._buffers: + self.bias = bias_tensor + else: + self.register_buffer("bias", bias_tensor) + else: + self.bias = None + + def list_buffers(self): + buffers = [] + if hasattr(self, "weight") and self.weight is not None: + buffers.append(self.weight) + if hasattr(self, "weight_scale") and self.weight_scale is not None: + buffers.append(self.weight_scale) + if hasattr(self, "weight_scale_2") and self.weight_scale_2 is not None: + buffers.append(self.weight_scale_2) + if hasattr(self, "bias") and self.bias is not None: + buffers.append(self.bias) + return buffers + + def extra_repr(self) -> str: + return ( + f"in_features={self.in_features}, out_features={self.out_features}, " + f"bias={self.bias is not None}, weight_block_size={self.weight_block_size}" + ) + + def _weight_to_matrix(self, linear: nn.Module) -> torch.Tensor: + return _weight_to_matrix(linear) + + def pack(self, linear: nn.Module, scales: torch.Tensor, zeros: torch.Tensor, g_idx: torch.Tensor = None): + self.pack_original(linear=linear, scales=scales, zeros=zeros, g_idx=g_idx) + + def pack_block( + self, + linear: nn.Module, + scales: torch.Tensor, + zeros: torch.Tensor, + g_idx: torch.Tensor = None, + block_in: int = 8192, + workers: int = 1, + ): + del block_in, workers + self.pack_original(linear=linear, scales=scales, zeros=zeros, g_idx=g_idx) + + def pack_gpu( + self, + linear: nn.Module, + scales: torch.Tensor, + zeros: torch.Tensor, + g_idx: torch.Tensor = None, + *, + block_in: int = 8192, + device: torch.device | None = None, + ): + del block_in, device + self.pack_original(linear=linear, scales=scales, zeros=zeros, g_idx=g_idx) + + def _pack_with_gptq_scales( + self, weight: torch.Tensor, scales: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Pack a weight already on the NVFP4 grid using the loop-produced group scales. + + GPTQ error feedback mutates columns after each group scale is fixed, so the + final weight must be encoded against those exact scales — re-deriving scales + from the weight (RTN-style) would land some blocks on a different grid. + + That constraint also dictates how `weight_scale_2` is chosen here. The loop + already fixed each block scale `s_i`, so the two-level form must satisfy + `block_e4m3_i * ws2 == s_i` exactly, or the packed codes stop matching their + scales. A **power-of-two** `ws2` makes that exact: dividing by it only shifts + the exponent, leaving the E4M3 mantissa untouched. A greedy `amax/(6*448)` + would not, so it is used only on the RTN path where scales are free. + """ + + if f32_to_f4_unpacked is None or pack_uint4 is None: + raise RuntimeError("NVFP4 packing requires torchao with NVFP4 support.") + + scale = scales.to(device="cpu", dtype=torch.float32) + expected_groups = self.in_features // self.weight_block_size + if scale.shape != (self.out_features, expected_groups): + raise ValueError( + f"NVFP4 pack expects scales of shape {(self.out_features, expected_groups)}, " + f"got {tuple(scale.shape)}." + ) + + expanded = scale.repeat_interleave(self.weight_block_size, dim=1) + # Reciprocal-multiply, not divide: matches torchao's nvfp4_quantize numerics + # (see NVFP4Quantizer.quantize). A 1-ULP difference can flip an E2M1 code. + scaled = torch.clamp(weight * expanded.reciprocal(), -6.0, 6.0) + codes = f32_to_f4_unpacked(scaled.contiguous()) + packed = pack_uint4(codes) + if packed.dtype is not torch.uint8: + packed = packed.view(torch.uint8) + + # ws2 = 2^ceil(log2(max(s) / 448)): the smallest power of two that brings + # every block scale within E4M3's range. + smax = scale.max() + if smax > 0: + exponent = torch.ceil(torch.log2(smax / _F8_E4M3_MAX)) + global_scale = torch.exp2(exponent).to(torch.float32).reshape(()) + else: + global_scale = torch.ones((), dtype=torch.float32) + + block_scale = (scale * global_scale.reciprocal()).to(torch.float8_e4m3fn) + # The exactness the docstring depends on is cheap to verify, and a silent + # mismatch here means every packed code is read against the wrong scale. + recovered = block_scale.to(torch.float32) * global_scale + if not torch.equal(recovered, scale): + bad = int((recovered != scale).sum()) + raise RuntimeError( + f"NVFP4 two-level split changed {bad} of {scale.numel()} block scales; " + "packed codes would no longer match their scales." + ) + + return packed.contiguous(), block_scale.contiguous(), global_scale.contiguous() + + @torch.inference_mode() + def pack_original( + self, + linear: nn.Module, + scales: torch.Tensor, + zeros: torch.Tensor, + g_idx: torch.Tensor = None, + *, + smooth=None, + ): + del zeros, g_idx + + weight = self._weight_to_matrix(linear).to(device="cpu", dtype=torch.float32) + if isinstance(scales, torch.Tensor) and scales.numel() > 0: + qweight, weight_scale, weight_scale_2 = self._pack_with_gptq_scales(weight, scales) + else: + weight = _apply_optional_smoother( + weight, + smooth=smooth, + group_size=self.smooth_block_size(), + ) + qweight, weight_scale, weight_scale_2 = quantize_nvfp4_weight( + weight, + weight_block_size=self.weight_block_size, + ) + + if "weight" in self._buffers: + self.weight = qweight + else: + self.register_buffer("weight", qweight) + + if "weight_scale" in self._buffers: + self.weight_scale = weight_scale + else: + self.register_buffer("weight_scale", weight_scale) + + if "weight_scale_2" in self._buffers: + self.weight_scale_2 = weight_scale_2 + else: + self.register_buffer("weight_scale_2", weight_scale_2) + + if linear.bias is not None: + bias = linear.bias.detach().to(device="cpu", dtype=torch.float16) + if "bias" in self._buffers: + self.bias = bias + else: + self.register_buffer("bias", bias) + else: + self.bias = None + + self._native_linear_hard_disabled = False + + def dequantize_weight( + self, + *, + device: torch.device | str | None = None, + dtype: torch.dtype | None = None, + ) -> torch.Tensor: + target_device = self.weight.device if device is None else torch.device(device) + target_dtype = torch.float32 if dtype is None else dtype + + # dequantize_f4_e2m1 handles both the fused CPU path (bf16/fp16 targets) + # and the reference torchao path for other targets/devices. + dequant_dtype = target_dtype if target_dtype in (torch.bfloat16, torch.float16) else torch.bfloat16 + source = self.weight if self.weight.device == target_device else self.weight.to(device=target_device) + scale = ( + self.weight_scale + if self.weight_scale.device == target_device + else self.weight_scale.to(device=target_device) + ) + # Fold the per-tensor scale into the block scales rather than scaling the + # dequantized weight afterwards: one pass, and it keeps the fp32 product + # exact instead of rounding through the E4M3 block scale twice. + scale_2 = getattr(self, "weight_scale_2", None) + if isinstance(scale_2, torch.Tensor): + scale = scale.to(dtype=torch.float32) * scale_2.to(device=target_device, dtype=torch.float32) + weight = dequantize_f4_e2m1( + source, + scale=scale, + axis=None, + target_dtype=dequant_dtype, + ) + if weight.dtype != target_dtype: + weight = weight.to(dtype=target_dtype) + return weight.transpose(0, 1).contiguous() + + def _native_weight(self) -> "NVFP4Tensor": + if NVFP4Tensor is None: + raise RuntimeError("TorchNVFP4Linear requires torchao NVFP4Tensor support.") + return NVFP4Tensor( + self.weight, + self.weight_scale, + block_size=self.weight_block_size, + orig_dtype=torch.bfloat16, + per_tensor_scale=getattr(self, "weight_scale_2", None), + ) + + def _can_use_native_linear(self, x_flat: torch.Tensor) -> bool: + return ( + not self._native_linear_hard_disabled + and NVFP4Tensor is not None + and x_flat.device.type == "cuda" + and x_flat.dtype is torch.bfloat16 + and device_supports_native_fp4(x_flat.device) + ) + + def _forward_dequant_matmul(self, x_flat: torch.Tensor) -> torch.Tensor: + weight = self.dequantize_weight(device=x_flat.device, dtype=x_flat.dtype) + return torch.matmul(x_flat, weight) + + def forward(self, x: torch.Tensor): + original_shape = x.shape[:-1] + (self.out_features,) + x_flat = x.reshape(-1, x.shape[-1]) + + if self._can_use_native_linear(x_flat): + try: + output = F.linear(x_flat, self._native_weight(), None) + except Exception: + self._native_linear_hard_disabled = True + output = self._forward_dequant_matmul(x_flat) + else: + output = self._forward_dequant_matmul(x_flat) + + if self.bias is not None: + bias = self.bias + if bias.device != output.device or bias.dtype != output.dtype: + bias = bias.to(device=output.device, dtype=output.dtype) + output = output + bias + + if self.adapter: + output = self.adapter.apply(x=x_flat, out=output) + + return output.reshape(original_shape) + + +__all__ = [ + "TorchFP4Linear", + "TorchNVFP4Linear", + "quantize_nvfp4_weight", + "compute_nvfp4_global_scale", +] diff --git a/gptqmodel/quantization/__init__.py b/gptqmodel/quantization/__init__.py index c2261e779..d1fa8c5f6 100644 --- a/gptqmodel/quantization/__init__.py +++ b/gptqmodel/quantization/__init__.py @@ -27,6 +27,7 @@ GPTAQConfig, GPTQConfig, HessianConfig, + NVFP4Config, ParoConfig, PreProcessorCode, PreProcessorConfig, diff --git a/gptqmodel/quantization/config.py b/gptqmodel/quantization/config.py index 658c59126..13a3effe4 100644 --- a/gptqmodel/quantization/config.py +++ b/gptqmodel/quantization/config.py @@ -114,6 +114,7 @@ class FORMAT(str, Enum): GPTQ_P = "gptq_p" GGUF = "gguf" FP8 = "fp8" + NVFP4 = "nvfp4" BITSANDBYTES = "bitsandbytes" MARLIN = "marlin" BITBLAS = "bitblas" @@ -134,6 +135,7 @@ class METHOD(str, Enum): GPTQ = "gptq" GGUF = "gguf" FP8 = "fp8" + NVFP4 = "nvfp4" BITSANDBYTES = "bitsandbytes" QQQ = "qqq" AWQ = "awq" @@ -652,6 +654,8 @@ def resolve_quant_format( return FORMAT.GGUF if method == METHOD.FP8: return FORMAT.FP8 + if method == METHOD.NVFP4: + return FORMAT.NVFP4 if method == METHOD.BITSANDBYTES: return FORMAT.BITSANDBYTES if method == METHOD.EXL3: @@ -1404,10 +1408,14 @@ def to_dict(self) -> Dict[str, Any]: FORMAT.GPTQ_P, FORMAT.MARLIN, FORMAT.BITBLAS, + FORMAT.NVFP4, }, METHOD.FP8: { FORMAT.FP8, }, + METHOD.NVFP4: { + FORMAT.NVFP4, + }, METHOD.BITSANDBYTES: { FORMAT.BITSANDBYTES, }, @@ -1439,6 +1447,7 @@ def to_dict(self) -> Dict[str, Any]: FORMAT.GPTQ_P, FORMAT.MARLIN, FORMAT.BITBLAS, + FORMAT.NVFP4, ) AWQ_EXPORT_FORMATS: Tuple[FORMAT, ...] = ( FORMAT.GEMM, @@ -1461,6 +1470,22 @@ def to_dict(self) -> Dict[str, Any]: FP8_EXPORT_FORMATS: Tuple[FORMAT, ...] = ( FORMAT.FP8, ) +NVFP4_EXPORT_FORMATS: Tuple[FORMAT, ...] = ( + FORMAT.NVFP4, +) +# NVFP4 fixes 16 FP4 values per FP8 E4M3 scale along the in-feature axis. +NVFP4_WEIGHT_BLOCK_SIZE = 16 + + +def _normalize_nvfp4_format(value: Optional[Union[str, FORMAT]]) -> str: + """Normalize accepted NVFP4 format aliases (`nvfp4`, `fp4`) to `nvfp4`.""" + + if value is None: + return FORMAT.NVFP4.value + normalized = str(getattr(value, "value", value)).strip().lower() + if normalized in {"nvfp4", "fp4"}: + return FORMAT.NVFP4.value + raise ValueError(f"NVFP4Config: unsupported `format`: `{value}`. Expected `nvfp4`.") BITSANDBYTES_EXPORT_FORMATS: Tuple[FORMAT, ...] = ( FORMAT.BITSANDBYTES, ) @@ -1485,6 +1510,7 @@ def to_dict(self) -> Dict[str, Any]: FORMAT.GPTQ_V2: METHOD.GPTQ, FORMAT.GPTQ_P: METHOD.GPTQ, FORMAT.FP8: METHOD.FP8, + FORMAT.NVFP4: METHOD.NVFP4, FORMAT.BITSANDBYTES: METHOD.BITSANDBYTES, FORMAT.EXL3: METHOD.EXL3, FORMAT.GGUF: METHOD.GGUF, @@ -2258,6 +2284,17 @@ def _normalize_fp8_kwargs(payload: Dict[str, Any]) -> Dict[str, Any]: return normalized +def _normalize_nvfp4_kwargs(payload: Dict[str, Any]) -> Dict[str, Any]: + normalized = dict(payload) + weight_only = normalized.pop("weight_only", None) + + if "smoother" not in normalized and "smooth" not in normalized: + normalized["smoother"] = _extract_weight_only_smooth(weight_only) + + normalized[FORMAT_FIELD_CODE] = _normalize_nvfp4_format(normalized.get(FORMAT_FIELD_CODE)) + return normalized + + def _normalize_bitsandbytes_kwargs(payload: Dict[str, Any]) -> Dict[str, Any]: normalized = dict(payload) weight_only = normalized.pop("weight_only", None) @@ -2303,6 +2340,8 @@ def _normalize_quantize_config_payload_for_target_cls(target_cls, payload: Dict[ expected_method = METHOD.AWQ elif target_cls is FP8Config: expected_method = METHOD.FP8 + elif target_cls is NVFP4Config: + expected_method = METHOD.NVFP4 elif target_cls is BitsAndBytesConfig: expected_method = METHOD.BITSANDBYTES elif target_cls is EXL3Config: @@ -2384,6 +2423,8 @@ def _prepare_target_quantize_config_kwargs(target_cls, payload: Dict[str, Any]) normalized = _normalize_gguf_kwargs(normalized) elif target_cls is FP8Config: normalized = _normalize_fp8_kwargs(normalized) + elif target_cls is NVFP4Config: + normalized = _normalize_nvfp4_kwargs(normalized) elif target_cls is BitsAndBytesConfig: normalized = _normalize_bitsandbytes_kwargs(normalized) return _filter_quantize_config_payload_for_target_cls(target_cls, normalized) @@ -2913,7 +2954,17 @@ def from_quant_config(cls, quantize_cfg, format: str = None): }: normalized[key] = val else: - normalized[key] = _normalize_format(val) + raw_method = quantize_cfg.get(METHOD_FIELD_CODE, quantize_cfg.get(QUANT_METHOD_FIELD)) + try: + method_hint = _normalize_quant_method(raw_method) if raw_method is not None else None + except ValueError: + method_hint = None + if method_hint == METHOD.NVFP4: + # NVFP4 payloads accept the `fp4` format alias, same as the + # constructor path (`_normalize_nvfp4_kwargs`). + normalized[key] = FORMAT(_normalize_nvfp4_format(val)) + else: + normalized[key] = _normalize_format(val) elif key in field_names: normalized[key] = val else: @@ -2997,6 +3048,8 @@ def from_quant_config(cls, quantize_cfg, format: str = None): normalized = _normalize_gguf_kwargs(normalized) elif target_cls is FP8Config: normalized = _normalize_fp8_kwargs(normalized) + elif target_cls is NVFP4Config: + normalized = _normalize_nvfp4_kwargs(normalized) elif target_cls is BitsAndBytesConfig: normalized = _normalize_bitsandbytes_kwargs(normalized) @@ -3212,6 +3265,46 @@ class GPTQConfig(PreProcessorConfig): metadata={"help": "Skip heavy computations for fast model loading validation"}, ) hessian: Optional[HessianConfig] = field(default_factory=HessianConfig) + act_format: Optional[FORMAT] = field( + default=None, + metadata={ + "help": "Activation format the deployed kernel will use (w4a4). When set, the " + "GPTQ Hessian is accumulated on Q(X) instead of X, so the solved weights " + "are optimal for the inputs the kernel actually sees. None = weight-only." + }, + ) + act_group_size: int = field( + default=NVFP4_WEIGHT_BLOCK_SIZE, + metadata={"help": "Activation quantization block size, along the last dim."}, + ) + act_amax_path: Optional[str] = field( + default=None, + metadata={ + "help": "Path to an activation amax sidecar (JSON) used to derive NVFP4's static " + "per-tensor global scale as amax/6/448. Accepts a flat {key: amax} mapping " + "or {'amax': {key: amax}}. Without it only block scales are applied, which " + "will not match a kernel that uses a static global activation scale." + }, + ) + nvfp4_scale_sweep: int = field( + default=8, + metadata={ + "help": "Per-block E4M3 scale search width for format=nvfp4. For each 16-wide " + "block, evaluate this many neighbouring E4M3 scale codes and keep the " + "MSE-minimal one; amax-derived scales are only optimal for a block's " + "single largest value (measured ~13% reconstruction error vs the " + "ModelOpt-RTN reference). 0 = amax-derived scales only." + }, + ) + act_amax_key_rules: Optional[List[Tuple[str, str]]] = field( + default=None, + metadata={ + "help": "Ordered (regex, replacement) pairs mapping a module name to its sidecar " + "key, first match wins; default is to look up the module name itself. " + "Modules sharing an input must map to the same key -- e.g. w1/w3 both to " + "`layer{N}.moe_input` -- because they share one activation scale." + }, + ) def allowed_quant_methods(self) -> Tuple[METHOD, ...]: return (METHOD.GPTQ,) @@ -3249,6 +3342,71 @@ def __post_init__(self): if self.act_group_aware and self.desc_act: raise ValueError("QuantizeConfig:: `act_group_aware` == `True` requires `desc_act` == `False`.") + if self.format == FORMAT.NVFP4: + # NVFP4 export stores contiguous 16-wide FP4 blocks with E4M3 scales, so the + # GPTQ loop must quantize on that exact grid without column reordering. + if quant_bits_width(self.bits) != 4: + raise ValueError("QuantizeConfig: `format=nvfp4` requires `bits=4`.") + if self.group_size != NVFP4_WEIGHT_BLOCK_SIZE: + raise ValueError( + f"QuantizeConfig: `format=nvfp4` requires `group_size={NVFP4_WEIGHT_BLOCK_SIZE}`." + ) + if not self.sym: + raise ValueError("QuantizeConfig: `format=nvfp4` requires `sym=True`.") + if self.desc_act: + raise ValueError("QuantizeConfig: `format=nvfp4` does not support `desc_act=True`.") + if self.act_group_aware: + if act_group_aware_user_value: + raise ValueError("QuantizeConfig: `format=nvfp4` does not support `act_group_aware=True`.") + self.act_group_aware = False + if float(self.mse or 0.0) > 0.0: + raise ValueError("QuantizeConfig: `format=nvfp4` does not support `mse` scale search.") + if self.fallback is not None: + # Only the RTN fallback routes through NVFP4Quantizer; every other + # strategy quantizes on an asymmetric integer grid whose zero-points + # the NVFP4 packer cannot represent. + if self.fallback.strategy != FallbackStrategy.RTN: + raise ValueError( + "QuantizeConfig: `format=nvfp4` requires `fallback.strategy=rtn`." + ) + # Smoothing multiplies the block scale by a per-block factor, so the + # stored scale is no longer E4M3-exact and packing stops being lossless. + if self.fallback.smooth is not None: + raise ValueError( + "QuantizeConfig: `format=nvfp4` does not support `fallback.smooth`." + ) + if self.mock_quantization: + # The mock loop rounds on an integer grid and applies one scale per + # 128-column block, neither of which lands on the NVFP4 grid. + raise ValueError( + "QuantizeConfig: `format=nvfp4` does not support `mock_quantization=True`." + ) + + if self.act_format is not None: + # w4a4: the Hessian is built from Q(X), so the activation grid has to be one + # we can actually reproduce here. + if self.act_format != FORMAT.NVFP4: + raise ValueError( + f"QuantizeConfig: `act_format={self.act_format}` is not supported; " + "only `act_format=nvfp4` is implemented." + ) + if self.format != FORMAT.NVFP4: + raise ValueError( + "QuantizeConfig: `act_format=nvfp4` requires `format=nvfp4`." + ) + if self.act_group_size <= 0: + raise ValueError("QuantizeConfig: `act_group_size` must be positive.") + if self.mock_quantization: + raise ValueError( + "QuantizeConfig: `act_format` does not support `mock_quantization=True`." + ) + elif self.act_amax_path is not None or self.act_amax_key_rules is not None: + # A sidecar with no activation quantization silently does nothing; that is + # far more likely a misconfiguration than an intent. + raise ValueError( + "QuantizeConfig: `act_amax_path`/`act_amax_key_rules` require `act_format` to be set." + ) + def _resolve_activation_ordering( self, desc_act_user_value: Optional[bool], @@ -3651,6 +3809,90 @@ def _update_output_payload(self, out: Dict[str, Any]) -> None: def uses_weight_only_lifecycle(self) -> bool: return True +@dataclass +class NVFP4Config(PreProcessorConfig): + """Weight-only NVFP4 (FP4 E2M1 + per-16-block FP8 E4M3 scales) export config.""" + + bits: int = field(default=4, metadata={"choices": [4]}) + method: METHOD = field(default=METHOD.NVFP4) + format: Optional[str] = field(default=FORMAT.NVFP4) + group_size: int = field(default=-1) + desc_act: Optional[bool] = field(default=False) + sym: bool = field(default=True) + # NVFP4 fixes the scaling block to 16 values along in-features. + weight_block_size: int = field(default=NVFP4_WEIGHT_BLOCK_SIZE) + + def _resolve_checkpoint_format(self) -> FORMAT: + self.format = FORMAT.NVFP4 + return FORMAT.NVFP4 + + def allowed_quant_methods(self) -> Tuple[METHOD, ...]: + return (METHOD.NVFP4,) + + def supported_export_formats(self) -> Tuple[FORMAT, ...]: + return NVFP4_EXPORT_FORMATS + + def default_desc_act(self) -> bool: + return False + + def __post_init__(self): + self._normalize_preprocessor_state() + super().__post_init__() + + if self.bits != 4: + raise ValueError("NVFP4Config: `bits` must be `4`.") + + if self.method != METHOD.NVFP4: + raise ValueError("NVFP4Config: `method` must be `nvfp4`.") + + self.group_size = -1 + self.desc_act = False + self.sym = True + self.format = _normalize_nvfp4_format(self.format) + + if int(self.weight_block_size) != NVFP4_WEIGHT_BLOCK_SIZE: + raise ValueError( + f"NVFP4Config: `weight_block_size` must be `{NVFP4_WEIGHT_BLOCK_SIZE}`." + ) + self.weight_block_size = NVFP4_WEIGHT_BLOCK_SIZE + + if self.dynamic is not None: + self.dynamic = { + **{k: v for k, v in self.dynamic.items() if k.startswith('-')}, + **{k: v for k, v in self.dynamic.items() if not k.startswith('-')}, + } + for layer, layer_dict in self.dynamic.items(): + self._normalize_dynamic_layer_config(layer, layer_dict) + + def _normalize_dynamic_layer_config( + self, + layer_name: str, + layer_dict: Dict[str, Any], + ) -> None: + if "bits" in layer_dict and int(layer_dict["bits"]) != 4: + raise ValueError(f"NVFP4Config: layer `{layer_name}` only supports 4-bit NVFP4 weights.") + if "group_size" in layer_dict and layer_dict["group_size"] not in (-1, None): + raise ValueError("NVFP4Config: `group_size` is not used; keep it at `-1`.") + if "weight_block_size" in layer_dict and int(layer_dict["weight_block_size"]) != NVFP4_WEIGHT_BLOCK_SIZE: + raise ValueError( + f"NVFP4Config: `weight_block_size` must be `{NVFP4_WEIGHT_BLOCK_SIZE}`." + ) + raw_format = layer_dict.get(FORMAT_FIELD_CODE) + if raw_format is not None: + layer_dict[FORMAT_FIELD_CODE] = _normalize_nvfp4_format(raw_format) + + def quant_linear_init_kwargs(self) -> Dict[str, Any]: + return { + "weight_block_size": self.weight_block_size, + } + + def _update_output_payload(self, out: Dict[str, Any]) -> None: + out[FORMAT_FIELD_CODE] = FORMAT.NVFP4.value + out["weight_block_size"] = self.weight_block_size + + def uses_weight_only_lifecycle(self) -> bool: + return True + @dataclass class BitsAndBytesConfig(PreProcessorConfig): bits: int = field(default=4, metadata={"choices": [4, 8]}) @@ -4078,9 +4320,9 @@ def uses_weight_only_lifecycle(self) -> bool: return True def clone_weight_only_config_for_module( - qcfg: Union[RTNConfig, GGUFConfig, FP8Config, BitsAndBytesConfig], + qcfg: Union[RTNConfig, GGUFConfig, FP8Config, NVFP4Config, BitsAndBytesConfig], module_full_name: str, -) -> Optional[Union[RTNConfig, GGUFConfig, FP8Config, BitsAndBytesConfig]]: +) -> Optional[Union[RTNConfig, GGUFConfig, FP8Config, NVFP4Config, BitsAndBytesConfig]]: if qcfg.dynamic_get(layer_name=module_full_name) is False: return None @@ -4225,11 +4467,12 @@ def _resolve_quantize_config_class(payload: Dict[str, Any]) -> type[BaseQuantize WeightOnlyMethod.RTN, WeightOnlyMethod.GGUF, WeightOnlyMethod.FP8, + WeightOnlyMethod.NVFP4, WeightOnlyMethod.BITSANDBYTES, }: raise ValueError( "QuantizeConfig: unsupported weight-only config. Weight-only export currently supports " - "`rtn`, `gguf`, `fp8`, and `bitsandbytes`." + "`rtn`, `gguf`, `fp8`, `nvfp4`, and `bitsandbytes`." ) if ( format_value == FORMAT.GGUF @@ -4240,6 +4483,8 @@ def _resolve_quantize_config_class(payload: Dict[str, Any]) -> type[BaseQuantize return GGUFConfig if weight_only_method == WeightOnlyMethod.FP8: return FP8Config + if weight_only_method == WeightOnlyMethod.NVFP4: + return NVFP4Config if weight_only_method == WeightOnlyMethod.BITSANDBYTES: return BitsAndBytesConfig if weight_only_method == WeightOnlyMethod.RTN: @@ -4248,6 +4493,12 @@ def _resolve_quantize_config_class(payload: Dict[str, Any]) -> type[BaseQuantize return RTNConfig if method == METHOD.FP8 or format_value == FORMAT.FP8 or _looks_like_fp8_fmt(fp8_storage_fmt): return FP8Config + if method == METHOD.NVFP4: + return NVFP4Config + # `format=nvfp4` with the default/explicit GPTQ method runs the GPTQ algorithm + # on the NVFP4 grid; only `method=nvfp4` selects the direct weight-only path. + if format_value == FORMAT.NVFP4 and method != METHOD.GPTQ: + return NVFP4Config if method == METHOD.BITSANDBYTES or format_value == FORMAT.BITSANDBYTES or _looks_like_bitsandbytes_format(raw_format_value): return BitsAndBytesConfig if method == METHOD.EXL3 or format_value == FORMAT.EXL3: @@ -4276,6 +4527,7 @@ def _known_quantize_config_field_names() -> set[str]: ParoConfig, QQQConfig, FP8Config, + NVFP4Config, BitsAndBytesConfig, EXL3Config, RTNConfig, diff --git a/gptqmodel/quantization/dtype.py b/gptqmodel/quantization/dtype.py index 15e34f7a7..6ce2fd641 100644 --- a/gptqmodel/quantization/dtype.py +++ b/gptqmodel/quantization/dtype.py @@ -42,6 +42,10 @@ "dequantize_fp8", "dequantize_f8_e4m3", "dequantize_f4_e2m1", + "dequantize_block_fp8", + "block_scale_multiplier", + "fake_quantize_nvfp4", + "fp8_storage_dtypes", "is_fp4_packed_dtype", ] @@ -612,6 +616,18 @@ def dequantize_f8_e4m3( if scale is not None and scale_inv is not None: raise ValueError("Provide either scale or scale_inv, not both") + + # MX checkpoints (e.g. mxfp8's UE8M0 `weight_scale_inv`) store block scales as + # uint8 exponent bytes: the multiplier is 2^(b - 127). A genuine float scale is + # never serialized as uint8, so route uint8 scales through the block dequant + # helper unconditionally. Feeding the raw bytes onward is wrong twice over: + # magnitudes ~110-130 exceed 1, so `_fast_scale_arg` flips scale_inv into + # divide mode (observed absmax 3.9 vs correct 0.28 on Minimax-M3-0602), and + # the fast CPU path does not understand 2D block grids in the first place. + mx_scale = scale if scale is not None else scale_inv + if mx_scale is not None and mx_scale.dtype == torch.uint8: + return dequantize_block_fp8(tensor, mx_scale, target_dtype=target_dtype) + if tensor.dtype in _FLOAT8_DTYPES and _can_use_fast_path( tensor, scale if scale is not None else scale_inv, @@ -681,6 +697,108 @@ def dequantize_fp8( ) +# E8M0 is a scale format, never a weight payload, so it is excluded here: a +# tensor stored as e8m0 is metadata and must not be dequantized as if it were +# quantized data. +_FP8_STORAGE_DTYPE_NAMES = ("float8_e4m3fn", "float8_e5m2", "float8_e4m3fnuz", "float8_e5m2fnuz") +_FP8_STORAGE_DTYPES = tuple( + getattr(torch, name) for name in _FP8_STORAGE_DTYPE_NAMES if hasattr(torch, name) +) + + +def fp8_storage_dtypes() -> tuple[torch.dtype, ...]: + """FP8 dtypes that can carry quantized weight payloads (excludes E8M0 scales).""" + + return _FP8_STORAGE_DTYPES + + +def fake_quantize_nvfp4( + x: torch.Tensor, + *, + block_size: int = 16, + per_tensor_scale: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Quantize to the NVFP4 grid and immediately dequantize, preserving shape/dtype. + + Used to make calibration statistics describe the activations a w4a4 kernel + actually consumes. Blocks run along the last dim, matching the weight path in + ``quantize_to_nvfp4``; ``per_tensor_scale`` is the optional static global + scale (omitted, as in the weight path, when the grid has block scales only). + """ + + if NVFP4Tensor is None or nvfp4_quantize is None: + raise RuntimeError("fake_quantize_nvfp4 requires torchao with NVFP4 support.") + + if x.shape[-1] % block_size != 0: + raise ValueError( + f"NVFP4 activation quantization needs the last dim ({x.shape[-1]}) " + f"to be a multiple of block_size ({block_size})." + ) + + orig_dtype = x.dtype + scales, packed = nvfp4_quantize( + x.to(torch.float32).contiguous(), + block_size=block_size, + per_tensor_scale=per_tensor_scale, + ) + tensor = NVFP4Tensor( + packed, + scales, + block_size, + orig_dtype, + per_tensor_scale=per_tensor_scale, + ) + return tensor.dequantize(orig_dtype) + + +def block_scale_multiplier(scale: torch.Tensor, target_shape: torch.Size) -> torch.Tensor: + """Expand a block scale so it multiplies elementwise against ``target_shape``. + + ``uint8`` scales are MX E8M0 exponents, whose multiplier is ``2**(e - 127)``. + Float scales are already multipliers (block-fp8 checkpoints store them that + way). Each dim is repeated by ``target_dim // scale_dim``, so ``[1, 32]`` + and ``[128, 128]`` block layouts are both handled without special-casing. + """ + + if scale.dtype == torch.uint8: + mult = torch.exp2(scale.to(torch.float32) - 127.0) + else: + mult = scale.to(torch.float32) + + while mult.ndim < len(target_shape): + mult = mult.unsqueeze(-1) + if mult.ndim != len(target_shape): + raise ValueError( + f"block scale rank {mult.ndim} exceeds weight rank {len(target_shape)}" + ) + + for dim, (want, have) in enumerate(zip(target_shape, mult.shape)): + if have == want: + continue + if have == 0 or want % have != 0: + raise ValueError( + f"block scale dim {dim} ({have}) does not evenly divide weight dim ({want})" + ) + mult = mult.repeat_interleave(want // have, dim=dim) + return mult + + +def dequantize_block_fp8( + tensor: torch.Tensor, + scale: torch.Tensor, + *, + target_dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: + """Dequantize a block-scaled FP8 weight to ``target_dtype``. + + Unlike :func:`dequantize_fp8`, this applies a *block* scale (one entry per + fixed-size group of elements) and decodes E8M0 exponent bytes, which is what + mxfp8 checkpoints store alongside each weight. + """ + + return (tensor.to(torch.float32) * block_scale_multiplier(scale, tensor.shape)).to(target_dtype) + + def dequantize_f4_e2m1( tensor: torch.Tensor, *, diff --git a/gptqmodel/quantization/foem.py b/gptqmodel/quantization/foem.py index b17002068..7c070282b 100644 --- a/gptqmodel/quantization/foem.py +++ b/gptqmodel/quantization/foem.py @@ -20,6 +20,7 @@ from ..quantization import QuantizeConfig from ..utils.torch import TORCH_GTE_28, torch_compile, torch_sync from .gptq import GPTQ +from .quantizer import NVFP4Quantizer class FOEM(GPTQ): @@ -155,6 +156,8 @@ def quantize( W = self.module_copy self.module_copy = None + if isinstance(self.quantizer, NVFP4Quantizer): + self.quantizer.prime_global_scale(W) self.quantizer.find_params(W, weight=True) H = self.H diff --git a/gptqmodel/quantization/gptaq.py b/gptqmodel/quantization/gptaq.py index 79ef7f0dd..4472f7f65 100644 --- a/gptqmodel/quantization/gptaq.py +++ b/gptqmodel/quantization/gptaq.py @@ -20,6 +20,7 @@ from ..quantization import QuantizeConfig from ..utils.torch import TORCH_GTE_28, torch_compile, torch_sync from .gptq import GPTQ +from .quantizer import NVFP4Quantizer class GPTAQ(GPTQ): @@ -140,6 +141,8 @@ def quantize( W = self.module_copy self.module_copy = None + if isinstance(self.quantizer, NVFP4Quantizer): + self.quantizer.prime_global_scale(W) self.quantizer.find_params(W, weight=True) H = self.H diff --git a/gptqmodel/quantization/gptq.py b/gptqmodel/quantization/gptq.py index af04c8d49..d3cee8c85 100644 --- a/gptqmodel/quantization/gptq.py +++ b/gptqmodel/quantization/gptq.py @@ -6,12 +6,15 @@ # Based on original gptq algorithm and code from https://github.com/IST-DASLab/gptq import contextlib +import functools +import json import math import os +import re import sys import threading import time -from typing import Dict, Optional, Tuple +from typing import Callable, Dict, Optional, Tuple import numpy as np import torch @@ -21,10 +24,11 @@ from ..looper.named_module import NamedModule from ..quantization import QuantizeConfig -from ..quantization.config import FallbackStrategy, SmoothMSE +from ..quantization.config import FORMAT, FallbackStrategy, SmoothMSE from ..utils.device import get_device from ..utils.logger import setup_logger from ..utils.torch import torch_sync +from .dtype import fake_quantize_nvfp4 from .fallback_smooth import mse_optimal_quant, smooth_block from .gar import ( compose_final_perm, @@ -34,7 +38,7 @@ invert_perm, ) from .npu_linalg import npu_inverse_cholesky_factor -from .quantizer import HF_OPTIMUM, Quantizer +from .quantizer import HF_OPTIMUM, NVFP4Quantizer, Quantizer log = setup_logger() @@ -129,6 +133,28 @@ def _device_supports_bfloat16(device: torch.device) -> bool: return support +@functools.lru_cache(maxsize=8) +def _load_act_amax_table(path: str) -> Dict[str, float]: + """Load an activation amax sidecar, cached so thousands of modules read it once. + + Accepts a flat ``{key: amax}`` mapping or ``{"amax": {key: amax}}``, the shape + the NVFP4 conversion pipeline emits. + """ + + with open(path, "r") as handle: + payload = json.load(handle) + + if isinstance(payload, dict) and isinstance(payload.get("amax"), dict): + payload = payload["amax"] + + if not isinstance(payload, dict): + raise ValueError( + f"GPTQ: activation amax sidecar '{path}' must be a mapping, or contain an 'amax' mapping." + ) + + return {str(k): float(v) for k, v in payload.items()} + + def get_number_of_rows_and_cols(layer: nn.Module): # return layer.weight.shape[0], np.prod(layer.weight.shape[1:]) if isinstance(layer, NamedModule): @@ -215,6 +241,7 @@ def __init__(self, module: nn.Module, qcfg: Optional[QuantizeConfig] = None): self.nsamples = 0 self.quantizer = self.create_quantizer(name=self.name) + self._act_quantizer = self._create_act_quantizer() # fwd counter self.fwd_counter = 0 @@ -269,7 +296,90 @@ def validate_module(module): # def has_hessian_issues(self) -> bool: # return any([self.issue_zero_samples, self.issue_nan_hessian, self.issue_non_invertible]) + def _resolve_act_global_scale(self) -> Optional[torch.Tensor]: + """NVFP4's static per-tensor activation scale, from the amax sidecar. + + 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 here, so the global scale is read from a sidecar produced by + an earlier calibration pass (`amax / 6 / 448`, matching the serving + pipeline). Returns ``None`` when no sidecar is configured, leaving block + scales only. + + Keys are resolved through ``act_amax_key_rules`` so modules sharing an + input map to one entry — w1/w3 see the same MoE input and must therefore + share a scale, as must q/k/v. + """ + + path = getattr(self.qcfg, "act_amax_path", None) + if not path: + return None + + table = _load_act_amax_table(path) + + # `self.name` is layer-relative (`mlp.experts.1.gate_proj`) -- it carries no + # layer index for the key rules to capture. NamedModule.full_name is the + # full dotted path (`...layers.3.mlp.experts.1.gate_proj`), so resolve + # against that when available. + name = getattr(self._named_module, "full_name", None) or self.name + + key = name + for pattern, replacement in (getattr(self.qcfg, "act_amax_key_rules", None) or []): + if re.search(pattern, name): + key = re.sub(pattern, replacement, name) + break + + if key not in table: + # Failing loud beats silently falling back to block-only scales, which + # would quantize against a different grid than the kernel uses. + raise KeyError( + f"GPTQ: activation amax for module '{name}' (key '{key}') not found in " + f"'{path}'. Add the entry or extend `act_amax_key_rules`." + ) + + amax = float(table[key]) + if not math.isfinite(amax) or amax <= 0.0: + raise ValueError( + f"GPTQ: activation amax for '{key}' must be finite and positive, got {amax}." + ) + + # NVFP4 second-level scale: amax / F4_E2M1_MAX / F8_E4M3_MAX. + return torch.tensor(amax / 6.0 / 448.0, dtype=torch.float32) + + def _create_act_quantizer(self) -> Optional[Callable[[torch.Tensor], torch.Tensor]]: + """Build the activation fake-quantizer for w4a4, or ``None`` for weight-only. + + Under w4a4 the kernel consumes ``Q(X)``, so the Hessian must be built from + ``Q(X)`` rather than ``X`` — otherwise GPTQ solves for inputs that never + occur. Returning ``None`` preserves the weight-only behaviour exactly. + """ + + act_format = getattr(self.qcfg, "act_format", None) + if act_format is None: + return None + + if act_format != FORMAT.NVFP4: + raise ValueError(f"GPTQ: unsupported `act_format={act_format}`.") + + block_size = int(getattr(self.qcfg, "act_group_size", 16)) + # Resolved on first use, not here: `self.name` is not final at construction + # (the looper names modules afterwards), and a name-keyed sidecar lookup done + # too early silently reads the wrong entry. + resolved: Dict[str, Optional[torch.Tensor]] = {} + + def _quantize_act(x: torch.Tensor) -> torch.Tensor: + if "scale" not in resolved: + resolved["scale"] = self._resolve_act_global_scale() + scale = resolved["scale"] + if scale is not None: + scale = scale.to(x.device) + return fake_quantize_nvfp4(x, block_size=block_size, per_tensor_scale=scale) + + return _quantize_act + def create_quantizer(self, name: str) -> Quantizer: + if getattr(self.qcfg, "format", None) == FORMAT.NVFP4: + return NVFP4Quantizer(qcfg=self.qcfg, name=name) return Quantizer(qcfg=self.qcfg, name=name) def shape(self): @@ -528,6 +638,13 @@ def process_batch(self, inp: torch.Tensor) -> Tuple[int, Optional[torch.Tensor], # Delay dtype conversion until we materialize Hessian chunks to avoid unnecessary temporaries reshaped_inp = reshaped_inp.contiguous() + + # w4a4: the kernel consumes Q(X), so the Hessian must describe Q(X), not X. + # Applied before the TP zero-padding below so the pad columns never enter a + # block's amax, and after any input-scale transform baked into `inp`. + if self._act_quantizer is not None: + reshaped_inp = self._act_quantizer(reshaped_inp) + if self._tp_pad_cols: pad = reshaped_inp.new_zeros((reshaped_inp.shape[0], self._tp_pad_cols)) reshaped_inp = torch.cat((reshaped_inp, pad), dim=1) @@ -679,6 +796,11 @@ def _fallback_quantize(self, strategy: FallbackStrategy, blocksize: int): scale_chunks = [] zero_chunks = [] + if isinstance(self.quantizer, NVFP4Quantizer): + # Same two-level priming as the main path: the RTN fallback quantizes + # group by group and would otherwise floor-clamp every block scale. + self.quantizer.prime_global_scale(W) + for start in range(0, self.columns, effective_group_size): end = min(start + effective_group_size, self.columns) block = W[:, start:end] @@ -948,8 +1070,11 @@ def quantize( use_hessian = False threshold_text = str(getattr(self.fallback, "threshold", None)) threshold_info = f", threshold_raw={threshold_raw}" if threshold_raw is not None and is_percent else "" + # full_name where available: `self.name` is layer-relative and repeats + # across layers, which makes fallback logs ambiguous in multi-layer runs. + fallback_log_name = getattr(self._named_module, "full_name", None) or self.name log.warn( - f"Quantization: Module `{self.name}` -> " + f"Quantization: Module `{fallback_log_name}` -> " f"Using `{resolved_strategy.value}` fallback quantization (observed {self.nsamples} samples, threshold={threshold_text}{threshold_info}, max_total={self.expected_nsamples})." ) self.H = self.create_H(target_device=target_device) @@ -983,6 +1108,10 @@ def quantize( W = self.module_copy.to(device=self.H.device) del self.module_copy + if isinstance(self.quantizer, NVFP4Quantizer): + # Two-level NVFP4 needs the per-tensor global scale fixed from the FULL + # weight before any per-group find_params -- groups alone floor-clamp. + self.quantizer.prime_global_scale(W) self.quantizer.find_params(W, weight=True) # H = self.H.to(device=self.H.device) @@ -1244,6 +1373,12 @@ def quantize( if math.isnan(avg_loss): print("Losses sum item:", torch.sum(Losses).item()) if fallback_configured: + if getattr(self.qcfg, "format", None) == FORMAT.NVFP4: + # The mock retry quantizes on an integer grid and reuses one + # scale per block, both off the NVFP4 grid; the RTN fallback + # goes through NVFP4Quantizer and stays on it. + log.info(f"Quantization: Failed due to `NaN` loss for `{self.name}`, using `{resolved_strategy.value}` fallback quantization for `{self.name}`") + return self._fallback_quantize(resolved_strategy, blocksize) log.info(f"Quantization: Failed due to `NaN` loss for `{self.name}`, use mock quantization retry for `{self.name}`") self.qcfg.mock_quantization = True return self.quantize(blocksize=blocksize) diff --git a/gptqmodel/quantization/quantizer.py b/gptqmodel/quantization/quantizer.py index fa99c5542..39ff6e21b 100644 --- a/gptqmodel/quantization/quantizer.py +++ b/gptqmodel/quantization/quantizer.py @@ -5,6 +5,8 @@ # adapted from @qwopqwop200 's [GPTQ-for-LLaMa](https://github.com/qwopqwop200/GPTQ-for-LLaMa/tree/cuda), which itself is based on [gptq](https://github.com/IST-DASLab/gptq) +from typing import Optional + import torch import torch.nn as nn @@ -173,4 +175,135 @@ class QQQQuantizer(Quantizer): def requires_groupwise_processing(self) -> bool: return self.qcfg.group_size == -1 and self.qcfg.sym -__all__ = ["Quantizer"] + +try: + from torchao.prototype.mx_formats.kernels import f4_unpacked_to_f32, f32_to_f4_unpacked +except Exception: + f4_unpacked_to_f32 = None + f32_to_f4_unpacked = None + +_F4_E2M1_MAX = 6.0 +_F8_E4M3_MAX = 448.0 + + +class NVFP4Quantizer(Quantizer): + """Quantizer whose grid is NVFP4: E2M1 values scaled per 16-wide block by E4M3 scales. + + Scale derivation and value rounding mirror torchao's ``nvfp4_quantize`` exactly so + the GPTQ-corrected weight is bit-identical to the packed NVFP4 checkpoint layout. + """ + + def configure(self, *args, **kwargs): + if f32_to_f4_unpacked is None or f4_unpacked_to_f32 is None: + raise RuntimeError("NVFP4Quantizer requires torchao with NVFP4 support.") + super().configure(*args, **kwargs) + self.global_scale: Optional[torch.Tensor] = None + + def prime_global_scale(self, weight: torch.Tensor) -> torch.Tensor: + """Fix the per-tensor second-level scale from the FULL weight, before the loop. + + NVFP4 is two-level: block scales are E4M3 *relative to* a per-tensor global + scale. Without one, ``find_params`` derives absolute block scales + ``amax/6`` -- for realistic weights those sit below E4M3's smallest normal + (2^-6) and get floor-clamped (measured 96.8% of blocks on Minimax-M3 + experts), compressing the grid and costing ~1.4x reconstruction error. + + The global scale is snapped to a power of two, ``2^ceil(log2(amax/(6*448)))`` + (tore-quant's lossless-conversion snap), for exactness downstream: the + largest block's E4M3 ratio lands in (224, 448] by construction, so + ``_pack_with_gptq_scales`` re-derives exactly this value from the effective + scales and the two-level split reproduces the loop's grid bit-for-bit. + + find_params only ever sees one 16-wide group at a time, so this must be + called once per module with the full weight; an unprimed find_params raises. + """ + + amax = weight.detach().abs().amax().to(torch.float32) + if amax > 0: + exponent = torch.ceil(torch.log2(amax / (_F4_E2M1_MAX * _F8_E4M3_MAX))) + self.global_scale = torch.exp2(exponent).reshape(()) + else: + self.global_scale = torch.ones((), dtype=torch.float32) + return self.global_scale + + def _sweep_block_scales( + self, x_flat: torch.Tensor, base_ratio: torch.Tensor, gs: torch.Tensor, steps: int + ) -> torch.Tensor: + """Per-block MSE search over neighbouring E4M3 scale codes. + + The amax-derived scale is only optimal for a block's single largest value; + a slightly smaller scale clips that one value but resolves the other 15 + finer, and usually wins on MSE (the ModelOpt-RTN reference checkpoints do + exactly this: ~75% of their block scales differ from amax-derived, worth + ~13% reconstruction error). Positive-E4M3 bit patterns are monotonic, so + neighbouring codes are integer offsets on the uint8 view. Candidates span + [code-(steps-2), code+1]; the swept scale is still e4m3 * gs, so pack's + two-level split stays exact. + """ + + codes = base_ratio.view(torch.uint8).to(torch.int16) # (rows,) + offsets = torch.arange(-(steps - 2), 2, device=codes.device, dtype=torch.int16) + cand = torch.clamp(codes.unsqueeze(1) + offsets, 1, 126).to(torch.uint8) # (rows, C) + cand_scale = cand.view(torch.float8_e4m3fn).to(torch.float32) * gs # (rows, C) + + # Quantize every block against every candidate with the exact production + # rounding path (reciprocal multiply -> clamp -> E2M1 round-trip). + scaled = torch.clamp( + x_flat.unsqueeze(1) * cand_scale.unsqueeze(-1).reciprocal(), + -_F4_E2M1_MAX, _F4_E2M1_MAX, + ) # (rows, C, 16) + deq = f4_unpacked_to_f32(f32_to_f4_unpacked(scaled.contiguous())) * cand_scale.unsqueeze(-1) + err = (deq - x_flat.unsqueeze(1)).pow(2).sum(dim=-1) # (rows, C) + best = err.argmin(dim=1) # (rows,) + return cand_scale.gather(1, best.unsqueeze(1)).squeeze(1) + + def find_params(self, x, weight=False): + if not weight: + raise NotImplementedError("NVFP4Quantizer only supports weight quantization.") + + dev = x.device + x_flat = x.flatten(1).to(torch.float32) + + amax = x_flat.abs().amax(dim=1) + gs = getattr(self, "global_scale", None) + if gs is None: + # No single-level fallback: absolute amax/6 block scales sit below + # E4M3's 2^-6 normal floor on realistic weights (measured 96.8% of + # blocks clamped, 1.4x reconstruction error). An unprimed caller is a + # bug, not a mode. + raise RuntimeError( + "NVFP4Quantizer.find_params called without a primed global scale; " + "call prime_global_scale(full_weight) before per-group find_params." + ) + # Two-level (torchao nvfp4_quantize with per_tensor_scale): block ratio + # amax/(6*gs) uses E4M3's full range; effective scale = e4m3 * gs. + gs_dev = gs.to(amax.device) + ratio = amax / (_F4_E2M1_MAX * gs_dev) + ratio = torch.clamp(ratio, min=torch.finfo(torch.float8_e4m3fn).tiny, max=_F8_E4M3_MAX) + ratio_e4m3 = ratio.to(torch.float8_e4m3fn) + + sweep = int(getattr(self.qcfg, "nvfp4_scale_sweep", 0) or 0) + if sweep > 1: + scale = self._sweep_block_scales(x_flat, ratio_e4m3, gs_dev, sweep) + else: + scale = ratio_e4m3.to(torch.float32) * gs_dev + + shape = [-1] + [1] * (len(x.shape) - 1) + self.scale = scale.reshape(shape).to(device=dev) + self.zero = torch.zeros_like(self.scale) + + def quantize(self, x): + orig_dtype = x.dtype + # Multiply by the reciprocal rather than dividing, matching torchao's + # nvfp4_quantize (and the MSLK triton kernel it tracks). `x / s` and + # `x * (1/s)` differ by up to 1 ULP in fp32, and with only 16 E2M1 codes a + # value on a rounding boundary can land on a different code. + scaled = torch.clamp( + x.to(torch.float32) * self.scale.reciprocal(), -_F4_E2M1_MAX, _F4_E2M1_MAX + ) + codes = f32_to_f4_unpacked(scaled.contiguous()) + dequant = f4_unpacked_to_f32(codes) * self.scale + return dequant.to(orig_dtype) + + +__all__ = ["Quantizer", "QQQQuantizer", "NVFP4Quantizer"] diff --git a/gptqmodel/utils/backend.py b/gptqmodel/utils/backend.py index 9d4dde801..301f779e8 100644 --- a/gptqmodel/utils/backend.py +++ b/gptqmodel/utils/backend.py @@ -50,6 +50,9 @@ class BACKEND(str, Enum): # FP8 kernels FP8_TORCH = "fp8_torch" + # NVFP4 kernels + NVFP4_TORCH = "nvfp4_torch" + # GGUF kernels / engines GGUF_TORCH = "gguf_torch" GGUF_TRITON = "gguf_triton" @@ -130,6 +133,9 @@ class PROFILE(str, Enum): "fp8": { BACKEND.TORCH: BACKEND.FP8_TORCH, }, + "nvfp4": { + BACKEND.TORCH: BACKEND.NVFP4_TORCH, + }, "exl3": { BACKEND.EXLLAMA_V3: BACKEND.EXL3_EXLLAMA_V3, BACKEND.TORCH: BACKEND.EXL3_TORCH, diff --git a/gptqmodel/utils/model.py b/gptqmodel/utils/model.py index e58e04f14..88692d74e 100644 --- a/gptqmodel/utils/model.py +++ b/gptqmodel/utils/model.py @@ -40,6 +40,7 @@ DEVICE, EXPERT_INDEX_PLACEHOLDER, SUPPORTS_MODULE_TYPES, + normalize_device, ) from ..nn_modules.qlinear import BaseQuantLinear, GPTQQuantLinear from ..nn_modules.qlinear.exllamav2 import ExllamaV2Linear @@ -622,7 +623,7 @@ def create_quant_module( dtype=dtype, in_features=in_features, out_features=out_features, - device=DEVICE(device) if isinstance(device, str) else device, + device=normalize_device(device) if device is not None else None, adapter=adapter, # TODO FIX ME..need to pass Eora if loaded ) if err is not None: diff --git a/gptqmodel/utils/structure.py b/gptqmodel/utils/structure.py index 6edf2ef33..bd4d3077c 100644 --- a/gptqmodel/utils/structure.py +++ b/gptqmodel/utils/structure.py @@ -44,6 +44,7 @@ from safetensors import safe_open from torch import nn +from ..quantization.dtype import dequantize_block_fp8, fp8_storage_dtypes from ..utils.logger import setup_logger @@ -2294,6 +2295,58 @@ def _load_checkpoint_tensors_for_module_path( tensors[rel_name] = handler.get_tensor(full_name) return tensors + def _materialize_dtype(self) -> torch.dtype: + """Compute dtype for dequantized checkpoint tensors.""" + + for cfg in (getattr(self.config, "text_config", None), self.config): + dtype = getattr(cfg, "torch_dtype", None) or getattr(cfg, "dtype", None) + if isinstance(dtype, torch.dtype): + return dtype + if isinstance(dtype, str) and hasattr(torch, dtype): + resolved = getattr(torch, dtype) + if isinstance(resolved, torch.dtype): + return resolved + return torch.bfloat16 + + def _scale_tensor_name(self, full_name: str) -> Optional[str]: + """Checkpoint name of the block scale paired with ``full_name``, if any.""" + + candidate = f"{full_name}_scale_inv" + return candidate if candidate in self._weight_map else None + + def _read_checkpoint_tensor(self, handler, shard_path: str, full_name: str) -> torch.Tensor: + """Read a checkpoint tensor, applying block scales to FP8 payloads. + + An FP8 checkpoint stores only the mantissa payload; the magnitude lives + in a sibling ``*_scale_inv`` tensor. Copying the payload into a bf16 + parameter with a plain dtype cast silently drops that scale, so the + dequant has to happen here, before the tensor reaches the target param. + + Layers the exporter left unquantized (``ignored_layers``) carry no scale + tensor, so the lookup misses and they keep the plain-cast path. + """ + + tensor = handler.get_tensor(full_name) + if tensor.dtype not in fp8_storage_dtypes(): + return tensor + + scale_name = self._scale_tensor_name(full_name) + if scale_name is None: + return tensor + + scale_shard = self._weight_map.get(scale_name) + if scale_shard is None: + return tensor + + scale_path = os.path.join(self.model_local_path, scale_shard) + if scale_path == shard_path: + scale = handler.get_tensor(scale_name) + else: + with safe_open(scale_path, framework="pt", device="cpu") as scale_handler: + scale = scale_handler.get_tensor(scale_name) + + return dequantize_block_fp8(tensor, scale, target_dtype=self._materialize_dtype()) + def _copy_checkpoint_tensors_into_submodule( self, *, @@ -2414,7 +2467,9 @@ def _copy_checkpoint_tensors_into_submodule( ) shard_path = os.path.join(self.model_local_path, shard) with safe_open(shard_path, framework="pt", device="cpu") as handler: - parts.append(handler.get_tensor(full_name)) + parts.append( + self._read_checkpoint_tensor(handler, shard_path, full_name) + ) try: tensor = torch.cat(parts, dim=concat_dim).contiguous() @@ -2511,7 +2566,9 @@ def _copy_checkpoint_tensors_into_submodule( rel_name=rel_name, modules_by_name=modules_by_name, ) - checkpoint_tensor = handler.get_tensor(full_name) + checkpoint_tensor = self._read_checkpoint_tensor( + handler, shard_path, full_name + ) tensor = self._transform_checkpoint_tensor( checkpoint_tensor, expert_index=expert_index, @@ -2759,6 +2816,15 @@ def _materialize_direct_meta_tensors( ) -> int: synced = 0 + # Quantized linear tensors are produced by quantization and restored from the + # quant offload directory, never from the source checkpoint. Kernels whose + # buffer names shadow checkpoint tensors (e.g. NVFP4/FP8 `weight`) must not be + # re-aliased from the original dense weights here. + from ..nn_modules.qlinear import BaseQuantLinear + + if isinstance(shell_sub, BaseQuantLinear): + return synced + with torch.inference_mode(): for name, shell_param in dict(shell_sub.named_parameters(recurse=False)).items(): if not _is_meta_tensor(shell_param): @@ -2782,7 +2848,9 @@ def _materialize_direct_meta_tensors( )) source_path = os.path.join(self.model_local_path, shard) with safe_open(source_path, framework="pt", device="cpu") as handler: - parts.append(handler.get_tensor(full_name)) + parts.append( + self._read_checkpoint_tensor(handler, source_path, full_name) + ) try: source_param = torch.cat(parts, dim=concat_dim).contiguous() @@ -2844,7 +2912,9 @@ def _materialize_direct_meta_tensors( source_path = os.path.join(self.model_local_path, shard) with safe_open(source_path, framework="pt", device="cpu") as handler: - checkpoint_param = handler.get_tensor(full_name) + checkpoint_param = self._read_checkpoint_tensor( + handler, source_path, full_name + ) source_param = self._transform_checkpoint_tensor( checkpoint_param, expert_index=expert_index, @@ -2919,7 +2989,9 @@ def _materialize_direct_meta_tensors( )) source_path = os.path.join(self.model_local_path, shard) with safe_open(source_path, framework="pt", device="cpu") as handler: - parts.append(handler.get_tensor(full_name)) + parts.append( + self._read_checkpoint_tensor(handler, source_path, full_name) + ) try: source_buffer = torch.cat(parts, dim=concat_dim).contiguous() @@ -2979,7 +3051,9 @@ def _materialize_direct_meta_tensors( source_path = os.path.join(self.model_local_path, shard) with safe_open(source_path, framework="pt", device="cpu") as handler: - checkpoint_buffer = handler.get_tensor(full_name) + checkpoint_buffer = self._read_checkpoint_tensor( + handler, source_path, full_name + ) source_buffer = self._transform_checkpoint_tensor( checkpoint_buffer, expert_index=expert_index, diff --git a/scripts/minimax_m3_nvfp4/amax_per_layer.json b/scripts/minimax_m3_nvfp4/amax_per_layer.json new file mode 100644 index 000000000..9a844320a --- /dev/null +++ b/scripts/minimax_m3_nvfp4/amax_per_layer.json @@ -0,0 +1,130 @@ +{ + "amax": { + "layer3.moe_input": 2.453125, + "layer3.w2_input": 17.625, + "layer4.moe_input": 3.140625, + "layer4.w2_input": 24.75, + "layer5.moe_input": 3.203125, + "layer5.w2_input": 24.875, + "layer6.moe_input": 2.734375, + "layer6.w2_input": 27.25, + "layer7.moe_input": 2.984375, + "layer7.w2_input": 28.75, + "layer8.moe_input": 2.8125, + "layer8.w2_input": 41.5, + "layer9.moe_input": 3.171875, + "layer9.w2_input": 52.75, + "layer10.moe_input": 2.875, + "layer10.w2_input": 56.0, + "layer11.moe_input": 3.03125, + "layer11.w2_input": 49.5, + "layer12.moe_input": 3.203125, + "layer12.w2_input": 56.0, + "layer13.moe_input": 3.125, + "layer13.w2_input": 56.0, + "layer14.moe_input": 2.96875, + "layer14.w2_input": 56.0, + "layer15.moe_input": 2.875, + "layer15.w2_input": 56.0, + "layer16.moe_input": 3.109375, + "layer16.w2_input": 56.0, + "layer17.moe_input": 2.625, + "layer17.w2_input": 56.0, + "layer18.moe_input": 2.484375, + "layer18.w2_input": 56.0, + "layer19.moe_input": 2.890625, + "layer19.w2_input": 56.0, + "layer20.moe_input": 2.75, + "layer20.w2_input": 56.0, + "layer21.moe_input": 3.046875, + "layer21.w2_input": 56.0, + "layer22.moe_input": 4.28125, + "layer22.w2_input": 56.0, + "layer23.moe_input": 4.28125, + "layer23.w2_input": 56.0, + "layer24.moe_input": 5.1875, + "layer24.w2_input": 56.0, + "layer25.moe_input": 5.1875, + "layer25.w2_input": 56.0, + "layer26.moe_input": 4.5625, + "layer26.w2_input": 56.0, + "layer27.moe_input": 4.875, + "layer27.w2_input": 56.0, + "layer28.moe_input": 3.96875, + "layer28.w2_input": 56.0, + "layer29.moe_input": 3.96875, + "layer29.w2_input": 56.0, + "layer30.moe_input": 4.28125, + "layer30.w2_input": 56.0, + "layer31.moe_input": 3.5625, + "layer31.w2_input": 56.0, + "layer32.moe_input": 4.28125, + "layer32.w2_input": 56.0, + "layer33.moe_input": 4.21875, + "layer33.w2_input": 56.0, + "layer34.moe_input": 4.96875, + "layer34.w2_input": 56.0, + "layer35.moe_input": 5.78125, + "layer35.w2_input": 56.0, + "layer36.moe_input": 5.28125, + "layer36.w2_input": 56.0, + "layer37.moe_input": 4.9375, + "layer37.w2_input": 56.0, + "layer38.moe_input": 4.875, + "layer38.w2_input": 56.0, + "layer39.moe_input": 4.8125, + "layer39.w2_input": 56.0, + "layer40.moe_input": 4.5, + "layer40.w2_input": 56.0, + "layer41.moe_input": 4.5, + "layer41.w2_input": 56.0, + "layer42.moe_input": 4.53125, + "layer42.w2_input": 56.0, + "layer43.moe_input": 4.65625, + "layer43.w2_input": 56.0, + "layer44.moe_input": 4.59375, + "layer44.w2_input": 56.0, + "layer45.moe_input": 4.84375, + "layer45.w2_input": 56.0, + "layer46.moe_input": 5.15625, + "layer46.w2_input": 56.0, + "layer47.moe_input": 5.0, + "layer47.w2_input": 56.0, + "layer48.moe_input": 6.625, + "layer48.w2_input": 56.0, + "layer49.moe_input": 6.3125, + "layer49.w2_input": 56.0, + "layer50.moe_input": 6.90625, + "layer50.w2_input": 56.0, + "layer51.moe_input": 8.625, + "layer51.w2_input": 56.0, + "layer52.moe_input": 10.0625, + "layer52.w2_input": 56.0, + "layer53.moe_input": 11.75, + "layer53.w2_input": 56.0, + "layer54.moe_input": 15.0625, + "layer54.w2_input": 56.0, + "layer55.moe_input": 15.9375, + "layer55.w2_input": 56.0, + "layer56.moe_input": 18.0, + "layer56.w2_input": 56.0, + "layer57.moe_input": 18.25, + "layer57.w2_input": 56.0, + "layer58.moe_input": 23.625, + "layer58.w2_input": 56.0, + "layer59.moe_input": 31.875, + "layer59.w2_input": 56.0 + }, + "meta": { + "model": "/data/huggingface/Minimax-M3-0602", + "samples": 64, + "max_tokens": 8192, + "seed": 42, + "total_tokens": 470987, + "method": "plain max over |x|, reduced across experts (peer-max)", + "sites": { + "moe_input": "gate_proj/up_proj input", + "w2_input": "down_proj input" + } + } +} \ No newline at end of file diff --git a/scripts/minimax_m3_nvfp4/build_variant.py b/scripts/minimax_m3_nvfp4/build_variant.py new file mode 100644 index 000000000..abde6c3b3 --- /dev/null +++ b/scripts/minimax_m3_nvfp4/build_variant.py @@ -0,0 +1,126 @@ +"""Build serving variants of the w4a4 -final checkpoint. + +Two orthogonal transforms, matching the team's existing checkpoint families: + +* ``input_scale`` -> 1.0 everywhere (the `-inscale1` convention; the serving + stack's dynamic activation quantization ignores/conflicts with static scales, + and the nvidia donor ships 1.0). +* Optionally swap whole MoE layers back to source MXFP8 (the alt3x6 / midsplit2 + placement recipes): those layers' experts become the raw source + {weight F8_E4M3, weight_scale_inv U8} pairs, dropping the NVFP4 quad entirely. + + FIN=/data/huggingface/Minimax-M3-0602-NVFP4-GPTQ-w4a4-final \ + SRC=/data/huggingface/Minimax-M3-0602 \ + OUT=... MXFP8_LAYERS="22-28,32-38" \ + python scripts/minimax_m3_nvfp4/build_variant.py +""" + +import json +import os +import re +import shutil +import time + +import torch +from safetensors import safe_open +from safetensors.torch import save_file + +FIN = os.environ["FIN"] +SRC = os.environ.get("SRC", "/data/huggingface/Minimax-M3-0602") +OUT = os.environ["OUT"] +MXFP8_LAYERS = os.environ.get("MXFP8_LAYERS", "").strip() +SHARD_BYTES = int(os.environ.get("SHARD_BYTES", str(8 << 30))) + +t0 = time.time() +def log(m): print(f"[variant +{time.time()-t0:7.1f}s] {m}", flush=True) + +def parse_layers(spec): + out = set() + for part in spec.split(","): + part = part.strip() + if not part: continue + if "-" in part: + a, b = part.split("-"); out |= set(range(int(a), int(b) + 1)) + else: + out.add(int(part)) + return out + +mx_layers = parse_layers(MXFP8_LAYERS) +log(f"input_scale -> 1.0 everywhere; mxfp8 layers: {sorted(mx_layers) or 'none'}") + +fin_idx = json.load(open(f"{FIN}/model.safetensors.index.json"))["weight_map"] +src_idx = json.load(open(f"{SRC}/model.safetensors.index.json"))["weight_map"] + +EXPERT_RE = re.compile(r"^language_model\.model\.layers\.(\d+)\.block_sparse_moe\.experts\.(\d+)\.(w[123])\.") +ONE = torch.ones((), dtype=torch.float32) + +# plan in FIN order; for mxfp8-layer expert modules, emit the source pair at +# first encounter and skip the rest of the quad +plan, emitted = [], set() +for name in fin_idx: + m = EXPERT_RE.match(name) + if m and int(m.group(1)) in mx_layers: + base = name.rsplit(".", 1)[0] if not name.endswith("_scale_2") and not name.endswith("_scale") and not name.endswith("input_scale") else None + base = f"language_model.model.layers.{m.group(1)}.block_sparse_moe.experts.{m.group(2)}.{m.group(3)}" + if base in emitted: + continue + emitted.add(base) + for suffix in (".weight", ".weight_scale_inv"): + if base + suffix not in src_idx: + raise SystemExit(f"source missing {base}{suffix}") + plan.append((base + suffix, "src", base + suffix)) + elif name.endswith(".input_scale"): + plan.append((name, "one", None)) + else: + plan.append((name, "fin", name)) + +n_src = sum(1 for _, o, _ in plan if o == "src") +n_one = sum(1 for _, o, _ in plan if o == "one") +log(f"plan: {len(plan)} tensors ({n_src} mxfp8-source, {n_one} input_scale->1.0, " + f"{len(plan) - n_src - n_one} copied)") + +os.makedirs(OUT, exist_ok=True) +handles = {} +def read(origin, key): + if origin == "one": + return ONE.clone() + root, index = (SRC, src_idx) if origin == "src" else (FIN, fin_idx) + ck = (root, index[key]) + if ck not in handles: + handles.clear() + handles[ck] = safe_open(os.path.join(root, index[key]), framework="pt") + return handles[ck].get_tensor(key) + +weight_map, buffer, buf_bytes, shard_id = {}, {}, 0, 0 +def flush(): + global buffer, buf_bytes, shard_id + if not buffer: return + shard_id += 1 + fname = f"model-{shard_id:05d}.safetensors" + save_file(buffer, os.path.join(OUT, fname)) + for k in buffer: weight_map[k] = fname + buffer, buf_bytes = {}, 0 + +for out_name, origin, key in plan: + t = read(origin, key) + buffer[out_name] = t + buf_bytes += t.numel() * t.element_size() + if buf_bytes >= SHARD_BYTES: + flush() +flush() + +final = {f"model-{i:05d}.safetensors": f"model-{i:05d}-of-{shard_id:05d}.safetensors" + for i in range(1, shard_id + 1)} +for old, new in final.items(): + os.replace(os.path.join(OUT, old), os.path.join(OUT, new)) +weight_map = {k: final[v] for k, v in weight_map.items()} +with open(f"{OUT}/model.safetensors.index.json", "w") as f: + json.dump({"metadata": {"total_size": sum( + os.path.getsize(os.path.join(OUT, s)) for s in set(weight_map.values()))}, + "weight_map": weight_map}, f, indent=1) + +for fname in os.listdir(FIN): + if fname.endswith((".json", ".py", ".jinja", ".txt")) and fname != "model.safetensors.index.json": + shutil.copy2(os.path.join(FIN, fname), os.path.join(OUT, fname)) +log(f"wrote {shard_id} shards, {len(weight_map)} tensors -> {OUT}") +log("DONE") diff --git a/scripts/minimax_m3_nvfp4/coverage_layer_expert.json b/scripts/minimax_m3_nvfp4/coverage_layer_expert.json new file mode 100644 index 000000000..f6648442a --- /dev/null +++ b/scripts/minimax_m3_nvfp4/coverage_layer_expert.json @@ -0,0 +1 @@ +{"model": "/data/huggingface/Minimax-M3-0602", "samples": 64, "max_tokens": 8192, "seed": 42, "total_tokens": 470987, "top_k": 4, "counts": {"3": [48, 9713, 33972, 20998, 35310, 20325, 14853, 25354, 8256, 24547, 4756, 15106, 19268, 40840, 3699, 755, 8893, 11940, 2025, 19313, 2380, 3070, 6007, 27750, 173, 20587, 15621, 66, 8998, 6936, 97, 6707, 343, 4614, 2499, 1542, 2394, 19831, 314, 44826, 20011, 29111, 15465, 15165, 23308, 283251, 5271, 3311, 2571, 100, 14042, 2546, 11704, 15353, 13562, 13629, 56, 774, 5459, 13213, 7367, 20864, 28, 12044, 8119, 7953, 18966, 30388, 1312, 4, 14799, 125, 1150, 15735, 0, 1104, 364, 5719, 34892, 9792, 8231, 33741, 3731, 11818, 4909, 14112, 10302, 8936, 9850, 2901, 7016, 1426, 34032, 588, 23837, 2218, 3921, 178520, 3974, 15146, 978, 2424, 18262, 9534, 2595, 345, 35893, 64, 29770, 1249, 4188, 8168, 19375, 902, 637, 2316, 24332, 7038, 26771, 6416, 32750, 13609, 2895, 10903, 1563, 58480, 7929, 2000], "4": [6288, 4776, 1199, 35970, 1769, 116, 326, 14913, 25277, 186, 259034, 23180, 1747, 38658, 28831, 10289, 15118, 250, 444, 2156, 515, 331, 1801, 7178, 233, 14692, 23784, 8076, 2252, 14356, 344, 6184, 37154, 731, 28732, 1936, 4214, 1659, 377, 8141, 136, 36308, 6146, 18525, 31597, 4685, 2160, 885, 353, 23512, 7431, 44910, 5409, 1943, 783, 2338, 626, 3750, 6082, 8186, 5668, 2680, 785, 562, 7230, 1644, 1068, 8844, 76368, 2145, 2608, 7587, 3, 28847, 51735, 6105, 55006, 5848, 9496, 23, 7757, 2926, 90, 17396, 5, 20760, 26098, 3603, 55001, 60, 12, 7091, 9412, 6285, 232868, 1265, 15072, 14914, 2299, 236, 33403, 22066, 4931, 54097, 13037, 1628, 366, 27445, 4395, 26916, 3967, 12225, 40793, 5400, 2053, 5133, 6583, 8444, 5346, 6216, 37834, 2856, 5133, 12914, 992, 2898, 294, 199], "5": [20247, 14495, 11870, 3134, 6137, 4003, 25463, 14471, 22651, 360, 44326, 65258, 3612, 28043, 29264, 196, 388, 3109, 5322, 12835, 10032, 15814, 6199, 2828, 24859, 3185, 33957, 8863, 2780, 20839, 7236, 19155, 9, 10773, 3391, 17793, 11100, 13963, 323, 11879, 11547, 12823, 26193, 403, 1814, 14702, 4350, 1373, 35914, 5315, 23290, 9782, 5214, 14471, 24000, 5216, 25503, 4792, 312, 17800, 9088, 18856, 25762, 2641, 329536, 8889, 627, 8270, 9481, 9145, 12676, 1415, 12989, 7473, 14363, 33450, 2825, 21381, 17548, 14375, 3703, 615, 1264, 8662, 19652, 7120, 29861, 9962, 355, 153, 3862, 2025, 6794, 15856, 3446, 16787, 30899, 11951, 3791, 14457, 2771, 27299, 5685, 3532, 4483, 781, 155129, 13676, 2112, 1523, 763, 9313, 3395, 3415, 14874, 4024, 19027, 3934, 10657, 625, 3471, 7129, 10791, 11956, 3291, 8928, 13648, 735], "6": [14338, 6964, 11085, 31597, 26411, 11231, 11150, 35818, 11402, 29760, 7797, 1944, 13866, 7471, 21079, 7546, 29814, 10939, 51327, 1464, 12529, 7253, 36162, 30660, 35188, 11144, 10777, 2564, 5230, 39685, 5855, 4667, 8066, 11349, 29098, 11089, 9302, 1914, 16006, 6284, 14518, 37296, 16026, 6182, 7705, 10992, 6160, 37803, 4649, 35879, 23223, 4120, 55544, 53, 3731, 4577, 15537, 8519, 14641, 13582, 3403, 1716, 11310, 776, 35997, 5729, 12843, 46003, 32972, 15478, 793, 8853, 29422, 6136, 2625, 6001, 8954, 13460, 14454, 2951, 222, 11385, 19947, 43656, 1898, 953, 28757, 3925, 2183, 11788, 52992, 1075, 23095, 467, 2020, 10923, 7812, 8407, 3442, 18195, 9210, 7149, 30629, 16823, 464, 17922, 18417, 4611, 3763, 16959, 3778, 2662, 4894, 13999, 104270, 18208, 685, 931, 3985, 463, 987, 4017, 33248, 29793, 29891, 15928, 13609, 5073], "7": [8503, 22718, 1218, 12764, 18329, 13000, 63879, 6956, 2920, 8147, 7941, 43842, 217893, 642, 10315, 3429, 621, 6219, 7766, 34997, 39646, 8541, 20355, 30709, 592, 42331, 12382, 14779, 4700, 27435, 1471, 772, 2123, 5776, 15230, 6800, 16432, 3401, 538, 9645, 1237, 3019, 18165, 19172, 20176, 299, 11896, 673, 2944, 587, 206382, 2061, 18352, 1165, 4374, 18814, 64563, 4889, 501, 4337, 13616, 9215, 2184, 37880, 231, 7545, 6215, 4324, 18157, 3142, 8286, 8051, 7516, 1073, 9302, 11501, 3563, 1677, 6, 5835, 1877, 4351, 33219, 9586, 12749, 4977, 85, 11966, 4704, 12635, 29408, 17719, 819, 1374, 10928, 16752, 8458, 2373, 4412, 10340, 272, 2954, 19572, 19200, 941, 16037, 5188, 15056, 15352, 23580, 42534, 338, 8343, 16306, 7582, 10695, 9768, 14010, 19031, 3221, 17088, 45732, 10071, 22980, 412, 1716, 2761, 15794], "8": [13363, 51120, 9531, 19345, 14423, 9119, 9687, 45415, 1424, 2060, 26419, 5632, 18271, 16040, 6565, 13969, 9692, 1692, 24633, 15069, 3099, 25864, 34010, 26863, 6054, 48854, 36367, 10612, 3091, 5296, 3138, 16120, 2, 1317, 9178, 7358, 8073, 189, 48216, 4015, 10446, 268, 2396, 29818, 27088, 4672, 481, 6330, 2441, 14753, 7797, 32936, 19646, 4397, 11493, 5950, 21069, 2522, 2474, 6878, 2470, 6859, 3035, 1571, 14540, 4896, 7106, 0, 51, 7334, 2640, 8980, 2566, 23014, 17651, 8295, 22096, 1112, 2091, 7217, 12062, 6947, 2434, 9995, 11589, 596, 9578, 3742, 2414, 2132, 49494, 19057, 604, 827, 4900, 2197, 1605, 851, 147636, 363, 497, 48925, 12267, 1016, 12030, 4184, 12365, 7122, 5094, 880, 12, 16302, 32281, 2195, 3785, 135969, 5535, 15364, 14484, 870, 7508, 34978, 34872, 7524, 177418, 5532, 17585, 9767], "9": [13671, 24064, 40360, 6329, 8790, 112261, 7384, 5205, 1999, 19567, 157127, 1560, 2427, 1141, 9968, 19925, 20665, 69408, 15306, 239, 1758, 13744, 10890, 23410, 16155, 5498, 3689, 7203, 29908, 30882, 3095, 3209, 7243, 30330, 8685, 9867, 12920, 57, 1440, 41316, 9720, 1998, 5918, 59886, 431, 2902, 816, 19479, 909, 19077, 3792, 440, 27600, 17820, 39425, 1201, 2491, 10516, 20071, 7201, 13489, 1734, 3268, 2912, 17726, 8857, 14768, 3711, 5126, 2079, 58924, 22400, 363, 7021, 7027, 21, 2214, 9640, 5423, 244, 17897, 19670, 5375, 59368, 40351, 5110, 1734, 40, 16264, 671, 1009, 2869, 35288, 91, 48882, 15918, 1593, 109997, 1822, 3760, 16305, 20937, 3618, 6910, 14561, 5309, 11930, 941, 415, 10345, 3900, 2647, 10202, 3383, 30520, 13214, 16195, 7222, 1799, 8288, 11969, 11221, 29126, 2863, 9361, 2595, 2624, 504], "10": [42641, 1593, 4107, 187, 963, 27, 20913, 19102, 13547, 10068, 3, 41402, 12452, 25860, 30655, 36185, 7108, 22152, 1016, 1296, 7535, 26229, 32949, 38414, 22712, 2980, 16500, 4094, 74728, 5387, 10604, 21238, 5652, 9970, 8889, 586, 17431, 51043, 2372, 3880, 21896, 204, 618, 6895, 2648, 4602, 5650, 46146, 4638, 242, 39046, 33015, 2741, 5708, 1914, 1267, 2438, 7346, 5537, 2314, 3986, 160, 18377, 8290, 14000, 1054, 10287, 10546, 2065, 3172, 3907, 5699, 2619, 5004, 12055, 16920, 25630, 185, 1123, 21481, 5110, 7865, 32473, 4383, 45664, 9758, 10239, 46214, 32572, 726, 19756, 59029, 5029, 326, 31144, 9866, 2392, 8557, 39799, 52129, 780, 940, 15266, 11763, 13618, 23747, 16179, 5473, 21151, 36828, 12129, 39298, 10787, 2454, 25219, 17895, 6701, 46220, 18918, 33103, 1202, 8952, 10478, 8868, 2816, 7771, 11975, 34491], "11": [1976, 12987, 10869, 12781, 16126, 3660, 175539, 1908, 5516, 5235, 30853, 16903, 53448, 14383, 28348, 22365, 395, 1156, 5967, 3663, 214, 2960, 3583, 2892, 31552, 7285, 14200, 2795, 3478, 5372, 6190, 10409, 11201, 5161, 6171, 33277, 6761, 12444, 785, 19282, 8813, 844, 14063, 7601, 676, 149424, 7394, 8385, 4879, 28052, 11736, 6146, 27120, 414, 46857, 252, 934, 60858, 6828, 8465, 37725, 6742, 21759, 6406, 65863, 1305, 5805, 1617, 1055, 15176, 4474, 1773, 3716, 1, 11524, 15960, 45990, 2007, 1822, 613, 5971, 418, 81, 2598, 2610, 6371, 12045, 6521, 8598, 6077, 25224, 13182, 1372, 6936, 27534, 2770, 7466, 11695, 3093, 67211, 5212, 11322, 68353, 8232, 35901, 31957, 2348, 3560, 64595, 697, 1315, 9640, 343, 5209, 3526, 3999, 729, 3460, 386, 1688, 10508, 16912, 22778, 6206, 53536, 760, 13048, 18791], "12": [63773, 15060, 3305, 3643, 561, 2940, 3895, 4413, 28262, 137013, 17595, 7256, 5077, 566, 9681, 39132, 17971, 13630, 28158, 14900, 181, 5926, 2359, 26268, 11683, 4804, 13920, 12085, 6888, 12614, 1563, 11419, 1911, 54352, 29444, 52585, 4418, 19123, 9331, 14049, 19514, 12543, 351, 26648, 8918, 2004, 3615, 3391, 10917, 2716, 24765, 70426, 7, 12406, 810, 1731, 2483, 6892, 13347, 13272, 2124, 5131, 6446, 179013, 464, 6174, 2401, 5368, 720, 1277, 6009, 40400, 346, 1294, 23720, 804, 6954, 20669, 10606, 6395, 3086, 2050, 4028, 1945, 8235, 36793, 21433, 4251, 5284, 26939, 11699, 378, 2622, 6471, 80550, 23976, 5617, 12604, 64394, 35778, 3304, 7697, 13, 23888, 1489, 2409, 1960, 16427, 131, 1021, 15278, 15524, 16816, 13310, 9077, 113, 4158, 6012, 2524, 20573, 9981, 27082, 13026, 7794, 4944, 13503, 7823, 5113], "13": [34458, 37489, 3881, 70950, 12052, 14988, 1590, 3787, 4371, 48330, 5977, 14254, 30363, 7095, 5531, 1881, 2692, 1014, 4922, 6379, 29839, 11936, 20636, 211, 13077, 12667, 3748, 4185, 481, 36741, 13852, 7468, 44091, 19793, 2852, 2698, 4809, 95286, 268, 3373, 72175, 1892, 12732, 8393, 13342, 15312, 3397, 3517, 1247, 26696, 3, 6810, 2881, 4092, 8680, 121, 20037, 12902, 24566, 11391, 5244, 21861, 11516, 2022, 4738, 7560, 8534, 3093, 4322, 5918, 9166, 25603, 479, 8752, 476, 3410, 6046, 2742, 5498, 8724, 15614, 21047, 14535, 11836, 4747, 4763, 1026, 11561, 10079, 29591, 24180, 46365, 10015, 16612, 2051, 11809, 6567, 7364, 40307, 50412, 8999, 673, 8234, 5714, 9466, 1199, 1175, 36653, 7850, 15601, 33942, 5996, 524, 2372, 5410, 40156, 22141, 7730, 8507, 37085, 18556, 144163, 22523, 2121, 21194, 813, 27081, 11684], "14": [1060, 7572, 7888, 1707, 25935, 9900, 28269, 1458, 14280, 11354, 29522, 21061, 181702, 5713, 11342, 5454, 3111, 23651, 5814, 5099, 15729, 25844, 17691, 5315, 3233, 592, 1105, 27, 16986, 1062, 54505, 69071, 10659, 16811, 14969, 16252, 4, 772, 1931, 22608, 23837, 4406, 1512, 4819, 23654, 31513, 11683, 1617, 32959, 6198, 14057, 18299, 16350, 371, 39741, 3146, 8204, 17345, 11473, 1242, 30743, 18748, 219284, 2810, 1316, 3908, 3747, 5760, 13083, 498, 17680, 1834, 1565, 6938, 423, 2960, 232, 151, 4532, 109, 20908, 52805, 33736, 4890, 23835, 501, 24061, 661, 9309, 13638, 9653, 2608, 13773, 3313, 13627, 17702, 13825, 2872, 18405, 13244, 22531, 13208, 19652, 49035, 12953, 2051, 3759, 25310, 2569, 1314, 2773, 4065, 12106, 9409, 2065, 3415, 5853, 8350, 1047, 38318, 726, 804, 19432, 120, 1080, 6677, 6840, 9315], "15": [1956, 79126, 155, 22478, 9456, 320, 3711, 7051, 3026, 4827, 6994, 17007, 1492, 1684, 20755, 525, 4220, 2131, 2694, 12587, 11757, 5641, 12420, 26363, 30846, 9209, 35142, 14255, 5402, 10371, 8689, 92375, 4129, 7322, 1722, 1809, 42516, 10705, 2240, 544, 59923, 2543, 3090, 6121, 3886, 10934, 3601, 351, 6057, 12628, 258, 20956, 15076, 54537, 11750, 19370, 562, 18620, 604, 53216, 100, 2365, 1395, 25438, 35145, 21410, 2364, 24047, 527, 35434, 2727, 3147, 457, 187472, 8857, 6622, 6734, 1058, 6440, 2085, 105, 20577, 34305, 238, 1478, 67269, 4319, 2275, 2043, 1221, 2193, 174, 1754, 0, 10312, 8423, 3064, 2861, 3199, 4496, 18211, 55363, 5901, 497, 14443, 35078, 28470, 6736, 67283, 8408, 713, 7734, 9171, 1769, 10069, 35145, 2450, 32138, 37581, 12620, 22851, 5055, 48123, 16826, 5563, 2622, 4579, 4684], "16": [9774, 6938, 2525, 2515, 2514, 10136, 9269, 4655, 3088, 12541, 5189, 1790, 366, 62853, 5396, 3493, 5092, 3347, 5224, 1674, 19844, 3898, 392, 2078, 33051, 1650, 3395, 19820, 16496, 290, 136351, 47882, 2990, 4691, 758, 7227, 14505, 6527, 1500, 4255, 863, 22419, 6889, 3240, 295, 1672, 21856, 3085, 15402, 71338, 7704, 1, 345, 17916, 1391, 9328, 8445, 545, 3074, 2019, 778, 14980, 6796, 3052, 1420, 1579, 20150, 163327, 7977, 44985, 750, 38, 19560, 2081, 343, 14336, 2244, 3086, 4972, 4027, 7207, 6097, 4607, 1038, 9152, 3047, 5785, 24407, 11411, 50654, 371, 2334, 9201, 18559, 1492, 14020, 898, 14307, 12291, 41612, 12153, 4471, 4643, 34805, 1946, 2672, 18080, 3122, 41204, 10782, 8001, 182988, 23733, 6261, 100980, 10926, 9874, 2200, 1501, 26715, 23595, 11791, 1141, 8715, 20696, 59884, 12188, 104], "17": [7696, 5668, 2106, 21078, 2631, 7904, 27130, 133, 5760, 1223, 12313, 9709, 47333, 2233, 4738, 6, 21326, 2638, 15039, 632, 48321, 1274, 34854, 22869, 3795, 9287, 13794, 5113, 737, 37107, 12787, 20627, 115149, 11483, 757, 15337, 4977, 9198, 809, 2998, 301, 5693, 1946, 30963, 99716, 13609, 3230, 1424, 4327, 70598, 28402, 4185, 11121, 175, 2449, 15568, 8165, 19229, 5055, 1586, 35, 12472, 3442, 4324, 1298, 3832, 2889, 25943, 3895, 6646, 95, 2117, 2682, 6855, 10156, 9886, 17656, 9404, 8896, 22984, 392, 3909, 1038, 1008, 713, 8335, 20491, 8096, 14924, 1205, 12033, 6363, 4127, 6829, 170603, 4095, 54069, 596, 11639, 12970, 2757, 1274, 3944, 1720, 113, 27976, 13896, 119896, 2558, 5722, 12706, 93376, 31208, 2726, 3530, 19588, 29, 19716, 522, 70633, 4104, 616, 157, 5867, 1829, 38187, 631, 3414], "18": [7282, 6897, 20319, 6611, 87502, 3485, 5313, 1694, 4455, 6707, 10542, 41566, 1522, 3127, 6436, 23048, 50116, 51645, 5267, 1173, 3743, 3338, 17, 15446, 63382, 52803, 24355, 212, 3932, 2483, 5195, 1474, 5068, 28545, 103, 1544, 6751, 1808, 4671, 278, 2407, 304, 1483, 1214, 183555, 116, 38230, 42993, 666, 72756, 142662, 1670, 2667, 119, 278, 25094, 2238, 206, 624, 23910, 6214, 99, 4519, 15870, 3501, 1567, 9235, 640, 1343, 1040, 14077, 22369, 2267, 4499, 12, 240, 4537, 18280, 1027, 3311, 45770, 2802, 31153, 6353, 1936, 5494, 6, 4589, 7159, 47323, 6243, 5232, 2478, 295, 10537, 1720, 1237, 22, 31073, 4153, 82469, 701, 18871, 2125, 2120, 752, 1008, 11658, 2495, 24853, 12263, 21926, 5922, 1044, 6861, 54275, 41505, 16094, 15895, 3578, 70627, 59, 1222, 1654, 15282, 2989, 6024, 16407], "19": [5657, 43214, 269, 14895, 21, 4756, 14311, 14609, 22, 6637, 75, 1689, 20080, 38, 641, 40, 5992, 272880, 1225, 14470, 42034, 31340, 152, 1164, 4302, 882, 3321, 3509, 43705, 2050, 10940, 25702, 53479, 12143, 5753, 16597, 529, 8000, 2009, 427, 42896, 8414, 2, 16046, 750, 20908, 7676, 5151, 1181, 214, 2776, 1397, 740, 13786, 5733, 21133, 2343, 1894, 3866, 4486, 7306, 2441, 2604, 16142, 33034, 3014, 7099, 3236, 12755, 9049, 5039, 7557, 59571, 1207, 2981, 13207, 10872, 2962, 50651, 2151, 13304, 67691, 34682, 35482, 3129, 5669, 599, 3324, 27099, 11473, 9627, 2952, 11096, 4111, 14787, 428, 10972, 3865, 3608, 1150, 13064, 8178, 7901, 65, 3599, 6764, 5338, 77, 16526, 8424, 39003, 1517, 3450, 2846, 1719, 409, 1314, 2277, 2831, 209645, 36693, 4991, 13175, 7316, 33403, 4035, 74845, 5666], "20": [11170, 2006, 30833, 306, 4901, 589, 2763, 25076, 34894, 4337, 2830, 137, 1102, 2672, 3996, 59819, 531, 17814, 1383, 2027, 158745, 495, 6098, 3497, 6720, 70250, 784, 345, 75577, 1, 4528, 130, 398, 10161, 15876, 2430, 32039, 4490, 285, 758, 7879, 13211, 10058, 5494, 4678, 5067, 38926, 9983, 10000, 4783, 12521, 13181, 92477, 1082, 1047, 14589, 3227, 7110, 7217, 1639, 2293, 267, 2389, 7169, 2422, 5489, 1900, 86047, 1147, 37547, 3256, 16643, 5901, 6477, 6767, 19732, 1829, 6347, 6703, 17955, 2469, 62742, 4984, 3192, 1859, 80880, 44259, 496, 154752, 1347, 13856, 1397, 4283, 5740, 54961, 2901, 3619, 43025, 15101, 980, 4014, 1741, 379, 17007, 18258, 875, 17040, 72891, 345, 1752, 2375, 3320, 7183, 4431, 3816, 3887, 1, 6778, 9834, 4246, 256, 2607, 32352, 12633, 1869, 1758, 1324, 42891], "21": [10791, 587, 7923, 39, 839, 6715, 1992, 263, 3820, 3, 3215, 53307, 17923, 30528, 121, 7419, 7522, 145, 16159, 27753, 10142, 21506, 24, 9177, 2424, 303, 2437, 32646, 2268, 218014, 1305, 1612, 51951, 2278, 90495, 8549, 6005, 1535, 2042, 322, 13531, 43369, 2294, 6068, 15137, 31926, 5231, 1591, 1935, 1946, 28, 4400, 13391, 19520, 1459, 14448, 11296, 64455, 2507, 1439, 3494, 1633, 10202, 13315, 9591, 17957, 31327, 10391, 28083, 1971, 8452, 2448, 79750, 3722, 9509, 4994, 6441, 1204, 5864, 3423, 13394, 20293, 104937, 19435, 5996, 7682, 4029, 1970, 16727, 48, 3138, 4570, 5565, 4198, 3618, 1876, 10881, 2506, 1674, 4250, 1246, 5230, 23407, 68, 64, 16317, 1247, 310, 6568, 57392, 12208, 56033, 9894, 8559, 9926, 10330, 12357, 136, 89, 1253, 4184, 3214, 189858, 9292, 1471, 23350, 8288, 3029], "22": [3030, 32063, 15189, 1293, 17339, 6635, 1482, 14770, 1601, 18, 10883, 2642, 3087, 199, 4848, 1900, 70085, 41877, 1970, 56, 8164, 716, 9786, 5610, 9696, 154474, 3765, 974, 182, 247, 8309, 13373, 1162, 15791, 3168, 43966, 18359, 3933, 299, 23, 27866, 6921, 737, 341, 106816, 9, 11521, 7680, 3897, 4342, 1630, 2522, 12000, 2810, 32151, 16378, 5973, 7184, 1928, 59584, 830, 25031, 5000, 2019, 3028, 2008, 39802, 34658, 4902, 12287, 7343, 29942, 12819, 2028, 103, 2267, 2747, 2706, 5035, 7352, 1527, 3741, 2449, 30399, 10356, 25149, 4666, 2162, 62400, 543, 4280, 2870, 24257, 110, 647, 1732, 0, 45707, 12670, 66134, 4559, 53465, 5516, 13381, 9352, 1753, 227659, 4697, 1436, 1109, 6709, 1439, 8174, 6235, 873, 2732, 8199, 1724, 692, 2681, 61722, 487, 7906, 42383, 14085, 34588, 5891, 5541], "23": [2446, 4030, 171, 3752, 10591, 644, 10368, 14765, 2057, 34048, 638, 23029, 3650, 1836, 393, 522, 2078, 13965, 96, 1790, 35813, 5285, 64519, 18034, 548, 3509, 5098, 10749, 1362, 5061, 13224, 62529, 70306, 3259, 17239, 13896, 9112, 12841, 845, 1261, 48095, 36, 8374, 3462, 185, 1458, 5001, 4440, 26760, 2466, 59, 2171, 10785, 747, 8557, 1627, 968, 16733, 1464, 2294, 1526, 2438, 9, 9416, 286, 6826, 336, 1849, 1278, 5487, 11338, 179, 16112, 158, 140725, 19916, 6907, 2267, 656, 12222, 239171, 16506, 158374, 3137, 921, 5931, 1269, 7461, 3212, 1456, 40969, 9216, 10518, 15622, 2554, 2460, 4250, 819, 68, 4479, 1112, 14365, 68162, 507, 36475, 17373, 22194, 5947, 800, 383, 30, 11669, 23183, 3764, 12090, 2, 2070, 26, 36494, 3280, 1975, 2595, 8, 169660, 7505, 39595, 1864, 1455], "24": [123, 28194, 884, 10685, 3837, 184, 12732, 16, 16556, 16152, 2207, 5735, 69526, 743, 19864, 2092, 1110, 5394, 27818, 1629, 22422, 29935, 81, 5799, 1146, 4556, 861, 23137, 226, 13349, 234501, 6349, 107403, 1499, 347, 6594, 19177, 21766, 3768, 8453, 4140, 4092, 193, 14984, 4383, 1086, 23028, 7996, 3093, 20, 46836, 377, 2808, 10054, 134, 1850, 2930, 997, 8986, 393, 225, 75196, 179, 3012, 412, 59, 512, 403, 1998, 20116, 126510, 53867, 65690, 10784, 1425, 1990, 4887, 2547, 33396, 5790, 45931, 1501, 19, 4720, 2485, 1240, 3782, 6215, 27122, 4472, 1000, 68, 2174, 10523, 2588, 8812, 29593, 69162, 7238, 4313, 4828, 4458, 831, 103984, 1338, 2480, 15774, 18939, 2619, 1148, 95, 7737, 17128, 1705, 2630, 2930, 21996, 22617, 1584, 11020, 33727, 953, 12934, 31598, 12891, 3119, 6584, 19115], "25": [174, 3459, 14708, 47546, 4370, 1771, 4862, 2014, 34893, 1691, 2266, 4281, 67951, 16987, 12024, 2384, 9787, 10721, 5254, 970, 7288, 2011, 55, 1636, 2988, 53757, 13378, 3, 9456, 28, 5487, 25642, 1618, 19782, 340, 2023, 3202, 10534, 4616, 341, 71894, 1993, 3, 13566, 34003, 4794, 1335, 16688, 6919, 82867, 1354, 15450, 4130, 3291, 20838, 851, 3877, 18320, 1039, 589, 66859, 1235, 33461, 4586, 332, 18420, 2214, 2100, 16385, 2794, 8079, 104137, 955, 100032, 37800, 33688, 8722, 3886, 2270, 2834, 11801, 3464, 7390, 5088, 5111, 49665, 9321, 12218, 134633, 574, 5996, 7419, 3315, 5426, 3933, 24235, 51329, 51255, 3, 3980, 427, 47039, 10657, 5315, 5450, 1269, 6054, 5800, 2066, 693, 5471, 1345, 3500, 4104, 16260, 2302, 7657, 1092, 2421, 4737, 4495, 6912, 138720, 4190, 8034, 5530, 19347, 17], "26": [5667, 1120, 35706, 12024, 1251, 827, 48, 12224, 6639, 2296, 859, 1460, 37, 15343, 34, 3311, 39153, 3206, 7560, 4031, 2872, 6768, 8917, 2941, 881, 44042, 11405, 112, 4589, 5477, 445, 15552, 167, 12372, 21363, 28789, 3992, 1354, 33, 1278, 3115, 40659, 66267, 10142, 5618, 154727, 12483, 4631, 695, 13138, 12765, 5341, 21713, 2256, 6871, 41026, 3460, 7224, 2062, 1719, 1189, 2, 3378, 1448, 16288, 9712, 3732, 181572, 13799, 12315, 1170, 26345, 7035, 683, 2574, 649, 4852, 15018, 3519, 1011, 7766, 2682, 40019, 6376, 25261, 247, 1091, 9529, 2816, 4703, 3414, 94815, 7574, 40955, 0, 1646, 29060, 7, 12944, 57048, 5991, 2752, 16892, 23032, 8303, 4033, 21892, 38047, 44664, 8872, 4039, 2917, 1278, 836, 573, 3787, 7766, 13164, 2953, 59191, 9902, 9497, 4281, 136254, 20433, 15583, 13520, 7195], "27": [2902, 6770, 2071, 2913, 2947, 11352, 1456, 4568, 1033, 1743, 17150, 79185, 442, 485, 106947, 393, 5716, 3555, 1890, 15874, 1344, 669, 60959, 55, 170, 9563, 1778, 4166, 5309, 16339, 2566, 30326, 37, 4039, 3560, 1383, 31364, 1522, 4880, 3609, 1144, 4262, 58657, 5253, 407, 765, 409, 16802, 10996, 2777, 448, 3798, 25, 6067, 23, 9239, 2747, 2387, 2419, 95253, 7962, 10979, 15358, 1142, 1540, 6526, 403, 21400, 3408, 20285, 212545, 1723, 23577, 17224, 971, 6615, 78, 6915, 39500, 6420, 65221, 33952, 1958, 727, 4057, 3986, 1567, 10373, 6978, 2689, 65, 5250, 737, 33205, 2021, 1570, 66402, 38, 10802, 5406, 27849, 12529, 12504, 51, 4987, 2695, 1048, 7634, 1405, 1093, 13649, 6790, 15257, 3047, 37820, 19899, 33519, 22892, 18548, 96386, 64, 2360, 778, 94184, 31224, 44214, 16149, 16890], "28": [2877, 6089, 574, 4520, 1563, 93953, 28996, 438, 11621, 7852, 2062, 5084, 2022, 11591, 1681, 2575, 7393, 1415, 36893, 11761, 2802, 2107, 4760, 12057, 4, 13968, 2990, 3411, 5983, 6116, 69462, 91668, 11537, 2584, 4644, 10369, 137, 26686, 3862, 7704, 32817, 18129, 529, 104257, 22272, 31873, 90, 37451, 2349, 517, 2414, 5406, 307, 15470, 464, 11270, 267, 61, 5832, 879, 1348, 74, 21937, 17550, 145, 4, 17625, 614, 8146, 1777, 4112, 18141, 2333, 2958, 927, 3189, 5294, 540, 4463, 97142, 34021, 6926, 14810, 93618, 15366, 4384, 19514, 12700, 18700, 19999, 11535, 1601, 5425, 5148, 83226, 211, 4099, 777, 755, 7296, 5046, 11339, 7160, 416, 13520, 11114, 80487, 1494, 12468, 18064, 3974, 10334, 10853, 2316, 3977, 45287, 1150, 64, 15340, 9973, 40320, 36730, 13385, 744, 4214, 77047, 61245, 992], "29": [9653, 3726, 742, 1421, 1118, 2104, 5255, 4927, 466, 2969, 88828, 5351, 10493, 1267, 37959, 6796, 1677, 22538, 2987, 5163, 15703, 4757, 29982, 718, 1172, 1127, 1442, 12907, 22545, 1656, 184, 8408, 1842, 41482, 10370, 1192, 15205, 14192, 15487, 83, 7088, 202, 2418, 7783, 7, 7913, 3742, 4386, 14793, 3008, 2748, 18918, 7011, 48748, 15804, 3891, 6049, 28779, 88, 147334, 12712, 3047, 40798, 200370, 2615, 991, 27242, 38934, 6412, 7953, 4103, 18923, 11333, 31327, 4477, 2771, 4361, 53601, 336, 9067, 266, 4996, 6706, 6789, 5302, 850, 5810, 11455, 24, 35444, 39782, 33939, 67, 5513, 1517, 25421, 893, 10094, 4142, 777, 47441, 746, 10767, 6849, 57447, 1682, 32965, 470, 13306, 7286, 116938, 32709, 916, 5901, 2413, 8384, 72, 1260, 1547, 338, 1957, 1184, 5393, 39, 74155, 9933, 10090, 3996], "30": [549, 19090, 5738, 17634, 15788, 41273, 15407, 1268, 5995, 5054, 175, 24703, 1357, 6290, 31394, 55732, 1926, 4386, 3952, 336, 2719, 6928, 2318, 1110, 42619, 5041, 696, 3169, 10076, 2579, 12866, 20751, 5894, 29, 1576, 6818, 74041, 5652, 5786, 33201, 4762, 806, 2754, 980, 1573, 9899, 26418, 3635, 17997, 7091, 2908, 19998, 5911, 154561, 17797, 43109, 12913, 70021, 3162, 573, 23, 5138, 19448, 18834, 2254, 10658, 475, 23020, 5449, 9912, 392, 3617, 3931, 6096, 14657, 16584, 36992, 27216, 400, 94, 916, 2131, 1859, 6427, 21684, 2261, 6287, 28663, 15303, 4790, 5191, 5526, 15641, 53337, 41527, 10172, 19309, 428, 2812, 36377, 22268, 44696, 6258, 5845, 1533, 1882, 4091, 1230, 35524, 16371, 33260, 73509, 1078, 16593, 476, 19705, 6608, 17789, 9990, 17715, 3094, 7048, 7258, 11475, 4981, 109280, 7193, 2583], "31": [8519, 4215, 9482, 2, 224, 41, 6911, 9310, 11424, 11290, 6429, 2880, 116, 10259, 2419, 6502, 43434, 11790, 2070, 1319, 23746, 12315, 7167, 2365, 3531, 7713, 16372, 465, 5194, 178, 1114, 1092, 15171, 3484, 12863, 9025, 14979, 41737, 14859, 27805, 126, 3164, 1255, 6658, 14243, 36449, 485, 16325, 6928, 1512, 502, 18462, 39579, 70117, 9462, 6670, 1025, 66156, 15462, 27301, 51450, 10031, 7318, 53016, 1888, 1515, 3199, 3, 34592, 2976, 11828, 3252, 2414, 9953, 42725, 11526, 9332, 17875, 11131, 20471, 62675, 220, 9156, 2050, 6591, 1866, 32483, 2200, 906, 19337, 4979, 26766, 26199, 12992, 7721, 4421, 3280, 1470, 11556, 7262, 508, 8019, 24345, 8334, 16595, 108013, 5791, 88530, 88682, 9835, 23535, 41672, 2378, 31367, 2985, 520, 10086, 63934, 1431, 2121, 5640, 1299, 280, 6881, 12478, 32690, 23460, 222], "32": [3001, 3553, 9638, 45079, 5188, 23601, 104024, 2352, 8103, 20176, 22331, 5147, 2663, 6856, 1333, 5231, 6854, 4536, 7773, 27690, 5784, 5219, 9177, 5072, 118, 8350, 55, 8637, 2727, 11447, 5308, 1506, 1238, 5316, 21055, 5031, 947, 9149, 833, 478, 5586, 51364, 4546, 3941, 888, 16306, 162151, 4087, 8356, 21499, 22, 1353, 2172, 7376, 47458, 1628, 8019, 12110, 6069, 1558, 1401, 1245, 80614, 13055, 26577, 59384, 5779, 39934, 5684, 19017, 4581, 18140, 59383, 1304, 55812, 4945, 54750, 2735, 490, 13182, 10322, 7470, 15619, 12468, 934, 1459, 9738, 18868, 7447, 42321, 15888, 10942, 24050, 6002, 54540, 1094, 11653, 4795, 28958, 3949, 738, 1122, 73871, 49248, 97, 34970, 6949, 9985, 12253, 276, 4958, 14837, 24664, 7202, 181, 15777, 9118, 24508, 3689, 250, 16740, 14, 5696, 1482, 8053, 8355, 170, 3151], "33": [6781, 3825, 33268, 4447, 6373, 9484, 356, 1409, 38840, 7609, 3520, 22548, 3919, 34005, 6352, 583, 1786, 35930, 2380, 1109, 27248, 34933, 454, 16502, 101720, 10307, 11577, 15954, 4134, 7605, 47601, 8750, 66, 51211, 2, 23835, 30685, 16560, 8204, 5382, 5479, 967, 435, 3545, 4742, 139321, 7677, 32426, 45220, 1132, 3138, 13444, 18931, 1123, 2890, 1271, 16077, 8386, 756, 10900, 14443, 3952, 4921, 19426, 379, 22985, 1566, 2003, 7849, 1891, 74347, 7198, 10483, 26308, 12001, 4706, 55247, 4012, 50652, 11907, 12106, 48, 162421, 17160, 1327, 7938, 5201, 10440, 720, 0, 6201, 9220, 941, 7425, 2653, 9359, 2508, 8301, 2920, 13857, 24, 2420, 2132, 6465, 5780, 14294, 24586, 6062, 10066, 5647, 73589, 29730, 532, 1699, 1016, 25, 17587, 17900, 14581, 136, 15833, 1894, 9106, 8, 8813, 8629, 22439, 819], "34": [542, 987, 10956, 6201, 15404, 70669, 39233, 2616, 13903, 7324, 6782, 21621, 6910, 6185, 6868, 1917, 4262, 157, 2825, 31414, 26007, 11227, 7698, 8582, 604, 7648, 12055, 0, 83350, 512, 158, 5466, 1705, 117, 308, 36656, 23916, 9673, 305, 7146, 61, 36374, 3048, 29887, 2072, 20150, 4246, 6151, 4339, 2060, 5493, 15469, 1526, 92161, 24195, 7786, 3379, 9435, 4670, 195, 1286, 10385, 34202, 7219, 31082, 844, 352, 5069, 2454, 4810, 64884, 38721, 6590, 30395, 64849, 1028, 6944, 133, 27180, 148, 37050, 13118, 14677, 751, 17807, 51606, 10778, 29442, 7407, 38206, 6055, 9791, 2890, 58632, 8737, 22427, 1490, 774, 1851, 2733, 204, 21617, 35099, 33007, 4585, 1920, 14474, 95669, 10574, 3966, 1896, 8446, 7369, 47, 1332, 126, 4494, 671, 65217, 4075, 8476, 10648, 1405, 13745, 53410, 20446, 13492, 135], "35": [106, 1066, 18099, 44365, 8195, 18912, 44034, 22674, 3802, 4799, 2701, 1186, 2511, 2002, 32885, 6466, 71332, 17346, 30537, 4417, 86704, 13891, 87, 11, 23459, 1838, 15717, 12611, 18278, 596, 10554, 6046, 215, 2959, 7078, 22290, 3838, 24486, 36458, 11597, 2684, 3241, 1622, 26250, 16741, 44, 2051, 5465, 2431, 5263, 1119, 1234, 4494, 24987, 5459, 5435, 3575, 2072, 32026, 571, 16, 9120, 62859, 1041, 23254, 8420, 2025, 416, 3136, 2830, 11571, 1407, 537, 26471, 3331, 643, 1444, 7028, 22, 8452, 36695, 14861, 7178, 10671, 12410, 2577, 1399, 42758, 6980, 43158, 31780, 19755, 12757, 29215, 3415, 3011, 8514, 340, 4383, 38, 3870, 4000, 1820, 57024, 47118, 2030, 296, 67157, 5413, 10018, 8191, 34173, 1753, 13134, 9374, 121185, 1971, 6518, 41655, 8247, 7544, 7122, 60828, 8020, 49684, 28359, 673, 27941], "36": [574, 4868, 5786, 20659, 1841, 5283, 9701, 1409, 98289, 18588, 7447, 14073, 9094, 32375, 40071, 16481, 20, 4744, 4312, 32263, 7823, 2727, 7403, 5221, 14614, 2247, 48448, 9256, 814, 2478, 3058, 29818, 19135, 321, 14702, 30429, 2605, 13, 5629, 35549, 3620, 4608, 5453, 1867, 6228, 1361, 42098, 237, 42174, 18025, 7472, 4545, 29263, 34860, 12547, 158, 93, 1889, 35695, 20741, 7159, 1415, 217, 2375, 661, 3678, 7821, 2931, 35801, 5694, 661, 2963, 1658, 5388, 19493, 6830, 3343, 17294, 2949, 2570, 41319, 8244, 52061, 4020, 6729, 76291, 74092, 20749, 1993, 1126, 8837, 24721, 10775, 9676, 538, 3314, 742, 21666, 3232, 533, 267, 6044, 5573, 57592, 1177, 5370, 36362, 35860, 3204, 9687, 3766, 3198, 20746, 16024, 1932, 733, 13174, 17243, 37128, 3068, 41372, 69788, 4788, 62773, 1475, 64143, 5804, 10996], "37": [32349, 28323, 19268, 12652, 10918, 3037, 4567, 117, 8278, 4783, 830, 4821, 4970, 3100, 26312, 10221, 9879, 14231, 11263, 5330, 2207, 8230, 1127, 3632, 2680, 42637, 104, 68757, 33272, 47830, 13319, 6358, 3068, 69, 22610, 3339, 5378, 20530, 24925, 11695, 5738, 8528, 21779, 22615, 11317, 3905, 6403, 897, 8997, 18026, 11410, 4348, 61270, 63999, 23799, 10347, 18607, 26071, 13989, 3482, 26070, 4589, 2608, 10007, 5097, 43916, 3029, 201, 8105, 625, 10138, 2303, 3586, 24616, 5180, 19938, 498, 47871, 10617, 12014, 8215, 2050, 13719, 5878, 36754, 1927, 8804, 4948, 7188, 3397, 24264, 7977, 992, 2533, 50321, 31863, 6450, 1461, 5460, 5651, 1745, 276, 43827, 17230, 7565, 585, 16068, 103, 12771, 3862, 7336, 355, 74823, 49648, 2221, 94599, 45, 6327, 20163, 138078, 5540, 8259, 1713, 9191, 2401, 3507, 2029, 278], "38": [12606, 31, 7746, 350, 2553, 5779, 630, 8606, 440, 19480, 40165, 5735, 5900, 5518, 21377, 288, 33068, 9526, 9080, 22063, 17344, 3606, 81623, 22403, 11386, 3368, 1158, 1211, 73, 16246, 12, 5309, 4521, 12574, 4422, 21122, 24008, 31741, 39313, 3286, 4806, 43107, 29079, 27833, 1652, 1845, 7350, 837, 1359, 7938, 7623, 5038, 13546, 53300, 26655, 43835, 1006, 7690, 2874, 2752, 747, 19137, 7386, 86, 15448, 442, 5686, 55310, 342, 12826, 7576, 10504, 3413, 131486, 22599, 487, 2058, 99358, 12222, 22662, 1551, 2571, 1066, 36291, 6096, 46335, 6985, 810, 11293, 8534, 4190, 13006, 11466, 16950, 3675, 2241, 1231, 2355, 7742, 11555, 3704, 14245, 452, 230, 7777, 1880, 34181, 9158, 18028, 28862, 890, 6992, 1433, 3943, 161117, 45204, 2696, 7007, 4193, 61990, 4447, 258, 1899, 113, 8599, 1682, 4771, 16687], "39": [7819, 4552, 9293, 28103, 8569, 4559, 344, 2820, 7099, 1812, 12503, 7639, 37281, 4871, 3097, 111, 4940, 21184, 12103, 2975, 4701, 11690, 1923, 1300, 4, 7044, 10869, 4829, 1707, 33015, 12296, 2262, 7923, 10899, 632, 10155, 3193, 11242, 34556, 32164, 42347, 3278, 86, 42, 1657, 4079, 52692, 3895, 6750, 13772, 121101, 1903, 1172, 4735, 8962, 28163, 3115, 14641, 18751, 855, 69, 40763, 6521, 42199, 11903, 4274, 109869, 251, 3343, 3218, 3975, 15442, 4486, 20042, 39462, 19222, 33497, 73239, 663, 13156, 3499, 1698, 27281, 155182, 9450, 1659, 8581, 4827, 21907, 6244, 18028, 34897, 4125, 7423, 9574, 1495, 5293, 13785, 3898, 2752, 5151, 14169, 12455, 148261, 3041, 12538, 124, 7766, 115, 2591, 6238, 26610, 180, 2568, 26418, 4304, 2191, 1704, 21150, 35865, 4435, 12192, 10620, 10593, 5694, 383, 1477, 1849], "40": [7320, 3029, 3029, 4162, 90, 9000, 3944, 9140, 7577, 12633, 1451, 2686, 4657, 15822, 26550, 1492, 14565, 846, 118, 12411, 3866, 71550, 5474, 335, 5579, 2134, 18094, 8860, 31116, 1101, 7855, 8257, 62892, 2819, 11359, 4072, 88448, 19850, 17395, 18939, 1057, 2274, 16102, 5193, 4967, 16605, 33, 2123, 19100, 15319, 10717, 85022, 7784, 79, 8872, 5836, 1290, 2787, 23895, 3378, 23192, 7185, 9442, 2915, 9877, 347, 22609, 15221, 3113, 9594, 7910, 3072, 219842, 2250, 6403, 7489, 2628, 6000, 16986, 69864, 8010, 31560, 8048, 21631, 11794, 16531, 93, 8667, 17923, 806, 4753, 14933, 10413, 208, 5759, 6404, 23846, 4901, 117637, 33445, 3798, 6540, 10059, 43, 53422, 1592, 62239, 27469, 2324, 9881, 11476, 13188, 997, 5103, 14351, 1144, 579, 196, 26255, 2774, 8540, 27451, 5482, 14030, 4548, 5693, 3541, 4982], "41": [8156, 26596, 4517, 65626, 379, 27503, 2316, 4664, 9477, 36289, 7817, 9022, 377, 4072, 8319, 22763, 7758, 67123, 7089, 80124, 1011, 5011, 4528, 4279, 4101, 6111, 140, 4240, 9295, 12769, 29803, 69, 950, 18708, 4994, 11370, 3218, 4949, 1940, 923, 9962, 2467, 5125, 15722, 3243, 4443, 37350, 168, 440, 65107, 11402, 4834, 17229, 18497, 18983, 4227, 18746, 112776, 13602, 3526, 27818, 10354, 5233, 7611, 11710, 7287, 22150, 13138, 494, 4209, 22891, 90327, 9144, 46662, 1119, 6360, 12534, 10709, 2171, 1326, 12971, 21860, 11427, 507, 7359, 29137, 19911, 9182, 35751, 602, 30992, 22379, 26244, 5360, 7978, 2648, 1260, 1662, 1802, 49533, 4065, 2312, 10811, 35651, 6974, 10603, 10191, 4823, 7957, 1223, 135, 34700, 3155, 2163, 75272, 203, 29150, 124, 36742, 9157, 6253, 7255, 35068, 15730, 5651, 5515, 1020, 13988], "42": [2636, 5544, 3930, 4282, 3764, 9171, 26840, 10150, 29254, 5251, 21595, 26431, 2769, 19772, 2042, 1103, 6615, 19478, 6522, 7130, 1025, 4442, 4739, 47074, 2776, 50770, 5042, 9460, 3194, 3697, 3173, 2895, 19035, 1833, 10205, 6260, 14654, 1055, 1033, 38484, 1743, 9210, 2502, 29316, 7320, 6186, 50819, 10781, 9, 158692, 71177, 11091, 91846, 928, 4215, 15130, 20826, 2126, 1242, 10089, 4430, 182, 74, 10234, 2946, 25932, 6708, 21433, 24659, 245, 46612, 15775, 7194, 3051, 11963, 14860, 2870, 4645, 2959, 4877, 4624, 65, 1210, 2477, 390, 1594, 11540, 1782, 157029, 4920, 5586, 7220, 1909, 3680, 7627, 479, 37712, 22065, 23931, 14016, 6914, 15862, 8337, 14376, 4747, 849, 8156, 36671, 548, 5273, 26917, 416, 74988, 11133, 10188, 19254, 5945, 12220, 1869, 12169, 22591, 14068, 6878, 10638, 6470, 32381, 4939, 37273], "43": [278, 1748, 2939, 9757, 288, 30440, 25597, 67046, 66664, 16815, 9984, 9882, 42585, 74212, 5487, 12599, 91096, 2707, 15955, 43442, 15691, 207, 4526, 26807, 10110, 41743, 23435, 7645, 11150, 4706, 1258, 13908, 3169, 822, 1596, 6440, 8094, 8540, 19508, 23286, 9157, 640, 26504, 13822, 16123, 9048, 9003, 5605, 4452, 21704, 3788, 18086, 6915, 4111, 19589, 28, 426, 23980, 4267, 899, 7703, 9164, 15439, 9993, 2787, 6724, 8791, 2876, 877, 2575, 8687, 114963, 5727, 5, 16692, 3593, 203, 10609, 41463, 16563, 208, 17432, 14845, 13396, 23596, 5196, 363, 8407, 1264, 20688, 1775, 5165, 15484, 1845, 5911, 3268, 7732, 14778, 7610, 4094, 3363, 16099, 133053, 17789, 1991, 18653, 1816, 22106, 4490, 6708, 49802, 10024, 15444, 2640, 1377, 59933, 6185, 9983, 1534, 5506, 3377, 2175, 12848, 9614, 4778, 5664, 2597, 35599], "44": [4308, 5061, 20160, 9561, 13669, 14091, 1543, 5138, 8239, 1408, 8293, 9755, 8230, 2194, 4557, 5404, 1613, 19967, 31371, 7477, 3629, 4783, 5728, 29390, 45219, 8854, 18033, 1342, 1366, 11472, 2842, 19117, 18847, 137, 51814, 7171, 4759, 8667, 12858, 66091, 13412, 1593, 13669, 1560, 11088, 2587, 9589, 7003, 723, 934, 28567, 1799, 13341, 27, 1936, 52011, 81, 6, 18039, 35490, 3868, 15268, 32862, 1270, 5811, 40641, 14, 6652, 140795, 8428, 13831, 1038, 8995, 16254, 40909, 60048, 1439, 12988, 36289, 9845, 5748, 7752, 43707, 52437, 11647, 10575, 1870, 7934, 20013, 37357, 8595, 1130, 1733, 50163, 14035, 21738, 2631, 49372, 1173, 5793, 2360, 157, 7783, 13348, 2113, 3608, 1725, 53901, 43410, 1273, 8293, 3634, 13613, 14607, 484, 12437, 7496, 29131, 8846, 39733, 34686, 23208, 390, 2534, 5954, 731, 9722, 8510], "45": [43760, 9474, 13786, 21035, 5451, 559, 31883, 546, 4531, 79, 1230, 13426, 226, 3424, 3232, 6144, 21606, 35253, 3891, 53881, 10192, 6181, 4012, 5348, 12797, 31794, 3217, 1633, 7136, 55192, 1573, 5260, 7138, 58157, 16865, 109766, 14268, 2909, 11584, 23579, 3810, 7983, 9064, 9942, 10291, 26379, 517, 72581, 46470, 6288, 3510, 23869, 2890, 5302, 26000, 972, 2572, 25848, 3993, 16613, 886, 6280, 64175, 3236, 23096, 31865, 128923, 1072, 3435, 11125, 5561, 20156, 16228, 3790, 8797, 1185, 27684, 22281, 3681, 11686, 200, 2619, 17700, 7885, 6275, 10179, 10295, 5581, 5920, 6218, 763, 11274, 1957, 3826, 404, 2472, 1267, 6607, 22243, 4407, 20080, 28599, 7426, 12560, 17936, 5521, 49360, 18260, 41683, 10808, 10745, 2703, 42, 7551, 56188, 3651, 1025, 7422, 882, 18, 18563, 17699, 6257, 461, 17079, 13660, 5342, 12281], "46": [4133, 6250, 100, 48870, 32617, 1387, 20478, 5387, 2342, 21798, 7790, 9826, 15976, 61557, 23846, 21133, 11745, 16454, 54364, 8791, 14459, 58093, 8437, 8746, 1088, 8388, 8254, 1277, 18398, 43165, 295, 9887, 43223, 105, 5210, 4232, 27, 2250, 10480, 14818, 6069, 18337, 6308, 23646, 4950, 5308, 1301, 2830, 96, 39, 14204, 60, 17345, 2762, 35168, 4098, 36115, 37095, 7875, 12784, 743, 761, 20842, 3536, 63062, 18774, 11759, 64, 593, 13120, 23247, 21299, 20172, 25264, 5805, 11549, 7377, 2276, 10664, 275, 26492, 8246, 11703, 426, 16989, 34862, 29920, 2641, 32937, 517, 4981, 8272, 24023, 2329, 1122, 15915, 5352, 3183, 1754, 7266, 228, 9293, 1529, 29283, 7445, 2207, 5051, 12878, 8356, 8859, 23672, 166757, 1262, 9956, 3408, 63655, 27663, 32101, 10979, 10809, 10569, 8167, 2205, 7280, 10039, 7079, 13162, 17878], "47": [101420, 10658, 11884, 3607, 27982, 13018, 6624, 9771, 24672, 14210, 11575, 21141, 14553, 10759, 25832, 12111, 1590, 7272, 1813, 1557, 7641, 4147, 8108, 3434, 13768, 17941, 1117, 12055, 30, 633, 2775, 956, 23493, 148972, 8129, 11416, 1343, 22995, 5642, 22543, 11082, 2613, 138, 3581, 78060, 55638, 10196, 171, 5799, 24460, 11990, 47, 11260, 10541, 3164, 20765, 5968, 6549, 6470, 38883, 1570, 470, 16923, 63, 14114, 4053, 5898, 30527, 105888, 12683, 7651, 12111, 44862, 18801, 621, 2816, 998, 1536, 1288, 29856, 8403, 6664, 5395, 3942, 2775, 24491, 8423, 18437, 10588, 15206, 994, 16703, 3071, 8890, 22436, 5839, 4093, 3821, 1562, 9360, 59, 3048, 11797, 164, 6520, 4201, 15371, 27316, 151355, 9814, 6589, 18707, 14563, 41781, 2515, 20748, 104, 4897, 20201, 3625, 11627, 30392, 17003, 573, 2445, 1547, 6570, 32], "48": [181, 65845, 3449, 1544, 25504, 6139, 4875, 4324, 4613, 897, 284, 3884, 31650, 13963, 20972, 14203, 32145, 6449, 394, 5575, 1686, 9821, 38968, 175329, 1492, 28282, 4332, 16744, 5605, 6843, 135, 140, 3809, 22488, 19287, 43669, 14935, 10501, 28755, 10431, 22664, 13217, 11214, 225, 2726, 186, 7091, 121048, 7318, 18951, 13754, 3007, 232, 2267, 24170, 10256, 13502, 15709, 3087, 23414, 804, 74547, 24156, 5652, 55436, 4446, 418, 6474, 1741, 9355, 1911, 5018, 17166, 13205, 2689, 875, 14600, 4457, 5154, 1078, 6832, 2283, 11797, 11071, 1627, 11444, 106, 746, 2318, 7238, 13822, 5049, 146582, 12641, 867, 3520, 1412, 18085, 14753, 2539, 28099, 9186, 7188, 4318, 1070, 1553, 2871, 12609, 27665, 55436, 49738, 137, 35772, 13560, 20871, 908, 5347, 799, 4911, 1617, 7285, 5432, 36589, 4762, 3890, 7747, 977, 5517], "49": [41649, 12308, 750, 30928, 4629, 11196, 847, 6806, 37731, 6600, 5001, 26222, 3980, 23537, 796, 13265, 5576, 3590, 4771, 19420, 41836, 12267, 8853, 207, 1060, 1272, 14575, 7211, 6674, 12234, 10246, 20830, 1131, 6448, 7175, 26212, 10377, 7394, 10902, 851, 14734, 4872, 1118, 14656, 6344, 1778, 13640, 3804, 9162, 22674, 41946, 16630, 415, 3150, 12228, 9587, 550, 12543, 39911, 10342, 7684, 2013, 1696, 78, 3711, 8739, 27675, 9278, 3508, 3696, 6834, 91719, 5947, 6436, 14780, 4595, 2087, 4522, 31680, 49911, 367, 6869, 2792, 68000, 67649, 66, 86813, 7224, 19914, 27, 2174, 20427, 20014, 8047, 4531, 16275, 34516, 42647, 10604, 44, 14535, 4778, 4179, 4717, 20986, 10068, 5884, 2853, 20100, 1042, 47560, 1105, 33447, 13127, 8311, 48579, 8985, 5821, 1401, 3486, 156026, 2422, 4127, 245, 11805, 10303, 10615, 14411], "50": [363, 5103, 106352, 18333, 1093, 13120, 6125, 1212, 1995, 16812, 98, 3671, 14921, 8648, 8411, 57746, 32021, 7368, 160760, 4393, 1990, 7475, 18594, 16389, 8981, 43076, 305, 18379, 6842, 29198, 8232, 94, 5279, 5519, 10907, 646, 1335, 2086, 28153, 3092, 3406, 10136, 56360, 1776, 19669, 5867, 33659, 23446, 9043, 5951, 7011, 87153, 17, 1089, 5108, 12624, 134, 6268, 4518, 14574, 2144, 30160, 2973, 40580, 8405, 287, 5996, 34535, 11005, 1631, 6118, 34304, 1175, 579, 8153, 8092, 2070, 32561, 7871, 8919, 6621, 9920, 403, 17027, 3618, 31436, 12656, 64085, 3915, 922, 772, 11709, 34673, 1383, 1732, 3545, 28126, 4549, 7638, 7215, 1320, 1602, 12057, 79599, 9375, 2084, 18010, 76, 19596, 20363, 53133, 8556, 3028, 7326, 232, 1916, 14238, 3273, 7506, 10261, 460, 4317, 79064, 8684, 7825, 9430, 871, 7317], "51": [13790, 3562, 880, 4192, 8100, 4563, 9747, 8202, 15623, 42162, 257, 21047, 7404, 9980, 2416, 18600, 957, 24010, 58641, 7537, 19, 4089, 2582, 27775, 9377, 20924, 5809, 16030, 991, 32305, 36122, 138, 27068, 273, 5694, 6659, 59479, 4443, 12167, 3941, 27449, 11292, 4534, 3469, 8805, 12322, 1379, 9649, 6577, 4432, 18463, 2466, 14074, 18275, 12431, 7757, 141445, 694, 8268, 11183, 1143, 24524, 5501, 9056, 10155, 3375, 8935, 154, 636, 437, 22500, 396, 7536, 22405, 4716, 2041, 3188, 3316, 12008, 35544, 5849, 77, 22069, 2763, 121538, 1130, 32289, 77145, 1384, 22695, 152106, 51777, 11124, 2705, 6463, 17343, 5526, 51537, 34430, 126, 5258, 1365, 8534, 75, 6478, 10660, 22364, 2909, 5729, 2788, 13943, 527, 5545, 22685, 3926, 22379, 12527, 9194, 119, 6047, 437, 7961, 11597, 2098, 11260, 14444, 11601, 1337], "52": [5896, 1419, 5838, 7195, 1184, 11255, 47125, 1, 8295, 3820, 296, 7747, 4125, 143008, 5010, 26568, 29642, 33492, 1256, 15703, 5539, 162398, 5409, 10911, 96518, 502, 11206, 6128, 36, 556, 874, 5407, 3149, 4535, 38699, 20307, 151468, 5760, 3850, 1640, 266, 268, 5494, 5359, 12925, 53798, 3845, 2892, 34301, 1337, 4797, 11544, 10438, 22141, 13500, 6240, 14679, 6505, 141, 18083, 59358, 10967, 25235, 4256, 7656, 10501, 2954, 6994, 23539, 5386, 384, 666, 9370, 9893, 14918, 1216, 6768, 11108, 1665, 4745, 12857, 7711, 5601, 15864, 6301, 2562, 7928, 23872, 8225, 4201, 367, 2648, 11071, 6107, 6861, 4174, 8600, 1675, 9189, 147, 36581, 11535, 8825, 8817, 10270, 1860, 83, 402, 12019, 39185, 6958, 8627, 44775, 14072, 4815, 8795, 10202, 21595, 12573, 15850, 19131, 63727, 19013, 11337, 5255, 1583, 126, 77], "53": [18556, 6540, 38373, 24531, 2914, 2701, 185, 661, 8149, 19370, 3400, 41440, 11563, 2699, 147, 8572, 49685, 14766, 54486, 7533, 5558, 3190, 82788, 31536, 2732, 374, 60438, 109001, 298, 9176, 30657, 209, 4261, 9087, 18496, 30128, 34299, 6032, 23079, 801, 76981, 1139, 4550, 4292, 4684, 62, 21704, 11582, 9961, 28674, 27872, 7818, 889, 25079, 3593, 4312, 8549, 2229, 13453, 326, 43194, 24524, 6515, 62213, 14818, 11332, 3477, 5964, 22222, 2716, 2518, 4080, 4157, 487, 1643, 296, 3346, 4054, 8, 5892, 30104, 100, 407, 28325, 2910, 8888, 45386, 17953, 7181, 67749, 1592, 1276, 9704, 429, 10143, 451, 15590, 5266, 13096, 2156, 5634, 8346, 32150, 4020, 33092, 2478, 1233, 4973, 5128, 37474, 114, 2052, 16309, 27852, 23698, 5598, 73308, 13766, 230, 5818, 609, 7392, 10402, 3886, 13354, 5652, 7806, 3252], "54": [51023, 3978, 6913, 51668, 6389, 6090, 22071, 4633, 7740, 9737, 5990, 1343, 4760, 5758, 64407, 15247, 5688, 4925, 900, 16723, 11690, 561, 14609, 12786, 3217, 3236, 10190, 3355, 6551, 22194, 46320, 1521, 188, 5191, 15711, 246, 13914, 71670, 25430, 34629, 600, 11286, 2210, 33220, 31282, 1932, 14679, 4887, 232, 13972, 24506, 216, 4390, 75736, 7114, 26252, 9761, 14127, 14707, 39842, 5536, 7739, 7438, 8071, 1378, 12026, 44, 7899, 10391, 46548, 42625, 17759, 5182, 19594, 15168, 43030, 5296, 64682, 25906, 48863, 2869, 887, 772, 16199, 2033, 10112, 6765, 50, 3331, 2890, 3502, 5568, 10257, 13500, 4818, 387, 1705, 1733, 61664, 35023, 16465, 3092, 4917, 24364, 2345, 647, 2256, 7120, 23536, 371, 3888, 5393, 4873, 19443, 547, 32927, 40446, 2851, 7, 36916, 3002, 3, 3722, 5904, 86510, 25430, 6570, 11020], "55": [26946, 722, 64487, 6520, 21012, 20577, 9499, 16031, 13469, 3492, 1684, 12357, 2656, 813, 9918, 4201, 4557, 14732, 2842, 4216, 72023, 77, 508, 650, 30069, 7724, 14743, 1743, 3499, 717, 5008, 18041, 4445, 7476, 14993, 34333, 11300, 104069, 6313, 10020, 76, 5524, 7154, 3086, 5490, 9191, 1932, 22397, 30261, 46546, 2197, 4124, 11664, 17545, 57498, 1046, 17579, 102, 2517, 1739, 101787, 17891, 3005, 9664, 9616, 13744, 9941, 6762, 970, 5606, 689, 17863, 1050, 7360, 7494, 15035, 8547, 1951, 10478, 4373, 39200, 1122, 9147, 11027, 28004, 4944, 2587, 18887, 17882, 78617, 47993, 6829, 9338, 47150, 92943, 955, 12211, 7625, 3193, 24851, 17239, 700, 629, 10195, 9559, 990, 6896, 2859, 4354, 9208, 15, 11337, 1539, 6101, 15499, 20, 2345, 14193, 3961, 6819, 15827, 4648, 21061, 93872, 91, 121, 10984, 58405], "56": [85770, 9091, 19427, 14311, 1770, 378, 5424, 29357, 1246, 11375, 21797, 11, 2730, 375, 4808, 9783, 45779, 44994, 1262, 3401, 8422, 5635, 3369, 5831, 3662, 19930, 63985, 14915, 65, 817, 16664, 5587, 864, 2, 5406, 4486, 29623, 342, 5463, 29960, 22111, 11246, 7854, 2223, 7451, 18667, 2861, 371, 6157, 1792, 329, 4852, 19188, 1070, 4954, 5652, 5333, 15517, 13004, 6291, 8703, 71240, 5064, 12533, 350, 16032, 17126, 81187, 25848, 1439, 14092, 116972, 6525, 70815, 15678, 1095, 13478, 1587, 317, 1840, 10548, 15310, 1371, 2794, 21908, 16426, 11391, 15846, 327, 25466, 13684, 7574, 142, 9720, 16597, 32535, 12105, 2415, 22939, 9624, 566, 22907, 358, 3722, 13905, 4443, 16022, 11378, 6767, 13963, 67471, 19845, 37236, 473, 11675, 23618, 8693, 21201, 847, 12319, 54994, 10159, 1180, 17648, 16873, 10262, 45543, 4197], "57": [23849, 23860, 18781, 10902, 41576, 6241, 3054, 1644, 6488, 18410, 24402, 3911, 34421, 3638, 8108, 15354, 7884, 1599, 6223, 7228, 34648, 9792, 11218, 1312, 19331, 6653, 65, 12657, 5681, 5757, 13994, 50, 560, 41018, 6648, 32051, 181, 16041, 1255, 137114, 8224, 9110, 450, 49969, 65083, 10548, 3509, 865, 9688, 6225, 4914, 8176, 25622, 16335, 234, 1136, 1728, 4414, 51340, 9306, 1785, 5041, 6736, 3067, 26908, 4762, 7085, 2946, 9043, 24156, 4194, 2736, 13, 36876, 1951, 10394, 7824, 1744, 25625, 12806, 11946, 1110, 295, 1545, 42149, 157, 2132, 22365, 4666, 27071, 11503, 2415, 12828, 2722, 25137, 2244, 7541, 4889, 21131, 2507, 1631, 6247, 1055, 1718, 5280, 8374, 4229, 9442, 967, 339, 8517, 137841, 3184, 51575, 252, 59656, 26681, 109010, 3756, 3729, 5601, 2538, 2902, 53491, 1416, 20665, 27152, 14180], "58": [3539, 7555, 11907, 33940, 11816, 8560, 14418, 24087, 13557, 4239, 3139, 112167, 7389, 19107, 253, 1973, 10323, 6711, 8567, 3023, 16776, 5949, 5924, 7535, 32033, 3427, 6691, 1449, 1796, 4069, 7543, 10858, 925, 8834, 13351, 19989, 2035, 11701, 5610, 2500, 16841, 61594, 4675, 9157, 8538, 29168, 25664, 34336, 4411, 56245, 13722, 24666, 27954, 7894, 279, 532, 26250, 9347, 99661, 2714, 10794, 1479, 40711, 108, 636, 66, 17987, 6355, 84000, 2618, 27047, 0, 85686, 10743, 7990, 6404, 10964, 1273, 1315, 44670, 6560, 2317, 1208, 474, 26549, 7155, 30820, 2494, 4744, 1686, 4492, 3363, 17091, 3153, 100254, 1848, 14526, 7205, 751, 812, 4050, 3996, 23162, 110, 21964, 12758, 34347, 91, 16098, 748, 6465, 9299, 9706, 9122, 31634, 5215, 17687, 10541, 7544, 5439, 339, 1357, 12032, 64788, 14785, 7892, 9992, 5526], "59": [19174, 33738, 13975, 2984, 8682, 35339, 5006, 38176, 8990, 27352, 519, 1752, 4656, 47377, 52439, 3221, 23552, 4462, 41662, 1067, 30994, 300, 48927, 5186, 776, 2896, 5624, 4426, 5601, 14323, 8664, 22209, 26391, 8477, 3860, 3592, 4746, 1234, 4186, 398, 11087, 43378, 97717, 3349, 31572, 4488, 61933, 3288, 21144, 5296, 5923, 4259, 9360, 23475, 61757, 1883, 9261, 13726, 3162, 7820, 42, 2489, 73864, 5505, 17124, 57722, 23234, 6091, 160, 639, 1981, 14396, 24639, 442, 589, 18360, 11371, 9616, 1171, 3662, 4604, 3449, 9196, 410, 1778, 235, 6885, 5068, 23439, 263, 20174, 13468, 13767, 33466, 4897, 30, 13862, 2335, 1959, 31662, 6888, 7289, 5605, 9415, 4008, 35833, 3476, 8723, 27637, 24370, 15046, 18672, 38858, 788, 21535, 46, 27430, 2907, 3439, 68211, 412, 459, 17356, 2339, 59898, 1, 1484, 6968]}} \ No newline at end of file diff --git a/scripts/minimax_m3_nvfp4/debug_expert_capture.py b/scripts/minimax_m3_nvfp4/debug_expert_capture.py new file mode 100644 index 000000000..097eec7f3 --- /dev/null +++ b/scripts/minimax_m3_nvfp4/debug_expert_capture.py @@ -0,0 +1,93 @@ +"""Diagnose why layer-3 expert capture saw ~0 tokens during the pilot. + +Loads the model the same way the pilot does, attaches plain forward hooks to all +layer-3 expert Linears, runs ONE 8k calibration sample through the model, and +reports per-expert token counts plus which experts implementation is active. +""" + +import json +import random + +import torch +import torch.nn as nn +from transformers import AutoTokenizer + +MODEL = "/data/sgambhira/models/Minimax-M3-0602" +SNAP = ( + "/data/sgambhira/hf-home/hub/datasets--togethercomputer--m3-quant-eval-runs/" + "snapshots/27f2a493e10922e878840d27924aabd0bbf5d745" +) + +from gptqmodel import BACKEND, GPTQModel, QuantizeConfig # noqa: E402 + +qcfg = QuantizeConfig( + bits=4, + group_size=16, + format="nvfp4", + offload_to_disk_path="/data/sgambhira/gptqmodel-offload/debug-capture", +) +model = GPTQModel.load(MODEL, quantize_config=qcfg, backend=BACKEND.TORCH) +print("DBG model loaded", flush=True) + +cfg = model.model.config +text_cfg = getattr(cfg, "text_config", cfg) +print("DBG top config _experts_implementation:", getattr(cfg, "_experts_implementation", ""), flush=True) +print("DBG text config _experts_implementation:", getattr(text_cfg, "_experts_implementation", ""), flush=True) + +# Materialize only what a 4-layer forward needs (weights are lazy/meta otherwise), +# using the same API the looper uses. +lm = model.model.model.language_model +device = torch.device("cpu") +for sub in [lm.embed_tokens, lm.rotary_emb, lm.norm, *lm.layers[:4]]: + model.shell_module_materialize(sub, device) +print("DBG materialized embed/rotary/norm/layers0-3", flush=True) +# truncate the stack so the forward stops after layer 3 +lm.layers = nn.ModuleList(list(lm.layers[:4])) + +layer3 = lm.layers[3] +experts_mod = layer3.mlp.experts +print("DBG experts module class:", type(experts_mod).__name__, flush=True) +print("DBG experts forward qualname:", type(experts_mod).forward.__qualname__, flush=True) +has_3d = any(p.ndim == 3 for p in experts_mod.parameters(recurse=False)) +print("DBG experts has 3D params:", has_3d, flush=True) + +counts = {} + +def make_hook(name): + def hook(mod, inp, out): + counts[name] = counts.get(name, 0) + inp[0].shape[0] + return hook + +hooks = [] +for e in range(128): + child = getattr(experts_mod, str(e), None) + if child is None: + continue + for proj in ("gate_proj", "up_proj", "down_proj"): + m = getattr(child, proj, None) + if isinstance(m, nn.Linear): + hooks.append(m.register_forward_hook(make_hook(f"{e}.{proj}"))) +print("DBG hooks attached:", len(hooks), flush=True) + +# one 8k calibration sample +tokenizer = AutoTokenizer.from_pretrained(MODEL) +rows = [json.loads(l) for l in open(f"{SNAP}/hle-mxfp8-0602-text2158/responses.jsonl")] +r = random.Random(42).sample(rows, 64)[0] +messages = [ + {"role": "user", "content": r["question"]}, + {"role": "assistant", "content": r["response"], "reasoning_content": r["reasoning"]}, +] +text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False) +ids = tokenizer(text, add_special_tokens=False, return_tensors="pt")["input_ids"][:, :8192] +print("DBG sample tokens:", ids.shape[1], flush=True) + +with torch.inference_mode(): + lm(input_ids=ids) + +gate_counts = {k: v for k, v in counts.items() if k.endswith("gate_proj")} +total = sum(gate_counts.values()) +nonzero = len(gate_counts) +print(f"DBG experts hit: {nonzero}/128 | total routed tokens (gate): {total:,} | expected ~{ids.shape[1]*4:,}", flush=True) +top = sorted(gate_counts.items(), key=lambda kv: -kv[1])[:8] +print("DBG top experts:", top, flush=True) +print("DBG DONE", flush=True) diff --git a/scripts/minimax_m3_nvfp4/emit_act_amax.py b/scripts/minimax_m3_nvfp4/emit_act_amax.py new file mode 100644 index 000000000..9e108bf07 --- /dev/null +++ b/scripts/minimax_m3_nvfp4/emit_act_amax.py @@ -0,0 +1,276 @@ +"""Step 0 for w4a4: emit the per-layer activation amax sidecar. + +NVFP4 activations carry a *static* per-tensor global scale (calibrated, frozen at +serving) plus dynamic per-16-block E4M3 scales. Only the block scales can be +derived while GPTQ accumulates its Hessian, so the global scale has to come from +an earlier pass -- this one. Consumed via `ACT_AMAX=` in run.sh, which turns it +into `global_scale = amax / 6 / 448` (see docs/w4a4_hessian.md). + +Scales are keyed by *input site*, not by module, because modules sharing an input +share a scale: + + layer{N}.moe_input -> input to routed-expert gate_proj/up_proj (w1/w3) + layer{N}.w2_input -> input to routed-expert down_proj (w2), post-SwiGLU + +A plain max is taken over abs, matching tore-quant's `save_amax_sidecar` +(no percentile or MSE clipping), and reduced across every expert of a layer -- +the peer-max sync `fixup_moe_expert_amax` performs. + +Streams the model window-by-window through one GPU exactly as +p1b_expert_coverage.py does (materialize -> forward all samples -> free), since +the checkpoint is far larger than any single device. +""" + +import json +import os +import random +import time + +import torch +import torch.nn as nn +from transformers import AutoTokenizer + +MODEL = os.environ.get("MODEL", "/data/huggingface/Minimax-M3-0602") +OUT_JSON = os.environ.get( + "OUT_JSON", "/data/sgambhira/GPTQModel/scripts/minimax_m3_nvfp4/amax_per_layer.json" +) +SNAP = os.environ.get( + "SNAP", + "/data/sgambhira/hf-home/hub/datasets--togethercomputer--m3-quant-eval-runs/" + "snapshots/27f2a493e10922e878840d27924aabd0bbf5d745", +) +OFFLOAD = os.environ.get("OFFLOAD", "/data/sgambhira/gptqmodel-offload/amax") +SEED = int(os.environ.get("SEED", "42")) +NUM_SAMPLES = int(os.environ.get("NUM_SAMPLES", "64")) +MAX_TOKENS = int(os.environ.get("MAX_TOKENS", "8192")) +WINDOW = int(os.environ.get("WINDOW", "5")) +DEVICE = torch.device(os.environ.get("DEVICE", "cuda:0")) +CKPT = os.environ.get("CKPT", "/data/sgambhira/gptqmodel-offload/amax_resume.pt") + +t0 = time.time() + + +def log(msg: str) -> None: + print(f"[amax +{time.time() - t0:7.1f}s] {msg}", flush=True) + + +from gptqmodel import BACKEND, GPTQModel, QuantizeConfig # noqa: E402 + +qcfg = QuantizeConfig(bits=4, group_size=16, format="nvfp4", offload_to_disk_path=OFFLOAD) +model = GPTQModel.load(MODEL, quantize_config=qcfg, backend=BACKEND.TORCH) +log("model loaded") + +lm = model.model.model.language_model +all_layers = list(lm.layers) +num_layers = len(all_layers) + +# --- calibration inputs (identical to the coverage / GPTQ calibration set) --- +tokenizer = AutoTokenizer.from_pretrained(MODEL) +rows = [json.loads(line) for line in open(f"{SNAP}/hle-mxfp8-0602-text2158/responses.jsonl")] +sample = random.Random(SEED).sample(rows, NUM_SAMPLES) +inputs = [] +for r in sample: + messages = [ + {"role": "user", "content": r["question"]}, + {"role": "assistant", "content": r["response"], "reasoning_content": r["reasoning"]}, + ] + text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False) + ids = tokenizer(text, add_special_tokens=False, return_tensors="pt")["input_ids"][:, :MAX_TOKENS] + inputs.append(ids) +total_tokens = sum(i.shape[1] for i in inputs) +log(f"calibration: {len(inputs)} samples, {total_tokens:,} tokens") + +# --- amax capture ----------------------------------------------------------- +amax: dict[str, float] = {} + + +def _update(key: str, tensor: torch.Tensor) -> None: + if not torch.is_tensor(tensor) or tensor.numel() == 0: + return + v = tensor.detach().abs().max().to(torch.float32).item() + if v > amax.get(key, 0.0): + amax[key] = v + + +def make_moe_input_hook(layer_idx: int): + # Pre-hook on the experts module: args[0] is the hidden state entering w1/w3. + def hook(module, args, kwargs): + if args and torch.is_tensor(args[0]): + _update(f"layer{layer_idx}.moe_input", args[0]) + return hook + + +def make_apply_gate_patch(layer_idx: int, original): + """Capture w2_input by wrapping the fused experts' `_apply_gate`. + + MiniMaxM3VLExperts stores experts as 3D nn.Parameters -- `gate_up_proj` and + `down_proj` -- and runs them with F.linear inside a Python loop, so there is + no down_proj submodule to hook. Its forward is: + + current = self._apply_gate(F.linear(hidden[idx], gate_up_proj[e])) + current = F.linear(current, down_proj[e]) * top_k_weights[...] + + so `_apply_gate`'s return value *is* the down_proj input -- the post-SwiGLU + activation we need. Wrapping it captures that exactly, with no + reimplementation of the SwiGLU to drift out of sync. It fires once per hit + expert, so the max is already reduced across the layer's experts (the + peer-max sync). + """ + + def patched(gate_up: torch.Tensor) -> torch.Tensor: + out = original(gate_up) + _update(f"layer{layer_idx}.w2_input", out) + return out + + return patched + + +def make_w2_input_hook(layer_idx: int): + # Pre-hook on a per-expert down_proj: its input IS the post-SwiGLU activation. + # Fires once per expert, so the max is already reduced across the layer. + def hook(module, args): + if args and torch.is_tensor(args[0]): + _update(f"layer{layer_idx}.w2_input", args[0]) + return hook + + +# The experts tree differs by loader, so support both rather than betting on one: +# * GPTQModel.load splits the checkpoint's fused 3D tensors into 128 per-expert +# nn.Linear modules (gate_proj / up_proj / down_proj) -- hook down_proj. +# * Stock transformers keeps MiniMaxM3VLExperts fused as 3D nn.Parameters with +# no submodules -- wrap `_apply_gate`, whose return value is the down_proj +# input. +hook_handles = [] +patched_experts = [] +moe_layers = [] +n_down = 0 +for li, layer in enumerate(all_layers): + experts = getattr(getattr(layer, "mlp", None), "experts", None) + if not isinstance(experts, nn.Module): + continue + moe_layers.append(li) + hook_handles.append( + experts.register_forward_pre_hook(make_moe_input_hook(li), with_kwargs=True) + ) + + downs = [m for name, m in experts.named_modules() + if name.endswith("down_proj") and isinstance(m, nn.Linear)] + if downs: + for m in downs: + hook_handles.append(m.register_forward_pre_hook(make_w2_input_hook(li))) + n_down += len(downs) + continue + + original_apply_gate = getattr(experts, "_apply_gate", None) + if original_apply_gate is None: + raise SystemExit( + f"layer {li}: {type(experts).__name__} exposes neither down_proj Linears nor " + "`_apply_gate`; w2_input cannot be captured. Adapt the capture to this " + "architecture rather than emitting a sidecar missing every down_proj scale." + ) + # Instance attribute shadows the bound method, so `self._apply_gate(x)` hits it. + experts._apply_gate = make_apply_gate_patch(li, original_apply_gate) + patched_experts.append(experts) + +log(f"hooked {len(moe_layers)} MoE layers ({moe_layers[:3]}...{moe_layers[-2:]}): " + f"moe_input via experts pre-hook; w2_input via {n_down} down_proj hooks " + f"+ {len(patched_experts)} _apply_gate wrappers") + +# --- windowed streaming forward (same structure as p1b_expert_coverage.py) --- +model.shell_module_materialize(lm.rotary_emb, DEVICE) +model.shell_module_materialize(lm.norm, DEVICE) + +start_layer = 0 +if os.path.exists(CKPT): + state = torch.load(CKPT, map_location="cpu", weights_only=False) + hidden = state["hidden"] + amax.update(state["amax"]) + start_layer = state["next_layer"] + log(f"resumed from checkpoint at layer {start_layer}") +else: + embed = model.shell_module_materialize(lm.embed_tokens, DEVICE) + hidden = [] + with torch.inference_mode(): + for ids in inputs: + hidden.append(embed(ids.to(DEVICE)).to("cpu")) + log("embedded all samples") + +orig_layers = lm.layers + +for w_start in range(start_layer, num_layers, WINDOW): + w_layers = all_layers[w_start : w_start + WINDOW] + lm.layers = orig_layers + for layer in w_layers: + model.shell_module_materialize(layer, DEVICE) + lm.layers = nn.ModuleList(w_layers) + + window_out = {} + + def tail_hook(module, args, output): + window_out["h"] = (output[0] if isinstance(output, tuple) else output).detach() + + tail = w_layers[-1].register_forward_hook(tail_hook) + + with torch.inference_mode(): + for si in range(len(hidden)): + h = hidden[si].to(DEVICE) + pos = torch.arange(h.shape[1], device=DEVICE).unsqueeze(0) + lm(inputs_embeds=h, position_ids=pos, use_cache=False) + hidden[si] = window_out["h"].to("cpu") + + tail.remove() + for layer in w_layers: + layer.to_empty(device="meta") + torch.cuda.empty_cache() + log(f"window done: layers {w_start}-{w_start + len(w_layers) - 1}") + + next_layer = w_start + len(w_layers) + if next_layer < num_layers: + tmp = CKPT + ".tmp" + torch.save({"hidden": hidden, "amax": amax, "next_layer": next_layer}, tmp) + os.replace(tmp, CKPT) + +lm.layers = orig_layers +for h in hook_handles: + h.remove() +for experts in patched_experts: + # Drop the instance attribute so the class's bound method is live again. + del experts._apply_gate + +# --- emit ------------------------------------------------------------------- +missing = [ + f"layer{li}.{site}" + for li in moe_layers + for site in ("moe_input", "w2_input") + if f"layer{li}.{site}" not in amax +] +if missing: + raise SystemExit(f"missing amax entries, refusing to write a partial sidecar: {missing[:8]}") + +payload = { + "amax": {k: amax[k] for k in sorted(amax, key=lambda s: (int(s.split(".")[0][5:]), s))}, + "meta": { + "model": MODEL, + "samples": NUM_SAMPLES, + "max_tokens": MAX_TOKENS, + "seed": SEED, + "total_tokens": total_tokens, + "method": "plain max over |x|, reduced across experts (peer-max)", + "sites": {"moe_input": "gate_proj/up_proj input", "w2_input": "down_proj input"}, + }, +} +os.makedirs(os.path.dirname(OUT_JSON) or ".", exist_ok=True) +with open(OUT_JSON, "w") as f: + json.dump(payload, f, indent=1) +log(f"wrote {OUT_JSON} ({len(payload['amax'])} entries)") + +if os.path.exists(CKPT): + os.remove(CKPT) + +vals = list(payload["amax"].values()) +log(f"amax range: min={min(vals):.4g} max={max(vals):.4g}") +log(f"implied global_scale range: {min(vals) / 6 / 448:.4g} .. {max(vals) / 6 / 448:.4g}") +for li in moe_layers[:3] + moe_layers[-2:]: + log(f" layer{li}: moe_input={amax[f'layer{li}.moe_input']:.4g} " + f"w2_input={amax[f'layer{li}.w2_input']:.4g}") +log("DONE") diff --git a/scripts/minimax_m3_nvfp4/generate_serving_config.py b/scripts/minimax_m3_nvfp4/generate_serving_config.py new file mode 100644 index 000000000..4fb6d4376 --- /dev/null +++ b/scripts/minimax_m3_nvfp4/generate_serving_config.py @@ -0,0 +1,127 @@ +"""Generate the serving quantization_config for w4a4 variant checkpoints. + +The spliced/variant checkpoints inherit the SOURCE config.json verbatim, whose +quantization_config says mxfp8 -- wrong for NVFP4(+mxfp8-swap) experts and +unservable. The serving stack expects the ModelOpt layout the reference +checkpoints carry: config.json quantization_config with quant_algo=NVFP4 and an +exclude_modules/ignore list naming everything NOT NVFP4, plus an +hf_quant_config.json sidecar with the same list. + +Rather than hand-deriving the lists, the base (non-swapped exclusions: attention, +index projections, shared experts, dense MLP, routers, vision, ...) and the +per-swapped-layer template are EXTRACTED from the midsplit2 reference and +re-instantiated for this variant's layer set. Generator output is verified by +exact set-equality against BOTH reference checkpoints before anything is written. + + OUT=/data/huggingface/Minimax-M3-0602-NVFP4-GPTQ-w4a4-inscale1-skip14mid \ + MXFP8_LAYERS="22-28,32-38" \ + python scripts/minimax_m3_nvfp4/generate_serving_config.py +""" + +import json +import os +import re + +MID = "/data/huggingface/MiniMax-M3-NVFP4-mxfp8skip14-midsplit2" +ALT = "/data/huggingface/MiniMax-M3-NVFP4-alt3x6-0708" +MID_LAYERS = set(range(22, 29)) | set(range(32, 39)) +ALT_LAYERS = {3,4,5,12,13,14,21,22,23,30,31,32,39,40,41,48,49,50,57,58,59} + +OUT = os.environ["OUT"] +MXFP8_LAYERS = os.environ.get("MXFP8_LAYERS", "").strip() + +def parse_layers(spec): + out = set() + for part in spec.split(","): + part = part.strip() + if not part: continue + if "-" in part: + a, b = part.split("-"); out |= set(range(int(a), int(b) + 1)) + else: out.add(int(part)) + return out + +variant_layers = parse_layers(MXFP8_LAYERS) + +def load_exclude(root): + return set(json.load(open(f"{root}/config.json"))["quantization_config"]["exclude_modules"]) + +mid_ex, alt_ex = load_exclude(MID), load_exclude(ALT) + +LAYER_RE = re.compile(r"^language_model\.model\.layers\.(\d+)(?=$|\.)") +def layer_of(entry): + m = LAYER_RE.match(entry) + return int(m.group(1)) if m else None + +def swap_entries(exclude, layers): + """Entries belonging to swapped layers: expert modules, the `layers.N.*` + wildcard, and the bare `layers.N` entry the references also carry.""" + out = set() + for k in exclude: + li = layer_of(k) + if li in layers and ( + ".block_sparse_moe.experts." in k + or k.endswith(f".layers.{li}.*") + or k.endswith(f".layers.{li}") + ): + out.add(k) + return out + +# base = midsplit2's excludes minus its swapped-layer entries +base = mid_ex - swap_entries(mid_ex, MID_LAYERS) + +# per-layer template from midsplit2's layer 22, parameterized on the layer number +template = sorted(re.sub(r"\.layers\.22(?=$|\.)", ".layers.{L}", k) + for k in swap_entries(mid_ex, {22})) +assert template, "no layer-22 swap entries found in midsplit2" + +def build(layers): + ex = set(base) + for li in sorted(layers): + ex |= {t.replace("{L}", str(li)) for t in template} + return ex + +# --- verify the generator reproduces BOTH references exactly ----------------- +for name, ref_ex, ref_layers in (("midsplit2", mid_ex, MID_LAYERS), ("alt3x6", alt_ex, ALT_LAYERS)): + got = build(ref_layers) + if got != ref_ex: + miss, extra = sorted(ref_ex - got)[:5], sorted(got - ref_ex)[:5] + raise SystemExit(f"generator does not reproduce {name}: missing {len(ref_ex-got)} {miss}, " + f"extra {len(got-ref_ex)} {extra}") + print(f"[cfg] generator reproduces {name} exactly ({len(ref_ex)} entries)") + +# --- build this variant's config --------------------------------------------- +exclude = sorted(build(variant_layers)) +print(f"[cfg] variant layers {sorted(variant_layers) or 'none'} -> {len(exclude)} exclude entries") + +ref_cfg = json.load(open(f"{MID}/config.json"))["quantization_config"] +qc = { + "quant_algo": ref_cfg["quant_algo"], + "kv_cache_quant_algo": ref_cfg["kv_cache_quant_algo"], + "group_size": ref_cfg["group_size"], + "exclude_modules": exclude, + "quant_method": ref_cfg["quant_method"], + "ignore": exclude, +} + +cfg_path = f"{OUT}/config.json" +cfg = json.load(open(cfg_path)) +cfg["quantization_config"] = qc +with open(cfg_path, "w") as f: + json.dump(cfg, f, indent=2) +print(f"[cfg] wrote quantization_config -> {cfg_path}") + +hq_ref = json.load(open(f"{MID}/hf_quant_config.json")) +hq = { + "producer": {"name": hq_ref["producer"]["name"], "version": "minimax-m3-w4a4-gptq-nvfp4"}, + "quant_method": hq_ref["quant_method"], + "quantization": { + "quant_algo": hq_ref["quantization"]["quant_algo"], + "kv_cache_quant_algo": hq_ref["quantization"]["kv_cache_quant_algo"], + "group_size": hq_ref["quantization"]["group_size"], + "exclude_modules": exclude, + }, +} +with open(f"{OUT}/hf_quant_config.json", "w") as f: + json.dump(hq, f, indent=1) +print(f"[cfg] wrote {OUT}/hf_quant_config.json") +print("[cfg] DONE") diff --git a/scripts/minimax_m3_nvfp4/p03_load_dryrun.py b/scripts/minimax_m3_nvfp4/p03_load_dryrun.py new file mode 100644 index 000000000..b30ba5e52 --- /dev/null +++ b/scripts/minimax_m3_nvfp4/p03_load_dryrun.py @@ -0,0 +1,116 @@ +"""Phase 0.3: GPTQModel.load dry run for Minimax-M3-0602 NVFP4 plan. + +Verifies: (a) load routes to the quantize path and MXFP8 decodes to bf16, +(b) runtime module naming (block_sparse_moe/w1 vs mlp/gate_proj), (c) the +experts-only dynamic rules select the expected 16,512 modules. + +Run: .venv/bin/python scripts/minimax_m3_nvfp4/p03_load_dryrun.py +See: docs/minimax_m3_nvfp4_quantization_plan.md +""" + +import json + +import torch +import torch.nn as nn + +MODEL = "/data/sgambhira/models/Minimax-M3-0602" +SKIP_LAYERS = "22|23|24|25|26|27|28|32|33|34|35|36|37|38" + +from gptqmodel import BACKEND, GPTQModel, QuantizeConfig # noqa: E402 + +qcfg = QuantizeConfig( + bits=4, + group_size=16, + format="nvfp4", + dynamic={ + r"-:.*self_attn\..*": {}, + r"-:.*shared_experts\..*": {}, + r"-:.*layers\.(0|1|2)\.mlp\..*": {}, + rf"-:.*layers\.({SKIP_LAYERS})\..*": {}, + }, + moe_vram_strategy="balanced", +) +print("P03 qcfg class:", type(qcfg).__name__, "format:", qcfg.format, "device:", qcfg.device, flush=True) + +# Non-quantized models always load CPU-side (lazy); GPUs are used per-layer +# during quantization. Do not pass device= when qcfg.device is set/auto. +model = GPTQModel.load(MODEL, quantize_config=qcfg, backend=BACKEND.TORCH) +print("P03 loaded model class:", type(model.model).__name__, flush=True) + +# (b) runtime module names for layer 3 +names_l3 = [ + (n, type(m).__name__) + for n, m in model.model.named_modules() + if ".layers.3." in n and isinstance(m, nn.Linear) +] +print("P03 layer-3 linear count:", len(names_l3), flush=True) +for n, t in names_l3[:12]: + print("P03 l3 module:", n, t, flush=True) +expert_markers = { + "mlp.experts": sum(1 for n, _ in names_l3 if "mlp.experts." in n), + "block_sparse_moe.experts": sum(1 for n, _ in names_l3 if "block_sparse_moe.experts." in n), +} +print("P03 expert naming counts (layer 3):", expert_markers, flush=True) + +# (a) dequant spot check on one expert projection +expert_name = next((n for n, _ in names_l3 if ".experts.0." in n), None) +print("P03 spot-check module:", expert_name, flush=True) +mod = dict(model.model.named_modules())[expert_name] +w = mod.weight +print("P03 module weight dtype/shape/device:", w.dtype, tuple(w.shape), w.device, flush=True) + +if w.device.type != "meta": + from safetensors import safe_open + + index = json.load(open(f"{MODEL}/model.safetensors.index.json"))["weight_map"] + cands = [ + k + for k in index + if ".layers.3." in k + and ".experts.0." in k + and k.endswith(".weight") + and "scale" not in k + and "shared" not in k + ] + print("P03 checkpoint candidates:", cands[:3], flush=True) + matched = False + for ck in cands: + with safe_open(f"{MODEL}/{index[ck]}", framework="pt", device="cpu") as f: + fp8 = f.get_tensor(ck) + scale_u8 = f.get_tensor(ck.replace(".weight", ".weight_scale_inv")) + if tuple(fp8.shape) != tuple(w.shape): + continue + scale = torch.pow(2.0, scale_u8.to(torch.float32) - 127.0) + manual = fp8.to(torch.float32) * scale.repeat_interleave(32, dim=1) + diff = (w.detach().to(torch.float32).cpu() - manual).abs().max().item() + print(f"P03 dequant max abs diff vs manual mxfp8 ({ck}):", diff, flush=True) + matched = True + break + if not matched: + print("P03 no shape-matching checkpoint candidate found", flush=True) + +# (c) dynamic selection counts over all language-model layers +included, excluded = [], [] +for n, m in model.model.named_modules(): + if not isinstance(m, nn.Linear) or ".layers." not in n: + continue + if "visual" in n or "vision" in n: + continue + if qcfg.dynamic_get(layer_name=n) is False: + excluded.append(n) + else: + included.append(n) +experts_included = [n for n in included if ".experts." in n and "shared" not in n] +non_expert_included = [n for n in included if n not in set(experts_included)] +print("P03 total linears:", len(included) + len(excluded), flush=True) +print("P03 included expert linears:", len(experts_included), "(expect 16512)", flush=True) +print("P03 included NON-expert linears:", len(non_expert_included), flush=True) +for n in non_expert_included[:10]: + print("P03 non-expert included (should be empty or router-only):", n, flush=True) +skip_hits = [ + n + for n in included + if any(f".layers.{x}." in n for x in "22 23 24 25 26 27 28 32 33 34 35 36 37 38".split()) +] +print("P03 included modules in skip layers (expect 0):", len(skip_hits), flush=True) +print("P03 DONE", flush=True) diff --git a/scripts/minimax_m3_nvfp4/p1_pilot_layer3.py b/scripts/minimax_m3_nvfp4/p1_pilot_layer3.py new file mode 100644 index 000000000..1087e9abf --- /dev/null +++ b/scripts/minimax_m3_nvfp4/p1_pilot_layer3.py @@ -0,0 +1,123 @@ +"""Phase 1: single-layer (layer 3) GPTQ->NVFP4 pilot for Minimax-M3-0602. + +Quantizes ONLY layer-3 routed experts (384 modules) with the full 64-trajectory +HLE calibration set. Gates checked here before the full run: + - per-expert calibration coverage / RTN-fallback rate under native routing + - loss sanity (also validates the MXFP8->bf16 decode numerically) + - saved tensor layout (uint8 packed weight + float8_e4m3fn weight_scale) + - what non-quantized modules are written as (plan open-decision #1) + +Run: .venv/bin/python scripts/minimax_m3_nvfp4/p1_pilot_layer3.py +See: docs/minimax_m3_nvfp4_quantization_plan.md +""" + +import json +import random + +import torch +from transformers import AutoTokenizer + +MODEL = "/data/sgambhira/models/Minimax-M3-0602" +OUT = "/data/sgambhira/models/Minimax-M3-0602-NVFP4-pilot-l3" +SNAP = ( + "/data/sgambhira/hf-home/hub/datasets--togethercomputer--m3-quant-eval-runs/" + "snapshots/27f2a493e10922e878840d27924aabd0bbf5d745" +) +SEED = 42 +NUM_SAMPLES = 64 +# 80GB H100 headroom (user-confirmed): keep the FIRST 8k tokens per sample. The sdpa +# sparse-attention path materializes dense per-head masks (64 x n^2 x 2B) plus a bool +# intermediate — ~9GB + ~4GB at 8k vs 32GB + 17GB at 16k, which OOMed. +MAX_TOKENS_PER_SAMPLE = 8192 + +from gptqmodel import BACKEND, GPTQModel, QuantizeConfig # noqa: E402 + +qcfg = QuantizeConfig( + bits=4, + group_size=16, + format="nvfp4", + dynamic={ + r"-:.*self_attn\..*": {}, + r"-:.*shared_experts\..*": {}, + r"-:.*layers\.(0|1|2)\.mlp\..*": {}, + # pilot: skip every layer except 3 + r"-:.*layers\.([4-9]|[1-5][0-9])\..*": {}, + }, + moe_vram_strategy="balanced", + # /tmp is on the ~full 124G root disk; module offload must go to /data (15T free). + offload_to_disk_path="/data/sgambhira/gptqmodel-offload/pilot-l3", +) +print("P1 qcfg:", type(qcfg).__name__, qcfg.format, flush=True) + +# --- calibration: 64 seed-42 HLE trajectories, pre-tokenized (matches P0.2) --- +tokenizer = AutoTokenizer.from_pretrained(MODEL) +rows = [json.loads(l) for l in open(f"{SNAP}/hle-mxfp8-0602-text2158/responses.jsonl")] +sample = random.Random(SEED).sample(rows, NUM_SAMPLES) + +calibration = [] +total_tokens = 0 +for r in sample: + messages = [ + {"role": "user", "content": r["question"]}, + {"role": "assistant", "content": r["response"], "reasoning_content": r["reasoning"]}, + ] + text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False) + enc = tokenizer(text, add_special_tokens=False, return_tensors="pt") + input_ids = enc["input_ids"][:, :MAX_TOKENS_PER_SAMPLE] + attention_mask = enc["attention_mask"][:, :MAX_TOKENS_PER_SAMPLE] + calibration.append({"input_ids": input_ids, "attention_mask": attention_mask}) + total_tokens += input_ids.shape[1] +truncated = sum(1 for c in calibration if c["input_ids"].shape[1] == MAX_TOKENS_PER_SAMPLE) +print( + f"P1 calibration: {len(calibration)} samples, {total_tokens:,} tokens " + f"({truncated} truncated to first {MAX_TOKENS_PER_SAMPLE})", + flush=True, +) + +model = GPTQModel.load(MODEL, quantize_config=qcfg, backend=BACKEND.TORCH) +print("P1 model loaded", flush=True) + +model.quantize(calibration, batch_size=1, backend=BACKEND.TORCH) +print("P1 quantize complete", flush=True) + +model.save(OUT) +print("P1 saved to", OUT, flush=True) + +# --- post-checks on the saved checkpoint --- +from safetensors import safe_open # noqa: E402 + +index = json.load(open(f"{OUT}/model.safetensors.index.json"))["weight_map"] + +def _tensor(name): + with safe_open(f"{OUT}/{index[name]}", framework="pt", device="cpu") as f: + return f.get_tensor(name) + +l3_keys = sorted(k for k in index if ".layers.3.mlp.experts.0." in k) +print("P1 saved layer-3 expert-0 tensors:", l3_keys, flush=True) +for k in l3_keys: + t = _tensor(k) + print(f"P1 saved {k}: {t.dtype} {tuple(t.shape)}", flush=True) + +# non-quantized passthrough: layer-4 expert + layer-3 attention +for probe in [ + next((k for k in index if ".layers.4.mlp.experts.0." in k and k.endswith(".weight")), None), + next((k for k in index if ".layers.3.self_attn.q_proj." in k and k.endswith(".weight")), None), +]: + if probe: + t = _tensor(probe) + print(f"P1 passthrough {probe}: {t.dtype} {tuple(t.shape)}", flush=True) + +# dequant round-trip vs torchao for one quantized expert +wk = next((k for k in index if ".layers.3.mlp.experts.0.down_proj.weight" == k.split("model.", 1)[-1] or k.endswith("layers.3.mlp.experts.0.down_proj.weight")), None) +sk = wk.replace(".weight", ".weight_scale") if wk else None +if wk and sk in index: + from torchao.prototype.mx_formats.nvfp4_tensor import NVFP4Tensor + + w, s = _tensor(wk), _tensor(sk) + nv = NVFP4Tensor(w, s, block_size=16, orig_dtype=torch.bfloat16) + deq = nv.dequantize(torch.bfloat16) if not hasattr(nv, "to_dtype") else nv.to_dtype(torch.bfloat16) + print(f"P1 nvfp4 round-trip ok: {wk} -> dequant {tuple(deq.shape)} finite={torch.isfinite(deq).all().item()}", flush=True) +else: + print("P1 WARNING: could not locate quantized down_proj weight/scale pair", flush=True) + +print("P1 DONE", flush=True) diff --git a/scripts/minimax_m3_nvfp4/p1b_expert_coverage.py b/scripts/minimax_m3_nvfp4/p1b_expert_coverage.py new file mode 100644 index 000000000..9c8a01c78 --- /dev/null +++ b/scripts/minimax_m3_nvfp4/p1b_expert_coverage.py @@ -0,0 +1,208 @@ +"""Phase 1b: per-expert routing coverage across all MoE layers of Minimax-M3-0602. + +Streams the model layer-window by layer-window through one GPU (materialize -> +forward all samples -> free), capturing each MoE layer's routed expert +assignments (top_k_index) at the experts-module boundary. Native top-4 routing, +64 seed-42 HLE samples truncated to their first 8k tokens — identical to the +Phase 1 pilot calibration. + +Output: JSON matrix [57 layers x 128 experts] of routed token counts plus a +per-layer concentration summary. +""" + +import json +import os +import random +import time + +import torch +import torch.nn as nn +from transformers import AutoTokenizer + +# /scratch/tonyzhang/models/Minimax-M3-0602 now returns EIO on config.json and +# every safetensors shard (btrfs read errors); this copy is intact. +MODEL = "/data/huggingface/Minimax-M3-0602" +SNAP = ( + "/data/sgambhira/hf-home/hub/datasets--togethercomputer--m3-quant-eval-runs/" + "snapshots/27f2a493e10922e878840d27924aabd0bbf5d745" +) +OUT_JSON = "/data/sgambhira/GPTQModel/scripts/minimax_m3_nvfp4/coverage_layer_expert.json" +SEED = 42 +NUM_SAMPLES = 64 +MAX_TOKENS = 8192 +# A recurring co-tenant pipeline used to claim GPUs 0-3 for ~150GB in bursts, +# which forced WINDOW=1 (~23GB peak). With the GPU idle, WINDOW=5 amortizes the +# per-window materialize/free cost over 5 layers instead of 1. +WINDOW = 5 +NUM_EXPERTS = 128 +DEVICE = torch.device("cuda:1") +# Co-tenant bursts OOM-killed two previous runs; checkpoint the CPU-side hidden +# states + counts so a restart resumes instead of redoing completed layers. +CKPT = "/data/sgambhira/gptqmodel-offload/coverage_resume.pt" +CKPT_EVERY = 1 # windows between checkpoint saves +OOM_RETRIES = 20 +OOM_WAIT_S = 90 + + +def _retry_oom(fn, what, max_tries=OOM_RETRIES, wait_s=OOM_WAIT_S): + for i in range(max_tries): + try: + return fn() + except torch.OutOfMemoryError: + torch.cuda.empty_cache() + print(f"COV OOM during {what} (attempt {i + 1}/{max_tries}); waiting {wait_s}s", flush=True) + time.sleep(wait_s) + raise RuntimeError(f"COV {what} failed after {max_tries} OOM retries") + +from gptqmodel import BACKEND, GPTQModel, QuantizeConfig # noqa: E402 + +qcfg = QuantizeConfig( + bits=4, + group_size=16, + format="nvfp4", + offload_to_disk_path="/data/sgambhira/gptqmodel-offload/coverage", +) +model = GPTQModel.load(MODEL, quantize_config=qcfg, backend=BACKEND.TORCH) +print("COV model loaded", flush=True) + +lm = model.model.model.language_model +num_layers = len(lm.layers) +all_layers = list(lm.layers) + +# --- calibration inputs (identical to pilot) --- +tokenizer = AutoTokenizer.from_pretrained(MODEL) +rows = [json.loads(l) for l in open(f"{SNAP}/hle-mxfp8-0602-text2158/responses.jsonl")] +sample = random.Random(SEED).sample(rows, NUM_SAMPLES) +inputs = [] +for r in sample: + messages = [ + {"role": "user", "content": r["question"]}, + {"role": "assistant", "content": r["response"], "reasoning_content": r["reasoning"]}, + ] + text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False) + ids = tokenizer(text, add_special_tokens=False, return_tensors="pt")["input_ids"][:, :MAX_TOKENS] + inputs.append(ids) +total_tokens = sum(i.shape[1] for i in inputs) +print(f"COV calibration: {len(inputs)} samples, {total_tokens:,} tokens", flush=True) + +# --- routed-assignment capture at the experts-module boundary --- +counts = {} # layer_idx -> LongTensor[NUM_EXPERTS] + +def make_experts_hook(layer_idx): + def pre_hook(module, args, kwargs): + idx = None + if len(args) >= 2 and torch.is_tensor(args[1]): + idx = args[1] + elif "top_k_index" in kwargs: + idx = kwargs["top_k_index"] + if idx is None: + return + c = torch.bincount(idx.reshape(-1).to("cpu", torch.int64), minlength=NUM_EXPERTS) + counts[layer_idx] = counts.get(layer_idx, torch.zeros(NUM_EXPERTS, dtype=torch.long)) + c + return pre_hook + +hook_handles = [] +moe_layers = [] +for li, layer in enumerate(all_layers): + experts = getattr(getattr(layer, "mlp", None), "experts", None) + if isinstance(experts, nn.Module): + moe_layers.append(li) + hook_handles.append(experts.register_forward_pre_hook(make_experts_hook(li), with_kwargs=True)) +print(f"COV hooked {len(moe_layers)} MoE layers: {moe_layers[:4]}...{moe_layers[-2:]}", flush=True) + +# --- windowed streaming forward --- +# hidden states enter each window via inputs_embeds; rotary/causal-mask are +# position-derived so per-window model forwards are exact. +model.shell_module_materialize(lm.rotary_emb, DEVICE) +model.shell_module_materialize(lm.norm, DEVICE) + +start_layer = 0 +if os.path.exists(CKPT): + state = torch.load(CKPT, map_location="cpu") + hidden = state["hidden"] + counts.update(state["counts"]) + start_layer = state["next_layer"] + print(f"COV resumed from checkpoint: continuing at layer {start_layer}", flush=True) +else: + embed = model.shell_module_materialize(lm.embed_tokens, DEVICE) + hidden = [] + with torch.inference_mode(): + for ids in inputs: + hidden.append(embed(ids.to(DEVICE)).to("cpu")) + print("COV embedded all samples", flush=True) + +orig_layers = lm.layers + +for w_start in range(start_layer, num_layers, WINDOW): + w_layers = all_layers[w_start : w_start + WINDOW] + # materialization resolves module paths against the model tree, so the full + # layer list must be attached while materializing; truncate only for forward. + lm.layers = orig_layers + for layer in w_layers: + model.shell_module_materialize(layer, DEVICE) + lm.layers = nn.ModuleList(w_layers) + + # capture the last window layer's output as the next window's input + window_out = {} + def tail_hook(module, args, output): + window_out["h"] = (output[0] if isinstance(output, tuple) else output).detach() + tail = w_layers[-1].register_forward_hook(tail_hook) + + with torch.inference_mode(): + for si in range(len(hidden)): + h = hidden[si].to(DEVICE) + pos = torch.arange(h.shape[1], device=DEVICE).unsqueeze(0) + lm(inputs_embeds=h, position_ids=pos, use_cache=False) + hidden[si] = window_out["h"].to("cpu") + + tail.remove() + for layer in w_layers: + layer.to_empty(device="meta") + torch.cuda.empty_cache() + print(f"COV window done: layers {w_start}-{w_start + len(w_layers) - 1}", flush=True) + + next_layer = w_start + len(w_layers) + if ((w_start - start_layer) // WINDOW + 1) % CKPT_EVERY == 0 and next_layer < num_layers: + tmp = CKPT + ".tmp" + torch.save({"hidden": hidden, "counts": counts, "next_layer": next_layer}, tmp) + os.replace(tmp, CKPT) + print(f"COV checkpoint saved through layer {next_layer - 1}", flush=True) + +lm.layers = orig_layers + +for h in hook_handles: + h.remove() + +# --- summarize --- +matrix = {str(li): counts.get(li, torch.zeros(NUM_EXPERTS, dtype=torch.long)).tolist() for li in moe_layers} +payload = { + "model": MODEL, + "samples": NUM_SAMPLES, + "max_tokens": MAX_TOKENS, + "seed": SEED, + "total_tokens": total_tokens, + "top_k": 4, + "counts": matrix, +} +with open(OUT_JSON, "w") as f: + json.dump(payload, f) +print(f"COV wrote {OUT_JSON}", flush=True) + +if os.path.exists(CKPT): + os.remove(CKPT) # a finished run must not seed the next fresh run + +threshold = 0.005 * total_tokens +print("COV layer | experts>0 | experts>=fallback_thr | top1_share | top8_share", flush=True) +for li in moe_layers: + c = torch.tensor(matrix[str(li)], dtype=torch.float64) + tot = c.sum().item() + if tot == 0: + print(f"COV {li:>3} | ZERO CAPTURE", flush=True) + continue + srt = c.sort(descending=True).values + print( + f"COV {li:>3} | {(c > 0).sum().item():>3} | {(c >= threshold).sum().item():>3} | " + f"{srt[0].item() / tot:5.1%} | {srt[:8].sum().item() / tot:5.1%}", + flush=True, + ) +print("COV DONE", flush=True) diff --git a/scripts/minimax_m3_nvfp4/quantize_w4a4.py b/scripts/minimax_m3_nvfp4/quantize_w4a4.py new file mode 100644 index 000000000..19a88a802 --- /dev/null +++ b/scripts/minimax_m3_nvfp4/quantize_w4a4.py @@ -0,0 +1,208 @@ +"""GPTQ -> NVFP4 quantization of Minimax-M3-0602, calibrated for w4a4. + +Quantizes the routed experts of all 57 MoE layers to NVFP4. Everything else -- +attention, index projections, shared experts, the 3 dense MLP layers, routers, +norms, embeddings, lm_head, vision tower -- is left untouched and passes through +in the source checkpoint's precision, matching the format map in +tore-quant/recipes/minimax-m3-nvfp4/README.md. + +When an activation amax sidecar is supplied, the GPTQ Hessian is accumulated on +Q(X) rather than X, so the weights are solved for the activations a w4a4 kernel +actually consumes (docs/w4a4_hessian.md). Without one it degrades to a +weight-only calibration and says so. + +Driven by run.sh; all knobs are environment variables. +""" + +import json +import os +import random +import sys +import time + +import torch +from transformers import AutoTokenizer + +MODEL = os.environ.get("MODEL", "/data/huggingface/Minimax-M3-0602") +OUT = os.environ.get("OUT", "/data/huggingface/Minimax-M3-0602-NVFP4-GPTQ-w4a4") +SNAP = os.environ.get( + "SNAP", + "/data/sgambhira/hf-home/hub/datasets--togethercomputer--m3-quant-eval-runs/" + "snapshots/27f2a493e10922e878840d27924aabd0bbf5d745", +) +OFFLOAD = os.environ.get("OFFLOAD", "/data/sgambhira/gptqmodel-offload/w4a4") +ACT_AMAX = os.environ.get("ACT_AMAX", "").strip() +SEED = int(os.environ.get("SEED", "42")) +NUM_SAMPLES = int(os.environ.get("NUM_SAMPLES", "64")) +MAX_TOKENS = int(os.environ.get("MAX_TOKENS", "8192")) +BATCH_SIZE = int(os.environ.get("BATCH_SIZE", "1")) +# Fallback threshold, in TOKENS seen by a module (not samples). +# +# The library default "0.5%" is a fixed quantile of the routing distribution and +# is scale-invariant -- an expert is RTN'd iff its routing share < pct/(100*top_k), +# with total tokens cancelling -- so more calibration data does not reduce it. +# Worse, at 2,355 tokens it sits below the input dim, so ~22% of the modules it +# lets through are GPTQ-solved on a singular Hessian held up only by damping. +# +# H = XᵀX is (columns x columns) and rank(XᵀX) = rank(X) <= min(n_tokens, columns), +# so a module needs at least `columns` (= in_features) tokens to be full rank. +# Default here is per-projection thresholds at exactly that bound. +FALLBACK_THRESHOLD = os.environ.get("FALLBACK_THRESHOLD", "").strip() +# Multiplier on in_features. 1.0 = bare full rank (well-defined but ill-conditioned +# at exactly n == columns); >1 buys conditioning margin at the cost of more RTN. +RANK_MARGIN = float(os.environ.get("RANK_MARGIN", "1.0")) +LAYERS = os.environ.get("LAYERS", "").strip() # e.g. "3" or "3-5"; empty = all + +from gptqmodel import BACKEND, GPTQModel, QuantizeConfig # noqa: E402 +from gptqmodel.quantization.config import FORMAT, Fallback, FallbackStrategy # noqa: E402 + +t0 = time.time() + + +def log(msg: str) -> None: + print(f"[w4a4 +{time.time() - t0:7.1f}s] {msg}", flush=True) + + +# --- what to quantize ------------------------------------------------------- +# Only routed experts. `-:` prefixes a skip rule. +dynamic = { + r"-:.*self_attn\..*": {}, + r"-:.*shared_experts\..*": {}, + r"-:.*\.layers\.(0|1|2)\.mlp\..*": {}, + r"-:.*mlp\.gate\..*": {}, # routers stay F32 + r"-:.*(embed_tokens|lm_head).*": {}, + r"-:.*(vision|visual|mm_projector|patch_merge).*": {}, +} +if LAYERS: + lo, _, hi = LAYERS.partition("-") + keep = set(range(int(lo), int(hi or lo) + 1)) + skip = [str(i) for i in range(60) if i not in keep] + dynamic[rf"-:.*\.layers\.({'|'.join(skip)})\..*"] = {} + log(f"LAYERS={LAYERS}: quantizing only {sorted(keep)}") + +# --- RTN fallback threshold, keyed to each projection's Hessian rank --------- +# All three projections of one expert see the same routed tokens, but they have +# different in_features, so a single global threshold either admits singular +# solves (too low) or RTNs down_proj modules that were perfectly well-posed +# (too high). `dynamic` supports a per-module `fallback` override, so key it to +# the actual dimension: gate/up_proj are (3072, 6144) and down_proj (6144, 3072). +IN_FEATURES = {"gate_proj": 6144, "up_proj": 6144, "down_proj": 3072} + +if FALLBACK_THRESHOLD: + default_fallback = Fallback(strategy=FallbackStrategy.RTN, threshold=FALLBACK_THRESHOLD) + log(f"fallback: global threshold {FALLBACK_THRESHOLD} (overriding rank-based default)") +else: + # Anything not matched below keeps the strictest bound. + default_fallback = Fallback(strategy=FallbackStrategy.RTN, + threshold=int(max(IN_FEATURES.values()) * RANK_MARGIN)) + for proj, cols in IN_FEATURES.items(): + thr = int(cols * RANK_MARGIN) + dynamic[rf".*mlp\.experts\.\d+\.{proj}$"] = { + "fallback": {"strategy": "rtn", "threshold": thr} + } + log("fallback: per-projection rank thresholds " + + ", ".join(f"{p}={int(c * RANK_MARGIN)}" for p, c in IN_FEATURES.items()) + + f" (RANK_MARGIN={RANK_MARGIN}) -- every GPTQ solve gets a full-rank Hessian") + +# --- activation quantization (w4a4) ---------------------------------------- +act_kwargs = {} +if ACT_AMAX: + if not os.path.isfile(ACT_AMAX): + sys.exit(f"ACT_AMAX sidecar not found: {ACT_AMAX}") + act_kwargs = dict( + act_format=FORMAT.NVFP4, + act_group_size=16, + act_amax_path=ACT_AMAX, + # w1/w3 (gate_proj/up_proj) share the MoE input, so they share one amax + # entry; w2 (down_proj) sees the post-SwiGLU activation and gets its own. + act_amax_key_rules=[ + (r".*\.layers\.(\d+)\.mlp\.experts\.\d+\.(?:gate_proj|up_proj)$", r"layer\1.moe_input"), + (r".*\.layers\.(\d+)\.mlp\.experts\.\d+\.down_proj$", r"layer\1.w2_input"), + ], + ) + log(f"w4a4: Hessian will be accumulated on Q(X), amax sidecar {ACT_AMAX}") +else: + log("NO ACT_AMAX: weight-only calibration (Hessian on clean X). NOT w4a4-optimal.") + +qcfg = QuantizeConfig( + bits=4, + group_size=16, + format="nvfp4", + sym=True, + desc_act=False, + dynamic=dynamic, + fallback=default_fallback, + moe_vram_strategy="balanced", + offload_to_disk_path=OFFLOAD, + **act_kwargs, +) +log(f"qcfg format={qcfg.format} bits={qcfg.bits} group_size={qcfg.group_size}") + +# --- calibration: 64 seed-42 HLE trajectories, first 8k tokens -------------- +tokenizer = AutoTokenizer.from_pretrained(MODEL) +rows = [json.loads(line) for line in open(f"{SNAP}/hle-mxfp8-0602-text2158/responses.jsonl")] +sample = random.Random(SEED).sample(rows, NUM_SAMPLES) + +calibration = [] +total_tokens = 0 +for r in sample: + messages = [ + {"role": "user", "content": r["question"]}, + {"role": "assistant", "content": r["response"], "reasoning_content": r["reasoning"]}, + ] + text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False) + enc = tokenizer(text, add_special_tokens=False, return_tensors="pt") + calibration.append({ + "input_ids": enc["input_ids"][:, :MAX_TOKENS], + "attention_mask": enc["attention_mask"][:, :MAX_TOKENS], + }) + total_tokens += calibration[-1]["input_ids"].shape[1] + +log(f"calibration: {len(calibration)} samples, {total_tokens:,} tokens") + +# --- run -------------------------------------------------------------------- +model = GPTQModel.load(MODEL, quantize_config=qcfg, backend=BACKEND.TORCH) +log("model loaded") + +model.quantize(calibration, batch_size=BATCH_SIZE, backend=BACKEND.TORCH) +log("quantize complete") + +model.save(OUT) +log(f"saved -> {OUT}") + +# --- post-check: the export must actually be NVFP4 -------------------------- +# A silently-passed-through expert is the failure mode that looks like success, +# so assert the layout rather than trusting the run log. +from safetensors import safe_open # noqa: E402 + +index = json.load(open(f"{OUT}/model.safetensors.index.json"))["weight_map"] +# Probe INSIDE the quantized layer set: with LAYERS set, most layers are +# passthrough and legitimately lack NVFP4 scales, so a layer-blind probe fails +# on a perfectly good checkpoint (it did: it hit passthrough layer 10 first). +probe_layer = LAYERS.partition("-")[0] if LAYERS else "3" +probe = next( + (k for k in index + if k.endswith(f".layers.{probe_layer}.mlp.experts.0.down_proj.weight")), + None, +) +if probe is None: + sys.exit(f"post-check FAILED: no routed-expert weight for layer {probe_layer} in the saved index") + +with safe_open(f"{OUT}/{index[probe]}", framework="pt", device="cpu") as f: + w = f.get_tensor(probe) +scale_key = probe.replace(".weight", ".weight_scale") +if scale_key not in index: + sys.exit(f"post-check FAILED: {probe} has no weight_scale -- not NVFP4") +scale2_key = probe.replace(".weight", ".weight_scale_2") +if scale2_key not in index: + sys.exit(f"post-check FAILED: {probe} has no weight_scale_2 -- one-level export") +with safe_open(f"{OUT}/{index[scale_key]}", framework="pt", device="cpu") as f: + s = f.get_tensor(scale_key) + +log(f"post-check {probe}: weight {w.dtype} {tuple(w.shape)} | scale {s.dtype} {tuple(s.shape)}") +assert w.dtype == torch.uint8, f"expected packed uint8 weight, got {w.dtype}" +assert s.dtype == torch.float8_e4m3fn, f"expected float8_e4m3fn scale, got {s.dtype}" + +n_scales = sum(1 for k in index if k.endswith(".weight_scale")) +log(f"post-check OK: {n_scales:,} NVFP4 weight_scale tensors in the checkpoint") +log("DONE") diff --git a/scripts/minimax_m3_nvfp4/run.sh b/scripts/minimax_m3_nvfp4/run.sh new file mode 100755 index 000000000..1779bc041 --- /dev/null +++ b/scripts/minimax_m3_nvfp4/run.sh @@ -0,0 +1,191 @@ +#!/usr/bin/env bash +set -euo pipefail + +# GPTQ -> NVFP4 quantization of Minimax-M3-0602, calibrated for w4a4. +# +# Quantizes the routed experts of all 57 MoE layers (57 x 128 x 3 = 21,888 +# modules). Attention, shared experts, dense MLP 0-2, routers, norms, embeddings +# and the vision tower pass through untouched -- the format map from +# tore-quant/recipes/minimax-m3-nvfp4. +# +# bash scripts/minimax_m3_nvfp4/run.sh # preflight + full run +# PREFLIGHT_ONLY=1 bash scripts/minimax_m3_nvfp4/run.sh # checks only +# LAYERS=3 bash scripts/minimax_m3_nvfp4/run.sh # single-layer smoke test +# ACT_AMAX=/path/amax_per_layer.json bash .../run.sh # true w4a4 calibration +# +# READ BEFORE RUNNING -- two things this does NOT do: +# +# 1. Without ACT_AMAX this is a *weight-only* calibration. The Hessian is built +# on clean X, so the weights are solved for activations a w4a4 kernel will +# not feed them. See docs/w4a4_hessian.md. Produce a sidecar with +# tore-quant `ptq.py --emit-amax-sidecar`. +# +# 2. The export writes `weight` (packed uint8) + `weight_scale` (float8_e4m3fn) +# + `weight_scale_2` (fp32 per-tensor). It does NOT write `input_scale`; +# a serving stack expecting the full ModelOpt layout needs that added as a +# post-step (derivable from the ACT_AMAX sidecar: amax/6/448 per input site). +# ACT_AMAX changes what the weights are optimized for, not the tensor set. + +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO="$(cd "${HERE}/../.." && pwd)" +PY="${PYTHON_BIN:-${REPO}/.venv/bin/python}" + +# /scratch/tonyzhang/models/Minimax-M3-0602 returns EIO on config.json and every +# safetensors shard (btrfs read errors). This copy is intact. +MODEL="${MODEL:-/data/huggingface/Minimax-M3-0602}" +OUT="${OUT:-/data/huggingface/Minimax-M3-0602-NVFP4-GPTQ-w4a4}" +SNAP="${SNAP:-/data/sgambhira/hf-home/hub/datasets--togethercomputer--m3-quant-eval-runs/snapshots/27f2a493e10922e878840d27924aabd0bbf5d745}" +OFFLOAD="${OFFLOAD:-/data/sgambhira/gptqmodel-offload/w4a4}" +ACT_AMAX="${ACT_AMAX:-}" +LAYERS="${LAYERS:-}" +NUM_SAMPLES="${NUM_SAMPLES:-64}" +MAX_TOKENS="${MAX_TOKENS:-8192}" +# Empty = per-projection rank thresholds (gate/up 6144, down 3072). Set to a +# number or "N%" to override globally. +FALLBACK_THRESHOLD="${FALLBACK_THRESHOLD:-}" +RANK_MARGIN="${RANK_MARGIN:-1.0}" +LOG="${LOG:-${HERE}/w4a4_run.log}" +PREFLIGHT_ONLY="${PREFLIGHT_ONLY:-0}" + +# Weights alone are ~211 GiB (57 x 3.80 GiB NVFP4 experts) plus ~16.5 GiB of +# passthrough; leave room for the offload dir on the same filesystem. +MIN_FREE_GIB="${MIN_FREE_GIB:-600}" + +say() { printf '\n\033[1m== %s\033[0m\n' "$*"; } +die() { printf '\033[31mFAIL: %s\033[0m\n' "$*" >&2; exit 1; } + +# Affinity watchdog: a host agent episodically pins spawned python processes to a +# single core, and wekanode busy-polls cores 1/28/56/84 at 100% -- a process +# parked there is starved (state R, frozen CPU counters, undeliverable signals). +# Re-widen within 2s whenever the driver or splice collapses to one core. +# Machine-specific mask (224-core B200 box); harmless elsewhere. +WEKA_SAFE_MASK="0,2-27,29-55,57-83,85-223" +start_watchdog() { + command -v taskset >/dev/null || return 0 + ( + while :; do + for P in $(pgrep -f "quantize_w4a4.py|splice_source_passthrough.py" 2>/dev/null); do + A=$(awk '/Cpus_allowed_list/{print $2}' "/proc/${P}/status" 2>/dev/null) + case "${A}" in + *[,-]*|"") : ;; + *) taskset -pc "${WEKA_SAFE_MASK}" "${P}" >/dev/null 2>&1 && echo "watchdog: re-pinned pid ${P} from core ${A}" ;; + esac + done + sleep 2 + done + ) & + WATCHDOG_PID=$! + trap 'kill "${WATCHDOG_PID}" 2>/dev/null || true' EXIT +} + +say "Preflight" + +[[ -x "${PY}" ]] || die "python not executable: ${PY} (set PYTHON_BIN)" + +# Read CONTENT, not just stat: /scratch's copy passes `ls` and fails on read. +[[ -d "${MODEL}" ]] || die "model dir missing: ${MODEL}" +head -c 1 "${MODEL}/config.json" >/dev/null 2>&1 || die "cannot READ ${MODEL}/config.json (I/O error?)" +shards=$(ls "${MODEL}"/*.safetensors 2>/dev/null | wc -l) +[[ "${shards}" -gt 0 ]] || die "no safetensors shards in ${MODEL}" +first_shard=$(ls "${MODEL}"/*.safetensors | head -1) +dd if="${first_shard}" of=/dev/null bs=1M count=1 status=none 2>/dev/null \ + || die "cannot READ ${first_shard} -- checkpoint is corrupt" +echo " model ${MODEL} (${shards} shards, readable)" + +[[ -f "${SNAP}/hle-mxfp8-0602-text2158/responses.jsonl" ]] \ + || die "calibration responses.jsonl missing under ${SNAP}" +echo " calib data ${SNAP}" + +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)')" \ + || die "ACT_AMAX is not valid JSON" +else + printf '\033[33m act amax NOT SET -- weight-only calibration, not w4a4-optimal\033[0m\n' +fi + +mkdir -p "${OFFLOAD}" "$(dirname "${OUT}")" +free_gib=$(df -PBG "$(dirname "${OUT}")" | awk 'NR==2{gsub(/G/,"",$4); print $4}') +[[ "${free_gib}" -ge "${MIN_FREE_GIB}" ]] \ + || die "only ${free_gib} GiB free at $(dirname "${OUT}"); need >= ${MIN_FREE_GIB}" +echo " disk ${free_gib} GiB free" + +"${PY}" - <<'PYCHK' || die "no usable CUDA device" +import torch, sys +if not torch.cuda.is_available(): + sys.exit(1) +free, total = torch.cuda.mem_get_info(0) +print(f" gpu {torch.cuda.device_count()}x {torch.cuda.get_device_name(0)}, " + f"{free/2**30:.0f}/{total/2**30:.0f} GiB free on cuda:0") +PYCHK + +if [[ -e "${OUT}" ]]; then + printf '\033[33m WARNING %s already exists; model.save() will overwrite it\033[0m\n' "${OUT}" +fi + +[[ "${PREFLIGHT_ONLY}" == "1" ]] && { say "Preflight only -- stopping"; exit 0; } + +say "Expected RTN fallback rate" +# A module seeing fewer tokens than its own in_features cannot form a full-rank +# Hessian, so it is RTN-quantized rather than GPTQ-solved. Surface the rate up +# front rather than discovering it in the log. +COVERAGE="${HERE}/coverage_layer_expert.json" +if [[ -f "${COVERAGE}" ]]; then + "${PY}" - <5} tok -> RTN {r:,}/{tot:,} = {r/tot:.1%}") +print(" Thresholds are each projection's in_features: H = XtX is (cols x cols) and") +print(" rank(XtX) <= min(n_tokens, cols), so fewer tokens than cols gives a singular H.") +print(" NOTE: more calibration data does NOT lower this. A percentage threshold is") +print(" scale-invariant (RTN iff share < pct/(100*top_k)); only the threshold moves it.") +PYRTN +else + echo " (no coverage_layer_expert.json; run p1b_expert_coverage.py to predict)" +fi + +start_watchdog + +say "Quantizing" +echo " out ${OUT}" +echo " log ${LOG}" +[[ -n "${LAYERS}" ]] && echo " layers ${LAYERS} (subset)" + +MODEL="${MODEL}" OUT="${OUT}" SNAP="${SNAP}" OFFLOAD="${OFFLOAD}" \ +ACT_AMAX="${ACT_AMAX}" LAYERS="${LAYERS}" NUM_SAMPLES="${NUM_SAMPLES}" \ +MAX_TOKENS="${MAX_TOKENS}" FALLBACK_THRESHOLD="${FALLBACK_THRESHOLD}" \ +RANK_MARGIN="${RANK_MARGIN}" \ +PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}" \ +COLUMNS="${COLUMNS:-360}" \ + "${PY}" -u "${HERE}/quantize_w4a4.py" 2>&1 | tee "${LOG}" + +# Phase 3: splice quantized experts into the SOURCE checkpoint so passthrough +# modules keep their original serialization (MXFP8/BF16/F32, source naming) and +# experts gain input_scale from the amax sidecar -- the serving format map. +SPLICE="${SPLICE:-1}" +OUT_FINAL="${OUT_FINAL:-${OUT}-final}" +if [[ "${SPLICE}" == "1" ]]; then + say "Splicing source passthrough -> ${OUT_FINAL}" + SRC="${MODEL}" QNT="${OUT}" OUT="${OUT_FINAL}" ACT_AMAX="${ACT_AMAX}" "${PY}" -u "${HERE}/splice_source_passthrough.py" 2>&1 | tee -a "${LOG}" +fi + +say "Done" +# Roster: which modules were RTN'd (named in WARN lines) vs GPTQ-solved (the rest). +ROSTER="${OUT}/quant_roster.json" +"${PY}" - "${LOG}" "${ROSTER}" <<'PYROSTER' || true +import json, re, sys +log, out = sys.argv[1], sys.argv[2] +rtn = {} +for m in re.finditer(r"Module `([\w.]+)` -> Using `rtn` fallback quantization \(observed (\d+) samples, threshold=(\d+)", open(log, errors="replace").read()): + rtn[m.group(1)] = {"observed_tokens": int(m.group(2)), "threshold": int(m.group(3))} +json.dump({"rtn_modules": rtn, "note": "GPTQ-solved = quantized modules not listed here"}, open(out, "w"), indent=1) +print(f" roster {out} ({len(rtn)} RTN modules)") +PYROSTER +echo " checkpoint ${OUT}" +echo " next validate + add input_scale/weight_scale_2 if your serving" +echo " stack expects the ModelOpt NVFP4 layout (see header note 2)" diff --git a/scripts/minimax_m3_nvfp4/splice_source_passthrough.py b/scripts/minimax_m3_nvfp4/splice_source_passthrough.py new file mode 100644 index 000000000..10feb1604 --- /dev/null +++ b/scripts/minimax_m3_nvfp4/splice_source_passthrough.py @@ -0,0 +1,180 @@ +"""Splice GPTQ-NVFP4 experts into the source checkpoint (passthrough = source bytes). + +GPTQModel's save materializes every module, so passthrough tensors come out as +dequantized BF16 -- numerically right, but not the serving format map: attention, +index projections, shared experts, dense MLP 0-2 and non-quantized expert layers +must stay byte-identical MXFP8/BF16/F32 from the source release +(tore-quant/recipes/minimax-m3-nvfp4 format map). + +This builds the final checkpoint the way tore-quant's build_nvfp4_mxfp8_mix does: +walk the SOURCE tensor list, and only for routed experts of quantized layers swap +in the GPTQ output's NVFP4 tensors (renamed to source naming), plus the +`input_scale` the ModelOpt layout expects, derived from the activation amax +sidecar (amax / 6 / 448 per input site -- the same numbers the Hessian was +calibrated with). + + SRC=/data/huggingface/Minimax-M3-0602 \ + QNT=/data/huggingface/Minimax-M3-0602-NVFP4-w4a4-smoke-l3 \ + OUT=/data/huggingface/Minimax-M3-0602-NVFP4-w4a4-smoke-l3-final \ + ACT_AMAX=scripts/minimax_m3_nvfp4/amax_per_layer.json \ + python scripts/minimax_m3_nvfp4/splice_source_passthrough.py +""" + +import json +import os +import re +import shutil +import time + +import torch +from safetensors import safe_open +from safetensors.torch import save_file + +SRC = os.environ.get("SRC", "/data/huggingface/Minimax-M3-0602") +QNT = os.environ["QNT"] +OUT = os.environ["OUT"] +ACT_AMAX = os.environ.get("ACT_AMAX", "").strip() +SHARD_BYTES = int(os.environ.get("SHARD_BYTES", str(8 << 30))) + +t0 = time.time() + + +def log(msg: str) -> None: + print(f"[splice +{time.time() - t0:7.1f}s] {msg}", flush=True) + + +# --- naming: source <-> gptqmodel-save --------------------------------------- +# language_model.model.layers.N.block_sparse_moe.experts.E.{w1,w2,w3} +# <-> model.language_model.layers.N.mlp.experts.E.{gate_proj,down_proj,up_proj} +PROJ_TO_QNT = {"w1": "gate_proj", "w2": "down_proj", "w3": "up_proj"} +SRC_EXPERT_RE = re.compile( + r"^language_model\.model\.layers\.(\d+)\.block_sparse_moe\.experts\.(\d+)\.(w[123])\." +) + + +def qnt_base(layer: str, expert: str, proj: str) -> str: + return f"model.language_model.layers.{layer}.mlp.experts.{expert}.{PROJ_TO_QNT[proj]}" + + +src_index = json.load(open(f"{SRC}/model.safetensors.index.json"))["weight_map"] +qnt_index = json.load(open(f"{QNT}/model.safetensors.index.json"))["weight_map"] + +# Quantized modules are exactly those the GPTQ output gave a weight_scale_2. +quantized_layers = sorted({ + int(m.group(1)) + for k in qnt_index if k.endswith(".weight_scale_2") + for m in [re.search(r"\.layers\.(\d+)\.mlp\.experts\.", k)] if m +}) +log(f"quantized layers in {QNT}: {quantized_layers}") +if not quantized_layers: + raise SystemExit("no quantized expert layers found in QNT -- nothing to splice") + +amax = {} +if ACT_AMAX: + payload = json.load(open(ACT_AMAX)) + amax = payload.get("amax", payload) + log(f"input_scale source: {ACT_AMAX} ({len(amax)} entries)") +else: + log("WARNING: no ACT_AMAX -- input_scale tensors will NOT be written") + +# --- plan the output tensor list, in source order ----------------------------- +# Per quantized expert projection: drop source {weight, weight_scale_inv}; emit +# {weight, weight_scale, weight_scale_2[, input_scale]} from QNT under source names. +PROJ_SITE = {"w1": "moe_input", "w3": "moe_input", "w2": "w2_input"} + +plan = [] # (out_name, origin, origin_key) origin in {"src", "qnt", "amax"} +emitted = set() +for name in src_index: + m = SRC_EXPERT_RE.match(name) + if m and int(m.group(1)) in quantized_layers: + layer, expert, proj = m.group(1), m.group(2), m.group(3) + base_src = f"language_model.model.layers.{layer}.block_sparse_moe.experts.{expert}.{proj}" + if base_src in emitted: + continue # second source tensor (weight_scale_inv) of the same module + emitted.add(base_src) + base_q = qnt_base(layer, expert, proj) + for suffix in (".weight", ".weight_scale", ".weight_scale_2"): + if base_q + suffix not in qnt_index: + raise SystemExit(f"QNT missing {base_q}{suffix} -- incomplete quantized module") + plan.append((base_src + suffix, "qnt", base_q + suffix)) + if amax: + key = f"layer{layer}.{PROJ_SITE[proj]}" + if key not in amax: + raise SystemExit(f"amax sidecar missing {key}") + plan.append((base_src + ".input_scale", "amax", key)) + else: + plan.append((name, "src", name)) + +log(f"plan: {len(plan)} tensors ({sum(1 for _, o, _ in plan if o != 'src')} spliced, " + f"{sum(1 for _, o, _ in plan if o == 'src')} source passthrough)") + +# --- stream tensors into fresh shards ----------------------------------------- +os.makedirs(OUT, exist_ok=True) +handles = {} + + +def read_tensor(origin: str, key: str) -> torch.Tensor: + if origin == "amax": + return torch.tensor(float(amax[key]) / 6.0 / 448.0, dtype=torch.float32) + root, index = (SRC, src_index) if origin == "src" else (QNT, qnt_index) + shard = index[key] + cache_key = (root, shard) + if cache_key not in handles: + handles.clear() # one open shard per side at a time; source order keeps locality + handles[cache_key] = safe_open(os.path.join(root, shard), framework="pt") + return handles[cache_key].get_tensor(key) + + +weight_map = {} +buffer, buf_bytes, shard_id = {}, 0, 0 +total = 0 + + +def flush(): + global buffer, buf_bytes, shard_id + if not buffer: + return + shard_id += 1 + fname = f"model-{shard_id:05d}.safetensors" + save_file(buffer, os.path.join(OUT, fname)) + for k in buffer: + weight_map[k] = fname + buffer, buf_bytes = {}, 0 + + +for out_name, origin, key in plan: + t = read_tensor(origin, key) + buffer[out_name] = t + buf_bytes += t.numel() * t.element_size() + total += 1 + if buf_bytes >= SHARD_BYTES: + flush() + if shard_id % 8 == 0: + log(f" wrote shard {shard_id} ({total}/{len(plan)} tensors)") +flush() +log(f"wrote {shard_id} shards, {len(weight_map)} tensors") + +# rename shards to the canonical model-XXXXX-of-YYYYY pattern +final_names = {} +for i in range(1, shard_id + 1): + old = f"model-{i:05d}.safetensors" + new = f"model-{i:05d}-of-{shard_id:05d}.safetensors" + os.replace(os.path.join(OUT, old), os.path.join(OUT, new)) + final_names[old] = new +weight_map = {k: final_names[v] for k, v in weight_map.items()} + +with open(f"{OUT}/model.safetensors.index.json", "w") as f: + json.dump({"metadata": {"total_size": sum( + os.path.getsize(os.path.join(OUT, s)) for s in set(weight_map.values()) + )}, "weight_map": weight_map}, f, indent=1) + +# --- config + aux files: source's, with quantized experts removed from the +# mxfp8 ignore/exclude story. Serving-side quantization_config for the mixed +# NVFP4+MXFP8 layout follows the MiniMax-M3-NVFP4-mxfp8skip14 recipe and is left +# to the serving prep step; here we copy source aux files verbatim. +for fname in os.listdir(SRC): + if fname.endswith((".json", ".py", ".jinja", ".txt")) and fname != "model.safetensors.index.json": + shutil.copy2(os.path.join(SRC, fname), os.path.join(OUT, fname)) +log("copied source aux/config files verbatim (serving quantization_config is a separate prep step)") + +log(f"DONE -> {OUT}") diff --git a/tests/test_fp8_e8m0_dequant.py b/tests/test_fp8_e8m0_dequant.py new file mode 100644 index 000000000..0fc55ad76 --- /dev/null +++ b/tests/test_fp8_e8m0_dequant.py @@ -0,0 +1,89 @@ +# SPDX-FileCopyrightText: 2026 ModelCloud.ai +# SPDX-License-Identifier: Apache-2.0 + +"""uint8 scales entering dequantize_fp8 are UE8M0 exponent bytes, not magnitudes. + +MX checkpoints (mxfp8) serialize `weight_scale_inv` as uint8 E8M0: multiplier +2^(b - 127). Feeding the raw bytes into the float scale arithmetic is the bug +this pins: byte values ~110-130 exceed 1, so `_fast_scale_arg` flipped scale_inv +into divide mode and every dequantized weight came out wrong by a block-dependent +factor (observed absmax 3.9 instead of 0.28 on Minimax-M3-0602). This was the +weight source for GPTQ's quant_source path, so it would have silently poisoned +whole quantization runs. +""" + +import os + + +os.environ.setdefault("CUDA_VISIBLE_DEVICES", "") + +import pytest +import torch + +from gptqmodel.quantization.dtype import dequantize_block_fp8, dequantize_fp8 + + +BLOCK = 32 # mxfp8 weight_block_size [1, 32] + + +def _mx_tensor(rows=8, cols=128, seed=0): + """A synthetic mxfp8 pair: fp8 payload + uint8 E8M0 block scales.""" + + torch.manual_seed(seed) + payload = (torch.randn(rows, cols) * 100).to(torch.float8_e4m3fn) + # exponents spanning well below and above the bias, like real checkpoints + scale = torch.randint(110, 135, (rows, cols // BLOCK), dtype=torch.uint8) + return payload, scale + + +def test_uint8_scale_inv_decodes_as_e8m0(): + payload, scale = _mx_tensor() + out = dequantize_fp8(payload, scale_inv=scale, axis=None, target_dtype=torch.bfloat16) + ref = dequantize_block_fp8(payload, scale, target_dtype=torch.bfloat16) + assert torch.equal(out, ref), "uint8 scale_inv not decoded as 2^(b-127)" + + +def test_uint8_decode_is_multiplicative_not_divisive(): + """The historic failure mode: bytes > 1 flipped the path into divide mode.""" + + payload = torch.full((1, BLOCK), 4.0).to(torch.float8_e4m3fn) + scale = torch.tensor([[127 - 3]], dtype=torch.uint8) # 2^-3 = 0.125 + out = dequantize_fp8(payload, scale_inv=scale, axis=None, target_dtype=torch.float32) + assert torch.allclose(out, torch.full((1, BLOCK), 0.5)), ( + f"expected 4.0 * 2^-3 = 0.5, got {out[0, 0].item()} " + "(4.0 / 124 would indicate divide-by-raw-bytes has returned)" + ) + + +def test_float_scale_inv_behavior_unchanged(): + """Float scales must not take the uint8/E8M0 branch: their multiplier + semantics are pre-existing and pinned here with a full-shape scale, which is + unambiguous on every internal path.""" + + payload, _ = _mx_tensor(seed=1) + scale_f32 = torch.full(payload.shape, 0.125, dtype=torch.float32) + out = dequantize_fp8(payload, scale_inv=scale_f32, axis=None, target_dtype=torch.float32) + ref = payload.to(torch.float32) * 0.125 + assert torch.allclose(out, ref) + + +@pytest.mark.skipif( + not os.path.isdir("/data/huggingface/Minimax-M3-0602"), reason="M3 checkpoint not present" +) +def test_real_minimax_m3_tensor_decodes_to_sane_magnitude(): + """End-to-end pin against the real checkpoint that exposed the bug.""" + + import json + + from safetensors import safe_open + + d = "/data/huggingface/Minimax-M3-0602" + idx = json.load(open(f"{d}/model.safetensors.index.json"))["weight_map"] + name = "language_model.model.layers.3.self_attn.q_proj.weight" + with safe_open(os.path.join(d, idx[name]), framework="pt") as f: + w = f.get_tensor(name) + with safe_open(os.path.join(d, idx[name + "_scale_inv"]), framework="pt") as f: + s = f.get_tensor(name + "_scale_inv") + out = dequantize_fp8(w, scale_inv=s, axis=None, target_dtype=torch.bfloat16) + absmax = out.float().abs().max().item() + assert 0.1 < absmax < 1.0, f"absmax {absmax}: ~0.28 is correct, 3.9 is divide-mode, 448 is raw" diff --git a/tests/test_nvfp4_export.py b/tests/test_nvfp4_export.py new file mode 100644 index 000000000..5f48daeed --- /dev/null +++ b/tests/test_nvfp4_export.py @@ -0,0 +1,467 @@ +# SPDX-FileCopyrightText: 2026 ModelCloud.ai +# SPDX-License-Identifier: Apache-2.0 + +"""NVFP4 export coverage: GPTQ-on-NVFP4-grid and weight-only NVFP4 packing. + +1. `NVFP4Quantizer` must be bit-identical to torchao's ``nvfp4_quantize`` grid. +2. `TorchNVFP4Linear` packing must be lossless against the scales produced by + the GPTQ loop (error feedback mutates columns after group scales are fixed, + so export must reuse the loop's scales, never re-derive them). +3. End-to-end: GPTQ quantize -> save -> reload -> forward on a tiny model, + with the checkpoint stored in the NVFP4 layout (packed uint8 weight + + float8_e4m3fn `weight_scale`). +""" + +import os +from pathlib import Path + + +# Keep this test on CPU so it works on CPU-only runners. +os.environ.setdefault("CUDA_VISIBLE_DEVICES", "") + +import pytest +import torch +import torch.nn as nn +from safetensors import safe_open +from tokenizers import Tokenizer +from tokenizers.models import WordLevel +from tokenizers.pre_tokenizers import Whitespace +from tokenizers.trainers import WordLevelTrainer +from transformers import PreTrainedTokenizerFast +from transformers.models.llama.modeling_llama import LlamaConfig, LlamaForCausalLM + +from gptqmodel import BACKEND, GPTQModel, QuantizeConfig +from gptqmodel.nn_modules.qlinear.fp4 import TorchNVFP4Linear, quantize_nvfp4_weight +from gptqmodel.quantization.config import FORMAT, METHOD, GPTQConfig, NVFP4Config +from gptqmodel.quantization.quantizer import NVFP4Quantizer + + +try: + from torchao.prototype.mx_formats.nvfp4_tensor import NVFP4Tensor, nvfp4_quantize +except Exception: + NVFP4Tensor = None + nvfp4_quantize = None + +pytestmark = [ + pytest.mark.cpu, + pytest.mark.skipif(NVFP4Tensor is None, reason="torchao NVFP4 support required"), +] + +BLOCK = 16 + + +def _nvfp4_dequant(packed: torch.Tensor, scales: torch.Tensor, dtype=torch.float32) -> torch.Tensor: + nv = NVFP4Tensor(packed, scales, block_size=BLOCK, orig_dtype=dtype) + to_dtype = getattr(nv, "to_dtype", None) + if callable(to_dtype): + return to_dtype(dtype) + return nv.dequantize(dtype) + + +# --------------------------------------------------------------------------- +# Config resolution +# --------------------------------------------------------------------------- + + +def test_format_nvfp4_resolves_to_gptq_config(): + cfg = QuantizeConfig(bits=4, group_size=16, format="nvfp4") + assert isinstance(cfg, GPTQConfig) + assert cfg.method == METHOD.GPTQ + assert cfg.format == FORMAT.NVFP4 + assert cfg.sym is True + assert cfg.desc_act is False + assert cfg.act_group_aware is False + + payload = cfg.to_dict() + reloaded = QuantizeConfig.from_quant_config(payload) + assert isinstance(reloaded, GPTQConfig) + assert reloaded.format == FORMAT.NVFP4 + + +def test_method_nvfp4_resolves_to_weight_only_config(): + cfg = QuantizeConfig(quant_method="nvfp4") + assert isinstance(cfg, NVFP4Config) + assert cfg.uses_weight_only_lifecycle() is True + assert cfg.bits == 4 + assert cfg.weight_block_size == BLOCK + + reloaded = QuantizeConfig.from_quant_config(cfg.to_dict()) + assert isinstance(reloaded, NVFP4Config) + + shorthand = QuantizeConfig(weight_only="nvfp4") + assert isinstance(shorthand, NVFP4Config) + + +def test_from_quant_config_accepts_fp4_alias(): + cfg = QuantizeConfig.from_quant_config({"quant_method": "nvfp4", "format": "fp4", "bits": 4}) + assert isinstance(cfg, NVFP4Config) + assert cfg.format == FORMAT.NVFP4 + + +@pytest.mark.parametrize( + "kwargs", + [ + {"bits": 8, "group_size": 16}, + {"bits": 4, "group_size": 128}, + {"bits": 4, "group_size": 16, "sym": False}, + {"bits": 4, "group_size": 16, "desc_act": True}, + {"bits": 4, "group_size": 16, "act_group_aware": True}, + # non-RTN fallback strategies quantize off the NVFP4 grid + {"bits": 4, "group_size": 16, "fallback": {"strategy": "mean"}}, + # fallback smoothing breaks E4M3-exact scales, so packing stops being lossless + {"bits": 4, "group_size": 16, "fallback": {"strategy": "rtn", "smooth": "mse"}}, + # the mock loop rounds on an integer grid, off the NVFP4 grid + {"bits": 4, "group_size": 16, "mock_quantization": True}, + ], +) +def test_format_nvfp4_rejects_incompatible_gptq_settings(kwargs): + with pytest.raises(ValueError): + QuantizeConfig(format="nvfp4", **kwargs) + + +# --------------------------------------------------------------------------- +# Quantizer grid exactness +# --------------------------------------------------------------------------- + + +def test_nvfp4_quantizer_requires_primed_global_scale(): + """Single-level (unprimed) scale derivation was removed: absolute amax/6 + block scales floor-clamp at E4M3's 2^-6 on realistic weights. An unprimed + find_params is a caller bug and must raise, not silently degrade.""" + + qcfg = QuantizeConfig(bits=4, group_size=16, format="nvfp4") + quantizer = NVFP4Quantizer(qcfg=qcfg) + quantizer.configure(perchannel=True) + + with pytest.raises(RuntimeError, match="prime_global_scale"): + quantizer.find_params(torch.randn(64, BLOCK), weight=True) + + +# --------------------------------------------------------------------------- +# Kernel packing +# --------------------------------------------------------------------------- + + +def test_direct_pack_matches_torchao(): + """RTN pack must match torchao's TWO-level path: per-tensor scale amax/(6*448) + plus E4M3 block scales, matching the ModelOpt/Quark/tore-quant layout.""" + + torch.manual_seed(1) + lin = nn.Linear(128, 64, bias=True) + + module = TorchNVFP4Linear( + bits=4, + group_size=-1, + sym=True, + desc_act=False, + in_features=128, + out_features=64, + bias=True, + register_buffers=True, + ) + module.pack_original(linear=lin, scales=None, zeros=None) + + w = lin.weight.detach().float() + gs_ref = (w.abs().max() / (6.0 * 448.0)).reshape(()) + scales_ref, packed_ref = nvfp4_quantize(w.contiguous(), block_size=BLOCK, per_tensor_scale=gs_ref) + assert torch.equal(module.weight, packed_ref.view(torch.uint8)) + assert torch.equal(module.weight_scale.view(torch.uint8), scales_ref.view(torch.uint8)) + assert torch.equal(module.weight_scale_2, gs_ref) + + deq = module.dequantize_weight(dtype=torch.float32).T + ref = NVFP4Tensor( + packed_ref, scales_ref, BLOCK, torch.float32, per_tensor_scale=gs_ref + ).dequantize(torch.float32) + assert torch.allclose(deq, ref, atol=1e-3) + + +def test_gptq_scale_pack_is_lossless(): + """Error feedback shrinks/grows columns after the group scale is fixed; packing + against the loop's scales must reproduce the GPTQ output weight exactly.""" + + torch.manual_seed(2) + IN, OUT = 128, 64 + qcfg = QuantizeConfig(bits=4, group_size=16, format="nvfp4") + quantizer = NVFP4Quantizer(qcfg=qcfg) + quantizer.configure(perchannel=True) + + W = torch.randn(OUT, IN, dtype=torch.float32) + quantizer.prime_global_scale(W) + Q = torch.empty_like(W) + scales = [] + for g in range(0, IN, BLOCK): + block = W[:, g : g + BLOCK] + quantizer.find_params(block, weight=True) + # simulate GPTQ error feedback mutating the block after scale is fixed + block = block + 0.05 * torch.randn_like(block) + Q[:, g : g + BLOCK] = quantizer.quantize(block) + scales.append(quantizer.scale) + scales = torch.cat(scales, dim=1) + + lin = nn.Linear(IN, OUT, bias=False) + lin.weight.data = Q.clone() + module = TorchNVFP4Linear( + bits=4, + group_size=16, + sym=True, + desc_act=False, + in_features=IN, + out_features=OUT, + bias=False, + register_buffers=True, + ) + module.pack_original(linear=lin, scales=scales, zeros=None) + + deq = module.dequantize_weight(dtype=torch.float32).T + assert torch.equal(deq, Q) + + +def test_quantize_nvfp4_weight_rejects_bad_shapes(): + with pytest.raises(ValueError): + quantize_nvfp4_weight(torch.randn(4, 24)) # 24 % 16 != 0 + with pytest.raises(ValueError): + quantize_nvfp4_weight(torch.randn(4, 16, 2)) + + +# --------------------------------------------------------------------------- +# End-to-end: GPTQ quantize -> save -> reload -> forward +# --------------------------------------------------------------------------- + +_CALIBRATION_TEXTS = [ + "tiny nvfp4 calibration sample one with enough tokens to survive minimum length filtering", + "tiny nvfp4 calibration sample two exercising the gptq loop over the fp4 grid cleanly", + "another synthetic calibration example that is intentionally verbose so filtering keeps it", +] * 2 + + +def _build_local_tokenizer(model_dir: Path) -> PreTrainedTokenizerFast: + tokenizer = Tokenizer(WordLevel(unk_token="[UNK]")) + tokenizer.pre_tokenizer = Whitespace() + trainer = WordLevelTrainer(special_tokens=["[PAD]", "[UNK]", "[BOS]", "[EOS]"]) + tokenizer.train_from_iterator(_CALIBRATION_TEXTS, trainer=trainer) + + fast_tokenizer = PreTrainedTokenizerFast( + tokenizer_object=tokenizer, + bos_token="[BOS]", + eos_token="[EOS]", + unk_token="[UNK]", + pad_token="[PAD]", + ) + fast_tokenizer.save_pretrained(model_dir) + return fast_tokenizer + + +def _build_tiny_llama_fixture(model_dir: Path) -> tuple[LlamaConfig, PreTrainedTokenizerFast]: + config = LlamaConfig( + num_hidden_layers=2, + hidden_size=64, + intermediate_size=128, + num_attention_heads=4, + num_key_value_heads=4, + vocab_size=128, + pad_token_id=0, + bos_token_id=1, + eos_token_id=2, + ) + model = LlamaForCausalLM(config) + model.save_pretrained(model_dir) + tokenizer = _build_local_tokenizer(model_dir) + return config, tokenizer + + +def _build_calibration_dataset(tokenizer: PreTrainedTokenizerFast) -> list[dict[str, object]]: + dataset = [] + for text in _CALIBRATION_TEXTS: + encoded = tokenizer(text, return_tensors="pt") + dataset.append( + { + "input_ids": encoded["input_ids"], + "attention_mask": encoded["attention_mask"], + } + ) + return dataset + + +@pytest.mark.slow +def test_gptq_nvfp4_end_to_end(tmp_path: Path): + model_dir = tmp_path / "native" + quantized_dir = tmp_path / "quantized" + + config, tokenizer = _build_tiny_llama_fixture(model_dir) + calibration = _build_calibration_dataset(tokenizer) + + quantize_config = QuantizeConfig( + bits=4, + group_size=16, + format="nvfp4", + device="cpu", + ) + assert isinstance(quantize_config, GPTQConfig) + + model = GPTQModel.load( + str(model_dir), + quantize_config=quantize_config, + backend=BACKEND.TORCH, + ) + model.quantize( + calibration, + batch_size=1, + backend=BACKEND.TORCH, + calibration_data_min_length=1, + ) + model.save(quantized_dir) + + # Checkpoint layout: packed uint8 fp4 weight + float8_e4m3fn weight_scale. + shards = sorted(quantized_dir.glob("*.safetensors")) + assert shards, "no safetensors shard written" + weight_key = "model.layers.0.self_attn.q_proj.weight" + scale_key = "model.layers.0.self_attn.q_proj.weight_scale" + with safe_open(shards[0], framework="pt", device="cpu") as f: + keys = set(f.keys()) + assert weight_key in keys and scale_key in keys, sorted(keys) + weight = f.get_tensor(weight_key) + scale = f.get_tensor(scale_key) + hidden = config.hidden_size + assert weight.dtype is torch.uint8 + assert weight.shape == (hidden, hidden // 2) + assert scale.dtype is torch.float8_e4m3fn + assert scale.shape == (hidden, hidden // BLOCK) + + # The saved tensors must decode with the torchao NVFP4 layout. + _ = _nvfp4_dequant(weight, scale) + + reloaded = GPTQModel.load(str(quantized_dir), device="cpu") + modules = dict(reloaded.named_modules()) + for suffix in ("q_proj", "k_proj", "v_proj", "o_proj"): + name = f"model.model.layers.0.self_attn.{suffix}" + assert isinstance(modules[name], TorchNVFP4Linear), name + for suffix in ("gate_proj", "up_proj", "down_proj"): + name = f"model.model.layers.0.mlp.{suffix}" + assert isinstance(modules[name], TorchNVFP4Linear), name + + encoded = tokenizer("nvfp4 forward smoke", return_tensors="pt") + with torch.inference_mode(): + out = reloaded.model(input_ids=encoded["input_ids"]) + assert out.logits.shape[-1] == config.vocab_size + assert torch.isfinite(out.logits).all() + + +@pytest.mark.slow +def test_weight_only_nvfp4_end_to_end(tmp_path: Path): + """RTN-style direct NVFP4 export (`method=nvfp4`) without calibration.""" + + model_dir = tmp_path / "native" + quantized_dir = tmp_path / "quantized" + + config, tokenizer = _build_tiny_llama_fixture(model_dir) + + quantize_config = QuantizeConfig(quant_method="nvfp4", device="cpu") + assert isinstance(quantize_config, NVFP4Config) + + model = GPTQModel.load( + str(model_dir), + quantize_config=quantize_config, + backend=BACKEND.TORCH, + ) + model.quantize(calibration=None, backend=BACKEND.TORCH) + model.save(quantized_dir) + + reloaded = GPTQModel.load(str(quantized_dir), device="cpu") + modules = dict(reloaded.named_modules()) + name = "model.model.layers.0.mlp.down_proj" + assert isinstance(modules[name], TorchNVFP4Linear), name + + encoded = tokenizer("nvfp4 weight only smoke", return_tensors="pt") + with torch.inference_mode(): + out = reloaded.model(input_ids=encoded["input_ids"]) + assert torch.isfinite(out.logits).all() + + +def test_primed_quantizer_matches_torchao_two_level_grid(): + """With a primed global scale, find_params must land on torchao's two-level + grid -- the fix for 96.8% of block scales floor-clamping at E4M3's 2^-6 on + realistic weight magnitudes (1.4x reconstruction error vs the donor RTN).""" + + torch.manual_seed(3) + qcfg = QuantizeConfig(bits=4, group_size=16, format="nvfp4", nvfp4_scale_sweep=0) + quantizer = NVFP4Quantizer(qcfg=qcfg) + quantizer.configure(perchannel=True) + + W = torch.randn(32, 64, dtype=torch.float32) * 0.02 + gs = quantizer.prime_global_scale(W) + assert float(gs) > 0 and (torch.log2(gs) == torch.log2(gs).round()), "global scale must be a power of two" + + # collect the loop's per-group effective scales + ours = [] + for g in range(0, 64, BLOCK): + quantizer.find_params(W[:, g : g + BLOCK], weight=True) + ours.append(quantizer.scale.flatten()) + ours = torch.stack(ours, dim=1) + + ref_scales, _ = nvfp4_quantize(W.contiguous(), block_size=BLOCK, per_tensor_scale=gs.reshape(())) + ref = ref_scales.to(torch.float32) * gs + assert torch.equal(ours, ref), "primed find_params diverged from torchao two-level scales" + + # the floor-clamp regression itself: on these magnitudes the one-level path + # pins ~100% of blocks at 2^-6; two-level must not. + floored = (ours == torch.finfo(torch.float8_e4m3fn).tiny).float().mean() + assert floored < 0.5, f"{floored:.0%} of scales at the E4M3 floor -- still single-level" + + +def test_primed_scales_pack_losslessly_with_recovered_global_scale(): + """pack must re-derive exactly the primed global scale from the effective + scales, so the two-level split reproduces the loop's grid bit-for-bit.""" + + torch.manual_seed(4) + IN, OUT = 128, 32 + qcfg = QuantizeConfig(bits=4, group_size=16, format="nvfp4") + quantizer = NVFP4Quantizer(qcfg=qcfg) + quantizer.configure(perchannel=True) + + W = torch.randn(OUT, IN, dtype=torch.float32) * 0.02 + gs = quantizer.prime_global_scale(W) + scales = [] + for g in range(0, IN, BLOCK): + quantizer.find_params(W[:, g : g + BLOCK], weight=True) + scales.append(quantizer.scale.reshape(-1, 1)) + scales = torch.cat(scales, dim=1) + + module = TorchNVFP4Linear( + bits=4, group_size=BLOCK, sym=True, desc_act=False, + in_features=IN, out_features=OUT, bias=False, + ) + packed, block_scale, ws2 = module._pack_with_gptq_scales(W, scales) + # The invariant is exactness of the split, not equality of the decomposition: + # if the max block ratio E4M3-rounds down to exactly 224, pack legitimately + # recovers gs/2 -- the product is identical either way. + assert torch.equal(block_scale.to(torch.float32) * ws2, scales), "two-level split not exact" + lg = torch.log2(ws2.reshape(())) + assert lg == lg.round(), "recovered global scale must stay a power of two" + assert float(lg) in (float(torch.log2(gs)), float(torch.log2(gs)) - 1.0), ( + "recovered global scale should be gs or gs/2 (boundary rounding)" + ) + + +def test_scale_sweep_beats_amax_scales(): + """The per-block E4M3 scale sweep must never lose to amax-derived scales + (amax is in the candidate set) and should measurably win on realistic + magnitudes -- the ModelOpt-RTN reference gets ~13% from exactly this.""" + + torch.manual_seed(5) + W = torch.randn(64, 128, dtype=torch.float32) * 0.02 + + errs = {} + for sweep in (0, 8): + qcfg = QuantizeConfig(bits=4, group_size=16, format="nvfp4", nvfp4_scale_sweep=sweep) + quantizer = NVFP4Quantizer(qcfg=qcfg) + quantizer.configure(perchannel=True) + quantizer.prime_global_scale(W) + out = torch.empty_like(W) + for g in range(0, 128, BLOCK): + block = W[:, g : g + BLOCK] + quantizer.find_params(block, weight=True) + out[:, g : g + BLOCK] = quantizer.quantize(block) + errs[sweep] = ((out - W).norm() / W.norm()).item() + + assert errs[8] <= errs[0] + 1e-9, f"sweep lost: {errs}" + assert errs[8] < errs[0] * 0.97, f"sweep gained <3%: {errs}" diff --git a/tests/test_w4a4_hessian.py b/tests/test_w4a4_hessian.py new file mode 100644 index 000000000..97ffac8ee --- /dev/null +++ b/tests/test_w4a4_hessian.py @@ -0,0 +1,292 @@ +# SPDX-FileCopyrightText: 2026 ModelCloud.ai +# SPDX-License-Identifier: Apache-2.0 + +"""w4a4: the GPTQ Hessian must be accumulated on quantized activations. + +Under w4a4 the kernel consumes ``Q(X)``, so a Hessian built from ``X`` solves for +inputs that never occur. The load-bearing test here is +``test_hessian_differs_from_clean_activations``: the natural way to get this +wrong is to capture activations with a ``forward_pre_hook``, which fires before +any input quantizer and therefore silently accumulates clean activations. That +bug leaves every config flag looking correct, so it is only detectable by +asserting the Hessian actually changed. + +See `docs/w4a4_hessian.md`. +""" + +import json +import os + + +# Keep this test on CPU so it works on CPU-only runners. +os.environ.setdefault("CUDA_VISIBLE_DEVICES", "") + +import pytest +import torch +import torch.nn as nn + +from gptqmodel.quantization import QuantizeConfig +from gptqmodel.quantization.config import FORMAT +from gptqmodel.quantization.dtype import fake_quantize_nvfp4 +from gptqmodel.quantization.gptq import GPTQ + + +IN_FEATURES = 64 +OUT_FEATURES = 32 +BLOCK = 16 + + +def _qcfg(act: bool) -> QuantizeConfig: + kwargs = dict(bits=4, group_size=BLOCK, format="nvfp4", sym=True, desc_act=False) + if act: + kwargs.update(act_format=FORMAT.NVFP4, act_group_size=BLOCK) + return QuantizeConfig(**kwargs) + + +def _hessian_for(act: bool, inp: torch.Tensor) -> torch.Tensor: + torch.manual_seed(0) + module = nn.Linear(IN_FEATURES, OUT_FEATURES, bias=False) + gptq = GPTQ(module=module, qcfg=_qcfg(act)) + gptq.add_batch(inp, torch.empty(0)) + return gptq.finalize_hessian().clone() + + +@pytest.fixture(scope="module") +def calib_input() -> torch.Tensor: + torch.manual_seed(1234) + return torch.randn(4, 128, IN_FEATURES) + + +def test_fake_quantize_nvfp4_round_trip(): + """Quantize+dequantize preserves layout but lands values on the FP4 grid.""" + + torch.manual_seed(0) + x = torch.randn(8, IN_FEATURES) + q = fake_quantize_nvfp4(x, block_size=BLOCK) + + assert q.shape == x.shape + assert q.dtype == x.dtype + assert torch.isfinite(q).all() + assert not torch.equal(q, x), "fake quantization was a no-op" + # E2M1 has 16 codes, so one block can never hold more distinct values. + assert len(torch.unique(q[0, :BLOCK])) <= 16 + + +def test_fake_quantize_nvfp4_rejects_ragged_last_dim(): + """A last dim that is not block-divisible must raise, not silently mis-block.""" + + with pytest.raises(ValueError, match="multiple of block_size"): + fake_quantize_nvfp4(torch.randn(4, BLOCK + 1), block_size=BLOCK) + + +def test_weight_only_config_leaves_activations_untouched(): + """Without `act_format`, no activation quantizer exists at all.""" + + gptq = GPTQ(module=nn.Linear(IN_FEATURES, OUT_FEATURES, bias=False), qcfg=_qcfg(act=False)) + assert gptq._act_quantizer is None + + +def test_act_format_installs_quantizer(): + gptq = GPTQ(module=nn.Linear(IN_FEATURES, OUT_FEATURES, bias=False), qcfg=_qcfg(act=True)) + assert gptq._act_quantizer is not None + + +def test_hessian_differs_from_clean_activations(calib_input): + """The regression that matters: `act_format` must actually change the Hessian. + + A hook-based implementation would see pre-quantizer activations and produce a + Hessian identical to the weight-only one while appearing configured. + """ + + h_clean = _hessian_for(act=False, inp=calib_input) + h_quant = _hessian_for(act=True, inp=calib_input) + + assert h_clean.shape == h_quant.shape + assert not torch.allclose(h_clean, h_quant), ( + "Hessian is identical with and without activation quantization — " + "activations are not being quantized before accumulation" + ) + rel = ((h_clean - h_quant).norm() / h_clean.norm()).item() + assert rel > 1e-3, f"Hessian changed by only {rel:.2%}; expected a visible FP4 effect" + + +def test_quantized_hessian_is_well_formed(calib_input): + """Quantizing activations must not break the properties the solve relies on.""" + + h = _hessian_for(act=True, inp=calib_input) + + assert torch.isfinite(h).all() + assert torch.allclose(h, h.T, atol=1e-4), "Hessian lost symmetry" + assert (torch.diag(h) > 0).all(), "Hessian gained a dead column under activation quantization" + + +def test_hessian_is_deterministic(calib_input): + """Same inputs must give the same Hessian — the solve is not robust to drift.""" + + assert torch.equal(_hessian_for(act=True, inp=calib_input), _hessian_for(act=True, inp=calib_input)) + + +@pytest.mark.parametrize( + "kwargs, match", + [ + (dict(bits=4, group_size=128, format="gptq", act_format=FORMAT.NVFP4), "requires `format=nvfp4`"), + (dict(bits=4, group_size=BLOCK, format="nvfp4", act_format=FORMAT.NVFP4, act_group_size=0), "act_group_size"), + (dict(bits=4, group_size=BLOCK, format="nvfp4", act_amax_path="x.json"), "require `act_format`"), + ], +) +def test_config_rejects_unsupported_activation_settings(kwargs, match): + with pytest.raises(ValueError, match=match): + QuantizeConfig(**kwargs) + + +# --------------------------------------------------------------- amax sidecar -- + + +def _sidecar(tmp_path, payload) -> str: + path = tmp_path / "amax_per_layer.json" + path.write_text(json.dumps(payload)) + return str(path) + + +def _gptq_named(name: str, qcfg: QuantizeConfig) -> GPTQ: + torch.manual_seed(0) + module = nn.Linear(IN_FEATURES, OUT_FEATURES, bias=False) + gptq = GPTQ(module=module, qcfg=qcfg) + gptq.name = name + return gptq + + +def _qcfg_sidecar(path, rules=None) -> QuantizeConfig: + return QuantizeConfig( + bits=4, group_size=BLOCK, format="nvfp4", sym=True, desc_act=False, + act_format=FORMAT.NVFP4, act_group_size=BLOCK, + act_amax_path=path, act_amax_key_rules=rules, + ) + + +def test_global_scale_matches_pipeline_formula(tmp_path): + """global_scale must be amax / 6 / 448 — the NVFP4 second-level scale.""" + + amax = 12.0 + path = _sidecar(tmp_path, {"layer0.moe_input": amax}) + gptq = _gptq_named("layer0.moe_input", _qcfg_sidecar(path)) + + scale = gptq._resolve_act_global_scale() + assert scale is not None + assert scale.item() == pytest.approx(amax / 6.0 / 448.0) + + +def test_sidecar_accepts_nested_amax_mapping(tmp_path): + """The pipeline emits {'amax': {...}}; a flat mapping must work too.""" + + path = _sidecar(tmp_path, {"amax": {"layer0.moe_input": 6.0}, "meta": {"ignored": 1}}) + gptq = _gptq_named("layer0.moe_input", _qcfg_sidecar(path)) + assert gptq._resolve_act_global_scale().item() == pytest.approx(6.0 / 6.0 / 448.0) + + +def test_key_rules_share_one_scale_across_sibling_modules(tmp_path): + """w1 and w3 consume the same MoE input, so they must resolve to one entry. + + Getting this wrong gives siblings different activation scales than the kernel + applies, which is exactly the mismatch the sidecar exists to prevent. + """ + + path = _sidecar(tmp_path, {"layer7.moe_input": 8.0, "layer7.w2_input": 2.0}) + rules = [ + (r"^layers\.(\d+)\..*\.(?:w1|w3)$", r"layer\1.moe_input"), + (r"^layers\.(\d+)\..*\.w2$", r"layer\1.w2_input"), + ] + cfg = _qcfg_sidecar(path, rules) + + w1 = _gptq_named("layers.7.ffn.experts.3.w1", cfg)._resolve_act_global_scale() + w3 = _gptq_named("layers.7.ffn.experts.3.w3", cfg)._resolve_act_global_scale() + w2 = _gptq_named("layers.7.ffn.experts.3.w2", cfg)._resolve_act_global_scale() + + assert w1.item() == w3.item(), "w1 and w3 share an input and must share a scale" + assert w2.item() != w1.item(), "w2 sees the post-SwiGLU activation and needs its own" + assert w2.item() == pytest.approx(2.0 / 6.0 / 448.0) + + +def test_missing_amax_key_fails_loud(tmp_path): + """A missing entry must raise, not silently fall back to block-only scales.""" + + path = _sidecar(tmp_path, {"layer0.moe_input": 6.0}) + gptq = _gptq_named("layers.99.ffn.experts.0.w1", _qcfg_sidecar(path)) + with pytest.raises(KeyError, match="not found"): + gptq._resolve_act_global_scale() + + +@pytest.mark.parametrize("bad", [0.0, -1.0, float("inf"), float("nan")]) +def test_invalid_amax_rejected(tmp_path, bad): + path = _sidecar(tmp_path, {"k": bad}) + gptq = _gptq_named("k", _qcfg_sidecar(path)) + with pytest.raises(ValueError, match="finite and positive"): + gptq._resolve_act_global_scale() + + +def test_key_rules_resolve_against_named_module_full_name(tmp_path): + """The looper names modules layer-relative (`mlp.experts.1.gate_proj`); + only NamedModule.full_name carries the layer index the key rules capture. + Resolving against `self.name` fails on every real looper run.""" + + from types import SimpleNamespace + + path = _sidecar(tmp_path, {"layer3.moe_input": 6.0}) + rules = [(r".*\.layers\.(\d+)\.mlp\.experts\.\d+\.(?:gate_proj|up_proj)$", r"layer\1.moe_input")] + gptq = _gptq_named("mlp.experts.1.gate_proj", _qcfg_sidecar(path, rules)) + gptq._named_module = SimpleNamespace( + full_name="model.language_model.layers.3.mlp.experts.1.gate_proj" + ) + assert gptq._resolve_act_global_scale().item() == pytest.approx(6.0 / 6.0 / 448.0) + + +def test_no_sidecar_leaves_global_scale_unset(): + gptq = _gptq_named("anything", _qcfg(act=True)) + assert gptq._resolve_act_global_scale() is None + + +def test_global_scale_changes_the_hessian(tmp_path, calib_input): + """The sidecar must actually reach the quantizer, not just parse.""" + + torch.manual_seed(0) + inp = calib_input + + def hessian(cfg, name="k"): + torch.manual_seed(0) + g = GPTQ(module=nn.Linear(IN_FEATURES, OUT_FEATURES, bias=False), qcfg=cfg) + g.name = name + g._act_quantizer = g._create_act_quantizer() + g.add_batch(inp, torch.empty(0)) + return g.finalize_hessian().clone() + + block_only = hessian(_qcfg(act=True)) + # An amax far from the data's own range forces a visibly different grid. + with_scale = hessian(_qcfg_sidecar(_sidecar(tmp_path, {"k": 40.0}))) + + assert not torch.allclose(block_only, with_scale), ( + "global scale had no effect on the Hessian — sidecar is not reaching the quantizer" + ) + + +def test_dynamic_fallback_override_survives_caller_baseline(): + """base.py threads the global quantize_config.fallback into loop(); it must be + the baseline, not a clobber -- per-module dynamic `fallback` overrides (the + per-projection rank thresholds) previously never took effect.""" + + from gptqmodel.looper.gptq_processor import clone_gptq_config_for_module + from gptqmodel.quantization.config import Fallback, FallbackStrategy + + qcfg = QuantizeConfig( + bits=4, group_size=BLOCK, format="nvfp4", sym=True, desc_act=False, + dynamic={r".*\.down_proj$": {"fallback": {"strategy": "rtn", "threshold": 3072}}}, + fallback=Fallback(strategy=FallbackStrategy.RTN, threshold=6144), + ) + caller_baseline = Fallback(strategy=FallbackStrategy.RTN, threshold=6144) + + clone = clone_gptq_config_for_module( + qcfg, "model.layers.3.mlp.experts.0.down_proj", fallback=caller_baseline) + assert clone.fallback.threshold == 3072, "dynamic override clobbered by caller fallback" + + clone2 = clone_gptq_config_for_module( + qcfg, "model.layers.3.mlp.experts.0.gate_proj", fallback=caller_baseline) + assert clone2.fallback.threshold == 6144, "non-overridden module must keep the baseline"