[shengliangxu]

Note

Re-parented on 2026-10-30 onto #3136, #2741 [2/3] (the expert-parallel loader). #2741 was split into three PRs on review: #3131 [1/3], #3136 [2/3] and #2741 [3/3], its post-load refusals. This PR needs only the loader, so it sits on [2/3] and no longer waits for [3/3]. The series first became a tree on 2026-10-27 (this PR was [10/21]).

Base: #3136. The diff below is this PR's alone, and it merges after #3136.

Re-parent: this PR's diff is unchanged, except that its dense-part load re-anchors around #3136's load guard. from_pretrained, with this PR's warmup and dense-meta contexts, runs inside the guard. The dense part's FSDP2 + DCP load, which runs collectives, comes after it. The PR's original 37 commits are condensed into one; the review-round commits are kept.

Series layout

• main ([OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745, [OMNIML-6064] Load each rank's checkpoint shard straight into a meta-built model #2739, [OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736, [OMNIML-6054] Decide how a model is distributed, from the model built on the meta device #2737, [OMNIML-6082] [1/3] Let a reverse conversion record a merge rule no key matched #3100, [OMNIML-6082] [2/3] Move each tensor of a sharded model whole to one owner rank #3101, [OMNIML-6081] [1/2] Pack sharded weights in FSDP2 weight-update windows, and NVFP4 DTensor shards #3099, [OMNIML-6073] [1/2] Calibrate and fold expert-parallel fused experts on each rank's own experts #3122, [OMNIML-6060] [1/3] Add require_transformers, and build a VLM without its LM head #3131 [1/3] and [OMNIML-6082] [3/3] Write a sharded model's HF checkpoint shards from each tensor's owner rank #2751 have merged)
  • [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 → {[OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067, [OMNIML-6071] Calibrate a distributed model on per-rank data #2740}; [OMNIML-6054] [1/2] Move fsdp2_wrap out of distributed.py into fsdp2_utils #3140 → [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779
  • shengliangx/hfep-base-6a-ep-loader = [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 + [OMNIML-6071] Calibrate a distributed model on per-rank data #2740
    • [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 [OMNIML-6060] [3/3] Refuse an expert-parallel load that would calibrate the experts wrong #2741 [2/3] expert-parallel loader
      • [OMNIML-6060] [3/3] Refuse an expert-parallel load that would calibrate the experts wrong #2741 [3/3] the loader's post-load refusals
      • [OMNIML-6061] Load the dense part sharded, and size the load warmup to each rank #2742 dense sharding (this PR)
        • [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 (merges after [OMNIML-6060] [3/3] Refuse an expert-parallel load that would calibrate the experts wrong #2741 [3/3])
        • [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-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 → [OMNIML-6084] Export an expert-parallel model through the no-gather writer #2753 expert-parallel export

What does this PR do?

Type of change: New feature

#2741 replicates the dense (non-expert) part of an expert-parallel model on every rank. That does not fit when the dense part alone is bigger than a GPU, which the plan from #2737 already recognizes (dense="fsdp").

• The dense part is FSDP2-sharded over every rank when the plan says so (dense="fsdp", with the experts sharded over every rank too), and nothing beside the plan decides it:

  • it is left on meta during from_pretrained;
  • each decoder layer is then fully_shard-ed around its experts, which stay expert-parallel. The blocks come from Hugging Face's own _no_split_modules (_fsdp2_wrap_modules), looked up on the language backbone rather than on a VLM wrapper that would also list its vision blocks. A VLM whose language model shows no blocks (no _no_split_modules match, no detected decoder layers) is refused, naming MODELOPT_FSDP_SIZE=1, instead of wrapped by the VLM's own _no_split_modules, which can name vision blocks;
  • each rank loads only its own shards of the dense weights from the checkpoint through DCP ([OMNIML-6064] Load each rank's checkpoint shard straight into a meta-built model #2739).

  No rank ever holds the whole dense part. A dense param the loader already placed as a replicated DTensor is re-registered as a plain local parameter first, so FSDP2 shards it like any other.

• A block-FP8 checkpoint meets the loader's format refusal here, whatever its dense layout. This load does not dequantize block FP8; [OMNIML-6062] Dequantize block-FP8 checkpoints as they load expert-parallel #2743 adds that, with the rule that keeps a block-FP8 dense part replicated.

• A dense part FSDP2-sharded within several expert-data-parallel replicas is refused before loading. The no-gather export cannot write that layout, so load_and_prepare_ep refuses a plan with dense="fsdp" and edp > 1 before it looks at anything, rather than after the load and calibration it would otherwise run. The error says to set MODELOPT_EP_DEGREE to the world size, or MODELOPT_FSDP_SIZE=1. A replicated dense part at edp > 1 stays supported. The decision knows the layout too: left to decide, it replicates such a dense part and warns that it may run out of memory, naming the same way out; asked to shard it, it refuses with the same advice.

• The load-time allocator warmup is sized to the real per-rank footprint. transformers' caching_allocator_warmup otherwise reserves about the whole device on every rank, and torch keeps that reservation, starving the all-to-all's NCCL buffers. Now:

  • expert params count their 1/ep share;
  • dense params left on meta are not counted at all.

  The replacement forwards whatever else transformers passes the hook, and a hook whose leading parameters are not (model, accelerator_device_map, hf_quantizer) keeps the stock sizing. Warmup sizes measured while building this (moved here from the source):

  Model, layout	Sized before	Sized with this PR
  Nemotron-3-Ultra-550B, EP32	~275 GiB per rank (experts at full size)	~91 GiB per rank
  Kimi-K3, EP64, dense part sharded	185.8 GiB on a 184.31 GiB card (dense part on meta counted at full size)	dense part on meta not counted

A dense parameter off the model's dtype (an fp32 router) stays replicated, as in fsdp2_wrap, and the dense meta load refuses a transformers without the get_device hook it patches.

• A rank that fails before the DCP load fails every rank. _meta_dcp_load (main's, from [OMNIML-6064] Load each rank's checkpoint shard straight into a meta-built model #2739; this PR's dense load is its only caller) runs its per-rank preparation (the checkpoint's metadata read, the mapping checks, to_empty) through the failure-carrying all_reduce before dcp.load, so one rank failing there no longer leaves the others waiting in the load. Main's FSDP2 load (parallel_load_and_prepare_fsdp2) fills its shards through set_model_state_dict, not this function, so the fix reaches only this PR's dense load.
• Every alias of a tied parameter is loaded. The dense load names both aliases of a tied pair, so a tied checkpoint that stores only lm_head.weight fills both, as from_pretrained does.

Why shard the dense part, and why separate helpers

A model that needs it. Kimi-K3 at EP64 on GB200s (184.31 GiB each). With the dense part replicated, a rank holds the dense part whole and its 1/64 share of the routed experts. transformers' load warmup sized that at 185.8 GiB, more than the card before any activation. With the dense part sharded, 79.2 GiB of weights stayed resident per rank. The NVFP4 run over 16 GB200 nodes finished in about 10 minutes at a peak of 155 GiB per GPU, on an earlier build of this series.

What it costs where it is not needed. Qwen3-30B-A3B at EP16 on 4 nodes × 4 GB200s (transformers 5.9.0, torch 2.11). Both arms ran hf_ptq.py --recipe general/ptq/nvfp4_default-kv_fp8 --batch_size 2 --calib_size 512 --calib_seq 2048 --skip_generate with MODELOPT_EP_DEGREE=16, back to back in one allocation:

Per rank (max over the 16 ranks)	Replicated (MODELOPT_FSDP_SIZE=1)	FSDP2-sharded (MODELOPT_FSDP_SIZE=16)
Peak allocated after load	6.35 GiB	4.64 GiB
Peak allocated after calibration	8.41 GiB	6.83 GiB
Peak reserved after calibration	17.36 GiB	16.90 GiB
Load time	6.2 s	5.8 s
Calibration time	20.7 s	21.2 s
Export time	4.2 s	4.5 s

Sharding saves 1.7 GiB of resident dense weights per rank and costs about 0.5 s (2%) of calibration in the dense all-gathers. This model does not need it: its dense part is about 1.8 GiB. With MODELOPT_FSDP_SIZE unset, auto chose replicate and reproduced the replicated arm to within 0.01 GiB and 0.1 s. The two layouts export the same checkpoint (see Testing).

The 60% boundary is a heuristic. auto compares what a rank would hold with the dense part replicated (the dense part whole, plus its 1/ep share of the routed experts) with 60% of its GPU's memory. The other 40% is left for calibration's activations and the communication buffers, the all-to-all's among them. The two measured models sit far on either side of the line: Kimi-K3 just over 100% of the card before any activation, and Qwen3-30B-A3B at a few percent. No sweep across models and GPU sizes has placed the line itself. Where it picks wrong, MODELOPT_FSDP_SIZE=1 replicates the dense part and MODELOPT_FSDP_SIZE=<world size> shards it.

Why separate helpers, rather than extending the existing paths:

• parallel_load_and_prepare_fsdp2 builds the model from its config, FSDP2-shards every decoder layer with its experts, and has each layer's owner read the layer and broadcast it to every rank. Under expert parallelism the routed experts are not FSDP2-sharded. from_pretrained places each rank's share of them on the expert-parallel mesh through transformers' tensor-parallel plan, and the all-to-all forward routes tokens to them. The hybrid keeps from_pretrained for the experts and shards only the rest; reusing the FSDP2 loader would mean re-implementing that placement in it.

• fsdp2_wrap wraps each detected decoder layer whole, and the root by default, and refuses a model its detection misses. The hybrid has different wrap units:

  • each block is wrapped around its experts, which stay out of the FSDP2 groups;
  • the root stays replicated;
  • blocks come from the language backbone's _no_split_modules, so a VLM's vision blocks stay replicated.

  Folding that into fsdp2_wrap would change the wrap units and the accepted models of the existing FSDP2 load, so its callers keep it as it is.

• Accelerate offload cannot combine with this load. transformers refuses a device_map beside the tensor-parallel placement the expert-parallel load requests: "tp_plan and device_map are mutually exclusive" on 5.5.4, and the same refusal on 5.18. Offload would also keep a whole copy of the dense part per rank in host memory and move it to the GPU for every calibration forward. Sharding holds 1/world of it and all-gathers one block at a time.

Usage

plan = DistributionPlan(size(), ep_size=size(), dense="fsdp")
model = load_and_prepare_ep(ckpt_path, plan, hf_config=config)

Testing

• In tests/unit/torch/utils/test_model_load_utils.py (first written as test_ep_dense_sharding.py), on CPU under gloo:

  • telling expert params from dense ones through the EP plan, under any prefix, with routers kept replicated;
  • a block-FP8 checkpoint with a sharded dense part refused for its format;
  • a dense part sharded within several expert-data-parallel replicas refused before anything loads (test_ep_loader.py);
  • dense weights resolving to meta during the load, with the device hook and the warmup's sizing restored even after an error;
  • each decoder layer sharded around its experts, and the blocks to wrap: from _no_split_modules on the language backbone, falling back to decoder-layer detection, and refused for a model without blocks and a VLM whose language model shows none (test_meta_dcp_load.py);
  • a replicated DTensor param sharded like a plain one;
  • the shard-then-load order.

• In tests/unit/torch/utils/test_model_load_utils.py (first written as test_ep_load_warmup.py):

  • the per-rank sizing of experts, meta dense weights, buffers and quantized elements;
  • transformers reserving exactly that footprint;
  • falling back to the stock sizing.

• Mutation-checked: each test fails on at least one deliberate break of the code.

• The expert-matcher cases are one table of 7 rows, and other duplicates become cases of the tests that cover them. Each drop is checked by mutation.

• The decision replicating, or refusing, a dense part within several expert-parallel groups; an fp32 router left replicated, with every layer still unsharding; the missing-get_device refusal.

• 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.

• GPU results at a named head are under "GPU runs, 2026-10-30" at the end of this section; the other GPU results here were measured on earlier revisions of this branch.

• tests/gpu/torch/utils/test_expert_parallel_load.py now also runs with the dense part sharded, at EP over every GPU (sharded within several EP groups, the loader refuses it): each rank's experts and the logits match the whole model. At 8192a75 the whole file passed: 12 on 4 GPUs; 8 on 2 GPUs, where the 4 cases with EP over half the GPUs skip.

• Also in tests/gpu/torch/utils/test_expert_parallel_load.py, QAT under expert parallelism: an FP8 or NVFP4 model calibrated expert-parallel must reach the whole model's scales, with its experts over every rank or half of them and the dense part replicated or FSDP2-sharded, and its gradients through the fake quantizers must match. These tests came with the fused-experts quantization ([OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745) before the series was reordered; they land here, the first PR with an FSDP2-sharded dense part. At 8192a75, where the whole file passed:

  Dense layers, EP span	2 GPUs: FP8, NVFP4	4 GPUs: FP8, NVFP4
  replicated, all ranks	pass, pass	pass, pass
  replicated, half the ranks	skipped	pass, pass
  FSDP2-sharded, all ranks	pass, pass	pass, pass

• End to end with the export fixes in [OMNIML-6081] [2/2] Fold fused experts' quantizers into their weights shard-locally #2750, the sharded and replicated dense layouts export the same checkpoint:

  • tiny Qwen3-MoE, EP=2 on 2 GPUs: byte-identical for NVFP4 and FP8;
  • Qwen3-30B-A3B, EP16 over 4 nodes: byte-identical weights and weight scales for NVFP4 and FP8, with only calibration-derived input scales differing.

  Both were measured before [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074's expert weight-scale fix, when both layouts inflated expert scales the same way: the agreement shows the two layouts match, not that those scales were right.

• 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 dense sharding that finds blocks but shards none is refused, pointing at MODELOPT_FSDP_SIZE=1; a TP-planned param keeps its share before the warmup's meta skip, and the TP-plan lookup tries transformers.distributed first. 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.

• Review fixes, 2026-10-27, round 2: the warmup replacement forwards whatever transformers passes the hook and leaves a hook of another signature alone; replicating a dense part within several expert-parallel groups warns that it may run out of memory, with the way out; the dense-sharding comments and progress line say where it wraps and that it shards over every rank. 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).

• Rebased over [OMNIML-6060] [3/3] Refuse an expert-parallel load that would calibrate the experts wrong #2741's last round, which prints the loader's notes through print_rank_0 and no longer imports is_master: this PR's dense-sharding notes import it themselves.

• Review suggestion, 2026-10-28: two of the dense-sharding notes print through print_rank_0; is_master stays where it keeps the meta-parameter listing off the other ranks.

• Review suggestion, 2026-10-28, second round: the EP-aware warmup-sizing note prints through print_rank_0 too; this PR's MoE-helper import joins [OMNIML-6054] Decide how a model is distributed, from the model built on the meta device #2737's module-level one.

• Review suggestion, 2026-10-28, third round: the warmup sizing credits a TP-planned param 1/ep_degree, its share of the ep submesh from_pretrained places it over, rather than 1/world_size. The test's world now holds two ep submeshes and fails with the old divisor.

• Declined review suggestion, 2026-10-28, third round: alias-safe re-registration of replicated parameters tied inside one decoder block, which no model this load supports has; the embedding-to-lm_head tie sits at the root, which is not wrapped.

• Review suggestion, 2026-10-28, fourth round: the warmup's docstring says only that no expert param is in _tp_plan under the EP load, which is what the sizing relies on.

• Review suggestion, 2026-10-28, fifth round: the expert-param matcher is the dot-padded containment test alone, which the exact and suffix matches were special cases of.

• Review suggestions, 2026-10-28, sixth round: the expert matcher's docstring states its one containment test, and the meta-placement context and the sharded dense load each have their own summary.

• Re-cut on 2026-10-29 (see the note at the top): rebased onto the re-cut [OMNIML-6060] [3/3] Refuse an expert-parallel load that would calibrate the experts wrong #2741 with no source conflict. The edits of the deleted test_resolve_distribution_plan.py (the env-knob table, and the dense part the expert-parallel load cannot shard) moved to test_model_load_utils.py. Pre-commit, mypy included, passes at this head over the whole change against main. On CPU at this head (torch 2.14), this PR's tests with test_ep_loader.py, test_meta_dcp_load.py and test_model_load_utils.py pass on transformers 5.18 and 5.5.4.

• Review fixes, 2026-10-29 (cjluo-nv's review bot):

  • A rank failing in _meta_dcp_load's preparation now fails every rank before dcp.load. A bounded two-rank test fails rank 1 either in its metadata read or with to_empty out of memory; with the fix reverted, rank 0 enters the load. (Corrected 2026-10-30: main's FSDP2 load does not call _meta_dcp_load, so the fix reaches only this PR's dense load.)
  • The dense load names both aliases of a tied parameter. A tied tiny Qwen3-MoE storing only lm_head.weight failed through this caller on the CPU before the fix, and now loads both aliases. The load also takes its device from the mesh.
  • Imports moved to module level, and the helper docs compressed; the benchmark numbers moved above.
  • Pre-commit, mypy included, passes against main. tests/unit/torch/utils passes on the CPU at this head on transformers 5.18 and 5.5.4. GPU runs of that head did not happen; see 8192a75 below.

• Re-parented 2026-10-30 onto [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 ([OMNIML-6060] [3/3] Refuse an expert-parallel load that would calibrate the experts wrong #2741 [2/3]):

  • from_pretrained re-anchored inside the load guard; the dense part's load after it.
  • Claude's three suggestions taken: the GPU test checks every rank for a missing gradient before its collective comparisons; _meta_dcp_load hands its prepared state over as a named tuple; replace_layer_number_by_wildcard is imported from either transformers location.
  • Pre-commit, mypy included, passes against main.
  • tests/unit/torch/utils passes on CPU at this head, on transformers 5.18 and 5.5.4.
  • The GPU tests did not run at that head; see 8192a75 below.

• Review fixes, 2026-10-30 (Claude's re-review): with the dense part sharded, a load in which no dense param came back on meta is refused, since every rank would then have loaded the whole dense part (test_a_dense_part_from_pretrained_did_not_leave_on_meta_is_refused fails without the check); the EP-aware warmup sizing warns when it falls back to transformers' own sizing (the fallback test asserts the warning and fails without it). Rebased with [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 onto main 99c3834; the rest of this PR's diff is unchanged. Pre-commit, mypy included, passes against main, and tests/unit/torch/utils passes on CPU at this head, on transformers 5.18 and 5.5.4.

• Review fixes, 2026-10-30 (cjluo-nv's review bot):

  • _fsdp2_wrap_modules' two imports stay in the function, as fsdp2_wrap's does, with a comment on why (6388dc4; superseded below, they are now at module scope): the export and quantization packages sit above this module, and the export package imports it. Moving them would not cycle today: with both at module scope, every entry order imports on transformers 5.18 and 5.5.4 and on torch 2.8, and import modelopt.torch already loads both modules. So the reason is layering, not cost.
  • The description now carries the Qwen3-30B-A3B costs, the Kimi-K3 case, the 60% rule and the reasons for separate helpers. It also corrects two claims: main's FSDP2 load does not call _meta_dcp_load, and block-FP8 handling lands in [OMNIML-6062] Dequantize block-FP8 checkpoints as they load expert-parallel #2743.
  • Pre-commit, mypy included, passes against main, and tests/unit/torch/utils passes on CPU at this head, on transformers 5.18 and 5.5.4.
  • Rebased onto [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 96be72d, which now refuses an expert-parallel plan matching no module, or styling an expert weight other than grouped_gemm. The dense-meta refusal test's unplaced-keys stand-in returns a list, as the function does.
  • The expert-parallel GPU tests ran at 8192a75 on 2 and 4 GPUs (below).

• Review fixes, 2026-10-30, round 2 (cjluo-nv's review bot):

  • _fsdp2_wrap_modules' two imports, and fsdp2_wrap's, are at module scope (c2f1bf9). Each of model_load_utils, modelopt.torch.export and modelopt.torch.quantization, imported first in a fresh interpreter, loads cleanly on transformers 5.18, 5.9.0 and 5.5.4.
  • The EP-aware warmup sizing also warns when transformers' hook takes other leading parameters, the way a transformers upgrade would reach its fallback; the hook-signature test asserts the warning and fails without it. The refusal of a dense part sharded within several expert-data-parallel replicas also names MODELOPT_EP_DEGREE=1, which keeps the dense part sharded by dropping expert parallelism (201e2a5).
  • Pre-commit, mypy included, passes against main, and tests/unit/torch/utils passes on CPU at this head, on transformers 5.18, 5.9.0 and 5.5.4.

• GPU runs, 2026-10-30, at 8192a75: NVIDIA GB200, torch 2.11.0a0 (nv26.02), transformers 5.9.0, pytest with -p no:rerunfailures.

  Test file	2 GPUs	4 GPUs
  tests/gpu/torch/utils/test_expert_parallel_load.py (EP logits, backward, FP8/NVFP4)	8 passed, 4 skipped (25 s)	12 passed (94 s)
  tests/gpu/torch/utils/test_model_load_utils.py (main's FSDP2 load)	6 passed (26 s)	not run
  tests/gpu/torch/export/test_fsdp2_export.py (main's FSDP2 load)	39 passed (147 s)	not run

  These results were taken at 8192a75. Up to aa85e91 this PR's diff since then was the module-scope imports, the hook-signature warning, the MODELOPT_EP_DEGREE=1 remedy text and two docstrings; the 2026-10-30 rebase below then changed the base under it, and the file was re-run on the rebased stack (the last entry here).

• Review suggestions, 2026-10-30 (Claude's review at 201e2a5, approved): two docstrings no longer over-claim. _fsdp2_wrap_modules leaves vision blocks unwrapped only for a VLM get_language_model_from_vl recognizes, and the dense meta load leaves buffers alone unless they are named .weight/.bias (aa85e91). Pre-commit, mypy included, passes against main, and tests/unit/torch/utils passes on CPU at this head, on transformers 5.18, 5.9.0 and 5.5.4.

• Review fix, 2026-10-30 (cjluo-nv's review bot): test_ep_dense_sharding.py and test_ep_load_warmup.py, which covered model_load_utils alone, moved into test_model_load_utils.py with their fixtures and parametrization, as CONTRIBUTING.md's naming rule asks (004de2a). No test was added or dropped: tests/unit/torch/utils gives the same counts on CPU on transformers 5.18, 5.9.0 and 5.5.4, and pre-commit, mypy included, passes against main.

• Rebased 2026-10-30 with [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 onto the rebuilt shengliangx/hfep-base-6a-ep-loader, which after [OMNIML-6060] [1/3] Add require_transformers, and build a VLM without its LM head #3131 merged holds [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779 (on [OMNIML-6054] [1/2] Move fsdp2_wrap out of distributed.py into fsdp2_utils #3140, which moves fsdp2_wrap into modelopt/torch/utils/fsdp2_utils.py), [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 on main 62a5c48. With fsdp2_wrap in its own module, this PR's paragraph in its docstring moved to fsdp2_utils.py, fsdp2_wrap keeps its local quantization import there (fsdp2_utils imports no Hugging Face or quantization package), and model_load_utils imports _off_dtype_params and fully_shard for the dense sharding itself. The rest of this PR's diff is unchanged. Pre-commit, mypy included, passes against main, and tests/unit/torch/utils passes on CPU at this head, on transformers 5.18, 5.9.0 and 5.5.4.

• GPU run, 2026-10-30, at [OMNIML-6062] Dequantize block-FP8 checkpoints as they load expert-parallel #2743's head 9c917ae, which contains this PR's head f991fa5: tests/gpu/torch/utils/test_expert_parallel_load.py on transformers 5.18.0, the test file unchanged since the run at 8192a75.

  GPUs	Passed	Skipped	Failed
  2	8	4	0
  4	12	0	0

  On 1 GPU every case skips, so it is no evidence. On 2 GPUs the four cases with expert parallelism over half the GPUs skip; on 4 GPUs all 12 run: each rank's experts and the logits against the whole model, and FP8 and NVFP4 quantized training (scales and gradients), with the experts over every GPU or half of them and the dense part replicated or FSDP2-sharded.

• Rebased 2026-10-30 onto [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 e251f6c (meenchen's review round), whose module-scope imports replace the load's local import block: the comment on the load's guard keeps [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136's note on what it rests on, and this PR's own diff is otherwise unchanged. Pre-commit, mypy included, passes against main, and tests/unit/torch/utils passes on CPU at this head, on transformers 5.18, 5.9.0 and 5.5.4.

• Rebased 2026-10-30 onto [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 1491b66 (its MXFP4 refusal and review fixes); this PR's own diff is unchanged. Pre-commit, mypy included, passes against main, and tests/unit/torch/utils passes on CPU at this head, on transformers 5.18, 5.9.0 and 5.5.4.

• GPU run, 2026-10-30, at [OMNIML-6062] Dequantize block-FP8 checkpoints as they load expert-parallel #2743's head 7192e05, which contains this PR at a4e2e1d: tests/gpu/torch/utils/test_expert_parallel_load.py on transformers 5.18.0.

  GPUs	Passed	Skipped	Failed
  1	0	12	0
  2	8	4	0
  4	12	0	0

  On 4 GPUs every case runs, the four with expert parallelism over half the GPUs included, as at the run before. 1 GPU skips everything, so it is no evidence.

• GPU run, 2026-10-30, at [OMNIML-6062] Dequantize block-FP8 checkpoints as they load expert-parallel #2743's head 7c83230, which contains this PR's head 4abc058: tests/gpu/torch/utils/test_expert_parallel_load.py on transformers 5.18.0.

  GPUs	Passed	Skipped	Failed
  1	0	12	0
  2	8	4	0
  4	12	0	0

  The same result as the run before, with the same test file, so it speaks for the code underneath. The file loads an unquantized checkpoint on transformers 5.18, so it cannot show the MXFP4 refusal or the block-FP8 floor of 5.8 firing: [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136's config-refusal table covers the MXFP4 refusal, and [OMNIML-6062] Dequantize block-FP8 checkpoints as they load expert-parallel #2743's 5.7.0 test the block-FP8 floor, on CPU.

• Rebased 2026-10-30 onto [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 ad9559d (meenchen's review, round 3, on the base rebuilt with [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779 47ca6d1); this PR's own diff is unchanged. Pre-commit, mypy included, passes against main, and tests/unit/torch/utils passes on CPU at this head, on transformers 5.18, 5.9.0 and 5.5.4. The GPU tests have not run on this head.

• Rebased 2026-10-30 onto main 059588b with [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 eb6c382, on shengliangx/hfep-base-6a-ep-loader rebuilt there from [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779 3b814a6, [OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067 fbff3f3 and [OMNIML-6071] Calibrate a distributed model on per-rank data #2740 186efae. [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 now also refuses a rank not yet bound to its GPU. This PR's own diff is unchanged. Pre-commit, mypy included, passes against main, and tests/unit/torch/utils passes on CPU at this head, on transformers 5.18, 5.9.0 and 5.5.4. The GPU tests have not run on this head.

• Rebased 2026-10-30 onto [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 56353ef, whose new commits change only tests and one comment; this PR's own diff is unchanged. Pre-commit, mypy included, passes against main, and tests/unit/torch/utils passes on CPU at this head, on transformers 5.18, 5.9.0 and 5.5.4.

• GPU run, 2026-10-30, at [OMNIML-6062] Dequantize block-FP8 checkpoints as they load expert-parallel #2743's head 7540a2e, which contains this PR at 6415c94: tests/gpu/torch/utils/test_expert_parallel_load.py on transformers 5.18.0.

  GPUs	Passed	Skipped	Failed
  1	0	12	0
  2	8	4	0
  4	12	0	0

  The same result as the runs before. On 4 GPUs every case runs, the four with expert parallelism over half the GPUs included; 1 GPU skips everything, so it is no evidence. Since that head, [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 changed only tests and one comment, so the result speaks for this head's code.

• Rebased 2026-10-30 onto [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 75bc2d6, which trims the docstring and comments this PR edits. This PR's own diff is unchanged but for its docstring paragraph on the sharded dense part, now one sentence on the trimmed docstring. Pre-commit, mypy included, passes against main, and tests/unit/torch/utils passes on CPU at this head, on transformers 5.18, 5.9.0 and 5.5.4. The GPU tests have not run on this head.

• Rebased 2026-10-30 with [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 ab06860 onto main d464d30, on shengliangx/hfep-base-6a-ep-loader rebuilt there from [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779 8f9d877, [OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067 d452768 and [OMNIML-6071] Calibrate a distributed model on per-rank data #2740 9a87d15; this PR's own diff is unchanged. Pre-commit, mypy included, passes against main, and tests/unit/torch/utils passes on CPU at this head, on transformers 5.18, 5.9.0 and 5.5.4.

• GPU runs, 2026-10-30, at an earlier head with the same own diff ([OMNIML-6062] Dequantize block-FP8 checkpoints as they load expert-parallel #2743's 6f5391d, which contains this PR at 72f1e49), on transformers 5.18.0:

  • tests/gpu/torch/utils/test_expert_parallel_load.py:

  GPUs	Passed	Skipped	Failed
  1	0	12	0
  2	8	4	0
  4	12	0	0

  On 4 GPUs every case runs, the four with expert parallelism over half the GPUs included; 1 GPU skips everything, so it is no evidence.

  • tests/unit/torch/utils/test_ep_loader.py on a host with one visible GPU: 59 of 59 passed.

• Review fixes, 2026-10-31 (cjluo-nv's review bot): the sharded dense part materializes on the GPU each rank is bound to, not the one its local rank numbers (f0c2a2b; the new test, local rank 1 bound to cuda:0, fails without it), and the FSDP2 wrapping runs in the load's failure-carrying guard, so a rank failing to wrap fails every rank before DCP's collectives (0143941; the new bounded two-rank test fails without it). Pre-commit, mypy included, passes against main, and tests/unit/torch/utils passes on CPU at this head, on transformers 5.18, 5.9.0 and 5.5.4. The GPU tests have not run on this head.

• Rebased 2026-10-31 onto [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 7747c5d, which moves test_ep_loader.py's cases into test_model_load_utils.py and the expert_parallel_mesh ones into test_moe_a2a.py. This PR's test edits moved with those cases, and its source diff is unchanged. Its dense-sharding tests' experts class is now _LayerExperts, because [OMNIML-6060] [3/3] Refuse an expert-parallel load that would calibrate the experts wrong #2741's tests define an _Experts in the same file. Pre-commit, mypy included, passes against main, and tests/unit/torch/utils collects and passes the same cases as before on CPU, on transformers 5.18, 5.9.0 and 5.5.4.

• GPU run, 2026-10-31, at 0143941, which has this head's source: tests/gpu/torch/utils/test_expert_parallel_load.py on transformers 5.18.0.

  GPUs	Passed	Skipped	Failed
  1	0	12	0
  2	8	4	0
  4	12	0	0

  On 4 GPUs every case runs, the dense-sharded arms included. The file has no dense-sharded arm with experts over half the GPUs, by design: this PR refuses a dense part FSDP2-sharded within several expert-data-parallel replicas (edp > 1), a layout the export cannot write yet.

• Review suggestion, 2026-10-31 (Claude): a transformers without get_total_byte_count now warns that the load warmup keeps its own sizing, as the other two fallbacks do (2860c0c); the test of that release expects the warning and fails without it. Pre-commit, mypy included, passes against main, and tests/unit/torch/utils passes on CPU at this head, on transformers 5.18, 5.9.0 and 5.5.4.

• GPU check of the bound-device fix (f0c2a2b), 2026-10-31: 2 GPUs, one visible per process, nccl, LOCAL_RANK 0/1, DistributionPlan(2, ep_size=2, dense="fsdp"). The dense-sharded path ran on both ranks in both runs.

  Head	Rank 0	Rank 1
  2860c0c (fixed)	loads	loads: its local rank is 1, and it materializes on its bound GPU, cuda:0
  591b986 (before the fix)	refuses, naming rank 1's failure	fails: "CUDA error: invalid device ordinal" on cuda:1

• Rebased 2026-10-31 with [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 930bb0f onto main e352cac, on shengliangx/hfep-base-6a-ep-loader rebuilt there from [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779 c78d2e5, [OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067 adb9757 and [OMNIML-6071] Calibrate a distributed model on per-rank data #2740 dea9b79; this PR's own diff is unchanged. Pre-commit, mypy included, passes against main, and tests/unit/torch/utils passes on CPU at this head, on transformers 5.18, 5.9.0 and 5.5.4.

• GPU run, 2026-10-31, at [OMNIML-6062] Dequantize block-FP8 checkpoints as they load expert-parallel #2743's head 939bf6e on main e352cac, which contains this PR's head 7ea4ac7: tests/gpu/torch/utils/test_expert_parallel_load.py on transformers 5.18.0.

  GPUs	Passed	Skipped	Failed
  1	0	12	0
  2	8	4	0
  4	12	0	0

  The same 12 cases as before, now over the rebuilt base, which includes main's fused-MoE GPTQ change. On 4 GPUs every case runs, the dense-sharded ones included: each rank's experts and the logits against the whole model, and both quantized formats' training. 1 GPU skips everything, so it is no evidence.

Before your PR is "Ready for review"

• Is this change backward compatible?: ✅ Additive; the dense part stays replicated unless asked.
• 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?: N/A. The feature's entry lands with [OMNIML-6085] [3/3] Run expert-parallel PTQ from hf_ptq.py #2754.
• Did you get Claude approval on this PR?: Approved 2026-10-28 at 05e21c2ba84, and again on 2026-10-30 at 32d57f5c81c, after the review fixes in Testing, which change main's _meta_dcp_load.

Additional Information

Part 10 of 21 of expert-parallel PTQ for Hugging Face MoE models; merges after #2741.

Summary by CodeRabbit

• New Features
  • Expert-parallel model loading supports dense weights that are replicated or FSDP-sharded, offering more flexible memory and distribution options.
  • FSDP2 wrapping follows model-defined layer boundaries, including language-model layers in vision-language models.
  • Block-FP8 checkpoints are refused by this PR's loader, whatever the dense layout; [OMNIML-6062] Dequantize block-FP8 checkpoints as they load expert-parallel #2743 adds their dequantization and keeps their dense part replicated.
• Bug Fixes
  • Improved loading and training consistency for expert-parallel models, including quantized configurations.

[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: 5e3b73d
• Latest commit, not reviewed: 7ea4ac7

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

• Configuration used: Repository: NVIDIA/Model-Optimizer/.coderabbit.yaml
• Review profile: CHILL
• Plan: Enterprise
• Run ID: e0f3ea3d-383e-4fd6-b3c5-2ac1a0842a5a

📥 Commits

Reviewing files that changed from the base of the PR and between dacd135 and 5e3b73d.

📒 Files selected for processing: 1

• tests/gpu/torch/utils/test_expert_parallel_load.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

EP loading now supports FSDP2 sharding for dense parameters while expert parameters remain expert-parallel. Block-FP8 checkpoints reject explicit dense FSDP loading. Automatic planning can select replicated dense weights. Tests cover planning, loading, sharding, forward behavior, gradients, and quantization.

Changes: Expert-parallel dense loading

Layer / File(s)	Summary
Distribution planning and warmup sizing modelopt/torch/utils/plugins/model_load_utils.py, tests/unit/torch/utils/test_resolve_distribution_plan.py, tests/unit/torch/utils/test_ep_load_warmup.py	Block-FP8 automatic EP plans use replicated dense weights, while explicit dense FSDP sharding raises NotImplementedError. EP-aware warmup sizing accounts for expert shards and can exclude dense parameters on meta.
Wrap boundaries and dense parameter selection modelopt/torch/utils/plugins/model_load_utils.py, tests/unit/torch/utils/test_meta_dcp_load.py, tests/unit/torch/utils/test_ep_dense_sharding.py	FSDP2 wrapping uses detected model boundaries, with decoder-layer fallback and language-backbone selection for VLMs. Dense sharding excludes EP expert parameters.
Dense parameter loading modelopt/torch/utils/plugins/model_load_utils.py, tests/unit/torch/utils/test_ep_dense_sharding.py, tests/unit/torch/utils/test_ep_loader.py	The loader leaves selected dense weights and biases on meta, then FSDP2-shards and DCP-loads them. EP plans with dense FSDP proceed through this loading path.
Distributed forward and training validation tests/gpu/torch/utils/test_expert_parallel_load.py	GPU tests compare forward outputs, gradients, and quantizer scales against whole-model references for replicated and FSDP2-sharded dense parameters.

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

Change: Feature

Sequence Diagram(s)

sequenceDiagram
  participant load_and_prepare_ep
  participant from_pretrained
  participant _load_ep_dense_distributed
  participant _meta_dcp_load
  load_and_prepare_ep->>from_pretrained: Load experts and leave selected dense weights on meta
  load_and_prepare_ep->>_load_ep_dense_distributed: Pass loaded model and dense FSDP mesh
  _load_ep_dense_distributed->>_meta_dcp_load: Load selected dense parameters from checkpoint shards

Loading

Merge Risk: Minimal · up to 5e3b7

The FSDP-sharded dense expert-parallel load path refuses unsupported layouts before it loads anything, and the GPU tests no longer exercise those layouts. No merge-blocking risk remains.

Pre-merge checks: 5

Check name	Status	Explanation
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns		PASS. The reviewed diff adds no torch.load(..., weights_only=False), numpy.load(..., allow_pickle=True), hardcoded trust_remote_code=True, external-input eval()/exec(), or # nosec comments…
Title check		The title "[OMNIML-6061] Load the dense part sharded, and size the load warmup to each rank" names both main changes in the PR: optional FSDP2 sharding of the dense portion during expert-parallel load…
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.

Full details: Security Anti-Patterns

Explanation

PASS. The reviewed diff adds no torch.load(..., weights_only=False), numpy.load(..., allow_pickle=True), hardcoded trust_remote_code=True, external-input eval()/exec(), or # nosec comments in modelopt or examples. The diff also changes no pyproject.toml or requirements file. SECURITY.md excludes tests from these rules, and the remaining changes are tests.

Full details: Title check

Explanation

The title "[OMNIML-6061] Load the dense part sharded, and size the load warmup to each rank" names both main changes in the PR: optional FSDP2 sharding of the dense portion during expert-parallel loading, and per-rank sizing of the Transformers allocator warmup. The changeset adds FSDP2 dense sharding (dense="fsdp", edp=1), block-FP8 handling, and EP-aware warmup sizing, so the title matches the main change. The ticket prefix is a tracker reference, not noise that hides meaning. The title is 80 characters and 14 words, which is longer than ideal but still readable and specific.

Finishing Touches

Generate docstrings

• 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 91.52542% with 15 lines in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (shengliangx/hfep-6a-ep-loader@ee17253). Learn more about missing BASE report.

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

Additional details and impacted files

@@                       Coverage Diff                        @@
##             shengliangx/hfep-6a-ep-loader    #2742   +/-   ##
================================================================
  Coverage                                 ?   61.83%           
================================================================
  Files                                    ?      649           
  Lines                                    ?    72969           
  Branches                                 ?        0           
================================================================
  Hits                                     ?    45117           
  Misses                                   ?    27852           
  Partials                                 ?        0           

Flag	Coverage Δ
unit	61.83% <91.52%> (?)

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.

[github-actions]

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

[shengliangxu]

/ok to test d6cb219

[shengliangxu]

/ok to test a5aecef

[shengliangxu]

/ok to test ffbb850

[shengliangxu]

/ok to test 88ebf30

[shengliangxu]

/ok to test 426a9d1

[shengliangxu]

/ok to test 591b986

[cjluo-nv]

Bot review (gpt-6.1-sol) — DM the bot to share feedback.

Comment: the prior fixes are verified, but bound-device handling and unguarded dense wrapping can still fail valid distributed loads.

Needs action:

• Guard rank-local dense wrapping before the next collective in _load_ep_dense_distributed; add a bounded two-rank wrapping-failure regression (inline).
• Use the caller-bound CUDA device for dense materialization and cover LOCAL_RANK != current_device() (inline).

No action needed:

• ✔️ Resolved since the last review: module-based test naming. Earlier alias, dense-meta, and DCP-preparation fixes remain intact.
• Existing refusal, environment-knob, and GPU-layout test edits are justified by the supported/refused layouts.

[cjluo-nv]

Bot comment.

[P1] Agree on success of the local wrapping stage before entering _meta_dcp_load. _fsdp2_shard_dense runs outside both _refused_together and the new preparation all_reduce. A rank-local exception here—for example an OOM while fully_shard allocates local storage for already-materialized non-.weight/.bias dense parameters—therefore leaves successful peers waiting in _meta_dcp_load's preparation gather. Structural wrap decisions being identical does not make allocation failures identical. Could the non-collective wrapping/name-collection stage use the existing failure-carrying guard too, keeping DCP itself outside it? Extend the bounded two-rank coverage to fail wrapping on one rank and assert every rank raises the originating error before DCP, rather than a collective timeout.

[shengliangxu]

Fixed in 0143941. The FSDP2 wrapping, and the dense names it collects, now run inside the load's failure-carrying guard, with DCP left outside it. A rank that fails to wrap (out of memory, say) makes every rank raise that failure before DCP's collectives. The new bounded two-rank test fails the wrapping on rank 1, stands DCP in by a bounded barrier, and asserts both ranks raise the originating error. It fails without the guard.

[cjluo-nv]

Bot comment.

[P2] Use the bound CUDA device, not dist.local_rank(), for materialization. load_and_prepare_ep's GPU preflight and _experts_sharded_over explicitly honor torch.cuda.current_device(), including a caller-owned group where each process sees one distinct physical GPU as cuda:0 while LOCAL_RANK is 0/1. That layout passes preflight and loads experts on cuda:0, but rank 1 now passes cuda:1 to _meta_dcp_prepare/to_empty, which fails with an invalid device ordinal. With multiple visible GPUs and a remapped binding, it instead materializes dense weights on a different GPU from the experts/FSDP compute device. Keep the CPU branch, but select CUDA through current_device() (as _move_to_fsdp_device does). Extend the device assertion to cover a local rank different from the bound device; the current CPU-only assertion and rank-bound GPU workers miss this.

[shengliangxu]

Fixed in f0c2a2b. The dense part now materializes on the GPU the rank is bound to: _own_gpu_index(), which the GPU check and the load routing share, keeping the CPU branch. The new test runs local rank 1 bound to cuda:0, as a launch showing each process one GPU has it, and asserts DCP gets cuda:0. It fails without the change, which passed cuda:1.

[shengliangxu]

Checked on GPUs since: 2 GPUs, one visible per process, nccl, LOCAL_RANK 0/1, with a dense-sharded expert-parallel plan. At this fix both ranks load, rank 1 on its bound cuda:0. Before it, rank 1 failed with "CUDA error: invalid device ordinal" on cuda:1, and rank 0 refused, naming rank 1's failure.

[shengliangxu]

/ok to test 0143941

[shengliangxu]

/claude review

[claude]

[SUGGESTION] This is now the only one of the three fallbacks that reverts to transformers' full-device sizing without saying so. The signature-mismatch branch below (:1845-1850) and the except Exception branch (:1908-1913) both warn precisely because that sizing "may reserve most of the device" and a later NCCL/OOM failure would then have no stated cause — but a transformers release that removes get_total_byte_count (renames it, moves it out of modeling_utils) lands here and is indistinguishable from a healthy load.

That is also the more likely of the three to happen: the function is private to modeling_utils, and _ep_dense_meta_load already had to grow a two-module lookup for get_device for exactly this reason.

Worth the same warn_rank_0 the siblings get, naming the version:

    orig = getattr(hf_modeling_utils, "get_total_byte_count", None)
    if orig is None:
        warn_rank_0(
            f"transformers {transformers.__version__} has no get_total_byte_count in "
            "modeling_utils to size the load warmup by; using its own sizing, which may reserve "
            "most of the device."
        )
        yield
        return

test_a_transformers_without_the_sizing_hook_is_left_alone would then gain a pytest.warns, matching the two sibling tests. Non-blocking.

[shengliangxu]

Taken in 2860c0c, with your wording. A release without get_total_byte_count now warns, naming the transformers version, as the other two fallbacks do. test_a_transformers_without_the_sizing_hook_is_left_alone expects the warning, and fails without it.

[shengliangxu]

/ok to test 00a592d

[shengliangxu]

/ok to test 2860c0c

[shengliangxu]

/claude review

[shengliangxu]

/ok to test 7ea4ac7