Repository navigation
fix: load modelopt NVFP4 per-tensor scales at native dtype - #4221
Merged
Merged
Conversation
quant_method: modelopt has no transformers quantizer, so the conversion loader cast every source to the bf16 skeleton dtype before the NVFP4 converters ran. uint8 qdata and e4m3 scales survive that, but the fp32 weight_scale_2 was rounded (up to ~0.4%), so the experts' per_tensor_scale, the dequantized non-routed linears and _nvfp4_pts no longer matched the grid the modelopt merge reads from the file. Keys routed to the NVFP4 converters now keep their checkpoint dtype for scales and uint8/fp8 weights on every load path, with or without the rank0-only patch. Float weights are still cast to the skeleton dtype. This replaces the U8/F8 special case in the rank0-only materialize patch.
The CPU lanes ignore tests/integrations/kernels/, so the test only ran on the GPU kernels lane even though it needs no GPU. tests/integrations/ kernels_cpu/ is collected by the existing tests/integrations/ argument.
Contributor
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the
 |
| Layer / File(s) | Summary |
|---|---|
Track NVFP4 sources and preserve native dtype src/axolotl/integrations/kernels/libs/scattermoe_lora/nvfp4_weight_converter.py |
The loader tracks source patterns that route to NVFP4 converters and materializes those sources with dtype=None, except ordinary float .weight tensors. Converter registration filters existing converters by their operations and installs the patch. |
Validate ModelOpt checkpoint loading tests/integrations/kernels_cpu/test_nvfp4_modelopt_load.py |
A parameterized integration test checks checkpoint loading with synchronous or threaded loading and rank-0-only conversion enabled or disabled. It verifies fused expert tensors, scales, quantized data, linear dequantization, and selected parameter values. |
Priority:
Estimated code review effort: 3 (Moderate) | ~25 minutes
Change: Bug fix
Merge Risk: 
2807a
-
No actionable issue is established that would prevent merging after normal checks.
Pre-merge checks: 
-
Check name Status Explanation Resolution Docstring Coverage Docstring coverage is 33.33% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documenta… Write docstrings for the functions missing them to satisfy the coverage threshold. -
Check name Status Explanation Description Check Check skipped - CodeRabbit’s high-level summary is enabled. Title check The title clearly and concisely describes the main change: loading ModelOpt NVFP4 per-tensor scales at their native dtype. Linked Issues check Check skipped because no linked issues were found for this pull request. Out of Scope Changes check Check skipped because no linked issues were found for this pull request.
Full details: Docstring Coverage
-
Explanation
Docstring coverage is 33.33% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 15 functions across 2 files.
Resolution
Write docstrings for the functions missing them to satisfy the coverage threshold.
-
Finishing Touches: 💡 1
-
📝 Generate docstrings: 💡
-
- Commit to this branch
- Create a new PR
🧪 Generate unit tests (beta)
-
- Commit to this branch
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.
Comment @coderabbitai help to get the list of available commands.
Contributor
|
📖 Documentation Preview: https://6ae303ce5c99db40ad50c63a--resonant-treacle-0fd729.netlify.app Deployed on Netlify from commit 2fd577b |
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
CI picked up pytest 9.2.0, which breaks two CPU tests that pass on 9.1.1: - test_trajectory_tracks_dist_muon: pytest.approx now iterates the expected value, and the test passes 0-d tensors. - test_chunked_forward: torch._dynamo.config.recompile_limit is gone by the time it runs, after test_encoder_canvas monkeypatched it on the same xdist worker. Pin every place that installs pytest until those tests are fixed.
winglian
added a commit
that referenced
this pull request
Oct 29, 2026
The 9.1.1 pin from #4221 was a stopgap for the two tests fixed in this branch. With them passing under pytest 9.2, CI goes back to the latest pytest.
2 tasks done
3 tasks done
NanoCode012
pushed a commit
to NanoCode012/axolotl
that referenced
this pull request
Oct 29, 2026
…i-cloud#4221) * fix: load modelopt NVFP4 per-tensor scales at native dtype quant_method: modelopt has no transformers quantizer, so the conversion loader cast every source to the bf16 skeleton dtype before the NVFP4 converters ran. uint8 qdata and e4m3 scales survive that, but the fp32 weight_scale_2 was rounded (up to ~0.4%), so the experts' per_tensor_scale, the dequantized non-routed linears and _nvfp4_pts no longer matched the grid the modelopt merge reads from the file. Keys routed to the NVFP4 converters now keep their checkpoint dtype for scales and uint8/fp8 weights on every load path, with or without the rank0-only patch. Float weights are still cast to the skeleton dtype. This replaces the U8/F8 special case in the rank0-only materialize patch. * test: move CPU-only NVFP4 modelopt load test to kernels_cpu The CPU lanes ignore tests/integrations/kernels/, so the test only ran on the GPU kernels lane even though it needs no GPU. tests/integrations/ kernels_cpu/ is collected by the existing tests/integrations/ argument. * chore: pin pytest to 9.1.1 CI picked up pytest 9.2.0, which breaks two CPU tests that pass on 9.1.1: - test_trajectory_tracks_dist_muon: pytest.approx now iterates the expected value, and the test passes 0-d tensors. - test_chunked_forward: torch._dynamo.config.recompile_limit is gone by the time it runs, after test_encoder_canvas monkeypatched it on the same xdist worker. Pin every place that installs pytest until those tests are fixed. (cherry picked from commit 7286a73)
winglian
added a commit
that referenced
this pull request
Oct 29, 2026
* Muon 8-bit test: compare the losses as floats for pytest 9.2 pytest 9.2's approx iterates the expected value to check for Decimal tolerances, which raises on the 0-d tensors _step returns. * Diffusion LM test: patch dynamo's recompile_limit with config.patch From pytest 9.2, monkeypatch reads the old value statically, finds nothing on torch._dynamo.config (whose keys go through __getattr__) and undoes the patch with delattr, which hides the key: every later compile on that xdist worker failed with "torch._dynamo.config.recompile_limit does not exist" (test_chunked_xentropy on CI). Dynamo's own config.patch restores it. * chore: unpin pytest The 9.1.1 pin from #4221 was a stopgap for the two tests fixed in this branch. With them passing under pytest 9.2, CI goes back to the latest pytest.
winglian
added a commit
that referenced
this pull request
Oct 29, 2026
* fix: pass chat_template_kwargs to the tokenizer as plain dicts * chore: bump mistral-common to 1.13.0 * feat: add Mistral Large 4 NVFP4 example * fix: use mistral-common tokenizer and packing in Mistral Large 4 example * fix: keep weight renamings when registering NVFP4 expert converters * feat: dequantize modelopt 4D block FP8 scales in Nvfp4LinearDequantize * feat: map Mistral-native w1/w3/w2 routed NVFP4 expert names * refactor: add NVFP4 MoE adapter hooks for routed proj names, conversion model type and layout * feat: resolve the NVFP4 safetensors index from config transformers_weights * feat: add Mistral Large 4 NVFP4 kernels adapter * feat: add mistral_large4 model support descriptor * fix: force non-reentrant gradient checkpointing for Mistral Large 4 * feat: add Mistral Large 4 MoE block to MOE_ARCH_BLOCK * docs: trim Mistral Large 4 README to the standard example format * chore: rename Mistral Large 4 example to nvfp4-lora * chore: shorten Mistral Large 4 example header comment * chore: drop lora kernel flags from Mistral Large 4 example * feat: fall back to the consolidated index for NVFP4 repos without config.json * feat: expose CCE hidden-state fp32 accumulation and gradient filtering options * chore: use exact CCE gradients in the Mistral Large 4 example * fix: find LoRA merge shards and write the index named by transformers_weights * fix: keep unanchored scoped weight renamings in the LoRA merge * fix: match dot-terminated weight renamings against module paths in the LoRA merge * feat: merge LoRA into modelopt FP8 weights * fix: reuse modelopt FP8 scales when re-quantizing merged weights * fix: read modelopt FP8 scales stored in another shard during the LoRA merge * fix: bound NVFP4 re-quantization code search memory * perf: snap merge-aware NVFP4 expert weights in bounded expert chunks * feat: train the Mistral Large 4 example with merge-aware NVFP4 LoRA over EP * chore: update docs * chore: add caveat in doc * docs: add Mistral Large 4 to latest updates * feat: take the base config from base_model_config in the LoRA merge * feat: ship the Mistral Large 4 NVFP4 config with the example * feat: add Mistral Large 4 NVFP4 vision LoRA example * chore: note the data_ptr cache keys to verify under FSDP2 on Blackwell * chore: pin cut-cross-entropy to v0.2.1 for Mistral Large 4 * fix: match torch's scalar-division rounding in the NVFP4 fake-quant kernel * perf: fuse the merge-aware NVFP4 expert snap into one Triton kernel * fix: disable expert parallelism when merging a LoRA trained with EP * docs: add the merge step and config note to the Mistral Large 4 README * docs: correct the dense weight placement note in the Mistral Large 4 examples * perf: cache Marlin expert weights under DDP where experts stay rank-resident * feat: target only merge-exact layers in the Mistral Large 4 LoRA examples * fix: make the Mistral Large 4 MoE weight renamings idempotent so PEFT loads expert LoRA * fix: only patch the Mistral Large 4 renamings when transformers matches the known version * fix: run the efficient LoRA merge with IEEE fp32 matmuls so tf32 configs merge exactly * fix: load modelopt NVFP4 per-tensor scales at native dtype (#4221) * fix: load modelopt NVFP4 per-tensor scales at native dtype quant_method: modelopt has no transformers quantizer, so the conversion loader cast every source to the bf16 skeleton dtype before the NVFP4 converters ran. uint8 qdata and e4m3 scales survive that, but the fp32 weight_scale_2 was rounded (up to ~0.4%), so the experts' per_tensor_scale, the dequantized non-routed linears and _nvfp4_pts no longer matched the grid the modelopt merge reads from the file. Keys routed to the NVFP4 converters now keep their checkpoint dtype for scales and uint8/fp8 weights on every load path, with or without the rank0-only patch. Float weights are still cast to the skeleton dtype. This replaces the U8/F8 special case in the rank0-only materialize patch. * test: move CPU-only NVFP4 modelopt load test to kernels_cpu The CPU lanes ignore tests/integrations/kernels/, so the test only ran on the GPU kernels lane even though it needs no GPU. tests/integrations/ kernels_cpu/ is collected by the existing tests/integrations/ argument. * chore: pin pytest to 9.1.1 CI picked up pytest 9.2.0, which breaks two CPU tests that pass on 9.1.1: - test_trajectory_tracks_dist_muon: pytest.approx now iterates the expected value, and the test passes 0-d tensors. - test_chunked_forward: torch._dynamo.config.recompile_limit is gone by the time it runs, after test_encoder_canvas monkeypatched it on the same xdist worker. Pin every place that installs pytest until those tests are fixed. (cherry picked from commit 7286a73) * chore: upgrade to transformers 5.21.0 transformers 5.21 calls `hf_quantizer.update_weight_conversions(weight_conversions, model)`, so the NVFP4-aware FP8 quantizer's override raised a TypeError for every fp8 checkpoint once installed (use_dsv4_kernels + use_scattermoe). It now forwards the extra arguments, which keeps 5.20.x working too. The `$` anchors on the expert converters stay: besides keeping `.weight` off the `.weight_scale*` keys on 5.20.x, they stop 5.21's base `update_weight_conversions` from rewriting the NVFP4 expert converters into Fp8Dequantize + `weight_scale_inv` converters. Adds a CPU test that runs a tiny DeepSeek-V4 NVFP4 meta model through `get_model_conversion_mapping` with the quantizer. (cherry picked from commit 1c76380) * chore: pin transformers 5.21.1 --------- Co-authored-by: Wing Lian <wing@axolotl.ai>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
quant_method: modelopthas no transformers quantizer, so transformers' conversion loader casts every checkpoint tensor to the bf16 skeleton dtype before Axolotl's NVFP4WeightConverters run. uint8 qdata and e4m3 group scales survive that round trip exactly. The fp32 per-tensorweight_scale_2does not: it is rounded to bf16, which is up to ~0.4% relative error in the per-tensor scale (0.25% measured on synthetic values). That rounded value ends up in:NVFP4Tensor'sper_tensor_scale;_nvfp4_pts, which merge-aware LoRA uses for its requant.The modelopt merge in
lora_merge.pyreads the exact fp32weight_scale_2from the file, so merge-aware training simulated a slightly different grid from the one the merge writes.Fix:
patch_conversion_loader_native_dtype()innvfp4_weight_converter.py, installed byregister_nvfp4_expert_converters. It wrapscore_model_loading.rename_source_keyto note when a key goes to one of our NVFP4 converters, andspawn_materialize(called next for that key) to skip the cast for that key's scales and uint8/fp8 weights..weightsources (an unquantized layer under a converter suffix) are still cast, so ordinary params keep the dtype they had before._materialize_copypatch. That patch's U8/F8 native-dtype clause is removed as redundant.weight_scalevalues are now exact as well.The test lives in a new
tests/integrations/kernels_cpu/directory because the CPU lanes ignoretests/integrations/kernels/.Motivation and Context
Maintainer-directed. Found while checking whether transformers 5.21 changes how modelopt NVFP4 MoE checkpoints load. It doesn't; this rounding is pre-existing and happens on both 5.20.1 (the pin) and 5.21.0.
How has this been tested?
tests/integrations/kernels_cpu/test_nvfp4_modelopt_load.py: a tiny synthetic modelopt NVFP4BeamConfigcheckpoint (2 layers, dense + sparse, 8 experts) loaded in four cases: with and without the rank0-only patch, each with sync and threaded loading. It asserts:per_tensor_scale, qdata and group-scale bytes;_nvfp4_pts;per_tensor_scale0.0601 vs 0.0600).-m "not gpu and not slow and not distributed_cpu" --ignore=tests/integrations/kernels/.test_nvfp4_linear_dequant,test_merge_aware_linear,test_kernels_adapters,test_nvfp4_layout_offline,test_weight_conversions,test_lora_merge_aware.Not run: torchao 0.18 (local macOS has 0.16; the CPU lane installs 0.18), the full CPU suite, and any GPU or multi-rank FSDP load. Local runs were on macOS with Python 3.11, using an uncommitted pytest shim (a
sys.monitoringstub, plus a stub for the triton-importingscattermoe_lorapackage__init__).AI Usage Disclaimer
Do not open this PR, including as a draft, until a human has reviewed the complete diff. The contributor may perform this review; a separate reviewer is not required. This also applies to maintainer-directed work.
Yes: Claude Code implemented the fix and test and ran the validation above under maintainer direction.
Types of changes
Summary by CodeRabbit