[shengliangxu]

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 [6/21]. Rebased onto main 7c3c868 later the same day, after #2739 merged, and onto main f03651e that night.

[2/2] of a split (2026-10-30), as review asked to keep each PR under the source-line budget: #3140 ([1/2]) moves fsdp2_wrap out of modelopt.torch.utils.distributed into modelopt.torch.utils.fsdp2_utils, and this PR routes load_pretrained by the distribution plan.

Source lines: 496 against #3140, under the ~500 budget: trimmed on 2026-10-31 on the user's decision not to split.

Base: #3140 (shengliangx/hfep-13a-move-fsdp2-wrap). The diff below is this PR's alone. Merge order: #3140 first, then this PR, which then retargets to main; until then GitHub runs no unit, GPU, example or regression tests on it (they run against main).

Re-cut: this PR's diff is unchanged; only its base moved. Review fixes since are listed at the end of Testing.

Series layout

• main
  • [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745 fused experts under grouped_mm (merged)
    • [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 calibration on grouped_mm
      • [OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067 Nemotron-H's experts under grouped_mm
      • [OMNIML-6071] Calibrate a distributed model on per-rank data #2740 calibration on per-rank data
  • [OMNIML-6064] Load each rank's checkpoint shard straight into a meta-built model #2739 meta-DCP checkpoint load (merged)
  • [OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736 distribution plan, run group, collectives (merged)
    • [OMNIML-6081] [2/2] Fold fused experts' quantizers into their weights shard-locally #2750 keep-fused fold
    • [OMNIML-6082] [3/3] Write a sharded model's HF checkpoint shards from each tensor's owner rank #2751 distributed-writer pieces
      • [OMNIML-6083] Write a sharded model's HF checkpoint without gathering it #2752 no-gather export (on [OMNIML-6082] [3/3] Write a sharded model's HF checkpoint shards from each tensor's owner rank #2751, [OMNIML-6081] [2/2] Fold fused experts' quantizers into their weights shard-locally #2750, [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779 and [OMNIML-6071] Calibrate a distributed model on per-rank data #2740)
        • [OMNIML-6084] Export an expert-parallel model through the no-gather writer #2753 expert-parallel export
    • [OMNIML-6054] Decide how a model is distributed, from the model built on the meta device #2737 plan decision (merged)
      • [OMNIML-6054] [1/2] Move fsdp2_wrap out of distributed.py into fsdp2_utils #3140 fsdp2_wrap move ([1/2])
        • [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779 load routing (this PR, [2/2])
        • [OMNIML-6060] [3/3] Refuse an expert-parallel load that would calibrate the experts wrong #2741 expert-parallel loader (on [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779, [OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067 and [OMNIML-6071] Calibrate a distributed model on per-rank data #2740)
          • [OMNIML-6061] Load the dense part sharded, and size the load warmup to each rank #2742 dense sharding
            • [OMNIML-6062] Dequantize block-FP8 checkpoints as they load expert-parallel #2743 block-FP8 dequantization
              • [OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780 expert-parallel load routing
                • [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746 expert parallelism in the model specs
                  • [OMNIML-6063] [3/3] Load a checkpoint its model spec decompresses through the expert-parallel loader #2747 on-the-fly decompression
                  • [OMNIML-6078] [2/2] Add the Kimi-K2.5 expert-parallel adapter #2748 Kimi-K2.5 and Kimi-VL adapters
                  • [OMNIML-6079] [2/2] Add the Kimi-K3 expert-parallel adapter #2749 Kimi-K3 adapter
                  • [OMNIML-6085] [3/3] Run expert-parallel PTQ from hf_ptq.py #2754 expert-parallel PTQ from hf_ptq.py (on [OMNIML-6076] [3/3] Apply the model spec when deciding and loading expert-parallel #2746 and [OMNIML-6084] Export an expert-parallel model through the no-gather writer #2753)

What does this PR do?

Type of change: New feature

load_pretrained (#2733) loads a model whole. A distributed load, such as parallel_load_and_prepare_fsdp2 (main), is a separate entry point, and every caller picks one itself, as hf_ptq.py --use_fsdp2 does. 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_pretrained the one loader, whole or distributed:

• distribution is a DistributionPlan; by default it is the plan resolve_distribution_plan ([OMNIML-6054] Decide how a model is distributed, from the model built on the meta device #2737) decides.
  • dense="fsdp": parallel_load_and_prepare_fsdp2 over the process group. A single process launched with torchrun --nproc_per_node=1 has no group, so the load sets one up, of one rank, once nothing can refuse the load anymore; a process started without a launcher cannot set one up and is refused. distributed_run ([OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736) destroys it with the run. One rank of a multi-rank launch whose group is not set up yet is refused too, told to set it up first with distributed_run(), rather than shard over the launch's ranks under a one-rank plan; the plan-size refusal says the same when that is why the sizes differ.
  • Otherwise from_pretrained. On a distributed run every rank gets the whole model, onto its own GPU when it has one and no device_map is passed.
  • A plan that shards the routed experts raises here; [OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780 adds that load, once the expert-parallel loader ([OMNIML-6060] [3/3] Refuse an expert-parallel load that would calibrate the experts wrong #2741, [OMNIML-6061] Load the dense part sharded, and size the load warmup to each rank #2742 and [OMNIML-6062] Dequantize block-FP8 checkpoints as they load expert-parallel #2743) is in. Until then a plan this PR's load decides keeps the EP degree at 1, so an MoE loads without expert parallelism by default, whole or FSDP2-sharded as the dense layout says. An EP degree MODELOPT_EP_DEGREE asks for, or a handed plan with ep_size > 1, is refused before anything loads, and the refusal names MODELOPT_EP_DEGREE=1.
  • resolve_distribution_plan takes the class the load builds with (auto_class, which load_pretrained passes). With a class the sharded loads do not build and no knob asking to shard, it plans the whole model on every rank without a meta-device build, so an encoder-decoder loaded through AutoModel, for example, loads multi-rank as before.
• The plan is recorded on the model and every submodule (set_distribution_plan, [OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736), where calibration and export read it. load_pretrained records it exactly once, at the end of every load; the loaders it calls record none. fsdp2_wrap records 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. hf_ptq.py --use_fsdp2 loads through parallel_load_and_prepare_fsdp2 directly and records no plan in this PR. Recording it there lands with [OMNIML-6083] Write a sharded model's HF checkpoint without gathering it #2752, the first PR whose export reads it on that path, and [OMNIML-6085] [3/3] Run expert-parallel PTQ from hf_ptq.py #2754 moves that load onto load_pretrained.
• fsdp2_wrap moved into modelopt.torch.utils.fsdp2_utils, which imports no Hugging Face package and which the loaders import it from, in [OMNIML-6054] [1/2] Move fsdp2_wrap out of distributed.py into fsdp2_utils #3140 ([1/2]), which keeps its old path. This PR makes it refuse a model that, or any of whose submodules, records a plan, as a model load_pretrained returned does: a plan is recorded once, so a model to shard yourself is one loaded with from_pretrained.
• Every load makes a Hub ID local and records the unplaced source keys ([OMNIML-6080] Derive the unplaced source keys once, at load, through a core load_pretrained #2733). The FSDP2 load builds its model from the config without the Hub options, so custom code an auto_map names comes with the environment's Hub settings.
• Every rank loads the same plan. Every rank takes rank 0's distribution, a plan or none, through one broadcast (broadcast_from, [OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736) before it branches, and loads a plan rank 0 was handed just as it would one it decided. Ranks that each loaded the plan they were handed could build mismatched meshes and hang, and ranks passing different kinds would wait in each other's collectives. Rank 0 checks what it was handed before it sends it, a non-plan included, so a distribution it refuses raises on every rank instead of leaving the others waiting in the broadcast.
• Refusals, before anything loads, and on every rank. Each rank checks its own arguments and reads the config itself, so each refusal runs through the failure-carrying all_reduce ([OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736): a refusal on any rank raises on every rank, before the checkpoint's fetch for the arguments, after it for a rank that cannot see the checkpoint rank 0 fetched, and before the sharded load's collectives for what the config says, instead of leaving the other ranks waiting there.
  • A sharded load cannot run the model type's loading hook, so it refuses a checkpoint the hook would build modules for: the MTP head above. Main's hf_ptq.py --use_fsdp2 refuses Nemotron-H MTP on the same loader (validate_fsdp2_supported), and this applies its rule to every caller. The hook could wrap the FSDP2 load's meta-device build, but the head's calibration forward inside the root FSDP2 unit and its quantized export from FSDP2 shards have never run on that path: supporting MTP under FSDP2 is a change to that loader and to hf_ptq, validated on a real checkpoint, not part of this routing.
  • It refuses from_pretrained's placement arguments (device_map, max_memory, ...) and an auto_class other than AutoModelForCausalLM (leave it out: the load picks the class), rather than ignore them, and an explicit dtype, saying why it does not apply: the load keeps the checkpoint's dtype, and FSDP2's mp_policy sets the dtype it computes in. An option passed as None, from_pretrained's own default, asks nothing and is not refused.
  • The FSDP2 load refuses a quantized checkpoint, whose tensors it would read as stored.
  • The FSDP2 load builds a text-only VLM it supports (Qwen3.5) through its image-text class, as hf_ptq's --use_fsdp2 does, and any other model through AutoModelForCausalLM. It refuses a config _auto_class would not build through that class, as most VLMs (which AutoModelForCausalLM refuses deep in the loader, or builds as the language model alone), before it sets up a group or reads anything, and the refusal says how to load the model whole on every rank: its model class as auto_class, with MODELOPT_FSDP_SIZE=1 (for a plan the load decided) or a replicated plan.
  • A plan made for another world size is refused, and so is anything other than a plan.
  • hf_config without 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. resolve_distribution_plan's docstring says to hand the config it returns to load_pretrained only with a plan that shards the model: a whole-model load reads the checkpoint's config itself and refuses hf_config.

The FSDP2 route's refusals name their ways out (MODELOPT_FSDP_SIZE=1 for a plan the load decided, or a replicated plan), and a config handed with a plan is copied before the load builds from it. With a plan handed and no config, the FSDP2 route reads the checkpoint's config once and hands it to the load, rather than reading it twice. Deciding a plan builds the meta-device model from a copy of the config, so the config returned with the plan, which a sharded load builds from, is the same on every rank: the build writes its dtype and its attention and experts implementations, and only rank 0 builds.

Routing details are in the docs page this PR adds, docs/source/guides/_distributed_loading.rst ("Loading Hugging Face models across ranks", under the quantization guide), which load_pretrained's docstring links.

Usage

# torchrun --nproc_per_node 8 quantize.py
from modelopt.torch.utils.distributed import distributed_run
from modelopt.torch.utils.plugins.model_load_utils import load_pretrained

with distributed_run():
    model = load_pretrained("meta-llama/Llama-3.1-8B")  # a whole copy on every rank's GPU
    # With MODELOPT_FSDP_SIZE=world: FSDP2-sharded over the 8 ranks instead.

Testing

• Tests in tests/unit/torch/utils/test_model_load_utils.py, beside the plan decision's (the module's test file, as CONTRIBUTING asks). They run on CPU, and stand-in loaders record their calls. They cover:

  • each load: the arguments it gets and the plan it records;
  • a decided FSDP2 load, handed the config the plan was decided from instead of reading it again, and a handed FSDP2 plan's load reading the checkpoint's config once (that check fails without its fix);
  • a single process's FSDP2 load setting up its group of one, and leaving none behind when the load is refused;
  • every refusal, raised before any loader runs, the expert-parallel plan's included; test_a_load_refuses_what_it_cannot_load has a VLM row (a Llava config), which fails without the refusal;
  • on two real gloo ranks (cpu_dist_workers), rank 0 handed a replicated plan and rank 1 an FSDP2 one, or none: 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. Rank 0 handed a non-plan fails both ranks;
  • on two real gloo ranks, rank 1 alone refused, by two configs or a dtype before the checkpoint is fetched or by a quantized hf_config before the FSDP2 load: both ranks raise rank 1's refusal. The fetch and the load stand in as collectives on a group bounded to 10 s, so without the change the test fails fast (rank 0 gets a dropped connection) instead of hanging;
  • a caller's own model class: loaded whole on every rank, planned without building the model, with no sharding knob, and refused by the sharded load with one.

• With [OMNIML-6054] Decide how a model is distributed, from the model built on the meta device #2737's plan-decision tests: the config returned with the plan is the one every rank read, unchanged by rank 0's meta-device build (new test_the_config_returned_with_the_plan_is_as_every_rank_read_it, which fails without its fix).

• A single process's load, placed as from_pretrained places 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 in tests/unit/torch/utils/test_unplaced_source_keys.py, which now also checks the plan.

• fsdp2_wrap leaving the plan to the load, and refusing a model recorded as replicated: CPU unit tests in tests/unit/torch/utils/test_fsdp2_utils.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: [OMNIML-6054] [1/2] Move fsdp2_wrap out of distributed.py into fsdp2_utils #3140's test_fsdp2_wrap_keeps_its_old_path.

• Mutation-checked: 17 mutants, each failing a test. They drop the rank-GPU default, pin a single process's model, drop the config handoff or the world-size check, load an expert-parallel or FSDP2 plan whole, accept a string, let each rank load the plan it was handed, record no plan in the load or record one in fsdp2_wrap too, set no group up for a single process's FSDP2 load, set one up over a live group, or have the old fsdp2_wrap path drop an argument or the result.

• New tests in tests/gpu/torch/utils/test_model_load_utils.py run load_pretrained left to decide, on every GPU of the node:

  • a dense model loads whole on each rank, and with MODELOPT_FSDP_SIZE=world FSDP2-sharded;
  • each rank's weights match the checkpoint;
  • the load records its unplaced keys and its plan.
  • It needs 2 GPUs, so this PR's GPU CI lane, on one GPU, skips it, and main's nightly GPU run, on two, runs it. It passed on 2 and 4 GPUs on 2026-10-29 (the GPU check below).

• tests/examples/hf_ptq/test_example_utils.py: get_model's stand-in model is a module now, since load_pretrained records its plan on it.

• The load refusals are one parametrized test, and the single-process load test gives way to [OMNIML-6080] Derive the unplaced source keys once, at load, through a core load_pretrained #2733's real load. Each drop is checked by mutation.

• A single process without a launcher is refused an FSDP2 load, without stubbing the setup, and a handed config is left as the caller passed it. A rank of a multi-rank launch whose process group is not set up is refused any load, before anything is fetched. The refusal table has an explicit dtype.

• Locally on this revision, the stack rebased on main c6fa613: pre-commit over the whole PR against main, at every PR of the stack; and at the stack's tip the full tests/unit on 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.

• Re-cut on 2026-10-27 (see the note at the top): pre-commit, mypy included, passes at this head over the whole change against main, and the series' PRs merged together give a tree byte-identical to the former stack's tip. This PR's own unit tests pass at this head on CPU (torch 2.14, transformers 5.18); CI runs the full unit suite.

• Review fixes, 2026-10-27: a rank of a multi-rank launch without its process group is refused on every route, before anything is fetched; a caller's model class with a knob asking to shard is refused before the model is built; a changelog entry for fsdp2_wrap's move. Each behaviour fix has a test that fails without it; pre-commit, mypy included, and this PR's own unit tests pass at this head.

• Rebased on 2026-10-27 onto main 7c3c868, after [OMNIML-6064] Load each rank's checkpoint shard straight into a meta-built model #2739 merged: pre-commit, mypy included, and this PR's own unit tests pass at this head, and the series' PRs merged together add no unit-test failure over main (tests/unit/torch utils, export, quantization and models, on transformers 5.18 and 5.5.4 and on torch 2.8). [OMNIML-6064] Load each rank's checkpoint shard straight into a meta-built model #2739's FSDPModule import, now on main, and this PR's CPUOffloadPolicy, fully_shard are one import line, and fsdp2_wrap sits ahead of the meta-DCP loader, as before.

• Review fixes, 2026-10-27 (Claude's review after the rebase): a sharded load takes dtype=None and torch_dtype as from_pretrained does. None and "auto" keep the checkpoint's dtype under either name, and any other dtype is refused with the dtype message, where torch_dtype got the placement one. The new cases fail without the change.

• Review suggestions, 2026-10-27 evening: a launch whose processes set no group up (hf_ptq under torchrun without --use_fsdp2) loads the whole model on each process again, refusing there only a plan for several ranks and a sharded load; a handed plan is checked before the checkpoint is fetched; an FSDP2 layout decided from rank 0's GPU memory says so when it refuses placement options; fsdp2_wrap imports the layer detection where it uses it.

• Review fixes, 2026-10-27 night: a launch without its process group refuses MODELOPT_EP_DEGREE above 1 or MODELOPT_FSDP_SIZE naming the world before anything is fetched, and the sharded-load refusal names the plan to pass. The refusal test's new rows fail without the change.

• Review suggestions, 2026-10-27, last round: on a launch without its group, a MODELOPT_EP_DEGREE that is not an EP degree is refused with the plan decision's message, and a padded value is read as its number; load_pretrained's docstring says the automatic layout weighs the model in the checkpoint's own dtype. The refusal test's new row fails without the change.

• Review fix, 2026-10-27, last round: a one-rank sharded plan in a launch without its group is refused before the checkpoint is fetched, and the advice for a model with no plan names load_pretrained again, which records the plan from this PR on.

• Review fixes, 2026-10-27, last round: hf_config with a plan that loads the model whole is refused before the checkpoint is fetched, with the no-plan case's message, and the new pre-fetch test row fails without the change; fsdp2_wrap's docstring says the FSDP2 load wraps with it.

• Review suggestions, 2026-10-28: the refusals that depend only on a handed plan's arguments (auto_class, dtype, options the sharded load does not take) run before the checkpoint is fetched, and the new test row fails without the change; a test docstring no longer says fsdp2_wrap loads.

• Review suggestions, 2026-10-28, fifth round: the handed-plan test pins the auto_class and dtype arms of the pre-fetch refusal, each failing with its arm left out of the pre-fetch call; the four function-local imports of modelopt.torch.utils.distributed go, now that [OMNIML-6054] Decide how a model is distributed, from the model built on the meta device #2737 imports it at module top.

• Declined review suggestion, 2026-10-28, fifth round: one helper for the kwargs view the pre-fetch refusal and _load_sharded read; the two agree, and the handed-plan test pins each pre-fetch arm.

• Rebased on 2026-10-29 onto [OMNIML-6054] Decide how a model is distributed, from the model built on the meta device #2737's current head 2a6290b470a (main 2b900c8fa41). This PR's own diff is unchanged, except that its test additions follow the module-mirroring test files: the plan-decision test goes to test_model_load_utils.py ([OMNIML-6054] Decide how a model is distributed, from the model built on the meta device #2737 deleted test_resolve_distribution_plan.py), and the fsdp2_wrap old-path test to test_distributed.py (main renamed test_distributed_run.py). One test fixture changed: main's broadcast_from now picks its sender by modelopt.torch.utils.distributed.rank(), which the load-routing tests' process_group fixture sets to 1 for the FSDP2 load, so run_as now runs the real broadcast with the deciding rank as its sender. No source change. Pre-commit passes against main; tests/unit/torch/utils and the hf_ptq argument and example-utils tests pass on CPU with transformers 5.18 and 5.5.4.

• 2026-10-29, at 59e83993426: tests/unit/torch/export and tests/examples/hf_ptq/test_carry_over_layouts.py also pass on CPU with transformers 5.18 and 5.5.4, apart from 7 test_export_diffusers.py cases that fail the same way on [OMNIML-6054] Decide how a model is distributed, from the model built on the meta device #2737's head in this environment; GitHub runs no unit tests on this PR until it targets main.

• GPU check, 2026-10-29, at 59e83993426, the first run of this PR's multi-GPU tests, on GB200 (torch 2.11.0a0 from NGC PyTorch 26.02, transformers 5.9.0): the new test_load_pretrained_distribution.py, the two files whose tests now record the plan by hand (test_model_load_utils.py, test_export_carry_over.py), and test_distributed.py, whose fsdp2_wrap tests import it through its old path:

  Run at 59e83993426	2 GPUs	4 GPUs
  tests/gpu/torch/utils/test_load_pretrained_distribution.py	2 passed, 20.9 s	2 passed, 23.2 s
  tests/gpu/torch/utils/test_model_load_utils.py	4 passed, 25.6 s	4 passed, 26.9 s
  tests/gpu/torch/export/test_export_carry_over.py	2 passed, 22.8 s	2 passed, 29.2 s
  tests/gpu/torch/utils/test_distributed.py	8 passed, 21.9 s	8 passed, 26.6 s

• Rebased on 2026-10-29 onto main 70147bcbfa1, [OMNIML-6054] Decide how a model is distributed, from the model built on the meta device #2737's squash, after [OMNIML-6054] Decide how a model is distributed, from the model built on the meta device #2737 merged; this PR now targets main. Its own diff is unchanged. The one conflict was the changelog: main's Transformer Engine 1.x line and this PR's fsdp2_wrap line are both kept. Pre-commit, mypy included, passes against main; tests/unit/torch/utils, tests/unit/torch/export and the hf_ptq tests that load through load_pretrained pass on CPU with transformers 5.18 and 5.5.4, apart from the same 7 test_export_diffusers.py cases as above. GitHub now runs the full unit, GPU, example and regression CI on it.

• Review fixes, 2026-10-29 (cjluo-nv's review): a refusal on any rank fails every rank, of the arguments before the checkpoint is fetched and of what the config says before the sharded load (the argument checks now run after the plan's broadcast; the sharded load's checks moved into _sharded_load_kwargs, run inside the all_reduce). The new two-rank test fails on all three of its cases with the change reverted. The routing tests moved into the unit and GPU test_model_load_utils.py files, test bodies unchanged; the changelog entry is two sentences; the MTP refusal's reason is above. Pre-commit, mypy included, passes against main; tests/unit/torch/utils, tests/unit/torch/export and the hf_ptq tests that load through load_pretrained pass on CPU with transformers 5.18 and 5.5.4, apart from the same 7 test_export_diffusers.py cases as above.

• Review fix, 2026-10-29, cjluo-nv's second round: a rank that cannot see the checkpoint rank 0 fetched fails every rank, the fetch running inside the same all_reduce as the other refusals. The two-rank visibility test (test_a_rank_that_cannot_see_the_checkpoint_fails_every_rank) also loads by a handed plan, on the test's own gloo group with the error caught on both ranks, and fails in about 8 s with the change reverted. Pre-commit, mypy included, passes against main; tests/unit/torch/utils, tests/unit/torch/export and the hf_ptq tests that load through load_pretrained pass on CPU with transformers 5.18 and 5.5.4, apart from the same 7 test_export_diffusers.py cases.

• Rebased on 2026-10-30 onto main 99c3834f62c (after [OMNIML-6081] [1/2] Pack sharded weights in FSDP2 weight-update windows, and NVFP4 DTensor shards #3099 and [OMNIML-6081] [2/2] Fold fused experts' quantizers into their weights shard-locally #2750 merged), without conflicts; this PR's own diff is unchanged. Pre-commit, mypy included, passes against main.

• Review fix, 2026-10-30 (cjluo-nv): config= given as a path, as from_pretrained takes it, is read as the load reads it, and the plan decision, the loading hook and a sharded load use that config, so a weights-only checkpoint with its config in another directory loads again (it had failed on the decision's read of the checkpoint's own config.json). test_a_config_passed_for_the_load_needs_none_on_disk now also loads one through LlamaForCausalLM, with and without model_type; both cases fail without the change. Pre-commit, mypy included, passes against main; tests/unit/torch/utils and the hf_ptq example tests pass on CPU with transformers 5.18 and 5.5.4.

• Review suggestion, 2026-10-30 (Claude): the config= path is read through the model class's own config_class, as from_pretrained reads it, and through AutoConfig for the auto classes; a new row loads a model class of the caller's own whose config class AutoConfig cannot resolve, and fails without the change. Pre-commit passes against main; tests/unit/torch/utils and the hf_ptq example tests pass on CPU with transformers 5.18 and 5.5.4.

• Review fix, 2026-10-30 (Claude): a process without a group, or in a group of one, takes DistributionPlan() without deciding or reading a config, as before plans existed, so a model class of the caller's own with its model_type given loads again whatever AutoConfig knows (test_a_model_class_of_the_callers_own_reads_its_config_itself, which fails without the change); a model class's own config class reads a config= path without trust_remote_code. Pre-commit passes against main; tests/unit/torch/utils and the hf_ptq example tests pass on CPU with transformers 5.18 and 5.5.4.

• Review fix, 2026-10-30 (Claude), the several-rank arm: the plan decision reads the checkpoint's config as the load does, through the model class's own config_class or, for an auto class, AutoConfig, so a model class of the caller's own loads whole on every rank whatever AutoConfig knows; test_a_model_class_of_the_callers_own_loads_whole has a two-rank row with such a config on disk, which fails without the change. The other AutoConfig reads on the load path (a model_type left to read, which main reads too; the FSDP2 route's, which only auto classes reach) refuse nothing main loaded. Pre-commit passes against main; tests/unit/torch/utils and the hf_ptq example tests pass on CPU with transformers 5.18 and 5.5.4.

• Review fixes, 2026-10-30 (Claude): on a plan that loads the model whole, the options only the FSDP2 load takes (cpu_offload, mp_policy, broadcast_chunk_size) are refused with the other load checks, rather than reach from_pretrained; MODELOPT_EP_DEGREE has one parse, the plan decision's, which a launch without its group now uses too and which takes surrounding spaces and refuses a degree below 1. New rows (an FSDP2 option on a replicated plan, a padded auto, a degree of -2) fail without the change. Source insertions were then 509, a few over the ~500 budget, for these fixes. Pre-commit passes against main; tests/unit/torch/utils and the hf_ptq example tests pass on CPU with transformers 5.18 and 5.5.4.

• Review suggestions, 2026-10-30 (Claude's approving review at cb7fb1e8495): a distributed whole-model load defaults each rank's device_map to its GPU only where CUDA is available, so a run without GPUs (a gloo group on CPU) loads whole again; a model type left to read is read as the model class reads its config; MODELOPT_FSDP_SIZE takes surrounding spaces as MODELOPT_EP_DEGREE does. New rows (two gloo ranks loading whole without a device_map, a caller's model class whose model type is read, a padded MODELOPT_FSDP_SIZE) fail without the change. Source insertions are 508. Pre-commit passes against main; tests/unit/torch/utils and the hf_ptq example tests pass on CPU with transformers 5.18 and 5.5.4.

• 2026-10-30, line budget (cjluo-nv): the routing's comments and docstrings are tightened, with no change in what a load does or refuses, bringing source insertions to 488; the automatic layout's placement refusal reads MODELOPT_FSDP_SIZE through its one parse, _dense_from_env (Claude's last suggestion). Pre-commit passes against main; tests/unit/torch/utils and the hf_ptq example tests pass on CPU with transformers 5.18 and 5.5.4.

• Split on 2026-10-30 (cjluo-nv's review): the fsdp2_wrap move, its old path, its changelog line and its test went to [OMNIML-6054] [1/2] Move fsdp2_wrap out of distributed.py into fsdp2_utils #3140 ([1/2], 54 source lines); this PR is rebased onto it and keeps the rest (434 source lines). The two merged together give exactly the tree of this PR's previous head a51881aa6a8 on main cb9ab4de069; the rebase dropped three commits whose whole change [OMNIML-6054] [1/2] Move fsdp2_wrap out of distributed.py into fsdp2_utils #3140 now carries, and resolved the commits that had moved or reworded fsdp2_wrap to [OMNIML-6054] [1/2] Move fsdp2_wrap out of distributed.py into fsdp2_utils #3140's version. Pre-commit, mypy included, passes against main at both heads; tests/unit/torch/utils and the hf_ptq example tests pass on CPU with transformers 5.18 and 5.5.4 at both heads.

• 2026-10-30, after [OMNIML-6054] [1/2] Move fsdp2_wrap out of distributed.py into fsdp2_utils #3140's review fixes: rebased onto [OMNIML-6054] [1/2] Move fsdp2_wrap out of distributed.py into fsdp2_utils #3140 6c44b40df74 (it now also carries fsdp2_wrap's dtype helpers, 114 source lines); this PR's own diff is unchanged. Review suggestions (Claude's approving review): load_pretrained's docstring says that a single process, or a group of one, reads no knob and takes the FSDP2 load only from a plan passed as distribution, and that a launch without its group loads with no device_map default. Source insertions: 435. Pre-commit passes against main; tests/unit/torch/utils and the hf_ptq example tests pass on CPU with transformers 5.18 and 5.5.4.

• Review fixes, 2026-10-30 (Claude): the whole-model load runs inside the same failure-carrying all_reduce as the refusals, so a rank whose load fails alone (out of memory on its GPU, say) fails every rank rather than leave the others in the next collective; test_a_whole_load_that_fails_on_one_rank_fails_every_rank (two gloo ranks) fails without it. set_distribution_plan's docstring wording. Source insertions: 439. Pre-commit passes against main; tests/unit/torch/utils and the hf_ptq example tests pass on CPU with transformers 5.18 and 5.5.4.

• Rebased on 2026-10-30 onto [OMNIML-6054] [1/2] Move fsdp2_wrap out of distributed.py into fsdp2_utils #3140 48891f691fb (its old path's narrower import error); this PR's own diff is unchanged. Pre-commit passes against main; tests/unit/torch/utils and the hf_ptq example tests pass on CPU with transformers 5.18 and 5.5.4.

• Review fix, 2026-10-30 (Claude), rebased onto [OMNIML-6054] [1/2] Move fsdp2_wrap out of distributed.py into fsdp2_utils #3140 26bf59ace7c: a model FSDP2-sharded after load_pretrained recorded a replicating plan is refused by distribution_plan(), rather than read as replicated by the export; record-once stays, and the load_pretrained and fsdp2_wrap docstrings say to wrap a model loaded with from_pretrained and record its plan after. The two-rank test_a_model_with_no_plan_raises_in_a_run_of_several_ranks covers it and fails without the change. Source insertions: 453. Pre-commit, mypy included, passes against main; tests/unit/torch/utils and the hf_ptq example tests pass on CPU with transformers 5.18 and 5.5.4.

• Review fixes, 2026-10-30 (cjluo-nv, Claude), rebased onto [OMNIML-6054] [1/2] Move fsdp2_wrap out of distributed.py into fsdp2_utils #3140 a4bdab004de: the FSDP2 load is refused, with the other load checks and so on every rank, on a rank without a GPU, which in a group the caller set up used to fail inside the load while the others waited in its broadcasts (test_a_load_refused_on_one_rank_fails_every_rank's new two-rank no-gpu row, which fails without it); config, variant or hf_config passed as None, their default, is taken as not passed instead of refused (three new rows of test_a_load_it_decides_reads_the_config_once, which fail without it). Source insertions: 466. Pre-commit, mypy included, passes against main; tests/unit/torch/utils and the hf_ptq example tests pass on CPU with transformers 5.18 and 5.5.4.

• Review fixes, 2026-10-30 (Claude), rebased onto [OMNIML-6054] [1/2] Move fsdp2_wrap out of distributed.py into fsdp2_utils #3140 a627d7d6477: a distributed whole-model load defaults a rank's device_map only to a GPU of its own (its local rank below the visible GPU count), so a host with fewer GPUs than ranks no longer hands a rank a device it lacks, and the FSDP2 load's refusal asks the same; the CHANGELOG says a model loaded whole on several ranks goes onto each rank's own GPU; the two-rank gloo load test masks CUDA, so it no longer depends on the host's GPUs; a model class given by name is refused by name; an ep_size that is no whole number ([2], 2.9) is refused rather than raise a TypeError or be truncated. New rows (no_gpu_of_its_own, fewer-gpus-than-ranks, the by-name and whole-number tests) fail without their change. Source insertions: 478. Pre-commit, mypy included, passes against main; tests/unit/torch/utils and the hf_ptq example tests pass on CPU with transformers 5.18 and 5.5.4, and tests/unit/torch/utils without transformers (339 passed, 35 skipped).

• Review fix, 2026-10-30 (Claude): the check that a model recorded as replicated was not sharded after its plan was recorded moved from distribution_plan(), where it walked the model on every read of a replicated plan (the export reads it per weight), to fsdp2_wrap, which now refuses such a model; distribution_plan() is an attribute read again. test_fsdp2_wrap_refuses_a_model_recorded_as_replicated fails without the guard. The single-process FSDP2 load's comment says the group it sets up outlives the load outside distributed_run. Source insertions: 478. Pre-commit, mypy included, passes against main; tests/unit/torch/utils and the hf_ptq example tests pass on CPU with transformers 5.18 and 5.5.4, and tests/unit/torch/utils without transformers.

• Review fixes, 2026-10-30 (Claude): a load option outside the sharded load's own set passed as None, from_pretrained's default (device_map=None as hf_ptq passes it, max_memory, offload_folder, quantization_config), is taken as not passed instead of refused by the FSDP2 load or the automatic layout (new placement-none and none rows, which fail without it); the CHANGELOG files what changes for a load in a group of several ranks under Backward Breaking Changes, with MODELOPT_FSDP_SIZE=1 / MODELOPT_EP_DEGREE=1 or a replicated plan to keep it whole; the refusal tests match the option, knob or value a refusal names rather than its wording (tests only). Source insertions: 487. Pre-commit, mypy included, passes against main; tests/unit/torch/utils and the hf_ptq example tests pass on CPU with transformers 5.18 and 5.5.4, and tests/unit/torch/utils without transformers.

• Review fixes, 2026-10-30 (Claude): the docstring says a rank's whole model goes onto its own GPU only when no device_map is passed, device_map=None keeping from_pretrained's load on the CPU for the caller to place (a row now pins it); mp_policy or cpu_offload passed as None no longer refuses a whole-model load, which made the same call load or raise by rank 0's GPU memory (new row, which fails without it). Source insertions: 494. Pre-commit, mypy included, passes against main; tests/unit/torch/utils and the hf_ptq example tests pass on CPU with transformers 5.18 and 5.5.4, and tests/unit/torch/utils without transformers.

• Rebased on 2026-10-30 onto [OMNIML-6054] [1/2] Move fsdp2_wrap out of distributed.py into fsdp2_utils #3140 ad004655ca9 (cjluo-nv's review there: fsdp2_wrap lives in modelopt.torch.utils.fsdp2_utils, free of the Hugging Face packages). This PR's own changes to the wrap, its docstring and its refusal of a model recorded as replicated, moved with it, and their tests into tests/unit/torch/utils/test_fsdp2_utils.py; the docstrings' references to fsdp2_wrap point there. Source insertions: 496. Pre-commit, mypy included, passes against main; tests/unit/torch/utils and the hf_ptq example tests pass on CPU with transformers 5.18 and 5.5.4, and tests/unit/torch/utils without transformers (344 passed, 33 skipped).

• Review fix, 2026-10-30 (Claude): the FSDP2 load refuses a VLM whose language model is quantized, its quantization_config under text_config or llm_config as Kimi-K2.5 keeps it, rather than only a top-level one; its remote code builds it as a causal LM, so the load had read the packed weights as stored. The new quantized-language-model row fails without the change. Source insertions: 498. Pre-commit, mypy included, passes against main; tests/unit/torch/utils and the hf_ptq example tests pass on CPU with transformers 5.18 and 5.5.4, and tests/unit/torch/utils without transformers.

• Review suggestions taken on the user's decision, 2026-10-30 (Claude): fsdp2_wrap refuses a model that, or any of whose submodules, records any distribution plan, not only a replicating one, naming the plan; a rank of a replicated load that has no GPU of its own, and no device_map, warns on that rank (not through warn_rank_0) that it loads where from_pretrained places it. New rows (a sharded plan, a plan on a submodule) and test_a_rank_without_a_gpu_of_its_own_says_so fail without the change, the last also with the warning sent through warn_rank_0. Source insertions: 511, past the budget by the user's decision. Pre-commit, mypy included, passes against main; tests/unit/torch/utils and the hf_ptq example tests pass on CPU with transformers 5.18 and 5.5.4, and tests/unit/torch/utils without transformers.

• Review fixes, 2026-10-30 (Claude): the FSDP2 load refuses a checkpoint only when it would not build the class _auto_class names, so a Qwen3.5 VLM, which it builds through its image-text class as hf_ptq's --use_fsdp2 does, reaches it (test_a_text_only_vlm_the_fsdp2_load_builds_reaches_it, which fails without the change); its refusals offer MODELOPT_FSDP_SIZE=1 only for a plan the load decided (test_a_sharded_load_offers_the_knob_only_for_a_plan_it_decided); the changelog files fsdp2_wrap's refusal of a model load_pretrained returned as breaking; the no-GPU warning no longer claims a rank's peers have GPUs; and parallel_load_and_prepare_fsdp2's docstring says a direct caller records its own plan. Source insertions: 529. Pre-commit, mypy included, passes against main; tests/unit/torch/utils and the hf_ptq example tests pass on CPU with transformers 5.18 and 5.5.4, and tests/unit/torch/utils without transformers (346 passed, 33 skipped).

• Review suggestion, 2026-10-30 (Claude's approving review at c49fd54f6ee): set_distribution_plan and fsdp2_wrap find a recorded plan with one walk, distributed._recorded_plan, and fsdp2_wrap refuses with a ValueError, as set_distribution_plan does for the same record, rather than a RuntimeError; its test expects the ValueError. Source insertions: 543. Pre-commit, mypy included, passes against main; tests/unit/torch/utils and the hf_ptq example tests pass on CPU with transformers 5.18 and 5.5.4, and tests/unit/torch/utils without transformers (346 passed, 33 skipped).

• Review suggestions, 2026-10-30 (Claude's approving review at 4ffccd8088a): a whole-model load takes cpu_offload=False, the FSDP2 load's own default and what hf_ptq passes without --cpu_offload, as asking nothing (new cpu_offload_at_its_default row, which fails without the change); the refusal of placement options under a layout rank 0's GPU memory chose moved into _refuse_sharded_load_options, which already decides which options place a whole model, and uses _whole_instead for its way out (the refusals are unchanged). Source insertions: 538. Pre-commit, mypy included, passes against main; tests/unit/torch/utils and the hf_ptq example tests pass on CPU with transformers 5.18 and 5.5.4, and tests/unit/torch/utils without transformers (346 passed, 33 skipped).

• Review fixes, 2026-10-30 (Claude): under a layout rank 0's GPU memory chose, the FSDP2 load's refusals of a quantized checkpoint and of one carrying an MTP head name the way to load it whole, as its other refusals do, and the changelog lists them (test_a_measured_layout_names_the_way_back_when_it_refuses_a_checkpoint, which fails without the change); the docstring says which FSDP2-only options a whole load takes as asking nothing; the dtype check is written once. Source insertions: 540. Pre-commit, mypy included, passes against main; tests/unit/torch/utils and the hf_ptq example tests pass on CPU with transformers 5.18 and 5.5.4, and tests/unit/torch/utils without transformers (346 passed, 33 skipped).

• Fix, 2026-10-30 (found on [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 by meenchen's review, agreed with that PR): a rank's own GPU is the one its process is bound to once CUDA is initialized, as torch's DeviceMesh takes it, else the one its local rank numbers (_own_gpu_index(), with _has_own_gpu() over it). A launch that exposes one GPU to each process, each bound to cuda:0, no longer reads every rank but the first as GPU-less: the FSDP2 load no longer refuses it, a whole load no longer leaves it on the CPU with a warning, and both place onto the bound device (test_a_rank_bound_to_its_one_visible_gpu_loads_onto_it, and the two-rank test_one_visible_gpu_per_process_shards_on_every_rank, both failing without the change). Source insertions: 556. Pre-commit, mypy included, passes against main; tests/unit/torch/utils and the hf_ptq example tests pass on CPU with transformers 5.18 and 5.5.4, and tests/unit/torch/utils without transformers (346 passed, 33 skipped).

• Review fixes, 2026-10-30 (Claude): an EP degree left to the plan decision is decided as 1, so a multi-rank load of an MoE with no knob set loads without expert parallelism instead of being refused, while a degree MODELOPT_EP_DEGREE asks for is still refused (test_a_load_it_decides_keeps_the_experts_whole, whose left-to-decide row fails without the change); a whole-model load takes broadcast_chunk_size=8, the FSDP2 load's own default, as asking nothing (new row). Source insertions: 563. Pre-commit, mypy included, passes against main; tests/unit/torch/utils and the hf_ptq example tests pass on CPU with transformers 5.18 and 5.5.4, and tests/unit/torch/utils without transformers (346 passed, 33 skipped).

• Rebased on 2026-10-30 onto [OMNIML-6054] [1/2] Move fsdp2_wrap out of distributed.py into fsdp2_utils #3140 74a636f02c0, itself on main 059588b89a5, without conflicts; this PR's own diff is unchanged. Pre-commit, mypy included, passes against main; tests/unit/torch/utils and the hf_ptq example tests pass on CPU with transformers 5.18 and 5.5.4, and tests/unit/torch/utils without transformers.

• Review suggestions, 2026-10-30 (Claude, at 47ca6d1c422): MODELOPT_EP_DEGREE or MODELOPT_FSDP_SIZE set to whitespace only reads as unset (auto), as an empty one does, instead of being refused: env.py strips both knobs, and the two parses no longer strip their own (new exported-blank rows, which fail without the change). The docstring says a sharded load builds a text-only VLM through its image-text class itself and takes no auto_class but AutoModelForCausalLM, so leave it out for one. Source insertions: 567. Pre-commit, mypy included, passes against main; tests/unit/torch/utils and the hf_ptq example tests pass on CPU with transformers 5.18 and 5.5.4, and tests/unit/torch/utils without transformers (345 passed, 35 skipped).

• Review suggestions, 2026-10-30 (Claude, at e78cbe68469): the docstring says an option at from_pretrained's own default, device_map=None included, places nothing and does not keep a load whole. A launch without a group that MODELOPT_EP_DEGREE asks to shard is told to set MODELOPT_EP_DEGREE=1 as well, since no load here shards the routed experts. A sharded load's refusal of an auto_class names its way out (leave it out, or load the model whole with it), and both auto_class refusals say which class the load builds. The ep-knob and model_class test rows now expect the way out and fail without the change. Messages and docstring only. Source insertions: 577. Pre-commit, mypy included, passes against main; tests/unit/torch/utils and the hf_ptq example tests pass on CPU with transformers 5.18 and 5.5.4, and tests/unit/torch/utils without transformers (345 passed, 35 skipped).

• Rebased on 2026-10-30 onto [OMNIML-6054] [1/2] Move fsdp2_wrap out of distributed.py into fsdp2_utils #3140 b4445d14a8a, itself on main d464d30b120, without conflicts; this PR's own diff is unchanged. Pre-commit, mypy included, passes against main; tests/unit/torch/utils and the hf_ptq example tests pass on CPU with transformers 5.18 (778 passed) and 5.5.4 (757 passed, 21 skipped), and tests/unit/torch/utils without transformers (346 passed, 36 skipped).

• Review fixes, 2026-10-31 (cjluo-nv): a load whose plan the load decides parses MODELOPT_EP_DEGREE in the load-arguments check, so a knob one rank cannot parse fails every rank instead of leaving rank 0 in the plan decision's collectives (test_a_knob_one_rank_cannot_parse_fails_every_rank, two ranks, which fails on the group's bound without the change). The unit tests pin the CUDA state the GPU check reads (_pin_cuda), so a host process that set CUDA up no longer hands them cuda:0; with CUDA set up beforehand, five of them failed before. load_pretrained's docstring is shortened to the routes, the record and the refusal guarantees, its extended guidance moved to Routing details above, and the routing's comments are tightened, with no change in behaviour: source insertions 483 against [OMNIML-6054] [1/2] Move fsdp2_wrap out of distributed.py into fsdp2_utils #3140. Pre-commit, mypy included, passes against main; tests/unit/torch/utils and the hf_ptq example tests pass on CPU with transformers 5.18 (779 passed) and 5.5.4 (758 passed, 21 skipped), and tests/unit/torch/utils without transformers (346 passed, 36 skipped).

• Review suggestions, 2026-10-31 (Claude, at 54ae57e1cfd): the refusal under a layout rank 0's GPU memory chose says the options it refuses do not apply to a sharded load, rather than that they place a whole model, as they are any keyword the sharded load does not take (its placed row now matches the new wording and fails without the change). The docstring says to pass device_map=None where the model is placed otherwise, as under init_quantized_weights, which refuses the per-rank device_map a distributed whole-model load defaults to. Source insertions: 484 against [OMNIML-6054] [1/2] Move fsdp2_wrap out of distributed.py into fsdp2_utils #3140. Pre-commit, mypy included, passes against main; tests/unit/torch/utils and the hf_ptq example tests pass on CPU with transformers 5.18 (779 passed) and 5.5.4 (758 passed, 21 skipped), and tests/unit/torch/utils without transformers (346 passed, 36 skipped).

• Review fixes, 2026-10-31 (cjluo-nv, at 32179e5ad47): a sharded load is refused inside init_quantized_weights, on every rank and before anything loads, since the FSDP2 load builds the model with from_config and would skip the quantization and compression the context's from_pretrained applies (test_a_sharded_load_is_refused_inside_init_quantized_weights, with the automatic layout forced to FSDP2 inside the real context and with a handed sharded plan, fails without the change). The routing guidance is published as a docs page, which load_pretrained's docstring links, and the docstring and the page say a sharded load's from_config takes no Hub options. Source insertions: 494 against [OMNIML-6054] [1/2] Move fsdp2_wrap out of distributed.py into fsdp2_utils #3140. Pre-commit, mypy included, passes against main; the docs build passes with warnings as errors; tests/unit/torch/utils, the hf_ptq example tests and the accelerate plugin tests pass on CPU with transformers 5.18 (784 passed) and 5.5.4 (763 passed, 21 skipped), and tests/unit/torch/utils without transformers (346 passed, 36 skipped).

• Review fix, 2026-10-31 (Claude, at 6ed77327b96): inside init_quantized_weights, a multi-rank load whose plan the load decides gives every rank the whole model when nothing asks to shard, without measuring it, and leaves the device map to the context, which computes it and refuses one passed; before, it was refused when the layout sharded and failed on an injected device_map when it replicated. A knob or a handed plan that asks to shard is still refused. test_a_load_it_decides_inside_init_quantized_weights_is_whole_and_unplaced (rank 0's GPU memory fitting the model or not) fails without the change. Source insertions: 495 against [OMNIML-6054] [1/2] Move fsdp2_wrap out of distributed.py into fsdp2_utils #3140. Pre-commit, mypy included, passes against main; tests/unit/torch/utils, the hf_ptq example tests and the accelerate plugin tests pass on CPU with transformers 5.18 (786 passed) and 5.5.4 (765 passed, 21 skipped), and tests/unit/torch/utils without transformers (346 passed, 36 skipped).

• Rebased on 2026-10-31 onto [OMNIML-6054] [1/2] Move fsdp2_wrap out of distributed.py into fsdp2_utils #3140 30ab95bdb53, itself on main e352caca3c3, without conflicts; this PR's own diff is unchanged. Pre-commit, mypy included, passes against main; tests/unit/torch/utils, the hf_ptq example tests and the accelerate plugin tests pass on CPU with transformers 5.18 (786 passed) and 5.5.4 (765 passed, 21 skipped), and tests/unit/torch/utils without transformers (346 passed, 36 skipped).

• Review suggestion, 2026-10-31 (Claude, at c78d2e50ed6): the changelog's backward-breaking entry lists what the sharded load refuses, the class refusal of a layout rank 0's GPU memory chose included, in two sentences, and the whole load's per-rank GPU placement is its own bullet. Text only. Source insertions: 496 against [OMNIML-6054] [1/2] Move fsdp2_wrap out of distributed.py into fsdp2_utils #3140. Pre-commit, mypy included, passes against main.

Before your PR is "Ready for review"

• Is this change backward compatible?: ❌ Three changes are filed under Backward Breaking Changes. A multi-rank load_pretrained now decides how to distribute the model rather than load it whole on every rank (MODELOPT_FSDP_SIZE=1 and MODELOPT_EP_DEGREE=1, or a replicated plan, keep it whole), and a model it loads whole goes onto each rank's own GPU unless device_map is passed. load_pretrained records the plan on the model it returns, even in a single process, and fsdp2_wrap now refuses a model that records one (wrap a model loaded with from_pretrained). Otherwise a single process loads as before unless a plan it is handed shards the model, and distribution is new. fsdp2_wrap's old import path still works ([OMNIML-6054] [1/2] Move fsdp2_wrap out of distributed.py into fsdp2_utils #3140).
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A
• Did you write any new necessary tests?: ✅
• Did you update Changelog?: ✅ 0.49.0: Misc, and three Backward Breaking Changes entries.
• Did you get Claude approval on this PR?: ✅ Approved 2026-10-29 at 1064e3d09ce, after the fetch fix; its suggestions are answered in their threads. Approved again at 3160a0fb703 after the config-path fix; its suggestion is taken in 6c36e2f8a88; its reviews since are answered in f065849c295, c53b4863b21 and cb7fb1e8495; approved again at cb7fb1e8495 and d9d394b8f55, their suggestions taken in 1b308a07425 and a51881aa6a8, re-requested.

Additional Information

Part 6 of 21 of expert-parallel PTQ for Hugging Face MoE models, now [2/2] after #3140; #2737, below it, merged on 2026-10-29. The expert-parallel half of this routing is #2780.

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

  • Distributed loading supports replicated or FSDP2-sharded models through configurable distribution plans, which are recorded on loaded models. Automatic plan selection and Hub authentication and caching options are supported.

• Documentation

  • Added guidance for distributed loading and distribution-plan recording.

[copy-pr-bot]

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.

[coderabbitai]

Note

Reviews paused

It 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 reviews.auto_review.auto_pause_after_reviewed_commits setting.

• Last reviewed commit: fb8bf38
• Latest commit, not reviewed: 0d3d093

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

No actionable comments were generated in the recent review.

ℹ️ Recent review info

⚙️ Run configuration

• Source mode: Auto
• Repository: NVIDIA/Model-Optimizer
• Branch: shengliangx/hfep-13b-load-routing
• File: .coderabbit.yaml
• Commit: fb8bf38
• Layers: Repository YAML
• Review profile: CHILL
• Plan: Enterprise
• Run ID: 45c74fdb-8dbc-4269-9909-f3d32a1010a4

📥 Commits

Reviewing files that changed from the base of the PR and between 6150cee and fb8bf38.

📒 Files selected for processing: 4

• CHANGELOG.rst
• modelopt/torch/utils/plugins/model_load_utils.py
• tests/unit/torch/utils/test_load_pretrained_distribution.py
• tests/unit/torch/utils/test_unplaced_source_keys.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 10 remain after this review.

Walkthrough

load_pretrained now accepts or resolves a distribution plan and selects whole-model or FSDP2-sharded loading. It validates inputs, records load metadata, and provides FSDP2 wrapping through the model-loading plugin.

Changes: Distributed model loading

Layer / File(s)	Summary
Plan-based load routing modelopt/torch/utils/plugins/model_load_utils.py, tests/unit/torch/utils/test_load_pretrained_distribution.py, tests/gpu/torch/utils/test_load_pretrained_distribution.py, tests/examples/hf_ptq/test_example_utils.py, CHANGELOG.rst, tests/unit/torch/utils/test_unplaced_source_keys.py	load_pretrained resolves or accepts a distribution plan and selects whole-model or FSDP2-sharded loading. Sharded loads validate supported arguments and checkpoints. The loader records the plan and unplaced keys. Unit and GPU tests cover routing, validation, and placement.
FSDP2 wrapping and delegation modelopt/torch/utils/plugins/model_load_utils.py, modelopt/torch/utils/distributed.py, tests/unit/torch/utils/test_distributed_run.py, tests/unit/torch/utils/test_load_pretrained_distribution.py	The plugin implements FSDP2 wrapping, and distributed.py delegates to it. Tests check argument forwarding and confirm that wrapping alone does not record a distribution plan.
Distribution metadata in distributed tests tests/gpu/torch/export/test_export_carry_over.py, tests/gpu/torch/utils/test_model_load_utils.py	Distributed test models record an FSDP distribution plan before export-related checks.

Estimated code review effort: 4 (Complex) | ~45 minutes

Change: Feature · Severity of issue fixed: Medium

Sequence Diagram(s)

sequenceDiagram
  participant Caller
  participant load_pretrained
  participant _route
  participant _load_sharded
  participant parallel_load_and_prepare_fsdp2
  participant from_pretrained
  Caller->>load_pretrained: Provide checkpoint and optional distribution plan
  load_pretrained->>_route: Validate plan and select load path
  alt Dense FSDP plan
    load_pretrained->>_load_sharded: Load checkpoint with config
    _load_sharded->>parallel_load_and_prepare_fsdp2: Load on local CUDA device
  else Whole-model plan
    load_pretrained->>from_pretrained: Load model with device map
  end
  load_pretrained-->>Caller: Return model with plan and unplaced keys recorded

Loading

Suggested reviewers: kevalmorabia97, sugunav14

Merge Risk: Minimal · up to fb8bf

The documented distributed loading paths are ready for normal checks. Users selecting an unsupported expert-parallel plan must choose a supported loading option.

Pre-merge checks: 5

Check name	Status	Explanation
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 modelopt Python files, torch.load uses weights_only=True; no added weights_only=False, allow_pickle=True, ha…
Title check		The title clearly and concisely describes the main change: routing load_pretrained based on the distribution plan.
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.

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 modelopt Python files, torch.load uses weights_only=True; no added weights_only=False, allow_pickle=True, hardcoded trust_remote_code=True, eval()/exec() call, or # nosec comment exists. The pull request changes no pyproject.toml or requirements*.txt file, so it adds no dependency requiring license review.

Finishing Touches

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

[codecov]

Codecov Report

❌ Patch coverage is 98.82353% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.75%. Comparing base (99c3834) to head (a51881a).
⚠️ Report is 3 commits behind head on main.

Files with missing lines	Patch %	Lines
modelopt/torch/utils/plugins/model_load_utils.py	98.80%	2 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##             main    #2779       +/-   ##
===========================================
+ Coverage   71.55%   81.75%   +10.19%     
===========================================
  Files         664      657        -7     
  Lines       76659    73928     -2731     
===========================================
+ Hits        54856    60442     +5586     
+ Misses      21803    13486     -8317     

Flag	Coverage Δ
examples	?
unit	62.51% <97.05%> (+0.18%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[shengliangxu]

/ok to test 3ac7c7b

[shengliangxu]

/ok to test df76b36

[github-actions]

PR Preview Action v1.8.1
🚀 View preview at https://NVIDIA.github.io/Model-Optimizer/pr-preview/pr-2779/
Built to branch gh-pages at 2026-10-30 07:08 UTC. Preview will be ready when the GitHub Pages deployment is complete.

[shengliangxu]

/ok to test aa3615e

[shengliangxu]

/ok to test 168c90d

[shengliangxu]

/ok to test 537ff8d

[shengliangxu]

/ok to test aa2a065

[claude]

[SUGGESTION] This entry's class-refusal clause is narrower than the code, and the entry is well past CONTRIBUTING's "one or two sentences … for external users".

Accuracy. "a model class other than a causal LM is refused when MODELOPT_FSDP_SIZE or MODELOPT_EP_DEGREE asks to shard" describes _decide_distribution's auto_class refusal (model_load_utils.py:1441-1448) only. A measured layout refuses a class too: with distribution=None, no knob set and auto_class left out, _decide_distribution builds on meta, dense="auto" picks "fsdp" for a checkpoint over 60% of rank 0's GPU, and _sharded_load_kwargs then raises NotImplementedError because _auto_class(hf_config) is not builds (:588-602) — e.g. a large Phi-4-multimodal-style config AutoModelForCausalLM accepts but _auto_class routes to AutoModelForImageTextToText. The refusal itself is fine (it ends with _whole_instead(decided)), but a reader of this entry concludes their knob-free multi-rank load of a non-causal checkpoint is unaffected.

Length. The first sentence is ~110 words carrying five distinct behaviour changes. docs/source/guides/_distributed_loading.rst is now the right home for the detail; the changelog needs what users must do.

Suggested replacement for the entry:

- ``load_pretrained`` called in a process group of several ranks now decides how to distribute the
  model rather than load it whole on every rank: a model too big for one GPU is FSDP2-sharded, and
  the sharded load refuses placement arguments, a ``dtype``, a quantized checkpoint, one carrying an
  MTP head, a model class it does not build, and a call inside ``init_quantized_weights``; an MoE
  loads without expert parallelism until the expert-parallel load lands. Set ``MODELOPT_FSDP_SIZE=1``
  and ``MODELOPT_EP_DEGREE=1``, or pass ``distribution=DistributionPlan(<ranks>, dense="replicate")``,
  to keep loading it whole on every rank; see :doc:`/guides/_distributed_loading`.

(The device_map note and the "whole load goes onto each rank's own GPU" behaviour read better as their own short bullet than inside this one.)

[shengliangxu]

Taken in 0d3d0937c05, close to your wording, so the entry now lists every class the sharded load refuses and not only the knob case.

• The entry lists what the sharded load refuses: placement arguments, a dtype, a quantized checkpoint, one carrying an MTP head, a model class it does not build, and a call inside init_quantized_weights. The class refusal now covers a layout rank 0's GPU memory chose.
• It says an MoE loads without expert parallelism, which MODELOPT_EP_DEGREE cannot ask for until the expert-parallel load lands.
• It names the settings that keep loading whole, and points to the new guide by its title. I used the title because the changelog has no :doc: links and is also read on GitHub.
• The per-rank GPU placement of a whole load is now a short bullet of its own.

[shengliangxu]

/ok to test c78d2e5

[shengliangxu]

/ok to test 0d3d093