Repository navigation
[OMNIML-6054] [6/21] Route load_pretrained by the distribution plan - #2779
shengliangxu wants to merge 23 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:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info
📝 Walkthrough
Merge Risk: ⚪ Minimal · up to
|
| Check name | Status | Explanation | Resolution |
|---|---|---|---|
| Docstring Coverage |  |
Docstring coverage is 68.42% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 57 functions across 10 files. (1 skipped:… | 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 names the main change: load_pretrained now picks its loading route (whole-model or FSDP2) from a DistributionPlan. The changeset matches this. The title does not mention the moved fsdp2_wrap… Linked Issues check The only linked issue, #2774, is closed and completed. The issue states it is historical context only. No active directly linked target remains, so no linked-issue coding requirements apply to this pu…Out of Scope Changes check The whole-PR summary (mergeBase 4e10244 to currentHead fb8bf38) shows changes that match the PR description. In modelopt/torch/utils/plugins/model_load_utils.py, load_pretrained gains a distribution… Security Anti-Patterns No listed security anti-pattern is introduced by the pull request. In the changed modeloptPython files,torch.loadusesweights_only=True; no addedweights_only=False,allow_pickle=True, ha…Full details: Title check
-
Explanation
The title names the main change: load_pretrained now picks its loading route (whole-model or FSDP2) from a DistributionPlan. The changeset matches this. The title does not mention the moved fsdp2_wrap or the new Hub options, but it does not need to cover every detail. The bracketed prefix [6/21] is a series marker and does not affect the assessment.
Full details: Linked Issues check
-
Explanation
The only linked issue,
#2774, is closed and completed. The issue states it is historical context only. No active directly linked target remains, so no linked-issue coding requirements apply to this pull request. The prior PASS result still holds.
Full details: Out of Scope Changes check
-
Explanation
The whole-PR summary (mergeBase 4e10244 to currentHead fb8bf38) shows changes that match the PR description. In modelopt/torch/utils/plugins/model_load_utils.py, load_pretrained gains a distribution parameter, resolve_distribution_plan gains config and Hub access options, and fsdp2_wrap moves here. In modelopt/torch/utils/distributed.py, fsdp2_wrap forwards to the new location and its signature is unchanged. The Hub access options support the PR's Hub ID localization. The test changes cover the same functions: test_example_utils.py (FakeModel now a Module so the plan can be recorded), test_export_carry_over.py and test_model_load_utils.py (manual plan recording after fsdp2_wrap), test_distributed_run.py (forwarding), test_unplaced_source_keys.py (plan check), test_resolve_distribution_plan.py (Hub options), and the new unit and GPU test modules. The CHANGELOG.rst entry documents the same behavior. No file in the summary is shown to be unrelated. The assessment rests on the AI-generated summary; the source hunks were not read in this review.
Full details: Security Anti-Patterns
-
Explanation
No listed security anti-pattern is introduced by the pull request. In the changed
modeloptPython files,torch.loadusesweights_only=True; no addedweights_only=False,allow_pickle=True, hardcodedtrust_remote_code=True,eval()/exec()call, or# noseccomment exists. The pull request changes nopyproject.tomlorrequirements*.txtfile, so it adds no dependency requiring license review.
-
Full details: Docstring Coverage
-
Explanation
Docstring coverage is 68.42% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 57 functions across 10 files. (1 skipped: 1 unsupported.)
Resolution
Write docstrings for the functions missing them to satisfy the coverage threshold.
✨ Finishing Touches: 💡 2
-
📝 Generate docstrings: 💡
-
- Commit to this branch
- Create a new PR
🛠️ Fix failing CI checks: 💡
-
- 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
Comment @coderabbitai help to get the list of available commands.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## shengliangx/hfep-2-plan-decision #2779 +/- ##
===================================================================
Coverage ? 61.36%
===================================================================
Files ? 648
Lines ? 71985
Branches ? 0
===================================================================
Hits ? 44173
Misses ? 27812
Partials ? 0
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 3ac7c7b |
3ac7c7b to
df76b36
Compare
shengliangxu
changed the title
|
/ok to test df76b36 |
df76b36 to
aa3615e
Compare
|
|
/ok to test aa3615e |
aa3615e to
168c90d
Compare
|
/ok to test 168c90d |
168c90d to
537ff8d
Compare
shengliangxu
changed the title
|
/ok to test 537ff8d |
465f1eb to
aa2a065
Compare
|
/ok to test aa2a065 |
|
@coderabbitai review |
✅ Action performedReview finished.
|
|
/ok to test ec4159a |
|
Question: does It looks like it only lets a sharded load reuse the config |
|
/ok to test 17cfb96 |
|
@Fridah-nv, it lets a caller that decides the plan itself load by that plan without the config being prepared twice.
Carrying the prepared config inside the plan would remove the argument. But a |
load_pretrained becomes the one loader for a model ModelOpt will optimize and export, whole or distributed. A new distribution argument picks the load: "sequential" (from_pretrained as before), "fsdp2" (parallel_load_and_prepare_fsdp2), or a DistributionPlan -- FSDP2 over the world, or a whole copy on every rank's own GPU. A plan that shards the routed experts raises: no load here does that yet. Left at None, the active ModelOptSession decides: a single process loads whole, a distributed session decides the plan with the new resolve_distribution_plan, through decide_distribution_plan with the MODELOPT_EP_DEGREE and MODELOPT_FSDP_SIZE knobs (modelopt.torch.utils.env). The plan a load runs is kept as the session's synced plan, the same on every rank. Every load makes a Hub ID local and records the unplaced source keys. The FSDP2 load cannot run the model type's loading hook, so it refuses a checkpoint the hook would build modules for (a Nemotron-H MTP head), and it refuses from_pretrained arguments it would otherwise ignore, a quantized checkpoint, whose tensors it would read as stored, and a partial fsdp_size, which it cannot shard. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
load_pretrained kept a plan the caller passed as session.plan, which holds rank 0's, but routed and loaded with this rank's own: ranks handed different plans built mismatched meshes and hung while the session said they agreed. Every rank now loads the plan the session holds, once its own plan passed the checks, so a refused plan still leaves the session as it was. A plan the session decides itself is already held, and is no longer sent a second time. The test that rank 0 alone works out the unplaced keys answers the plan the load now decides first. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Every rank loads the plan rank 0 was handed, now set with session.plan.decide(lambda: handed) rather than an assignment, and read back with session.plan.value. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Rank and world size come from the process group, and a plan handed to every rank goes out from rank 0 through broadcast_from. Rank 0 now checks that plan before it sends, so a plan it refuses raises on every rank instead of leaving the others in the broadcast. Nothing keeps the plan after the load. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
load_pretrained's distribution is now a DistributionPlan or None, which decides through resolve_distribution_plan; the 'sequential' and 'fsdp2' strings that picked a load without a plan are gone. A single process loads by a single process's plan, so every load has one. Whatever the route, the plan is recorded on the model and every submodule, where calibration and export read it back instead of inspecting the model. parallel_load_and_prepare_fsdp2 records its own plan as well, for a caller that uses it directly. The decision's copy here goes too: it lives in resolve_distribution_plan now. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
fsdp2_wrap FSDP2-shards a model over the process group, so it records that plan on the model itself. A caller that wraps a model with it -- the FSDP2 load among them, which no longer records its own -- exports and calibrates it by the plan without a further call. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…n on load_pretrained records the plan on every module of the model it returns, so the stand-in model these tests have from_pretrained return is a torch module now. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
FSDP2 shards through a process group even on one rank, and the FSDP2 load refused to start without one, leaving every caller to set it up first. load_pretrained sets up the group of one itself, after anything that could refuse the load, so a refused load leaves none behind; the run (distributed_run) destroys it with the run. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
fsdp2_wrap shards a model as part of loading it, the one load that calls it lives in model_load_utils, and so does every other loader. It moves there; the FSDP2 helpers it uses stay in distributed.py, and modelopt.torch.utils.distributed.fsdp2_wrap forwards to it for callers of the old path. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
fsdp2_wrap recorded a plan it derived from the process group, and load_pretrained then recorded the one it loaded by over it: two writers, one silently replacing the other. load_pretrained is the one that knows the plan, so it alone records it, once, at the end of every load. fsdp2_wrap records nothing, as on main; the tests that wrap or load a model through it directly record the plan themselves, as any caller distributing a model itself must. 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>
resolve_distribution_plan reads the checkpoint through the failure-carrying all_reduce, which gathers from every rank of a launch. The run_as fixture answers that gather as it answers the plan's broadcast: every rank read the checkpoint as this one did. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The test played rank 1 against a stand-in broadcast that answered with rank 0's plan, and checked which stub loader ran. It now runs on the session's two gloo workers: rank 0 is handed a replicated plan and rank 1 an FSDP2 one, both load a real checkpoint onto the CPU, and both must load it whole and record rank 0's plan. A rank that loaded its own plan would refuse the placement an FSDP2 load cannot honor, so a regression fails at once instead of waiting in a collective. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Fold near-duplicates into one test per behavior: - test_a_load_refuses_what_it_cannot_load takes the another-world, quantized-checkpoint, not-a-plan and hf-config-on-a-whole-load refusals; each case still fails if its refusal goes. - test_a_single_process_loads_the_model_whole_as_asked goes, and its plan check moves to test_unplaced_source_keys.py's real single-process load. That load fails if a single process is placed on the rank GPU, loses the caller's kwargs, records another plan or reaches the FSDP2 loader. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
test_fsdp2_wrap_leaves_the_plan_to_the_load imported it inside the test,
though it is neither optional nor circular. At module scope it sits after the
file's importorskip("transformers"), so the transformers-free lane still skips
the module first.
Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
load_pretrained now resolves the distribution plan before it loads, and resolve_distribution_plan resolved the checkpoint and read its config without token, cache_dir or local_files_only. That read reaches the Hub for custom code an auto_map names in another repo, as the loads after it do. resolve_distribution_plan takes the three options, and load_pretrained passes on the ones it was given. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The load_pretrained test cannot see whether the plan's read resolves the checkpoint with token, cache_dir and local_files_only: load_pretrained has made the checkpoint local by then. Drive resolve_distribution_plan with a Hub ID and check that the resolution and the config read both get them. (Test from the epdp port of this change.) Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
load_pretrained keeps token, cache_dir and local_files_only out of the keywords it routes, and a sharded load given a plan but no config read the config without them. That read reaches the Hub for custom code an auto_map names in another repo, as the whole-model load's do. _load_sharded now takes the Hub options for its config read. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
load_pretrained honored a config passed for the load only in the whole-model from_pretrained: the plan was still decided from the checkpoint's config.json, so a directory without one failed before loading, and a sharded load refused config as an option it does not take. resolve_distribution_plan takes the caller's config and decides from a copy of it (deciding sets fields on it) without reading the checkpoint's; load_pretrained passes it on, and a sharded load builds from it as its hf_config. Passing both config and hf_config is refused. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
With a plan handed in, nothing is decided, so the model type comes from the passed config itself rather than from the plan's copy of it. The test only covered a plan load_pretrained decides, and a load that ignored the passed config there (model type from hf_config alone) still passed it. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
A sharded load builds from hf_config as given. Without a plan handed in, load_pretrained decided one from the checkpoint's own config and then loaded the caller's hf_config, so the model loaded was not the one the plan was decided for, nor one the model spec's checks had seen. hf_config is the config resolve_distribution_plan returns with its plan; without that plan it is now refused, pointing at config= for a config the plan should be decided from. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…oint Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
|
/ok to test fb9d332 |
This branch has not been deployed
2 participants
What does this PR do?
Type of change: New feature
load_pretrained(#2733) loads a model whole. A distributed load, such asparallel_load_and_prepare_fsdp2(main), is a separate entry point, and every caller picks one itself, ashf_ptq.py --use_fsdp2does. A caller that picks a loader also gets what that loader skips: a sharded load cannot run the model type's loading hook, so a Nemotron-H checkpoint's MTP head is not built and its tensors go to the export unquantized.This PR makes
load_pretrainedthe one loader, whole or distributed:distributionis aDistributionPlan; by default it is the planresolve_distribution_plan([5/21]) decides.dense="fsdp":parallel_load_and_prepare_fsdp2over the process group. A single process has no group, so the load sets one up, of one rank, once nothing can refuse the load anymore.distributed_run([4/21]) destroys it with the run.from_pretrained. On a distributed run every rank gets the whole model, on its own GPU.set_distribution_plan, [4/21]), where calibration and export read it.load_pretrainedrecords it exactly once, at the end of every load; the loaders it calls record none.fsdp2_wraprecords nothing, as on main, so a caller that wraps or loads a model through it directly records the plan itself, as the error for a missing plan says.fsdp2_wrapmoves tomodel_load_utils, with the load that calls it and every other loader.modelopt.torch.utils.distributed.fsdp2_wrapforwards to it, so callers of the old path keep working, and the FSDP2 helpers it uses stay indistributed.broadcast_from, [4/21]), just as it would one it decided. Ranks that each loaded the plan they were handed could build mismatched meshes and hang. Rank 0 checks the plan before it sends it, so a plan it refuses raises on every rank instead of leaving the others waiting in the broadcast.from_pretrained's placement arguments (device_map,max_memory, ...) and a model class it does not build, rather than ignore them.hf_configwithout the plan it was resolved with is refused. A sharded load builds the model from it as given, so a plan decided here, from the checkpoint's own config, would not be the one it came with.config=is for a config the plan should be decided from.Usage
Testing
tests/unit/torch/utils/test_load_pretrained_distribution.py. It runs on CPU, and stand-in loaders record their calls. It covers:cpu_dist_workers), rank 0 handed a replicated plan and rank 1 an FSDP2 one: both load whole and record rank 0's plan. A rank that loads the plan it was handed fails at once, on the placement refusal.from_pretrainedplaces it, with the caller's arguments and a single process's plan: [OMNIML-6080] Derive the unplaced source keys once, at load, through a core load_pretrained #2733's real CPU load intests/unit/torch/utils/test_unplaced_source_keys.py, which now also checks the plan.fsdp2_wrapleaving the plan to the load: a CPU unit test intests/unit/torch/utils/test_load_pretrained_distribution.py. [OMNIML-6080] Derive the unplaced source keys once, at load, through a core load_pretrained #2733's GPU tests that wrap or load through it directly (test_export_carry_over.py,test_model_load_utils.py) record the plan themselves. Its old path forwarding every argument:tests/unit/torch/utils/test_distributed_run.py.fsdp2_wraptoo, set no group up for a single process's FSDP2 load, set one up over a live group, or have the oldfsdp2_wrappath drop an argument or the result.tests/gpu/torch/utils/test_load_pretrained_distribution.pyrunsload_pretrainedleft to decide, on every GPU of the node:MODELOPT_FSDP_SIZE=worldFSDP2-sharded;tests/examples/hf_ptq/test_example_utils.py:get_model's stand-in model is a module now, sinceload_pretrainedrecords its plan on it.Before your PR is "Ready for review"
load_pretrainedkeeps its behavior in a single process, anddistributionis new.fsdp2_wrapbehaves as on main and records no plan; its old import path still works.CONTRIBUTING.md: N/AAdditional Information
Part 6 of 21 of expert-parallel PTQ for Hugging Face MoE models; merges after [5/21]. The expert-parallel half of this routing is [12/21].
Replaces #2774, which carried this routing whole as [19/20]. Moving its branch down the stack made the branch part of #2774's old base, so GitHub closed #2774 as merged into that base; nothing reached main.
Summary by CodeRabbit
New Features
Documentation