Repository navigation
[OMNIML-5899] Export IQ checkpoints from HF and Megatron - #2447
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThe export system now supports GGML IQ1_S and IQ2_XS formats. It validates IQ metadata, packs Hugging Face and Megatron weights, rejects unsupported layouts, documents payload contracts, and adds unit and GPU coverage. ChangesGGML IQ export
Priority: ⬇️ Low Estimated code review effort: 3 (Moderate) | ~25 minutes Change: Feature Sequence Diagram(s)sequenceDiagram
participant MegatronExporter
participant CUDAIQQuantizer
participant ExportState
MegatronExporter->>CUDAIQQuantizer: Pack IQ1_S or IQ2_XS weight
CUDAIQQuantizer->>ExportState: Return CPU uint8 payload
MegatronExporter->>ExportState: Store packed weight and IQ metadata
Merge Risk: 🟠 High · up to Do not merge yet: the export package cannot load until the GGML provider is included, and mixed-format tensor-parallel exports can produce unsupported IQ payloads instead of being rejected. 🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
Full details: Docstring CoverageExplanation Docstring coverage is 31.78% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 129 functions across 25 files. (1 skipped: 1 unsupported.) ✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Comment |
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| # Save the merged weights | ||
| if merged_weight_scale is None: | ||
| if qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): | ||
| self._state_dict.update(self._get_iq_weight_state(prefix, merged_weight, qformat)) |
There was a problem hiding this comment.
[CRITICAL Algorithm] IQ blocks are packed along the output axis here, not the input axis.
What's wrong: three lines above (line 1887) merged_weight is transposed from Megatron's [E, out, in] to HF's [E, in, out]. _get_iq_weight_state → quantize_iq{1_s,2_xs} always forms its 256-element GGML blocks along weight.shape[-1] (see quantize_iq2_xs: blocks = weight.contiguous().reshape(-1, 256), packed_shape = (*weight.shape[:-1], weight.shape[-1] // 256, bytes)). After the transpose that last axis is out_features, so each super-block groups 256 output channels instead of 256 contiguous input elements.
Same defect at line 2023 in _pack_name_remapping_gpt_oss, where the transpose at line 1963 is unconditional (and for linear_fc1 the gate/up interleave then permutes the last axis again before packing).
Why it matters:
- It silently disagrees with calibration.
iq2_xs_fake_quant/iq1_s_fake_quantcallquantize_iq*(inputs)on the module's own[out, in]weight, so PTQ evaluated blocks alongin. The exported checkpoint therefore does not reproduce the model that was measured — accuracy loss with no error raised wheneverout_features % 256 == 0. - When
out_features % 256 != 0(e.g. anffn_hidden_sizeof 1408 or 5120·k that isn't a multiple of 256)validate_weightraises a confusing "requires the last weight dimension to be divisible by 256" error naming a shape the user never configured. - The dense paths (
_populate_state_dict, the qkv/gate-up splits) pack the untransposed[out, in]weight, so fused-MoE experts end up with a different block axis than every other layer in the same checkpoint.
Note the HF exporter already establishes the right convention for exactly this case: _export_quantized_weight wraps BMM-expert packing in maybe_transpose_expert_weight_dimensions(...) → quantize → transpose back, specifically to keep blocks on the contraction axis.
Suggested fix: pack before the layout transpose, i.e. in both methods pack the stacked [E, out, in] tensor and only then move the payload into HF order (the packed tensor is [E, out, in//256, bytes], so the logical transpose has to be applied to merged_weight first and the packing deferred, or the pack has to be done on the pre-transpose tensor and the consumer told the blocks live on in). Concretely, for _pack_name_remapping:
merged_weight = torch.stack(weight_list, dim=0) # [E, out, in]
if qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS):
# GGML blocks must run along the contraction (in_features) axis, which is
# merged_weight's last dim *before* the HF [E, in, out] transpose. Pack here
# and skip the transpose; a packed payload cannot be transposed afterwards.
self._state_dict.update(self._get_iq_weight_state(prefix, merged_weight, qformat))
return
# Megatron is [num_experts, out, in]; most HF layouts want [num_experts, in, out].
if transpose:
merged_weight = merged_weight.transpose(-2, -1).contiguous()
...and mirror it in _pack_name_remapping_gpt_oss (bias/input_scale handling still needs to run, so hoist the pack rather than early-returning there). Whichever convention you land on, please state the block axis for fused MoE experts explicitly in docs/source/deployment/3_unified_hf.rst — the current "IQ weight representation" section only describes the dense [out, in] case.
There was a problem hiding this comment.
Fixed in f2d3a834a. Fused experts are now packed per expert while each logical weight is still [out, in], so 256-value blocks remain on the contraction axis. Packed CPU payloads are stacked without a post-pack transpose. Exact packed-byte coverage was added for the standard and GPT-OSS expert paths.
| delattr(sub_module, weight_name) | ||
| sub_module.register_buffer("weight", packed_weight) |
There was a problem hiding this comment.
[IMPORTANT Compatibility] This is the only export format that turns weight from a Parameter into a buffer.
What's wrong: every other branch of this function ends at line 846 with setattr(sub_module, weight_name, nn.Parameter(quantized_weight, requires_grad=False)), and that already covers non-float packed payloads — pack_int4_in_uint8 returns uint8 and INT4_AWQ stores it as a Parameter (integer dtypes are legal for nn.Parameter as long as requires_grad=False). The IQ path instead does delattr + register_buffer, silently changing the module contract for this one format.
Why it matters: it breaks any consumer that enumerates weights via named_parameters() / isinstance(m.weight, nn.Parameter) rather than state_dict(). The PR already had to add a compensating branch in moe_utils.py:221 for exactly this reason, and the same guard will be needed at every future call site. TiedWeightMap and sync_tied_input_amax (quant_utils.py:1852) both gate on isinstance(m.weight, torch.nn.Parameter); they happen to be built in _prepare_model_for_export before packing today, so nothing is broken right now, but the invariant is one reordering away from a silent tied-weight dedup miss.
Suggested fix: keep the weight a Parameter so the module contract is uniform across formats, and revert the moe_utils.py special case:
packed_weight, _ = quantize_iq(weight.to(dtype))
setattr(sub_module, weight_name, nn.Parameter(packed_weight, requires_grad=False))
maybe_clear_cuda_cache()
return(This also drops the delattr/hardcoded-"weight" asymmetry — the weight_name != "weight" guard above then only needs to exist for the quantizer_attr_names bookkeeping, not for the assignment.)
There was a problem hiding this comment.
Fixed in f2d3a834a. IQ export now replaces weight with nn.Parameter(..., requires_grad=False), matching the other packed formats. The fused-expert buffer special case and its test were removed.
| module: torch.nn.Module, | ||
| dtype: torch.dtype = torch.float16, | ||
| name_to_value: dict[str, torch.Tensor] | None = None, | ||
| keep_weight_device: bool = False, |
There was a problem hiding this comment.
[IMPORTANT Performance] keep_weight_device=True moves the fused-MoE stack from CPU onto the GPU and can add several GB to the export rank's peak.
What's wrong: for IQ, _get_quantized_state now returns an on-GPU weight (which is correct for the dense paths — _get_iq_weight_state packs and immediately .detach().cpu()s, so the GPU copy is transient). But _pack_name_remapping / _pack_name_remapping_gpt_oss accumulate weight_list across all experts and then do torch.stack(...) followed by .transpose(-2, -1).contiguous(). Previously every entry was a CPU copy and both temporaries lived in host memory; now the stack and the contiguous transposed copy are two full-size CUDA allocations on top of the expert weights that are already resident as model parameters (module.weight.to(dtype) is a no-op alias when the dtype already matches, so the list itself is free — the two temporaries are the cost).
Why it matters: for a large MoE decoder layer this is ~2× the fused expert tensor in extra device memory, per layer, on the last-PP-stage rank that is also holding the model. E.g. 128 experts × 1408 × 7168 in bf16 is ≈2.6 GB fused, so ≈5 GB of new transient device memory — enough to OOM an export that previously fit. Since save_pretrained already restricts IQ to TP=1, this rank has no TP sharding to shrink it.
Suggested fix: two options, either is fine:
- Pack per expert before stacking (each
[out, in]expert packs independently to[out, in//256, bytes], thentorch.stackthe small uint8 payloads on CPU) — this composes naturally with the block-axis fix on the_get_iq_weight_statecall and removes the transposed float copy entirely. - Or scope
keep_weight_deviceto the callers that actually benefit (the dense_populate_state_dict/ split paths) and leave the packed-expert paths on CPU, since_pack_name_remappingis the one place where the on-device weight is held rather than consumed immediately.
Worth a note in the keep_weight_device docstring either way: as written it reads as a neutral device toggle, but it shifts a whole-layer allocation from host to device for MoE.
There was a problem hiding this comment.
Fixed in f2d3a834a. Each expert is packed immediately and moved to CPU before stacking. The export no longer builds a full fused floating-point stack and contiguous transpose on GPU.
| # (pre_quant_scale is the AWQ / NVFP4_AWQ / SVDQuant companion, renamed in the KV-cache pass.) | ||
| weight_suffixes = ( | ||
| "weight", | ||
| "weight_shape", |
There was a problem hiding this comment.
[SUGGESTION] weight_shape is added to the tied-weight suffix list but IQ export never emits it, so this entry is dead for the format that motivated it.
_get_iq_weight_state and the HF _export_quantized_weight IQ branch both discard the second return of quantize_iq* (packed_weight, _ = ...), and the docs section added in this PR states that no shape tensor is stored. The only producer of a weight_shape key is the compressed_tensors CompressedLinear path in quantization/plugins/huggingface.py:1282. That makes the addition a harmless consistency fix for that path, but it reads as if IQ exports a shape companion.
Related, and the reason it's worth a look: dequantize_iq1_s / dequantize_iq2_xs take weight_shape as a required argument, so ModelOpt's own decoder cannot read back a checkpoint that ModelOpt just wrote without the caller re-deriving the shape by hand. Consider adding a small helper next to the packers, e.g.
def iq_logical_shape(packed_weights: torch.Tensor) -> torch.Tensor:
"""Recover the logical weight shape from a packed IQ payload."""
return torch.tensor(
(*packed_weights.shape[:-2], packed_weights.shape[-2] * GGML_BLOCK_SIZE),
dtype=torch.int64,
)so the documented recovery rule lives in one place that both the docs and dequantize_iq* callers can point at, instead of being restated prose-only in 3_unified_hf.rst. (If you'd rather just drop the "weight_shape" line here, that's fine too — but then the comment above about extending weight_suffixes should say which path produces it.)
There was a problem hiding this comment.
The dead weight_shape suffix was removed in f2d3a834a. The proposed logical-shape helper changes the codec API and belongs in #2446, so it is intentionally kept out of this export PR.
| elif quant_algo in ("IQ1_S", "IQ2_XS"): | ||
| effective_bits, payload_bytes = (1.5625, 50) if quant_algo == "IQ1_S" else (2.3125, 74) | ||
| return { | ||
| "weights": { | ||
| "dynamic": False, | ||
| "num_bits": 1 if quant_algo == "IQ1_S" else 2, | ||
| "effective_bits": effective_bits, | ||
| "type": "int", | ||
| "group_size": 256, | ||
| "packing": "ggml", | ||
| "block_payload_bytes": payload_bytes, | ||
| } |
There was a problem hiding this comment.
[SUGGESTION] The IQ weights group carries three keys that aren't part of the schema this function is documented to emit, and it describes IQ as group-wise int quantization.
Per the docstring above, config_groups.group_0.* mirrors compressed-tensors' QuantizationArgs; every other branch here sticks to dynamic / num_bits / type / group_size / strategy. This branch adds effective_bits, packing, and block_payload_bytes.
Two things worth reconsidering:
"type": "int", "num_bits": 1, "group_size": 256claims a group-wise affine int scheme, but the checkpoint has noweight_scale/weight_zero_pointcompanion (all block metadata is inside the uint8 payload). A loader that dispatches onconfig_groups— rather than on the ModelOpt-specificquant_algo: "IQ1_S"— will accept this group and then try to decompressweightas group-quantized int, which cannot work. Sinceconvert_hf_quant_config_formatwrites this intoconfig.jsonasquantization_config, that's the config transformers actually sees.num_bitsandeffective_bitsdisagree by design (1 vs 1.5625). Anything that sizes buffers fromnum_bits × numelwill under-allocate.
Nothing here is load-breaking today given quant_algo is present and distinct, so this is non-blocking — but a short comment in this branch stating that IQ groups are not compressed-tensors-decodable and that packing: "ggml" is the discriminator a consumer must check would save the next reader a trip through quantize_iq2_xs. Nesting the ModelOpt-only fields (e.g. under a single "ggml": {...} sub-dict) would make that structural rather than a comment.
There was a problem hiding this comment.
Fixed in f2d3a834a. IQ metadata no longer presents the payload as a compressed-tensors integer weights group. Uniform exports carry ModelOpt-owned top-level format metadata, while mixed exports use an IQ group without a weights schema. Tests cover both forms.
There was a problem hiding this comment.
Claude review — IQ checkpoint export
Reviewed all 7 changed files (175 additions): the 6 modelopt/torch/export/ modules in full, plus the docs/source/deployment/3_unified_hf.rst contract section. Traced the new IQ format constants and _get_iq_weight_state through every _get_quantized_state caller in unified_export_megatron.py, and read the quantize_iq* / validate_weight implementations from the stacked quantization PR to check the packing contract. No prior Claude review on this PR.
Findings: CRITICAL: 1, IMPORTANT: 2, SUGGESTION: 2
Most impactful
1. Fused-MoE experts pack IQ blocks along the wrong axis (CRITICAL). _pack_name_remapping and _pack_name_remapping_gpt_oss transpose merged_weight from Megatron [E, out, in] to HF [E, in, out] and then call _get_iq_weight_state. quantize_iq* always blocks along shape[-1], so each 256-element GGML super-block groups output channels instead of contiguous input elements. Calibration (iq2_xs_fake_quant) packs the module's own [out, in] weight along in, so the exported checkpoint does not reproduce the model that was measured — silent accuracy loss whenever out_features % 256 == 0, and a misleading divisibility error when it isn't. The dense paths in this same file pack the untransposed weight, so fused experts also disagree with every other layer in the same checkpoint. The HF exporter already sets the correct precedent with maybe_transpose_expert_weight_dimensions (transpose, quantize, transpose back) for exactly this case.
2. IQ is the only format that demotes weight from Parameter to buffer (IMPORTANT). Every other branch of _export_quantized_weight ends with nn.Parameter(quantized_weight, requires_grad=False), and that already handles uint8 payloads (INT4_AWQ does it). The delattr + register_buffer here is what forced the compensating isinstance(wrapper.weight, nn.Parameter) branch in moe_utils.py; keeping it a Parameter makes that patch unnecessary and keeps named_parameters()-based consumers (TiedWeightMap, sync_tied_input_amax) working by construction rather than by ordering luck.
3. keep_weight_device=True shifts a whole-layer MoE allocation from host to device (IMPORTANT). Transient for the dense paths, but _pack_name_remapping holds weight_list across all experts and then allocates a GPU torch.stack plus a transposed .contiguous() copy — roughly 2x the fused expert tensor in new device memory per layer, on a rank that already holds the model and (per the new TP=1 restriction) has no sharding to shrink it.
The two SUGGESTIONs cover the unused weight_shape suffix vs. dequantize_iq*'s required weight_shape argument, and the non-standard keys in the compressed-tensors-shaped config_groups entry.
Verified as correct
quant_algoplumbing round-trips end to end:get_quantization_formatproduces"iq1_s",process_layer_quant_configuppercases to"IQ1_S", and the newconvert_hf_quant_config_formatbranch matches. The Megatron writer reaches the same code viaprocess_layer_quant_config(combined_layer_config_dict), sogroup_size/packing/block_payload_bytesland there too.- All seven
_get_quantized_statecall sites have a matching IQ branch — no path leaves an unpacked (or CUDA-resident) IQ weight in_state_dict. - The early return in
_get_quantized_statecorrectly skips amax/scale collection, andFUSION_FREE_FORMATSmembership is right for a weight-only format with no cross-module scales. - Excluded modules are safe:
qformat is Nonemakesis_iqfalse, sokeep_bf16/ excluded weights never hit the packer. Thekeep_bf16handling in the GatedDeltaNetin_projpath is the only site that needs it and it has it. - The documented shape-recovery rule
[*shape[:-2], shape[-2] * 256]is self-consistent for 1-D, 2-D and 3-D logical weights.
Minor note (not counted)
The TP guard in save_pretrained derives quantization_format from self._get_quantization_format(self.model), which is per-rank. If a PP stage ever returns None while another returns an IQ format, the non-IQ ranks skip the raise and proceed to torch.distributed.barrier() while the others have already thrown — an NCCL timeout instead of the clean NotImplementedError. Every PP stage of a transformer has linear layers in practice, so this is theoretical; gating on get_tensor_model_parallel_world_size() != 1 first, or all-reducing the format, would remove the possibility.
Risk
Moderate-to-high for MoE, low for dense. Finding 1 makes fused-expert IQ checkpoints numerically wrong without any error, and MoE is the main use case for aggressive 1-2 bit formats. The dense TP=1 path looks sound. Everything is additive and gated behind the two new format constants, so no existing format or checkpoint is affected.
🤖 Generated with Claude Code
c994ca2 to
c93a0ad
Compare
89b12cd to
d516738
Compare
c93a0ad to
305ad2e
Compare
d516738 to
0dda727
Compare
305ad2e to
1b4e81a
Compare
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (claude-opus-5) — DM the bot to share feedback.
Commenting: the Megatron packed-expert paths pack IQ blocks after the [E, out, in] -> [E, in, out] transpose, so blocks form along the wrong axis, and five of the six new Megatron branches have no test.
Needs action:
- Confirm the block axis in
_pack_name_remapping/_pack_name_remapping_gpt_oss(unified_export_megatron.py:1914): packing after the transpose groups 256 weights along the output dim, unlike fake-quant and the HF path. Pack before the transpose or explain why not. - Add tests for the untested IQ branches:
_gated_mlp_slicing,_grouped_mlp_slicing,_qkv_slicing,_gated_delta_net_slicing,_pack_name_remapping— only_name_remappingis covered. - Replace the hardcoded
50/74/1.5625/2.3125/256inquant_utils.py:733andconvert_hf_config.py:121withIQ*_BLOCK_BYTES/IQ*_EFFECTIVE_BITS/GGML_BLOCK_SIZEfrommodelopt.torch.quantization.ggml. - Drop the
"weight_shape"entry added toweight_suffixesinquant_utils.py:1170—_BASE_SKIP_KEYSalready removes anyweight_shapekey, and IQ export never emits one.
No action needed:
- Merge order: this stack sits on #2446, whose llama.cpp/ggml MIT provenance is still awaiting OSRB sign-off.
| @@ -1846,7 +1910,9 @@ def _pack_name_remapping(self, module, prefix, layer_type=None, is_mtp=False, tr | |||
| merged_input_scale = None | |||
|
|
|||
| # Save the merged weights | |||
| if merged_weight_scale is None: | |||
| if qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): | |||
| self._state_dict.update(self._get_iq_weight_state(prefix, merged_weight, qformat)) | |||
There was a problem hiding this comment.
Bot comment.
merged_weight has already been transposed to [E, in, out] above, so quantize_iq* blocks the last dim — the output features — into 256-element GGML blocks. Everywhere else (HF _export_quantized_weight, _name_remapping, _gated_mlp_slicing) the blocks run along the input dim, which is also what the weight_quantizer used during calibration/fake-quant. That means the exported experts are quantized on a different axis than the model was calibrated for, and out % 256 != 0 would raise from validate_weight even when in is a valid multiple. Please pack before the transpose (and re-pack/reshape into the HF layout), or document why the output-dim grouping is correct here. Same applies to _pack_name_remapping_gpt_oss, where the gate/up interleave also happens on that last dim before packing.
There was a problem hiding this comment.
Fixed in f2d3a834a by packing each expert before any layout transpose. For GPT-OSS, gate/up interleaving happens on the logical output rows before packing, then the CPU payloads are stacked. Exact-byte tests cover both expert paths.
| @@ -1152,6 +1167,7 @@ def _export_key(key: str) -> str: | |||
| # (pre_quant_scale is the AWQ / NVFP4_AWQ / SVDQuant companion, renamed in the KV-cache pass.) | |||
| weight_suffixes = ( | |||
| "weight", | |||
| "weight_shape", | |||
There was a problem hiding this comment.
Bot comment.
_BASE_SKIP_KEYS already contains "weight_shape", so any key containing it is dropped before the tied-weight dedup runs and can never appear in post_state_dict — the members filter if ak in post_state_dict makes this entry dead. IQ export also discards the logical-shape tensor (packed_weight, _ = quantize_iq(...)), which the new test asserts. Suggest removing this line.
There was a problem hiding this comment.
Removed in f2d3a834a. IQ export does not emit a separate weight_shape tensor, so the tied-weight suffix entry was dead.
| @@ -722,6 +729,14 @@ def process_layer_quant_config(layer_config_dict): | |||
| "quant_algo": "MXFP8", | |||
| "group_size": block_size_value, | |||
| } | |||
| elif v in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): | |||
| payload_bytes = 50 if v == QUANTIZATION_IQ1_S else 74 | |||
There was a problem hiding this comment.
Bot comment.
50 / 74 / 256 are re-derived here, again in convert_hf_config._quant_algo_to_group_config, and again in the docs. modelopt.torch.quantization.ggml already exports IQ1_S_BLOCK_BYTES, IQ2_XS_BLOCK_BYTES, IQ1_S_EFFECTIVE_BITS, IQ2_XS_EFFECTIVE_BITS and GGML_BLOCK_SIZE — importing them keeps the metadata in sync with the codec if a payload layout ever changes. Also note group_size is hardcoded to 256 here while the sibling branches use block_size_value from the recorded awq_block_size; asserting they agree would catch a misconfigured quantizer.
There was a problem hiding this comment.
Fixed in f2d3a834a. Export metadata now imports the per-format block-byte, block-size, and effective-bit constants from the codec package. It also verifies that the recorded block size matches the format constant, with a mismatch regression test.
### What does this PR do? Type of change: Code refactoring `#2448` added the GGML IQ packing kernels as **two** torch extensions, `modelopt_cuda_ext_iq1_s` and `modelopt_cuda_ext_iq2_xs`. This merges them into one, `modelopt_cuda_ext_ggml`. The existing per-extension split in `extensions.py` exists for reasons that don't apply to the IQ formats: `get_cuda_ext` gates on CUDA `>=11` while `_fp8`/`_mx` gate on `>=11.8`, and `_mx` needs `--use_fast_math`, which must not reach the base `tensor_quant` kernels. `get_cuda_ext_iq1_s` and `get_cuda_ext_iq2_xs` differed in none of that — same `>=11.8` gate, same `-O3` flags, same `common.cuh` — so the split only compiled the shared header twice, ran nvcc twice, and grew the loader, `__getattr__`, and `precompile()` once per format. With IQ2_XXS / IQ3_S / IQ4_NL plausibly following, that scales badly. Changes: - New `ggml/ggml.cpp` holds both host-side validation wrappers and the single `PYBIND11_MODULE`, binding `iq1_s_pack` and `iq2_xs_pack` (previously each module exported a bare `pack`). Deletes `ggml/iq1_s.cpp` and `ggml/iq2_xs.cpp`; the validation logic and docstrings carry over unchanged. - `get_cuda_ext_iq1_s` + `get_cuda_ext_iq2_xs` → `get_cuda_ext_ggml`, which builds `ggml.cpp`, `iq1_s.cu`, and `iq2_xs.cu` together. The retry-on-`raise_if_failed` semantics of the old getters are preserved. - Each format keeps its kernels in its own translation unit, so adding a format is a new `.cu` plus one `module.def` — no new extension, loader, or `precompile()` line. No caller outside `extensions.py` and its tests referenced the old getters on `main`, so nothing else changes. **Note for the follow-up PRs in the `#2448` series (`#2446`/`#2447`/`#2449`): the codec layer should call `get_cuda_ext_ggml().iq1_s_pack(...)` / `.iq2_xs_pack(...)` instead of `get_cuda_ext_iq1_s().pack(...)` / `get_cuda_ext_iq2_xs().pack(...)`.** ### Usage ```python from modelopt.torch.quantization.extensions import get_cuda_ext_ggml ext = get_cuda_ext_ggml(raise_if_failed=True) iq1_s_payload = ext.iq1_s_pack(weight, iq1s_grid) # uint8 [numel / 256, 50] iq2_xs_payload = ext.iq2_xs_pack(weight, iq2xs_grid, scales) # uint8 [numel / 256, 74] ``` ### Testing Ran on a single H200 NVL (TRT-LLM `1.3.0rc27.dev202609170000` container), building the merged extension from scratch: - `pytest tests/gpu/_extensions/test_torch_extensions.py` — **24 passed** (6:44). This is the full existing IQ suite (zero-block layout, encode, dtype rejection, row-straddling rejection, invalid/negative-zero scales, byte-exact dtype equivalence, and the brute-force optimality round-trip) reparametrized onto the merged module, plus the untouched `modelopt_cuda_ext` / `_fp8` / `_mx` load tests. - Verified `precompile()` loads all four extensions and that the merged module exports exactly `iq1_s_pack` and `iq2_xs_pack` with the expected arities. - Off-GPU: compiled the three sources directly and linked them into one `.so` to confirm no duplicate-symbol collisions between the two `.cu` translation units. - `pre-commit run --files ...` passes on all changed files (ruff, mypy, clang-format, bandit, license headers). ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ — the removed getters were added in `#2448` (merged today, unreleased) and have no callers outside this file's own tests. - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A — no new code or dependencies; the moved wrappers keep their original attribution. - Did you write any new necessary tests?: ✅ — existing coverage reparametrized onto the merged module; no behavior change to test. - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A — internal refactor of an unreleased, not-yet-wired-up API. - Did you get Claude approval on this PR?: ❌ — not yet run. ### Additional Information Follow-up to #2448. Merge before the remaining PRs in that series (#2446, #2447, #2449) land, so the codec layer is written against `get_cuda_ext_ggml` and no rename is needed afterwards. 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Added IQ1_S packing support through the GGML CUDA extension. - Added a unified GGML extension loader for IQ1_S and IQ2_XS packing. - Improved extension loading reliability when a cached extension is unavailable. - **Changes** - Renamed the IQ2_XS packing binding from `pack` to `iq2_xs_pack`. - Consolidated IQ1_S and IQ2_XS extension access under the shared GGML loader. - Updated GPU validation and coverage to use the unified extension interface. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
## Summary - add IQ1_S and IQ2_XS reference codecs and a weight-only fake-quant backend - register and export both formats from the quantization package - cache compact packed weights across unchanged forwards and invalidate on tensor or config changes - use one Python-side IQ2_XS FP16 scale predictor for both reference and CUDA packing - validate packed payload metadata, normalize CUDA cache keys, and define a shared non-finite policy ## PR split This work is split into four focused PRs. Each PR targets `main` and owns a disjoint file set: 1. **Kernel** — [#2448: Add CUDA kernels for IQ packing](#2448) 2. **Quantization** — [#2446: Add IQ quantization codecs and backend](#2446) 3. **Export** — [#2447: Export IQ checkpoints from HF and Megatron](#2447) 4. **Recipes** — [#2449: Add IQ post-training quantization recipes](#2449) The required merge order is #2448, #2446, #2447, then #2449. ## Scope This PR owns the Python codecs, backend dispatch, package registration, license attribution, CPU codec/backend tests, and CUDA numerical/reference-path tests. The native CUDA layer and direct extension tests remain in #2448; export and recipes remain in their own PRs. ## Why the codecs are separate from `qtensor` The new `ggml/` package contains stateless reference codecs and fake-quant backend functions. They transform ordinary tensors into packed format payloads and reconstruct tensors for fake quantization; they do not define persistent runtime quantized-tensor objects. `BaseQuantizedTensor` subclasses under `qtensor/` own runtime tensor objects and execution dispatch. Keeping the codecs separate avoids claiming a runtime tensor contract that these formats do not yet provide. A `qtensor` type can be added later if a runtime execution path requires one. ## Compatibility boundary The Python encoders intentionally use fixed-scale, unweighted searches. They are not intended to reproduce another encoder's bytes for every input when that encoder performs iterative scale refinement or importance weighting. Compatibility is defined by the canonical codebooks, 50/74-byte payload layouts, and pinned dequantization formulas. IQ2_XS computes the FP16 superblock scale once in the Python predictor and passes it to the CUDA packer. This removes a duplicate floating-point reduction and makes native/reference byte parity use the same scale. Non-finite input elements are treated as zero during packing in both implementations. The unit tests construct nonzero payload fields independently and validate metadata, signs, local scales, and global scales. The CUDA tests compare native packed bytes with this Python reference encoder. ## Test coverage - [IQ1_S CPU codec tests](https://github.com/NVIDIA/Model-Optimizer/blob/e8d937081d8cd01cf8e44d43915df443b79deb17/tests/unit/torch/quantization/test_iq1_s.py) - [IQ2_XS CPU codec tests](https://github.com/NVIDIA/Model-Optimizer/blob/e8d937081d8cd01cf8e44d43915df443b79deb17/tests/unit/torch/quantization/test_iq2_xs.py) - [registered backend and cache tests](https://github.com/NVIDIA/Model-Optimizer/blob/e8d937081d8cd01cf8e44d43915df443b79deb17/tests/unit/torch/quantization/test_ggml_backend.py) - [IQ1_S CUDA byte-parity, numerical, non-finite, zero-payload, and fallback tests](https://github.com/NVIDIA/Model-Optimizer/blob/e8d937081d8cd01cf8e44d43915df443b79deb17/tests/gpu/torch/quantization/test_iq1_s_cuda.py) - [IQ2_XS CUDA byte-parity, numerical, non-finite, zero-payload, underflow, and fallback tests](https://github.com/NVIDIA/Model-Optimizer/blob/e8d937081d8cd01cf8e44d43915df443b79deb17/tests/gpu/torch/quantization/test_iq2_xs_cuda.py) ## Licensing The embedded codebook data cites the pinned upstream MIT source, carries its license notice, and uses the repository's third-party license mechanism. Human OSRB/code-owner confirmation is still required; this PR does not claim that approval. ## Validation - focused lint, format, and type checks pass for all changed Python files - 36 focused CPU codec and backend tests pass locally - all 20 direct-extension and CUDA integration test cases collect locally; runtime CUDA execution remains delegated to GPU CI - restricted-term scan passes <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Added GGML quantization support for IQ1_S and IQ2_XS formats. - Added quantization, dequantization, and fake-quantization workflows with pass-through gradients. - Added CPU fallback when CUDA acceleration is unavailable. - Added validation for packed weights, tensor shapes, formats, and backend options. - Added configurable chunk processing and caching for repeated quantization. - **Tests** - Added comprehensive CPU and CUDA coverage for accuracy, validation, caching, fallback behavior, and edge cases. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com> Co-authored-by: Chenjie Luo <chenjiel@nvidia.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (claude-opus-5) — DM the bot to share feedback.
Commenting: all four items from the last round are still present in the code at this commit — the rebase onto a main that now carries modelopt/torch/quantization/ggml/ is the only visible change.
Needs action:
- Harden the TP>1 IQ guard at
unified_export_megatron.py:319:_get_quantization_format(self.model)returns only the first non-NONEformat, so a mixed-format model whose IQ layers come after e.g. an FP8 layer bypasses it and packs TP-sharded weights. Scan per-layer formats, or state why mixed IQ cannot occur. - Reject an enabled
input_quantizer/_pre_quant_scalein the IQ branch of_get_quantization_from_layer(quant_utils.py:488); both export paths return beforeinput_scalecollection and the existingpre_quant_scaleguard, so W-IQ + A-FP8 exports silently as weight-only. - Forward the caller's
group_sizeinto_quant_algo_to_group_configatconvert_hf_config.py:242so an invalid value raises instead of being rewritten to 256; add a test alongsidetest_iq_quantization_config_rejects_mismatched_block_size. - Reword the
IQ1_S/IQ2_XSbullets in3_unified_hf.rst:53— "importance-aware" implies imatrix support these encoders do not have.
No action needed:
- ✔️ Resolved earlier and still good: fused-MoE IQ rejection with tests on both entry points, and
ggml/now resolving on this branch so the export imports load.
Two of these let a wrong checkpoint out the door silently, which is the worst failure shape for a format nothing downstream validates. **TP>1 guard could be bypassed.** The Megatron guard keyed on ``_get_quantization_format(self.model)``, which by its own docstring returns only the first non-NONE format in the tree. A mixed-format model whose IQ layers follow, say, an FP8 one slipped past and packed TP-sharded weights as if they were whole. Add ``uses_iq_quantization``, which scans every layer, and guard on that. It reads ``num_bits`` directly rather than resolving each layer's full format, so an unrelated unsupported quantizer elsewhere in the model cannot turn the guard into an error. **W-IQ + A-FP8 exported as weight-only.** The IQ branch of ``_get_quantization_from_layer`` returned as soon as it validated the backend, while the neighbouring INT8 branch consults ``input_quantizer`` to tell SQ from WO. Both exporters then return before ``input_scale`` collection and before the ``pre_quant_scale`` handling, so an enabled activation quantizer vanished without a trace. Reject it instead; the GGML block payload has nowhere to put an activation scale. **group_size was not forwarded.** ``_quant_algo_to_group_config`` already validates it, and a test already covered that, but the top-level IQ branch in ``convert_hf_config`` called it without one, so a mismatched value was rewritten to the block size rather than rejected. The MIXED_PRECISION branch beside it was already correct. **Docs overstated the encoders.** "importance-aware" implies imatrix support these encoders do not have -- they neither refine the scale iteratively nor weight by importance, unlike upstream ``quantize_row_iq1_s_impl``. Call it codebook quantization. Four tests added. The two behavioural ones were mutation-checked: each fails with its fix reverted. That caught a first attempt at the group_size test which passed either way, because it exercised the MIXED_PRECISION path instead of the branch actually changed. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (claude-opus-5) — DM the bot to share feedback.
Commenting: all four items from the last round are fixed, but the new uses_iq_quantization helper reads num_bits off a quantizer that may be a SequentialQuantizer, which has no such attribute.
Needs action:
- Unwrap
SequentialQuantizer(or usegetattr(..., "num_bits", None)) inuses_iq_quantization(quant_utils.py:468) —save_pretrainednow calls it on every Megatron export, so a W4A8_AWQ model raisesAttributeErrorbefore any format dispatch. See inline. - Add a non-IQ
SequentialQuantizercase totest_uses_iq_quantization_false_without_iq_layers; the current coverage is FP8/IQTensorQuantizeronly. - Confirm or fix rank-uniformity of the new guard (
unified_export_megatron.py:321):uses_iq_quantization(self.model)is rank-local, so a PP stage with no IQ layer skips the raise and blocks in_gather_exclude_modules()while its peers exit. Checking TP size first, or all-reducing the flag, closes it.
No action needed:
- ✔️ Resolved since the last review: the TP>1 mixed-format bypass, the W-IQ + enabled-activation-quantizer silent weight-only export (plus
pre_quant_scale),group_sizeforwarding atconvert_hf_config.py:242with a test, and the "importance-aware" wording in3_unified_hf.rst.
| if ( | ||
| weight_quantizer is not None | ||
| and weight_quantizer.is_enabled | ||
| and weight_quantizer.num_bits in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS) |
There was a problem hiding this comment.
Bot comment.
representative_weight_quantizer can return a SequentialQuantizer (W4A8_AWQ), and that class only delegates fake_quant / is_enabled / amax to its members — num_bits falls through to nn.Module.__getattr__ and raises AttributeError. is_enabled delegates to self[0], so the and chain does reach .num_bits for an enabled sequential quantizer.
This is now on the unconditional path: save_pretrained calls uses_iq_quantization(self.model) for every Megatron export, and W4A8_AWQ is otherwise supported there (process_layer_quant_config has a w4a8_awq branch, to_quantized_weight packs it). So a W4A8_AWQ Megatron export that works today would fail with an unrelated AttributeError.
The sibling helper in this file already handles it (get_weight_block_size: if isinstance(weight_quantizer, SequentialQuantizer): weight_quantizer = weight_quantizer[0]). Mirroring that, or getattr(weight_quantizer, "num_bits", None) in (...), is enough — IQ is never a sequential format, so either reads as "not IQ". Worth a regression test with a W4A8-style config alongside test_uses_iq_quantization_false_without_iq_layers.
…nk-uniform The helper added in the previous commit read ``weight_quantizer.num_bits`` directly. A ``SequentialQuantizer`` has ``is_enabled`` but no ``num_bits``, so this raised ``AttributeError`` on any W4A8_AWQ model -- and since ``save_pretrained`` calls the guard on every Megatron export, it raised before any format dispatch. A regression introduced by 106686d; read the attribute defensively instead. A SequentialQuantizer is never IQ, which is a single quantizer with backend="ggml", so ``None`` is the right answer for it. The guard was also rank-local, as it was before 106686d: under pipeline parallelism a stage holding no IQ layer skipped the raise and then blocked in ``_gather_exclude_modules`` while its peers exited -- a hang rather than a clean error, the same failure shape ``_check_weight_quantization_took_effect`` already warns about. Agree across ranks first via ``all_gather_object``, mirroring ``_gather_exclude_modules``, including its ``is_initialized`` guard for single-process export. Tests: a SequentialQuantizer case, mutation-checked -- with the defensive read reverted it fails with the exact AttributeError above. 30 tests in test_get_quantization.py, 246 across the export and ggml unit suites. The rank-agreement path is not covered here: it needs Megatron and more than one rank, neither available in this environment. ``_reject_unsupported_fused_iq_export`` still raises from rank-local inspection and has the same hazard; left alone as it is outside this round's findings. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (claude-opus-5) — DM the bot to share feedback.
Nudge: all three items from the last round are fixed, but the fused-MoE rejection still has the rank-divergence problem the author just solved for the TP guard.
Needs action:
- Make
_reject_unsupported_fused_iq_exportrank-uniform inunified_export_megatron.py: it raises from inside the per-expert loop, so a PP stage (or EP rank) owning no fused expert skips the raise and blocks in_gather_exclude_modules()while peers exit. Reuse the_any_rank_uses_iq_quantization()pattern, or hoist it next to the TP guard. - Confirm
save_pretrainedis entered by every rank before the new unconditionalall_gather_objectin_any_rank_uses_iq_quantization()— it now runs on every Megatron export, IQ or not. - Report GPU/Megatron CI green on the rebased branch, and confirm the #2448 → #2446 → this merge order now that
quantization/ggml/is onmain.
No action needed:
- ✔️ Resolved since the last review: the
SequentialQuantizerAttributeErrorinuses_iq_quantization(nowgetattr(..., "num_bits", None)), its regression test, the rank-local TP>1 guard,group_sizeforwarding inconvert_hf_config.py, and the "importance-aware" doc wording.
_reject_unsupported_fused_iq_export raises from inside the two per-expert loops, so a pipeline stage or EP rank owning no fused expert skipped it and then blocked in the collectives in save_pretrained while its peers exited -- the same divergence just fixed for the TP>1 guard, one layer down. Decide it up front instead. Whether experts are packed into one fused tensor is a property of the architecture's export rule table, not of which layers a rank happens to own, so it is identical everywhere: pair it with the already rank-agreed IQ flag and raise beside the TP check. The in-loop rejection stays as a backstop for paths that do not enter save_pretrained. The predicate is a module-level pure function over a mapping dict rather than a method, so it is testable without Megatron or a second rank -- four parametrized cases plus a non-CustomModuleMapping entry. Also confirmed, no change needed: save_pretrained has no early return before the all_gather_object added in 070baa5, and already contained _gather_exclude_modules, _gather_layer_config_dict, three barriers and an all_gather_object, one of which carries a comment about stranding peers. It already required every rank to enter; this adds no new constraint, it only moves the first collective earlier. 35 tests in test_get_quantization.py, 251 across the export and ggml unit suites. The rank-agreement path itself remains uncovered here: it needs Megatron and more than one rank. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (claude-opus-5) — DM the bot to share feedback.
Nudge: the rank-divergence fix from the last round landed cleanly, but the fused-expert guard is now broader than the problem it solves and CI/merge order still needs confirming.
Needs action:
- Confirm the over-rejection in
save_pretrainedis intended:_mappings_pack_fused_experts(all_mcore_hf_export_mapping[self.arch])keys off the architecture's rule table, so an IQ dense-layer export on any fused-expert arch (Llama4, GPT-OSS) raises even when the experts themselves are BF16/FP8. Narrow it to the experts' own format, or say why mixed IQ+fused cannot occur. - Add a test that
save_pretrainedraises the fused-MoENotImplementedError— only_mappings_pack_fused_expertsand the per-expert_reject_unsupported_fused_iq_exportpaths are covered today. - Report GPU/Megatron CI green on the rebased branch and confirm the #2448 → #2446 → this merge order.
No action needed:
- ✔️ Resolved since the last review: the rank-divergent fused-MoE rejection (now hoisted next to the TP=1 guard on the rank-uniform rule table, with unit tests), and the
_any_rank_uses_iq_quantization()all_gather —save_pretrainedalready runs unconditional barriers/all-gathers, so every rank must enter it. - ggml/llama.cpp OSRB sign-off sits in #2446, outside this file set.
a232c1a hoisted the fused-MoE rejection beside the TP guard to make it rank-uniform, keying it off the architecture's export rule table. That is the wrong key: the table says whether an architecture *has* fused experts, not whether *those experts* are IQ. So an IQ dense-layer export whose experts are BF16 or FP8 was rejected outright on every fused architecture -- measured, that is GptOssForCausalLM and Llama4ForConditionalGeneration, exactly the two the review named. Blocking a working configuration is worse than the hang it was meant to prevent, which only affects a configuration that is unsupported anyway. Restore the per-expert rejection as the sole mechanism. It is precise: it fires on the expert's own qformat. It is still rank-local, so the divergence the previous round asked about remains open, and the docstring now records both the hazard and why the architecture rule table cannot close it -- a correct fix needs a rank-uniform way to know the experts themselves are IQ, which means resolving the fused rule keys against the local model rather than reading a static table. I do not have a way to test that here: it needs Megatron and more than one rank, and my two attempts at closing it from static information have each traded one failure mode for another. Worth the author's judgement on whether mixed IQ + fused-expert models can occur at all, which would settle it without any detection. 30 tests in test_get_quantization.py, 246 across the export and ggml unit suites. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (claude-opus-5) — DM the bot to share feedback.
Nudge: the fused-expert over-rejection is gone, but moving the guard back into the per-expert loop re-opens the rank-divergence hang the previous round had closed.
Needs action:
- 💬 Author reverted to a per-expert
_reject_unsupported_fused_iq_exportand documented why the arch rule table alone can't gate it — still needs rank-uniformity: an EP/PP rank owning no fused expert never enters the loop, skips the raise, and blocks in_gather_exclude_modules()while peers exit. Gate on the rank-uniformuse_packed_local_experts/pack_name_remapping*rule AND anall_gathered "experts are IQ" flag (the_any_rank_uses_iq_quantization()pattern), or state in the PR body that the hang is acceptable. - Report GPU/Megatron CI green on the rebased branch and confirm the #2448 → #2446 → this merge order.
No action needed:
- ✔️ Resolved since the last review: IQ dense-layer export on fused-expert architectures (Llama4/GPT-OSS) is no longer rejected, the guard now keys off the experts' own format, and both fused entry points are covered by tests asserting
_state_dict == {}. The earlier ask for asave_pretrained-level fused-MoE test is moot now that the check no longer lives there. - ggml/llama.cpp OSRB sign-off sits in #2446, outside this file set.
…iform The fused-MoE rejection raises from inside the per-expert loops, so it only runs on ranks that own an expert -- a stage without one skipped it and blocked in save_pretrained's collectives while its peers exited. My two attempts at detecting the condition up front each traded that hang for a worse problem, the second rejecting IQ dense exports on every fused architecture. Constrain the topology instead. IQ export already required TP=1; require PP=1 as well. Both sizes are identical on every rank and the IQ flag is already agreed across them, so these raise everywhere or nowhere. With PP=1 every rank holds the same layers and reaches the same loops, and expert parallelism shards a set of experts quantized alike, so the per-expert rejection fires on all ranks together. No detection needed, and nothing over-rejected. The residual gap is a rank holding no local expert at all, which needs expert-parallel size to exceed the expert count; recorded in the docstring rather than guarded, since no supported topology reaches it. TP=1 was already load-bearing for correctness, not just uniformity: packing happens during export, so a tensor-parallel shard would be packed as if it were a whole weight. The deployment doc now states both constraints. 246 tests across the export and ggml unit suites. The guards themselves stay uncovered here -- they need Megatron and more than one rank. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
|
/ok to test 5d0bbab |
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (claude-opus-5) — DM the bot to share feedback.
Nudge: last round's fused-expert rank-divergence hang is properly closed by the new rank-uniform PP=1/TP=1 guards, but the guard's detector looks blind to grouped-MoE quantizers.
Needs action:
- Confirm or fix
uses_iq_quantizationinquant_utils.pyforTEGroupedLinear: it iteratesweight_attr_names, which yields nothing for a module carryingweight0..Nplus aGroupedQuantizer, so an experts-only IQ model returnsFalseand TP>1 is never rejected. See inline. - Add a test for the new PP>1
NotImplementedErrorintests/gpu_megatron/torch/export/test_unified_export_megatron.py; onlytest_megatron_iq_export_rejects_tensor_parallelismcovers the TP branch. - Report GPU/Megatron CI green on the rebased branch and confirm the #2448 → #2446 → this merge order now that
quantization/ggml/is onmain.
No action needed:
- ✔️ Resolved since the last review: the fused-MoE rank divergence —
_any_rank_uses_iq_quantization()now decides TP=1 and the new PP=1 requirement uniformly before any collective, and the residualep_size > num_expertscase is documented in_reject_unsupported_fused_iq_export.
| # getattr: a SequentialQuantizer has is_enabled but no num_bits, and is never IQ -- | ||
| # IQ is a single quantizer with backend="ggml". | ||
| if ( | ||
| weight_quantizer is not None |
There was a problem hiding this comment.
Bot comment.
This walk may not see grouped-MoE experts. weight_attr_names(module) yields "weight" only when module.weight is not None, and for other names only when representative_weight_quantizer(module, name) finds <name>_weight_quantizer. A TEGroupedLinear normally has no weight attribute (_grouped_mlp_slicing does has_weight = hasattr(module, "weight") and later delattrs the one it temporarily assigns) and holds weight0..weightN with a single weight_quantizer that is a GroupedQuantizer — so weight_attr_names yields nothing here, and the recursion into weight_quantizer / its child TensorQuantizers has no weight either.
Net effect: a model whose IQ quantizers live only on grouped MoE experts (an expert-only recipe is a plausible shape for 1–2 bit formats) returns False, save_pretrained skips both the TP=1 and PP=1 raises, and _grouped_mlp_slicing then packs TP/ETP-sharded expert weights as if they were whole — the exact failure this guard exists to prevent.
representative_weight_quantizer already unwraps a GroupedQuantizer to q[0], so handling this is mostly a matter of also considering weight0 (or iterating the GroupedQuantizer members directly) when no standard weight attr is present. If grouped experts can never be the only IQ layers in a shipped recipe, a one-line comment saying so would close this out instead.
test_megatron_iq_export_rejects_pipeline_parallelism mirrors the existing TP
test: it patches the parallel-size helpers, so it needs no real Megatron model.
Verified that the construction those tests use does trip the guard --
is_enabled True, num_bits "iq2_xs", uses_iq_quantization True -- which also
confirms the existing TP test still passes now that the guard routes through
_any_rank_uses_iq_quantization.
On the TEGroupedLinear finding: confirmed, and broader than the IQ guard.
Reproduced with a module shaped like one -- weight0..N parameters and a single
GroupedQuantizer under weight_quantizer:
weight_attr_names yields: []
get_quantization_format: None
uses_iq_quantization: False
So this is not a divergence between the guard and the exporter; both are blind
to that layout for every format, and an experts-only model reports no format at
all. representative_weight_quantizer handles the GroupedQuantizer correctly --
the yield in weight_attr_names is gated on a plain ``weight`` parameter that a
TEGroupedLinear does not have.
Patching only uses_iq_quantization would make it see IQ where the rest of the
export sees nothing, rejecting TP>1 for a model the exporter would then treat
as unquantized. The fix belongs in weight_attr_names, which every format's
detection shares and which I cannot exercise here -- it needs Transformer
Engine and Megatron. Recorded in the docstring instead.
Whether this is reachable at all is worth the author's judgement: TEGroupedLinear
is the grouped-expert layout, and fused-expert IQ export is rejected outright.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
|
/ok to test 68b7230 |
Both statements were written while #2447 was still open and are stale as of its merge in ed5c5ed. ptq.md claimed these recipes "configure simulated weight quantization only; packed checkpoint export is added separately". Unified HF export now writes the packed GGML blocks, so say that instead, along with the two limits a recipe user will otherwise discover by hitting an exception: Megatron export requires tensor and pipeline parallel sizes of 1, and does not support fused-MoE experts. The CHANGELOG entry said "fake quantization". This PR carries the only entry for the whole feature -- #2446, #2447 and #2448 each added none -- so it has to describe the shipped surface, which now includes export. Verified against main rather than the PR descriptions: the TP and PP guards and the fused-MoE rejection are present in unified_export_megatron.py, and unified_export_hf.py calls quantize_iq1_s/quantize_iq2_xs. No rebase was needed -- the branch already carries main at b311c05. The recipe configs check out: the validator loads both recipes cleanly and the 12 preset tests pass, including the one asserting the recipe matches IQ1_S_BLOCK_SIZE / IQ1_S_EFFECTIVE_BITS. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
## Summary - add numerics configs for IQ1_S and IQ2_XS - add model presets and general PTQ recipes - document the supported weight-shape requirement - reuse the packed IQ weight across forwards instead of re-encoding it every time - add end-to-end `hf_ptq` coverage for both formats - add the release-note entry ## PR split This work is split into four focused PRs. Each PR targets `main` and owns a disjoint file set: 1. **Kernel** — [#2448: Add CUDA kernels for IQ packing](#2448) 2. **Quantization** — [#2446: Add IQ quantization codecs and backend](#2446) 3. **Export** — [#2447: Export IQ checkpoints from HF and Megatron](#2447) 4. **Recipes** — [#2449: Add IQ post-training quantization recipes](#2449) The required merge order is #2448, #2446, #2447, then #2449. ## Scope Sixteen files. The PR started as eight recipe config, documentation and changelog files; the end-to-end test added for them surfaced two performance bugs in the already-merged codec, and fixing those pulled in the codec files and their tests. Beyond the original recipe set it now touches four files owned by #2446 — `ggml/common.py`, `ggml/iq1_s.py`, `ggml/iq2_xs.py` and `ggml/backend.py` — plus three test files. It still contains no kernel or export files. Keeping those fixes here rather than moving them to #2446 is deliberate and confirmed with the stack owner: #2446 is already merged, and both bugs are only observable through the end-to-end test this PR adds, so splitting them would separate each fix from the test that demonstrates it. ## Packed-weight cache fix Adding the end-to-end test made the cost visible: a TinyLlama IQ1_S `hf_ptq` run spent **499 of its 537 seconds inside the IQ1_S encoder**, packing the same 154 weights 15400 times — 100 times each. The repacking is not calibration. These recipes set `algorithm: null` and `hf_ptq` logs `Dynamic quantization. Calibration skipped.`. The 100 passes are the sample `generate()` calls `hf_ptq` makes before and after quantization: one decode step re-runs weight fake-quant on every linear, and the packed payload was thrown away each time. `_PackedWeightCache` was already there to prevent exactly this, and it never hit. It keyed on the identity of the tensor the backend was handed, but `TensorQuantizer` passes a fresh *view* of the weight on every forward, so the identity check never matched twice. The fix anchors the entry to `inputs._base` — the parameter the view is taken from — held as a weakref. The parameter is stable across forwards, so the cache hits; the reference stays weak, so the payload is released with the weight and offloaded/meta-device flows are unaffected. (A strong reference does make the cache hit, but pins full-precision storage for the life of the quantizer, which is the opposite of what those flows need.) Measured on TinyLlama IQ1_S `hf_ptq`, 2×H100, same command before and after: | | packer calls | packing time | wall clock | |---|---|---|---| | before | 15400 | 499.0 s | 8m57s | | after | 154 (one per weight) | 6.2 s | 3m21s | `test_ggml_weight_is_packed_once_across_forwards` pins this: it counts encoder calls across five forwards under `torch.inference_mode()` (what `generate()` runs under) and asserts exactly one. ## Decode chunk fix With packing cached, the end-to-end cost moved entirely into the decode, and IQ2_XS was still 4x slower than IQ1_S (1081s vs 260s per case). Instrumenting both showed packing was no longer the cost at all — IQ2_XS packs *faster*: | | pack calls | packing time | wall clock | |---|---|---|---| | IQ1_S | 154 | 6.2 s | 3m21s | | IQ2_XS | 154 | 1.6 s | 17m59s | The cause was one constant serving two loops with opposite characteristics. `_DEFAULT_BLOCK_CHUNK_SIZE` bounds the torch encode fallback, which holds the large codebook-search temporaries and runs once per weight; IQ2_XS sets it to 256 rather than IQ1_S's 1024 because its search sweeps sixteen local scales per grid tile. But the *decode* shared it — and the decode has tiny temporaries, runs on every forward, and is never cached, so a small chunk only multiplies kernel launches. Decoding a 2048x5632 weight: | chunk | IQ2_XS decode | transient peak | |---|---|---| | 256 (was) | 91.4 ms | +24 MiB | | 1024 | 22.9 ms | +31 MiB | | 4096 (now) | 5.8 ms | +56 MiB | The decode now takes its own `_DEFAULT_DECODE_CHUNK_SIZE`, threaded through `fake_quantize_with_cache`. The encode bounds are untouched, so the memory ceiling stays where it was aimed. The seven-iteration sign-parity loop in `dequantize_iq2_xs` is also folded into three XOR steps, off the same per-forward path. Per case in `tests/examples/hf_ptq` on 2xH100: | | before | after | |---|---|---| | IQ1_S | 259.65 s | 94.64 s | | IQ2_XS | 1081.05 s | 102.77 s | Both now fit the 300s `tests/examples` default, so the cases carry no explicit timeout. ## Integration contracts - `block_sizes: {-1: 256}` records the native packed-block contract; it does not drive the GGML fake-quant scale search. The export path reads `TensorQuantizer.block_sizes[-1]` through `get_weight_block_size`, validates it against the format block size, and records `group_size: 256` in checkpoint metadata. - The model presets intentionally expose `--qformat iq1_s` and `--qformat iq2_xs` in `hf_ptq`. ## Testing - `tests/unit/torch/quantization/ -k 'ggml or iq1 or iq2'` — 52 passed, 1 skipped - `tests/unit/recipe/test_presets.py` — both shipped IQ recipes load with the expected backend and no unused search option - `tests/examples/hf_ptq/test_llm_ptq.py -k 'iq1_s or iq2_xs'` — 2 passed on 2×H100 (TinyLlama, both formats end to end through export), 3m17s for the pair - pre-commit hooks pass on all changed files - larger-model sanity check outside CI: Qwen3.8-27B (2256 quantizers) quantizes and exports with `general/ptq/iq1_s` on 2xH100 in 66m46s, export itself 192s --------- Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com> Co-authored-by: Chenjie Luo <chenjiel@nvidia.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
### What does this PR do? Type of change: new feature llama.cpp defines five GGML IQ formats at one and two bits; we ship two. This adds **IQ2_XXS** at 2.0625 bits per weight, between IQ1_S and IQ2_XS, and is the **first of three**. On a real mixed-precision checkpoint (`unsloth/Qwen3.8-27B-GGUF`, `Qwen3.8-27B-UD-IQ1_S.gguf`) IQ2_XXS alone covers **59 tensors and 2.84 B parameters — 10.6% of the file**, which a reader limited to IQ1_S/IQ2_XS cannot consume. Across all three PRs the missing formats account for 17.3%. | format | bpw | bytes/256 | codebook | | |---|---|---|---|---| | `iq1_s` | 1.5625 | 50 | `iq1s_grid` (2048) | existing | | **`iq2_xxs`** | **2.0625** | **66** | **`iq2xxs_grid` (256)** | **this PR** | | `iq2_xs` | 2.3125 | 74 | `iq2xs_grid` (512) | existing | The encoder follows the existing single-pass grid search at a fixed anchored super-block scale, and the CUDA kernel the existing per-block structure. IQ2_XXS reuses IQ2_XS's even-parity sign rule but packs a 4-bit sub-block scale into the same 32-bit word as four 7-bit sign indices, and its 256-entry grid needs no high index bits. ### Groundwork the next two reuse Two things land here because IQ2_XXS is the first format to need them: - **Export registry.** The IQ family was spelled as a two-element tuple at **nine** sites across `quant_utils.py`, `unified_export_hf.py` and `unified_export_megatron.py`. Those become an `IQ_FORMATS` frozenset plus per-format packer and block-geometry tables, so a format is a row rather than a sweep through the exporters. - **Shared test contract.** The per-format test files had drifted apart — each of `iq1_s` and `iq2_xs` tested things the other did not. They become one parametrized module per layer (unit and CUDA), so every format is held to the same contract and a new one inherits it. ### Usage ```bash python examples/hf_ptq/hf_ptq.py --pyt_ckpt_path <model> --recipe general/ptq/iq2_xxs ``` ### Testing **The decoder is validated against llama.cpp's own output, not just round-tripped.** Every IQ2_XXS tensor in the checkpoint above, compared against `dequantize_row_iq2_xxs` from `ggml-quants.c`: ``` IQ2_XXS: 59 tensors, 11,100,160 blocks → 0 mismatched, max|diff| 0.0 ``` The new codebook matches the `ggml-common.h` table entry for entry, as does the `ksigns_iq2xs` sign table. Blocks lifted from that checkpoint ship as conformance vectors so CI keeps checking bytes we did not produce; mutation testing confirms they catch a wrong sign-field width. The CUDA encoder is byte-identical to the PyTorch reference on a fixed input and runs at **1047.9 M elem/s against the torch search's 10.7** on a 5632×2048 weight. - `tests/unit/torch/quantization/ -k 'ggml or iq1 or iq2 or iq_'` — 99 passed - `tests/gpu/torch/quantization/test_iq_formats_cuda.py` — 21 passed (7 checks × 3 formats) - `tests/unit/recipe/test_presets.py` — passing; `general/ptq` now holds 29 recipes, `ptq.md` updated - reconstruction error decreases monotonically with bit width, pinned by a test Pre-existing failures in `tests/unit/torch/export/` and `test_autoquant.py` are `transformers`/`torchvision` import problems in my environment — identical counts with and without this change. ### A finding about already-merged code Checking the new kernel against its PyTorch reference at 4096 blocks showed that **CUDA and torch encoders disagree on roughly 1 block in 6000 — including the already-merged `iq2_xs`**, at 0.0163% against IQ2_XXS's 0.0000%. Root cause: both compute `xnorm − 2·scale·dot + scale²·qnorm`, but CUDA fuses it with `fmaf` while torch uses separate ops; where two local scales fall within a float32 ULP the roundings pick different sides. Adjudicated against float64, neither path is better (5 to 6). Worst-case cost is **1.48e-08** relative reconstruction error, and run-to-run determinism on a given device holds. This is pre-existing, not introduced here — `test_iq2_xs_cuda.py` asserts exact byte parity but on a 16-block weight where ties essentially never arise. I have **not** changed that test; rewording a guarantee on merged code belongs in its own change. The new shared GPU tests assert exact parity on a small fixed input and compare reconstruction error at scale. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: ✅ — the new codebook is a GGML table, carried in `codebooks.py` beside the existing ones so the MIT-licensed surface stays in that one file, with the source revision recorded. No new dependencies. - Did you write any new necessary tests?: ✅ - Did you update Changelog?: ✅ - Did you get Claude approval on this PR?: ❌ — not yet run ### Additional Information First of three; **IQ2_S** and **IQ1_M** follow and build on this branch. Replaces #2505, which carried all three at once. Follows #2446 / #2447 / #2448 / #2449, which landed IQ1_S and IQ2_XS. 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Added IQ2_XXS weight-only quantization, including CUDA acceleration and support for Hugging Face and Megatron exports. - Added the `general/ptq/iq2_xxs` recipe. It requires no calibration data and supports eligible layers with a weight dimension divisible by 256. - Updated the PTQ recipe catalog to list IQ1_S, IQ2_XXS, and IQ2_XS at approximately 1.56, 2.06, and 2.31 bits per weight, respectively. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Chenjie Luo <chenjiel@nvidia.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
2 participants
Summary
uint8weight contract and the fused-expert boundaryPR split
This work is split into four focused PRs. Each PR targets
mainand owns a disjoint file set:The required merge order is #2448, #2446, #2447, then #2449.
Scope
This PR owns only export code, deployment documentation, and export tests. It targets
mainand should merge after #2448 and #2446. It does not contain kernel, codec/backend, or recipe files.Deployment consumer boundary
Dense weights and individually named expert weights use the documented shaped
uint8contract. Megatron fused-MoE IQ export is intentionally rejected withNotImplementedError: its payload would have shape[num_experts, out_features, in_features // 256, payload_bytes], and no deployment loader in this stack currently owns that layout. Support should be enabled only with a loader integration test.Test coverage
Validation
Summary by CodeRabbit
New Features
Limitations
weightattributes in Hugging Face models.