From e124449e6e1fe14a0b8d87efa10ba6b6f239f6cd Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Sat, 10 Oct 2026 21:23:15 +0000
Subject: [PATCH 01/38] Export an expert-parallel model through the no-gather
 writer

An EP-sharded model could not be exported at all. transformers pre-slices each
expert parameter across the EP group, so save_pretrained sees a
tensor-parallel model and refuses outright. Routing now sends an
expert-parallel model to the no-gather writer unconditionally, since it has no
fallback; a model sharded by FSDP2 alone keeps its opt-in.

What expert parallelism changes in the writer and around it:

- The global expert index. Under FSDP2 alone a per-expert key's index is the
  DTensor's own dim-0 offset, but under EP transformers has already sliced the
  parameter to this rank's group, so that offset is group-relative and every
  group would write experts 0..n-1 over each other -- a silently wrong
  checkpoint. The coordinate now comes from the routed experts' ParallelState
  (expert_parallel_state), the same state calibration reads.
- Per-expert weight quantizers. The fold indexed them by the DTensor's global
  dim-0 offset, the FSDP2 convention where the quantizer list stays global;
  under EP the list holds only this rank's experts. What tells the two apart is
  the DTensor's global dim-0 extent -- equal to the list's length under FSDP2,
  larger under EP -- and an assert pins the local invariant. (Seen as an
  IndexError on every rank exporting GLM-5.3-Flash under EP16.)
- The rank. It was get_rank() only under FSDP2 and 0 otherwise, so under EP
  every process took the rank-0 branch and rewrote config.json at once, and
  the closing barrier was skipped, so the job hung instead of exiting.
- Carried weights. Under EP every other rank's experts read as unplaced on
  each rank, and every rank believed it was the single sender of the carried
  set; the writer's point-to-point protocol then deadlocked (GLM-5.3-Flash,
  EP16: 15 of 16 ranks parked in the barrier). Rank 0 alone writes them now.
- transformers stamps a DistributedConfig on the config when sharding; it is
  not JSON-serializable and describes this run's mesh, so it is dropped first.
- A rank whose export raises prints its traceback before the finally clause's
  barrier, where it would otherwise block on ranks that never fail and the
  job would hang with no error in the log.

_hf_shard_flags derives (fsdp2_sharded, ep_sharded) once for all of these
gates, kept separate because the FSDP2-dense + EP-expert hybrid is both at
once. _process_hf_model_for_export names the in-place processing half the
writer path runs: prepare, pack quantized weights with sharded experts left
fused, return the quant config.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/export/distributed_writer.py   |  73 +++++++++--
 modelopt/torch/export/unified_export_hf.py    | 121 +++++++++++++++---
 modelopt/torch/utils/distributed.py           |  12 ++
 .../torch/export/test_unified_export_hf.py    |  26 ++++
 tests/unit/torch/utils/test_distributed.py    |  28 ++++
 5 files changed, 227 insertions(+), 33 deletions(-)

diff --git a/modelopt/torch/export/distributed_writer.py b/modelopt/torch/export/distributed_writer.py
index 21ad6da9f34..51b9365d807 100644
--- a/modelopt/torch/export/distributed_writer.py
+++ b/modelopt/torch/export/distributed_writer.py
@@ -17,7 +17,8 @@
 
 Every tensor is sent to exactly one owner rank over batched point-to-point sends and
 written from there, so no rank ever materializes the whole model and no unit is
-replicated across ranks. Selected by ``MODELOPT_FSDP2_EXPORT=nogather``.
+replicated across ranks. Used for every expert-parallel model; a model sharded by FSDP2 alone opts in
+with ``MODELOPT_FSDP2_EXPORT=nogather``.
 """
 
 import json
@@ -690,7 +691,7 @@ def _distributed_save_hf_checkpoint_impl(
 
       1. Fused expert weights (when present): each rank splits its LOCAL shard into per-expert keys
          with global indices and keeps them in place -- no communication. Experts are therefore
-         already disjoint across ranks.
+         already disjoint across ranks. dp-replica ranks under DP x EP skip this.
       2. Dense / non-expert weights: each fqn is assigned to exactly ONE owner rank and moved there
          (:func:`_materialize_owned`), so every tensor becomes whole on its owner while the set stays
          spread across ranks. Replicated and plain values need no transport; dim-0 shards travel
@@ -716,6 +717,50 @@ def _distributed_save_hf_checkpoint_impl(
 
     export_dir = Path(export_dir)
     rank, world = dist.get_rank(), dist.get_world_size()
+
+    # Where this rank sits in expert parallelism, if any.
+    #
+    # Under FSDP2 alone the fused 3-D expert weight is sharded on dim 0, so each per-expert key's
+    # global index is just the DTensor's own dim-0 offset and the answer below stays (0, 0).
+    #
+    # Under expert parallelism it is not: transformers pre-slices each expert param to this rank's
+    # EP group, so the EP split is INVISIBLE to the per-param DTensor -- its dim-0 offset is relative
+    # to the group, not global. Without ep_rank every group would write experts 0..n-1 over each
+    # other and the checkpoint would be quietly wrong rather than broken. The coordinate comes from
+    # the routed experts' ParallelState -- the same source calibration reads, not transformers'
+    # internal ``model._device_mesh``. An expert's data-parallel group is its expert-data-parallel one.
+    from modelopt.torch.utils.distributed import expert_parallel_state
+
+    ep_rank, dp_idx = 0, 0
+    parallel_state = expert_parallel_state(model)
+    if parallel_state is not None:
+        ep_rank = parallel_state.expert_model_parallel_group.rank()
+        edp_group = parallel_state.data_parallel_group
+        if edp_group.is_initialized():
+            # Experts are replicated across EDP, so only the first replica writes them.
+            dp_idx = edp_group.rank()
+    # Total config expert count, to tell a per-EP-group expert tensor (size < total) from one that
+    # spans all experts (classic DP x EP) when computing the per-expert global index offset.
+    _cfg = getattr(model, "config", None)
+
+    def _experts_of(c):
+        if c is None:
+            return 0
+        return int(
+            getattr(c, "num_experts", 0)
+            or getattr(c, "num_local_experts", 0)
+            or getattr(c, "n_routed_experts", 0)  # nemotron_h
+            or 0
+        )
+
+    # VLMs nest the text-model fields under ``config.text_config`` (the quantized backbone is
+    # ``model.language_model``), so fall back to it when the top-level config has no expert count.
+    total_experts = _experts_of(_cfg) or _experts_of(getattr(_cfg, "text_config", None))
+    # DP x EP: experts are dp-REPLICATED across the EDP group (each holds a full EP-sharded copy). Only
+    # the first replica (dp_idx == 0) writes them; replica ranks (dp_idx > 0) skip the per-expert split.
+    # dp_idx is the EDP-group rank resolved above -- 0 when experts are not EDP-replicated (dense /
+    # FSDP2 / EP-without-DP), so every such rank writes.
+    write_experts = dp_idx == 0
     sharded_sd = get_model_state_dict(model, options=StateDictOptions(full_state_dict=False))
 
     # Re-materialize every inference tensor as a NORMAL one UP FRONT. A caller in inference mode hands
@@ -789,17 +834,21 @@ def _is_expert_key(key: str) -> bool:
         for j, k in enumerate(in_keys):
             sharded_sd[k] = stacked[j].clone()
     # (1) Experts: split this rank's local shard into per-expert keys (global idx); stays local.
-    # FSDP2 shards a fused 3-D expert weight on dim 0, the expert axis, so each per-expert key's global
-    # index is the DTensor's own dim-0 offset, which _split_local_fused_module reads.
+    # Skipped on dp-replica ranks under DP x EP (their experts are written by dp group 0).
     local_sd: dict = {}
-    for prefix in prefixes:
-        local_sd.update(_split_local_fused_module(prefix, sharded_sd))
-    local_sd = postprocess_state_dict(
-        local_sd, maxbound, kv_cache_format, is_modelopt_qlora, tied_map=tied_map
-    )
-    # In place: the split's tensors are its own copies, so one made contiguous replaces the old.
-    for _k in list(local_sd.keys()):
-        local_sd[_k] = local_sd[_k].detach().contiguous()
+    if write_experts:
+        for prefix in prefixes:
+            local_sd.update(
+                _split_local_fused_module(
+                    prefix, sharded_sd, ep_rank=ep_rank, total_experts=total_experts
+                )
+            )
+        local_sd = postprocess_state_dict(
+            local_sd, maxbound, kv_cache_format, is_modelopt_qlora, tied_map=tied_map
+        )
+        # In place: the split's tensors are its own copies, so one made contiguous replaces the old.
+        for _k in list(local_sd.keys()):
+            local_sd[_k] = local_sd[_k].detach().contiguous()
     # At this point local_sd holds ONLY this rank's (already split, per-expert) experts -- capture the
     # set so _materialize_owned can tell disjoint experts from replicated dense.
     expert_key_set = set(local_sd)
diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py
index 9c63130bfd1..a072e9dfa69 100644
--- a/modelopt/torch/export/unified_export_hf.py
+++ b/modelopt/torch/export/unified_export_hf.py
@@ -22,6 +22,7 @@
 import os
 import re
 import tempfile
+import traceback
 import warnings
 from builtins import ValueError
 from collections import defaultdict
@@ -1087,6 +1088,40 @@ def pack_quantized_weights(
     _reconstruct_fused_moe_linear(model)
 
 
+def _hf_shard_flags(model: nn.Module) -> tuple[bool, bool]:
+    """Return ``(fsdp2_sharded, ep_sharded)`` for an HF model about to be exported.
+
+    Expert parallelism shards the experts too, just not by FSDP2: transformers pre-slices the expert
+    params across the EP group. The two are reported separately because a model can be both at once,
+    FSDP2-dense with expert-parallel experts. Neither holds without an initialized process group.
+    """
+    from modelopt.torch.utils.distributed import expert_parallel_state
+
+    if not (torch.distributed.is_available() and torch.distributed.is_initialized()):
+        return False, False
+    fsdp2_sharded = is_fsdp2_model(model)
+    ep_sharded = torch.distributed.get_world_size() > 1 and expert_parallel_state(model) is not None
+    return fsdp2_sharded, ep_sharded
+
+
+def _process_hf_model_for_export(
+    model: nn.Module, dtype: torch.dtype | None = None, is_modelopt_qlora: bool = False
+) -> dict[str, Any]:
+    """Fold a quantized HF model's quantizers into its weights in place; return its quant config.
+
+    The processing half of an export, for a caller that serializes the model itself: the no-gather
+    writer, which writes each rank's shard, and the speculative-decoding exporters, which keep only
+    the draft. Afterwards the model's own state dict holds the deployment weights. Nothing is
+    gathered, so it runs on a live FSDP2 or expert-parallel model; there the fused MoE experts stay
+    fused, since each rank folds only its local shard and the writer splits them after
+    consolidation.
+    """
+    dtype, _, quant_config = _prepare_model_for_export(model, dtype, is_modelopt_qlora)
+    keep_fused_experts = any(_hf_shard_flags(model))
+    pack_quantized_weights(model, dtype, is_modelopt_qlora, keep_fused_experts=keep_fused_experts)
+    return quant_config
+
+
 def _export_transformers_checkpoint(
     model: nn.Module,
     dtype: torch.dtype | None = None,
@@ -1933,6 +1968,21 @@ def read_unplaced_weights(model: nn.Module, *, keys_only: bool = False) -> dict[
     return out
 
 
+def _write_base_config(model: nn.Module, export_dir: Path) -> None:
+    """Write the no-gather export's non-weight files, once every rank has written its shards.
+
+    ``save_non_weight_artifacts`` writes them, as it does for the FSDP2 streaming writer; this
+    keeps the config attributes an expert-parallel run sets out of config.json.
+    """
+    # transformers stamps a DistributedConfig on the config when it shards the model. It is not
+    # JSON-serializable, and it describes THIS RUN's mesh -- the EP degree and tp plan the export
+    # happened to use -- not a property of the checkpoint, so baking it in would mislead whoever
+    # loads it next even if it could be written.
+    if getattr(model.config, "distributed_config", None) is not None:
+        model.config.distributed_config = None
+    save_non_weight_artifacts(model, export_dir)
+
+
 def _no_gather_export_asked() -> bool:
     """Whether ``MODELOPT_FSDP2_EXPORT`` picks the no-gather writer for an FSDP2-sharded export.
 
@@ -1995,10 +2045,20 @@ def export_hf_checkpoint(
     # emits it from rank 0 alone, so on the other ranks a full read is host memory spent and thrown
     # away -- and the carried set is the large stuff. The KEY list is still resolved everywhere,
     # because get_quant_config runs per rank and the configs have to agree.
-    # How the model is distributed decides how it is exported. A diffusers pipeline is not a module,
-    # and nothing here distributes one.
-    plan = distribution_plan(model) if isinstance(model, nn.Module) else DistributionPlan()
-    _writes_extra = not plan.is_fsdp_dense or torch.distributed.get_rank() == 0
+    # Expert parallelism shards the experts too, just not by FSDP2 -- transformers pre-slices the
+    # expert params across the EP group. Computed here rather than with the other predicates below
+    # because the carry-over gate needs it and runs ahead of them (the layerwise and diffusers paths
+    # return before those are reached).
+    _fsdp2_sharded, _ep_sharded = _hf_shard_flags(model)
+    # Either form of sharding, including the hybrid where both hold: a model can be FSDP2-dense AND
+    # expert-parallel, so these are tracked separately rather than as an either/or.
+    _sharded_export = _ep_sharded or _fsdp2_sharded
+    # A sharded export writes the carried weights from rank 0 alone. This has to cover EP as well as
+    # FSDP2: under EP every OTHER rank's experts are absent from this rank's model, so they read as
+    # unplaced and each rank would carry them -- 889 tensors per rank on GLM-5.3-Flash. Every rank
+    # then hands them to the writer, whose point-to-point protocol expects one sender, and the export
+    # deadlocks with rank 0 inside the writer and the rest parked in the barrier below.
+    _writes_extra = not _sharded_export or torch.distributed.get_rank() == 0
     _carried = read_unplaced_weights(model, keys_only=not _writes_extra)
     if _writes_extra and _carried:
         extra_state_dict = {**_carried, **(extra_state_dict or {})}
@@ -2044,14 +2104,27 @@ def export_hf_checkpoint(
         )
         return
 
-    is_fsdp2_sharded = plan.is_fsdp_dense
-    # Rank 0's choice: the two writers issue different collectives, so ranks that read their
-    # environments differently would wait in collectives the others never join.
-    use_no_gather_writer = is_fsdp2_sharded and broadcast_from(
-        "MODELOPT_FSDP2_EXPORT", _no_gather_export_asked
+    is_fsdp2_sharded = _fsdp2_sharded
+
+    # An expert-parallel model is sharded too, just not by FSDP2: transformers pre-slices the expert
+    # params across the EP group, so save_pretrained sees a tensor-parallel model and refuses it
+    # outright. There is no gather fallback for that shape, so EP always takes the no-gather writer,
+    # FSDP2-sharded dense included (the hybrid is both at once). A model sharded by FSDP2 alone keeps
+    # the existing path unless MODELOPT_FSDP2_EXPORT=nogather opts it in, as rank 0 reads it: the
+    # two writers issue different collectives, so ranks that read their environments differently
+    # would wait in collectives the others never join. Derived from _ep_sharded above so the
+    # carry-over gate and this cannot disagree about what EP is.
+    is_ep_sharded = _ep_sharded
+    use_no_gather_writer = is_ep_sharded or (
+        is_fsdp2_sharded and broadcast_from("MODELOPT_FSDP2_EXPORT", _no_gather_export_asked)
     )
-    # Not the global rank: a non-FSDP2 export writes the whole checkpoint from every process.
-    rank = torch.distributed.get_rank() if is_fsdp2_sharded else 0
+
+    # The rank used to gate the single-writer steps below. It must be the real rank whenever the
+    # checkpoint is written cooperatively -- under expert parallelism every process would otherwise
+    # read 0 here, all of them would rewrite config.json at once, and one would truncate the file
+    # another was reading. A non-sharded export does write the whole checkpoint from every process,
+    # which is what the 0 is for.
+    rank = torch.distributed.get_rank() if (is_fsdp2_sharded or is_ep_sharded) else 0
     # Offloaded models take the streaming path: it materializes one layer at a time and
     # writes each straight to a shard file, so peak memory is one layer plus one shard
     # buffer instead of the whole quantized state dict.
@@ -2089,12 +2162,9 @@ def export_hf_checkpoint(
             from .distributed_writer import distributed_save_hf_checkpoint
 
             _qlora = kwargs.get("is_modelopt_qlora", False)
-            # The writer builds its own name-based tied-weight map, so this one is not handed down.
-            _dtype, _, hf_quant_config = _prepare_model_for_export(model, dtype, _qlora)
-            # Preparing the model does not pack its weights, so pack them here. keep_fused_experts
-            # leaves the MoE experts fused and sharded through the fold, so each rank folds only
-            # its local experts; the writer then splits each rank's experts into per-expert tensors.
-            pack_quantized_weights(model, _dtype, _qlora, keep_fused_experts=True)
+            # Folds in place; the model is sharded here, so the MoE experts stay fused and each rank
+            # folds only its LOCAL experts. The writer splits them after consolidation.
+            hf_quant_config = _process_hf_model_for_export(model, dtype, _qlora)
             if getattr(model, "hf_quantizer", None) is not None:
                 model.hf_quantizer = None
             reverted = distributed_save_hf_checkpoint(
@@ -2109,7 +2179,7 @@ def export_hf_checkpoint(
                 extra_state_dict=extra_state_dict,
             )
             if rank == 0:
-                save_non_weight_artifacts(model, export_dir)
+                _write_base_config(model, export_dir)
                 # As the streaming path below: every FSDP2 rank holds the same ModelOpt state.
                 if save_modelopt_state and ModeloptStateManager.is_converted(model):
                     torch.save(modelopt_state(model), export_dir / _MODELOPT_STATE_SAVE_NAME)
@@ -2206,16 +2276,25 @@ def export_hf_checkpoint(
         if rank == 0:
             _write_hf_export_config(model, hf_quant_config, export_dir)
             copy_off_index_safetensors(_source_checkpoint(model), export_dir)
-            if not (_offloaded or is_fsdp2_sharded):
-                # The streaming paths already did, from save_non_weight_artifacts.
+            if not (_offloaded or is_fsdp2_sharded or use_no_gather_writer):
+                # The streaming paths and the no-gather writer already did, from
+                # save_non_weight_artifacts.
                 _copy_non_model_files_from_source(model, export_dir)
 
     except Exception as e:
+        # Print (not just warnings.warn, which the default filter can dedupe/suppress) BEFORE the
+        # barrier below: under FSDP2/EP every other rank is still waiting there, so a failing rank
+        # that raises straight into `finally` blocks on that barrier forever without this rank's
+        # traceback ever reaching the log -- the job just hangs with no visible error.
+        _rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0
+        print(f"[export_hf_checkpoint] rank {_rank} failed:\n{traceback.format_exc()}", flush=True)
         warnings.warn(
             "Cannot export model to the model_config. The modelopt-optimized model state_dict"
             " can be saved with torch.save for further inspection."
         )
         raise e
     finally:
-        if is_fsdp2_sharded:
+        # Every rank waits for the writer to finish. Without this the ranks that wrote nothing exit
+        # first and the job hangs on the ones still writing.
+        if is_fsdp2_sharded or is_ep_sharded:
             torch.distributed.barrier()
diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py
index 2a82598dbb4..facaf25fe44 100644
--- a/modelopt/torch/utils/distributed.py
+++ b/modelopt/torch/utils/distributed.py
@@ -57,6 +57,7 @@
     "distributed_forward_loop",
     "distributed_run",
     "distribution_plan",
+    "expert_parallel_state",
     "fsdp2_wrap",
     "gather_bytes",
     "gather_state_dicts",
@@ -690,6 +691,17 @@ def distributed_run(
             _destroy_group(exc)
 
 
+def expert_parallel_state(model: torch.nn.Module) -> "ParallelState | None":
+    """Return the ``ParallelState`` of the model's routed experts, or ``None`` if they are not expert-parallel.
+
+    The expert-parallel loader seeds every routed-expert module with a ``ParallelState`` that carries the
+    expert-model-parallel group (``load_and_prepare_ep``); no other module of a HuggingFace model has one.
+    """
+    for module in model.modules():
+        state = getattr(module, "_parallel_state", None)
+        if state is not None and state.expert_model_parallel_group.is_initialized():
+            return state
+    return None
 class _StridedIterableShard(IterableDataset):
     """Rank-strided shard of a streaming ``IterableDataset`` (no ``len()`` required).
 
diff --git a/tests/unit/torch/export/test_unified_export_hf.py b/tests/unit/torch/export/test_unified_export_hf.py
index ef113dbcc4d..cd79f26bb4a 100644
--- a/tests/unit/torch/export/test_unified_export_hf.py
+++ b/tests/unit/torch/export/test_unified_export_hf.py
@@ -1107,3 +1107,29 @@ def test_source_checkpoint_is_none_for_a_model_built_from_config(tmp_path, monke
 
     assert _source_checkpoint(model) is None
     assert not any(export_dir.iterdir())
+
+
+def test_process_hf_model_for_export_folds_in_place_and_returns_the_quant_config():
+    """The processing half the speculative-decoding exporters call before keeping only the draft."""
+    from transformers import LlamaConfig, LlamaForCausalLM
+
+    from modelopt.torch.export.unified_export_hf import _process_hf_model_for_export
+
+    config = LlamaConfig(
+        vocab_size=64,
+        hidden_size=32,
+        intermediate_size=64,
+        num_hidden_layers=1,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+    )
+    model = LlamaForCausalLM(config).eval()
+    mtq.quantize(model, mtq.FP8_DEFAULT_CFG, lambda m: m(torch.randint(0, 64, (1, 8))))
+
+    quant_config = _process_hf_model_for_export(model, torch.float16)
+
+    assert quant_config["quantization"]["quant_algo"] == "FP8"
+    # Folded in place: the model's own state dict is now the deployment one.
+    q_proj = model.model.layers[0].self_attn.q_proj
+    assert q_proj.weight.dtype == torch.float8_e4m3fn
+    assert "model.layers.0.self_attn.q_proj.weight_scale" in model.state_dict()
diff --git a/tests/unit/torch/utils/test_distributed.py b/tests/unit/torch/utils/test_distributed.py
index f292843c89a..54811fb924f 100644
--- a/tests/unit/torch/utils/test_distributed.py
+++ b/tests/unit/torch/utils/test_distributed.py
@@ -22,6 +22,7 @@
 from contextlib import nullcontext
 from datetime import timedelta
 from functools import partial
+from types import SimpleNamespace
 from unittest import mock
 
 import pytest
@@ -43,6 +44,7 @@
     broadcast_from,
     distributed_run,
     distribution_plan,
+    expert_parallel_state,
     set_distribution_plan,
 )
 
@@ -548,3 +550,29 @@ def test_the_run_leaves_a_group_it_found_alone(monkeypatch, recwarn, world_size,
         if raises:
             raise ValueError("calibration failed")
     assert "timeout=1:00:00) does not apply" in str(recwarn.pop(UserWarning).message)
+
+
+# ---------------------------------------------------------------------------- expert parallel state
+
+
+def _state(expert_parallel: bool):
+    return SimpleNamespace(
+        expert_model_parallel_group=SimpleNamespace(is_initialized=lambda: expert_parallel)
+    )
+
+
+def test_the_routed_experts_state_is_found_by_its_expert_parallel_group():
+    model = torch.nn.Module()
+    model.attn = torch.nn.Linear(2, 2)
+    model.experts = torch.nn.Module()
+    for module in model.modules():
+        module._parallel_state = _state(expert_parallel=False)
+    model.experts._parallel_state = expert = _state(expert_parallel=True)
+    assert expert_parallel_state(model) is expert
+
+
+def test_a_model_without_expert_parallel_experts_has_no_expert_state():
+    model = torch.nn.Sequential(torch.nn.Linear(2, 2))
+    assert expert_parallel_state(model) is None
+    model[0]._parallel_state = _state(expert_parallel=False)
+    assert expert_parallel_state(model) is None

From 3d3a9eb410f9c3063c00c292abce070fbc2dc74b Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Sun, 11 Oct 2026 02:07:56 +0000
Subject: [PATCH 02/38] Keep the EP run's DistributedConfig and torch_a2a plan
 out of config.json

A distributed export wrote a config.json that differed from a single-GPU
export of the same checkpoint in two keys:

- base_model_ep_plan, which the EP loader swaps to the torch_a2a styles.
  Those styles exist only while ModelOpt is imported, so loading the export
  expert-parallel anywhere else fails on an unknown style. transformers drops
  the tp and pp plans when it serializes a config, but not the EP plan.
- distributed_config: null. The writer nulled the DistributedConfig that
  transformers stamps on the config (it is not JSON-serializable), and the
  null was written out. Nulling it in place also left transformers' tp_plan
  property raising on the live model, since it reads
  distributed_config.enable_expert_parallel whenever the attribute exists.

While config.json is written, each config -- the root and every nested one,
where a VLM keeps its plan -- now holds the value its source config.json
declares, or loses the attribute when the source has none. The run's values
come back afterwards.

Found comparing a Qwen3-30B-A3B export from the EP path with a single-GPU
export of the same checkpoint.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/export/unified_export_hf.py    | 64 +++++++++++++++++--
 .../torch/export/test_unified_export_hf.py    | 36 +++++++++++
 2 files changed, 93 insertions(+), 7 deletions(-)

diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py
index a072e9dfa69..a2b800d2096 100644
--- a/modelopt/torch/export/unified_export_hf.py
+++ b/modelopt/torch/export/unified_export_hf.py
@@ -1968,19 +1968,69 @@ def read_unplaced_weights(model: nn.Module, *, keys_only: bool = False) -> dict[
     return out
 
 
+# Config attributes a distributed load sets to describe THIS RUN rather than the checkpoint: the
+# DistributedConfig transformers stamps on the config when it shards the model (the EP degree and tp
+# plan the export happened to use; not JSON-serializable either), and the base_model_ep_plan the EP
+# loader swapped to the torch_a2a styles, which exist only while ModelOpt is imported.
+_DISTRIBUTED_RUN_CONFIG_ATTRS = ("distributed_config", "base_model_ep_plan")
+
+
+@contextlib.contextmanager
+def _source_distributed_config_attrs(model: nn.Module):
+    """Within the block, every config holds the run attributes its source ``config.json`` declared.
+
+    transformers drops ``base_model_tp_plan`` and ``base_model_pp_plan`` when it serializes a config,
+    but not ``base_model_ep_plan``, so without this the export ships a plan whose styles nothing but
+    ModelOpt can load. Each config -- the root and every nested one, e.g. a VLM's ``text_config``,
+    where the plan lives -- gets the value its source ``config.json`` holds at the same place, or
+    loses the attribute when the source has none: the config a single-GPU export writes. A source
+    that cannot be read counts as having none. The run's values come back on exit, since the
+    sharded model still reads them (transformers' ``tp_plan`` checks ``distributed_config``).
+    """
+    from transformers import PretrainedConfig
+
+    source: Any = {}
+    ckpt = _source_checkpoint(model)
+    if ckpt and Path(ckpt, "config.json").is_file():
+        with contextlib.suppress(OSError, ValueError):
+            source = json.loads(Path(ckpt, "config.json").read_text())
+
+    _missing = object()
+    saved: list[tuple[Any, str, Any]] = []
+
+    def use_source(config: Any, source: Any) -> None:
+        source = source if isinstance(source, dict) else {}
+        for attr in _DISTRIBUTED_RUN_CONFIG_ATTRS:
+            if attr in vars(config) or attr in source:
+                saved.append((config, attr, vars(config).get(attr, _missing)))
+                if attr in source:
+                    setattr(config, attr, source[attr])
+                elif attr in vars(config):
+                    delattr(config, attr)
+        for key, sub_config in list(vars(config).items()):
+            if isinstance(sub_config, PretrainedConfig):
+                use_source(sub_config, source.get(key))
+
+    try:
+        use_source(model.config, source)
+        yield
+    finally:
+        for config, attr, value in reversed(saved):
+            if value is _missing:
+                if attr in vars(config):
+                    delattr(config, attr)
+            else:
+                setattr(config, attr, value)
+
+
 def _write_base_config(model: nn.Module, export_dir: Path) -> None:
     """Write the no-gather export's non-weight files, once every rank has written its shards.
 
     ``save_non_weight_artifacts`` writes them, as it does for the FSDP2 streaming writer; this
     keeps the config attributes an expert-parallel run sets out of config.json.
     """
-    # transformers stamps a DistributedConfig on the config when it shards the model. It is not
-    # JSON-serializable, and it describes THIS RUN's mesh -- the EP degree and tp plan the export
-    # happened to use -- not a property of the checkpoint, so baking it in would mislead whoever
-    # loads it next even if it could be written.
-    if getattr(model.config, "distributed_config", None) is not None:
-        model.config.distributed_config = None
-    save_non_weight_artifacts(model, export_dir)
+    with _source_distributed_config_attrs(model):
+        save_non_weight_artifacts(model, export_dir)
 
 
 def _no_gather_export_asked() -> bool:
diff --git a/tests/unit/torch/export/test_unified_export_hf.py b/tests/unit/torch/export/test_unified_export_hf.py
index cd79f26bb4a..a26112b1241 100644
--- a/tests/unit/torch/export/test_unified_export_hf.py
+++ b/tests/unit/torch/export/test_unified_export_hf.py
@@ -1133,3 +1133,39 @@ def test_process_hf_model_for_export_folds_in_place_and_returns_the_quant_config
     q_proj = model.model.layers[0].self_attn.q_proj
     assert q_proj.weight.dtype == torch.float8_e4m3fn
     assert "model.layers.0.self_attn.q_proj.weight_scale" in model.state_dict()
+
+
+@pytest.mark.parametrize("source_plan", [None, {"layers.*.feed_forward.experts": "grouped_gemm"}])
+def test_distributed_export_writes_the_config_a_single_gpu_export_writes(tmp_path, source_plan):
+    """The EP load's DistributedConfig and torch_a2a plan stay out of config.json, the source's own
+    plan survives, and the live configs keep the run's values afterwards."""
+    from transformers import AutoConfig, Llama4Config
+    from transformers.distributed import DistributedConfig
+
+    from modelopt.torch.export.unified_export_hf import _write_base_config
+
+    source = tmp_path / "source"
+    Llama4Config().save_pretrained(source)
+    if source_plan is not None:
+        config_json = json.loads((source / "config.json").read_text())
+        config_json["text_config"]["base_model_ep_plan"] = source_plan
+        (source / "config.json").write_text(json.dumps(config_json))
+    AutoConfig.from_pretrained(source).save_pretrained(tmp_path / "single_gpu")
+
+    # What the EP load leaves behind: the stamped DistributedConfig and the swapped plan.
+    config = AutoConfig.from_pretrained(source)
+    distributed_config = DistributedConfig(enable_expert_parallel=True)
+    swapped_plan = {"layers.*.feed_forward.experts": "torch_a2a_experts"}
+    config.distributed_config = distributed_config
+    config.text_config.base_model_ep_plan = swapped_plan
+    model = SimpleNamespace(
+        config=config, _modelopt_source_checkpoint=str(source), can_generate=lambda: False
+    )
+
+    _write_base_config(model, tmp_path / "export")
+
+    written = json.loads((tmp_path / "export" / "config.json").read_text())
+    assert written == json.loads((tmp_path / "single_gpu" / "config.json").read_text())
+    assert written["text_config"].get("base_model_ep_plan") == source_plan
+    assert config.distributed_config is distributed_config
+    assert config.text_config.base_model_ep_plan is swapped_plan

From 9b9f580ee34ba3ee6dbe37a35943933bee75b284 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Sun, 11 Oct 2026 03:17:09 +0000
Subject: [PATCH 03/38] Describe _process_hf_model_for_export's one caller

Its docstring and test named the speculative-decoding exporters as a second
caller; in this stack the no-gather writer is the only one.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/export/unified_export_hf.py        | 9 ++++-----
 tests/unit/torch/export/test_unified_export_hf.py | 2 +-
 2 files changed, 5 insertions(+), 6 deletions(-)

diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py
index a2b800d2096..08e6ade6332 100644
--- a/modelopt/torch/export/unified_export_hf.py
+++ b/modelopt/torch/export/unified_export_hf.py
@@ -1110,11 +1110,10 @@ def _process_hf_model_for_export(
     """Fold a quantized HF model's quantizers into its weights in place; return its quant config.
 
     The processing half of an export, for a caller that serializes the model itself: the no-gather
-    writer, which writes each rank's shard, and the speculative-decoding exporters, which keep only
-    the draft. Afterwards the model's own state dict holds the deployment weights. Nothing is
-    gathered, so it runs on a live FSDP2 or expert-parallel model; there the fused MoE experts stay
-    fused, since each rank folds only its local shard and the writer splits them after
-    consolidation.
+    writer, which writes each rank's shard. Afterwards the model's own state dict holds the
+    deployment weights. Nothing is gathered, so it runs on a live FSDP2 or expert-parallel model;
+    there the fused MoE experts stay fused, since each rank folds only its local shard and the
+    writer splits them after consolidation.
     """
     dtype, _, quant_config = _prepare_model_for_export(model, dtype, is_modelopt_qlora)
     keep_fused_experts = any(_hf_shard_flags(model))
diff --git a/tests/unit/torch/export/test_unified_export_hf.py b/tests/unit/torch/export/test_unified_export_hf.py
index a26112b1241..4a0cb931018 100644
--- a/tests/unit/torch/export/test_unified_export_hf.py
+++ b/tests/unit/torch/export/test_unified_export_hf.py
@@ -1110,7 +1110,7 @@ def test_source_checkpoint_is_none_for_a_model_built_from_config(tmp_path, monke
 
 
 def test_process_hf_model_for_export_folds_in_place_and_returns_the_quant_config():
-    """The processing half the speculative-decoding exporters call before keeping only the draft."""
+    """The processing half the no-gather writer runs before writing each rank's shard."""
     from transformers import LlamaConfig, LlamaForCausalLM
 
     from modelopt.torch.export.unified_export_hf import _process_hf_model_for_export

From 1e2dc893de1d2ba6218e4d922989f161170af281 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Sun, 11 Oct 2026 04:42:29 +0000
Subject: [PATCH 04/38] Keep the load-time checkpoint renames out of
 config.json too

Kimi-K3's expert-parallel prepare hook installs _checkpoint_conversion_mapping
on the root and text configs, renaming the source's per-expert w1/w3/w2 so
transformers can fuse them on load. Those regexes name the source's tensors,
not the export's, and the distributed writer serialized them into config.json
at both levels: the 16-node Kimi-K3 export carried them while its
base_model_ep_plan and distributed_config were already gone. Restore them
with the other run attributes, so a reader of the export never sees renames
meant for a different checkpoint.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/export/unified_export_hf.py        | 12 +++++++++---
 tests/unit/torch/export/test_unified_export_hf.py | 12 +++++++++---
 2 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py
index 08e6ade6332..1fee18347e7 100644
--- a/modelopt/torch/export/unified_export_hf.py
+++ b/modelopt/torch/export/unified_export_hf.py
@@ -1969,9 +1969,15 @@ def read_unplaced_weights(model: nn.Module, *, keys_only: bool = False) -> dict[
 
 # Config attributes a distributed load sets to describe THIS RUN rather than the checkpoint: the
 # DistributedConfig transformers stamps on the config when it shards the model (the EP degree and tp
-# plan the export happened to use; not JSON-serializable either), and the base_model_ep_plan the EP
-# loader swapped to the torch_a2a styles, which exist only while ModelOpt is imported.
-_DISTRIBUTED_RUN_CONFIG_ATTRS = ("distributed_config", "base_model_ep_plan")
+# plan the export happened to use; not JSON-serializable either), the base_model_ep_plan the EP
+# loader swapped to the torch_a2a styles, which exist only while ModelOpt is imported, and the
+# _checkpoint_conversion_mapping a model's expert-parallel prepare hook installs to rename the
+# source's tensors on load (Kimi-K3's per-expert w1/w3/w2), which names nothing in the export.
+_DISTRIBUTED_RUN_CONFIG_ATTRS = (
+    "distributed_config",
+    "base_model_ep_plan",
+    "_checkpoint_conversion_mapping",
+)
 
 
 @contextlib.contextmanager
diff --git a/tests/unit/torch/export/test_unified_export_hf.py b/tests/unit/torch/export/test_unified_export_hf.py
index 4a0cb931018..d55e1a6775d 100644
--- a/tests/unit/torch/export/test_unified_export_hf.py
+++ b/tests/unit/torch/export/test_unified_export_hf.py
@@ -1137,8 +1137,8 @@ def test_process_hf_model_for_export_folds_in_place_and_returns_the_quant_config
 
 @pytest.mark.parametrize("source_plan", [None, {"layers.*.feed_forward.experts": "grouped_gemm"}])
 def test_distributed_export_writes_the_config_a_single_gpu_export_writes(tmp_path, source_plan):
-    """The EP load's DistributedConfig and torch_a2a plan stay out of config.json, the source's own
-    plan survives, and the live configs keep the run's values afterwards."""
+    """The EP load's DistributedConfig, torch_a2a plan and load-time renames stay out of config.json,
+    the source's own plan survives, and the live configs keep the run's values afterwards."""
     from transformers import AutoConfig, Llama4Config
     from transformers.distributed import DistributedConfig
 
@@ -1152,12 +1152,16 @@ def test_distributed_export_writes_the_config_a_single_gpu_export_writes(tmp_pat
         (source / "config.json").write_text(json.dumps(config_json))
     AutoConfig.from_pretrained(source).save_pretrained(tmp_path / "single_gpu")
 
-    # What the EP load leaves behind: the stamped DistributedConfig and the swapped plan.
+    # What the EP load leaves behind: the stamped DistributedConfig, the swapped plan, and the
+    # renames a prepare hook installs on both configs (as Kimi-K3's does).
     config = AutoConfig.from_pretrained(source)
     distributed_config = DistributedConfig(enable_expert_parallel=True)
     swapped_plan = {"layers.*.feed_forward.experts": "torch_a2a_experts"}
+    renames = {r"experts\.(\d+)\.w1\.": r"experts.\1.gate_proj."}
     config.distributed_config = distributed_config
     config.text_config.base_model_ep_plan = swapped_plan
+    config._checkpoint_conversion_mapping = renames
+    config.text_config._checkpoint_conversion_mapping = renames
     model = SimpleNamespace(
         config=config, _modelopt_source_checkpoint=str(source), can_generate=lambda: False
     )
@@ -1169,3 +1173,5 @@ def test_distributed_export_writes_the_config_a_single_gpu_export_writes(tmp_pat
     assert written["text_config"].get("base_model_ep_plan") == source_plan
     assert config.distributed_config is distributed_config
     assert config.text_config.base_model_ep_plan is swapped_plan
+    assert config._checkpoint_conversion_mapping is renames
+    assert config.text_config._checkpoint_conversion_mapping is renames

From 8e74d3f0309938308ed820ce8fb1caa07faa4d5c Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Sun, 11 Oct 2026 07:39:42 +0000
Subject: [PATCH 05/38] Name experts back through converters that loading
 rewrote

The export reverts in-memory names to the checkpoint's by replaying the
conversions from_pretrained recorded. A load that rewrote a converter lists
every checkpoint tensor it consumes as a source: a dequantizing FP8 load adds
each weight's weight_scale_inv, and an on-the-fly compressed-tensors
decompressor turns each .weight into weight_packed / weight_scale /
weight_global_scale. The same expert leaf then repeats, the expert reverse
bailed with an arity mismatch, and the whole reverse was skipped, leaving the
export on in-memory names.

On Kimi-K3 that wrote the routed experts as
block_sparse_moe.experts.N.gate_proj/up_proj/down_proj, while the model's own
modeling code builds them as w1/w3/w2 (gate_proj and friends belong to its
dense MLP class), so a stock load matched none of the 989,184 expert tensors.
A converter rewritten by an FP8 dequantizing load fails the same way.

Read a module's leaf through the regex anchor and the compressed-tensors
sibling suffixes, and keep each leaf once, in order. The dense split rule
gets the same treatment, for a fused dense converter a dequantizing load
rewrote.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 .../torch/export/quant_aware_conversion.py    | 31 ++++++-
 .../export/test_quant_aware_conversion.py     | 82 +++++++++++++++++++
 2 files changed, 109 insertions(+), 4 deletions(-)

diff --git a/modelopt/torch/export/quant_aware_conversion.py b/modelopt/torch/export/quant_aware_conversion.py
index 1005a3a904f..2e69233a0b9 100644
--- a/modelopt/torch/export/quant_aware_conversion.py
+++ b/modelopt/torch/export/quant_aware_conversion.py
@@ -98,6 +98,18 @@
 )
 _LEAF_SUFFIXES = (".weight", *_QUANT_STATE_SUFFIXES, ".bias")
 
+# Checkpoint tensors a load-time rewrite folds into a weight -- compressed-tensors' stored
+# parameters, FP8's activation scale -- and so lists as sources next to it. They name the same
+# module as that weight; here they only matter for reading the module's leaf off a pattern.
+_LOAD_SIBLING_SUFFIXES = (
+    ".weight_packed",
+    ".weight_global_scale",
+    ".weight_shape",
+    ".weight_zero_point",
+    ".weight_g_idx",
+    ".activation_scale",
+)
+
 # Leaves that are per-tensor scalars (0-d) and must be *duplicated*, not split, when
 # a fused module is un-fused.
 _SCALAR_LEAF_SUFFIXES = (".weight_scale_2", ".input_scale")
@@ -884,7 +896,7 @@ def _expert_leaf_renames(rev) -> list[RenameRule]:
     fused leaf as per-expert parts, mapped back to the hub leaves positionally.
     """
     src_leaf = _leaf(_as_list(rev.source_patterns)[0])
-    hub_leaves = [_leaf(t) for t in _as_list(rev.target_patterns)]
+    hub_leaves = _distinct_leaves(_as_list(rev.target_patterns))
     part_leaves = _FUSED_EXPERT_PART_NAMES.get(src_leaf, [src_leaf])
     if len(part_leaves) != len(hub_leaves):
         raise QuantConversionUnsupportedError(
@@ -903,7 +915,7 @@ def _expert_leaf_renames(rev) -> list[RenameRule]:
 def _dense_split_rule(rev, ops) -> SplitRule:
     """Un-fuse a dense (non-expert) fused linear that survives in the state dict."""
     fused = _leaf_suffix(_as_list(rev.source_patterns)[0])
-    parts = tuple(_leaf_suffix(t) for t in _as_list(rev.target_patterns))
+    parts = tuple("." + leaf for leaf in _distinct_leaves(_as_list(rev.target_patterns)))
     dim = next((op.dim for op in ops if hasattr(op, "dim")), 0)
     return SplitRule(fused_suffix=fused, part_suffixes=parts, dim=dim)
 
@@ -912,10 +924,21 @@ def _as_list(x) -> list:
     return list(x) if isinstance(x, (list, tuple)) else [x]
 
 
+def _distinct_leaves(patterns: list[str]) -> list[str]:
+    """The module leaves ``patterns`` name, each once, in order.
+
+    A converter that loading rewrote lists every checkpoint tensor it consumes per module: a
+    dequantizing quantizer adds each weight's ``weight_scale_inv``, an on-the-fly decompressor the
+    stored ``weight_packed`` / ``weight_scale`` / ``weight_global_scale``. Those all name the same
+    module, so repeated leaves collapse to it.
+    """
+    return list(dict.fromkeys(_leaf(p) for p in patterns))
+
+
 def _leaf(pattern: str) -> str:
     """Bare leaf name from a conversion pattern, e.g. ``.experts.*.w1.weight`` -> ``w1``."""
-    p = pattern
-    for suffix in _LEAF_SUFFIXES:
+    p = pattern.removesuffix("$")
+    for suffix in (*_LEAF_SUFFIXES, *_LOAD_SIBLING_SUFFIXES):
         if p.endswith(suffix):
             p = p[: -len(suffix)]
             break
diff --git a/tests/unit/torch/export/test_quant_aware_conversion.py b/tests/unit/torch/export/test_quant_aware_conversion.py
index c584c085b8d..9235dc3967e 100644
--- a/tests/unit/torch/export/test_quant_aware_conversion.py
+++ b/tests/unit/torch/export/test_quant_aware_conversion.py
@@ -1281,3 +1281,85 @@ def test_revert_quant_config_names_mapper():
     q2 = {"exclude_modules": ["x*"]}
     revert_quant_config_names(q2, None)
     assert q2["exclude_modules"] == ["x*"]
+
+
+def _as_loading_rewrote(conversions, siblings):
+    """``conversions`` with every expert converter's ``.weight`` sources expanded the way a
+    load-time rewrite expands them: one anchored source per checkpoint tensor it consumes."""
+    from transformers.core_model_loading import WeightConverter
+
+    rewritten = []
+    for conv in conversions:
+        sources = conv.source_patterns if isinstance(conv.source_patterns, list) else []
+        if not isinstance(conv, WeightConverter) or not any(".experts." in s for s in sources):
+            rewritten.append(conv)
+            continue
+        stems = [s.removesuffix("$").removesuffix(".weight") for s in sources]
+        rewritten.append(
+            WeightConverter(
+                source_patterns=[f"{stem}.{leaf}$" for leaf in siblings for stem in stems],
+                target_patterns=getattr(conv, "_original_target_patterns", conv.target_patterns),
+                operations=list(conv.operations),
+            )
+        )
+    return rewritten
+
+
+@pytest.mark.parametrize(
+    "siblings",
+    [("weight", "weight_scale_inv"), ("weight_packed", "weight_scale", "weight_global_scale")],
+    ids=["fp8-dequantize", "compressed-tensors-decompress"],
+)
+def test_experts_revert_through_converters_loading_rewrote(siblings):
+    """The export replays the conversions ``from_pretrained`` recorded. After a dequantizing or
+    decompressing load, each expert converter lists every checkpoint tensor it consumed, so the
+    same leaf repeats; the expert reverse used to bail on that and leave every name in-memory."""
+    pytest.importorskip("transformers")
+    from _test_utils.torch.transformers_models import get_tiny_mixtral
+
+    try:
+        from transformers.conversion_mapping import get_model_conversion_mapping
+    except ImportError:
+        pytest.skip("transformers build has no conversion_mapping API")
+    model = get_tiny_mixtral(**_MIXTRAL_KWARGS)
+    stock = get_model_conversion_mapping(model, add_legacy=False)
+    if not any(".experts." in str(c.source_patterns) for c in stock):
+        pytest.skip("transformers build has no mixtral expert conversion")
+    model._weight_conversions = _as_loading_rewrote(stock, siblings)
+
+    p = "model.layers.0"
+    sd = {f"{p}.mlp.gate.weight": torch.randn(2, 32)}
+    for e in range(2):
+        sd.update(_nvfp4_linear(f"{p}.mlp.experts.{e}.gate_proj", 64, 32))
+        sd.update(_nvfp4_linear(f"{p}.mlp.experts.{e}.up_proj", 64, 32))
+        sd.update(_nvfp4_linear(f"{p}.mlp.experts.{e}.down_proj", 32, 64))
+
+    out = revert_weight_conversion_quant_aware(model, sd)
+
+    for e in range(2):
+        base = f"{p}.block_sparse_moe.experts.{e}"
+        assert f"{base}.w1.weight" in out
+        assert f"{base}.w3.weight" in out
+        assert f"{base}.w2.weight" in out
+        assert f"{base}.w1.weight_scale_2" in out
+    assert not any(".mlp.experts." in k for k in out)
+
+
+def test_dense_split_names_each_part_once_through_load_siblings():
+    """A dense fused converter that a dequantizing load rewrote still un-fuses into one part per
+    module, not one per checkpoint tensor the load consumed."""
+    pytest.importorskip("transformers")
+    from transformers.core_model_loading import Concatenate, WeightConverter
+
+    from modelopt.torch.export.quant_aware_conversion import _dense_split_rule
+
+    modules = ("self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj")
+    leaves = ("weight", "weight_scale_inv")
+    conv = WeightConverter(
+        source_patterns=[f"{m}.{leaf}$" for leaf in leaves for m in modules],
+        target_patterns="self_attn.qkv_proj",
+        operations=[Concatenate(dim=0)],
+    )
+    rev = conv.reverse_transform()
+    rule = _dense_split_rule(rev, list(rev.operations))
+    assert rule.part_suffixes == (".q_proj", ".k_proj", ".v_proj")

From 848a5abd79b2c394dcca7aebe530085aab9c3717 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Sun, 11 Oct 2026 11:11:48 +0000
Subject: [PATCH 06/38] Stop restoring load-time checkpoint renames around
 config.json

They were kept out of config.json because Kimi-K3's prepare hook installed
_checkpoint_conversion_mapping on its configs. The hook registers its
per-expert conversion with transformers' conversion registry instead now, and
nothing in the expert-parallel path sets that attribute any more.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/export/unified_export_hf.py        | 12 +++---------
 tests/unit/torch/export/test_unified_export_hf.py | 12 +++---------
 2 files changed, 6 insertions(+), 18 deletions(-)

diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py
index 1fee18347e7..08e6ade6332 100644
--- a/modelopt/torch/export/unified_export_hf.py
+++ b/modelopt/torch/export/unified_export_hf.py
@@ -1969,15 +1969,9 @@ def read_unplaced_weights(model: nn.Module, *, keys_only: bool = False) -> dict[
 
 # Config attributes a distributed load sets to describe THIS RUN rather than the checkpoint: the
 # DistributedConfig transformers stamps on the config when it shards the model (the EP degree and tp
-# plan the export happened to use; not JSON-serializable either), the base_model_ep_plan the EP
-# loader swapped to the torch_a2a styles, which exist only while ModelOpt is imported, and the
-# _checkpoint_conversion_mapping a model's expert-parallel prepare hook installs to rename the
-# source's tensors on load (Kimi-K3's per-expert w1/w3/w2), which names nothing in the export.
-_DISTRIBUTED_RUN_CONFIG_ATTRS = (
-    "distributed_config",
-    "base_model_ep_plan",
-    "_checkpoint_conversion_mapping",
-)
+# plan the export happened to use; not JSON-serializable either), and the base_model_ep_plan the EP
+# loader swapped to the torch_a2a styles, which exist only while ModelOpt is imported.
+_DISTRIBUTED_RUN_CONFIG_ATTRS = ("distributed_config", "base_model_ep_plan")
 
 
 @contextlib.contextmanager
diff --git a/tests/unit/torch/export/test_unified_export_hf.py b/tests/unit/torch/export/test_unified_export_hf.py
index d55e1a6775d..4a0cb931018 100644
--- a/tests/unit/torch/export/test_unified_export_hf.py
+++ b/tests/unit/torch/export/test_unified_export_hf.py
@@ -1137,8 +1137,8 @@ def test_process_hf_model_for_export_folds_in_place_and_returns_the_quant_config
 
 @pytest.mark.parametrize("source_plan", [None, {"layers.*.feed_forward.experts": "grouped_gemm"}])
 def test_distributed_export_writes_the_config_a_single_gpu_export_writes(tmp_path, source_plan):
-    """The EP load's DistributedConfig, torch_a2a plan and load-time renames stay out of config.json,
-    the source's own plan survives, and the live configs keep the run's values afterwards."""
+    """The EP load's DistributedConfig and torch_a2a plan stay out of config.json, the source's own
+    plan survives, and the live configs keep the run's values afterwards."""
     from transformers import AutoConfig, Llama4Config
     from transformers.distributed import DistributedConfig
 
@@ -1152,16 +1152,12 @@ def test_distributed_export_writes_the_config_a_single_gpu_export_writes(tmp_pat
         (source / "config.json").write_text(json.dumps(config_json))
     AutoConfig.from_pretrained(source).save_pretrained(tmp_path / "single_gpu")
 
-    # What the EP load leaves behind: the stamped DistributedConfig, the swapped plan, and the
-    # renames a prepare hook installs on both configs (as Kimi-K3's does).
+    # What the EP load leaves behind: the stamped DistributedConfig and the swapped plan.
     config = AutoConfig.from_pretrained(source)
     distributed_config = DistributedConfig(enable_expert_parallel=True)
     swapped_plan = {"layers.*.feed_forward.experts": "torch_a2a_experts"}
-    renames = {r"experts\.(\d+)\.w1\.": r"experts.\1.gate_proj."}
     config.distributed_config = distributed_config
     config.text_config.base_model_ep_plan = swapped_plan
-    config._checkpoint_conversion_mapping = renames
-    config.text_config._checkpoint_conversion_mapping = renames
     model = SimpleNamespace(
         config=config, _modelopt_source_checkpoint=str(source), can_generate=lambda: False
     )
@@ -1173,5 +1169,3 @@ def test_distributed_export_writes_the_config_a_single_gpu_export_writes(tmp_pat
     assert written["text_config"].get("base_model_ep_plan") == source_plan
     assert config.distributed_config is distributed_config
     assert config.text_config.base_model_ep_plan is swapped_plan
-    assert config._checkpoint_conversion_mapping is renames
-    assert config.text_config._checkpoint_conversion_mapping is renames

From 8a27ccce80f7f860f15983ea51e65ea9e10ded5e Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Sun, 11 Oct 2026 11:34:57 +0000
Subject: [PATCH 07/38] Reverse the conversions around a dequantizing FP8 load

A FineGrainedFP8 dequantize=True load adds a converter that decodes each
weight in place, and its reverse re-encodes it (Fp8Quantize), which the
quant-aware reverse does not support. It then gave up on the whole mapping,
so MiniMax-M2.7 exported its routed experts under in-memory names
(mlp.experts.N.gate_proj) instead of the checkpoint's
(block_sparse_moe.experts.N.w1).

Such a converter renames nothing, and the export writes ModelOpt's own
quantized weights and scales rather than re-encoding the source's format, so
it contributes no reverse rule now.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 .../torch/export/quant_aware_conversion.py    | 29 +++++++++++++++++++
 .../export/test_quant_aware_conversion.py     | 28 ++++++++++++++++++
 2 files changed, 57 insertions(+)

diff --git a/modelopt/torch/export/quant_aware_conversion.py b/modelopt/torch/export/quant_aware_conversion.py
index 2e69233a0b9..02cb7b0915e 100644
--- a/modelopt/torch/export/quant_aware_conversion.py
+++ b/modelopt/torch/export/quant_aware_conversion.py
@@ -757,6 +757,31 @@ def _drop_shadowed_prefix_renames(model, rules: list[RenameRule]) -> list[Rename
     return kept
 
 
+def _quantization_codec_ops() -> tuple[type, ...]:
+    """Transformers' ops that decode or encode a weight's on-disk quantization in place."""
+    try:
+        from transformers.integrations import finegrained_fp8
+    except Exception:  # transformers without the FP8 integration
+        return ()
+    names = ("Fp8Dequantize", "Fp8Quantize", "Fp8DecodeScale", "Fp8EncodeScale")
+    return tuple(op for op in (getattr(finegrained_fp8, n, None) for n in names) if op is not None)
+
+
+def _decodes_in_place(rev, ops) -> bool:
+    """Whether a reversed converter only re-encodes a weight's quantization, under the same name.
+
+    A dequantizing load (FineGrainedFP8 ``dequantize=True``) decodes each weight in place, and
+    its reverse would re-encode the weight in the source's format. The export never wants that:
+    it writes ModelOpt's own quantized weights and scales. Such a converter renames nothing.
+    """
+    codec_ops = _quantization_codec_ops()
+    if not ops or not codec_ops or not all(isinstance(op, codec_ops) for op in ops):
+        return False
+    sources = [p.strip("^$") for p in _as_list(rev.source_patterns)]
+    targets = [p.strip("^$") for p in _as_list(rev.target_patterns)]
+    return bool(sources) and sources[:1] == targets[:1]
+
+
 def _build_reverse_rules(
     model,
 ) -> tuple[list[SplitRule], list[_MergeRule], list[RenameRule], list[str], set[str] | None]:
@@ -779,6 +804,8 @@ def _build_reverse_rules(
       survives in the state dict, so it is un-fused via a :class:`SplitRule`.
     * Dense splitting ``WeightConverter`` (reverse is ``Concatenate`` only): the split
       tensors are re-fused under the transform's original sub-model scope.
+    * A dequantizing load's ``WeightConverter`` (quantization codec ops only, same name): no
+      rule -- the export writes ModelOpt's own quantization, not the source's.
 
     Raises :class:`QuantConversionUnsupportedError` for any op shape not covered, so
     the caller falls back to the legacy (in-memory-name) behavior.
@@ -854,6 +881,8 @@ def _build_reverse_rules(
                         dim=ops[0].dim,
                     )
                 )
+            elif _decodes_in_place(rev, ops):
+                continue
             else:
                 raise QuantConversionUnsupportedError(
                     f"unsupported reverse ops: {[type(o).__name__ for o in ops]}"
diff --git a/tests/unit/torch/export/test_quant_aware_conversion.py b/tests/unit/torch/export/test_quant_aware_conversion.py
index 9235dc3967e..e9156e91998 100644
--- a/tests/unit/torch/export/test_quant_aware_conversion.py
+++ b/tests/unit/torch/export/test_quant_aware_conversion.py
@@ -705,6 +705,34 @@ def test_build_reverse_rules_orders_prefix_reorder_after_container():
     assert not any(".mlp.experts." in k for k in out)
 
 
+def test_a_dequantizing_fp8_load_still_reverses_the_other_conversions():
+    """A FineGrainedFP8 ``dequantize=True`` load adds a converter that only decodes each weight in
+    place. It renames nothing, so it must neither abort the reverse of the model's real
+    conversions (MiniMax-M2.7 kept its in-memory expert names) nor re-encode a weight."""
+    pytest.importorskip("transformers.core_model_loading")
+    fp8 = pytest.importorskip("transformers.integrations.finegrained_fp8")
+    from transformers.core_model_loading import WeightConverter, WeightRenaming
+
+    conversions = [
+        WeightRenaming(".block_sparse_moe.experts.", ".mlp.experts."),
+        WeightConverter(
+            source_patterns=["weight$", "weight_scale_inv", "activation_scale"],
+            target_patterns="weight",
+            operations=[fp8.Fp8Dequantize(None)],
+        ),
+    ]
+    model = types.SimpleNamespace(_weight_conversions=conversions)
+    sd = _nvfp4_linear("model.layers.0.mlp.experts.0.w1", 8, 16)
+    sd["model.layers.0.self_attn.q_proj.weight"] = torch.randn(8, 16, dtype=torch.bfloat16)
+
+    out = revert_weight_conversion_quant_aware(model, sd)
+
+    expert = "model.layers.0.block_sparse_moe.experts.0.w1"
+    assert {f"{expert}.weight", f"{expert}.weight_scale", f"{expert}.weight_scale_2"} <= set(out)
+    assert out["model.layers.0.self_attn.q_proj.weight"].dtype == torch.bfloat16
+    assert not any(".mlp.experts." in k or k.endswith("weight_scale_inv") for k in out)
+
+
 def test_nested_text_prefix_reverse_does_not_capture_vlm_siblings():
     """A nested text-model conversion must not rewrite the full VLM namespace."""
     pytest.importorskip("transformers.core_model_loading")

From 55dc3857b457c9bdb776236510e1d7845dbccbbe Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Sun, 11 Oct 2026 13:02:27 +0000
Subject: [PATCH 08/38] Recognize the identity reverse older transformers give
 a dequantization

transformers 5.5 and 5.6 reverse Fp8Dequantize into their identity op rather
than Fp8Quantize, so on CI's minimum transformers the dequantizing converter
still aborted the reverse.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/export/quant_aware_conversion.py | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/modelopt/torch/export/quant_aware_conversion.py b/modelopt/torch/export/quant_aware_conversion.py
index 02cb7b0915e..e6dcd965bd7 100644
--- a/modelopt/torch/export/quant_aware_conversion.py
+++ b/modelopt/torch/export/quant_aware_conversion.py
@@ -758,13 +758,20 @@ def _drop_shadowed_prefix_renames(model, rules: list[RenameRule]) -> list[Rename
 
 
 def _quantization_codec_ops() -> tuple[type, ...]:
-    """Transformers' ops that decode or encode a weight's on-disk quantization in place."""
+    """Transformers' ops that decode or encode a weight's on-disk quantization in place.
+
+    Includes the identity op, which is what older releases (5.5, 5.6) reverse ``Fp8Dequantize``
+    into; later ones reverse it into ``Fp8Quantize``.
+    """
     try:
+        from transformers import core_model_loading
         from transformers.integrations import finegrained_fp8
     except Exception:  # transformers without the FP8 integration
         return ()
     names = ("Fp8Dequantize", "Fp8Quantize", "Fp8DecodeScale", "Fp8EncodeScale")
-    return tuple(op for op in (getattr(finegrained_fp8, n, None) for n in names) if op is not None)
+    ops = [getattr(finegrained_fp8, n, None) for n in names]
+    ops.append(getattr(core_model_loading, "_IdentityOp", None))
+    return tuple(op for op in ops if op is not None)
 
 
 def _decodes_in_place(rev, ops) -> bool:

From e59190bcc7f6d7f16e1623376d6b14364ca9b213 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Sun, 11 Oct 2026 16:39:16 +0000
Subject: [PATCH 09/38] Skip the export config tests where transformers is not
 installed

CI's partial-install torch lane runs tests/unit/torch without the hf extras. These two tests import transformers inside the test body, so they failed there instead of skipping like the rest of the file.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 tests/unit/torch/export/test_unified_export_hf.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tests/unit/torch/export/test_unified_export_hf.py b/tests/unit/torch/export/test_unified_export_hf.py
index 4a0cb931018..05bbdc3c083 100644
--- a/tests/unit/torch/export/test_unified_export_hf.py
+++ b/tests/unit/torch/export/test_unified_export_hf.py
@@ -1111,6 +1111,7 @@ def test_source_checkpoint_is_none_for_a_model_built_from_config(tmp_path, monke
 
 def test_process_hf_model_for_export_folds_in_place_and_returns_the_quant_config():
     """The processing half the no-gather writer runs before writing each rank's shard."""
+    pytest.importorskip("transformers")
     from transformers import LlamaConfig, LlamaForCausalLM
 
     from modelopt.torch.export.unified_export_hf import _process_hf_model_for_export
@@ -1139,6 +1140,7 @@ def test_process_hf_model_for_export_folds_in_place_and_returns_the_quant_config
 def test_distributed_export_writes_the_config_a_single_gpu_export_writes(tmp_path, source_plan):
     """The EP load's DistributedConfig and torch_a2a plan stay out of config.json, the source's own
     plan survives, and the live configs keep the run's values afterwards."""
+    pytest.importorskip("transformers")
     from transformers import AutoConfig, Llama4Config
     from transformers.distributed import DistributedConfig
 

From 59077aefa489bab9fae31b8652ae8a0865eeadc1 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Wed, 14 Oct 2026 17:41:58 +0000
Subject: [PATCH 10/38] Fail a sharded export on every rank when any rank's
 export fails

A barrier alone let the ranks whose writes succeeded return, and the run carried on as if the
checkpoint were whole while another rank's export raised. The sharded export now ends by
gathering every rank's outcome (all_reduce), so every rank raises with the failing rank's error.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/export/unified_export_hf.py    | 22 +++++++++++++-
 .../torch/export/test_unified_export_hf.py    | 30 +++++++++++++++++++
 2 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py
index 08e6ade6332..f05f4b0042f 100644
--- a/modelopt/torch/export/unified_export_hf.py
+++ b/modelopt/torch/export/unified_export_hf.py
@@ -99,6 +99,7 @@
 from modelopt.torch.utils.dataset_utils import _disable_use_cache
 from modelopt.torch.utils.distributed import (
     DistributionPlan,
+    all_reduce,
     broadcast_from,
     distribution_plan,
     is_fsdp2_model,
@@ -2179,6 +2180,7 @@ def export_hf_checkpoint(
     # buffer instead of the whole quantized state dict.
     _offloaded = has_accelerate_offload(model)
 
+    error: Exception | None = None
     try:
         if _offloaded:
             # Imported here rather than at module scope: the streaming exporter imports the
@@ -2331,6 +2333,7 @@ def export_hf_checkpoint(
                 _copy_non_model_files_from_source(model, export_dir)
 
     except Exception as e:
+        error = e
         # Print (not just warnings.warn, which the default filter can dedupe/suppress) BEFORE the
         # barrier below: under FSDP2/EP every other rank is still waiting there, so a failing rank
         # that raises straight into `finally` blocks on that barrier forever without this rank's
@@ -2346,4 +2349,21 @@ def export_hf_checkpoint(
         # Every rank waits for the writer to finish. Without this the ranks that wrote nothing exit
         # first and the job hangs on the ones still writing.
         if is_fsdp2_sharded or is_ep_sharded:
-            torch.distributed.barrier()
+            _finish_export_together(error)
+
+
+def _finish_export_together(error: Exception | None) -> None:
+    """Wait for every rank to finish writing, and raise on all of them if any failed.
+
+    A barrier alone lets the ranks whose writes succeeded return, and the run carries on as if the
+    checkpoint were whole while another rank's export raised. Every rank's outcome is gathered
+    instead (:func:`~modelopt.torch.utils.distributed.all_reduce`), so every rank raises with the
+    failing rank's error.
+    """
+
+    def exported() -> bool:
+        if error is not None:
+            raise error
+        return True
+
+    all_reduce("export_ok", exported, all)
diff --git a/tests/unit/torch/export/test_unified_export_hf.py b/tests/unit/torch/export/test_unified_export_hf.py
index 05bbdc3c083..685755c15ca 100644
--- a/tests/unit/torch/export/test_unified_export_hf.py
+++ b/tests/unit/torch/export/test_unified_export_hf.py
@@ -18,10 +18,12 @@
 import importlib.util
 import json
 from copy import deepcopy
+from functools import partial
 from types import SimpleNamespace
 
 import pytest
 import torch
+from _test_utils.torch.distributed.utils import spawn_multiprocess_job
 from _test_utils.torch.quantization.tied_modules import (
     make_tied_linear_pair,
     wrap_in_parent_with_tied_keys,
@@ -1171,3 +1173,31 @@ def test_distributed_export_writes_the_config_a_single_gpu_export_writes(tmp_pat
     assert written["text_config"].get("base_model_ep_plan") == source_plan
     assert config.distributed_config is distributed_config
     assert config.text_config.base_model_ep_plan is swapped_plan
+
+
+def _finish_export_job(failing_rank, rank, size):
+    from modelopt.torch.export.unified_export_hf import _finish_export_together
+
+    if failing_rank is None:
+        _finish_export_together(None)
+    elif rank == failing_rank:
+        with pytest.raises(OSError, match="disk full"):
+            _finish_export_together(OSError("disk full"))
+    else:
+        with pytest.raises(RuntimeError, match=rf"rank {failing_rank} failed.*disk full"):
+            _finish_export_together(None)
+
+
+@pytest.mark.parametrize("failing_rank", [None, 0, 1])
+def test_every_rank_fails_an_export_that_any_rank_failed(failing_rank):
+    """A barrier alone lets the ranks that wrote their shards carry on as if the export succeeded."""
+    spawn_multiprocess_job(2, partial(_finish_export_job, failing_rank), backend="gloo")
+
+
+def test_a_single_process_export_ends_on_its_own_outcome():
+    """Nothing to wait for: the export's own error, if any, is all that is raised."""
+    from modelopt.torch.export.unified_export_hf import _finish_export_together
+
+    _finish_export_together(None)
+    with pytest.raises(OSError, match="disk full"):
+        _finish_export_together(OSError("disk full"))

From 2839200223fbb899e82f42d941256762e52828d4 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Wed, 14 Oct 2026 20:41:19 +0000
Subject: [PATCH 11/38] Read an expert-parallel export's sharding from the
 model's distribution plan

_hf_shard_flags reports FSDP2 and expert-parallel sharding from the model's recorded plan, rather
than from an FSDP2 module or a seeded expert-parallel group, so the export's route, its single
writer of the carried weights and its fail-together finish all follow the plan. The writer takes
its EP coordinate from the experts' group only when the plan is expert-parallel, and raises if
the experts carry none: without it every EP rank would write experts 0..n-1 over the others'.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/export/distributed_writer.py   | 11 +++++++---
 modelopt/torch/export/unified_export_hf.py    | 15 ++++++-------
 modelopt/torch/utils/distributed.py           |  2 +-
 .../torch/export/test_distributed_writer.py   | 12 +++++++++++
 .../torch/export/test_unified_export_hf.py    | 21 +++++++++++++++++++
 5 files changed, 48 insertions(+), 13 deletions(-)

diff --git a/modelopt/torch/export/distributed_writer.py b/modelopt/torch/export/distributed_writer.py
index 51b9365d807..c6e20cd183a 100644
--- a/modelopt/torch/export/distributed_writer.py
+++ b/modelopt/torch/export/distributed_writer.py
@@ -729,11 +729,16 @@ def _distributed_save_hf_checkpoint_impl(
     # other and the checkpoint would be quietly wrong rather than broken. The coordinate comes from
     # the routed experts' ParallelState -- the same source calibration reads, not transformers'
     # internal ``model._device_mesh``. An expert's data-parallel group is its expert-data-parallel one.
-    from modelopt.torch.utils.distributed import expert_parallel_state
+    from modelopt.torch.utils.distributed import distribution_plan, expert_parallel_state
 
     ep_rank, dp_idx = 0, 0
-    parallel_state = expert_parallel_state(model)
-    if parallel_state is not None:
+    if distribution_plan(model).is_expert_parallel:
+        parallel_state = expert_parallel_state(model)
+        if parallel_state is None:
+            raise RuntimeError(
+                "The model's plan is expert-parallel, but no routed expert carries its expert-parallel "
+                "group: load it with load_and_prepare_ep."
+            )
         ep_rank = parallel_state.expert_model_parallel_group.rank()
         edp_group = parallel_state.data_parallel_group
         if edp_group.is_initialized():
diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py
index f05f4b0042f..0d819e6ab98 100644
--- a/modelopt/torch/export/unified_export_hf.py
+++ b/modelopt/torch/export/unified_export_hf.py
@@ -98,7 +98,6 @@
 from modelopt.torch.utils import print_rank_0
 from modelopt.torch.utils.dataset_utils import _disable_use_cache
 from modelopt.torch.utils.distributed import (
-    DistributionPlan,
     all_reduce,
     broadcast_from,
     distribution_plan,
@@ -1090,19 +1089,17 @@ def pack_quantized_weights(
 
 
 def _hf_shard_flags(model: nn.Module) -> tuple[bool, bool]:
-    """Return ``(fsdp2_sharded, ep_sharded)`` for an HF model about to be exported.
+    """Return ``(fsdp2_sharded, ep_sharded)`` for an HF model about to be exported, from its plan.
 
     Expert parallelism shards the experts too, just not by FSDP2: transformers pre-slices the expert
     params across the EP group. The two are reported separately because a model can be both at once,
-    FSDP2-dense with expert-parallel experts. Neither holds without an initialized process group.
+    FSDP2-dense with expert-parallel experts. A diffusers pipeline is not a module, and nothing here
+    distributes one.
     """
-    from modelopt.torch.utils.distributed import expert_parallel_state
-
-    if not (torch.distributed.is_available() and torch.distributed.is_initialized()):
+    if not isinstance(model, nn.Module):
         return False, False
-    fsdp2_sharded = is_fsdp2_model(model)
-    ep_sharded = torch.distributed.get_world_size() > 1 and expert_parallel_state(model) is not None
-    return fsdp2_sharded, ep_sharded
+    plan = distribution_plan(model)
+    return plan.is_fsdp_dense, plan.is_expert_parallel
 
 
 def _process_hf_model_for_export(
diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py
index facaf25fe44..3a0d96458a0 100644
--- a/modelopt/torch/utils/distributed.py
+++ b/modelopt/torch/utils/distributed.py
@@ -692,7 +692,7 @@ def distributed_run(
 
 
 def expert_parallel_state(model: torch.nn.Module) -> "ParallelState | None":
-    """Return the ``ParallelState`` of the model's routed experts, or ``None`` if they are not expert-parallel.
+    """Return the ``ParallelState`` of the model's routed experts, or ``None`` if none has one.
 
     The expert-parallel loader seeds every routed-expert module with a ``ParallelState`` that carries the
     expert-model-parallel group (``load_and_prepare_ep``); no other module of a HuggingFace model has one.
diff --git a/tests/unit/torch/export/test_distributed_writer.py b/tests/unit/torch/export/test_distributed_writer.py
index a35e47086c9..99d0df37604 100644
--- a/tests/unit/torch/export/test_distributed_writer.py
+++ b/tests/unit/torch/export/test_distributed_writer.py
@@ -177,6 +177,18 @@ def _fsdp2_shard(model, mesh, wrap_root=True):
     return model
 
 
+def test_an_expert_parallel_plan_needs_its_experts_expert_parallel_group(
+    gloo_world_of_one, tmp_path
+):
+    """Without the group the writer could not offset each EP rank's experts, and every rank would
+    write experts 0..n-1 over the others'."""
+    model = _load(_save(_llama_config(), tmp_path / "src"))
+    set_distribution_plan(model, DistributionPlan(2, ep_size=2))
+
+    with pytest.raises(RuntimeError, match="no routed expert carries its expert-parallel group"):
+        _write(model, tmp_path / "export")
+
+
 # FP8 weights, with an FP8 KV cache in layer 0 and an NVFP4 one in layer 1: a mixed-precision KV
 # cache, whose per-layer formats the export reads from the quantization config.
 _FP8_MIXED_KV_CFG = copy.deepcopy(mtq.FP8_DEFAULT_CFG)
diff --git a/tests/unit/torch/export/test_unified_export_hf.py b/tests/unit/torch/export/test_unified_export_hf.py
index 685755c15ca..4adea2e9a81 100644
--- a/tests/unit/torch/export/test_unified_export_hf.py
+++ b/tests/unit/torch/export/test_unified_export_hf.py
@@ -1201,3 +1201,24 @@ def test_a_single_process_export_ends_on_its_own_outcome():
     _finish_export_together(None)
     with pytest.raises(OSError, match="disk full"):
         _finish_export_together(OSError("disk full"))
+
+
+@pytest.mark.parametrize(
+    ("dense", "ep_size", "flags"),
+    [
+        ("replicate", 1, (False, False)),
+        ("fsdp", 1, (True, False)),
+        ("replicate", 2, (False, True)),
+        ("fsdp", 2, (True, True)),
+    ],
+    ids=["replicated", "fsdp2", "expert-parallel", "both"],
+)
+def test_the_export_takes_how_the_model_is_sharded_from_its_plan(dense, ep_size, flags):
+    """A model can be FSDP2-dense and expert-parallel at once, so the two are read separately."""
+    from modelopt.torch.export.unified_export_hf import _hf_shard_flags
+    from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan
+
+    model = torch.nn.Linear(2, 2)
+    set_distribution_plan(model, DistributionPlan(2, ep_size, dense))
+
+    assert _hf_shard_flags(model) == flags

From 87311facbc1e2ba1ba7c006f647db5ad59ec3dcf Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Wed, 14 Oct 2026 23:46:16 +0000
Subject: [PATCH 12/38] Keep two blank lines after expert_parallel_state

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/utils/distributed.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py
index 3a0d96458a0..3ccc6022ff1 100644
--- a/modelopt/torch/utils/distributed.py
+++ b/modelopt/torch/utils/distributed.py
@@ -702,6 +702,8 @@ def expert_parallel_state(model: torch.nn.Module) -> "ParallelState | None":
         if state is not None and state.expert_model_parallel_group.is_initialized():
             return state
     return None
+
+
 class _StridedIterableShard(IterableDataset):
     """Rank-strided shard of a streaming ``IterableDataset`` (no ``len()`` required).
 

From ad1ad14c71bccc085fd420524a8fa1839084c0f9 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Thu, 15 Oct 2026 22:31:35 +0000
Subject: [PATCH 13/38] Run the every-rank-fails export test on the shared gloo
 workers

Each of its three cases spawned two ranks, about 14 s of the linux unit job apiece.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 tests/unit/torch/export/test_unified_export_hf.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/tests/unit/torch/export/test_unified_export_hf.py b/tests/unit/torch/export/test_unified_export_hf.py
index 4adea2e9a81..cd15da1f2b3 100644
--- a/tests/unit/torch/export/test_unified_export_hf.py
+++ b/tests/unit/torch/export/test_unified_export_hf.py
@@ -23,7 +23,6 @@
 
 import pytest
 import torch
-from _test_utils.torch.distributed.utils import spawn_multiprocess_job
 from _test_utils.torch.quantization.tied_modules import (
     make_tied_linear_pair,
     wrap_in_parent_with_tied_keys,
@@ -1189,9 +1188,9 @@ def _finish_export_job(failing_rank, rank, size):
 
 
 @pytest.mark.parametrize("failing_rank", [None, 0, 1])
-def test_every_rank_fails_an_export_that_any_rank_failed(failing_rank):
+def test_every_rank_fails_an_export_that_any_rank_failed(cpu_dist_workers, failing_rank):
     """A barrier alone lets the ranks that wrote their shards carry on as if the export succeeded."""
-    spawn_multiprocess_job(2, partial(_finish_export_job, failing_rank), backend="gloo")
+    cpu_dist_workers.run(partial(_finish_export_job, failing_rank))
 
 
 def test_a_single_process_export_ends_on_its_own_outcome():

From d820cf3b912cd4a1a1f8c9c1963925ff9fd21154 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Thu, 15 Oct 2026 20:57:19 +0000
Subject: [PATCH 14/38] Refuse an expert-parallel export whose FSDP2 dense
 shards are replicated

With fewer expert-parallel ranks than the world, FSDP2 shards the dense part
within each expert-parallel group and replicates it across the groups. The
no-gather writer, which every expert-parallel export takes, writes a dense tensor
only from a dim-0 shard over the whole world, so it refused that layout -- after
the export had folded the quantizers. The export now refuses the plan before it
folds anything and names the knobs that avoid it.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/export/unified_export_hf.py    | 14 +++++++++++++
 .../torch/export/test_unified_export_hf.py    | 20 +++++++++++++++++++
 2 files changed, 34 insertions(+)

diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py
index 0d819e6ab98..e0bdd9b2396 100644
--- a/modelopt/torch/export/unified_export_hf.py
+++ b/modelopt/torch/export/unified_export_hf.py
@@ -2097,6 +2097,20 @@ def export_hf_checkpoint(
     # because the carry-over gate needs it and runs ahead of them (the layerwise and diffusers paths
     # return before those are reached).
     _fsdp2_sharded, _ep_sharded = _hf_shard_flags(model)
+    if _fsdp2_sharded and _ep_sharded:
+        plan = distribution_plan(model)
+        # With fewer expert-parallel ranks than the world, FSDP2 shards the dense part within each
+        # expert-parallel group and replicates it across them, and the no-gather writer every
+        # expert-parallel export takes assembles a dense tensor only from a dim-0 shard over the
+        # whole world. Refused here, before anything is folded, not after the export has begun.
+        if plan.edp > 1:
+            raise NotImplementedError(
+                f"The export cannot write this plan's dense part: with ep_size={plan.ep_size} of "
+                f"world_size={plan.world_size} and dense='fsdp', it is sharded within each of the "
+                f"{plan.edp} expert-parallel groups and replicated across them. Set "
+                f"MODELOPT_EP_DEGREE={plan.world_size} if that divides the expert count, or "
+                "MODELOPT_FSDP_SIZE=1 to replicate the dense part."
+            )
     # Either form of sharding, including the hybrid where both hold: a model can be FSDP2-dense AND
     # expert-parallel, so these are tracked separately rather than as an either/or.
     _sharded_export = _ep_sharded or _fsdp2_sharded
diff --git a/tests/unit/torch/export/test_unified_export_hf.py b/tests/unit/torch/export/test_unified_export_hf.py
index cd15da1f2b3..5796d4ce321 100644
--- a/tests/unit/torch/export/test_unified_export_hf.py
+++ b/tests/unit/torch/export/test_unified_export_hf.py
@@ -1221,3 +1221,23 @@ def test_the_export_takes_how_the_model_is_sharded_from_its_plan(dense, ep_size,
     set_distribution_plan(model, DistributionPlan(2, ep_size, dense))
 
     assert _hf_shard_flags(model) == flags
+
+
+def test_an_expert_parallel_plan_with_replicated_dense_shards_is_refused_before_folding(tmp_path):
+    """Below the world size, EP leaves FSDP2's dense shards replicated across the expert-parallel
+    groups, which the no-gather writer cannot write: refused before the export folds anything."""
+    from modelopt.torch.export.unified_export_hf import export_hf_checkpoint
+    from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan
+
+    model = torch.nn.Sequential(torch.nn.Linear(16, 16))
+    mtq.quantize(model, mtq.FP8_DEFAULT_CFG, lambda m: m(torch.randn(2, 16)))
+    set_distribution_plan(model, DistributionPlan(4, ep_size=2, dense="fsdp"))
+    before = {k: v.clone() for k, v in model.state_dict().items()}
+
+    with pytest.raises(NotImplementedError, match=r"MODELOPT_EP_DEGREE=4 .* MODELOPT_FSDP_SIZE=1"):
+        export_hf_checkpoint(model, export_dir=tmp_path / "export")
+
+    after = model.state_dict()
+    assert after.keys() == before.keys()
+    assert all(torch.equal(after[k], v) for k, v in before.items())
+    assert not (tmp_path / "export").exists()

From cfe2c840d81709af2d7c1061744d5efa856e9814 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Thu, 15 Oct 2026 20:59:59 +0000
Subject: [PATCH 15/38] Stop every rank before the no-gather write when one
 rank's processing failed

Each rank folds only its own shard, so processing can fail on one rank alone,
for instance when that rank's experts mix quantized and unquantized ones. That
rank went on to the closing export_ok consensus while its peers started the
writer's collectives, and the two did not match. Processing now ends in its own
consensus, so every rank raises before any of them writes.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/export/unified_export_hf.py    | 21 ++++++++++-
 .../torch/export/test_unified_export_hf.py    | 37 +++++++++++++++++++
 2 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py
index e0bdd9b2396..e3ff60c4bac 100644
--- a/modelopt/torch/export/unified_export_hf.py
+++ b/modelopt/torch/export/unified_export_hf.py
@@ -2226,7 +2226,7 @@ def export_hf_checkpoint(
             _qlora = kwargs.get("is_modelopt_qlora", False)
             # Folds in place; the model is sharded here, so the MoE experts stay fused and each rank
             # folds only its LOCAL experts. The writer splits them after consolidation.
-            hf_quant_config = _process_hf_model_for_export(model, dtype, _qlora)
+            hf_quant_config = _process_together(model, dtype, _qlora)
             if getattr(model, "hf_quantizer", None) is not None:
                 model.hf_quantizer = None
             reverted = distributed_save_hf_checkpoint(
@@ -2363,6 +2363,25 @@ def export_hf_checkpoint(
             _finish_export_together(error)
 
 
+def _process_together(
+    model: nn.Module, dtype: torch.dtype | None, is_modelopt_qlora: bool
+) -> dict[str, Any]:
+    """:func:`_process_hf_model_for_export` on every rank, raising on all of them if any failed.
+
+    Each rank folds only its own shard, so processing can fail on one rank alone -- one rank's
+    experts mixing quantized and unquantized ones, say. That rank would leave for the closing
+    consensus while its peers start the writer's collectives, and the two would not match.
+    """
+    quant_config: list[dict[str, Any]] = []
+
+    def processed() -> bool:
+        quant_config.append(_process_hf_model_for_export(model, dtype, is_modelopt_qlora))
+        return True
+
+    all_reduce("export processed", processed, all)
+    return quant_config[0]
+
+
 def _finish_export_together(error: Exception | None) -> None:
     """Wait for every rank to finish writing, and raise on all of them if any failed.
 
diff --git a/tests/unit/torch/export/test_unified_export_hf.py b/tests/unit/torch/export/test_unified_export_hf.py
index 5796d4ce321..cac21a9eb14 100644
--- a/tests/unit/torch/export/test_unified_export_hf.py
+++ b/tests/unit/torch/export/test_unified_export_hf.py
@@ -1193,6 +1193,43 @@ def test_every_rank_fails_an_export_that_any_rank_failed(cpu_dist_workers, faili
     cpu_dist_workers.run(partial(_finish_export_job, failing_rank))
 
 
+def _one_rank_fails_processing_job(rank, size, tmp_dir):
+    from unittest import mock
+
+    from modelopt.torch.export import distributed_writer, unified_export_hf
+    from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan
+
+    model = torch.nn.Linear(2, 2)
+    set_distribution_plan(model, DistributionPlan(size, ep_size=size))
+    failure = NotImplementedError("rank 1's experts mix quantized and unquantized ones")
+
+    def process(*args):
+        if rank == 1:
+            raise failure
+        return {}
+
+    expected = (
+        pytest.raises(NotImplementedError, match="rank 1's experts")
+        if rank == 1
+        else pytest.raises(RuntimeError, match="rank 1 failed to compute export processed")
+    )
+    with (
+        mock.patch.object(unified_export_hf, "_process_hf_model_for_export", process),
+        mock.patch.object(distributed_writer, "distributed_save_hf_checkpoint") as write,
+        expected,
+    ):
+        unified_export_hf.export_hf_checkpoint(model, export_dir=tmp_dir)
+    write.assert_not_called()
+
+
+def test_every_rank_stops_before_the_write_when_one_ranks_processing_failed(
+    cpu_dist_workers, tmp_path
+):
+    """Each rank folds only its own shard, so its processing can fail alone; the others must not
+    go on into the writer's collectives without it."""
+    cpu_dist_workers.run(_one_rank_fails_processing_job, tmp_dir=tmp_path)
+
+
 def test_a_single_process_export_ends_on_its_own_outcome():
     """Nothing to wait for: the export's own error, if any, is all that is raised."""
     from modelopt.torch.export.unified_export_hf import _finish_export_together

From 0877280a3c2c593c81bc8f977a55e666b9d5d7db Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Thu, 15 Oct 2026 23:30:45 +0000
Subject: [PATCH 16/38] Agree across ranks on fused experts' quantization
 before an expert-parallel export

Each rank holds only its own experts' weight quantizers, so the keep-fused
fold's all-or-none check saw only that rank's share. A rank whose experts were
all unquantized and one whose experts were all quantized both passed, folded
differently, and only the quantized one entered the writer's input-scale
all-reduce, out of step with its peer's ownership gather. An expert-parallel
export now gathers every rank's per-projection state before anything is folded,
and every rank refuses a projection whose experts are not all quantized or all
unquantized.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/export/unified_export_hf.py    | 46 +++++++++++++
 .../torch/export/test_unified_export_hf.py    | 64 +++++++++++++++++++
 2 files changed, 110 insertions(+)

diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py
index e3ff60c4bac..f259cbbedad 100644
--- a/modelopt/torch/export/unified_export_hf.py
+++ b/modelopt/torch/export/unified_export_hf.py
@@ -2111,6 +2111,9 @@ def export_hf_checkpoint(
                 f"MODELOPT_EP_DEGREE={plan.world_size} if that divides the expert count, or "
                 "MODELOPT_FSDP_SIZE=1 to replicate the dense part."
             )
+    if _ep_sharded:
+        # Each rank holds only its own experts' quantizers, so they are agreed on before any folds.
+        _agree_on_fused_expert_quantization(model)
     # Either form of sharding, including the hybrid where both hold: a model can be FSDP2-dense AND
     # expert-parallel, so these are tracked separately rather than as an either/or.
     _sharded_export = _ep_sharded or _fsdp2_sharded
@@ -2363,6 +2366,49 @@ def export_hf_checkpoint(
             _finish_export_together(error)
 
 
+def _agree_on_fused_expert_quantization(model: nn.Module) -> None:
+    """Refuse, on every rank, a fused-experts projection quantized for some experts and not others.
+
+    The no-gather export writes each projection's experts as one fused tensor, so they must be all
+    quantized or all not. Under expert parallelism a rank holds only its own experts' quantizers
+    and sees only its share: ranks that disagree would fold differently, and only the quantized
+    ones would join the writer's input-scale all-reduce. So every rank's view is gathered, and every
+    rank raises, before any of them folds. Collective.
+    """
+
+    def states() -> dict[str, str]:
+        found: dict[str, str] = {}
+        for name, module in model.named_modules():
+            if _hf_export_handlers._has_fused_experts_quantizers(module):
+                for proj in (getattr(module, "_first_proj_attr", "gate_up_proj"), "down_proj"):
+                    enabled = [q.is_enabled for q in getattr(module, f"{proj}_weight_quantizers")]
+                    found[f"{name}.{proj}"] = (
+                        "quantized" if all(enabled) else "mixed" if any(enabled) else "unquantized"
+                    )
+        return found
+
+    def agree(per_rank: list[dict[str, str]]) -> dict[str, str]:
+        # A projection's state where every rank reports the same one, else what each rank reported.
+        merged: dict[str, str] = {}
+        for key in sorted(set().union(*per_rank)):
+            seen = [found.get(key, "absent") for found in per_rank]
+            merged[key] = (
+                seen[0]
+                if len(set(seen)) == 1
+                else ", ".join(f"rank {r}: {state}" for r, state in enumerate(seen))
+            )
+        return merged
+
+    agreed = all_reduce("fused expert quantization", states, agree)
+    unsupported = {k: v for k, v in agreed.items() if v not in ("quantized", "unquantized")}
+    if unsupported:
+        raise NotImplementedError(
+            "The expert-parallel export writes each fused-experts projection as one tensor, so its "
+            "experts must be all quantized or all unquantized, on every rank: "
+            + "; ".join(f"{key} ({state})" for key, state in unsupported.items())
+        )
+
+
 def _process_together(
     model: nn.Module, dtype: torch.dtype | None, is_modelopt_qlora: bool
 ) -> dict[str, Any]:
diff --git a/tests/unit/torch/export/test_unified_export_hf.py b/tests/unit/torch/export/test_unified_export_hf.py
index cac21a9eb14..7a4c57be060 100644
--- a/tests/unit/torch/export/test_unified_export_hf.py
+++ b/tests/unit/torch/export/test_unified_export_hf.py
@@ -1230,6 +1230,70 @@ def test_every_rank_stops_before_the_write_when_one_ranks_processing_failed(
     cpu_dist_workers.run(_one_rank_fails_processing_job, tmp_dir=tmp_path)
 
 
+class _FusedExpertsQuantizers(torch.nn.Module):
+    """What the export reads of a fused-experts container: each projection's weight quantizers."""
+
+    def __init__(self, enabled):
+        super().__init__()
+        self.gate_up_proj_weight_quantizers = [SimpleNamespace(is_enabled=e) for e in enabled]
+        self.down_proj_weight_quantizers = [SimpleNamespace(is_enabled=e) for e in enabled]
+
+
+def _expert_quantization_job(rank, size, tmp_dir, enabled_per_rank, refused):
+    from unittest import mock
+
+    from modelopt.torch.export import distributed_writer, unified_export_hf
+    from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan
+
+    model = torch.nn.Module()
+    model.experts = _FusedExpertsQuantizers(enabled_per_rank[rank])
+    set_distribution_plan(model, DistributionPlan(size, ep_size=size))
+    raised = None
+    with (
+        mock.patch.object(
+            unified_export_hf, "_process_hf_model_for_export", return_value={}
+        ) as fold,
+        mock.patch.object(distributed_writer, "distributed_save_hf_checkpoint") as write,
+        mock.patch.object(unified_export_hf, "_write_base_config"),
+        mock.patch.object(unified_export_hf, "_write_hf_export_config"),
+    ):
+        # Not pytest.raises: when nothing is raised it fails with a BaseException, which would take
+        # the shared worker down instead of failing this test alone.
+        try:
+            unified_export_hf.export_hf_checkpoint(model, export_dir=tmp_dir)
+        except NotImplementedError as e:
+            raised = str(e)
+    if refused is None:
+        assert raised is None, raised
+    else:
+        assert raised is not None and f"experts.gate_up_proj ({refused})" in raised, raised
+    # Refused before any rank folds or writes; agreed states export as usual.
+    assert fold.called == write.called == (refused is None)
+
+
+@pytest.mark.parametrize(
+    ("enabled_per_rank", "refused"),
+    [
+        (((False, False), (True, True)), "rank 0: unquantized, rank 1: quantized"),
+        (((True, True), (True, False)), "rank 0: quantized, rank 1: mixed"),
+        (((True, True), (True, True)), None),
+        (((False, False), (False, False)), None),
+    ],
+    ids=["ranks-disagree", "one-rank-mixed", "all-quantized", "all-unquantized"],
+)
+def test_expert_parallel_ranks_agree_on_their_experts_quantization_before_folding(
+    cpu_dist_workers, tmp_path, enabled_per_rank, refused
+):
+    """Each rank sees only its own experts' quantizers. Ranks that disagree would fold differently,
+    and only the quantized ones would join the writer's input-scale all-reduce."""
+    cpu_dist_workers.run(
+        _expert_quantization_job,
+        tmp_dir=tmp_path,
+        enabled_per_rank=enabled_per_rank,
+        refused=refused,
+    )
+
+
 def test_a_single_process_export_ends_on_its_own_outcome():
     """Nothing to wait for: the export's own error, if any, is all that is raised."""
     from modelopt.torch.export.unified_export_hf import _finish_export_together

From 32e4ff98c64df64af5aad46fb3251e0f5e96b6f4 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 16 Oct 2026 00:28:29 +0000
Subject: [PATCH 17/38] Test the expert-quantization agreement on every-rank
 mixes and down_proj

The two-rank test gave both projections the same states and never had every rank
mixed, so an agreement that read only the first projection, or accepted a
projection mixed on every rank, still passed. Both cases are now refused before
processing, as the agreement intends.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 .../torch/export/test_unified_export_hf.py    | 48 ++++++++++++++-----
 1 file changed, 36 insertions(+), 12 deletions(-)

diff --git a/tests/unit/torch/export/test_unified_export_hf.py b/tests/unit/torch/export/test_unified_export_hf.py
index 7a4c57be060..28d1fd723ce 100644
--- a/tests/unit/torch/export/test_unified_export_hf.py
+++ b/tests/unit/torch/export/test_unified_export_hf.py
@@ -1233,20 +1233,22 @@ def test_every_rank_stops_before_the_write_when_one_ranks_processing_failed(
 class _FusedExpertsQuantizers(torch.nn.Module):
     """What the export reads of a fused-experts container: each projection's weight quantizers."""
 
-    def __init__(self, enabled):
+    def __init__(self, enabled, down_enabled=None):
         super().__init__()
+        down_enabled = enabled if down_enabled is None else down_enabled
         self.gate_up_proj_weight_quantizers = [SimpleNamespace(is_enabled=e) for e in enabled]
-        self.down_proj_weight_quantizers = [SimpleNamespace(is_enabled=e) for e in enabled]
+        self.down_proj_weight_quantizers = [SimpleNamespace(is_enabled=e) for e in down_enabled]
 
 
-def _expert_quantization_job(rank, size, tmp_dir, enabled_per_rank, refused):
+def _expert_quantization_job(rank, size, tmp_dir, enabled_per_rank, down_enabled_per_rank, refused):
     from unittest import mock
 
     from modelopt.torch.export import distributed_writer, unified_export_hf
     from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan
 
     model = torch.nn.Module()
-    model.experts = _FusedExpertsQuantizers(enabled_per_rank[rank])
+    down_enabled = None if down_enabled_per_rank is None else down_enabled_per_rank[rank]
+    model.experts = _FusedExpertsQuantizers(enabled_per_rank[rank], down_enabled)
     set_distribution_plan(model, DistributionPlan(size, ep_size=size))
     raised = None
     with (
@@ -1266,23 +1268,44 @@ def _expert_quantization_job(rank, size, tmp_dir, enabled_per_rank, refused):
     if refused is None:
         assert raised is None, raised
     else:
-        assert raised is not None and f"experts.gate_up_proj ({refused})" in raised, raised
+        assert raised is not None and refused in raised, raised
     # Refused before any rank folds or writes; agreed states export as usual.
     assert fold.called == write.called == (refused is None)
 
 
 @pytest.mark.parametrize(
-    ("enabled_per_rank", "refused"),
+    ("enabled_per_rank", "down_enabled_per_rank", "refused"),
     [
-        (((False, False), (True, True)), "rank 0: unquantized, rank 1: quantized"),
-        (((True, True), (True, False)), "rank 0: quantized, rank 1: mixed"),
-        (((True, True), (True, True)), None),
-        (((False, False), (False, False)), None),
+        (
+            ((False, False), (True, True)),
+            None,
+            "experts.gate_up_proj (rank 0: unquantized, rank 1: quantized)",
+        ),
+        (
+            ((True, True), (True, False)),
+            None,
+            "experts.gate_up_proj (rank 0: quantized, rank 1: mixed)",
+        ),
+        (((True, False), (False, True)), None, "experts.gate_up_proj (mixed)"),
+        (
+            ((True, True), (True, True)),
+            ((False, False), (True, True)),
+            "experts.down_proj (rank 0: unquantized, rank 1: quantized)",
+        ),
+        (((True, True), (True, True)), None, None),
+        (((False, False), (False, False)), None, None),
+    ],
+    ids=[
+        "ranks-disagree",
+        "one-rank-mixed",
+        "every-rank-mixed",
+        "down-proj-disagrees",
+        "all-quantized",
+        "all-unquantized",
     ],
-    ids=["ranks-disagree", "one-rank-mixed", "all-quantized", "all-unquantized"],
 )
 def test_expert_parallel_ranks_agree_on_their_experts_quantization_before_folding(
-    cpu_dist_workers, tmp_path, enabled_per_rank, refused
+    cpu_dist_workers, tmp_path, enabled_per_rank, down_enabled_per_rank, refused
 ):
     """Each rank sees only its own experts' quantizers. Ranks that disagree would fold differently,
     and only the quantized ones would join the writer's input-scale all-reduce."""
@@ -1290,6 +1313,7 @@ def test_expert_parallel_ranks_agree_on_their_experts_quantization_before_foldin
         _expert_quantization_job,
         tmp_dir=tmp_path,
         enabled_per_rank=enabled_per_rank,
+        down_enabled_per_rank=down_enabled_per_rank,
         refused=refused,
     )
 

From 080c998cf3808e9a9057fb6f3ade3f32dc926be3 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 16 Oct 2026 07:07:00 +0000
Subject: [PATCH 18/38] Test that a failed export fails every rank through the
 export itself

The tests of the closing consensus called _finish_export_together
directly, so they re-tested all_reduce: had export_hf_checkpoint gone
back to a plain barrier, every test would still have passed. One pooled
export now fails on rank 0 alone, after every rank wrote its shards, and
both ranks must raise with its error. Its collectives run on a group
that gives up in seconds, so a rank left waiting fails the test instead
of hanging it.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 .../torch/export/test_unified_export_hf.py    | 61 ++++++++++++-------
 1 file changed, 38 insertions(+), 23 deletions(-)

diff --git a/tests/unit/torch/export/test_unified_export_hf.py b/tests/unit/torch/export/test_unified_export_hf.py
index 28d1fd723ce..bf839a234a1 100644
--- a/tests/unit/torch/export/test_unified_export_hf.py
+++ b/tests/unit/torch/export/test_unified_export_hf.py
@@ -18,6 +18,7 @@
 import importlib.util
 import json
 from copy import deepcopy
+from datetime import timedelta
 from functools import partial
 from types import SimpleNamespace
 
@@ -1174,23 +1175,46 @@ def test_distributed_export_writes_the_config_a_single_gpu_export_writes(tmp_pat
     assert config.text_config.base_model_ep_plan is swapped_plan
 
 
-def _finish_export_job(failing_rank, rank, size):
-    from modelopt.torch.export.unified_export_hf import _finish_export_together
+def _one_rank_fails_after_the_write_job(rank, size, tmp_dir):
+    from unittest import mock
 
-    if failing_rank is None:
-        _finish_export_together(None)
-    elif rank == failing_rank:
-        with pytest.raises(OSError, match="disk full"):
-            _finish_export_together(OSError("disk full"))
-    else:
-        with pytest.raises(RuntimeError, match=rf"rank {failing_rank} failed.*disk full"):
-            _finish_export_together(None)
+    from modelopt.torch.export import distributed_writer, unified_export_hf
+    from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan
+
+    model = torch.nn.Linear(2, 2)
+    set_distribution_plan(model, DistributionPlan(size, ep_size=size))
+    # A rank left waiting in a collective its peer never joins gives up in seconds on this group.
+    group = torch.distributed.new_group(timeout=timedelta(seconds=10))
+    gather = partial(torch.distributed.all_gather_object, group=group)
+    raised = None
+    with (
+        mock.patch.object(torch.distributed, "all_gather_object", gather),
+        mock.patch.object(
+            torch.distributed, "barrier", partial(torch.distributed.barrier, group=group)
+        ),
+        mock.patch.object(unified_export_hf, "_process_hf_model_for_export", return_value={}),
+        mock.patch.object(distributed_writer, "distributed_save_hf_checkpoint") as write,
+        # Rank 0 alone writes config.json, once every rank has written its shards.
+        mock.patch.object(
+            unified_export_hf, "_write_base_config", side_effect=OSError("disk full")
+        ),
+    ):
+        # Not pytest.raises: when nothing is raised it fails with a BaseException, which would take
+        # the shared worker down instead of failing this test alone.
+        try:
+            unified_export_hf.export_hf_checkpoint(model, export_dir=tmp_dir)
+        except (OSError, RuntimeError) as e:
+            raised = e
+    assert write.called
+    assert raised is not None and "disk full" in str(raised), raised
 
 
-@pytest.mark.parametrize("failing_rank", [None, 0, 1])
-def test_every_rank_fails_an_export_that_any_rank_failed(cpu_dist_workers, failing_rank):
-    """A barrier alone lets the ranks that wrote their shards carry on as if the export succeeded."""
-    cpu_dist_workers.run(partial(_finish_export_job, failing_rank))
+def test_every_rank_fails_an_export_that_one_rank_failed_after_the_write(
+    cpu_dist_workers, tmp_path
+):
+    """The ranks that wrote their shards must not return as if the checkpoint were whole while rank 0
+    failed to finish it: a barrier alone would let them."""
+    cpu_dist_workers.run(_one_rank_fails_after_the_write_job, tmp_dir=tmp_path)
 
 
 def _one_rank_fails_processing_job(rank, size, tmp_dir):
@@ -1318,15 +1342,6 @@ def test_expert_parallel_ranks_agree_on_their_experts_quantization_before_foldin
     )
 
 
-def test_a_single_process_export_ends_on_its_own_outcome():
-    """Nothing to wait for: the export's own error, if any, is all that is raised."""
-    from modelopt.torch.export.unified_export_hf import _finish_export_together
-
-    _finish_export_together(None)
-    with pytest.raises(OSError, match="disk full"):
-        _finish_export_together(OSError("disk full"))
-
-
 @pytest.mark.parametrize(
     ("dense", "ep_size", "flags"),
     [

From b1e789afe8014f590b686747601a7151429828d5 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 16 Oct 2026 18:29:32 +0000
Subject: [PATCH 19/38] Trim the expert-parallel export's tests: merge the
 lookups, drop covered ones

Merged:
- test_a_dequantizing_fp8_load_still_reverses_the_other_conversions becomes
  the fp8-dequantize case of test_experts_revert_through_converters_loading_rewrote,
  as a real dequantizing load leaves both converters: the case appends the
  in-place Fp8Dequantize converter and checks that a bf16 weight is not
  re-encoded. Aborting the reverse on that converter, and on transformers 5.5
  not recognizing its identity reverse, still fail the case.
- test_a_model_without_expert_parallel_experts_has_no_expert_state folds into
  test_the_routed_experts_state_is_found_by_its_expert_parallel_group.

Dropped, each caught by a remaining test when its code is broken:
- test_process_hf_model_for_export_folds_in_place_and_returns_the_quant_config:
  returning no quant config, folding nothing, or never keeping experts fused
  fails test_nogather_export_writes_what_the_single_process_export_writes.
- test_the_export_takes_how_the_model_is_sharded_from_its_plan: dropping either
  flag, swapping them, or reading a pipeline's plan fails the after-the-write
  failure test, the replicated-dense refusal or
  test_export_diffusers_pipeline_config.
- The one-rank-mixed case of the expert-quantization agreement: every mutation
  of the agreement it caught, the ranks-disagree and every-rank-mixed cases
  catch too.

The agreement cases now name the refused projection instead of pinning each
rank's state in the message, and the processing-failure job checks that rank 0's
error names rank 1. Both jobs, and the after-the-write job in place of its own
group, run their collectives inside bounded_collectives, and the
processing-failure job uses try/except rather than pytest.raises.

Kept, since nothing else fails when their code is broken: the agreement's
down-proj-disagrees case (an agreement that reads only the first projection)
and its all-unquantized case (one that refuses an agreed unquantized
projection).

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 .../export/test_quant_aware_conversion.py     |  54 ++++------
 .../torch/export/test_unified_export_hf.py    | 100 ++++--------------
 tests/unit/torch/utils/test_distributed.py    |   9 +-
 3 files changed, 44 insertions(+), 119 deletions(-)

diff --git a/tests/unit/torch/export/test_quant_aware_conversion.py b/tests/unit/torch/export/test_quant_aware_conversion.py
index e9156e91998..791efd2a5dc 100644
--- a/tests/unit/torch/export/test_quant_aware_conversion.py
+++ b/tests/unit/torch/export/test_quant_aware_conversion.py
@@ -705,34 +705,6 @@ def test_build_reverse_rules_orders_prefix_reorder_after_container():
     assert not any(".mlp.experts." in k for k in out)
 
 
-def test_a_dequantizing_fp8_load_still_reverses_the_other_conversions():
-    """A FineGrainedFP8 ``dequantize=True`` load adds a converter that only decodes each weight in
-    place. It renames nothing, so it must neither abort the reverse of the model's real
-    conversions (MiniMax-M2.7 kept its in-memory expert names) nor re-encode a weight."""
-    pytest.importorskip("transformers.core_model_loading")
-    fp8 = pytest.importorskip("transformers.integrations.finegrained_fp8")
-    from transformers.core_model_loading import WeightConverter, WeightRenaming
-
-    conversions = [
-        WeightRenaming(".block_sparse_moe.experts.", ".mlp.experts."),
-        WeightConverter(
-            source_patterns=["weight$", "weight_scale_inv", "activation_scale"],
-            target_patterns="weight",
-            operations=[fp8.Fp8Dequantize(None)],
-        ),
-    ]
-    model = types.SimpleNamespace(_weight_conversions=conversions)
-    sd = _nvfp4_linear("model.layers.0.mlp.experts.0.w1", 8, 16)
-    sd["model.layers.0.self_attn.q_proj.weight"] = torch.randn(8, 16, dtype=torch.bfloat16)
-
-    out = revert_weight_conversion_quant_aware(model, sd)
-
-    expert = "model.layers.0.block_sparse_moe.experts.0.w1"
-    assert {f"{expert}.weight", f"{expert}.weight_scale", f"{expert}.weight_scale_2"} <= set(out)
-    assert out["model.layers.0.self_attn.q_proj.weight"].dtype == torch.bfloat16
-    assert not any(".mlp.experts." in k or k.endswith("weight_scale_inv") for k in out)
-
-
 def test_nested_text_prefix_reverse_does_not_capture_vlm_siblings():
     """A nested text-model conversion must not rewrite the full VLM namespace."""
     pytest.importorskip("transformers.core_model_loading")
@@ -1341,7 +1313,12 @@ def _as_loading_rewrote(conversions, siblings):
 def test_experts_revert_through_converters_loading_rewrote(siblings):
     """The export replays the conversions ``from_pretrained`` recorded. After a dequantizing or
     decompressing load, each expert converter lists every checkpoint tensor it consumed, so the
-    same leaf repeats; the expert reverse used to bail on that and leave every name in-memory."""
+    same leaf repeats; the expert reverse used to bail on that and leave every name in-memory.
+
+    A FineGrainedFP8 ``dequantize=True`` load also adds a converter that only decodes each weight
+    in place. It renames nothing, so it must neither abort the reverse of the model's real
+    conversions (MiniMax-M2.7 kept its in-memory expert names) nor re-encode a weight.
+    """
     pytest.importorskip("transformers")
     from _test_utils.torch.transformers_models import get_tiny_mixtral
 
@@ -1354,9 +1331,23 @@ def test_experts_revert_through_converters_loading_rewrote(siblings):
     if not any(".experts." in str(c.source_patterns) for c in stock):
         pytest.skip("transformers build has no mixtral expert conversion")
     model._weight_conversions = _as_loading_rewrote(stock, siblings)
+    if "weight_scale_inv" in siblings:
+        fp8 = pytest.importorskip("transformers.integrations.finegrained_fp8")
+        from transformers.core_model_loading import WeightConverter
+
+        model._weight_conversions.append(
+            WeightConverter(
+                source_patterns=["weight$", "weight_scale_inv", "activation_scale"],
+                target_patterns="weight",
+                operations=[fp8.Fp8Dequantize(None)],
+            )
+        )
 
     p = "model.layers.0"
-    sd = {f"{p}.mlp.gate.weight": torch.randn(2, 32)}
+    sd = {
+        f"{p}.mlp.gate.weight": torch.randn(2, 32),
+        f"{p}.self_attn.q_proj.weight": torch.randn(8, 16, dtype=torch.bfloat16),
+    }
     for e in range(2):
         sd.update(_nvfp4_linear(f"{p}.mlp.experts.{e}.gate_proj", 64, 32))
         sd.update(_nvfp4_linear(f"{p}.mlp.experts.{e}.up_proj", 64, 32))
@@ -1370,7 +1361,8 @@ def test_experts_revert_through_converters_loading_rewrote(siblings):
         assert f"{base}.w3.weight" in out
         assert f"{base}.w2.weight" in out
         assert f"{base}.w1.weight_scale_2" in out
-    assert not any(".mlp.experts." in k for k in out)
+    assert out[f"{p}.self_attn.q_proj.weight"].dtype == torch.bfloat16
+    assert not any(".mlp.experts." in k or k.endswith("weight_scale_inv") for k in out)
 
 
 def test_dense_split_names_each_part_once_through_load_siblings():
diff --git a/tests/unit/torch/export/test_unified_export_hf.py b/tests/unit/torch/export/test_unified_export_hf.py
index bf839a234a1..a364a6ba017 100644
--- a/tests/unit/torch/export/test_unified_export_hf.py
+++ b/tests/unit/torch/export/test_unified_export_hf.py
@@ -18,12 +18,11 @@
 import importlib.util
 import json
 from copy import deepcopy
-from datetime import timedelta
-from functools import partial
 from types import SimpleNamespace
 
 import pytest
 import torch
+from _test_utils.torch.distributed.utils import bounded_collectives
 from _test_utils.torch.quantization.tied_modules import (
     make_tied_linear_pair,
     wrap_in_parent_with_tied_keys,
@@ -1111,33 +1110,6 @@ def test_source_checkpoint_is_none_for_a_model_built_from_config(tmp_path, monke
     assert not any(export_dir.iterdir())
 
 
-def test_process_hf_model_for_export_folds_in_place_and_returns_the_quant_config():
-    """The processing half the no-gather writer runs before writing each rank's shard."""
-    pytest.importorskip("transformers")
-    from transformers import LlamaConfig, LlamaForCausalLM
-
-    from modelopt.torch.export.unified_export_hf import _process_hf_model_for_export
-
-    config = LlamaConfig(
-        vocab_size=64,
-        hidden_size=32,
-        intermediate_size=64,
-        num_hidden_layers=1,
-        num_attention_heads=4,
-        num_key_value_heads=2,
-    )
-    model = LlamaForCausalLM(config).eval()
-    mtq.quantize(model, mtq.FP8_DEFAULT_CFG, lambda m: m(torch.randint(0, 64, (1, 8))))
-
-    quant_config = _process_hf_model_for_export(model, torch.float16)
-
-    assert quant_config["quantization"]["quant_algo"] == "FP8"
-    # Folded in place: the model's own state dict is now the deployment one.
-    q_proj = model.model.layers[0].self_attn.q_proj
-    assert q_proj.weight.dtype == torch.float8_e4m3fn
-    assert "model.layers.0.self_attn.q_proj.weight_scale" in model.state_dict()
-
-
 @pytest.mark.parametrize("source_plan", [None, {"layers.*.feed_forward.experts": "grouped_gemm"}])
 def test_distributed_export_writes_the_config_a_single_gpu_export_writes(tmp_path, source_plan):
     """The EP load's DistributedConfig and torch_a2a plan stay out of config.json, the source's own
@@ -1183,15 +1155,9 @@ def _one_rank_fails_after_the_write_job(rank, size, tmp_dir):
 
     model = torch.nn.Linear(2, 2)
     set_distribution_plan(model, DistributionPlan(size, ep_size=size))
-    # A rank left waiting in a collective its peer never joins gives up in seconds on this group.
-    group = torch.distributed.new_group(timeout=timedelta(seconds=10))
-    gather = partial(torch.distributed.all_gather_object, group=group)
     raised = None
     with (
-        mock.patch.object(torch.distributed, "all_gather_object", gather),
-        mock.patch.object(
-            torch.distributed, "barrier", partial(torch.distributed.barrier, group=group)
-        ),
+        bounded_collectives(),
         mock.patch.object(unified_export_hf, "_process_hf_model_for_export", return_value={}),
         mock.patch.object(distributed_writer, "distributed_save_hf_checkpoint") as write,
         # Rank 0 alone writes config.json, once every rank has written its shards.
@@ -1232,17 +1198,23 @@ def process(*args):
             raise failure
         return {}
 
-    expected = (
-        pytest.raises(NotImplementedError, match="rank 1's experts")
-        if rank == 1
-        else pytest.raises(RuntimeError, match="rank 1 failed to compute export processed")
-    )
+    raised = None
     with (
+        bounded_collectives(),
         mock.patch.object(unified_export_hf, "_process_hf_model_for_export", process),
         mock.patch.object(distributed_writer, "distributed_save_hf_checkpoint") as write,
-        expected,
     ):
-        unified_export_hf.export_hf_checkpoint(model, export_dir=tmp_dir)
+        # Not pytest.raises: when nothing is raised it fails with a BaseException, which would take
+        # the shared worker down instead of failing this test alone.
+        try:
+            unified_export_hf.export_hf_checkpoint(model, export_dir=tmp_dir)
+        except (NotImplementedError, RuntimeError) as e:
+            raised = e
+    # Rank 1 raises its own error, and rank 0 one that names rank 1, before either writes.
+    if rank == 1:
+        assert raised is failure, raised
+    else:
+        assert isinstance(raised, RuntimeError) and "rank 1" in str(raised), raised
     write.assert_not_called()
 
 
@@ -1282,6 +1254,7 @@ def _expert_quantization_job(rank, size, tmp_dir, enabled_per_rank, down_enabled
         mock.patch.object(distributed_writer, "distributed_save_hf_checkpoint") as write,
         mock.patch.object(unified_export_hf, "_write_base_config"),
         mock.patch.object(unified_export_hf, "_write_hf_export_config"),
+        bounded_collectives(),
     ):
         # Not pytest.raises: when nothing is raised it fails with a BaseException, which would take
         # the shared worker down instead of failing this test alone.
@@ -1300,28 +1273,14 @@ def _expert_quantization_job(rank, size, tmp_dir, enabled_per_rank, down_enabled
 @pytest.mark.parametrize(
     ("enabled_per_rank", "down_enabled_per_rank", "refused"),
     [
-        (
-            ((False, False), (True, True)),
-            None,
-            "experts.gate_up_proj (rank 0: unquantized, rank 1: quantized)",
-        ),
-        (
-            ((True, True), (True, False)),
-            None,
-            "experts.gate_up_proj (rank 0: quantized, rank 1: mixed)",
-        ),
-        (((True, False), (False, True)), None, "experts.gate_up_proj (mixed)"),
-        (
-            ((True, True), (True, True)),
-            ((False, False), (True, True)),
-            "experts.down_proj (rank 0: unquantized, rank 1: quantized)",
-        ),
+        (((False, False), (True, True)), None, "experts.gate_up_proj"),
+        (((True, False), (False, True)), None, "experts.gate_up_proj"),
+        (((True, True), (True, True)), ((False, False), (True, True)), "experts.down_proj"),
         (((True, True), (True, True)), None, None),
         (((False, False), (False, False)), None, None),
     ],
     ids=[
         "ranks-disagree",
-        "one-rank-mixed",
         "every-rank-mixed",
         "down-proj-disagrees",
         "all-quantized",
@@ -1342,27 +1301,6 @@ def test_expert_parallel_ranks_agree_on_their_experts_quantization_before_foldin
     )
 
 
-@pytest.mark.parametrize(
-    ("dense", "ep_size", "flags"),
-    [
-        ("replicate", 1, (False, False)),
-        ("fsdp", 1, (True, False)),
-        ("replicate", 2, (False, True)),
-        ("fsdp", 2, (True, True)),
-    ],
-    ids=["replicated", "fsdp2", "expert-parallel", "both"],
-)
-def test_the_export_takes_how_the_model_is_sharded_from_its_plan(dense, ep_size, flags):
-    """A model can be FSDP2-dense and expert-parallel at once, so the two are read separately."""
-    from modelopt.torch.export.unified_export_hf import _hf_shard_flags
-    from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan
-
-    model = torch.nn.Linear(2, 2)
-    set_distribution_plan(model, DistributionPlan(2, ep_size, dense))
-
-    assert _hf_shard_flags(model) == flags
-
-
 def test_an_expert_parallel_plan_with_replicated_dense_shards_is_refused_before_folding(tmp_path):
     """Below the world size, EP leaves FSDP2's dense shards replicated across the expert-parallel
     groups, which the no-gather writer cannot write: refused before the export folds anything."""
diff --git a/tests/unit/torch/utils/test_distributed.py b/tests/unit/torch/utils/test_distributed.py
index 54811fb924f..7dddb8e2379 100644
--- a/tests/unit/torch/utils/test_distributed.py
+++ b/tests/unit/torch/utils/test_distributed.py
@@ -565,14 +565,9 @@ def test_the_routed_experts_state_is_found_by_its_expert_parallel_group():
     model = torch.nn.Module()
     model.attn = torch.nn.Linear(2, 2)
     model.experts = torch.nn.Module()
+    assert expert_parallel_state(model) is None
     for module in model.modules():
         module._parallel_state = _state(expert_parallel=False)
+    assert expert_parallel_state(model) is None
     model.experts._parallel_state = expert = _state(expert_parallel=True)
     assert expert_parallel_state(model) is expert
-
-
-def test_a_model_without_expert_parallel_experts_has_no_expert_state():
-    model = torch.nn.Sequential(torch.nn.Linear(2, 2))
-    assert expert_parallel_state(model) is None
-    model[0]._parallel_state = _state(expert_parallel=False)
-    assert expert_parallel_state(model) is None

From 95734a5c7a7b1f586304aada7ef3a371f92bfefb Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 16 Oct 2026 21:57:54 +0000
Subject: [PATCH 20/38] Import the dense split rule at the conversion tests'
 module scope

test_dense_split_names_each_part_once_through_load_siblings imported
_dense_split_rule inside the test, next to the file's module-level
quant_aware_conversion imports. Its transformers imports stay local, with the
note the file already uses: CI's transformers-free lane collects this file.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 tests/unit/torch/export/test_quant_aware_conversion.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/tests/unit/torch/export/test_quant_aware_conversion.py b/tests/unit/torch/export/test_quant_aware_conversion.py
index 791efd2a5dc..ae55f94f63c 100644
--- a/tests/unit/torch/export/test_quant_aware_conversion.py
+++ b/tests/unit/torch/export/test_quant_aware_conversion.py
@@ -40,6 +40,7 @@
     SplitRule,
     _assert_experts_pre_expanded,
     _build_reverse_rules,
+    _dense_split_rule,
     _source_checkpoint_keys,
     _SourceAwareRenamer,
     apply_reverse_rules,
@@ -1286,6 +1287,7 @@ def test_revert_quant_config_names_mapper():
 def _as_loading_rewrote(conversions, siblings):
     """``conversions`` with every expert converter's ``.weight`` sources expanded the way a
     load-time rewrite expands them: one anchored source per checkpoint tensor it consumes."""
+    # Local import: optional dependency; every caller importorskips transformers first.
     from transformers.core_model_loading import WeightConverter
 
     rewritten = []
@@ -1320,6 +1322,7 @@ def test_experts_revert_through_converters_loading_rewrote(siblings):
     conversions (MiniMax-M2.7 kept its in-memory expert names) nor re-encode a weight.
     """
     pytest.importorskip("transformers")
+    # Local import: optional dependency, guarded by the importorskip above.
     from _test_utils.torch.transformers_models import get_tiny_mixtral
 
     try:
@@ -1369,10 +1372,9 @@ def test_dense_split_names_each_part_once_through_load_siblings():
     """A dense fused converter that a dequantizing load rewrote still un-fuses into one part per
     module, not one per checkpoint tensor the load consumed."""
     pytest.importorskip("transformers")
+    # Local import: optional dependency, guarded by the importorskip above.
     from transformers.core_model_loading import Concatenate, WeightConverter
 
-    from modelopt.torch.export.quant_aware_conversion import _dense_split_rule
-
     modules = ("self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj")
     leaves = ("weight", "weight_scale_inv")
     conv = WeightConverter(

From 670002e991c218fdb4ec8165143917e0f525c27c Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 16 Oct 2026 21:58:03 +0000
Subject: [PATCH 21/38] Import the expert-parallel export tests' ModelOpt
 modules at module scope

The expert-parallel export tests imported mock, the distributed writer,
unified_export_hf and the distribution-plan helpers inside each test. Move
those to the module's import block; the transformers imports stay local, with
the file's existing note, since the transformers-free lane collects it.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 .../torch/export/test_unified_export_hf.py    | 26 +++++--------------
 1 file changed, 6 insertions(+), 20 deletions(-)

diff --git a/tests/unit/torch/export/test_unified_export_hf.py b/tests/unit/torch/export/test_unified_export_hf.py
index a364a6ba017..d0b408ad722 100644
--- a/tests/unit/torch/export/test_unified_export_hf.py
+++ b/tests/unit/torch/export/test_unified_export_hf.py
@@ -19,6 +19,7 @@
 import json
 from copy import deepcopy
 from types import SimpleNamespace
+from unittest import mock
 
 import pytest
 import torch
@@ -30,6 +31,7 @@
 from safetensors.torch import save_file
 
 import modelopt.torch.quantization as mtq
+from modelopt.torch.export import distributed_writer, unified_export_hf
 from modelopt.torch.export.model_utils import (
     TiedWeightMap,
     get_language_model_from_vl,
@@ -45,10 +47,13 @@
     _fuse_shared_input_modules,
     _resolve_export_dtype,
     _source_checkpoint,
+    _write_base_config,
     collect_shared_input_modules,
+    export_hf_checkpoint,
     read_unplaced_weights,
 )
 from modelopt.torch.quantization.nn import TensorQuantizer
+from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan
 
 
 def test_multimodal_detection_accepts_null_architectures():
@@ -1115,11 +1120,10 @@ def test_distributed_export_writes_the_config_a_single_gpu_export_writes(tmp_pat
     """The EP load's DistributedConfig and torch_a2a plan stay out of config.json, the source's own
     plan survives, and the live configs keep the run's values afterwards."""
     pytest.importorskip("transformers")
+    # Local import: optional dependency, guarded by the importorskip above.
     from transformers import AutoConfig, Llama4Config
     from transformers.distributed import DistributedConfig
 
-    from modelopt.torch.export.unified_export_hf import _write_base_config
-
     source = tmp_path / "source"
     Llama4Config().save_pretrained(source)
     if source_plan is not None:
@@ -1148,11 +1152,6 @@ def test_distributed_export_writes_the_config_a_single_gpu_export_writes(tmp_pat
 
 
 def _one_rank_fails_after_the_write_job(rank, size, tmp_dir):
-    from unittest import mock
-
-    from modelopt.torch.export import distributed_writer, unified_export_hf
-    from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan
-
     model = torch.nn.Linear(2, 2)
     set_distribution_plan(model, DistributionPlan(size, ep_size=size))
     raised = None
@@ -1184,11 +1183,6 @@ def test_every_rank_fails_an_export_that_one_rank_failed_after_the_write(
 
 
 def _one_rank_fails_processing_job(rank, size, tmp_dir):
-    from unittest import mock
-
-    from modelopt.torch.export import distributed_writer, unified_export_hf
-    from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan
-
     model = torch.nn.Linear(2, 2)
     set_distribution_plan(model, DistributionPlan(size, ep_size=size))
     failure = NotImplementedError("rank 1's experts mix quantized and unquantized ones")
@@ -1237,11 +1231,6 @@ def __init__(self, enabled, down_enabled=None):
 
 
 def _expert_quantization_job(rank, size, tmp_dir, enabled_per_rank, down_enabled_per_rank, refused):
-    from unittest import mock
-
-    from modelopt.torch.export import distributed_writer, unified_export_hf
-    from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan
-
     model = torch.nn.Module()
     down_enabled = None if down_enabled_per_rank is None else down_enabled_per_rank[rank]
     model.experts = _FusedExpertsQuantizers(enabled_per_rank[rank], down_enabled)
@@ -1304,9 +1293,6 @@ def test_expert_parallel_ranks_agree_on_their_experts_quantization_before_foldin
 def test_an_expert_parallel_plan_with_replicated_dense_shards_is_refused_before_folding(tmp_path):
     """Below the world size, EP leaves FSDP2's dense shards replicated across the expert-parallel
     groups, which the no-gather writer cannot write: refused before the export folds anything."""
-    from modelopt.torch.export.unified_export_hf import export_hf_checkpoint
-    from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan
-
     model = torch.nn.Sequential(torch.nn.Linear(16, 16))
     mtq.quantize(model, mtq.FP8_DEFAULT_CFG, lambda m: m(torch.randn(2, 16)))
     set_distribution_plan(model, DistributionPlan(4, ep_size=2, dense="fsdp"))

From d6fd0dd988323dacddbeb642806c16c46d2048b7 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Tue, 27 Oct 2026 01:24:08 +0000
Subject: [PATCH 22/38] Skip a fused-experts projection without weight
 quantizers when the ranks agree

The expert-parallel export's agreement step read down_proj_weight_quantizers unguarded, while every
other walk of the pair treats the second projection as optional: a layout that registers only the
first projection's quantizers raised AttributeError inside the collective, on every rank.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/export/unified_export_hf.py        | 6 +++++-
 tests/unit/torch/export/test_unified_export_hf.py | 5 ++++-
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py
index f259cbbedad..8a0c77773b6 100644
--- a/modelopt/torch/export/unified_export_hf.py
+++ b/modelopt/torch/export/unified_export_hf.py
@@ -2381,7 +2381,11 @@ def states() -> dict[str, str]:
         for name, module in model.named_modules():
             if _hf_export_handlers._has_fused_experts_quantizers(module):
                 for proj in (getattr(module, "_first_proj_attr", "gate_up_proj"), "down_proj"):
-                    enabled = [q.is_enabled for q in getattr(module, f"{proj}_weight_quantizers")]
+                    # As iter_weights_for_calibration does: a layout may register only the first.
+                    quantizers = getattr(module, f"{proj}_weight_quantizers", None)
+                    if quantizers is None:
+                        continue
+                    enabled = [q.is_enabled for q in quantizers]
                     found[f"{name}.{proj}"] = (
                         "quantized" if all(enabled) else "mixed" if any(enabled) else "unquantized"
                     )
diff --git a/tests/unit/torch/export/test_unified_export_hf.py b/tests/unit/torch/export/test_unified_export_hf.py
index d0b408ad722..18aeeddb53a 100644
--- a/tests/unit/torch/export/test_unified_export_hf.py
+++ b/tests/unit/torch/export/test_unified_export_hf.py
@@ -1227,7 +1227,8 @@ def __init__(self, enabled, down_enabled=None):
         super().__init__()
         down_enabled = enabled if down_enabled is None else down_enabled
         self.gate_up_proj_weight_quantizers = [SimpleNamespace(is_enabled=e) for e in enabled]
-        self.down_proj_weight_quantizers = [SimpleNamespace(is_enabled=e) for e in down_enabled]
+        if down_enabled != "absent":  # a layout may register only the first projection's
+            self.down_proj_weight_quantizers = [SimpleNamespace(is_enabled=e) for e in down_enabled]
 
 
 def _expert_quantization_job(rank, size, tmp_dir, enabled_per_rank, down_enabled_per_rank, refused):
@@ -1265,6 +1266,7 @@ def _expert_quantization_job(rank, size, tmp_dir, enabled_per_rank, down_enabled
         (((False, False), (True, True)), None, "experts.gate_up_proj"),
         (((True, False), (False, True)), None, "experts.gate_up_proj"),
         (((True, True), (True, True)), ((False, False), (True, True)), "experts.down_proj"),
+        (((True, True), (True, True)), ("absent", "absent"), None),
         (((True, True), (True, True)), None, None),
         (((False, False), (False, False)), None, None),
     ],
@@ -1272,6 +1274,7 @@ def _expert_quantization_job(rank, size, tmp_dir, enabled_per_rank, down_enabled
         "ranks-disagree",
         "every-rank-mixed",
         "down-proj-disagrees",
+        "no-down-proj-quantizers",
         "all-quantized",
         "all-unquantized",
     ],

From 041c2e99db306b04cef721110f39424de6406caf Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Tue, 27 Oct 2026 01:24:08 +0000
Subject: [PATCH 23/38] Fold the no-gather export's experts fused without
 reading the plan again

Its only caller is the no-gather writer, on a sharded model, which has already decided the route; a
second read of the plan that disagreed would have folded the experts split while the writer
expected them fused.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/export/unified_export_hf.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py
index 8a0c77773b6..d7118c2826c 100644
--- a/modelopt/torch/export/unified_export_hf.py
+++ b/modelopt/torch/export/unified_export_hf.py
@@ -1114,8 +1114,9 @@ def _process_hf_model_for_export(
     writer splits them after consolidation.
     """
     dtype, _, quant_config = _prepare_model_for_export(model, dtype, is_modelopt_qlora)
-    keep_fused_experts = any(_hf_shard_flags(model))
-    pack_quantized_weights(model, dtype, is_modelopt_qlora, keep_fused_experts=keep_fused_experts)
+    # Only the no-gather writer calls this, and only for a sharded model, so the experts stay fused
+    # without reading the plan a second time.
+    pack_quantized_weights(model, dtype, is_modelopt_qlora, keep_fused_experts=True)
     return quant_config
 
 

From cd45930ed2932e3703b16d02aed425030c25f52a Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Tue, 27 Oct 2026 01:24:08 +0000
Subject: [PATCH 24/38] Cache the transformers ops that re-encode a weight's
 quantization in place

They are a constant of the installed transformers, and were looked up again for every converter.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/export/quant_aware_conversion.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/modelopt/torch/export/quant_aware_conversion.py b/modelopt/torch/export/quant_aware_conversion.py
index e6dcd965bd7..914d938d19f 100644
--- a/modelopt/torch/export/quant_aware_conversion.py
+++ b/modelopt/torch/export/quant_aware_conversion.py
@@ -62,6 +62,7 @@
 (in-memory-name) behavior rather than emit a silently-wrong checkpoint.
 """
 
+import functools
 import logging
 import re
 from collections.abc import Iterable
@@ -757,6 +758,7 @@ def _drop_shadowed_prefix_renames(model, rules: list[RenameRule]) -> list[Rename
     return kept
 
 
+@functools.cache
 def _quantization_codec_ops() -> tuple[type, ...]:
     """Transformers' ops that decode or encode a weight's on-disk quantization in place.
 

From eba585a5d99e46cbe338791f99bc42242e75c425 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Tue, 27 Oct 2026 01:49:04 +0000
Subject: [PATCH 25/38] Refuse to save the ModelOpt state from an
 expert-parallel export

The no-gather export now saves the ModelOpt state from rank 0, as the streaming export does. An
expert-parallel model always takes that writer, and there rank 0 holds only its own experts' state,
so the save would write an incomplete one. The export refuses it up front, before any collective.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/export/unified_export_hf.py        |  7 +++++++
 tests/unit/torch/export/test_unified_export_hf.py | 13 +++++++++++++
 2 files changed, 20 insertions(+)

diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py
index d7118c2826c..45707d36e52 100644
--- a/modelopt/torch/export/unified_export_hf.py
+++ b/modelopt/torch/export/unified_export_hf.py
@@ -2098,6 +2098,13 @@ def export_hf_checkpoint(
     # because the carry-over gate needs it and runs ahead of them (the layerwise and diffusers paths
     # return before those are reached).
     _fsdp2_sharded, _ep_sharded = _hf_shard_flags(model)
+    # Under expert parallelism each rank holds only its own experts' ModelOpt state, so the
+    # rank-0 save would write an incomplete one. Refused before any collective, alike on every rank.
+    if _ep_sharded and save_modelopt_state:
+        raise NotImplementedError(
+            "save_modelopt_state=True is not supported for an expert-parallel export: each rank "
+            "holds only its own experts' ModelOpt state."
+        )
     if _fsdp2_sharded and _ep_sharded:
         plan = distribution_plan(model)
         # With fewer expert-parallel ranks than the world, FSDP2 shards the dense part within each
diff --git a/tests/unit/torch/export/test_unified_export_hf.py b/tests/unit/torch/export/test_unified_export_hf.py
index 18aeeddb53a..c35a9aec2ad 100644
--- a/tests/unit/torch/export/test_unified_export_hf.py
+++ b/tests/unit/torch/export/test_unified_export_hf.py
@@ -1293,6 +1293,19 @@ def test_expert_parallel_ranks_agree_on_their_experts_quantization_before_foldin
     )
 
 
+def test_an_expert_parallel_export_refuses_to_save_the_modelopt_state(tmp_path):
+    """Each expert-parallel rank holds only its own experts' ModelOpt state, so saving rank 0's
+    would write an incomplete one: refused before anything is written."""
+    model = torch.nn.Sequential(torch.nn.Linear(16, 16))
+    mtq.quantize(model, mtq.FP8_DEFAULT_CFG, lambda m: m(torch.randn(2, 16)))
+    set_distribution_plan(model, DistributionPlan(2, ep_size=2))
+
+    with pytest.raises(NotImplementedError, match="expert-parallel export"):
+        export_hf_checkpoint(model, export_dir=tmp_path / "export", save_modelopt_state=True)
+
+    assert not (tmp_path / "export").exists()
+
+
 def test_an_expert_parallel_plan_with_replicated_dense_shards_is_refused_before_folding(tmp_path):
     """Below the world size, EP leaves FSDP2's dense shards replicated across the expert-parallel
     groups, which the no-gather writer cannot write: refused before the export folds anything."""

From 57c76b15b3d9163822c4859869a5d23a60b851a6 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Tue, 27 Oct 2026 18:57:52 +0000
Subject: [PATCH 26/38] Look up the identity decode apart from transformers'
 FP8 integration

transformers 5.5 imports triton atop its FP8 integration, so without
triton that import fails, and with it the shared import of
core_model_loading: the identity op 5.5 and 5.6 reverse Fp8Dequantize
into was lost too, and the reverse of a dequantizing load aborted, the
MiniMax-M2.7 failure. Each source of codec ops is now imported on its
own.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 .../torch/export/quant_aware_conversion.py    | 19 ++++++++++++-------
 .../export/test_quant_aware_conversion.py     | 19 +++++++++++++++++++
 2 files changed, 31 insertions(+), 7 deletions(-)

diff --git a/modelopt/torch/export/quant_aware_conversion.py b/modelopt/torch/export/quant_aware_conversion.py
index 914d938d19f..8d94eee3197 100644
--- a/modelopt/torch/export/quant_aware_conversion.py
+++ b/modelopt/torch/export/quant_aware_conversion.py
@@ -62,6 +62,7 @@
 (in-memory-name) behavior rather than emit a silently-wrong checkpoint.
 """
 
+import contextlib
 import functools
 import logging
 import re
@@ -765,14 +766,18 @@ def _quantization_codec_ops() -> tuple[type, ...]:
     Includes the identity op, which is what older releases (5.5, 5.6) reverse ``Fp8Dequantize``
     into; later ones reverse it into ``Fp8Quantize``.
     """
-    try:
-        from transformers import core_model_loading
+    ops: list[Any] = []
+    # Imported apart: 5.5 imports triton atop its FP8 integration, so without triton that import
+    # fails, while the identity op those releases reverse into lives in core_model_loading.
+    with contextlib.suppress(Exception):  # transformers without the FP8 integration
         from transformers.integrations import finegrained_fp8
-    except Exception:  # transformers without the FP8 integration
-        return ()
-    names = ("Fp8Dequantize", "Fp8Quantize", "Fp8DecodeScale", "Fp8EncodeScale")
-    ops = [getattr(finegrained_fp8, n, None) for n in names]
-    ops.append(getattr(core_model_loading, "_IdentityOp", None))
+
+        names = ("Fp8Dequantize", "Fp8Quantize", "Fp8DecodeScale", "Fp8EncodeScale")
+        ops += [getattr(finegrained_fp8, n, None) for n in names]
+    with contextlib.suppress(Exception):
+        from transformers import core_model_loading
+
+        ops.append(getattr(core_model_loading, "_IdentityOp", None))
     return tuple(op for op in ops if op is not None)
 
 
diff --git a/tests/unit/torch/export/test_quant_aware_conversion.py b/tests/unit/torch/export/test_quant_aware_conversion.py
index ae55f94f63c..e32a709c642 100644
--- a/tests/unit/torch/export/test_quant_aware_conversion.py
+++ b/tests/unit/torch/export/test_quant_aware_conversion.py
@@ -23,6 +23,7 @@
 """
 
 import json
+import sys
 import types
 import warnings
 from fnmatch import fnmatchcase
@@ -1368,6 +1369,24 @@ def test_experts_revert_through_converters_loading_rewrote(siblings):
     assert not any(".mlp.experts." in k or k.endswith("weight_scale_inv") for k in out)
 
 
+def test_the_identity_decode_is_known_without_the_fp8_integration(monkeypatch):
+    """transformers 5.5 imports triton atop its FP8 integration, so without triton that import
+    fails. The identity op 5.5 and 5.6 reverse ``Fp8Dequantize`` into is ``core_model_loading``'s,
+    and must still read as a decode in place, or the whole reverse of a dequantizing load aborts."""
+    core = pytest.importorskip("transformers.core_model_loading")
+    if not hasattr(core, "_IdentityOp"):
+        pytest.skip("transformers build has no _IdentityOp")
+    integrations = pytest.importorskip("transformers.integrations")
+    monkeypatch.delattr(integrations, "finegrained_fp8", raising=False)
+    monkeypatch.setitem(sys.modules, "transformers.integrations.finegrained_fp8", None)
+    quant_aware_conversion._quantization_codec_ops.cache_clear()
+    try:
+        rev = types.SimpleNamespace(source_patterns=["weight$"], target_patterns="weight")
+        assert quant_aware_conversion._decodes_in_place(rev, [core._IdentityOp()])
+    finally:
+        quant_aware_conversion._quantization_codec_ops.cache_clear()
+
+
 def test_dense_split_names_each_part_once_through_load_siblings():
     """A dense fused converter that a dequantizing load rewrote still un-fuses into one part per
     module, not one per checkpoint tensor the load consumed."""

From a54c18de7b3d9b1947eaa9b672d99868cc682a40 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Tue, 27 Oct 2026 19:50:42 +0000
Subject: [PATCH 27/38] Fail every rank of a sharded export that one rank
 exited or was interrupted in

The closing consensus learned of a rank's failure only through
`except Exception`, so a SystemExit or KeyboardInterrupt on one rank
reached it as a success: the other ranks returned as if the checkpoint
were whole while that rank's export had stopped. The export now hands
any BaseException to the consensus, which, as all_reduce does, exits
every rank with an exit's code and fails them on anything else.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/export/unified_export_hf.py    |  9 +++++++--
 .../torch/export/test_unified_export_hf.py    | 20 ++++++++++---------
 2 files changed, 18 insertions(+), 11 deletions(-)

diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py
index 45707d36e52..f32e436b3fb 100644
--- a/modelopt/torch/export/unified_export_hf.py
+++ b/modelopt/torch/export/unified_export_hf.py
@@ -2202,7 +2202,7 @@ def export_hf_checkpoint(
     # buffer instead of the whole quantized state dict.
     _offloaded = has_accelerate_offload(model)
 
-    error: Exception | None = None
+    error: BaseException | None = None
     try:
         if _offloaded:
             # Imported here rather than at module scope: the streaming exporter imports the
@@ -2367,6 +2367,11 @@ def export_hf_checkpoint(
             " can be saved with torch.save for further inspection."
         )
         raise e
+    except BaseException as e:
+        # An exit or an interrupt on one rank must fail the others too, or they return as if the
+        # checkpoint were whole.
+        error = e
+        raise
     finally:
         # Every rank waits for the writer to finish. Without this the ranks that wrote nothing exit
         # first and the job hangs on the ones still writing.
@@ -2440,7 +2445,7 @@ def processed() -> bool:
     return quant_config[0]
 
 
-def _finish_export_together(error: Exception | None) -> None:
+def _finish_export_together(error: BaseException | None) -> None:
     """Wait for every rank to finish writing, and raise on all of them if any failed.
 
     A barrier alone lets the ranks whose writes succeeded return, and the run carries on as if the
diff --git a/tests/unit/torch/export/test_unified_export_hf.py b/tests/unit/torch/export/test_unified_export_hf.py
index c35a9aec2ad..a095fd006c2 100644
--- a/tests/unit/torch/export/test_unified_export_hf.py
+++ b/tests/unit/torch/export/test_unified_export_hf.py
@@ -1151,7 +1151,7 @@ def test_distributed_export_writes_the_config_a_single_gpu_export_writes(tmp_pat
     assert config.text_config.base_model_ep_plan is swapped_plan
 
 
-def _one_rank_fails_after_the_write_job(rank, size, tmp_dir):
+def _one_rank_fails_after_the_write_job(rank, size, tmp_dir, failure):
     model = torch.nn.Linear(2, 2)
     set_distribution_plan(model, DistributionPlan(size, ep_size=size))
     raised = None
@@ -1160,26 +1160,28 @@ def _one_rank_fails_after_the_write_job(rank, size, tmp_dir):
         mock.patch.object(unified_export_hf, "_process_hf_model_for_export", return_value={}),
         mock.patch.object(distributed_writer, "distributed_save_hf_checkpoint") as write,
         # Rank 0 alone writes config.json, once every rank has written its shards.
-        mock.patch.object(
-            unified_export_hf, "_write_base_config", side_effect=OSError("disk full")
-        ),
+        mock.patch.object(unified_export_hf, "_write_base_config", side_effect=failure),
     ):
         # Not pytest.raises: when nothing is raised it fails with a BaseException, which would take
         # the shared worker down instead of failing this test alone.
         try:
             unified_export_hf.export_hf_checkpoint(model, export_dir=tmp_dir)
-        except (OSError, RuntimeError) as e:
+        except (OSError, RuntimeError, SystemExit) as e:
             raised = e
     assert write.called
-    assert raised is not None and "disk full" in str(raised), raised
+    if isinstance(failure, SystemExit):  # an exit exits every rank, with its code
+        assert isinstance(raised, SystemExit) and raised.code == failure.code, raised
+    else:
+        assert raised is not None and "disk full" in str(raised), raised
 
 
+@pytest.mark.parametrize("failure", [OSError("disk full"), SystemExit(3)], ids=["error", "exit"])
 def test_every_rank_fails_an_export_that_one_rank_failed_after_the_write(
-    cpu_dist_workers, tmp_path
+    cpu_dist_workers, tmp_path, failure
 ):
     """The ranks that wrote their shards must not return as if the checkpoint were whole while rank 0
-    failed to finish it: a barrier alone would let them."""
-    cpu_dist_workers.run(_one_rank_fails_after_the_write_job, tmp_dir=tmp_path)
+    failed to finish it: a barrier alone would let them, and so would an exit, which is no Exception."""
+    cpu_dist_workers.run(_one_rank_fails_after_the_write_job, tmp_dir=tmp_path, failure=failure)
 
 
 def _one_rank_fails_processing_job(rank, size, tmp_dir):

From c96a024aa40e2e88e6442b27799937402be60a01 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Thu, 29 Oct 2026 21:57:32 +0000
Subject: [PATCH 28/38] Name no loader the expert-parallel export does not
 build on

The writer's refusal of an expert-parallel model with no routed expert's
ParallelState, and expert_parallel_state's docstring, named
load_and_prepare_ep, which is not below this change. They now say what
such a load does instead: it seeds each routed expert's ParallelState.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/export/distributed_writer.py | 3 ++-
 modelopt/torch/utils/distributed.py         | 4 ++--
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/modelopt/torch/export/distributed_writer.py b/modelopt/torch/export/distributed_writer.py
index c6e20cd183a..824b17ff458 100644
--- a/modelopt/torch/export/distributed_writer.py
+++ b/modelopt/torch/export/distributed_writer.py
@@ -737,7 +737,8 @@ def _distributed_save_hf_checkpoint_impl(
         if parallel_state is None:
             raise RuntimeError(
                 "The model's plan is expert-parallel, but no routed expert carries its expert-parallel "
-                "group: load it with load_and_prepare_ep."
+                "group: load it with an expert-parallel loader, which seeds each routed expert's "
+                "ParallelState."
             )
         ep_rank = parallel_state.expert_model_parallel_group.rank()
         edp_group = parallel_state.data_parallel_group
diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py
index 3ccc6022ff1..bb3dca0b260 100644
--- a/modelopt/torch/utils/distributed.py
+++ b/modelopt/torch/utils/distributed.py
@@ -694,8 +694,8 @@ def distributed_run(
 def expert_parallel_state(model: torch.nn.Module) -> "ParallelState | None":
     """Return the ``ParallelState`` of the model's routed experts, or ``None`` if none has one.
 
-    The expert-parallel loader seeds every routed-expert module with a ``ParallelState`` that carries the
-    expert-model-parallel group (``load_and_prepare_ep``); no other module of a HuggingFace model has one.
+    An expert-parallel load seeds every routed-expert module with a ``ParallelState`` that carries the
+    expert-model-parallel group; no other module of a HuggingFace model has one.
     """
     for module in model.modules():
         state = getattr(module, "_parallel_state", None)

From 5f62b529710b93a66e5da7b6ee607ec32cc20d02 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Thu, 29 Oct 2026 23:38:13 +0000
Subject: [PATCH 29/38] Count the experts with moe's helper, and import the
 writer's EP names at the top

The writer read the expert count off the config with its own lookup,
which missed num_routed_experts and the llm_config a NemotronH-Omni
config nests its text model under.
modelopt.torch.utils.moe._num_experts_from_config reads every spelling
and both nestings, and returns 0 for a model that is not an MoE, as the
lookup did, so the writer uses it.

distribution_plan and expert_parallel_state were imported inside the
function. Nothing forces that: the writer already imports
modelopt.torch.utils.distributed at module scope, and moe.py imports only
torch. Both now come from the top.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/export/distributed_writer.py | 21 +++------------------
 1 file changed, 3 insertions(+), 18 deletions(-)

diff --git a/modelopt/torch/export/distributed_writer.py b/modelopt/torch/export/distributed_writer.py
index 824b17ff458..9915edc7d0f 100644
--- a/modelopt/torch/export/distributed_writer.py
+++ b/modelopt/torch/export/distributed_writer.py
@@ -36,7 +36,8 @@
 from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict
 from torch.distributed.tensor import DTensor
 
-from modelopt.torch.utils.distributed import all_reduce
+from modelopt.torch.utils.distributed import all_reduce, distribution_plan, expert_parallel_state
+from modelopt.torch.utils.moe import _num_experts_from_config
 
 from .model_utils import TiedWeightMap
 from .moe_utils import _FUSED_PROJ, _fused_experts_prefixes, _split_local_fused_module
@@ -729,8 +730,6 @@ def _distributed_save_hf_checkpoint_impl(
     # other and the checkpoint would be quietly wrong rather than broken. The coordinate comes from
     # the routed experts' ParallelState -- the same source calibration reads, not transformers'
     # internal ``model._device_mesh``. An expert's data-parallel group is its expert-data-parallel one.
-    from modelopt.torch.utils.distributed import distribution_plan, expert_parallel_state
-
     ep_rank, dp_idx = 0, 0
     if distribution_plan(model).is_expert_parallel:
         parallel_state = expert_parallel_state(model)
@@ -747,21 +746,7 @@ def _distributed_save_hf_checkpoint_impl(
             dp_idx = edp_group.rank()
     # Total config expert count, to tell a per-EP-group expert tensor (size < total) from one that
     # spans all experts (classic DP x EP) when computing the per-expert global index offset.
-    _cfg = getattr(model, "config", None)
-
-    def _experts_of(c):
-        if c is None:
-            return 0
-        return int(
-            getattr(c, "num_experts", 0)
-            or getattr(c, "num_local_experts", 0)
-            or getattr(c, "n_routed_experts", 0)  # nemotron_h
-            or 0
-        )
-
-    # VLMs nest the text-model fields under ``config.text_config`` (the quantized backbone is
-    # ``model.language_model``), so fall back to it when the top-level config has no expert count.
-    total_experts = _experts_of(_cfg) or _experts_of(getattr(_cfg, "text_config", None))
+    total_experts = _num_experts_from_config(getattr(model, "config", None))
     # DP x EP: experts are dp-REPLICATED across the EDP group (each holds a full EP-sharded copy). Only
     # the first replica (dp_idx == 0) writes them; replica ranks (dp_idx > 0) skip the per-expert split.
     # dp_idx is the EDP-group rank resolved above -- 0 when experts are not EDP-replicated (dense /

From 092c784043ccabb8c06813c13c15a3f10c3f7788 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Thu, 29 Oct 2026 23:38:23 +0000
Subject: [PATCH 30/38] Drop the can_generate stub the config test no longer
 needs

_write_base_config asked the model whether it can generate before it
wrote generation_config.json. It now writes the non-weight files through
save_non_weight_artifacts, which does not ask, so the test's stub model
drops the answer.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 tests/unit/torch/export/test_unified_export_hf.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/tests/unit/torch/export/test_unified_export_hf.py b/tests/unit/torch/export/test_unified_export_hf.py
index a095fd006c2..472e3e747cb 100644
--- a/tests/unit/torch/export/test_unified_export_hf.py
+++ b/tests/unit/torch/export/test_unified_export_hf.py
@@ -1138,9 +1138,7 @@ def test_distributed_export_writes_the_config_a_single_gpu_export_writes(tmp_pat
     swapped_plan = {"layers.*.feed_forward.experts": "torch_a2a_experts"}
     config.distributed_config = distributed_config
     config.text_config.base_model_ep_plan = swapped_plan
-    model = SimpleNamespace(
-        config=config, _modelopt_source_checkpoint=str(source), can_generate=lambda: False
-    )
+    model = SimpleNamespace(config=config, _modelopt_source_checkpoint=str(source))
 
     _write_base_config(model, tmp_path / "export")
 

From 423c0a0b51b8142f9148187b2d9019bc4071748d Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 00:07:10 +0000
Subject: [PATCH 31/38] Number a rank's experts across expert-parallel groups
 in the local split

The no-gather writer's local split now numbers experts by their FSDP2
dim-0 offset alone. Under expert parallelism transformers slices each
expert parameter to the rank's group, so that offset is group-relative:
the split takes the rank's expert-parallel coordinate and the config's
expert count again, adds the group's base when the tensor holds one
group's experts, and keeps the test of the numbering under each layout.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/export/moe_utils.py            | 27 ++++++++++++++++---
 .../export/test_fused_experts_keep_fused.py   | 26 ++++++++++++++++++
 2 files changed, 49 insertions(+), 4 deletions(-)

diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py
index e12ed147abf..50b8d909f3b 100644
--- a/modelopt/torch/export/moe_utils.py
+++ b/modelopt/torch/export/moe_utils.py
@@ -804,17 +804,29 @@ def _dtensor_dim0_offset(dt) -> int:
         return 0
 
 
-def _split_local_fused_module(prefix: str, sharded_sd: dict) -> dict:
+def _split_local_fused_module(
+    prefix: str, sharded_sd: dict, ep_rank: int = 0, total_experts: int = 0
+) -> dict:
     """Split this rank's local shard of a fused-experts module into per-expert keys.
 
     Split THIS rank's *local* shard of one fused-experts module into per-expert keys with GLOBAL
     expert indices (plain, local tensors -- no all-gather).
 
-    FSDP2 shards the fused 3-D expert weight on dim 0 (the expert axis), so every expert this rank
+    FSDP2/EP shard the fused 3-D expert weight on dim 0 (the expert axis), so every expert this rank
     owns is whole and local. We reuse :func:`split_fused_experts_state_dict` on the ``to_local()``
     shard (which numbers experts ``0..E_local-1``) and then shift each expert index by the rank's
     global dim-0 offset, so the resulting per-expert keys match the single-process export and stay
     on this rank for a direct distributed write (no consolidation).
+
+    The fused expert tensor can arrive in two EP layouts, distinguished by its dim-0 size vs
+    ``total_experts`` (the config expert count):
+      * **per-EP-group** (size < total): transformers pre-sliced the param to this rank's EP group
+        (FSDP x EP -> a dp-sharded DTensor of the group; or a plain ep-shard). ``_dtensor_dim0_offset``
+        only gives the WITHIN-group offset, so we add the EP block base ``ep_rank * group_size``.
+      * **all-experts** (size == total): an EP DTensor sharded across the ep mesh spanning all experts
+        (classic DP x EP). ``_dtensor_dim0_offset`` already returns the TRUE global offset -> no base.
+    ``total_experts == 0``, a count the caller could not read off the model's config, takes the
+    tensor for a per-EP-group slice.
     """
     gated = f"{prefix}.gate_up_proj" in sharded_sd
     first = "gate_up_proj" if gated else "up_proj"
@@ -822,6 +834,13 @@ def _split_local_fused_module(prefix: str, sharded_sd: dict) -> dict:
     # dp_offset: this rank's dim-0 start WITHIN its DTensor; n_local: experts it owns.
     dp_offset = _dtensor_dim0_offset(fw) if isinstance(fw, DTensor) else 0
     n_local = fw.to_local().shape[0] if isinstance(fw, DTensor) else fw.shape[0]
+    group_size = int(fw.shape[0])
+    # Add the EP block base only when the tensor is a per-EP-group slice (size < total_experts) or
+    # total is unknown (0). When it spans all experts, dp_offset is already global.
+    add_ep_base = (
+        bool(ep_rank) and group_size > 0 and (total_experts == 0 or group_size < total_experts)
+    )
+    offset = dp_offset + (ep_rank * group_size if add_ep_base else 0)
 
     def _loc(key):
         v = sharded_sd.get(key)
@@ -864,12 +883,12 @@ def _loc_scale(key):
     split_local = split_fused_experts_state_dict(
         local_fused
     )  # keys {prefix}.{i}.* (i in 0..E_local-1)
-    if dp_offset == 0:
+    if offset == 0:
         return split_local
 
     plen = len(prefix) + 1
     out: dict = {}
     for key, val in split_local.items():
         i_str, tail = key[plen:].split(".", 1)
-        out[f"{prefix}.{int(i_str) + dp_offset}.{tail}"] = val
+        out[f"{prefix}.{int(i_str) + offset}.{tail}"] = val
     return out
diff --git a/tests/unit/torch/export/test_fused_experts_keep_fused.py b/tests/unit/torch/export/test_fused_experts_keep_fused.py
index a646c3743a3..6ed2b81ad57 100644
--- a/tests/unit/torch/export/test_fused_experts_keep_fused.py
+++ b/tests/unit/torch/export/test_fused_experts_keep_fused.py
@@ -565,6 +565,32 @@ def test_fused_experts_prefixes_lists_each_fused_module_once_in_order():
     ]
 
 
+@pytest.mark.parametrize(
+    ("ep_rank", "total_experts", "first_index"),
+    [
+        (0, 0, 0),
+        # This EP group's slice of a larger expert set: shift by the group base.
+        (1, 2 * LOCAL_EXPERTS, LOCAL_EXPERTS),
+        (3, 4 * LOCAL_EXPERTS, 3 * LOCAL_EXPERTS),
+        # The tensor already spans every expert, so its indices are global.
+        (1, LOCAL_EXPERTS, 0),
+        # Total unknown: the legacy per-group shift.
+        (1, 0, LOCAL_EXPERTS),
+    ],
+)
+def test_split_local_numbers_experts_globally(ep_rank, total_experts, first_index):
+    sd = _fused_state_dict(gated=True)
+    unshifted = split_fused_experts_state_dict(dict(sd))
+
+    out = _split_local_fused_module(PREFIX, sd, ep_rank=ep_rank, total_experts=total_experts)
+
+    want = {}
+    for key, value in unshifted.items():
+        expert, tail = key[len(PREFIX) + 1 :].split(".", 1)
+        want[f"{PREFIX}.{int(expert) + first_index}.{tail}"] = value
+    _assert_byte_identical(out, want)
+
+
 # ------------------------------------------------------------------ sharded (DTensor) experts
 
 

From 6317c6c4a399e7934dcbaf207efa89b8a23c0318 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 00:09:02 +0000
Subject: [PATCH 32/38] Check the routed experts' groups against the plan
 before writing them

The writer reads each rank's expert-parallel coordinate and expert
replica index off the routed experts' ParallelState. A ParallelState
left with its default data-parallel group holds the whole world there:
taken as the replica index, it would leave every rank but global rank 0
out of the expert write, and the checkpoint would miss experts with no
error. The writer now refuses groups whose sizes are not the plan's
ep_size and edp, on every rank alike, as the sizes are the same on every
rank. A two-rank test seeds the default group and the plan's own.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/export/distributed_writer.py   | 18 ++++++++--
 .../torch/export/test_distributed_writer.py   | 35 ++++++++++++++++++-
 2 files changed, 50 insertions(+), 3 deletions(-)

diff --git a/modelopt/torch/export/distributed_writer.py b/modelopt/torch/export/distributed_writer.py
index 9915edc7d0f..d380b3c7735 100644
--- a/modelopt/torch/export/distributed_writer.py
+++ b/modelopt/torch/export/distributed_writer.py
@@ -731,7 +731,8 @@ def _distributed_save_hf_checkpoint_impl(
     # the routed experts' ParallelState -- the same source calibration reads, not transformers'
     # internal ``model._device_mesh``. An expert's data-parallel group is its expert-data-parallel one.
     ep_rank, dp_idx = 0, 0
-    if distribution_plan(model).is_expert_parallel:
+    plan = distribution_plan(model)
+    if plan.is_expert_parallel:
         parallel_state = expert_parallel_state(model)
         if parallel_state is None:
             raise RuntimeError(
@@ -739,8 +740,21 @@ def _distributed_save_hf_checkpoint_impl(
                 "group: load it with an expert-parallel loader, which seeds each routed expert's "
                 "ParallelState."
             )
-        ep_rank = parallel_state.expert_model_parallel_group.rank()
+        ep_group = parallel_state.expert_model_parallel_group
         edp_group = parallel_state.data_parallel_group
+        # The coordinates index the plan's own groups. A ParallelState left with its default
+        # data-parallel group holds the whole world there: taken as the replica index, it would leave
+        # every rank but global rank 0 out of the expert write, with no error. The groups have the
+        # same sizes on every rank, so every rank refuses alike.
+        edp_size = edp_group.world_size() if edp_group.is_initialized() else 1
+        if ep_group.world_size() != plan.ep_size or edp_size != plan.edp:
+            raise RuntimeError(
+                f"The routed experts' expert-parallel group holds {ep_group.world_size()} ranks and "
+                f"their data-parallel group {edp_size}, but the model's plan has "
+                f"ep_size={plan.ep_size} and edp={plan.edp}, so the writer cannot tell which rank "
+                "writes which experts."
+            )
+        ep_rank = ep_group.rank()
         if edp_group.is_initialized():
             # Experts are replicated across EDP, so only the first replica writes them.
             dp_idx = edp_group.rank()
diff --git a/tests/unit/torch/export/test_distributed_writer.py b/tests/unit/torch/export/test_distributed_writer.py
index 99d0df37604..810eff08870 100644
--- a/tests/unit/torch/export/test_distributed_writer.py
+++ b/tests/unit/torch/export/test_distributed_writer.py
@@ -64,7 +64,7 @@
     _export_transformers_checkpoint,
     export_hf_checkpoint,
 )
-from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan
+from modelopt.torch.utils.distributed import DistributionPlan, ParallelState, set_distribution_plan
 
 _WRITER = "modelopt.torch.export.distributed_writer"
 
@@ -189,6 +189,39 @@ def test_an_expert_parallel_plan_needs_its_experts_expert_parallel_group(
         _write(model, tmp_path / "export")
 
 
+def _expert_parallel_groups_job(rank, world_size, tmp_dir):
+    for name, data_parallel_group, refused in (
+        # Left at its default, the data-parallel group is the whole world.
+        ("default-data-parallel-group", None, True),
+        # The plan's single replica leaves the experts' data-parallel group unused.
+        ("unused-data-parallel-group", -1, False),
+    ):
+        model = _load(tmp_dir / "src")
+        set_distribution_plan(model, DistributionPlan(world_size, ep_size=world_size))
+        model.model.layers[0].mlp._parallel_state = ParallelState(
+            data_parallel_group=data_parallel_group, expert_model_parallel_group=dist.group.WORLD
+        )
+        try:
+            with bounded_collectives():
+                _write(model, tmp_dir / name)
+        except RuntimeError as error:
+            raised = str(error)
+        else:
+            raised = "no error"
+        assert ("ep_size=2 and edp=1" in raised) if refused else raised == "no error", raised
+
+
+def test_an_expert_parallel_export_checks_its_experts_groups_against_the_plan(
+    cpu_dist_workers, tmp_path
+):
+    """Each rank's expert coordinates come from the routed experts' groups. Taken as the replica
+    index, a data-parallel group left at its default, the whole world, would leave every rank but
+    rank 0 out of the expert write with no error, so every rank refuses it; the plan's own groups
+    pass."""
+    _save(_llama_config(), tmp_path / "src")
+    cpu_dist_workers.run(_expert_parallel_groups_job, tmp_dir=tmp_path)
+
+
 # FP8 weights, with an FP8 KV cache in layer 0 and an NVFP4 one in layer 1: a mixed-precision KV
 # cache, whose per-layer formats the export reads from the quantization config.
 _FP8_MIXED_KV_CFG = copy.deepcopy(mtq.FP8_DEFAULT_CFG)

From 4687e268ffe9e8633889ad6b729136c086a00a0f Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 00:07:46 +0000
Subject: [PATCH 33/38] Fail every rank of a sharded export when rank 0 cannot
 read the carried weights

A sharded export reads the carried tensors off the source on rank 0
alone, and its peers read only their names. That read sat ahead of every
collective the export joins, so a failure there left rank 0 raising out
of the export while its peers waited in the first collective for it.

The read now runs inside an all_reduce for a sharded export, so a
failure on rank 0 raises on every rank, before any of them goes on. The
names every rank seeds get_quant_config with are rank 0's, the same on
every rank. A non-sharded export reads as before, with no collective.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/export/unified_export_hf.py    | 27 +++++++++----
 .../torch/export/test_unified_export_hf.py    | 38 +++++++++++++++++++
 2 files changed, 58 insertions(+), 7 deletions(-)

diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py
index f32e436b3fb..70660a1e831 100644
--- a/modelopt/torch/export/unified_export_hf.py
+++ b/modelopt/torch/export/unified_export_hf.py
@@ -2131,15 +2131,28 @@ def export_hf_checkpoint(
     # then hands them to the writer, whose point-to-point protocol expects one sender, and the export
     # deadlocks with rank 0 inside the writer and the rest parked in the barrier below.
     _writes_extra = not _sharded_export or torch.distributed.get_rank() == 0
-    _carried = read_unplaced_weights(model, keys_only=not _writes_extra)
+    _carried: dict = {}
+
+    def carried_over_names() -> list[str]:
+        # Everything the export writes in original precision straight from the source, by either
+        # mechanism: tensors carried here, and the off-index weight files copied verbatim alongside.
+        # get_quant_config reads this to seed exclude_modules; recorded here because it runs before
+        # that, and because only this point knows what was actually written rather than what was
+        # merely unplaced.
+        _carried.update(read_unplaced_weights(model, keys_only=not _writes_extra))
+        return sorted({*_carried, *off_index_tensor_names(model)})
+
+    if _sharded_export:
+        # Rank 0 reads more of the source here than its peers, so a failure there alone must fail
+        # every rank rather than leave the others waiting in the export's first collective. Every
+        # rank then seeds get_quant_config with rank 0's names.
+        model._modelopt_carried_over_names = all_reduce(
+            "carried weights", carried_over_names, lambda names: names[0]
+        )
+    else:
+        model._modelopt_carried_over_names = carried_over_names()
     if _writes_extra and _carried:
         extra_state_dict = {**_carried, **(extra_state_dict or {})}
-    # Everything the export writes in original precision straight from the source, by either
-    # mechanism: tensors carried above, and the off-index weight files copied verbatim alongside.
-    # get_quant_config reads this to seed exclude_modules; recorded here because it runs before
-    # that, and because only this point knows what was actually written rather than what was
-    # merely unplaced.
-    model._modelopt_carried_over_names = sorted({*_carried, *off_index_tensor_names(model)})
 
     from .layerwise_export import LAYERWISE_EXPORTER_ATTR
 
diff --git a/tests/unit/torch/export/test_unified_export_hf.py b/tests/unit/torch/export/test_unified_export_hf.py
index 472e3e747cb..0273ffd31ce 100644
--- a/tests/unit/torch/export/test_unified_export_hf.py
+++ b/tests/unit/torch/export/test_unified_export_hf.py
@@ -1182,6 +1182,44 @@ def test_every_rank_fails_an_export_that_one_rank_failed_after_the_write(
     cpu_dist_workers.run(_one_rank_fails_after_the_write_job, tmp_dir=tmp_path, failure=failure)
 
 
+def _rank_0_cannot_read_the_carried_weights_job(rank, size, tmp_dir):
+    model = torch.nn.Linear(2, 2)
+    set_distribution_plan(model, DistributionPlan(size, ep_size=size))
+    read = unified_export_hf.read_unplaced_weights
+
+    def unreadable_on_rank_0(*args, **kwargs):
+        if rank == 0:
+            raise OSError("a carried shard of the source is unreadable")
+        return read(*args, **kwargs)
+
+    raised = None
+    with (
+        bounded_collectives(),
+        mock.patch.object(unified_export_hf, "read_unplaced_weights", unreadable_on_rank_0),
+        mock.patch.object(unified_export_hf, "_process_hf_model_for_export", return_value={}),
+        mock.patch.object(distributed_writer, "distributed_save_hf_checkpoint") as write,
+    ):
+        # Not pytest.raises: when nothing is raised it fails with a BaseException, which would take
+        # the shared worker down instead of failing this test alone.
+        try:
+            unified_export_hf.export_hf_checkpoint(model, export_dir=tmp_dir)
+        except (OSError, RuntimeError) as e:
+            raised = e
+    # Rank 0 raises its own error, and rank 1 one that names it, before either writes.
+    assert raised is not None and "a carried shard of the source is unreadable" in str(raised), (
+        raised
+    )
+    write.assert_not_called()
+
+
+def test_every_rank_fails_an_export_whose_carried_weights_rank_0_cannot_read(
+    cpu_dist_workers, tmp_path
+):
+    """Rank 0 alone reads the carried tensors off the source, so its failure there must fail every
+    rank, not leave the others waiting in the export's first collective."""
+    cpu_dist_workers.run(_rank_0_cannot_read_the_carried_weights_job, tmp_dir=tmp_path)
+
+
 def _one_rank_fails_processing_job(rank, size, tmp_dir):
     model = torch.nn.Linear(2, 2)
     set_distribution_plan(model, DistributionPlan(size, ep_size=size))

From b0b1f3b83cb2d1204a684cbeb935ba286f85eb79 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 00:10:18 +0000
Subject: [PATCH 34/38] Skip a codec converter as a decode in place only when
 it names one module

A converter whose reverse only re-encodes a weight's quantization is
skipped, and its reverse rules with it, because it renames nothing. The
check compared only the first source and target, so a converter with a
later target naming another module would have been skipped too, leaving
that rename out of the exported names with no error. Every pattern must
now name a tensor of the same module, the weight or a scale the load
folded into it; another converter is reversed or refused, as before.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 .../torch/export/quant_aware_conversion.py    |  7 +++++--
 .../export/test_quant_aware_conversion.py     | 21 +++++++++++++++++++
 2 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/modelopt/torch/export/quant_aware_conversion.py b/modelopt/torch/export/quant_aware_conversion.py
index 8d94eee3197..36bcb7522f4 100644
--- a/modelopt/torch/export/quant_aware_conversion.py
+++ b/modelopt/torch/export/quant_aware_conversion.py
@@ -786,14 +786,17 @@ def _decodes_in_place(rev, ops) -> bool:
 
     A dequantizing load (FineGrainedFP8 ``dequantize=True``) decodes each weight in place, and
     its reverse would re-encode the weight in the source's format. The export never wants that:
-    it writes ModelOpt's own quantized weights and scales. Such a converter renames nothing.
+    it writes ModelOpt's own quantized weights and scales. Such a converter renames nothing: the
+    weight maps onto itself, and every pattern names a tensor of the same module, the weight or a
+    scale the load folded into it. One naming another module would rename it, so it is not skipped.
     """
     codec_ops = _quantization_codec_ops()
     if not ops or not codec_ops or not all(isinstance(op, codec_ops) for op in ops):
         return False
     sources = [p.strip("^$") for p in _as_list(rev.source_patterns)]
     targets = [p.strip("^$") for p in _as_list(rev.target_patterns)]
-    return bool(sources) and sources[:1] == targets[:1]
+    modules = {p.rpartition(".")[0] for p in (*sources, *targets)}
+    return bool(sources) and sources[:1] == targets[:1] and len(modules) == 1
 
 
 def _build_reverse_rules(
diff --git a/tests/unit/torch/export/test_quant_aware_conversion.py b/tests/unit/torch/export/test_quant_aware_conversion.py
index e32a709c642..d6feac8799b 100644
--- a/tests/unit/torch/export/test_quant_aware_conversion.py
+++ b/tests/unit/torch/export/test_quant_aware_conversion.py
@@ -1387,6 +1387,27 @@ def test_the_identity_decode_is_known_without_the_fp8_integration(monkeypatch):
         quant_aware_conversion._quantization_codec_ops.cache_clear()
 
 
+@pytest.mark.parametrize(
+    ("sources", "targets", "in_place"),
+    [
+        # A dequantizing FP8 load, reversed: the weight and the scales folded into it.
+        (["weight$"], ["weight$", "weight_scale_inv", "activation_scale"], True),
+        (["mlp.up_proj.weight$"], ["mlp.up_proj.weight$", "mlp.up_proj.weight_scale_inv"], True),
+        # A later target names another module, so the converter renames it.
+        (["mlp.up_proj.weight$"], ["mlp.up_proj.weight$", "mlp.gate_proj.weight"], False),
+    ],
+    ids=["bare-siblings", "module-siblings", "another-module"],
+)
+def test_a_codec_converter_decodes_in_place_only_if_it_names_one_module(sources, targets, in_place):
+    """A converter skipped as a decode in place gives up its reverse rules, so it must rename
+    nothing: a target naming another module is kept, to be reversed or refused."""
+    core = pytest.importorskip("transformers.core_model_loading")
+    if not hasattr(core, "_IdentityOp"):
+        pytest.skip("transformers build has no _IdentityOp")
+    rev = types.SimpleNamespace(source_patterns=sources, target_patterns=targets)
+    assert quant_aware_conversion._decodes_in_place(rev, [core._IdentityOp()]) is in_place
+
+
 def test_dense_split_names_each_part_once_through_load_siblings():
     """A dense fused converter that a dequantizing load rewrote still un-fuses into one part per
     module, not one per checkpoint tensor the load consumed."""

From 33ba11619391b0ecd9b77197fa38acc8fe975718 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 00:33:29 +0000
Subject: [PATCH 35/38] Send rank 0's carried names with broadcast_from, and
 print a traceback only when sharded

Only rank 0's list of carried names is used, so its peers no longer read
the source for names they would throw away: broadcast_from runs the read
on rank 0 alone, sends the list, and raises on every rank when that read
fails, as the all-reduce did.

The traceback print before the closing consensus is for a sharded export,
where the other ranks wait in it. A single process prints the traceback
once again when the error reaches the top, so it no longer prints it here.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/export/unified_export_hf.py | 20 +++++++++++---------
 1 file changed, 11 insertions(+), 9 deletions(-)

diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py
index 70660a1e831..67aef86d7ca 100644
--- a/modelopt/torch/export/unified_export_hf.py
+++ b/modelopt/torch/export/unified_export_hf.py
@@ -2143,12 +2143,9 @@ def carried_over_names() -> list[str]:
         return sorted({*_carried, *off_index_tensor_names(model)})
 
     if _sharded_export:
-        # Rank 0 reads more of the source here than its peers, so a failure there alone must fail
-        # every rank rather than leave the others waiting in the export's first collective. Every
-        # rank then seeds get_quant_config with rank 0's names.
-        model._modelopt_carried_over_names = all_reduce(
-            "carried weights", carried_over_names, lambda names: names[0]
-        )
+        # Rank 0 reads the source here and sends every rank its names, so a failure of that read
+        # fails every rank rather than leave the others waiting in the export's first collective.
+        model._modelopt_carried_over_names = broadcast_from("carried weights", carried_over_names)
     else:
         model._modelopt_carried_over_names = carried_over_names()
     if _writes_extra and _carried:
@@ -2372,9 +2369,14 @@ def carried_over_names() -> list[str]:
         # Print (not just warnings.warn, which the default filter can dedupe/suppress) BEFORE the
         # barrier below: under FSDP2/EP every other rank is still waiting there, so a failing rank
         # that raises straight into `finally` blocks on that barrier forever without this rank's
-        # traceback ever reaching the log -- the job just hangs with no visible error.
-        _rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0
-        print(f"[export_hf_checkpoint] rank {_rank} failed:\n{traceback.format_exc()}", flush=True)
+        # traceback ever reaching the log -- the job just hangs with no visible error. A single
+        # process has no barrier, and the interpreter prints the traceback.
+        if is_fsdp2_sharded or is_ep_sharded:
+            print(
+                f"[export_hf_checkpoint] rank {torch.distributed.get_rank()} failed:\n"
+                f"{traceback.format_exc()}",
+                flush=True,
+            )
         warnings.warn(
             "Cannot export model to the model_config. The modelopt-optimized model state_dict"
             " can be saved with torch.save for further inspection."

From ace58051c75c40d99492450a04c4c1f1a18218d9 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 18:57:06 +0000
Subject: [PATCH 36/38] Import PretrainedConfig with the module's guarded
 transformers import

The config check in _source_distributed_config_attrs imported PretrainedConfig inside the function, with no stated reason. transformers is optional for this module, and the module already imports it once, guarded, at the top. PretrainedConfig joins that import, and the check reads _TRANSFORMERS_CONFIG_TYPES, which is empty without transformers, as _TRANSFORMERS_MODEL_TYPES is.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/export/unified_export_hf.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py
index 67aef86d7ca..a92b33c313c 100644
--- a/modelopt/torch/export/unified_export_hf.py
+++ b/modelopt/torch/export/unified_export_hf.py
@@ -67,11 +67,13 @@
     HAS_DIFFUSERS = False
 
 try:
-    from transformers import PreTrainedModel
+    from transformers import PretrainedConfig, PreTrainedModel
 
     _TRANSFORMERS_MODEL_TYPES: tuple[type, ...] = (PreTrainedModel,)
+    _TRANSFORMERS_CONFIG_TYPES: tuple[type, ...] = (PretrainedConfig,)
 except ImportError:
     _TRANSFORMERS_MODEL_TYPES = ()
+    _TRANSFORMERS_CONFIG_TYPES = ()
 
 
 from modelopt.torch.opt.conversion import ModeloptStateManager, modelopt_state
@@ -1985,8 +1987,6 @@ def _source_distributed_config_attrs(model: nn.Module):
     that cannot be read counts as having none. The run's values come back on exit, since the
     sharded model still reads them (transformers' ``tp_plan`` checks ``distributed_config``).
     """
-    from transformers import PretrainedConfig
-
     source: Any = {}
     ckpt = _source_checkpoint(model)
     if ckpt and Path(ckpt, "config.json").is_file():
@@ -2006,7 +2006,7 @@ def use_source(config: Any, source: Any) -> None:
                 elif attr in vars(config):
                     delattr(config, attr)
         for key, sub_config in list(vars(config).items()):
-            if isinstance(sub_config, PretrainedConfig):
+            if isinstance(sub_config, _TRANSFORMERS_CONFIG_TYPES):
                 use_source(sub_config, source.get(key))
 
     try:

From 57f6e2072fdc66ad773e63bc5d08ae1eb916397a Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Sat, 31 Oct 2026 00:10:32 +0000
Subject: [PATCH 37/38] Drop the consolidation step from the export's comments

The no-gather export has no consolidation step: each rank's writer splits
its own fused experts into per-expert tensors. #2752 fixed its own copy of
this wording; the docstring and comment this PR adds said the same. Comments
and a docstring only.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/export/unified_export_hf.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py
index a92b33c313c..12bb9e3ec3d 100644
--- a/modelopt/torch/export/unified_export_hf.py
+++ b/modelopt/torch/export/unified_export_hf.py
@@ -1112,8 +1112,8 @@ def _process_hf_model_for_export(
     The processing half of an export, for a caller that serializes the model itself: the no-gather
     writer, which writes each rank's shard. Afterwards the model's own state dict holds the
     deployment weights. Nothing is gathered, so it runs on a live FSDP2 or expert-parallel model;
-    there the fused MoE experts stay fused, since each rank folds only its local shard and the
-    writer splits them after consolidation.
+    there the fused MoE experts stay fused, since each rank folds only its local shard, and the
+    writer then splits each rank's experts into per-expert tensors.
     """
     dtype, _, quant_config = _prepare_model_for_export(model, dtype, is_modelopt_qlora)
     # Only the no-gather writer calls this, and only for a sharded model, so the experts stay fused
@@ -2246,7 +2246,8 @@ def carried_over_names() -> list[str]:
 
             _qlora = kwargs.get("is_modelopt_qlora", False)
             # Folds in place; the model is sharded here, so the MoE experts stay fused and each rank
-            # folds only its LOCAL experts. The writer splits them after consolidation.
+            # folds only its local experts. The writer then splits each rank's experts into
+            # per-expert tensors.
             hf_quant_config = _process_together(model, dtype, _qlora)
             if getattr(model, "hf_quantizer", None) is not None:
                 model.hf_quantizer = None

From 72a687c50c039f3675a4da3901caf46fd6461fca Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Sat, 31 Oct 2026 22:16:40 +0000
Subject: [PATCH 38/38] Give only the single-sender reason for carrying weights
 from rank 0

The comment above _writes_extra said that under expert parallelism every
other rank's experts read as unplaced, which is wrong: the 889 tensors on
GLM-5.3-Flash are its MTP layer, and other ranks' experts map to the fused
expert parameters every rank holds. The gate stands on one reason: both
sharded writers take extra_state_dict from rank 0 alone. The comment now
says so, and that the names, unlike the tensors, reach every rank. Comment
only.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/export/unified_export_hf.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py
index 12bb9e3ec3d..357cb994ac8 100644
--- a/modelopt/torch/export/unified_export_hf.py
+++ b/modelopt/torch/export/unified_export_hf.py
@@ -2125,11 +2125,11 @@ def export_hf_checkpoint(
     # Either form of sharding, including the hybrid where both hold: a model can be FSDP2-dense AND
     # expert-parallel, so these are tracked separately rather than as an either/or.
     _sharded_export = _ep_sharded or _fsdp2_sharded
-    # A sharded export writes the carried weights from rank 0 alone. This has to cover EP as well as
-    # FSDP2: under EP every OTHER rank's experts are absent from this rank's model, so they read as
-    # unplaced and each rank would carry them -- 889 tensors per rank on GLM-5.3-Flash. Every rank
-    # then hands them to the writer, whose point-to-point protocol expects one sender, and the export
-    # deadlocks with rank 0 inside the writer and the rest parked in the barrier below.
+    # A sharded export reads and hands over the carried weights on rank 0 alone, since both sharded
+    # writers take extra_state_dict from that one sender: on any other rank, reading them (889
+    # tensors on GLM-5.3-Flash, its MTP layer) would be host memory spent for nothing. This gates
+    # the tensors only. Every rank still resolves their names, which _modelopt_carried_over_names
+    # below records alike on every rank for the quant config.
     _writes_extra = not _sharded_export or torch.distributed.get_rank() == 0
     _carried: dict = {}