[OMNIML-6084] Export an expert-parallel model through the no-gather writer

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
@@ -35,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
@@ -690,7 +692,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 +718,54 @@ 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.
+    ep_rank, dp_idx = 0, 0
+    plan = distribution_plan(model)
+    if plan.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 an expert-parallel loader, which seeds each routed expert's "
+                "ParallelState."
+            )
+        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()
+    # 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.
+    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 /
+    # 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 +839,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)

modelopt/torch/export/moe_utils.py

@@ -804,24 +804,43 @@ 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"
     fw = sharded_sd[f"{prefix}.{first}"]
     # 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

modelopt/torch/export/quant_aware_conversion.py

@@ -62,6 +62,8 @@
 (in-memory-name) behavior rather than emit a silently-wrong checkpoint.
 """
 
+import contextlib
+import functools
 import logging
 import re
 from collections.abc import Iterable
@@ -98,6 +100,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")
@@ -745,6 +759,46 @@ 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.
+
+    Includes the identity op, which is what older releases (5.5, 5.6) reverse ``Fp8Dequantize``
+    into; later ones reverse it into ``Fp8Quantize``.
+    """
+    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
+
+        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)
+
+
+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: 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)]
+    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(
     model,
 ) -> tuple[list[SplitRule], list[_MergeRule], list[RenameRule], list[str], set[str] | None]:
@@ -767,6 +821,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.
@@ -842,6 +898,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]}"
@@ -884,7 +942,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 +961,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 +970,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