Skip to content

feat(kimi_k25): implement multimodal get_inputs_embeds for Megatron training #106

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
zhihanliu-collab wants to merge 2 commits into
base: main
Choose a base branch
from
Draft

feat(kimi_k25): implement multimodal get_inputs_embeds for Megatron training #106

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
d03aefe
feat(kimi_k25): implement multimodal get_inputs_embeds for Megatron t…
zhihanliu-collab Jun 3, 2026
2772934
fix: guard plain-text dummy vision pass against NaN propagation
zhihanliu-collab Jun 3, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 35 additions & 4 deletions src/mcore_bridge/model/mm_gpts/kimi_vl.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -72,7 +72,6 @@ class KimiK25Vit(HuggingFaceVit):
module_mapping = {'vision_tower': 'vision_tower', 'mm_projector': 'mm_projector'}
_vision_tower = ['vision_tower']
_aligner = ['mm_projector']
test_mm_type = 'text'

def prepare_model(self, hf_config: PretrainedConfig):
output = []
Expand All @@ -85,10 +84,42 @@ def prepare_model(self, hf_config: PretrainedConfig):
self.vision_tower = MoonViT3dPretrainedModel._from_config(vit_config)
self.mm_projector = PatchMergerMLP(proj_config).to(self.vision_tower.dtype)

def _encode_images(self, pixel_values, grid_thws):
# vision_tower returns a list of un-projected feature tensors; mm_projector
# (PatchMergerMLP) maps them to the language hidden size. Mirrors
# KimiK25ForConditionalGeneration.forward.
image_features = self.vision_tower(pixel_values, grid_thws)
image_features = self.mm_projector(image_features)
return torch.cat(image_features, dim=0)

def get_inputs_embeds(self, inputs_embeds, **kwargs):
pixel_values = kwargs.pop('pixel_values', None)
if pixel_values is not None:
raise NotImplementedError('Kimi-K25 currently only supports plain text training.')
input_ids = kwargs['input_ids']
pixel_values = kwargs.get('pixel_values')
grid_thws = kwargs.get('grid_thws')
dtype = next(self.vision_tower.parameters()).dtype
if pixel_values is not None and pixel_values.size(0) > 0:
if grid_thws is None:
raise KeyError('pixel_values present in inputs but grid_thws is missing')
pixel_values = pixel_values.to(device=inputs_embeds.device, dtype=dtype)
grid_thws = grid_thws.to(inputs_embeds.device)
image_features = self._encode_images(pixel_values, grid_thws)
image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype)
media_token_id = self.hf_config.media_placeholder_token_id
image_mask = (input_ids == media_token_id).unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_features)
else:
# plain-text batch: still run the vision graph on a dummy image so that
# DP ranks stay in sync during gradient all-reduce.
vision_config = self.hf_config.vision_config
patch_size = vision_config.patch_size
merge_kernel_size = vision_config.merge_kernel_size
kernel = merge_kernel_size[0] if isinstance(merge_kernel_size, (list, tuple)) else merge_kernel_size
h = w = kernel * 2
dummy_pixels = torch.zeros((h * w, 3, patch_size, patch_size), dtype=dtype, device=inputs_embeds.device)
dummy_grid = input_ids.new_tensor([[1, h, w]])
image_features = self._encode_images(dummy_pixels, dummy_grid)
inputs_embeds = inputs_embeds + image_features.mean().to(
device=inputs_embeds.device, dtype=inputs_embeds.dtype) * 0.

@gemini-code-assist gemini-code-assist Bot Jun 3, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

high

Numerical Instability and Robustness Safeguards

  1. NaN Propagation Prevention (High Severity): Running a dummy forward pass on an all-zero tensor (torch.zeros) can sometimes lead to numerical instabilities (e.g., division by zero or zero-variance in normalization layers) in the vision tower, resulting in NaN values in image_features. If image_features contains NaN, image_features.mean() * 0. will evaluate to NaN (since NaN * 0 is NaN in IEEE 754). This will propagate NaN to inputs_embeds, corrupting the entire plain-text batch and causing training divergence. Using torch.nan_to_num on the zero-multiplier term ensures that any numerical instability in the dummy pass is safely zeroed out and does not affect text-only training.

  2. Defensive Attribute Access (Medium Severity): Accessing vision_config.patch_size and vision_config.merge_kernel_size directly can raise AttributeError or TypeError if they are missing or None. Using getattr with safe fallbacks is more robust. Additionally, checking len(merge_kernel_size) > 0 prevents IndexError if merge_kernel_size is an empty list or tuple.

Suggested change
vision_config = self.hf_config.vision_config
patch_size = vision_config.patch_size
merge_kernel_size = vision_config.merge_kernel_size
kernel = merge_kernel_size[0] if isinstance(merge_kernel_size, (list, tuple)) else merge_kernel_size
h = w = kernel * 2
dummy_pixels = torch.zeros((h * w, 3, patch_size, patch_size), dtype=dtype, device=inputs_embeds.device)
dummy_grid = input_ids.new_tensor([[1, h, w]])
image_features = self._encode_images(dummy_pixels, dummy_grid)
inputs_embeds = inputs_embeds + image_features.mean().to(
device=inputs_embeds.device, dtype=inputs_embeds.dtype) * 0.
vision_config = self.hf_config.vision_config
patch_size = getattr(vision_config, "patch_size", 14)
merge_kernel_size = getattr(vision_config, "merge_kernel_size", 2) or 2
kernel = merge_kernel_size[0] if isinstance(merge_kernel_size, (list, tuple)) and len(merge_kernel_size) > 0 else merge_kernel_size
h = w = kernel * 2
dummy_pixels = torch.zeros((h * w, 3, patch_size, patch_size), dtype=dtype, device=inputs_embeds.device)
dummy_grid = input_ids.new_tensor([[1, h, w]])
image_features = self._encode_images(dummy_pixels, dummy_grid)
zero_loss_term = torch.nan_to_num(image_features.mean() * 0.)
inputs_embeds = inputs_embeds + zero_loss_term.to(
device=inputs_embeds.device, dtype=inputs_embeds.dtype)

return inputs_embeds


Expand Down
5 changes: 5 additions & 0 deletions tests/test_mllm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -87,6 +87,10 @@ def test_kimi_vl():
_test_model('moonshotai/Kimi-VL-A3B-Thinking-2506')


def test_kimi_k25():
_test_model('moonshotai/Kimi-K2.6')


def test_qwen3_vl():
_test_model('Qwen/Qwen3-VL-4B-Instruct')

Expand Down Expand Up @@ -134,6 +138,7 @@ def test_gemma4():
# test_glm4_6v_flash()
# test_ovis2_5()
# test_kimi_vl()
# test_kimi_k25()
# test_qwen3_vl()
# test_qwen3_vl_moe()
# test_qwen3_omni()
Expand Down