bug fix 6567139 - #2152
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:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (4)
🚧 Files skipped from review as they are similar to previous changes (2)
📝 WalkthroughWalkthroughFSDP2 now keeps non-dominant-dtype parameters replicated and moves them to the compute device. FSDP2 export now builds one reusable parameter index per mapping, preserving tied-weight resolution and reducing repeated parameter scans. ChangesFSDP2 mixed-dtype handling
FSDP2 export parameter indexing
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant Model
participant fsdp2_wrap
participant fully_shard
participant create_fsdp_param_mapping
participant build_param_index
participant get_prefixed_param_names
Model->>fsdp2_wrap: provide mixed-dtype parameters
fsdp2_wrap->>fully_shard: pass off-dtype parameters as ignored_params
fully_shard-->>fsdp2_wrap: create sharded parameters
fsdp2_wrap->>Model: move ignored parameters to the FSDP2 device
create_fsdp_param_mapping->>build_param_index: build one parameter index
build_param_index-->>create_fsdp_param_mapping: return parameter metadata
create_fsdp_param_mapping->>get_prefixed_param_names: resolve FSDP parameters with the index
get_prefixed_param_names-->>create_fsdp_param_mapping: return deterministic names
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 inconclusive)
✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
modelopt/torch/utils/distributed.py (1)
243-248: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winEmit the mixed-dtype warning only on rank 0.
fsdp2_wrapruns on every rank, butwarnings.warnis process-local. Guard the call withis_master()or use the existingwarn_rank_0helper.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelopt/torch/utils/distributed.py` around lines 243 - 248, Restrict the mixed-parameter-dtype warning in fsdp2_wrap to rank 0 by guarding the warn call with is_master() or reusing the existing warn_rank_0 helper. Preserve the current warning message and calculations.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@modelopt/torch/utils/distributed.py`:
- Around line 243-248: Restrict the mixed-parameter-dtype warning in fsdp2_wrap
to rank 0 by guarding the warn call with is_master() or reusing the existing
warn_rank_0 helper. Preserve the current warning message and calculations.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 339be468-54f7-443f-84a6-2a5d49963514
📒 Files selected for processing (4)
CHANGELOG.rstmodelopt/torch/utils/distributed.pytests/gpu/torch/utils/test_distributed.pytests/unit/torch/utils/test_distributed.py
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #2152 +/- ##
==========================================
- Coverage 78.66% 78.18% -0.49%
==========================================
Files 522 522
Lines 60420 60454 +34
==========================================
- Hits 47532 47265 -267
- Misses 12888 13189 +301
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:
|
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-5) — DM the bot to share feedback.
Small, focused fix (+198/-0 across 4 files): fsdp2_wrap now collects params whose dtype differs from the model's dominant (by element count) dtype and hands them to fully_shard(ignored_params=...), so mixed-dtype checkpoints (Nemotron-3-Nano's fp32 MoE router gates) no longer trip FSDP2's _init_mp_dtypes uniform-dtype assertion. The root cause and rejected alternatives (uniform cast, mp_policy, nested FSDP group for the off-dtype params) are well argued in the PR body, no existing helper is duplicated (no dominant-dtype utility exists in modelopt/torch/utils), the dominant-dtype computation is rank-deterministic (param iteration order is identical everywhere, so ties resolve the same on all ranks), and both a unit test for _off_dtype_params and a GPU test for the wrap are included. New files carry the project's standard NVIDIA Apache-2.0 header (year 2024, matching neighboring files rather than LICENSE_HEADER's 2026) — no licensing action needed. Overall it looks right; I'm nudging rather than approving for the verification gaps and one latent API hazard below.
Points for the owner:
- The author states the new GPU test
tests/gpu/torch/utils/test_distributed.py::test_fsdp2_wrap_mixed_dtypeshas never been run on a GPU (executed only on a 2-rank gloo/CPU stand-in), and that the original 2-node × 8-GPU Nemotron-3-Nano repro is unverified end-to-end. This is the only coverage of the actual fix, so it should be run on real NCCL before merge. fully_sharddoes not moveignored_paramsto the compute device (only managed states go through_move_states_to_device). Today's only in-repo caller isparallel_load_and_prepare_fsdp2, which materializes on device via_materialize_meta_model(and_promote_non_dtensor_to_gpuundercpu_offload), so this is fine — butfsdp2_wrapis exported in__all__, and any caller passing a CPU-resident model will now silently leave the off-dtype params on CPU and fail at forward. Worth a docstring note or an explicit.to(device); the new GPU test can't catch it because it pre-moves the model with.to(device).- The heuristic is global, not per shard group: if an entire decoder layer happened to be off-dtype it would be excluded from sharding altogether (silent memory regression, only surfaced via the
warnpercentage). Fine for the router-gate case that motivates the PR; a threshold/error above some share of elements may be worth considering. Also, thewarnfires on every rank in a torchrun job. - The CodeRabbit "Summary by CodeRabbit" block in the PR body describes many unrelated features (NVFP4 headroom calibration, TE MoE per-expert quantization, MLflow, EAGLE-3 CP fix) that are not in this 4-file diff — stale/misleading; consider trimming so reviewers aren't misled about scope.
|
Thanks — three of the four were right and are addressed in 682ee4e. Rundown: GPU test never run on a GPU. Closed. This box has sm_120 Blackwell GPUs and only one env has a torch built for them (2.8.0+cu128); that env has an unrelated broken sklearn/numpy pairing that kills pytest collection, so I drove the test function directly under
Your point that the existing test could not catch this was exactly right, so Warning fires on every rank. Fixed — now Global heuristic / threshold. Not taken, deliberately. I surveyed the affected population first: the only two families are Nemotron-3-Nano and ERNIE 4.5 MoE ( Worth noting for the archive: CodeRabbit block. Removed — you were right, it had summarized the whole 0.47 changelog section because the diff touches |
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-5) — DM the bot to share feedback.
Re-review of PR #2152 (fsdp2_wrap mixed-dtype fix, now 4 files +247/-0). The three actionable items from the previous round are addressed in code, and the new _move_to_fsdp_device + second GPU test look correct to me:
- Warning now goes through
warn_rank_0(fires once per job); the local import ofmodelopt.torch.utils.loggingis genuinely required —logging.pydoesfrom . import distributed as distat module top, so a top-level import here would be circular, and the comment says so. Fine per conventions. ignored_paramsdevice move: verified against the loader path —parallel_load_and_prepare_fsdp2builds on meta,_move_to_fsdp_deviceskips meta params, and_materialize_meta_model/_promote_non_dtensor_to_gpustill materialize them afterwards, so deferred init is unaffected.ignored_paramsis only consumed by_get_managed_statesat wrap time, so the new Parameter objects created by_applyare harmless.- Global-vs-per-shard-group heuristic: declined with a concrete survey (Nemotron-3-Nano + ERNIE 4.5 MoE routers only;
_keep_in_fp32_modulesis afrom_pretrained-only path and does not widen the population). Reasonable to leave as-is.
Points I'd still like a human to sign off on:
- 💬 Author replied: the GPU test was run on 2x RTX PRO 6000 (NCCL) in both directions, but the env's sklearn/numpy pairing breaks pytest collection, so the test function was driven directly under
torchrun— still warrants sign-off because the new tests have therefore never executed through thedist_workersfixture/pytest harness they're written against (tests/gpu/conftest.pybuilds a session-scopedDistributedWorkerPool(world_size=torch.cuda.device_count())), and the original 2-node × 8-GPU Nemotron-3-Nano repro is still unverified end-to-end. A single CI/GPU-node run oftests/gpu/torch/utils/test_distributed.pywould close this. _move_to_fsdp_devicedepends on two private torch internals (fully_shard.state(module)andstate._fsdp_param_group.device) and silently returns without moving anything if it can't find a param group — a caller with a CPU model would then fail later at forward with an opaque device-mismatch error. Consider deriving the device from a managed DTensor instead (next(p for p in model.parameters() if isinstance(p, DTensor)).to_local().device, using the already-importedDTensor), or at least warning on thedevice is Nonefall-through.- Minor: the "Summary by CodeRabbit" block the author says was removed is still present in the PR body snapshot I received and still describes unrelated 0.47 features (NVFP4 headroom, TE MoE, MLflow, EAGLE-3 CP) — worth confirming it's actually gone so reviewers aren't misled about scope.
No licensing action needed: the two new test files carry the project's standard NVIDIA Apache-2.0 header matching neighboring test files, and no third-party code is vendored.
|
All three addressed in fd24ef4. Never ran through pytest/ That is the actual Private internals + silent fall-through. Fixed, though not the way suggested. Both private accesses are gone and the fall-through now warns with the parameter count instead of returning quietly. I did not use So the device comes from CodeRabbit block. It had been removed; CodeRabbit regenerated it on the previous push, which is why your snapshot still showed it. Removed again now that the pushes are done — it is a bot artifact rather than an authored summary, so it will likely return if anyone pushes again. |
realAsma
left a comment
realAsma
left a comment
There was a problem hiding this comment.
Thanks @sugunav14 for the fix.
Note that Pytorch nightly already support mixed dtype FSDP2 wrapping. Once it lands in stable release we could skip calling _off_dtype_params.
I added a related note here:
|
RB: Can you please explain more about this FSDP2 mixed data format support in Pytorch main: #2152 (review) |
|
@sugunav14 Can we add a TODO note for #2152 (review) |
🐝 I’m checking the upstream FSDP2 change and will summarize the implications shortly. |
🐝 Follow-up on the PyTorch-main change (“Support mixed trainable dtypes in FSDP2 groups”): it does not cast the model to one dtype or leave the FP32 gates replicated. During That lets BF16 weights and FP32 router gates remain their original dtypes while each group has a single dtype for its all-gather/reduce-scatter and gradient cast-back. The FP32 gates can therefore be FSDP-managed/DTensor-sharded too; unlike this PR’s workaround, they need not be passed as So, once the PyTorch release containing that upstream change is the minimum supported version, |
|
Thanks @realAsma, @kevalmorabia97. @kevalmorabia97 — changelog to 0.46. Already on the branch; you moved it in dcc9d54 while I had the same edit staged locally, so I dropped mine rather than conflict with it. Confirmed the entry is under 0.46's Bug Fixes. @realAsma — upstream FSDP2 mixed-dtype support. Good to know, and it is the better long-term answer than this workaround. Recorded as a TODO on TODO: Drop this and shard the off-dtype params once a stable PyTorch release includes FSDP2
mixed-precision parameter dtype support (already on nightly).Worth noting for whoever picks that up: dropping Status on the one open verification item ( — zero occurrences of the |
|
Folded the FSDP2 export fix into this PR (2b3eca1 + changelog in 5ee5d6d), since it is unreachable without the dtype fix — every prior run on this checkpoint died at calibration, so the export path was never exercised.
The change is small — 32 source lines, ~8 of actual logic — plus a 121-line test file:
Measured at that scale: 1151 ms -> 5.1 ms per call (227x), i.e. ~1.9 h -> ~31 s of export. Deliberately not cached across calls — Tests: 7 new unit tests including a linear-scan oracle asserting identical results for every module, tied-weight ordering, and a regression guard counting parameter walks (verified it fails against the old behaviour). Full GPU re-run on the merged tree: 16 passed across |
FSDP2 requires every parameter in a shard group to share one dtype, but HF
models routinely pin a few parameters to fp32 for stability. Nemotron-3-Nano's
remote modeling code declares its MoE router gates float32 while the rest of the
checkpoint is bfloat16, so each decoder layer's shard group mixed dtypes and
`--use_fsdp2` PTQ died on the first calibration forward with:
AssertionError: FSDP expects uniform original parameter dtype
but got {torch.bfloat16, torch.float32}
`fsdp2_wrap` now finds parameters whose dtype differs from the model's dominant
one (by element count) and passes them to `fully_shard(ignored_params=...)`.
They stay replicated in their original dtype rather than being cast, so router
precision and the exported checkpoint are unchanged; a warning names them and
reports their share of the model. `mp_policy` is not an alternative here --
FSDP2 asserts on original dtypes regardless of it.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
Review follow-ups on the mixed-dtype FSDP2 wrap. `fully_shard` filters ignored_params out in `_get_managed_states`, so they never reach `_move_states_to_device` and stay on whatever device the caller built the model on. The in-repo caller materializes on device already, but `fsdp2_wrap` is public, so a CPU-resident model would have silently kept its off-dtype params on CPU and failed at forward. `fsdp2_wrap` now moves them onto the shards' device, read from the FSDP param group rather than guessed; meta params are skipped so the deferred-init path in `parallel_load_and_prepare_fsdp2` is unaffected. The new `test_fsdp2_wrap_moves_ignored_params_to_device` covers it by building on CPU, which the existing test cannot do since it pre-moves the model. Also switch the mixed-dtype warning to `warn_rank_0` so it fires once per job rather than once per rank (lazy import: logging imports this module). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
Review follow-up. `_move_to_fsdp_device` read `fully_shard.state(m)._fsdp_param_group.device`, two private torch internals, and silently did nothing if no param group turned up — a caller with a CPU-resident model would then have hit an opaque device mismatch at forward. It now takes the device from a sharded param's `DTensor.device_mesh`, which is public, and warns instead of returning quietly when nothing was sharded. Deliberately not the local shard's device: under `cpu_offload` the shard rests on CPU while compute stays on the accelerator, so `to_local().device` would strand the ignored params on CPU. Measured on 2x RTX PRO 6000: cpu_offload=False -> to_local()=cuda:0, device_mesh.device_type=cuda cpu_offload=True -> to_local()=cpu, device_mesh.device_type=cuda `test_fsdp2_wrap_moves_ignored_params_to_device` is parametrized over `cpu_offload` to pin this down; the local-shard variant fails the offload case. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Review follow-up from @realAsma: PyTorch nightly already supports mixed-dtype FSDP2 wrapping, so record a TODO to drop this workaround and shard the off-dtype params once that reaches a stable release. Worded to match the existing note in tensor_quantizer.py's enable_lsq. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
`create_fsdp_param_mapping` resolved each FSDPParam's module by walking every `model.named_parameters()`, and export calls it once per quantized module. That is quadratic in (parameters x modules), which is harmless for dense models and intractable for a large MoE. Exporting Nemotron-3-Nano-30B-A3B (6,243 parameter tensors, 6,004 quantized modules, ~256 FSDPParams per MoE layer group) produced no output for 3h35m with every GPU at 0% util. py-spy showed every sample `active+gil` in get_prefixed_param_names, so it was never a stalled collective -- just ~9.6e9 Python-level id comparisons. `build_param_index` now maps `id(param) -> (position, name)` once per mapping call, and `get_prefixed_param_names` takes it as an argument. Position ordering preserves the old "first in named_parameters() order" result, which matters for tied weights. At the above scale this is 1151 ms -> 5.1 ms per call, i.e. 1.9h -> 31s of export. The index is deliberately not cached across calls: `fsdp2_aware_weight_update` swaps in quantized parameters and rebuilds FSDPParams between them, so a shared map would go stale. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
5ee5d6d to
fdb3dc4
Compare
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
kevalmorabia97
added
the
cherry-pick-done
## Cherry-picked PRs - #2172 - #2087 - #2152 - #2060 - #2008 - #2194 <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added streaming Hugging Face checkpoint export for disk- and CPU-offloaded models, including sharded safetensors output. * Added a PTQ recipe for NVFP4 expert quantization with FP8 KV-cache support and layerwise offload. * Added support for additional Nemotron-H model layouts and more reliable conversation input handling. * **Bug Fixes** * Improved FSDP2 handling of mixed parameter data types. * Fixed tied-weight deduplication and checkpoint export consistency. * **Documentation** * Added a unified deployment support matrix with updated framework requirements, model coverage, quantization guidance, and hardware notes. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Signed-off-by: Zhiyu Cheng <zhiyuc@nvidia.com> Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com> Signed-off-by: Jennifer Chen <jennifchen@nvidia.com> Signed-off-by: Fridah-nv <fridah@nvidia.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Signed-off-by: Juhi Mittal <juhim@nvidia.com> Co-authored-by: Zhiyu <zhiyuc@nvidia.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: sugunav14 <178320438+sugunav14@users.noreply.github.com> Co-authored-by: Jenny Chen <jennifchen@nvidia.com> Co-authored-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com> Co-authored-by: Juhi Mittal <39641197+juhi10071998@users.noreply.github.com>
4 participants
What does this PR do?
Type of change: Bug fix
Fixes nvbug 6567139:
--use_fsdp2PTQ of Nemotron-3-Nano-30B-A3B dies on the first calibration forward withRoot cause.
fsdp2_wrapcallsfully_shardon each decoder layer, and FSDP2 requires every parameter in a shard group to share one dtype. Nemotron-3-Nano's remote modeling code pins the MoE router gate to fp32 while the rest of the checkpoint is bf16:so every MoE layer's param group is
{bf16, fp32}andFSDPParamGroup._init_mp_dtypesasserts. The crash surfaces during calibration only because FSDP2'slazy_initruns at the first forward — the model is already unwrappable atfully_shardtime. Nothing about quantization is involved;--use_fsdp2on this checkpoint fails regardless of recipe. This only bites under--trust_remote_code: transformers' nativenemotron_hbuilds fully bf16.Fix.
fsdp2_wrapnow finds parameters whose dtype differs from the model's dominant one (by element count) and passes them tofully_shard(ignored_params=...). They stay replicated in their original dtype instead of being cast, so router precision and the exported checkpoint are unchanged. A warning names them and reports their share of the model — for Nemotron-3-Nano that is 23 fp32 MoE router gates (backbone.layers.N.mixer.gate.weight, 128 experts x 2688 hidden), ~30 MB replicated per rank against a 30B model.Casting to a uniform dtype was rejected because it would change both calibration routing and the exported weights.
mp_policyis not an alternative: FSDP2 asserts on original dtypes regardless of it.Replication was chosen over giving the off-dtype params their own nested FSDP group (which would shard them) because router gates are
[n_experts, hidden]— ~2 MiB each, 30-100 MiB total across every affected model — so sharding them would add a latency-bound all-gather per MoE layer on a tensor whose dim-0 (64-128 experts) cannot even split cleanly across ranks.fully_shardfilters ignored params out in_get_managed_states, so it never moves them to the compute device.fsdp2_wraptherefore moves them itself, reading the device off the FSDP param group rather than guessing; meta params are skipped soparallel_load_and_prepare_fsdp2's deferred init is unaffected. The mixed-dtype warning goes throughwarn_rank_0so it fires once per job, not once per rank.Second fix: FSDP2 export param mapping
Folding in a second, independent bug found while verifying the first.
create_fsdp_param_mappingresolved eachFSDPParam's module by scanning all ofmodel.named_parameters(), and export calls it once per quantized module — quadratic in (parameters x modules). Harmless for dense models, intractable for a large MoE.Exporting Nemotron-3-Nano-30B-A3B (6,243 parameter tensors, 6,004 quantized modules, ~256 FSDPParams per MoE layer group) produced no output for 3h35m with every GPU at 0% util. py-spy showed every sample
active+gilinget_prefixed_param_names, so it was CPU burn, not a stalled collective — roughly 9.6e9 Python-level id comparisons.build_param_indexnow mapsid(param) -> (position, name)once per mapping call. The position ordering preserves the previous "first innamed_parameters()order" result, which matters for tied weights. Measured at that scale: 1151 ms -> 5.1 ms per call (227x), i.e. ~1.9 h -> ~31 s of export.The index is deliberately not cached across calls:
fsdp2_aware_weight_updateswaps in quantized parameters and rebuildsFSDPParams between them, so a shared map would go stale.This was unreachable before the dtype fix — every prior run on this checkpoint died at calibration — which is why the two ship together.
Usage
No API change — the existing command now works:
Testing
tests/unit/torch/utils/test_distributed.py(new): 4 tests for_off_dtype_params, including that "dominant" is by element count rather than parameter count. Passing.tests/gpu/torch/utils/test_distributed.py(new):test_fsdp2_wrap_mixed_dtypeswraps a model carrying an fp32 router gate, forwards it, and checks the fp32 parameter stays non-DTensor, fp32, and on the compute device alongside the shards.test_fsdp2_wrap_moves_ignored_params_to_devicebuilds the model on CPU and checks the ignored params are moved onto the shards' device.dist_workersfixture: 3 passed, includingcpu_offload=True. Without the fix the wrap raisesAssertionError: FSDP expects uniform original parameter dtype.set_model_state_dict(..., full_state_dict=True)into a layer mixing sharded DTensor and ignored plain params writes both correctly, which is what_broadcast_load_groupdoes per decoder layer.tests/unit/torch/utils/: 155 passed.pre-commitclean.Before your PR is "Ready for review"
CONTRIBUTING.md: N/AAdditional Information
fully_shard(ignored_params=...)requires torch >= 2.7; the repo already pinstorch>=2.8.🤖 Generated with Claude Code
Summary by CodeRabbit
Bug Fixes
Performance
Tests