[OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time

modelopt/torch/quantization/plugins/huggingface.py

@@ -41,7 +41,7 @@
 from modelopt.torch.kernels.quantization.gemm import IS_AVAILABLE as IS_TRITON_AVAILABLE
 from modelopt.torch.opt.dynamic import DynamicModule
 from modelopt.torch.utils import warn_rank_0
-from modelopt.torch.utils.distributed import ParallelState
+from modelopt.torch.utils.distributed import ParallelState, recorded_distribution_plan
 
 from ..algorithms import AutoQuantizeGradientSearcher
 from ..conversion import register
@@ -596,7 +596,9 @@ def convert_hf_parallel_linears_on_the_fly(model):
         elif HFRowParallelLinear.is_compatible(module, tp_mesh):
             HFRowParallelLinear.convert(module)
             converted = True
-    if tp_mesh is not None and not converted:
+    # An expert-parallel load keeps its experts' mesh there too, and shards no Linear by design.
+    plan = recorded_distribution_plan(model)
+    if tp_mesh is not None and not converted and not (plan is not None and plan.is_expert_parallel):
         warnings.warn(
             "Found a transformers tensor-parallel mesh but no TP-sharded linear layer; they will be"
             " quantized as unsharded layers."

modelopt/torch/utils/moe.py

@@ -15,25 +15,25 @@
 
 """HuggingFace-coupled MoE expert-parallel (EP) helpers.
 
-Finding a model's routed experts, choosing the EP degree they are sharded over, and the DP=EP
-all-to-all expert transport (``torch_a2a``).
+Finding a model's routed experts, choosing the EP degree they are sharded over, the DP=EP
+all-to-all expert transport (``torch_a2a``). The expert-parallel model loader is
+:func:`~modelopt.torch.utils.plugins.model_load_utils.load_and_prepare_ep`.
 
 The all-to-all transport is plain ``torch.distributed.all_to_all_single`` (NCCL) wrapped around
-transformers' quantizable ``_grouped_linear``. ``register_torch_a2a()`` registers it into
-transformers' ``ExpertsInterface`` lazily -- never at import time, so importing this module has no
-transformers-``ExpertsInterface`` side effect.
+transformers' quantizable ``_grouped_linear``. ``register_torch_a2a()`` registers
+it into transformers' ``ExpertsInterface`` (lazily, from ``load_and_prepare_ep``'s a2a path -- never at
+import time, so importing this module has no transformers-``ExpertsInterface`` side effect).
 """
 
 import warnings
-from typing import TYPE_CHECKING
 
 import torch
 import torch.distributed as dist
+import torch.nn as nn
 from packaging.version import parse
 from torch.distributed.nn.functional import all_to_all_single as autograd_all_to_all_single
 
-if TYPE_CHECKING:
-    from torch import nn
+from .distributed import distribution_plan
 
 # The transformers release the expert-parallel load needs, for every model: 5.6 builds a model only
 # with the experts implementations it ships, never torch_a2a.
@@ -559,7 +559,9 @@ def register_torch_a2a():
     These two styles are functionally transformers' ``megamoe_router`` / ``megamoe_experts`` (added
     in 5.10.4 for DeepGEMM Mega-MoE), reimplemented here so the transport works below 5.10.4 too.
     The base classes (``RouterParallel``, ``TensorParallelLayer``) are stock wherever the EP infra
-    exists (transformers >= 5.6); ``grouped_gemm`` (expert weight sharding) is stock too.
+    exists (transformers >= 5.6); ``grouped_gemm`` (expert weight sharding) is stock too. A model
+    builds with ``torch_a2a`` only from 5.7, the floor ``load_and_prepare_ep`` enforces: 5.6 accepts
+    only the experts implementations it ships.
     """
     global _TORCH_A2A_REGISTERED
     if _TORCH_A2A_REGISTERED:
@@ -626,3 +628,23 @@ def shard_tensor(self, param, tensor_idx=None, device=None, dtype=None):
     ALL_PARALLEL_STYLES.register("torch_a2a_router", TorchA2ARouter())
     ALL_PARALLEL_STYLES.register("torch_a2a_experts", TorchA2AExperts())
     _TORCH_A2A_REGISTERED = True
+
+
+def expert_parallel_mesh(model: nn.Module):
+    """The mesh ``model``'s routed experts are sharded over, or ``None`` if its plan shards none.
+
+    :func:`~modelopt.torch.utils.plugins.model_load_utils.load_and_prepare_ep` loads the model over
+    its expert-parallel submesh, and keeps that submesh on the model as ``_device_mesh``.
+    Calibration takes the mesh to re-install the all-to-all forwards
+    (:func:`~modelopt.torch.utils.distributed.distributed_forward_loop`).
+    """
+    if not distribution_plan(model).is_expert_parallel:
+        return None
+    mesh = getattr(model, "_device_mesh", None)
+    if mesh is None:
+        raise RuntimeError(
+            f"{type(model).__name__} is planned expert-parallel but carries no device mesh. Pass "
+            "the model load_and_prepare_ep returned, not a submodule of it: the loader keeps the "
+            "mesh it shards the experts over on that model alone."
+        )
+    return mesh