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

[yushengsu-thu]

What does this PR do ?

Extends multi-LoRA (#4218) to grouped MoE expert linears, bringing the multi-adapter path to parity with single-adapter LoRA for grouped experts.

Stacked on #4218 — this branch includes #4218's commits as its base, so the Files tab shows both until #4218 merges. Only the last commit is new here. Ported from radixark#23, where it runs in production on the bridge branch.

Why

MultiLoRA skips every MoE expert linear, so on MoE models the experts — most of the trainable capacity — silently get no adapter.

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

On 8x H200 (radixark fork, same file contents as this PR): test_multi_lora_moe.py 17 passed; all multi-LoRA tests 70 passed. 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.

Changelog

• Add MultiLoRAGroupedExpertLinear and MoE slot routing (install_moe_slot_routing) to megatron.bridge.peft.multi_lora_layers
• MultiLoRA now wraps grouped expert linears; sequential (SequentialMLP) experts are still skipped with a clearer warning
• MultiLoRAGroupedExpertLinear.reset_adapter forks the expert RNG tracker (EDP-consistent slot reuse)
• New unit tests: tests/unit_tests/peft/test_multi_lora_moe.py (CPU routing math + single-GPU integration)

GitHub Actions CI

See the CI section in the Contributing doc. A Nvidia developer will need to approve and trigger the CI for external contributors.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?
• Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  • Reviewer: Does the PR have correct import guards for all optional libraries?

Additional Information

• Depends on [peft] Support multi-lora for Megatron models #4218

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[Copilot]

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

[yaoyu-33]

Hi @yushengsu-thu — this needs a rebase before it can be reviewed/merged.

Why it conflicts: #4218 was squash-merged into main as 1f12931e2. This branch still carries #4218's six original (unsquashed) commits as its base, so git sees them as conflicting with the squashed version plus the later lora_layers.py / utils.py refactors on main.

Fix: drop the #4218 base and keep only your one new commit on top of current main. I verified locally that it cherry-picks cleanly — no manual conflict resolution needed:

git fetch upstream main
git checkout yusheng/multilora-moe-experts
git reset --hard upstream/main
git cherry-pick 5fde419eae1cd7f67136f283ac6965dbdf6a69dc
git push --force-with-lease

I also confirmed on the rebased result that every symbol this PR imports still exists after main's peft refactor (GroupedExpertLinearAdapter, all2all_hp2sp, is_grouped_expert_linear, init_method_normal, sort_chunks_by_idxs), and that ruff check + ruff format --check pass.

I would have pushed the rebase for you, but the fork is owned by the radixark org and GitHub's "allow edits by maintainers" does not apply to org-owned forks, so my push is rejected.

Note #5156 is stacked on this branch and will need rebasing after this one lands — see my comment there.