From eecb2464d4d7e9a548674e11b61fc3d103afd1b7 Mon Sep 17 00:00:00 2001
From: Yusheng Su <yushengsu.thu@gmail.com>
Date: Fri, 24 Jul 2026 20:20:01 -0700
Subject: [PATCH 1/6] feat(multi-lora): enable and validate MoE expert adapters
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Training-side support for putting multi-LoRA adapters on MoE experts, on top of
MultiLoRAGroupedExpertLinear in megatron-bridge.

Force moe_permute_fusion off for multi-LoRA runs that target expert leaves. Most
bridge MoE providers default it on (Qwen3-MoE among them), and with it on the
dispatcher records TE's row_id_map instead of a token gather index, which expert
adapters cannot replay. Disabling it only costs the fused permute kernel, so this
turns off an optimization rather than refusing to build.

Fix two things that would have corrupted or lost expert adapters:

* slice_lora_to_rank addressed the rank axis from the front, so a packed
  grouped-expert export ([num_experts, rank, in] / [num_experts, out, rank]) was
  sliced on its expert or output axis. Address it from the end instead.
* Megatron-native adapter shards were named by (tp, pp) only and written by DP
  rank 0. Expert-parallel ranks sharing a (tp, pp) coordinate hold different
  local experts, so their shards collided and a resume loaded one EP rank's
  experts onto every rank. Shards are now keyed by (tp, pp, ep) — the suffix is
  omitted at ep_size 1 so existing checkpoints stay loadable — and both the
  writer election and the completeness check come from one cached gloo all-gather
  of the realized coordinates. Neither can be derived locally: no single group
  rank elects one writer per coordinate, and the realized coordinates are not the
  cross product of the group sizes once expert TP is smaller than tensor TP,
  which expert multi-LoRA requires.

Validate the configurations the adapter cannot serve. Model-dependent checks live
in _validate_multi_lora_moe_support, where the provider values are concrete:
expert TP must be 1, experts must be grouped (SequentialMLP linears are skipped,
so the experts would train nothing), no fp8/fp4 experts, no capacity padding, and
both expert projections must be targeted since sglang applies the expert delta
after gate_up and after down together. Launch-time checks cover pipeline size 1
and --qkv-format thd; the latter closes a pre-existing hole, as per-slot token
spans assume samples pack contiguously in the sequence-major flattening, which
bshd's [b, s] batch does not.

Expert TP is deliberately NOT checked against the CLI: its default is None and
Megatron only resolves it to tensor_model_parallel_size in its own validate_args,
which runs after miles' — so comparing the raw value rejected runs that simply
omitted the flag.

Also stop forwarding exclude_modules to MultiLoRA: ModuleMatcher asserts it is
empty whenever target_modules is set, and the exclusion has already been
subtracted from target_modules during argument validation.
---
 .../megatron_utils/bridge_lora_helpers.py     |  81 +++++++++-
 .../megatron_utils/multi_lora_utils.py        | 151 ++++++++++++++----
 miles/utils/multi_lora.py                     |  52 ++++++
 .../test_multi_lora_checkpoint_naming.py      |  57 +++++++
 .../megatron_utils/test_slice_lora_to_rank.py |  37 +++++
 tests/fast/utils/test_arguments.py            |  37 +++++
 .../fast/utils/test_targets_expert_leaves.py  |  42 +++++
 7 files changed, 422 insertions(+), 35 deletions(-)
 create mode 100644 tests/fast/backends/megatron_utils/test_multi_lora_checkpoint_naming.py
 create mode 100644 tests/fast/utils/test_targets_expert_leaves.py

diff --git a/miles/backends/megatron_utils/bridge_lora_helpers.py b/miles/backends/megatron_utils/bridge_lora_helpers.py
index 9114a14c5b5..979aa91c8ed 100644
--- a/miles/backends/megatron_utils/bridge_lora_helpers.py
+++ b/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,67 @@ 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.
+
+    These depend on the resolved Megatron config rather than the CLI, so they are
+    checked here instead of in ``validate_multi_lora_args``. Each would otherwise
+    fail deep in the first forward, or — for a non-grouped expert implementation
+    — silently leave the experts without adapters.
+    """
+    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 against the provider rather than the CLI: --expert-tensor-parallel-size
+    # defaults to None and is only resolved to tensor_model_parallel_size later, so a
+    # launch-time check on args would reject runs that simply omit the flag.
+    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}): expert TP would shard the adapter rank axis and need ETP collectives "
+        f"between the two adapter GEMMs, which the multi-slot grouped-expert layer does not "
+        f"implement. Set --expert-tensor-parallel-size 1."
+    )
+    assert getattr(provider, "moe_grouped_gemm", False), (
+        "Multi-LoRA on MoE experts requires a grouped expert implementation "
+        "(moe_grouped_gemm=True). With SequentialMLP the per-expert linears are skipped, "
+        "so the experts would train no adapter at all."
+    )
+    assert not getattr(provider, "fp8", None) and not getattr(provider, "fp4", None), (
+        "Multi-LoRA on MoE experts does not support fp8/fp4 expert quantization: the grouped "
+        "MLP pads its input to the quantization alignment, which desynchronizes the dispatched "
+        "token order the adapter's slot routing is built from."
+    )
+    # sglang applies the expert LoRA delta after gate_up and after down together, so
+    # it wraps a fused MoE layer only when both projections are served. Training just
+    # one of them would leave those expert weights loaded and never applied — the
+    # rollout policy would silently differ from the trained one.
+    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)}): sglang cannot apply an "
+            f"expert delta for only one of the two expert projections, so the missing side's "
+            f"trained weights would be 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 "
+        "(the drop-and-pad dispatch path has no slot-routing implementation)."
+    )
+    # Set above, before finalize(); assert here so a provider default that survives
+    # (or a future finalize that re-enables it) fails at build time rather than
+    # silently mis-routing every expert token.
+    assert not getattr(provider, "moe_permute_fusion", False), (
+        "Multi-LoRA on MoE experts requires moe_permute_fusion=False: the fused permute "
+        "records a row_id_map instead of a token gather index, so expert adapters cannot "
+        "follow the dispatcher's permutation."
+    )
+
+
 def _setup_lora_model_via_bridge(args: Namespace) -> list:
     """Build Megatron model with LoRA using Megatron-Bridge.
 
@@ -104,6 +168,17 @@ 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):
+        # Multi-LoRA expert adapters follow the dispatcher's token permutation, and
+        # the fused permute records TE's row_id_map instead of a token gather index
+        # (so the adapter cannot replay it). Most bridge MoE providers default this
+        # on — e.g. Qwen3-MoE — so turn it off here rather than refusing 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:
@@ -113,6 +188,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)
diff --git a/miles/backends/megatron_utils/multi_lora_utils.py b/miles/backends/megatron_utils/multi_lora_utils.py
index 707541fa4e7..3336303dff9 100644
--- a/miles/backends/megatron_utils/multi_lora_utils.py
+++ b/miles/backends/megatron_utils/multi_lora_utils.py
@@ -32,6 +32,9 @@ def create_multi_lora_instance(args: Namespace):
 
         lora_cls = LoRA
 
+    # exclude_modules is deliberately not forwarded: ModuleMatcher.match asserts it
+    # is empty whenever target_modules is set, and --exclude-modules has already
+    # been subtracted from args.target_modules during argument 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 +46,87 @@ 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:
+    """Name of the Megatron-native adapter shard for one set of parallel coordinates.
+
+    Expert-parallel ranks sharing a ``(tp, pp)`` coordinate hold *different* local
+    experts, so their adapter shards are distinct files. The suffix is omitted
+    when ``ep_size == 1`` so checkpoints written before expert adapters existed
+    stay loadable.
+    """
+    name = f"adapter_megatron_tp{tp_rank}_pp{pp_rank}"
+    if ep_size > 1:
+        name += f"_ep{ep_rank}"
+    return name + ".pt"
+
+
+_shard_topology: tuple[bool, tuple[tuple[int, int, int], ...]] | None = None
+
+
+def adapter_shard_topology() -> tuple[bool, tuple[tuple[int, int, int], ...]]:
+    """Elect one writer per adapter shard and enumerate the shards that exist.
+
+    Returns ``(this_rank_writes_its_shard, all_realized_(tp, pp, ep)_coordinates)``.
+
+    Adapter params are replicated across the data-parallel and CP dimensions but
+    not across TP, PP, or EP, so exactly one rank per ``(tp, pp, ep)`` coordinate
+    must write, and resume must wait for exactly those coordinates. Neither can be
+    derived locally:
+
+    * No single group rank elects one writer per coordinate.
+      ``intra_dp_cp.rank == 0`` is one rank for the whole DP group, so it covers
+      only one EP rank; expert-DP rank 0 covers only one TP rank, because ranks
+      sharing a coordinate differ in both CP and expert-DP index.
+    * The realized coordinates are not the ``tp x pp x ep`` cross product. When
+      expert tensor parallelism is smaller than tensor parallelism — which expert
+      multi-LoRA requires, since it needs ETP=1 — expert-parallel ranks are carved
+      out of the tensor-parallel dimension, so only a subset of ``(tp, ep)`` pairs
+      is occupied. Enumerating the cross product would wait forever for shards no
+      rank ever writes.
+
+    One gloo all-gather, cached: the topology is fixed for the run. Every caller
+    must therefore be on a rank-uniform path.
+    """
+    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
+
+    from miles.utils.distributed_utils import get_gloo_group
+
+    my_rank = dist.get_rank()
+    group = get_gloo_group()
+    gathered: list[object] = [None] * dist.get_world_size(group=group)
+    dist.all_gather_object(gathered, (coords, my_rank), group=group)
+    is_writer = my_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)
+    # Checkpoints written before expert adapters existed have no ep suffix, and all
+    # EP ranks may read the same file: without expert adapters every param in the
+    # shard is replicated across EP.
+    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 +138,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 +188,33 @@ 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 down to ``adapter_rank``.
+
+    The rank axis is the second-to-last dim of ``lora_A`` and the last dim of
+    ``lora_B``, addressed from the end so a packed grouped-expert export
+    (``[num_experts, rank, in]`` / ``[num_experts, out, rank]``) is sliced on its
+    rank axis rather than on the expert or output 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 +229,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 +243,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 +268,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}")
 
diff --git a/miles/utils/multi_lora.py b/miles/utils/multi_lora.py
index a195f731691..795de528577 100644
--- a/miles/utils/multi_lora.py
+++ b/miles/utils/multi_lora.py
@@ -46,6 +46,25 @@ def define_new_adapter_metrics(snapshot: dict) -> None:
         define_step_key_metric_group(prefix=f"{name}/perf", step_key="rollout/step")
 
 
+# MLP leaf names match both the dense MLP and the MoE experts, so their presence
+# in target_modules is the signal that adapters may land inside the experts.
+_EXPERT_LEAF_NAMES = ("linear_fc1", "linear_fc2", "gate_proj", "up_proj", "down_proj")
+# Bulk selectors that expand to include those leaves. "all-linear" is expanded to
+# concrete names during argument validation, but "all" is only resolved later by
+# the target-module conversion, so match it here too.
+_ALL_MODULE_ALIASES = ("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) for tm in (target_modules or [])]
+    if any(entry.strip().lower() in _ALL_MODULE_ALIASES for entry in entries):
+        return True
+    return any(leaf in entry for entry in entries for leaf in _EXPERT_LEAF_NAMES)
+
+
 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 +87,39 @@ 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"
+    # Multi-LoRA expert adapters use the per-expert layout only: every slot shares
+    # one max-rank buffer and an adapter's tensors are all sliced to a single rank,
+    # which the shared-outer layout's [1, r, dim] outer weights do not fit. Without
+    # this check --experts-shared-outer-loras still forces
+    # --sglang-experts-shared-outer-loras, so serving would expect a layout training
+    # never produces.
+    # Enforced at launch rather than only when the weight-sync mixin is built, and
+    # load-bearing beyond sync: per-forward adapter routing is read from state set
+    # before the forward, which an activation recompute only sees correctly while
+    # no later micro-batch's forward has run in between — true without pipelining.
+    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 address the micro-batch flattened the way Megatron
+    # flattens hidden states, which is sequence-major. Only 'thd' packs samples
+    # end-to-end so that sorting samples by slot makes each slot's tokens
+    # contiguous there; with 'bshd' the [b, s] batch interleaves samples in that
+    # flattening and every slot's span would cover the wrong tokens.
+    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 NOT validated here: --expert-tensor-parallel-size
+    # defaults to None and Megatron only resolves it to tensor_model_parallel_size
+    # in its own validate_args, which runs after this. Those checks live in
+    # _validate_multi_lora_moe_support, where the provider values are concrete and
+    # the model is known to be MoE.
     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"
diff --git a/tests/fast/backends/megatron_utils/test_multi_lora_checkpoint_naming.py b/tests/fast/backends/megatron_utils/test_multi_lora_checkpoint_naming.py
new file mode 100644
index 00000000000..7fb24ef207d
--- /dev/null
+++ b/tests/fast/backends/megatron_utils/test_multi_lora_checkpoint_naming.py
@@ -0,0 +1,57 @@
+"""Megatron-native adapter shards are keyed by parallel coordinates.
+
+Expert-parallel ranks that share a ``(tp, pp)`` coordinate hold *different* local
+experts, so their adapter shards are different files. Without the ep suffix they
+overwrite each other and a resume loads one EP rank's experts onto every rank.
+
+Resume must also wait for exactly the shards that exist. The realized
+``(tp, pp, ep)`` coordinates are not the cross product of the group sizes: with
+expert tensor parallelism smaller than tensor parallelism — what expert multi-LoRA
+requires — only a subset of ``(tp, ep)`` pairs is occupied, so enumerating the
+cross product would wait for shards no rank ever writes.
+"""
+
+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 with ETP=1: expert-parallel ranks are carved out of the
+    # tensor-parallel dimension, so (tp=0, ep=1) and (tp=1, ep=0) do not 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))
diff --git a/tests/fast/backends/megatron_utils/test_slice_lora_to_rank.py b/tests/fast/backends/megatron_utils/test_slice_lora_to_rank.py
index 6b226657662..be241bfdf17 100644
--- a/tests/fast/backends/megatron_utils/test_slice_lora_to_rank.py
+++ b/tests/fast/backends/megatron_utils/test_slice_lora_to_rank.py
@@ -46,3 +46,40 @@ def test_full_rank_tensor_passes_through():
 def test_non_lora_names_pass_through():
     tensor = torch.ones(32, 8)
     assert slice_lora_to_rank("x.some_other.weight", tensor, 16) is tensor
+
+
+# A packed grouped-expert export is [num_experts, rank, in] / [num_experts, out, rank].
+# Slicing dim 0 / dim 1 there would cut the expert or output axis instead of the
+# rank axis — dropping whole experts while reporting the requested shape.
+
+
+def test_packed_expert_lora_a_is_sliced_on_the_rank_dim():
+    # 4 experts, max rank 32, in_features 8; only the first 16 ranks are live.
+    tensor = torch.zeros(4, 32, 8)
+    tensor[:, :16] = 1.0
+    out = slice_lora_to_rank("base_model.experts.gate_proj.lora_A.weight", tensor, 16)
+    assert out.shape == (4, 16, 8)
+    assert torch.equal(out, tensor[:, :16])
+
+
+def test_packed_expert_lora_b_is_sliced_on_the_rank_dim():
+    tensor = torch.zeros(4, 8, 32)
+    tensor[..., :16] = 1.0
+    out = slice_lora_to_rank("base_model.experts.gate_proj.lora_B.weight", tensor, 16)
+    assert out.shape == (4, 8, 16)
+    assert torch.equal(out, tensor[..., :16])
+
+
+def test_packed_expert_nonzero_padding_is_rejected():
+    tensor = torch.ones(4, 32, 8)
+    with pytest.raises(AssertionError, match="padded dims are non-zero"):
+        slice_lora_to_rank("x.experts.gate_proj.lora_A.weight", tensor, 16)
+
+
+def test_packed_expert_fewer_experts_than_rank_is_not_confused():
+    # 2 experts with max rank 32 sliced to rank 4: slicing dim 0 would return
+    # only 2 experts' worth of rows and silently pass the shape check.
+    tensor = torch.zeros(2, 32, 8)
+    tensor[:, :4] = 1.0
+    out = slice_lora_to_rank("x.experts.gate_proj.lora_A.weight", tensor, 4)
+    assert out.shape == (2, 4, 8)
diff --git a/tests/fast/utils/test_arguments.py b/tests/fast/utils/test_arguments.py
index 2dbec7238f1..b9e7dcd12eb 100644
--- a/tests/fast/utils/test_arguments.py
+++ b/tests/fast/utils/test_arguments.py
@@ -251,6 +251,43 @@ def test_rejects_experimental_ft_trainer(self, monkeypatch):
         with pytest.raises(AssertionError, match="MILES_EXPERIMENTAL_FT_TRAINER"):
             miles_validate_args(args)
 
+    def test_rejects_pipeline_parallelism(self):
+        # No single rank holds a complete adapter to push, and a pipelined schedule
+        # would recompute activations against a later micro-batch's adapter routing.
+        args = self._parse([])
+        args.pipeline_model_parallel_size = 2
+        with pytest.raises(AssertionError, match="pipeline-model-parallel-size 1"):
+            miles_validate_args(args)
+
+    def test_rejects_bshd_qkv_format(self):
+        # Per-slot token spans address the micro-batch as Megatron flattens hidden
+        # states (sequence-major); a [b, s] batch interleaves samples there, so every
+        # span would cover the wrong tokens.
+        args = self._parse([])
+        args.qkv_format = "bshd"
+        with pytest.raises(AssertionError, match="qkv-format thd"):
+            miles_validate_args(args)
+
+    def test_rejects_shared_outer_expert_loras(self):
+        # Multi-LoRA expert adapters are per-expert only; the flag would otherwise
+        # still switch sglang to a layout training never produces.
+        args = self._parse([])
+        args.experts_shared_outer_loras = True
+        with pytest.raises(AssertionError, match="experts-shared-outer-loras"):
+            miles_validate_args(args)
+
+    def test_accepts_expert_leaf_targets_without_expert_tp_flag(self):
+        # --expert-tensor-parallel-size defaults to None and is only resolved to
+        # tensor_model_parallel_size by Megatron's own validate_args, which runs
+        # after this. Comparing the raw value here rejected every run that simply
+        # omitted the flag, including dense ones.
+        args = self._parse(["--target-modules", "gate_proj,up_proj,down_proj"])
+        args.expert_tensor_parallel_size = None
+
+        miles_validate_args(args)
+
+        assert args.multi_lora is True
+
 
 class TestResolveFtComponents:
     def test_disabled_with_no_components_returns_empty_without_warning(self, caplog) -> None:
diff --git a/tests/fast/utils/test_targets_expert_leaves.py b/tests/fast/utils/test_targets_expert_leaves.py
new file mode 100644
index 00000000000..517cd3dd491
--- /dev/null
+++ b/tests/fast/utils/test_targets_expert_leaves.py
@@ -0,0 +1,42 @@
+"""targets_expert_leaves decides whether adapters can land inside MoE experts.
+
+It gates the MoE-specific multi-LoRA handling (turning off permute fusion, the
+expert-parallel validations), so a false negative means those are silently skipped
+and expert tokens get routed against a permutation the adapter cannot replay.
+"""
+
+from miles.utils.multi_lora import targets_expert_leaves
+
+
+def test_mlp_leaf_names_target_experts():
+    # These names match the dense MLP and the routed experts alike.
+    assert targets_expert_leaves(["gate_proj", "up_proj", "down_proj"])
+    assert targets_expert_leaves(["linear_fc1"])
+    assert targets_expert_leaves(["linear_fc2"])
+
+
+def test_expert_scoped_wildcards_target_experts():
+    assert targets_expert_leaves(["*.layers.*.mlp.experts.linear_fc1"])
+
+
+def test_attention_only_targets_do_not():
+    assert not targets_expert_leaves(["linear_qkv", "linear_proj"])
+    assert not targets_expert_leaves(["q_proj", "k_proj", "v_proj", "o_proj"])
+
+
+def test_bulk_aliases_target_experts():
+    # "all-linear" is expanded to concrete names during argument validation, but
+    # "all" is only resolved later by the target-module conversion, so the alias
+    # itself has to count.
+    for alias in ("all", "all-linear", "all_linear", "ALL"):
+        assert targets_expert_leaves([alias]), alias
+
+
+def test_bare_string_is_accepted():
+    assert targets_expert_leaves("gate_proj")
+    assert not targets_expert_leaves("linear_qkv")
+
+
+def test_empty_targets_do_not():
+    assert not targets_expert_leaves(None)
+    assert not targets_expert_leaves([])

From 5f8ca7bb4f8a5e7ba379b8a3fcbb4febba9d3a70 Mon Sep 17 00:00:00 2001
From: Yusheng Su <yushengsu.thu@gmail.com>
Date: Sun, 26 Jul 2026 10:10:35 -0700
Subject: [PATCH 2/6] docs: trim multi-LoRA comments to 1-2 lines per review

Comment-only change addressing the 19 review comments: every flagged
docstring/comment block keeps its load-bearing constraint and drops the
derivation. No code change; the four touched fast-test files pass 48/48.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 .../megatron_utils/multi_lora_utils.py        | 52 ++++---------------
 miles/utils/multi_lora.py                     | 32 +++---------
 .../test_multi_lora_checkpoint_naming.py      | 19 ++-----
 .../megatron_utils/test_slice_lora_to_rank.py |  8 ++-
 tests/fast/utils/test_arguments.py            | 16 ++----
 .../fast/utils/test_targets_expert_leaves.py  | 12 ++---
 6 files changed, 31 insertions(+), 108 deletions(-)

diff --git a/miles/backends/megatron_utils/multi_lora_utils.py b/miles/backends/megatron_utils/multi_lora_utils.py
index 3336303dff9..dafa099e993 100644
--- a/miles/backends/megatron_utils/multi_lora_utils.py
+++ b/miles/backends/megatron_utils/multi_lora_utils.py
@@ -32,9 +32,8 @@ def create_multi_lora_instance(args: Namespace):
 
         lora_cls = LoRA
 
-    # exclude_modules is deliberately not forwarded: ModuleMatcher.match asserts it
-    # is empty whenever target_modules is set, and --exclude-modules has already
-    # been subtracted from args.target_modules during argument validation.
+    # exclude_modules was already subtracted from target_modules during arg validation
+    # (ModuleMatcher asserts it is empty when target_modules is set).
     return MultiLoRA(
         target_modules=convert_target_modules_to_megatron(args.target_modules, lora_type=lora_cls),
         n_adapters=args.multi_lora_n_adapters,
@@ -47,13 +46,8 @@ def create_multi_lora_instance(args: Namespace):
 
 
 def megatron_shard_name(tp_rank: int, pp_rank: int, ep_rank: int, ep_size: int) -> str:
-    """Name of the Megatron-native adapter shard for one set of parallel coordinates.
-
-    Expert-parallel ranks sharing a ``(tp, pp)`` coordinate hold *different* local
-    experts, so their adapter shards are distinct files. The suffix is omitted
-    when ``ep_size == 1`` so checkpoints written before expert adapters existed
-    stay loadable.
-    """
+    """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}"
@@ -64,29 +58,8 @@ def megatron_shard_name(tp_rank: int, pp_rank: int, ep_rank: int, ep_size: int)
 
 
 def adapter_shard_topology() -> tuple[bool, tuple[tuple[int, int, int], ...]]:
-    """Elect one writer per adapter shard and enumerate the shards that exist.
-
-    Returns ``(this_rank_writes_its_shard, all_realized_(tp, pp, ep)_coordinates)``.
-
-    Adapter params are replicated across the data-parallel and CP dimensions but
-    not across TP, PP, or EP, so exactly one rank per ``(tp, pp, ep)`` coordinate
-    must write, and resume must wait for exactly those coordinates. Neither can be
-    derived locally:
-
-    * No single group rank elects one writer per coordinate.
-      ``intra_dp_cp.rank == 0`` is one rank for the whole DP group, so it covers
-      only one EP rank; expert-DP rank 0 covers only one TP rank, because ranks
-      sharing a coordinate differ in both CP and expert-DP index.
-    * The realized coordinates are not the ``tp x pp x ep`` cross product. When
-      expert tensor parallelism is smaller than tensor parallelism — which expert
-      multi-LoRA requires, since it needs ETP=1 — expert-parallel ranks are carved
-      out of the tensor-parallel dimension, so only a subset of ``(tp, ep)`` pairs
-      is occupied. Enumerating the cross product would wait forever for shards no
-      rank ever writes.
-
-    One gloo all-gather, cached: the topology is fixed for the run. Every caller
-    must therefore be on a rank-uniform path.
-    """
+    """Return ``(this_rank_writes_its_shard, realized (tp, pp, ep) coords)`` from one cached
+    gloo all-gather — with ETP < TP the realized set is not the tp x pp x ep cross product."""
     global _shard_topology
     if _shard_topology is not None:
         return _shard_topology
@@ -122,9 +95,7 @@ def find_latest_checkpoint(ckpt_dir: Path) -> tuple[Path | None, int]:
 
     expected = {megatron_shard_name(*coord, ep_size) for coord in coords}
     my_shard = megatron_shard_name(*my_coords, ep_size)
-    # Checkpoints written before expert adapters existed have no ep suffix, and all
-    # EP ranks may read the same file: without expert adapters every param in the
-    # shard is replicated across EP.
+    # 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)
 
@@ -188,13 +159,8 @@ 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:
-    """Trim a max-rank-padded LoRA tensor down to ``adapter_rank``.
-
-    The rank axis is the second-to-last dim of ``lora_A`` and the last dim of
-    ``lora_B``, addressed from the end so a packed grouped-expert export
-    (``[num_experts, rank, in]`` / ``[num_experts, out, rank]``) is sliced on its
-    rank axis rather than on the expert or output axis.
-    """
+    """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]:
diff --git a/miles/utils/multi_lora.py b/miles/utils/multi_lora.py
index 795de528577..07b1c350133 100644
--- a/miles/utils/multi_lora.py
+++ b/miles/utils/multi_lora.py
@@ -46,12 +46,9 @@ def define_new_adapter_metrics(snapshot: dict) -> None:
         define_step_key_metric_group(prefix=f"{name}/perf", step_key="rollout/step")
 
 
-# MLP leaf names match both the dense MLP and the MoE experts, so their presence
-# in target_modules is the signal that adapters may land inside the experts.
+# MLP leaf names match the dense MLP and the MoE experts alike; the bulk aliases expand
+# to them ("all" is only resolved by the later target-module conversion, so match it here).
 _EXPERT_LEAF_NAMES = ("linear_fc1", "linear_fc2", "gate_proj", "up_proj", "down_proj")
-# Bulk selectors that expand to include those leaves. "all-linear" is expanded to
-# concrete names during argument validation, but "all" is only resolved later by
-# the target-module conversion, so match it here too.
 _ALL_MODULE_ALIASES = ("all", "all-linear", "all_linear")
 
 
@@ -87,26 +84,14 @@ 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"
-    # Multi-LoRA expert adapters use the per-expert layout only: every slot shares
-    # one max-rank buffer and an adapter's tensors are all sliced to a single rank,
-    # which the shared-outer layout's [1, r, dim] outer weights do not fit. Without
-    # this check --experts-shared-outer-loras still forces
-    # --sglang-experts-shared-outer-loras, so serving would expect a layout training
-    # never produces.
-    # Enforced at launch rather than only when the weight-sync mixin is built, and
-    # load-bearing beyond sync: per-forward adapter routing is read from state set
-    # before the forward, which an activation recompute only sees correctly while
-    # no later micro-batch's forward has run in between — true without pipelining.
+    # Per-forward adapter routing is only recompute-safe without pipelining, so enforce
+    # at launch rather than when the weight-sync mixin is built.
     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 address the micro-batch flattened the way Megatron
-    # flattens hidden states, which is sequence-major. Only 'thd' packs samples
-    # end-to-end so that sorting samples by slot makes each slot's tokens
-    # contiguous there; with 'bshd' the [b, s] batch interleaves samples in that
-    # flattening and every slot's span would cover the wrong tokens.
+    # 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})."
@@ -115,11 +100,8 @@ def validate_multi_lora_args(args: Any) -> None:
         "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 NOT validated here: --expert-tensor-parallel-size
-    # defaults to None and Megatron only resolves it to tensor_model_parallel_size
-    # in its own validate_args, which runs after this. Those checks live in
-    # _validate_multi_lora_moe_support, where the provider values are concrete and
-    # the model is known to be MoE.
+    # 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"
diff --git a/tests/fast/backends/megatron_utils/test_multi_lora_checkpoint_naming.py b/tests/fast/backends/megatron_utils/test_multi_lora_checkpoint_naming.py
index 7fb24ef207d..97ab9f3c5ec 100644
--- a/tests/fast/backends/megatron_utils/test_multi_lora_checkpoint_naming.py
+++ b/tests/fast/backends/megatron_utils/test_multi_lora_checkpoint_naming.py
@@ -1,15 +1,5 @@
-"""Megatron-native adapter shards are keyed by parallel coordinates.
-
-Expert-parallel ranks that share a ``(tp, pp)`` coordinate hold *different* local
-experts, so their adapter shards are different files. Without the ep suffix they
-overwrite each other and a resume loads one EP rank's experts onto every rank.
-
-Resume must also wait for exactly the shards that exist. The realized
-``(tp, pp, ep)`` coordinates are not the cross product of the group sizes: with
-expert tensor parallelism smaller than tensor parallelism — what expert multi-LoRA
-requires — only a subset of ``(tp, ep)`` pairs is occupied, so enumerating the
-cross product would wait for shards no rank ever writes.
-"""
+"""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
 
@@ -42,9 +32,8 @@ def test_completeness_check_requires_every_realized_shard(tmp_path):
 
 
 def test_completeness_ignores_unrealized_coordinates(tmp_path):
-    # TP=2, EP=2 with ETP=1: expert-parallel ranks are carved out of the
-    # tensor-parallel dimension, so (tp=0, ep=1) and (tp=1, ep=0) do not exist.
-    # A cross-product check would demand four shards and never resume.
+    # 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()
diff --git a/tests/fast/backends/megatron_utils/test_slice_lora_to_rank.py b/tests/fast/backends/megatron_utils/test_slice_lora_to_rank.py
index be241bfdf17..2df14f54b1f 100644
--- a/tests/fast/backends/megatron_utils/test_slice_lora_to_rank.py
+++ b/tests/fast/backends/megatron_utils/test_slice_lora_to_rank.py
@@ -48,9 +48,8 @@ def test_non_lora_names_pass_through():
     assert slice_lora_to_rank("x.some_other.weight", tensor, 16) is tensor
 
 
-# A packed grouped-expert export is [num_experts, rank, in] / [num_experts, out, rank].
-# Slicing dim 0 / dim 1 there would cut the expert or output axis instead of the
-# rank axis — dropping whole experts while reporting the requested shape.
+# Packed grouped-expert exports ([E, rank, in] / [E, out, rank]) must be sliced on the
+# rank axis — dim 0/1 slicing would drop experts while reporting the requested shape.
 
 
 def test_packed_expert_lora_a_is_sliced_on_the_rank_dim():
@@ -77,8 +76,7 @@ def test_packed_expert_nonzero_padding_is_rejected():
 
 
 def test_packed_expert_fewer_experts_than_rank_is_not_confused():
-    # 2 experts with max rank 32 sliced to rank 4: slicing dim 0 would return
-    # only 2 experts' worth of rows and silently pass the shape check.
+    # 2 experts, max rank 32, sliced to 4: dim-0 slicing would silently pass the shape check.
     tensor = torch.zeros(2, 32, 8)
     tensor[:, :4] = 1.0
     out = slice_lora_to_rank("x.experts.gate_proj.lora_A.weight", tensor, 4)
diff --git a/tests/fast/utils/test_arguments.py b/tests/fast/utils/test_arguments.py
index b9e7dcd12eb..c2c2c232fd9 100644
--- a/tests/fast/utils/test_arguments.py
+++ b/tests/fast/utils/test_arguments.py
@@ -252,35 +252,29 @@ def test_rejects_experimental_ft_trainer(self, monkeypatch):
             miles_validate_args(args)
 
     def test_rejects_pipeline_parallelism(self):
-        # No single rank holds a complete adapter to push, and a pipelined schedule
-        # would recompute activations against a later micro-batch's adapter routing.
+        # Adapter routing is not recompute-safe under a pipelined schedule.
         args = self._parse([])
         args.pipeline_model_parallel_size = 2
         with pytest.raises(AssertionError, match="pipeline-model-parallel-size 1"):
             miles_validate_args(args)
 
     def test_rejects_bshd_qkv_format(self):
-        # Per-slot token spans address the micro-batch as Megatron flattens hidden
-        # states (sequence-major); a [b, s] batch interleaves samples there, so every
-        # span would cover the wrong tokens.
+        # bshd interleaves samples in the sequence-major flattening the spans assume.
         args = self._parse([])
         args.qkv_format = "bshd"
         with pytest.raises(AssertionError, match="qkv-format thd"):
             miles_validate_args(args)
 
     def test_rejects_shared_outer_expert_loras(self):
-        # Multi-LoRA expert adapters are per-expert only; the flag would otherwise
-        # still switch sglang to a layout training never produces.
+        # Per-expert layout only; the flag would switch sglang to a layout training never produces.
         args = self._parse([])
         args.experts_shared_outer_loras = True
         with pytest.raises(AssertionError, match="experts-shared-outer-loras"):
             miles_validate_args(args)
 
     def test_accepts_expert_leaf_targets_without_expert_tp_flag(self):
-        # --expert-tensor-parallel-size defaults to None and is only resolved to
-        # tensor_model_parallel_size by Megatron's own validate_args, which runs
-        # after this. Comparing the raw value here rejected every run that simply
-        # omitted the flag, including dense ones.
+        # --expert-tensor-parallel-size stays None until Megatron's own validate_args;
+        # comparing the raw value here rejected every run that omitted the flag.
         args = self._parse(["--target-modules", "gate_proj,up_proj,down_proj"])
         args.expert_tensor_parallel_size = None
 
diff --git a/tests/fast/utils/test_targets_expert_leaves.py b/tests/fast/utils/test_targets_expert_leaves.py
index 517cd3dd491..9b2844a0e70 100644
--- a/tests/fast/utils/test_targets_expert_leaves.py
+++ b/tests/fast/utils/test_targets_expert_leaves.py
@@ -1,9 +1,5 @@
-"""targets_expert_leaves decides whether adapters can land inside MoE experts.
-
-It gates the MoE-specific multi-LoRA handling (turning off permute fusion, the
-expert-parallel validations), so a false negative means those are silently skipped
-and expert tokens get routed against a permutation the adapter cannot replay.
-"""
+"""targets_expert_leaves gates the MoE-specific multi-LoRA handling (permute-fusion
+off, expert validations); a false negative silently skips those."""
 
 from miles.utils.multi_lora import targets_expert_leaves
 
@@ -25,9 +21,7 @@ def test_attention_only_targets_do_not():
 
 
 def test_bulk_aliases_target_experts():
-    # "all-linear" is expanded to concrete names during argument validation, but
-    # "all" is only resolved later by the target-module conversion, so the alias
-    # itself has to count.
+    # "all" is only resolved by the later target-module conversion, so the alias itself counts.
     for alias in ("all", "all-linear", "all_linear", "ALL"):
         assert targets_expert_leaves([alias]), alias
 

From 658e1349d3b3d287e1420c0976702b30ce3c1d8e Mon Sep 17 00:00:00 2001
From: Yusheng Su <yushengsu.thu@gmail.com>
Date: Sun, 26 Jul 2026 10:19:47 -0700
Subject: [PATCH 3/6] docs: reduce bridge_lora_helpers comments and assert
 messages to one sentence per review

---
 .../megatron_utils/bridge_lora_helpers.py     | 55 ++++++-------------
 1 file changed, 17 insertions(+), 38 deletions(-)

diff --git a/miles/backends/megatron_utils/bridge_lora_helpers.py b/miles/backends/megatron_utils/bridge_lora_helpers.py
index 979aa91c8ed..a6e2558ad4e 100644
--- a/miles/backends/megatron_utils/bridge_lora_helpers.py
+++ b/miles/backends/megatron_utils/bridge_lora_helpers.py
@@ -70,63 +70,44 @@ def _get_model_config_from_wrapped(model):
 
 
 def _validate_multi_lora_moe_support(args: Namespace, provider) -> None:
-    """Reject MoE configs the multi-slot grouped-expert adapter cannot serve.
-
-    These depend on the resolved Megatron config rather than the CLI, so they are
-    checked here instead of in ``validate_multi_lora_args``. Each would otherwise
-    fail deep in the first forward, or — for a non-grouped expert implementation
-    — silently leave the experts without adapters.
-    """
+    """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 against the provider rather than the CLI: --expert-tensor-parallel-size
-    # defaults to None and is only resolved to tensor_model_parallel_size later, so a
-    # launch-time check on args would reject runs that simply omit the flag.
+    # 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}): expert TP would shard the adapter rank axis and need ETP collectives "
-        f"between the two adapter GEMMs, which the multi-slot grouped-expert layer does not "
-        f"implement. Set --expert-tensor-parallel-size 1."
+        f"{expert_tp}); set --expert-tensor-parallel-size 1."
     )
     assert getattr(provider, "moe_grouped_gemm", False), (
-        "Multi-LoRA on MoE experts requires a grouped expert implementation "
-        "(moe_grouped_gemm=True). With SequentialMLP the per-expert linears are skipped, "
-        "so the experts would train no adapter at all."
+        "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 expert quantization: the grouped "
-        "MLP pads its input to the quantization alignment, which desynchronizes the dispatched "
-        "token order the adapter's slot routing is built from."
+        "Multi-LoRA on MoE experts does not support fp8/fp4 experts (quantization padding "
+        "desynchronizes the dispatched token order)."
     )
-    # sglang applies the expert LoRA delta after gate_up and after down together, so
-    # it wraps a fused MoE layer only when both projections are served. Training just
-    # one of them would leave those expert weights loaded and never applied — the
-    # rollout policy would silently differ from the trained one.
+    # 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)}): sglang cannot apply an "
-            f"expert delta for only one of the two expert projections, so the missing side's "
-            f"trained weights would be dropped at rollout time."
+            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 "
-        "(the drop-and-pad dispatch path has no slot-routing implementation)."
+        "Multi-LoRA on MoE experts does not support --moe-pad-expert-input-to-capacity."
     )
-    # Set above, before finalize(); assert here so a provider default that survives
-    # (or a future finalize that re-enables it) fails at build time rather than
-    # silently mis-routing every expert token.
+    # Set before finalize(); assert so a surviving provider default fails at build time.
     assert not getattr(provider, "moe_permute_fusion", False), (
-        "Multi-LoRA on MoE experts requires moe_permute_fusion=False: the fused permute "
-        "records a row_id_map instead of a token gather index, so expert adapters cannot "
-        "follow the dispatcher's permutation."
+        "Multi-LoRA on MoE experts requires moe_permute_fusion=False (the fused permute's "
+        "row_id_map is not a token gather index)."
     )
 
 
@@ -169,10 +150,8 @@ def _setup_lora_model_via_bridge(args: Namespace) -> list:
     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):
-        # Multi-LoRA expert adapters follow the dispatcher's token permutation, and
-        # the fused permute records TE's row_id_map instead of a token gather index
-        # (so the adapter cannot replay it). Most bridge MoE providers default this
-        # on — e.g. Qwen3-MoE — so turn it off here rather than refusing to build.
+        # 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 "

From fc9c5218f4f5205862822cd3e93be32e7d3b7b02 Mon Sep 17 00:00:00 2001
From: Yusheng Su <yushengsu.thu@gmail.com>
Date: Sun, 26 Jul 2026 10:33:02 -0700
Subject: [PATCH 4/6] style: apply black to bridge_lora_helpers.py

---
 miles/backends/megatron_utils/bridge_lora_helpers.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/miles/backends/megatron_utils/bridge_lora_helpers.py b/miles/backends/megatron_utils/bridge_lora_helpers.py
index a6e2558ad4e..d7a3a42687b 100644
--- a/miles/backends/megatron_utils/bridge_lora_helpers.py
+++ b/miles/backends/megatron_utils/bridge_lora_helpers.py
@@ -101,9 +101,9 @@ def _validate_multi_lora_moe_support(args: Namespace, provider) -> None:
             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_pad_expert_input_to_capacity", False
+    ), "Multi-LoRA on MoE experts does not support --moe-pad-expert-input-to-capacity."
     # Set before finalize(); assert so a surviving provider default fails at build time.
     assert not getattr(provider, "moe_permute_fusion", False), (
         "Multi-LoRA on MoE experts requires moe_permute_fusion=False (the fused permute's "

From 7d907ac6b3875b10e9c7724e490192c07af41e7f Mon Sep 17 00:00:00 2001
From: Yusheng Su <yushengsu.thu@gmail.com>
Date: Sun, 26 Jul 2026 12:39:00 -0700
Subject: [PATCH 5/6] =?UTF-8?q?refactor:=20address=20review=20=E2=80=94=20?=
 =?UTF-8?q?top-level=20import/global,=20current=5Frank=20rename,=20leaf-na?=
 =?UTF-8?q?me=20mapping,=20squashed=20comments?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../megatron_utils/bridge_lora_helpers.py      |  5 ++---
 .../megatron_utils/multi_lora_utils.py         | 18 ++++++++----------
 miles/utils/multi_lora.py                      | 18 +++++++++---------
 3 files changed, 19 insertions(+), 22 deletions(-)

diff --git a/miles/backends/megatron_utils/bridge_lora_helpers.py b/miles/backends/megatron_utils/bridge_lora_helpers.py
index d7a3a42687b..f40a39fdf0d 100644
--- a/miles/backends/megatron_utils/bridge_lora_helpers.py
+++ b/miles/backends/megatron_utils/bridge_lora_helpers.py
@@ -104,10 +104,9 @@ def _validate_multi_lora_moe_support(args: Namespace, provider) -> None:
     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."
-    # Set before finalize(); assert so a surviving provider default fails at build time.
     assert not getattr(provider, "moe_permute_fusion", False), (
-        "Multi-LoRA on MoE experts requires moe_permute_fusion=False (the fused permute's "
-        "row_id_map is not a token gather index)."
+        "Multi-LoRA on MoE experts requires moe_permute_fusion=False (set off before "
+        "finalize(); the fused permute's row_id_map is not a token gather index)."
     )
 
 
diff --git a/miles/backends/megatron_utils/multi_lora_utils.py b/miles/backends/megatron_utils/multi_lora_utils.py
index dafa099e993..bbee1886a20 100644
--- a/miles/backends/megatron_utils/multi_lora_utils.py
+++ b/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,8 +36,7 @@ def create_multi_lora_instance(args: Namespace):
 
         lora_cls = LoRA
 
-    # exclude_modules was already subtracted from target_modules during arg validation
-    # (ModuleMatcher asserts it is empty when target_modules is set).
+    # 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,
@@ -54,9 +57,6 @@ def megatron_shard_name(tp_rank: int, pp_rank: int, ep_rank: int, ep_size: int)
     return name + ".pt"
 
 
-_shard_topology: tuple[bool, tuple[tuple[int, int, int], ...]] | None = None
-
-
 def adapter_shard_topology() -> tuple[bool, tuple[tuple[int, int, int], ...]]:
     """Return ``(this_rank_writes_its_shard, realized (tp, pp, ep) coords)`` from one cached
     gloo all-gather — with ETP < TP the realized set is not the tp x pp x ep cross product."""
@@ -69,13 +69,11 @@ def adapter_shard_topology() -> tuple[bool, tuple[tuple[int, int, int], ...]]:
         _shard_topology = (True, (coords,))
         return _shard_topology
 
-    from miles.utils.distributed_utils import get_gloo_group
-
-    my_rank = dist.get_rank()
+    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, my_rank), group=group)
-    is_writer = my_rank == min(rank for entry_coords, rank in gathered if entry_coords == coords)
+    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
 
diff --git a/miles/utils/multi_lora.py b/miles/utils/multi_lora.py
index 07b1c350133..67ff62d15ff 100644
--- a/miles/utils/multi_lora.py
+++ b/miles/utils/multi_lora.py
@@ -46,20 +46,21 @@ def define_new_adapter_metrics(snapshot: dict) -> None:
         define_step_key_metric_group(prefix=f"{name}/perf", step_key="rollout/step")
 
 
-# MLP leaf names match the dense MLP and the MoE experts alike; the bulk aliases expand
-# to them ("all" is only resolved by the later target-module conversion, so match it here).
-_EXPERT_LEAF_NAMES = ("linear_fc1", "linear_fc2", "gate_proj", "up_proj", "down_proj")
-_ALL_MODULE_ALIASES = ("all", "all-linear", "all_linear")
+# 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) for tm in (target_modules or [])]
-    if any(entry.strip().lower() in _ALL_MODULE_ALIASES for entry in entries):
+    entries = [str(tm).strip().lower() for tm in (target_modules or [])]
+    if any(entry in _ALL_MODULE_ALIASES for entry in entries):
         return True
-    return any(leaf in entry for entry in entries for leaf in _EXPERT_LEAF_NAMES)
+    # 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:
@@ -84,8 +85,7 @@ 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"
-    # Per-forward adapter routing is only recompute-safe without pipelining, so enforce
-    # at launch rather than when the weight-sync mixin is built.
+    # 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 "

From 8176f5dab1bc1a1526a6203d37889192cee0405d Mon Sep 17 00:00:00 2001
From: Yusheng Su <yushengsu.thu@gmail.com>
Date: Sun, 26 Jul 2026 12:46:35 -0700
Subject: [PATCH 6/6] docs: squash remaining multi-LoRA comments per review

---
 miles/backends/megatron_utils/bridge_lora_helpers.py | 7 +++----
 miles/backends/megatron_utils/multi_lora_utils.py    | 3 +--
 2 files changed, 4 insertions(+), 6 deletions(-)

diff --git a/miles/backends/megatron_utils/bridge_lora_helpers.py b/miles/backends/megatron_utils/bridge_lora_helpers.py
index f40a39fdf0d..ae2e6a880ec 100644
--- a/miles/backends/megatron_utils/bridge_lora_helpers.py
+++ b/miles/backends/megatron_utils/bridge_lora_helpers.py
@@ -104,10 +104,9 @@ def _validate_multi_lora_moe_support(args: Namespace, provider) -> None:
     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 (set off before "
-        "finalize(); the fused permute's row_id_map is not a token gather index)."
-    )
+    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:
diff --git a/miles/backends/megatron_utils/multi_lora_utils.py b/miles/backends/megatron_utils/multi_lora_utils.py
index bbee1886a20..4d6d7689098 100644
--- a/miles/backends/megatron_utils/multi_lora_utils.py
+++ b/miles/backends/megatron_utils/multi_lora_utils.py
@@ -58,8 +58,7 @@ def megatron_shard_name(tp_rank: int, pp_rank: int, ep_rank: int, ep_size: int)
 
 
 def adapter_shard_topology() -> tuple[bool, tuple[tuple[int, int, int], ...]]:
-    """Return ``(this_rank_writes_its_shard, realized (tp, pp, ep) coords)`` from one cached
-    gloo all-gather — with ETP < TP the realized set is not the tp x pp x ep cross product."""
+    """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