Repository navigation
[OMNIML-6083] [16/21] Write a sharded model's HF checkpoint without gathering it - #2752
shengliangxu wants to merge 19 commits into
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. |
|
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:
Walkthrough
Merge Risk:
 |
| Check name | Status | Explanation | Resolution |
|---|---|---|---|
| Docstring Coverage |  |
Docstring coverage is 58.33% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 48 functions across 6 files. | Write docstrings for the functions missing them to satisfy the coverage threshold. |
-
Check name Status Explanation 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. Security Anti-Patterns PASS. The PR changes only modelopt/torch/export/distributed_writer.pyandmodelopt/torch/export/unified_export_hf.pyin the package; their added lines contain no `torch.load(..., weights_only=Fals…Title check The title clearly and concisely describes the main change: exporting a sharded model's Hugging Face checkpoint without gathering it. Description Check Check skipped - CodeRabbit’s high-level summary is enabled. Full details: Security Anti-Patterns
-
Explanation
PASS. The PR changes only
modelopt/torch/export/distributed_writer.pyandmodelopt/torch/export/unified_export_hf.pyin the package; their added lines contain notorch.load(..., weights_only=False), pickle-enabled NumPy load, hardcodedtrust_remote_code=True,eval,exec, or# nosec. The PR does not modifypyproject.tomlor requirements files, and it adds no new PIP dependency.
-
✨ Finishing Touches: 💡 2
-
📝 Generate docstrings: 💡
-
- Commit to this branch
- Create a new PR
🧪 Generate unit tests (beta)
-
- Commit to this branch
- Create a new PR
🛠️ Fix failing CI checks: 💡
-
- Commit to this branch
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
4810e33 to
72eab4a
Compare
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## shengliangx/hfep-12b-writer-pieces #2752 +/- ##
======================================================================
+ Coverage 62.05% 62.13% +0.08%
======================================================================
Files 656 656
Lines 73481 73584 +103
======================================================================
+ Hits 45601 45725 +124
+ Misses 27880 27859 -21
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:
|
|
/ok to test 72eab4a |
72eab4a to
a30a335
Compare
|
/ok to test a30a335 |
shengliangxu
changed the title
a30a335 to
751b1d1
Compare
751b1d1 to
717b73e
Compare
|
/ok to test 717b73e |
shengliangxu
changed the title
717b73e to
446370c
Compare
shengliangxu
changed the title
446370c to
755ebf2
Compare
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Claude review — [16/21] no-gather distributed HF checkpoint writer (third pass)
Findings: CRITICAL 0 · IMPORTANT 1 · SUGGESTION 2
Scope: bare /claude review, so a full review. 8 files changed (991+/116−). Reviewed in full: modelopt/torch/export/distributed_writer.py, modelopt/torch/export/unified_export_hf.py, all three test diffs (tests/unit/torch/export/test_distributed_writer.py, tests/gpu/torch/export/test_distributed_hf_export.py, test_unified_export_hf.py / test_fused_experts_keep_fused.py), CHANGELOG.rst, examples/hf_ptq/README.md. Supporting context read outside the diff: DistributionPlan/distribution_plan/broadcast_from/all_reduce in modelopt/torch/utils/distributed.py, resolve_distribution_plan in model_load_utils.py, _materialize_owned/_revert_whole_sd/_write_even_shards, _split_local_fused_module and the consumed set in moe_utils.py, _prepare_model_for_export, save_non_weight_artifacts, _revert_quant_config_names_best_effort. No prompt-injection or review-steering text in the diff.
Resolved since the previous round
- The reverse-conversion flag is threaded out and the quant-config revert is gated on it (
if reverted:), so weight names andhf_quant_config.jsoncan no longer disagree.test_nogather_export_names_modules_in_the_quant_config_as_the_weights_doassertsexclude_modules ⊆the written module names. - The expert
_input_scalereduce dropped its backend-string device branch and now reduces on the tensors own device, andtest_two_ranks_write_every_expert_with_the_largest_input_scalecovers it under gloo — the code/test contradiction is gone. - The
isinstance(model, nn.Module)guard keeps a diffusers pipeline offdistribution_plan(). - The GPU test docstring no longer describes a DCP staging step, and the unfailable
sharded/assertion is gone.
Most impactful (new this round)
[IMPORTANT Performance] _to_normal copies every DTensor unconditionally, so peak is the sharded model twice (inline) — the plain branch is guarded by v.is_inference(), the DTensor branch is not, and the copy is additive rather than a replacement: get_model_state_dict returns detach()-ed views (as the wrapper docstring itself notes), so models parameters still own the originals and overwriting sharded_sd[_k] frees only the view. That makes the in-place-vs-comprehension rationale below it self-defeating — both forms are ~2x, and the copy this loop creates is the second one — and contradicts the functions "Peak memory per rank is its owned set, not the model". On the documented GPU usage a model sized to fit through calibration now needs 2x its shard at the first step of the export. Guarding the DTensor branch removes the cost outright for the non-inference-mode callers (every _nogather_export unit test here, and any QAT flow); for the inference-mode case most of the copy is redundant with the clones _split_local_fused_module and _gather_shard0_to_owners already make, and at minimum the docstring and comment should state the bound the code actually has.
Suggestions (non-blocking)
- The route is selected on
is_fsdp_densealone, so an expert-parallel plan reaches a writer that declares EP out of scope (inline) —resolve_distribution_plancan returnep_size > 1withdense="fsdp", andep_rank/dp_idxare hardcoded to 0. It does fail rather than mis-write (the per-expert keys collide and_write_even_shardsraises on every rank), but incidentally and with a message that names the writer rather than the plan. - Unused
pi/_tied(inline).
Carried over, unaddressed — raise again only if you intend to act
_write_base_configstill duplicatessave_non_weight_artifactsminus the trailing copy; the two concrete deltas are the_patch_revert_weight_conversionwrapper and thecan_generate()guard.- Dead EP bookkeeping (
ep_rank, dp_idx = 0, 0,write_experts) — the inline suggestion above is the other half of this one. - The expert
_input_scaleall_reduceis still sized from per-rank keys behindif in_keys:, ahead of_materialize_owneds divergence check; unreachable under theis_fsdp_denseplan this writer targets.
What checked out
- Mixed-precision KV cache goes through
_get_kv_cache_postprocess_configandmaxbound=448matcheskv_cache_max_boundon the gather and both streaming paths;_FP8_MIXED_KV_CFGcovers it end to end. _is_expert_keys suffix tuple matchesmoe_utils.py:672sconsumedset exactly, for both the gated (gate_up_proj) and ungated (up_proj)first, so no fused tensor is written twice under its fused name and none is dropped. Keys outside both sets (e.g. a fused_bias) land indense_sdunder the fused name, which is what the single-process path does too.- The input-scale reduce precedes the fold-sensitive steps and only touches activation scales, so no weight is quantized against a scale the checkpoint does not carry. Dense scales get no reduce, but owner-takes-all for those is pre-existing to the FSDP2 streaming writer, not new here.
_assert_byte_identicalcompares dtype and shape as well as bytes, so thereshape(())/.float()in the reduce is covered against the single-process reference rather than slipping through a raw-byte comparison._revert_quant_config_names_best_effortreads onlymodel.named_modules()— no collective — so calling it on rank 0 alone is safe, matching the streaming path.extra_state_dictsingle ownership:_writes_extramatches the writers rank-0 merge,check_no_overlapruns insideall_reduceso a rank-0 collision raises everywhere, and the merge happens after the reverse so hub-named carried tensors are not renamed._modelopt_carried_over_namesis resolved on every rank, soget_quant_configstill agrees.broadcast_fromis short-circuited byis_fsdp2_sharded, which comes from the recorded plan and is therefore identical on every rank — no rank enters a broadcast the others skip.- The
save_modelopt_staterefusal is raised beforepack_quantized_weightsmutates the model, and before thetry, so a refusal leaves the model as it was and no rank is left in a collective. - File-write ordering is safe:
copy_non_model_filesskips_EXPORT_OWNED_FILESand files the export already wrote, so_copy_non_model_files_from_sourcecannot clobber the embeddedquantization_config. - The GPU test reads safetensors headers directly (argument evaluation order in the length-then-header read is correct), expects per-expert shapes derived from the source, halves the last dim for packed U8, and asserts no
lm_head.weightfor the tied case via theunexpectedcheck. - Dropping
test_fused_experts_fold_in_a_weight_window...for the Qwen2-MoE CPU case is a net gain: the new test asserts the observable effect (experts written quantized with the shared expert packed after them) rather than the monkeypatched call sequence. - PR size within budget: 352 insertions outside
tests/anddocs/. CHANGELOG entries are user-facing and filed under the right sub-sections, including theset_distribution_planrequirement under Backward Breaking Changes.
Risk
Moderate-low. The writer stays opt-in behind MODELOPT_FSDP2_EXPORT=nogather and the default export is untouched, and the structural failure modes (dropped key, mis-assembled shard, divergent ranks, duplicate name, missing index) are each asserted or explicitly refused — the correctness picture is in good shape after two rounds of fixes. The remaining IMPORTANT is not a wrong checkpoint but a memory bound: it doubles per-rank resident weights in the one path whose reason to exist is fitting a model that the gather path cannot, and the surrounding comments assert the opposite, so it is worth closing (or at least restating accurately) before merge.
🤖 Generated with Claude Code
|
/ok to test cc5eba8 |
|
/ok to test 8b2d1df |
|
/claude review |
| ep_rank, dp_idx = 0, 0 | ||
| # Total config expert count, to tell a per-EP-group expert tensor (size < total) from one that | ||
| # spans all experts (classic DP x EP) when computing the per-expert global index offset. | ||
| _cfg = getattr(model, "config", None) | ||
|
|
||
| def _experts_of(c): | ||
| if c is None: | ||
| return 0 | ||
| return int( | ||
| getattr(c, "num_experts", 0) | ||
| or getattr(c, "num_local_experts", 0) | ||
| or getattr(c, "n_routed_experts", 0) # nemotron_h | ||
| or 0 | ||
| ) | ||
|
|
||
| # VLMs nest the text-model fields under ``config.text_config`` (the quantized backbone is | ||
| # ``model.language_model``), so fall back to it when the top-level config has no expert count. | ||
| total_experts = _experts_of(_cfg) or _experts_of(getattr(_cfg, "text_config", None)) | ||
| # DP x EP: experts are dp-REPLICATED across the EDP group (each holds a full EP-sharded copy). Only | ||
| # the first replica (dp_idx == 0) writes them; replica ranks (dp_idx > 0) skip the per-expert split. | ||
| # dp_idx is the EDP-group rank resolved above -- 0 when experts are not EDP-replicated (dense / | ||
| # FSDP2 / EP-without-DP), so every such rank writes. | ||
| write_experts = dp_idx == 0 |
There was a problem hiding this comment.
[SUGGESTION] This EP scaffolding has no runtime effect on this path, and one of its comments states something that is not true.
ep_rank and dp_idx are literals, so:
write_experts = dp_idx == 0is alwaysTrue— theif write_experts:guard below and the "Skipped on dp-replica ranks under DP x EP" note in the docstring describe a branch that cannot be taken here.total_expertsis never read._split_local_fused_moduleonly consults it insideadd_ep_base = bool(ep_rank) and group_size > 0 and (total_experts == 0 or group_size < total_experts), andbool(0)short-circuits that toFalseunconditionally. So_cfg,_experts_of, thetext_configVLM fallback and the two comment blocks explaining the per-EP-group vs all-experts distinction are all dead on this path.- The comment "dp_idx is the EDP-group rank resolved above" is wrong — nothing resolves it; it is the literal two lines up.
Why it matters: ~25 lines of code and comments describe an EP-hybrid routing that the function's own docstring says is out of scope, so a reader has to re-derive that ep_rank == 0 kills all of it. _experts_of is also a latent trap — int(getattr(c, "num_experts", 0) or ...) raises on a config where num_experts is a list, for a value nothing consumes.
Suggested fix: drop _cfg/_experts_of/total_experts/write_experts/dp_idx, call _split_local_fused_module(prefix, sharded_sd) with its defaults, and keep one sentence saying that this writer is FSDP2-only so the global expert index comes from the DTensor's own dim-0 offset. The EP-hybrid resolution can arrive with [17/21], where it is actually exercised.
| in_keys = sorted(k for k in sharded_sd if k.endswith("_input_scale") and _is_expert_key(k)) | ||
| if in_keys: | ||
| stacked = torch.stack([sharded_sd[k].detach().reshape(()).float() for k in in_keys]) | ||
| dist.all_reduce(stacked, op=dist.ReduceOp.MAX) | ||
| for j, k in enumerate(in_keys): | ||
| sharded_sd[k] = stacked[j].clone() |
There was a problem hiding this comment.
[SUGGESTION] This is the only collective in the writer whose participation and payload size are decided from rank-local state, and it runs ahead of the divergence diagnostics that exist precisely to turn this class of mistake into an error.
in_keys is derived from this rank's own sharded_sd keys via _is_expert_key, which depends on prefix_set — itself built locally by _fused_experts_prefixes(sharded_sd). So:
- a rank with an empty
in_keyswhile its peers have some skips theall_reduceentirely and walks into_materialize_owned'sall_gather_object, while the peers sit in theall_reduce— a hang with no message; - a rank whose
in_keyshas a different length reduces a differently-shapedstacked— again a hang or an opaque NCCL/gloo size error.
_materialize_owned already catches exactly this ("dense key sets differ across ranks, so the export would deadlock", "ranks disagree on which tensors are dim-0 sharded"), and because expert keys are the complement of dense keys, a divergent expert classification would be named by that check — but only after this all_reduce has already hung.
In practice get_model_state_dict(full_state_dict=False) gives every FSDP2 rank the same key set, so this is unlikely to fire today. It is still the one place in this file that can hang instead of diagnosing, which is the standard the rest of the module holds itself to (and CLAUDE.md's pre-push rule on rank-dependent work starting a collective).
Suggested fix: fold in_keys into the single all_gather_object metadata exchange that _materialize_owned already performs — or, if the reduce has to stay here, gather the key lists first, raise on disagreement, and then reduce over the agreed list unconditionally (an empty list on every rank skipping it is fine, since that decision is then global).
| # How the model is distributed decides how it is exported. A diffusers pipeline is not a module, | ||
| # and nothing here distributes one. | ||
| plan = distribution_plan(model) if isinstance(model, nn.Module) else DistributionPlan() | ||
| _writes_extra = not plan.is_fsdp_dense or torch.distributed.get_rank() == 0 |
There was a problem hiding this comment.
[SUGGESTION] Switching from is_fsdp2_model(model) to plan.is_fsdp_dense drops the is_available() and is_initialized() guard that used to protect this get_rank() call (and the one at line 2015).
DistributionPlan.__post_init__ accepts DistributionPlan(dense="fsdp") with the default world_size=1, and set_distribution_plan records it without consulting the process group. A model carrying such a plan with no live default group makes export_hf_checkpoint die on this line with RuntimeError: Default process group has not been initialized, before it has done anything — rather than exporting as a single process.
Realistically unreachable today (FSDP2 fully_shard needs a mesh, so a genuinely sharded model has a group), but the old expression could not fail this way and the new one can. broadcast_from two lines down already handles the non-distributed case gracefully, so the asymmetry is easy to close:
_distributed = torch.distributed.is_available() and torch.distributed.is_initialized()
_writes_extra = not (plan.is_fsdp_dense and _distributed) or torch.distributed.get_rank() == 0Alternatively, have DistributionPlan.__post_init__ (or set_distribution_plan) reject dense="fsdp" outside a distributed run, so an impossible plan is refused where it is recorded instead of surfacing at export.
|
Claude review — [16/21] No-gather sharded HF checkpoint write Scope: full review (the trigger comment carried no scoping instructions). 8 changed files, +1011/−117. Reviewed Findings — CRITICAL: 0 · IMPORTANT: 0 · SUGGESTION: 3
All three are non-blocking. What checked out
One note out of scope The plan requirement is a real break for anyone using the ModelOpt Risk: low. The new path is opt-in behind 🤖 Generated with Claude Code |
|
/ok to test 0b46535 |
distributed_save_hf_checkpoint writes an already-processed sharded model as HF safetensors with no rank-0 gather: fused experts are split shard-locally, every other tensor moves whole to one owner rank, the reverse conversion runs there, and each rank writes its own shards. No rank holds the whole model. A model sharded by FSDP2 alone opts in with MODELOPT_FSDP2_EXPORT=nogather; the expert-parallel export routes to it in the next change. On that path export_hf_checkpoint packs with keep_fused_experts; writes config.json, a sanitized generation_config.json and hf_quant_config.json on rank 0; and copies the source's non-model files (tokenizer, processor, remote code, chat templates), which the writer never sees. It rejects save_modelopt_state, which the path cannot honor without save_pretrained, rather than drop it silently. A model built from a config has an empty config._name_or_path, which as a path is the working directory, so the export copied the working directory's tokenizer, processor and code files as if they were the model's. An empty name now means no source checkpoint. The output matches the single-process export tensor for tensor (CPU tests: FP8 and NVFP4, one and two ranks, root-wrapped and layers-only FSDP2). A GPU test PTQs a tiny Qwen3 (dense and MoE) under FSDP2 with FP8 and NVFP4, opting each export into the writer, and checks the checkpoint against the safetensors headers. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
export_hf_checkpoint takes the FSDP2 routes, streaming or no-gather, when the model's recorded plan FSDP2-shards it, and reads the carried-over weights on rank 0 alone for the same reason, rather than looking for an FSDP2 module. A diffusers pipeline is not a module and is never distributed, so it is exported without a plan, as a new test checks. The writer tests record the plan where they shard a model with fully_shard themselves. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
fsdp2_wrap records no plan; load_pretrained does. These tests wrap the model by hand, and the export picks its route from the plan, so each records the FSDP2 plan it wrapped the model by. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Batch sizing and the HF export read how a model is distributed from the plan load_pretrained records, and raise for an FSDP2-sharded or multi-rank model without one, where they used to look at the model and the process group. Users who distribute a model themselves, fsdp2_wrap included, now record its plan. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
main opened 0.49.0 for what merges after 0.48.0's cut, which this PR does. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
It spawned two ranks, about 14 s of the linux unit job. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…iter test The writer now raises it on every rank, so rank 0 no longer has to release rank 1 from the closing barrier; the dedicated test bounds that case, so the two-rank job no longer repeats it. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The window around the fold was guarded only by a test of two stubs' call order. The MoE no-gather export test used Qwen3-MoE, whose layers pack nothing after their experts, so it exported them quantized with or without the window. The window matters when a dense weight in the same FSDP2 unit is packed after the fold: that weight's own window unshards and reshards the unit, which puts FSDP2's original shards back over a fold FSDP2 does not track. Qwen2-MoE's shared expert is such a weight, so the test now exports Qwen2-MoE, which fails without the window, and the stub test goes. The GPU MoE export now requires every expert to be exported quantized, not just some weight. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Every GPU test of the no-gather export needs two GPUs, which the PR runner does not have. What they checked runs on CPU instead, through the same writer and export. The two-rank write is now handed tensors the model does not own on every rank, and must write them once. A linear left unquantized already is checked: every no-gather equivalence case exports the lm_head unquantized, byte for byte as the single-process export does. The GPU completeness test ties the dense model's embeddings, so its expected keys also check that the alias is dropped. The GPU tests these cover go, with the GPU refusal of save_modelopt_state, which the unit test of the same refusal covers. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…vered ones Merged: - The inference-mode and extra_state_dict writes become cases of test_distributed_save_writes_the_hub_checkpoint_of_a_sharded_moe. - test_a_model_built_from_a_config_names_no_source_checkpoint folds into test_source_checkpoint_is_none_for_a_model_built_from_config, which now also checks that nothing is copied from the working directory. - test_materialize_owned_leaves_every_tensor_whole_and_plain moves into the two-rank job, where ownership is real: its tensors cycle through dim-0 shards, replicated DTensors and plain tensors, plus a zero-row tensor. Dropping, duplicating or not unwrapping any of them, treating experts as dense, or writing the zero-row tensor as float32 fails the job. Dropped, each caught by a remaining test when its code is broken: - test_revert_whole_sd_keeps_a_state_dict_with_nothing_to_revert: wiping a dense state dict fails test_nogather_export_writes_what_the_single_process_export_writes. - test_revert_whole_sd_renames_to_hub_names_in_place: discarding the reverse fails the hub-checkpoint test. - test_a_pipeline_is_exported_without_a_distribution_plan: reading a pipeline's plan fails test_export_diffusers_pipeline_config on a real tiny pipeline. The two-rank job checks the consistency errors with try/except inside bounded_collectives rather than pytest.raises, which takes the shared worker down when nothing is raised. Kept, since no remaining test fails when their code is broken: the split's storage, dense pass-through and prefix tests and the handler test in test_fused_experts_keep_fused.py, the mesh-memo test, and test_write_even_shards_writes_even_shards_and_an_index_of_every_tensor, the only check that written tensors leave the rank's set, of the file count, and of a non-contiguous tensor. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The no-gather writer took its device from torch.cuda.current_device() on every call, though only the expert input-scale all-reduce needs it. On CPU torch that raised, so every writer unit test patched current_device. Take the device from the backend: the CPU under a gloo-only group, the current CUDA device otherwise (distributed_run sets up cpu:gloo,cuda:nccl). The tests drop the patch, and fail on CPU torch without the change. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
distributed_save_hf_checkpoint checked extra_state_dict against the model's keys on rank 0 only, outside any collective. A collision raised on rank 0 while the other ranks went on into the write's first collective and waited there until the timeout, then raised an unrelated timeout. Run the check inside all_reduce, so every rank raises rank 0's error at once. The two-rank writer job covers it, and now checks rank 0's checkpoint bytes after its last collective, so a failed check cannot leave rank 1 waiting either. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
_write_hf_export_config rewrote generation_config.json, which every export route has already written by then: save_pretrained on the plain path, save_non_weight_artifacts on the streaming paths and _write_base_config, behind can_generate(), on the no-gather path. The layerwise path returns before it. Main's helper never wrote it. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…r not The overlap check ran its all_reduce only where extra_state_dict was non-empty, but a sharded export hands the tensors it carries over (an MTP head, an auxiliary tower) to rank 0 alone. Rank 0 then entered the overlap collective while the others went on to "shard counts", and every sharded export carrying weights failed with "Ranks reduced out of step". Every rank joins now; only rank 0, which holds the extras, checks and adds them. The two-rank writer job hands its extras to rank 0 only, as the export does. Found by the epdp port. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The comment credited is_fsdp2_model with choosing the writer; the recorded plan and the MODELOPT_FSDP2_EXPORT opt-in choose it. The assertion stays: it checks the model is really sharded, so the test exercises the shard writer. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…rank 0's writer choice The no-gather export passed the raw kv_cache_quant_algo to the writer, so a mixed-precision KV cache reached postprocess_state_dict as the "MIXED_PRECISION" sentinel and failed its format assert. It now resolves the per-layer map through _get_kv_cache_postprocess_config, as the gather and streaming exports do; the single-process comparison covers an FP8 and NVFP4 KV cache. Which writer runs is rank 0's MODELOPT_FSDP2_EXPORT, broadcast: the two writers issue different collectives. The two-rank export sets the knob on rank 0 only. The no-gather branch drops a save_modelopt_state write it refused earlier, and _write_base_config says why a config save needs no revert_weight_conversion patch. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The no-gather writer keeps the in-memory tensor names on every rank when any rank cannot reverse the conversion mapping (a merge group split across ranks), but export_hf_checkpoint renamed the quant config to hub names regardless, so exclude_modules named modules the checkpoint did not hold and a loader would read an excluded layer as quantized. distributed_save_hf_checkpoint now returns whether the names were reverted, and the config is renamed only when they were. The experts' input-scale all-reduce runs on the scales' own device instead of one picked from the backend name, and the CPU tests now cover it: the one-rank Qwen2-MoE export quantizes expert inputs (its comment called the reduce CUDA-only), and a two-rank export whose ranks calibrated different input amaxes must write the largest on every expert. Also drops a comment that described no code, and the GPU test's description of a staging step this writer never had, with the assertion that nothing left that staging directory behind. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The no-gather writer copied every DTensor's local shard up front to drop the inference flag, though only an inference tensor carries it, and its comments called the copy free: the state dict's values are views of the model's own parameters, so each copy adds to the rank's memory instead of replacing a value. A DTensor is now copied only when its local shard is an inference tensor, as a plain tensor already was, so a caller outside inference mode (hf_ptq exports under no_grad) copies nothing. On two gloo ranks exporting a 223 MiB bf16 Llama, the writer's peak above each rank's 112 MiB shard drops from 278 and 246 MiB to 224 and 217 MiB. The docstrings and comments now state the bound the code has. Also drops an unused loop index. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Since [8/21], get_max_batch_size measures a model with no recorded plan per rank, as it did before, so only export_hf_checkpoint raises for a sharded or multi-rank model that carries no plan. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
|
/ok to test 3b132ec |
What does this PR do?
Type of change: New feature, with a bug fix for models built from a config
distributed_save_hf_checkpointwrites an already-processed sharded model as HF safetensors with no rank-0 gather:No rank ever holds the whole model. On top of its own shard, which the model keeps, a rank holds what it writes; only an inference tensor is copied up front, so an export under
no_grad, as hf_ptq runs it, copies nothing. A model sharded by FSDP2 alone opts in withMODELOPT_FSDP2_EXPORT=nogather; the expert-parallel export routes to it unconditionally in [17/21].Two gloo ranks exporting a 223 MiB bf16 Llama, the writer's peak RSS above each rank's 112 MiB shard:
export_hf_checkpointon this path:keep_fused_experts;config.json, a sanitizedgeneration_config.jsonandhf_quant_config.jsonon rank 0.hf_quant_config.jsonnames modules as the weights do: when any rank cannot reverse the conversion mapping, every rank's weights keep their in-memory names, and the quant config keeps them too (distributed_save_hf_checkpointreturns whether the names were reverted);It rejects
save_modelopt_state, which this path cannot honor because it never runssave_pretrained, rather than dropping it silently.The export reads how the model is sharded from its distribution plan ([4/21]), not from an FSDP2 module in it: the FSDP2 routes, streaming or no-gather, and which rank reads the carried-over weights. A model FSDP2-sharded with
fsdp2_wrapor loaded throughload_pretrainedcarries its plan. One wrapped by hand records it withset_distribution_plan, and without it the export raises rather than treat it as whole. A diffusers pipeline is not a module and is exported without a plan.The KV-cache format goes through
_get_kv_cache_postprocess_config, so a mixed-precision KV cache exports as on the other paths, and which writer runs is rank 0'sMODELOPT_FSDP2_EXPORT.Usage
Testing
tests/unit/torch/export/test_distributed_writer.pyadds end-to-end cases on CPU under gloo:save_pretrained, tied and untied;extra_state_dictkept under its own names;export_hf_checkpointon the no-gather path, root-wrapped and layers-only: byte-identical to the single-process export, plus config files and the copied source sidecar, for FP8 and NVFP4;exclude_modulesas in the weights;save_modelopt_staterejected.tests/unit/torch/export/test_unified_export_hf.py: a model built from a config, which main's Speed up CPU unit tests to stay under the 20 min CI limit #2823 gives no source checkpoint, copies nothing from the working directory at export.test_export_diffusers_pipeline_configon a real tiny pipeline. Only rank 0 reading the carried weights is not unit-tested, since it needs two FSDP2 ranks; the GPU carry-over tests cover it.tests/gpu/torch/export/test_distributed_hf_export.py: PTQs a tiny Qwen3 (dense, with tied embeddings, and MoE) under FSDP2 with FP8 and NVFP4, opting each export into the writer, and checks the checkpoint against the safetensors headers: every routed expert exported quantized, and the tied head written once. Its earlier version passed on a 4-GPU node, 8 tests, including checks of extras, ties, refusals and unquantized layers that now run on CPU above, where every PR runs them. On 1 GPU, all skip. Its docstring no longer describes a DCP staging step this writer never had, and an assertion on asharded/staging directory, which could not fail, is dropped; not re-run on GPU after that change.test_export_diffusers.pytest included. Each drop is checked by mutation.tests/unit/torch/utils,export,models, the fused-experts quantization tests and the hf_ptq example tests) on transformers 5.18 and 5.5 and on torch 2.8. They pass apart from the diffusers-export and GGML tests, which fail the same way on main in that environment. CI runs the full unit suite.Before your PR is "Ready for review"
load_pretrained,fsdp2_wrapincluded, must record its plan withset_distribution_planbeforeexport_hf_checkpoint; without it the export raises.CONTRIBUTING.md: N/Aget_max_batch_sizeneed a distributed model's plan recorded.Additional Information
Part 16 of 21 of expert-parallel PTQ for Hugging Face MoE models; merges after [15/21].
Summary by CodeRabbit
MODELOPT_FSDP2_EXPORT=nogather, ranks write checkpoint shards in parallel without gathering the full model on one host; MoE experts remain sharded.