4/5 support kimi 2.5 full + lora: shared-outer grouped-expert LoRA - #1222
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for Kimi-K2.5 LoRA GRPO training and serving, including a new example script, DeepSeek/Kimi MLA target module mappings, and integration with SGLang's shared-outer grouped-expert LoRA. It also optimizes weight synchronization by skipping base weight sync during LoRA distributed updates and sending full adapters in a single call. The review feedback identifies three key issues: a potential IndexError when gathering serialized tensors from ranks without LoRA weights, a typo in the PYTHONBUFFERED environment variable in the example script, and a potential AttributeError when validating SGLang arguments if the shared-outer LoRA attribute is missing.
nanjiangwill
changed the title
There was a problem hiding this comment.
Discussed in this PR: #1057. So I approved to unblocked ci first.
yueming-yuan
left a comment
yueming-yuan
left a comment
There was a problem hiding this comment.
general file changes (update weight, sglang_engine) LGTM, @yushengsu-thu @jybsuper to check lora specific logics
Add --experts-shared-outer-loras: gate_up lora_A and down lora_B are shared across experts (expert_dim=1), matching SGLang PR #21466's serving contract. Maps DeepSeek/Kimi MLA HF names (q_a_proj, kv_a_proj_with_mqa, ...) to Megatron linear_* and drops q_b/kv_b from the SGLang rollout adapter config (unsupported by default get_hidden_dim). Switch the LoRA weight sync to per-TP-rank serialized_named_tensors and accumulate the adapter into a single load_lora_adapter_from_tensors call; skip base-weight sync under distributed update (base stays on the rollout GPU). Adds examples/lora/run-kimi-k25-megatron-lora.sh.
yueming-yuan
left a comment
yueming-yuan
left a comment
There was a problem hiding this comment.
approved after offline discussion
4f67038 to
4c033c6
Compare
3 participants
Part 4/5 of splitting #1057 (Kimi K2.5 full-param + LoRA support) into reviewable PRs, rebased on latest main.
Add
--experts-shared-outer-loras: gate_up lora_A and down lora_B are shared across experts (expert_dim=1), matching SGLang PR #21466's serving contract. Maps DeepSeek/Kimi MLA HF names (q_a_proj,kv_a_proj_with_mqa, ...) to Megatronlinear_*and dropsq_b/kv_bfrom the SGLang rollout adapter config (unsupported by the default get_hidden_dim).Switch LoRA weight sync to per-TP-rank
serialized_named_tensorsand accumulate the adapter into a singleload_lora_adapter_from_tensorscall; skip base-weight sync under distributed update (base stays on the rollout GPU). Addsexamples/lora/run-kimi-k25-megatron-lora.sh.Builds on #1221 (Kimi full-param) at runtime; file set is disjoint from #1219-#1221 so it merges independently. #1223 stacks on this one.