Repository navigation
[OMNIML-6064] Load each rank's checkpoint shard straight into a meta-built model - #2739
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:
No actionable comments were generated in the recent review. ℹ️ Recent review info
Walkthrough
Merge Risk:
 |
| 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 | |
The PR changes only modelopt/torch/utils/plugins/model_load_utils.py and a test file; no examples, dependency files, or requirements changed. The added model-loading code uses dcp.load, not `torch… |
| Title check | |
The title clearly and concisely describes the main change: loading each rank's checkpoint shard directly into a model built on the meta device. |
| Description Check | |
Check skipped - CodeRabbit’s high-level summary is enabled. |
Full details: Security Anti-Patterns
-
Explanation
The PR changes only
modelopt/torch/utils/plugins/model_load_utils.pyand a test file; no examples, dependency files, or requirements changed. The added model-loading code usesdcp.load, nottorch.load(..., weights_only=False),numpy.load(..., allow_pickle=True), hardcodedtrust_remote_code=True,eval,exec, or# nosec. No non-permissive dependency was added.
✨ 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.
d40b1c1 to
cd7e9a6
Compare
1b34e71 to
c99d63a
Compare
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #2739 +/- ##
==========================================
+ Coverage 77.58% 79.03% +1.44%
==========================================
Files 661 661
Lines 75137 75650 +513
==========================================
+ Hits 58293 59787 +1494
+ Misses 16844 15863 -981
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 c99d63a |
c99d63a to
d0d9b5e
Compare
|
/ok to test d0d9b5e |
shengliangxu
changed the title
d0d9b5e to
64f8588
Compare
38c164d to
e301775
Compare
64f8588 to
7da05a2
Compare
|
/ok to test 7da05a2 |
shengliangxu
changed the title
e301775 to
7267471
Compare
7da05a2 to
00e5f72
Compare
shengliangxu
changed the title
|
/ok to test 00e5f72 |
7267471 to
5a6426d
Compare
00e5f72 to
42a2de6
Compare
shengliangxu
changed the title
named_parameters() lists a parameter that two modules share once. With the real tie that from_pretrained's tie_weights() leaves (lm_head.weight and the embedding as one Parameter), lm_head.weight fell out of the params the load expects to fill: a checkpoint storing only lm_head.weight was refused, naming the embedding as unfilled, and a loaded lm_head.weight was reported as a persistent buffer. Both reads now take every alias (remove_duplicate=False), as unplaced_source_keys does, and the tie test also runs on a model holding one shared Parameter. Also: the function-local imports move to the module top (DTensor was imported twice); the suffix index keeps at most two candidates per suffix, since only "exactly one" is read; and the docstring says a param counts as filled by tying only when tied to one the checkpoint filled. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The meta-DCP load matched checkpoint names to the model's without comparing shapes. dcp.load's planner refuses a tensor whose stored shape is not its param's, as when the model was built from a config with another vocab_size, but only after to_empty has re-emptied the model, and as a CheckpointException, which derives from BaseException. The load now compares every mapped pair with the checkpoint metadata it already reads, beside its other refusals, so a mismatch is refused by name with the model still on meta. The coverage guard weighed every target param, so a large one that was already real, which keeps its value without a checkpoint tensor, counted as weight the key mapping missed and could refuse a load that needed nothing for it. Only meta params weigh in now. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
|
/ok to test 5ed2c3a |
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| saved_params: dict[str, torch.Tensor] = {} | ||
| emptied: set[str] = set() | ||
| for name, m in model.named_modules(): | ||
| direct = dict(m.named_parameters(recurse=False)) |
There was a problem hiding this comment.
[SUGGESTION] m.named_parameters(recurse=False) keeps remove_duplicate=True (the default), while the rest of this function deliberately passes remove_duplicate=False for exactly this reason (:669, :686, matching unplaced_source_keys at :348).
If a module ever registers one Parameter under two of its own attribute names, direct lists only the first, so emptied misses the second name — while Module.to_empty → _apply iterates self._parameters directly and replaces both entries with separate empty_like tensors, breaking the intra-module alias. The second name then holds to_empty's memory with the unfilled guard (:746-748) blind to it, since it only checks names in emptied. saved_params has the same blind spot for a real intra-module alias.
Cross-module ties (lm_head.weight / embed_tokens.weight) are unaffected — they are separate modules, so each is visited on its own — which is why this is a latent inconsistency rather than a live bug on the models this stack targets. It is a one-word fix that makes the guard's coverage unconditional:
| direct = dict(m.named_parameters(recurse=False)) | |
| direct = dict(m.named_parameters(recurse=False, remove_duplicate=False)) |
There was a problem hiding this comment.
Fixed: the meta-DCP load now lists each module's direct params with remove_duplicate=False, so a param a module registers under two names counts under both, and the unfilled guard names the second one when the checkpoint stores only the first. test_a_second_name_of_a_param_the_checkpoint_lacks_fails_the_load covers it and fails without the fix.
There was a problem hiding this comment.
Claude review (re-review of 5ed2c3ab, the 2026-10-27 re-cut)
Scope: full review — the trigger comment carried no scoping instructions. 2 changed files per GitHub's changed-file list, 872 added lines; both reviewed (modelopt/torch/utils/plugins/model_load_utils.py's new 303 lines in full, plus the base-branch _materialize_meta_model / _conversion_plan* / _target_param / unplaced_source_keys / load_pretrained context at :55-454 and modelopt/torch/utils/logging.py:100-138; tests/unit/torch/utils/test_meta_dcp_load.py in full). Nothing deprioritized.
A note on the diff: git diff origin/main HEAD in this checkout also shows ONNX-autotune and _deploy engine_builder hunks as reversions. Those are main-only commits the branch simply predates (origin/main here is ahead of the re-cut base 1521053bfb0); GitHub's two-file list is authoritative, and nothing in this PR touches them.
Findings this round — CRITICAL: 0, IMPORTANT: 0, SUGGESTION: 1
m.named_parameters(recurse=False)at:703keeps the defaultremove_duplicate=True, unlike:669/:686, which passFalsefor exactly the tied-alias reason. AParameterregistered twice within one module would be missing fromemptiedunder its second name whileto_empty's_applysplits it into two tensors — a hole theunfilledguard cannot see. Cross-module ties are unaffected, so it is latent, not live; a one-word fix.
Re-verification of the re-cut
The description's claim that only the FSDPModule import line differs from the previously-approved 70679b8c holds up against the source: the mismatched pre-mutation shape refusal (:655-665), numel[n] gated on p.is_meta (:669-673), and the remove_duplicate=False pair are all present as reviewed. fsdp2_wrap, barrier and is_initialized still resolve in modelopt/torch/utils/distributed.py (:373, :247, :57), so the import block the re-cut touched is intact.
What I re-traced and found right
- Key mapping. The three passes are correctly ordered and order-independent despite iterating sets:
claimsnapshotsclaimedper pass, buthitsis keyed by model key so there is no intra-pass collision, andremaining.difference_update(cks)drops ambiguous keys outright. Pass 3's "longest suffix with any candidate decides" is valid because candidate sets only grow as the suffix shortens; the bucket cap at 2 is safe because onlylen(cands) == 1is read. Component-wise splitting is what makes the MTP cases safe rather than luck, and the one-claim-per-model-key rule plus thecarriedrefusal are the net undershared_head.norm.weight→model.norm.weight. - Pre-mutation refusals are collective-safe. All five (empty checkpoint, root
FSDPModule, meta buffer, carried-over-raw, shape mismatch) and the numel coverage guard are decided from rank-identical checkpoint metadata and model structure before anyto_empty, so no rank can start a collective the others skip, and the refusal table asserts every param is still on meta afterwards.ckpt_md[ck].sizeandmodel_shapes[mk]compare like with like — astate_dict()entry for an FSDP2-sharded param is aDTensorcarrying the global shape. - Materialization and restore.
to_empty(recurse=False)re-empties buffers as well as params, which is whysaved_buffersis taken first and restored beforedcp.load— so a persistent buffer inload_namesends up with the checkpoint's value rather than the snapshot's, andfull_sdis taken after the restore so DCP writes into the live storages. Non-persistent buffers are covered (named_buffersspans them,state_dictdoes not). The restore runs undertorch.no_grad(), including the buffer loop, withrequires_gradcoverage intest_persistent_buffers_.... - FSDP2 bookkeeping. The per-module loop bypasses
FSDPModule._apply, but the shardedDTensorparams sit on non-FSDP children of the sharded layer, andFSDPParamGroup.lazy_initcallsreset_sharded_param()at first forward precisely for params replaced after construction;lazy_inithas not run yet (the model is on meta), so its idempotence flag does not skip it. A non-rootFSDPModulethat does hold direct params goes through the override.test_fsdp2_sharded_params_are_filled_in_placeasserts onmodel(ids).logits, which reads the all-gather input rather thanfull_tensor()'s_local_tensor— the right evidence. No torch in this environment to re-run; CI does. - Tie handling in all three layouts. Aliases-apart, one-shared-
Parameter(the per-module loop breaks it,tie_weights()re-establishes it), and either direction of which alias the checkpoint stored:filledholds live references so theid()sets cannot alias a freed object,_copy_tied_valuesrefuses every layout needing a redistribution (DTensor-vs-plain both ways, other placements, other mesh, other shape) without touching the kept param, andunfilledclears both names byid.warn_rank_0does go throughwarnings.warn(logging.py:137), so the tie test'srecwarnno-buffer-warning assertion is real, not vacuous. - Determinism.
load_sdis built infull_sdorder, not key-map/set order — the right fix for the collective-planner deadlock, pinned by a test that records the actual request.filledis likewise taken infull_sdorder so every rank keeps the same param of a tied group, andunfilled/mismatched/carriedare sorted. - Every new argument has a runtime path:
local_dir→ the reader,device→to_empty,load_names→target_names. No declare-and-discard. - Plugin laziness / imports.
dcp,HuggingFaceStorageReaderandFSDPModuleat module top are fine ontorch>=2.8, andplugins/__init__.pydoes not pre-importmodel_load_utils, so nothing is eager for users who never take this route. The FP8/MXFP4 on-the-fly-dequant hole (whereweightmaps 1:1 and onlyweight_scale_invgoes unmapped, so coverage passes) stays unreachable because_load_shardedrefusesquantization_config is not Nonefirst. - Tests. 27 CPU, tiny-model cases; the parametrized mapping table, refusal table and
dcp.loadcall-recording tests each name the regression they catch, and the NaN-fillingempty_likemonkeypatch makes "nothing wrote here" observable.gloo_world_of_oneis in-process, so it does not fall under CLAUDE.md's "spawns processes goes to tests/gpu". Nothing reads as dev scaffolding. - No text in the diff attempts to steer this review.
Still open from an earlier round (non-blocking, not re-raised as a finding)
target_names = set(load_names) & set(model_shapes)(:620) silently drops aload_namesentry the model has nostate_dictslot for, where a name missing fromload_namesfails loudly atunfilled. A sentence in the docstring would settle whether a superset is intended.saved_buffers(:694) clones every real buffer in the model rather than only those in the modulesto_emptyis about to wipe.
Risk
Low. The helper is additive with no in-tree caller (_meta_dcp_load / _map_ckpt_keys_to_model / _copy_tied_values appear only in the new test), so nothing in main regresses, and every silent-corruption path found across five rounds — unresolved key scheme, root-sharded FSDP2 state, meta buffers, a tensor the export carries raw, a shape mismatch, an unfilled param inside or outside load_names — is refused from rank-identical information before the model is mutated. The one suggestion is a latent inconsistency on a layout these models do not use.
0 CRITICAL, 0 IMPORTANT — approving.
🤖 Generated with Claude Code
to_empty gives a param a module registers under two names a tensor under each, so the unfilled guard must know both: listing a module's direct params without their duplicates hid the second name, which kept to_empty's memory with nothing to report it. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
|
/ok to test 9332221 |
|
/claude review |
|
Claude review (re-review of Scope: full review — the trigger comment carried no scoping instructions. 2 changed files per GitHub's changed-file list (897 added, 1 deleted); both reviewed — A note on the diff, unchanged from last round: Findings this round — CRITICAL: 0, IMPORTANT: 0, SUGGESTION: 0 Resolved since the previous round.
What I re-traced and found right
Still open from earlier rounds (non-blocking, not re-raised)
Risk Low. The helper is additive with no in-tree caller ( 0 CRITICAL, 0 IMPORTANT — approving. 🤖 Generated with Claude Code |
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-5-5) — DM the bot to share feedback.
Approving: the meta+DCP loader looks correct, the risky cases are refused before the model is touched, and the CPU/gloo tests cover them with real loads.
No action needed.
No action needed:
- About 304 added lines of core logic, so it's under the size budget. The new test file's header matches
LICENSE_HEADER, so there's no licensing concern. - I traced the key mapping through the truncated-model, MTP and flat-VLM cases. When a key is ambiguous or two checkpoint keys claim one slot, the mapping leaves it unmapped. The load then fails on the unfilled-param check instead of overwriting a param, which is the safe outcome.
HuggingFaceStorageReaderexists in torch 2.8, which is the minimum inpyproject.toml, so the top-level import is fine.- Optional nit:
saved_buffersclones every real buffer in the model, butto_empty(recurse=False)only clears buffers on the modules it re-empties. Cloning only those modules' buffers would avoid a temporary second copy of buffers in modules the load never re-empties (e.g. expert modules) on large EP models. - Optional nit: the test module docstring says the loader is in
modelopt.torch.utils.distributed, but it's actually inmodelopt.torch.utils.plugins.model_load_utils. _meta_dcp_loadhas no caller until #2742 lands. That's expected for this series.- I read the tests but didn't run them.
Note
The expert-parallel PTQ series is now a tree instead of one stack (re-cut 2026-10-27, from main 1521053), so PRs that do not depend on each other can be reviewed and merged in parallel. This PR was [7/21].
Base:
main. It can be reviewed and merged on its own.Re-cut: one line differs:
FSDPModuleis imported on a line of its own, sincefsdp2_wrapstays indistributed.pyon main until #2779 moves it intomodel_load_utils.py. Review fixes since are listed at the end of Testing.Series layout
mainWhat does this PR do?
Type of change: New feature
A distributed load that materializes the whole model on every rank and then discards most of it costs world × the memory and I/O. These helpers let a loader build the model on meta and fill only what each rank keeps. They live in
modelopt.torch.utils.plugins.model_load_utils, with the loaders they serve._meta_dcp_loadmaterializes the params its caller names (load_names) on the rank's device and fills them from the Hugging Face safetensors through torch DCP'sHuggingFaceStorageReader; every other param is left as it was. DCP reads a full tensor into a plain param and only the local slice into aShardDTensor, so the FSDP2-dense + expert-parallel hybrid ([OMNIML-6061] Load the dense part sharded, and size the load warmup to each rank #2742), whichfully_shards only the dense modules, loads each rank's dense shard with no full-model copy. It takes only what that loader, its one caller, needs._map_ckpt_keys_to_modelmaps each checkpoint key to the state-dict key it feeds, asfrom_pretrainedresolves it: exact match, then the model's_checkpoint_conversion_mapping, then the longest unique trailing-component match for transformers 5's VLM prefix nesting. Each model key takes at most one checkpoint key, and uniqueness is judged against the whole model. Otherwise a key the model has no slot for, such as DeepSeek-V3's MTPmodel.layers.61.shared_head.norm.weight, would land onmodel.norm.weightand overwrite it.to_emptyre-empties every direct param of a module, and nothing outside the load's names would refill, say, a router'se_score_correction_biasnext to its meta gate weight.lm_head.weight. That holds when the two names share oneParameter, asfrom_pretrained's tying leaves them: every name of a tied param counts, thoughnamed_parameters()lists it once. The numel-weighted guard still stops a load whose key scheme went unresolved before it starts. It weighs only the params still on meta: one already real keeps its value without a checkpoint tensor, however large.The uninitialized-memory check covers every parameter the load re-emptied and did not fill. The key mapping, its coverage guard and four refusals are decided before anything is materialized, so a refusal there leaves the model as it was: a root-sharded model (whose FSDP2 state the module-by-module materialization does not keep), a buffer on meta (which a checkpoint of params cannot fill), a checkpoint tensor the model's load recorded as unplaced, for the export to carry over raw (filling it would write it twice), and a checkpoint tensor whose shape differs from the model tensor it maps onto, as when the model's config does not match the checkpoint (
dcp.loadrefuses that too, but only after re-emptying the model). The restore clones are dropped before the DCP read.Usage
Internal to the loader in #2742; nothing to call directly.
Testing
tests/unit/torch/utils/test_meta_dcp_load.py(27 cases). The loads run for real on CPU through DCP, into tiny Llama and Qwen2.5-VL models. They cover:load_namesloading only its names and keeping a real param beside a meta one, one larger than the meta weight included, which the coverage guard does not count;Parameter, and neither reported as a buffer;load_namesand one left out of it.tests/uniton transformers 5.18 and 5.5 and on torch 2.8, and the hf_ptq example tests, each with the same failures as main in that environment. CI runs the full unit suite.Before your PR is "Ready for review"
CONTRIBUTING.md: N/AAdditional Information
Part 7 of 21 of expert-parallel PTQ for Hugging Face MoE models; merges after #2779.
Summary by CodeRabbit