diff --git a/src/megatron/bridge/peft/multi_lora.py b/src/megatron/bridge/peft/multi_lora.py index 441def32c6..422848814d 100644 --- a/src/megatron/bridge/peft/multi_lora.py +++ b/src/megatron/bridge/peft/multi_lora.py @@ -29,14 +29,18 @@ from megatron.bridge.peft.base import PEFT from megatron.bridge.peft.module_matcher import ModuleMatcher -from megatron.bridge.peft.multi_lora_layers import MultiLoRALinear -from megatron.bridge.peft.utils import is_expert_linear +from megatron.bridge.peft.multi_lora_layers import ( + MultiLoRAGroupedExpertLinear, + MultiLoRALinear, + install_moe_slot_routing, +) +from megatron.bridge.peft.utils import is_expert_linear, is_grouped_expert_linear logger = logging.getLogger(__name__) -# One-shot flag so the "expert modules are skipped" warning fires once per -# process rather than once per (many) expert linear. +# One-shot flag so the "sequential expert modules are skipped" warning fires once +# per process rather than once per (many) expert linear. _EXPERT_SKIP_WARNED = False @@ -55,6 +59,9 @@ class MultiLoRA(PEFT, ModuleMatcher): lora_B_init_method: Initialisation method for the B matrix. a2a_experimental: Enable experimental all-to-all communication. lora_dtype: Data type for adapter weights. + normalize_moe_lora: Unsupported for multi-LoRA; see :meth:`__call__`. + share_expert_adapters: Unsupported for multi-LoRA; see :meth:`__call__`. + experts_shared_outer_loras: Unsupported for multi-LoRA; see :meth:`__call__`. """ target_modules: List[str] = field( @@ -69,6 +76,31 @@ class MultiLoRA(PEFT, ModuleMatcher): lora_B_init_method: str = "zero" a2a_experimental: bool = False lora_dtype: Optional[torch.dtype] = None + # Accepted (rather than rejected as unknown kwargs) so callers that share an + # argument surface with single-LoRA get an explicit error instead of a + # silently different adapter layout. Validated in __call__. + normalize_moe_lora: bool = False + share_expert_adapters: bool = False + experts_shared_outer_loras: bool = False + + def __call__(self, model, training: bool = True): + """Apply multi-LoRA, then install MoE slot routing for wrapped expert linears.""" + # Every slot shares one max-rank buffer and consumers slice all of an + # adapter's tensors to a single rank, so an expert-specific rank + # (normalize_moe_lora) or a layout that changes the exported tensor + # count per expert would break that contract rather than the layers. + for unsupported in ("normalize_moe_lora", "share_expert_adapters", "experts_shared_outer_loras"): + if getattr(self, unsupported): + raise NotImplementedError( + f"MultiLoRA does not support {unsupported}=True; expert adapters use the " + f"per-expert layout at the same max rank as every other target module." + ) + + model = super().__call__(model, training=training) + hooked = install_moe_slot_routing(model) + if hooked: + logger.info("Installed multi-lora MoE slot routing on %d MoE layer(s)", hooked) + return model def transform(self, module: nn.Module, name: Optional[str] = None, prefix: Optional[str] = None) -> nn.Module: if isinstance(module, MultiLoRALinear): @@ -78,20 +110,37 @@ def transform(self, module: nn.Module, name: Optional[str] = None, prefix: Optio (match, full_name) = ans if is_expert_linear(full_name): - # MoE expert linears are not supported by the grouped-GEMM - # multi-LoRA layer; skipping keeps them from crashing, but on an - # MoE model that silently drops the experts (most of the - # trainable capacity) from LoRA. Warn once so it isn't invisible. - global _EXPERT_SKIP_WARNED - if not _EXPERT_SKIP_WARNED: - logger.warning( - "MultiLoRA does not support MoE expert linears; skipping all expert " - "modules (e.g. %s). On MoE models only non-expert (attention/dense) " - "layers will receive LoRA adapters.", - full_name, - ) - _EXPERT_SKIP_WARNED = True - return module + if not is_grouped_expert_linear(full_name): + # SequentialMLP keeps one linear per expert + # (mlp.experts.local_experts.N.linear_fc*), each seeing only + # its own routed tokens in dispatcher order. Neither the + # dense per-slot spans nor the grouped (slot, expert) + # routing segments that, so skip with a one-shot warning. + global _EXPERT_SKIP_WARNED + if not _EXPERT_SKIP_WARNED: + logger.warning( + "MultiLoRA does not support sequential MoE expert linears; skipping " + "them (e.g. %s). Use a grouped expert implementation " + "(moe_grouped_gemm=True) to put adapters on experts.", + full_name, + ) + _EXPERT_SKIP_WARNED = True + return module + + logger.info(f"Adding multi-lora ({self.n_adapters} adapters) to expert: {full_name}") + + return MultiLoRAGroupedExpertLinear( + to_wrap=module, + n_adapters=self.n_adapters, + dim=self.dim, + alpha=self.alpha, + full_name=full_name, + num_local_experts=module.num_gemms, + column_init_method=self.lora_A_init_method, + row_init_method=self.lora_B_init_method, + dropout=self.dropout, + dropout_position=self.dropout_position, + ) if isinstance(module, TopKRouter): return module diff --git a/src/megatron/bridge/peft/multi_lora_layers.py b/src/megatron/bridge/peft/multi_lora_layers.py index 5236c518fc..f5d4e484cb 100644 --- a/src/megatron/bridge/peft/multi_lora_layers.py +++ b/src/megatron/bridge/peft/multi_lora_layers.py @@ -21,22 +21,38 @@ Forward stacks the raw weights of all adapters and uses ``torch._grouped_mm`` for a single fused kernel; TP/SP collectives are issued once around the two GEMMs to match the layout of the wrapped base linear. + +:class:`MultiLoRAGroupedExpertLinear` is the MoE counterpart, wrapping a grouped +expert linear (``mlp.experts.linear_fc{1,2}`` of a ``TEGroupedMLP``) with one +low-rank pair per (adapter slot, local expert). Inside the experts the token +order is the dispatcher's expert-major permutation rather than the micro-batch's +adapter-major order, so ``tokens_per_adapter`` alone cannot segment it; +:func:`install_moe_slot_routing` co-permutes a per-token slot id through the +dispatcher to recover the per-(slot, expert) segmentation. """ -from typing import Any, Dict, Optional, Tuple +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Sequence, Tuple, Union import torch import torch.nn as nn from megatron.core import parallel_state from megatron.core.tensor_parallel.mappings import ( + all_to_all, gather_from_sequence_parallel_region, gather_from_tensor_model_parallel_region, reduce_from_tensor_model_parallel_region, scatter_to_sequence_parallel_region, ) +from megatron.core.transformer.moe.moe_utils import sort_chunks_by_idxs from megatron.bridge.peft.adapter_wrapper import AdapterWrapper -from megatron.bridge.peft.utils import ParallelLinearAdapter, all2all_hp2sp, get_adapter_attributes_from_linear +from megatron.bridge.peft.utils import ( + GroupedExpertLinearAdapter, + ParallelLinearAdapter, + all2all_hp2sp, + get_adapter_attributes_from_linear, +) class MultiLoRALinear(AdapterWrapper): @@ -273,11 +289,275 @@ def sharded_state_dict( return sharded_sd +@dataclass +class ExpertSlotRouting: + """Per-forward mapping from dispatched expert tokens to adapter slots. + + Built once per MoE layer forward by :func:`install_moe_slot_routing`'s hook + and shared by that layer's ``linear_fc1``/``linear_fc2`` adapters, which see + the same rows in the same order. + + Attributes: + sort_idx: Permutation ordering the dispatched tokens by + ``(slot, local_expert)``, so one grouped GEMM covers every + (slot, expert) pair. + inverse_idx: Inverse of ``sort_idx``, restoring the base layer's order. + group_offsets: Inclusive cumsum of the ``n_adapters * num_local_experts`` + group sizes, in ``slot``-major order (group ``s * E + e``). + slot_token_counts: Tokens per slot, for per-token alpha/rank scaling. + num_tokens: Row count the routing was built for; guards against a base + layer that pads its input (e.g. fp8 quantization padding). + """ + + sort_idx: torch.Tensor + inverse_idx: torch.Tensor + group_offsets: torch.Tensor + slot_token_counts: torch.Tensor + num_tokens: int + + +class MultiLoRAGroupedExpertLinear(MultiLoRALinear): + """Grouped MoE expert linear wrapped with *N* concurrent LoRA adapters. + + One :class:`GroupedExpertLinearAdapter` per slot, i.e. an independent + low-rank pair per (slot, local expert). Reusing the single-LoRA adapter + class keeps the packed ``[num_local_experts, ...]`` weight layout that the + bridge's grouped-expert export and distributed checkpointing already + understand; this class only owns the multi-slot forward. + + Subclassing :class:`MultiLoRALinear` is deliberate: the slot lifecycle + helpers here and the ``isinstance``-based multi-LoRA discovery in downstream + consumers (per-slot optimizers, adapter-state zeroing) then pick expert + layers up with no changes. + + Unlike the dense layer, the wrapped base sits *inside* the experts, where + rows are the dispatcher's expert-major permutation of tokens from every EP + rank — ``tokens_per_adapter`` does not segment it. The per-forward + :class:`ExpertSlotRouting` supplies that segmentation instead. + """ + + def __init__( + self, + to_wrap: nn.Module, + n_adapters: int, + dim: int, + alpha: float, + full_name: str, + num_local_experts: int, + column_init_method: str = "xavier", + row_init_method: str = "zero", + dropout: float = 0.0, + dropout_position: str = "pre", + ) -> None: + nn.Module.__init__(self) + # Same reason as the dense layer: the grouped-GEMM forward never runs an + # adapter's own forward, so dropout would be silently dropped. + assert dropout == 0.0, ( + f"MultiLoRAGroupedExpertLinear grouped-GEMM path does not apply adapter dropout " + f"(got dropout={dropout}); set dropout/--lora-dropout to 0." + ) + self.to_wrap = to_wrap + self._adapter_enabled = True + self.n_adapters = n_adapters + self.max_rank = dim + self.base_linear_name = full_name + self.num_local_experts = num_local_experts + self._column_init_method = column_init_method + self._row_init_method = row_init_method + + # Not defensive: the dispatch requirements below are the only thing standing + # between an unsupported MoE config and silently mis-routed expert tokens, so + # a missing config must fail rather than let every check read its default. + config = to_wrap.config + expert_tp_size = parallel_state.get_expert_tensor_parallel_world_size() or getattr( + config, "expert_tensor_parallel_size", 1 + ) + # Expert TP would shard the adapter's rank axis and require ETP + # collectives between the two GEMMs. Refuse rather than inherit the + # single-LoRA grouped adapter's per-shard factorization, which is not + # equivalent to the full LoRA product for a row-parallel base. + if expert_tp_size and expert_tp_size > 1: + raise NotImplementedError( + f"Multi-LoRA on MoE experts requires expert_tensor_parallel_size=1 " + f"(got {expert_tp_size}) for {full_name}." + ) + # TEGroupedMLP pads its input to the quantization alignment before + # calling the grouped linears, which would desynchronize the row order + # the slot routing was built for. + if getattr(config, "fp8", None) or getattr(config, "fp4", None): + raise NotImplementedError( + f"Multi-LoRA on MoE experts does not support fp8/fp4 expert quantization " + f"(quantization padding changes the dispatched token order) for {full_name}." + ) + # The remaining dispatch requirements are checked here, at model build + # time, rather than in the routing hook: a hook that raises on only some + # ranks would leave its peers waiting in the companion all-to-all. + dispatcher_type = getattr(config, "moe_token_dispatcher_type", None) + if dispatcher_type != "alltoall": + raise NotImplementedError( + f"Multi-LoRA on MoE experts requires moe_token_dispatcher_type='alltoall' " + f"(got {dispatcher_type!r}) for {full_name}: the slot routing replays that " + f"dispatcher's permutation stages." + ) + # With fusion on, the dispatcher's recorded permutation is TE's row_id_map + # — a per-(expert, token) destination table with -1 for unrouted pairs — + # which is only meaningful to the matching fused unpermute, not as a + # gather index over local tokens. + if getattr(config, "moe_permute_fusion", False): + raise NotImplementedError( + f"Multi-LoRA on MoE experts requires moe_permute_fusion=False (got True) for " + f"{full_name}: the fused permute records a row_id_map instead of a token gather " + f"index, so the adapter cannot follow the dispatcher's permutation." + ) + if getattr(config, "moe_pad_expert_input_to_capacity", False): + raise NotImplementedError( + f"Multi-LoRA on MoE experts does not support moe_pad_expert_input_to_capacity " + f"for {full_name} (the drop-and-pad dispatch has no slot-routing implementation)." + ) + + attrs = get_adapter_attributes_from_linear(to_wrap, is_expert=True) + self.input_is_parallel = attrs.input_is_parallel + # With expert TP disabled the adapter needs no TP/SP collectives at all: + # the base grouped linear issues none, and the dispatched rows it is + # handed are already gathered to full sequence. + self.disable_sequence_parallel_comm = True + self.use_a2a = False + self._gather_output = False + + first_param = next(to_wrap.parameters()) + self.adapters = nn.ModuleList( + [ + GroupedExpertLinearAdapter( + attrs.in_features, + attrs.out_features, + dim, + num_local_experts=num_local_experts, + base_linear_name=full_name, + activation="identity", + column_init_method=column_init_method, + row_init_method=row_init_method, + input_is_parallel=attrs.input_is_parallel, + dropout=dropout, + dropout_position=dropout_position, + model_parallel_config=config, + alpha=alpha, + base_linear_is_parallel=attrs.base_linear_is_parallel, + params_device=first_param.device, + params_dtype=first_param.dtype, + ) + for _ in range(n_adapters) + ] + ) + + self.tokens_per_adapter: Optional[torch.Tensor] = None + # Republished by the MoE-layer forward pre-hook on every experts forward + # (including recompute replays), so it is never stale when read; the None + # here only guards a forward that runs before install_moe_slot_routing. + self.expert_slot_routing: Optional[ExpertSlotRouting] = None + device = first_param.device + dtype = first_param.dtype + self.register_buffer("alpha_values", torch.ones(n_adapters, dtype=dtype, device=device), persistent=False) + self.register_buffer( + "rank_values", torch.full((n_adapters,), dim, dtype=dtype, device=device), persistent=False + ) + + def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + linear_output, bias, layernorm_output = self.base_linear_forward(x, *args, **kwargs) + + if not self._adapter_enabled: + return linear_output, bias + + routing = self.expert_slot_routing + if routing is None: + raise RuntimeError( + f"{self.base_linear_name}: no expert slot routing for this forward. " + f"install_moe_slot_routing(model) must run after MultiLoRA is applied (it is " + f"called by MultiLoRA.__call__), and set_tokens_per_adapter_slot(model, counts) " + f"before every forward." + ) + + x_flat = layernorm_output.reshape(-1, layernorm_output.shape[-1]) + + # Every slot's weights must stay in the autograd graph even when this + # rank's experts received no tokens: Megatron's grad buffers expect a + # grad hook per trainable parameter, and a slot missing from the graph + # would leave its bucket incomplete. torch.stack below provides that in + # the normal path; the empty path needs an explicit zero term. + if routing.num_tokens == 0: + zero_term = sum((a.linear_in.weight.sum() + a.linear_out.weight.sum()) for a in self.adapters) * 0.0 + return linear_output + zero_term, bias + + if x_flat.shape[0] != routing.num_tokens: + raise RuntimeError( + f"{self.base_linear_name}: base layer received {x_flat.shape[0]} rows but the " + f"slot routing was built for {routing.num_tokens}. The base grouped linear must " + f"not pad or reorder its input between dispatch and the expert GEMMs." + ) + + # [n_adapters, num_local_experts, ...] -> one group per (slot, expert), + # slot-major so group index s*E+e matches the sort key below. The stack + # is a copy, but it is what keeps every slot in the graph (see above). + stacked_A = torch.stack([a.linear_in.weight for a in self.adapters]) + stacked_B = torch.stack([a.linear_out.weight for a in self.adapters]) + num_groups = stacked_A.shape[0] * stacked_A.shape[1] + grouped_A = stacked_A.reshape(num_groups, *stacked_A.shape[2:]) + grouped_B = stacked_B.reshape(num_groups, *stacked_B.shape[2:]) + + x_sorted = x_flat.index_select(0, routing.sort_idx) + mid = torch._grouped_mm(x_sorted, grouped_A.transpose(-2, -1), routing.group_offsets) + out = torch._grouped_mm(mid, grouped_B.transpose(-2, -1), routing.group_offsets) + + # Scaling is applied in sorted (slot-major) order, before unsorting. + # Same dtype caveat as the dense layer: the ratio is rounded to the + # activation dtype while sglang multiplies an exact fp32 ratio. + scaling = self.alpha_values / self.rank_values + out = out * torch.repeat_interleave(scaling, routing.slot_token_counts).unsqueeze(-1) + out = out.index_select(0, routing.inverse_idx) + + return linear_output + out.reshape(linear_output.shape), bias + + def reset_adapter(self, idx: int) -> None: + # Same model-parallel RNG discipline as the dense layer, but forking the + # *expert* tracker: its seed is constant across expert-data-parallel peers + # and varies across ep/etp ranks, matching how Megatron seeds the base + # expert weights. Forking the dense tracker instead would make EDP + # replicas of a reused slot diverge, since the dense tracker's seed + # varies along the dimension EDP peers differ in. + from megatron.core.tensor_parallel.random import ( + get_cuda_rng_tracker, + get_expert_parallel_rng_tracker_name, + ) + + adapter = self.adapters[idx] + col_fn = ParallelLinearAdapter._get_init_fn(None, self._column_init_method) + row_fn = ParallelLinearAdapter._get_init_fn(None, self._row_init_method) + with get_cuda_rng_tracker().fork(get_expert_parallel_rng_tracker_name()): + col_fn(adapter.linear_in.weight.data) + row_fn(adapter.linear_out.weight.data) + + def _apply_rank_mask(self, idx: int) -> None: + """Zero the padded rank rows of A and rank columns of B for slot ``idx``. + + Packed grouped-expert weights are ``[num_local_experts, rank, in]`` and + ``[num_local_experts, out, rank]``, so the rank axis is 1 and -1 + respectively, for every local expert at once. ``__init__`` rejects + expert TP, so the rank axis is never sharded and needs no local + remapping (unlike the dense layer). + """ + actual_rank = int(self.rank_values[idx].item()) + if actual_rank >= self.max_rank: + return + adapter = self.adapters[idx] + with torch.no_grad(): + adapter.linear_in.weight.data[:, actual_rank:, :].zero_() + adapter.linear_out.weight.data[..., actual_rank:].zero_() + + # ================================================================== # Standalone functions # ================================================================== -_MULTI_LORA_TYPES = MultiLoRALinear +_MULTI_LORA_TYPES = (MultiLoRALinear,) def _iter_multi_lora_modules(model): @@ -299,6 +579,210 @@ def set_tokens_per_adapter_slot(model, tokens_per_adapter: torch.Tensor) -> None module.tokens_per_adapter = tokens_per_adapter +def _split_sizes_to_list(splits) -> Optional[List[int]]: + """Normalize a dispatcher split spec to the list form ``all_to_all`` expects.""" + if splits is None: + return None + if isinstance(splits, torch.Tensor): + return splits.tolist() + return [int(s) for s in splits] + + +def _co_permute_slot_ids(dispatcher, slot_ids: torch.Tensor, num_local_experts: int) -> torch.Tensor: + """Apply the dispatcher's token permutation to a per-token adapter-slot vector. + + Mirrors :class:`MoEAlltoAllTokenDispatcher`'s dispatch stages in order — + local expert-major permute, EP all-to-all, then the local-expert sort — using + the dispatcher's own recorded metadata, so the result is aligned row-for-row + with the ``permuted_local_hidden_states`` the experts receive. + + The companion all-to-all carries one int32 per dispatched token, i.e. + ``2/hidden_size`` of the hidden-state exchange it shadows. + """ + from megatron.core.transformer.moe.token_dispatcher import MoEAlltoAllTokenDispatcher + + # Checked by class, not by attribute: the all-gather dispatcher records a + # permutation mapping too, but its stages (TP*EP all-gather, no EP + # all-to-all, no local-expert chunk sort) are not the ones replayed here. + if not isinstance(dispatcher, MoEAlltoAllTokenDispatcher): + raise NotImplementedError( + f"Multi-LoRA on MoE experts replays MoEAlltoAllTokenDispatcher's permutation " + f"stages; got {type(dispatcher).__name__}. Use moe_token_dispatcher_type='alltoall'." + ) + slot_ids = slot_ids.index_select(0, dispatcher.reversed_local_input_permutation_mapping) + + if getattr(dispatcher, "ep_size", 1) > 1: + slot_ids = all_to_all( + dispatcher.ep_group, + slot_ids, + _split_sizes_to_list(dispatcher.output_splits), + _split_sizes_to_list(dispatcher.input_splits), + ) + + # Expert TP is rejected at layer construction; assert here too because the + # dispatcher would otherwise have gathered rows this vector does not cover. + if getattr(dispatcher, "tp_size", 1) > 1: + raise NotImplementedError( + "Multi-LoRA on MoE experts requires expert_tensor_parallel_size=1, but the token " + f"dispatcher reports expert TP size {dispatcher.tp_size}." + ) + + if num_local_experts > 1: + slot_ids, _ = sort_chunks_by_idxs( + slot_ids, + dispatcher.num_global_tokens_per_local_expert.ravel(), + dispatcher.sort_input_by_local_experts, + fused=False, + ) + return slot_ids + + +def _build_expert_slot_routing( + dispatcher, + tokens_per_expert: Union[torch.Tensor, Sequence[int]], + tokens_per_adapter: torch.Tensor, + n_adapters: int, + num_local_experts: int, + device: torch.device, +) -> ExpertSlotRouting: + """Derive the per-(slot, expert) segmentation of one MoE layer's dispatched tokens. + + The checks below run before the companion all-to-all in + :func:`_co_permute_slot_ids`, so they can only fail together on all ranks — + they depend on the micro-batch and on group sizes, which are uniform within a + tensor-parallel group. A rank-local data corruption that tripped one of them + on a single rank would hang its expert-parallel peers in that all-to-all + rather than surfacing the error; every configuration-level requirement is + therefore checked at layer construction instead, where failure is uniform. + """ + slot_ids = torch.repeat_interleave( + torch.arange(n_adapters, device=device, dtype=torch.int32), + tokens_per_adapter.to(device=device), + ) + + # ``tokens_per_adapter`` counts the whole micro-batch, while a MoE layer's + # input is sequence-parallel sharded. The shard is contiguous in the same + # flattened order the dense layer's spans assume, so slicing by TP rank + # recovers this rank's slot ids. + local_tokens = int(dispatcher.hidden_shape_before_permute[0]) + if slot_ids.shape[0] != local_tokens: + tp_size = parallel_state.get_tensor_model_parallel_world_size() + if slot_ids.shape[0] != local_tokens * tp_size: + raise RuntimeError( + f"Cannot map {slot_ids.shape[0]} adapter-slot token ids onto {local_tokens} " + f"local MoE tokens with tensor_model_parallel_size={tp_size}. Check that " + f"set_tokens_per_adapter_slot() was given this micro-batch's token counts." + ) + tp_rank = parallel_state.get_tensor_model_parallel_rank() + slot_ids = slot_ids.narrow(0, tp_rank * local_tokens, local_tokens) + + slot_ids = _co_permute_slot_ids(dispatcher, slot_ids, num_local_experts) + + if isinstance(tokens_per_expert, torch.Tensor): + per_expert = tokens_per_expert.to(device=device, dtype=torch.long) + else: + per_expert = torch.tensor(list(tokens_per_expert), device=device, dtype=torch.long) + expert_ids = torch.repeat_interleave(torch.arange(num_local_experts, device=device, dtype=torch.long), per_expert) + if expert_ids.shape[0] != slot_ids.shape[0]: + raise RuntimeError( + f"Dispatched token count mismatch: tokens_per_expert sums to {expert_ids.shape[0]} " + f"but the co-permuted slot ids cover {slot_ids.shape[0]} tokens." + ) + + # Slot-major key so a stable sort yields contiguous (slot, expert) groups + # while preserving the dispatcher's expert order inside each slot. + keys = slot_ids.long() * num_local_experts + expert_ids + sort_idx = torch.argsort(keys, stable=True) + inverse_idx = torch.empty_like(sort_idx) + inverse_idx.scatter_(0, sort_idx, torch.arange(sort_idx.shape[0], device=device)) + + counts = torch.bincount(keys, minlength=n_adapters * num_local_experts) + return ExpertSlotRouting( + sort_idx=sort_idx, + inverse_idx=inverse_idx, + group_offsets=counts.cumsum(dim=0, dtype=torch.int32), + slot_token_counts=counts.view(n_adapters, num_local_experts).sum(dim=1), + num_tokens=int(sort_idx.shape[0]), + ) + + +def _make_slot_routing_hook(moe_layer: nn.Module, expert_layers: List[MultiLoRAGroupedExpertLinear]): + """Build the forward pre-hook that publishes slot routing to one MoE layer's adapters. + + The routing is rebuilt from the layer's current ``tokens_per_adapter``, so an + activation recompute must happen while that still describes the micro-batch + being recomputed. That holds without pipelining (each micro-batch's backward + immediately follows its forward) — which is the only supported multi-LoRA + configuration, since weight sync also requires + ``pipeline_model_parallel_size == 1``. A pipelined schedule would interleave a + later micro-batch's forward before the earlier one's recompute and would need + the counts carried on the graph instead. + """ + + def hook(module: nn.Module, args: Tuple[Any, ...]) -> None: + if not any(layer._adapter_enabled for layer in expert_layers): + return None + if len(args) < 2: + raise RuntimeError( + f"Expected the experts module to be called as (hidden_states, tokens_per_expert, " + f"...); got {len(args)} positional argument(s)." + ) + hidden_states, tokens_per_expert = args[0], args[1] + reference = expert_layers[0] + tokens_per_adapter = reference.tokens_per_adapter + if tokens_per_adapter is None: + raise RuntimeError( + "set_tokens_per_adapter_slot(model, adapter_token_counts) must run before every " + "forward when MoE experts carry multi-LoRA adapters." + ) + routing = _build_expert_slot_routing( + moe_layer.token_dispatcher, + tokens_per_expert, + tokens_per_adapter, + reference.n_adapters, + reference.num_local_experts, + hidden_states.device, + ) + for layer in expert_layers: + layer.expert_slot_routing = routing + return None + + return hook + + +def install_moe_slot_routing(model) -> int: + """Install per-MoE-layer slot routing for wrapped grouped expert linears. + + A forward pre-hook on each MoE layer's ``experts`` module runs after the + token dispatcher has permuted and exchanged tokens but before the expert + GEMMs, which is the only point where both the dispatcher's permutation + metadata and the final row order are available. + + Idempotent, and a no-op on models whose expert linears carry no adapters. + Returns the number of MoE layers hooked. + """ + from megatron.core.transformer.moe.moe_layer import BaseMoELayer + + installed = 0 + for model_chunk in model if isinstance(model, list) else [model]: + for module in model_chunk.modules(): + if not isinstance(module, BaseMoELayer): + continue + experts = getattr(module, "experts", None) + if experts is None: + continue + expert_layers = [m for m in experts.modules() if isinstance(m, MultiLoRAGroupedExpertLinear)] + if not expert_layers: + continue + if getattr(experts, "_multi_lora_slot_routing_handle", None) is not None: + continue + experts._multi_lora_slot_routing_handle = experts.register_forward_pre_hook( + _make_slot_routing_hook(module, expert_layers) + ) + installed += 1 + return installed + + def init_adapter_slot(model, idx: int, rank: int, alpha: float) -> None: """Claim slot ``idx`` across every multi-LoRA layer for an adapter. diff --git a/tests/unit_tests/peft/test_multi_lora.py b/tests/unit_tests/peft/test_multi_lora.py index acc8ffb489..119e575760 100644 --- a/tests/unit_tests/peft/test_multi_lora.py +++ b/tests/unit_tests/peft/test_multi_lora.py @@ -58,9 +58,27 @@ def __init__(self, to_wrap: nn.Module, **kwargs) -> None: self.init_kwargs = kwargs +class FakeMultiLoRAGroupedExpertLinear(nn.Module): + """Stand-in for ``MultiLoRAGroupedExpertLinear`` that records constructor kwargs.""" + + def __init__(self, to_wrap: nn.Module, **kwargs) -> None: + super().__init__() + self.to_wrap = to_wrap + self.init_kwargs = kwargs + + def multi_lora_linear_patch(): - """Patch ``MultiLoRALinear`` in the transform module with a recording fake.""" - return patch.object(multi_lora_module, "MultiLoRALinear", FakeMultiLoRALinear) + """Patch both multi-LoRA layer types in the transform module with recording fakes. + + ``MultiLoRA.transform`` short-circuits on ``isinstance(module, MultiLoRALinear)`` + for idempotency, and the real grouped-expert class subclasses the dense one, so + the fakes are patched together to keep that relationship out of the matching tests. + """ + return patch.multiple( + multi_lora_module, + MultiLoRALinear=FakeMultiLoRALinear, + MultiLoRAGroupedExpertLinear=FakeMultiLoRAGroupedExpertLinear, + ) def multi_lora_topk_router_patch(router_cls: type): @@ -107,8 +125,16 @@ def __init__(self) -> None: ) +class _FakeGroupedExpertLinear(nn.Linear): + """Grouped expert linear stand-in: a linear that also reports ``num_gemms``.""" + + def __init__(self, in_features: int, out_features: int, num_gemms: int = 4) -> None: + super().__init__(in_features, out_features) + self.num_gemms = num_gemms + + class MoEModel(nn.Module): - """Model with a dense MLP linear and an expert linear of the same name.""" + """Dense MLP, grouped expert, and sequential expert linears of the same name.""" def __init__(self) -> None: super().__init__() @@ -116,9 +142,13 @@ def __init__(self) -> None: self.decoder.layers = nn.ModuleList([nn.Module()]) layer = self.decoder.layers[0] layer.mlp = nn.Module() - layer.mlp.linear_fc1 = nn.Linear(32, 64) # dense -> should be wrapped + layer.mlp.linear_fc1 = nn.Linear(32, 64) # dense -> dense multi-LoRA layer.mlp.experts = nn.Module() - layer.mlp.experts.linear_fc1 = nn.Linear(32, 64) # expert -> should be skipped + # grouped (TEGroupedMLP-style) -> grouped-expert multi-LoRA + layer.mlp.experts.linear_fc1 = _FakeGroupedExpertLinear(32, 64) + # sequential (SequentialMLP-style) -> skipped + layer.mlp.experts.local_experts = nn.ModuleList([nn.Module()]) + layer.mlp.experts.local_experts[0].linear_fc1 = nn.Linear(32, 64) class _DummyTopKRouter(nn.Module): @@ -277,16 +307,36 @@ def test_transform_wildcard_matching(self) -> None: assert isinstance(transformed.layers[1]["attention"]["linear_qkv"], nn.Linear) assert isinstance(transformed.layers[1]["mlp"]["linear_fc2"], nn.Linear) - def test_transform_skips_expert_linear(self) -> None: + def test_transform_routes_expert_linears_by_implementation(self) -> None: model = MoEModel() peft = MultiLoRA(target_modules=["linear_fc1"]) transformed = peft(model, training=True) layer = transformed.decoder.layers[0] - # Dense MLP linear is wrapped; the routed-expert linear of the same name is skipped. + # Dense MLP linear gets the dense layer; the grouped expert linear of the + # same name gets the grouped-expert layer; the sequential per-expert + # linear is left alone (its token order is not slot-segmentable). assert isinstance(layer.mlp.linear_fc1, FakeMultiLoRALinear) - assert isinstance(layer.mlp.experts.linear_fc1, nn.Linear) + assert isinstance(layer.mlp.experts.linear_fc1, FakeMultiLoRAGroupedExpertLinear) + assert isinstance(layer.mlp.experts.local_experts[0].linear_fc1, nn.Linear) + + def test_transform_passes_local_expert_count(self) -> None: + model = MoEModel() + peft = MultiLoRA(target_modules=["linear_fc1"], n_adapters=3, dim=8, alpha=16) + + transformed = peft(model, training=True) + + wrapped = transformed.decoder.layers[0].mlp.experts.linear_fc1 + assert wrapped.init_kwargs["num_local_experts"] == 4 + assert wrapped.init_kwargs["n_adapters"] == 3 + assert wrapped.init_kwargs["dim"] == 8 + + @pytest.mark.parametrize("flag", ["normalize_moe_lora", "share_expert_adapters", "experts_shared_outer_loras"]) + def test_unsupported_expert_layouts_raise(self, flag: str) -> None: + peft = MultiLoRA(target_modules=["linear_fc1"], **{flag: True}) + with pytest.raises(NotImplementedError, match=flag): + peft(MoEModel(), training=True) def test_transform_skips_topk_router(self) -> None: router = _DummyTopKRouter() diff --git a/tests/unit_tests/peft/test_multi_lora_layers.py b/tests/unit_tests/peft/test_multi_lora_layers.py index 46b6dabf55..4486803c5a 100644 --- a/tests/unit_tests/peft/test_multi_lora_layers.py +++ b/tests/unit_tests/peft/test_multi_lora_layers.py @@ -27,7 +27,8 @@ Mock-level (no distributed): * B7: grouped-GEMM path rejects adapter dropout > 0 (it cannot apply it) * B2: expose_adapter_slot / hide_adapters restore the ModuleList even if the body raises - * B8: MoE expert linears are skipped with a one-time warning, not silently + * B8: sequential MoE expert linears are skipped with a one-time warning, not + silently; grouped MoE expert linears are wrapped with the grouped-expert layer * B9: load_adapter raises on a checkpoint/model mismatch in either direction (params missing from the checkpoint, or checkpoint tensors no param consumed) @@ -563,27 +564,54 @@ def test_expose_adapter_slot_restores_on_success(): # --------------------------------------------------------------------------- # -# B8: expert linears are skipped, but with a one-time warning (not silently). +# B8: sequential expert linears are skipped, but with a one-time warning +# (grouped expert linears are wrapped — see test_grouped_expert_linear_wrapped). # --------------------------------------------------------------------------- # -def test_expert_skip_warns_once(): +def test_sequential_expert_skip_warns_once(): multi_lora_mod._EXPERT_SKIP_WARNED = False mlora = MultiLoRA(target_modules=["linear_fc1"], n_adapters=2, dim=8, alpha=16) module = nn.Linear(4, 4) - full = "decoder.layers.0.mlp.experts.linear_fc1" + prefix = "decoder.layers.0.mlp.experts.local_experts.0." + full = prefix + "linear_fc1" with ( - patch.object(multi_lora_mod, "is_expert_linear", return_value=True), patch.object(mlora, "match", return_value=(MagicMock(), full)), patch.object(multi_lora_mod, "logger") as logmock, ): - out1 = mlora.transform(module, name="linear_fc1", prefix="decoder.layers.0.mlp.experts.") - out2 = mlora.transform(module, name="linear_fc1", prefix="decoder.layers.0.mlp.experts.") + out1 = mlora.transform(module, name="linear_fc1", prefix=prefix) + out2 = mlora.transform(module, name="linear_fc1", prefix=prefix) - # expert modules are returned unwrapped... + # sequential expert modules are returned unwrapped... assert out1 is module and out2 is module # ...and the warning fires exactly once across both skips assert logmock.warning.call_count == 1 +def test_grouped_expert_linear_wrapped(): + """A grouped expert linear gets the multi-slot grouped-expert layer.""" + mlora = MultiLoRA(target_modules=["linear_fc1"], n_adapters=2, dim=8, alpha=16) + module = nn.Linear(4, 4) + module.num_gemms = 3 + prefix = "decoder.layers.0.mlp.experts." + full = prefix + "linear_fc1" + + recorded = {} + + class _Fake(nn.Module): + def __init__(self, **kwargs): + super().__init__() + recorded.update(kwargs) + + with ( + patch.object(mlora, "match", return_value=(MagicMock(), full)), + patch.object(multi_lora_mod, "MultiLoRAGroupedExpertLinear", _Fake), + ): + out = mlora.transform(module, name="linear_fc1", prefix=prefix) + + assert isinstance(out, _Fake) + assert recorded["num_local_experts"] == 3 + assert recorded["full_name"] == full + + # --------------------------------------------------------------------------- # # B9: load_adapter raises on a checkpoint/model mismatch in either direction. # --------------------------------------------------------------------------- # diff --git a/tests/unit_tests/peft/test_multi_lora_moe.py b/tests/unit_tests/peft/test_multi_lora_moe.py new file mode 100644 index 0000000000..6f5087936f --- /dev/null +++ b/tests/unit_tests/peft/test_multi_lora_moe.py @@ -0,0 +1,574 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for multi-adapter LoRA on grouped MoE expert linears. + +Two layers of coverage: + +CPU (fake token dispatcher, EP=1, no collectives): + * the co-permuted slot ids match the permutation the dispatcher applied + * ``sort_idx``/``inverse_idx``/``group_offsets``/``slot_token_counts`` agree with a + reference stable sort by ``(slot, local_expert)`` + * the sequence-parallel narrow picks this rank's shard of the micro-batch + * unsupported dispatch configurations raise instead of silently mis-routing + * ``install_moe_slot_routing`` is idempotent and skips adapter-free MoE layers + +Single-GPU integration (needs CUDA + model-parallel init): + * the grouped forward equals a per-(slot, expert) reference LoRA computation + * every slot's parameters receive a gradient, including slots that got no tokens + and the all-empty case (Megatron's grad buffers expect one hook per parameter) +""" + +import os +from unittest.mock import patch + +import pytest +import torch +import torch.nn as nn + +from megatron.bridge.peft.multi_lora_layers import ( + MultiLoRAGroupedExpertLinear, + _build_expert_slot_routing, + _co_permute_slot_ids, + install_moe_slot_routing, +) + + +# ====================================================================== +# Test doubles +# ====================================================================== + + +def _fake_dispatcher_cls(): + """Subclass the real dispatcher so the routing's isinstance gate is satisfied. + + Only the attributes the slot routing reads are populated; the real ``__init__`` + needs a full transformer config and process groups. + """ + from megatron.core.transformer.moe.token_dispatcher import MoEAlltoAllTokenDispatcher + + class _FakeAlltoAllDispatcher(MoEAlltoAllTokenDispatcher): + def __init__( + self, + sorted_indices: torch.Tensor, + num_local_tokens: int, + *, + num_local_experts: int = 1, + tokens_per_expert: torch.Tensor | None = None, + ep_size: int = 1, + tp_size: int = 1, + ) -> None: + self.reversed_local_input_permutation_mapping = sorted_indices + self.hidden_shape_before_permute = (num_local_tokens, 8) + self.ep_size = ep_size + self.tp_size = tp_size + self.num_local_experts = num_local_experts + # With EP=1 there is a single source chunk per local expert, so the + # local-expert sort is the identity — as it is in Megatron for ep*tp == 1. + if tokens_per_expert is not None: + self.num_global_tokens_per_local_expert = tokens_per_expert.view(1, -1) + self.sort_input_by_local_experts = torch.arange(num_local_experts) + + return _FakeAlltoAllDispatcher + + +def _FakeDispatcher(*args, **kwargs): # noqa: N802 - reads as a class at call sites + return _fake_dispatcher_cls()(*args, **kwargs) + + +def _expert_major_permutation(token_experts: list[int], num_experts: int) -> torch.Tensor: + """Indices that group tokens by expert, mirroring ``moe_utils.permute`` for topk=1. + + ``permute`` builds these with ``token_indices.masked_select(routing_map.T)``, + i.e. expert-major with token order preserved inside each expert. + """ + order: list[int] = [] + for expert in range(num_experts): + order.extend(i for i, e in enumerate(token_experts) if e == expert) + return torch.tensor(order, dtype=torch.long) + + +def _reference_routing(slot_ids: list[int], expert_ids: list[int], n_adapters: int, num_local_experts: int): + """Reference (slot, expert) segmentation computed with plain Python.""" + keys = [s * num_local_experts + e for s, e in zip(slot_ids, expert_ids)] + order = sorted(range(len(keys)), key=lambda i: (keys[i], i)) # stable + counts = [0] * (n_adapters * num_local_experts) + for k in keys: + counts[k] += 1 + return order, counts + + +# ====================================================================== +# CPU: slot-id co-permutation and segmentation +# ====================================================================== + + +def test_co_permute_follows_dispatcher_permutation(): + # 6 tokens, slots [0,0,0,1,1,1], routed to experts [1,0,1,0,1,0]. + token_experts = [1, 0, 1, 0, 1, 0] + sorted_indices = _expert_major_permutation(token_experts, num_experts=2) + slot_ids = torch.tensor([0, 0, 0, 1, 1, 1], dtype=torch.int32) + dispatcher = _FakeDispatcher( + sorted_indices, + num_local_tokens=6, + num_local_experts=2, + tokens_per_expert=torch.tensor([3, 3]), + ) + + permuted = _co_permute_slot_ids(dispatcher, slot_ids, num_local_experts=2) + + # Expert 0 receives tokens 1,3,5 (slots 0,1,1); expert 1 receives 0,2,4 (slots 0,0,1). + assert permuted.tolist() == [0, 1, 1, 0, 0, 1] + + +def test_routing_matches_reference_segmentation(): + token_experts = [1, 0, 1, 0, 1, 0, 0, 1] + sorted_indices = _expert_major_permutation(token_experts, num_experts=2) + tokens_per_expert = torch.tensor([4, 4]) + dispatcher = _FakeDispatcher( + sorted_indices, + num_local_tokens=8, + num_local_experts=2, + tokens_per_expert=tokens_per_expert, + ) + tokens_per_adapter = torch.tensor([3, 5], dtype=torch.int32) + + routing = _build_expert_slot_routing( + dispatcher, + tokens_per_expert, + tokens_per_adapter, + n_adapters=2, + num_local_experts=2, + device=torch.device("cpu"), + ) + + dispatched_slots = [0, 0, 0, 1, 1, 1, 1, 1] + permuted_slots = [dispatched_slots[i] for i in sorted_indices.tolist()] + permuted_experts = [token_experts[i] for i in sorted_indices.tolist()] + expected_order, expected_counts = _reference_routing(permuted_slots, permuted_experts, 2, 2) + + assert routing.num_tokens == 8 + assert routing.sort_idx.tolist() == expected_order + assert routing.group_offsets.tolist() == torch.tensor(expected_counts).cumsum(0).tolist() + assert routing.slot_token_counts.tolist() == [3, 5] + assert routing.group_offsets.dtype == torch.int32 + + +def test_inverse_index_round_trips(): + token_experts = [0, 1, 0, 1] + sorted_indices = _expert_major_permutation(token_experts, num_experts=2) + tokens_per_expert = torch.tensor([2, 2]) + routing = _build_expert_slot_routing( + _FakeDispatcher(sorted_indices, num_local_tokens=4, num_local_experts=2, tokens_per_expert=tokens_per_expert), + tokens_per_expert, + torch.tensor([2, 2], dtype=torch.int32), + n_adapters=2, + num_local_experts=2, + device=torch.device("cpu"), + ) + + rows = torch.arange(4 * 3, dtype=torch.float32).reshape(4, 3) + assert torch.equal(rows.index_select(0, routing.sort_idx).index_select(0, routing.inverse_idx), rows) + + +def test_sequence_parallel_narrow_selects_this_rank_shard(): + """With SP the MoE input is a contiguous shard of the micro-batch's tokens.""" + # 4 local tokens out of an 8-token micro-batch: this is TP rank 1's half, + # whose slot ids are the second half of [0,0,0,0,1,1,1,1]. + token_experts = [0, 0, 1, 1] + sorted_indices = _expert_major_permutation(token_experts, num_experts=2) + tokens_per_expert = torch.tensor([2, 2]) + dispatcher = _FakeDispatcher( + sorted_indices, num_local_tokens=4, num_local_experts=2, tokens_per_expert=tokens_per_expert + ) + + with ( + patch("megatron.core.parallel_state.get_tensor_model_parallel_world_size", return_value=2), + patch("megatron.core.parallel_state.get_tensor_model_parallel_rank", return_value=1), + ): + routing = _build_expert_slot_routing( + dispatcher, + tokens_per_expert, + torch.tensor([4, 4], dtype=torch.int32), + n_adapters=2, + num_local_experts=2, + device=torch.device("cpu"), + ) + + # All four local tokens belong to slot 1 (the second half of the micro-batch). + assert routing.slot_token_counts.tolist() == [0, 4] + + +def test_token_count_mismatch_raises(): + token_experts = [0, 1] + dispatcher = _FakeDispatcher( + _expert_major_permutation(token_experts, num_experts=2), num_local_tokens=2, num_local_experts=1 + ) + with ( + patch("megatron.core.parallel_state.get_tensor_model_parallel_world_size", return_value=1), + pytest.raises(RuntimeError, match="adapter-slot token ids"), + ): + _build_expert_slot_routing( + dispatcher, + torch.tensor([2]), + torch.tensor([2, 3], dtype=torch.int32), # sums to 5, not 2 + n_adapters=2, + num_local_experts=1, + device=torch.device("cpu"), + ) + + +def test_expert_tensor_parallel_dispatch_raises(): + slot_ids = torch.zeros(4, dtype=torch.int32) + expert_tp = _FakeDispatcher(torch.arange(4), num_local_tokens=4, tp_size=2) + with pytest.raises(NotImplementedError, match="expert_tensor_parallel_size=1"): + _co_permute_slot_ids(expert_tp, slot_ids, num_local_experts=1) + + +def test_non_alltoall_dispatcher_raises(): + """The all-gather dispatcher records a permutation mapping too, so the gate is by type. + + Its dispatch stages differ completely (TP*EP all-gather, no EP all-to-all, no + local-expert chunk sort), so replaying the all-to-all stages against it would + silently pair tokens with the wrong slots. + """ + + class _AllGatherLike: + # Same attribute MoEAllGatherTokenDispatcher sets. + reversed_local_input_permutation_mapping = torch.arange(4) + ep_size = 1 + tp_size = 1 + num_local_experts = 1 + + with pytest.raises(NotImplementedError, match="MoEAlltoAllTokenDispatcher"): + _co_permute_slot_ids(_AllGatherLike(), torch.zeros(4, dtype=torch.int32), num_local_experts=1) + + +# ====================================================================== +# CPU: hook installation +# ====================================================================== + + +class _FakeExperts(nn.Module): + def __init__(self, child: nn.Module | None) -> None: + super().__init__() + if child is not None: + self.linear_fc1 = child + + def forward(self, hidden, tokens_per_expert, probs): # pragma: no cover - never called + return hidden, None + + +class _FakeMoELayer(nn.Module): + """Stands in for the isinstance check only; BaseMoELayer needs a full config.""" + + def __init__(self, child: nn.Module | None) -> None: + super().__init__() + self.experts = _FakeExperts(child) + self.token_dispatcher = None + + +def test_install_moe_slot_routing_is_idempotent_and_selective(): + wrapped = MultiLoRAGroupedExpertLinear.__new__(MultiLoRAGroupedExpertLinear) + nn.Module.__init__(wrapped) + + with_adapter = _FakeMoELayer(wrapped) + without_adapter = _FakeMoELayer(nn.Linear(4, 4)) + model = nn.ModuleList([with_adapter, without_adapter]) + + # install_moe_slot_routing imports BaseMoELayer at call time, so patching the + # source module is enough to make the fake layers match. + with patch("megatron.core.transformer.moe.moe_layer.BaseMoELayer", _FakeMoELayer): + assert install_moe_slot_routing([model]) == 1 + # A second pass must not stack a duplicate hook on the same experts module. + assert install_moe_slot_routing([model]) == 0 + + assert len(with_adapter.experts._forward_pre_hooks) == 1 + assert len(without_adapter.experts._forward_pre_hooks) == 0 + + +# ====================================================================== +# Single-GPU integration +# ====================================================================== + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU + model-parallel init") +class TestMultiLoRAGroupedExpertLinearGPU: + @pytest.fixture(autouse=True) + def _mp(self): + import megatron.core.parallel_state as parallel_state + import torch.distributed as dist + + if not dist.is_initialized(): + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", "29557") + os.environ.setdefault("RANK", "0") + os.environ.setdefault("LOCAL_RANK", "0") + os.environ.setdefault("WORLD_SIZE", "1") + torch.cuda.set_device(0) + dist.init_process_group(backend="nccl", world_size=1, rank=0) + if not parallel_state.model_parallel_is_initialized(): + parallel_state.initialize_model_parallel(tensor_model_parallel_size=1, pipeline_model_parallel_size=1) + from megatron.core.process_groups_config import ProcessGroupCollection + + from megatron.bridge.training.initialize import _set_random_seed + + _set_random_seed( + seed_=1234, + data_parallel_random_init=False, + te_rng_tracker=True, + inference_rng_tracker=False, + pg_collection=ProcessGroupCollection.use_mpu_process_groups(), + ) + yield + try: + if parallel_state.model_parallel_is_initialized(): + parallel_state.destroy_model_parallel() + if dist.is_initialized(): + dist.destroy_process_group() + except Exception: + pass + + def _build_base(self, *, hidden=16, ffn=32, num_local_experts=2, **config_overrides): + """A grouped expert linear as TEGroupedMLP builds its ``linear_fc1``.""" + from megatron.core.extensions.transformer_engine import TEColumnParallelGroupedLinear + from megatron.core.transformer.transformer_config import TransformerConfig + + from megatron.bridge.peft.utils import init_method_normal + + config = TransformerConfig( + num_layers=1, + hidden_size=hidden, + ffn_hidden_size=ffn, + num_attention_heads=1, + num_moe_experts=num_local_experts, + moe_grouped_gemm=True, + # Not the TransformerConfig default ('allgather'), and required: the + # slot routing replays this dispatcher's permutation stages. + moe_token_dispatcher_type="alltoall", + sequence_parallel=False, + tensor_model_parallel_size=1, + expert_tensor_parallel_size=1, + bf16=True, + params_dtype=torch.bfloat16, + **config_overrides, + ) + base = TEColumnParallelGroupedLinear( + num_gemms=num_local_experts, + input_size=hidden, + output_size=ffn, + config=config, + init_method=init_method_normal(0.02), + bias=False, + skip_bias_add=True, + is_expert=True, + tp_comm_buffer_name="fc1", + ).cuda() + return base, config + + def _build(self, *, hidden=16, ffn=32, num_local_experts=2, n_adapters=2, dim=8, alpha=16): + base, config = self._build_base(hidden=hidden, ffn=ffn, num_local_experts=num_local_experts) + layer = MultiLoRAGroupedExpertLinear( + to_wrap=base, + n_adapters=n_adapters, + dim=dim, + alpha=alpha, + full_name="decoder.layers.0.mlp.experts.linear_fc1", + num_local_experts=num_local_experts, + ) + layer.adapters.to(device="cuda", dtype=torch.bfloat16) + return layer, config + + @staticmethod + def _routing_for(slots: list[int], experts: list[int], n_adapters: int, num_local_experts: int): + """Build ExpertSlotRouting directly for a known (slot, expert) assignment.""" + from megatron.bridge.peft.multi_lora_layers import ExpertSlotRouting + + keys = torch.tensor( + [s * num_local_experts + e for s, e in zip(slots, experts)], device="cuda", dtype=torch.long + ) + sort_idx = torch.argsort(keys, stable=True) + inverse_idx = torch.empty_like(sort_idx) + inverse_idx.scatter_(0, sort_idx, torch.arange(keys.numel(), device="cuda")) + counts = torch.bincount(keys, minlength=n_adapters * num_local_experts) + return ExpertSlotRouting( + sort_idx=sort_idx, + inverse_idx=inverse_idx, + group_offsets=counts.cumsum(dim=0, dtype=torch.int32), + slot_token_counts=counts.view(n_adapters, num_local_experts).sum(dim=1), + num_tokens=keys.numel(), + ) + + def test_forward_matches_per_token_reference(self): + layer, _ = self._build() + # Tokens are expert-major (as they arrive from the dispatcher) with slots + # interleaved inside each expert — the case a contiguous per-slot span + # cannot express. + experts = [0, 0, 0, 1, 1, 1] + slots = [0, 1, 0, 1, 1, 0] + for idx, rank, alpha in ((0, 8, 16.0), (1, 4, 8.0)): + layer.init_adapter_slot(idx, rank=rank, alpha=alpha) + # B is zero-initialised, so give it content or every delta is zero. + for adapter in layer.adapters: + with torch.no_grad(): + adapter.linear_out.weight.normal_(std=0.02) + layer._apply_rank_mask(1) + + x = torch.randn(len(slots), 16, dtype=torch.bfloat16, device="cuda") + layer.expert_slot_routing = self._routing_for(slots, experts, 2, 2) + # m_splits is how TEGroupedMLP passes tokens_per_expert to the grouped linear. + out, _ = layer(x, [3, 3]) + + base_out, _ = layer.to_wrap(x, [3, 3]) + expected = base_out.clone() + for i, (slot, expert) in enumerate(zip(slots, experts)): + adapter = layer.adapters[slot] + scale = layer.alpha_values[slot] / layer.rank_values[slot] + mid = x[i] @ adapter.linear_in.weight[expert].T + expected[i] += (mid @ adapter.linear_out.weight[expert].T) * scale + + torch.testing.assert_close(out, expected, rtol=2e-2, atol=2e-2) + + def test_every_slot_gets_a_gradient(self): + """Slot 1 receives no tokens, but Megatron's grad buffers still need its hook.""" + layer, _ = self._build() + for idx in (0, 1): + layer.init_adapter_slot(idx, rank=8, alpha=16.0) + for adapter in layer.adapters: + with torch.no_grad(): + adapter.linear_out.weight.normal_(std=0.02) + + x = torch.randn(4, 16, dtype=torch.bfloat16, device="cuda") + layer.expert_slot_routing = self._routing_for([0, 0, 0, 0], [0, 0, 1, 1], 2, 2) + out, _ = layer(x, [2, 2]) + out.sum().backward() + + for idx, adapter in enumerate(layer.adapters): + assert adapter.linear_in.weight.grad is not None, f"slot {idx} linear_in got no grad" + assert adapter.linear_out.weight.grad is not None, f"slot {idx} linear_out got no grad" + + def test_no_tokens_still_produces_gradients(self): + layer, config = self._build() + for idx in (0, 1): + layer.init_adapter_slot(idx, rank=8, alpha=16.0) + + # The base grouped linear has its own all-empty handling; stub it out so + # this exercises only the adapter's empty-batch path. + class _EmptyBase(nn.Module): + def forward(self, x, *args, **kwargs): + return x.new_zeros((x.shape[0], config.ffn_hidden_size), requires_grad=True), None + + layer.to_wrap = _EmptyBase() + + x = torch.zeros(0, 16, dtype=torch.bfloat16, device="cuda") + layer.expert_slot_routing = self._routing_for([], [], 2, 2) + out, _ = layer(x, [0, 0]) + out.sum().backward() + + for idx, adapter in enumerate(layer.adapters): + assert adapter.linear_in.weight.grad is not None, f"slot {idx} linear_in got no grad" + assert adapter.linear_out.weight.grad is not None, f"slot {idx} linear_out got no grad" + + def test_rank_mask_zeroes_padded_rank_on_every_expert(self): + layer, _ = self._build(dim=8) + with torch.no_grad(): + for adapter in layer.adapters: + adapter.linear_in.weight.fill_(1.0) + adapter.linear_out.weight.fill_(1.0) + + layer.init_adapter_slot(0, rank=3, alpha=6.0) + + adapter = layer.adapters[0] + assert adapter.linear_in.weight[:, 3:, :].abs().max().item() == 0 + assert adapter.linear_out.weight[..., 3:].abs().max().item() == 0 + assert adapter.linear_in.weight[:, :3, :].abs().min().item() == 1 + assert adapter.linear_out.weight[..., :3].abs().min().item() == 1 + + def test_reset_adapter_draws_from_the_expert_rng_stream(self): + """Expert adapters must re-init from the expert tracker, not the dense one. + + The dense model-parallel seed varies with the *dense* TP rank, and + expert-data-parallel peers differ in exactly that rank whenever + tensor_model_parallel_size != expert_tensor_parallel_size (miles' MoE + recipes run TP>1 with ETP=1). Drawing from the dense stream there would + give each EDP replica of a reused slot different weights. + """ + import megatron.core.tensor_parallel.random as mcore_random + + layer, _ = self._build() + real_tracker = mcore_random.get_cuda_rng_tracker() + forked: list[object] = [] + + class _RecordingTracker: + def fork(self, name=None): + forked.append(name) + return real_tracker.fork(name) if name is not None else real_tracker.fork() + + with patch.object(mcore_random, "get_cuda_rng_tracker", lambda: _RecordingTracker()): + layer.clear_adapter_slot(0) + + assert forked == [mcore_random.get_expert_parallel_rng_tracker_name()], ( + f"expert adapter re-init forked {forked}, not the expert rng stream" + ) + # The stream must also actually exist, or the fork would raise at runtime. + assert mcore_random.get_expert_parallel_rng_tracker_name() in real_tracker.get_states() + + @pytest.mark.parametrize( + "field, value, message", + [ + # Most bridge MoE providers default fusion on (e.g. Qwen3-MoE), and the + # fused permute records TE's row_id_map rather than a token gather index. + ("moe_permute_fusion", True, "moe_permute_fusion=False"), + ("moe_token_dispatcher_type", "allgather", "alltoall"), + ("moe_pad_expert_input_to_capacity", True, "moe_pad_expert_input_to_capacity"), + ], + ) + def test_unsupported_dispatch_configs_are_rejected_at_build_time(self, field, value, message): + """Config-level requirements must fail uniformly at wrap time, not inside the hook. + + A hook that raised on only some ranks would leave its expert-parallel peers + waiting in the companion all-to-all. + + The field is set on the resolved config rather than passed to + ``TransformerConfig``, whose own validation would otherwise reject some of + these combinations before the adapter is ever constructed. + """ + base, config = self._build_base() + setattr(config, field, value) + with pytest.raises(NotImplementedError, match=message): + MultiLoRAGroupedExpertLinear( + to_wrap=base, + n_adapters=2, + dim=8, + alpha=16, + full_name="decoder.layers.0.mlp.experts.linear_fc1", + num_local_experts=2, + ) + + def test_expert_tensor_parallel_is_rejected(self): + # Build the base with real (ETP=1) state, then construct the adapter layer + # under a patched ETP size so only the layer's own guard is exercised. + base, _ = self._build_base() + with ( + patch("megatron.core.parallel_state.get_expert_tensor_parallel_world_size", return_value=2), + pytest.raises(NotImplementedError, match="expert_tensor_parallel_size=1"), + ): + MultiLoRAGroupedExpertLinear( + to_wrap=base, + n_adapters=2, + dim=8, + alpha=16, + full_name="decoder.layers.0.mlp.experts.linear_fc1", + num_local_experts=2, + )