Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
|
Important Draft PR not reviewedDraft PRs are not automatically reviewed by default.
To automatically review draft PRs, update your CodeRabbit configuration: reviews:
auto_review:
drafts: trueComment |
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #2381 +/- ##
==========================================
- Coverage 70.96% 69.80% -1.17%
==========================================
Files 542 595 +53
Lines 63784 68763 +4979
==========================================
+ Hits 45266 47998 +2732
- Misses 18518 20765 +2247
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
ChenhanYu
changed the title
286f594 to
075e991
Compare
|
/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.
[IMPORTANT Export] _get_iq_weight_state appends a . to prefix, but at this call site (and at the _pack_name_remapping_gpt_oss twin around line 2025) prefix is a complete tensor key, not a module prefix.
The non-IQ branches right below treat it that way:
elif merged_weight_scale is None:
self._state_dict[prefix] = merged_weight
else:
self._state_dict[prefix] = to_quantized_weight(...)
self._state_dict[prefix + "_weight_scale"] = merged_weight_scaleand the rule book confirms it — mcore_llama.py:72 maps local_experts.linear_fc1 to "language_model.model.layers.{}.feed_forward.experts.gate_up_proj" and mcore_gptoss.py:41 to "model.layers.{}.mlp.experts.gate_up_proj", neither with a trailing dot. Since _get_iq_weight_state does prefix = prefix.rstrip(".") + ".", packed MoE experts emit
...experts.gate_up_proj.packed_weights...experts.gate_up_proj.weight_shape
next to siblings named ...experts.gate_up_proj_weight_scale / ..._input_scale. Every other IQ call site (1232, 1280/1283, 1486, 1653, 1775) passes a real module prefix ending in ., so those come out right — only the two pack_name_remapping paths (Llama4 / GPT-OSS MoE) produce this inconsistent key shape. A consumer that resolves gate_up_proj as a parameter name plus suffixes will not find these tensors.
Suggested fix: don't normalize inside the helper; take the two conventions explicitly at the call site, e.g.
# in _get_iq_weight_state, replace the rstrip/append with:
# (no prefix rewriting; caller supplies the exact base)
return {
f"{prefix}_packed_weights": packed_weights.detach().cpu(),
f"{prefix}_weight_shape": weight_shape.detach().cpu(),
}and have the module-prefix call sites pass prefix + "weight" so they still yield ...weight_packed_weights-style keys — or, if the dotted form is the intended wire format, pass prefix + "gate_up_proj." style bases here instead of the bare parameter key. Either way the two packed-expert paths should match the naming the rest of the exporter uses for that tensor.
|
|
||
| @torch.no_grad() | ||
| def quantize_iq1_s( | ||
| weight: torch.Tensor, *, block_chunk_size: int = 4 |
There was a problem hiding this comment.
[IMPORTANT Performance] block_chunk_size=4 makes the CPU/eager encoder launch-bound by ~16×. _encode_blocks runs a fixed Python nest of 16 grid tiles × 2 shifts × 8 local scales = 256 iterations plus the matmuls, per chunk, independent of chunk size. With 4 blocks per chunk, a 4096×4096 weight (65 536 blocks) needs 16 384 chunks ⇒ ~4.2 M tiny tensor ops; at 64 blocks it's 1 024 chunks for the same arithmetic.
Memory isn't the constraint: the largest temporary is error/shifted_dot at chunk × 32 × 16 floats, ~1 MB at chunk 64 — and quantize_iq2_xs in this same PR already defaults to 64 with a larger grid (512×8 with per-entry sign search). The asymmetry looks unintentional.
| weight: torch.Tensor, *, block_chunk_size: int = 4 | |
| weight: torch.Tensor, *, block_chunk_size: int = 64 |
| error = ( | ||
| xnorm.unsqueeze(-1) - 2 * scale * shifted_dot + scale.square() * shifted_norm | ||
| ) | ||
| tile_error, tile_index = error.min(dim=-1) | ||
| replace = tile_error < best_error[:, :, choice] |
There was a problem hiding this comment.
[IMPORTANT Compatibility] The IQ1_S eager error is left unclamped, while iq1_s.cu clamps it in both search stages (fmaxf(error, 0.0f)) and the sibling iq2_xs.py:186 applies .clamp_min_(0). So the two IQ1_S implementations minimize different objectives.
The clamp is not cosmetic for argmin: xnorm - 2·scale·dot + scale²·norm is an exact-arithmetic-nonnegative expression evaluated in floating point, so near-perfect fits land at small values of either sign. When one candidate rounds to -1e-7 and another to +1e-7, unclamped CPU prefers the negative one and clamped CUDA sees a tie broken by lowest entry index — different entry, different payload byte. The same applies at the group_error.argmin stage, where 4 clamped-vs-unclamped terms are summed.
Fix: mirror the kernel (and IQ2_XS) here —
error = (
xnorm.unsqueeze(-1)
- 2 * scale * shifted_dot
+ scale.square() * shifted_norm
).clamp_min_(0)Two related gaps worth closing in the same pass (details in a top-level comment): the CPU↔CUDA payloads are never asserted equal by any test, and get_cuda_ext_iq1_s/get_cuda_ext_iq2_xs build with --use_fast_math, which is a second independent source of divergence.
| packed, shape = quantize_iq1_s(weight) | ||
| packed_again, _ = quantize_iq1_s(weight) | ||
| reconstructed = dequantize_iq1_s(packed, shape) | ||
|
|
||
| assert packed.shape == (8, 1, 50) | ||
| assert torch.equal(packed, packed_again) | ||
| normalized_mse = ( | ||
| reconstructed.float() - weight.float() | ||
| ).square().mean() / weight.float().square().mean() | ||
| assert normalized_mse < 0.25 |
There was a problem hiding this comment.
[IMPORTANT Compatibility] These GPU tests only check that the CUDA encoder agrees with itself (packed == packed_again) and lands under an MSE bound. Nothing in the suite asserts that the CUDA payload equals the eager payload for the same input — yet the two are presented as interchangeable implementations behind search_impl: auto, so which one produced a checkpoint depends on where the export ran.
Three concrete divergence sources, all live today:
- The eager error at
iq1_s.py:187is unclamped;iq1_s.cuappliesfmaxf(..., 0.0f)in both stages (flagged inline). extensions.pybuilds both IQ kernels with--use_fast_math, i.e. imprecisediv/sqrtand denormal flush-to-zero.find_scalecomputesdand rounds it to fp16 — a single ULP of drift inamax / _IQ1_S_NATIVE_MAX * 0.61before the__float2halfcan tipdto the neighbouring fp16 value, after which the entire 50-byte block differs, not one field.- Reduction order differs by construction: the kernel accumulates
dotwith sequentialfmafacross a warp shuffle tree, the eager path usesmatmul/sum(dim=-1).
Suggested test to add here (the useful assertion is bitwise, since the payload is what ships):
def test_iq1_s_cuda_pack_matches_eager_pack():
generator = torch.Generator().manual_seed(1234)
weight = torch.randn((8, 512), generator=generator, dtype=torch.bfloat16)
eager, eager_shape = quantize_iq1_s(weight, search_impl="eager")
cuda, cuda_shape = quantize_iq1_s(weight.cuda())
assert torch.equal(cuda.cpu(), eager)
assert torch.equal(cuda_shape.cpu(), eager_shape)(adjust to whatever knob forces the eager path — backend_extra_args={"search_impl": ...} today only reaches it through the quantizer). If exact equality is not a goal, that's a legitimate design choice, but then it should be stated in the module docstring, because it means the exported GGML bytes are host-dependent; a tolerance-based check on the dequantized values plus a documented note would be the alternative.
| def test_iq1_s_dequantizes_ggml_metadata_bit_fields(): | ||
| packed = torch.zeros((1, 1, 50), dtype=torch.uint8) | ||
| d = torch.tensor([2.0], dtype=torch.float16).view(torch.uint8) | ||
| packed[0, 0, :2] = d | ||
| entries = torch.tensor([0, 256, 511, 2047], dtype=torch.int64) | ||
| packed[0, 0, 2:6] = (entries & 0xFF).to(torch.uint8) | ||
| qh = ( | ||
| ((entries[0] >> 8) & 7) | ||
| | (((entries[1] >> 8) & 7) << 3) | ||
| | (((entries[2] >> 8) & 7) << 6) | ||
| | (((entries[3] >> 8) & 7) << 9) | ||
| | (3 << 12) | ||
| | (1 << 15) | ||
| ) | ||
| packed[0, 0, 34] = (qh & 0xFF).to(torch.uint8) | ||
| packed[0, 0, 35] = (qh >> 8).to(torch.uint8) | ||
|
|
||
| decoded = dequantize_iq1_s(packed, torch.tensor([1, 256]), dtype=torch.float32) | ||
| expected = (iq1_s_grid()[entries] - 0.125) * 14.0 | ||
|
|
||
| assert torch.equal(decoded[0, :32].reshape(4, 8), expected) |
There was a problem hiding this comment.
[SUGGESTION] This is the closest thing to a format-compatibility test, and it's self-referential: expected is computed with iq1_s_grid() and the same (g - 0.125) * d formula the decoder uses, so an error in the grid table itself or in the scale ladder passes. The whole point of IQ1_S/IQ2_XS is that llama.cpp can read the bytes.
Worth adding a handful of golden vectors — a fixed input array with the 50-byte (resp. 74-byte) payload produced by llama-quantize/ggml_quantize_chunk, checked in as a literal, plus the expected dequantize_row_iq1_s output. That pins the grid contents, the 2*ls+1 ladder, the ±0.125 delta, the qh field packing, and (for IQ2_XS) the ksigns_iq2xs even-parity convention against the reference in one assertion. Byte-exact encoder equality is a strong ask given the search heuristics differ; asserting the decoder reproduces reference GGML blocks is cheap and hermetic (tests/unit stays offline — just constants).
| ) | ||
| if not weight.is_floating_point(): | ||
| raise TypeError(f"{format_name} requires a floating-point weight, got {weight.dtype}") | ||
| if not torch.isfinite(weight).all(): |
There was a problem hiding this comment.
[SUGGESTION] torch.isfinite(weight).all() in a Python if forces a device→host sync on every quantize_iq* call, and it reads the full tensor. CONTRIBUTING's coding standards call out avoiding CPU-GPU syncs specifically. It's once per weight rather than per block, so the cost is bounded — but the check also doesn't buy much: amax of a tensor containing inf/NaN already produces a non-finite d that clamp(max=65504) turns into either 65504 or NaN, and a NaN weight is a bug upstream. Either drop it or gate it behind a debug flag, so the hot export path doesn't serialize on it.
| 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] This branch hardcodes "group_size": 256 and silently ignores the group_size argument (gs) that every other branch honours. That's correct for GGML — the 256-value block is baked into the 50/74-byte layout — but a caller who passes group_size=128 gets a config claiming 256 with no warning, and the mismatch surfaces only when a consumer tries to unpack. Worth asserting the contract instead of dropping the argument:
elif quant_algo in ("IQ1_S", "IQ2_XS"):
if gs not in (None, 256):
warnings.warn(f"{quant_algo} uses a fixed GGML block of 256; ignoring group_size={gs}")
effective_bits, payload_bytes = (1.5625, 50) if quant_algo == "IQ1_S" else (2.3125, 74)Also "type": "int" with "num_bits": 1 is a lossy description of the format for anything reading this generically — the values are ternary {-1,0,1} (IQ1_S) / magnitude-grid codes (IQ2_XS) with a shared fp16 block scale, not 1-bit integers. effective_bits and packing: "ggml" carry the real information; consider a comment noting num_bits is nominal here so a future reader doesn't try to derive the size from it.
Claude review — full review of
|
There was a problem hiding this comment.
Claude review complete — 0 CRITICAL, 4 IMPORTANT, 3 SUGGESTION. Not approving; the IMPORTANT findings should be resolved first. Full write-up in the summary comment above.
IMPORTANT:
- Export —
_get_iq_weight_statenormalizesprefixto end in., but the twopack_name_remappingcall sites (unified_export_megatron.py:1912,:2025) pass a complete tensor key, not a module prefix. Llama4 / GPT-OSS packed MoE experts end up with...experts.gate_up_proj.packed_weightsnext to...experts.gate_up_proj_weight_scale. The other six IQ call sites are fine. - Compatibility — the IQ1_S eager search error is unclamped (
ggml/iq1_s.py:187) whileiq1_s.cuappliesfmaxf(error, 0.0f)in both stages andiq2_xs.py:186uses.clamp_min_(0). The two IQ1_S implementations minimize different objectives, which changes argmin outcomes near zero error. - Compatibility — nothing asserts CPU/CUDA payload equality;
test_iq1_s_cuda.pyonly checks self-determinism. Combined with--use_fast_mathon both kernels and the fp16 rounding ofd, exported GGML bytes can depend on which host produced them. - Performance —
quantize_iq1_sdefaults toblock_chunk_size=4, ~16× more launches than needed for identical arithmetic;quantize_iq2_xsalready uses 64 with a larger grid.
SUGGESTION (non-blocking): no llama.cpp golden byte vectors pinning GGML compatibility; validate_weight's torch.isfinite(...).all() forces a CPU-GPU sync; the IQ branch of _quant_algo_to_group_config silently ignores its group_size argument.
The format work itself — both bit layouts, scale ladders, the IQ2_XS parity convention, the state-dict rename-after-skip mechanism, and both kernels' shared-memory and synchronization structure — checked out against llama.cpp and against itself. The open items are in export plumbing and test coverage, not in the quantizers.
|
Signed-off-by: Chenhan D. Yu <5185878+ChenhanYu@users.noreply.github.com>
Signed-off-by: Chenhan D. Yu <5185878+ChenhanYu@users.noreply.github.com>
Add TP=1 Megatron unified checkpoint export for IQ1_S and IQ2_XS, including final-layout packing and coverage for the scale-free payload schema. Signed-off-by: Chenhan D. Yu <5185878+ChenhanYu@users.noreply.github.com>
Signed-off-by: Chenhan D. Yu <5185878+ChenhanYu@users.noreply.github.com>
Signed-off-by: Chenhan D. Yu <5185878+ChenhanYu@users.noreply.github.com>
Signed-off-by: Chenhan D. Yu <5185878+ChenhanYu@users.noreply.github.com>
Store IQ payloads under canonical weight keys and remove redundant logical-shape tensors.\n\nSigned-off-by: Chenhan D. Yu <5185878+ChenhanYu@users.noreply.github.com> Signed-off-by: Chenhan D. Yu <5185878+ChenhanYu@users.noreply.github.com>
Register packed quantized expert weights as buffers so unified export retains them. Add regression coverage for buffer-backed fused-expert projections. Signed-off-by: Chenhan D. Yu <5185878+ChenhanYu@users.noreply.github.com>
4b12d33 to
73baffe
Compare
|
This large PR is superseded by three focused draft PRs:
The export and CUDA PRs will be retargeted to main after #2443 merges. |
2 participants
What does this PR do?
Type of change: new feature
Adds IQ1_S and IQ2_XS weight quantization and unified Hugging Face and Megatron Core checkpoint export (TP=1). The implementation provides deterministic CPU reference encoders/decoders, optional CUDA encoders, the canonical format codebooks, and GGML-compatible 50-byte IQ1_S or 74-byte IQ2_XS blocks for every 256 logical weights. Unified safetensors export writes each supported weight as
packed_weightsplusweight_shape, with format metadata and reusable PTQ recipes.Usage
Use
general/ptq/iq2_xsfor IQ2_XS.Testing
uv run --frozen pytest -q tests/unit/torch/quantization/test_iq1_s.py tests/unit/torch/quantization/test_iq2_xs.py tests/unit/torch/export/test_export_weight.py tests/unit/torch/export/test_get_quantization.py tests/unit/recipe/test_recipe_docs.py— 35 passed.uv run --frozen pytest -q tests/unit/recipe/test_recipe_docs.py tests/unit/recipe/test_loader.py— 337 passed.uvx --from pre-commit==4.6.0 pre-commit run --from-ref origin/main --to-ref HEAD— all applicable hooks passed.nvcc.Before your PR is "Ready for review"
Make sure you read and follow Contributor guidelines and your commits are signed (
git commit -s -S).Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded
trust_remote_code=True,torch.load(..., weights_only=False),pickle, etc.).CONTRIBUTING.md: ❌ — llama.cpp MIT attribution and third-party notices are included; the internal OSRB/NVBug record remains to be updated before ready-for-review./claude review.Additional Information
The canonical IQ1_S and IQ2_XS codebooks are adapted from llama.cpp at revision 9b05354e under its MIT license. The search behavior is aligned with the PSX-LUTS
search_impl="auto"path used for QAT. Runtime integration in vLLM is intentionally out of scope for this PR.