[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: main, where #3136 and #2780 have merged (2026-11-03). The diff below is this PR's alone. #2780 merged first, so this PR carries what the two need together: it lifts #2780's refusal of an expert-parallel plan with a sharded dense part, and the plan decision's give-way to whole-model FSDP2.

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-6060] [1/3] Add require_transformers, and build a VLM without its LM head #3131 [1/3], [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 [2/3] and [OMNIML-6060] Load an expert-parallel plan through load_pretrained, and apply the model spec when deciding and loading it #2780 have merged)
  • [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

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.

• Rebased 2026-11-01 with [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 4dd38a3 onto main cf38ec3, on shengliangx/hfep-base-6a-ep-loader rebuilt there from [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779 c42cb61, [OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067 cb83366 and [OMNIML-6071] Calibrate a distributed model on per-rank data #2740 ff372fb; 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.

• Rebased 2026-11-01 with [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 35f1ac0 onto main 51645db, where [OMNIML-6054] [1/2] Move fsdp2_wrap out of distributed.py into fsdp2_utils #3140 merged, on shengliangx/hfep-base-6a-ep-loader rebuilt there from [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779 1aee07f, [OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067 495f0be and [OMNIML-6071] Calibrate a distributed model on per-rank data #2740 6a06a54; this PR's own diff is unchanged. The tree is identical to the previous head, as [OMNIML-6054] [1/2] Move fsdp2_wrap out of distributed.py into fsdp2_utils #3140 merged as its own content. Pre-commit, mypy included, passes against main.

• Rebased 2026-11-02 with [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 39f55f2 onto main 2d4e8d0, on shengliangx/hfep-base-6a-ep-loader rebuilt there from [OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067 e6e422d and [OMNIML-6071] Calibrate a distributed model on per-rank data #2740 5829a85 ([OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779 is now part of main); 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.

• Rebased 2026-11-02 with [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 c0041d0 onto main 4f08a12, on shengliangx/hfep-base-6a-ep-loader rebuilt there from [OMNIML-6071] Calibrate a distributed model on per-rank data #2740 5829a85 ([OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067 is now part of main); 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.

• Rebased 2026-11-02 with [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 d87f5bd onto main 653ef61, where [OMNIML-6071] Calibrate a distributed model on per-rank data #2740 merged; 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.

• Rebased 2026-11-03 onto main a6cc167, where [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 merged; this PR's own diff is unchanged. Pre-commit, mypy included, passes against main.

• Review fix, 2026-11-03 (kevalmorabia97's review): test_a_tied_head_the_checkpoint_stores_alone_reaches_both_aliases writes copies of the tensors it read with load_file before save_file overwrites their file (ab58a1a). With safetensors 0.8.0-rc.0 on transformers 5.9.0 the test failed before the change, every loaded tensor reading as zeros, and passes after it. Test only. Pre-commit, mypy included, passes against main, and tests/unit/torch/utils and the HF plugin tests pass on CPU at this head, on transformers 5.18, 5.9.0 and 5.5.4.

• Rebased 2026-11-03 onto main 7d44246, where [OMNIML-6060] Load an expert-parallel plan through load_pretrained, and apply the model spec when deciding and loading it #2780 merged, with the junction the hf_ptq helper branch carried until now. In this PR's first commit, the plan decision keeps [OMNIML-6060] Load an expert-parallel plan through load_pretrained, and apply the model spec when deciding and loading it #2780's refusal warning and drops its give-way to whole-model FSDP2, for this PR's dense-part rule. A commit on top (9768e7a) lifts _route's refusal of an expert-parallel plan with dense="fsdp", says in load_pretrained's docstring and the distributed-loading guide that the dense part is replicated or FSDP2-sharded as the plan says, updates [OMNIML-6060] Load an expert-parallel plan through load_pretrained, and apply the model spec when deciding and loading it #2780's plan tests to match (an MoE left to decide keeps its EP degree and FSDP2-shards its dense part; an expert-parallel plan with dense="fsdp" loads through load_and_prepare_ep), and gives the dense-sharding tests' config doubles a model type. The rest of this PR's diff is unchanged. Pre-commit, mypy included, passes against main, and tests/unit/torch/utils, tests/unit/torch/export and the HF plugin tests pass on CPU at this head, on transformers 5.18, 5.9.0 and 5.5.4, apart from 7 diffusers-export cases that fail on main itself in these environments.

• Rebased 2026-11-03 onto main 5870698, where [OMNIML-6063] [2/2] Decompress compressed-tensors MoE experts as they load, expert-parallel loads included #3137 merged. In the GPU test file, [OMNIML-6063] [2/2] Decompress compressed-tensors MoE experts as they load, expert-parallel loads included #3137's compressed-experts tests and this PR's dense-sharded cases sit side by side. This PR's diff is otherwise unchanged. Pre-commit, mypy included, passes against main, the GPU file collects its 14 tests, and tests/unit/torch/utils, tests/unit/torch/export, tests/unit/torch/models and the HF plugin tests pass on CPU at this head, on transformers 5.18 and 5.5.4, apart from 7 diffusers-export cases that fail on main itself in these environments.

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 of expert-parallel PTQ for Hugging Face MoE models.

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: c66e2ff

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.

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

CodeRabbit Verify

• Run CodeRabbit Verify to test the changes in this pull request.

ℹ️ 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.

[codecov]

Codecov Report

❌ Patch coverage is 95.95960% with 8 lines in your changes missing coverage. Please review.
✅ Project coverage is 80.32%. Comparing base (5870698) to head (c66e2ff).

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

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2742      +/-   ##
==========================================
+ Coverage   72.12%   80.32%   +8.19%     
==========================================
  Files         677      677              
  Lines       78261    78443     +182     
==========================================
+ Hits        56447    63010    +6563     
+ Misses      21814    15433    -6381     

Flag	Coverage Δ
examples-diffusers	19.61% <0.00%> (-0.20%)	⬇️
examples-gpt-oss	13.30% <0.00%> (-0.17%)	⬇️
examples-hf_ptq	23.43% <15.15%> (-0.06%)	⬇️
examples-llm_distill	13.36% <0.00%> (-0.15%)	⬇️
examples-llm_eval	17.28% <15.15%> (-0.09%)	⬇️
examples-llm_qat	17.18% <0.00%> (-0.22%)	⬇️
examples-llm_sparsity	16.64% <15.15%> (-0.09%)	⬇️
examples-megatron_bridge	26.02% <0.00%> (-0.61%)	⬇️
examples-specdec_bench	13.08% <0.00%> (-0.15%)	⬇️
examples-speculative_decoding	17.52% <15.15%> (-0.04%)	⬇️
examples-torch_onnx	20.86% <0.00%> (-0.26%)	⬇️
examples-torch_trt	15.02% <0.00%> (-0.17%)	⬇️
examples-vllm_serve	17.39% <15.15%> (-0.10%)	⬇️
gpu	58.36% <15.15%> (+36.52%)	⬆️
regression	14.80% <0.00%> (-0.17%)	⬇️
unit	63.18% <95.95%> (+0.08%)	⬆️

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
Preview removed because the pull request was closed.
2026-11-03 18:32 UTC

[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 ab58a1a

[shengliangxu]

@kevalmorabia97 Fixed in ab58a1a. test_a_tied_head_the_checkpoint_stores_alone_reaches_both_aliases now writes copies of the tensors it read ({k: v.clone() for k, v in load_file(weights).items()}), so save_file no longer overwrites the file they map. With safetensors 0.8.0-rc.0 on transformers 5.9.0, the test fails before the change and passes after it. No other test in #2742 or #2743 rewrites a file it read with load_file.

[kevalmorabia97]

LGTM. The tied-head test now clones the tensors before rewriting the file; no other changes since my last review. Still open for the author: the 2-GPU dense backward/fp8 QAT cases fail on Ada with transformers 5.9 (layer-1 grads off 3-13%, replicated layout too) - please confirm whether that is Ada numerics or a transformers-5.9 issue.

[shengliangxu]

/ok to test 9768e7a

[shengliangxu]

/ok to test c66e2ff

[shengliangxu]

@kevalmorabia97 Thanks for this. We reproduced it on 2 RTX 6000 Ada GPUs, and it comes from routing, not from a wrong gradient: one token in layer 1 goes to a different pair of experts in the sharded model than in the whole model.

We ran tests/gpu/torch/utils/test_expert_parallel_load.py -k dense on 2 GPUs. The test uses a tiny Qwen3-MoE (8 experts, 2 layers), with the experts sharded over both GPUs and the dense part replicated or FSDP2-sharded. On Ada we used torch 2.11.0 (CUDA 13.0, NCCL 2.28.9); there, transformers runs the experts' grouped GEMM through torch.nn.functional.grouped_mm.

GPU	Code	transformers	Backward cases	Largest gradient error	Tokens routed differently (layer 0 / layer 1)
RTX 6000 Ada	this PR	5.9.0	2 failed	0.156	0 / 1 of 64
RTX 6000 Ada	this PR	5.18.0	2 failed	0.156	0 / 1 of 64
RTX 6000 Ada	main with this PR's test file	5.9.0 and 5.18.0	replicated failed	0.156	0 / 1 of 64
GB200	this PR	5.9.0 and 5.18.0	2 passed	0.003	0 / 0
GB200	main with this PR's test file	5.9.0	replicated passed	0.003	0 / 0
a second GPU type	this PR, and main with its test file	5.18.0	passed	0.003	0 / 0

Everything else passed wherever it could run: the forward cases, and the fp8 and nvfp4 QAT cases. The ep-half cases need 4 GPUs and skipped. On main, the dense-sharded cases fail by design, since main's loader refuses a sharded dense part.

So on Ada:

• The failure doesn't depend on transformers: 5.9.0 and 5.18.0 give the same numbers to four digits. Nor does it depend on this PR: main's replicated path fails the same way.
• Layer 0 routes identically, and its gradients stay within 0.009. In layer 1, one of 64 tokens takes a different pair of experts, and only layer 1's gradients fail: the experts by 9-14%, the router gate by 3-4% and the post-attention norm by 12-16%.

That supports the first explanation. Ada's bf16 GEMMs put layer 0's output a little further from the whole model's (0.009 against 0.003), which is enough to flip a near-tied route in layer 1. A token that takes other experts sends its gradient to other weights, and no tolerance should absorb that.

So I'd change neither the tolerance nor the transformers floor. Fixed the test in #3198: the tiny model gets a router that decides clearly, so bf16 noise cannot flip a route, and the test checks that both models route every token alike before it compares gradients, so a flip is reported as one.

On the same 2 RTX 6000 Ada GPUs, the fixed file passes on transformers 5.9.0 and 5.18.0 (10 passed, 4 skipped), with the largest gradient error at 0.011.

Your fp8 cases failed too. Ours passed on Ada, but only without ModelOpt's FP8 CUDA extension, which our CUDA 13.3 toolchain could not build against torch 2.11. If your run had the extension, its FP8 rounding could flip a route the same way.