5/5 support kimi 2.5 full + lora: LoRA base-weight CPU backup - #1223
Merged
yueming-yuan merged 1 commit intoMay 29, 2026
Merged
Conversation
nanjiangwill
requested review from
Zhichenzzz,
fzyzcjy,
guapisolo,
jybsuper,
maocheng23,
yueming-yuan and
yushengsu-thu
as code owners
May 26, 2026 23:42
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces support for Kimi-K2.5 LoRA GRPO training, adding arguments for CPU backup of base weights and shared-outer grouped-expert LoRA. It optimizes weight syncing by skipping base weight transfers when CPU backup is enabled and updates target module mappings between Megatron and Hugging Face formats. The review feedback highlights three key issues: a potential IndexError when a pipeline parallel rank has no LoRA weights, a restrictive dotted-path check during module name conversion, and a potential AttributeError when accessing SGLang-specific arguments directly.
nanjiangwill
changed the title
yushengsu-thu
approved these changes
May 27, 2026
Collaborator
There was a problem hiding this comment.
Discussed in this PR: #1057. So I approved to unblocked ci first.
nanjiangwill
force-pushed
the
kimi25-pr6-lora-cpu-backup
branch
from
May 28, 2026 23:47
df9ed15 to
39dd23a
Compare
Frankstein73
pushed a commit
to sii-avalanche/miles
that referenced
this pull request
Jun 1, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Part 5/5 of splitting #1057 (Kimi K2.5 full-param + LoRA support) into reviewable PRs, rebased on latest main.
Add
--lora-base-cpu-backup: with--colocate, keep a host-RAM mirror of the base weights in SGLang (enable_weights_cpu_backup) so they survivetorch_memory_saver.pause()across rollout/training swaps, and have the trainer skip the per-step base weight sync (and the surrounding restore / re-quantize calls) entirely. Trades host RAM for faster onload/offload; ignored unless--colocateand LoRA are both on.Stacks on #1222 (shared files: arguments.py, lora_utils.py, sglang_engine.py, update_weight_from_tensor.py). Until #1222 merges, the diff here also shows that PR's changes; review the last commit (
perf: LoRA base-weight CPU backup) in isolation, or after #1222 lands.