feat(multi-lora): enable and validate MoE expert adapters

miles/backends/megatron_utils/bridge_lora_helpers.py

@@ -6,15 +6,18 @@
 
 from __future__ import annotations
 
+import logging
 from argparse import Namespace
 from dataclasses import dataclass
 
 from megatron.core.utils import get_attr_wrapped_model
 
 from miles.utils.hf_config import load_hf_config
-from miles.utils.multi_lora import is_multi_lora_enabled
+from miles.utils.multi_lora import is_multi_lora_enabled, targets_expert_leaves
 
-from .lora_utils import patch_param_grad_buffer_for_colocate_mode_lora
+from .lora_utils import convert_target_modules_to_hf, patch_param_grad_buffer_for_colocate_mode_lora
+
+logger = logging.getLogger(__name__)
 
 
 @dataclass
@@ -66,6 +69,46 @@ def _get_model_config_from_wrapped(model):
     return get_attr_wrapped_model(model, "config", allow_none=False)
 
 
+def _validate_multi_lora_moe_support(args: Namespace, provider) -> None:
+    """Reject MoE configs the multi-slot grouped-expert adapter cannot serve (checked
+    post-finalize because they depend on the resolved provider, not the CLI)."""
+    if not getattr(provider, "num_moe_experts", None):
+        return
+    if not targets_expert_leaves(args.target_modules):
+        logger.info("[multilora] MoE model with no expert leaves in --target-modules; experts stay frozen")
+        return
+
+    # Checked on the provider: --expert-tensor-parallel-size stays None until Megatron resolves it.
+    expert_tp = getattr(provider, "expert_tensor_parallel_size", 1) or 1
+    assert expert_tp == 1, (
+        f"Multi-LoRA on MoE experts requires expert_tensor_parallel_size=1 (resolved to "
+        f"{expert_tp}); set --expert-tensor-parallel-size 1."
+    )
+    assert getattr(provider, "moe_grouped_gemm", False), (
+        "Multi-LoRA on MoE experts requires moe_grouped_gemm=True (SequentialMLP expert "
+        "linears are skipped, so the experts would train no adapter)."
+    )
+    assert not getattr(provider, "fp8", None) and not getattr(provider, "fp4", None), (
+        "Multi-LoRA on MoE experts does not support fp8/fp4 experts (quantization padding "
+        "desynchronizes the dispatched token order)."
+    )
+    # sglang only wraps a fused MoE layer when both expert projections are targeted.
+    served = set(convert_target_modules_to_hf(list(args.target_modules)))
+    expert_pair = {"gate_proj", "up_proj", "down_proj"}
+    if served & expert_pair:
+        assert expert_pair <= served, (
+            f"Multi-LoRA on MoE experts requires all of {sorted(expert_pair)} in "
+            f"--target-modules (got {sorted(served & expert_pair)}); a one-sided expert "
+            f"target is dropped at rollout time."
+        )
+    assert not getattr(
+        provider, "moe_pad_expert_input_to_capacity", False
+    ), "Multi-LoRA on MoE experts does not support --moe-pad-expert-input-to-capacity."
+    assert not getattr(
+        provider, "moe_permute_fusion", False
+    ), "Multi-LoRA on MoE experts requires moe_permute_fusion=False."
+
+
 def _setup_lora_model_via_bridge(args: Namespace) -> list:
     """Build Megatron model with LoRA using Megatron-Bridge.
 
@@ -105,6 +148,15 @@ def _setup_lora_model_via_bridge(args: Namespace) -> list:
     provider.variable_seq_lengths = True
     provider.moe_token_dispatcher_type = "alltoall"
     provider.moe_router_load_balancing_type = "none"
+    if is_multi_lora_enabled(args) and targets_expert_leaves(args.target_modules):
+        # Expert adapters cannot replay the fused permute's row_id_map, and most bridge
+        # MoE providers default the fusion on — so turn it off rather than refuse to build.
+        if getattr(provider, "moe_permute_fusion", False):
+            logger.info(
+                "[multilora] disabling moe_permute_fusion: expert adapters replay the "
+                "dispatcher's permutation, which the fused kernel does not expose"
+            )
+        provider.moe_permute_fusion = False
     if getattr(args, "decoder_first_pipeline_num_layers", None) is not None:
         provider.num_layers_in_first_pipeline_stage = args.decoder_first_pipeline_num_layers
     if getattr(args, "decoder_last_pipeline_num_layers", None) is not None:
@@ -114,6 +166,8 @@ def _setup_lora_model_via_bridge(args: Namespace) -> list:
     provider.finalize()
 
     if is_multi_lora_enabled(args):
+        _validate_multi_lora_moe_support(args, provider)
+
         from miles.backends.megatron_utils.multi_lora_utils import create_multi_lora_instance
 
         lora = create_multi_lora_instance(args)

miles/backends/megatron_utils/multi_lora_utils.py

@@ -12,9 +12,13 @@
 from miles.backends.training_utils.parallel import get_parallel_state
 from miles.ray.multi_lora.controller import get_multi_lora_controller
 from miles.utils.adapter_config import AdapterRun
+from miles.utils.distributed_utils import get_gloo_group
 
 logger = logging.getLogger(__name__)
 
+# Cached by adapter_shard_topology(); the topology is fixed for the run.
+_shard_topology: tuple[bool, tuple[tuple[int, int, int], ...]] | None = None
+
 
 def create_multi_lora_instance(args: Namespace):
     """Create a MultiLoRA instance from training args."""
@@ -32,6 +36,7 @@ def create_multi_lora_instance(args: Namespace):
 
         lora_cls = LoRA
 
+    # exclude_modules was already folded into target_modules during arg validation.
     return MultiLoRA(
         target_modules=convert_target_modules_to_megatron(args.target_modules, lora_type=lora_cls),
         n_adapters=args.multi_lora_n_adapters,
@@ -43,21 +48,53 @@ def create_multi_lora_instance(args: Namespace):
     )
 
 
-def all_megatron_checkpoints_exist(step_dir: Path, tp_size, pp_size) -> bool:
-    return all(
-        (step_dir / f"adapter_megatron_tp{tp}_pp{pp}.pt").exists() for tp in range(tp_size) for pp in range(pp_size)
-    )
+def megatron_shard_name(tp_rank: int, pp_rank: int, ep_rank: int, ep_size: int) -> str:
+    """Adapter shard name for one (tp, pp, ep) coordinate; EP ranks hold different local
+    experts. The ep suffix is omitted at ep_size == 1 so legacy checkpoints stay loadable."""
+    name = f"adapter_megatron_tp{tp_rank}_pp{pp_rank}"
+    if ep_size > 1:
+        name += f"_ep{ep_rank}"
+    return name + ".pt"
+
+
+def adapter_shard_topology() -> tuple[bool, tuple[tuple[int, int, int], ...]]:
+    """Return ``(this_rank_writes_its_shard, realized (tp, pp, ep) coords)`` via one cached gloo all-gather."""
+    global _shard_topology
+    if _shard_topology is not None:
+        return _shard_topology
+    parallel_state = get_parallel_state()
+    coords = (parallel_state.tp.rank, parallel_state.pp.rank, parallel_state.ep.rank)
+    if not dist.is_initialized():
+        _shard_topology = (True, (coords,))
+        return _shard_topology
+
+    current_rank = dist.get_rank()
+    group = get_gloo_group()
+    gathered: list[object] = [None] * dist.get_world_size(group=group)
+    dist.all_gather_object(gathered, (coords, current_rank), group=group)
+    is_writer = current_rank == min(rank for entry_coords, rank in gathered if entry_coords == coords)
+    _shard_topology = (is_writer, tuple(sorted({entry_coords for entry_coords, _ in gathered})))
+    return _shard_topology
+
+
+def all_megatron_checkpoints_exist(step_dir: Path, shard_names) -> bool:
+    return all((step_dir / name).exists() for name in shard_names)
 
 
 def find_latest_checkpoint(ckpt_dir: Path) -> tuple[Path | None, int]:
+    _, coords = adapter_shard_topology()
     if not ckpt_dir.exists():
         return None, 0
 
     parallel_state = get_parallel_state()
-    tp_size = parallel_state.tp.size
-    pp_size = parallel_state.pp.size
-    tp_rank = parallel_state.tp.rank
-    pp_rank = parallel_state.pp.rank
+    ep_size = parallel_state.ep.size
+    my_coords = (parallel_state.tp.rank, parallel_state.pp.rank, parallel_state.ep.rank)
+
+    expected = {megatron_shard_name(*coord, ep_size) for coord in coords}
+    my_shard = megatron_shard_name(*my_coords, ep_size)
+    # Legacy pre-expert-adapter layout: no ep suffix; safe for all EP ranks to read (EP-replicated).
+    legacy = {megatron_shard_name(tp, pp, 0, 1) for tp, pp, _ in coords}
+    my_legacy = megatron_shard_name(my_coords[0], my_coords[1], 0, 1)
 
     def get_step(d):
         return int(d.name.split("_")[1])
@@ -69,8 +106,11 @@ def get_step(d):
     )
     for step_dir in step_dirs:
         step = get_step(step_dir)
-        if all_megatron_checkpoints_exist(step_dir, tp_size, pp_size):
-            return step_dir / f"adapter_megatron_tp{tp_rank}_pp{pp_rank}.pt", step
+        if all_megatron_checkpoints_exist(step_dir, expected):
+            return step_dir / my_shard, step
+        if ep_size > 1 and all_megatron_checkpoints_exist(step_dir, legacy):
+            logger.info(f"[multilora] resuming from pre-expert-parallel shard layout in {step_dir}")
+            return step_dir / my_legacy, step
 
     return None, 0
 
@@ -116,20 +156,28 @@ def zero_optimizer_state_for_adapter(optimizer, model, idx: int) -> None:
 
 
 def slice_lora_to_rank(hf_name: str, tensor: torch.Tensor, adapter_rank: int) -> torch.Tensor:
-    if "lora_A" in hf_name and adapter_rank < tensor.shape[0]:
-        remainder = tensor[adapter_rank:]
-        assert remainder.abs().max() == 0, (
-            f"lora_A padded dims are non-zero: {hf_name}, "
-            f"max={remainder.abs().max().item():.6e}, shape={tensor.shape}, rank={adapter_rank}"
-        )
-        return tensor[:adapter_rank]
-    if "lora_B" in hf_name and adapter_rank < tensor.shape[1]:
-        remainder = tensor[:, adapter_rank:]
-        assert remainder.abs().max() == 0, (
-            f"lora_B padded dims are non-zero: {hf_name}, "
-            f"max={remainder.abs().max().item():.6e}, shape={tensor.shape}, rank={adapter_rank}"
-        )
-        return tensor[:, :adapter_rank]
+    """Trim a max-rank-padded LoRA tensor to ``adapter_rank`` on the rank axis, addressed
+    from the end so packed grouped-expert exports are not sliced on the expert axis."""
+    if "lora_A" in hf_name:
+        rank_dim = tensor.ndim - 2
+        if adapter_rank < tensor.shape[rank_dim]:
+            remainder = tensor.narrow(rank_dim, adapter_rank, tensor.shape[rank_dim] - adapter_rank)
+            assert remainder.abs().max() == 0, (
+                f"lora_A padded dims are non-zero: {hf_name}, "
+                f"max={remainder.abs().max().item():.6e}, shape={tensor.shape}, rank={adapter_rank}"
+            )
+            return tensor.narrow(rank_dim, 0, adapter_rank)
+        return tensor
+    if "lora_B" in hf_name:
+        rank_dim = tensor.ndim - 1
+        if adapter_rank < tensor.shape[rank_dim]:
+            remainder = tensor.narrow(rank_dim, adapter_rank, tensor.shape[rank_dim] - adapter_rank)
+            assert remainder.abs().max() == 0, (
+                f"lora_B padded dims are non-zero: {hf_name}, "
+                f"max={remainder.abs().max().item():.6e}, shape={tensor.shape}, rank={adapter_rank}"
+            )
+            return tensor.narrow(rank_dim, 0, adapter_rank)
+        return tensor
     return tensor
 
 
@@ -144,7 +192,7 @@ def save_multi_lora_checkpoints(
     Layout (per adapter)::
 
         {adapter.save}/checkpoints/step_{iteration}/
-        ├── adapter_megatron_tp{tp}_pp{pp}.pt   ← per-rank shard, fast resume
+        ├── adapter_megatron_tp{tp}_pp{pp}[_ep{ep}].pt  ← per-rank shard, fast resume
         ├── adapter_model.safetensors           ← gathered HF, inference / external
         └── adapter_config.json                 ← HF PEFT metadata (r, alpha, ...)
     """
@@ -158,11 +206,11 @@ def save_multi_lora_checkpoints(
     parallel_state = get_parallel_state()
     tp_rank = parallel_state.tp.rank
     pp_rank = parallel_state.pp.rank
-    # One writer per (tp, pp) shard: LoRA params are replicated across DP AND
-    # CP, so gate on the combined dp×cp group. Gating on intra_dp alone left
-    # every CP rank writing the same shard file and racing the os.replace.
-    is_dp_cp_rank_0 = parallel_state.intra_dp_cp.rank == 0
-    is_global_writer = is_dp_cp_rank_0 and tp_rank == 0 and pp_rank == 0
+    ep_rank = parallel_state.ep.rank
+    ep_size = parallel_state.ep.size
+    # Exactly one writer per (tp, pp, ep) shard; see adapter_shard_topology.
+    is_shard_writer, _ = adapter_shard_topology()
+    is_global_writer = is_shard_writer and tp_rank == 0 and pp_rank == 0 and ep_rank == 0
 
     target_modules_hf = (
         convert_target_modules_to_hf(list(args.target_modules))
@@ -183,21 +231,21 @@ def save_multi_lora_checkpoints(
 
         final_dir = config.save / "checkpoints" / f"step_{iteration}"
         tmp_dir = config.save / "checkpoints" / f"_tmp_step_{iteration}"
-        if is_dp_cp_rank_0:
+        if is_shard_writer:
             tmp_dir.mkdir(parents=True, exist_ok=True)
         if dist.is_initialized():
             dist.barrier()
 
         with expose_adapter_slot(model, adapter.slot):
             # Megatron checkpoints
-            if is_dp_cp_rank_0:
+            if is_shard_writer:
                 shard: dict[str, torch.Tensor] = {
                     name: param.data.cpu()
                     for batch in model
                     for name, param in batch.named_parameters()
                     if ".adapter." in name
                 }
-                native_path = tmp_dir / f"adapter_megatron_tp{tp_rank}_pp{pp_rank}.pt"
+                native_path = tmp_dir / megatron_shard_name(tp_rank, pp_rank, ep_rank, ep_size)
                 torch.save(shard, native_path)
                 logger.info(f"{log_prefix} saved Megatron shard " f"({len(shard)} tensors) to {native_path}")
 

miles/utils/multi_lora.py

@@ -46,6 +46,23 @@ def define_new_adapter_metrics(snapshot: dict) -> None:
         define_step_key_metric_group(prefix=f"{name}/perf", step_key="rollout/step")
 
 
+# Leaf module names that can live inside MoE experts (they also name the dense MLP
+# projections); the bulk aliases expand to them during target-module resolution.
+_EXPERT_LEAF_NAMES = frozenset({"linear_fc1", "linear_fc2", "gate_proj", "up_proj", "down_proj"})
+_ALL_MODULE_ALIASES = frozenset({"all", "all-linear", "all_linear"})
+
+
+def targets_expert_leaves(target_modules: Any) -> bool:
+    """Whether ``target_modules`` can put adapters on MoE expert linears."""
+    if isinstance(target_modules, str):
+        target_modules = [target_modules]
+    entries = [str(tm).strip().lower() for tm in (target_modules or [])]
+    if any(entry in _ALL_MODULE_ALIASES for entry in entries):
+        return True
+    # Map each entry (possibly a dotted or wildcard path) to its leaf module name.
+    return any(entry.split(".")[-1] in _EXPERT_LEAF_NAMES for entry in entries)
+
+
 def validate_multi_lora_args(args: Any) -> None:
     """Set ``args.multi_lora``, then validate and default the multi-LoRA arg
     surface. Called from ``miles_validate_args``; a no-op for normal runs."""
@@ -68,6 +85,23 @@ def validate_multi_lora_args(args: Any) -> None:
     assert args.lora_rank > 0, "--lora-rank must be set when --multi-lora-n-adapters > 0"
     assert args.target_modules is not None, "--target-modules must be set when --multi-lora-n-adapters > 0"
     assert args.train_backend == "megatron", "Multi-LoRA currently requires --train-backend megatron"
+    # Adapter routing is only recompute-safe without pipelining; enforce at launch.
+    assert getattr(args, "pipeline_model_parallel_size", 1) == 1, (
+        "Multi-LoRA requires --pipeline-model-parallel-size 1: no single rank holds a "
+        "complete adapter to push to the rollout engines, and a pipelined schedule would "
+        "recompute activations against a later micro-batch's adapter routing."
+    )
+    # Per-slot token spans assume sequence-major contiguous sample packing, which only 'thd' provides.
+    assert getattr(args, "qkv_format", "thd") == "thd", (
+        "Multi-LoRA requires --qkv-format thd: per-adapter token spans assume the "
+        f"micro-batch packs samples contiguously, which bshd does not (got {args.qkv_format!r})."
+    )
+    assert not getattr(args, "experts_shared_outer_loras", False), (
+        "Multi-LoRA does not support --experts-shared-outer-loras; MoE expert adapters "
+        "use the per-expert layout. Drop the flag (and --sglang-experts-shared-outer-loras)."
+    )
+    # Expert-parallel sizes are checked post-finalize in _validate_multi_lora_moe_support:
+    # --expert-tensor-parallel-size stays None until Megatron's own validate_args resolves it.
     assert "muon" not in str(getattr(args, "optimizer", "")).lower(), (
         "Multi-LoRA does not support Muon: per-adapter decoupled stepping is only "
         "implemented for Adam-family per-slot optimizers"

tests/fast/backends/megatron_utils/test_multi_lora_checkpoint_naming.py

@@ -0,0 +1,46 @@
+"""Adapter shards are keyed by (tp, pp, ep): EP ranks hold different local experts, and
+the realized coordinates are not the tp x pp x ep cross product when ETP < TP."""
+
+from miles.backends.megatron_utils.multi_lora_utils import all_megatron_checkpoints_exist, megatron_shard_name
+
+
+def _names(coords, ep_size):
+    return {megatron_shard_name(*coord, ep_size) for coord in coords}
+
+
+def test_shard_name_omits_ep_suffix_without_expert_parallelism():
+    # Checkpoints written before expert adapters existed must stay loadable.
+    assert megatron_shard_name(0, 0, 0, ep_size=1) == "adapter_megatron_tp0_pp0.pt"
+    assert megatron_shard_name(1, 2, 0, ep_size=1) == "adapter_megatron_tp1_pp2.pt"
+
+
+def test_shard_name_is_unique_per_expert_parallel_rank():
+    names = {megatron_shard_name(0, 0, ep, ep_size=4) for ep in range(4)}
+    assert len(names) == 4
+    assert megatron_shard_name(0, 0, 2, ep_size=4) == "adapter_megatron_tp0_pp0_ep2.pt"
+
+
+def test_completeness_check_requires_every_realized_shard(tmp_path):
+    coords = [(0, 0, 0), (0, 0, 1), (0, 0, 2)]
+    for coord in coords[:2]:
+        (tmp_path / megatron_shard_name(*coord, 3)).touch()
+
+    assert not all_megatron_checkpoints_exist(tmp_path, _names(coords, 3))
+
+    (tmp_path / megatron_shard_name(*coords[2], 3)).touch()
+    assert all_megatron_checkpoints_exist(tmp_path, _names(coords, 3))
+
+
+def test_completeness_ignores_unrealized_coordinates(tmp_path):
+    # TP=2, EP=2, ETP=1: only (0,0,0) and (1,0,1) exist; a cross-product check
+    # would demand four shards and never resume.
+    coords = [(0, 0, 0), (1, 0, 1)]
+    for coord in coords:
+        (tmp_path / megatron_shard_name(*coord, 2)).touch()
+
+    assert all_megatron_checkpoints_exist(tmp_path, _names(coords, 2))
+
+
+def test_completeness_check_with_single_shard(tmp_path):
+    (tmp_path / "adapter_megatron_tp0_pp0.pt").touch()
+    assert all_megatron_checkpoints_exist(tmp_path, _names([(0, 0, 0)], 1))