feat(peft): support multi-LoRA on grouped MoE expert linears

[yushengsu-thu]

Why

MultiLoRA skipped every MoE expert linear, so on MoE models the experts — most of the trainable capacity — silently got no adapter. This brings the multi-adapter path to parity with single-adapter LoRA for grouped experts.

What

• MultiLoRAGroupedExpertLinear wraps mlp.experts.linear_fc{1,2} of a TEGroupedMLP with one low-rank pair per (slot, local expert). Each slot is an existing GroupedExpertLinearAdapter, so the packed [num_local_experts, ...] export and checkpoint layout is reused unchanged; subclassing MultiLoRALinear keeps isinstance-based multi-LoRA discovery working with no downstream changes.
• Slot routing. Inside the experts, rows are the dispatcher's expert-major permutation, so tokens_per_adapter no longer segments them. install_moe_slot_routing puts a forward pre-hook on each MoE layer's experts module that co-permutes an int32 slot-id vector through the dispatcher's own stages (local permute, EP all-to-all, local-expert sort), then stable-sorts by (slot, expert) so one grouped GEMM covers every pair.
• Every slot's weights stay in the autograd graph (stacked weights; explicit zero term on the empty-batch path) — Megatron's grad buffers expect one grad hook per trainable param.

Rejected configurations

Raised at layer construction, not in the hook (a hook raising on some ranks would strand peers in the companion all-to-all): expert TP > 1, moe_permute_fusion (records TE's row_id_map, not a token gather index), fp8/fp4 experts (input padding desyncs the row order), capacity padding, non-alltoall dispatchers. SequentialMLP experts stay skipped with a clearer warning; normalize_moe_lora / share_expert_adapters / experts_shared_outer_loras raise NotImplementedError.

Bug fix

reset_adapter forked the dense RNG tracker, whose seed varies with the dense TP rank — expert-DP peers differ in exactly that rank whenever TP != ETP, so replicas of a reused slot diverged. It now forks the expert tracker.

Testing

8x H200: test_multi_lora_moe.py 17 passed; all multi-LoRA tests 70 passed (9 failures elsewhere in tests/unit_tests/peft/ are pre-existing single-LoRA default mismatches on bridge). Real-model EP=2 integration (Qwen3-MoE-shaped, real MoEAlltoAllTokenDispatcher): all expert adapter params get grads, and routing tokens to different slots gives different losses. Export through AutoBridge.export_adapter_weights is byte-identical (names and shapes) to single-adapter LoRA.

Merge order

First of three: this → radixark/miles#1794 → sgl-project/sglang#32376.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[yushengsu-thu]

Export contract: identical to the single-adapter path

Since the whole point of reusing GroupedExpertLinearAdapter per slot is that the existing
grouped-expert export keeps working, that was checked directly rather than argued: a wrapped MoE model
was exported through AutoBridge.export_adapter_weights (the call miles' weight sync uses), once with
MultiLoRA and once with single-adapter LoRA, on the same gpt-oss config with 4 experts.

Both produce byte-identical names and shapes:

model.layers.0.mlp.experts.gate_up_proj.lora_A.weight  (1, 4, 8, 2880)
model.layers.0.mlp.experts.gate_up_proj.lora_B.weight  (1, 4, 5760, 8)
model.layers.0.mlp.experts.down_proj.lora_A.weight     (1, 4, 8, 2880)
model.layers.0.mlp.experts.down_proj.lora_B.weight     (1, 4, 2880, 8)

So a multi-LoRA slot exported through expose_adapter_slot is indistinguishable, downstream, from a
single-adapter export — the payload shape the serving side already handles in CI.

Worth recording for anyone reading the layout: the per-expert index lives in the tensor's expert axis
here, not in the key, because gpt-oss's HF format stacks experts into one tensor per projection. The
naming is therefore architecture-dependent (a Qwen3-MoE-style HF layout with experts.{i}.gate_proj
modules would key per expert instead), which is a property of the architecture bridge, not of this PR.

[yushengsu-thu]

Pushed 592f5e2: a GLM-5.2_5layer multi-LoRA e2e smoke (the first MLA model through this path — TP=EP=4, ETP=1, SP, thd, dp-attention rollout) caught two defects in the dense MultiLoRALinear on replicated base linears (TELinear parallel_mode='duplicated', i.e. MLA q/kv down-projections):

1. Per-slot token spans covered the full micro-batch while a replicated base under sequence parallelism consumes the SP shard directly — the grouped GEMM read out of bounds (device-side assert). Spans are now intersected with the rank's contiguous token window, same invariant as the MoE slot routing's SP narrow. 4 new CPU tests cover the window math.
2. _gather_output keyed only on input_is_parallel, leaving the adapter output column-sharded while a replicated base outputs full width. Now mirrors ParallelLinearAdapter's lin_out_gather_output.

Qwen3-style models never reach either path (all their targets are TP-parallel), which is why the existing e2e missed them. After the fix the GLM-5.2 smoke completes 3 training steps end-to-end (per-step distributed upserts to both dp-attention engines, (tp, pp, ep) checkpoint shards at the realized coordinates), run under CUDA_LAUNCH_BLOCKING=1; the three multi-LoRA test files pass 74/74.