Skip to content

3/5 support kimi 2.5 full + lora: Kimi K2.5 VL full-parameter support - #1221

Merged
yueming-yuan merged 5 commits into
radixark:mainfrom
nanjiangwill:kimi25-pr4-fullparam
May 28, 2026
Merged

3/5 support kimi 2.5 full + lora: Kimi K2.5 VL full-parameter support#1221
yueming-yuan merged 5 commits into
radixark:mainfrom
nanjiangwill:kimi25-pr4-fullparam

Conversation

@nanjiangwill

@nanjiangwill nanjiangwill commented May 26, 2026

Copy link
Copy Markdown
Contributor

Part 3/5 of splitting #1057 (Kimi K2.5 full-param + LoRA support) into reviewable PRs, rebased on latest main.

🙏 The original Kimi K2.5 full-parameter implementation is the work of @GeLee-Q — thanks for making it work! This PR is the reviewable split of that work (credited via Co-authored-by on the commit).

Megatron->HF weight conversion for Kimi K2.5 / VL (megatron_to_hf/kimi_vl.py + dispatch in megatron_to_hf/__init__.py) and language_model.-prefixed decoder layer matching in the weight iterator (update_weight/common.py). Media placeholder tokens are expanded to grid-derived counts at train time (training_utils/data.py). General colocate weight-sync robustness used by this path: IPC collect after MultiprocessingSerializer sync (memory_utils + actor) and a reloadable broadcast_object_list. Adds scripts/run-kimi-k25.sh.

Uses the processor unification from #1219 at runtime; file sets are disjoint so it merges independently.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces support for Kimi VL and Kimi-K2.5 models, including weight conversion utilities, multimodal rollout data expansion, and training scripts. The review feedback highlights a critical runtime error where reversed() is called directly on a PyTorch tensor, which must be converted to a list first. Additionally, the environment variable PYTHONBUFFERED in the run script should be corrected to PYTHONUNBUFFERED. Other recommendations include using getattr for safer configuration retrieval, enforcing strict length checks in zip operations, and wrapping process group destruction in a try...finally block to guarantee cleanup.

Comment thread miles/backends/training_utils/data.py Outdated
Comment thread scripts/run-kimi-k25.sh
Comment thread miles/backends/megatron_utils/megatron_to_hf/kimi_vl.py
Comment thread miles/backends/megatron_utils/megatron_to_hf/kimi_vl.py
Comment thread miles/backends/training_utils/data.py Outdated
Comment thread miles/backends/training_utils/data.py Outdated
Comment thread miles/backends/megatron_utils/actor.py Outdated
@nanjiangwill nanjiangwill changed the title 4/6 support kimi 2.5 full + lora: Kimi K2.5 VL full-parameter support 3/5 support kimi 2.5 full + lora: Kimi K2.5 VL full-parameter support May 26, 2026
Add Megatron->HF weight conversion for Kimi K2.5 / VL (kimi_vl.py +
dispatch in megatron_to_hf/__init__) and the language_model.-prefixed
decoder layer matching in the weight iterator. Expand media placeholder
tokens to grid-derived counts at train time (training_utils/data.py).
General colocate weight-sync robustness used by the kimi path: IPC
collect after MultiprocessingSerializer sync (memory_utils + actor) and
a reloadable broadcast_object_list. Adds scripts/run-kimi-k25.sh.

The original Kimi K2.5 full-parameter implementation is the work of
JiLi (@GeLee-Q); this commit is the reviewable split of that work.

Co-authored-by: JiLi <22428217+GeLee-Q@users.noreply.github.com>
@nanjiangwill
nanjiangwill force-pushed the kimi25-pr4-fullparam branch from aa6712c to dca1bd9 Compare May 26, 2026 23:53
Comment thread miles/backends/megatron_utils/actor.py Outdated
Comment thread miles/backends/training_utils/data.py Outdated

@yueming-yuan yueming-yuan left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

overall LGTM with nits comments

Zhichenzzz and others added 2 commits May 28, 2026 04:12
data.py was growing with VL-only logic; move the 5 helpers and
expand_multimodal_rollout_data_in_place into mm_data.py so generic
batching stays out of multimodal preprocessing's way and future VL
models can extend mm_data.py without inflating data.py.

Pure mechanical refactor: function bodies are unchanged. Hardcoded
media token id 163605 and 2x2 patch merger dims are promoted to
KIMI_VL_MEDIA_TOKEN_ID / _KIMI_VL_MERGE_H / _KIMI_VL_MERGE_W constants.
data.py keeps a re-export of expand_multimodal_rollout_data_in_place
via 'from .mm_data import ...', so the two call sites in get_batch /
get_data_iterator are unchanged.
@nanjiangwill
nanjiangwill force-pushed the kimi25-pr4-fullparam branch from 8f03759 to c1ae7d3 Compare May 28, 2026 04:27

@Zhichenzzz Zhichenzzz left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

lgtm! Thanks @nanjiangwill

@yueming-yuan
yueming-yuan merged commit 408721a into radixark:main May 28, 2026
49 of 52 checks passed
@nanjiangwill
nanjiangwill deleted the kimi25-pr4-fullparam branch May 28, 2026 23:44
Frankstein73 pushed a commit to sii-avalanche/miles that referenced this pull request Jun 1, 2026
…radixark#1221)

Co-authored-by: JiLi <22428217+GeLee-Q@users.noreply.github.com>
Co-authored-by: Zhichenzzz <zczeng@uw.edu>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants