feat(multi-lora): enable and validate MoE expert adapters - #1794
feat(multi-lora): enable and validate MoE expert adapters#1794yushengsu-thu wants to merge 9 commits into
Conversation
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
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.
a98e4c3 to
eecb246
Compare
|
Description updated with two follow-up PRs found by a GLM-5.2_5layer multi-LoRA e2e (the first MLA + dp-attention model through this path; TP=EP=4, ETP=1, thd, disaggregated 4 train + 4 rollout):
With both applied on the current |
yushengsu-thu
left a comment
yushengsu-thu
left a comment
There was a problem hiding this comment.
remove the comments or reduce to 1 or 2 lines
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>
yushengsu-thu
left a comment
yushengsu-thu
left a comment
There was a problem hiding this comment.
reduce to 1 sentence
…me, leaf-name mapping, squashed comments
|
Need this: #1803 
|
CI triage: both failures are environment issues, not this PR
Pre-existing image breakage; it reproduces on plain Not rerunnable: the repo classifies Not reachable from this diff either — the failing run logs
HuggingFace rate limit, not connectivity: |
miles modification - support moe lora layer
megatron-bridge:
enable and validate MoE expert adapters: radixark/Megatron-Bridge#23
fix(peft): correct dense MultiLoRALinear on replicated base linears: radixark/Megatron-Bridge#25 — required for MLA models (GLM-5 / DeepSeek): their q/kv down-projections are replicated linears, whose multi-LoRA token spans and output gather were wrong under sequence parallelism
sglang:
Warn instead of silently dropping one-sided MoE expert LoRA targets: sgl-project/sglang#32376
Allow load_lora_adapter_from_distributed under dp attention: sgl-project/sglang#32421 — required for dp-attention serving recipes (GLM-5.2's standard config): the distributed-load route was the only LoRA endpoint still asserting dp_size == 1
Both follow-up fixes were found by a GLM-5.2_5layer multi-LoRA e2e (TP=EP=4, ETP=1, thd, dp-attention rollout); with them applied on the current sglang-miles / bridge tips the run completes 3 training steps end-to-end (per-step distributed upserts,
(tp, pp, ep)checkpoint shards, clean adapter retire).What
moe_permute_fusionoff when multi-LoRA targets expert leaves. Most bridge MoE providers default it on (Qwen3-MoE included), and the fused permute records TE'srow_id_map, which expert adapters cannot replay. Costs only the fused permute kernel.slice_lora_to_rank: address the rank axis from the end. A packed grouped-expert export is 4-D (lora_A (1, E, rank, in)/lora_B (1, E, out, rank)); front-addressing sliced the singleton/expert axis instead —lora_A's rank was never trimmed andlora_Bsilently dropped experts. Reachable whenever an adapter's rank is below--lora-rank.(tp, pp, ep). EP ranks sharing a(tp, pp)coordinate hold different local experts, so shards collided and a resume loaded one EP rank's experts onto every rank. Writer election and the completeness check both come from one cached gloo all-gather of the realized coordinates — neither is derivable locally, and the realized set is not thetp × pp × epcross product once ETP < TP (which expert multi-LoRA requires). Suffix omitted atep_size == 1with a legacy fallback, so existing checkpoints stay loadable.finalize()(_validate_multi_lora_moe_support): ETP = 1, grouped experts, no fp8/fp4, no capacity padding, and both expert projections targeted — sglang applies the expert delta after gate_up and after down together, so a one-sided target would be silently dropped at rollout. Launch-time: PP = 1 and--qkv-format thd(per-slot token spans need contiguous sample packing;bshdviolates it). Expert TP is deliberately not checked at the CLI: itsNonedefault is only resolved by Megatron aftermiles_validate_args.exclude_modulesis no longer forwarded toMultiLoRA(already subtracted;ModuleMatcherasserts it empty).Testing
8x H200 with radixark/Megatron-Bridge#23:
tests/fast -k "lora or adapter"→ 226 passed. New tests cover expert-leaf gating, shard naming and completeness with unrealized coordinates, packed-expert slicing, and the validation cases (incl. the unresolved expert-TP default). Writer election exercised undertorchrun --nproc_per_node=4(TP=2, EP=2, ETP=1): exactly one writer per realized coordinate, 2 shards where the cross product would demand 4. Dense DP=2 e2e results in a comment below.Merge order
radixark/Megatron-Bridge#23 (merged) → this → sgl-project/sglang#32376 (merged).
For GLM-5.2-class models (MLA + dp-attention serving), additionally: radixark/Megatron-Bridge#25 and sgl-project/sglang#32421 — independent of this PR and of each other, mergeable in any order.