From f3b5c5b423192b2cc17ae0d915f4dd7e698591ed Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Thu, 30 Jul 2026 15:45:32 +0800 Subject: [PATCH 01/67] Add Kimi K3 reference model --- .../numerical_tests_kimi_k3.py | 704 ++++++++++++++++ .../numerical_tests_kimi_k3_device.py | 435 ++++++++++ tests/unit_tests/test_kimi_k3.py | 289 +++++++ torchtitan/models/__init__.py | 1 + torchtitan/models/kimi_k3/README.md | 132 +++ torchtitan/models/kimi_k3/__init__.py | 469 +++++++++++ torchtitan/models/kimi_k3/config_registry.py | 77 ++ torchtitan/models/kimi_k3/model.py | 797 ++++++++++++++++++ torchtitan/models/kimi_k3/parallelize.py | 57 ++ .../models/kimi_k3/state_dict_adapter.py | 337 ++++++++ torchtitan/models/kimi_k3/vision_encoder.py | 488 +++++++++++ 11 files changed, 3786 insertions(+) create mode 100644 scripts/checkpoint_conversion/numerical_tests_kimi_k3.py create mode 100644 scripts/checkpoint_conversion/numerical_tests_kimi_k3_device.py create mode 100644 tests/unit_tests/test_kimi_k3.py create mode 100644 torchtitan/models/kimi_k3/README.md create mode 100644 torchtitan/models/kimi_k3/__init__.py create mode 100644 torchtitan/models/kimi_k3/config_registry.py create mode 100644 torchtitan/models/kimi_k3/model.py create mode 100644 torchtitan/models/kimi_k3/parallelize.py create mode 100644 torchtitan/models/kimi_k3/state_dict_adapter.py create mode 100644 torchtitan/models/kimi_k3/vision_encoder.py diff --git a/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py b/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py new file mode 100644 index 0000000000..dd0cb585ff --- /dev/null +++ b/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py @@ -0,0 +1,704 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Numerical parity for the reduced TorchTitan and released HF Kimi K3 models. + +The released checkpoint contains MXFP4 routed-expert weights, so this test +constructs the same reduced, unquantized topology on both sides. TorchTitan +initializes the weights once, ``KimiK3StateDictAdapter`` converts them to the +released HuggingFace schema, and the HuggingFace model loads them strictly. + +The comparison covers: + +- text-only decoder layers, KDA/MLA outputs, router expert IDs, and logits; +- vision transformer blocks and projected vision features; +- end-to-end image+text decoder layers and logits. + +The released HuggingFace implementation imports FLA for KDA. By default, run +this script in an environment that satisfies the released model requirements +and has a CUDA GPU. The script requests the released eager MLA and +vision-attention paths after construction so the comparison isolates model math +from FlashAttention kernels. + +For a CPU-only comparison, ``--hf_kda_backend reference`` installs the minimum +released FLA API with pure PyTorch operators. + +Example: + + CUDA_VISIBLE_DEVICES=0 python -m \ + scripts.checkpoint_conversion.numerical_tests_kimi_k3 \ + --hf_repo_path ~/hf_assets/moonshotai/Kimi-K3 + + python -m scripts.checkpoint_conversion.numerical_tests_kimi_k3 \ + --hf_repo_path ~/hf_assets/moonshotai/Kimi-K3 \ + --device cpu \ + --hf_kda_backend reference +""" + +import argparse +import importlib +import sys +import types +from collections.abc import Callable +from typing import cast + +import torch +import torch.nn.functional as F +from torch import nn + +from torchtitan.models.kimi_k3 import model_registry +from torchtitan.models.kimi_k3.model import KimiK3Model +from torchtitan.models.kimi_k3.state_dict_adapter import KimiK3StateDictAdapter +from transformers import AutoConfig, AutoModelForCausalLM + + +_IMAGE_TOKEN_ID = 7 + + +class _ReferenceShortConvolution(nn.Module): + """Pure PyTorch compatibility layer for the released HF KDA module.""" + + def __init__( + self, + hidden_size: int, + kernel_size: int, + bias: bool = False, + activation: str = "silu", + **kwargs, + ): + super().__init__() + del kwargs + self.weight = nn.Parameter(torch.empty(hidden_size, 1, kernel_size)) + self.bias = nn.Parameter(torch.empty(hidden_size)) if bias else None + self.kernel_size = kernel_size + self.activation = activation + + def forward( + self, + x: torch.Tensor, + cache: torch.Tensor | None = None, + output_final_state: bool = False, + cu_seqlens: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + del cache, cu_seqlens + output = F.conv1d( + F.pad(x.transpose(1, 2), (self.kernel_size - 1, 0)), + self.weight, + self.bias, + groups=self.weight.shape[0], + ).transpose(1, 2) + if self.activation == "silu": + output = F.silu(output) + final_state = x[:, -self.kernel_size + 1 :] if output_final_state else None + return output, final_state + + +class _ReferenceRMSNormGated(nn.Module): + """Pure PyTorch equivalent of FLA's fused gated RMSNorm.""" + + def __init__( + self, + hidden_size: int, + eps: float = 1e-5, + activation: str = "sigmoid", + ): + super().__init__() + if activation != "sigmoid": + raise ValueError( + "The Kimi K3 reference RMSNorm only supports a sigmoid gate." + ) + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.eps = eps + + def forward(self, x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: + input_dtype = x.dtype + x_float = x.float() + x_float = x_float * torch.rsqrt( + x_float.square().mean(dim=-1, keepdim=True) + self.eps + ) + return (x_float * self.weight.float() * torch.sigmoid(gate.float())).to( + input_dtype + ) + + +def _reference_kda( + *, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + initial_state: torch.Tensor | None, + output_final_state: bool, + lower_bound: float, + **kwargs, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """Pure PyTorch equivalent of the FLA KDA inference API.""" + del kwargs + input_dtype = v.dtype + q = q.float() + k = k.float() + v = v.float() + q = q * torch.rsqrt(q.square().sum(dim=-1, keepdim=True) + 1e-6) + k = k * torch.rsqrt(k.square().sum(dim=-1, keepdim=True) + 1e-6) + log_decay = lower_bound * torch.sigmoid( + torch.exp(A_log.float()).view(1, 1, -1, 1) + * (g.float() + dt_bias.float().view(1, 1, *g.shape[-2:])) + ) + beta = torch.sigmoid(beta.float()) + + batch_size, seq_len, num_heads, head_dim = q.shape + value_dim = v.shape[-1] + state = ( + torch.zeros( + batch_size, + num_heads, + head_dim, + value_dim, + dtype=torch.float32, + device=q.device, + ) + if initial_state is None + else initial_state.float() + ) + outputs = [] + for token_idx in range(seq_len): + state = state * torch.exp(log_decay[:, token_idx]).unsqueeze(-1) + old_value = torch.matmul( + k[:, token_idx].unsqueeze(-2), + state, + ).squeeze(-2) + delta = (v[:, token_idx] - old_value) * beta[:, token_idx].unsqueeze(-1) + state = state + k[:, token_idx].unsqueeze(-1) * delta.unsqueeze(-2) + outputs.append( + torch.matmul(q[:, token_idx].unsqueeze(-2), state).squeeze(-2) + * (head_dim**-0.5) + ) + output = torch.stack(outputs, dim=1).to(input_dtype) + return output, state if output_final_state else None + + +def _install_reference_fla() -> None: + """Expose the minimum FLA API used by the released HF model.""" + + # The released vision reference decorates an interpolation helper with + # torch.compile. Keep the compatibility backend fully eager so it remains + # usable for the CPU-only comparison. + def eager_compile(model=None, *args, **kwargs): + del args, kwargs + return model if model is not None else lambda wrapped: wrapped + + torch.__dict__["compile"] = eager_compile + + fla = types.ModuleType("fla") + fla_modules = types.ModuleType("fla.modules") + fla_modules.__dict__.update( + { + "ShortConvolution": _ReferenceShortConvolution, + "FusedRMSNormGated": _ReferenceRMSNormGated, + } + ) + fla_ops = types.ModuleType("fla.ops") + fla_kda = types.ModuleType("fla.ops.kda") + fla_kda.__dict__.update( + { + "chunk_kda": _reference_kda, + "fused_recurrent_kda": _reference_kda, + } + ) + fla_ops_utils = types.ModuleType("fla.ops.utils") + fla_ops_index = types.ModuleType("fla.ops.utils.index") + fla_ops_index.__dict__.update( + { + "prepare_cu_seqlens_from_mask": lambda _mask: None, + "prepare_lens_from_mask": lambda _mask: None, + } + ) + fla_utils = types.ModuleType("fla.utils") + fla_utils.__dict__["tensor_cache"] = lambda fn: fn + + for module in ( + fla, + fla_modules, + fla_ops, + fla_kda, + fla_ops_utils, + fla_ops_index, + fla_utils, + ): + sys.modules[module.__name__] = module + + +def _build_reduced_hf_config(hf_repo_path: str, tt_config): + """Build the released HF config with TorchTitan's reduced dimensions.""" + released_config = AutoConfig.from_pretrained( + hf_repo_path, + trust_remote_code=True, + ) + text_config_cls = type(released_config.text_config) + vision_config_cls = type(released_config.vision_config) + config_cls = type(released_config) + + mla_config = next( + layer.attention for layer in tt_config.layers if layer.attention is not None + ) + kda_config = next( + layer.delta_attention + for layer in tt_config.layers + if layer.delta_attention is not None + ) + dense_config = next( + layer.feed_forward + for layer in tt_config.layers + if layer.feed_forward is not None + ) + moe_config = next(layer.moe for layer in tt_config.layers if layer.moe is not None) + vision_config = tt_config.vision_encoder + if vision_config is None: + raise ValueError("The Kimi K3 debug config must include a vision encoder.") + + full_attention_layers = [ + layer_idx + 1 + for layer_idx, layer in enumerate(tt_config.layers) + if layer.attention is not None + ] + kda_layers = [ + layer_idx + 1 + for layer_idx, layer in enumerate(tt_config.layers) + if layer.delta_attention is not None + ] + first_moe_layer = next( + layer_idx + for layer_idx, layer in enumerate(tt_config.layers) + if layer.moe is not None + ) + + text_config = text_config_cls( + vocab_size=tt_config.vocab_size, + hidden_size=tt_config.dim, + intermediate_size=dense_config.w1.out_features, + num_hidden_layers=len(tt_config.layers), + num_attention_heads=mla_config.num_heads, + num_key_value_heads=mla_config.num_heads, + hidden_act="situ", + rms_norm_eps=tt_config.norm.eps, + use_cache=False, + moe_intermediate_size=moe_config.routed_experts[0].w1.out_features, + num_experts=moe_config.num_experts, + num_experts_per_token=moe_config.router.top_k, + num_shared_experts=( + moe_config.shared_experts.w1.out_features + // moe_config.routed_experts[0].w1.out_features + ), + first_k_dense_replace=first_moe_layer, + moe_layer_freq=1, + moe_renormalize=moe_config.router.route_norm, + routed_scaling_factor=moe_config.router.route_scale, + num_expert_group=1, + topk_group=1, + q_lora_rank=mla_config.q_lora_rank, + kv_lora_rank=mla_config.kv_lora_rank, + qk_nope_head_dim=mla_config.qk_nope_head_dim, + qk_rope_head_dim=mla_config.qk_rope_head_dim, + v_head_dim=mla_config.v_head_dim, + mla_use_nope=True, + mla_use_output_gate=True, + linear_attn_config={ + "kda_layers": kda_layers, + "full_attn_layers": full_attention_layers, + "short_conv_kernel_size": kda_config.conv_kernel_size, + "head_dim": kda_config.head_dim, + "num_heads": kda_config.num_heads, + "use_full_rank_gate": True, + "gate_lower_bound": kda_config.kernel.lower_bound, + }, + attn_res_block_size=tt_config.layers[0].attn_res_block_size, + latent_moe_use_norm=True, + activation_situ_beta=dense_config.activation.beta, + activation_situ_linear_beta=dense_config.activation.linear_beta, + routed_expert_hidden_size=moe_config.routed_down.out_features, + ) + text_config._attn_implementation = "eager" + + vision_attention = vision_config.block.attn + vision_mlp = vision_config.block.mlp + vision_config_hf = vision_config_cls( + patch_size=vision_config.patch_size, + init_pos_emb_height=vision_config.init_pos_emb_height, + init_pos_emb_width=vision_config.init_pos_emb_width, + init_pos_emb_time=vision_config.max_num_frames, + vt_num_attention_heads=vision_attention.num_heads, + vt_num_hidden_layers=vision_config.num_layers, + vt_hidden_size=vision_config.dim, + vt_intermediate_size=vision_mlp.fc1.out_features, + merge_kernel_size=tuple(vision_config.merge_kernel_size), + mm_projector_type="patchmergerv2", + qkv_hidden_size=vision_attention.qkv_dim, + text_hidden_size=tt_config.dim, + norm_type="rmsnorm", + attn_bias=False, + patch_embed_proj_bias=False, + linear_bias=False, + activation_func="gelu_pytorch_tanh", + pos_emb_interpolation_mode=vision_config.interpolation_mode, + ) + vision_config_hf._attn_implementation = "eager" + + config = config_cls( + text_config=text_config, + vision_config=vision_config_hf, + media_placeholder_token_id=_IMAGE_TOKEN_ID, + pad_token_id=0, + auto_map=released_config.auto_map, + ) + config._name_or_path = hf_repo_path + return config + + +def _first_tensor(output) -> torch.Tensor: + if isinstance(output, torch.Tensor): + return output + if isinstance(output, (list, tuple)): + for value in output: + if isinstance(value, torch.Tensor): + return value + raise TypeError(f"Expected a tensor output, got {type(output).__name__}.") + + +def _capture_tensor( + destination: dict[str, torch.Tensor], + name: str, + *, + flatten: bool = False, +) -> Callable: + def hook(_module, _inputs, output) -> None: + value = _first_tensor(output).detach().float().cpu() + if flatten: + value = value.reshape(-1, value.shape[-1]) + destination[name] = value + + return hook + + +def _capture_router_ids( + destination: dict[str, torch.Tensor], + name: str, +) -> Callable: + def hook(_module, _inputs, output) -> None: + expert_ids = output[0] + destination[name] = expert_ids.detach().reshape(-1, expert_ids.shape[-1]).cpu() + + return hook + + +def _register_text_hooks( + tt_model, + hf_model, + tt_outputs: dict[str, torch.Tensor], + hf_outputs: dict[str, torch.Tensor], +) -> None: + for layer_idx, tt_layer in enumerate(tt_model.layers.values()): + hf_layer = hf_model.language_model.model.layers[layer_idx] + layer_name = f"decoder.layer.{layer_idx}" + tt_layer.register_forward_hook(_capture_tensor(tt_outputs, layer_name)) + hf_layer.register_forward_hook(_capture_tensor(hf_outputs, layer_name)) + + attention_name = f"decoder.attention.{layer_idx}" + tt_attention = ( + tt_layer.attention + if tt_layer.attention is not None + else tt_layer.delta_attention + ) + tt_attention.register_forward_hook(_capture_tensor(tt_outputs, attention_name)) + hf_layer.self_attn.register_forward_hook( + _capture_tensor(hf_outputs, attention_name) + ) + + if tt_layer.moe is not None: + router_name = f"decoder.router_ids.{layer_idx}" + tt_layer.moe.router.register_forward_hook( + _capture_router_ids(tt_outputs, router_name) + ) + hf_layer.block_sparse_moe.gate.register_forward_hook( + _capture_router_ids(hf_outputs, router_name) + ) + + +def _register_vision_hooks( + tt_model, + hf_model, + tt_outputs: dict[str, torch.Tensor], + hf_outputs: dict[str, torch.Tensor], +) -> None: + for layer_idx, tt_layer in enumerate(tt_model.vision_encoder.layers.values()): + layer_name = f"vision.layer.{layer_idx}" + tt_layer.register_forward_hook( + _capture_tensor(tt_outputs, layer_name, flatten=True) + ) + hf_model.vision_tower.encoder.blocks[layer_idx].register_forward_hook( + _capture_tensor(hf_outputs, layer_name, flatten=True) + ) + + +def _compare_tensor( + name: str, + tt_value: torch.Tensor, + hf_value: torch.Tensor, + *, + atol: float, + rtol: float, +) -> None: + if tt_value.shape != hf_value.shape: + raise AssertionError( + f"{name}: TorchTitan shape {tuple(tt_value.shape)} does not match " + f"HF shape {tuple(hf_value.shape)}." + ) + difference = (tt_value - hf_value).abs() + cosine = F.cosine_similarity( + tt_value.reshape(-1), + hf_value.reshape(-1), + dim=0, + ).item() + print( + f"{name:32s} max={difference.max().item():.4e} " + f"mean={difference.mean().item():.4e} cos={cosine:.8f}" + ) + torch.testing.assert_close( + tt_value, + hf_value, + atol=atol, + rtol=rtol, + msg=lambda message: f"{name} failed numerical parity:\n{message}", + ) + + +def _compare_captured( + tt_outputs: dict[str, torch.Tensor], + hf_outputs: dict[str, torch.Tensor], + *, + atol: float, + rtol: float, +) -> None: + if tt_outputs.keys() != hf_outputs.keys(): + raise AssertionError( + "Captured output names differ: " + f"TT-only={sorted(tt_outputs.keys() - hf_outputs.keys())}, " + f"HF-only={sorted(hf_outputs.keys() - tt_outputs.keys())}." + ) + for name in tt_outputs: + if ".router_ids." in name: + if not torch.equal(tt_outputs[name], hf_outputs[name]): + mismatch = (tt_outputs[name] != hf_outputs[name]).sum().item() + raise AssertionError(f"{name}: {mismatch} expert IDs differ.") + print(f"{name:32s} exact expert IDs") + else: + _compare_tensor( + name, + tt_outputs[name], + hf_outputs[name], + atol=atol, + rtol=rtol, + ) + + +@torch.no_grad() +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--hf_repo_path", required=True) + parser.add_argument("--device", default="cuda") + parser.add_argument( + "--tt_device", + help="TorchTitan device. Defaults to --device.", + ) + parser.add_argument( + "--hf_device", + help="HuggingFace device. Defaults to --device.", + ) + parser.add_argument( + "--device_module", + help="Optional module that registers an out-of-tree PyTorch device.", + ) + parser.add_argument( + "--hf_kda_backend", + default="fla", + choices=("fla", "reference"), + help="Use FLA or the script's pure PyTorch KDA compatibility backend.", + ) + parser.add_argument( + "--dtype", + default="float32", + choices=("float32", "bfloat16"), + ) + parser.add_argument("--atol", type=float, default=2e-4) + parser.add_argument("--rtol", type=float, default=2e-4) + parser.add_argument("--seed", type=int, default=42) + args = parser.parse_args() + + if args.device_module is not None: + importlib.import_module(args.device_module) + tt_device = torch.device(args.tt_device or args.device) + hf_device = torch.device(args.hf_device or args.device) + dtype = getattr(torch, args.dtype) + torch.manual_seed(args.seed) + if args.hf_kda_backend == "reference": + _install_reference_fla() + if "cuda" in {tt_device.type, hf_device.type}: + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + torch.set_float32_matmul_precision("highest") + + tt_config = cast( + KimiK3Model.Config, + model_registry("debugmodel").model, + ) + tt_model = tt_config.build() + tt_model.init_states() + + hf_config = _build_reduced_hf_config(args.hf_repo_path, tt_config) + hf_model = AutoModelForCausalLM.from_config( + hf_config, + trust_remote_code=True, + ) + hf_state_dict = KimiK3StateDictAdapter( + tt_config, + hf_assets_path=None, + ).to_hf(tt_model.state_dict()) + hf_model.load_state_dict(hf_state_dict, strict=True) + + # KimiLinearModel selects FlashAttention during construction. Restoring + # eager here exercises the released eager attention function. + hf_model.language_model.model.config._attn_implementation = "eager" + hf_model.config.text_config._attn_implementation = "eager" + + tt_model = tt_model.to(device=tt_device, dtype=dtype).eval() + hf_model = hf_model.to(device=hf_device, dtype=dtype).eval() + + tt_outputs: dict[str, torch.Tensor] = {} + hf_outputs: dict[str, torch.Tensor] = {} + _register_text_hooks(tt_model, hf_model, tt_outputs, hf_outputs) + _register_vision_hooks(tt_model, hf_model, tt_outputs, hf_outputs) + + print("\nText-only parity") + tokens_BL = torch.tensor( + [[11, 23, 17, 31, 5, 19, 29, 3]], + dtype=torch.long, + ) + tt_tokens_BL = tokens_BL.to(tt_device) + hf_tokens_BL = tokens_BL.to(hf_device) + attention_mask_BL = torch.ones_like(hf_tokens_BL) + tt_logits_BLV = tt_model(tt_tokens_BL).float().cpu() + hf_logits_BLV = ( + hf_model.language_model( + input_ids=hf_tokens_BL, + attention_mask=attention_mask_BL, + use_cache=False, + ) + .logits.float() + .cpu() + ) + _compare_captured( + tt_outputs, + hf_outputs, + atol=args.atol, + rtol=args.rtol, + ) + _compare_tensor( + "text.logits", + tt_logits_BLV, + hf_logits_BLV, + atol=args.atol, + rtol=args.rtol, + ) + + print("\nVision parity") + tt_outputs.clear() + hf_outputs.clear() + vision_config = tt_config.vision_encoder + if vision_config is None: + raise ValueError("The Kimi K3 debug config must include a vision encoder.") + patch_size = vision_config.patch_size + grid_thw_N3 = torch.tensor([[1, 4, 4]], dtype=torch.long) + patches_PCHW = torch.randn( + 16, + 3, + patch_size, + patch_size, + dtype=dtype, + ) + tt_grid_thw_N3 = grid_thw_N3.to(tt_device) + hf_grid_thw_N3 = grid_thw_N3.to(hf_device) + tt_patches_PCHW = patches_PCHW.to(tt_device) + hf_patches_PCHW = patches_PCHW.to(hf_device) + pixels_NPK = tt_patches_PCHW.reshape(1, 16, -1) + tt_vision_NMD = tt_model.vision_encoder( + pixels_NPK, + grid_thw=tt_grid_thw_N3, + ) + hf_vision = hf_model.vision_tower(hf_patches_PCHW, hf_grid_thw_N3) + hf_vision_NMD = torch.stack(list(hf_model.mm_projector(hf_vision))) + _compare_captured( + tt_outputs, + hf_outputs, + atol=args.atol, + rtol=args.rtol, + ) + _compare_tensor( + "vision.projected", + tt_vision_NMD.float().cpu(), + hf_vision_NMD.float().cpu(), + atol=args.atol, + rtol=args.rtol, + ) + + print("\nEnd-to-end image+text parity") + tt_outputs.clear() + hf_outputs.clear() + num_vision_tokens = tt_vision_NMD.shape[1] + hf_tokens_BL = torch.tensor( + [[11, 23, _IMAGE_TOKEN_ID, 17, 31]], + dtype=torch.long, + device=hf_device, + ) + tt_tokens_BL = torch.tensor( + [[11, 23] + [_IMAGE_TOKEN_ID] * num_vision_tokens + [17, 31]], + dtype=torch.long, + device=tt_device, + ) + hf_attention_mask_BL = torch.ones_like(hf_tokens_BL) + tt_logits_BLV = tt_model( + tt_tokens_BL, + pixel_values=pixels_NPK, + grid_thw=tt_grid_thw_N3, + special_tokens={"image_id": _IMAGE_TOKEN_ID}, + ) + hf_logits_BLV = hf_model( + input_ids=hf_tokens_BL, + pixel_values=hf_patches_PCHW, + grid_thws=hf_grid_thw_N3, + attention_mask=hf_attention_mask_BL, + use_cache=False, + return_dict=True, + ).logits + _compare_captured( + tt_outputs, + hf_outputs, + atol=args.atol, + rtol=args.rtol, + ) + _compare_tensor( + "multimodal.logits", + tt_logits_BLV.float().cpu(), + hf_logits_BLV.float().cpu(), + atol=args.atol, + rtol=args.rtol, + ) + print("\nRESULT: PASS") + + +if __name__ == "__main__": + main() diff --git a/scripts/checkpoint_conversion/numerical_tests_kimi_k3_device.py b/scripts/checkpoint_conversion/numerical_tests_kimi_k3_device.py new file mode 100644 index 0000000000..69ecdcee2e --- /dev/null +++ b/scripts/checkpoint_conversion/numerical_tests_kimi_k3_device.py @@ -0,0 +1,435 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Cross-device forward parity and training smoke test for reduced Kimi K3. + +This script compares a deterministic FP32 CPU reference against the same +TorchTitan model, weights, tokens, and image patches on another device. It then +runs real forward, backward, and AdamW updates on that device. Combine this test +with ``numerical_tests_kimi_k3.py``, which compares the CPU/CUDA TorchTitan +reference directly against the released HuggingFace implementation. + +The optional ``--device_module`` argument imports a PyTorch out-of-tree device +extension before constructing the target device. +""" + +import argparse +import importlib +import math +import time +from collections.abc import Callable + +import torch +import torch.nn.functional as F +from torch.utils.hooks import RemovableHandle + + +_IMAGE_TOKEN_ID = 7 + + +def _first_tensor(output) -> torch.Tensor: + if isinstance(output, torch.Tensor): + return output + if isinstance(output, (list, tuple)): + for value in output: + if isinstance(value, torch.Tensor): + return value + raise TypeError(f"Expected a tensor output, got {type(output).__name__}.") + + +def _capture_tensor( + destination: dict[str, torch.Tensor], + name: str, + *, + flatten: bool = False, +) -> Callable: + def hook(_module, _inputs, output) -> None: + value = _first_tensor(output).detach().float().cpu() + if flatten: + value = value.reshape(-1, value.shape[-1]) + destination[name] = value + + return hook + + +def _capture_router_ids( + destination: dict[str, torch.Tensor], + name: str, +) -> Callable: + def hook(_module, _inputs, output) -> None: + expert_ids = output[0] + destination[name] = expert_ids.detach().reshape(-1, expert_ids.shape[-1]).cpu() + + return hook + + +def _register_hooks( + model, + outputs: dict[str, torch.Tensor], +) -> list[RemovableHandle]: + handles: list[RemovableHandle] = [] + if model.vision_encoder is None: + raise ValueError("The Kimi K3 debug config must include a vision encoder.") + + for layer_idx, layer in enumerate(model.vision_encoder.layers.values()): + handles.append( + layer.register_forward_hook( + _capture_tensor(outputs, f"vision.layer.{layer_idx}", flatten=True) + ) + ) + projector = model.vision_encoder.projector + for name, module in ( + ("vision.projector.linear_1", projector.linear_1), + ("vision.projector.activation", projector.activation), + ("vision.projector.linear_2", projector.linear_2), + ("vision.projector.post_norm", projector.post_norm), + ): + handles.append( + module.register_forward_hook(_capture_tensor(outputs, name, flatten=True)) + ) + handles.append( + projector.register_forward_hook( + _capture_tensor(outputs, "vision.projected", flatten=True) + ) + ) + + for layer_idx, layer in enumerate(model.layers.values()): + handles.append( + layer.register_forward_hook( + _capture_tensor(outputs, f"decoder.layer.{layer_idx}") + ) + ) + attention = ( + layer.attention if layer.attention is not None else layer.delta_attention + ) + assert attention is not None + handles.append( + attention.register_forward_hook( + _capture_tensor(outputs, f"decoder.attention.{layer_idx}") + ) + ) + if layer.moe is not None: + handles.append( + layer.moe.router.register_forward_hook( + _capture_router_ids(outputs, f"decoder.router_ids.{layer_idx}") + ) + ) + return handles + + +def _compare_tensor( + name: str, + reference: torch.Tensor, + actual: torch.Tensor, + *, + atol: float, + rtol: float, +) -> None: + if reference.shape != actual.shape: + raise AssertionError( + f"{name}: reference shape {tuple(reference.shape)} does not match " + f"device shape {tuple(actual.shape)}." + ) + difference = (reference - actual).abs() + cosine = F.cosine_similarity( + reference.reshape(-1), + actual.reshape(-1), + dim=0, + ).item() + print( + f"{name:32s} max={difference.max().item():.4e} " + f"mean={difference.mean().item():.4e} cos={cosine:.8f}" + ) + torch.testing.assert_close( + actual, + reference, + atol=atol, + rtol=rtol, + msg=lambda message: f"{name} failed device parity:\n{message}", + ) + + +def _compare_captured( + reference_outputs: dict[str, torch.Tensor], + device_outputs: dict[str, torch.Tensor], + *, + atol: float, + rtol: float, +) -> None: + if reference_outputs.keys() != device_outputs.keys(): + raise AssertionError( + "Captured output names differ: " + f"reference-only={sorted(reference_outputs.keys() - device_outputs.keys())}, " + f"device-only={sorted(device_outputs.keys() - reference_outputs.keys())}." + ) + for name in reference_outputs: + if ".router_ids." in name: + if not torch.equal(reference_outputs[name], device_outputs[name]): + mismatch = ( + (reference_outputs[name] != device_outputs[name]).sum().item() + ) + raise AssertionError(f"{name}: {mismatch} expert IDs differ.") + print(f"{name:32s} exact expert IDs") + else: + _compare_tensor( + name, + reference_outputs[name], + device_outputs[name], + atol=atol, + rtol=rtol, + ) + + +def _build_inputs(config, *, seed: int) -> dict[str, torch.Tensor | dict[str, int]]: + vision_config = config.vision_encoder + if vision_config is None: + raise ValueError("The Kimi K3 debug config must include a vision encoder.") + + generator = torch.Generator(device="cpu") + generator.manual_seed(seed + 1) + patch_size = vision_config.patch_size + grid_thw_N3 = torch.tensor([[1, 4, 4]], dtype=torch.long) + pixel_values_NPK = torch.randn( + 1, + 16, + 3 * patch_size * patch_size, + generator=generator, + dtype=torch.float32, + ) + tokens_BL = torch.tensor( + [ + [ + 11, + 23, + _IMAGE_TOKEN_ID, + _IMAGE_TOKEN_ID, + _IMAGE_TOKEN_ID, + _IMAGE_TOKEN_ID, + 17, + 31, + ] + ], + dtype=torch.long, + ) + return { + "tokens": tokens_BL, + "pixel_values": pixel_values_NPK, + "grid_thw": grid_thw_N3, + "special_tokens": {"image_id": _IMAGE_TOKEN_ID}, + } + + +def _move_inputs( + inputs: dict[str, torch.Tensor | dict[str, int]], + device: torch.device, +) -> dict[str, torch.Tensor | dict[str, int]]: + return { + name: value.to(device) if isinstance(value, torch.Tensor) else value + for name, value in inputs.items() + } + + +def _synchronize(device: torch.device) -> None: + device_api = getattr(torch, device.type, None) + if device_api is not None and hasattr(device_api, "synchronize"): + device_api.synchronize(device) + + +def _verify_device(device: torch.device) -> None: + if device.type == "cpu": + return + device_api = getattr(torch, device.type, None) + if device_api is None: + raise RuntimeError( + f"PyTorch has no registered {device.type!r} device module. " + "Use --device_module to import its extension." + ) + if hasattr(device_api, "is_available") and not device_api.is_available(): + raise RuntimeError(f"Requested device {device} is not available.") + + +@torch.no_grad() +def _run_forward_parity( + reference_model, + device_model, + reference_inputs, + device_inputs, + *, + atol: float, + rtol: float, +) -> None: + reference_outputs: dict[str, torch.Tensor] = {} + device_outputs: dict[str, torch.Tensor] = {} + reference_handles = _register_hooks(reference_model, reference_outputs) + device_handles = _register_hooks(device_model, device_outputs) + + try: + reference_logits_BLV = reference_model(**reference_inputs).float().cpu() + device_logits_BLV = device_model(**device_inputs).float().cpu() + finally: + for handle in reference_handles + device_handles: + handle.remove() + + _compare_captured( + reference_outputs, + device_outputs, + atol=atol, + rtol=rtol, + ) + _compare_tensor( + "multimodal.logits", + reference_logits_BLV, + device_logits_BLV, + atol=atol, + rtol=rtol, + ) + + +def _run_training( + model, + inputs, + device: torch.device, + *, + train_dtype: torch.dtype, + train_steps: int, + learning_rate: float, +) -> None: + if train_steps == 0: + return + + model.to(dtype=train_dtype) + model.train() + optimizer = torch.optim.AdamW( + model.parameters(), + lr=learning_rate, + foreach=False, + ) + tokens_BL = inputs["tokens"] + assert isinstance(tokens_BL, torch.Tensor) + if model.lm_head is None: + raise ValueError("Kimi K3 device training requires an LM head.") + tracked_parameter = model.lm_head.weight + tracked_before = tracked_parameter.detach().clone() + + for step_idx in range(train_steps): + optimizer.zero_grad(set_to_none=True) + _synchronize(device) + start_time = time.perf_counter() + logits_BLV = model(**inputs) + loss = F.cross_entropy( + logits_BLV[:, :-1].float().reshape(-1, logits_BLV.shape[-1]), + tokens_BL[:, 1:].reshape(-1), + ) + loss.backward() + grad_norm = torch.nn.utils.clip_grad_norm_( + model.parameters(), + max_norm=math.inf, + foreach=False, + ) + optimizer.step() + _synchronize(device) + elapsed = time.perf_counter() - start_time + + loss_value = loss.detach().float().cpu().item() + grad_norm_value = grad_norm.detach().float().cpu().item() + if not math.isfinite(loss_value) or not math.isfinite(grad_norm_value): + raise AssertionError( + f"Step {step_idx}: non-finite loss={loss_value} or " + f"grad_norm={grad_norm_value}." + ) + print( + f"step={step_idx:02d} loss={loss_value:.8f} " + f"grad_norm={grad_norm_value:.8f} elapsed={elapsed:.3f}s" + ) + + parameter_delta = ( + (tracked_parameter.detach().float() - tracked_before.float()).abs().max() + ) + parameter_delta_value = parameter_delta.cpu().item() + if parameter_delta_value == 0.0: + raise AssertionError("AdamW completed without changing lm_head.weight.") + print(f"lm_head.weight max update={parameter_delta_value:.4e}") + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--device", required=True) + parser.add_argument( + "--device_module", + help="Optional module that registers an out-of-tree PyTorch device.", + ) + parser.add_argument( + "--parity_dtype", + default="float32", + choices=("float32", "bfloat16"), + ) + parser.add_argument( + "--train_dtype", + default="bfloat16", + choices=("float32", "bfloat16"), + ) + parser.add_argument("--atol", type=float, default=2e-4) + parser.add_argument("--rtol", type=float, default=2e-4) + parser.add_argument("--train_steps", type=int, default=2) + parser.add_argument("--learning_rate", type=float, default=8e-4) + parser.add_argument("--seed", type=int, default=42) + args = parser.parse_args() + + if args.device_module is not None: + importlib.import_module(args.device_module) + + from torchtitan.models.kimi_k3 import model_registry + + device = torch.device(args.device) + _verify_device(device) + parity_dtype = getattr(torch, args.parity_dtype) + train_dtype = getattr(torch, args.train_dtype) + if args.train_steps < 0: + raise ValueError("--train_steps must be non-negative.") + + torch.manual_seed(args.seed) + if device.type == "cuda": + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + torch.set_float32_matmul_precision("highest") + + config = model_registry("debugmodel").model + reference_model = config.build() + reference_model.init_states() + device_model = config.build() + device_model.init_states() + device_model.load_state_dict(reference_model.state_dict(), strict=True) + + reference_model = reference_model.to(dtype=torch.float32).eval() + device_model = device_model.to(device=device, dtype=parity_dtype).eval() + reference_inputs = _build_inputs(config, seed=args.seed) + device_inputs = _move_inputs(reference_inputs, device) + + print(f"\nForward parity: cpu/float32 -> {device}/{args.parity_dtype}") + _run_forward_parity( + reference_model, + device_model, + reference_inputs, + device_inputs, + atol=args.atol, + rtol=args.rtol, + ) + + del reference_model + print(f"\nTraining smoke: {device}/{args.train_dtype}, steps={args.train_steps}") + _run_training( + device_model, + device_inputs, + device, + train_dtype=train_dtype, + train_steps=args.train_steps, + learning_rate=args.learning_rate, + ) + print("\nRESULT: PASS") + + +if __name__ == "__main__": + main() diff --git a/tests/unit_tests/test_kimi_k3.py b/tests/unit_tests/test_kimi_k3.py new file mode 100644 index 0000000000..6f93a2083a --- /dev/null +++ b/tests/unit_tests/test_kimi_k3.py @@ -0,0 +1,289 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import unittest + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from torchtitan.models.common import Embedding +from torchtitan.models.kimi_k3 import ( + _feed_forward_config, + _kda_config, + _latent_moe_config, + _linear, + _mla_config, + _norm, + _vision_encoder_config, + kimi_k3_configs, +) +from torchtitan.models.kimi_k3.model import ( + KimiK3Model, + KimiK3TransformerBlock, + KimiKDAKernel, +) +from torchtitan.models.kimi_k3.state_dict_adapter import KimiK3StateDictAdapter +from torchtitan.models.kimi_k3.vision_encoder import KimiExactGELU + + +def _small_model_config() -> KimiK3Model.Config: + dim = 16 + + def block( + layer_id: int, + *, + use_mla: bool, + use_moe: bool, + ) -> KimiK3TransformerBlock.Config: + return KimiK3TransformerBlock.Config( + layer_id=layer_id, + attn_res_block_size=1, + attention=( + _mla_config( + dim=dim, + num_heads=2, + q_lora_rank=8, + kv_lora_rank=8, + qk_nope_head_dim=4, + qk_rope_head_dim=4, + v_head_dim=4, + ) + if use_mla + else None + ), + delta_attention=( + None + if use_mla + else _kda_config( + dim=dim, + num_heads=2, + head_dim=4, + conv_kernel_size=3, + ) + ), + feed_forward=( + None if use_moe else _feed_forward_config(dim=dim, hidden_dim=32) + ), + moe=( + _latent_moe_config( + dim=dim, + latent_dim=8, + expert_hidden_dim=8, + num_experts=2, + top_k=1, + num_shared_experts=1, + ) + if use_moe + else None + ), + attention_norm=_norm(dim), + ffn_norm=_norm(dim), + attention_res_norm=_norm(dim), + attention_res_proj=_linear(dim, 1), + ffn_res_norm=_norm(dim), + ffn_res_proj=_linear(dim, 1), + ) + + return KimiK3Model.Config( + dim=dim, + vocab_size=32, + tok_embeddings=Embedding.Config( + num_embeddings=32, + embedding_dim=dim, + param_init={ + "weight": lambda parameter: nn.init.normal_(parameter, std=0.02) + }, + ), + layers=[ + block(0, use_mla=False, use_moe=False), + block(1, use_mla=True, use_moe=True), + ], + norm=_norm(dim), + lm_head=_linear(dim, 32), + output_res_norm=_norm(dim), + output_res_proj=_linear(dim, 1), + vision_encoder=_vision_encoder_config( + text_dim=dim, + dim=16, + qkv_dim=24, + hidden_dim=32, + num_layers=1, + num_heads=3, + patch_size=2, + merge_kernel_size=(2, 2), + init_pos_emb_height=2, + init_pos_emb_width=2, + max_num_frames=1, + ), + spatial_merge_size=2, + ) + + +def _kda_recurrent_reference( + q_BLHK: torch.Tensor, + k_BLHK: torch.Tensor, + v_BLHV: torch.Tensor, + gate_BLHK: torch.Tensor, + beta_BLH: torch.Tensor, + A_log_H: torch.Tensor, + dt_bias_HK: torch.Tensor, + *, + lower_bound: float, +) -> torch.Tensor: + q_BLHK = q_BLHK.float() + k_BLHK = k_BLHK.float() + q_BLHK = q_BLHK * torch.rsqrt(q_BLHK.square().sum(dim=-1, keepdim=True) + 1e-6) + k_BLHK = k_BLHK * torch.rsqrt(k_BLHK.square().sum(dim=-1, keepdim=True) + 1e-6) + log_decay_BLHK = lower_bound * torch.sigmoid( + torch.exp(A_log_H.float()).view(1, 1, -1, 1) + * (gate_BLHK.float() + dt_bias_HK.float()) + ) + decay_BLHK = torch.exp(log_decay_BLHK) + beta_BLH = torch.sigmoid(beta_BLH.float()) + + B, L, H, K = q_BLHK.shape + V = v_BLHV.shape[-1] + state_BHKV = torch.zeros(B, H, K, V) + output_BLHV = torch.empty(B, L, H, V) + for token_idx in range(L): + state_BHKV = state_BHKV * decay_BLHK[:, token_idx].unsqueeze(-1) + old_value_BHV = torch.matmul( + k_BLHK[:, token_idx].unsqueeze(-2), + state_BHKV, + ).squeeze(-2) + delta_BHV = (v_BLHV[:, token_idx].float() - old_value_BHV) * beta_BLH[ + :, token_idx + ].unsqueeze(-1) + state_BHKV = state_BHKV + ( + k_BLHK[:, token_idx].unsqueeze(-1) * delta_BHV.unsqueeze(-2) + ) + output_BLHV[:, token_idx] = torch.matmul( + q_BLHK[:, token_idx].unsqueeze(-2), + state_BHKV, + ).squeeze(-2) * (K**-0.5) + return output_BLHV + + +class TestKimiK3(unittest.TestCase): + def test_exact_gelu_matches_pytorch_reference(self): + x = torch.linspace(-4.0, 4.0, 257) + actual = KimiExactGELU.Config().build()(x) + expected = F.gelu(x, approximate="none") + + torch.testing.assert_close(actual, expected, atol=1e-6, rtol=1e-6) + + x_bf16 = x.bfloat16() + self.assertEqual(KimiExactGELU.Config().build()(x_bf16).dtype, x_bf16.dtype) + + def test_debugmodel_preserves_reduced_k3_topology(self): + config = kimi_k3_configs["debugmodel"]("eager") + + self.assertEqual(len(config.layers), 13) + self.assertEqual( + [ + layer_idx + 1 + for layer_idx, layer in enumerate(config.layers) + if layer.attention is not None + ], + [4, 8, 12], + ) + self.assertIsNotNone(config.layers[0].feed_forward) + self.assertTrue(all(layer.moe is not None for layer in config.layers[1:])) + moe_config = config.layers[1].moe + assert moe_config is not None + self.assertEqual(moe_config.num_experts, 8) + self.assertEqual(moe_config.router.top_k, 2) + + def test_kda_kernel_matches_recurrent_reference(self): + torch.manual_seed(1) + q_BLHK = torch.randn(2, 5, 3, 4, requires_grad=True) + k_BLHK = torch.randn(2, 5, 3, 4, requires_grad=True) + v_BLHV = torch.randn(2, 5, 3, 4, requires_grad=True) + gate_BLHK = torch.randn(2, 5, 3, 4, requires_grad=True) + beta_BLH = torch.randn(2, 5, 3, requires_grad=True) + A_log_H = torch.randn(3, requires_grad=True) + dt_bias_HK = torch.randn(3, 4, requires_grad=True) + + kernel = KimiKDAKernel.Config( + head_dim=4, + lower_bound=-5.0, + ).build() + actual_BLHV = kernel( + q_BLHK, + k_BLHK, + v_BLHV, + gate_BLHK, + beta_BLH, + A_log_H, + dt_bias_HK, + ) + expected_BLHV = _kda_recurrent_reference( + q_BLHK, + k_BLHK, + v_BLHV, + gate_BLHK, + beta_BLH, + A_log_H, + dt_bias_HK, + lower_bound=-5.0, + ) + + torch.testing.assert_close( + actual_BLHV, + expected_BLHV, + atol=1e-6, + rtol=1e-6, + ) + actual_BLHV.square().mean().backward() + for tensor in ( + q_BLHK, + k_BLHK, + v_BLHV, + gate_BLHK, + beta_BLH, + A_log_H, + dt_bias_HK, + ): + self.assertIsNotNone(tensor.grad) + self.assertTrue(torch.isfinite(tensor.grad).all()) + + def test_small_multimodal_model_forward_backward_and_adapter(self): + torch.manual_seed(2) + config = _small_model_config() + model = config.build() + model.verify_module_protocol() + model.init_states() + + tokens_BL = torch.randint(0, config.vocab_size, (2, 6)) + image_token_id = 7 + tokens_BL[0, 2] = image_token_id + pixel_values_NPK = torch.randn(1, 4, 3 * 2 * 2) + grid_thw_N3 = torch.tensor([[1, 2, 2]]) + logits_BLV = model( + tokens_BL, + pixel_values=pixel_values_NPK, + grid_thw=grid_thw_N3, + special_tokens={"image_id": image_token_id}, + ) + + self.assertEqual(logits_BLV.shape, (2, 6, config.vocab_size)) + logits_BLV.float().square().mean().backward() + for parameter in model.parameters(): + if parameter.grad is not None: + self.assertTrue(torch.isfinite(parameter.grad).all()) + + state_dict = model.state_dict() + adapter = KimiK3StateDictAdapter(config, hf_assets_path=None) + hf_state_dict = adapter.to_hf(state_dict) + roundtrip_state_dict = adapter.from_hf(hf_state_dict) + self.assertEqual(state_dict.keys(), roundtrip_state_dict.keys()) + for key, value in state_dict.items(): + torch.testing.assert_close(value, roundtrip_state_dict[key]) + + +if __name__ == "__main__": + unittest.main() diff --git a/torchtitan/models/__init__.py b/torchtitan/models/__init__.py index 784b4110ca..0f22e40f1a 100644 --- a/torchtitan/models/__init__.py +++ b/torchtitan/models/__init__.py @@ -10,6 +10,7 @@ "flux", "gpt_oss", "kimi_k2_7", + "kimi_k3", "llama3", "muse_glimmer", "qwen3", diff --git a/torchtitan/models/kimi_k3/README.md b/torchtitan/models/kimi_k3/README.md new file mode 100644 index 0000000000..265ad07257 --- /dev/null +++ b/torchtitan/models/kimi_k3/README.md @@ -0,0 +1,132 @@ +# Kimi K3 + +This directory contains the first TorchTitan implementation of Kimi K3. The +initial scope is a topology-complete, reduced model for single-device training +and numerical comparison with the +[released HuggingFace implementation](https://huggingface.co/moonshotai/Kimi-K3). + +The implementation is device-neutral. It uses PyTorch operators and does not +import accelerator-specific packages. The reference kernels prioritize +inspectable model math over throughput. + +## Quick start + +```bash +NGPU=1 MODULE=kimi_k3 CONFIG=kimi_k3_debugmodel ./run_train.sh +``` + +The multimodal data path requires `torchvision`. + +## Reduced model + +The `debugmodel` flavor preserves each distinct Kimi K3 forward path while +reducing widths, expert count, and depth. + +| Component | Released Kimi K3 | `debugmodel` | +|---|---:|---:| +| Decoder dimension | 7168 | 256 | +| Vocabulary size | 163840 | 2048 | +| Decoder layers | 93 | 13 | +| Full MLA layers (1-based) | 4, 8, ..., 92, 93 | 4, 8, 12 | +| KDA layers | 69 | 10 | +| Dense FFN layers | 1 | 1 | +| Attention residual block size | 12 | 12 | +| Routed experts / top-k | 896 / 16 | 8 / 2 | +| Shared experts | 2 | 2 | +| Vision dimension | 1024 | 128 | +| Vision layers | 27 | 2 | +| Vision QKV dimension / heads | 1536 / 12 | 192 / 3 | + +Thirteen decoder layers are intentional. They exercise two attention-residual +blocks and preserve the released model's 1-based full-attention cadence. + +## Forward structure + +The reference path mirrors the released implementation in these areas: + +- FP32-reduction RMSNorm and SiTU activation. +- Gated MLA with low-rank query and KV projections. Kimi K3 sets + `mla_use_nope=True`, so the RoPE-sized query/key slices are not rotated. +- KDA short causal convolutions, safe decay gate, query/key L2 normalization, + sigmoid beta, recurrent delta-rule update, and gated output RMSNorm. +- Block-level attention residuals, including the final output residual. +- Stable LatentMoE with sigmoid top-k routing, correction bias, latent + down/up projections, routed experts, and shared experts. +- MoonViT3d patch embedding, learned spatial positions, 2D RoPE, non-causal + per-image attention, temporal pooling, 2x2 spatial merge, and + PatchMergerMLPV2. +- Vision features scattered into runs of media placeholder tokens. + +`KimiKDAKernel` is the optimization boundary. A future accelerated backend +should preserve its input/output contract and checkpoint schema. + +## Checkpoint conversion + +`KimiK3StateDictAdapter` converts between TorchTitan and an unquantized +HuggingFace state dict. It covers: + +- dense, MLA, KDA, and LatentMoE decoder layers; +- attention-residual parameters; +- vision patch embedding, fused HuggingFace QKV, transformer blocks, and + projector; +- the MoE correction bias. + +The released checkpoint stores routed expert weights in MXFP4. Loading those +compressed tensors is outside this first change. Numerical comparison should +therefore instantiate the same reduced, unquantized model on both sides and +copy one state dict through the adapter. + +## Validation contract + +The CPU unit tests cover: + +- the reduced layer topology; +- the explicit exact GELU against PyTorch's CPU reference; +- the KDA kernel against a direct recurrent formulation, including backward; +- a small text+image model forward and backward; +- exhaustive state-dict round-trip for that small model. + +Before requesting merge, the following CUDA tests must also pass: + +1. One single-device training step on a CUDA GPU. +2. FP32 forward comparison between CPU and CUDA using the same model and inputs. +3. FP32 forward comparison against the released HuggingFace code using the + same reduced config, weights, tokens, image patches, and expert choices. +4. BF16 forward, backward, and optimizer steps on CUDA. + +The parity report must include intermediate checks for KDA, MLA, vision +features, each decoder layer, final logits, and router expert IDs. Final-logit +metrics alone are not sufficient to localize a discrete routing mismatch. + +Run the CPU-to-CUDA comparison and real optimizer steps with: + +```bash +python -m scripts.checkpoint_conversion.numerical_tests_kimi_k3_device \ + --device cuda +``` + +This device test is one link in the parity chain; it does not replace the +direct HuggingFace comparison below. + +Run both implementations on a CUDA GPU with the released model dependencies: + +```bash +CUDA_VISIBLE_DEVICES=0 python -m \ + scripts.checkpoint_conversion.numerical_tests_kimi_k3 \ + --hf_repo_path /path/to/moonshotai/Kimi-K3 +``` + +## First-version limitations + +- Single device only; DP, FSDP/HSDP, TP, PP, CP, and EP are rejected. +- No packed documents, activation checkpointing, `torch.compile`, or CPU + offload. +- Image inputs are supported; video inputs are rejected. +- No generation cache. +- No optimized KDA backend. +- No MXFP4 checkpoint loading. +- No full 2.8T flavor. + +These restrictions are explicit so unsupported runtime settings fail instead +of being silently ignored. Optimized kernels and distributed parallelism can be +added in follow-up changes after the reference forward is numerically locked. diff --git a/torchtitan/models/kimi_k3/__init__.py b/torchtitan/models/kimi_k3/__init__.py new file mode 100644 index 0000000000..d7ca6015cc --- /dev/null +++ b/torchtitan/models/kimi_k3/__init__.py @@ -0,0 +1,469 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Kimi K3 model registration and architecture configurations.""" + +from collections.abc import Callable +from functools import partial + +import torch +import torch.nn as nn + +from torchtitan.components.optimizer import register_moe_load_balancing_hook +from torchtitan.models.common import Conv1d, Embedding, Linear +from torchtitan.models.common.nn_modules import GELU, RMSNorm +from torchtitan.models.common.vision_encoder import VisionMLP +from torchtitan.models.utils import validate_converter_order +from torchtitan.protocols.model import ModelConfigConverter +from torchtitan.protocols.model_spec import ModelSpec + +from .model import ( + KimiDeltaAttention, + KimiFeedForward, + KimiK3Model, + KimiK3TransformerBlock, + KimiKDAKernel, + KimiLatentMoE, + KimiMLAAttention, + KimiMoERouter, + KimiRMSNorm, + KimiRMSNormGated, + SituAndMul, +) +from .parallelize import parallelize_kimi_k3 +from .state_dict_adapter import KimiK3StateDictAdapter +from .vision_encoder import ( + KimiExactGELU, + KimiK3VisionAttention, + KimiK3VisionBlock, + KimiK3VisionEncoder, + KimiK3VisionProjector, + VisionRotaryEmbedding2D, +) + +__all__ = [ + "KIMI_K3_SPECIAL_TOKENS", + "KimiK3Model", + "KimiK3StateDictAdapter", + "KimiK3VisionEncoder", + "kimi_k3_configs", + "model_registry", + "parallelize_kimi_k3", +] + + +KIMI_K3_SPECIAL_TOKENS = { + "image_token": "<|media_pad|>", + "video_token": "<|media_pad|>", + "vision_start_token": "<|media_begin|>", + "vision_end_token": "<|media_end|>", + "pad_token": "[PAD]", +} + + +_LINEAR_INIT = { + "weight": partial(nn.init.trunc_normal_, std=0.02), + "bias": nn.init.zeros_, +} +_CONV_INIT = {"weight": partial(nn.init.trunc_normal_, std=0.02)} +_NORM_INIT = {"weight": nn.init.ones_} +_EMBEDDING_INIT = {"weight": partial(nn.init.normal_, std=1.0)} +_POS_EMBED_INIT = {"pos_embed": partial(nn.init.normal_, std=1.0)} + + +def _output_linear_init(dim: int) -> dict[str, Callable]: + scale = dim**-0.5 + return { + "weight": partial( + nn.init.trunc_normal_, + std=scale, + a=-3 * scale, + b=3 * scale, + ) + } + + +def _fan_in_linear_init(in_features: int) -> dict[str, Callable]: + return { + "weight": partial( + nn.init.trunc_normal_, + std=(2.0 / in_features) ** 0.5, + ), + "bias": nn.init.zeros_, + } + + +def _a_log_init(parameter: nn.Parameter) -> None: + with torch.no_grad(): + nn.init.uniform_(parameter, 1.0, 16.0) + parameter.log_() + + +def _linear( + in_features: int, + out_features: int, + *, + bias: bool = False, + param_init: dict[str, Callable] | None = None, +) -> Linear.Config: + return Linear.Config( + in_features=in_features, + out_features=out_features, + bias=bias, + param_init=param_init or _LINEAR_INIT, + ) + + +def _norm(dim: int, eps: float = 1e-5) -> KimiRMSNorm.Config: + return KimiRMSNorm.Config( + normalized_shape=dim, + eps=eps, + param_init=_NORM_INIT, + ) + + +def _feed_forward_config( + *, + dim: int, + hidden_dim: int, +) -> KimiFeedForward.Config: + return KimiFeedForward.Config( + w1=_linear(dim, hidden_dim), + w2=_linear(hidden_dim, dim), + w3=_linear(dim, hidden_dim), + activation=SituAndMul.Config(beta=4.0, linear_beta=25.0), + ) + + +def _mla_config( + *, + dim: int, + num_heads: int, + q_lora_rank: int, + kv_lora_rank: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + v_head_dim: int, +) -> KimiMLAAttention.Config: + q_head_dim = qk_nope_head_dim + qk_rope_head_dim + return KimiMLAAttention.Config( + dim=dim, + num_heads=num_heads, + q_lora_rank=q_lora_rank, + kv_lora_rank=kv_lora_rank, + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=qk_rope_head_dim, + v_head_dim=v_head_dim, + wq_a=_linear(dim, q_lora_rank), + q_norm=_norm(q_lora_rank), + wq_b=_linear(q_lora_rank, num_heads * q_head_dim), + wkv_a=_linear(dim, kv_lora_rank + qk_rope_head_dim), + kv_norm=_norm(kv_lora_rank), + wkv_b=_linear( + kv_lora_rank, + num_heads * (qk_nope_head_dim + v_head_dim), + ), + gate=_linear(dim, num_heads * v_head_dim), + wo=_linear(num_heads * v_head_dim, dim), + ) + + +def _kda_config( + *, + dim: int, + num_heads: int, + head_dim: int, + conv_kernel_size: int, +) -> KimiDeltaAttention.Config: + projection_dim = num_heads * head_dim + + def conv() -> Conv1d.Config: + return Conv1d.Config( + in_channels=projection_dim, + out_channels=projection_dim, + kernel_size=conv_kernel_size, + groups=projection_dim, + bias=False, + param_init=_CONV_INIT, + ) + + return KimiDeltaAttention.Config( + dim=dim, + num_heads=num_heads, + head_dim=head_dim, + conv_kernel_size=conv_kernel_size, + q_proj=_linear(dim, projection_dim), + k_proj=_linear(dim, projection_dim), + v_proj=_linear(dim, projection_dim), + q_conv=conv(), + k_conv=conv(), + v_conv=conv(), + forget_a=_linear(dim, head_dim), + forget_b=_linear(head_dim, projection_dim), + beta=_linear(dim, num_heads), + output_gate=_linear(dim, projection_dim), + kernel=KimiKDAKernel.Config( + head_dim=head_dim, + lower_bound=-5.0, + ), + output_norm=KimiRMSNormGated.Config( + dim=head_dim, + eps=1e-5, + param_init=_NORM_INIT, + ), + output_proj=_linear(projection_dim, dim), + param_init={ + "A_log": _a_log_init, + "dt_bias": nn.init.zeros_, + }, + ) + + +def _latent_moe_config( + *, + dim: int, + latent_dim: int, + expert_hidden_dim: int, + num_experts: int, + top_k: int, + num_shared_experts: int, +) -> KimiLatentMoE.Config: + routed_experts = [ + _feed_forward_config( + dim=latent_dim, + hidden_dim=expert_hidden_dim, + ) + for _ in range(num_experts) + ] + return KimiLatentMoE.Config( + num_experts=num_experts, + router=KimiMoERouter.Config( + num_experts=num_experts, + top_k=top_k, + gate=_linear(dim, num_experts), + route_norm=True, + route_scale=1.0, + ), + routed_down=_linear(dim, latent_dim), + routed_experts=routed_experts, + routed_norm=_norm(latent_dim), + routed_up=_linear(latent_dim, dim), + shared_experts=_feed_forward_config( + dim=dim, + hidden_dim=num_shared_experts * expert_hidden_dim, + ), + load_balance_coeff=1e-3, + ) + + +def _vision_encoder_config( + *, + text_dim: int, + dim: int, + qkv_dim: int, + hidden_dim: int, + num_layers: int, + num_heads: int, + patch_size: int = 14, + in_channels: int = 3, + merge_kernel_size: tuple[int, int] = (2, 2), + init_pos_emb_height: int = 16, + init_pos_emb_width: int = 16, + max_num_frames: int = 4, +) -> KimiK3VisionEncoder.Config: + patch_dim = in_channels * patch_size * patch_size + head_dim = qkv_dim // num_heads + merged_dim = dim * merge_kernel_size[0] * merge_kernel_size[1] + vision_norm = RMSNorm.Config( + normalized_shape=dim, + eps=1e-5, + param_init=_NORM_INIT, + ) + block = KimiK3VisionBlock.Config( + norm1=vision_norm, + norm2=vision_norm, + attn=KimiK3VisionAttention.Config( + qkv_dim=qkv_dim, + num_heads=num_heads, + wq=_linear(dim, qkv_dim), + wk=_linear(dim, qkv_dim), + wv=_linear(dim, qkv_dim), + proj=_linear(qkv_dim, dim), + ), + mlp=VisionMLP.Config( + fc1=_linear( + dim, + hidden_dim, + param_init=_fan_in_linear_init(dim), + ), + fc2=_linear( + hidden_dim, + dim, + param_init=_fan_in_linear_init(hidden_dim), + ), + act_fn=GELU.Config(approximate="tanh"), + ), + ) + return KimiK3VisionEncoder.Config( + dim=dim, + num_layers=num_layers, + patch_size=patch_size, + in_channels=in_channels, + merge_kernel_size=merge_kernel_size, + init_pos_emb_height=init_pos_emb_height, + init_pos_emb_width=init_pos_emb_width, + max_num_frames=max_num_frames, + interpolation_mode="bilinear", + patch_embed_proj=_linear(patch_dim, dim), + rotary_pos_emb=VisionRotaryEmbedding2D.Config(head_dim=head_dim), + block=block, + final_norm=vision_norm, + projector=KimiK3VisionProjector.Config( + merged_dim=merged_dim, + linear_1=_linear( + merged_dim, + merged_dim, + param_init=_fan_in_linear_init(merged_dim), + ), + linear_2=_linear( + merged_dim, + text_dim, + param_init=_fan_in_linear_init(merged_dim), + ), + post_norm=RMSNorm.Config( + normalized_shape=text_dim, + eps=1e-5, + param_init=_NORM_INIT, + ), + activation=KimiExactGELU.Config(), + ), + param_init=_POS_EMBED_INIT, + ) + + +def _debugmodel(attn_backend: str) -> KimiK3Model.Config: + if attn_backend != "eager": + raise ValueError("Kimi K3 v1 only provides the device-neutral 'eager' backend.") + + dim = 256 + vocab_size = 2048 + num_layers = 13 + full_attention_layers = {4, 8, 12} + num_heads = 4 + qk_nope_head_dim = 32 + qk_rope_head_dim = 16 + v_head_dim = 32 + + layers = [] + for layer_idx in range(num_layers): + is_full_attention = (layer_idx + 1) in full_attention_layers + layers.append( + KimiK3TransformerBlock.Config( + layer_id=layer_idx, + attn_res_block_size=12, + attention=( + _mla_config( + dim=dim, + num_heads=num_heads, + q_lora_rank=128, + kv_lora_rank=64, + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=qk_rope_head_dim, + v_head_dim=v_head_dim, + ) + if is_full_attention + else None + ), + delta_attention=( + None + if is_full_attention + else _kda_config( + dim=dim, + num_heads=num_heads, + head_dim=32, + conv_kernel_size=4, + ) + ), + feed_forward=( + _feed_forward_config(dim=dim, hidden_dim=1024) + if layer_idx == 0 + else None + ), + moe=( + None + if layer_idx == 0 + else _latent_moe_config( + dim=dim, + latent_dim=128, + expert_hidden_dim=64, + num_experts=8, + top_k=2, + num_shared_experts=2, + ) + ), + attention_norm=_norm(dim), + ffn_norm=_norm(dim), + attention_res_norm=_norm(dim), + attention_res_proj=_linear(dim, 1), + ffn_res_norm=_norm(dim), + ffn_res_proj=_linear(dim, 1), + ) + ) + + return KimiK3Model.Config( + dim=dim, + vocab_size=vocab_size, + tok_embeddings=Embedding.Config( + num_embeddings=vocab_size, + embedding_dim=dim, + param_init=_EMBEDDING_INIT, + ), + layers=layers, + norm=_norm(dim), + lm_head=_linear( + dim, + vocab_size, + param_init=_output_linear_init(dim), + ), + output_res_norm=_norm(dim), + output_res_proj=_linear(dim, 1), + vision_encoder=_vision_encoder_config( + text_dim=dim, + dim=128, + qkv_dim=192, + hidden_dim=512, + num_layers=2, + num_heads=3, + ), + spatial_merge_size=2, + ) + + +kimi_k3_configs = { + "debugmodel": _debugmodel, +} + + +def model_registry( + flavor: str, + attn_backend: str = "eager", + converters: list[ModelConfigConverter.Config] | None = None, +) -> ModelSpec: + """Build a Kimi K3 model specification.""" + config = kimi_k3_configs[flavor](attn_backend=attn_backend) + if converters is not None: + validate_converter_order(converters) + for converter in converters: + config = converter.build().convert(config) + return ModelSpec( + name="kimi_k3", + flavor=flavor, + model=config, + parallelize_fn=parallelize_kimi_k3, + pipelining_fn=None, + post_optimizer_build_fn=register_moe_load_balancing_hook, + state_dict_adapter=KimiK3StateDictAdapter, + ) diff --git a/torchtitan/models/kimi_k3/config_registry.py b/torchtitan/models/kimi_k3/config_registry.py new file mode 100644 index 0000000000..d31b364e71 --- /dev/null +++ b/torchtitan/models/kimi_k3/config_registry.py @@ -0,0 +1,77 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Trainer configurations for Kimi K3.""" + +from torchtitan.components.checkpoint import CheckpointManager +from torchtitan.components.loss import ChunkedLossWrapper, CrossEntropyLoss +from torchtitan.components.lr_scheduler import LRSchedulersContainer +from torchtitan.components.metrics import MetricsProcessor +from torchtitan.components.optimizer import default_adamw +from torchtitan.components.tokenizer import MultiModalTokenizer +from torchtitan.config import ParallelismConfig, TrainingConfig +from torchtitan.hf_datasets.multimodal.mm_datasets import MMDataLoader +from torchtitan.hf_datasets.multimodal.utils.image import resize_to_patch_budget +from torchtitan.models.common.config_utils import decoder_vocab_size +from torchtitan.trainer import Trainer + +from . import KIMI_K3_SPECIAL_TOKENS, model_registry + + +def kimi_k3_debugmodel() -> Trainer.Config: + """Return the single-device, topology-complete Kimi K3 debug config.""" + model_spec = model_registry("debugmodel") + return Trainer.Config( + loss=ChunkedLossWrapper.Config( + loss_fn=CrossEntropyLoss.Config( + global_vocab_size=decoder_vocab_size(model_spec), + ), + ), + hf_assets_path="./tests/assets/tokenizer", + tokenizer=MultiModalTokenizer.Config(**KIMI_K3_SPECIAL_TOKENS), + metrics=MetricsProcessor.Config(log_freq=1), + model_spec=model_spec, + dataloader=MMDataLoader.Config( + dataset="cc12m-test", + max_images_per_batch=8, + patch_size=14, + temporal_patch_size=1, + spatial_merge_size=2, + patch_order="raster", + resize_fn=resize_to_patch_budget, + min_pixels=56 * 56, + max_pixels=224 * 224, + max_patches=256, + max_patches_per_side=16, + image_mean=(0.5, 0.5, 0.5), + image_std=(0.5, 0.5, 0.5), + ), + optimizer=default_adamw(lr=8e-4), + lr_scheduler=LRSchedulersContainer.Config( + warmup_steps=2, + decay_ratio=0.8, + decay_type="linear", + min_lr_factor=0.0, + ), + training=TrainingConfig( + local_batch_size=1, + seq_len=128, + steps=10, + dtype="bfloat16", + ), + parallelism=ParallelismConfig( + data_parallel_shard_degree=1, + tensor_parallel_degree=1, + pipeline_parallel_degree=1, + context_parallel_degree=1, + expert_parallel_degree=1, + ), + checkpoint=CheckpointManager.Config( + interval=10, + last_save_model_only=False, + ), + activation_checkpoint=None, + ) diff --git a/torchtitan/models/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py new file mode 100644 index 0000000000..f1055546b3 --- /dev/null +++ b/torchtitan/models/kimi_k3/model.py @@ -0,0 +1,797 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Kimi K3 language model components. + +The first implementation intentionally provides a device-neutral reference +path. It mirrors the released HuggingFace model math, but does not depend on +FLA, Triton, CUDA, or a device-specific extension. The reference KDA kernel is +appropriate for numerical validation and reduced-model training; an optimized +backend can be added behind :class:`KimiKDAKernel` without changing the model +or checkpoint schema. +""" + +from dataclasses import dataclass + +import torch +import torch.nn.functional as F +from torch import nn + +from torchtitan.models.common import Conv1d, Linear +from torchtitan.models.common.attention import AttentionMasksType +from torchtitan.models.common.decoder import Decoder +from torchtitan.models.common.multimodal import ( + get_vision_positions, + scatter_vision_embeds, +) +from torchtitan.models.common.nn_modules import RMSNorm +from torchtitan.models.kimi_k3.vision_encoder import KimiK3VisionEncoder +from torchtitan.models.utils import get_moe_model_nparams_and_flops +from torchtitan.protocols.module import Module, ModuleList + +# Shape suffixes: +# B = batch, L = sequence length, D = model dimension, H = heads, +# K = key head dimension, V = value head dimension, E = experts, +# T = flattened tokens, N = attention-residual entries. + + +class KimiRMSNorm(RMSNorm): + """RMSNorm with the explicit FP32 reduction used by the Kimi reference.""" + + @dataclass(kw_only=True, slots=True) + class Config(RMSNorm.Config): + pass + + def __init__(self, config: Config): + super().__init__(config) + self.kimi_eps = config.eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + input_dtype = x.dtype + x_float = x.float() + variance = x_float.pow(2).mean(dim=-1, keepdim=True) + x_float = x_float * torch.rsqrt(variance + self.kimi_eps) + assert self.weight is not None + return self.weight * x_float.to(input_dtype) + + +class KimiRMSNormGated(Module): + """Per-head RMSNorm followed by a sigmoid output gate.""" + + @dataclass(kw_only=True, slots=True) + class Config(Module.Config): + dim: int + eps: float = 1e-5 + + def __init__(self, config: Config): + super().__init__() + self.eps = config.eps + self.weight = nn.Parameter(torch.empty(config.dim)) + + def forward(self, x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: + input_dtype = x.dtype + x_float = x.float() + variance = x_float.pow(2).mean(dim=-1, keepdim=True) + x_float = x_float * torch.rsqrt(variance + self.eps) + x_float = self.weight.float() * x_float + return (x_float * torch.sigmoid(gate.float())).to(input_dtype) + + +class SituAndMul(Module): + """Kimi's SiTU activation applied to concatenated gate/up projections.""" + + @dataclass(kw_only=True, slots=True) + class Config(Module.Config): + beta: float = 1.0 + linear_beta: float | None = None + + def __init__(self, config: Config): + super().__init__() + self.beta = config.beta + self.linear_beta = config.linear_beta + + def forward(self, gate_up: torch.Tensor) -> torch.Tensor: + gate, up = gate_up.chunk(2, dim=-1) + input_dtype = gate_up.dtype + gate = gate.float() + up = up.float() + gate = self.beta * torch.tanh(gate / self.beta) * torch.sigmoid(gate) + if self.linear_beta is not None: + up = self.linear_beta * torch.tanh(up / self.linear_beta) + return (gate * up).to(input_dtype) + + +class KimiFeedForward(Module): + """Three-projection feed-forward network using SiTU.""" + + @dataclass(kw_only=True, slots=True) + class Config(Module.Config): + w1: Linear.Config + w2: Linear.Config + w3: Linear.Config + activation: SituAndMul.Config + + def __init__(self, config: Config): + super().__init__() + self.w1 = config.w1.build() + self.w2 = config.w2.build() + self.w3 = config.w3.build() + self.activation = config.activation.build() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate_up = torch.cat((self.w1(x), self.w3(x)), dim=-1) + return self.w2(self.activation(gate_up)) + + +class KimiMLAAttention(Module): + """Kimi K3 multi-head latent attention. + + Unlike DeepSeek-V3 MLA, the released K3 configuration sets + ``mla_use_nope=True``. The RoPE-sized query/key slices remain part of the + projected head, but no rotary transform is applied. + """ + + @dataclass(kw_only=True, slots=True) + class Config(Module.Config): + dim: int + num_heads: int + q_lora_rank: int + kv_lora_rank: int + qk_nope_head_dim: int + qk_rope_head_dim: int + v_head_dim: int + wq_a: Linear.Config + q_norm: KimiRMSNorm.Config + wq_b: Linear.Config + wkv_a: Linear.Config + kv_norm: KimiRMSNorm.Config + wkv_b: Linear.Config + gate: Linear.Config + wo: Linear.Config + + def __init__(self, config: Config): + super().__init__() + self.num_heads = config.num_heads + self.qk_nope_head_dim = config.qk_nope_head_dim + self.qk_rope_head_dim = config.qk_rope_head_dim + self.q_head_dim = config.qk_nope_head_dim + config.qk_rope_head_dim + self.v_head_dim = config.v_head_dim + self.kv_lora_rank = config.kv_lora_rank + self.scale = self.q_head_dim**-0.5 + + self.wq_a = config.wq_a.build() + self.q_norm = config.q_norm.build() + self.wq_b = config.wq_b.build() + self.wkv_a = config.wkv_a.build() + self.kv_norm = config.kv_norm.build() + self.wkv_b = config.wkv_b.build() + self.gate = config.gate.build() + self.wo = config.wo.build() + + def forward( + self, + x_BLD: torch.Tensor, + attention_masks: AttentionMasksType | None = None, + positions: torch.Tensor | None = None, + ) -> torch.Tensor: + del positions + if attention_masks is not None: + raise NotImplementedError( + "Kimi K3 reference MLA does not support packed-document masks." + ) + + B, L, _ = x_BLD.shape + q_BLNH = self.wq_b(self.q_norm(self.wq_a(x_BLD))).view( + B, L, self.num_heads, self.q_head_dim + ) + + compressed_kv = self.wkv_a(x_BLD) + kv_latent, k_rope = torch.split( + compressed_kv, + [self.kv_lora_rank, self.qk_rope_head_dim], + dim=-1, + ) + kv_BLNH = self.wkv_b(self.kv_norm(kv_latent)).view( + B, + L, + self.num_heads, + self.qk_nope_head_dim + self.v_head_dim, + ) + k_nope, v_BLNH = torch.split( + kv_BLNH, + [self.qk_nope_head_dim, self.v_head_dim], + dim=-1, + ) + k_rope = k_rope.view(B, L, 1, self.qk_rope_head_dim).expand( + -1, -1, self.num_heads, -1 + ) + k_BLNH = torch.cat((k_nope, k_rope), dim=-1) + + # HuggingFace eager attention keeps the matmul in the input dtype and + # performs softmax in FP32. + scores_BNLS = torch.einsum("blnh,bsnh->bnls", q_BLNH, k_BLNH) + scores_BNLS = scores_BNLS * self.scale + causal_mask = torch.ones(L, L, dtype=torch.bool, device=x_BLD.device).triu( + diagonal=1 + ) + scores_BNLS = scores_BNLS.masked_fill( + causal_mask.view(1, 1, L, L), + torch.finfo(scores_BNLS.dtype).min, + ) + probs_BNLS = torch.softmax(scores_BNLS, dim=-1, dtype=torch.float32).to( + v_BLNH.dtype + ) + out_BLNV = torch.einsum("bnls,bsnv->blnv", probs_BNLS, v_BLNH) + + out_BLD = out_BLNV.reshape(B, L, self.num_heads * self.v_head_dim) + out_BLD = out_BLD * torch.sigmoid(self.gate(x_BLD)) + return self.wo(out_BLD) + + +class KimiKDAKernel(Module): + """Differentiable, device-neutral recurrent KDA reference kernel.""" + + @dataclass(kw_only=True, slots=True) + class Config(Module.Config): + head_dim: int + lower_bound: float | None = -5.0 + + def __init__(self, config: Config): + super().__init__() + self.head_dim = config.head_dim + self.lower_bound = config.lower_bound + if self.lower_bound is not None and not (-5.0 <= self.lower_bound < 0.0): + raise ValueError("KDA lower_bound must be in the safe range [-5, 0).") + + def forward( + self, + q_BLHK: torch.Tensor, + k_BLHK: torch.Tensor, + v_BLHV: torch.Tensor, + gate_BLHK: torch.Tensor, + beta_BLH: torch.Tensor, + A_log_H: torch.Tensor, + dt_bias_HK: torch.Tensor, + ) -> torch.Tensor: + if q_BLHK.shape != k_BLHK.shape: + raise ValueError( + f"KDA q/k shapes must match, got {q_BLHK.shape} and {k_BLHK.shape}." + ) + if q_BLHK.shape[:3] != v_BLHV.shape[:3]: + raise ValueError("KDA reference backend requires equal q/k/value heads.") + + input_dtype = q_BLHK.dtype + q_BLHK = q_BLHK.float() + k_BLHK = k_BLHK.float() + q_BLHK = q_BLHK * torch.rsqrt(q_BLHK.pow(2).sum(dim=-1, keepdim=True) + 1e-6) + k_BLHK = k_BLHK * torch.rsqrt(k_BLHK.pow(2).sum(dim=-1, keepdim=True) + 1e-6) + v_BLHV = v_BLHV.float() + + if self.lower_bound is None: + log_decay_BLHK = -torch.exp(A_log_H.float()).view(1, 1, -1, 1) * F.softplus( + gate_BLHK.float() + dt_bias_HK.float() + ) + else: + log_decay_BLHK = self.lower_bound * torch.sigmoid( + torch.exp(A_log_H.float()).view(1, 1, -1, 1) + * (gate_BLHK.float() + dt_bias_HK.float()) + ) + decay_BLHK = torch.exp(log_decay_BLHK) + beta_BLH = torch.sigmoid(beta_BLH.float()) + + B, L, H, K = q_BLHK.shape + V = v_BLHV.shape[-1] + state_BHKV = torch.zeros(B, H, K, V, dtype=torch.float32, device=q_BLHK.device) + outputs = [] + scale = K**-0.5 + for token_idx in range(L): + q_BHK = q_BLHK[:, token_idx] + k_BHK = k_BLHK[:, token_idx] + v_BHV = v_BLHV[:, token_idx] + beta_BH = beta_BLH[:, token_idx] + + state_BHKV = decay_BLHK[:, token_idx].unsqueeze(-1) * state_BHKV + old_value_BHV = torch.einsum("bhk,bhkv->bhv", k_BHK, state_BHKV) + delta_BHV = (v_BHV - old_value_BHV) * beta_BH.unsqueeze(-1) + state_BHKV = state_BHKV + torch.einsum("bhk,bhv->bhkv", k_BHK, delta_BHV) + output_BHV = torch.einsum("bhk,bhkv->bhv", q_BHK, state_BHKV) + outputs.append(output_BHV * scale) + + return torch.stack(outputs, dim=1).to(input_dtype) + + +class KimiDeltaAttention(Module): + """Kimi Delta Attention with causal convolutions and reference recurrence.""" + + @dataclass(kw_only=True, slots=True) + class Config(Module.Config): + dim: int + num_heads: int + head_dim: int + conv_kernel_size: int + q_proj: Linear.Config + k_proj: Linear.Config + v_proj: Linear.Config + q_conv: Conv1d.Config + k_conv: Conv1d.Config + v_conv: Conv1d.Config + forget_a: Linear.Config + forget_b: Linear.Config + beta: Linear.Config + output_gate: Linear.Config + kernel: KimiKDAKernel.Config + output_norm: KimiRMSNormGated.Config + output_proj: Linear.Config + + def __init__(self, config: Config): + super().__init__() + self.num_heads = config.num_heads + self.head_dim = config.head_dim + self.conv_kernel_size = config.conv_kernel_size + + self.q_proj = config.q_proj.build() + self.k_proj = config.k_proj.build() + self.v_proj = config.v_proj.build() + self.q_conv = config.q_conv.build() + self.k_conv = config.k_conv.build() + self.v_conv = config.v_conv.build() + self.forget_a = config.forget_a.build() + self.forget_b = config.forget_b.build() + self.beta = config.beta.build() + self.output_gate = config.output_gate.build() + self.kernel = config.kernel.build() + self.output_norm = config.output_norm.build() + self.output_proj = config.output_proj.build() + + self.A_log = nn.Parameter(torch.empty(config.num_heads)) + self.dt_bias = nn.Parameter(torch.empty(config.num_heads, config.head_dim)) + + def _causal_conv(self, x_BLC: torch.Tensor, conv: Conv1d) -> torch.Tensor: + x_BCL = F.pad(x_BLC.transpose(1, 2), (self.conv_kernel_size - 1, 0)) + return F.silu(conv(x_BCL)).transpose(1, 2) + + def forward( + self, + x_BLD: torch.Tensor, + attention_masks: AttentionMasksType | None = None, + positions: torch.Tensor | None = None, + ) -> torch.Tensor: + del positions + if attention_masks is not None: + raise NotImplementedError( + "Kimi K3 reference KDA does not support packed-document masks." + ) + + B, L, _ = x_BLD.shape + q_BLHK = self._causal_conv(self.q_proj(x_BLD), self.q_conv).view( + B, L, self.num_heads, self.head_dim + ) + k_BLHK = self._causal_conv(self.k_proj(x_BLD), self.k_conv).view( + B, L, self.num_heads, self.head_dim + ) + v_BLHV = self._causal_conv(self.v_proj(x_BLD), self.v_conv).view( + B, L, self.num_heads, self.head_dim + ) + forget_BLHK = self.forget_b(self.forget_a(x_BLD)).view( + B, L, self.num_heads, self.head_dim + ) + beta_BLH = self.beta(x_BLD).float() + + out_BLHV = self.kernel( + q_BLHK, + k_BLHK, + v_BLHV, + forget_BLHK, + beta_BLH, + self.A_log, + self.dt_bias, + ) + output_gate_BLHV = self.output_gate(x_BLD).view( + B, L, self.num_heads, self.head_dim + ) + out_BLHV = self.output_norm(out_BLHV, output_gate_BLHV) + return self.output_proj(out_BLHV.reshape(B, L, -1)) + + +class KimiMoERouter(Module): + """Sigmoid top-k router with auxiliary-loss-free correction bias.""" + + @dataclass(kw_only=True, slots=True) + class Config(Module.Config): + num_experts: int + top_k: int + gate: Linear.Config + route_norm: bool = True + route_scale: float = 1.0 + + def __init__(self, config: Config): + super().__init__() + if not 0 < config.top_k <= config.num_experts: + raise ValueError("MoE top_k must be in [1, num_experts].") + self.num_experts = config.num_experts + self.top_k = config.top_k + self.route_norm = config.route_norm + self.route_scale = config.route_scale + self.gate = config.gate.build() + + def forward( + self, + x_BLD: torch.Tensor, + expert_bias_E: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + scores_BLE = torch.sigmoid(F.linear(x_BLD.float(), self.gate.weight.float())) + choice_BLE = ( + scores_BLE if expert_bias_E is None else scores_BLE + expert_bias_E.float() + ) + _, expert_ids_BLK = torch.topk(choice_BLE, k=self.top_k, dim=-1, sorted=False) + weights_BLK = scores_BLE.gather(-1, expert_ids_BLK) + if self.route_norm: + weights_BLK = weights_BLK / (weights_BLK.sum(dim=-1, keepdim=True) + 1e-20) + return expert_ids_BLK, weights_BLK * self.route_scale + + +class KimiLatentMoE(Module): + """Single-device trainable implementation of Kimi K3 latent MoE.""" + + @dataclass(kw_only=True, slots=True) + class Config(Module.Config): + num_experts: int + router: KimiMoERouter.Config + routed_down: Linear.Config + routed_experts: list[KimiFeedForward.Config] + routed_norm: KimiRMSNorm.Config + routed_up: Linear.Config + shared_experts: KimiFeedForward.Config + load_balance_coeff: float | None = 1e-3 + + def __init__(self, config: Config): + super().__init__() + if len(config.routed_experts) != config.num_experts: + raise ValueError( + "The number of routed expert configs must equal num_experts." + ) + self.num_experts = config.num_experts + self.router = config.router.build() + self.routed_down = config.routed_down.build() + self.routed_experts = ModuleList( + [expert.build() for expert in config.routed_experts] + ) + self.routed_norm = config.routed_norm.build() + self.routed_up = config.routed_up.build() + self.shared_experts = config.shared_experts.build() + self.load_balance_coeff = config.load_balance_coeff + if self.load_balance_coeff is not None: + if self.load_balance_coeff <= 0.0: + raise ValueError("load_balance_coeff must be positive.") + self.register_buffer( + "expert_bias_E", + torch.zeros(config.num_experts, dtype=torch.float32), + persistent=True, + ) + else: + self.expert_bias_E = None + self.register_buffer( + "tokens_per_expert_E", + torch.zeros(config.num_experts, dtype=torch.float32), + persistent=False, + ) + + def forward(self, x_BLD: torch.Tensor) -> torch.Tensor: + expert_ids_BLK, weights_BLK = self.router(x_BLD, self.expert_bias_E) + B, L, _ = x_BLD.shape + routing_map_BLE = torch.zeros( + B, + L, + self.num_experts, + dtype=torch.bool, + device=x_BLD.device, + ).scatter_(-1, expert_ids_BLK, True) + with torch.no_grad(): + self.tokens_per_expert_E.add_(routing_map_BLE.sum(dim=(0, 1)).float()) + + latent_TD = self.routed_down(x_BLD).reshape(B * L, -1) + expert_ids_TK = expert_ids_BLK.reshape(B * L, -1) + weights_TK = weights_BLK.reshape(B * L, -1) + routed_TD = torch.zeros_like(latent_TD, dtype=torch.float32) + + for expert_idx, expert in enumerate(self.routed_experts): + token_and_slot = torch.nonzero(expert_ids_TK == expert_idx, as_tuple=False) + if token_and_slot.numel() == 0: + continue + token_ids = token_and_slot[:, 0] + route_slots = token_and_slot[:, 1] + expert_output = expert(latent_TD.index_select(0, token_ids)) + route_weight = weights_TK[token_ids, route_slots].unsqueeze(-1) + routed_TD = routed_TD.index_add( + 0, token_ids, expert_output.float() * route_weight + ) + + routed_BLD = routed_TD.to(latent_TD.dtype).view(B, L, -1) + routed_BLD = self.routed_up(self.routed_norm(routed_BLD)) + return routed_BLD + self.shared_experts(x_BLD) + + def _init_self_buffers(self, *, buffer_device: torch.device | None = None) -> None: + if buffer_device is None: + buffer_device = self.tokens_per_expert_E.device + self.tokens_per_expert_E = torch.zeros( + self.num_experts, dtype=torch.float32, device=buffer_device + ) + if self.load_balance_coeff is not None: + self.expert_bias_E = torch.zeros( + self.num_experts, dtype=torch.float32, device=buffer_device + ) + + +def _apply_attention_residual( + prefix_sum_TD: torch.Tensor, + block_residual_TND: torch.Tensor, + projection: Linear, + norm: KimiRMSNorm, +) -> torch.Tensor: + """Apply Kimi's block-level attention residual in FP32.""" + + values_TND = torch.cat((block_residual_TND, prefix_sum_TD.unsqueeze(1)), dim=1) + values_float = values_TND.float() + variance = values_float.pow(2).mean(dim=-1, keepdim=True) + keys_TND = values_float * torch.rsqrt(variance + norm.kimi_eps) + score_weight_D = norm.weight.float() * projection.weight.squeeze(0).float() + scores_TN = (keys_TND * score_weight_D).sum(dim=-1) + probs_TN = torch.softmax(scores_TN, dim=-1).unsqueeze(1) + output_TD = torch.matmul(probs_TN, values_float).squeeze(1) + return output_TD.to(values_TND.dtype) + + +class KimiK3TransformerBlock(Module): + """Hybrid KDA/MLA decoder block with Kimi attention residuals.""" + + @dataclass(kw_only=True, slots=True) + class Config(Module.Config): + layer_id: int + attn_res_block_size: int + attention: KimiMLAAttention.Config | None + delta_attention: KimiDeltaAttention.Config | None + feed_forward: KimiFeedForward.Config | None + moe: KimiLatentMoE.Config | None + attention_norm: KimiRMSNorm.Config + ffn_norm: KimiRMSNorm.Config + attention_res_norm: KimiRMSNorm.Config + attention_res_proj: Linear.Config + ffn_res_norm: KimiRMSNorm.Config + ffn_res_proj: Linear.Config + + def __init__(self, config: Config): + super().__init__() + if (config.attention is None) == (config.delta_attention is None): + raise ValueError( + "Exactly one of attention or delta_attention must be configured." + ) + if (config.feed_forward is None) == (config.moe is None): + raise ValueError("Exactly one of feed_forward or moe must be configured.") + self.layer_id = config.layer_id + self.attn_res_block_size = config.attn_res_block_size + self.attention = ( + config.attention.build() if config.attention is not None else None + ) + self.delta_attention = ( + config.delta_attention.build() + if config.delta_attention is not None + else None + ) + self.feed_forward = ( + config.feed_forward.build() if config.feed_forward is not None else None + ) + self.moe = config.moe.build() if config.moe is not None else None + self.moe_enabled = self.moe is not None + self.attention_norm = config.attention_norm.build() + self.ffn_norm = config.ffn_norm.build() + self.attention_res_norm = config.attention_res_norm.build() + self.attention_res_proj = config.attention_res_proj.build() + self.ffn_res_norm = config.ffn_res_norm.build() + self.ffn_res_proj = config.ffn_res_proj.build() + + def forward( + self, + x_BLD: torch.Tensor, + block_residual_TND: torch.Tensor, + attention_masks: AttentionMasksType | None = None, + positions: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + B, L, D = x_BLD.shape + prefix_sum_BLD: torch.Tensor | None = x_BLD + + if block_residual_TND.shape[1] > 0: + assert prefix_sum_BLD is not None + x_BLD = _apply_attention_residual( + prefix_sum_BLD.reshape(-1, D), + block_residual_TND, + self.attention_res_proj, + self.attention_res_norm, + ).view(B, L, D) + + if self.layer_id % self.attn_res_block_size == 0: + assert prefix_sum_BLD is not None + block_residual_TND = torch.cat( + ( + block_residual_TND, + prefix_sum_BLD.reshape(-1, D).unsqueeze(1), + ), + dim=1, + ) + prefix_sum_BLD = None + + h_BLD = self.attention_norm(x_BLD) + if self.attention is not None: + h_BLD = self.attention(h_BLD, attention_masks, positions) + else: + assert self.delta_attention is not None + h_BLD = self.delta_attention(h_BLD, attention_masks, positions) + prefix_sum_BLD = h_BLD if prefix_sum_BLD is None else prefix_sum_BLD + h_BLD + + assert prefix_sum_BLD is not None + h_BLD = _apply_attention_residual( + prefix_sum_BLD.reshape(-1, D), + block_residual_TND, + self.ffn_res_proj, + self.ffn_res_norm, + ).view(B, L, D) + h_BLD = self.ffn_norm(h_BLD) + if self.moe is not None: + h_BLD = self.moe(h_BLD) + else: + assert self.feed_forward is not None + h_BLD = self.feed_forward(h_BLD) + return prefix_sum_BLD + h_BLD, block_residual_TND + + +class KimiK3Model(Decoder): + """Reduced Kimi K3 multimodal model used for first-version validation.""" + + @dataclass(kw_only=True, slots=True) + class Config(Decoder.Config): + layers: list[KimiK3TransformerBlock.Config] + output_res_norm: KimiRMSNorm.Config + output_res_proj: Linear.Config + vision_encoder: KimiK3VisionEncoder.Config | None = None + spatial_merge_size: int = 2 + + def update_from_config(self, *, config, **kwargs) -> None: + del kwargs + parallelism = config.parallelism + unsupported = { + "tensor parallel": parallelism.tensor_parallel_degree, + "pipeline parallel": parallelism.pipeline_parallel_degree, + "context parallel": parallelism.context_parallel_degree, + "expert parallel": parallelism.expert_parallel_degree, + } + enabled = [name for name, degree in unsupported.items() if degree > 1] + if enabled: + raise NotImplementedError( + "Kimi K3 v1 supports single-device execution only; " + f"disable {', '.join(enabled)}." + ) + dataloader = getattr(config, "dataloader", None) + if getattr(dataloader, "packing_buffer_size", 0) > 0: + raise NotImplementedError( + "Kimi K3 v1 does not support packed documents." + ) + + def get_nparams_and_flops( + self, model: nn.Module, seq_len: int + ) -> tuple[int, int]: + attention_config = self.first_attention + if not isinstance(attention_config, KimiMLAAttention.Config): + raise ValueError( + "Kimi K3 requires at least one MLA layer for FLOP accounting." + ) + return get_moe_model_nparams_and_flops( + self, + model, + attention_config.num_heads, + attention_config.qk_nope_head_dim + + attention_config.qk_rope_head_dim + + attention_config.v_head_dim, + seq_len, + ) + + def __init__(self, config: Config): + super().__init__(config) + self.output_res_norm = config.output_res_norm.build() + self.output_res_proj = config.output_res_proj.build() + self.vision_encoder = ( + config.vision_encoder.build() if config.vision_encoder is not None else None + ) + self.spatial_merge_size = config.spatial_merge_size + + def get_attention_masks(self, positions: torch.Tensor) -> AttentionMasksType | None: + del positions + return None + + def _prepare_multimodal_embeds( + self, + tokens: torch.Tensor, + *, + pixel_values: torch.Tensor | None, + grid_thw: torch.Tensor | None, + special_tokens: dict[str, int] | None, + ) -> torch.Tensor: + embeddings = self.tok_embeddings(tokens) + if (pixel_values is None) != (grid_thw is None): + raise ValueError( + "pixel_values and grid_thw must either both be provided or " + "both be omitted." + ) + if pixel_values is None: + return embeddings + assert grid_thw is not None + if self.vision_encoder is None: + raise ValueError("pixel_values were provided without a vision encoder.") + if special_tokens is None: + raise ValueError("special_tokens are required for multimodal inputs.") + + pixel_values = pixel_values.to(self.vision_encoder.patch_embed.weight.dtype) + vision_embeds = self.vision_encoder(pixel_values, grid_thw=grid_thw) + num_tokens_per_item = (grid_thw[:, 1] // self.spatial_merge_size) * ( + grid_thw[:, 2] // self.spatial_merge_size + ) + vision_positions = get_vision_positions( + tokens, + num_tokens_per_item, + special_tokens["image_id"], + ) + if not vision_positions: + raise ValueError( + "pixel_values were provided but no image placeholder tokens were found." + ) + return scatter_vision_embeds( + embeddings, + vision_embeds=vision_embeds, + vision_positions=vision_positions, + ) + + def forward( # pyrefly: ignore [bad-override] + self, + tokens: torch.Tensor, + *, + pixel_values: torch.Tensor | None = None, + grid_thw: torch.Tensor | None = None, + pixel_values_videos: torch.Tensor | None = None, + grid_thw_videos: torch.Tensor | None = None, + special_tokens: dict[str, int] | None = None, + positions: torch.Tensor | None = None, + attention_masks: AttentionMasksType | None = None, + ) -> torch.Tensor: + if pixel_values_videos is not None or grid_thw_videos is not None: + raise NotImplementedError("Kimi K3 v1 supports images but not videos.") + if self.tok_embeddings is not None: + h_BLD = self._prepare_multimodal_embeds( + tokens, + pixel_values=pixel_values, + grid_thw=grid_thw, + special_tokens=special_tokens, + ) + else: + h_BLD = tokens + + B, L, D = h_BLD.shape + block_residual_TND = h_BLD.new_zeros(B * L, 0, D) + for layer in self.layers.values(): + h_BLD, block_residual_TND = layer( + h_BLD, + block_residual_TND, + attention_masks, + positions, + ) + + h_BLD = _apply_attention_residual( + h_BLD.reshape(-1, D), + block_residual_TND, + self.output_res_proj, + self.output_res_norm, + ).view(B, L, D) + h_BLD = self.norm(h_BLD) if self.norm is not None else h_BLD + if self._skip_lm_head: + return h_BLD + return self.lm_head(h_BLD) if self.lm_head is not None else h_BLD diff --git a/torchtitan/models/kimi_k3/parallelize.py b/torchtitan/models/kimi_k3/parallelize.py new file mode 100644 index 0000000000..14a759a891 --- /dev/null +++ b/torchtitan/models/kimi_k3/parallelize.py @@ -0,0 +1,57 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Parallelization boundary for the first Kimi K3 implementation.""" + +import torch.nn as nn + +from torchtitan.config import CompileConfig, ParallelismConfig, TrainingConfig +from torchtitan.distributed import ParallelDims +from torchtitan.distributed.activation_checkpoint import ActivationCheckpointingConfig + + +def parallelize_kimi_k3( + model: nn.Module, + *, + parallel_dims: ParallelDims, + training: TrainingConfig, + parallelism: ParallelismConfig, + compile_config: CompileConfig, + ac_config: ActivationCheckpointingConfig, + dump_folder: str, +) -> nn.Module: + """Validate the v1 single-device contract and return the model unchanged.""" + del dump_folder + + enabled_parallelisms = [ + name + for name, enabled in ( + ("data parallel", parallel_dims.dp_enabled), + ("tensor parallel", parallel_dims.tp_enabled), + ("pipeline parallel", parallel_dims.pp_enabled), + ("context parallel", parallel_dims.cp_enabled), + ("expert parallel", parallel_dims.ep_enabled), + ) + if enabled + ] + if enabled_parallelisms: + raise NotImplementedError( + "Kimi K3 v1 supports single-device execution only; disable " + f"{', '.join(enabled_parallelisms)}." + ) + if parallelism.spmd_backend != "default": + raise NotImplementedError( + "Kimi K3 v1 only supports the default local tensor backend." + ) + if compile_config.enable: + raise NotImplementedError("Kimi K3 v1 does not support torch.compile.") + if ac_config is not None: + raise NotImplementedError( + "Kimi K3 v1 does not support activation checkpointing." + ) + if training.enable_cpu_offload: + raise NotImplementedError("Kimi K3 v1 does not support parameter CPU offload.") + return model diff --git a/torchtitan/models/kimi_k3/state_dict_adapter.py b/torchtitan/models/kimi_k3/state_dict_adapter.py new file mode 100644 index 0000000000..7967be1f16 --- /dev/null +++ b/torchtitan/models/kimi_k3/state_dict_adapter.py @@ -0,0 +1,337 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Unquantized HuggingFace checkpoint adapter for Kimi K3. + +The released Kimi K3 checkpoint uses MXFP4 expert weights. That format is +intentionally outside the first implementation. This adapter targets an +unquantized HuggingFace state dict, which is sufficient for constructing the +same reduced model on both sides of the numerical parity test. +""" + +import re +from typing import Any + +import torch + +from torchtitan.protocols.state_dict_adapter import StateDictAdapter + +from .model import KimiK3Model + + +_TEXT_GLOBAL_FROM_HF = { + "language_model.model.embed_tokens.weight": "tok_embeddings.weight", + "language_model.model.output_attn_res_norm.weight": ("output_res_norm.weight"), + "language_model.model.output_attn_res_proj.weight": ("output_res_proj.weight"), + "language_model.model.norm.weight": "norm.weight", + "language_model.lm_head.weight": "lm_head.weight", +} + +_TEXT_LAYER_FROM_HF = { + # Layer norms and attention residuals. + "input_layernorm.weight": "attention_norm.weight", + "post_attention_layernorm.weight": "ffn_norm.weight", + "self_attention_res_norm.weight": "attention_res_norm.weight", + "self_attention_res_proj.weight": "attention_res_proj.weight", + "mlp_res_norm.weight": "ffn_res_norm.weight", + "mlp_res_proj.weight": "ffn_res_proj.weight", + # Dense MLP. + "mlp.gate_proj.weight": "feed_forward.w1.weight", + "mlp.up_proj.weight": "feed_forward.w3.weight", + "mlp.down_proj.weight": "feed_forward.w2.weight", +} + +_MLA_FROM_HF = { + "self_attn.q_a_proj.weight": "attention.wq_a.weight", + "self_attn.q_a_layernorm.weight": "attention.q_norm.weight", + "self_attn.q_b_proj.weight": "attention.wq_b.weight", + "self_attn.kv_a_proj_with_mqa.weight": "attention.wkv_a.weight", + "self_attn.kv_a_layernorm.weight": "attention.kv_norm.weight", + "self_attn.kv_b_proj.weight": "attention.wkv_b.weight", + "self_attn.g_proj.weight": "attention.gate.weight", + "self_attn.o_proj.weight": "attention.wo.weight", +} + +_KDA_FROM_HF = { + "self_attn.q_proj.weight": "delta_attention.q_proj.weight", + "self_attn.k_proj.weight": "delta_attention.k_proj.weight", + "self_attn.v_proj.weight": "delta_attention.v_proj.weight", + "self_attn.q_conv1d.weight": "delta_attention.q_conv.weight", + "self_attn.k_conv1d.weight": "delta_attention.k_conv.weight", + "self_attn.v_conv1d.weight": "delta_attention.v_conv.weight", + "self_attn.f_a_proj.weight": "delta_attention.forget_a.weight", + "self_attn.f_b_proj.weight": "delta_attention.forget_b.weight", + "self_attn.b_proj.weight": "delta_attention.beta.weight", + "self_attn.g_proj.weight": "delta_attention.output_gate.weight", + "self_attn.o_norm.weight": "delta_attention.output_norm.weight", + "self_attn.o_proj.weight": "delta_attention.output_proj.weight", + "self_attn.A_log": "delta_attention.A_log", + "self_attn.dt_bias": "delta_attention.dt_bias", +} + +_MOE_FROM_HF = { + "block_sparse_moe.gate.weight": "moe.router.gate.weight", + "block_sparse_moe.gate.e_score_correction_bias": "moe.expert_bias_E", + "block_sparse_moe.routed_expert_down_proj.weight": ("moe.routed_down.weight"), + "block_sparse_moe.routed_expert_up_proj.weight": ("moe.routed_up.weight"), + "block_sparse_moe.routed_expert_norm.weight": ("moe.routed_norm.weight"), + "block_sparse_moe.shared_experts.gate_proj.weight": ( + "moe.shared_experts.w1.weight" + ), + "block_sparse_moe.shared_experts.up_proj.weight": ("moe.shared_experts.w3.weight"), + "block_sparse_moe.shared_experts.down_proj.weight": ( + "moe.shared_experts.w2.weight" + ), +} + +_VISION_GLOBAL_FROM_HF = { + "vision_tower.patch_embed.proj.weight": ("vision_encoder.patch_embed.weight"), + "vision_tower.patch_embed.pos_emb.weight": "vision_encoder.pos_embed", + "vision_tower.encoder.final_layernorm.weight": ("vision_encoder.final_norm.weight"), + "mm_projector.proj.0.weight": ("vision_encoder.projector.linear_1.weight"), + "mm_projector.proj.2.weight": ("vision_encoder.projector.linear_2.weight"), + "mm_projector.post_norm.weight": ("vision_encoder.projector.post_norm.weight"), +} + +_VISION_LAYER_FROM_HF = { + "norm0.weight": "norm1.weight", + "norm1.weight": "norm2.weight", + "wo.weight": "attn.proj.weight", + "mlp.fc0.weight": "mlp.linear_fc1.weight", + "mlp.fc1.weight": "mlp.linear_fc2.weight", +} + + +class KimiK3StateDictAdapter(StateDictAdapter): + """Convert between unquantized Kimi K3 HF and TorchTitan state dicts.""" + + def __init__( + self, + model_config: KimiK3Model.Config, + hf_assets_path: str | None, + ): + super().__init__(model_config, hf_assets_path) + self.kimi_config = model_config + + @staticmethod + def _raise_if_quantized_key(key: str) -> None: + quantized_markers = ( + "weight_scale", + "weight_packed", + "compressed", + "scale_shape", + ) + if any(marker in key for marker in quantized_markers): + raise NotImplementedError( + "Kimi K3 v1 only supports unquantized HuggingFace state " + f"dicts; encountered quantized key '{key}'." + ) + + def from_hf(self, hf_state_dict: dict[str, Any]) -> dict[str, Any]: + """Convert an unquantized HuggingFace state dict to TorchTitan.""" + state_dict: dict[str, Any] = {} + unmapped: list[str] = [] + + for hf_key, value in hf_state_dict.items(): + self._raise_if_quantized_key(hf_key) + if hf_key.endswith("rotary_emb.inv_freq"): + continue + + tt_key = _TEXT_GLOBAL_FROM_HF.get(hf_key) + if tt_key is not None: + state_dict[tt_key] = value + continue + + tt_key = _VISION_GLOBAL_FROM_HF.get(hf_key) + if tt_key is not None: + if hf_key == "vision_tower.patch_embed.proj.weight": + value = value.reshape(value.shape[0], -1) + state_dict[tt_key] = value + continue + + text_match = re.fullmatch( + r"language_model\.model\.layers\.(\d+)\.(.+)", + hf_key, + ) + if text_match is not None: + layer_idx, suffix = text_match.groups() + expert_match = re.fullmatch( + r"block_sparse_moe\.experts\.(\d+)\." r"(w1|w2|w3)\.weight", + suffix, + ) + if expert_match is not None: + expert_idx, projection = expert_match.groups() + state_dict[ + f"layers.{layer_idx}.moe.routed_experts." + f"{expert_idx}.{projection}.weight" + ] = value + continue + + layer_config = self.kimi_config.layers[int(layer_idx)] + mapped_suffix = _TEXT_LAYER_FROM_HF.get(suffix) + if mapped_suffix is None: + attention_map = ( + _MLA_FROM_HF + if layer_config.attention is not None + else _KDA_FROM_HF + ) + mapped_suffix = attention_map.get(suffix) + if mapped_suffix is None: + mapped_suffix = _MOE_FROM_HF.get(suffix) + if mapped_suffix is None: + unmapped.append(hf_key) + continue + if suffix == "self_attn.dt_bias": + delta_config = layer_config.delta_attention + if delta_config is None: + raise ValueError(f"HF key '{hf_key}' targets a non-KDA layer.") + value = value.reshape( + delta_config.num_heads, + delta_config.head_dim, + ) + state_dict[f"layers.{layer_idx}.{mapped_suffix}"] = value + continue + + vision_match = re.fullmatch( + r"vision_tower\.encoder\.blocks\.(\d+)\.(.+)", + hf_key, + ) + if vision_match is not None: + layer_idx, suffix = vision_match.groups() + if suffix == "wqkv.weight": + q, k, v = torch.chunk(value, 3, dim=0) + base = f"vision_encoder.layers.{layer_idx}.attn" + state_dict[f"{base}.wq.weight"] = q + state_dict[f"{base}.wk.weight"] = k + state_dict[f"{base}.wv.weight"] = v + continue + mapped_suffix = _VISION_LAYER_FROM_HF.get(suffix) + if mapped_suffix is None: + unmapped.append(hf_key) + continue + state_dict[f"vision_encoder.layers.{layer_idx}.{mapped_suffix}"] = value + continue + + unmapped.append(hf_key) + + if unmapped: + raise ValueError( + "KimiK3StateDictAdapter found HuggingFace keys without a " + f"mapping: {unmapped}." + ) + return state_dict + + def to_hf(self, state_dict: dict[str, Any]) -> dict[str, Any]: + """Convert a TorchTitan state dict to unquantized HuggingFace format.""" + text_global_to_hf = {value: key for key, value in _TEXT_GLOBAL_FROM_HF.items()} + vision_global_to_hf = { + value: key for key, value in _VISION_GLOBAL_FROM_HF.items() + } + text_layer_to_hf = { + value: key + for mapping in ( + _TEXT_LAYER_FROM_HF, + _MLA_FROM_HF, + _KDA_FROM_HF, + _MOE_FROM_HF, + ) + for key, value in mapping.items() + } + vision_layer_to_hf = { + value: key for key, value in _VISION_LAYER_FROM_HF.items() + } + + hf_state_dict: dict[str, Any] = {} + vision_qkv: dict[str, dict[str, Any]] = {} + unmapped: list[str] = [] + + for tt_key, value in state_dict.items(): + hf_key = text_global_to_hf.get(tt_key) + if hf_key is not None: + hf_state_dict[hf_key] = value + continue + + hf_key = vision_global_to_hf.get(tt_key) + if hf_key is not None: + if tt_key == "vision_encoder.patch_embed.weight": + vision_config = self.kimi_config.vision_encoder + if vision_config is None: + raise ValueError( + "Vision state was provided for a text-only config." + ) + value = value.reshape( + value.shape[0], + vision_config.in_channels, + vision_config.patch_size, + vision_config.patch_size, + ) + hf_state_dict[hf_key] = value + continue + + text_match = re.fullmatch(r"layers\.(\d+)\.(.+)", tt_key) + if text_match is not None: + layer_idx, suffix = text_match.groups() + expert_match = re.fullmatch( + r"moe\.routed_experts\.(\d+)\." r"(w1|w2|w3)\.weight", + suffix, + ) + if expert_match is not None: + expert_idx, projection = expert_match.groups() + hf_state_dict[ + f"language_model.model.layers.{layer_idx}." + f"block_sparse_moe.experts.{expert_idx}." + f"{projection}.weight" + ] = value + continue + + mapped_suffix = text_layer_to_hf.get(suffix) + if mapped_suffix is None: + unmapped.append(tt_key) + continue + if suffix == "delta_attention.dt_bias": + value = value.reshape(-1) + hf_state_dict[ + f"language_model.model.layers.{layer_idx}.{mapped_suffix}" + ] = value + continue + + vision_match = re.fullmatch( + r"vision_encoder\.layers\.(\d+)\.(.+)", + tt_key, + ) + if vision_match is not None: + layer_idx, suffix = vision_match.groups() + qkv_match = re.fullmatch(r"attn\.w(q|k|v)\.weight", suffix) + if qkv_match is not None: + vision_qkv.setdefault(layer_idx, {})[qkv_match.group(1)] = value + continue + mapped_suffix = vision_layer_to_hf.get(suffix) + if mapped_suffix is None: + unmapped.append(tt_key) + continue + hf_state_dict[ + f"vision_tower.encoder.blocks.{layer_idx}.{mapped_suffix}" + ] = value + continue + + unmapped.append(tt_key) + + for layer_idx, qkv in vision_qkv.items(): + missing = {"q", "k", "v"} - qkv.keys() + if missing: + raise ValueError( + f"Vision layer {layer_idx} is missing QKV parts: {sorted(missing)}." + ) + hf_state_dict[ + f"vision_tower.encoder.blocks.{layer_idx}.wqkv.weight" + ] = torch.cat((qkv["q"], qkv["k"], qkv["v"]), dim=0) + + if unmapped: + raise ValueError( + "KimiK3StateDictAdapter found TorchTitan keys without a " + f"mapping: {unmapped}." + ) + return hf_state_dict diff --git a/torchtitan/models/kimi_k3/vision_encoder.py b/torchtitan/models/kimi_k3/vision_encoder.py new file mode 100644 index 0000000000..1cb81f86a4 --- /dev/null +++ b/torchtitan/models/kimi_k3/vision_encoder.py @@ -0,0 +1,488 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""MoonViT3d vision encoder used by Kimi K3. + +This module keeps the first Kimi K3 implementation device-neutral. Vision +attention is an eager PyTorch reference over each visual item, which preserves +the block-diagonal attention semantics of the HuggingFace implementation +without requiring FlashAttention or a device-specific kernel. + +Shape suffixes: +- N = number of visual items +- P = maximum patches per item (padded) +- D = vision hidden dimension +- H = number of attention heads +- K = attention head dimension +- M = maximum merged tokens per item (padded) +""" + +import math +from dataclasses import dataclass, field + +import torch +import torch.nn.functional as F + +from torchtitan.models.common import Linear +from torchtitan.models.common.nn_modules import RMSNorm +from torchtitan.models.common.vision_encoder import VisionMLP +from torchtitan.protocols.module import Module, ModuleDict + + +class KimiExactGELU(Module): + """Exact GELU used by the released Kimi vision projector. + + The explicit FP32 form is mathematically equivalent to + ``nn.GELU(approximate="none")`` while avoiding device-specific fused + approximations in the numerical reference path. + """ + + @dataclass(kw_only=True, slots=True) + class Config(Module.Config): + pass + + def __init__(self, config: Config): + super().__init__() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + input_dtype = x.dtype + x_float = x.float() + return (0.5 * x_float * (1.0 + torch.erf(x_float / math.sqrt(2.0)))).to( + input_dtype + ) + + +def _get_temporal_pos_embed( + num_frames: int, + embed_dim: int, + *, + device: torch.device, +) -> torch.Tensor: + """Return fixed 1D sinusoidal embeddings for video frame positions.""" + grid = torch.arange(num_frames, dtype=torch.float32, device=device) + omega = torch.arange(embed_dim // 2, dtype=torch.float32, device=device) / ( + embed_dim / 2.0 + ) + omega = 1.0 / 10000.0**omega + angles = torch.outer(grid, omega) + return torch.cat((angles.sin(), angles.cos()), dim=-1) + + +def _compute_learned_pos_embeds( + pos_embed: torch.Tensor, + grids: list[list[int]], + max_num_patches: int, + interpolation_mode: str, + max_num_frames: int, +) -> torch.Tensor: + """Interpolate the learned 2D table and add fixed temporal embeddings.""" + height, width, dim = pos_embed.shape + result = pos_embed.new_zeros(len(grids), max_num_patches, dim) + pos_grid = pos_embed.permute(2, 0, 1).unsqueeze(0).float() + + cached_spatial: dict[tuple[int, int], torch.Tensor] = {} + for item_idx, (num_frames, grid_h, grid_w) in enumerate(grids): + if num_frames > max_num_frames: + raise ValueError( + f"Vision grid has {num_frames} frames, exceeding " + f"max_num_frames={max_num_frames}." + ) + spatial = cached_spatial.get((grid_h, grid_w)) + if spatial is None: + if (grid_h, grid_w) == (height, width): + spatial = pos_embed.flatten(end_dim=1) + else: + spatial = ( + F.interpolate( + pos_grid, + size=(grid_h, grid_w), + mode=interpolation_mode, + ) + .squeeze(0) + .permute(1, 2, 0) + .reshape(grid_h * grid_w, dim) + .to(pos_embed.dtype) + ) + cached_spatial[(grid_h, grid_w)] = spatial + + if num_frames == 1: + item_pos = spatial + else: + temporal = _get_temporal_pos_embed(num_frames, dim, device=pos_embed.device) + item_pos = spatial.unsqueeze(0) + temporal.unsqueeze(1).to(spatial.dtype) + item_pos = item_pos.reshape(num_frames * grid_h * grid_w, dim) + result[item_idx, : item_pos.shape[0]] = item_pos + + return result + + +def _compute_2d_rope_cache( + freq_table: torch.Tensor, + grids: list[list[int]], + max_num_patches: int, + head_dim: int, +) -> torch.Tensor: + """Build the real-valued 2D RoPE cache in raster patch order.""" + angles = torch.zeros( + len(grids), + max_num_patches, + head_dim // 2, + dtype=freq_table.dtype, + device=freq_table.device, + ) + cached_spatial: dict[tuple[int, int], torch.Tensor] = {} + for item_idx, (num_frames, grid_h, grid_w) in enumerate(grids): + spatial = cached_spatial.get((grid_h, grid_w)) + if spatial is None: + flat = torch.arange(grid_h * grid_w, device=freq_table.device) + x_angles = freq_table[flat % grid_w] + y_angles = freq_table[flat // grid_w] + spatial = torch.stack((x_angles, y_angles), dim=-1).reshape( + grid_h * grid_w, head_dim // 2 + ) + cached_spatial[(grid_h, grid_w)] = spatial + item_angles = spatial.repeat(num_frames, 1) + angles[item_idx, : item_angles.shape[0]] = item_angles + + return torch.stack((angles.cos(), angles.sin()), dim=-1).unsqueeze(2) + + +def _apply_2d_rope( + q_NPHK: torch.Tensor, + k_NPHK: torch.Tensor, + rope_cache_NP1C2: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Apply 2D RoPE using the real form of complex multiplication.""" + + cos_NP1C = rope_cache_NP1C2[..., 0] + sin_NP1C = rope_cache_NP1C2[..., 1] + + def rotate(x_NPHK: torch.Tensor) -> torch.Tensor: + x_NPHC2 = x_NPHK.float().reshape(*x_NPHK.shape[:-1], -1, 2) + real_NPHC = x_NPHC2[..., 0] * cos_NP1C - x_NPHC2[..., 1] * sin_NP1C + imag_NPHC = x_NPHC2[..., 0] * sin_NP1C + x_NPHC2[..., 1] * cos_NP1C + return torch.stack((real_NPHC, imag_NPHC), dim=-1).flatten(-2) + + return rotate(q_NPHK).to(q_NPHK.dtype), rotate(k_NPHK).to(k_NPHK.dtype) + + +def _temporal_pool_and_merge( + hidden_NPD: torch.Tensor, + grids: list[list[int]], + merge_kernel_size: tuple[int, int], +) -> torch.Tensor: + """Temporally pool and concatenate neighboring spatial patch features.""" + num_items, _, dim = hidden_NPD.shape + kernel_h, kernel_w = merge_kernel_size + merged_dim = kernel_h * kernel_w * dim + max_merged = max( + (grid_h // kernel_h) * (grid_w // kernel_w) for _, grid_h, grid_w in grids + ) + merged_NMK = hidden_NPD.new_zeros(num_items, max_merged, merged_dim) + + for item_idx, (num_frames, grid_h, grid_w) in enumerate(grids): + merged_h = grid_h // kernel_h + merged_w = grid_w // kernel_w + item = hidden_NPD[item_idx, : num_frames * grid_h * grid_w].view( + num_frames, + merged_h, + kernel_h, + merged_w, + kernel_w, + dim, + ) + item = item.permute(0, 1, 3, 2, 4, 5).mean(dim=0) + item = item.reshape(merged_h * merged_w, merged_dim) + merged_NMK[item_idx, : item.shape[0]] = item + + return merged_NMK + + +class VisionRotaryEmbedding2D(Module): + """Per-axis frequency table for MoonViT's interleaved 2D RoPE.""" + + @dataclass(kw_only=True, slots=True) + class Config(Module.Config): + head_dim: int + theta: float = 10000.0 + + def __init__(self, config: Config): + super().__init__() + if config.head_dim % 4 != 0: + raise ValueError( + "Vision 2D RoPE head_dim must be divisible by 4, " + f"got {config.head_dim}." + ) + self.head_dim = config.head_dim + self.theta = config.theta + self.register_buffer("inv_freq", self._compute_inv_freq(), persistent=False) + + def _compute_inv_freq(self, *, device: torch.device | None = None) -> torch.Tensor: + return 1.0 / ( + self.theta + ** ( + torch.arange( + 0, + self.head_dim, + 4, + dtype=torch.float32, + device=device, + ) + / self.head_dim + ) + ) + + def _init_self_buffers(self, *, buffer_device: torch.device | None = None) -> None: + device = buffer_device or self.inv_freq.device + self.inv_freq = self._compute_inv_freq(device=device) + + def forward(self, seqlen: int) -> torch.Tensor: + positions = torch.arange( + seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype + ) + return torch.outer(positions, self.inv_freq) + + +class KimiK3VisionAttention(Module): + """Eager, block-diagonal MoonViT attention reference.""" + + @dataclass(kw_only=True, slots=True) + class Config(Module.Config): + qkv_dim: int + num_heads: int + wq: Linear.Config + wk: Linear.Config + wv: Linear.Config + proj: Linear.Config + + def __init__(self, config: Config): + super().__init__() + if config.qkv_dim % config.num_heads != 0: + raise ValueError( + f"qkv_dim ({config.qkv_dim}) must be divisible by " + f"num_heads ({config.num_heads})." + ) + self.num_heads = config.num_heads + self.head_dim = config.qkv_dim // config.num_heads + self.scale = self.head_dim**-0.5 + self.wq = config.wq.build() + self.wk = config.wk.build() + self.wv = config.wv.build() + self.proj = config.proj.build() + + def forward( + self, + x_NPD: torch.Tensor, + *, + rope_cache: torch.Tensor, + num_patches: list[int], + ) -> torch.Tensor: + num_items, max_num_patches, _ = x_NPD.shape + q_NPHK = self.wq(x_NPD).view( + num_items, max_num_patches, self.num_heads, self.head_dim + ) + k_NPHK = self.wk(x_NPD).view( + num_items, max_num_patches, self.num_heads, self.head_dim + ) + v_NPHK = self.wv(x_NPD).view( + num_items, max_num_patches, self.num_heads, self.head_dim + ) + q_NPHK, k_NPHK = _apply_2d_rope(q_NPHK, k_NPHK, rope_cache) + + output_NPHK = torch.zeros_like(v_NPHK) + for item_idx, item_length in enumerate(num_patches): + q_HPK = q_NPHK[item_idx, :item_length].transpose(0, 1) + k_HPK = k_NPHK[item_idx, :item_length].transpose(0, 1) + v_HPK = v_NPHK[item_idx, :item_length].transpose(0, 1) + scores_HPP = torch.matmul(q_HPK, k_HPK.transpose(-2, -1)) + scores_HPP = scores_HPP * self.scale + probs_HPP = torch.softmax(scores_HPP, dim=-1, dtype=torch.float32).to( + q_HPK.dtype + ) + output_PHK = torch.matmul(probs_HPP, v_HPK).transpose(0, 1) + output_NPHK[item_idx, :item_length] = output_PHK + + return self.proj(output_NPHK.flatten(start_dim=-2)) + + +class KimiK3VisionBlock(Module): + """MoonViT pre-norm attention and MLP block.""" + + @dataclass(kw_only=True, slots=True) + class Config(Module.Config): + norm1: RMSNorm.Config + norm2: RMSNorm.Config + attn: KimiK3VisionAttention.Config + mlp: VisionMLP.Config + + def __init__(self, config: Config): + super().__init__() + self.norm1 = config.norm1.build() + self.norm2 = config.norm2.build() + self.attn = config.attn.build() + self.mlp = config.mlp.build() + + def forward( + self, + x_NPD: torch.Tensor, + *, + rope_cache: torch.Tensor, + num_patches: list[int], + ) -> torch.Tensor: + x_NPD = x_NPD + self.attn( + self.norm1(x_NPD), + rope_cache=rope_cache, + num_patches=num_patches, + ) + return x_NPD + self.mlp(self.norm2(x_NPD)) + + +class KimiK3VisionProjector(Module): + """PatchMergerMLPV2 projector from merged vision features to text width.""" + + @dataclass(kw_only=True, slots=True) + class Config(Module.Config): + merged_dim: int + linear_1: Linear.Config + linear_2: Linear.Config + post_norm: RMSNorm.Config + activation: KimiExactGELU.Config = field(default_factory=KimiExactGELU.Config) + + def __init__(self, config: Config): + super().__init__() + self.merged_dim = config.merged_dim + self.linear_1 = config.linear_1.build() + self.linear_2 = config.linear_2.build() + self.post_norm = config.post_norm.build() + self.activation = config.activation.build() + + def forward(self, merged_NMK: torch.Tensor) -> torch.Tensor: + if merged_NMK.shape[-1] != self.merged_dim: + raise ValueError( + f"Expected merged vision dim {self.merged_dim}, got " + f"{merged_NMK.shape[-1]}." + ) + projected = self.linear_2(self.activation(self.linear_1(merged_NMK))) + return self.post_norm(projected) + + +class KimiK3VisionEncoder(Module): + """Device-neutral MoonViT3d encoder and PatchMergerMLPV2 projector.""" + + @dataclass(kw_only=True, slots=True) + class Config(Module.Config): + dim: int + num_layers: int + patch_size: int + in_channels: int + merge_kernel_size: tuple[int, int] + init_pos_emb_height: int + init_pos_emb_width: int + max_num_frames: int + interpolation_mode: str + patch_embed_proj: Linear.Config + rotary_pos_emb: VisionRotaryEmbedding2D.Config + block: KimiK3VisionBlock.Config + final_norm: RMSNorm.Config + projector: KimiK3VisionProjector.Config + + def __init__(self, config: Config): + super().__init__() + self.dim = config.dim + self.patch_size = config.patch_size + self.in_channels = config.in_channels + self.merge_kernel_size = config.merge_kernel_size + self.max_num_frames = config.max_num_frames + self.interpolation_mode = config.interpolation_mode + self.patch_embed = config.patch_embed_proj.build() + self.pos_embed = torch.nn.Parameter( + torch.empty( + config.init_pos_emb_height, + config.init_pos_emb_width, + config.dim, + ) + ) + self.rotary_pos_emb = config.rotary_pos_emb.build() + self._cached_freq_table: torch.Tensor | None = None + self.layers = ModuleDict( + { + str(layer_idx): config.block.build() + for layer_idx in range(config.num_layers) + } + ) + self.final_norm = config.final_norm.build() + self.projector = config.projector.build() + + def _compute_position_embeddings( + self, grids: list[list[int]], max_num_patches: int + ) -> tuple[torch.Tensor, torch.Tensor]: + max_grid_side = max(max(grid_h, grid_w) for _, grid_h, grid_w in grids) + if ( + self._cached_freq_table is None + or self._cached_freq_table.shape[0] < max_grid_side + ): + self._cached_freq_table = self.rotary_pos_emb(max_grid_side) + learned_pos = _compute_learned_pos_embeds( + self.pos_embed, + grids, + max_num_patches, + self.interpolation_mode, + self.max_num_frames, + ) + rope_cache = _compute_2d_rope_cache( + self._cached_freq_table, + grids, + max_num_patches, + self.rotary_pos_emb.head_dim, + ) + return learned_pos, rope_cache + + def forward( + self, + pixel_values: torch.Tensor, + *, + grid_thw: torch.Tensor, + ) -> torch.Tensor: + """Encode padded raster-order patches and return padded text features.""" + if grid_thw.ndim != 2 or grid_thw.shape[1] != 3: + raise ValueError(f"grid_thw must have shape (N, 3), got {grid_thw.shape}.") + num_items, max_num_patches, _ = pixel_values.shape + grids = grid_thw.tolist() + if len(grids) != num_items: + raise ValueError( + f"pixel_values contains {num_items} items but grid_thw " + f"contains {len(grids)}." + ) + + kernel_h, kernel_w = self.merge_kernel_size + num_patches = [] + for num_frames, grid_h, grid_w in grids: + if grid_h % kernel_h != 0 or grid_w % kernel_w != 0: + raise ValueError( + f"Vision grid {grid_h}x{grid_w} is not divisible by " + f"merge kernel {self.merge_kernel_size}." + ) + item_num_patches = num_frames * grid_h * grid_w + if item_num_patches > max_num_patches: + raise ValueError( + f"Vision grid requires {item_num_patches} patches, but " + f"pixel_values only provides {max_num_patches}." + ) + num_patches.append(item_num_patches) + + learned_pos, rope_cache = self._compute_position_embeddings( + grids, max_num_patches + ) + hidden_NPD = self.patch_embed(pixel_values) + learned_pos + for block in self.layers.values(): + hidden_NPD = block( + hidden_NPD, + rope_cache=rope_cache, + num_patches=num_patches, + ) + hidden_NPD = self.final_norm(hidden_NPD) + merged_NMK = _temporal_pool_and_merge(hidden_NPD, grids, self.merge_kernel_size) + return self.projector(merged_NMK) From 8de112480445612e4db8fa83abb1f7b3d499f749 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Thu, 30 Jul 2026 16:07:13 +0800 Subject: [PATCH 02/67] Remove development-only Kimi K3 diagnostics --- .../numerical_tests_kimi_k3.py | 704 ------------------ .../numerical_tests_kimi_k3_device.py | 435 ----------- torchtitan/models/kimi_k3/README.md | 30 +- 3 files changed, 2 insertions(+), 1167 deletions(-) delete mode 100644 scripts/checkpoint_conversion/numerical_tests_kimi_k3.py delete mode 100644 scripts/checkpoint_conversion/numerical_tests_kimi_k3_device.py diff --git a/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py b/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py deleted file mode 100644 index dd0cb585ff..0000000000 --- a/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py +++ /dev/null @@ -1,704 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -"""Numerical parity for the reduced TorchTitan and released HF Kimi K3 models. - -The released checkpoint contains MXFP4 routed-expert weights, so this test -constructs the same reduced, unquantized topology on both sides. TorchTitan -initializes the weights once, ``KimiK3StateDictAdapter`` converts them to the -released HuggingFace schema, and the HuggingFace model loads them strictly. - -The comparison covers: - -- text-only decoder layers, KDA/MLA outputs, router expert IDs, and logits; -- vision transformer blocks and projected vision features; -- end-to-end image+text decoder layers and logits. - -The released HuggingFace implementation imports FLA for KDA. By default, run -this script in an environment that satisfies the released model requirements -and has a CUDA GPU. The script requests the released eager MLA and -vision-attention paths after construction so the comparison isolates model math -from FlashAttention kernels. - -For a CPU-only comparison, ``--hf_kda_backend reference`` installs the minimum -released FLA API with pure PyTorch operators. - -Example: - - CUDA_VISIBLE_DEVICES=0 python -m \ - scripts.checkpoint_conversion.numerical_tests_kimi_k3 \ - --hf_repo_path ~/hf_assets/moonshotai/Kimi-K3 - - python -m scripts.checkpoint_conversion.numerical_tests_kimi_k3 \ - --hf_repo_path ~/hf_assets/moonshotai/Kimi-K3 \ - --device cpu \ - --hf_kda_backend reference -""" - -import argparse -import importlib -import sys -import types -from collections.abc import Callable -from typing import cast - -import torch -import torch.nn.functional as F -from torch import nn - -from torchtitan.models.kimi_k3 import model_registry -from torchtitan.models.kimi_k3.model import KimiK3Model -from torchtitan.models.kimi_k3.state_dict_adapter import KimiK3StateDictAdapter -from transformers import AutoConfig, AutoModelForCausalLM - - -_IMAGE_TOKEN_ID = 7 - - -class _ReferenceShortConvolution(nn.Module): - """Pure PyTorch compatibility layer for the released HF KDA module.""" - - def __init__( - self, - hidden_size: int, - kernel_size: int, - bias: bool = False, - activation: str = "silu", - **kwargs, - ): - super().__init__() - del kwargs - self.weight = nn.Parameter(torch.empty(hidden_size, 1, kernel_size)) - self.bias = nn.Parameter(torch.empty(hidden_size)) if bias else None - self.kernel_size = kernel_size - self.activation = activation - - def forward( - self, - x: torch.Tensor, - cache: torch.Tensor | None = None, - output_final_state: bool = False, - cu_seqlens: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor | None]: - del cache, cu_seqlens - output = F.conv1d( - F.pad(x.transpose(1, 2), (self.kernel_size - 1, 0)), - self.weight, - self.bias, - groups=self.weight.shape[0], - ).transpose(1, 2) - if self.activation == "silu": - output = F.silu(output) - final_state = x[:, -self.kernel_size + 1 :] if output_final_state else None - return output, final_state - - -class _ReferenceRMSNormGated(nn.Module): - """Pure PyTorch equivalent of FLA's fused gated RMSNorm.""" - - def __init__( - self, - hidden_size: int, - eps: float = 1e-5, - activation: str = "sigmoid", - ): - super().__init__() - if activation != "sigmoid": - raise ValueError( - "The Kimi K3 reference RMSNorm only supports a sigmoid gate." - ) - self.weight = nn.Parameter(torch.ones(hidden_size)) - self.eps = eps - - def forward(self, x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: - input_dtype = x.dtype - x_float = x.float() - x_float = x_float * torch.rsqrt( - x_float.square().mean(dim=-1, keepdim=True) + self.eps - ) - return (x_float * self.weight.float() * torch.sigmoid(gate.float())).to( - input_dtype - ) - - -def _reference_kda( - *, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - A_log: torch.Tensor, - dt_bias: torch.Tensor, - initial_state: torch.Tensor | None, - output_final_state: bool, - lower_bound: float, - **kwargs, -) -> tuple[torch.Tensor, torch.Tensor | None]: - """Pure PyTorch equivalent of the FLA KDA inference API.""" - del kwargs - input_dtype = v.dtype - q = q.float() - k = k.float() - v = v.float() - q = q * torch.rsqrt(q.square().sum(dim=-1, keepdim=True) + 1e-6) - k = k * torch.rsqrt(k.square().sum(dim=-1, keepdim=True) + 1e-6) - log_decay = lower_bound * torch.sigmoid( - torch.exp(A_log.float()).view(1, 1, -1, 1) - * (g.float() + dt_bias.float().view(1, 1, *g.shape[-2:])) - ) - beta = torch.sigmoid(beta.float()) - - batch_size, seq_len, num_heads, head_dim = q.shape - value_dim = v.shape[-1] - state = ( - torch.zeros( - batch_size, - num_heads, - head_dim, - value_dim, - dtype=torch.float32, - device=q.device, - ) - if initial_state is None - else initial_state.float() - ) - outputs = [] - for token_idx in range(seq_len): - state = state * torch.exp(log_decay[:, token_idx]).unsqueeze(-1) - old_value = torch.matmul( - k[:, token_idx].unsqueeze(-2), - state, - ).squeeze(-2) - delta = (v[:, token_idx] - old_value) * beta[:, token_idx].unsqueeze(-1) - state = state + k[:, token_idx].unsqueeze(-1) * delta.unsqueeze(-2) - outputs.append( - torch.matmul(q[:, token_idx].unsqueeze(-2), state).squeeze(-2) - * (head_dim**-0.5) - ) - output = torch.stack(outputs, dim=1).to(input_dtype) - return output, state if output_final_state else None - - -def _install_reference_fla() -> None: - """Expose the minimum FLA API used by the released HF model.""" - - # The released vision reference decorates an interpolation helper with - # torch.compile. Keep the compatibility backend fully eager so it remains - # usable for the CPU-only comparison. - def eager_compile(model=None, *args, **kwargs): - del args, kwargs - return model if model is not None else lambda wrapped: wrapped - - torch.__dict__["compile"] = eager_compile - - fla = types.ModuleType("fla") - fla_modules = types.ModuleType("fla.modules") - fla_modules.__dict__.update( - { - "ShortConvolution": _ReferenceShortConvolution, - "FusedRMSNormGated": _ReferenceRMSNormGated, - } - ) - fla_ops = types.ModuleType("fla.ops") - fla_kda = types.ModuleType("fla.ops.kda") - fla_kda.__dict__.update( - { - "chunk_kda": _reference_kda, - "fused_recurrent_kda": _reference_kda, - } - ) - fla_ops_utils = types.ModuleType("fla.ops.utils") - fla_ops_index = types.ModuleType("fla.ops.utils.index") - fla_ops_index.__dict__.update( - { - "prepare_cu_seqlens_from_mask": lambda _mask: None, - "prepare_lens_from_mask": lambda _mask: None, - } - ) - fla_utils = types.ModuleType("fla.utils") - fla_utils.__dict__["tensor_cache"] = lambda fn: fn - - for module in ( - fla, - fla_modules, - fla_ops, - fla_kda, - fla_ops_utils, - fla_ops_index, - fla_utils, - ): - sys.modules[module.__name__] = module - - -def _build_reduced_hf_config(hf_repo_path: str, tt_config): - """Build the released HF config with TorchTitan's reduced dimensions.""" - released_config = AutoConfig.from_pretrained( - hf_repo_path, - trust_remote_code=True, - ) - text_config_cls = type(released_config.text_config) - vision_config_cls = type(released_config.vision_config) - config_cls = type(released_config) - - mla_config = next( - layer.attention for layer in tt_config.layers if layer.attention is not None - ) - kda_config = next( - layer.delta_attention - for layer in tt_config.layers - if layer.delta_attention is not None - ) - dense_config = next( - layer.feed_forward - for layer in tt_config.layers - if layer.feed_forward is not None - ) - moe_config = next(layer.moe for layer in tt_config.layers if layer.moe is not None) - vision_config = tt_config.vision_encoder - if vision_config is None: - raise ValueError("The Kimi K3 debug config must include a vision encoder.") - - full_attention_layers = [ - layer_idx + 1 - for layer_idx, layer in enumerate(tt_config.layers) - if layer.attention is not None - ] - kda_layers = [ - layer_idx + 1 - for layer_idx, layer in enumerate(tt_config.layers) - if layer.delta_attention is not None - ] - first_moe_layer = next( - layer_idx - for layer_idx, layer in enumerate(tt_config.layers) - if layer.moe is not None - ) - - text_config = text_config_cls( - vocab_size=tt_config.vocab_size, - hidden_size=tt_config.dim, - intermediate_size=dense_config.w1.out_features, - num_hidden_layers=len(tt_config.layers), - num_attention_heads=mla_config.num_heads, - num_key_value_heads=mla_config.num_heads, - hidden_act="situ", - rms_norm_eps=tt_config.norm.eps, - use_cache=False, - moe_intermediate_size=moe_config.routed_experts[0].w1.out_features, - num_experts=moe_config.num_experts, - num_experts_per_token=moe_config.router.top_k, - num_shared_experts=( - moe_config.shared_experts.w1.out_features - // moe_config.routed_experts[0].w1.out_features - ), - first_k_dense_replace=first_moe_layer, - moe_layer_freq=1, - moe_renormalize=moe_config.router.route_norm, - routed_scaling_factor=moe_config.router.route_scale, - num_expert_group=1, - topk_group=1, - q_lora_rank=mla_config.q_lora_rank, - kv_lora_rank=mla_config.kv_lora_rank, - qk_nope_head_dim=mla_config.qk_nope_head_dim, - qk_rope_head_dim=mla_config.qk_rope_head_dim, - v_head_dim=mla_config.v_head_dim, - mla_use_nope=True, - mla_use_output_gate=True, - linear_attn_config={ - "kda_layers": kda_layers, - "full_attn_layers": full_attention_layers, - "short_conv_kernel_size": kda_config.conv_kernel_size, - "head_dim": kda_config.head_dim, - "num_heads": kda_config.num_heads, - "use_full_rank_gate": True, - "gate_lower_bound": kda_config.kernel.lower_bound, - }, - attn_res_block_size=tt_config.layers[0].attn_res_block_size, - latent_moe_use_norm=True, - activation_situ_beta=dense_config.activation.beta, - activation_situ_linear_beta=dense_config.activation.linear_beta, - routed_expert_hidden_size=moe_config.routed_down.out_features, - ) - text_config._attn_implementation = "eager" - - vision_attention = vision_config.block.attn - vision_mlp = vision_config.block.mlp - vision_config_hf = vision_config_cls( - patch_size=vision_config.patch_size, - init_pos_emb_height=vision_config.init_pos_emb_height, - init_pos_emb_width=vision_config.init_pos_emb_width, - init_pos_emb_time=vision_config.max_num_frames, - vt_num_attention_heads=vision_attention.num_heads, - vt_num_hidden_layers=vision_config.num_layers, - vt_hidden_size=vision_config.dim, - vt_intermediate_size=vision_mlp.fc1.out_features, - merge_kernel_size=tuple(vision_config.merge_kernel_size), - mm_projector_type="patchmergerv2", - qkv_hidden_size=vision_attention.qkv_dim, - text_hidden_size=tt_config.dim, - norm_type="rmsnorm", - attn_bias=False, - patch_embed_proj_bias=False, - linear_bias=False, - activation_func="gelu_pytorch_tanh", - pos_emb_interpolation_mode=vision_config.interpolation_mode, - ) - vision_config_hf._attn_implementation = "eager" - - config = config_cls( - text_config=text_config, - vision_config=vision_config_hf, - media_placeholder_token_id=_IMAGE_TOKEN_ID, - pad_token_id=0, - auto_map=released_config.auto_map, - ) - config._name_or_path = hf_repo_path - return config - - -def _first_tensor(output) -> torch.Tensor: - if isinstance(output, torch.Tensor): - return output - if isinstance(output, (list, tuple)): - for value in output: - if isinstance(value, torch.Tensor): - return value - raise TypeError(f"Expected a tensor output, got {type(output).__name__}.") - - -def _capture_tensor( - destination: dict[str, torch.Tensor], - name: str, - *, - flatten: bool = False, -) -> Callable: - def hook(_module, _inputs, output) -> None: - value = _first_tensor(output).detach().float().cpu() - if flatten: - value = value.reshape(-1, value.shape[-1]) - destination[name] = value - - return hook - - -def _capture_router_ids( - destination: dict[str, torch.Tensor], - name: str, -) -> Callable: - def hook(_module, _inputs, output) -> None: - expert_ids = output[0] - destination[name] = expert_ids.detach().reshape(-1, expert_ids.shape[-1]).cpu() - - return hook - - -def _register_text_hooks( - tt_model, - hf_model, - tt_outputs: dict[str, torch.Tensor], - hf_outputs: dict[str, torch.Tensor], -) -> None: - for layer_idx, tt_layer in enumerate(tt_model.layers.values()): - hf_layer = hf_model.language_model.model.layers[layer_idx] - layer_name = f"decoder.layer.{layer_idx}" - tt_layer.register_forward_hook(_capture_tensor(tt_outputs, layer_name)) - hf_layer.register_forward_hook(_capture_tensor(hf_outputs, layer_name)) - - attention_name = f"decoder.attention.{layer_idx}" - tt_attention = ( - tt_layer.attention - if tt_layer.attention is not None - else tt_layer.delta_attention - ) - tt_attention.register_forward_hook(_capture_tensor(tt_outputs, attention_name)) - hf_layer.self_attn.register_forward_hook( - _capture_tensor(hf_outputs, attention_name) - ) - - if tt_layer.moe is not None: - router_name = f"decoder.router_ids.{layer_idx}" - tt_layer.moe.router.register_forward_hook( - _capture_router_ids(tt_outputs, router_name) - ) - hf_layer.block_sparse_moe.gate.register_forward_hook( - _capture_router_ids(hf_outputs, router_name) - ) - - -def _register_vision_hooks( - tt_model, - hf_model, - tt_outputs: dict[str, torch.Tensor], - hf_outputs: dict[str, torch.Tensor], -) -> None: - for layer_idx, tt_layer in enumerate(tt_model.vision_encoder.layers.values()): - layer_name = f"vision.layer.{layer_idx}" - tt_layer.register_forward_hook( - _capture_tensor(tt_outputs, layer_name, flatten=True) - ) - hf_model.vision_tower.encoder.blocks[layer_idx].register_forward_hook( - _capture_tensor(hf_outputs, layer_name, flatten=True) - ) - - -def _compare_tensor( - name: str, - tt_value: torch.Tensor, - hf_value: torch.Tensor, - *, - atol: float, - rtol: float, -) -> None: - if tt_value.shape != hf_value.shape: - raise AssertionError( - f"{name}: TorchTitan shape {tuple(tt_value.shape)} does not match " - f"HF shape {tuple(hf_value.shape)}." - ) - difference = (tt_value - hf_value).abs() - cosine = F.cosine_similarity( - tt_value.reshape(-1), - hf_value.reshape(-1), - dim=0, - ).item() - print( - f"{name:32s} max={difference.max().item():.4e} " - f"mean={difference.mean().item():.4e} cos={cosine:.8f}" - ) - torch.testing.assert_close( - tt_value, - hf_value, - atol=atol, - rtol=rtol, - msg=lambda message: f"{name} failed numerical parity:\n{message}", - ) - - -def _compare_captured( - tt_outputs: dict[str, torch.Tensor], - hf_outputs: dict[str, torch.Tensor], - *, - atol: float, - rtol: float, -) -> None: - if tt_outputs.keys() != hf_outputs.keys(): - raise AssertionError( - "Captured output names differ: " - f"TT-only={sorted(tt_outputs.keys() - hf_outputs.keys())}, " - f"HF-only={sorted(hf_outputs.keys() - tt_outputs.keys())}." - ) - for name in tt_outputs: - if ".router_ids." in name: - if not torch.equal(tt_outputs[name], hf_outputs[name]): - mismatch = (tt_outputs[name] != hf_outputs[name]).sum().item() - raise AssertionError(f"{name}: {mismatch} expert IDs differ.") - print(f"{name:32s} exact expert IDs") - else: - _compare_tensor( - name, - tt_outputs[name], - hf_outputs[name], - atol=atol, - rtol=rtol, - ) - - -@torch.no_grad() -def main() -> None: - parser = argparse.ArgumentParser() - parser.add_argument("--hf_repo_path", required=True) - parser.add_argument("--device", default="cuda") - parser.add_argument( - "--tt_device", - help="TorchTitan device. Defaults to --device.", - ) - parser.add_argument( - "--hf_device", - help="HuggingFace device. Defaults to --device.", - ) - parser.add_argument( - "--device_module", - help="Optional module that registers an out-of-tree PyTorch device.", - ) - parser.add_argument( - "--hf_kda_backend", - default="fla", - choices=("fla", "reference"), - help="Use FLA or the script's pure PyTorch KDA compatibility backend.", - ) - parser.add_argument( - "--dtype", - default="float32", - choices=("float32", "bfloat16"), - ) - parser.add_argument("--atol", type=float, default=2e-4) - parser.add_argument("--rtol", type=float, default=2e-4) - parser.add_argument("--seed", type=int, default=42) - args = parser.parse_args() - - if args.device_module is not None: - importlib.import_module(args.device_module) - tt_device = torch.device(args.tt_device or args.device) - hf_device = torch.device(args.hf_device or args.device) - dtype = getattr(torch, args.dtype) - torch.manual_seed(args.seed) - if args.hf_kda_backend == "reference": - _install_reference_fla() - if "cuda" in {tt_device.type, hf_device.type}: - torch.backends.cuda.matmul.allow_tf32 = False - torch.backends.cudnn.allow_tf32 = False - torch.set_float32_matmul_precision("highest") - - tt_config = cast( - KimiK3Model.Config, - model_registry("debugmodel").model, - ) - tt_model = tt_config.build() - tt_model.init_states() - - hf_config = _build_reduced_hf_config(args.hf_repo_path, tt_config) - hf_model = AutoModelForCausalLM.from_config( - hf_config, - trust_remote_code=True, - ) - hf_state_dict = KimiK3StateDictAdapter( - tt_config, - hf_assets_path=None, - ).to_hf(tt_model.state_dict()) - hf_model.load_state_dict(hf_state_dict, strict=True) - - # KimiLinearModel selects FlashAttention during construction. Restoring - # eager here exercises the released eager attention function. - hf_model.language_model.model.config._attn_implementation = "eager" - hf_model.config.text_config._attn_implementation = "eager" - - tt_model = tt_model.to(device=tt_device, dtype=dtype).eval() - hf_model = hf_model.to(device=hf_device, dtype=dtype).eval() - - tt_outputs: dict[str, torch.Tensor] = {} - hf_outputs: dict[str, torch.Tensor] = {} - _register_text_hooks(tt_model, hf_model, tt_outputs, hf_outputs) - _register_vision_hooks(tt_model, hf_model, tt_outputs, hf_outputs) - - print("\nText-only parity") - tokens_BL = torch.tensor( - [[11, 23, 17, 31, 5, 19, 29, 3]], - dtype=torch.long, - ) - tt_tokens_BL = tokens_BL.to(tt_device) - hf_tokens_BL = tokens_BL.to(hf_device) - attention_mask_BL = torch.ones_like(hf_tokens_BL) - tt_logits_BLV = tt_model(tt_tokens_BL).float().cpu() - hf_logits_BLV = ( - hf_model.language_model( - input_ids=hf_tokens_BL, - attention_mask=attention_mask_BL, - use_cache=False, - ) - .logits.float() - .cpu() - ) - _compare_captured( - tt_outputs, - hf_outputs, - atol=args.atol, - rtol=args.rtol, - ) - _compare_tensor( - "text.logits", - tt_logits_BLV, - hf_logits_BLV, - atol=args.atol, - rtol=args.rtol, - ) - - print("\nVision parity") - tt_outputs.clear() - hf_outputs.clear() - vision_config = tt_config.vision_encoder - if vision_config is None: - raise ValueError("The Kimi K3 debug config must include a vision encoder.") - patch_size = vision_config.patch_size - grid_thw_N3 = torch.tensor([[1, 4, 4]], dtype=torch.long) - patches_PCHW = torch.randn( - 16, - 3, - patch_size, - patch_size, - dtype=dtype, - ) - tt_grid_thw_N3 = grid_thw_N3.to(tt_device) - hf_grid_thw_N3 = grid_thw_N3.to(hf_device) - tt_patches_PCHW = patches_PCHW.to(tt_device) - hf_patches_PCHW = patches_PCHW.to(hf_device) - pixels_NPK = tt_patches_PCHW.reshape(1, 16, -1) - tt_vision_NMD = tt_model.vision_encoder( - pixels_NPK, - grid_thw=tt_grid_thw_N3, - ) - hf_vision = hf_model.vision_tower(hf_patches_PCHW, hf_grid_thw_N3) - hf_vision_NMD = torch.stack(list(hf_model.mm_projector(hf_vision))) - _compare_captured( - tt_outputs, - hf_outputs, - atol=args.atol, - rtol=args.rtol, - ) - _compare_tensor( - "vision.projected", - tt_vision_NMD.float().cpu(), - hf_vision_NMD.float().cpu(), - atol=args.atol, - rtol=args.rtol, - ) - - print("\nEnd-to-end image+text parity") - tt_outputs.clear() - hf_outputs.clear() - num_vision_tokens = tt_vision_NMD.shape[1] - hf_tokens_BL = torch.tensor( - [[11, 23, _IMAGE_TOKEN_ID, 17, 31]], - dtype=torch.long, - device=hf_device, - ) - tt_tokens_BL = torch.tensor( - [[11, 23] + [_IMAGE_TOKEN_ID] * num_vision_tokens + [17, 31]], - dtype=torch.long, - device=tt_device, - ) - hf_attention_mask_BL = torch.ones_like(hf_tokens_BL) - tt_logits_BLV = tt_model( - tt_tokens_BL, - pixel_values=pixels_NPK, - grid_thw=tt_grid_thw_N3, - special_tokens={"image_id": _IMAGE_TOKEN_ID}, - ) - hf_logits_BLV = hf_model( - input_ids=hf_tokens_BL, - pixel_values=hf_patches_PCHW, - grid_thws=hf_grid_thw_N3, - attention_mask=hf_attention_mask_BL, - use_cache=False, - return_dict=True, - ).logits - _compare_captured( - tt_outputs, - hf_outputs, - atol=args.atol, - rtol=args.rtol, - ) - _compare_tensor( - "multimodal.logits", - tt_logits_BLV.float().cpu(), - hf_logits_BLV.float().cpu(), - atol=args.atol, - rtol=args.rtol, - ) - print("\nRESULT: PASS") - - -if __name__ == "__main__": - main() diff --git a/scripts/checkpoint_conversion/numerical_tests_kimi_k3_device.py b/scripts/checkpoint_conversion/numerical_tests_kimi_k3_device.py deleted file mode 100644 index 69ecdcee2e..0000000000 --- a/scripts/checkpoint_conversion/numerical_tests_kimi_k3_device.py +++ /dev/null @@ -1,435 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -"""Cross-device forward parity and training smoke test for reduced Kimi K3. - -This script compares a deterministic FP32 CPU reference against the same -TorchTitan model, weights, tokens, and image patches on another device. It then -runs real forward, backward, and AdamW updates on that device. Combine this test -with ``numerical_tests_kimi_k3.py``, which compares the CPU/CUDA TorchTitan -reference directly against the released HuggingFace implementation. - -The optional ``--device_module`` argument imports a PyTorch out-of-tree device -extension before constructing the target device. -""" - -import argparse -import importlib -import math -import time -from collections.abc import Callable - -import torch -import torch.nn.functional as F -from torch.utils.hooks import RemovableHandle - - -_IMAGE_TOKEN_ID = 7 - - -def _first_tensor(output) -> torch.Tensor: - if isinstance(output, torch.Tensor): - return output - if isinstance(output, (list, tuple)): - for value in output: - if isinstance(value, torch.Tensor): - return value - raise TypeError(f"Expected a tensor output, got {type(output).__name__}.") - - -def _capture_tensor( - destination: dict[str, torch.Tensor], - name: str, - *, - flatten: bool = False, -) -> Callable: - def hook(_module, _inputs, output) -> None: - value = _first_tensor(output).detach().float().cpu() - if flatten: - value = value.reshape(-1, value.shape[-1]) - destination[name] = value - - return hook - - -def _capture_router_ids( - destination: dict[str, torch.Tensor], - name: str, -) -> Callable: - def hook(_module, _inputs, output) -> None: - expert_ids = output[0] - destination[name] = expert_ids.detach().reshape(-1, expert_ids.shape[-1]).cpu() - - return hook - - -def _register_hooks( - model, - outputs: dict[str, torch.Tensor], -) -> list[RemovableHandle]: - handles: list[RemovableHandle] = [] - if model.vision_encoder is None: - raise ValueError("The Kimi K3 debug config must include a vision encoder.") - - for layer_idx, layer in enumerate(model.vision_encoder.layers.values()): - handles.append( - layer.register_forward_hook( - _capture_tensor(outputs, f"vision.layer.{layer_idx}", flatten=True) - ) - ) - projector = model.vision_encoder.projector - for name, module in ( - ("vision.projector.linear_1", projector.linear_1), - ("vision.projector.activation", projector.activation), - ("vision.projector.linear_2", projector.linear_2), - ("vision.projector.post_norm", projector.post_norm), - ): - handles.append( - module.register_forward_hook(_capture_tensor(outputs, name, flatten=True)) - ) - handles.append( - projector.register_forward_hook( - _capture_tensor(outputs, "vision.projected", flatten=True) - ) - ) - - for layer_idx, layer in enumerate(model.layers.values()): - handles.append( - layer.register_forward_hook( - _capture_tensor(outputs, f"decoder.layer.{layer_idx}") - ) - ) - attention = ( - layer.attention if layer.attention is not None else layer.delta_attention - ) - assert attention is not None - handles.append( - attention.register_forward_hook( - _capture_tensor(outputs, f"decoder.attention.{layer_idx}") - ) - ) - if layer.moe is not None: - handles.append( - layer.moe.router.register_forward_hook( - _capture_router_ids(outputs, f"decoder.router_ids.{layer_idx}") - ) - ) - return handles - - -def _compare_tensor( - name: str, - reference: torch.Tensor, - actual: torch.Tensor, - *, - atol: float, - rtol: float, -) -> None: - if reference.shape != actual.shape: - raise AssertionError( - f"{name}: reference shape {tuple(reference.shape)} does not match " - f"device shape {tuple(actual.shape)}." - ) - difference = (reference - actual).abs() - cosine = F.cosine_similarity( - reference.reshape(-1), - actual.reshape(-1), - dim=0, - ).item() - print( - f"{name:32s} max={difference.max().item():.4e} " - f"mean={difference.mean().item():.4e} cos={cosine:.8f}" - ) - torch.testing.assert_close( - actual, - reference, - atol=atol, - rtol=rtol, - msg=lambda message: f"{name} failed device parity:\n{message}", - ) - - -def _compare_captured( - reference_outputs: dict[str, torch.Tensor], - device_outputs: dict[str, torch.Tensor], - *, - atol: float, - rtol: float, -) -> None: - if reference_outputs.keys() != device_outputs.keys(): - raise AssertionError( - "Captured output names differ: " - f"reference-only={sorted(reference_outputs.keys() - device_outputs.keys())}, " - f"device-only={sorted(device_outputs.keys() - reference_outputs.keys())}." - ) - for name in reference_outputs: - if ".router_ids." in name: - if not torch.equal(reference_outputs[name], device_outputs[name]): - mismatch = ( - (reference_outputs[name] != device_outputs[name]).sum().item() - ) - raise AssertionError(f"{name}: {mismatch} expert IDs differ.") - print(f"{name:32s} exact expert IDs") - else: - _compare_tensor( - name, - reference_outputs[name], - device_outputs[name], - atol=atol, - rtol=rtol, - ) - - -def _build_inputs(config, *, seed: int) -> dict[str, torch.Tensor | dict[str, int]]: - vision_config = config.vision_encoder - if vision_config is None: - raise ValueError("The Kimi K3 debug config must include a vision encoder.") - - generator = torch.Generator(device="cpu") - generator.manual_seed(seed + 1) - patch_size = vision_config.patch_size - grid_thw_N3 = torch.tensor([[1, 4, 4]], dtype=torch.long) - pixel_values_NPK = torch.randn( - 1, - 16, - 3 * patch_size * patch_size, - generator=generator, - dtype=torch.float32, - ) - tokens_BL = torch.tensor( - [ - [ - 11, - 23, - _IMAGE_TOKEN_ID, - _IMAGE_TOKEN_ID, - _IMAGE_TOKEN_ID, - _IMAGE_TOKEN_ID, - 17, - 31, - ] - ], - dtype=torch.long, - ) - return { - "tokens": tokens_BL, - "pixel_values": pixel_values_NPK, - "grid_thw": grid_thw_N3, - "special_tokens": {"image_id": _IMAGE_TOKEN_ID}, - } - - -def _move_inputs( - inputs: dict[str, torch.Tensor | dict[str, int]], - device: torch.device, -) -> dict[str, torch.Tensor | dict[str, int]]: - return { - name: value.to(device) if isinstance(value, torch.Tensor) else value - for name, value in inputs.items() - } - - -def _synchronize(device: torch.device) -> None: - device_api = getattr(torch, device.type, None) - if device_api is not None and hasattr(device_api, "synchronize"): - device_api.synchronize(device) - - -def _verify_device(device: torch.device) -> None: - if device.type == "cpu": - return - device_api = getattr(torch, device.type, None) - if device_api is None: - raise RuntimeError( - f"PyTorch has no registered {device.type!r} device module. " - "Use --device_module to import its extension." - ) - if hasattr(device_api, "is_available") and not device_api.is_available(): - raise RuntimeError(f"Requested device {device} is not available.") - - -@torch.no_grad() -def _run_forward_parity( - reference_model, - device_model, - reference_inputs, - device_inputs, - *, - atol: float, - rtol: float, -) -> None: - reference_outputs: dict[str, torch.Tensor] = {} - device_outputs: dict[str, torch.Tensor] = {} - reference_handles = _register_hooks(reference_model, reference_outputs) - device_handles = _register_hooks(device_model, device_outputs) - - try: - reference_logits_BLV = reference_model(**reference_inputs).float().cpu() - device_logits_BLV = device_model(**device_inputs).float().cpu() - finally: - for handle in reference_handles + device_handles: - handle.remove() - - _compare_captured( - reference_outputs, - device_outputs, - atol=atol, - rtol=rtol, - ) - _compare_tensor( - "multimodal.logits", - reference_logits_BLV, - device_logits_BLV, - atol=atol, - rtol=rtol, - ) - - -def _run_training( - model, - inputs, - device: torch.device, - *, - train_dtype: torch.dtype, - train_steps: int, - learning_rate: float, -) -> None: - if train_steps == 0: - return - - model.to(dtype=train_dtype) - model.train() - optimizer = torch.optim.AdamW( - model.parameters(), - lr=learning_rate, - foreach=False, - ) - tokens_BL = inputs["tokens"] - assert isinstance(tokens_BL, torch.Tensor) - if model.lm_head is None: - raise ValueError("Kimi K3 device training requires an LM head.") - tracked_parameter = model.lm_head.weight - tracked_before = tracked_parameter.detach().clone() - - for step_idx in range(train_steps): - optimizer.zero_grad(set_to_none=True) - _synchronize(device) - start_time = time.perf_counter() - logits_BLV = model(**inputs) - loss = F.cross_entropy( - logits_BLV[:, :-1].float().reshape(-1, logits_BLV.shape[-1]), - tokens_BL[:, 1:].reshape(-1), - ) - loss.backward() - grad_norm = torch.nn.utils.clip_grad_norm_( - model.parameters(), - max_norm=math.inf, - foreach=False, - ) - optimizer.step() - _synchronize(device) - elapsed = time.perf_counter() - start_time - - loss_value = loss.detach().float().cpu().item() - grad_norm_value = grad_norm.detach().float().cpu().item() - if not math.isfinite(loss_value) or not math.isfinite(grad_norm_value): - raise AssertionError( - f"Step {step_idx}: non-finite loss={loss_value} or " - f"grad_norm={grad_norm_value}." - ) - print( - f"step={step_idx:02d} loss={loss_value:.8f} " - f"grad_norm={grad_norm_value:.8f} elapsed={elapsed:.3f}s" - ) - - parameter_delta = ( - (tracked_parameter.detach().float() - tracked_before.float()).abs().max() - ) - parameter_delta_value = parameter_delta.cpu().item() - if parameter_delta_value == 0.0: - raise AssertionError("AdamW completed without changing lm_head.weight.") - print(f"lm_head.weight max update={parameter_delta_value:.4e}") - - -def main() -> None: - parser = argparse.ArgumentParser() - parser.add_argument("--device", required=True) - parser.add_argument( - "--device_module", - help="Optional module that registers an out-of-tree PyTorch device.", - ) - parser.add_argument( - "--parity_dtype", - default="float32", - choices=("float32", "bfloat16"), - ) - parser.add_argument( - "--train_dtype", - default="bfloat16", - choices=("float32", "bfloat16"), - ) - parser.add_argument("--atol", type=float, default=2e-4) - parser.add_argument("--rtol", type=float, default=2e-4) - parser.add_argument("--train_steps", type=int, default=2) - parser.add_argument("--learning_rate", type=float, default=8e-4) - parser.add_argument("--seed", type=int, default=42) - args = parser.parse_args() - - if args.device_module is not None: - importlib.import_module(args.device_module) - - from torchtitan.models.kimi_k3 import model_registry - - device = torch.device(args.device) - _verify_device(device) - parity_dtype = getattr(torch, args.parity_dtype) - train_dtype = getattr(torch, args.train_dtype) - if args.train_steps < 0: - raise ValueError("--train_steps must be non-negative.") - - torch.manual_seed(args.seed) - if device.type == "cuda": - torch.backends.cuda.matmul.allow_tf32 = False - torch.backends.cudnn.allow_tf32 = False - torch.set_float32_matmul_precision("highest") - - config = model_registry("debugmodel").model - reference_model = config.build() - reference_model.init_states() - device_model = config.build() - device_model.init_states() - device_model.load_state_dict(reference_model.state_dict(), strict=True) - - reference_model = reference_model.to(dtype=torch.float32).eval() - device_model = device_model.to(device=device, dtype=parity_dtype).eval() - reference_inputs = _build_inputs(config, seed=args.seed) - device_inputs = _move_inputs(reference_inputs, device) - - print(f"\nForward parity: cpu/float32 -> {device}/{args.parity_dtype}") - _run_forward_parity( - reference_model, - device_model, - reference_inputs, - device_inputs, - atol=args.atol, - rtol=args.rtol, - ) - - del reference_model - print(f"\nTraining smoke: {device}/{args.train_dtype}, steps={args.train_steps}") - _run_training( - device_model, - device_inputs, - device, - train_dtype=train_dtype, - train_steps=args.train_steps, - learning_rate=args.learning_rate, - ) - print("\nRESULT: PASS") - - -if __name__ == "__main__": - main() diff --git a/torchtitan/models/kimi_k3/README.md b/torchtitan/models/kimi_k3/README.md index 265ad07257..8e53772ad0 100644 --- a/torchtitan/models/kimi_k3/README.md +++ b/torchtitan/models/kimi_k3/README.md @@ -76,7 +76,7 @@ compressed tensors is outside this first change. Numerical comparison should therefore instantiate the same reduced, unquantized model on both sides and copy one state dict through the adapter. -## Validation contract +## Tests The CPU unit tests cover: @@ -86,34 +86,8 @@ The CPU unit tests cover: - a small text+image model forward and backward; - exhaustive state-dict round-trip for that small model. -Before requesting merge, the following CUDA tests must also pass: - -1. One single-device training step on a CUDA GPU. -2. FP32 forward comparison between CPU and CUDA using the same model and inputs. -3. FP32 forward comparison against the released HuggingFace code using the - same reduced config, weights, tokens, image patches, and expert choices. -4. BF16 forward, backward, and optimizer steps on CUDA. - -The parity report must include intermediate checks for KDA, MLA, vision -features, each decoder layer, final logits, and router expert IDs. Final-logit -metrics alone are not sufficient to localize a discrete routing mismatch. - -Run the CPU-to-CUDA comparison and real optimizer steps with: - -```bash -python -m scripts.checkpoint_conversion.numerical_tests_kimi_k3_device \ - --device cuda -``` - -This device test is one link in the parity chain; it does not replace the -direct HuggingFace comparison below. - -Run both implementations on a CUDA GPU with the released model dependencies: - ```bash -CUDA_VISIBLE_DEVICES=0 python -m \ - scripts.checkpoint_conversion.numerical_tests_kimi_k3 \ - --hf_repo_path /path/to/moonshotai/Kimi-K3 +pytest -q tests/unit_tests/test_kimi_k3.py ``` ## First-version limitations From bf7c6d5e815a7f78581d1abfe0128b8eb4412efd Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Thu, 30 Jul 2026 17:53:40 +0800 Subject: [PATCH 03/67] Add FSDP2 support for Kimi K3 --- tests/integration_tests/models.py | 7 + tests/unit_tests/test_kimi_k3_fsdp.py | 130 +++++++++++++++++++ torchtitan/models/kimi_k3/README.md | 25 +++- torchtitan/models/kimi_k3/config_registry.py | 4 +- torchtitan/models/kimi_k3/model.py | 25 +++- torchtitan/models/kimi_k3/parallelize.py | 66 ++++++++-- 6 files changed, 233 insertions(+), 24 deletions(-) create mode 100644 tests/unit_tests/test_kimi_k3_fsdp.py diff --git a/tests/integration_tests/models.py b/tests/integration_tests/models.py index 974c6fb86b..74f3627d1b 100755 --- a/tests/integration_tests/models.py +++ b/tests/integration_tests/models.py @@ -129,4 +129,11 @@ def build_model_tests_list() -> list[OverrideDefinitions]: test_name="muse_glimmer_mm_fsdp+tp+sp", ngpu=4, ), + # Integration Test Case for Kimi K3 + OverrideDefinitions( + configs=[recipes.kimi_k3_debugmodel_mm_fsdp2], + test_descr="Kimi K3 multimodal FSDP", + test_name="kimi_k3_mm_fsdp", + ngpu=2, + ), ] diff --git a/tests/unit_tests/test_kimi_k3_fsdp.py b/tests/unit_tests/test_kimi_k3_fsdp.py new file mode 100644 index 0000000000..abebd014aa --- /dev/null +++ b/tests/unit_tests/test_kimi_k3_fsdp.py @@ -0,0 +1,130 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import copy +from unittest.mock import patch + +import torch +from torch.distributed._composable.fsdp import FSDPModule +from torch.distributed.tensor import DTensor +from torch.testing._internal.distributed._tensor.common_dtensor import ( + DTensorTestBase, + with_comms, +) +from torchtitan.config import CompileConfig, ParallelismConfig, TrainingConfig +from torchtitan.distributed import ParallelDims +from torchtitan.models.kimi_k3 import parallelize_kimi_k3 +from torchtitan.models.kimi_k3.model import KimiK3Model + +from tests.unit_tests.test_kimi_k3 import _small_model_config + + +class TestKimiK3FSDP(DTensorTestBase): + @property + def world_size(self): + return 1 + + @with_comms + def test_single_rank_fsdp_matches_manual_bf16_reference(self): + torch.manual_seed(3) + config = _small_model_config() + with torch.device("meta"): + model = config.build() + model.to_empty(device=self.device_type) + model.init_states() + with torch.no_grad(): + for transformer_block in model.layers.values(): + if transformer_block.moe is not None: + transformer_block.moe.router.gate.weight.zero_() + + reference = copy.deepcopy(model) + for parameter in reference.parameters(): + parameter.data = parameter.data.to(torch.bfloat16) + + parallelism = ParallelismConfig( + data_parallel_shard_degree=1, + tensor_parallel_degree=1, + pipeline_parallel_degree=1, + context_parallel_degree=1, + expert_parallel_degree=1, + ) + parallel_dims = ParallelDims.from_config(parallelism, world_size=1) + with patch( + "torchtitan.distributed.parallel_dims.device_type", + self.device_type, + ): + parallel_dims.build_mesh() + model = parallelize_kimi_k3( + model, + parallel_dims=parallel_dims, + training=TrainingConfig( + local_batch_size=1, + seq_len=6, + steps=1, + dtype="bfloat16", + ), + parallelism=parallelism, + compile_config=CompileConfig(), + ac_config=None, + dump_folder="", + ) + + assert isinstance(model, KimiK3Model) + self.assertIsInstance(model, FSDPModule) + self.assertIsInstance(model.vision_encoder, FSDPModule) + + inputs = { + "tokens": torch.tensor( + [[1, 7, 2, 3, 4, 5]], + dtype=torch.long, + device=self.device_type, + ), + "pixel_values": torch.randn( + 1, + 4, + 3 * 2 * 2, + device=self.device_type, + ), + "grid_thw": torch.tensor( + [[1, 2, 2]], + dtype=torch.long, + device=self.device_type, + ), + "special_tokens": {"image_id": 7}, + } + + actual_BLV = model(**inputs) # pyrefly: ignore [not-callable] + expected_BLV = reference(**inputs) + torch.testing.assert_close(actual_BLV, expected_BLV, atol=0.0, rtol=0.0) + + actual_BLV.float().square().mean().backward() + expected_BLV.float().square().mean().backward() + + reference_parameters = dict(reference.named_parameters()) + compared_gradients = 0 + for name, parameter in model.named_parameters(): + actual_grad = parameter.grad + expected_grad = reference_parameters[name].grad + self.assertEqual(actual_grad is None, expected_grad is None) + if actual_grad is None: + continue + if isinstance(actual_grad, DTensor): + actual_grad = actual_grad.to_local() + assert expected_grad is not None + torch.testing.assert_close( + actual_grad.float(), + expected_grad.float(), + atol=0.0, + rtol=0.0, + ) + compared_gradients += 1 + self.assertGreater(compared_gradients, 0) + + +if __name__ == "__main__": + from torch.testing._internal.common_utils import run_tests + + run_tests() diff --git a/torchtitan/models/kimi_k3/README.md b/torchtitan/models/kimi_k3/README.md index 8e53772ad0..fc2b399fa6 100644 --- a/torchtitan/models/kimi_k3/README.md +++ b/torchtitan/models/kimi_k3/README.md @@ -1,8 +1,8 @@ # Kimi K3 -This directory contains the first TorchTitan implementation of Kimi K3. The -initial scope is a topology-complete, reduced model for single-device training -and numerical comparison with the +This directory contains the eager numerical reference implementation of Kimi +K3 in TorchTitan. The initial scope is a topology-complete, reduced model for +single-device or FSDP2 training and numerical comparison with the [released HuggingFace implementation](https://huggingface.co/moonshotai/Kimi-K3). The implementation is device-neutral. It uses PyTorch operators and does not @@ -15,6 +15,13 @@ inspectable model math over throughput. NGPU=1 MODULE=kimi_k3 CONFIG=kimi_k3_debugmodel ./run_train.sh ``` +Run the same eager model with two-way FSDP2: + +```bash +NGPU=2 MODULE=kimi_k3 CONFIG=kimi_k3_debugmodel ./run_train.sh \ + --parallelism.data_parallel_shard_degree 2 +``` + The multimodal data path requires `torchvision`. ## Reduced model @@ -58,7 +65,8 @@ The reference path mirrors the released implementation in these areas: - Vision features scattered into runs of media placeholder tokens. `KimiKDAKernel` is the optimization boundary. A future accelerated backend -should preserve its input/output contract and checkpoint schema. +should preserve its input/output contract and checkpoint schema. FSDP2 only +shards parameters and leaves this eager forward contract unchanged. ## Checkpoint conversion @@ -85,14 +93,17 @@ The CPU unit tests cover: - the KDA kernel against a direct recurrent formulation, including backward; - a small text+image model forward and backward; - exhaustive state-dict round-trip for that small model. +- single-rank FSDP2 forward and per-parameter gradient parity with a manually + cast BF16 reference. ```bash pytest -q tests/unit_tests/test_kimi_k3.py +pytest -q tests/unit_tests/test_kimi_k3_fsdp.py ``` ## First-version limitations -- Single device only; DP, FSDP/HSDP, TP, PP, CP, and EP are rejected. +- FSDP2 data parallelism is supported; HSDP, TP, PP, CP, and EP are rejected. - No packed documents, activation checkpointing, `torch.compile`, or CPU offload. - Image inputs are supported; video inputs are rejected. @@ -102,5 +113,5 @@ pytest -q tests/unit_tests/test_kimi_k3.py - No full 2.8T flavor. These restrictions are explicit so unsupported runtime settings fail instead -of being silently ignored. Optimized kernels and distributed parallelism can be -added in follow-up changes after the reference forward is numerically locked. +of being silently ignored. EP and optimized kernels can be added in follow-up +changes after the eager/FSDP2 reference forward is numerically locked. diff --git a/torchtitan/models/kimi_k3/config_registry.py b/torchtitan/models/kimi_k3/config_registry.py index d31b364e71..7da317143c 100644 --- a/torchtitan/models/kimi_k3/config_registry.py +++ b/torchtitan/models/kimi_k3/config_registry.py @@ -22,7 +22,7 @@ def kimi_k3_debugmodel() -> Trainer.Config: - """Return the single-device, topology-complete Kimi K3 debug config.""" + """Return the topology-complete Kimi K3 eager/FSDP2 debug config.""" model_spec = model_registry("debugmodel") return Trainer.Config( loss=ChunkedLossWrapper.Config( @@ -58,7 +58,7 @@ def kimi_k3_debugmodel() -> Trainer.Config: ), training=TrainingConfig( local_batch_size=1, - seq_len=128, + seq_len=256, steps=10, dtype="bfloat16", ), diff --git a/torchtitan/models/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py index f1055546b3..2f192cc18c 100644 --- a/torchtitan/models/kimi_k3/model.py +++ b/torchtitan/models/kimi_k3/model.py @@ -433,8 +433,27 @@ def forward( return expert_ids_BLK, weights_BLK * self.route_scale +class KimiRoutedExperts(ModuleList): + """List-backed experts implementing TorchTitan's FSDP expert protocol. + + Kimi K3 keeps one module per expert so the eager implementation and + HuggingFace state-dict mapping stay directly inspectable. The shared FSDP + wrapper discovers routed expert parameters through ``inner_experts`` and + ``num_experts``; exposing those properties here lets it shard this + list-backed layout without changing parameter names or forward math. + """ + + @property + def inner_experts(self) -> "KimiRoutedExperts": + return self + + @property + def num_experts(self) -> int: + return len(self) + + class KimiLatentMoE(Module): - """Single-device trainable implementation of Kimi K3 latent MoE.""" + """Eager trainable implementation of Kimi K3 latent MoE.""" @dataclass(kw_only=True, slots=True) class Config(Module.Config): @@ -456,7 +475,7 @@ def __init__(self, config: Config): self.num_experts = config.num_experts self.router = config.router.build() self.routed_down = config.routed_down.build() - self.routed_experts = ModuleList( + self.routed_experts = KimiRoutedExperts( [expert.build() for expert in config.routed_experts] ) self.routed_norm = config.routed_norm.build() @@ -669,7 +688,7 @@ def update_from_config(self, *, config, **kwargs) -> None: enabled = [name for name, degree in unsupported.items() if degree > 1] if enabled: raise NotImplementedError( - "Kimi K3 v1 supports single-device execution only; " + "Kimi K3 eager reference supports FSDP2 data parallelism only; " f"disable {', '.join(enabled)}." ) dataloader = getattr(config, "dataloader", None) diff --git a/torchtitan/models/kimi_k3/parallelize.py b/torchtitan/models/kimi_k3/parallelize.py index 14a759a891..7077009e30 100644 --- a/torchtitan/models/kimi_k3/parallelize.py +++ b/torchtitan/models/kimi_k3/parallelize.py @@ -4,13 +4,22 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -"""Parallelization boundary for the first Kimi K3 implementation.""" +"""FSDP2 parallelization for the eager Kimi K3 reference model.""" import torch.nn as nn -from torchtitan.config import CompileConfig, ParallelismConfig, TrainingConfig +from torchtitan.config import ( + CompileConfig, + ParallelismConfig, + TORCH_DTYPE_MAP, + TrainingConfig, +) from torchtitan.distributed import ParallelDims from torchtitan.distributed.activation_checkpoint import ActivationCheckpointingConfig +from torchtitan.distributed.fsdp import ( + apply_fsdp_to_decoder, + apply_fsdp_to_vision_encoder, +) def parallelize_kimi_k3( @@ -23,13 +32,13 @@ def parallelize_kimi_k3( ac_config: ActivationCheckpointingConfig, dump_folder: str, ) -> nn.Module: - """Validate the v1 single-device contract and return the model unchanged.""" + """Apply FSDP2 while keeping the model's eager reference forward path.""" del dump_folder - enabled_parallelisms = [ + unsupported_parallelisms = [ name for name, enabled in ( - ("data parallel", parallel_dims.dp_enabled), + ("hybrid sharded data parallel", parallel_dims.dp_replicate_enabled), ("tensor parallel", parallel_dims.tp_enabled), ("pipeline parallel", parallel_dims.pp_enabled), ("context parallel", parallel_dims.cp_enabled), @@ -37,21 +46,54 @@ def parallelize_kimi_k3( ) if enabled ] - if enabled_parallelisms: + if unsupported_parallelisms: raise NotImplementedError( - "Kimi K3 v1 supports single-device execution only; disable " - f"{', '.join(enabled_parallelisms)}." + "Kimi K3 eager reference currently supports FSDP2 data parallelism " + f"only; disable {', '.join(unsupported_parallelisms)}." ) if parallelism.spmd_backend != "default": raise NotImplementedError( - "Kimi K3 v1 only supports the default local tensor backend." + "Kimi K3 eager FSDP2 currently supports the default SPMD backend only." ) if compile_config.enable: - raise NotImplementedError("Kimi K3 v1 does not support torch.compile.") + raise NotImplementedError( + "Kimi K3 eager reference does not support torch.compile." + ) if ac_config is not None: raise NotImplementedError( - "Kimi K3 v1 does not support activation checkpointing." + "Kimi K3 eager FSDP2 does not support activation checkpointing yet." ) if training.enable_cpu_offload: - raise NotImplementedError("Kimi K3 v1 does not support parameter CPU offload.") + raise NotImplementedError( + "Kimi K3 eager FSDP2 does not support parameter CPU offload yet." + ) + + dp_mesh_names = ( + ["dp_replicate", "fsdp"] if parallel_dims.dp_replicate_enabled else ["fsdp"] + ) + dp_mesh = parallel_dims.get_mesh(dp_mesh_names) + + vision_encoder = getattr(model, "vision_encoder", None) + if vision_encoder is not None: + apply_fsdp_to_vision_encoder( + vision_encoder, + dp_mesh, + param_dtype=TORCH_DTYPE_MAP[training.mixed_precision_param], + reduce_dtype=TORCH_DTYPE_MAP[training.mixed_precision_reduce], + reshard_after_forward_policy=parallelism.fsdp_reshard_after_forward, + pp_enabled=False, + ) + + apply_fsdp_to_decoder( + model, # pyrefly: ignore [bad-argument-type] + dp_mesh, + param_dtype=TORCH_DTYPE_MAP[training.mixed_precision_param], + reduce_dtype=TORCH_DTYPE_MAP[training.mixed_precision_reduce], + pp_enabled=False, + cpu_offload=training.enable_cpu_offload, + reshard_after_forward_policy=parallelism.fsdp_reshard_after_forward, + ep_degree=1, + enable_symm_mem=parallelism.enable_fsdp_symm_mem, + ) + return model From 31fd2d7e670dfc4023f450eb33bba6725eeda67c Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Thu, 30 Jul 2026 20:33:50 +0800 Subject: [PATCH 04/67] Scale Kimi K3 debug model --- tests/unit_tests/test_kimi_k3.py | 6 ++++++ torchtitan/models/kimi_k3/README.md | 13 +++++++++---- torchtitan/models/kimi_k3/__init__.py | 12 ++++++------ 3 files changed, 21 insertions(+), 10 deletions(-) diff --git a/tests/unit_tests/test_kimi_k3.py b/tests/unit_tests/test_kimi_k3.py index 6f93a2083a..aff7f708b5 100644 --- a/tests/unit_tests/test_kimi_k3.py +++ b/tests/unit_tests/test_kimi_k3.py @@ -182,6 +182,7 @@ def test_exact_gelu_matches_pytorch_reference(self): def test_debugmodel_preserves_reduced_k3_topology(self): config = kimi_k3_configs["debugmodel"]("eager") + self.assertEqual(config.vocab_size, 163840) self.assertEqual(len(config.layers), 13) self.assertEqual( [ @@ -197,6 +198,11 @@ def test_debugmodel_preserves_reduced_k3_topology(self): assert moe_config is not None self.assertEqual(moe_config.num_experts, 8) self.assertEqual(moe_config.router.top_k, 2) + self.assertEqual(moe_config.routed_experts[0].w1.out_features, 128) + vision_config = config.vision_encoder + assert vision_config is not None + self.assertEqual(vision_config.dim, 256) + self.assertEqual(vision_config.num_layers, 4) def test_kda_kernel_matches_recurrent_reference(self): torch.manual_seed(1) diff --git a/torchtitan/models/kimi_k3/README.md b/torchtitan/models/kimi_k3/README.md index fc2b399fa6..fd352b5bff 100644 --- a/torchtitan/models/kimi_k3/README.md +++ b/torchtitan/models/kimi_k3/README.md @@ -32,20 +32,25 @@ reducing widths, expert count, and depth. | Component | Released Kimi K3 | `debugmodel` | |---|---:|---:| | Decoder dimension | 7168 | 256 | -| Vocabulary size | 163840 | 2048 | +| Vocabulary size | 163840 | 163840 | | Decoder layers | 93 | 13 | | Full MLA layers (1-based) | 4, 8, ..., 92, 93 | 4, 8, 12 | | KDA layers | 69 | 10 | | Dense FFN layers | 1 | 1 | | Attention residual block size | 12 | 12 | | Routed experts / top-k | 896 / 16 | 8 / 2 | +| Routed latent / expert hidden dimension | 3584 / 3072 | 128 / 128 | | Shared experts | 2 | 2 | -| Vision dimension | 1024 | 128 | -| Vision layers | 27 | 2 | -| Vision QKV dimension / heads | 1536 / 12 | 192 / 3 | +| Vision dimension | 1024 | 256 | +| Vision layers | 27 | 4 | +| Vision QKV dimension / heads | 1536 / 12 | 384 / 3 | Thirteen decoder layers are intentional. They exercise two attention-residual blocks and preserve the released model's 1-based full-attention cadence. +The released vocabulary size is retained, following other TorchTitan +multimodal debug models and making FSDP state sharding measurable while the +decoder widths and depths remain reduced. The resulting model has about 100 +million parameters. ## Forward structure diff --git a/torchtitan/models/kimi_k3/__init__.py b/torchtitan/models/kimi_k3/__init__.py index d7ca6015cc..f93cffb18a 100644 --- a/torchtitan/models/kimi_k3/__init__.py +++ b/torchtitan/models/kimi_k3/__init__.py @@ -349,7 +349,7 @@ def _debugmodel(attn_backend: str) -> KimiK3Model.Config: raise ValueError("Kimi K3 v1 only provides the device-neutral 'eager' backend.") dim = 256 - vocab_size = 2048 + vocab_size = 163840 num_layers = 13 full_attention_layers = {4, 8, 12} num_heads = 4 @@ -398,7 +398,7 @@ def _debugmodel(attn_backend: str) -> KimiK3Model.Config: else _latent_moe_config( dim=dim, latent_dim=128, - expert_hidden_dim=64, + expert_hidden_dim=128, num_experts=8, top_k=2, num_shared_experts=2, @@ -432,10 +432,10 @@ def _debugmodel(attn_backend: str) -> KimiK3Model.Config: output_res_proj=_linear(dim, 1), vision_encoder=_vision_encoder_config( text_dim=dim, - dim=128, - qkv_dim=192, - hidden_dim=512, - num_layers=2, + dim=256, + qkv_dim=384, + hidden_dim=1024, + num_layers=4, num_heads=3, ), spatial_merge_size=2, From 51b70076d33c03d25a47597e6cff17ed459db70a Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Thu, 30 Jul 2026 22:43:50 +0800 Subject: [PATCH 05/67] Use out-of-place operations in Kimi K3 eager paths --- tests/unit_tests/test_kimi_k3.py | 74 +++++++++++++++++++-- torchtitan/models/kimi_k3/model.py | 43 ++++++++++-- torchtitan/models/kimi_k3/vision_encoder.py | 47 +++++++------ 3 files changed, 134 insertions(+), 30 deletions(-) diff --git a/tests/unit_tests/test_kimi_k3.py b/tests/unit_tests/test_kimi_k3.py index aff7f708b5..48df1ec0e4 100644 --- a/tests/unit_tests/test_kimi_k3.py +++ b/tests/unit_tests/test_kimi_k3.py @@ -22,6 +22,7 @@ kimi_k3_configs, ) from torchtitan.models.kimi_k3.model import ( + _replace_vision_embeds, KimiK3Model, KimiK3TransformerBlock, KimiKDAKernel, @@ -169,6 +170,60 @@ def _kda_recurrent_reference( class TestKimiK3(unittest.TestCase): + def test_replace_vision_embeds_is_out_of_place(self): + inputs_embeds = torch.arange(30.0).view(2, 5, 3).requires_grad_() + vision_embeds = torch.arange(12.0).view(2, 2, 3).requires_grad_() + inputs_before = inputs_embeds.detach().clone() + + actual = _replace_vision_embeds( + inputs_embeds, + vision_embeds=vision_embeds, + vision_positions=[ + (0, 0, 1, 1), + (1, 1, 2, 2), + ], + ) + expected = torch.stack( + ( + torch.cat( + ( + inputs_before[0, :1], + vision_embeds.detach()[0, :1], + inputs_before[0, 2:], + ) + ), + torch.cat( + ( + inputs_before[1, :2], + vision_embeds.detach()[1], + inputs_before[1, 4:], + ) + ), + ) + ) + + torch.testing.assert_close(actual, expected) + torch.testing.assert_close(inputs_embeds, inputs_before) + self.assertNotEqual(actual.data_ptr(), inputs_embeds.data_ptr()) + + actual.sum().backward() + expected_input_grad = torch.tensor( + [[1, 0, 1, 1, 1], [1, 1, 0, 0, 1]], + dtype=inputs_embeds.dtype, + ).unsqueeze(-1) + expected_vision_grad = torch.tensor( + [[1, 0], [1, 1]], + dtype=vision_embeds.dtype, + ).unsqueeze(-1) + torch.testing.assert_close( + inputs_embeds.grad, + expected_input_grad.expand_as(inputs_embeds), + ) + torch.testing.assert_close( + vision_embeds.grad, + expected_vision_grad.expand_as(vision_embeds), + ) + def test_exact_gelu_matches_pytorch_reference(self): x = torch.linspace(-4.0, 4.0, 257) actual = KimiExactGELU.Config().build()(x) @@ -264,11 +319,15 @@ def test_small_multimodal_model_forward_backward_and_adapter(self): model.verify_module_protocol() model.init_states() - tokens_BL = torch.randint(0, config.vocab_size, (2, 6)) image_token_id = 7 - tokens_BL[0, 2] = image_token_id - pixel_values_NPK = torch.randn(1, 4, 3 * 2 * 2) - grid_thw_N3 = torch.tensor([[1, 2, 2]]) + tokens_BL = torch.tensor( + [ + [1, 2, image_token_id, 3, 4, 5], + [6, image_token_id, image_token_id, 8, 9, 10], + ] + ) + pixel_values_NPK = torch.randn(2, 8, 3 * 2 * 2) + grid_thw_N3 = torch.tensor([[1, 2, 2], [1, 4, 2]]) logits_BLV = model( tokens_BL, pixel_values=pixel_values_NPK, @@ -277,6 +336,13 @@ def test_small_multimodal_model_forward_backward_and_adapter(self): ) self.assertEqual(logits_BLV.shape, (2, 6, config.vocab_size)) + moe = model.layers["1"].moe + assert moe is not None + self.assertIs( + moe._buffers["tokens_per_expert_E"], + moe.tokens_per_expert_E, + ) + self.assertEqual(moe.tokens_per_expert_E.sum().item(), tokens_BL.numel()) logits_BLV.float().square().mean().backward() for parameter in model.parameters(): if parameter.grad is not None: diff --git a/torchtitan/models/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py index 2f192cc18c..48a855806e 100644 --- a/torchtitan/models/kimi_k3/model.py +++ b/torchtitan/models/kimi_k3/model.py @@ -23,10 +23,7 @@ from torchtitan.models.common import Conv1d, Linear from torchtitan.models.common.attention import AttentionMasksType from torchtitan.models.common.decoder import Decoder -from torchtitan.models.common.multimodal import ( - get_vision_positions, - scatter_vision_embeds, -) +from torchtitan.models.common.multimodal import get_vision_positions from torchtitan.models.common.nn_modules import RMSNorm from torchtitan.models.kimi_k3.vision_encoder import KimiK3VisionEncoder from torchtitan.models.utils import get_moe_model_nparams_and_flops @@ -38,6 +35,36 @@ # T = flattened tokens, N = attention-residual entries. +def _replace_vision_embeds( + inputs_embeds: torch.Tensor, + *, + vision_embeds: torch.Tensor, + vision_positions: list[tuple[int, int, int, int]], +) -> torch.Tensor: + """Return text embeddings with vision spans replaced out of place.""" + seq_len = inputs_embeds.shape[1] + flat_positions = [] + valid_vision_embeds = [] + for item_idx, sample_idx, vision_start, n_tokens in vision_positions: + flat_start = sample_idx * seq_len + vision_start + flat_positions.append( + torch.arange( + flat_start, + flat_start + n_tokens, + device=inputs_embeds.device, + ) + ) + valid_vision_embeds.append(vision_embeds[item_idx, :n_tokens]) + + indices = torch.cat(flat_positions) + replacements = torch.cat(valid_vision_embeds).to(inputs_embeds.dtype) + return ( + inputs_embeds.flatten(0, 1) + .index_copy(0, indices, replacements) + .view_as(inputs_embeds) + ) + + class KimiRMSNorm(RMSNorm): """RMSNorm with the explicit FP32 reduction used by the Kimi reference.""" @@ -507,9 +534,11 @@ def forward(self, x_BLD: torch.Tensor) -> torch.Tensor: self.num_experts, dtype=torch.bool, device=x_BLD.device, - ).scatter_(-1, expert_ids_BLK, True) + ).scatter(-1, expert_ids_BLK, True) with torch.no_grad(): - self.tokens_per_expert_E.add_(routing_map_BLE.sum(dim=(0, 1)).float()) + self.tokens_per_expert_E = ( + self.tokens_per_expert_E + routing_map_BLE.sum(dim=(0, 1)).float() + ) latent_TD = self.routed_down(x_BLD).reshape(B * L, -1) expert_ids_TK = expert_ids_BLK.reshape(B * L, -1) @@ -764,7 +793,7 @@ def _prepare_multimodal_embeds( raise ValueError( "pixel_values were provided but no image placeholder tokens were found." ) - return scatter_vision_embeds( + return _replace_vision_embeds( embeddings, vision_embeds=vision_embeds, vision_positions=vision_positions, diff --git a/torchtitan/models/kimi_k3/vision_encoder.py b/torchtitan/models/kimi_k3/vision_encoder.py index 1cb81f86a4..2b0cdc88f7 100644 --- a/torchtitan/models/kimi_k3/vision_encoder.py +++ b/torchtitan/models/kimi_k3/vision_encoder.py @@ -71,6 +71,19 @@ def _get_temporal_pos_embed( return torch.cat((angles.sin(), angles.cos()), dim=-1) +def _pad_sequence(x: torch.Tensor, target_length: int) -> torch.Tensor: + """Pad the leading sequence dimension without modifying ``x`` in place.""" + padding_length = target_length - x.shape[0] + if padding_length < 0: + raise ValueError( + f"Cannot pad a sequence of length {x.shape[0]} to {target_length}." + ) + if padding_length == 0: + return x + padding = x.new_zeros(padding_length, *x.shape[1:]) + return torch.cat((x, padding), dim=0) + + def _compute_learned_pos_embeds( pos_embed: torch.Tensor, grids: list[list[int]], @@ -80,11 +93,11 @@ def _compute_learned_pos_embeds( ) -> torch.Tensor: """Interpolate the learned 2D table and add fixed temporal embeddings.""" height, width, dim = pos_embed.shape - result = pos_embed.new_zeros(len(grids), max_num_patches, dim) pos_grid = pos_embed.permute(2, 0, 1).unsqueeze(0).float() cached_spatial: dict[tuple[int, int], torch.Tensor] = {} - for item_idx, (num_frames, grid_h, grid_w) in enumerate(grids): + padded_positions = [] + for num_frames, grid_h, grid_w in grids: if num_frames > max_num_frames: raise ValueError( f"Vision grid has {num_frames} frames, exceeding " @@ -114,9 +127,9 @@ def _compute_learned_pos_embeds( temporal = _get_temporal_pos_embed(num_frames, dim, device=pos_embed.device) item_pos = spatial.unsqueeze(0) + temporal.unsqueeze(1).to(spatial.dtype) item_pos = item_pos.reshape(num_frames * grid_h * grid_w, dim) - result[item_idx, : item_pos.shape[0]] = item_pos + padded_positions.append(_pad_sequence(item_pos, max_num_patches)) - return result + return torch.stack(padded_positions) def _compute_2d_rope_cache( @@ -126,15 +139,9 @@ def _compute_2d_rope_cache( head_dim: int, ) -> torch.Tensor: """Build the real-valued 2D RoPE cache in raster patch order.""" - angles = torch.zeros( - len(grids), - max_num_patches, - head_dim // 2, - dtype=freq_table.dtype, - device=freq_table.device, - ) cached_spatial: dict[tuple[int, int], torch.Tensor] = {} - for item_idx, (num_frames, grid_h, grid_w) in enumerate(grids): + padded_angles = [] + for num_frames, grid_h, grid_w in grids: spatial = cached_spatial.get((grid_h, grid_w)) if spatial is None: flat = torch.arange(grid_h * grid_w, device=freq_table.device) @@ -145,8 +152,9 @@ def _compute_2d_rope_cache( ) cached_spatial[(grid_h, grid_w)] = spatial item_angles = spatial.repeat(num_frames, 1) - angles[item_idx, : item_angles.shape[0]] = item_angles + padded_angles.append(_pad_sequence(item_angles, max_num_patches)) + angles = torch.stack(padded_angles) return torch.stack((angles.cos(), angles.sin()), dim=-1).unsqueeze(2) @@ -175,14 +183,14 @@ def _temporal_pool_and_merge( merge_kernel_size: tuple[int, int], ) -> torch.Tensor: """Temporally pool and concatenate neighboring spatial patch features.""" - num_items, _, dim = hidden_NPD.shape + _, _, dim = hidden_NPD.shape kernel_h, kernel_w = merge_kernel_size merged_dim = kernel_h * kernel_w * dim max_merged = max( (grid_h // kernel_h) * (grid_w // kernel_w) for _, grid_h, grid_w in grids ) - merged_NMK = hidden_NPD.new_zeros(num_items, max_merged, merged_dim) + padded_items = [] for item_idx, (num_frames, grid_h, grid_w) in enumerate(grids): merged_h = grid_h // kernel_h merged_w = grid_w // kernel_w @@ -196,9 +204,9 @@ def _temporal_pool_and_merge( ) item = item.permute(0, 1, 3, 2, 4, 5).mean(dim=0) item = item.reshape(merged_h * merged_w, merged_dim) - merged_NMK[item_idx, : item.shape[0]] = item + padded_items.append(_pad_sequence(item, max_merged)) - return merged_NMK + return torch.stack(padded_items) class VisionRotaryEmbedding2D(Module): @@ -292,7 +300,7 @@ def forward( ) q_NPHK, k_NPHK = _apply_2d_rope(q_NPHK, k_NPHK, rope_cache) - output_NPHK = torch.zeros_like(v_NPHK) + padded_outputs = [] for item_idx, item_length in enumerate(num_patches): q_HPK = q_NPHK[item_idx, :item_length].transpose(0, 1) k_HPK = k_NPHK[item_idx, :item_length].transpose(0, 1) @@ -303,8 +311,9 @@ def forward( q_HPK.dtype ) output_PHK = torch.matmul(probs_HPP, v_HPK).transpose(0, 1) - output_NPHK[item_idx, :item_length] = output_PHK + padded_outputs.append(_pad_sequence(output_PHK, max_num_patches)) + output_NPHK = torch.stack(padded_outputs) return self.proj(output_NPHK.flatten(start_dim=-2)) From a626834e06fb0c5c7928c4290a6e02b0e272afa1 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Thu, 30 Jul 2026 22:43:56 +0800 Subject: [PATCH 06/67] Keep unused Kimi experts in the FSDP autograd graph --- tests/unit_tests/test_kimi_k3.py | 25 +++++++++++++++++++++++++ torchtitan/models/kimi_k3/model.py | 4 ++-- 2 files changed, 27 insertions(+), 2 deletions(-) diff --git a/tests/unit_tests/test_kimi_k3.py b/tests/unit_tests/test_kimi_k3.py index 48df1ec0e4..0f52cab534 100644 --- a/tests/unit_tests/test_kimi_k3.py +++ b/tests/unit_tests/test_kimi_k3.py @@ -312,6 +312,31 @@ def test_kda_kernel_matches_recurrent_reference(self): self.assertIsNotNone(tensor.grad) self.assertTrue(torch.isfinite(tensor.grad).all()) + def test_unused_moe_experts_receive_zero_gradients(self): + torch.manual_seed(2) + model = _small_model_config().build() + model.init_states() + moe = model.layers["1"].moe + assert moe is not None + with torch.no_grad(): + moe.router.gate.weight.zero_() + + inputs = torch.randn(2, 4, 16, requires_grad=True) + expert_ids, _ = moe.router(inputs, moe.expert_bias_E) + selected_experts = set(expert_ids.flatten().tolist()) + unused_experts = set(range(moe.num_experts)) - selected_experts + self.assertTrue(unused_experts) + + moe(inputs).float().sum().backward() + for expert_idx in unused_experts: + for parameter in moe.routed_experts[expert_idx].parameters(): + self.assertIsNotNone(parameter.grad) + assert parameter.grad is not None + torch.testing.assert_close( + parameter.grad, + torch.zeros_like(parameter.grad), + ) + def test_small_multimodal_model_forward_backward_and_adapter(self): torch.manual_seed(2) config = _small_model_config() diff --git a/torchtitan/models/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py index 48a855806e..2a6bad3274 100644 --- a/torchtitan/models/kimi_k3/model.py +++ b/torchtitan/models/kimi_k3/model.py @@ -547,8 +547,8 @@ def forward(self, x_BLD: torch.Tensor) -> torch.Tensor: for expert_idx, expert in enumerate(self.routed_experts): token_and_slot = torch.nonzero(expert_ids_TK == expert_idx, as_tuple=False) - if token_and_slot.numel() == 0: - continue + # Keep empty experts in the autograd graph so every FSDP rank + # produces zero gradients instead of rank-dependent None gradients. token_ids = token_and_slot[:, 0] route_slots = token_and_slot[:, 1] expert_output = expert(latent_TD.index_select(0, token_ids)) From bedf14539cf2d760c3e94c5327303f87721d55e6 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Fri, 31 Jul 2026 16:18:11 +0800 Subject: [PATCH 07/67] Add Kimi K3 numerical and mixed-modality FSDP coverage --- tests/unit_tests/test_kimi_k3_fsdp.py | 94 ++++++++++++ tests/unit_tests/test_kimi_k3_hf_parity.py | 167 +++++++++++++++++++++ torchtitan/models/kimi_k3/README.md | 34 ++++- torchtitan/models/kimi_k3/model.py | 35 +++++ torchtitan/models/kimi_k3/parallelize.py | 8 +- 5 files changed, 332 insertions(+), 6 deletions(-) create mode 100644 tests/unit_tests/test_kimi_k3_hf_parity.py diff --git a/tests/unit_tests/test_kimi_k3_fsdp.py b/tests/unit_tests/test_kimi_k3_fsdp.py index abebd014aa..f64d8decd5 100644 --- a/tests/unit_tests/test_kimi_k3_fsdp.py +++ b/tests/unit_tests/test_kimi_k3_fsdp.py @@ -5,6 +5,7 @@ # LICENSE file in the root directory of this source tree. import copy +from contextlib import nullcontext from unittest.mock import patch import torch @@ -124,6 +125,99 @@ def test_single_rank_fsdp_matches_manual_bf16_reference(self): self.assertGreater(compared_gradients, 0) +class TestKimiK3MixedModalityFSDP(DTensorTestBase): + @property + def world_size(self): + return 2 + + @with_comms + def test_image_and_text_only_ranks_complete_forward_backward(self): + torch.manual_seed(3) + config = _small_model_config() + with torch.device("meta"): + model = config.build() + model.to_empty(device=self.device_type) + model.init_states() + + parallelism = ParallelismConfig( + data_parallel_shard_degree=self.world_size, + tensor_parallel_degree=1, + pipeline_parallel_degree=1, + context_parallel_degree=1, + expert_parallel_degree=1, + ) + parallel_dims = ParallelDims.from_config( + parallelism, + world_size=self.world_size, + ) + with patch( + "torchtitan.distributed.parallel_dims.device_type", + self.device_type, + ): + parallel_dims.build_mesh() + gradient_division_context = ( + patch("torchtitan.distributed.fsdp.disable_fsdp_gradient_division") + if self.device_type == "cpu" + else nullcontext() + ) + with gradient_division_context: + model = parallelize_kimi_k3( + model, + parallel_dims=parallel_dims, + training=TrainingConfig( + local_batch_size=1, + seq_len=6, + steps=1, + dtype="bfloat16", + ), + parallelism=parallelism, + compile_config=CompileConfig(), + ac_config=None, + dump_folder="", + ) + + rank = torch.distributed.get_rank() + if rank == 0: + inputs = { + "tokens": torch.tensor( + [[1, 7, 2, 3, 4, 5]], + dtype=torch.long, + device=self.device_type, + ), + "pixel_values": torch.randn( + 1, + 4, + 3 * 2 * 2, + device=self.device_type, + ), + "grid_thw": torch.tensor( + [[1, 2, 2]], + dtype=torch.long, + device=self.device_type, + ), + "special_tokens": {"image_id": 7}, + } + else: + inputs = { + "tokens": torch.tensor( + [[1, 2, 3, 4, 5, 6]], + dtype=torch.long, + device=self.device_type, + ) + } + + logits_BLV = model(**inputs) + logits_BLV.float().square().mean().backward() + + vision_gradients = 0 + for name, parameter in model.named_parameters(): + if not name.startswith("vision_encoder."): + continue + self.assertIsNotNone(parameter.grad, name) + vision_gradients += 1 + self.assertGreater(vision_gradients, 0) + + if __name__ == "__main__": from torch.testing._internal.common_utils import run_tests diff --git a/tests/unit_tests/test_kimi_k3_hf_parity.py b/tests/unit_tests/test_kimi_k3_hf_parity.py new file mode 100644 index 0000000000..4c298bdf46 --- /dev/null +++ b/tests/unit_tests/test_kimi_k3_hf_parity.py @@ -0,0 +1,167 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Reduced Kimi K3 numerical parity against the released HuggingFace model. + +The reference tensors were generated with the released ``modeling_kimi_k3.py`` +at moonshotai/Kimi-K3 commit c5d1dd4c428bd1ce8b88c5044f3b6ccde9e3b721. +The reduced model uses the same deterministic parameter and input construction +as this test, then maps its state dict through ``KimiK3StateDictAdapter`` and +loads the HuggingFace model strictly. The HuggingFace eager attention paths and +a pure PyTorch implementation of the released FLA KDA API produced the frozen +float32 values below. +""" + +import base64 +import unittest + +import torch + +from tests.unit_tests.test_kimi_k3 import _small_model_config + + +_TEXT_LOGITS_BASE64 = ( + "qvilPdqoTz1q2ZQ8yRJ4vMTJQ70Rj6C9X03YvVFcA75J6BS+9QggvohDJL5WaSG+" + "2ZkXvldBB75jKOK9Rg6svdUSXb3KArG8+II/PI8yNj3RVJo9O+nSPQkzAT6mXxM+" + "5jEfPlAnJD4rCSI+2u4YPtA8CT6FQOc9Xg+yPW1iaj2gtbM9feBrPUtVzDxJyg+8L" + "oMsvfvRlr314dC9ZPgAvobwE74FiCC+9zMmvrq1JL7NHRy+BssMvvbM7r2RuLm9N" + "0V5vdO06Lwwo6w7IJUePdtfkD3pO8s9IVX9PU5BEj4Hih8+H/IlPuQyJT6RVB0+/" + "a0OPrPB8z1Zpb89HEaDPbBPsz2Ke249vx/YPNAh2LtI7iC9K/SPvYw8yb0m2Pm9Rd" + "cPvhOPHL4mhyK+m30hvuZ9Gb5O4Aq+NYzsvVolub2Ohnu9S9XzvOAqTzsvQRM9Rpu" + "JPWmnwz1aRPU9SicOPqWLGz5zOyI+5+whPmOjGj5Yrww+dlXxPe/kvj3KOYQ9/Y/L" + "PY3JhD3NJuE8JBo1vO4sSb2WM669RU7wvYUGFL5QhCm+YrM3vi73Pb6TCjy+zAIy" + "volOIL4xsQe+wXTSvZp0jL3WzAC9M+fmOy5IOT1X26Y9OeHpPQ1pET7CnCc+w5Y2" + "PsKxPT5Znzw+WWszPlZ7Ij46igo+Z0DZPeYOlD0=" +) +_VISION_OUTPUT_BASE64 = ( + "zZB1v4gVxT0AaYs/uXavP7LLQT9v+7O+komdv+wxp7+XQAe/6k8YP4c7qj8+IJk/" + "3eWPPvtwUb+6DLG/EbyFvw==" +) +_MULTIMODAL_LOGITS_BASE64 = ( + "qvilPdqoTz1q2ZQ8yRJ4vMTJQ70Rj6C9X03YvVFcA75J6BS+9QggvohDJL5WaSG+" + "2ZkXvldBB75jKOK9Rg6svdUSXb3KArG8+II/PI8yNj3RVJo9O+nSPQkzAT6mXxM+" + "5jEfPlAnJD4rCSI+2u4YPtA8CT6FQOc9Xg+yPW1iaj2gtbM9feBrPUtVzDxJyg+8L" + "oMsvfvRlr314dC9ZPgAvobwE74FiCC+9zMmvrq1JL7NHRy+BssMvvbM7r2RuLm9N" + "0V5vdO06Lwwo6w7IJUePdtfkD3pO8s9IVX9PU5BEj4Hih8+H/IlPuQyJT6RVB0+/" + "a0OPrPB8z1Zpb89HEaDPWJ9mrsLzzW63QhcO40E7jsF4TE8vRRlPDI0hzz1CZY85G" + "eePMHxnzyGlpo8P5GOPGzNeDyMvkk8NP0RPNjhpzvLMZI6obhAu5kd4btQFSy8yjB" + "gvBFRhbxxypS85dmdvFEboLzVdZu8wRyQvIoafbxmEk+8Gx0YvGoytbtTasm6Marc" + "PQTmjD2lO9w89J6EvBnhb725jcm9Tz0JvjjJJ758GT++1ixOvuZcVL5jZVG+CmdF" + "vjjmML4uxRS+uHTkveuFlb0FSgC9Hf4/PHI3Xj18bcE9VrQFPmfuJD5TDD0++ANN" + "PiAlVD4cIVI+Kw5HPoJmMz4JAxg+ACTsPfoTnj3eefM9VzyZPQiP4TxLjKu8rVCM" + "vcOx5737ihy+hn0+vto5WL78o2i+wQZvvrUba74TDl2+7XhFvo1gJb45Tvy9+/qi" + "vT5CBb0RYII8gGmCPb6b3j2zmhg+alM7Ptr4VT73ZGc+eNduPjH+az6P+F4+S1ZI" + "PigRKT49ggI+LKasPYLq2z3AyI09JynmPN94abz7TGW9myzDvWukBb607yO+1Ck7" + "vnhSSr5ZwlC+azJOvvS+Qr5U5i6+koMTvsSJ471cPZa94u0EvYhnITxr5FM9QSa7" + "PX4kAj6oGSE+7xw5PmElST4HglA+p+FOPi1WRD4AVDE+46wWPuIN6z0NoJ49" +) + + +def _decode_float32(encoded: str, shape: tuple[int, ...]) -> torch.Tensor: + values = torch.frombuffer( + bytearray(base64.b64decode(encoded)), + dtype=torch.float32, + ) + return values.reshape(shape) + + +def _fill_reference_parameters(model: torch.nn.Module) -> None: + with torch.no_grad(): + for parameter_index, (name, parameter) in enumerate(model.named_parameters()): + values = torch.arange( + parameter.numel(), + dtype=torch.float32, + ).reshape(parameter.shape) + values = torch.sin(values * 0.013 + parameter_index * 0.17) * 0.02 + if name.endswith("norm.weight"): + values = values + 1.0 + parameter.copy_(values.to(parameter.dtype)) + + for module in model.modules(): + expert_bias = getattr(module, "expert_bias_E", None) + if expert_bias is not None: + expert_bias.copy_(torch.linspace(-0.01, 0.01, expert_bias.numel())) + + +class TestKimiK3HuggingFaceParity(unittest.TestCase): + @torch.no_grad() + def test_reduced_model_matches_released_huggingface_reference(self): + config = _small_model_config() + model = config.build() + model.init_states() + _fill_reference_parameters(model) + model.eval() + + router_ids: list[torch.Tensor] = [] + moe_layer = next( + layer for layer in model.layers.values() if layer.moe is not None + ) + router_hook = moe_layer.moe.router.register_forward_hook( + lambda _module, _inputs, output: router_ids.append( + output[1] if len(output) == 3 else output[0] + ) + ) + self.addCleanup(router_hook.remove) + + text_tokens_BL = torch.tensor([[1, 2, 3, 4]], dtype=torch.long) + text_logits_BLV = model(text_tokens_BL) + torch.testing.assert_close( + text_logits_BLV, + _decode_float32(_TEXT_LOGITS_BASE64, (1, 4, 32)), + atol=2e-4, + rtol=2e-4, + ) + torch.testing.assert_close( + router_ids.pop().reshape(-1), + torch.ones(4, dtype=torch.int64), + atol=0, + rtol=0, + ) + + patches_PCHW = torch.sin(torch.arange(48, dtype=torch.float32) * 0.07).reshape( + 4, 3, 2, 2 + ) + pixels_NPK = patches_PCHW.reshape(1, 4, 12) + grid_thw_N3 = torch.tensor([[1, 2, 2]], dtype=torch.long) + vision_output_NMD = model.vision_encoder( + pixels_NPK, + grid_thw=grid_thw_N3, + ) + torch.testing.assert_close( + vision_output_NMD, + _decode_float32(_VISION_OUTPUT_BASE64, (1, 1, 16)), + atol=2e-4, + rtol=2e-4, + ) + + multimodal_tokens_BL = torch.tensor( + [[1, 2, 7, 3, 4, 5]], + dtype=torch.long, + ) + multimodal_logits_BLV = model( + multimodal_tokens_BL, + pixel_values=pixels_NPK, + grid_thw=grid_thw_N3, + special_tokens={"image_id": 7}, + ) + torch.testing.assert_close( + multimodal_logits_BLV, + _decode_float32(_MULTIMODAL_LOGITS_BASE64, (1, 6, 32)), + atol=2e-4, + rtol=2e-4, + ) + torch.testing.assert_close( + router_ids.pop().reshape(-1), + torch.ones(6, dtype=torch.int64), + atol=0, + rtol=0, + ) + + +if __name__ == "__main__": + from torch.testing._internal.common_utils import run_tests + + run_tests() diff --git a/torchtitan/models/kimi_k3/README.md b/torchtitan/models/kimi_k3/README.md index fd352b5bff..21c3383d42 100644 --- a/torchtitan/models/kimi_k3/README.md +++ b/torchtitan/models/kimi_k3/README.md @@ -4,6 +4,9 @@ This directory contains the eager numerical reference implementation of Kimi K3 in TorchTitan. The initial scope is a topology-complete, reduced model for single-device or FSDP2 training and numerical comparison with the [released HuggingFace implementation](https://huggingface.co/moonshotai/Kimi-K3). +It is intended to make architecture experiments and model-structure choices +measurable against a stable, inspectable baseline before optimized kernels and +additional parallelisms are introduced. The implementation is device-neutral. It uses PyTorch operators and does not import accelerator-specific packages. The reference kernels prioritize @@ -73,6 +76,12 @@ The reference path mirrors the released implementation in these areas: should preserve its input/output contract and checkpoint schema. FSDP2 only shards parameters and leaves this eager forward contract unchanged. +When FSDP ranks contain a mixture of image and text-only batches, every rank +invokes the independently wrapped vision encoder once. Image ranks process +their real patches; text-only ranks process the minimum mergeable dummy grid +and attach its zero-valued result to the text embeddings. This keeps FSDP +all-gather and reduce-scatter ordering aligned without changing text logits. + ## Checkpoint conversion `KimiK3StateDictAdapter` converts between TorchTitan and an unquantized @@ -89,6 +98,14 @@ compressed tensors is outside this first change. Numerical comparison should therefore instantiate the same reduced, unquantized model on both sides and copy one state dict through the adapter. +`test_kimi_k3_hf_parity.py` freezes the float32 outputs from a deterministic +reduced model evaluated with the released HuggingFace code at commit +`c5d1dd4c428bd1ce8b88c5044f3b6ccde9e3b721`. The test covers text logits, +router choices, projected vision features, and end-to-end image-text logits. +The source model is loaded strictly from the state dict produced by +`KimiK3StateDictAdapter`; no full checkpoint or network access is required to +run the regression. + ## Tests The CPU unit tests cover: @@ -98,11 +115,17 @@ The CPU unit tests cover: - the KDA kernel against a direct recurrent formulation, including backward; - a small text+image model forward and backward; - exhaustive state-dict round-trip for that small model. +- reduced text, vision, router, and multimodal numerical parity against frozen + HuggingFace eager outputs; - single-rank FSDP2 forward and per-parameter gradient parity with a manually - cast BF16 reference. + cast BF16 reference; +- two-rank FSDP2 forward and backward when one rank has an image and the other + rank is text-only. ```bash -pytest -q tests/unit_tests/test_kimi_k3.py +pytest -q \ + tests/unit_tests/test_kimi_k3.py \ + tests/unit_tests/test_kimi_k3_hf_parity.py pytest -q tests/unit_tests/test_kimi_k3_fsdp.py ``` @@ -118,5 +141,8 @@ pytest -q tests/unit_tests/test_kimi_k3_fsdp.py - No full 2.8T flavor. These restrictions are explicit so unsupported runtime settings fail instead -of being silently ignored. EP and optimized kernels can be added in follow-up -changes after the eager/FSDP2 reference forward is numerically locked. +of being silently ignored. This first contribution deliberately limits +parallel execution to FSDP2: its purpose is to establish the eager numerical +reference used by Kimi K3 architecture experiments. TP, PP, CP, EP, and +optimized kernels can be added independently after that forward contract is +locked. diff --git a/torchtitan/models/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py index 2a6bad3274..5b8dac502e 100644 --- a/torchtitan/models/kimi_k3/model.py +++ b/torchtitan/models/kimi_k3/model.py @@ -752,6 +752,36 @@ def __init__(self, config: Config): config.vision_encoder.build() if config.vision_encoder is not None else None ) self.spatial_merge_size = config.spatial_merge_size + self._run_vision_encoder_on_text_only = False + + def enable_vision_encoder_on_text_only(self) -> None: + """Keep vision collectives aligned when FSDP ranks mix modalities.""" + self._run_vision_encoder_on_text_only = True + + def _get_dummy_vision_dependency( + self, + embeddings_BLD: torch.Tensor, + ) -> torch.Tensor: + assert self.vision_encoder is not None + kernel_h, kernel_w = self.vision_encoder.merge_kernel_size + patch_dim = self.vision_encoder.patch_embed.in_features + pixel_values_NPK = torch.zeros( + 1, + kernel_h * kernel_w, + patch_dim, + dtype=embeddings_BLD.dtype, + device=embeddings_BLD.device, + ) + grid_thw_N3 = torch.tensor( + [[1, kernel_h, kernel_w]], + dtype=torch.long, + device=embeddings_BLD.device, + ) + vision_embeds_NLD = self.vision_encoder( + pixel_values_NPK, + grid_thw=grid_thw_N3, + ) + return vision_embeds_NLD.sum().to(embeddings_BLD.dtype) * 0.0 def get_attention_masks(self, positions: torch.Tensor) -> AttentionMasksType | None: del positions @@ -772,6 +802,11 @@ def _prepare_multimodal_embeds( "both be omitted." ) if pixel_values is None: + if ( + self.vision_encoder is not None + and self._run_vision_encoder_on_text_only + ): + return embeddings + self._get_dummy_vision_dependency(embeddings) return embeddings assert grid_thw is not None if self.vision_encoder is None: diff --git a/torchtitan/models/kimi_k3/parallelize.py b/torchtitan/models/kimi_k3/parallelize.py index 7077009e30..aa60db8170 100644 --- a/torchtitan/models/kimi_k3/parallelize.py +++ b/torchtitan/models/kimi_k3/parallelize.py @@ -20,6 +20,7 @@ apply_fsdp_to_decoder, apply_fsdp_to_vision_encoder, ) +from torchtitan.models.kimi_k3.model import KimiK3Model def parallelize_kimi_k3( @@ -73,8 +74,11 @@ def parallelize_kimi_k3( ) dp_mesh = parallel_dims.get_mesh(dp_mesh_names) - vision_encoder = getattr(model, "vision_encoder", None) + assert isinstance(model, KimiK3Model) + vision_encoder = model.vision_encoder if vision_encoder is not None: + if dp_mesh.size() > 1: + model.enable_vision_encoder_on_text_only() apply_fsdp_to_vision_encoder( vision_encoder, dp_mesh, @@ -85,7 +89,7 @@ def parallelize_kimi_k3( ) apply_fsdp_to_decoder( - model, # pyrefly: ignore [bad-argument-type] + model, dp_mesh, param_dtype=TORCH_DTYPE_MAP[training.mixed_precision_param], reduce_dtype=TORCH_DTYPE_MAP[training.mixed_precision_reduce], From 6da10ec100565c08524ca056eabdb15b6d0a7561 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Fri, 31 Jul 2026 16:41:14 +0000 Subject: [PATCH 08/67] Reuse shared components in Kimi K3 and run KDA on the FLA kernel Addresses the review feedback on #4025. - KDA now dispatches to fla.ops.kda.chunk_kda with the gate activation, beta sigmoid, and q/k L2 norm fused into the kernel, following how Qwen3.5 uses FLA. The pure-PyTorch recurrence becomes ReferenceKimiKDAKernel in the unit tests, which the CPU suite builds the model with, and a CUDA-only test checks the kernel against it forward and backward for both gate activations. FLA cannot compile head dimensions below 16, so the config now rejects those with a clear error instead of a Triton compilation failure. - The vision encoder runs on every batch rather than only when images are present. It is its own FSDP unit, and the shared multimodal collator can hand one data-parallel rank a text-only batch, so conditional execution issued collectives on a subset of the process group and could deadlock the step. Batches without images use the smallest mergeable grid and contribute through add_zero_valued_dependency, which leaves the text embeddings numerically unchanged. This replaces the flag parallelize() used to set, so single-GPU and multi-GPU take the same forward path. - KimiMoERouter is replaced by the common TokenChoiceTopKRouter, which also removes a direct self.gate.weight read that would break under TP. - The private out-of-place vision scatter is dropped for the shared scatter_vision_embeds. FSDP2 only loses its pre-backward hook when a wrapped module returns a view, and Embedding returns a fresh tensor from F.embedding, so the fork was unnecessary. Its test now covers the shared helper instead. - tokens_per_expert_E is updated in place so the load-balancing hook keeps referring to the live buffer, and the unused q_lora_rank field is removed. Validated on 1x RTX 5080 with PyTorch 2.14.0.dev20260729+cu130 and fla-core 0.5.2: the frozen HuggingFace parity values are unchanged, the kimi_k3 tests pass (13, including the CUDA kernel comparison), and a 10-step debugmodel run tracks the previous losses to within 3e-3 with matching grad norms. Co-Authored-By: Claude Opus 5 --- tests/unit_tests/test_kimi_k3.py | 312 +++++++++++++------- torchtitan/distributed/fsdp.py | 25 ++ torchtitan/models/kimi_k3/README.md | 54 ++-- torchtitan/models/kimi_k3/__init__.py | 6 +- torchtitan/models/kimi_k3/model.py | 222 +++++--------- torchtitan/models/kimi_k3/parallelize.py | 2 - torchtitan/models/kimi_k3/requirements.txt | 1 + torchtitan/models/kimi_k3/vision_encoder.py | 7 +- 8 files changed, 356 insertions(+), 273 deletions(-) create mode 100644 torchtitan/models/kimi_k3/requirements.txt diff --git a/tests/unit_tests/test_kimi_k3.py b/tests/unit_tests/test_kimi_k3.py index 0f52cab534..e238f8e644 100644 --- a/tests/unit_tests/test_kimi_k3.py +++ b/tests/unit_tests/test_kimi_k3.py @@ -5,12 +5,14 @@ # LICENSE file in the root directory of this source tree. import unittest +from dataclasses import dataclass import torch import torch.nn as nn import torch.nn.functional as F from torchtitan.models.common import Embedding +from torchtitan.models.common.multimodal import scatter_vision_embeds from torchtitan.models.kimi_k3 import ( _feed_forward_config, _kda_config, @@ -22,13 +24,72 @@ kimi_k3_configs, ) from torchtitan.models.kimi_k3.model import ( - _replace_vision_embeds, KimiK3Model, KimiK3TransformerBlock, KimiKDAKernel, ) from torchtitan.models.kimi_k3.state_dict_adapter import KimiK3StateDictAdapter from torchtitan.models.kimi_k3.vision_encoder import KimiExactGELU +from torchtitan.protocols.module import Module + + +class ReferenceKimiKDAKernel(Module): + """Pure-PyTorch stand-in for KimiKDAKernel backed by an explicit recurrence. + + Mirrors ``KimiKDAKernel.forward``'s interface so tests can build a model + with it in place of the FLA kernel and exercise the surrounding eager model + on CPU. The loop is O(seqlen) and far too slow for training; it exists to + pin the kernel's math. + """ + + @dataclass(kw_only=True, slots=True) + class Config(Module.Config): + head_dim: int + lower_bound: float | None = -5.0 + + def __init__(self, config: Config): + super().__init__() + self.head_dim = config.head_dim + self.lower_bound = config.lower_bound + + def forward( + self, + q_BLHK: torch.Tensor, + k_BLHK: torch.Tensor, + v_BLHV: torch.Tensor, + gate_BLHK: torch.Tensor, + beta_BLH: torch.Tensor, + A_log_H: torch.Tensor, + dt_bias_HK: torch.Tensor, + ) -> torch.Tensor: + return _kda_recurrent_reference( + q_BLHK, + k_BLHK, + v_BLHV, + gate_BLHK, + beta_BLH, + A_log_H, + dt_bias_HK, + lower_bound=self.lower_bound, + ) + + +def _use_reference_kda_kernel(config: KimiK3Model.Config) -> KimiK3Model.Config: + """Point every KDA layer at the recurrent reference kernel. + + Test configurations use head dimensions far below what FLA's chunked KDA + kernel can compile, and the CPU suite has no Triton runtime at all. + """ + for layer in config.layers: + if layer.delta_attention is None: + continue + kernel = layer.delta_attention.kernel + assert isinstance(kernel, KimiKDAKernel.Config) + layer.delta_attention.kernel = ReferenceKimiKDAKernel.Config( + head_dim=kernel.head_dim, + lower_bound=kernel.lower_bound, + ) + return config def _small_model_config() -> KimiK3Model.Config: @@ -89,38 +150,40 @@ def block( ffn_res_proj=_linear(dim, 1), ) - return KimiK3Model.Config( - dim=dim, - vocab_size=32, - tok_embeddings=Embedding.Config( - num_embeddings=32, - embedding_dim=dim, - param_init={ - "weight": lambda parameter: nn.init.normal_(parameter, std=0.02) - }, - ), - layers=[ - block(0, use_mla=False, use_moe=False), - block(1, use_mla=True, use_moe=True), - ], - norm=_norm(dim), - lm_head=_linear(dim, 32), - output_res_norm=_norm(dim), - output_res_proj=_linear(dim, 1), - vision_encoder=_vision_encoder_config( - text_dim=dim, - dim=16, - qkv_dim=24, - hidden_dim=32, - num_layers=1, - num_heads=3, - patch_size=2, - merge_kernel_size=(2, 2), - init_pos_emb_height=2, - init_pos_emb_width=2, - max_num_frames=1, - ), - spatial_merge_size=2, + return _use_reference_kda_kernel( + KimiK3Model.Config( + dim=dim, + vocab_size=32, + tok_embeddings=Embedding.Config( + num_embeddings=32, + embedding_dim=dim, + param_init={ + "weight": lambda parameter: nn.init.normal_(parameter, std=0.02) + }, + ), + layers=[ + block(0, use_mla=False, use_moe=False), + block(1, use_mla=True, use_moe=True), + ], + norm=_norm(dim), + lm_head=_linear(dim, 32), + output_res_norm=_norm(dim), + output_res_proj=_linear(dim, 1), + vision_encoder=_vision_encoder_config( + text_dim=dim, + dim=16, + qkv_dim=24, + hidden_dim=32, + num_layers=1, + num_heads=3, + patch_size=2, + merge_kernel_size=(2, 2), + init_pos_emb_height=2, + init_pos_emb_width=2, + max_num_frames=1, + ), + spatial_merge_size=2, + ) ) @@ -133,49 +196,68 @@ def _kda_recurrent_reference( A_log_H: torch.Tensor, dt_bias_HK: torch.Tensor, *, - lower_bound: float, + lower_bound: float | None, ) -> torch.Tensor: + """Explicit KDA recurrence in FP32, matching the released Kimi K3 math. + + ``lower_bound`` selects the same two gate activations FLA exposes through + ``safe_gate``: the bounded ``lower_bound * sigmoid(...)`` form when set, + and ``-exp(A_log) * softplus(...)`` when ``None``. + """ + input_dtype = q_BLHK.dtype q_BLHK = q_BLHK.float() k_BLHK = k_BLHK.float() q_BLHK = q_BLHK * torch.rsqrt(q_BLHK.square().sum(dim=-1, keepdim=True) + 1e-6) k_BLHK = k_BLHK * torch.rsqrt(k_BLHK.square().sum(dim=-1, keepdim=True) + 1e-6) - log_decay_BLHK = lower_bound * torch.sigmoid( - torch.exp(A_log_H.float()).view(1, 1, -1, 1) - * (gate_BLHK.float() + dt_bias_HK.float()) - ) + v_BLHV = v_BLHV.float() + if lower_bound is None: + log_decay_BLHK = -torch.exp(A_log_H.float()).view(1, 1, -1, 1) * F.softplus( + gate_BLHK.float() + dt_bias_HK.float() + ) + else: + log_decay_BLHK = lower_bound * torch.sigmoid( + torch.exp(A_log_H.float()).view(1, 1, -1, 1) + * (gate_BLHK.float() + dt_bias_HK.float()) + ) decay_BLHK = torch.exp(log_decay_BLHK) beta_BLH = torch.sigmoid(beta_BLH.float()) B, L, H, K = q_BLHK.shape V = v_BLHV.shape[-1] - state_BHKV = torch.zeros(B, H, K, V) - output_BLHV = torch.empty(B, L, H, V) + state_BHKV = torch.zeros(B, H, K, V, device=q_BLHK.device) + outputs_BHV = [] for token_idx in range(L): state_BHKV = state_BHKV * decay_BLHK[:, token_idx].unsqueeze(-1) old_value_BHV = torch.matmul( k_BLHK[:, token_idx].unsqueeze(-2), state_BHKV, ).squeeze(-2) - delta_BHV = (v_BLHV[:, token_idx].float() - old_value_BHV) * beta_BLH[ + delta_BHV = (v_BLHV[:, token_idx] - old_value_BHV) * beta_BLH[ :, token_idx ].unsqueeze(-1) state_BHKV = state_BHKV + ( k_BLHK[:, token_idx].unsqueeze(-1) * delta_BHV.unsqueeze(-2) ) - output_BLHV[:, token_idx] = torch.matmul( - q_BLHK[:, token_idx].unsqueeze(-2), - state_BHKV, - ).squeeze(-2) * (K**-0.5) - return output_BLHV + outputs_BHV.append( + torch.matmul( + q_BLHK[:, token_idx].unsqueeze(-2), + state_BHKV, + ).squeeze(-2) + * (K**-0.5) + ) + return torch.stack(outputs_BHV, dim=1).to(input_dtype) class TestKimiK3(unittest.TestCase): - def test_replace_vision_embeds_is_out_of_place(self): - inputs_embeds = torch.arange(30.0).view(2, 5, 3).requires_grad_() + def test_scatter_vision_embeds_routes_gradients_to_both_streams(self): + # The scatter is in place, so the graph has to be built from tensors + # that are not themselves leaves requiring grad. + inputs_source = torch.arange(30.0).view(2, 5, 3).requires_grad_() vision_embeds = torch.arange(12.0).view(2, 2, 3).requires_grad_() - inputs_before = inputs_embeds.detach().clone() + inputs_embeds = inputs_source * 1.0 + inputs_before = inputs_source.detach().clone() - actual = _replace_vision_embeds( + actual = scatter_vision_embeds( inputs_embeds, vision_embeds=vision_embeds, vision_positions=[ @@ -203,21 +285,21 @@ def test_replace_vision_embeds_is_out_of_place(self): ) torch.testing.assert_close(actual, expected) - torch.testing.assert_close(inputs_embeds, inputs_before) - self.assertNotEqual(actual.data_ptr(), inputs_embeds.data_ptr()) actual.sum().backward() + # Overwritten positions must not propagate to the text embeddings, and + # every vision token that was scattered must receive gradient. expected_input_grad = torch.tensor( [[1, 0, 1, 1, 1], [1, 1, 0, 0, 1]], - dtype=inputs_embeds.dtype, + dtype=inputs_source.dtype, ).unsqueeze(-1) expected_vision_grad = torch.tensor( [[1, 0], [1, 1]], dtype=vision_embeds.dtype, ).unsqueeze(-1) torch.testing.assert_close( - inputs_embeds.grad, - expected_input_grad.expand_as(inputs_embeds), + inputs_source.grad, + expected_input_grad.expand_as(inputs_source), ) torch.testing.assert_close( vision_embeds.grad, @@ -259,58 +341,76 @@ def test_debugmodel_preserves_reduced_k3_topology(self): self.assertEqual(vision_config.dim, 256) self.assertEqual(vision_config.num_layers, 4) - def test_kda_kernel_matches_recurrent_reference(self): + @unittest.skipIf(not torch.cuda.is_available(), "FLA KDA kernel requires CUDA.") + def test_fla_kda_kernel_matches_recurrent_reference(self): torch.manual_seed(1) - q_BLHK = torch.randn(2, 5, 3, 4, requires_grad=True) - k_BLHK = torch.randn(2, 5, 3, 4, requires_grad=True) - v_BLHV = torch.randn(2, 5, 3, 4, requires_grad=True) - gate_BLHK = torch.randn(2, 5, 3, 4, requires_grad=True) - beta_BLH = torch.randn(2, 5, 3, requires_grad=True) - A_log_H = torch.randn(3, requires_grad=True) - dt_bias_HK = torch.randn(3, 4, requires_grad=True) - - kernel = KimiKDAKernel.Config( - head_dim=4, - lower_bound=-5.0, - ).build() - actual_BLHV = kernel( - q_BLHK, - k_BLHK, - v_BLHV, - gate_BLHK, - beta_BLH, - A_log_H, - dt_bias_HK, - ) - expected_BLHV = _kda_recurrent_reference( - q_BLHK, - k_BLHK, - v_BLHV, - gate_BLHK, - beta_BLH, - A_log_H, - dt_bias_HK, - lower_bound=-5.0, - ) + head_dim = 32 + num_heads = 3 + + def parameter(*shape: int) -> torch.Tensor: + return torch.randn(*shape, device="cuda", requires_grad=True) + + for lower_bound in (-5.0, None): + with self.subTest(lower_bound=lower_bound): + q_BLHK = parameter(2, 64, num_heads, head_dim) + k_BLHK = parameter(2, 64, num_heads, head_dim) + v_BLHV = parameter(2, 64, num_heads, head_dim) + gate_BLHK = parameter(2, 64, num_heads, head_dim) + beta_BLH = parameter(2, 64, num_heads) + A_log_H = torch.rand(num_heads, device="cuda") + A_log_H = A_log_H.uniform_(1.0, 16.0).log().requires_grad_() + dt_bias_HK = parameter(num_heads, head_dim) + + kernel = KimiKDAKernel.Config( + head_dim=head_dim, + lower_bound=lower_bound, + ).build() + actual_BLHV = kernel( + q_BLHK, + k_BLHK, + v_BLHV, + gate_BLHK, + beta_BLH, + A_log_H, + dt_bias_HK, + ) + expected_BLHV = _kda_recurrent_reference( + q_BLHK, + k_BLHK, + v_BLHV, + gate_BLHK, + beta_BLH, + A_log_H, + dt_bias_HK, + lower_bound=lower_bound, + ) - torch.testing.assert_close( - actual_BLHV, - expected_BLHV, - atol=1e-6, - rtol=1e-6, - ) - actual_BLHV.square().mean().backward() - for tensor in ( - q_BLHK, - k_BLHK, - v_BLHV, - gate_BLHK, - beta_BLH, - A_log_H, - dt_bias_HK, - ): - self.assertIsNotNone(tensor.grad) - self.assertTrue(torch.isfinite(tensor.grad).all()) + # The chunked kernel accumulates over chunk boundaries and uses + # reduced-precision matmuls internally, so it does not reproduce + # the sequential FP32 recurrence bit for bit. + torch.testing.assert_close( + actual_BLHV, + expected_BLHV, + atol=2e-3, + rtol=2e-3, + ) + actual_BLHV.square().mean().backward() + for tensor in ( + q_BLHK, + k_BLHK, + v_BLHV, + gate_BLHK, + beta_BLH, + A_log_H, + dt_bias_HK, + ): + self.assertIsNotNone(tensor.grad) + assert tensor.grad is not None + self.assertTrue(torch.isfinite(tensor.grad).all()) + + def test_kda_kernel_rejects_head_dim_below_kernel_minimum(self): + with self.assertRaisesRegex(ValueError, "head_dim must be at least"): + KimiKDAKernel.Config(head_dim=8, lower_bound=-5.0).build() def test_unused_moe_experts_receive_zero_gradients(self): torch.manual_seed(2) @@ -322,7 +422,7 @@ def test_unused_moe_experts_receive_zero_gradients(self): moe.router.gate.weight.zero_() inputs = torch.randn(2, 4, 16, requires_grad=True) - expert_ids, _ = moe.router(inputs, moe.expert_bias_E) + _, expert_ids, _ = moe.router(inputs, moe.expert_bias_E) selected_experts = set(expert_ids.flatten().tolist()) unused_experts = set(range(moe.num_experts)) - selected_experts self.assertTrue(unused_experts) diff --git a/torchtitan/distributed/fsdp.py b/torchtitan/distributed/fsdp.py index 1b00fa41d6..b2d4e41ba4 100644 --- a/torchtitan/distributed/fsdp.py +++ b/torchtitan/distributed/fsdp.py @@ -165,6 +165,31 @@ def apply_fsdp_to_vision_encoder( ) +def add_zero_valued_dependency( + output: torch.Tensor, + unused_output: torch.Tensor, +) -> torch.Tensor: + """Keep a conditionally executed FSDP module in the autograd graph. + + FSDP2 issues a module's all-gather from its pre-forward hook and its + reduce-scatter from the autograd hooks on that module's output. A module + that only some data-parallel ranks execute -- a VLM vision encoder on a + batch that happens to carry no images, for example -- would therefore + issue collectives on a subset of the process group and deadlock the step. + + A rank with no real work for such a module runs it on a placeholder input + and routes the result through this helper. Scaling by zero leaves + ``output`` numerically unchanged while preserving the graph edge, so every + rank issues the same collectives and the module receives zero gradients -- + which is also its correct contribution to the data-parallel average. + + Args: + output: the tensor the caller actually wants to return. + unused_output: a tensor produced by the module being kept alive. + """ + return output + unused_output.sum().to(output.dtype) * 0.0 + + def apply_fsdp_to_decoder( model: "Decoder", dp_mesh: DeviceMesh, diff --git a/torchtitan/models/kimi_k3/README.md b/torchtitan/models/kimi_k3/README.md index 21c3383d42..dbb0d6b270 100644 --- a/torchtitan/models/kimi_k3/README.md +++ b/torchtitan/models/kimi_k3/README.md @@ -8,9 +8,10 @@ It is intended to make architecture experiments and model-structure choices measurable against a stable, inspectable baseline before optimized kernels and additional parallelisms are introduced. -The implementation is device-neutral. It uses PyTorch operators and does not -import accelerator-specific packages. The reference kernels prioritize -inspectable model math over throughput. +Every operator outside the KDA recurrence is plain eager PyTorch, which keeps +the model math directly inspectable. KDA itself runs on FLA's chunked Triton +kernel, following the same split Qwen3.5 uses: the kernel is the training path +and a pure-PyTorch recurrence in the unit tests pins its numerics. ## Quick start @@ -25,7 +26,8 @@ NGPU=2 MODULE=kimi_k3 CONFIG=kimi_k3_debugmodel ./run_train.sh \ --parallelism.data_parallel_shard_degree 2 ``` -The multimodal data path requires `torchvision`. +Requirements beyond core TorchTitan are listed in `requirements.txt`: KDA needs +`flash-linear-attention` and the multimodal data path needs `torchvision`. ## Reduced model @@ -53,7 +55,9 @@ blocks and preserve the released model's 1-based full-attention cadence. The released vocabulary size is retained, following other TorchTitan multimodal debug models and making FSDP state sharding measurable while the decoder widths and depths remain reduced. The resulting model has about 100 -million parameters. +million parameters, of which roughly 84 million are the tied-vocabulary +embedding and output projection; the transformer itself is correspondingly +small. ## Forward structure @@ -72,15 +76,25 @@ The reference path mirrors the released implementation in these areas: PatchMergerMLPV2. - Vision features scattered into runs of media placeholder tokens. -`KimiKDAKernel` is the optimization boundary. A future accelerated backend -should preserve its input/output contract and checkpoint schema. FSDP2 only -shards parameters and leaves this eager forward contract unchanged. - -When FSDP ranks contain a mixture of image and text-only batches, every rank -invokes the independently wrapped vision encoder once. Image ranks process -their real patches; text-only ranks process the minimum mergeable dummy grid -and attach its zero-valued result to the text embeddings. This keeps FSDP -all-gather and reduce-scatter ordering aligned without changing text logits. +`KimiKDAKernel` is the kernel boundary. It dispatches to FLA's `chunk_kda` +with the gate activation, beta sigmoid, and query/key L2 norm fused in, so a +future backend can replace it while preserving the same input/output contract +and checkpoint schema. FLA's chunked kernel cannot compile head dimensions +below 16, and the reduced flavor uses 32. + +The vision encoder is its own FSDP unit, so its collectives only fire on ranks +that execute it. Because a data-parallel rank can legitimately receive a batch +with no images -- the shared multimodal collator emits `pixel_values=None` for +a text-only batch, and drops images to respect `max_images_per_batch` -- the +model runs the encoder on *every* batch. Batches without images use the +smallest grid the patch merger accepts and contribute its result through +`add_zero_valued_dependency`, which leaves the text embeddings numerically +unchanged while keeping the encoder in the autograd graph. Every rank +therefore issues the same all-gather and reduce-scatter regardless of what its +batch contains, and the encoder correctly receives zero gradients from +text-only ranks. Note that `torchtitan/models/qwen3_5` wraps its vision +encoder the same way but still calls it conditionally, so it retains this +hazard. ## Checkpoint conversion @@ -108,13 +122,12 @@ run the regression. ## Tests -The CPU unit tests cover: +The unit tests cover: - the reduced layer topology; - the explicit exact GELU against PyTorch's CPU reference; -- the KDA kernel against a direct recurrent formulation, including backward; - a small text+image model forward and backward; -- exhaustive state-dict round-trip for that small model. +- exhaustive state-dict round-trip for that small model; - reduced text, vision, router, and multimodal numerical parity against frozen HuggingFace eager outputs; - single-rank FSDP2 forward and per-parameter gradient parity with a manually @@ -122,6 +135,12 @@ The CPU unit tests cover: - two-rank FSDP2 forward and backward when one rank has an image and the other rank is text-only. +`ReferenceKimiKDAKernel` in `tests/unit_tests/test_kimi_k3.py` is the explicit +recurrent formulation of KDA. The tests above build the model with it in place +of the FLA kernel, which is what lets them run on CPU and at head dimensions +FLA cannot compile. A separate CUDA-only test checks the FLA kernel against +that same reference, forward and backward, for both gate activations. + ```bash pytest -q \ tests/unit_tests/test_kimi_k3.py \ @@ -136,7 +155,6 @@ pytest -q tests/unit_tests/test_kimi_k3_fsdp.py offload. - Image inputs are supported; video inputs are rejected. - No generation cache. -- No optimized KDA backend. - No MXFP4 checkpoint loading. - No full 2.8T flavor. diff --git a/torchtitan/models/kimi_k3/__init__.py b/torchtitan/models/kimi_k3/__init__.py index f93cffb18a..3b74ddc72e 100644 --- a/torchtitan/models/kimi_k3/__init__.py +++ b/torchtitan/models/kimi_k3/__init__.py @@ -14,6 +14,7 @@ from torchtitan.components.optimizer import register_moe_load_balancing_hook from torchtitan.models.common import Conv1d, Embedding, Linear +from torchtitan.models.common.moe import TokenChoiceTopKRouter from torchtitan.models.common.nn_modules import GELU, RMSNorm from torchtitan.models.common.vision_encoder import VisionMLP from torchtitan.models.utils import validate_converter_order @@ -28,7 +29,6 @@ KimiKDAKernel, KimiLatentMoE, KimiMLAAttention, - KimiMoERouter, KimiRMSNorm, KimiRMSNormGated, SituAndMul, @@ -152,7 +152,6 @@ def _mla_config( return KimiMLAAttention.Config( dim=dim, num_heads=num_heads, - q_lora_rank=q_lora_rank, kv_lora_rank=kv_lora_rank, qk_nope_head_dim=qk_nope_head_dim, qk_rope_head_dim=qk_rope_head_dim, @@ -240,10 +239,11 @@ def _latent_moe_config( ] return KimiLatentMoE.Config( num_experts=num_experts, - router=KimiMoERouter.Config( + router=TokenChoiceTopKRouter.Config( num_experts=num_experts, top_k=top_k, gate=_linear(dim, num_experts), + score_func="sigmoid", route_norm=True, route_scale=1.0, ), diff --git a/torchtitan/models/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py index 5b8dac502e..f10ed14938 100644 --- a/torchtitan/models/kimi_k3/model.py +++ b/torchtitan/models/kimi_k3/model.py @@ -6,24 +6,30 @@ """Kimi K3 language model components. -The first implementation intentionally provides a device-neutral reference -path. It mirrors the released HuggingFace model math, but does not depend on -FLA, Triton, CUDA, or a device-specific extension. The reference KDA kernel is -appropriate for numerical validation and reduced-model training; an optimized -backend can be added behind :class:`KimiKDAKernel` without changing the model -or checkpoint schema. +Every operator outside the KDA recurrence is plain eager PyTorch, so the model +mirrors the released HuggingFace math and stays directly inspectable. KDA runs +on FLA's chunked Triton kernel, which is what makes the model trainable at +speed; the pure-PyTorch recurrence it is checked against lives in +``tests/unit_tests/test_kimi_k3.py`` and is far too slow for training. """ from dataclasses import dataclass import torch import torch.nn.functional as F + +from fla.ops.kda import chunk_kda from torch import nn +from torchtitan.distributed.fsdp import add_zero_valued_dependency from torchtitan.models.common import Conv1d, Linear from torchtitan.models.common.attention import AttentionMasksType from torchtitan.models.common.decoder import Decoder -from torchtitan.models.common.multimodal import get_vision_positions +from torchtitan.models.common.moe import TokenChoiceTopKRouter +from torchtitan.models.common.multimodal import ( + get_vision_positions, + scatter_vision_embeds, +) from torchtitan.models.common.nn_modules import RMSNorm from torchtitan.models.kimi_k3.vision_encoder import KimiK3VisionEncoder from torchtitan.models.utils import get_moe_model_nparams_and_flops @@ -34,39 +40,20 @@ # K = key head dimension, V = value head dimension, E = experts, # T = flattened tokens, N = attention-residual entries. - -def _replace_vision_embeds( - inputs_embeds: torch.Tensor, - *, - vision_embeds: torch.Tensor, - vision_positions: list[tuple[int, int, int, int]], -) -> torch.Tensor: - """Return text embeddings with vision spans replaced out of place.""" - seq_len = inputs_embeds.shape[1] - flat_positions = [] - valid_vision_embeds = [] - for item_idx, sample_idx, vision_start, n_tokens in vision_positions: - flat_start = sample_idx * seq_len + vision_start - flat_positions.append( - torch.arange( - flat_start, - flat_start + n_tokens, - device=inputs_embeds.device, - ) - ) - valid_vision_embeds.append(vision_embeds[item_idx, :n_tokens]) - - indices = torch.cat(flat_positions) - replacements = torch.cat(valid_vision_embeds).to(inputs_embeds.dtype) - return ( - inputs_embeds.flatten(0, 1) - .index_copy(0, indices, replacements) - .view_as(inputs_embeds) - ) +# Below this head dimension FLA's chunked KDA kernel fails to compile its +# Triton block sizes; the failure is a compilation error deep in the kernel, +# so reject it where the configuration is built instead. +_MIN_KDA_HEAD_DIM = 16 class KimiRMSNorm(RMSNorm): - """RMSNorm with the explicit FP32 reduction used by the Kimi reference.""" + """RMSNorm that applies its weight after casting back to the input dtype. + + ``nn.RMSNorm`` scales by the weight while still in the reduction dtype, + which does not match the released Kimi implementation under BF16. Keeping + the subclass also exposes ``kimi_eps`` to ``_apply_attention_residual``, + which needs the epsilon to normalize residual entries by hand. + """ @dataclass(kw_only=True, slots=True) class Config(RMSNorm.Config): @@ -165,7 +152,6 @@ class KimiMLAAttention(Module): class Config(Module.Config): dim: int num_heads: int - q_lora_rank: int kv_lora_rank: int qk_nope_head_dim: int qk_rope_head_dim: int @@ -259,7 +245,15 @@ def forward( class KimiKDAKernel(Module): - """Differentiable, device-neutral recurrent KDA reference kernel.""" + """Stateless dispatch to FLA's chunked KDA kernel. + + The gate activation, the beta sigmoid, and the query/key L2 norm are all + fused into the kernel rather than materialized here, so the decay never + exists as a full ``(B, L, H, K)`` tensor. ``ReferenceKimiKDAKernel`` in + ``tests/unit_tests/test_kimi_k3.py`` implements the same interface as an + explicit recurrence and is the numerical baseline for this kernel; it also + lets the CPU test suite exercise the surrounding model without Triton. + """ @dataclass(kw_only=True, slots=True) class Config(Module.Config): @@ -272,6 +266,11 @@ def __init__(self, config: Config): self.lower_bound = config.lower_bound if self.lower_bound is not None and not (-5.0 <= self.lower_bound < 0.0): raise ValueError("KDA lower_bound must be in the safe range [-5, 0).") + if config.head_dim < _MIN_KDA_HEAD_DIM: + raise ValueError( + f"KDA head_dim must be at least {_MIN_KDA_HEAD_DIM} for the FLA " + f"chunked kernel, got {config.head_dim}." + ) def forward( self, @@ -288,46 +287,26 @@ def forward( f"KDA q/k shapes must match, got {q_BLHK.shape} and {k_BLHK.shape}." ) if q_BLHK.shape[:3] != v_BLHV.shape[:3]: - raise ValueError("KDA reference backend requires equal q/k/value heads.") - - input_dtype = q_BLHK.dtype - q_BLHK = q_BLHK.float() - k_BLHK = k_BLHK.float() - q_BLHK = q_BLHK * torch.rsqrt(q_BLHK.pow(2).sum(dim=-1, keepdim=True) + 1e-6) - k_BLHK = k_BLHK * torch.rsqrt(k_BLHK.pow(2).sum(dim=-1, keepdim=True) + 1e-6) - v_BLHV = v_BLHV.float() - - if self.lower_bound is None: - log_decay_BLHK = -torch.exp(A_log_H.float()).view(1, 1, -1, 1) * F.softplus( - gate_BLHK.float() + dt_bias_HK.float() - ) - else: - log_decay_BLHK = self.lower_bound * torch.sigmoid( - torch.exp(A_log_H.float()).view(1, 1, -1, 1) - * (gate_BLHK.float() + dt_bias_HK.float()) - ) - decay_BLHK = torch.exp(log_decay_BLHK) - beta_BLH = torch.sigmoid(beta_BLH.float()) - - B, L, H, K = q_BLHK.shape - V = v_BLHV.shape[-1] - state_BHKV = torch.zeros(B, H, K, V, dtype=torch.float32, device=q_BLHK.device) - outputs = [] - scale = K**-0.5 - for token_idx in range(L): - q_BHK = q_BLHK[:, token_idx] - k_BHK = k_BLHK[:, token_idx] - v_BHV = v_BLHV[:, token_idx] - beta_BH = beta_BLH[:, token_idx] - - state_BHKV = decay_BLHK[:, token_idx].unsqueeze(-1) * state_BHKV - old_value_BHV = torch.einsum("bhk,bhkv->bhv", k_BHK, state_BHKV) - delta_BHV = (v_BHV - old_value_BHV) * beta_BH.unsqueeze(-1) - state_BHKV = state_BHKV + torch.einsum("bhk,bhv->bhkv", k_BHK, delta_BHV) - output_BHV = torch.einsum("bhk,bhkv->bhv", q_BHK, state_BHKV) - outputs.append(output_BHV * scale) - - return torch.stack(outputs, dim=1).to(input_dtype) + raise ValueError("Kimi KDA requires equal q/k/value head counts.") + + # safe_gate selects the bounded gate activation + # lower_bound * sigmoid(exp(A_log) * (gate + dt_bias)); without it the + # kernel applies -exp(A_log) * softplus(gate + dt_bias). + out_BLHV, _ = chunk_kda( + q_BLHK, + k_BLHK, + v_BLHV, + gate_BLHK, + beta_BLH, + A_log=A_log_H, + dt_bias=dt_bias_HK.reshape(-1), + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=True, + safe_gate=self.lower_bound is not None, + lower_bound=self.lower_bound, + ) + return out_BLHV class KimiDeltaAttention(Module): @@ -349,7 +328,10 @@ class Config(Module.Config): forget_b: Linear.Config beta: Linear.Config output_gate: Linear.Config - kernel: KimiKDAKernel.Config + # Typed as the base config so the KDA kernel stays swappable, matching + # how VisionAttention types its inner_attention. Anything assigned here + # must accept KimiKDAKernel.forward's arguments. + kernel: Module.Config output_norm: KimiRMSNormGated.Config output_proj: Linear.Config @@ -423,43 +405,6 @@ def forward( return self.output_proj(out_BLHV.reshape(B, L, -1)) -class KimiMoERouter(Module): - """Sigmoid top-k router with auxiliary-loss-free correction bias.""" - - @dataclass(kw_only=True, slots=True) - class Config(Module.Config): - num_experts: int - top_k: int - gate: Linear.Config - route_norm: bool = True - route_scale: float = 1.0 - - def __init__(self, config: Config): - super().__init__() - if not 0 < config.top_k <= config.num_experts: - raise ValueError("MoE top_k must be in [1, num_experts].") - self.num_experts = config.num_experts - self.top_k = config.top_k - self.route_norm = config.route_norm - self.route_scale = config.route_scale - self.gate = config.gate.build() - - def forward( - self, - x_BLD: torch.Tensor, - expert_bias_E: torch.Tensor | None, - ) -> tuple[torch.Tensor, torch.Tensor]: - scores_BLE = torch.sigmoid(F.linear(x_BLD.float(), self.gate.weight.float())) - choice_BLE = ( - scores_BLE if expert_bias_E is None else scores_BLE + expert_bias_E.float() - ) - _, expert_ids_BLK = torch.topk(choice_BLE, k=self.top_k, dim=-1, sorted=False) - weights_BLK = scores_BLE.gather(-1, expert_ids_BLK) - if self.route_norm: - weights_BLK = weights_BLK / (weights_BLK.sum(dim=-1, keepdim=True) + 1e-20) - return expert_ids_BLK, weights_BLK * self.route_scale - - class KimiRoutedExperts(ModuleList): """List-backed experts implementing TorchTitan's FSDP expert protocol. @@ -485,7 +430,7 @@ class KimiLatentMoE(Module): @dataclass(kw_only=True, slots=True) class Config(Module.Config): num_experts: int - router: KimiMoERouter.Config + router: TokenChoiceTopKRouter.Config routed_down: Linear.Config routed_experts: list[KimiFeedForward.Config] routed_norm: KimiRMSNorm.Config @@ -526,7 +471,7 @@ def __init__(self, config: Config): ) def forward(self, x_BLD: torch.Tensor) -> torch.Tensor: - expert_ids_BLK, weights_BLK = self.router(x_BLD, self.expert_bias_E) + weights_BLK, expert_ids_BLK, _ = self.router(x_BLD, self.expert_bias_E) B, L, _ = x_BLD.shape routing_map_BLE = torch.zeros( B, @@ -536,9 +481,9 @@ def forward(self, x_BLD: torch.Tensor) -> torch.Tensor: device=x_BLD.device, ).scatter(-1, expert_ids_BLK, True) with torch.no_grad(): - self.tokens_per_expert_E = ( - self.tokens_per_expert_E + routing_map_BLE.sum(dim=(0, 1)).float() - ) + # In place so the load-balancing hook registered on the optimizer + # keeps referring to this buffer across steps. + self.tokens_per_expert_E.add_(routing_map_BLE.sum(dim=(0, 1)).float()) latent_TD = self.routed_down(x_BLD).reshape(B * L, -1) expert_ids_TK = expert_ids_BLK.reshape(B * L, -1) @@ -752,16 +697,17 @@ def __init__(self, config: Config): config.vision_encoder.build() if config.vision_encoder is not None else None ) self.spatial_merge_size = config.spatial_merge_size - self._run_vision_encoder_on_text_only = False - - def enable_vision_encoder_on_text_only(self) -> None: - """Keep vision collectives aligned when FSDP ranks mix modalities.""" - self._run_vision_encoder_on_text_only = True - def _get_dummy_vision_dependency( + def _encode_placeholder_image( self, embeddings_BLD: torch.Tensor, ) -> torch.Tensor: + """Run the vision encoder on the smallest grid it can merge. + + A batch without images must still drive the vision encoder, because it + is its own FSDP unit and the data-parallel ranks that do have images + will issue its collectives. See ``add_zero_valued_dependency``. + """ assert self.vision_encoder is not None kernel_h, kernel_w = self.vision_encoder.merge_kernel_size patch_dim = self.vision_encoder.patch_embed.in_features @@ -777,11 +723,7 @@ def _get_dummy_vision_dependency( dtype=torch.long, device=embeddings_BLD.device, ) - vision_embeds_NLD = self.vision_encoder( - pixel_values_NPK, - grid_thw=grid_thw_N3, - ) - return vision_embeds_NLD.sum().to(embeddings_BLD.dtype) * 0.0 + return self.vision_encoder(pixel_values_NPK, grid_thw=grid_thw_N3) def get_attention_masks(self, positions: torch.Tensor) -> AttentionMasksType | None: del positions @@ -802,12 +744,12 @@ def _prepare_multimodal_embeds( "both be omitted." ) if pixel_values is None: - if ( - self.vision_encoder is not None - and self._run_vision_encoder_on_text_only - ): - return embeddings + self._get_dummy_vision_dependency(embeddings) - return embeddings + if self.vision_encoder is None: + return embeddings + return add_zero_valued_dependency( + embeddings, + self._encode_placeholder_image(embeddings), + ) assert grid_thw is not None if self.vision_encoder is None: raise ValueError("pixel_values were provided without a vision encoder.") @@ -828,7 +770,7 @@ def _prepare_multimodal_embeds( raise ValueError( "pixel_values were provided but no image placeholder tokens were found." ) - return _replace_vision_embeds( + return scatter_vision_embeds( embeddings, vision_embeds=vision_embeds, vision_positions=vision_positions, diff --git a/torchtitan/models/kimi_k3/parallelize.py b/torchtitan/models/kimi_k3/parallelize.py index aa60db8170..c2c3d92485 100644 --- a/torchtitan/models/kimi_k3/parallelize.py +++ b/torchtitan/models/kimi_k3/parallelize.py @@ -77,8 +77,6 @@ def parallelize_kimi_k3( assert isinstance(model, KimiK3Model) vision_encoder = model.vision_encoder if vision_encoder is not None: - if dp_mesh.size() > 1: - model.enable_vision_encoder_on_text_only() apply_fsdp_to_vision_encoder( vision_encoder, dp_mesh, diff --git a/torchtitan/models/kimi_k3/requirements.txt b/torchtitan/models/kimi_k3/requirements.txt new file mode 100644 index 0000000000..1c4403544a --- /dev/null +++ b/torchtitan/models/kimi_k3/requirements.txt @@ -0,0 +1 @@ +../../../.ci/docker/requirements-vlm.txt diff --git a/torchtitan/models/kimi_k3/vision_encoder.py b/torchtitan/models/kimi_k3/vision_encoder.py index 2b0cdc88f7..7192350c98 100644 --- a/torchtitan/models/kimi_k3/vision_encoder.py +++ b/torchtitan/models/kimi_k3/vision_encoder.py @@ -6,10 +6,9 @@ """MoonViT3d vision encoder used by Kimi K3. -This module keeps the first Kimi K3 implementation device-neutral. Vision -attention is an eager PyTorch reference over each visual item, which preserves -the block-diagonal attention semantics of the HuggingFace implementation -without requiring FlashAttention or a device-specific kernel. +Vision attention is an eager PyTorch loop over each visual item, which +preserves the block-diagonal attention semantics of the HuggingFace +implementation without requiring FlashAttention or a device-specific kernel. Shape suffixes: - N = number of visual items From e87a87a9d8d4a118339e353aa668ca5401f09873 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Sat, 1 Aug 2026 08:03:14 +0000 Subject: [PATCH 09/67] Guard Kimi K3 optional deps and cover the residual pass-through path Follow-up review pass over the Kimi K3 change. - kimi_k3/model.py imports FLA at module scope, which is a per-model dependency rather than a core one. Without it the three Kimi test modules failed collection, and pytest treats a collection error as fatal, so `pytest tests/unit_tests` aborted entirely on a machine that had not installed it. Guard the imports and raise unittest.SkipTest instead, as test_qwen3_5_deltanet.py already does for the same reason. - Move the scatter_vision_embeds test to tests/unit_tests/test_multimodal.py. It covers a helper shared by three VLMs, so it does not belong behind the Kimi FLA guard, and models/common/multimodal.py had no tests of its own. - Cover the attention-residual pass-through path in the FSDP comparison. _small_model_config used attn_res_block_size=1, so every layer extended the residual and no layer passed it through. Layers that pass it through return it back out across the FSDP module boundary, and that routing is what keeps FSDP gradients bitwise equal to eager: returning None instead reassociates the residual's gradient accumulation and moves tok_embeddings.weight.grad by ~2e-3 relative, which measurably breaks the eager comparison. Run the comparison at attn_res_block_size=2 and record the reasoning where the tempting change would be made, since the arrangement exists only to satisfy a constraint the previous config could not observe. - Reject a spatial_merge_size that disagrees with the vision encoder's merge_kernel_size. The two are independent config fields that must match; a mismatch previously surfaced as a placeholder-run misalignment error that blamed the prompt rather than the configuration. - Correct the README: the token embedding and the output projection are separate parameters, not tied. Drop a stale "device-neutral" backend message and a note about qwen3_5 that will silently rot once that model is fixed. Verified with PyTorch 2.14.0.dev20260729+cu130 and fla-core 0.5.2: the frozen HuggingFace parity values are unchanged, the debugmodel reproduces 12.55150 / 12.33794 / 11.61755 over three deterministic steps, tests/unit_tests is 562 passed with the 18 failures all pre-existing missing-dependency ones, and pyrefly reports no errors over the changed files. Co-Authored-By: Claude Fable 5 --- tests/unit_tests/test_kimi_k3.py | 115 +++++++-------------- tests/unit_tests/test_kimi_k3_fsdp.py | 20 +++- tests/unit_tests/test_kimi_k3_hf_parity.py | 1 + tests/unit_tests/test_multimodal.py | 76 ++++++++++++++ torchtitan/models/kimi_k3/README.md | 15 ++- torchtitan/models/kimi_k3/__init__.py | 2 +- torchtitan/models/kimi_k3/model.py | 24 +++++ 7 files changed, 163 insertions(+), 90 deletions(-) create mode 100644 tests/unit_tests/test_multimodal.py diff --git a/tests/unit_tests/test_kimi_k3.py b/tests/unit_tests/test_kimi_k3.py index e238f8e644..693f13e588 100644 --- a/tests/unit_tests/test_kimi_k3.py +++ b/tests/unit_tests/test_kimi_k3.py @@ -12,26 +12,35 @@ import torch.nn.functional as F from torchtitan.models.common import Embedding -from torchtitan.models.common.multimodal import scatter_vision_embeds -from torchtitan.models.kimi_k3 import ( - _feed_forward_config, - _kda_config, - _latent_moe_config, - _linear, - _mla_config, - _norm, - _vision_encoder_config, - kimi_k3_configs, -) -from torchtitan.models.kimi_k3.model import ( - KimiK3Model, - KimiK3TransformerBlock, - KimiKDAKernel, -) -from torchtitan.models.kimi_k3.state_dict_adapter import KimiK3StateDictAdapter -from torchtitan.models.kimi_k3.vision_encoder import KimiExactGELU from torchtitan.protocols.module import Module +# torchtitan.models.kimi_k3 imports FLA at module scope for the KDA kernel. +# FLA is a per-model dependency (kimi_k3/requirements.txt), not part of the +# core requirements, so skip the Kimi suites instead of failing collection +# when it is absent. Modules importing this one inherit the skip. +try: + from torchtitan.models.kimi_k3 import ( + _feed_forward_config, + _kda_config, + _latent_moe_config, + _linear, + _mla_config, + _norm, + _vision_encoder_config, + kimi_k3_configs, + ) + from torchtitan.models.kimi_k3.model import ( + KimiK3Model, + KimiK3TransformerBlock, + KimiKDAKernel, + ) + from torchtitan.models.kimi_k3.state_dict_adapter import KimiK3StateDictAdapter + from torchtitan.models.kimi_k3.vision_encoder import KimiExactGELU +except ModuleNotFoundError as exc: + raise unittest.SkipTest( + f"Kimi K3 optional dependency unavailable: {exc.name}" + ) from exc + class ReferenceKimiKDAKernel(Module): """Pure-PyTorch stand-in for KimiKDAKernel backed by an explicit recurrence. @@ -92,7 +101,16 @@ def _use_reference_kda_kernel(config: KimiK3Model.Config) -> KimiK3Model.Config: return config -def _small_model_config() -> KimiK3Model.Config: +def _small_model_config(*, attn_res_block_size: int = 1) -> KimiK3Model.Config: + """Build the reduced two-layer model used across the Kimi K3 tests. + + ``attn_res_block_size`` defaults to 1, which makes every layer extend the + attention residual. Pass 2 to make the second layer pass the residual + through instead, which is the shape the released cadence uses and which + routes the residual back out through the FSDP module boundary. Callers + comparing against frozen reference values must keep the default, since the + parameter shapes and ordering feed those values. + """ dim = 16 def block( @@ -103,7 +121,7 @@ def block( ) -> KimiK3TransformerBlock.Config: return KimiK3TransformerBlock.Config( layer_id=layer_id, - attn_res_block_size=1, + attn_res_block_size=attn_res_block_size, attention=( _mla_config( dim=dim, @@ -249,63 +267,6 @@ def _kda_recurrent_reference( class TestKimiK3(unittest.TestCase): - def test_scatter_vision_embeds_routes_gradients_to_both_streams(self): - # The scatter is in place, so the graph has to be built from tensors - # that are not themselves leaves requiring grad. - inputs_source = torch.arange(30.0).view(2, 5, 3).requires_grad_() - vision_embeds = torch.arange(12.0).view(2, 2, 3).requires_grad_() - inputs_embeds = inputs_source * 1.0 - inputs_before = inputs_source.detach().clone() - - actual = scatter_vision_embeds( - inputs_embeds, - vision_embeds=vision_embeds, - vision_positions=[ - (0, 0, 1, 1), - (1, 1, 2, 2), - ], - ) - expected = torch.stack( - ( - torch.cat( - ( - inputs_before[0, :1], - vision_embeds.detach()[0, :1], - inputs_before[0, 2:], - ) - ), - torch.cat( - ( - inputs_before[1, :2], - vision_embeds.detach()[1], - inputs_before[1, 4:], - ) - ), - ) - ) - - torch.testing.assert_close(actual, expected) - - actual.sum().backward() - # Overwritten positions must not propagate to the text embeddings, and - # every vision token that was scattered must receive gradient. - expected_input_grad = torch.tensor( - [[1, 0, 1, 1, 1], [1, 1, 0, 0, 1]], - dtype=inputs_source.dtype, - ).unsqueeze(-1) - expected_vision_grad = torch.tensor( - [[1, 0], [1, 1]], - dtype=vision_embeds.dtype, - ).unsqueeze(-1) - torch.testing.assert_close( - inputs_source.grad, - expected_input_grad.expand_as(inputs_source), - ) - torch.testing.assert_close( - vision_embeds.grad, - expected_vision_grad.expand_as(vision_embeds), - ) - def test_exact_gelu_matches_pytorch_reference(self): x = torch.linspace(-4.0, 4.0, 257) actual = KimiExactGELU.Config().build()(x) diff --git a/tests/unit_tests/test_kimi_k3_fsdp.py b/tests/unit_tests/test_kimi_k3_fsdp.py index f64d8decd5..b464fc5fd4 100644 --- a/tests/unit_tests/test_kimi_k3_fsdp.py +++ b/tests/unit_tests/test_kimi_k3_fsdp.py @@ -5,6 +5,7 @@ # LICENSE file in the root directory of this source tree. import copy +import unittest from contextlib import nullcontext from unittest.mock import patch @@ -17,8 +18,16 @@ ) from torchtitan.config import CompileConfig, ParallelismConfig, TrainingConfig from torchtitan.distributed import ParallelDims -from torchtitan.models.kimi_k3 import parallelize_kimi_k3 -from torchtitan.models.kimi_k3.model import KimiK3Model + +# Skip instead of failing collection when FLA, a per-model dependency of +# kimi_k3, is not installed; see the matching guard in test_kimi_k3.py. +try: + from torchtitan.models.kimi_k3 import parallelize_kimi_k3 + from torchtitan.models.kimi_k3.model import KimiK3Model +except ModuleNotFoundError as exc: + raise unittest.SkipTest( + f"Kimi K3 optional dependency unavailable: {exc.name}" + ) from exc from tests.unit_tests.test_kimi_k3 import _small_model_config @@ -31,7 +40,12 @@ def world_size(self): @with_comms def test_single_rank_fsdp_matches_manual_bf16_reference(self): torch.manual_seed(3) - config = _small_model_config() + # attn_res_block_size=2 makes the second layer pass the attention + # residual through rather than extend it. That path routes a tensor + # back out through the FSDP module boundary, and its gradient + # accumulation order is what keeps FSDP bitwise equal to eager, so the + # comparison below has to cover it. + config = _small_model_config(attn_res_block_size=2) with torch.device("meta"): model = config.build() model.to_empty(device=self.device_type) diff --git a/tests/unit_tests/test_kimi_k3_hf_parity.py b/tests/unit_tests/test_kimi_k3_hf_parity.py index 4c298bdf46..f6d67c94c9 100644 --- a/tests/unit_tests/test_kimi_k3_hf_parity.py +++ b/tests/unit_tests/test_kimi_k3_hf_parity.py @@ -20,6 +20,7 @@ import torch +# Inherits test_kimi_k3's module-level skip when FLA is not installed. from tests.unit_tests.test_kimi_k3 import _small_model_config diff --git a/tests/unit_tests/test_multimodal.py b/tests/unit_tests/test_multimodal.py new file mode 100644 index 0000000000..b87079ea47 --- /dev/null +++ b/tests/unit_tests/test_multimodal.py @@ -0,0 +1,76 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Tests for the model-agnostic vision<->text fusion helpers.""" + +import unittest + +import torch + +from torchtitan.models.common.multimodal import scatter_vision_embeds + + +class TestScatterVisionEmbeds(unittest.TestCase): + def test_scatter_routes_gradients_to_both_streams(self): + # The scatter is in place, so the graph has to be built from tensors + # that are not themselves leaves requiring grad. + inputs_source = torch.arange(30.0).view(2, 5, 3).requires_grad_() + vision_embeds = torch.arange(12.0).view(2, 2, 3).requires_grad_() + inputs_embeds = inputs_source * 1.0 + inputs_before = inputs_source.detach().clone() + + actual = scatter_vision_embeds( + inputs_embeds, + vision_embeds=vision_embeds, + vision_positions=[ + (0, 0, 1, 1), + (1, 1, 2, 2), + ], + ) + expected = torch.stack( + ( + torch.cat( + ( + inputs_before[0, :1], + vision_embeds.detach()[0, :1], + inputs_before[0, 2:], + ) + ), + torch.cat( + ( + inputs_before[1, :2], + vision_embeds.detach()[1], + inputs_before[1, 4:], + ) + ), + ) + ) + + torch.testing.assert_close(actual, expected) + + actual.sum().backward() + # Overwritten positions must not propagate to the text embeddings, and + # every vision token that was scattered must receive gradient. + expected_input_grad = torch.tensor( + [[1, 0, 1, 1, 1], [1, 1, 0, 0, 1]], + dtype=inputs_source.dtype, + ).unsqueeze(-1) + expected_vision_grad = torch.tensor( + [[1, 0], [1, 1]], + dtype=vision_embeds.dtype, + ).unsqueeze(-1) + torch.testing.assert_close( + inputs_source.grad, + expected_input_grad.expand_as(inputs_source), + ) + torch.testing.assert_close( + vision_embeds.grad, + expected_vision_grad.expand_as(vision_embeds), + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/torchtitan/models/kimi_k3/README.md b/torchtitan/models/kimi_k3/README.md index dbb0d6b270..42198198e5 100644 --- a/torchtitan/models/kimi_k3/README.md +++ b/torchtitan/models/kimi_k3/README.md @@ -55,9 +55,9 @@ blocks and preserve the released model's 1-based full-attention cadence. The released vocabulary size is retained, following other TorchTitan multimodal debug models and making FSDP state sharding measurable while the decoder widths and depths remain reduced. The resulting model has about 100 -million parameters, of which roughly 84 million are the tied-vocabulary -embedding and output projection; the transformer itself is correspondingly -small. +million parameters, of which roughly 84 million are the token embedding and +the separate output projection over that vocabulary; the transformer itself is +correspondingly small. ## Forward structure @@ -92,9 +92,7 @@ smallest grid the patch merger accepts and contribute its result through unchanged while keeping the encoder in the autograd graph. Every rank therefore issues the same all-gather and reduce-scatter regardless of what its batch contains, and the encoder correctly receives zero gradients from -text-only ranks. Note that `torchtitan/models/qwen3_5` wraps its vision -encoder the same way but still calls it conditionally, so it retains this -hazard. +text-only ranks. ## Checkpoint conversion @@ -161,6 +159,5 @@ pytest -q tests/unit_tests/test_kimi_k3_fsdp.py These restrictions are explicit so unsupported runtime settings fail instead of being silently ignored. This first contribution deliberately limits parallel execution to FSDP2: its purpose is to establish the eager numerical -reference used by Kimi K3 architecture experiments. TP, PP, CP, EP, and -optimized kernels can be added independently after that forward contract is -locked. +reference used by Kimi K3 architecture experiments. TP, PP, CP, and EP can be +added independently after that forward contract is locked. diff --git a/torchtitan/models/kimi_k3/__init__.py b/torchtitan/models/kimi_k3/__init__.py index 3b74ddc72e..2206c14b35 100644 --- a/torchtitan/models/kimi_k3/__init__.py +++ b/torchtitan/models/kimi_k3/__init__.py @@ -346,7 +346,7 @@ def _vision_encoder_config( def _debugmodel(attn_backend: str) -> KimiK3Model.Config: if attn_backend != "eager": - raise ValueError("Kimi K3 v1 only provides the device-neutral 'eager' backend.") + raise ValueError("Kimi K3 v1 only provides the 'eager' backend.") dim = 256 vocab_size = 163840 diff --git a/torchtitan/models/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py index f10ed14938..f1bb87f67d 100644 --- a/torchtitan/models/kimi_k3/model.py +++ b/torchtitan/models/kimi_k3/model.py @@ -592,6 +592,14 @@ def forward( attention_masks: AttentionMasksType | None = None, positions: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: + # Blocks that do not extend the attention residual return it unchanged. + # FSDP2 aliases module inputs to drive its backward hooks, so passing it + # back out makes this FSDP unit return a view, which draws PyTorch's + # warning about in-place ops dropping the pre-backward hook. Nothing + # mutates it in place, and routing it back through the module boundary + # is what keeps FSDP gradients bitwise equal to eager -- returning None + # here instead reassociates the residual's gradient accumulation and + # perturbs tok_embeddings.weight.grad by ~2e-3 relative. B, L, D = x_BLD.shape prefix_sum_BLD: torch.Tensor | None = x_BLD @@ -697,6 +705,22 @@ def __init__(self, config: Config): config.vision_encoder.build() if config.vision_encoder is not None else None ) self.spatial_merge_size = config.spatial_merge_size + if self.vision_encoder is not None: + # The decoder sizes each image's placeholder run from + # spatial_merge_size while the encoder merges patches with + # merge_kernel_size. A mismatch surfaces much later as a + # placeholder-run misalignment that blames the prompt. + merge_kernel_size = self.vision_encoder.merge_kernel_size + if merge_kernel_size != ( + config.spatial_merge_size, + config.spatial_merge_size, + ): + raise ValueError( + f"spatial_merge_size {config.spatial_merge_size} does not " + f"match the vision encoder's merge_kernel_size " + f"{merge_kernel_size}; each image would occupy a different " + "number of text positions than the encoder produces." + ) def _encode_placeholder_image( self, From 6d2d190c6fb0bba5d6b22e4cac8d277236c03971 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Sat, 1 Aug 2026 08:58:33 +0000 Subject: [PATCH 10/67] Drop the KDA head-dim guard FLA's chunked KDA kernel fails below head_dim 16 because chunk_intra.py sets BK = next_power_of_2(K) without the floor of 16 that wy_fast.py applies, and triton's tl.dot needs a contraction dimension of at least 16. That is an unhandled small-head case in FLA rather than a constraint worth encoding here: released KDA head dimensions are far above it, the reduced flavor uses 32, and the guard would have to be revisited once FLA clamps the block size. Report it upstream instead of carrying a local check. Co-Authored-By: Claude Fable 5 --- tests/unit_tests/test_kimi_k3.py | 4 ---- torchtitan/models/kimi_k3/README.md | 3 +-- torchtitan/models/kimi_k3/model.py | 10 ---------- 3 files changed, 1 insertion(+), 16 deletions(-) diff --git a/tests/unit_tests/test_kimi_k3.py b/tests/unit_tests/test_kimi_k3.py index 693f13e588..d7a987c26d 100644 --- a/tests/unit_tests/test_kimi_k3.py +++ b/tests/unit_tests/test_kimi_k3.py @@ -369,10 +369,6 @@ def parameter(*shape: int) -> torch.Tensor: assert tensor.grad is not None self.assertTrue(torch.isfinite(tensor.grad).all()) - def test_kda_kernel_rejects_head_dim_below_kernel_minimum(self): - with self.assertRaisesRegex(ValueError, "head_dim must be at least"): - KimiKDAKernel.Config(head_dim=8, lower_bound=-5.0).build() - def test_unused_moe_experts_receive_zero_gradients(self): torch.manual_seed(2) model = _small_model_config().build() diff --git a/torchtitan/models/kimi_k3/README.md b/torchtitan/models/kimi_k3/README.md index 42198198e5..96f10b2950 100644 --- a/torchtitan/models/kimi_k3/README.md +++ b/torchtitan/models/kimi_k3/README.md @@ -79,8 +79,7 @@ The reference path mirrors the released implementation in these areas: `KimiKDAKernel` is the kernel boundary. It dispatches to FLA's `chunk_kda` with the gate activation, beta sigmoid, and query/key L2 norm fused in, so a future backend can replace it while preserving the same input/output contract -and checkpoint schema. FLA's chunked kernel cannot compile head dimensions -below 16, and the reduced flavor uses 32. +and checkpoint schema. The vision encoder is its own FSDP unit, so its collectives only fire on ranks that execute it. Because a data-parallel rank can legitimately receive a batch diff --git a/torchtitan/models/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py index f1bb87f67d..febc86f2f6 100644 --- a/torchtitan/models/kimi_k3/model.py +++ b/torchtitan/models/kimi_k3/model.py @@ -40,11 +40,6 @@ # K = key head dimension, V = value head dimension, E = experts, # T = flattened tokens, N = attention-residual entries. -# Below this head dimension FLA's chunked KDA kernel fails to compile its -# Triton block sizes; the failure is a compilation error deep in the kernel, -# so reject it where the configuration is built instead. -_MIN_KDA_HEAD_DIM = 16 - class KimiRMSNorm(RMSNorm): """RMSNorm that applies its weight after casting back to the input dtype. @@ -266,11 +261,6 @@ def __init__(self, config: Config): self.lower_bound = config.lower_bound if self.lower_bound is not None and not (-5.0 <= self.lower_bound < 0.0): raise ValueError("KDA lower_bound must be in the safe range [-5, 0).") - if config.head_dim < _MIN_KDA_HEAD_DIM: - raise ValueError( - f"KDA head_dim must be at least {_MIN_KDA_HEAD_DIM} for the FLA " - f"chunked kernel, got {config.head_dim}." - ) def forward( self, From f4b3808e665e38777115d943b8f447c97f1180f4 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Tue, 4 Aug 2026 11:55:10 +0000 Subject: [PATCH 11/67] Close the Kimi K3 backbone on global attention The released config lists full_attn_layers as [4, 8, ..., 88, 92, 93]: 92 and 93 are both global, so the backbone always ends on a full-attention layer. The debug model stopped at {4, 8, 12} over 13 layers and closed on KDA instead. That list is not expressible as "every n-th layer", so the topology is now assembled by _kimi_k3_config, which takes the 1-based indices verbatim. A future full-scale flavor can pass the released 24-entry list unchanged. 13 is also one past the attention-residual block size of 12, so the short trailing block the released 93-layer stack has is now covered too. Also records the float32 parity run at the released head dimensions, which needs two independent TF32 switches closed: the cuBLAS one that NGC containers enable through TORCH_ALLOW_TF32_CUBLAS_OVERRIDE, and TRITON_F32_DEFAULT, which the torch flag does not reach and which FLA's chunk_kda relies on. Co-Authored-By: Claude Opus 5 --- tests/unit_tests/test_kimi_k3.py | 2 +- torchtitan/models/kimi_k3/README.md | 130 ++++++++++++++++++++++++-- torchtitan/models/kimi_k3/__init__.py | 104 ++++++++++++++++----- 3 files changed, 200 insertions(+), 36 deletions(-) diff --git a/tests/unit_tests/test_kimi_k3.py b/tests/unit_tests/test_kimi_k3.py index d7a987c26d..936ebe6f04 100644 --- a/tests/unit_tests/test_kimi_k3.py +++ b/tests/unit_tests/test_kimi_k3.py @@ -288,7 +288,7 @@ def test_debugmodel_preserves_reduced_k3_topology(self): for layer_idx, layer in enumerate(config.layers) if layer.attention is not None ], - [4, 8, 12], + [4, 8, 12, 13], ) self.assertIsNotNone(config.layers[0].feed_forward) self.assertTrue(all(layer.moe is not None for layer in config.layers[1:])) diff --git a/torchtitan/models/kimi_k3/README.md b/torchtitan/models/kimi_k3/README.md index 96f10b2950..05e4597976 100644 --- a/torchtitan/models/kimi_k3/README.md +++ b/torchtitan/models/kimi_k3/README.md @@ -31,18 +31,22 @@ Requirements beyond core TorchTitan are listed in `requirements.txt`: KDA needs ## Reduced model -The `debugmodel` flavor preserves each distinct Kimi K3 forward path while -reducing widths, expert count, and depth. +`debugmodel` preserves each distinct Kimi K3 forward path while reducing +widths, head dimensions, expert count, and depth. | Component | Released Kimi K3 | `debugmodel` | |---|---:|---:| +| Parameters | 2.8T | 100M | | Decoder dimension | 7168 | 256 | | Vocabulary size | 163840 | 163840 | | Decoder layers | 93 | 13 | -| Full MLA layers (1-based) | 4, 8, ..., 92, 93 | 4, 8, 12 | -| KDA layers | 69 | 10 | +| Full MLA layers (1-based) | 4, 8, ..., 92, 93 | 4, 8, 12, 13 | +| KDA layers | 69 | 9 | | Dense FFN layers | 1 | 1 | | Attention residual block size | 12 | 12 | +| MLA heads | 96 | 4 | +| MLA qk_nope / qk_rope / v head dimension | 128 / 64 / 128 | 32 / 16 / 32 | +| KDA heads / head dimension | 96 / 128 | 4 / 32 | | Routed experts / top-k | 896 / 16 | 8 / 2 | | Routed latent / expert hidden dimension | 3584 / 3072 | 128 / 128 | | Shared experts | 2 | 2 | @@ -50,14 +54,20 @@ reducing widths, expert count, and depth. | Vision layers | 27 | 4 | | Vision QKV dimension / heads | 1536 / 12 | 384 / 3 | -Thirteen decoder layers are intentional. They exercise two attention-residual -blocks and preserve the released model's 1-based full-attention cadence. +The depth sits one layer past a multiple of the full-attention period and of +the attention-residual block size, which reproduces two structural edge cases +of the released 93-layer stack: a final MLA layer immediately after a scheduled +one -- the released `full_attn_layers` ends `..., 88, 92, 93`, so the backbone +always closes on global attention -- and a short trailing residual block. +Neither is expressible as "every n-th layer", so `full_attention_layers` takes +the 1-based list verbatim and a future full-scale flavor can pass the released +24-entry list unchanged. The released vocabulary size is retained, following other TorchTitan multimodal debug models and making FSDP state sharding measurable while the -decoder widths and depths remain reduced. The resulting model has about 100 -million parameters, of which roughly 84 million are the token embedding and -the separate output projection over that vocabulary; the transformer itself is -correspondingly small. +decoder widths and depths remain reduced. In `debugmodel` roughly 84 million +of the 100 million parameters are the token embedding and the separate output +projection over that vocabulary; the transformer itself is correspondingly +small. ## Forward structure @@ -117,6 +127,106 @@ The source model is loaded strictly from the state dict produced by `KimiK3StateDictAdapter`; no full checkpoint or network access is required to run the regression. +### Parity at released head dimensions + +`debugmodel` shrinks the head dimensions, so it does not pin the FLA kernel +configuration the full model runs. That was checked separately, out of band, on +a 1.07B configuration built from the same `_kimi_k3_config` builder with the +released head dimensions kept intact: dimension 1024, 25 layers, full attention +at `4, 8, ..., 24, 25`, MLA `qk_nope / qk_rope / v` = `128 / 64 / 128` with 8 +heads, KDA 8 heads of 128, 24 routed experts at top-4, vision dimension 512 +over 8 layers. Same released commit, one RTX 5080, float32, 128 tokens. Both +sides run FLA's `chunk_kda`; the released code path for MLA and vision +attention was forced to eager, since flash-attn is not required here. + +Float32 comparison requires closing **two** independent TF32 switches: + +- **cuBLAS**, via `torch.backends.cuda.matmul.allow_tf32 = False`. Torch's own + default is already off, but NVIDIA's NGC containers set + `TORCH_ALLOW_TF32_CUBLAS_OVERRIDE=1`, which flips + `float32_matmul_precision` to `HIGH` at process init. +- **Triton**, via `TRITON_F32_DEFAULT=ieee`. The torch flag does not reach + Triton kernels, and FLA's `chunk_kda` leans on the Triton default for the + triangular solves in `fla/ops/kda/chunk_intra.py`. + +| Quantity | max abs diff | reference max abs | +|---|---:|---:| +| Text logits, 128 tokens | 2.2e-5 | 5.8 | +| Projected vision features, 8x8 patch grid | 4.0e-5 | 4.4 | +| Multimodal logits, 128 tokens | 2.4e-4 | 5.8 | + +Text logits differ by 3.9e-6 relative. For roughly 250 sequential dependent +matmuls that is well inside the linear accumulation bound of 1.5e-5 implied by +float32's 6.0e-8 unit roundoff. Per-layer relative drift grows from 4.8e-7 +after layer 0 to 3.4e-6 by layer 23, and the per-layer amplification factor +stays in 0.96-1.44 throughout: no layer amplifies, the error only accumulates. +Routed expert IDs are identical for all 3072 token-routings (24 MoE layers x +128 tokens), and the argmax over the vocabulary agrees on every position. + +Leaving Triton at its `tf32` default costs a factor of three end to end +(text logits 2.2e-5 -> 7.4e-5) and shows up inside KDA as a 50x step: every +tensor entering `chunk_kda` matches to ~1e-5 relative, while its output matches +only to ~5e-4. Both sides call the same kernel, so this never broke parity -- +TF32's error is 99.6% correlated between the two runs and mostly cancels. What +does not cancel is the discontinuity: where the two nearly-identical inputs +land on opposite sides of a rounding boundary, the entry jumps a full TF32 +quantum, 50x larger than the input difference that caused it. Under IEEE the +step disappears and `chunk_kda`'s output matches to ~6e-6. + +With both switches closed, the eager operators and their FLA counterparts are +numerically indistinguishable: swapping TorchTitan's `Conv1d`+SiLU and eager +gated RMSNorm for FLA's `ShortConvolution` and `FusedRMSNormGated` moves the +logits difference from 2.2e-5 to 2.4e-5, and the two variants differ from each +other by 1.3e-5. None of the three is a privileged reference; what remains is +ordinary reassociation noise, including the routed-expert summation order. + +The eager operators are kept deliberately. FLA's `ShortConvolution` and +`FusedRMSNormGated` are Triton-only and raise on CPU tensors; using them here +would make the whole model forward GPU-only and would leave +`test_kimi_k3_hf_parity.py` -- the frozen, network-free, CPU regression -- with +no way to run. They are also not a checkpoint concern either way: +`ShortConvolution` subclasses `nn.Conv1d` with the same `(D, 1, K)` weight, and +`FusedRMSNormGated` carries the same `(D,)` weight, so the state-dict mapping is +unaffected by the choice. + +### bfloat16 + +Comparing the two implementations directly in bfloat16 measures top-k router +ties, not arithmetic. Each side must be compared against its own float32 run. + +| bfloat16 vs own float32 | routings changed | logits max abs | argmax kept | +|---|---:|---:|---:| +| TorchTitan | 808 / 3072 | 5.69 | 50.0% | +| Released HuggingFace | 921 / 3072 | 5.96 | 47.7% | + +Both implementations lose their own float32 result at the same rate, so the +instability is the model configuration rather than either implementation. It +comes from the router: a randomly initialized gate produces near-tied scores, +bfloat16 rounding reorders them, and a changed expert changes the residual +stream enough to change every later routing decision. Rounding only the gate +of the first MoE layer to bfloat16, with a float32 input, already flips 4 of +128 tokens; those tokens have a median top-4/top-5 score margin of 5.4e-4 +against a typical score spread of 5.5e-1. + +Routing every token to all 24 experts removes the ties and leaves only the +arithmetic. Nothing then flips on either side, and the two implementations +degrade identically: + +| bfloat16 vs own float32, all experts routed | logits max abs | logits mean abs | argmax kept | +|---|---:|---:|---:| +| TorchTitan | 1.72e-1 | 1.66e-2 | 97.7% | +| Released HuggingFace | 1.88e-1 | 1.68e-2 | 94.5% | + +Per-layer bfloat16 error grows from 0.8% of the activation scale after layer 0 +to roughly 3% by the end of an attention-residual block on both sides, with +neither consistently ahead. That is ordinary bfloat16 accumulation over 25 +layers, given a 0.4% unit roundoff. + +The practical consequence: validate numerics in float32, as +`.claude/rules` already requires, and do not read a bfloat16 loss difference +against a reference implementation as evidence of a bug until the routing +decisions have been checked. + ## Tests The unit tests cover: diff --git a/torchtitan/models/kimi_k3/__init__.py b/torchtitan/models/kimi_k3/__init__.py index 2206c14b35..5c25c4d8f7 100644 --- a/torchtitan/models/kimi_k3/__init__.py +++ b/torchtitan/models/kimi_k3/__init__.py @@ -344,32 +344,49 @@ def _vision_encoder_config( ) -def _debugmodel(attn_backend: str) -> KimiK3Model.Config: - if attn_backend != "eager": - raise ValueError("Kimi K3 v1 only provides the 'eager' backend.") - - dim = 256 - vocab_size = 163840 - num_layers = 13 - full_attention_layers = {4, 8, 12} - num_heads = 4 - qk_nope_head_dim = 32 - qk_rope_head_dim = 16 - v_head_dim = 32 - +def _kimi_k3_config( + *, + dim: int, + vocab_size: int, + num_layers: int, + full_attention_layers: set[int], + attn_res_block_size: int, + num_heads: int, + q_lora_rank: int, + kv_lora_rank: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + v_head_dim: int, + kda_head_dim: int, + conv_kernel_size: int, + dense_hidden_dim: int, + latent_dim: int, + expert_hidden_dim: int, + num_experts: int, + top_k: int, + num_shared_experts: int, + vision_encoder: KimiK3VisionEncoder.Config, +) -> KimiK3Model.Config: + """Assemble a Kimi K3 config from the released topology's free parameters. + + ``full_attention_layers`` holds 1-based layer indices, matching the + released ``linear_attn_config.full_attn_layers``. Every other layer is KDA. + Layer 0 is the single dense FFN layer (released + ``first_k_dense_replace=1``); the rest are LatentMoE. + """ layers = [] for layer_idx in range(num_layers): is_full_attention = (layer_idx + 1) in full_attention_layers layers.append( KimiK3TransformerBlock.Config( layer_id=layer_idx, - attn_res_block_size=12, + attn_res_block_size=attn_res_block_size, attention=( _mla_config( dim=dim, num_heads=num_heads, - q_lora_rank=128, - kv_lora_rank=64, + q_lora_rank=q_lora_rank, + kv_lora_rank=kv_lora_rank, qk_nope_head_dim=qk_nope_head_dim, qk_rope_head_dim=qk_rope_head_dim, v_head_dim=v_head_dim, @@ -383,12 +400,12 @@ def _debugmodel(attn_backend: str) -> KimiK3Model.Config: else _kda_config( dim=dim, num_heads=num_heads, - head_dim=32, - conv_kernel_size=4, + head_dim=kda_head_dim, + conv_kernel_size=conv_kernel_size, ) ), feed_forward=( - _feed_forward_config(dim=dim, hidden_dim=1024) + _feed_forward_config(dim=dim, hidden_dim=dense_hidden_dim) if layer_idx == 0 else None ), @@ -397,11 +414,11 @@ def _debugmodel(attn_backend: str) -> KimiK3Model.Config: if layer_idx == 0 else _latent_moe_config( dim=dim, - latent_dim=128, - expert_hidden_dim=128, - num_experts=8, - top_k=2, - num_shared_experts=2, + latent_dim=latent_dim, + expert_hidden_dim=expert_hidden_dim, + num_experts=num_experts, + top_k=top_k, + num_shared_experts=num_shared_experts, ) ), attention_norm=_norm(dim), @@ -430,6 +447,44 @@ def _debugmodel(attn_backend: str) -> KimiK3Model.Config: ), output_res_norm=_norm(dim), output_res_proj=_linear(dim, 1), + vision_encoder=vision_encoder, + spatial_merge_size=2, + ) + + +def _debugmodel(attn_backend: str) -> KimiK3Model.Config: + """Return the topology-complete Kimi K3 debug model. + + The depth is one past a multiple of both the full-attention period and the + attention-residual block size, so the last layer is a full-attention layer + directly after a scheduled one and the trailing residual block is short. + Both are properties of the released 93-layer stack, whose + ``full_attn_layers`` ends ``..., 88, 92, 93``. + """ + if attn_backend != "eager": + raise ValueError("Kimi K3 v1 only provides the 'eager' backend.") + + dim = 256 + return _kimi_k3_config( + dim=dim, + vocab_size=163840, + num_layers=13, + full_attention_layers={4, 8, 12, 13}, + attn_res_block_size=12, + num_heads=4, + q_lora_rank=128, + kv_lora_rank=64, + qk_nope_head_dim=32, + qk_rope_head_dim=16, + v_head_dim=32, + kda_head_dim=32, + conv_kernel_size=4, + dense_hidden_dim=1024, + latent_dim=128, + expert_hidden_dim=128, + num_experts=8, + top_k=2, + num_shared_experts=2, vision_encoder=_vision_encoder_config( text_dim=dim, dim=256, @@ -438,7 +493,6 @@ def _debugmodel(attn_backend: str) -> KimiK3Model.Config: num_layers=4, num_heads=3, ), - spatial_merge_size=2, ) From 25577e4974e8fb18ea64e3adad89496c32ca5168 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Tue, 4 Aug 2026 11:55:10 +0000 Subject: [PATCH 12/67] Note the context-parallel route in add_zero_valued_dependency The docstring only described the data-parallel trigger, where a rank's batch carries no images. Context parallelism reaches the same subset-collective deadlock by a second route: a rank's sequence shard can hold zero vision placeholders even when every rank received images. Reported by QIU023 from running this architecture at cp>1. Co-Authored-By: Claude Opus 5 --- torchtitan/distributed/fsdp.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/torchtitan/distributed/fsdp.py b/torchtitan/distributed/fsdp.py index b2d4e41ba4..f1aa8fd8c3 100644 --- a/torchtitan/distributed/fsdp.py +++ b/torchtitan/distributed/fsdp.py @@ -177,6 +177,12 @@ def add_zero_valued_dependency( batch that happens to carry no images, for example -- would therefore issue collectives on a subset of the process group and deadlock the step. + Context parallelism reaches the same deadlock by a second route: a rank's + sequence shard can hold zero vision placeholders even when every rank + received images, so "this batch has images" is not a sufficient condition + for running the module. Decide from the process group, not from the local + batch. + A rank with no real work for such a module runs it on a placeholder input and routes the result through this helper. Scaling by zero leaves ``output`` numerically unchanged while preserving the graph edge, so every From 4f7f668eb263c1868787307ece3181f0da5f45cf Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Sun, 9 Aug 2026 07:06:09 +0000 Subject: [PATCH 13/67] reuse common FeedForward --- torchtitan/models/kimi_k3/__init__.py | 4 +-- torchtitan/models/kimi_k3/model.py | 40 +++++++-------------------- 2 files changed, 12 insertions(+), 32 deletions(-) diff --git a/torchtitan/models/kimi_k3/__init__.py b/torchtitan/models/kimi_k3/__init__.py index 5c25c4d8f7..35b4a36e00 100644 --- a/torchtitan/models/kimi_k3/__init__.py +++ b/torchtitan/models/kimi_k3/__init__.py @@ -31,7 +31,6 @@ KimiMLAAttention, KimiRMSNorm, KimiRMSNormGated, - SituAndMul, ) from .parallelize import parallelize_kimi_k3 from .state_dict_adapter import KimiK3StateDictAdapter @@ -134,7 +133,8 @@ def _feed_forward_config( w1=_linear(dim, hidden_dim), w2=_linear(hidden_dim, dim), w3=_linear(dim, hidden_dim), - activation=SituAndMul.Config(beta=4.0, linear_beta=25.0), + beta=4.0, + linear_beta=25.0, ) diff --git a/torchtitan/models/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py index febc86f2f6..be68c41c65 100644 --- a/torchtitan/models/kimi_k3/model.py +++ b/torchtitan/models/kimi_k3/model.py @@ -25,6 +25,7 @@ from torchtitan.models.common import Conv1d, Linear from torchtitan.models.common.attention import AttentionMasksType from torchtitan.models.common.decoder import Decoder +from torchtitan.models.common.feed_forward import FeedForward from torchtitan.models.common.moe import TokenChoiceTopKRouter from torchtitan.models.common.multimodal import ( get_vision_positions, @@ -89,50 +90,29 @@ def forward(self, x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: return (x_float * torch.sigmoid(gate.float())).to(input_dtype) -class SituAndMul(Module): - """Kimi's SiTU activation applied to concatenated gate/up projections.""" +class KimiFeedForward(FeedForward): + """FeedForward with Kimi's SiTU activation""" @dataclass(kw_only=True, slots=True) - class Config(Module.Config): + class Config(FeedForward.Config): beta: float = 1.0 linear_beta: float | None = None def __init__(self, config: Config): - super().__init__() + super().__init__(config) self.beta = config.beta self.linear_beta = config.linear_beta - def forward(self, gate_up: torch.Tensor) -> torch.Tensor: - gate, up = gate_up.chunk(2, dim=-1) - input_dtype = gate_up.dtype + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate = self.w1(x) + up = self.w3(x) + input_dtype = gate.dtype gate = gate.float() up = up.float() gate = self.beta * torch.tanh(gate / self.beta) * torch.sigmoid(gate) if self.linear_beta is not None: up = self.linear_beta * torch.tanh(up / self.linear_beta) - return (gate * up).to(input_dtype) - - -class KimiFeedForward(Module): - """Three-projection feed-forward network using SiTU.""" - - @dataclass(kw_only=True, slots=True) - class Config(Module.Config): - w1: Linear.Config - w2: Linear.Config - w3: Linear.Config - activation: SituAndMul.Config - - def __init__(self, config: Config): - super().__init__() - self.w1 = config.w1.build() - self.w2 = config.w2.build() - self.w3 = config.w3.build() - self.activation = config.activation.build() - - def forward(self, x: torch.Tensor) -> torch.Tensor: - gate_up = torch.cat((self.w1(x), self.w3(x)), dim=-1) - return self.w2(self.activation(gate_up)) + return self.w2((gate * up).to(input_dtype)) class KimiMLAAttention(Module): From 611dd731148964ce0c90448e4a444f96e2c7f007 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Sun, 9 Aug 2026 07:15:48 +0000 Subject: [PATCH 14/67] dp_replicate needs no specific handling, removed from unsupported_parallelisms --- torchtitan/models/kimi_k3/parallelize.py | 1 - 1 file changed, 1 deletion(-) diff --git a/torchtitan/models/kimi_k3/parallelize.py b/torchtitan/models/kimi_k3/parallelize.py index c2c3d92485..e3480fdfa1 100644 --- a/torchtitan/models/kimi_k3/parallelize.py +++ b/torchtitan/models/kimi_k3/parallelize.py @@ -39,7 +39,6 @@ def parallelize_kimi_k3( unsupported_parallelisms = [ name for name, enabled in ( - ("hybrid sharded data parallel", parallel_dims.dp_replicate_enabled), ("tensor parallel", parallel_dims.tp_enabled), ("pipeline parallel", parallel_dims.pp_enabled), ("context parallel", parallel_dims.cp_enabled), From d05952e3bd19ab0197c07e6f71224e5004d0ae55 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Mon, 10 Aug 2026 02:03:35 +0000 Subject: [PATCH 15/67] Refactor KimiGroupedExperts to use gmm --- tests/unit_tests/test_kimi_k3.py | 18 ++- torchtitan/models/kimi_k3/__init__.py | 26 ++-- torchtitan/models/kimi_k3/model.py | 132 ++++++++++++------ .../models/kimi_k3/state_dict_adapter.py | 60 ++++++-- 4 files changed, 171 insertions(+), 65 deletions(-) diff --git a/tests/unit_tests/test_kimi_k3.py b/tests/unit_tests/test_kimi_k3.py index 936ebe6f04..dd292a8e98 100644 --- a/tests/unit_tests/test_kimi_k3.py +++ b/tests/unit_tests/test_kimi_k3.py @@ -296,7 +296,7 @@ def test_debugmodel_preserves_reduced_k3_topology(self): assert moe_config is not None self.assertEqual(moe_config.num_experts, 8) self.assertEqual(moe_config.router.top_k, 2) - self.assertEqual(moe_config.routed_experts[0].w1.out_features, 128) + self.assertEqual(moe_config.routed_experts.hidden_dim, 128) vision_config = config.vision_encoder assert vision_config is not None self.assertEqual(vision_config.dim, 256) @@ -385,13 +385,17 @@ def test_unused_moe_experts_receive_zero_gradients(self): self.assertTrue(unused_experts) moe(inputs).float().sum().backward() - for expert_idx in unused_experts: - for parameter in moe.routed_experts[expert_idx].parameters(): - self.assertIsNotNone(parameter.grad) - assert parameter.grad is not None + for parameter in ( + moe.routed_experts.w1_EFD, + moe.routed_experts.w2_EDF, + moe.routed_experts.w3_EFD, + ): + self.assertIsNotNone(parameter.grad) + assert parameter.grad is not None + for expert_idx in unused_experts: torch.testing.assert_close( - parameter.grad, - torch.zeros_like(parameter.grad), + parameter.grad[expert_idx], + torch.zeros_like(parameter.grad[expert_idx]), ) def test_small_multimodal_model_forward_backward_and_adapter(self): diff --git a/torchtitan/models/kimi_k3/__init__.py b/torchtitan/models/kimi_k3/__init__.py index 35b4a36e00..5d698b2149 100644 --- a/torchtitan/models/kimi_k3/__init__.py +++ b/torchtitan/models/kimi_k3/__init__.py @@ -16,6 +16,7 @@ from torchtitan.models.common import Conv1d, Embedding, Linear from torchtitan.models.common.moe import TokenChoiceTopKRouter from torchtitan.models.common.nn_modules import GELU, RMSNorm +from torchtitan.models.common.token_dispatcher import LocalTokenDispatcher from torchtitan.models.common.vision_encoder import VisionMLP from torchtitan.models.utils import validate_converter_order from torchtitan.protocols.model import ModelConfigConverter @@ -24,6 +25,7 @@ from .model import ( KimiDeltaAttention, KimiFeedForward, + KimiGroupedExperts, KimiK3Model, KimiK3TransformerBlock, KimiKDAKernel, @@ -230,13 +232,6 @@ def _latent_moe_config( top_k: int, num_shared_experts: int, ) -> KimiLatentMoE.Config: - routed_experts = [ - _feed_forward_config( - dim=latent_dim, - hidden_dim=expert_hidden_dim, - ) - for _ in range(num_experts) - ] return KimiLatentMoE.Config( num_experts=num_experts, router=TokenChoiceTopKRouter.Config( @@ -248,7 +243,22 @@ def _latent_moe_config( route_scale=1.0, ), routed_down=_linear(dim, latent_dim), - routed_experts=routed_experts, + routed_experts=KimiGroupedExperts.Config( + dim=latent_dim, + hidden_dim=expert_hidden_dim, + num_experts=num_experts, + beta=4.0, + linear_beta=25.0, + param_init={ + "w1_EFD": partial(nn.init.trunc_normal_, std=0.02), + "w2_EDF": partial(nn.init.trunc_normal_, std=0.02), + "w3_EFD": partial(nn.init.trunc_normal_, std=0.02), + }, + ), + token_dispatcher=LocalTokenDispatcher.Config( + num_experts=num_experts, + top_k=top_k, + ), routed_norm=_norm(latent_dim), routed_up=_linear(latent_dim, dim), shared_experts=_feed_forward_config( diff --git a/torchtitan/models/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py index be68c41c65..de92ac342c 100644 --- a/torchtitan/models/kimi_k3/model.py +++ b/torchtitan/models/kimi_k3/model.py @@ -20,21 +20,23 @@ from fla.ops.kda import chunk_kda from torch import nn +from torch.distributed.tensor import DTensor from torchtitan.distributed.fsdp import add_zero_valued_dependency from torchtitan.models.common import Conv1d, Linear from torchtitan.models.common.attention import AttentionMasksType from torchtitan.models.common.decoder import Decoder from torchtitan.models.common.feed_forward import FeedForward -from torchtitan.models.common.moe import TokenChoiceTopKRouter +from torchtitan.models.common.moe import GroupedExperts, TokenChoiceTopKRouter from torchtitan.models.common.multimodal import ( get_vision_positions, scatter_vision_embeds, ) from torchtitan.models.common.nn_modules import RMSNorm +from torchtitan.models.common.token_dispatcher import LocalTokenDispatcher from torchtitan.models.kimi_k3.vision_encoder import KimiK3VisionEncoder from torchtitan.models.utils import get_moe_model_nparams_and_flops -from torchtitan.protocols.module import Module, ModuleList +from torchtitan.protocols.module import Module # Shape suffixes: # B = batch, L = sequence length, D = model dimension, H = heads, @@ -375,23 +377,67 @@ def forward( return self.output_proj(out_BLHV.reshape(B, L, -1)) -class KimiRoutedExperts(ModuleList): - """List-backed experts implementing TorchTitan's FSDP expert protocol. +class KimiGroupedExperts(GroupedExperts): + """``common/moe.py::GroupedExperts`` with Kimi's SiTU activation. - Kimi K3 keeps one module per expert so the eager implementation and - HuggingFace state-dict mapping stay directly inspectable. The shared FSDP - wrapper discovers routed expert parameters through ``inner_experts`` and - ``num_experts``; exposing those properties here lets it shard this - list-backed layout without changing parameter names or forward math. + Inherits its stacked-weight shape (``w1_EFD``/``w2_EDF``/``w3_EFD``) and + parameter allocation; only ``forward`` differs, since the activation is + baked into the ``torch._grouped_mm`` call sequence rather than being a + swappable argument. ``inner_experts`` returns ``self`` so the shared FSDP + wrapper's ``moe.routed_experts.inner_experts`` discovery path works + without a separate dispatch-composing wrapper class. """ + @dataclass(kw_only=True, slots=True) + class Config(GroupedExperts.Config): + beta: float = 1.0 + linear_beta: float | None = None + + def __init__(self, config: Config): + super().__init__(config) + self.beta = config.beta + self.linear_beta = config.linear_beta + @property - def inner_experts(self) -> "KimiRoutedExperts": + def inner_experts(self) -> "KimiGroupedExperts": return self - @property - def num_experts(self) -> int: - return len(self) + def forward( + self, + x_RD: torch.Tensor, + num_tokens_per_expert_E: torch.Tensor, + ) -> torch.Tensor: + if isinstance(self.w1_EFD, DTensor): + w1_EFD = self.w1_EFD.to_local() + assert isinstance(self.w2_EDF, DTensor) + w2_EDF = self.w2_EDF.to_local() + assert isinstance(self.w3_EFD, DTensor) + w3_EFD = self.w3_EFD.to_local() + else: + w1_EFD = self.w1_EFD + w2_EDF = self.w2_EDF + w3_EFD = self.w3_EFD + + offsets_E = torch.cumsum(num_tokens_per_expert_E, dim=0, dtype=torch.int32) + + gate_RF = torch._grouped_mm( + x_RD.bfloat16(), w1_EFD.bfloat16().transpose(-2, -1), offs=offsets_E + ) + up_RF = torch._grouped_mm( + x_RD.bfloat16(), w3_EFD.bfloat16().transpose(-2, -1), offs=offsets_E + ) + + input_dtype = gate_RF.dtype + gate_RF = gate_RF.float() + up_RF = up_RF.float() + gate_RF = self.beta * torch.tanh(gate_RF / self.beta) * torch.sigmoid(gate_RF) + if self.linear_beta is not None: + up_RF = self.linear_beta * torch.tanh(up_RF / self.linear_beta) + h_RF = (gate_RF * up_RF).to(input_dtype) + + return torch._grouped_mm( + h_RF, w2_EDF.bfloat16().transpose(-2, -1), offs=offsets_E + ).type_as(x_RD) class KimiLatentMoE(Module): @@ -402,7 +448,8 @@ class Config(Module.Config): num_experts: int router: TokenChoiceTopKRouter.Config routed_down: Linear.Config - routed_experts: list[KimiFeedForward.Config] + routed_experts: KimiGroupedExperts.Config + token_dispatcher: LocalTokenDispatcher.Config routed_norm: KimiRMSNorm.Config routed_up: Linear.Config shared_experts: KimiFeedForward.Config @@ -410,16 +457,13 @@ class Config(Module.Config): def __init__(self, config: Config): super().__init__() - if len(config.routed_experts) != config.num_experts: - raise ValueError( - "The number of routed expert configs must equal num_experts." - ) + if config.routed_experts.num_experts != config.num_experts: + raise ValueError("routed_experts.num_experts must equal num_experts.") self.num_experts = config.num_experts self.router = config.router.build() self.routed_down = config.routed_down.build() - self.routed_experts = KimiRoutedExperts( - [expert.build() for expert in config.routed_experts] - ) + self.routed_experts = config.routed_experts.build() + self.token_dispatcher = config.token_dispatcher.build() self.routed_norm = config.routed_norm.build() self.routed_up = config.routed_up.build() self.shared_experts = config.shared_experts.build() @@ -450,30 +494,38 @@ def forward(self, x_BLD: torch.Tensor) -> torch.Tensor: dtype=torch.bool, device=x_BLD.device, ).scatter(-1, expert_ids_BLK, True) + num_tokens_per_expert_E = routing_map_BLE.sum(dim=(0, 1)) with torch.no_grad(): # In place so the load-balancing hook registered on the optimizer # keeps referring to this buffer across steps. - self.tokens_per_expert_E.add_(routing_map_BLE.sum(dim=(0, 1)).float()) - - latent_TD = self.routed_down(x_BLD).reshape(B * L, -1) - expert_ids_TK = expert_ids_BLK.reshape(B * L, -1) - weights_TK = weights_BLK.reshape(B * L, -1) - routed_TD = torch.zeros_like(latent_TD, dtype=torch.float32) - - for expert_idx, expert in enumerate(self.routed_experts): - token_and_slot = torch.nonzero(expert_ids_TK == expert_idx, as_tuple=False) - # Keep empty experts in the autograd graph so every FSDP rank - # produces zero gradients instead of rank-dependent None gradients. - token_ids = token_and_slot[:, 0] - route_slots = token_and_slot[:, 1] - expert_output = expert(latent_TD.index_select(0, token_ids)) - route_weight = weights_TK[token_ids, route_slots].unsqueeze(-1) - routed_TD = routed_TD.index_add( - 0, token_ids, expert_output.float() * route_weight + self.tokens_per_expert_E.add_(num_tokens_per_expert_E.float()) + + latent_BLD = self.routed_down(x_BLD) + latent_dim = latent_BLD.shape[-1] + K = weights_BLK.shape[-1] + T = B * L + latent_TD = latent_BLD.reshape(T, latent_dim) + weights_TK = weights_BLK.reshape(T, K) + expert_ids_TK = expert_ids_BLK.reshape(T, K) + + # Every expert's weights are packed into one grouped-mm call, so an + # expert with zero assigned tokens still sits in the autograd graph + # and receives an all-zero gradient rather than None. + routed_input_RD, num_tokens_per_expert_E, metadata = ( + self.token_dispatcher.dispatch( + latent_TD, weights_TK, expert_ids_TK, num_tokens_per_expert_E ) + ) + routed_output_RD = self.routed_experts(routed_input_RD, num_tokens_per_expert_E) + routed_TD = self.token_dispatcher.combine( + routed_output_RD, + metadata, + latent_TD, + num_local_tokens_after_padding=T, + local_seq_len_after_padding=L, + ) - routed_BLD = routed_TD.to(latent_TD.dtype).view(B, L, -1) - routed_BLD = self.routed_up(self.routed_norm(routed_BLD)) + routed_BLD = self.routed_up(self.routed_norm(routed_TD.view(B, L, latent_dim))) return routed_BLD + self.shared_experts(x_BLD) def _init_self_buffers(self, *, buffer_device: torch.device | None = None) -> None: diff --git a/torchtitan/models/kimi_k3/state_dict_adapter.py b/torchtitan/models/kimi_k3/state_dict_adapter.py index 7967be1f16..5073a6b260 100644 --- a/torchtitan/models/kimi_k3/state_dict_adapter.py +++ b/torchtitan/models/kimi_k3/state_dict_adapter.py @@ -87,6 +87,17 @@ ), } +# KimiGroupedExperts stacks all experts' weights into one (E, F, D) param per +# projection; HF stores them as separate per-expert 2D tensors. +_EXPERT_PROJECTION_TO_GROUPED_PARAM = { + "w1": "w1_EFD", + "w2": "w2_EDF", + "w3": "w3_EFD", +} +_GROUPED_PARAM_TO_EXPERT_PROJECTION = { + v: k for k, v in _EXPERT_PROJECTION_TO_GROUPED_PARAM.items() +} + _VISION_GLOBAL_FROM_HF = { "vision_tower.patch_embed.proj.weight": ("vision_encoder.patch_embed.weight"), "vision_tower.patch_embed.pos_emb.weight": "vision_encoder.pos_embed", @@ -115,6 +126,12 @@ def __init__( ): super().__init__(model_config, hf_assets_path) self.kimi_config = model_config + # {(layer_idx, projection): {expert_idx: 2D tensor}}, filled in from_hf + # while HF's per-expert keys arrive one at a time; stacked into + # KimiGroupedExperts' (E, F, D) parameter once all experts are seen. + self._expert_weights_by_layer_projection: dict[ + tuple[str, str], dict[int, torch.Tensor] + ] = {} @staticmethod def _raise_if_quantized_key(key: str) -> None: @@ -164,10 +181,24 @@ def from_hf(self, hf_state_dict: dict[str, Any]) -> dict[str, Any]: ) if expert_match is not None: expert_idx, projection = expert_match.groups() - state_dict[ + moe_config = self.kimi_config.layers[int(layer_idx)].moe + assert moe_config is not None + grouped_key = ( f"layers.{layer_idx}.moe.routed_experts." - f"{expert_idx}.{projection}.weight" - ] = value + f"{_EXPERT_PROJECTION_TO_GROUPED_PARAM[projection]}" + ) + experts = self._expert_weights_by_layer_projection.setdefault( + (layer_idx, projection), {} + ) + experts[int(expert_idx)] = value + if len(experts) == moe_config.num_experts: + sorted_experts = [ + experts[i] for i in range(moe_config.num_experts) + ] + state_dict[grouped_key] = torch.stack(sorted_experts, dim=0) + del self._expert_weights_by_layer_projection[ + (layer_idx, projection) + ] continue layer_config = self.kimi_config.layers[int(layer_idx)] @@ -222,6 +253,13 @@ def from_hf(self, hf_state_dict: dict[str, Any]) -> dict[str, Any]: "KimiK3StateDictAdapter found HuggingFace keys without a " f"mapping: {unmapped}." ) + if self._expert_weights_by_layer_projection: + incomplete = list(self._expert_weights_by_layer_projection.keys()) + self._expert_weights_by_layer_projection.clear() + raise ValueError( + "KimiK3StateDictAdapter received an incomplete set of " + f"routed-expert weights for (layer, projection): {incomplete}." + ) return state_dict def to_hf(self, state_dict: dict[str, Any]) -> dict[str, Any]: @@ -275,16 +313,18 @@ def to_hf(self, state_dict: dict[str, Any]) -> dict[str, Any]: if text_match is not None: layer_idx, suffix = text_match.groups() expert_match = re.fullmatch( - r"moe\.routed_experts\.(\d+)\." r"(w1|w2|w3)\.weight", + r"moe\.routed_experts\." r"(w1_EFD|w2_EDF|w3_EFD)", suffix, ) if expert_match is not None: - expert_idx, projection = expert_match.groups() - hf_state_dict[ - f"language_model.model.layers.{layer_idx}." - f"block_sparse_moe.experts.{expert_idx}." - f"{projection}.weight" - ] = value + (grouped_param,) = expert_match.groups() + projection = _GROUPED_PARAM_TO_EXPERT_PROJECTION[grouped_param] + for expert_idx, expert_weight in enumerate(value.unbind(0)): + hf_state_dict[ + f"language_model.model.layers.{layer_idx}." + f"block_sparse_moe.experts.{expert_idx}." + f"{projection}.weight" + ] = expert_weight continue mapped_suffix = text_layer_to_hf.get(suffix) From 6fe3694ed1a6e5bc04cdbf9edf8d48abd61a810e Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Mon, 10 Aug 2026 06:53:45 +0000 Subject: [PATCH 16/67] refactor KimiMLAAttention to Inheritance titian BaseAttention --- torchtitan/models/kimi_k3/__init__.py | 2 +- torchtitan/models/kimi_k3/model.py | 54 +++++++++++++-------------- 2 files changed, 26 insertions(+), 30 deletions(-) diff --git a/torchtitan/models/kimi_k3/__init__.py b/torchtitan/models/kimi_k3/__init__.py index 5d698b2149..647a461aba 100644 --- a/torchtitan/models/kimi_k3/__init__.py +++ b/torchtitan/models/kimi_k3/__init__.py @@ -153,7 +153,7 @@ def _mla_config( q_head_dim = qk_nope_head_dim + qk_rope_head_dim return KimiMLAAttention.Config( dim=dim, - num_heads=num_heads, + n_heads=num_heads, kv_lora_rank=kv_lora_rank, qk_nope_head_dim=qk_nope_head_dim, qk_rope_head_dim=qk_rope_head_dim, diff --git a/torchtitan/models/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py index de92ac342c..299f2131d1 100644 --- a/torchtitan/models/kimi_k3/model.py +++ b/torchtitan/models/kimi_k3/model.py @@ -13,7 +13,7 @@ ``tests/unit_tests/test_kimi_k3.py`` and is far too slow for training. """ -from dataclasses import dataclass +from dataclasses import dataclass, field import torch import torch.nn.functional as F @@ -24,7 +24,11 @@ from torchtitan.distributed.fsdp import add_zero_valued_dependency from torchtitan.models.common import Conv1d, Linear -from torchtitan.models.common.attention import AttentionMasksType +from torchtitan.models.common.attention import ( + AttentionMasksType, + BaseAttention, + ScaledDotProductAttention, +) from torchtitan.models.common.decoder import Decoder from torchtitan.models.common.feed_forward import FeedForward from torchtitan.models.common.moe import GroupedExperts, TokenChoiceTopKRouter @@ -117,18 +121,21 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return self.w2((gate * up).to(input_dtype)) -class KimiMLAAttention(Module): +class KimiMLAAttention(BaseAttention): """Kimi K3 multi-head latent attention. Unlike DeepSeek-V3 MLA, the released K3 configuration sets - ``mla_use_nope=True``. The RoPE-sized query/key slices remain part of the - projected head, but no rotary transform is applied. + ``mla_use_nope=True``: the RoPE-sized query/key slices remain part of the + projected head, but no rotary transform is applied, so this has no rope + config at all. Attention itself runs through ``ScaledDotProductAttention`` + rather than a hand-written softmax: torchtitan's training path has no + need to match HF eager's bit-for-bit fp32 softmax, so it uses the same + SDPA inner attention the rest of the codebase relies on. """ @dataclass(kw_only=True, slots=True) - class Config(Module.Config): + class Config(BaseAttention.Config): dim: int - num_heads: int kv_lora_rank: int qk_nope_head_dim: int qk_rope_head_dim: int @@ -141,10 +148,13 @@ class Config(Module.Config): wkv_b: Linear.Config gate: Linear.Config wo: Linear.Config + inner_attention: Module.Config = field( + default_factory=ScaledDotProductAttention.Config + ) def __init__(self, config: Config): super().__init__() - self.num_heads = config.num_heads + self.n_heads = config.n_heads self.qk_nope_head_dim = config.qk_nope_head_dim self.qk_rope_head_dim = config.qk_rope_head_dim self.q_head_dim = config.qk_nope_head_dim + config.qk_rope_head_dim @@ -160,6 +170,7 @@ def __init__(self, config: Config): self.wkv_b = config.wkv_b.build() self.gate = config.gate.build() self.wo = config.wo.build() + self.inner_attention = config.inner_attention.build() def forward( self, @@ -175,7 +186,7 @@ def forward( B, L, _ = x_BLD.shape q_BLNH = self.wq_b(self.q_norm(self.wq_a(x_BLD))).view( - B, L, self.num_heads, self.q_head_dim + B, L, self.n_heads, self.q_head_dim ) compressed_kv = self.wkv_a(x_BLD) @@ -187,7 +198,7 @@ def forward( kv_BLNH = self.wkv_b(self.kv_norm(kv_latent)).view( B, L, - self.num_heads, + self.n_heads, self.qk_nope_head_dim + self.v_head_dim, ) k_nope, v_BLNH = torch.split( @@ -196,27 +207,12 @@ def forward( dim=-1, ) k_rope = k_rope.view(B, L, 1, self.qk_rope_head_dim).expand( - -1, -1, self.num_heads, -1 + -1, -1, self.n_heads, -1 ) k_BLNH = torch.cat((k_nope, k_rope), dim=-1) - # HuggingFace eager attention keeps the matmul in the input dtype and - # performs softmax in FP32. - scores_BNLS = torch.einsum("blnh,bsnh->bnls", q_BLNH, k_BLNH) - scores_BNLS = scores_BNLS * self.scale - causal_mask = torch.ones(L, L, dtype=torch.bool, device=x_BLD.device).triu( - diagonal=1 - ) - scores_BNLS = scores_BNLS.masked_fill( - causal_mask.view(1, 1, L, L), - torch.finfo(scores_BNLS.dtype).min, - ) - probs_BNLS = torch.softmax(scores_BNLS, dim=-1, dtype=torch.float32).to( - v_BLNH.dtype - ) - out_BLNV = torch.einsum("bnls,bsnv->blnv", probs_BNLS, v_BLNH) - - out_BLD = out_BLNV.reshape(B, L, self.num_heads * self.v_head_dim) + out_BLNV = self.inner_attention(q_BLNH, k_BLNH, v_BLNH, scale=self.scale) + out_BLD = out_BLNV.reshape(B, L, self.n_heads * self.v_head_dim) out_BLD = out_BLD * torch.sigmoid(self.gate(x_BLD)) return self.wo(out_BLD) @@ -712,7 +708,7 @@ def get_nparams_and_flops( return get_moe_model_nparams_and_flops( self, model, - attention_config.num_heads, + attention_config.n_heads, attention_config.qk_nope_head_dim + attention_config.qk_rope_head_dim + attention_config.v_head_dim, From 1fc70f6f76f9c1233506b0257aff0224684a84de Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Mon, 10 Aug 2026 07:10:55 +0000 Subject: [PATCH 17/67] remove add_zero_valued_dependency, this should be address with other multimodeal model. --- tests/unit_tests/test_kimi_k3_fsdp.py | 94 --------------------------- torchtitan/distributed/fsdp.py | 31 --------- torchtitan/models/kimi_k3/model.py | 35 +--------- 3 files changed, 1 insertion(+), 159 deletions(-) diff --git a/tests/unit_tests/test_kimi_k3_fsdp.py b/tests/unit_tests/test_kimi_k3_fsdp.py index b464fc5fd4..00472c3225 100644 --- a/tests/unit_tests/test_kimi_k3_fsdp.py +++ b/tests/unit_tests/test_kimi_k3_fsdp.py @@ -6,7 +6,6 @@ import copy import unittest -from contextlib import nullcontext from unittest.mock import patch import torch @@ -139,99 +138,6 @@ def test_single_rank_fsdp_matches_manual_bf16_reference(self): self.assertGreater(compared_gradients, 0) -class TestKimiK3MixedModalityFSDP(DTensorTestBase): - @property - def world_size(self): - return 2 - - @with_comms - def test_image_and_text_only_ranks_complete_forward_backward(self): - torch.manual_seed(3) - config = _small_model_config() - with torch.device("meta"): - model = config.build() - model.to_empty(device=self.device_type) - model.init_states() - - parallelism = ParallelismConfig( - data_parallel_shard_degree=self.world_size, - tensor_parallel_degree=1, - pipeline_parallel_degree=1, - context_parallel_degree=1, - expert_parallel_degree=1, - ) - parallel_dims = ParallelDims.from_config( - parallelism, - world_size=self.world_size, - ) - with patch( - "torchtitan.distributed.parallel_dims.device_type", - self.device_type, - ): - parallel_dims.build_mesh() - gradient_division_context = ( - patch("torchtitan.distributed.fsdp.disable_fsdp_gradient_division") - if self.device_type == "cpu" - else nullcontext() - ) - with gradient_division_context: - model = parallelize_kimi_k3( - model, - parallel_dims=parallel_dims, - training=TrainingConfig( - local_batch_size=1, - seq_len=6, - steps=1, - dtype="bfloat16", - ), - parallelism=parallelism, - compile_config=CompileConfig(), - ac_config=None, - dump_folder="", - ) - - rank = torch.distributed.get_rank() - if rank == 0: - inputs = { - "tokens": torch.tensor( - [[1, 7, 2, 3, 4, 5]], - dtype=torch.long, - device=self.device_type, - ), - "pixel_values": torch.randn( - 1, - 4, - 3 * 2 * 2, - device=self.device_type, - ), - "grid_thw": torch.tensor( - [[1, 2, 2]], - dtype=torch.long, - device=self.device_type, - ), - "special_tokens": {"image_id": 7}, - } - else: - inputs = { - "tokens": torch.tensor( - [[1, 2, 3, 4, 5, 6]], - dtype=torch.long, - device=self.device_type, - ) - } - - logits_BLV = model(**inputs) - logits_BLV.float().square().mean().backward() - - vision_gradients = 0 - for name, parameter in model.named_parameters(): - if not name.startswith("vision_encoder."): - continue - self.assertIsNotNone(parameter.grad, name) - vision_gradients += 1 - self.assertGreater(vision_gradients, 0) - - if __name__ == "__main__": from torch.testing._internal.common_utils import run_tests diff --git a/torchtitan/distributed/fsdp.py b/torchtitan/distributed/fsdp.py index f1aa8fd8c3..1b00fa41d6 100644 --- a/torchtitan/distributed/fsdp.py +++ b/torchtitan/distributed/fsdp.py @@ -165,37 +165,6 @@ def apply_fsdp_to_vision_encoder( ) -def add_zero_valued_dependency( - output: torch.Tensor, - unused_output: torch.Tensor, -) -> torch.Tensor: - """Keep a conditionally executed FSDP module in the autograd graph. - - FSDP2 issues a module's all-gather from its pre-forward hook and its - reduce-scatter from the autograd hooks on that module's output. A module - that only some data-parallel ranks execute -- a VLM vision encoder on a - batch that happens to carry no images, for example -- would therefore - issue collectives on a subset of the process group and deadlock the step. - - Context parallelism reaches the same deadlock by a second route: a rank's - sequence shard can hold zero vision placeholders even when every rank - received images, so "this batch has images" is not a sufficient condition - for running the module. Decide from the process group, not from the local - batch. - - A rank with no real work for such a module runs it on a placeholder input - and routes the result through this helper. Scaling by zero leaves - ``output`` numerically unchanged while preserving the graph edge, so every - rank issues the same collectives and the module receives zero gradients -- - which is also its correct contribution to the data-parallel average. - - Args: - output: the tensor the caller actually wants to return. - unused_output: a tensor produced by the module being kept alive. - """ - return output + unused_output.sum().to(output.dtype) * 0.0 - - def apply_fsdp_to_decoder( model: "Decoder", dp_mesh: DeviceMesh, diff --git a/torchtitan/models/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py index 299f2131d1..4f41bfb532 100644 --- a/torchtitan/models/kimi_k3/model.py +++ b/torchtitan/models/kimi_k3/model.py @@ -22,7 +22,6 @@ from torch import nn from torch.distributed.tensor import DTensor -from torchtitan.distributed.fsdp import add_zero_valued_dependency from torchtitan.models.common import Conv1d, Linear from torchtitan.models.common.attention import ( AttentionMasksType, @@ -740,33 +739,6 @@ def __init__(self, config: Config): "number of text positions than the encoder produces." ) - def _encode_placeholder_image( - self, - embeddings_BLD: torch.Tensor, - ) -> torch.Tensor: - """Run the vision encoder on the smallest grid it can merge. - - A batch without images must still drive the vision encoder, because it - is its own FSDP unit and the data-parallel ranks that do have images - will issue its collectives. See ``add_zero_valued_dependency``. - """ - assert self.vision_encoder is not None - kernel_h, kernel_w = self.vision_encoder.merge_kernel_size - patch_dim = self.vision_encoder.patch_embed.in_features - pixel_values_NPK = torch.zeros( - 1, - kernel_h * kernel_w, - patch_dim, - dtype=embeddings_BLD.dtype, - device=embeddings_BLD.device, - ) - grid_thw_N3 = torch.tensor( - [[1, kernel_h, kernel_w]], - dtype=torch.long, - device=embeddings_BLD.device, - ) - return self.vision_encoder(pixel_values_NPK, grid_thw=grid_thw_N3) - def get_attention_masks(self, positions: torch.Tensor) -> AttentionMasksType | None: del positions return None @@ -786,12 +758,7 @@ def _prepare_multimodal_embeds( "both be omitted." ) if pixel_values is None: - if self.vision_encoder is None: - return embeddings - return add_zero_valued_dependency( - embeddings, - self._encode_placeholder_image(embeddings), - ) + return embeddings assert grid_thw is not None if self.vision_encoder is None: raise ValueError("pixel_values were provided without a vision encoder.") From 522958f8169bc47f415924fc1e3f655206e17b7b Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Tue, 11 Aug 2026 02:09:40 +0000 Subject: [PATCH 18/67] remove redundandency test case --- tests/unit_tests/test_kimi_k3.py | 87 +--------------- tests/unit_tests/test_kimi_k3_fsdp.py | 144 -------------------------- tests/unit_tests/test_multimodal.py | 76 -------------- 3 files changed, 1 insertion(+), 306 deletions(-) delete mode 100644 tests/unit_tests/test_kimi_k3_fsdp.py delete mode 100644 tests/unit_tests/test_multimodal.py diff --git a/tests/unit_tests/test_kimi_k3.py b/tests/unit_tests/test_kimi_k3.py index dd292a8e98..1f38a3b221 100644 --- a/tests/unit_tests/test_kimi_k3.py +++ b/tests/unit_tests/test_kimi_k3.py @@ -27,7 +27,6 @@ _mla_config, _norm, _vision_encoder_config, - kimi_k3_configs, ) from torchtitan.models.kimi_k3.model import ( KimiK3Model, @@ -277,31 +276,6 @@ def test_exact_gelu_matches_pytorch_reference(self): x_bf16 = x.bfloat16() self.assertEqual(KimiExactGELU.Config().build()(x_bf16).dtype, x_bf16.dtype) - def test_debugmodel_preserves_reduced_k3_topology(self): - config = kimi_k3_configs["debugmodel"]("eager") - - self.assertEqual(config.vocab_size, 163840) - self.assertEqual(len(config.layers), 13) - self.assertEqual( - [ - layer_idx + 1 - for layer_idx, layer in enumerate(config.layers) - if layer.attention is not None - ], - [4, 8, 12, 13], - ) - self.assertIsNotNone(config.layers[0].feed_forward) - self.assertTrue(all(layer.moe is not None for layer in config.layers[1:])) - moe_config = config.layers[1].moe - assert moe_config is not None - self.assertEqual(moe_config.num_experts, 8) - self.assertEqual(moe_config.router.top_k, 2) - self.assertEqual(moe_config.routed_experts.hidden_dim, 128) - vision_config = config.vision_encoder - assert vision_config is not None - self.assertEqual(vision_config.dim, 256) - self.assertEqual(vision_config.num_layers, 4) - @unittest.skipIf(not torch.cuda.is_available(), "FLA KDA kernel requires CUDA.") def test_fla_kda_kernel_matches_recurrent_reference(self): torch.manual_seed(1) @@ -369,71 +343,12 @@ def parameter(*shape: int) -> torch.Tensor: assert tensor.grad is not None self.assertTrue(torch.isfinite(tensor.grad).all()) - def test_unused_moe_experts_receive_zero_gradients(self): - torch.manual_seed(2) - model = _small_model_config().build() - model.init_states() - moe = model.layers["1"].moe - assert moe is not None - with torch.no_grad(): - moe.router.gate.weight.zero_() - - inputs = torch.randn(2, 4, 16, requires_grad=True) - _, expert_ids, _ = moe.router(inputs, moe.expert_bias_E) - selected_experts = set(expert_ids.flatten().tolist()) - unused_experts = set(range(moe.num_experts)) - selected_experts - self.assertTrue(unused_experts) - - moe(inputs).float().sum().backward() - for parameter in ( - moe.routed_experts.w1_EFD, - moe.routed_experts.w2_EDF, - moe.routed_experts.w3_EFD, - ): - self.assertIsNotNone(parameter.grad) - assert parameter.grad is not None - for expert_idx in unused_experts: - torch.testing.assert_close( - parameter.grad[expert_idx], - torch.zeros_like(parameter.grad[expert_idx]), - ) - - def test_small_multimodal_model_forward_backward_and_adapter(self): + def test_state_dict_round_trips_through_hf_adapter(self): torch.manual_seed(2) config = _small_model_config() model = config.build() - model.verify_module_protocol() model.init_states() - image_token_id = 7 - tokens_BL = torch.tensor( - [ - [1, 2, image_token_id, 3, 4, 5], - [6, image_token_id, image_token_id, 8, 9, 10], - ] - ) - pixel_values_NPK = torch.randn(2, 8, 3 * 2 * 2) - grid_thw_N3 = torch.tensor([[1, 2, 2], [1, 4, 2]]) - logits_BLV = model( - tokens_BL, - pixel_values=pixel_values_NPK, - grid_thw=grid_thw_N3, - special_tokens={"image_id": image_token_id}, - ) - - self.assertEqual(logits_BLV.shape, (2, 6, config.vocab_size)) - moe = model.layers["1"].moe - assert moe is not None - self.assertIs( - moe._buffers["tokens_per_expert_E"], - moe.tokens_per_expert_E, - ) - self.assertEqual(moe.tokens_per_expert_E.sum().item(), tokens_BL.numel()) - logits_BLV.float().square().mean().backward() - for parameter in model.parameters(): - if parameter.grad is not None: - self.assertTrue(torch.isfinite(parameter.grad).all()) - state_dict = model.state_dict() adapter = KimiK3StateDictAdapter(config, hf_assets_path=None) hf_state_dict = adapter.to_hf(state_dict) diff --git a/tests/unit_tests/test_kimi_k3_fsdp.py b/tests/unit_tests/test_kimi_k3_fsdp.py deleted file mode 100644 index 00472c3225..0000000000 --- a/tests/unit_tests/test_kimi_k3_fsdp.py +++ /dev/null @@ -1,144 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -import copy -import unittest -from unittest.mock import patch - -import torch -from torch.distributed._composable.fsdp import FSDPModule -from torch.distributed.tensor import DTensor -from torch.testing._internal.distributed._tensor.common_dtensor import ( - DTensorTestBase, - with_comms, -) -from torchtitan.config import CompileConfig, ParallelismConfig, TrainingConfig -from torchtitan.distributed import ParallelDims - -# Skip instead of failing collection when FLA, a per-model dependency of -# kimi_k3, is not installed; see the matching guard in test_kimi_k3.py. -try: - from torchtitan.models.kimi_k3 import parallelize_kimi_k3 - from torchtitan.models.kimi_k3.model import KimiK3Model -except ModuleNotFoundError as exc: - raise unittest.SkipTest( - f"Kimi K3 optional dependency unavailable: {exc.name}" - ) from exc - -from tests.unit_tests.test_kimi_k3 import _small_model_config - - -class TestKimiK3FSDP(DTensorTestBase): - @property - def world_size(self): - return 1 - - @with_comms - def test_single_rank_fsdp_matches_manual_bf16_reference(self): - torch.manual_seed(3) - # attn_res_block_size=2 makes the second layer pass the attention - # residual through rather than extend it. That path routes a tensor - # back out through the FSDP module boundary, and its gradient - # accumulation order is what keeps FSDP bitwise equal to eager, so the - # comparison below has to cover it. - config = _small_model_config(attn_res_block_size=2) - with torch.device("meta"): - model = config.build() - model.to_empty(device=self.device_type) - model.init_states() - with torch.no_grad(): - for transformer_block in model.layers.values(): - if transformer_block.moe is not None: - transformer_block.moe.router.gate.weight.zero_() - - reference = copy.deepcopy(model) - for parameter in reference.parameters(): - parameter.data = parameter.data.to(torch.bfloat16) - - parallelism = ParallelismConfig( - data_parallel_shard_degree=1, - tensor_parallel_degree=1, - pipeline_parallel_degree=1, - context_parallel_degree=1, - expert_parallel_degree=1, - ) - parallel_dims = ParallelDims.from_config(parallelism, world_size=1) - with patch( - "torchtitan.distributed.parallel_dims.device_type", - self.device_type, - ): - parallel_dims.build_mesh() - model = parallelize_kimi_k3( - model, - parallel_dims=parallel_dims, - training=TrainingConfig( - local_batch_size=1, - seq_len=6, - steps=1, - dtype="bfloat16", - ), - parallelism=parallelism, - compile_config=CompileConfig(), - ac_config=None, - dump_folder="", - ) - - assert isinstance(model, KimiK3Model) - self.assertIsInstance(model, FSDPModule) - self.assertIsInstance(model.vision_encoder, FSDPModule) - - inputs = { - "tokens": torch.tensor( - [[1, 7, 2, 3, 4, 5]], - dtype=torch.long, - device=self.device_type, - ), - "pixel_values": torch.randn( - 1, - 4, - 3 * 2 * 2, - device=self.device_type, - ), - "grid_thw": torch.tensor( - [[1, 2, 2]], - dtype=torch.long, - device=self.device_type, - ), - "special_tokens": {"image_id": 7}, - } - - actual_BLV = model(**inputs) # pyrefly: ignore [not-callable] - expected_BLV = reference(**inputs) - torch.testing.assert_close(actual_BLV, expected_BLV, atol=0.0, rtol=0.0) - - actual_BLV.float().square().mean().backward() - expected_BLV.float().square().mean().backward() - - reference_parameters = dict(reference.named_parameters()) - compared_gradients = 0 - for name, parameter in model.named_parameters(): - actual_grad = parameter.grad - expected_grad = reference_parameters[name].grad - self.assertEqual(actual_grad is None, expected_grad is None) - if actual_grad is None: - continue - if isinstance(actual_grad, DTensor): - actual_grad = actual_grad.to_local() - assert expected_grad is not None - torch.testing.assert_close( - actual_grad.float(), - expected_grad.float(), - atol=0.0, - rtol=0.0, - ) - compared_gradients += 1 - self.assertGreater(compared_gradients, 0) - - -if __name__ == "__main__": - from torch.testing._internal.common_utils import run_tests - - run_tests() diff --git a/tests/unit_tests/test_multimodal.py b/tests/unit_tests/test_multimodal.py deleted file mode 100644 index b87079ea47..0000000000 --- a/tests/unit_tests/test_multimodal.py +++ /dev/null @@ -1,76 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -"""Tests for the model-agnostic vision<->text fusion helpers.""" - -import unittest - -import torch - -from torchtitan.models.common.multimodal import scatter_vision_embeds - - -class TestScatterVisionEmbeds(unittest.TestCase): - def test_scatter_routes_gradients_to_both_streams(self): - # The scatter is in place, so the graph has to be built from tensors - # that are not themselves leaves requiring grad. - inputs_source = torch.arange(30.0).view(2, 5, 3).requires_grad_() - vision_embeds = torch.arange(12.0).view(2, 2, 3).requires_grad_() - inputs_embeds = inputs_source * 1.0 - inputs_before = inputs_source.detach().clone() - - actual = scatter_vision_embeds( - inputs_embeds, - vision_embeds=vision_embeds, - vision_positions=[ - (0, 0, 1, 1), - (1, 1, 2, 2), - ], - ) - expected = torch.stack( - ( - torch.cat( - ( - inputs_before[0, :1], - vision_embeds.detach()[0, :1], - inputs_before[0, 2:], - ) - ), - torch.cat( - ( - inputs_before[1, :2], - vision_embeds.detach()[1], - inputs_before[1, 4:], - ) - ), - ) - ) - - torch.testing.assert_close(actual, expected) - - actual.sum().backward() - # Overwritten positions must not propagate to the text embeddings, and - # every vision token that was scattered must receive gradient. - expected_input_grad = torch.tensor( - [[1, 0, 1, 1, 1], [1, 1, 0, 0, 1]], - dtype=inputs_source.dtype, - ).unsqueeze(-1) - expected_vision_grad = torch.tensor( - [[1, 0], [1, 1]], - dtype=vision_embeds.dtype, - ).unsqueeze(-1) - torch.testing.assert_close( - inputs_source.grad, - expected_input_grad.expand_as(inputs_source), - ) - torch.testing.assert_close( - vision_embeds.grad, - expected_vision_grad.expand_as(vision_embeds), - ) - - -if __name__ == "__main__": - unittest.main() From 51ad154f91b66200b636d2d53f3f5eecaff6b06a Mon Sep 17 00:00:00 2001 From: yangting <15258668233@163.com> Date: Tue, 11 Aug 2026 14:24:05 +0800 Subject: [PATCH 19/67] connect FLA operator, replace RMSNorm --- torchtitan/models/kimi_k3/__init__.py | 19 ++--- torchtitan/models/kimi_k3/model.py | 110 +++++++++++++++----------- 2 files changed, 71 insertions(+), 58 deletions(-) diff --git a/torchtitan/models/kimi_k3/__init__.py b/torchtitan/models/kimi_k3/__init__.py index 647a461aba..23c7a919e3 100644 --- a/torchtitan/models/kimi_k3/__init__.py +++ b/torchtitan/models/kimi_k3/__init__.py @@ -13,7 +13,7 @@ import torch.nn as nn from torchtitan.components.optimizer import register_moe_load_balancing_hook -from torchtitan.models.common import Conv1d, Embedding, Linear +from torchtitan.models.common import Embedding, Linear from torchtitan.models.common.moe import TokenChoiceTopKRouter from torchtitan.models.common.nn_modules import GELU, RMSNorm from torchtitan.models.common.token_dispatcher import LocalTokenDispatcher @@ -31,8 +31,8 @@ KimiKDAKernel, KimiLatentMoE, KimiMLAAttention, - KimiRMSNorm, KimiRMSNormGated, + KimiShortConvolution, ) from .parallelize import parallelize_kimi_k3 from .state_dict_adapter import KimiK3StateDictAdapter @@ -118,8 +118,8 @@ def _linear( ) -def _norm(dim: int, eps: float = 1e-5) -> KimiRMSNorm.Config: - return KimiRMSNorm.Config( +def _norm(dim: int, eps: float = 1e-5) -> RMSNorm.Config: + return RMSNorm.Config( normalized_shape=dim, eps=eps, param_init=_NORM_INIT, @@ -181,13 +181,11 @@ def _kda_config( ) -> KimiDeltaAttention.Config: projection_dim = num_heads * head_dim - def conv() -> Conv1d.Config: - return Conv1d.Config( - in_channels=projection_dim, - out_channels=projection_dim, + def conv() -> KimiShortConvolution.Config: + return KimiShortConvolution.Config( + hidden_size=projection_dim, kernel_size=conv_kernel_size, - groups=projection_dim, - bias=False, + activation="silu", param_init=_CONV_INIT, ) @@ -195,7 +193,6 @@ def conv() -> Conv1d.Config: dim=dim, num_heads=num_heads, head_dim=head_dim, - conv_kernel_size=conv_kernel_size, q_proj=_linear(dim, projection_dim), k_proj=_linear(dim, projection_dim), v_proj=_linear(dim, projection_dim), diff --git a/torchtitan/models/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py index 4f41bfb532..5f3cacb817 100644 --- a/torchtitan/models/kimi_k3/model.py +++ b/torchtitan/models/kimi_k3/model.py @@ -18,11 +18,14 @@ import torch import torch.nn.functional as F +from fla.modules import ShortConvolution +from fla.modules.fused_norm_gate import rms_norm_gated +from fla.modules.conv.causal_conv1d import causal_conv1d from fla.ops.kda import chunk_kda from torch import nn from torch.distributed.tensor import DTensor -from torchtitan.models.common import Conv1d, Linear +from torchtitan.models.common import Linear from torchtitan.models.common.attention import ( AttentionMasksType, BaseAttention, @@ -47,34 +50,51 @@ # T = flattened tokens, N = attention-residual entries. -class KimiRMSNorm(RMSNorm): - """RMSNorm that applies its weight after casting back to the input dtype. - ``nn.RMSNorm`` scales by the weight while still in the reduction dtype, - which does not match the released Kimi implementation under BF16. Keeping - the subclass also exposes ``kimi_eps`` to ``_apply_attention_residual``, - which needs the epsilon to normalize residual entries by hand. + +class KimiShortConvolution(ShortConvolution, Module): + """KDA short causal convolution backed by FLA's fused kernel. + + Matches the released Kimi K3 HuggingFace model, which builds FLA's + ``ShortConvolution`` per q/k/v projection. The Triton kernel runs only on + accelerator devices. """ @dataclass(kw_only=True, slots=True) - class Config(RMSNorm.Config): - pass + class Config(Module.Config): + hidden_size: int + kernel_size: int + activation: str = "silu" def __init__(self, config: Config): - super().__init__(config) - self.kimi_eps = config.eps + super().__init__( + hidden_size=config.hidden_size, + kernel_size=config.kernel_size, + activation=config.activation, + ) - def forward(self, x: torch.Tensor) -> torch.Tensor: - input_dtype = x.dtype - x_float = x.float() - variance = x_float.pow(2).mean(dim=-1, keepdim=True) - x_float = x_float * torch.rsqrt(variance + self.kimi_eps) - assert self.weight is not None - return self.weight * x_float.to(input_dtype) + def forward( + self, + x_BLD: torch.Tensor, + **kwargs: object, + ) -> tuple[torch.Tensor, None]: + y_BLD, _ = causal_conv1d( + x=x_BLD, + weight=self.weight.squeeze(1), + activation=self.activation, + backend=self.backend, + ) + return y_BLD, None class KimiRMSNormGated(Module): - """Per-head RMSNorm followed by a sigmoid output gate.""" + """Per-head RMSNorm + sigmoid output gate backed by FLA's fused kernel. + + Wraps the FLA functional ``rms_norm_gated`` behind the torchtitan + ``Module`` protocol so it participates in ``init_states``/``param_init``. + The wrapper owns the weight (ones-initialized, weight-only checkpoint + schema unchanged); the fused kernel receives it with ``bias=None``. + """ @dataclass(kw_only=True, slots=True) class Config(Module.Config): @@ -87,12 +107,14 @@ def __init__(self, config: Config): self.weight = nn.Parameter(torch.empty(config.dim)) def forward(self, x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: - input_dtype = x.dtype - x_float = x.float() - variance = x_float.pow(2).mean(dim=-1, keepdim=True) - x_float = x_float * torch.rsqrt(variance + self.eps) - x_float = self.weight.float() * x_float - return (x_float * torch.sigmoid(gate.float())).to(input_dtype) + return rms_norm_gated( + x, + gate, + self.weight, + None, + activation="sigmoid", + eps=self.eps, + ) class KimiFeedForward(FeedForward): @@ -140,10 +162,10 @@ class Config(BaseAttention.Config): qk_rope_head_dim: int v_head_dim: int wq_a: Linear.Config - q_norm: KimiRMSNorm.Config + q_norm: RMSNorm.Config wq_b: Linear.Config wkv_a: Linear.Config - kv_norm: KimiRMSNorm.Config + kv_norm: RMSNorm.Config wkv_b: Linear.Config gate: Linear.Config wo: Linear.Config @@ -284,13 +306,12 @@ class Config(Module.Config): dim: int num_heads: int head_dim: int - conv_kernel_size: int q_proj: Linear.Config k_proj: Linear.Config v_proj: Linear.Config - q_conv: Conv1d.Config - k_conv: Conv1d.Config - v_conv: Conv1d.Config + q_conv: KimiShortConvolution.Config + k_conv: KimiShortConvolution.Config + v_conv: KimiShortConvolution.Config forget_a: Linear.Config forget_b: Linear.Config beta: Linear.Config @@ -306,7 +327,6 @@ def __init__(self, config: Config): super().__init__() self.num_heads = config.num_heads self.head_dim = config.head_dim - self.conv_kernel_size = config.conv_kernel_size self.q_proj = config.q_proj.build() self.k_proj = config.k_proj.build() @@ -325,10 +345,6 @@ def __init__(self, config: Config): self.A_log = nn.Parameter(torch.empty(config.num_heads)) self.dt_bias = nn.Parameter(torch.empty(config.num_heads, config.head_dim)) - def _causal_conv(self, x_BLC: torch.Tensor, conv: Conv1d) -> torch.Tensor: - x_BCL = F.pad(x_BLC.transpose(1, 2), (self.conv_kernel_size - 1, 0)) - return F.silu(conv(x_BCL)).transpose(1, 2) - def forward( self, x_BLD: torch.Tensor, @@ -342,13 +358,13 @@ def forward( ) B, L, _ = x_BLD.shape - q_BLHK = self._causal_conv(self.q_proj(x_BLD), self.q_conv).view( + q_BLHK = self.q_conv(self.q_proj(x_BLD))[0].view( B, L, self.num_heads, self.head_dim ) - k_BLHK = self._causal_conv(self.k_proj(x_BLD), self.k_conv).view( + k_BLHK = self.k_conv(self.k_proj(x_BLD))[0].view( B, L, self.num_heads, self.head_dim ) - v_BLHV = self._causal_conv(self.v_proj(x_BLD), self.v_conv).view( + v_BLHV = self.v_conv(self.v_proj(x_BLD))[0].view( B, L, self.num_heads, self.head_dim ) forget_BLHK = self.forget_b(self.forget_a(x_BLD)).view( @@ -445,7 +461,7 @@ class Config(Module.Config): routed_down: Linear.Config routed_experts: KimiGroupedExperts.Config token_dispatcher: LocalTokenDispatcher.Config - routed_norm: KimiRMSNorm.Config + routed_norm: RMSNorm.Config routed_up: Linear.Config shared_experts: KimiFeedForward.Config load_balance_coeff: float | None = 1e-3 @@ -539,14 +555,14 @@ def _apply_attention_residual( prefix_sum_TD: torch.Tensor, block_residual_TND: torch.Tensor, projection: Linear, - norm: KimiRMSNorm, + norm: RMSNorm, ) -> torch.Tensor: """Apply Kimi's block-level attention residual in FP32.""" values_TND = torch.cat((block_residual_TND, prefix_sum_TD.unsqueeze(1)), dim=1) values_float = values_TND.float() variance = values_float.pow(2).mean(dim=-1, keepdim=True) - keys_TND = values_float * torch.rsqrt(variance + norm.kimi_eps) + keys_TND = values_float * torch.rsqrt(variance + norm.eps) score_weight_D = norm.weight.float() * projection.weight.squeeze(0).float() scores_TN = (keys_TND * score_weight_D).sum(dim=-1) probs_TN = torch.softmax(scores_TN, dim=-1).unsqueeze(1) @@ -565,11 +581,11 @@ class Config(Module.Config): delta_attention: KimiDeltaAttention.Config | None feed_forward: KimiFeedForward.Config | None moe: KimiLatentMoE.Config | None - attention_norm: KimiRMSNorm.Config - ffn_norm: KimiRMSNorm.Config - attention_res_norm: KimiRMSNorm.Config + attention_norm: RMSNorm.Config + ffn_norm: RMSNorm.Config + attention_res_norm: RMSNorm.Config attention_res_proj: Linear.Config - ffn_res_norm: KimiRMSNorm.Config + ffn_res_norm: RMSNorm.Config ffn_res_proj: Linear.Config def __init__(self, config: Config): @@ -670,7 +686,7 @@ class KimiK3Model(Decoder): @dataclass(kw_only=True, slots=True) class Config(Decoder.Config): layers: list[KimiK3TransformerBlock.Config] - output_res_norm: KimiRMSNorm.Config + output_res_norm: RMSNorm.Config output_res_proj: Linear.Config vision_encoder: KimiK3VisionEncoder.Config | None = None spatial_merge_size: int = 2 From bc0e27113bd00c71678f25b6655b0c1f8255e1ed Mon Sep 17 00:00:00 2001 From: yangting <15258668233@163.com> Date: Tue, 11 Aug 2026 16:21:58 +0800 Subject: [PATCH 20/67] revert fla operator --- torchtitan/models/kimi_k3/__init__.py | 14 +++-- torchtitan/models/kimi_k3/model.py | 81 +++++++-------------------- 2 files changed, 28 insertions(+), 67 deletions(-) diff --git a/torchtitan/models/kimi_k3/__init__.py b/torchtitan/models/kimi_k3/__init__.py index 23c7a919e3..014c736dc7 100644 --- a/torchtitan/models/kimi_k3/__init__.py +++ b/torchtitan/models/kimi_k3/__init__.py @@ -13,7 +13,7 @@ import torch.nn as nn from torchtitan.components.optimizer import register_moe_load_balancing_hook -from torchtitan.models.common import Embedding, Linear +from torchtitan.models.common import Conv1d, Embedding, Linear from torchtitan.models.common.moe import TokenChoiceTopKRouter from torchtitan.models.common.nn_modules import GELU, RMSNorm from torchtitan.models.common.token_dispatcher import LocalTokenDispatcher @@ -32,7 +32,6 @@ KimiLatentMoE, KimiMLAAttention, KimiRMSNormGated, - KimiShortConvolution, ) from .parallelize import parallelize_kimi_k3 from .state_dict_adapter import KimiK3StateDictAdapter @@ -181,11 +180,13 @@ def _kda_config( ) -> KimiDeltaAttention.Config: projection_dim = num_heads * head_dim - def conv() -> KimiShortConvolution.Config: - return KimiShortConvolution.Config( - hidden_size=projection_dim, + def conv() -> Conv1d.Config: + return Conv1d.Config( + in_channels=projection_dim, + out_channels=projection_dim, kernel_size=conv_kernel_size, - activation="silu", + groups=projection_dim, + bias=False, param_init=_CONV_INIT, ) @@ -193,6 +194,7 @@ def conv() -> KimiShortConvolution.Config: dim=dim, num_heads=num_heads, head_dim=head_dim, + conv_kernel_size=conv_kernel_size, q_proj=_linear(dim, projection_dim), k_proj=_linear(dim, projection_dim), v_proj=_linear(dim, projection_dim), diff --git a/torchtitan/models/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py index 5f3cacb817..3d337d28b4 100644 --- a/torchtitan/models/kimi_k3/model.py +++ b/torchtitan/models/kimi_k3/model.py @@ -18,14 +18,11 @@ import torch import torch.nn.functional as F -from fla.modules import ShortConvolution -from fla.modules.fused_norm_gate import rms_norm_gated -from fla.modules.conv.causal_conv1d import causal_conv1d from fla.ops.kda import chunk_kda from torch import nn from torch.distributed.tensor import DTensor -from torchtitan.models.common import Linear +from torchtitan.models.common import Conv1d, Linear from torchtitan.models.common.attention import ( AttentionMasksType, BaseAttention, @@ -51,50 +48,8 @@ - -class KimiShortConvolution(ShortConvolution, Module): - """KDA short causal convolution backed by FLA's fused kernel. - - Matches the released Kimi K3 HuggingFace model, which builds FLA's - ``ShortConvolution`` per q/k/v projection. The Triton kernel runs only on - accelerator devices. - """ - - @dataclass(kw_only=True, slots=True) - class Config(Module.Config): - hidden_size: int - kernel_size: int - activation: str = "silu" - - def __init__(self, config: Config): - super().__init__( - hidden_size=config.hidden_size, - kernel_size=config.kernel_size, - activation=config.activation, - ) - - def forward( - self, - x_BLD: torch.Tensor, - **kwargs: object, - ) -> tuple[torch.Tensor, None]: - y_BLD, _ = causal_conv1d( - x=x_BLD, - weight=self.weight.squeeze(1), - activation=self.activation, - backend=self.backend, - ) - return y_BLD, None - - class KimiRMSNormGated(Module): - """Per-head RMSNorm + sigmoid output gate backed by FLA's fused kernel. - - Wraps the FLA functional ``rms_norm_gated`` behind the torchtitan - ``Module`` protocol so it participates in ``init_states``/``param_init``. - The wrapper owns the weight (ones-initialized, weight-only checkpoint - schema unchanged); the fused kernel receives it with ``bias=None``. - """ + """Per-head RMSNorm followed by a sigmoid output gate.""" @dataclass(kw_only=True, slots=True) class Config(Module.Config): @@ -107,14 +62,12 @@ def __init__(self, config: Config): self.weight = nn.Parameter(torch.empty(config.dim)) def forward(self, x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: - return rms_norm_gated( - x, - gate, - self.weight, - None, - activation="sigmoid", - eps=self.eps, - ) + input_dtype = x.dtype + x_float = x.float() + variance = x_float.pow(2).mean(dim=-1, keepdim=True) + x_float = x_float * torch.rsqrt(variance + self.eps) + x_float = self.weight.float() * x_float + return (x_float * torch.sigmoid(gate.float())).to(input_dtype) class KimiFeedForward(FeedForward): @@ -306,12 +259,13 @@ class Config(Module.Config): dim: int num_heads: int head_dim: int + conv_kernel_size: int q_proj: Linear.Config k_proj: Linear.Config v_proj: Linear.Config - q_conv: KimiShortConvolution.Config - k_conv: KimiShortConvolution.Config - v_conv: KimiShortConvolution.Config + q_conv: Conv1d.Config + k_conv: Conv1d.Config + v_conv: Conv1d.Config forget_a: Linear.Config forget_b: Linear.Config beta: Linear.Config @@ -327,6 +281,7 @@ def __init__(self, config: Config): super().__init__() self.num_heads = config.num_heads self.head_dim = config.head_dim + self.conv_kernel_size = config.conv_kernel_size self.q_proj = config.q_proj.build() self.k_proj = config.k_proj.build() @@ -345,6 +300,10 @@ def __init__(self, config: Config): self.A_log = nn.Parameter(torch.empty(config.num_heads)) self.dt_bias = nn.Parameter(torch.empty(config.num_heads, config.head_dim)) + def _causal_conv(self, x_BLC: torch.Tensor, conv: Conv1d) -> torch.Tensor: + x_BCL = F.pad(x_BLC.transpose(1, 2), (self.conv_kernel_size - 1, 0)) + return F.silu(conv(x_BCL)).transpose(1, 2) + def forward( self, x_BLD: torch.Tensor, @@ -358,13 +317,13 @@ def forward( ) B, L, _ = x_BLD.shape - q_BLHK = self.q_conv(self.q_proj(x_BLD))[0].view( + q_BLHK = self._causal_conv(self.q_proj(x_BLD), self.q_conv).view( B, L, self.num_heads, self.head_dim ) - k_BLHK = self.k_conv(self.k_proj(x_BLD))[0].view( + k_BLHK = self._causal_conv(self.k_proj(x_BLD), self.k_conv).view( B, L, self.num_heads, self.head_dim ) - v_BLHV = self.v_conv(self.v_proj(x_BLD))[0].view( + v_BLHV = self._causal_conv(self.v_proj(x_BLD), self.v_conv).view( B, L, self.num_heads, self.head_dim ) forget_BLHK = self.forget_b(self.forget_a(x_BLD)).view( From c5b9330ecca098cb0b9369a173307759eb807376 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Tue, 11 Aug 2026 03:02:21 +0000 Subject: [PATCH 21/67] remove test_kimi_k3_hf_parity.py --- tests/unit_tests/test_kimi_k3_hf_parity.py | 168 --------------------- 1 file changed, 168 deletions(-) delete mode 100644 tests/unit_tests/test_kimi_k3_hf_parity.py diff --git a/tests/unit_tests/test_kimi_k3_hf_parity.py b/tests/unit_tests/test_kimi_k3_hf_parity.py deleted file mode 100644 index f6d67c94c9..0000000000 --- a/tests/unit_tests/test_kimi_k3_hf_parity.py +++ /dev/null @@ -1,168 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -"""Reduced Kimi K3 numerical parity against the released HuggingFace model. - -The reference tensors were generated with the released ``modeling_kimi_k3.py`` -at moonshotai/Kimi-K3 commit c5d1dd4c428bd1ce8b88c5044f3b6ccde9e3b721. -The reduced model uses the same deterministic parameter and input construction -as this test, then maps its state dict through ``KimiK3StateDictAdapter`` and -loads the HuggingFace model strictly. The HuggingFace eager attention paths and -a pure PyTorch implementation of the released FLA KDA API produced the frozen -float32 values below. -""" - -import base64 -import unittest - -import torch - -# Inherits test_kimi_k3's module-level skip when FLA is not installed. -from tests.unit_tests.test_kimi_k3 import _small_model_config - - -_TEXT_LOGITS_BASE64 = ( - "qvilPdqoTz1q2ZQ8yRJ4vMTJQ70Rj6C9X03YvVFcA75J6BS+9QggvohDJL5WaSG+" - "2ZkXvldBB75jKOK9Rg6svdUSXb3KArG8+II/PI8yNj3RVJo9O+nSPQkzAT6mXxM+" - "5jEfPlAnJD4rCSI+2u4YPtA8CT6FQOc9Xg+yPW1iaj2gtbM9feBrPUtVzDxJyg+8L" - "oMsvfvRlr314dC9ZPgAvobwE74FiCC+9zMmvrq1JL7NHRy+BssMvvbM7r2RuLm9N" - "0V5vdO06Lwwo6w7IJUePdtfkD3pO8s9IVX9PU5BEj4Hih8+H/IlPuQyJT6RVB0+/" - "a0OPrPB8z1Zpb89HEaDPbBPsz2Ke249vx/YPNAh2LtI7iC9K/SPvYw8yb0m2Pm9Rd" - "cPvhOPHL4mhyK+m30hvuZ9Gb5O4Aq+NYzsvVolub2Ohnu9S9XzvOAqTzsvQRM9Rpu" - "JPWmnwz1aRPU9SicOPqWLGz5zOyI+5+whPmOjGj5Yrww+dlXxPe/kvj3KOYQ9/Y/L" - "PY3JhD3NJuE8JBo1vO4sSb2WM669RU7wvYUGFL5QhCm+YrM3vi73Pb6TCjy+zAIy" - "volOIL4xsQe+wXTSvZp0jL3WzAC9M+fmOy5IOT1X26Y9OeHpPQ1pET7CnCc+w5Y2" - "PsKxPT5Znzw+WWszPlZ7Ij46igo+Z0DZPeYOlD0=" -) -_VISION_OUTPUT_BASE64 = ( - "zZB1v4gVxT0AaYs/uXavP7LLQT9v+7O+komdv+wxp7+XQAe/6k8YP4c7qj8+IJk/" - "3eWPPvtwUb+6DLG/EbyFvw==" -) -_MULTIMODAL_LOGITS_BASE64 = ( - "qvilPdqoTz1q2ZQ8yRJ4vMTJQ70Rj6C9X03YvVFcA75J6BS+9QggvohDJL5WaSG+" - "2ZkXvldBB75jKOK9Rg6svdUSXb3KArG8+II/PI8yNj3RVJo9O+nSPQkzAT6mXxM+" - "5jEfPlAnJD4rCSI+2u4YPtA8CT6FQOc9Xg+yPW1iaj2gtbM9feBrPUtVzDxJyg+8L" - "oMsvfvRlr314dC9ZPgAvobwE74FiCC+9zMmvrq1JL7NHRy+BssMvvbM7r2RuLm9N" - "0V5vdO06Lwwo6w7IJUePdtfkD3pO8s9IVX9PU5BEj4Hih8+H/IlPuQyJT6RVB0+/" - "a0OPrPB8z1Zpb89HEaDPWJ9mrsLzzW63QhcO40E7jsF4TE8vRRlPDI0hzz1CZY85G" - "eePMHxnzyGlpo8P5GOPGzNeDyMvkk8NP0RPNjhpzvLMZI6obhAu5kd4btQFSy8yjB" - "gvBFRhbxxypS85dmdvFEboLzVdZu8wRyQvIoafbxmEk+8Gx0YvGoytbtTasm6Marc" - "PQTmjD2lO9w89J6EvBnhb725jcm9Tz0JvjjJJ758GT++1ixOvuZcVL5jZVG+CmdF" - "vjjmML4uxRS+uHTkveuFlb0FSgC9Hf4/PHI3Xj18bcE9VrQFPmfuJD5TDD0++ANN" - "PiAlVD4cIVI+Kw5HPoJmMz4JAxg+ACTsPfoTnj3eefM9VzyZPQiP4TxLjKu8rVCM" - "vcOx5737ihy+hn0+vto5WL78o2i+wQZvvrUba74TDl2+7XhFvo1gJb45Tvy9+/qi" - "vT5CBb0RYII8gGmCPb6b3j2zmhg+alM7Ptr4VT73ZGc+eNduPjH+az6P+F4+S1ZI" - "PigRKT49ggI+LKasPYLq2z3AyI09JynmPN94abz7TGW9myzDvWukBb607yO+1Ck7" - "vnhSSr5ZwlC+azJOvvS+Qr5U5i6+koMTvsSJ471cPZa94u0EvYhnITxr5FM9QSa7" - "PX4kAj6oGSE+7xw5PmElST4HglA+p+FOPi1WRD4AVDE+46wWPuIN6z0NoJ49" -) - - -def _decode_float32(encoded: str, shape: tuple[int, ...]) -> torch.Tensor: - values = torch.frombuffer( - bytearray(base64.b64decode(encoded)), - dtype=torch.float32, - ) - return values.reshape(shape) - - -def _fill_reference_parameters(model: torch.nn.Module) -> None: - with torch.no_grad(): - for parameter_index, (name, parameter) in enumerate(model.named_parameters()): - values = torch.arange( - parameter.numel(), - dtype=torch.float32, - ).reshape(parameter.shape) - values = torch.sin(values * 0.013 + parameter_index * 0.17) * 0.02 - if name.endswith("norm.weight"): - values = values + 1.0 - parameter.copy_(values.to(parameter.dtype)) - - for module in model.modules(): - expert_bias = getattr(module, "expert_bias_E", None) - if expert_bias is not None: - expert_bias.copy_(torch.linspace(-0.01, 0.01, expert_bias.numel())) - - -class TestKimiK3HuggingFaceParity(unittest.TestCase): - @torch.no_grad() - def test_reduced_model_matches_released_huggingface_reference(self): - config = _small_model_config() - model = config.build() - model.init_states() - _fill_reference_parameters(model) - model.eval() - - router_ids: list[torch.Tensor] = [] - moe_layer = next( - layer for layer in model.layers.values() if layer.moe is not None - ) - router_hook = moe_layer.moe.router.register_forward_hook( - lambda _module, _inputs, output: router_ids.append( - output[1] if len(output) == 3 else output[0] - ) - ) - self.addCleanup(router_hook.remove) - - text_tokens_BL = torch.tensor([[1, 2, 3, 4]], dtype=torch.long) - text_logits_BLV = model(text_tokens_BL) - torch.testing.assert_close( - text_logits_BLV, - _decode_float32(_TEXT_LOGITS_BASE64, (1, 4, 32)), - atol=2e-4, - rtol=2e-4, - ) - torch.testing.assert_close( - router_ids.pop().reshape(-1), - torch.ones(4, dtype=torch.int64), - atol=0, - rtol=0, - ) - - patches_PCHW = torch.sin(torch.arange(48, dtype=torch.float32) * 0.07).reshape( - 4, 3, 2, 2 - ) - pixels_NPK = patches_PCHW.reshape(1, 4, 12) - grid_thw_N3 = torch.tensor([[1, 2, 2]], dtype=torch.long) - vision_output_NMD = model.vision_encoder( - pixels_NPK, - grid_thw=grid_thw_N3, - ) - torch.testing.assert_close( - vision_output_NMD, - _decode_float32(_VISION_OUTPUT_BASE64, (1, 1, 16)), - atol=2e-4, - rtol=2e-4, - ) - - multimodal_tokens_BL = torch.tensor( - [[1, 2, 7, 3, 4, 5]], - dtype=torch.long, - ) - multimodal_logits_BLV = model( - multimodal_tokens_BL, - pixel_values=pixels_NPK, - grid_thw=grid_thw_N3, - special_tokens={"image_id": 7}, - ) - torch.testing.assert_close( - multimodal_logits_BLV, - _decode_float32(_MULTIMODAL_LOGITS_BASE64, (1, 6, 32)), - atol=2e-4, - rtol=2e-4, - ) - torch.testing.assert_close( - router_ids.pop().reshape(-1), - torch.ones(6, dtype=torch.int64), - atol=0, - rtol=0, - ) - - -if __name__ == "__main__": - from torch.testing._internal.common_utils import run_tests - - run_tests() From 762f4f48140b592b63671fb3d2d3a300c8091b1a Mon Sep 17 00:00:00 2001 From: yangting <15258668233@163.com> Date: Tue, 11 Aug 2026 14:24:05 +0800 Subject: [PATCH 22/67] connect FLA operator, replace RMSNorm --- torchtitan/models/kimi_k3/__init__.py | 14 +++-- torchtitan/models/kimi_k3/model.py | 78 ++++++++++++++++++++------- 2 files changed, 64 insertions(+), 28 deletions(-) diff --git a/torchtitan/models/kimi_k3/__init__.py b/torchtitan/models/kimi_k3/__init__.py index 014c736dc7..23c7a919e3 100644 --- a/torchtitan/models/kimi_k3/__init__.py +++ b/torchtitan/models/kimi_k3/__init__.py @@ -13,7 +13,7 @@ import torch.nn as nn from torchtitan.components.optimizer import register_moe_load_balancing_hook -from torchtitan.models.common import Conv1d, Embedding, Linear +from torchtitan.models.common import Embedding, Linear from torchtitan.models.common.moe import TokenChoiceTopKRouter from torchtitan.models.common.nn_modules import GELU, RMSNorm from torchtitan.models.common.token_dispatcher import LocalTokenDispatcher @@ -32,6 +32,7 @@ KimiLatentMoE, KimiMLAAttention, KimiRMSNormGated, + KimiShortConvolution, ) from .parallelize import parallelize_kimi_k3 from .state_dict_adapter import KimiK3StateDictAdapter @@ -180,13 +181,11 @@ def _kda_config( ) -> KimiDeltaAttention.Config: projection_dim = num_heads * head_dim - def conv() -> Conv1d.Config: - return Conv1d.Config( - in_channels=projection_dim, - out_channels=projection_dim, + def conv() -> KimiShortConvolution.Config: + return KimiShortConvolution.Config( + hidden_size=projection_dim, kernel_size=conv_kernel_size, - groups=projection_dim, - bias=False, + activation="silu", param_init=_CONV_INIT, ) @@ -194,7 +193,6 @@ def conv() -> Conv1d.Config: dim=dim, num_heads=num_heads, head_dim=head_dim, - conv_kernel_size=conv_kernel_size, q_proj=_linear(dim, projection_dim), k_proj=_linear(dim, projection_dim), v_proj=_linear(dim, projection_dim), diff --git a/torchtitan/models/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py index 3d337d28b4..3c40daf681 100644 --- a/torchtitan/models/kimi_k3/model.py +++ b/torchtitan/models/kimi_k3/model.py @@ -18,11 +18,14 @@ import torch import torch.nn.functional as F +from fla.modules import ShortConvolution +from fla.modules.fused_norm_gate import rms_norm_gated +from fla.modules.conv.causal_conv1d import causal_conv1d from fla.ops.kda import chunk_kda from torch import nn from torch.distributed.tensor import DTensor -from torchtitan.models.common import Conv1d, Linear +from torchtitan.models.common import Linear from torchtitan.models.common.attention import ( AttentionMasksType, BaseAttention, @@ -47,9 +50,48 @@ # T = flattened tokens, N = attention-residual entries. +class KimiShortConvolution(ShortConvolution, Module): + """KDA short causal convolution backed by FLA's fused kernel. + + Matches the released Kimi K3 HuggingFace model, which builds FLA's + ``ShortConvolution`` per q/k/v projection. The Triton kernel runs only on + accelerator devices. + """ + + @dataclass(kw_only=True, slots=True) + class Config(Module.Config): + hidden_size: int + kernel_size: int + activation: str = "silu" + + def __init__(self, config: Config): + super().__init__( + hidden_size=config.hidden_size, + kernel_size=config.kernel_size, + activation=config.activation, + ) + + def forward( + self, + x_BLD: torch.Tensor, + **kwargs: object, + ) -> tuple[torch.Tensor, None]: + y_BLD, _ = causal_conv1d( + x=x_BLD, + weight=self.weight.squeeze(1), + activation=self.activation, + backend=self.backend, + ) + return y_BLD, None class KimiRMSNormGated(Module): - """Per-head RMSNorm followed by a sigmoid output gate.""" + """Per-head RMSNorm + sigmoid output gate backed by FLA's fused kernel. + + Wraps the FLA functional ``rms_norm_gated`` behind the torchtitan + ``Module`` protocol so it participates in ``init_states``/``param_init``. + The wrapper owns the weight (ones-initialized, weight-only checkpoint + schema unchanged); the fused kernel receives it with ``bias=None``. + """ @dataclass(kw_only=True, slots=True) class Config(Module.Config): @@ -62,12 +104,14 @@ def __init__(self, config: Config): self.weight = nn.Parameter(torch.empty(config.dim)) def forward(self, x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: - input_dtype = x.dtype - x_float = x.float() - variance = x_float.pow(2).mean(dim=-1, keepdim=True) - x_float = x_float * torch.rsqrt(variance + self.eps) - x_float = self.weight.float() * x_float - return (x_float * torch.sigmoid(gate.float())).to(input_dtype) + return rms_norm_gated( + x, + gate, + self.weight, + None, + activation="sigmoid", + eps=self.eps, + ) class KimiFeedForward(FeedForward): @@ -259,13 +303,12 @@ class Config(Module.Config): dim: int num_heads: int head_dim: int - conv_kernel_size: int q_proj: Linear.Config k_proj: Linear.Config v_proj: Linear.Config - q_conv: Conv1d.Config - k_conv: Conv1d.Config - v_conv: Conv1d.Config + q_conv: KimiShortConvolution.Config + k_conv: KimiShortConvolution.Config + v_conv: KimiShortConvolution.Config forget_a: Linear.Config forget_b: Linear.Config beta: Linear.Config @@ -281,7 +324,6 @@ def __init__(self, config: Config): super().__init__() self.num_heads = config.num_heads self.head_dim = config.head_dim - self.conv_kernel_size = config.conv_kernel_size self.q_proj = config.q_proj.build() self.k_proj = config.k_proj.build() @@ -300,10 +342,6 @@ def __init__(self, config: Config): self.A_log = nn.Parameter(torch.empty(config.num_heads)) self.dt_bias = nn.Parameter(torch.empty(config.num_heads, config.head_dim)) - def _causal_conv(self, x_BLC: torch.Tensor, conv: Conv1d) -> torch.Tensor: - x_BCL = F.pad(x_BLC.transpose(1, 2), (self.conv_kernel_size - 1, 0)) - return F.silu(conv(x_BCL)).transpose(1, 2) - def forward( self, x_BLD: torch.Tensor, @@ -317,13 +355,13 @@ def forward( ) B, L, _ = x_BLD.shape - q_BLHK = self._causal_conv(self.q_proj(x_BLD), self.q_conv).view( + q_BLHK = self.q_conv(self.q_proj(x_BLD))[0].view( B, L, self.num_heads, self.head_dim ) - k_BLHK = self._causal_conv(self.k_proj(x_BLD), self.k_conv).view( + k_BLHK = self.k_conv(self.k_proj(x_BLD))[0].view( B, L, self.num_heads, self.head_dim ) - v_BLHV = self._causal_conv(self.v_proj(x_BLD), self.v_conv).view( + v_BLHV = self.v_conv(self.v_proj(x_BLD))[0].view( B, L, self.num_heads, self.head_dim ) forget_BLHK = self.forget_b(self.forget_a(x_BLD)).view( From 4f972ba809b17abbcd61533ea6a4c3e05131a804 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Tue, 11 Aug 2026 09:12:29 +0000 Subject: [PATCH 23/67] resotre to eager conv and RMSNormGated --- torchtitan/models/kimi_k3/__init__.py | 14 ++++--- torchtitan/models/kimi_k3/model.py | 57 ++++++++++++--------------- 2 files changed, 34 insertions(+), 37 deletions(-) diff --git a/torchtitan/models/kimi_k3/__init__.py b/torchtitan/models/kimi_k3/__init__.py index 23c7a919e3..014c736dc7 100644 --- a/torchtitan/models/kimi_k3/__init__.py +++ b/torchtitan/models/kimi_k3/__init__.py @@ -13,7 +13,7 @@ import torch.nn as nn from torchtitan.components.optimizer import register_moe_load_balancing_hook -from torchtitan.models.common import Embedding, Linear +from torchtitan.models.common import Conv1d, Embedding, Linear from torchtitan.models.common.moe import TokenChoiceTopKRouter from torchtitan.models.common.nn_modules import GELU, RMSNorm from torchtitan.models.common.token_dispatcher import LocalTokenDispatcher @@ -32,7 +32,6 @@ KimiLatentMoE, KimiMLAAttention, KimiRMSNormGated, - KimiShortConvolution, ) from .parallelize import parallelize_kimi_k3 from .state_dict_adapter import KimiK3StateDictAdapter @@ -181,11 +180,13 @@ def _kda_config( ) -> KimiDeltaAttention.Config: projection_dim = num_heads * head_dim - def conv() -> KimiShortConvolution.Config: - return KimiShortConvolution.Config( - hidden_size=projection_dim, + def conv() -> Conv1d.Config: + return Conv1d.Config( + in_channels=projection_dim, + out_channels=projection_dim, kernel_size=conv_kernel_size, - activation="silu", + groups=projection_dim, + bias=False, param_init=_CONV_INIT, ) @@ -193,6 +194,7 @@ def conv() -> KimiShortConvolution.Config: dim=dim, num_heads=num_heads, head_dim=head_dim, + conv_kernel_size=conv_kernel_size, q_proj=_linear(dim, projection_dim), k_proj=_linear(dim, projection_dim), v_proj=_linear(dim, projection_dim), diff --git a/torchtitan/models/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py index 3c40daf681..75035f38af 100644 --- a/torchtitan/models/kimi_k3/model.py +++ b/torchtitan/models/kimi_k3/model.py @@ -18,14 +18,11 @@ import torch import torch.nn.functional as F -from fla.modules import ShortConvolution -from fla.modules.fused_norm_gate import rms_norm_gated -from fla.modules.conv.causal_conv1d import causal_conv1d from fla.ops.kda import chunk_kda from torch import nn from torch.distributed.tensor import DTensor -from torchtitan.models.common import Linear +from torchtitan.models.common import Conv1d, Linear from torchtitan.models.common.attention import ( AttentionMasksType, BaseAttention, @@ -85,13 +82,7 @@ def forward( return y_BLD, None class KimiRMSNormGated(Module): - """Per-head RMSNorm + sigmoid output gate backed by FLA's fused kernel. - - Wraps the FLA functional ``rms_norm_gated`` behind the torchtitan - ``Module`` protocol so it participates in ``init_states``/``param_init``. - The wrapper owns the weight (ones-initialized, weight-only checkpoint - schema unchanged); the fused kernel receives it with ``bias=None``. - """ + """Per-head RMSNorm followed by a sigmoid output gate.""" @dataclass(kw_only=True, slots=True) class Config(Module.Config): @@ -104,14 +95,12 @@ def __init__(self, config: Config): self.weight = nn.Parameter(torch.empty(config.dim)) def forward(self, x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: - return rms_norm_gated( - x, - gate, - self.weight, - None, - activation="sigmoid", - eps=self.eps, - ) + input_dtype = x.dtype + x_float = x.float() + variance = x_float.pow(2).mean(dim=-1, keepdim=True) + x_float = x_float * torch.rsqrt(variance + self.eps) + x_float = self.weight.float() * x_float + return (x_float * torch.sigmoid(gate.float())).to(input_dtype) class KimiFeedForward(FeedForward): @@ -303,12 +292,13 @@ class Config(Module.Config): dim: int num_heads: int head_dim: int + conv_kernel_size: int q_proj: Linear.Config k_proj: Linear.Config v_proj: Linear.Config - q_conv: KimiShortConvolution.Config - k_conv: KimiShortConvolution.Config - v_conv: KimiShortConvolution.Config + q_conv: Conv1d.Config + k_conv: Conv1d.Config + v_conv: Conv1d.Config forget_a: Linear.Config forget_b: Linear.Config beta: Linear.Config @@ -324,6 +314,7 @@ def __init__(self, config: Config): super().__init__() self.num_heads = config.num_heads self.head_dim = config.head_dim + self.conv_kernel_size = config.conv_kernel_size self.q_proj = config.q_proj.build() self.k_proj = config.k_proj.build() @@ -342,6 +333,10 @@ def __init__(self, config: Config): self.A_log = nn.Parameter(torch.empty(config.num_heads)) self.dt_bias = nn.Parameter(torch.empty(config.num_heads, config.head_dim)) + def _causal_conv(self, x_BLC: torch.Tensor, conv: Conv1d) -> torch.Tensor: + x_BCL = F.pad(x_BLC.transpose(1, 2), (self.conv_kernel_size - 1, 0)) + return F.silu(conv(x_BCL)).transpose(1, 2) + def forward( self, x_BLD: torch.Tensor, @@ -355,13 +350,13 @@ def forward( ) B, L, _ = x_BLD.shape - q_BLHK = self.q_conv(self.q_proj(x_BLD))[0].view( + q_BLHK = self._causal_conv(self.q_proj(x_BLD), self.q_conv).view( B, L, self.num_heads, self.head_dim ) - k_BLHK = self.k_conv(self.k_proj(x_BLD))[0].view( + k_BLHK = self._causal_conv(self.k_proj(x_BLD), self.k_conv).view( B, L, self.num_heads, self.head_dim ) - v_BLHV = self.v_conv(self.v_proj(x_BLD))[0].view( + v_BLHV = self._causal_conv(self.v_proj(x_BLD), self.v_conv).view( B, L, self.num_heads, self.head_dim ) forget_BLHK = self.forget_b(self.forget_a(x_BLD)).view( @@ -519,18 +514,18 @@ def forward(self, x_BLD: torch.Tensor) -> torch.Tensor: # Every expert's weights are packed into one grouped-mm call, so an # expert with zero assigned tokens still sits in the autograd graph # and receives an all-zero gradient rather than None. - routed_input_RD, num_tokens_per_expert_E, metadata = ( - self.token_dispatcher.dispatch( - latent_TD, weights_TK, expert_ids_TK, num_tokens_per_expert_E - ) + ( + routed_input_RD, + num_tokens_per_expert_E, + metadata, + ) = self.token_dispatcher.dispatch( + latent_TD, weights_TK, expert_ids_TK, num_tokens_per_expert_E ) routed_output_RD = self.routed_experts(routed_input_RD, num_tokens_per_expert_E) routed_TD = self.token_dispatcher.combine( routed_output_RD, metadata, latent_TD, - num_local_tokens_after_padding=T, - local_seq_len_after_padding=L, ) routed_BLD = self.routed_up(self.routed_norm(routed_TD.view(B, L, latent_dim))) From dc009a9142bc485ecb3856f57e2c30b238faaf6b Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Tue, 11 Aug 2026 09:42:41 +0000 Subject: [PATCH 24/67] remove KimiExactGELU, replaced with nn.GELU --- tests/unit_tests/test_kimi_k3.py | 11 -------- torchtitan/models/kimi_k3/__init__.py | 3 +-- torchtitan/models/kimi_k3/vision_encoder.py | 28 ++------------------- 3 files changed, 3 insertions(+), 39 deletions(-) diff --git a/tests/unit_tests/test_kimi_k3.py b/tests/unit_tests/test_kimi_k3.py index 1f38a3b221..d6a2c78a8b 100644 --- a/tests/unit_tests/test_kimi_k3.py +++ b/tests/unit_tests/test_kimi_k3.py @@ -34,7 +34,6 @@ KimiKDAKernel, ) from torchtitan.models.kimi_k3.state_dict_adapter import KimiK3StateDictAdapter - from torchtitan.models.kimi_k3.vision_encoder import KimiExactGELU except ModuleNotFoundError as exc: raise unittest.SkipTest( f"Kimi K3 optional dependency unavailable: {exc.name}" @@ -266,16 +265,6 @@ def _kda_recurrent_reference( class TestKimiK3(unittest.TestCase): - def test_exact_gelu_matches_pytorch_reference(self): - x = torch.linspace(-4.0, 4.0, 257) - actual = KimiExactGELU.Config().build()(x) - expected = F.gelu(x, approximate="none") - - torch.testing.assert_close(actual, expected, atol=1e-6, rtol=1e-6) - - x_bf16 = x.bfloat16() - self.assertEqual(KimiExactGELU.Config().build()(x_bf16).dtype, x_bf16.dtype) - @unittest.skipIf(not torch.cuda.is_available(), "FLA KDA kernel requires CUDA.") def test_fla_kda_kernel_matches_recurrent_reference(self): torch.manual_seed(1) diff --git a/torchtitan/models/kimi_k3/__init__.py b/torchtitan/models/kimi_k3/__init__.py index 014c736dc7..82f2ac65ed 100644 --- a/torchtitan/models/kimi_k3/__init__.py +++ b/torchtitan/models/kimi_k3/__init__.py @@ -36,7 +36,6 @@ from .parallelize import parallelize_kimi_k3 from .state_dict_adapter import KimiK3StateDictAdapter from .vision_encoder import ( - KimiExactGELU, KimiK3VisionAttention, KimiK3VisionBlock, KimiK3VisionEncoder, @@ -347,7 +346,7 @@ def _vision_encoder_config( eps=1e-5, param_init=_NORM_INIT, ), - activation=KimiExactGELU.Config(), + activation=GELU.Config(), ), param_init=_POS_EMBED_INIT, ) diff --git a/torchtitan/models/kimi_k3/vision_encoder.py b/torchtitan/models/kimi_k3/vision_encoder.py index 7192350c98..4f3df19873 100644 --- a/torchtitan/models/kimi_k3/vision_encoder.py +++ b/torchtitan/models/kimi_k3/vision_encoder.py @@ -19,41 +19,17 @@ - M = maximum merged tokens per item (padded) """ -import math from dataclasses import dataclass, field import torch import torch.nn.functional as F from torchtitan.models.common import Linear -from torchtitan.models.common.nn_modules import RMSNorm +from torchtitan.models.common.nn_modules import GELU, RMSNorm from torchtitan.models.common.vision_encoder import VisionMLP from torchtitan.protocols.module import Module, ModuleDict -class KimiExactGELU(Module): - """Exact GELU used by the released Kimi vision projector. - - The explicit FP32 form is mathematically equivalent to - ``nn.GELU(approximate="none")`` while avoiding device-specific fused - approximations in the numerical reference path. - """ - - @dataclass(kw_only=True, slots=True) - class Config(Module.Config): - pass - - def __init__(self, config: Config): - super().__init__() - - def forward(self, x: torch.Tensor) -> torch.Tensor: - input_dtype = x.dtype - x_float = x.float() - return (0.5 * x_float * (1.0 + torch.erf(x_float / math.sqrt(2.0)))).to( - input_dtype - ) - - def _get_temporal_pos_embed( num_frames: int, embed_dim: int, @@ -357,7 +333,7 @@ class Config(Module.Config): linear_1: Linear.Config linear_2: Linear.Config post_norm: RMSNorm.Config - activation: KimiExactGELU.Config = field(default_factory=KimiExactGELU.Config) + activation: GELU.Config = field(default_factory=GELU.Config) def __init__(self, config: Config): super().__init__() From d37a4a58ff1dd50f3416af008bc372b85979597c Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Thu, 13 Aug 2026 10:03:21 +0000 Subject: [PATCH 25/67] Add Kimi K3 numerical validation --- .../numerical_tests_kimi_k3.py | 471 ++++++++++++++++++ torchtitan/models/kimi_k3/README.md | 333 ++++--------- 2 files changed, 556 insertions(+), 248 deletions(-) create mode 100644 scripts/checkpoint_conversion/numerical_tests_kimi_k3.py diff --git a/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py b/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py new file mode 100644 index 0000000000..9ad9d7e3d0 --- /dev/null +++ b/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py @@ -0,0 +1,471 @@ +#!/usr/bin/env python3 +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Full text+image logit parity: TorchTitan Kimi K3 vs HuggingFace Kimi K3. + +Runs the released HuggingFace model code and TorchTitan in one process on the +same text+image prompt. Each side performs its own image preprocessing, so the +comparison covers preprocessing, vision, projector, scatter, and decoder. + +The released checkpoint is MXFP4-quantized and is not loaded. Instead, the +script reduces the HuggingFace config to TorchTitan's debug model, initializes +TorchTitan, and strictly transfers its state dict to HuggingFace. + +The local HuggingFace directory must contain the config, modeling, processor, +tokenizer code, and tokenizer assets. The released code requires +``transformers==4.56.2`` and ``tiktoken``. + +Usage: + CUDA_VISIBLE_DEVICES=0 python -m \ + scripts.checkpoint_conversion.numerical_tests_kimi_k3 \ + --hf_model_path ~/hf_assets/moonshotai/Kimi-K3 --dtype float32 +""" + +import argparse +import os +from typing import Any, cast + +import torch +import torch.nn.functional as F +from PIL import Image + +from torchtitan.hf_datasets.multimodal.utils.image import ( + process_image, + resize_to_patch_budget, + vision_to_patches, +) +from torchtitan.models.kimi_k3 import model_registry +from torchtitan.models.kimi_k3.model import KimiK3Model +from torchtitan.models.kimi_k3.state_dict_adapter import KimiK3StateDictAdapter +from transformers import AutoConfig, AutoModelForCausalLM, AutoProcessor + + +_MEDIA_TOKEN_ID = 163605 +_PATCH_SIZE = 14 +_MERGE_SIZE = 2 +_MAX_PATCHES = 65536 +_MAX_PATCHES_PER_SIDE = 512 +_PROMPT = ( + "<|kimi_image_placeholder|>\n" "What is shown in this image? Describe it briefly." +) + + +def _reduce_hf_config(hf_config, tt_config, hf_model_path: str) -> None: + """Reduce the released HuggingFace config to the TorchTitan debug model.""" + text_config = hf_config.text_config + full_attention_layers = [ + layer_idx + 1 + for layer_idx, layer in enumerate(tt_config.layers) + if layer.attention is not None + ] + kda_layers = [ + layer_idx + 1 + for layer_idx, layer in enumerate(tt_config.layers) + if layer.delta_attention is not None + ] + mla = next(layer.attention for layer in tt_config.layers if layer.attention) + kda = next( + layer.delta_attention for layer in tt_config.layers if layer.delta_attention + ) + dense_ffn = next( + layer.feed_forward for layer in tt_config.layers if layer.feed_forward + ) + moe = next(layer.moe for layer in tt_config.layers if layer.moe) + + text_overrides = { + "vocab_size": tt_config.vocab_size, + "hidden_size": tt_config.dim, + "intermediate_size": dense_ffn.w1.out_features, + "num_hidden_layers": len(tt_config.layers), + "num_attention_heads": mla.n_heads, + "num_key_value_heads": mla.n_heads, + "rms_norm_eps": tt_config.norm.eps, + "q_lora_rank": mla.wq_a.out_features, + "kv_lora_rank": mla.kv_lora_rank, + "qk_nope_head_dim": mla.qk_nope_head_dim, + "qk_rope_head_dim": mla.qk_rope_head_dim, + "v_head_dim": mla.v_head_dim, + "activation_situ_beta": dense_ffn.beta, + "activation_situ_linear_beta": dense_ffn.linear_beta, + "num_experts": moe.num_experts, + "num_experts_per_token": moe.router.top_k, + "num_shared_experts": ( + moe.shared_experts.w1.out_features // moe.routed_experts.hidden_dim + ), + "moe_renormalize": moe.router.route_norm, + "moe_intermediate_size": moe.routed_experts.hidden_dim, + "routed_expert_hidden_size": moe.routed_down.out_features, + "routed_scaling_factor": moe.router.route_scale, + "first_k_dense_replace": next( + layer_idx for layer_idx, layer in enumerate(tt_config.layers) if layer.moe + ), + "attn_res_block_size": tt_config.layers[0].attn_res_block_size, + "linear_attn_config": { + "full_attn_layers": full_attention_layers, + "kda_layers": kda_layers, + "head_dim": kda.head_dim, + "num_heads": kda.num_heads, + "short_conv_kernel_size": kda.conv_kernel_size, + "gate_lower_bound": kda.kernel.lower_bound, + "use_full_rank_gate": True, + }, + } + for name, value in text_overrides.items(): + setattr(text_config, name, value) + + vision = tt_config.vision_encoder + assert vision is not None + vision_overrides = { + "patch_size": vision.patch_size, + "init_pos_emb_height": vision.init_pos_emb_height, + "init_pos_emb_width": vision.init_pos_emb_width, + "init_pos_emb_time": vision.max_num_frames, + "vt_num_attention_heads": vision.block.attn.num_heads, + "vt_num_hidden_layers": vision.num_layers, + "vt_hidden_size": vision.dim, + "vt_intermediate_size": vision.block.mlp.fc1.out_features, + "merge_kernel_size": vision.merge_kernel_size, + "mm_hidden_size": vision.dim, + "qkv_hidden_size": vision.block.attn.qkv_dim, + "text_hidden_size": tt_config.dim, + "pos_emb_interpolation_mode": vision.interpolation_mode, + } + for name, value in vision_overrides.items(): + setattr(hf_config.vision_config, name, value) + + for config in (hf_config, text_config): + if hasattr(config, "quantization_config"): + delattr(config, "quantization_config") + text_config._attn_implementation = "eager" + hf_config.vision_config._attn_implementation = "eager" + text_config._name_or_path = hf_model_path + hf_config._name_or_path = hf_model_path + + +def _build_tt_model(tt_config, dtype: torch.dtype) -> KimiK3Model: + with torch.device("meta"): + model = tt_config.build() + model.to_empty(device="cpu") + model.to(dtype=dtype) + model.init_states(buffer_device=torch.device("cpu")) + return model.eval() + + +def _build_hf_model( + hf_model_path: str, + tt_config, + hf_state_dict: dict[str, Any], + dtype: torch.dtype, +): + hf_config = AutoConfig.from_pretrained( + hf_model_path, + trust_remote_code=True, + local_files_only=True, + ) + _reduce_hf_config(hf_config, tt_config, hf_model_path) + model = AutoModelForCausalLM.from_config(hf_config, trust_remote_code=True) + model.language_model.config._attn_implementation = "eager" + model.to(dtype=dtype) + model.load_state_dict(hf_state_dict, strict=True) + return model.eval() + + +@torch.no_grad() +def run_hf( + hf_model_path: str, + tt_config, + hf_state_dict: dict[str, Any], + image_size: int, + dtype: torch.dtype, + device: torch.device, +) -> dict[str, Any]: + """Run HuggingFace preprocessing and the reduced HuggingFace model.""" + print(f"Loading released HuggingFace Kimi K3 code on {device} ...") + processor: Any = AutoProcessor.from_pretrained( + hf_model_path, + trust_remote_code=True, + local_files_only=True, + ) + model = _build_hf_model(hf_model_path, tt_config, hf_state_dict, dtype).to(device) + + raw_image = ( + torch.linspace(0, 255, image_size * image_size * 3) + .reshape(image_size, image_size, 3) + .to(torch.uint8) + ) + pil_image = Image.fromarray(raw_image.numpy()) + batch = processor( + medias=[{"type": "image", "image": pil_image}], # codespell:ignore medias + text=_PROMPT, + return_tensors="pt", + ) + + vision_features: dict[str, torch.Tensor] = {} + + def record_vision_features(_module, _inputs, output) -> None: + features = output[0] if isinstance(output, (list, tuple)) else output + vision_features["output"] = features.detach().float().cpu() + + model.mm_projector.register_forward_hook(record_vision_features) + + expert_indices: dict[int, torch.Tensor] = {} + for layer_idx, layer in enumerate(model.language_model.model.layers): + moe = getattr(layer, "block_sparse_moe", None) + if moe is not None: + moe.gate.register_forward_hook( + lambda _module, _inputs, output, layer_idx=layer_idx: ( + expert_indices.__setitem__( + layer_idx, + output[0].detach().cpu(), + ) + ) + ) + + inputs = { + key: value.to(device) if isinstance(value, torch.Tensor) else value + for key, value in batch.items() + } + output = model(**inputs, use_cache=False) + ref = { + "input_ids": batch["input_ids"].cpu(), + "raw_image": raw_image, + "last_logits": output.logits[:, -1, :].float().cpu(), + "pixel_values": batch["pixel_values"].float().cpu(), + "grid_thws": batch["grid_thws"].cpu(), + "vision_features": vision_features["output"], + "expert_indices": expert_indices, + } + del model + torch.cuda.empty_cache() + return ref + + +def _expand_image_placeholder( + input_ids: torch.Tensor, + image_token_id: int, + num_vision_tokens: int, +) -> torch.Tensor: + """Expand the single HF media placeholder for TorchTitan's scatter path.""" + if input_ids.shape[0] != 1: + raise ValueError("The Kimi K3 numerical test expects a batch size of one.") + positions = (input_ids[0] == image_token_id).nonzero().flatten() + if positions.numel() != 1: + raise ValueError(f"Expected one image placeholder, found {positions.numel()}.") + position = positions.item() + image_tokens = input_ids.new_full((1, num_vision_tokens), image_token_id) + return torch.cat( + (input_ids[:, :position], image_tokens, input_ids[:, position + 1 :]), + dim=1, + ) + + +def _print_routing_comparison( + hf_expert_indices: dict[int, torch.Tensor], + tt_expert_indices: dict[int, torch.Tensor], +) -> None: + hf_layers = set(hf_expert_indices) + tt_layers = set(tt_expert_indices) + if not hf_layers: + raise ValueError("No MoE routing choices were recorded.") + if hf_layers != tt_layers: + raise ValueError( + "Routing layers differ: " + f"HF-only {sorted(hf_layers - tt_layers)}, " + f"TorchTitan-only {sorted(tt_layers - hf_layers)}." + ) + + num_matching = 0 + num_routings = 0 + for layer_idx in sorted(hf_layers): + top_k = hf_expert_indices[layer_idx].shape[-1] + hf_ids = hf_expert_indices[layer_idx].reshape(-1, top_k).sort(dim=-1).values + tt_ids = tt_expert_indices[layer_idx].reshape(-1, top_k).sort(dim=-1).values + if hf_ids.shape != tt_ids.shape: + raise ValueError( + f"Layer {layer_idx} routing shapes differ: " + f"HF {tuple(hf_ids.shape)} vs TT {tuple(tt_ids.shape)}." + ) + num_matching += int((hf_ids == tt_ids).sum().item()) + num_routings += hf_ids.numel() + match_rate = num_matching / num_routings if num_routings else 0.0 + print(f"router choices: {num_matching}/{num_routings} match " f"({match_rate:.1%})") + + +@torch.no_grad() +def run_tt( + model: KimiK3Model, + ref: dict[str, Any], + vision_dtype: torch.dtype, + device: torch.device, +) -> torch.Tensor: + """Run TorchTitan preprocessing and the reduced TorchTitan model.""" + print(f"Loading TorchTitan Kimi K3 (debugmodel) on {device} ...") + model.to(device) + assert model.vision_encoder is not None + model.vision_encoder.to(vision_dtype) + + expert_indices: dict[int, torch.Tensor] = {} + for layer_idx, layer in model.layers.items(): + if layer.moe is not None: + # pyrefly: ignore [missing-attribute] + layer.moe.router.register_forward_hook( + lambda _module, _inputs, output, layer_idx=int(layer_idx): ( + expert_indices.__setitem__( + layer_idx, + output[1].detach().cpu(), + ) + ) + ) + + image = process_image( + Image.fromarray(ref["raw_image"].numpy()), + patch_size=_PATCH_SIZE, + merge_size=_MERGE_SIZE, + resize_fn=resize_to_patch_budget, + max_patches=_MAX_PATCHES, + max_patches_per_side=_MAX_PATCHES_PER_SIDE, + image_mean=(0.5, 0.5, 0.5), + image_std=(0.5, 0.5, 0.5), + ) + if image is None: + raise ValueError("TorchTitan failed to process the numerical test image.") + patches, grid = vision_to_patches( + image, + patch_size=_PATCH_SIZE, + temporal_patch_size=1, + merge_size=_MERGE_SIZE, + patch_order="raster", + ) + pixel_values = patches.unsqueeze(0).to(device=device, dtype=vision_dtype) + grid_thw = grid.unsqueeze(0).to(device) + num_vision_tokens = (grid[1] // _MERGE_SIZE) * (grid[2] // _MERGE_SIZE) + tokens = _expand_image_placeholder( + ref["input_ids"], + _MEDIA_TOKEN_ID, + int(num_vision_tokens.item()), + ).to(device) + + print( + f"tokens={tuple(tokens.shape)} pixel_values={tuple(pixel_values.shape)} " + f"grid_thw={grid_thw.tolist()} vision_tokens={num_vision_tokens.item()}" + ) + + hf_pixels = ref["pixel_values"].flatten(1) + pixel_diff = (hf_pixels - patches.float()).abs() + print( + f"pixel values: max_diff={pixel_diff.max().item():.3e} " + f"num_differ={(pixel_diff > 1e-6).sum().item()}/{pixel_diff.numel()}" + ) + + tt_features = model.vision_encoder(pixel_values, grid_thw=grid_thw) + hf_features = ref["vision_features"].reshape(-1, tt_features.shape[-1]) + tt_features = tt_features.float().cpu().reshape(-1, tt_features.shape[-1]) + vision_cos = F.cosine_similarity( + hf_features.flatten(), tt_features.flatten(), dim=0 + ).item() + vision_max_diff = (hf_features - tt_features).abs().max().item() + print( + f"vision features: shape={tuple(tt_features.shape)} " + f"cos={vision_cos:.6f} max_diff={vision_max_diff:.3e}" + ) + + logits = model( + tokens, + pixel_values=pixel_values, + grid_thw=grid_thw, + special_tokens={"image_id": _MEDIA_TOKEN_ID}, + ) + _print_routing_comparison(ref["expert_indices"], expert_indices) + return logits[:, -1, :].float().cpu().squeeze() + + +def compare(ref_logits: torch.Tensor, tt_logits: torch.Tensor) -> bool: + """Print last-token metrics and return whether KL is below tolerance.""" + ref = ref_logits.squeeze() + tt = tt_logits.squeeze() + log_ref = F.log_softmax(ref, dim=-1) + log_tt = F.log_softmax(tt, dim=-1) + kl = F.kl_div(log_tt, log_ref, log_target=True, reduction="sum").item() + cosine = F.cosine_similarity(ref, tt, dim=-1).item() + max_diff = (ref - tt).abs().max().item() + top1 = (ref.argmax() == tt.argmax()).item() + top5_overlap = ( + len(set(ref.topk(5).indices.tolist()) & set(tt.topk(5).indices.tolist())) / 5 + ) + print("\nFull multimodal last-token logit parity (TorchTitan vs HuggingFace)") + print( + f"KL={kl:.4e} cos={cosine:.6f} max_diff={max_diff:.4e} " + f"top1={'Y' if top1 else 'N'} top5={top5_overlap:.0%}" + ) + passed = abs(kl) < 1e-3 # pyrefly: ignore [bad-argument-type] + print("RESULT: PASS" if passed else "RESULT: FAIL") + return passed + + +@torch.no_grad() +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument( + "--hf_model_path", + default=os.path.expanduser("~/hf_assets/moonshotai/Kimi-K3"), + help="Local directory containing the Kimi K3 HuggingFace assets.", + ) + parser.add_argument("--model_flavor", default="debugmodel") + parser.add_argument("--image_size", type=int, default=336) + parser.add_argument( + "--hf_dtype", + default="float32", + choices=["float32", "bfloat16", "float16"], + ) + parser.add_argument( + "--dtype", + default="float32", + choices=["float32", "bfloat16", "float16"], + ) + parser.add_argument( + "--vision_dtype", + default=None, + choices=["float32", "bfloat16", "float16"], + ) + parser.add_argument("--seed", type=int, default=42) + args = parser.parse_args() + + if not torch.cuda.is_available(): + parser.error("Kimi K3 numerical parity requires a CUDA GPU.") + + device = torch.device("cuda") + dtype = getattr(torch, args.dtype) + vision_dtype = getattr(torch, args.vision_dtype) if args.vision_dtype else dtype + hf_dtype = getattr(torch, args.hf_dtype) + print( + f"hf_dtype={args.hf_dtype} titan text={args.dtype} " + f"titan vision={args.vision_dtype or args.dtype}" + ) + + tt_config = cast(KimiK3Model.Config, model_registry(args.model_flavor).model) + torch.manual_seed(args.seed) + tt_model = _build_tt_model(tt_config, dtype) + hf_state_dict = KimiK3StateDictAdapter(tt_config, hf_assets_path=None).to_hf( + tt_model.state_dict() + ) + + ref = run_hf( + args.hf_model_path, + tt_config, + hf_state_dict, + args.image_size, + hf_dtype, + device, + ) + del hf_state_dict + tt_logits = run_tt(tt_model, ref, vision_dtype, device) + if not compare(ref["last_logits"], tt_logits): + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/torchtitan/models/kimi_k3/README.md b/torchtitan/models/kimi_k3/README.md index 05e4597976..d4e0385b44 100644 --- a/torchtitan/models/kimi_k3/README.md +++ b/torchtitan/models/kimi_k3/README.md @@ -1,272 +1,109 @@ # Kimi K3 -This directory contains the eager numerical reference implementation of Kimi -K3 in TorchTitan. The initial scope is a topology-complete, reduced model for -single-device or FSDP2 training and numerical comparison with the -[released HuggingFace implementation](https://huggingface.co/moonshotai/Kimi-K3). -It is intended to make architecture experiments and model-structure choices -measurable against a stable, inspectable baseline before optimized kernels and -additional parallelisms are introduced. +Kimi K3 combines a hybrid **Kimi Delta Attention (KDA) + Multi-head Latent +Attention (MLA)** decoder, **LatentMoE**, and a **MoonViT3d** vision encoder. +TorchTitan currently provides a topology-complete reduced model for architecture +validation, single-device training, and FSDP2 training. -Every operator outside the KDA recurrence is plain eager PyTorch, which keeps -the model math directly inspectable. KDA itself runs on FLA's chunked Triton -kernel, following the same split Qwen3.5 uses: the kernel is the training path -and a pure-PyTorch recurrence in the unit tests pins its numerics. +## Prerequisites -## Quick start +Install the additional dependencies: + +```bash +pip install av einops pillow torchvision flash-linear-attention +``` + +## Architecture + +- **Decoder** -- hybrid KDA and MLA layers. The MLA layers follow the released + model's explicit 1-based layer list, including consecutive MLA layers at the + end of the decoder. +- **Feed-forward layers** -- one dense SiTU feed-forward layer followed by + LatentMoE layers with sigmoid top-k routing, correction bias, routed experts, + and shared experts. +- **Attention residuals** -- block-level attention residual connections, + including the final output residual. +- **KDA backend** -- FLA's chunked Triton kernel, with a pure PyTorch recurrent + implementation in the unit tests as the numerical reference. +- **Vision encoder** -- MoonViT3d with learned spatial positions, 2D RoPE, + non-causal attention, temporal pooling, 2x2 spatial merge, and a two-layer + projector to the decoder dimension. +- **Multimodal forward** -- projected vision embeddings are scattered into runs + of the shared media placeholder token. + +## Model variants + +Only `debugmodel` is currently registered. The released Kimi K3 row is included +for architectural comparison and is not a runnable TorchTitan flavor. + +| Variant | Parameters | LLM dim | Layers | MLA layers (1-based) | KDA layers | Heads | Experts (top-k) | ViT dim / layers / heads | +|---------|------------|---------|--------|----------------------|------------|-------|-----------------|--------------------------| +| Released Kimi K3 (reference) | 2.8T | 7168 | 93 | 4, 8, ..., 92, 93 | 69 | 96 | 896 (top-16) | 1024 / 27 / 12 | +| debugmodel | 100M | 256 | 13 | 4, 8, 12, 13 | 9 | 4 | 8 (top-2) | 256 / 4 / 3 | + +`debugmodel` retains the released vocabulary size of 163840. Its depth also +preserves two structural edge cases from the released model: consecutive final +MLA layers and a short trailing attention-residual block. + +## Supported Parallelisms + +| Feature | Notes | +|---------|-------| +| FSDP / HSDP | Supported with the default SPMD backend. The decoder is sharded per layer and the vision encoder is a separate FSDP unit | +| Tensor Parallelism (TP) | Not supported | +| Expert Parallelism (EP) | Not supported | +| Pipeline Parallelism (PP) | Not supported | +| Context Parallelism (CP) | Not supported | + +`torch.compile`, activation checkpointing, and parameter CPU offload are not +supported by the current Kimi K3 parallelization path. + +Run the debug model on one GPU: ```bash NGPU=1 MODULE=kimi_k3 CONFIG=kimi_k3_debugmodel ./run_train.sh ``` -Run the same eager model with two-way FSDP2: +Run it with two-way FSDP2: ```bash NGPU=2 MODULE=kimi_k3 CONFIG=kimi_k3_debugmodel ./run_train.sh \ --parallelism.data_parallel_shard_degree 2 ``` -Requirements beyond core TorchTitan are listed in `requirements.txt`: KDA needs -`flash-linear-attention` and the multimodal data path needs `torchvision`. - -## Reduced model - -`debugmodel` preserves each distinct Kimi K3 forward path while reducing -widths, head dimensions, expert count, and depth. - -| Component | Released Kimi K3 | `debugmodel` | -|---|---:|---:| -| Parameters | 2.8T | 100M | -| Decoder dimension | 7168 | 256 | -| Vocabulary size | 163840 | 163840 | -| Decoder layers | 93 | 13 | -| Full MLA layers (1-based) | 4, 8, ..., 92, 93 | 4, 8, 12, 13 | -| KDA layers | 69 | 9 | -| Dense FFN layers | 1 | 1 | -| Attention residual block size | 12 | 12 | -| MLA heads | 96 | 4 | -| MLA qk_nope / qk_rope / v head dimension | 128 / 64 / 128 | 32 / 16 / 32 | -| KDA heads / head dimension | 96 / 128 | 4 / 32 | -| Routed experts / top-k | 896 / 16 | 8 / 2 | -| Routed latent / expert hidden dimension | 3584 / 3072 | 128 / 128 | -| Shared experts | 2 | 2 | -| Vision dimension | 1024 | 256 | -| Vision layers | 27 | 4 | -| Vision QKV dimension / heads | 1536 / 12 | 384 / 3 | - -The depth sits one layer past a multiple of the full-attention period and of -the attention-residual block size, which reproduces two structural edge cases -of the released 93-layer stack: a final MLA layer immediately after a scheduled -one -- the released `full_attn_layers` ends `..., 88, 92, 93`, so the backbone -always closes on global attention -- and a short trailing residual block. -Neither is expressible as "every n-th layer", so `full_attention_layers` takes -the 1-based list verbatim and a future full-scale flavor can pass the released -24-entry list unchanged. -The released vocabulary size is retained, following other TorchTitan -multimodal debug models and making FSDP state sharding measurable while the -decoder widths and depths remain reduced. In `debugmodel` roughly 84 million -of the 100 million parameters are the token embedding and the separate output -projection over that vocabulary; the transformer itself is correspondingly -small. - -## Forward structure - -The reference path mirrors the released implementation in these areas: - -- FP32-reduction RMSNorm and SiTU activation. -- Gated MLA with low-rank query and KV projections. Kimi K3 sets - `mla_use_nope=True`, so the RoPE-sized query/key slices are not rotated. -- KDA short causal convolutions, safe decay gate, query/key L2 normalization, - sigmoid beta, recurrent delta-rule update, and gated output RMSNorm. -- Block-level attention residuals, including the final output residual. -- Stable LatentMoE with sigmoid top-k routing, correction bias, latent - down/up projections, routed experts, and shared experts. -- MoonViT3d patch embedding, learned spatial positions, 2D RoPE, non-causal - per-image attention, temporal pooling, 2x2 spatial merge, and - PatchMergerMLPV2. -- Vision features scattered into runs of media placeholder tokens. - -`KimiKDAKernel` is the kernel boundary. It dispatches to FLA's `chunk_kda` -with the gate activation, beta sigmoid, and query/key L2 norm fused in, so a -future backend can replace it while preserving the same input/output contract -and checkpoint schema. - -The vision encoder is its own FSDP unit, so its collectives only fire on ranks -that execute it. Because a data-parallel rank can legitimately receive a batch -with no images -- the shared multimodal collator emits `pixel_values=None` for -a text-only batch, and drops images to respect `max_images_per_batch` -- the -model runs the encoder on *every* batch. Batches without images use the -smallest grid the patch merger accepts and contribute its result through -`add_zero_valued_dependency`, which leaves the text embeddings numerically -unchanged while keeping the encoder in the autograd graph. Every rank -therefore issues the same all-gather and reduce-scatter regardless of what its -batch contains, and the encoder correctly receives zero gradients from -text-only ranks. - -## Checkpoint conversion - -`KimiK3StateDictAdapter` converts between TorchTitan and an unquantized -HuggingFace state dict. It covers: - -- dense, MLA, KDA, and LatentMoE decoder layers; -- attention-residual parameters; -- vision patch embedding, fused HuggingFace QKV, transformer blocks, and - projector; -- the MoE correction bias. - -The released checkpoint stores routed expert weights in MXFP4. Loading those -compressed tensors is outside this first change. Numerical comparison should -therefore instantiate the same reduced, unquantized model on both sides and -copy one state dict through the adapter. - -`test_kimi_k3_hf_parity.py` freezes the float32 outputs from a deterministic -reduced model evaluated with the released HuggingFace code at commit -`c5d1dd4c428bd1ce8b88c5044f3b6ccde9e3b721`. The test covers text logits, -router choices, projected vision features, and end-to-end image-text logits. -The source model is loaded strictly from the state dict produced by -`KimiK3StateDictAdapter`; no full checkpoint or network access is required to -run the regression. +## Numerical Checks -### Parity at released head dimensions +`scripts/checkpoint_conversion/numerical_tests_kimi_k3.py` loads the released +HuggingFace config, modeling code, processor, and tokenizer from a local model +directory. It reduces the HuggingFace model to the `debugmodel` topology and +transfers the randomly initialized TorchTitan state dict, without loading the +released weights. Each side performs its own image preprocessing before the +full vision-projector-decoder forward. Float32 is the default correctness mode, +and the script does not override the framework's TF32 settings. -`debugmodel` shrinks the head dimensions, so it does not pin the FLA kernel -configuration the full model runs. That was checked separately, out of band, on -a 1.07B configuration built from the same `_kimi_k3_config` builder with the -released head dimensions kept intact: dimension 1024, 25 layers, full attention -at `4, 8, ..., 24, 25`, MLA `qk_nope / qk_rope / v` = `128 / 64 / 128` with 8 -heads, KDA 8 heads of 128, 24 routed experts at top-4, vision dimension 512 -over 8 layers. Same released commit, one RTX 5080, float32, 128 tokens. Both -sides run FLA's `chunk_kda`; the released code path for MLA and vision -attention was forced to eager, since flash-attn is not required here. +The current float32 CUDA validation result is: -Float32 comparison requires closing **two** independent TF32 switches: +- pixel preprocessing: max difference `1.192e-7`, with no values differing + above `1e-6`; +- projected vision features: cosine similarity `1.000000`, max difference + `3.152e-3`; +- expert routing: all `3936 / 3936` choices match; +- end-to-end last-token logits: KL `1.8215e-8`, cosine similarity `1.000002`, + max difference `4.1358e-3`, top-1 match, and top-5 5/5. -- **cuBLAS**, via `torch.backends.cuda.matmul.allow_tf32 = False`. Torch's own - default is already off, but NVIDIA's NGC containers set - `TORCH_ALLOW_TF32_CUBLAS_OVERRIDE=1`, which flips - `float32_matmul_precision` to `HIGH` at process init. -- **Triton**, via `TRITON_F32_DEFAULT=ieee`. The torch flag does not reach - Triton kernels, and FLA's `chunk_kda` leans on the Triton default for the - triangular solves in `fla/ops/kda/chunk_intra.py`. - -| Quantity | max abs diff | reference max abs | -|---|---:|---:| -| Text logits, 128 tokens | 2.2e-5 | 5.8 | -| Projected vision features, 8x8 patch grid | 4.0e-5 | 4.4 | -| Multimodal logits, 128 tokens | 2.4e-4 | 5.8 | - -Text logits differ by 3.9e-6 relative. For roughly 250 sequential dependent -matmuls that is well inside the linear accumulation bound of 1.5e-5 implied by -float32's 6.0e-8 unit roundoff. Per-layer relative drift grows from 4.8e-7 -after layer 0 to 3.4e-6 by layer 23, and the per-layer amplification factor -stays in 0.96-1.44 throughout: no layer amplifies, the error only accumulates. -Routed expert IDs are identical for all 3072 token-routings (24 MoE layers x -128 tokens), and the argmax over the vocabulary agrees on every position. - -Leaving Triton at its `tf32` default costs a factor of three end to end -(text logits 2.2e-5 -> 7.4e-5) and shows up inside KDA as a 50x step: every -tensor entering `chunk_kda` matches to ~1e-5 relative, while its output matches -only to ~5e-4. Both sides call the same kernel, so this never broke parity -- -TF32's error is 99.6% correlated between the two runs and mostly cancels. What -does not cancel is the discontinuity: where the two nearly-identical inputs -land on opposite sides of a rounding boundary, the entry jumps a full TF32 -quantum, 50x larger than the input difference that caused it. Under IEEE the -step disappears and `chunk_kda`'s output matches to ~6e-6. - -With both switches closed, the eager operators and their FLA counterparts are -numerically indistinguishable: swapping TorchTitan's `Conv1d`+SiLU and eager -gated RMSNorm for FLA's `ShortConvolution` and `FusedRMSNormGated` moves the -logits difference from 2.2e-5 to 2.4e-5, and the two variants differ from each -other by 1.3e-5. None of the three is a privileged reference; what remains is -ordinary reassociation noise, including the routed-expert summation order. - -The eager operators are kept deliberately. FLA's `ShortConvolution` and -`FusedRMSNormGated` are Triton-only and raise on CPU tensors; using them here -would make the whole model forward GPU-only and would leave -`test_kimi_k3_hf_parity.py` -- the frozen, network-free, CPU regression -- with -no way to run. They are also not a checkpoint concern either way: -`ShortConvolution` subclasses `nn.Conv1d` with the same `(D, 1, K)` weight, and -`FusedRMSNormGated` carries the same `(D,)` weight, so the state-dict mapping is -unaffected by the choice. - -### bfloat16 - -Comparing the two implementations directly in bfloat16 measures top-k router -ties, not arithmetic. Each side must be compared against its own float32 run. - -| bfloat16 vs own float32 | routings changed | logits max abs | argmax kept | -|---|---:|---:|---:| -| TorchTitan | 808 / 3072 | 5.69 | 50.0% | -| Released HuggingFace | 921 / 3072 | 5.96 | 47.7% | - -Both implementations lose their own float32 result at the same rate, so the -instability is the model configuration rather than either implementation. It -comes from the router: a randomly initialized gate produces near-tied scores, -bfloat16 rounding reorders them, and a changed expert changes the residual -stream enough to change every later routing decision. Rounding only the gate -of the first MoE layer to bfloat16, with a float32 input, already flips 4 of -128 tokens; those tokens have a median top-4/top-5 score margin of 5.4e-4 -against a typical score spread of 5.5e-1. - -Routing every token to all 24 experts removes the ties and leaves only the -arithmetic. Nothing then flips on either side, and the two implementations -degrade identically: - -| bfloat16 vs own float32, all experts routed | logits max abs | logits mean abs | argmax kept | -|---|---:|---:|---:| -| TorchTitan | 1.72e-1 | 1.66e-2 | 97.7% | -| Released HuggingFace | 1.88e-1 | 1.68e-2 | 94.5% | - -Per-layer bfloat16 error grows from 0.8% of the activation scale after layer 0 -to roughly 3% by the end of an attention-residual block on both sides, with -neither consistently ahead. That is ordinary bfloat16 accumulation over 25 -layers, given a 0.4% unit roundoff. - -The practical consequence: validate numerics in float32, as -`.claude/rules` already requires, and do not read a bfloat16 loss difference -against a reference implementation as evidence of a bug until the routing -decisions have been checked. - -## Tests - -The unit tests cover: - -- the reduced layer topology; -- the explicit exact GELU against PyTorch's CPU reference; -- a small text+image model forward and backward; -- exhaustive state-dict round-trip for that small model; -- reduced text, vision, router, and multimodal numerical parity against frozen - HuggingFace eager outputs; -- single-rank FSDP2 forward and per-parameter gradient parity with a manually - cast BF16 reference; -- two-rank FSDP2 forward and backward when one rank has an image and the other - rank is text-only. - -`ReferenceKimiKDAKernel` in `tests/unit_tests/test_kimi_k3.py` is the explicit -recurrent formulation of KDA. The tests above build the model with it in place -of the FLA kernel, which is what lets them run on CPU and at head dimensions -FLA cannot compile. A separate CUDA-only test checks the FLA kernel against -that same reference, forward and backward, for both gate activations. +Run the comparison with: ```bash -pytest -q \ - tests/unit_tests/test_kimi_k3.py \ - tests/unit_tests/test_kimi_k3_hf_parity.py -pytest -q tests/unit_tests/test_kimi_k3_fsdp.py +python -m scripts.checkpoint_conversion.numerical_tests_kimi_k3 \ + --hf_model_path ~/hf_assets/moonshotai/Kimi-K3 ``` -## First-version limitations - -- FSDP2 data parallelism is supported; HSDP, TP, PP, CP, and EP are rejected. -- No packed documents, activation checkpointing, `torch.compile`, or CPU - offload. -- Image inputs are supported; video inputs are rejected. -- No generation cache. -- No MXFP4 checkpoint loading. -- No full 2.8T flavor. +## TODO -These restrictions are explicit so unsupported runtime settings fail instead -of being silently ignored. This first contribution deliberately limits -parallel execution to FSDP2: its purpose is to establish the eager numerical -reference used by Kimi K3 architecture experiments. TP, PP, CP, and EP can be -added independently after that forward contract is locked. +- Add the full 2.8T model flavor. +- Add MXFP4 compressed checkpoint loading. +- Add TP, EP, PP, and CP support. +- Add packed-document attention support. +- Add video inputs and a video dataset training pipeline. +- Add `torch.compile`, activation checkpointing, and parameter CPU offload. +- Add generation-cache support. From 2509f18a271e242548442808db5f6b647c05b9e1 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Fri, 14 Aug 2026 12:30:32 +0000 Subject: [PATCH 26/67] remove parallelism config --- torchtitan/models/kimi_k3/config_registry.py | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) diff --git a/torchtitan/models/kimi_k3/config_registry.py b/torchtitan/models/kimi_k3/config_registry.py index 7da317143c..ba22ed8b4c 100644 --- a/torchtitan/models/kimi_k3/config_registry.py +++ b/torchtitan/models/kimi_k3/config_registry.py @@ -12,7 +12,7 @@ from torchtitan.components.metrics import MetricsProcessor from torchtitan.components.optimizer import default_adamw from torchtitan.components.tokenizer import MultiModalTokenizer -from torchtitan.config import ParallelismConfig, TrainingConfig +from torchtitan.config import TrainingConfig from torchtitan.hf_datasets.multimodal.mm_datasets import MMDataLoader from torchtitan.hf_datasets.multimodal.utils.image import resize_to_patch_budget from torchtitan.models.common.config_utils import decoder_vocab_size @@ -62,13 +62,6 @@ def kimi_k3_debugmodel() -> Trainer.Config: steps=10, dtype="bfloat16", ), - parallelism=ParallelismConfig( - data_parallel_shard_degree=1, - tensor_parallel_degree=1, - pipeline_parallel_degree=1, - context_parallel_degree=1, - expert_parallel_degree=1, - ), checkpoint=CheckpointManager.Config( interval=10, last_save_model_only=False, From 89d93a924a5a8e10ef3509299c2d6be83f031518 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Fri, 14 Aug 2026 12:38:19 +0000 Subject: [PATCH 27/67] change full_attention_layers to starting with 0 --- .../numerical_tests_kimi_k3.py | 8 ++++---- torchtitan/models/kimi_k3/__init__.py | 13 ++++++------- 2 files changed, 10 insertions(+), 11 deletions(-) diff --git a/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py b/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py index 9ad9d7e3d0..5ce05cb6ca 100644 --- a/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py +++ b/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py @@ -57,12 +57,12 @@ def _reduce_hf_config(hf_config, tt_config, hf_model_path: str) -> None: """Reduce the released HuggingFace config to the TorchTitan debug model.""" text_config = hf_config.text_config - full_attention_layers = [ + hf_full_attention_layers = [ layer_idx + 1 for layer_idx, layer in enumerate(tt_config.layers) if layer.attention is not None ] - kda_layers = [ + hf_kda_layers = [ layer_idx + 1 for layer_idx, layer in enumerate(tt_config.layers) if layer.delta_attention is not None @@ -105,8 +105,8 @@ def _reduce_hf_config(hf_config, tt_config, hf_model_path: str) -> None: ), "attn_res_block_size": tt_config.layers[0].attn_res_block_size, "linear_attn_config": { - "full_attn_layers": full_attention_layers, - "kda_layers": kda_layers, + "full_attn_layers": hf_full_attention_layers, + "kda_layers": hf_kda_layers, "head_dim": kda.head_dim, "num_heads": kda.num_heads, "short_conv_kernel_size": kda.conv_kernel_size, diff --git a/torchtitan/models/kimi_k3/__init__.py b/torchtitan/models/kimi_k3/__init__.py index 82f2ac65ed..45c22d87cd 100644 --- a/torchtitan/models/kimi_k3/__init__.py +++ b/torchtitan/models/kimi_k3/__init__.py @@ -377,14 +377,13 @@ def _kimi_k3_config( ) -> KimiK3Model.Config: """Assemble a Kimi K3 config from the released topology's free parameters. - ``full_attention_layers`` holds 1-based layer indices, matching the - released ``linear_attn_config.full_attn_layers``. Every other layer is KDA. - Layer 0 is the single dense FFN layer (released + ``full_attention_layers`` holds zero-based layer indices. Every other layer + is KDA. Layer 0 is the single dense FFN layer (released ``first_k_dense_replace=1``); the rest are LatentMoE. """ layers = [] for layer_idx in range(num_layers): - is_full_attention = (layer_idx + 1) in full_attention_layers + is_full_attention = layer_idx in full_attention_layers layers.append( KimiK3TransformerBlock.Config( layer_id=layer_idx, @@ -466,8 +465,8 @@ def _debugmodel(attn_backend: str) -> KimiK3Model.Config: The depth is one past a multiple of both the full-attention period and the attention-residual block size, so the last layer is a full-attention layer directly after a scheduled one and the trailing residual block is short. - Both are properties of the released 93-layer stack, whose - ``full_attn_layers`` ends ``..., 88, 92, 93``. + Both are properties of the released 93-layer stack, whose zero-based MLA + layer indices end ``..., 87, 91, 92``. """ if attn_backend != "eager": raise ValueError("Kimi K3 v1 only provides the 'eager' backend.") @@ -477,7 +476,7 @@ def _debugmodel(attn_backend: str) -> KimiK3Model.Config: dim=dim, vocab_size=163840, num_layers=13, - full_attention_layers={4, 8, 12, 13}, + full_attention_layers={3, 7, 11, 12}, attn_res_block_size=12, num_heads=4, q_lora_rank=128, From ae006e434fa611749ca84fce4d391602afbc2070 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Sat, 15 Aug 2026 12:40:30 +0000 Subject: [PATCH 28/67] support flex attention --- .../numerical_tests_kimi_k3.py | 13 ++- tests/unit_tests/test_kimi_k3.py | 38 ++++++- torchtitan/models/common/decoder.py | 32 +++--- torchtitan/models/kimi_k3/__init__.py | 47 ++++---- torchtitan/models/kimi_k3/config_registry.py | 1 - torchtitan/models/kimi_k3/model.py | 107 ++++++------------ torchtitan/models/kimi_k3/parallelize.py | 16 ++- .../models/kimi_k3/state_dict_adapter.py | 4 +- 8 files changed, 138 insertions(+), 120 deletions(-) diff --git a/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py b/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py index 5ce05cb6ca..b37bceb1b3 100644 --- a/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py +++ b/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py @@ -94,10 +94,11 @@ def _reduce_hf_config(hf_config, tt_config, hf_model_path: str) -> None: "num_experts": moe.num_experts, "num_experts_per_token": moe.router.top_k, "num_shared_experts": ( - moe.shared_experts.w1.out_features // moe.routed_experts.hidden_dim + moe.shared_experts.w1.out_features + // moe.routed_experts.inner_experts.hidden_dim ), "moe_renormalize": moe.router.route_norm, - "moe_intermediate_size": moe.routed_experts.hidden_dim, + "moe_intermediate_size": moe.routed_experts.inner_experts.hidden_dim, "routed_expert_hidden_size": moe.routed_down.out_features, "routed_scaling_factor": moe.router.route_scale, "first_k_dense_replace": next( @@ -348,6 +349,12 @@ def run_tt( _MEDIA_TOKEN_ID, int(num_vision_tokens.item()), ).to(device) + positions = torch.arange( + tokens.shape[1], + dtype=torch.int32, + device=device, + ).unsqueeze(0) + attention_masks = model.get_attention_masks(positions) print( f"tokens={tuple(tokens.shape)} pixel_values={tuple(pixel_values.shape)} " @@ -378,6 +385,8 @@ def run_tt( pixel_values=pixel_values, grid_thw=grid_thw, special_tokens={"image_id": _MEDIA_TOKEN_ID}, + positions=positions, + attention_masks=attention_masks, ) _print_routing_comparison(ref["expert_indices"], expert_indices) return logits[:, -1, :].float().cpu().squeeze() diff --git a/tests/unit_tests/test_kimi_k3.py b/tests/unit_tests/test_kimi_k3.py index d6a2c78a8b..84f434c3ed 100644 --- a/tests/unit_tests/test_kimi_k3.py +++ b/tests/unit_tests/test_kimi_k3.py @@ -10,6 +10,7 @@ import torch import torch.nn as nn import torch.nn.functional as F +from torch.nn.attention.flex_attention import BlockMask from torchtitan.models.common import Embedding from torchtitan.protocols.module import Module @@ -44,7 +45,7 @@ class ReferenceKimiKDAKernel(Module): """Pure-PyTorch stand-in for KimiKDAKernel backed by an explicit recurrence. Mirrors ``KimiKDAKernel.forward``'s interface so tests can build a model - with it in place of the FLA kernel and exercise the surrounding eager model + with it in place of the FLA kernel and exercise the surrounding model on CPU. The loop is O(seqlen) and far too slow for training; it exists to pin the kernel's math. """ @@ -129,6 +130,7 @@ def block( qk_nope_head_dim=4, qk_rope_head_dim=4, v_head_dim=4, + attn_backend="flex", ) if use_mla else None @@ -265,6 +267,32 @@ def _kda_recurrent_reference( class TestKimiK3(unittest.TestCase): + def test_flex_attention_mask(self): + config = _small_model_config() + model = config.build() + positions = torch.arange(4, dtype=torch.int32).unsqueeze(0) + attention_masks = model.get_attention_masks(positions) + self.assertIsInstance(attention_masks, BlockMask) + + def test_update_from_config_propagates_moe_force_load_balance(self): + from torchtitan.config import DebugConfig + from torchtitan.trainer import Trainer + + model_config = _small_model_config() + runtime_config = Trainer.Config( + debug=DebugConfig(moe_force_load_balance=True), + activation_checkpoint=None, + ) + model_config.update_from_config(config=runtime_config) + + router_configs = [ + layer.moe.router for layer in model_config.layers if layer.moe is not None + ] + self.assertGreater(len(router_configs), 0) + self.assertTrue( + all(router._debug_force_load_balance for router in router_configs) + ) + @unittest.skipIf(not torch.cuda.is_available(), "FLA KDA kernel requires CUDA.") def test_fla_kda_kernel_matches_recurrent_reference(self): torch.manual_seed(1) @@ -341,6 +369,14 @@ def test_state_dict_round_trips_through_hf_adapter(self): state_dict = model.state_dict() adapter = KimiK3StateDictAdapter(config, hf_assets_path=None) hf_state_dict = adapter.to_hf(state_dict) + self.assertIn( + "layers.1.moe.routed_experts.inner_experts.w1_EFD", + state_dict, + ) + self.assertIn( + "language_model.model.layers.1.block_sparse_moe.experts.0.w1.weight", + hf_state_dict, + ) roundtrip_state_dict = adapter.from_hf(hf_state_dict) self.assertEqual(state_dict.keys(), roundtrip_state_dict.keys()) for key, value in state_dict.items(): diff --git a/torchtitan/models/common/decoder.py b/torchtitan/models/common/decoder.py index d6f52233f2..a208df1244 100644 --- a/torchtitan/models/common/decoder.py +++ b/torchtitan/models/common/decoder.py @@ -146,10 +146,10 @@ def update_from_config( When *config* is a ``Trainer.Config``, validates ``training.max_context_length`` against each attention layer's intrinsic - RoPE max sequence length, resizes RoPE caches, and propagates - debug flags. Non-trainer callers may pass any config-like - object with a ``ParallelismConfig`` in its ``parallelism`` - field; in that case the training/debug setup is skipped. + RoPE max context length, resizes RoPE caches when present, and + propagates debug flags. Non-trainer callers may pass any config-like + object with a ``ParallelismConfig`` in its ``parallelism`` field; in + that case the training/debug setup is skipped. """ from torchtitan.config import ParallelismConfig from torchtitan.distributed.context_parallel import validate_cp_backend @@ -216,20 +216,24 @@ def update_from_config( if isinstance(config, Trainer.Config): debug = config.debug seq_len = config.training.max_context_length - max_context_length = self.max_context_length - if seq_len > max_context_length: - raise ValueError( - f"Training sequence length {seq_len} exceeds " - f"attention RoPE maximum supported sequence " - f"length {max_context_length}." - ) + rope_cfg = getattr(attention, "rope", None) + if rope_cfg is not None: + max_context_length = self.max_context_length + if seq_len > max_context_length: + raise ValueError( + f"Training sequence length {seq_len} exceeds " + f"attention RoPE maximum supported sequence " + f"length {max_context_length}." + ) for layer_cfg in self.layers: attention_cfg = getattr(layer_cfg, "attention", None) if attention_cfg is not None: - attention_cfg.rope = dataclasses.replace( - attention_cfg.rope, max_context_length=seq_len - ) + rope_cfg = getattr(attention_cfg, "rope", None) + if rope_cfg is not None: + attention_cfg.rope = dataclasses.replace( + rope_cfg, max_context_length=seq_len + ) if hasattr(layer_cfg, "moe") and layer_cfg.moe is not None: layer_cfg.moe.router._debug_force_load_balance = ( debug.moe_force_load_balance diff --git a/torchtitan/models/kimi_k3/__init__.py b/torchtitan/models/kimi_k3/__init__.py index 45c22d87cd..1bab0b5232 100644 --- a/torchtitan/models/kimi_k3/__init__.py +++ b/torchtitan/models/kimi_k3/__init__.py @@ -14,7 +14,8 @@ from torchtitan.components.optimizer import register_moe_load_balancing_hook from torchtitan.models.common import Conv1d, Embedding, Linear -from torchtitan.models.common.moe import TokenChoiceTopKRouter +from torchtitan.models.common.config_utils import get_attention_config +from torchtitan.models.common.moe import RoutedExperts, TokenChoiceTopKRouter from torchtitan.models.common.nn_modules import GELU, RMSNorm from torchtitan.models.common.token_dispatcher import LocalTokenDispatcher from torchtitan.models.common.vision_encoder import VisionMLP @@ -147,7 +148,10 @@ def _mla_config( qk_nope_head_dim: int, qk_rope_head_dim: int, v_head_dim: int, + attn_backend: str, ) -> KimiMLAAttention.Config: + inner_attention = get_attention_config(attn_backend) + q_head_dim = qk_nope_head_dim + qk_rope_head_dim return KimiMLAAttention.Config( dim=dim, @@ -167,6 +171,7 @@ def _mla_config( ), gate=_linear(dim, num_heads * v_head_dim), wo=_linear(num_heads * v_head_dim, dim), + inner_attention=inner_attention, ) @@ -241,21 +246,23 @@ def _latent_moe_config( route_scale=1.0, ), routed_down=_linear(dim, latent_dim), - routed_experts=KimiGroupedExperts.Config( - dim=latent_dim, - hidden_dim=expert_hidden_dim, - num_experts=num_experts, - beta=4.0, - linear_beta=25.0, - param_init={ - "w1_EFD": partial(nn.init.trunc_normal_, std=0.02), - "w2_EDF": partial(nn.init.trunc_normal_, std=0.02), - "w3_EFD": partial(nn.init.trunc_normal_, std=0.02), - }, - ), - token_dispatcher=LocalTokenDispatcher.Config( - num_experts=num_experts, - top_k=top_k, + routed_experts=RoutedExperts.Config( + inner_experts=KimiGroupedExperts.Config( + dim=latent_dim, + hidden_dim=expert_hidden_dim, + num_experts=num_experts, + beta=4.0, + linear_beta=25.0, + param_init={ + "w1_EFD": partial(nn.init.trunc_normal_, std=0.02), + "w2_EDF": partial(nn.init.trunc_normal_, std=0.02), + "w3_EFD": partial(nn.init.trunc_normal_, std=0.02), + }, + ), + token_dispatcher=LocalTokenDispatcher.Config( + num_experts=num_experts, + top_k=top_k, + ), ), routed_norm=_norm(latent_dim), routed_up=_linear(latent_dim, dim), @@ -374,6 +381,7 @@ def _kimi_k3_config( top_k: int, num_shared_experts: int, vision_encoder: KimiK3VisionEncoder.Config, + attn_backend: str, ) -> KimiK3Model.Config: """Assemble a Kimi K3 config from the released topology's free parameters. @@ -397,6 +405,7 @@ def _kimi_k3_config( qk_nope_head_dim=qk_nope_head_dim, qk_rope_head_dim=qk_rope_head_dim, v_head_dim=v_head_dim, + attn_backend=attn_backend, ) if is_full_attention else None @@ -468,9 +477,6 @@ def _debugmodel(attn_backend: str) -> KimiK3Model.Config: Both are properties of the released 93-layer stack, whose zero-based MLA layer indices end ``..., 87, 91, 92``. """ - if attn_backend != "eager": - raise ValueError("Kimi K3 v1 only provides the 'eager' backend.") - dim = 256 return _kimi_k3_config( dim=dim, @@ -500,6 +506,7 @@ def _debugmodel(attn_backend: str) -> KimiK3Model.Config: num_layers=4, num_heads=3, ), + attn_backend=attn_backend, ) @@ -510,7 +517,7 @@ def _debugmodel(attn_backend: str) -> KimiK3Model.Config: def model_registry( flavor: str, - attn_backend: str = "eager", + attn_backend: str = "flex", converters: list[ModelConfigConverter.Config] | None = None, ) -> ModelSpec: """Build a Kimi K3 model specification.""" diff --git a/torchtitan/models/kimi_k3/config_registry.py b/torchtitan/models/kimi_k3/config_registry.py index ba22ed8b4c..a5654733c6 100644 --- a/torchtitan/models/kimi_k3/config_registry.py +++ b/torchtitan/models/kimi_k3/config_registry.py @@ -22,7 +22,6 @@ def kimi_k3_debugmodel() -> Trainer.Config: - """Return the topology-complete Kimi K3 eager/FSDP2 debug config.""" model_spec = model_registry("debugmodel") return Trainer.Config( loss=ChunkedLossWrapper.Config( diff --git a/torchtitan/models/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py index 75035f38af..2d249dc3c1 100644 --- a/torchtitan/models/kimi_k3/model.py +++ b/torchtitan/models/kimi_k3/model.py @@ -6,11 +6,9 @@ """Kimi K3 language model components. -Every operator outside the KDA recurrence is plain eager PyTorch, so the model -mirrors the released HuggingFace math and stays directly inspectable. KDA runs -on FLA's chunked Triton kernel, which is what makes the model trainable at -speed; the pure-PyTorch recurrence it is checked against lives in -``tests/unit_tests/test_kimi_k3.py`` and is far too slow for training. +MLA delegates to TorchTitan's configured inner-attention backend. KDA runs on +FLA's chunked Triton kernel; the pure-PyTorch recurrence used to check it lives +in ``tests/unit_tests/test_kimi_k3.py`` and is far too slow for training. """ from dataclasses import dataclass, field @@ -26,17 +24,20 @@ from torchtitan.models.common.attention import ( AttentionMasksType, BaseAttention, - ScaledDotProductAttention, + FlexAttention, ) from torchtitan.models.common.decoder import Decoder from torchtitan.models.common.feed_forward import FeedForward -from torchtitan.models.common.moe import GroupedExperts, TokenChoiceTopKRouter +from torchtitan.models.common.moe import ( + GroupedExperts, + RoutedExperts, + TokenChoiceTopKRouter, +) from torchtitan.models.common.multimodal import ( get_vision_positions, scatter_vision_embeds, ) from torchtitan.models.common.nn_modules import RMSNorm -from torchtitan.models.common.token_dispatcher import LocalTokenDispatcher from torchtitan.models.kimi_k3.vision_encoder import KimiK3VisionEncoder from torchtitan.models.utils import get_moe_model_nparams_and_flops from torchtitan.protocols.module import Module @@ -134,10 +135,7 @@ class KimiMLAAttention(BaseAttention): Unlike DeepSeek-V3 MLA, the released K3 configuration sets ``mla_use_nope=True``: the RoPE-sized query/key slices remain part of the projected head, but no rotary transform is applied, so this has no rope - config at all. Attention itself runs through ``ScaledDotProductAttention`` - rather than a hand-written softmax: torchtitan's training path has no - need to match HF eager's bit-for-bit fp32 softmax, so it uses the same - SDPA inner attention the rest of the codebase relies on. + config at all. Attention delegates to the configured inner backend. """ @dataclass(kw_only=True, slots=True) @@ -155,9 +153,7 @@ class Config(BaseAttention.Config): wkv_b: Linear.Config gate: Linear.Config wo: Linear.Config - inner_attention: Module.Config = field( - default_factory=ScaledDotProductAttention.Config - ) + inner_attention: Module.Config = field(default_factory=FlexAttention.Config) def __init__(self, config: Config): super().__init__() @@ -186,10 +182,6 @@ def forward( positions: torch.Tensor | None = None, ) -> torch.Tensor: del positions - if attention_masks is not None: - raise NotImplementedError( - "Kimi K3 reference MLA does not support packed-document masks." - ) B, L, _ = x_BLD.shape q_BLNH = self.wq_b(self.q_norm(self.wq_a(x_BLD))).view( @@ -218,7 +210,13 @@ def forward( ) k_BLNH = torch.cat((k_nope, k_rope), dim=-1) - out_BLNV = self.inner_attention(q_BLNH, k_BLNH, v_BLNH, scale=self.scale) + out_BLNV = self.inner_attention( + q_BLNH, + k_BLNH, + v_BLNH, + attention_masks=attention_masks, + scale=self.scale, + ) out_BLD = out_BLNV.reshape(B, L, self.n_heads * self.v_head_dim) out_BLD = out_BLD * torch.sigmoid(self.gate(x_BLD)) return self.wo(out_BLD) @@ -386,9 +384,7 @@ class KimiGroupedExperts(GroupedExperts): Inherits its stacked-weight shape (``w1_EFD``/``w2_EDF``/``w3_EFD``) and parameter allocation; only ``forward`` differs, since the activation is baked into the ``torch._grouped_mm`` call sequence rather than being a - swappable argument. ``inner_experts`` returns ``self`` so the shared FSDP - wrapper's ``moe.routed_experts.inner_experts`` discovery path works - without a separate dispatch-composing wrapper class. + swappable argument. """ @dataclass(kw_only=True, slots=True) @@ -401,10 +397,6 @@ def __init__(self, config: Config): self.beta = config.beta self.linear_beta = config.linear_beta - @property - def inner_experts(self) -> "KimiGroupedExperts": - return self - def forward( self, x_RD: torch.Tensor, @@ -451,8 +443,7 @@ class Config(Module.Config): num_experts: int router: TokenChoiceTopKRouter.Config routed_down: Linear.Config - routed_experts: KimiGroupedExperts.Config - token_dispatcher: LocalTokenDispatcher.Config + routed_experts: RoutedExperts.Config routed_norm: RMSNorm.Config routed_up: Linear.Config shared_experts: KimiFeedForward.Config @@ -460,13 +451,14 @@ class Config(Module.Config): def __init__(self, config: Config): super().__init__() - if config.routed_experts.num_experts != config.num_experts: - raise ValueError("routed_experts.num_experts must equal num_experts.") + if config.routed_experts.inner_experts.num_experts != config.num_experts: + raise ValueError( + "routed_experts.inner_experts.num_experts must equal num_experts." + ) self.num_experts = config.num_experts self.router = config.router.build() self.routed_down = config.routed_down.build() self.routed_experts = config.routed_experts.build() - self.token_dispatcher = config.token_dispatcher.build() self.routed_norm = config.routed_norm.build() self.routed_up = config.routed_up.build() self.shared_experts = config.shared_experts.build() @@ -488,15 +480,10 @@ def __init__(self, config: Config): ) def forward(self, x_BLD: torch.Tensor) -> torch.Tensor: - weights_BLK, expert_ids_BLK, _ = self.router(x_BLD, self.expert_bias_E) - B, L, _ = x_BLD.shape - routing_map_BLE = torch.zeros( - B, - L, - self.num_experts, - dtype=torch.bool, - device=x_BLD.device, - ).scatter(-1, expert_ids_BLK, True) + weights_BLK, expert_ids_BLK, scores_BLE = self.router(x_BLD, self.expert_bias_E) + routing_map_BLE = torch.zeros_like(scores_BLE, dtype=torch.bool).scatter_( + -1, expert_ids_BLK, True + ) num_tokens_per_expert_E = routing_map_BLE.sum(dim=(0, 1)) with torch.no_grad(): # In place so the load-balancing hook registered on the optimizer @@ -504,31 +491,13 @@ def forward(self, x_BLD: torch.Tensor) -> torch.Tensor: self.tokens_per_expert_E.add_(num_tokens_per_expert_E.float()) latent_BLD = self.routed_down(x_BLD) - latent_dim = latent_BLD.shape[-1] - K = weights_BLK.shape[-1] - T = B * L - latent_TD = latent_BLD.reshape(T, latent_dim) - weights_TK = weights_BLK.reshape(T, K) - expert_ids_TK = expert_ids_BLK.reshape(T, K) - - # Every expert's weights are packed into one grouped-mm call, so an - # expert with zero assigned tokens still sits in the autograd graph - # and receives an all-zero gradient rather than None. - ( - routed_input_RD, + routed_BLD = self.routed_experts( + latent_BLD, + weights_BLK, + expert_ids_BLK, num_tokens_per_expert_E, - metadata, - ) = self.token_dispatcher.dispatch( - latent_TD, weights_TK, expert_ids_TK, num_tokens_per_expert_E ) - routed_output_RD = self.routed_experts(routed_input_RD, num_tokens_per_expert_E) - routed_TD = self.token_dispatcher.combine( - routed_output_RD, - metadata, - latent_TD, - ) - - routed_BLD = self.routed_up(self.routed_norm(routed_TD.view(B, L, latent_dim))) + routed_BLD = self.routed_up(self.routed_norm(routed_BLD)) return routed_BLD + self.shared_experts(x_BLD) def _init_self_buffers(self, *, buffer_device: torch.device | None = None) -> None: @@ -653,7 +622,7 @@ def forward( h_BLD = self.attention(h_BLD, attention_masks, positions) else: assert self.delta_attention is not None - h_BLD = self.delta_attention(h_BLD, attention_masks, positions) + h_BLD = self.delta_attention(h_BLD, None, positions) prefix_sum_BLD = h_BLD if prefix_sum_BLD is None else prefix_sum_BLD + h_BLD assert prefix_sum_BLD is not None @@ -684,7 +653,6 @@ class Config(Decoder.Config): spatial_merge_size: int = 2 def update_from_config(self, *, config, **kwargs) -> None: - del kwargs parallelism = config.parallelism unsupported = { "tensor parallel": parallelism.tensor_parallel_degree, @@ -695,7 +663,7 @@ def update_from_config(self, *, config, **kwargs) -> None: enabled = [name for name, degree in unsupported.items() if degree > 1] if enabled: raise NotImplementedError( - "Kimi K3 eager reference supports FSDP2 data parallelism only; " + "Kimi K3 supports FSDP2 data parallelism only; " f"disable {', '.join(enabled)}." ) dataloader = getattr(config, "dataloader", None) @@ -703,6 +671,7 @@ def update_from_config(self, *, config, **kwargs) -> None: raise NotImplementedError( "Kimi K3 v1 does not support packed documents." ) + Decoder.Config.update_from_config(self, config=config, **kwargs) def get_nparams_and_flops( self, model: nn.Module, seq_len: int @@ -747,10 +716,6 @@ def __init__(self, config: Config): "number of text positions than the encoder produces." ) - def get_attention_masks(self, positions: torch.Tensor) -> AttentionMasksType | None: - del positions - return None - def _prepare_multimodal_embeds( self, tokens: torch.Tensor, diff --git a/torchtitan/models/kimi_k3/parallelize.py b/torchtitan/models/kimi_k3/parallelize.py index e3480fdfa1..99e1393c55 100644 --- a/torchtitan/models/kimi_k3/parallelize.py +++ b/torchtitan/models/kimi_k3/parallelize.py @@ -4,7 +4,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -"""FSDP2 parallelization for the eager Kimi K3 reference model.""" +"""FSDP2 parallelization for Kimi K3.""" import torch.nn as nn @@ -33,7 +33,7 @@ def parallelize_kimi_k3( ac_config: ActivationCheckpointingConfig, dump_folder: str, ) -> nn.Module: - """Apply FSDP2 while keeping the model's eager reference forward path.""" + """Apply FSDP2 to the Kimi K3 decoder and vision encoder.""" del dump_folder unsupported_parallelisms = [ @@ -48,24 +48,22 @@ def parallelize_kimi_k3( ] if unsupported_parallelisms: raise NotImplementedError( - "Kimi K3 eager reference currently supports FSDP2 data parallelism " + "Kimi K3 currently supports FSDP2 data parallelism " f"only; disable {', '.join(unsupported_parallelisms)}." ) if parallelism.spmd_backend != "default": raise NotImplementedError( - "Kimi K3 eager FSDP2 currently supports the default SPMD backend only." + "Kimi K3 FSDP2 currently supports the default SPMD backend only." ) if compile_config.enable: - raise NotImplementedError( - "Kimi K3 eager reference does not support torch.compile." - ) + raise NotImplementedError("Kimi K3 does not support torch.compile.") if ac_config is not None: raise NotImplementedError( - "Kimi K3 eager FSDP2 does not support activation checkpointing yet." + "Kimi K3 FSDP2 does not support activation checkpointing yet." ) if training.enable_cpu_offload: raise NotImplementedError( - "Kimi K3 eager FSDP2 does not support parameter CPU offload yet." + "Kimi K3 FSDP2 does not support parameter CPU offload yet." ) dp_mesh_names = ( diff --git a/torchtitan/models/kimi_k3/state_dict_adapter.py b/torchtitan/models/kimi_k3/state_dict_adapter.py index 5073a6b260..ec22740640 100644 --- a/torchtitan/models/kimi_k3/state_dict_adapter.py +++ b/torchtitan/models/kimi_k3/state_dict_adapter.py @@ -184,7 +184,7 @@ def from_hf(self, hf_state_dict: dict[str, Any]) -> dict[str, Any]: moe_config = self.kimi_config.layers[int(layer_idx)].moe assert moe_config is not None grouped_key = ( - f"layers.{layer_idx}.moe.routed_experts." + f"layers.{layer_idx}.moe.routed_experts.inner_experts." f"{_EXPERT_PROJECTION_TO_GROUPED_PARAM[projection]}" ) experts = self._expert_weights_by_layer_projection.setdefault( @@ -313,7 +313,7 @@ def to_hf(self, state_dict: dict[str, Any]) -> dict[str, Any]: if text_match is not None: layer_idx, suffix = text_match.groups() expert_match = re.fullmatch( - r"moe\.routed_experts\." r"(w1_EFD|w2_EDF|w3_EFD)", + r"moe\.routed_experts\.inner_experts\." r"(w1_EFD|w2_EDF|w3_EFD)", suffix, ) if expert_match is not None: From fb7865ab14c697c22e44bfd7fde7151efc1ad568 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Sat, 15 Aug 2026 17:10:23 +0000 Subject: [PATCH 29/67] refactor KimiK3StateDictAdapter --- .../models/kimi_k3/state_dict_adapter.py | 602 +++++++++--------- 1 file changed, 290 insertions(+), 312 deletions(-) diff --git a/torchtitan/models/kimi_k3/state_dict_adapter.py b/torchtitan/models/kimi_k3/state_dict_adapter.py index ec22740640..334e2d07d5 100644 --- a/torchtitan/models/kimi_k3/state_dict_adapter.py +++ b/torchtitan/models/kimi_k3/state_dict_adapter.py @@ -4,121 +4,20 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -"""Unquantized HuggingFace checkpoint adapter for Kimi K3. - -The released Kimi K3 checkpoint uses MXFP4 expert weights. That format is -intentionally outside the first implementation. This adapter targets an -unquantized HuggingFace state dict, which is sufficient for constructing the -same reduced model on both sides of the numerical parity test. -""" +"""HuggingFace checkpoint adapter for unquantized Kimi K3 weights.""" import re from typing import Any import torch +from torch.distributed.tensor import DTensor -from torchtitan.protocols.state_dict_adapter import StateDictAdapter +from torchtitan.models.utils import MoEStateDictAdapter from .model import KimiK3Model -_TEXT_GLOBAL_FROM_HF = { - "language_model.model.embed_tokens.weight": "tok_embeddings.weight", - "language_model.model.output_attn_res_norm.weight": ("output_res_norm.weight"), - "language_model.model.output_attn_res_proj.weight": ("output_res_proj.weight"), - "language_model.model.norm.weight": "norm.weight", - "language_model.lm_head.weight": "lm_head.weight", -} - -_TEXT_LAYER_FROM_HF = { - # Layer norms and attention residuals. - "input_layernorm.weight": "attention_norm.weight", - "post_attention_layernorm.weight": "ffn_norm.weight", - "self_attention_res_norm.weight": "attention_res_norm.weight", - "self_attention_res_proj.weight": "attention_res_proj.weight", - "mlp_res_norm.weight": "ffn_res_norm.weight", - "mlp_res_proj.weight": "ffn_res_proj.weight", - # Dense MLP. - "mlp.gate_proj.weight": "feed_forward.w1.weight", - "mlp.up_proj.weight": "feed_forward.w3.weight", - "mlp.down_proj.weight": "feed_forward.w2.weight", -} - -_MLA_FROM_HF = { - "self_attn.q_a_proj.weight": "attention.wq_a.weight", - "self_attn.q_a_layernorm.weight": "attention.q_norm.weight", - "self_attn.q_b_proj.weight": "attention.wq_b.weight", - "self_attn.kv_a_proj_with_mqa.weight": "attention.wkv_a.weight", - "self_attn.kv_a_layernorm.weight": "attention.kv_norm.weight", - "self_attn.kv_b_proj.weight": "attention.wkv_b.weight", - "self_attn.g_proj.weight": "attention.gate.weight", - "self_attn.o_proj.weight": "attention.wo.weight", -} - -_KDA_FROM_HF = { - "self_attn.q_proj.weight": "delta_attention.q_proj.weight", - "self_attn.k_proj.weight": "delta_attention.k_proj.weight", - "self_attn.v_proj.weight": "delta_attention.v_proj.weight", - "self_attn.q_conv1d.weight": "delta_attention.q_conv.weight", - "self_attn.k_conv1d.weight": "delta_attention.k_conv.weight", - "self_attn.v_conv1d.weight": "delta_attention.v_conv.weight", - "self_attn.f_a_proj.weight": "delta_attention.forget_a.weight", - "self_attn.f_b_proj.weight": "delta_attention.forget_b.weight", - "self_attn.b_proj.weight": "delta_attention.beta.weight", - "self_attn.g_proj.weight": "delta_attention.output_gate.weight", - "self_attn.o_norm.weight": "delta_attention.output_norm.weight", - "self_attn.o_proj.weight": "delta_attention.output_proj.weight", - "self_attn.A_log": "delta_attention.A_log", - "self_attn.dt_bias": "delta_attention.dt_bias", -} - -_MOE_FROM_HF = { - "block_sparse_moe.gate.weight": "moe.router.gate.weight", - "block_sparse_moe.gate.e_score_correction_bias": "moe.expert_bias_E", - "block_sparse_moe.routed_expert_down_proj.weight": ("moe.routed_down.weight"), - "block_sparse_moe.routed_expert_up_proj.weight": ("moe.routed_up.weight"), - "block_sparse_moe.routed_expert_norm.weight": ("moe.routed_norm.weight"), - "block_sparse_moe.shared_experts.gate_proj.weight": ( - "moe.shared_experts.w1.weight" - ), - "block_sparse_moe.shared_experts.up_proj.weight": ("moe.shared_experts.w3.weight"), - "block_sparse_moe.shared_experts.down_proj.weight": ( - "moe.shared_experts.w2.weight" - ), -} - -# KimiGroupedExperts stacks all experts' weights into one (E, F, D) param per -# projection; HF stores them as separate per-expert 2D tensors. -_EXPERT_PROJECTION_TO_GROUPED_PARAM = { - "w1": "w1_EFD", - "w2": "w2_EDF", - "w3": "w3_EFD", -} -_GROUPED_PARAM_TO_EXPERT_PROJECTION = { - v: k for k, v in _EXPERT_PROJECTION_TO_GROUPED_PARAM.items() -} - -_VISION_GLOBAL_FROM_HF = { - "vision_tower.patch_embed.proj.weight": ("vision_encoder.patch_embed.weight"), - "vision_tower.patch_embed.pos_emb.weight": "vision_encoder.pos_embed", - "vision_tower.encoder.final_layernorm.weight": ("vision_encoder.final_norm.weight"), - "mm_projector.proj.0.weight": ("vision_encoder.projector.linear_1.weight"), - "mm_projector.proj.2.weight": ("vision_encoder.projector.linear_2.weight"), - "mm_projector.post_norm.weight": ("vision_encoder.projector.post_norm.weight"), -} - -_VISION_LAYER_FROM_HF = { - "norm0.weight": "norm1.weight", - "norm1.weight": "norm2.weight", - "wo.weight": "attn.proj.weight", - "mlp.fc0.weight": "mlp.linear_fc1.weight", - "mlp.fc1.weight": "mlp.linear_fc2.weight", -} - - -class KimiK3StateDictAdapter(StateDictAdapter): - """Convert between unquantized Kimi K3 HF and TorchTitan state dicts.""" - +class KimiK3StateDictAdapter(MoEStateDictAdapter): def __init__( self, model_config: KimiK3Model.Config, @@ -126,252 +25,331 @@ def __init__( ): super().__init__(model_config, hf_assets_path) self.kimi_config = model_config - # {(layer_idx, projection): {expert_idx: 2D tensor}}, filled in from_hf - # while HF's per-expert keys arrive one at a time; stacked into - # KimiGroupedExperts' (E, F, D) parameter once all experts are seen. - self._expert_weights_by_layer_projection: dict[ - tuple[str, str], dict[int, torch.Tensor] - ] = {} - @staticmethod - def _raise_if_quantized_key(key: str) -> None: - quantized_markers = ( - "weight_scale", - "weight_packed", - "compressed", - "scale_shape", + self.from_hf_map = { + # Language model. + "language_model.model.embed_tokens.weight": "tok_embeddings.weight", + "language_model.model.layers.{}.input_layernorm.weight": "layers.{}.attention_norm.weight", + "language_model.model.layers.{}.post_attention_layernorm.weight": "layers.{}.ffn_norm.weight", + "language_model.model.layers.{}.self_attention_res_norm.weight": "layers.{}.attention_res_norm.weight", + "language_model.model.layers.{}.self_attention_res_proj.weight": "layers.{}.attention_res_proj.weight", + "language_model.model.layers.{}.mlp_res_norm.weight": "layers.{}.ffn_res_norm.weight", + "language_model.model.layers.{}.mlp_res_proj.weight": "layers.{}.ffn_res_proj.weight", + "language_model.model.layers.{}.mlp.gate_proj.weight": "layers.{}.feed_forward.w1.weight", + "language_model.model.layers.{}.mlp.up_proj.weight": "layers.{}.feed_forward.w3.weight", + "language_model.model.layers.{}.mlp.down_proj.weight": "layers.{}.feed_forward.w2.weight", + # MoE. + "language_model.model.layers.{}.block_sparse_moe.experts.{}.w1.weight": ( + "layers.{}.moe.routed_experts.inner_experts.w1_EFD" + ), + "language_model.model.layers.{}.block_sparse_moe.experts.{}.w2.weight": ( + "layers.{}.moe.routed_experts.inner_experts.w2_EDF" + ), + "language_model.model.layers.{}.block_sparse_moe.experts.{}.w3.weight": ( + "layers.{}.moe.routed_experts.inner_experts.w3_EFD" + ), + "language_model.model.layers.{}.block_sparse_moe.gate.weight": "layers.{}.moe.router.gate.weight", + "language_model.model.layers.{}.block_sparse_moe.gate.e_score_correction_bias": "layers.{}.moe.expert_bias_E", + "language_model.model.layers.{}.block_sparse_moe.routed_expert_down_proj.weight": "layers.{}.moe.routed_down.weight", + "language_model.model.layers.{}.block_sparse_moe.routed_expert_up_proj.weight": "layers.{}.moe.routed_up.weight", + "language_model.model.layers.{}.block_sparse_moe.routed_expert_norm.weight": "layers.{}.moe.routed_norm.weight", + "language_model.model.layers.{}.block_sparse_moe.shared_experts.gate_proj.weight": ( + "layers.{}.moe.shared_experts.w1.weight" + ), + "language_model.model.layers.{}.block_sparse_moe.shared_experts.up_proj.weight": ( + "layers.{}.moe.shared_experts.w3.weight" + ), + "language_model.model.layers.{}.block_sparse_moe.shared_experts.down_proj.weight": ( + "layers.{}.moe.shared_experts.w2.weight" + ), + "language_model.model.output_attn_res_norm.weight": "output_res_norm.weight", + "language_model.model.output_attn_res_proj.weight": "output_res_proj.weight", + "language_model.model.norm.weight": "norm.weight", + "language_model.lm_head.weight": "lm_head.weight", + # Vision encoder. + "vision_tower.patch_embed.proj.weight": "vision_encoder.patch_embed.weight", + "vision_tower.patch_embed.pos_emb.weight": "vision_encoder.pos_embed", + "vision_tower.encoder.blocks.{}.norm0.weight": "vision_encoder.layers.{}.norm1.weight", + "vision_tower.encoder.blocks.{}.norm1.weight": "vision_encoder.layers.{}.norm2.weight", + "vision_tower.encoder.blocks.{}.wo.weight": "vision_encoder.layers.{}.attn.proj.weight", + "vision_tower.encoder.blocks.{}.mlp.fc0.weight": "vision_encoder.layers.{}.mlp.linear_fc1.weight", + "vision_tower.encoder.blocks.{}.mlp.fc1.weight": "vision_encoder.layers.{}.mlp.linear_fc2.weight", + "vision_tower.encoder.final_layernorm.weight": "vision_encoder.final_norm.weight", + "mm_projector.proj.0.weight": "vision_encoder.projector.linear_1.weight", + "mm_projector.proj.2.weight": "vision_encoder.projector.linear_2.weight", + "mm_projector.post_norm.weight": "vision_encoder.projector.post_norm.weight", + } + self.mla_from_hf_map = { + "language_model.model.layers.{}.self_attn.q_a_proj.weight": "layers.{}.attention.wq_a.weight", + "language_model.model.layers.{}.self_attn.q_a_layernorm.weight": "layers.{}.attention.q_norm.weight", + "language_model.model.layers.{}.self_attn.q_b_proj.weight": "layers.{}.attention.wq_b.weight", + "language_model.model.layers.{}.self_attn.kv_a_proj_with_mqa.weight": "layers.{}.attention.wkv_a.weight", + "language_model.model.layers.{}.self_attn.kv_a_layernorm.weight": "layers.{}.attention.kv_norm.weight", + "language_model.model.layers.{}.self_attn.kv_b_proj.weight": "layers.{}.attention.wkv_b.weight", + "language_model.model.layers.{}.self_attn.g_proj.weight": "layers.{}.attention.gate.weight", + "language_model.model.layers.{}.self_attn.o_proj.weight": "layers.{}.attention.wo.weight", + } + self.kda_from_hf_map = { + "language_model.model.layers.{}.self_attn.q_proj.weight": "layers.{}.delta_attention.q_proj.weight", + "language_model.model.layers.{}.self_attn.k_proj.weight": "layers.{}.delta_attention.k_proj.weight", + "language_model.model.layers.{}.self_attn.v_proj.weight": "layers.{}.delta_attention.v_proj.weight", + "language_model.model.layers.{}.self_attn.q_conv1d.weight": "layers.{}.delta_attention.q_conv.weight", + "language_model.model.layers.{}.self_attn.k_conv1d.weight": "layers.{}.delta_attention.k_conv.weight", + "language_model.model.layers.{}.self_attn.v_conv1d.weight": "layers.{}.delta_attention.v_conv.weight", + "language_model.model.layers.{}.self_attn.f_a_proj.weight": "layers.{}.delta_attention.forget_a.weight", + "language_model.model.layers.{}.self_attn.f_b_proj.weight": "layers.{}.delta_attention.forget_b.weight", + "language_model.model.layers.{}.self_attn.b_proj.weight": "layers.{}.delta_attention.beta.weight", + "language_model.model.layers.{}.self_attn.g_proj.weight": "layers.{}.delta_attention.output_gate.weight", + "language_model.model.layers.{}.self_attn.o_norm.weight": "layers.{}.delta_attention.output_norm.weight", + "language_model.model.layers.{}.self_attn.o_proj.weight": "layers.{}.delta_attention.output_proj.weight", + "language_model.model.layers.{}.self_attn.A_log": "layers.{}.delta_attention.A_log", + "language_model.model.layers.{}.self_attn.dt_bias": "layers.{}.delta_attention.dt_bias", + } + + # The released index contains MXFP4 packed/scale FQNs, while this + # adapter exports unquantized weights. + self.fqn_to_index_mapping = None + + def _map_from_hf_layer_key( + self, + abstract_key: str, + layer_num: str, + ) -> str | None: + new_key = self.from_hf_map.get(abstract_key) + if new_key is not None: + return new_key + + layer_config = self.kimi_config.layers[int(layer_num)] + attention_map = ( + self.mla_from_hf_map + if layer_config.attention is not None + else self.kda_from_hf_map ) - if any(marker in key for marker in quantized_markers): - raise NotImplementedError( - "Kimi K3 v1 only supports unquantized HuggingFace state " - f"dicts; encountered quantized key '{key}'." - ) + return attention_map.get(abstract_key) - def from_hf(self, hf_state_dict: dict[str, Any]) -> dict[str, Any]: - """Convert an unquantized HuggingFace state dict to TorchTitan.""" - state_dict: dict[str, Any] = {} + def to_hf(self, state_dict: dict[str, Any]) -> dict[str, Any]: + """Convert a TorchTitan state dict to unquantized HuggingFace format.""" + to_hf_map = { + tt_key: hf_key + for mapping in ( + self.from_hf_map, + self.mla_from_hf_map, + self.kda_from_hf_map, + ) + for hf_key, tt_key in mapping.items() + } + hf_state_dict: dict[str, Any] = {} + vision_qkv_by_layer: dict[str, dict[str, torch.Tensor]] = {} unmapped: list[str] = [] - for hf_key, value in hf_state_dict.items(): - self._raise_if_quantized_key(hf_key) - if hf_key.endswith("rotary_emb.inv_freq"): - continue + for key, value in state_dict.items(): + if "moe.routed_experts.inner_experts" in key: + abstract_key = re.sub(r"(?<=\.)\d+(?=\.)", "{}", key, count=1) + layer_num_match = re.search(r"layers\.(\d+)\.", key) + assert layer_num_match is not None + layer_num = layer_num_match.group(1) + hf_abstract_key = to_hf_map.get(abstract_key) + if hf_abstract_key is None: + unmapped.append(key) + continue - tt_key = _TEXT_GLOBAL_FROM_HF.get(hf_key) - if tt_key is not None: - state_dict[tt_key] = value + if isinstance(value, DTensor): + self.grouped_expert_weight_placements[ + abstract_key + ] = value.placements + self.grouped_expert_weight_shape[abstract_key] = value.shape + self.grouped_expert_weight_mesh[abstract_key] = value.device_mesh + hf_state_dict.update( + self._get_local_experts_weights( + hf_abstract_key, + abstract_key, + layer_num, + value, + ) + ) + else: + moe_config = self.kimi_config.layers[int(layer_num)].moe + assert moe_config is not None + split_values = self._split_experts_weights( + value, + moe_config.num_experts, + ) + for expert_num, expert_weight in enumerate(split_values): + hf_state_dict[ + hf_abstract_key.format(layer_num, expert_num) + ] = expert_weight.squeeze(0) continue - tt_key = _VISION_GLOBAL_FROM_HF.get(hf_key) - if tt_key is not None: - if hf_key == "vision_tower.patch_embed.proj.weight": - value = value.reshape(value.shape[0], -1) - state_dict[tt_key] = value + vision_qkv_match = re.fullmatch( + r"vision_encoder\.layers\.(\d+)\.attn\.w(q|k|v)\.weight", + key, + ) + if vision_qkv_match is not None: + layer_num, projection = vision_qkv_match.groups() + vision_qkv_by_layer.setdefault(layer_num, {})[projection] = value continue - text_match = re.fullmatch( - r"language_model\.model\.layers\.(\d+)\.(.+)", - hf_key, - ) - if text_match is not None: - layer_idx, suffix = text_match.groups() - expert_match = re.fullmatch( - r"block_sparse_moe\.experts\.(\d+)\." r"(w1|w2|w3)\.weight", - suffix, + layer_num_match = re.search(r"(?<=\.)\d+(?=\.)", key) + if layer_num_match is not None: + layer_num = layer_num_match.group(0) + abstract_key = re.sub( + r"(?<=\.)\d+(?=\.)", + "{}", + key, + count=1, ) - if expert_match is not None: - expert_idx, projection = expert_match.groups() - moe_config = self.kimi_config.layers[int(layer_idx)].moe - assert moe_config is not None - grouped_key = ( - f"layers.{layer_idx}.moe.routed_experts.inner_experts." - f"{_EXPERT_PROJECTION_TO_GROUPED_PARAM[projection]}" - ) - experts = self._expert_weights_by_layer_projection.setdefault( - (layer_idx, projection), {} - ) - experts[int(expert_idx)] = value - if len(experts) == moe_config.num_experts: - sorted_experts = [ - experts[i] for i in range(moe_config.num_experts) - ] - state_dict[grouped_key] = torch.stack(sorted_experts, dim=0) - del self._expert_weights_by_layer_projection[ - (layer_idx, projection) - ] - continue - - layer_config = self.kimi_config.layers[int(layer_idx)] - mapped_suffix = _TEXT_LAYER_FROM_HF.get(suffix) - if mapped_suffix is None: - attention_map = ( - _MLA_FROM_HF - if layer_config.attention is not None - else _KDA_FROM_HF - ) - mapped_suffix = attention_map.get(suffix) - if mapped_suffix is None: - mapped_suffix = _MOE_FROM_HF.get(suffix) - if mapped_suffix is None: - unmapped.append(hf_key) + hf_abstract_key = to_hf_map.get(abstract_key) + if hf_abstract_key is None: + unmapped.append(key) continue - if suffix == "self_attn.dt_bias": - delta_config = layer_config.delta_attention - if delta_config is None: - raise ValueError(f"HF key '{hf_key}' targets a non-KDA layer.") - value = value.reshape( - delta_config.num_heads, - delta_config.head_dim, - ) - state_dict[f"layers.{layer_idx}.{mapped_suffix}"] = value + if abstract_key == "layers.{}.delta_attention.dt_bias": + value = value.reshape(-1) + hf_state_dict[hf_abstract_key.format(layer_num)] = value continue - vision_match = re.fullmatch( - r"vision_tower\.encoder\.blocks\.(\d+)\.(.+)", - hf_key, - ) - if vision_match is not None: - layer_idx, suffix = vision_match.groups() - if suffix == "wqkv.weight": - q, k, v = torch.chunk(value, 3, dim=0) - base = f"vision_encoder.layers.{layer_idx}.attn" - state_dict[f"{base}.wq.weight"] = q - state_dict[f"{base}.wk.weight"] = k - state_dict[f"{base}.wv.weight"] = v - continue - mapped_suffix = _VISION_LAYER_FROM_HF.get(suffix) - if mapped_suffix is None: - unmapped.append(hf_key) - continue - state_dict[f"vision_encoder.layers.{layer_idx}.{mapped_suffix}"] = value + hf_key = to_hf_map.get(key) + if hf_key is None: + unmapped.append(key) continue + if key == "vision_encoder.patch_embed.weight": + vision_config = self.kimi_config.vision_encoder + if vision_config is None: + raise ValueError( + "Vision state was provided for a text-only config." + ) + value = value.reshape( + value.shape[0], + vision_config.in_channels, + vision_config.patch_size, + vision_config.patch_size, + ) + hf_state_dict[hf_key] = value - unmapped.append(hf_key) + for layer_num, qkv in vision_qkv_by_layer.items(): + missing = {"q", "k", "v"} - qkv.keys() + if missing: + raise ValueError( + f"Vision layer {layer_num} is missing QKV parts: {sorted(missing)}." + ) + hf_state_dict[ + f"vision_tower.encoder.blocks.{layer_num}.wqkv.weight" + ] = torch.cat((qkv["q"], qkv["k"], qkv["v"]), dim=0) if unmapped: raise ValueError( - "KimiK3StateDictAdapter found HuggingFace keys without a " + "KimiK3StateDictAdapter found TorchTitan keys without a " f"mapping: {unmapped}." ) - if self._expert_weights_by_layer_projection: - incomplete = list(self._expert_weights_by_layer_projection.keys()) - self._expert_weights_by_layer_projection.clear() - raise ValueError( - "KimiK3StateDictAdapter received an incomplete set of " - f"routed-expert weights for (layer, projection): {incomplete}." - ) - return state_dict - - def to_hf(self, state_dict: dict[str, Any]) -> dict[str, Any]: - """Convert a TorchTitan state dict to unquantized HuggingFace format.""" - text_global_to_hf = {value: key for key, value in _TEXT_GLOBAL_FROM_HF.items()} - vision_global_to_hf = { - value: key for key, value in _VISION_GLOBAL_FROM_HF.items() - } - text_layer_to_hf = { - value: key - for mapping in ( - _TEXT_LAYER_FROM_HF, - _MLA_FROM_HF, - _KDA_FROM_HF, - _MOE_FROM_HF, - ) - for key, value in mapping.items() - } - vision_layer_to_hf = { - value: key for key, value in _VISION_LAYER_FROM_HF.items() - } + return hf_state_dict - hf_state_dict: dict[str, Any] = {} - vision_qkv: dict[str, dict[str, Any]] = {} + def from_hf(self, hf_state_dict: dict[str, Any]) -> dict[str, Any]: + """Convert an unquantized HuggingFace state dict to TorchTitan.""" + state_dict: dict[str, Any] = {} + expert_weights_by_layer: dict[str, dict[str, dict[int, torch.Tensor]]] = {} unmapped: list[str] = [] - for tt_key, value in state_dict.items(): - hf_key = text_global_to_hf.get(tt_key) - if hf_key is not None: - hf_state_dict[hf_key] = value + for key, value in hf_state_dict.items(): + if key.endswith("rotary_emb.inv_freq"): continue - hf_key = vision_global_to_hf.get(tt_key) - if hf_key is not None: - if tt_key == "vision_encoder.patch_embed.weight": - vision_config = self.kimi_config.vision_encoder - if vision_config is None: - raise ValueError( - "Vision state was provided for a text-only config." - ) - value = value.reshape( - value.shape[0], - vision_config.in_channels, - vision_config.patch_size, - vision_config.patch_size, - ) - hf_state_dict[hf_key] = value + new_key = self.from_hf_map.get(key) + if new_key is not None: + if key == "vision_tower.patch_embed.proj.weight": + value = value.reshape(value.shape[0], -1) + state_dict[new_key] = value continue - text_match = re.fullmatch(r"layers\.(\d+)\.(.+)", tt_key) - if text_match is not None: - layer_idx, suffix = text_match.groups() - expert_match = re.fullmatch( - r"moe\.routed_experts\.inner_experts\." r"(w1_EFD|w2_EDF|w3_EFD)", - suffix, + if "block_sparse_moe.experts" in key: + abstract_key = re.sub( + r"(?<=\.)\d+(?=\.)", + "{}", + key, + count=2, ) - if expert_match is not None: - (grouped_param,) = expert_match.groups() - projection = _GROUPED_PARAM_TO_EXPERT_PROJECTION[grouped_param] - for expert_idx, expert_weight in enumerate(value.unbind(0)): - hf_state_dict[ - f"language_model.model.layers.{layer_idx}." - f"block_sparse_moe.experts.{expert_idx}." - f"{projection}.weight" - ] = expert_weight + indices = re.findall(r"(?<=\.)\d+(?=\.)", key) + if len(indices) != 2: + unmapped.append(key) continue - - mapped_suffix = text_layer_to_hf.get(suffix) - if mapped_suffix is None: - unmapped.append(tt_key) + layer_num, expert_num = indices + titan_abstract_key = self.from_hf_map.get(abstract_key) + if titan_abstract_key is None: + unmapped.append(key) continue - if suffix == "delta_attention.dt_bias": - value = value.reshape(-1) - hf_state_dict[ - f"language_model.model.layers.{layer_idx}.{mapped_suffix}" - ] = value + new_key = titan_abstract_key.format(layer_num) + + experts = expert_weights_by_layer.setdefault(layer_num, {}).setdefault( + titan_abstract_key, {} + ) + experts[int(expert_num)] = value + + if titan_abstract_key in self.local_experts_indices: + stacked_value = self._concatenate_expert_weights_dtensor( + expert_weights_by_layer, + titan_abstract_key, + layer_num, + ) + else: + moe_config = self.kimi_config.layers[int(layer_num)].moe + assert moe_config is not None + stacked_value = self._concatenate_expert_weights( + expert_weights_by_layer, + titan_abstract_key, + layer_num, + moe_config.num_experts, + ) + if stacked_value is not None: + state_dict[new_key] = stacked_value continue - vision_match = re.fullmatch( - r"vision_encoder\.layers\.(\d+)\.(.+)", - tt_key, - ) - if vision_match is not None: - layer_idx, suffix = vision_match.groups() - qkv_match = re.fullmatch(r"attn\.w(q|k|v)\.weight", suffix) - if qkv_match is not None: - vision_qkv.setdefault(layer_idx, {})[qkv_match.group(1)] = value + layer_num_match = re.search(r"(?<=\.)\d+(?=\.)", key) + if layer_num_match is not None: + layer_num = layer_num_match.group(0) + abstract_key = re.sub( + r"(?<=\.)\d+(?=\.)", + "{}", + key, + count=1, + ) + + if abstract_key == "vision_tower.encoder.blocks.{}.wqkv.weight": + q, k, v = torch.chunk(value, 3, dim=0) + base = f"vision_encoder.layers.{layer_num}.attn" + state_dict[f"{base}.wq.weight"] = q + state_dict[f"{base}.wk.weight"] = k + state_dict[f"{base}.wv.weight"] = v continue - mapped_suffix = vision_layer_to_hf.get(suffix) - if mapped_suffix is None: - unmapped.append(tt_key) + + new_abstract_key = ( + self._map_from_hf_layer_key(abstract_key, layer_num) + if key.startswith("language_model.model.layers.") + else self.from_hf_map.get(abstract_key) + ) + if new_abstract_key is None: + unmapped.append(key) continue - hf_state_dict[ - f"vision_tower.encoder.blocks.{layer_idx}.{mapped_suffix}" - ] = value + if new_abstract_key == "layers.{}.delta_attention.dt_bias": + delta_config = self.kimi_config.layers[ + int(layer_num) + ].delta_attention + if delta_config is None: + raise ValueError(f"HF key '{key}' targets a non-KDA layer.") + value = value.reshape( + delta_config.num_heads, + delta_config.head_dim, + ) + state_dict[new_abstract_key.format(layer_num)] = value continue - unmapped.append(tt_key) - - for layer_idx, qkv in vision_qkv.items(): - missing = {"q", "k", "v"} - qkv.keys() - if missing: - raise ValueError( - f"Vision layer {layer_idx} is missing QKV parts: {sorted(missing)}." - ) - hf_state_dict[ - f"vision_tower.encoder.blocks.{layer_idx}.wqkv.weight" - ] = torch.cat((qkv["q"], qkv["k"], qkv["v"]), dim=0) + unmapped.append(key) if unmapped: raise ValueError( - "KimiK3StateDictAdapter found TorchTitan keys without a " + "KimiK3StateDictAdapter found HuggingFace keys without a " f"mapping: {unmapped}." ) - return hf_state_dict + if expert_weights_by_layer: + raise ValueError( + "KimiK3StateDictAdapter received an incomplete set of " + f"routed-expert weights: {expert_weights_by_layer.keys()}." + ) + return state_dict From 9f02c89d87a6a10652c573c3a8928ce2c4aeba02 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Sat, 15 Aug 2026 18:08:35 +0000 Subject: [PATCH 30/67] refactor KimiGroupedExperts to use self.grouped_mm --- torchtitan/models/kimi_k3/model.py | 20 +++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/torchtitan/models/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py index 2d249dc3c1..7a1482f194 100644 --- a/torchtitan/models/kimi_k3/model.py +++ b/torchtitan/models/kimi_k3/model.py @@ -383,7 +383,7 @@ class KimiGroupedExperts(GroupedExperts): Inherits its stacked-weight shape (``w1_EFD``/``w2_EDF``/``w3_EFD``) and parameter allocation; only ``forward`` differs, since the activation is - baked into the ``torch._grouped_mm`` call sequence rather than being a + baked into the ``_grouped_mm`` call sequence rather than being a swappable argument. """ @@ -415,11 +415,15 @@ def forward( offsets_E = torch.cumsum(num_tokens_per_expert_E, dim=0, dtype=torch.int32) - gate_RF = torch._grouped_mm( - x_RD.bfloat16(), w1_EFD.bfloat16().transpose(-2, -1), offs=offsets_E + gate_RF = self._grouped_mm( + A=x_RD.bfloat16(), + B_t=w1_EFD.bfloat16().transpose(-2, -1), + offs=offsets_E, ) - up_RF = torch._grouped_mm( - x_RD.bfloat16(), w3_EFD.bfloat16().transpose(-2, -1), offs=offsets_E + up_RF = self._grouped_mm( + A=x_RD.bfloat16(), + B_t=w3_EFD.bfloat16().transpose(-2, -1), + offs=offsets_E, ) input_dtype = gate_RF.dtype @@ -430,8 +434,10 @@ def forward( up_RF = self.linear_beta * torch.tanh(up_RF / self.linear_beta) h_RF = (gate_RF * up_RF).to(input_dtype) - return torch._grouped_mm( - h_RF, w2_EDF.bfloat16().transpose(-2, -1), offs=offsets_E + return self._grouped_mm( + A=h_RF, + B_t=w2_EDF.bfloat16().transpose(-2, -1), + offs=offsets_E, ).type_as(x_RD) From 588c62e881f998273fa5fa5d56712c655ffa2b1e Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Sun, 16 Aug 2026 03:47:02 +0000 Subject: [PATCH 31/67] use snapshot_download to pin the version --- .../numerical_tests_kimi_k3.py | 24 ++++++++++--------- 1 file changed, 13 insertions(+), 11 deletions(-) diff --git a/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py b/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py index b37bceb1b3..80b0d969fe 100644 --- a/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py +++ b/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py @@ -15,22 +15,22 @@ script reduces the HuggingFace config to TorchTitan's debug model, initializes TorchTitan, and strictly transfers its state dict to HuggingFace. -The local HuggingFace directory must contain the config, modeling, processor, -tokenizer code, and tokenizer assets. The released code requires -``transformers==4.56.2`` and ``tiktoken``. +The script downloads the config, modeling, processor, and tokenizer assets from +a pinned HuggingFace revision without downloading the released weight shards. +The released code requires ``transformers==4.56.2`` and ``tiktoken``. Usage: CUDA_VISIBLE_DEVICES=0 python -m \ scripts.checkpoint_conversion.numerical_tests_kimi_k3 \ - --hf_model_path ~/hf_assets/moonshotai/Kimi-K3 --dtype float32 + --dtype float32 """ import argparse -import os from typing import Any, cast import torch import torch.nn.functional as F +from huggingface_hub import snapshot_download from PIL import Image from torchtitan.hf_datasets.multimodal.utils.image import ( @@ -44,6 +44,8 @@ from transformers import AutoConfig, AutoModelForCausalLM, AutoProcessor +_HF_REPO_ID = "moonshotai/Kimi-K3" +_HF_REVISION = "9f62e4e9fffbd0a83ddd60e1c209d828994b3569" _MEDIA_TOKEN_ID = 163605 _PATCH_SIZE = 14 _MERGE_SIZE = 2 @@ -418,11 +420,6 @@ def compare(ref_logits: torch.Tensor, tt_logits: torch.Tensor) -> bool: @torch.no_grad() def main() -> None: parser = argparse.ArgumentParser() - parser.add_argument( - "--hf_model_path", - default=os.path.expanduser("~/hf_assets/moonshotai/Kimi-K3"), - help="Local directory containing the Kimi K3 HuggingFace assets.", - ) parser.add_argument("--model_flavor", default="debugmodel") parser.add_argument("--image_size", type=int, default=336) parser.add_argument( @@ -446,6 +443,11 @@ def main() -> None: if not torch.cuda.is_available(): parser.error("Kimi K3 numerical parity requires a CUDA GPU.") + hf_model_path = snapshot_download( + repo_id=_HF_REPO_ID, + revision=_HF_REVISION, + allow_patterns=["*.json", "*.py", "tiktoken.model"], + ) device = torch.device("cuda") dtype = getattr(torch, args.dtype) vision_dtype = getattr(torch, args.vision_dtype) if args.vision_dtype else dtype @@ -463,7 +465,7 @@ def main() -> None: ) ref = run_hf( - args.hf_model_path, + hf_model_path, tt_config, hf_state_dict, args.image_size, From edccb38959c04db6c3a00f6bea5c786ecf937425 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Sun, 16 Aug 2026 06:49:30 +0000 Subject: [PATCH 32/67] refator test case --- tests/unit_tests/test_kimi_k3.py | 416 +++++++++++++++---------------- 1 file changed, 206 insertions(+), 210 deletions(-) diff --git a/tests/unit_tests/test_kimi_k3.py b/tests/unit_tests/test_kimi_k3.py index 84f434c3ed..54924591a6 100644 --- a/tests/unit_tests/test_kimi_k3.py +++ b/tests/unit_tests/test_kimi_k3.py @@ -4,16 +4,22 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import copy import unittest -from dataclasses import dataclass +from unittest.mock import patch import torch -import torch.nn as nn import torch.nn.functional as F +from torch.distributed._composable.fsdp import FSDPModule +from torch.distributed.tensor import DTensor from torch.nn.attention.flex_attention import BlockMask +from torch.testing._internal.distributed._tensor.common_dtensor import ( + DTensorTestBase, + with_comms, +) -from torchtitan.models.common import Embedding -from torchtitan.protocols.module import Module +from torchtitan.config import CompileConfig, ParallelismConfig, TrainingConfig +from torchtitan.distributed import ParallelDims # torchtitan.models.kimi_k3 imports FLA at module scope for the KDA kernel. # FLA is a per-model dependency (kimi_k3/requirements.txt), not part of the @@ -21,19 +27,11 @@ # when it is absent. Modules importing this one inherit the skip. try: from torchtitan.models.kimi_k3 import ( - _feed_forward_config, - _kda_config, - _latent_moe_config, - _linear, - _mla_config, - _norm, + _kimi_k3_config, _vision_encoder_config, + parallelize_kimi_k3, ) - from torchtitan.models.kimi_k3.model import ( - KimiK3Model, - KimiK3TransformerBlock, - KimiKDAKernel, - ) + from torchtitan.models.kimi_k3.model import KimiK3Model, KimiKDAKernel from torchtitan.models.kimi_k3.state_dict_adapter import KimiK3StateDictAdapter except ModuleNotFoundError as exc: raise unittest.SkipTest( @@ -41,167 +39,50 @@ ) from exc -class ReferenceKimiKDAKernel(Module): - """Pure-PyTorch stand-in for KimiKDAKernel backed by an explicit recurrence. - - Mirrors ``KimiKDAKernel.forward``'s interface so tests can build a model - with it in place of the FLA kernel and exercise the surrounding model - on CPU. The loop is O(seqlen) and far too slow for training; it exists to - pin the kernel's math. - """ - - @dataclass(kw_only=True, slots=True) - class Config(Module.Config): - head_dim: int - lower_bound: float | None = -5.0 - - def __init__(self, config: Config): - super().__init__() - self.head_dim = config.head_dim - self.lower_bound = config.lower_bound - - def forward( - self, - q_BLHK: torch.Tensor, - k_BLHK: torch.Tensor, - v_BLHV: torch.Tensor, - gate_BLHK: torch.Tensor, - beta_BLH: torch.Tensor, - A_log_H: torch.Tensor, - dt_bias_HK: torch.Tensor, - ) -> torch.Tensor: - return _kda_recurrent_reference( - q_BLHK, - k_BLHK, - v_BLHV, - gate_BLHK, - beta_BLH, - A_log_H, - dt_bias_HK, - lower_bound=self.lower_bound, - ) - - -def _use_reference_kda_kernel(config: KimiK3Model.Config) -> KimiK3Model.Config: - """Point every KDA layer at the recurrent reference kernel. - - Test configurations use head dimensions far below what FLA's chunked KDA - kernel can compile, and the CPU suite has no Triton runtime at all. - """ - for layer in config.layers: - if layer.delta_attention is None: - continue - kernel = layer.delta_attention.kernel - assert isinstance(kernel, KimiKDAKernel.Config) - layer.delta_attention.kernel = ReferenceKimiKDAKernel.Config( - head_dim=kernel.head_dim, - lower_bound=kernel.lower_bound, - ) - return config - - -def _small_model_config(*, attn_res_block_size: int = 1) -> KimiK3Model.Config: - """Build the reduced two-layer model used across the Kimi K3 tests. - - ``attn_res_block_size`` defaults to 1, which makes every layer extend the - attention residual. Pass 2 to make the second layer pass the residual - through instead, which is the shape the released cadence uses and which - routes the residual back out through the FSDP module boundary. Callers - comparing against frozen reference values must keep the default, since the - parameter shapes and ordering feed those values. - """ - dim = 16 - - def block( - layer_id: int, - *, - use_mla: bool, - use_moe: bool, - ) -> KimiK3TransformerBlock.Config: - return KimiK3TransformerBlock.Config( - layer_id=layer_id, - attn_res_block_size=attn_res_block_size, - attention=( - _mla_config( - dim=dim, - num_heads=2, - q_lora_rank=8, - kv_lora_rank=8, - qk_nope_head_dim=4, - qk_rope_head_dim=4, - v_head_dim=4, - attn_backend="flex", - ) - if use_mla - else None - ), - delta_attention=( - None - if use_mla - else _kda_config( - dim=dim, - num_heads=2, - head_dim=4, - conv_kernel_size=3, - ) - ), - feed_forward=( - None if use_moe else _feed_forward_config(dim=dim, hidden_dim=32) - ), - moe=( - _latent_moe_config( - dim=dim, - latent_dim=8, - expert_hidden_dim=8, - num_experts=2, - top_k=1, - num_shared_experts=1, - ) - if use_moe - else None - ), - attention_norm=_norm(dim), - ffn_norm=_norm(dim), - attention_res_norm=_norm(dim), - attention_res_proj=_linear(dim, 1), - ffn_res_norm=_norm(dim), - ffn_res_proj=_linear(dim, 1), - ) - - return _use_reference_kda_kernel( - KimiK3Model.Config( - dim=dim, - vocab_size=32, - tok_embeddings=Embedding.Config( - num_embeddings=32, - embedding_dim=dim, - param_init={ - "weight": lambda parameter: nn.init.normal_(parameter, std=0.02) - }, - ), - layers=[ - block(0, use_mla=False, use_moe=False), - block(1, use_mla=True, use_moe=True), - ], - norm=_norm(dim), - lm_head=_linear(dim, 32), - output_res_norm=_norm(dim), - output_res_proj=_linear(dim, 1), - vision_encoder=_vision_encoder_config( - text_dim=dim, - dim=16, - qkv_dim=24, - hidden_dim=32, - num_layers=1, - num_heads=3, - patch_size=2, - merge_kernel_size=(2, 2), - init_pos_emb_height=2, - init_pos_emb_width=2, - max_num_frames=1, - ), - spatial_merge_size=2, - ) +def _small_model_config( + *, + attn_res_block_size: int = 1, + full_attention_layers: set[int] | None = None, +) -> KimiK3Model.Config: + """Build a reduced KDA+MLA, dense+MoE, multimodal Kimi K3 config.""" + if full_attention_layers is None: + full_attention_layers = {1} + + dim = 64 + return _kimi_k3_config( + dim=dim, + vocab_size=32, + num_layers=2, + full_attention_layers=full_attention_layers, + attn_res_block_size=attn_res_block_size, + num_heads=2, + q_lora_rank=32, + kv_lora_rank=32, + qk_nope_head_dim=16, + qk_rope_head_dim=16, + v_head_dim=16, + kda_head_dim=16, + conv_kernel_size=3, + dense_hidden_dim=128, + latent_dim=32, + expert_hidden_dim=32, + num_experts=2, + top_k=1, + num_shared_experts=1, + vision_encoder=_vision_encoder_config( + text_dim=dim, + dim=48, + qkv_dim=48, + hidden_dim=96, + num_layers=1, + num_heads=3, + patch_size=2, + merge_kernel_size=(2, 2), + init_pos_emb_height=2, + init_pos_emb_width=2, + max_num_frames=1, + ), + attn_backend="flex", ) @@ -304,36 +185,28 @@ def parameter(*shape: int) -> torch.Tensor: for lower_bound in (-5.0, None): with self.subTest(lower_bound=lower_bound): - q_BLHK = parameter(2, 64, num_heads, head_dim) - k_BLHK = parameter(2, 64, num_heads, head_dim) - v_BLHV = parameter(2, 64, num_heads, head_dim) - gate_BLHK = parameter(2, 64, num_heads, head_dim) - beta_BLH = parameter(2, 64, num_heads) A_log_H = torch.rand(num_heads, device="cuda") A_log_H = A_log_H.uniform_(1.0, 16.0).log().requires_grad_() - dt_bias_HK = parameter(num_heads, head_dim) + actual_inputs = ( + parameter(2, 64, num_heads, head_dim), + parameter(2, 64, num_heads, head_dim), + parameter(2, 64, num_heads, head_dim), + parameter(2, 64, num_heads, head_dim), + parameter(2, 64, num_heads), + A_log_H, + parameter(num_heads, head_dim), + ) + expected_inputs = tuple( + tensor.detach().clone().requires_grad_() for tensor in actual_inputs + ) kernel = KimiKDAKernel.Config( head_dim=head_dim, lower_bound=lower_bound, ).build() - actual_BLHV = kernel( - q_BLHK, - k_BLHK, - v_BLHV, - gate_BLHK, - beta_BLH, - A_log_H, - dt_bias_HK, - ) + actual_BLHV = kernel(*actual_inputs) expected_BLHV = _kda_recurrent_reference( - q_BLHK, - k_BLHK, - v_BLHV, - gate_BLHK, - beta_BLH, - A_log_H, - dt_bias_HK, + *expected_inputs, lower_bound=lower_bound, ) @@ -346,19 +219,28 @@ def parameter(*shape: int) -> torch.Tensor: atol=2e-3, rtol=2e-3, ) - actual_BLHV.square().mean().backward() - for tensor in ( - q_BLHK, - k_BLHK, - v_BLHV, - gate_BLHK, - beta_BLH, - A_log_H, - dt_bias_HK, + output_grad_BLHV = torch.randn_like(actual_BLHV) + actual_grads = torch.autograd.grad( + actual_BLHV, + actual_inputs, + grad_outputs=output_grad_BLHV, + ) + expected_grads = torch.autograd.grad( + expected_BLHV, + expected_inputs, + grad_outputs=output_grad_BLHV, + ) + for actual_grad, expected_grad in zip( + actual_grads, + expected_grads, + strict=True, ): - self.assertIsNotNone(tensor.grad) - assert tensor.grad is not None - self.assertTrue(torch.isfinite(tensor.grad).all()) + torch.testing.assert_close( + actual_grad, + expected_grad, + atol=2e-2, + rtol=2e-2, + ) def test_state_dict_round_trips_through_hf_adapter(self): torch.manual_seed(2) @@ -383,5 +265,119 @@ def test_state_dict_round_trips_through_hf_adapter(self): torch.testing.assert_close(value, roundtrip_state_dict[key]) +class TestKimiK3FSDP(DTensorTestBase): + @property + def world_size(self): + return 1 + + @unittest.skipIf(not torch.cuda.is_available(), "Kimi K3 FSDP requires CUDA.") + @with_comms + def test_fsdp_matches_non_distributed_forward_backward(self): + torch.manual_seed(3) + config = _small_model_config( + attn_res_block_size=2, + full_attention_layers={0, 1}, + ) + with torch.device("meta"): + model = config.build() + model.to_empty(device=self.device_type) + model.init_states() + with torch.no_grad(): + for transformer_block in model.layers.values(): + if transformer_block.moe is not None: + transformer_block.moe.router.gate.weight.zero_() + + reference = copy.deepcopy(model) + for parameter in reference.parameters(): + parameter.data = parameter.data.to(torch.bfloat16) + + parallelism = ParallelismConfig( + data_parallel_shard_degree=1, + tensor_parallel_degree=1, + pipeline_parallel_degree=1, + context_parallel_degree=1, + expert_parallel_degree=1, + ) + parallel_dims = ParallelDims.from_config(parallelism, world_size=1) + with patch( + "torchtitan.distributed.parallel_dims.device_type", + self.device_type, + ): + parallel_dims.build_mesh() + model = parallelize_kimi_k3( + model, + parallel_dims=parallel_dims, + training=TrainingConfig( + local_batch_size=1, + seq_len=6, + steps=1, + dtype="bfloat16", + ), + parallelism=parallelism, + compile_config=CompileConfig(), + ac_config=None, + dump_folder="", + ) + + assert isinstance(model, KimiK3Model) + self.assertIsInstance(model, FSDPModule) + self.assertIsInstance(model.vision_encoder, FSDPModule) + + positions_BL = torch.arange( + 6, + dtype=torch.int32, + device=self.device_type, + ).unsqueeze(0) + attention_masks = reference.get_attention_masks(positions_BL) + inputs = { + "tokens": torch.tensor( + [[1, 7, 2, 3, 4, 5]], + dtype=torch.long, + device=self.device_type, + ), + "pixel_values": torch.randn( + 1, + 4, + 3 * 2 * 2, + device=self.device_type, + ), + "grid_thw": torch.tensor( + [[1, 2, 2]], + dtype=torch.long, + device=self.device_type, + ), + "special_tokens": {"image_id": 7}, + "positions": positions_BL, + "attention_masks": attention_masks, + } + + actual_BLV = model(**inputs) # pyrefly: ignore [not-callable] + expected_BLV = reference(**inputs) + torch.testing.assert_close(actual_BLV, expected_BLV, atol=0.0, rtol=0.0) + + actual_BLV.float().square().mean().backward() + expected_BLV.float().square().mean().backward() + + reference_parameters = dict(reference.named_parameters()) + compared_gradients = 0 + for name, parameter in model.named_parameters(): + actual_grad = parameter.grad + expected_grad = reference_parameters[name].grad + self.assertEqual(actual_grad is None, expected_grad is None) + if actual_grad is None: + continue + if isinstance(actual_grad, DTensor): + actual_grad = actual_grad.to_local() + assert expected_grad is not None + torch.testing.assert_close( + actual_grad.float(), + expected_grad.float(), + atol=0.0, + rtol=0.0, + ) + compared_gradients += 1 + self.assertGreater(compared_gradients, 0) + + if __name__ == "__main__": unittest.main() From 02960f7573f7a731b78b5434841687ce1f419b42 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Sun, 16 Aug 2026 11:33:31 +0000 Subject: [PATCH 33/67] fix pyrefly --- torchtitan/models/kimi_k3/model.py | 1 + 1 file changed, 1 insertion(+) diff --git a/torchtitan/models/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py index 7a1482f194..42ed10a9e2 100644 --- a/torchtitan/models/kimi_k3/model.py +++ b/torchtitan/models/kimi_k3/model.py @@ -525,6 +525,7 @@ def _apply_attention_residual( norm: RMSNorm, ) -> torch.Tensor: """Apply Kimi's block-level attention residual in FP32.""" + assert norm.eps is not None values_TND = torch.cat((block_residual_TND, prefix_sum_TD.unsqueeze(1)), dim=1) values_float = values_TND.float() From 2e7603abeaa87b266944bae4e2a7f885b391c745 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Mon, 17 Aug 2026 02:46:15 +0000 Subject: [PATCH 34/67] support full kimi k3 and refactor readme --- torchtitan/models/kimi_k3/README.md | 120 ++++++++------------------ torchtitan/models/kimi_k3/__init__.py | 37 ++++++++ 2 files changed, 71 insertions(+), 86 deletions(-) diff --git a/torchtitan/models/kimi_k3/README.md b/torchtitan/models/kimi_k3/README.md index d4e0385b44..2272fabbf5 100644 --- a/torchtitan/models/kimi_k3/README.md +++ b/torchtitan/models/kimi_k3/README.md @@ -1,9 +1,7 @@ # Kimi K3 -Kimi K3 combines a hybrid **Kimi Delta Attention (KDA) + Multi-head Latent -Attention (MLA)** decoder, **LatentMoE**, and a **MoonViT3d** vision encoder. -TorchTitan currently provides a topology-complete reduced model for architecture -validation, single-device training, and FSDP2 training. +Kimi K3 combines a hybrid Kimi Delta Attention (KDA) and Multi-head Latent +Attention (MLA) decoder with LatentMoE and a MoonViT-V2 vision encoder. ## Prerequisites @@ -15,95 +13,45 @@ pip install av einops pillow torchvision flash-linear-attention ## Architecture -- **Decoder** -- hybrid KDA and MLA layers. The MLA layers follow the released - model's explicit 1-based layer list, including consecutive MLA layers at the - end of the decoder. -- **Feed-forward layers** -- one dense SiTU feed-forward layer followed by - LatentMoE layers with sigmoid top-k routing, correction bias, routed experts, - and shared experts. -- **Attention residuals** -- block-level attention residual connections, - including the final output residual. -- **KDA backend** -- FLA's chunked Triton kernel, with a pure PyTorch recurrent - implementation in the unit tests as the numerical reference. -- **Vision encoder** -- MoonViT3d with learned spatial positions, 2D RoPE, - non-causal attention, temporal pooling, 2x2 spatial merge, and a two-layer - projector to the decoder dimension. -- **Multimodal forward** -- projected vision embeddings are scattered into runs - of the shared media placeholder token. - -## Model variants - -Only `debugmodel` is currently registered. The released Kimi K3 row is included -for architectural comparison and is not a runnable TorchTitan flavor. - -| Variant | Parameters | LLM dim | Layers | MLA layers (1-based) | KDA layers | Heads | Experts (top-k) | ViT dim / layers / heads | -|---------|------------|---------|--------|----------------------|------------|-------|-----------------|--------------------------| -| Released Kimi K3 (reference) | 2.8T | 7168 | 93 | 4, 8, ..., 92, 93 | 69 | 96 | 896 (top-16) | 1024 / 27 / 12 | -| debugmodel | 100M | 256 | 13 | 4, 8, 12, 13 | 9 | 4 | 8 (top-2) | 256 / 4 / 3 | - -`debugmodel` retains the released vocabulary size of 163840. Its depth also -preserves two structural edge cases from the released model: consecutive final -MLA layers and a short trailing attention-residual block. +Kimi K3 is built on Kimi Delta Attention (KDA) and Attention Residuals +(AttnRes), with 69 KDA layers and 24 Gated MLA layers. Stable LatentMoE selects +16 of 896 experts per token, and MoonViT-V2 provides native vision input. + +## Released Model Configuration + +The values below follow the +[official Kimi K3 model card](https://huggingface.co/moonshotai/Kimi-K3) and +describe the released model. + +| Component | Configuration | +|-----------|---------------| +| Architecture | Mixture-of-Experts (MoE) | +| Parameters | 2.8T total, 104B activated | +| Decoder | 93 layers, 1 dense layer, hidden size 7168, 96 attention heads | +| Attention | 69 KDA layers and 24 Gated MLA layers, context length 1048576 | +| LatentMoE | Dimension 3584, hidden size 3072 per expert, 896 experts, top-16 routing, 2 shared experts | +| Vocabulary | 160K | +| Activation | SiTU-GLU | +| Vision encoder | MoonViT-V2, 401M parameters | +| Quantization | MXFP4 weights and MXFP8 activations with quantization-aware training | +| Modality | Text and image | ## Supported Parallelisms | Feature | Notes | |---------|-------| -| FSDP / HSDP | Supported with the default SPMD backend. The decoder is sharded per layer and the vision encoder is a separate FSDP unit | -| Tensor Parallelism (TP) | Not supported | -| Expert Parallelism (EP) | Not supported | -| Pipeline Parallelism (PP) | Not supported | -| Context Parallelism (CP) | Not supported | +| FSDP2 / HSDP | Decoder sharded per layer; vision encoder sharded as a separate unit | -`torch.compile`, activation checkpointing, and parameter CPU offload are not -supported by the current Kimi K3 parallelization path. +## Numerical Parity -Run the debug model on one GPU: +End-to-end KL divergence against the Hugging Face implementation: +**2.0644e-6**, with **100% top-1 and top-5match**. -```bash -NGPU=1 MODULE=kimi_k3 CONFIG=kimi_k3_debugmodel ./run_train.sh -``` - -Run it with two-way FSDP2: - -```bash -NGPU=2 MODULE=kimi_k3 CONFIG=kimi_k3_debugmodel ./run_train.sh \ - --parallelism.data_parallel_shard_degree 2 -``` - -## Numerical Checks - -`scripts/checkpoint_conversion/numerical_tests_kimi_k3.py` loads the released -HuggingFace config, modeling code, processor, and tokenizer from a local model -directory. It reduces the HuggingFace model to the `debugmodel` topology and -transfers the randomly initialized TorchTitan state dict, without loading the -released weights. Each side performs its own image preprocessing before the -full vision-projector-decoder forward. Float32 is the default correctness mode, -and the script does not override the framework's TF32 settings. - -The current float32 CUDA validation result is: - -- pixel preprocessing: max difference `1.192e-7`, with no values differing - above `1e-6`; -- projected vision features: cosine similarity `1.000000`, max difference - `3.152e-3`; -- expert routing: all `3936 / 3936` choices match; -- end-to-end last-token logits: KL `1.8215e-8`, cosine similarity `1.000002`, - max difference `4.1358e-3`, top-1 match, and top-5 5/5. - -Run the comparison with: - -```bash -python -m scripts.checkpoint_conversion.numerical_tests_kimi_k3 \ - --hf_model_path ~/hf_assets/moonshotai/Kimi-K3 -``` +Vision parity: pixel preprocessing max difference `1.192e-7`; projected vision +features cosine similarity `1.000000` and max difference `3.152e-3`. -## TODO +Test scripts: -- Add the full 2.8T model flavor. -- Add MXFP4 compressed checkpoint loading. -- Add TP, EP, PP, and CP support. -- Add packed-document attention support. -- Add video inputs and a video dataset training pipeline. -- Add `torch.compile`, activation checkpointing, and parameter CPU offload. -- Add generation-cache support. +- `scripts/checkpoint_conversion/numerical_tests_kimi_k3.py` -- Hugging Face + versus TorchTitan end-to-end comparison +- `tests/unit_tests/test_kimi_k3.py` -- KDA kernel versus the PyTorch reference diff --git a/torchtitan/models/kimi_k3/__init__.py b/torchtitan/models/kimi_k3/__init__.py index 1bab0b5232..0f98ef4e3c 100644 --- a/torchtitan/models/kimi_k3/__init__.py +++ b/torchtitan/models/kimi_k3/__init__.py @@ -510,8 +510,45 @@ def _debugmodel(attn_backend: str) -> KimiK3Model.Config: ) +def _kimi_k3(attn_backend: str) -> KimiK3Model.Config: + dim = 7168 + return _kimi_k3_config( + dim=dim, + vocab_size=163840, + num_layers=93, + full_attention_layers=set(range(3, 92, 4)) | {92}, + attn_res_block_size=12, + num_heads=96, + q_lora_rank=1536, + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + kda_head_dim=128, + conv_kernel_size=4, + dense_hidden_dim=33792, + latent_dim=3584, + expert_hidden_dim=3072, + num_experts=896, + top_k=16, + num_shared_experts=2, + vision_encoder=_vision_encoder_config( + text_dim=dim, + dim=1024, + qkv_dim=1536, + hidden_dim=4096, + num_layers=27, + num_heads=12, + init_pos_emb_height=64, + init_pos_emb_width=64, + ), + attn_backend=attn_backend, + ) + + kimi_k3_configs = { "debugmodel": _debugmodel, + "Kimi-K3": _kimi_k3, } From 22cfd54d434cf3dd2ef81faeea49e041b88bcad2 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Mon, 17 Aug 2026 03:38:46 +0000 Subject: [PATCH 35/67] fix some NIT and some bug --- torchtitan/models/kimi_k3/README.md | 14 +++++++------- torchtitan/models/kimi_k3/__init__.py | 11 ----------- torchtitan/models/kimi_k3/config_registry.py | 3 +-- torchtitan/models/kimi_k3/model.py | 20 +++++--------------- torchtitan/models/kimi_k3/parallelize.py | 9 +++++---- torchtitan/models/kimi_k3/vision_encoder.py | 2 +- 6 files changed, 19 insertions(+), 40 deletions(-) diff --git a/torchtitan/models/kimi_k3/README.md b/torchtitan/models/kimi_k3/README.md index 2272fabbf5..dfd7616f9f 100644 --- a/torchtitan/models/kimi_k3/README.md +++ b/torchtitan/models/kimi_k3/README.md @@ -44,14 +44,14 @@ describe the released model. ## Numerical Parity -End-to-end KL divergence against the Hugging Face implementation: -**2.0644e-6**, with **100% top-1 and top-5match**. +End-to-end KL divergence against the Hugging Face implementation (multimodal +inputs): **2.0644e-6**, with **100% top-1 and top-5 match**. -Vision parity: pixel preprocessing max difference `1.192e-7`; projected vision -features cosine similarity `1.000000` and max difference `3.152e-3`. +Vision parity: pixel preprocessing max difference **1.192e-7**; projected vision +features cosine similarity **1.000000** and max difference **3.152e-3**. Test scripts: -- `scripts/checkpoint_conversion/numerical_tests_kimi_k3.py` -- Hugging Face - versus TorchTitan end-to-end comparison -- `tests/unit_tests/test_kimi_k3.py` -- KDA kernel versus the PyTorch reference +- `scripts/checkpoint_conversion/numerical_tests_kimi_k3.py` -- Hugging Face vs. + TorchTitan comparison +- `tests/unit_tests/test_kimi_k3.py` -- KDA and FSDP2 correctness diff --git a/torchtitan/models/kimi_k3/__init__.py b/torchtitan/models/kimi_k3/__init__.py index 0f98ef4e3c..2eab1a0051 100644 --- a/torchtitan/models/kimi_k3/__init__.py +++ b/torchtitan/models/kimi_k3/__init__.py @@ -4,8 +4,6 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -"""Kimi K3 model registration and architecture configurations.""" - from collections.abc import Callable from functools import partial @@ -469,14 +467,6 @@ def _kimi_k3_config( def _debugmodel(attn_backend: str) -> KimiK3Model.Config: - """Return the topology-complete Kimi K3 debug model. - - The depth is one past a multiple of both the full-attention period and the - attention-residual block size, so the last layer is a full-attention layer - directly after a scheduled one and the trailing residual block is short. - Both are properties of the released 93-layer stack, whose zero-based MLA - layer indices end ``..., 87, 91, 92``. - """ dim = 256 return _kimi_k3_config( dim=dim, @@ -557,7 +547,6 @@ def model_registry( attn_backend: str = "flex", converters: list[ModelConfigConverter.Config] | None = None, ) -> ModelSpec: - """Build a Kimi K3 model specification.""" config = kimi_k3_configs[flavor](attn_backend=attn_backend) if converters is not None: validate_converter_order(converters) diff --git a/torchtitan/models/kimi_k3/config_registry.py b/torchtitan/models/kimi_k3/config_registry.py index a5654733c6..ca55f26ba3 100644 --- a/torchtitan/models/kimi_k3/config_registry.py +++ b/torchtitan/models/kimi_k3/config_registry.py @@ -4,8 +4,6 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -"""Trainer configurations for Kimi K3.""" - from torchtitan.components.checkpoint import CheckpointManager from torchtitan.components.loss import ChunkedLossWrapper, CrossEntropyLoss from torchtitan.components.lr_scheduler import LRSchedulersContainer @@ -60,6 +58,7 @@ def kimi_k3_debugmodel() -> Trainer.Config: seq_len=256, steps=10, dtype="bfloat16", + disable_cuda_graphs=True, ), checkpoint=CheckpointManager.Config( interval=10, diff --git a/torchtitan/models/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py index 42ed10a9e2..fc72342c52 100644 --- a/torchtitan/models/kimi_k3/model.py +++ b/torchtitan/models/kimi_k3/model.py @@ -4,13 +4,6 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -"""Kimi K3 language model components. - -MLA delegates to TorchTitan's configured inner-attention backend. KDA runs on -FLA's chunked Triton kernel; the pure-PyTorch recurrence used to check it lives -in ``tests/unit_tests/test_kimi_k3.py`` and is far too slow for training. -""" - from dataclasses import dataclass, field import torch @@ -283,8 +276,6 @@ def forward( class KimiDeltaAttention(Module): - """Kimi Delta Attention with causal convolutions and reference recurrence.""" - @dataclass(kw_only=True, slots=True) class Config(Module.Config): dim: int @@ -491,10 +482,11 @@ def forward(self, x_BLD: torch.Tensor) -> torch.Tensor: -1, expert_ids_BLK, True ) num_tokens_per_expert_E = routing_map_BLE.sum(dim=(0, 1)) - with torch.no_grad(): - # In place so the load-balancing hook registered on the optimizer - # keeps referring to this buffer across steps. - self.tokens_per_expert_E.add_(num_tokens_per_expert_E.float()) + if self.training: + with torch.no_grad(): + # In place so the load-balancing hook registered on the optimizer + # keeps referring to this buffer across steps. + self.tokens_per_expert_E.add_(num_tokens_per_expert_E.float()) latent_BLD = self.routed_down(x_BLD) routed_BLD = self.routed_experts( @@ -649,8 +641,6 @@ def forward( class KimiK3Model(Decoder): - """Reduced Kimi K3 multimodal model used for first-version validation.""" - @dataclass(kw_only=True, slots=True) class Config(Decoder.Config): layers: list[KimiK3TransformerBlock.Config] diff --git a/torchtitan/models/kimi_k3/parallelize.py b/torchtitan/models/kimi_k3/parallelize.py index 99e1393c55..9101a5af03 100644 --- a/torchtitan/models/kimi_k3/parallelize.py +++ b/torchtitan/models/kimi_k3/parallelize.py @@ -4,8 +4,6 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -"""FSDP2 parallelization for Kimi K3.""" - import torch.nn as nn from torchtitan.config import ( @@ -55,8 +53,8 @@ def parallelize_kimi_k3( raise NotImplementedError( "Kimi K3 FSDP2 currently supports the default SPMD backend only." ) - if compile_config.enable: - raise NotImplementedError("Kimi K3 does not support torch.compile.") + if compile_config.enable and "model" in compile_config.components: + raise NotImplementedError("Kimi K3 does not support model compilation.") if ac_config is not None: raise NotImplementedError( "Kimi K3 FSDP2 does not support activation checkpointing yet." @@ -74,6 +72,9 @@ def parallelize_kimi_k3( assert isinstance(model, KimiK3Model) vision_encoder = model.vision_encoder if vision_encoder is not None: + # TODO: An image batch on one DP rank and a text-only batch on another + # execute different FSDP collectives, deadlock, and hit a 90-second + # timeout. A general solution is needed. apply_fsdp_to_vision_encoder( vision_encoder, dp_mesh, diff --git a/torchtitan/models/kimi_k3/vision_encoder.py b/torchtitan/models/kimi_k3/vision_encoder.py index 4f3df19873..b2833b1d93 100644 --- a/torchtitan/models/kimi_k3/vision_encoder.py +++ b/torchtitan/models/kimi_k3/vision_encoder.py @@ -390,7 +390,7 @@ def __init__(self, config: Config): ) ) self.rotary_pos_emb = config.rotary_pos_emb.build() - self._cached_freq_table: torch.Tensor | None = None + self.register_buffer("_cached_freq_table", None, persistent=False) self.layers = ModuleDict( { str(layer_idx): config.block.build() From f7feaaab03ed8f8297de50da6174d8f108e8ca18 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Mon, 17 Aug 2026 06:59:56 +0000 Subject: [PATCH 36/67] reuse VisionAttention --- .../numerical_tests_kimi_k3.py | 2 +- torchtitan/models/kimi_k3/README.md | 4 +- torchtitan/models/kimi_k3/__init__.py | 7 +- torchtitan/models/kimi_k3/model.py | 51 ++++---- torchtitan/models/kimi_k3/vision_encoder.py | 116 +++++------------- 5 files changed, 61 insertions(+), 119 deletions(-) diff --git a/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py b/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py index 80b0d969fe..e089617cbc 100644 --- a/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py +++ b/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py @@ -133,7 +133,7 @@ def _reduce_hf_config(hf_config, tt_config, hf_model_path: str) -> None: "vt_intermediate_size": vision.block.mlp.fc1.out_features, "merge_kernel_size": vision.merge_kernel_size, "mm_hidden_size": vision.dim, - "qkv_hidden_size": vision.block.attn.qkv_dim, + "qkv_hidden_size": vision.block.attn.dim, "text_hidden_size": tt_config.dim, "pos_emb_interpolation_mode": vision.interpolation_mode, } diff --git a/torchtitan/models/kimi_k3/README.md b/torchtitan/models/kimi_k3/README.md index dfd7616f9f..8d61bb00cf 100644 --- a/torchtitan/models/kimi_k3/README.md +++ b/torchtitan/models/kimi_k3/README.md @@ -45,10 +45,10 @@ describe the released model. ## Numerical Parity End-to-end KL divergence against the Hugging Face implementation (multimodal -inputs): **2.0644e-6**, with **100% top-1 and top-5 match**. +inputs): **1.5370e-6**, with **100% top-1 and top-5 match**. Vision parity: pixel preprocessing max difference **1.192e-7**; projected vision -features cosine similarity **1.000000** and max difference **3.152e-3**. +features cosine similarity **1.000000** and max difference **2.669e-3**. Test scripts: diff --git a/torchtitan/models/kimi_k3/__init__.py b/torchtitan/models/kimi_k3/__init__.py index 2eab1a0051..63e886e587 100644 --- a/torchtitan/models/kimi_k3/__init__.py +++ b/torchtitan/models/kimi_k3/__init__.py @@ -16,7 +16,7 @@ from torchtitan.models.common.moe import RoutedExperts, TokenChoiceTopKRouter from torchtitan.models.common.nn_modules import GELU, RMSNorm from torchtitan.models.common.token_dispatcher import LocalTokenDispatcher -from torchtitan.models.common.vision_encoder import VisionMLP +from torchtitan.models.common.vision_encoder import VisionAttention, VisionMLP from torchtitan.models.utils import validate_converter_order from torchtitan.protocols.model import ModelConfigConverter from torchtitan.protocols.model_spec import ModelSpec @@ -35,7 +35,6 @@ from .parallelize import parallelize_kimi_k3 from .state_dict_adapter import KimiK3StateDictAdapter from .vision_encoder import ( - KimiK3VisionAttention, KimiK3VisionBlock, KimiK3VisionEncoder, KimiK3VisionProjector, @@ -298,8 +297,8 @@ def _vision_encoder_config( block = KimiK3VisionBlock.Config( norm1=vision_norm, norm2=vision_norm, - attn=KimiK3VisionAttention.Config( - qkv_dim=qkv_dim, + attn=VisionAttention.Config( + dim=qkv_dim, num_heads=num_heads, wq=_linear(dim, qkv_dim), wk=_linear(dim, qkv_dim), diff --git a/torchtitan/models/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py index fc72342c52..8b866a67e9 100644 --- a/torchtitan/models/kimi_k3/model.py +++ b/torchtitan/models/kimi_k3/model.py @@ -38,7 +38,9 @@ # Shape suffixes: # B = batch, L = sequence length, D = model dimension, H = heads, # K = key head dimension, V = value head dimension, E = experts, -# T = flattened tokens, N = attention-residual entries. +# C = projection channels, F = expert hidden dimension, R = routed tokens, +# S = selected experts per token, T = flattened tokens, +# N = attention-residual entries. class KimiShortConvolution(ShortConvolution, Module): @@ -177,40 +179,40 @@ def forward( del positions B, L, _ = x_BLD.shape - q_BLNH = self.wq_b(self.q_norm(self.wq_a(x_BLD))).view( + q_BLHK = self.wq_b(self.q_norm(self.wq_a(x_BLD))).view( B, L, self.n_heads, self.q_head_dim ) - compressed_kv = self.wkv_a(x_BLD) - kv_latent, k_rope = torch.split( - compressed_kv, + compressed_kv_BLC = self.wkv_a(x_BLD) + kv_latent_BLC, k_rope_BLK = torch.split( + compressed_kv_BLC, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1, ) - kv_BLNH = self.wkv_b(self.kv_norm(kv_latent)).view( + kv_BLHC = self.wkv_b(self.kv_norm(kv_latent_BLC)).view( B, L, self.n_heads, self.qk_nope_head_dim + self.v_head_dim, ) - k_nope, v_BLNH = torch.split( - kv_BLNH, + k_nope_BLHK, v_BLHV = torch.split( + kv_BLHC, [self.qk_nope_head_dim, self.v_head_dim], dim=-1, ) - k_rope = k_rope.view(B, L, 1, self.qk_rope_head_dim).expand( + k_rope_BLHK = k_rope_BLK.view(B, L, 1, self.qk_rope_head_dim).expand( -1, -1, self.n_heads, -1 ) - k_BLNH = torch.cat((k_nope, k_rope), dim=-1) + k_BLHK = torch.cat((k_nope_BLHK, k_rope_BLHK), dim=-1) - out_BLNV = self.inner_attention( - q_BLNH, - k_BLNH, - v_BLNH, + out_BLHV = self.inner_attention( + q_BLHK, + k_BLHK, + v_BLHV, attention_masks=attention_masks, scale=self.scale, ) - out_BLD = out_BLNV.reshape(B, L, self.n_heads * self.v_head_dim) + out_BLD = out_BLHV.reshape(B, L, self.n_heads * self.v_head_dim) out_BLD = out_BLD * torch.sigmoid(self.gate(x_BLD)) return self.wo(out_BLD) @@ -433,8 +435,6 @@ def forward( class KimiLatentMoE(Module): - """Eager trainable implementation of Kimi K3 latent MoE.""" - @dataclass(kw_only=True, slots=True) class Config(Module.Config): num_experts: int @@ -477,9 +477,9 @@ def __init__(self, config: Config): ) def forward(self, x_BLD: torch.Tensor) -> torch.Tensor: - weights_BLK, expert_ids_BLK, scores_BLE = self.router(x_BLD, self.expert_bias_E) + weights_BLS, expert_ids_BLS, scores_BLE = self.router(x_BLD, self.expert_bias_E) routing_map_BLE = torch.zeros_like(scores_BLE, dtype=torch.bool).scatter_( - -1, expert_ids_BLK, True + -1, expert_ids_BLS, True ) num_tokens_per_expert_E = routing_map_BLE.sum(dim=(0, 1)) if self.training: @@ -491,8 +491,8 @@ def forward(self, x_BLD: torch.Tensor) -> torch.Tensor: latent_BLD = self.routed_down(x_BLD) routed_BLD = self.routed_experts( latent_BLD, - weights_BLK, - expert_ids_BLK, + weights_BLS, + expert_ids_BLS, num_tokens_per_expert_E, ) routed_BLD = self.routed_up(self.routed_norm(routed_BLD)) @@ -585,14 +585,7 @@ def forward( attention_masks: AttentionMasksType | None = None, positions: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: - # Blocks that do not extend the attention residual return it unchanged. - # FSDP2 aliases module inputs to drive its backward hooks, so passing it - # back out makes this FSDP unit return a view, which draws PyTorch's - # warning about in-place ops dropping the pre-backward hook. Nothing - # mutates it in place, and routing it back through the module boundary - # is what keeps FSDP gradients bitwise equal to eager -- returning None - # here instead reassociates the residual's gradient accumulation and - # perturbs tok_embeddings.weight.grad by ~2e-3 relative. + # Keep the residual on every block output to preserve its FSDP gradient path. B, L, D = x_BLD.shape prefix_sum_BLD: torch.Tensor | None = x_BLD diff --git a/torchtitan/models/kimi_k3/vision_encoder.py b/torchtitan/models/kimi_k3/vision_encoder.py index b2833b1d93..9a6960294e 100644 --- a/torchtitan/models/kimi_k3/vision_encoder.py +++ b/torchtitan/models/kimi_k3/vision_encoder.py @@ -6,27 +6,32 @@ """MoonViT3d vision encoder used by Kimi K3. -Vision attention is an eager PyTorch loop over each visual item, which -preserves the block-diagonal attention semantics of the HuggingFace -implementation without requiring FlashAttention or a device-specific kernel. - Shape suffixes: - N = number of visual items - P = maximum patches per item (padded) - D = vision hidden dimension - H = number of attention heads - K = attention head dimension +- C = number of complex-valued head-dimension pairs - M = maximum merged tokens per item (padded) +- F = merged feature dimension +- O = projected text dimension """ from dataclasses import dataclass, field import torch import torch.nn.functional as F +from torch.nn.attention.flex_attention import BlockMask from torchtitan.models.common import Linear from torchtitan.models.common.nn_modules import GELU, RMSNorm -from torchtitan.models.common.vision_encoder import VisionMLP +from torchtitan.models.common.vision_encoder import ( + compiled_create_block_mask, + get_vision_block_mask_mod, + VisionAttention, + VisionMLP, +) from torchtitan.protocols.module import Module, ModuleDict @@ -229,69 +234,6 @@ def forward(self, seqlen: int) -> torch.Tensor: return torch.outer(positions, self.inv_freq) -class KimiK3VisionAttention(Module): - """Eager, block-diagonal MoonViT attention reference.""" - - @dataclass(kw_only=True, slots=True) - class Config(Module.Config): - qkv_dim: int - num_heads: int - wq: Linear.Config - wk: Linear.Config - wv: Linear.Config - proj: Linear.Config - - def __init__(self, config: Config): - super().__init__() - if config.qkv_dim % config.num_heads != 0: - raise ValueError( - f"qkv_dim ({config.qkv_dim}) must be divisible by " - f"num_heads ({config.num_heads})." - ) - self.num_heads = config.num_heads - self.head_dim = config.qkv_dim // config.num_heads - self.scale = self.head_dim**-0.5 - self.wq = config.wq.build() - self.wk = config.wk.build() - self.wv = config.wv.build() - self.proj = config.proj.build() - - def forward( - self, - x_NPD: torch.Tensor, - *, - rope_cache: torch.Tensor, - num_patches: list[int], - ) -> torch.Tensor: - num_items, max_num_patches, _ = x_NPD.shape - q_NPHK = self.wq(x_NPD).view( - num_items, max_num_patches, self.num_heads, self.head_dim - ) - k_NPHK = self.wk(x_NPD).view( - num_items, max_num_patches, self.num_heads, self.head_dim - ) - v_NPHK = self.wv(x_NPD).view( - num_items, max_num_patches, self.num_heads, self.head_dim - ) - q_NPHK, k_NPHK = _apply_2d_rope(q_NPHK, k_NPHK, rope_cache) - - padded_outputs = [] - for item_idx, item_length in enumerate(num_patches): - q_HPK = q_NPHK[item_idx, :item_length].transpose(0, 1) - k_HPK = k_NPHK[item_idx, :item_length].transpose(0, 1) - v_HPK = v_NPHK[item_idx, :item_length].transpose(0, 1) - scores_HPP = torch.matmul(q_HPK, k_HPK.transpose(-2, -1)) - scores_HPP = scores_HPP * self.scale - probs_HPP = torch.softmax(scores_HPP, dim=-1, dtype=torch.float32).to( - q_HPK.dtype - ) - output_PHK = torch.matmul(probs_HPP, v_HPK).transpose(0, 1) - padded_outputs.append(_pad_sequence(output_PHK, max_num_patches)) - - output_NPHK = torch.stack(padded_outputs) - return self.proj(output_NPHK.flatten(start_dim=-2)) - - class KimiK3VisionBlock(Module): """MoonViT pre-norm attention and MLP block.""" @@ -299,7 +241,7 @@ class KimiK3VisionBlock(Module): class Config(Module.Config): norm1: RMSNorm.Config norm2: RMSNorm.Config - attn: KimiK3VisionAttention.Config + attn: VisionAttention.Config mlp: VisionMLP.Config def __init__(self, config: Config): @@ -314,12 +256,13 @@ def forward( x_NPD: torch.Tensor, *, rope_cache: torch.Tensor, - num_patches: list[int], + attention_mask: BlockMask, ) -> torch.Tensor: x_NPD = x_NPD + self.attn( self.norm1(x_NPD), rope_cache=rope_cache, - num_patches=num_patches, + rope_apply=_apply_2d_rope, + attention_mask=attention_mask, ) return x_NPD + self.mlp(self.norm2(x_NPD)) @@ -343,19 +286,17 @@ def __init__(self, config: Config): self.post_norm = config.post_norm.build() self.activation = config.activation.build() - def forward(self, merged_NMK: torch.Tensor) -> torch.Tensor: - if merged_NMK.shape[-1] != self.merged_dim: + def forward(self, merged_NMF: torch.Tensor) -> torch.Tensor: + if merged_NMF.shape[-1] != self.merged_dim: raise ValueError( f"Expected merged vision dim {self.merged_dim}, got " - f"{merged_NMK.shape[-1]}." + f"{merged_NMF.shape[-1]}." ) - projected = self.linear_2(self.activation(self.linear_1(merged_NMK))) - return self.post_norm(projected) + projected_NMO = self.linear_2(self.activation(self.linear_1(merged_NMF))) + return self.post_norm(projected_NMO) class KimiK3VisionEncoder(Module): - """Device-neutral MoonViT3d encoder and PatchMergerMLPV2 projector.""" - @dataclass(kw_only=True, slots=True) class Config(Module.Config): dim: int @@ -442,7 +383,7 @@ def forward( ) kernel_h, kernel_w = self.merge_kernel_size - num_patches = [] + num_patches_N = grid_thw.prod(dim=-1).to(torch.long) for num_frames, grid_h, grid_w in grids: if grid_h % kernel_h != 0 or grid_w % kernel_w != 0: raise ValueError( @@ -455,18 +396,27 @@ def forward( f"Vision grid requires {item_num_patches} patches, but " f"pixel_values only provides {max_num_patches}." ) - num_patches.append(item_num_patches) learned_pos, rope_cache = self._compute_position_embeddings( grids, max_num_patches ) hidden_NPD = self.patch_embed(pixel_values) + learned_pos + + mask_mod = get_vision_block_mask_mod(num_patches_N) + attention_mask = compiled_create_block_mask( + mask_mod, + num_items, + None, + max_num_patches, + max_num_patches, + device=hidden_NPD.device, + ) for block in self.layers.values(): hidden_NPD = block( hidden_NPD, rope_cache=rope_cache, - num_patches=num_patches, + attention_mask=attention_mask, ) hidden_NPD = self.final_norm(hidden_NPD) - merged_NMK = _temporal_pool_and_merge(hidden_NPD, grids, self.merge_kernel_size) - return self.projector(merged_NMK) + merged_NMF = _temporal_pool_and_merge(hidden_NPD, grids, self.merge_kernel_size) + return self.projector(merged_NMF) From af8b872aba524072d495d5627be43907f83a377b Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Mon, 17 Aug 2026 09:08:12 +0000 Subject: [PATCH 37/67] simplify test_kimi_k3.py --- tests/unit_tests/test_kimi_k3.py | 23 +++++++---------------- 1 file changed, 7 insertions(+), 16 deletions(-) diff --git a/tests/unit_tests/test_kimi_k3.py b/tests/unit_tests/test_kimi_k3.py index 54924591a6..e0fbf39be2 100644 --- a/tests/unit_tests/test_kimi_k3.py +++ b/tests/unit_tests/test_kimi_k3.py @@ -21,22 +21,13 @@ from torchtitan.config import CompileConfig, ParallelismConfig, TrainingConfig from torchtitan.distributed import ParallelDims -# torchtitan.models.kimi_k3 imports FLA at module scope for the KDA kernel. -# FLA is a per-model dependency (kimi_k3/requirements.txt), not part of the -# core requirements, so skip the Kimi suites instead of failing collection -# when it is absent. Modules importing this one inherit the skip. -try: - from torchtitan.models.kimi_k3 import ( - _kimi_k3_config, - _vision_encoder_config, - parallelize_kimi_k3, - ) - from torchtitan.models.kimi_k3.model import KimiK3Model, KimiKDAKernel - from torchtitan.models.kimi_k3.state_dict_adapter import KimiK3StateDictAdapter -except ModuleNotFoundError as exc: - raise unittest.SkipTest( - f"Kimi K3 optional dependency unavailable: {exc.name}" - ) from exc +from torchtitan.models.kimi_k3 import ( + _kimi_k3_config, + _vision_encoder_config, + parallelize_kimi_k3, +) +from torchtitan.models.kimi_k3.model import KimiK3Model, KimiKDAKernel +from torchtitan.models.kimi_k3.state_dict_adapter import KimiK3StateDictAdapter def _small_model_config( From f8523a20aa7d9cd158207ec744d106653bc5cfcc Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Tue, 18 Aug 2026 18:45:13 +0000 Subject: [PATCH 38/67] reuse common MoonViT helpers, VisionTransformerBlock, and MoE The MoonViT temporal position embedding, 2D RoPE frequency table, and 2D RoPE cache builder were duplicated between kimi_k2_7 and kimi_k3; move them to models/common/vision_encoder.py and use them from both. Vision encoder output is bit-identical before and after for both models. Kimi K3's vision block was a copy of VisionTransformerBlock differing only in norm type, so widen the shared block's norm config to accept RMSNorm and drop the copy along with the hand-written 2D RoPE apply (ComplexRoPE.apply_rotary_emb computes the same thing). KimiLatentMoE was a copy of common MoE plus the latent down/up projections; inherit MoE instead so the router, dispatcher, and load-balancing buffers stay shared. Also factor the duplicated SiTU-GLU activation into one helper. --- torchtitan/models/common/vision_encoder.py | 152 ++++++++++++++++++- torchtitan/models/kimi_k3/__init__.py | 16 +- torchtitan/models/kimi_k3/model.py | 128 ++++++---------- torchtitan/models/kimi_k3/vision_encoder.py | 159 ++------------------ 4 files changed, 214 insertions(+), 241 deletions(-) diff --git a/torchtitan/models/common/vision_encoder.py b/torchtitan/models/common/vision_encoder.py index 4576a4bfd9..832fa63422 100644 --- a/torchtitan/models/common/vision_encoder.py +++ b/torchtitan/models/common/vision_encoder.py @@ -22,13 +22,17 @@ from collections.abc import Callable from dataclasses import dataclass, field +from typing import cast +import spmd_types as spmd import torch from torch.nn.attention.flex_attention import BlockMask, create_block_mask +from torchtitan.distributed.utils import get_spmd_backend from torchtitan.models.common import Linear from torchtitan.models.common.attention import FlexAttention, local_head_split -from torchtitan.models.common.nn_modules import GELU, LayerNorm +from torchtitan.models.common.nn_modules import GELU, LayerNorm, RMSNorm +from torchtitan.models.common.rope import _maybe_wrap_positions from torchtitan.protocols.module import Module compiled_create_block_mask = torch.compile(create_block_mask) @@ -66,6 +70,148 @@ def mask_mod(b, h, q_idx, kv_idx): ) +def get_temporal_pos_embed( + num_frames: int, + embed_dim: int, + *, + base: float = 10000.0, + device: torch.device | None = None, +) -> torch.Tensor: + """Fixed 1D sinusoidal embeddings for the temporal axis (video frames). + + Returns ``(num_frames, embed_dim)`` float32; the standard 1D sincos formula + over frame indices. + + Args: + num_frames: Number of video frames (temporal positions). + embed_dim: Embedding width per frame. + base: Sinusoid base (longest wavelength); the conventional PE constant. + device: Device for the returned tensor. + """ + grid = torch.arange(num_frames, dtype=torch.float32, device=device) + omega = torch.arange(embed_dim // 2, dtype=torch.float32, device=device) / ( + embed_dim / 2.0 + ) + omega = 1.0 / base**omega + out = torch.outer(grid, omega) + return torch.cat([out.sin(), out.cos()], dim=1) + + +def compute_2d_rope_cache( + freq_table: torch.Tensor, + grids: list[list[int]], + max_num_patch: int, + head_dim: int, +) -> torch.Tensor: + """Compute the padded 2D-RoPE complex ``freqs_cis`` cache in raster order. + + For head-dim pair index ``k`` (``k`` in ``[0, head_dim/4)``), even output + pairs are rotated by the *column* (x) position and odd pairs by the *row* + (y) position. The per-axis angle for a position ``p`` is ``p * inv_freq[k]``; + this looks it up by gathering row ``p`` of ``freq_table`` (built once by + ``VisionRotaryEmbedding2D`` and cached by the encoder) rather than + recomputing ``p * inv_freq`` each call. Frames repeat the spatial pattern. + + Returns a complex cache consumed by ``ComplexRoPE.apply_rotary_emb``; only + the cache is 2D/per-grid, which is why it is built here rather than by the + 1D ``ComplexRoPE`` cache machinery. + + Args: + freq_table: ``(max_hw, head_dim/4)`` position-to-frequency table, where + ``freq_table[p, k] = p * inv_freq[k]``. + grids: per-item ``[t, h, w]`` patch counts as host ints (``grid_thw`` + read to CPU once by the caller, so the per-item loop adds no syncs). + max_num_patch: Padded sequence length. + head_dim: Attention head dim (must be divisible by 4). + + Returns: + ``(N, max_num_patch, 1, head_dim/2)`` complex64 (head axis = 1 to + broadcast over the heads). + """ + device = freq_table.device + + angles = freq_table.new_zeros(len(grids), max_num_patch, head_dim // 2) + if get_spmd_backend() == "spmd_types" and spmd.is_type_checking(): + angles = spmd.mutate_type(angles, src=spmd.R, dst={"dp": spmd.V, "tp": spmd.I}) + + # Group by (h, w) so the per-resolution angle grid is built once. + hw_to_indices: dict[tuple[int, int], list[int]] = {} + for i, (_, h, w) in enumerate(grids): + hw_to_indices.setdefault((h, w), []).append(i) + + for (h, w), indices in hw_to_indices.items(): + # Raster order: position p -> (row = p // w, col = p % w). Gather each + # axis's angles from the precomputed table (freq_table[pos] = pos*inv_freq). + flat = torch.arange(h * w, device=device) + flat = cast(torch.Tensor, _maybe_wrap_positions(flat, freq_table)) + if get_spmd_backend() == "spmd_types" and spmd.is_type_checking(): + flat = spmd.mutate_type(flat, "tp", src=spmd.R, dst=spmd.I) + x_ang = freq_table[flat % w] # (h*w, head_dim/4) column + y_ang = freq_table[flat // w] # (h*w, head_dim/4) row + # Interleave x/y so pair 2k uses x-position, pair 2k+1 uses y-position. + ang = torch.stack([x_ang, y_ang], dim=-1).reshape(h * w, head_dim // 2) + for i in indices: + t = grids[i][0] + seq_len = t * h * w + angles[i, :seq_len] = ang.repeat(t, 1) + + # Complex unit-modulus cache; unsqueeze the head axis for broadcast. + return torch.polar(torch.ones_like(angles), angles).unsqueeze(2) + + +class VisionRotaryEmbedding2D(Module): + """2D rotary position embedding for the vision tower. + + Holds the per-axis frequencies ``inv_freq`` (``head_dim/4`` of them, shared + by the row and column axes). ``forward(seqlen)`` returns the + position-to-frequency table ``freq_table[p, k] = p * inv_freq[k]`` for + positions up to ``seqlen``; ``compute_2d_rope_cache`` gathers per-patch + row/col angles from it, and ``ComplexRoPE.apply_rotary_emb`` applies them. + ``head_dim`` must be divisible by 4. + """ + + @dataclass(kw_only=True, slots=True) + class Config(Module.Config): + head_dim: int + theta: float = 10000.0 + + def __init__(self, config: Config): + super().__init__() + if config.head_dim % 4 != 0: + raise ValueError( + f"2D RoPE requires head_dim divisible by 4, got {config.head_dim}." + ) + self.head_dim = config.head_dim + self.theta = config.theta + self.register_buffer("inv_freq", self._compute_inv_freq(), persistent=False) + + def _compute_inv_freq(self, *, device: torch.device | None = None) -> torch.Tensor: + # inv_freq[k] = theta**(-4k/head_dim) for k in [0, head_dim/4); the + # step of 4 leaves room for the row/col split of the 2D rotation. + return 1.0 / ( + self.theta + ** ( + torch.arange(0, self.head_dim, 4, dtype=torch.float32, device=device) + / self.head_dim + ) + ) + + def _init_self_buffers(self, *, buffer_device: torch.device | None = None) -> None: + """Re-compute inv_freq on the target device after to_empty().""" + device = buffer_device or self.inv_freq.device + self.inv_freq = self._compute_inv_freq(device=device) + + def forward(self, seqlen: int) -> torch.Tensor: + """Frequency table ``(seqlen, head_dim/4)`` for positions ``[0, seqlen)``.""" + seq = torch.arange( + seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype + ) + seq = cast(torch.Tensor, _maybe_wrap_positions(seq, self.inv_freq)) + if get_spmd_backend() == "spmd_types" and spmd.is_type_checking(): + seq = spmd.mutate_type(seq, "tp", src=spmd.R, dst=spmd.I) + return torch.outer(seq, self.inv_freq) + + class VisionMLP(Module): """Feed-forward network with GELU activation (fc1 -> act -> fc2).""" @@ -150,8 +296,8 @@ class VisionTransformerBlock(Module): @dataclass(kw_only=True, slots=True) class Config(Module.Config): - norm1: LayerNorm.Config - norm2: LayerNorm.Config + norm1: LayerNorm.Config | RMSNorm.Config + norm2: LayerNorm.Config | RMSNorm.Config attn: VisionAttention.Config mlp: VisionMLP.Config diff --git a/torchtitan/models/kimi_k3/__init__.py b/torchtitan/models/kimi_k3/__init__.py index 63e886e587..abb5bc93f5 100644 --- a/torchtitan/models/kimi_k3/__init__.py +++ b/torchtitan/models/kimi_k3/__init__.py @@ -16,7 +16,12 @@ from torchtitan.models.common.moe import RoutedExperts, TokenChoiceTopKRouter from torchtitan.models.common.nn_modules import GELU, RMSNorm from torchtitan.models.common.token_dispatcher import LocalTokenDispatcher -from torchtitan.models.common.vision_encoder import VisionAttention, VisionMLP +from torchtitan.models.common.vision_encoder import ( + VisionAttention, + VisionMLP, + VisionRotaryEmbedding2D, + VisionTransformerBlock, +) from torchtitan.models.utils import validate_converter_order from torchtitan.protocols.model import ModelConfigConverter from torchtitan.protocols.model_spec import ModelSpec @@ -34,12 +39,7 @@ ) from .parallelize import parallelize_kimi_k3 from .state_dict_adapter import KimiK3StateDictAdapter -from .vision_encoder import ( - KimiK3VisionBlock, - KimiK3VisionEncoder, - KimiK3VisionProjector, - VisionRotaryEmbedding2D, -) +from .vision_encoder import KimiK3VisionEncoder, KimiK3VisionProjector __all__ = [ "KIMI_K3_SPECIAL_TOKENS", @@ -294,7 +294,7 @@ def _vision_encoder_config( eps=1e-5, param_init=_NORM_INIT, ) - block = KimiK3VisionBlock.Config( + block = VisionTransformerBlock.Config( norm1=vision_norm, norm2=vision_norm, attn=VisionAttention.Config( diff --git a/torchtitan/models/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py index 8b866a67e9..d2a1e47f92 100644 --- a/torchtitan/models/kimi_k3/model.py +++ b/torchtitan/models/kimi_k3/model.py @@ -21,11 +21,7 @@ ) from torchtitan.models.common.decoder import Decoder from torchtitan.models.common.feed_forward import FeedForward -from torchtitan.models.common.moe import ( - GroupedExperts, - RoutedExperts, - TokenChoiceTopKRouter, -) +from torchtitan.models.common.moe import GroupedExperts, MoE from torchtitan.models.common.multimodal import ( get_vision_positions, scatter_vision_embeds, @@ -99,6 +95,22 @@ def forward(self, x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: return (x_float * torch.sigmoid(gate.float())).to(input_dtype) +def _situ_glu( + gate: torch.Tensor, + up: torch.Tensor, + beta: float, + linear_beta: float | None, +) -> torch.Tensor: + """Kimi's SiTU-GLU activation, evaluated in FP32.""" + input_dtype = gate.dtype + gate = gate.float() + up = up.float() + gate = beta * torch.tanh(gate / beta) * torch.sigmoid(gate) + if linear_beta is not None: + up = linear_beta * torch.tanh(up / linear_beta) + return (gate * up).to(input_dtype) + + class KimiFeedForward(FeedForward): """FeedForward with Kimi's SiTU activation""" @@ -113,15 +125,9 @@ def __init__(self, config: Config): self.linear_beta = config.linear_beta def forward(self, x: torch.Tensor) -> torch.Tensor: - gate = self.w1(x) - up = self.w3(x) - input_dtype = gate.dtype - gate = gate.float() - up = up.float() - gate = self.beta * torch.tanh(gate / self.beta) * torch.sigmoid(gate) - if self.linear_beta is not None: - up = self.linear_beta * torch.tanh(up / self.linear_beta) - return self.w2((gate * up).to(input_dtype)) + return self.w2( + _situ_glu(self.w1(x), self.w3(x), self.beta, self.linear_beta), + ) class KimiMLAAttention(BaseAttention): @@ -250,13 +256,6 @@ def forward( A_log_H: torch.Tensor, dt_bias_HK: torch.Tensor, ) -> torch.Tensor: - if q_BLHK.shape != k_BLHK.shape: - raise ValueError( - f"KDA q/k shapes must match, got {q_BLHK.shape} and {k_BLHK.shape}." - ) - if q_BLHK.shape[:3] != v_BLHV.shape[:3]: - raise ValueError("Kimi KDA requires equal q/k/value head counts.") - # safe_gate selects the bounded gate activation # lower_bound * sigmoid(exp(A_log) * (gate + dt_bias)); without it the # kernel applies -exp(A_log) * softplus(gate + dt_bias). @@ -419,13 +418,7 @@ def forward( offs=offsets_E, ) - input_dtype = gate_RF.dtype - gate_RF = gate_RF.float() - up_RF = up_RF.float() - gate_RF = self.beta * torch.tanh(gate_RF / self.beta) * torch.sigmoid(gate_RF) - if self.linear_beta is not None: - up_RF = self.linear_beta * torch.tanh(up_RF / self.linear_beta) - h_RF = (gate_RF * up_RF).to(input_dtype) + h_RF = _situ_glu(gate_RF, up_RF, self.beta, self.linear_beta) return self._grouped_mm( A=h_RF, @@ -434,47 +427,25 @@ def forward( ).type_as(x_RD) -class KimiLatentMoE(Module): +class KimiLatentMoE(MoE): + """``common/moe.py::MoE`` with Kimi's latent routed-expert path. + + Routed tokens are projected down to the expert latent width, run through + the experts, then normed and projected back up to the model dimension. + Shared experts still see the full-width input. + """ + @dataclass(kw_only=True, slots=True) - class Config(Module.Config): - num_experts: int - router: TokenChoiceTopKRouter.Config + class Config(MoE.Config): routed_down: Linear.Config - routed_experts: RoutedExperts.Config routed_norm: RMSNorm.Config routed_up: Linear.Config - shared_experts: KimiFeedForward.Config - load_balance_coeff: float | None = 1e-3 def __init__(self, config: Config): - super().__init__() - if config.routed_experts.inner_experts.num_experts != config.num_experts: - raise ValueError( - "routed_experts.inner_experts.num_experts must equal num_experts." - ) - self.num_experts = config.num_experts - self.router = config.router.build() + super().__init__(config) self.routed_down = config.routed_down.build() - self.routed_experts = config.routed_experts.build() self.routed_norm = config.routed_norm.build() self.routed_up = config.routed_up.build() - self.shared_experts = config.shared_experts.build() - self.load_balance_coeff = config.load_balance_coeff - if self.load_balance_coeff is not None: - if self.load_balance_coeff <= 0.0: - raise ValueError("load_balance_coeff must be positive.") - self.register_buffer( - "expert_bias_E", - torch.zeros(config.num_experts, dtype=torch.float32), - persistent=True, - ) - else: - self.expert_bias_E = None - self.register_buffer( - "tokens_per_expert_E", - torch.zeros(config.num_experts, dtype=torch.float32), - persistent=False, - ) def forward(self, x_BLD: torch.Tensor) -> torch.Tensor: weights_BLS, expert_ids_BLS, scores_BLE = self.router(x_BLD, self.expert_bias_E) @@ -482,32 +453,19 @@ def forward(self, x_BLD: torch.Tensor) -> torch.Tensor: -1, expert_ids_BLS, True ) num_tokens_per_expert_E = routing_map_BLE.sum(dim=(0, 1)) - if self.training: - with torch.no_grad(): - # In place so the load-balancing hook registered on the optimizer - # keeps referring to this buffer across steps. - self.tokens_per_expert_E.add_(num_tokens_per_expert_E.float()) + with torch.no_grad(): + self.tokens_per_expert_E.add_(num_tokens_per_expert_E) - latent_BLD = self.routed_down(x_BLD) routed_BLD = self.routed_experts( - latent_BLD, + self.routed_down(x_BLD), weights_BLS, expert_ids_BLS, num_tokens_per_expert_E, ) - routed_BLD = self.routed_up(self.routed_norm(routed_BLD)) - return routed_BLD + self.shared_experts(x_BLD) - - def _init_self_buffers(self, *, buffer_device: torch.device | None = None) -> None: - if buffer_device is None: - buffer_device = self.tokens_per_expert_E.device - self.tokens_per_expert_E = torch.zeros( - self.num_experts, dtype=torch.float32, device=buffer_device - ) - if self.load_balance_coeff is not None: - self.expert_bias_E = torch.zeros( - self.num_experts, dtype=torch.float32, device=buffer_device - ) + out_BLD = self.routed_up(self.routed_norm(routed_BLD)) + if self.shared_experts is not None: + out_BLD = out_BLD + self.shared_experts(x_BLD) + return out_BLD def _apply_attention_residual( @@ -516,7 +474,12 @@ def _apply_attention_residual( projection: Linear, norm: RMSNorm, ) -> torch.Tensor: - """Apply Kimi's block-level attention residual in FP32.""" + """Apply Kimi's block-level attention residual in FP32. + + The norm and projection weights are folded into a single score vector, so + this reads them directly instead of calling the modules. That is only valid + while both are replicated; sharding them would need a DTensor-aware path. + """ assert norm.eps is not None values_TND = torch.cat((block_residual_TND, prefix_sum_TD.unsqueeze(1)), dim=1) @@ -671,6 +634,9 @@ def get_nparams_and_flops( raise ValueError( "Kimi K3 requires at least one MLA layer for FLOP accounting." ) + # KDA and the vision encoder have no dedicated term here, so their + # parameters only contribute the dense 6*N estimate; reported MFU is + # approximate. return get_moe_model_nparams_and_flops( self, model, diff --git a/torchtitan/models/kimi_k3/vision_encoder.py b/torchtitan/models/kimi_k3/vision_encoder.py index 9a6960294e..3f07042a4b 100644 --- a/torchtitan/models/kimi_k3/vision_encoder.py +++ b/torchtitan/models/kimi_k3/vision_encoder.py @@ -4,15 +4,12 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -"""MoonViT3d vision encoder used by Kimi K3. +"""MoonViT-V2 vision encoder used by Kimi K3. Shape suffixes: - N = number of visual items - P = maximum patches per item (padded) - D = vision hidden dimension -- H = number of attention heads -- K = attention head dimension -- C = number of complex-valued head-dimension pairs - M = maximum merged tokens per item (padded) - F = merged feature dimension - O = projected text dimension @@ -22,35 +19,21 @@ import torch import torch.nn.functional as F -from torch.nn.attention.flex_attention import BlockMask from torchtitan.models.common import Linear from torchtitan.models.common.nn_modules import GELU, RMSNorm +from torchtitan.models.common.rope import ComplexRoPE from torchtitan.models.common.vision_encoder import ( compiled_create_block_mask, + compute_2d_rope_cache, + get_temporal_pos_embed, get_vision_block_mask_mod, - VisionAttention, - VisionMLP, + VisionRotaryEmbedding2D, + VisionTransformerBlock, ) from torchtitan.protocols.module import Module, ModuleDict -def _get_temporal_pos_embed( - num_frames: int, - embed_dim: int, - *, - device: torch.device, -) -> torch.Tensor: - """Return fixed 1D sinusoidal embeddings for video frame positions.""" - grid = torch.arange(num_frames, dtype=torch.float32, device=device) - omega = torch.arange(embed_dim // 2, dtype=torch.float32, device=device) / ( - embed_dim / 2.0 - ) - omega = 1.0 / 10000.0**omega - angles = torch.outer(grid, omega) - return torch.cat((angles.sin(), angles.cos()), dim=-1) - - def _pad_sequence(x: torch.Tensor, target_length: int) -> torch.Tensor: """Pad the leading sequence dimension without modifying ``x`` in place.""" padding_length = target_length - x.shape[0] @@ -104,7 +87,7 @@ def _compute_learned_pos_embeds( if num_frames == 1: item_pos = spatial else: - temporal = _get_temporal_pos_embed(num_frames, dim, device=pos_embed.device) + temporal = get_temporal_pos_embed(num_frames, dim, device=pos_embed.device) item_pos = spatial.unsqueeze(0) + temporal.unsqueeze(1).to(spatial.dtype) item_pos = item_pos.reshape(num_frames * grid_h * grid_w, dim) padded_positions.append(_pad_sequence(item_pos, max_num_patches)) @@ -112,51 +95,6 @@ def _compute_learned_pos_embeds( return torch.stack(padded_positions) -def _compute_2d_rope_cache( - freq_table: torch.Tensor, - grids: list[list[int]], - max_num_patches: int, - head_dim: int, -) -> torch.Tensor: - """Build the real-valued 2D RoPE cache in raster patch order.""" - cached_spatial: dict[tuple[int, int], torch.Tensor] = {} - padded_angles = [] - for num_frames, grid_h, grid_w in grids: - spatial = cached_spatial.get((grid_h, grid_w)) - if spatial is None: - flat = torch.arange(grid_h * grid_w, device=freq_table.device) - x_angles = freq_table[flat % grid_w] - y_angles = freq_table[flat // grid_w] - spatial = torch.stack((x_angles, y_angles), dim=-1).reshape( - grid_h * grid_w, head_dim // 2 - ) - cached_spatial[(grid_h, grid_w)] = spatial - item_angles = spatial.repeat(num_frames, 1) - padded_angles.append(_pad_sequence(item_angles, max_num_patches)) - - angles = torch.stack(padded_angles) - return torch.stack((angles.cos(), angles.sin()), dim=-1).unsqueeze(2) - - -def _apply_2d_rope( - q_NPHK: torch.Tensor, - k_NPHK: torch.Tensor, - rope_cache_NP1C2: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor]: - """Apply 2D RoPE using the real form of complex multiplication.""" - - cos_NP1C = rope_cache_NP1C2[..., 0] - sin_NP1C = rope_cache_NP1C2[..., 1] - - def rotate(x_NPHK: torch.Tensor) -> torch.Tensor: - x_NPHC2 = x_NPHK.float().reshape(*x_NPHK.shape[:-1], -1, 2) - real_NPHC = x_NPHC2[..., 0] * cos_NP1C - x_NPHC2[..., 1] * sin_NP1C - imag_NPHC = x_NPHC2[..., 0] * sin_NP1C + x_NPHC2[..., 1] * cos_NP1C - return torch.stack((real_NPHC, imag_NPHC), dim=-1).flatten(-2) - - return rotate(q_NPHK).to(q_NPHK.dtype), rotate(k_NPHK).to(k_NPHK.dtype) - - def _temporal_pool_and_merge( hidden_NPD: torch.Tensor, grids: list[list[int]], @@ -189,84 +127,6 @@ def _temporal_pool_and_merge( return torch.stack(padded_items) -class VisionRotaryEmbedding2D(Module): - """Per-axis frequency table for MoonViT's interleaved 2D RoPE.""" - - @dataclass(kw_only=True, slots=True) - class Config(Module.Config): - head_dim: int - theta: float = 10000.0 - - def __init__(self, config: Config): - super().__init__() - if config.head_dim % 4 != 0: - raise ValueError( - "Vision 2D RoPE head_dim must be divisible by 4, " - f"got {config.head_dim}." - ) - self.head_dim = config.head_dim - self.theta = config.theta - self.register_buffer("inv_freq", self._compute_inv_freq(), persistent=False) - - def _compute_inv_freq(self, *, device: torch.device | None = None) -> torch.Tensor: - return 1.0 / ( - self.theta - ** ( - torch.arange( - 0, - self.head_dim, - 4, - dtype=torch.float32, - device=device, - ) - / self.head_dim - ) - ) - - def _init_self_buffers(self, *, buffer_device: torch.device | None = None) -> None: - device = buffer_device or self.inv_freq.device - self.inv_freq = self._compute_inv_freq(device=device) - - def forward(self, seqlen: int) -> torch.Tensor: - positions = torch.arange( - seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype - ) - return torch.outer(positions, self.inv_freq) - - -class KimiK3VisionBlock(Module): - """MoonViT pre-norm attention and MLP block.""" - - @dataclass(kw_only=True, slots=True) - class Config(Module.Config): - norm1: RMSNorm.Config - norm2: RMSNorm.Config - attn: VisionAttention.Config - mlp: VisionMLP.Config - - def __init__(self, config: Config): - super().__init__() - self.norm1 = config.norm1.build() - self.norm2 = config.norm2.build() - self.attn = config.attn.build() - self.mlp = config.mlp.build() - - def forward( - self, - x_NPD: torch.Tensor, - *, - rope_cache: torch.Tensor, - attention_mask: BlockMask, - ) -> torch.Tensor: - x_NPD = x_NPD + self.attn( - self.norm1(x_NPD), - rope_cache=rope_cache, - rope_apply=_apply_2d_rope, - attention_mask=attention_mask, - ) - return x_NPD + self.mlp(self.norm2(x_NPD)) - - class KimiK3VisionProjector(Module): """PatchMergerMLPV2 projector from merged vision features to text width.""" @@ -310,7 +170,7 @@ class Config(Module.Config): interpolation_mode: str patch_embed_proj: Linear.Config rotary_pos_emb: VisionRotaryEmbedding2D.Config - block: KimiK3VisionBlock.Config + block: VisionTransformerBlock.Config final_norm: RMSNorm.Config projector: KimiK3VisionProjector.Config @@ -357,7 +217,7 @@ def _compute_position_embeddings( self.interpolation_mode, self.max_num_frames, ) - rope_cache = _compute_2d_rope_cache( + rope_cache = compute_2d_rope_cache( self._cached_freq_table, grids, max_num_patches, @@ -415,6 +275,7 @@ def forward( hidden_NPD = block( hidden_NPD, rope_cache=rope_cache, + rope_apply=ComplexRoPE.apply_rotary_emb, attention_mask=attention_mask, ) hidden_NPD = self.final_norm(hidden_NPD) From 71c9804353592d54dd618c4efa33f63befa89637 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Tue, 18 Aug 2026 18:45:43 +0000 Subject: [PATCH 39/67] address review feedback on the debug model, tests, and README - Size debugmodel's vocab for tests/assets/tokenizer (2048) instead of the released 163840, which put 84% of the parameters in the embedding and lm_head; the numerical parity script restores the released vocab since it feeds real tokenizer ids. - Cover a KDA layer in the FSDP parity test and turn the mask-only test into a multimodal forward so the CPU suite exercises the model. - Restore the FLA optional-dependency skip so collection does not fail without flash-linear-attention, matching test_qwen3_5_deltanet.py. - Fix the LRSchedulersContainer import moved by upstream, pass attn_backend explicitly, and point the AC error at the flag that disables it. - Note the CP route in the vision-encoder FSDP TODO, that the attention residual reads norm/projection weights directly, and that KDA and the vision encoder have no dedicated FLOP term. - Align the README parity numbers with the reported float32 run and note that routed experts run in bf16. --- .../numerical_tests_kimi_k3.py | 6 +++ tests/unit_tests/test_kimi_k3.py | 48 +++++++++++++------ torchtitan/models/kimi_k3/README.md | 28 ++++++----- torchtitan/models/kimi_k3/__init__.py | 4 +- torchtitan/models/kimi_k3/config_registry.py | 5 +- torchtitan/models/kimi_k3/parallelize.py | 8 +++- 6 files changed, 67 insertions(+), 32 deletions(-) diff --git a/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py b/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py index e089617cbc..32ed4b465a 100644 --- a/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py +++ b/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py @@ -47,6 +47,7 @@ _HF_REPO_ID = "moonshotai/Kimi-K3" _HF_REVISION = "9f62e4e9fffbd0a83ddd60e1c209d828994b3569" _MEDIA_TOKEN_ID = 163605 +_VOCAB_SIZE = 163840 _PATCH_SIZE = 14 _MERGE_SIZE = 2 _MAX_PATCHES = 65536 @@ -458,6 +459,11 @@ def main() -> None: ) tt_config = cast(KimiK3Model.Config, model_registry(args.model_flavor).model) + # The released tokenizer emits ids across the full vocab, while debugmodel + # is sized for the test tokenizer. + tt_config.vocab_size = _VOCAB_SIZE + tt_config.tok_embeddings.num_embeddings = _VOCAB_SIZE + tt_config.lm_head.out_features = _VOCAB_SIZE torch.manual_seed(args.seed) tt_model = _build_tt_model(tt_config, dtype) hf_state_dict = KimiK3StateDictAdapter(tt_config, hf_assets_path=None).to_hf( diff --git a/tests/unit_tests/test_kimi_k3.py b/tests/unit_tests/test_kimi_k3.py index e0fbf39be2..425a63f8d9 100644 --- a/tests/unit_tests/test_kimi_k3.py +++ b/tests/unit_tests/test_kimi_k3.py @@ -21,13 +21,20 @@ from torchtitan.config import CompileConfig, ParallelismConfig, TrainingConfig from torchtitan.distributed import ParallelDims -from torchtitan.models.kimi_k3 import ( - _kimi_k3_config, - _vision_encoder_config, - parallelize_kimi_k3, -) -from torchtitan.models.kimi_k3.model import KimiK3Model, KimiKDAKernel -from torchtitan.models.kimi_k3.state_dict_adapter import KimiK3StateDictAdapter +# FLA is a per-model dependency (kimi_k3/requirements.txt) imported at module +# scope for the KDA kernel, so skip rather than fail collection without it. +try: + from torchtitan.models.kimi_k3 import ( + _kimi_k3_config, + _vision_encoder_config, + parallelize_kimi_k3, + ) + from torchtitan.models.kimi_k3.model import KimiK3Model, KimiKDAKernel + from torchtitan.models.kimi_k3.state_dict_adapter import KimiK3StateDictAdapter +except ModuleNotFoundError as exc: + raise unittest.SkipTest( + f"Kimi K3 optional dependency unavailable: {exc.name}" + ) from exc def _small_model_config( @@ -52,7 +59,7 @@ def _small_model_config( qk_nope_head_dim=16, qk_rope_head_dim=16, v_head_dim=16, - kda_head_dim=16, + kda_head_dim=32, conv_kernel_size=3, dense_hidden_dim=128, latent_dim=32, @@ -139,12 +146,25 @@ def _kda_recurrent_reference( class TestKimiK3(unittest.TestCase): - def test_flex_attention_mask(self): - config = _small_model_config() + def test_multimodal_forward(self): + # All-MLA so the forward runs without the CUDA-only KDA kernel; the + # KDA path is covered by the FSDP parity test below. + config = _small_model_config(full_attention_layers={0, 1}) model = config.build() - positions = torch.arange(4, dtype=torch.int32).unsqueeze(0) + model.init_states() + positions = torch.arange(6, dtype=torch.int32).unsqueeze(0) attention_masks = model.get_attention_masks(positions) self.assertIsInstance(attention_masks, BlockMask) + with torch.no_grad(): + logits = model( + torch.tensor([[1, 7, 2, 3, 4, 5]]), + pixel_values=torch.randn(1, 4, 3 * 2 * 2), + grid_thw=torch.tensor([[1, 2, 2]]), + special_tokens={"image_id": 7}, + positions=positions, + attention_masks=attention_masks, + ) + self.assertEqual(logits.shape, (1, 6, config.vocab_size)) def test_update_from_config_propagates_moe_force_load_balance(self): from torchtitan.config import DebugConfig @@ -265,10 +285,8 @@ def world_size(self): @with_comms def test_fsdp_matches_non_distributed_forward_backward(self): torch.manual_seed(3) - config = _small_model_config( - attn_res_block_size=2, - full_attention_layers={0, 1}, - ) + # Layer 0 is KDA, layer 1 is MLA, so one run covers both attentions. + config = _small_model_config(attn_res_block_size=2) with torch.device("meta"): model = config.build() model.to_empty(device=self.device_type) diff --git a/torchtitan/models/kimi_k3/README.md b/torchtitan/models/kimi_k3/README.md index 8d61bb00cf..f440c6697f 100644 --- a/torchtitan/models/kimi_k3/README.md +++ b/torchtitan/models/kimi_k3/README.md @@ -44,14 +44,20 @@ describe the released model. ## Numerical Parity -End-to-end KL divergence against the Hugging Face implementation (multimodal -inputs): **1.5370e-6**, with **100% top-1 and top-5 match**. - -Vision parity: pixel preprocessing max difference **1.192e-7**; projected vision -features cosine similarity **1.000000** and max difference **2.669e-3**. - -Test scripts: - -- `scripts/checkpoint_conversion/numerical_tests_kimi_k3.py` -- Hugging Face vs. - TorchTitan comparison -- `tests/unit_tests/test_kimi_k3.py` -- KDA and FSDP2 correctness +`scripts/checkpoint_conversion/numerical_tests_kimi_k3.py` reduces the released +Hugging Face config to the debug model and compares both implementations on the +same text+image prompt, each side doing its own preprocessing. Float32 results: + +| Stage | Result | +|-------|--------| +| Pixel preprocessing | max difference `1.192e-7` | +| Projected vision features | cosine `1.000000`, max difference `3.152e-3` | +| MoE routing | `3936 / 3936` expert choices match | +| Last-token logits | KL `1.8215e-8`, top-1 match, top-5 `5 / 5` | + +Routed experts always run through the bf16 grouped GEMM, so the float32 logit +difference is bounded by bf16 rather than by float32. + +`tests/unit_tests/test_kimi_k3.py` covers the KDA kernel against a recurrent +reference, the HuggingFace state-dict round trip, a multimodal forward, and +FSDP2 forward/backward parity against non-distributed execution. diff --git a/torchtitan/models/kimi_k3/__init__.py b/torchtitan/models/kimi_k3/__init__.py index abb5bc93f5..51b1721a04 100644 --- a/torchtitan/models/kimi_k3/__init__.py +++ b/torchtitan/models/kimi_k3/__init__.py @@ -469,7 +469,9 @@ def _debugmodel(attn_backend: str) -> KimiK3Model.Config: dim = 256 return _kimi_k3_config( dim=dim, - vocab_size=163840, + # Sized for tests/assets/tokenizer, not the released 163840-token vocab, + # which would put 84% of the parameters in the embedding and lm_head. + vocab_size=2048, num_layers=13, full_attention_layers={3, 7, 11, 12}, attn_res_block_size=12, diff --git a/torchtitan/models/kimi_k3/config_registry.py b/torchtitan/models/kimi_k3/config_registry.py index ca55f26ba3..4737e7b905 100644 --- a/torchtitan/models/kimi_k3/config_registry.py +++ b/torchtitan/models/kimi_k3/config_registry.py @@ -6,9 +6,8 @@ from torchtitan.components.checkpoint import CheckpointManager from torchtitan.components.loss import ChunkedLossWrapper, CrossEntropyLoss -from torchtitan.components.lr_scheduler import LRSchedulersContainer from torchtitan.components.metrics import MetricsProcessor -from torchtitan.components.optimizer import default_adamw +from torchtitan.components.optimizer import default_adamw, LRSchedulersContainer from torchtitan.components.tokenizer import MultiModalTokenizer from torchtitan.config import TrainingConfig from torchtitan.hf_datasets.multimodal.mm_datasets import MMDataLoader @@ -20,7 +19,7 @@ def kimi_k3_debugmodel() -> Trainer.Config: - model_spec = model_registry("debugmodel") + model_spec = model_registry("debugmodel", attn_backend="flex") return Trainer.Config( loss=ChunkedLossWrapper.Config( loss_fn=CrossEntropyLoss.Config( diff --git a/torchtitan/models/kimi_k3/parallelize.py b/torchtitan/models/kimi_k3/parallelize.py index 9101a5af03..ce0d65fa8f 100644 --- a/torchtitan/models/kimi_k3/parallelize.py +++ b/torchtitan/models/kimi_k3/parallelize.py @@ -56,8 +56,10 @@ def parallelize_kimi_k3( if compile_config.enable and "model" in compile_config.components: raise NotImplementedError("Kimi K3 does not support model compilation.") if ac_config is not None: + # TODO: untested against the block's attention-residual tuple signature. raise NotImplementedError( - "Kimi K3 FSDP2 does not support activation checkpointing yet." + "Kimi K3 FSDP2 does not support activation checkpointing yet; " + "pass activation-checkpoint:none." ) if training.enable_cpu_offload: raise NotImplementedError( @@ -74,7 +76,9 @@ def parallelize_kimi_k3( if vision_encoder is not None: # TODO: An image batch on one DP rank and a text-only batch on another # execute different FSDP collectives, deadlock, and hit a 90-second - # timeout. A general solution is needed. + # timeout. Under CP the same deadlock is reachable even when every rank + # gets images, since a rank's sequence shard can hold no vision + # placeholders. A general solution is needed. apply_fsdp_to_vision_encoder( vision_encoder, dp_mesh, From 8d71077263b388efd3484ac19996098d036e9946 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Tue, 18 Aug 2026 18:48:23 +0000 Subject: [PATCH 40/67] register kimi_k3 in the CLI option freeze test_every_model_is_guarded requires every entry in _supported_models to have a configuration in _GUARDED_CONFIGS. The freeze snapshot passes, so kimi_k3 adds no new command-line options. --- tests/unit_tests/test_no_new_cli_options.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unit_tests/test_no_new_cli_options.py b/tests/unit_tests/test_no_new_cli_options.py index bbd8af5218..014d08b7ec 100644 --- a/tests/unit_tests/test_no_new_cli_options.py +++ b/tests/unit_tests/test_no_new_cli_options.py @@ -366,6 +366,7 @@ def _declared_cli_options( ("gpt_oss", "gpt_oss_debugmodel"), ("flux", "flux_debugmodel"), ("kimi_k2_7", "kimi_k2_5_debugmodel"), + ("kimi_k3", "kimi_k3_debugmodel"), ("muse_glimmer", "muse_glimmer_debugmodel_mm"), ) From e6e6e648fb205dd60163b200432f53447edb62b8 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Wed, 19 Aug 2026 02:47:58 +0000 Subject: [PATCH 41/67] fix unused change --- .../numerical_tests_kimi_k3.py | 6 - tests/unit_tests/test_kimi_k3.py | 48 ++---- torchtitan/models/common/vision_encoder.py | 152 +---------------- torchtitan/models/kimi_k3/README.md | 28 ++- torchtitan/models/kimi_k3/__init__.py | 20 +-- torchtitan/models/kimi_k3/config_registry.py | 2 +- torchtitan/models/kimi_k3/model.py | 5 +- torchtitan/models/kimi_k3/parallelize.py | 8 +- torchtitan/models/kimi_k3/vision_encoder.py | 159 ++++++++++++++++-- 9 files changed, 193 insertions(+), 235 deletions(-) diff --git a/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py b/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py index 32ed4b465a..e089617cbc 100644 --- a/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py +++ b/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py @@ -47,7 +47,6 @@ _HF_REPO_ID = "moonshotai/Kimi-K3" _HF_REVISION = "9f62e4e9fffbd0a83ddd60e1c209d828994b3569" _MEDIA_TOKEN_ID = 163605 -_VOCAB_SIZE = 163840 _PATCH_SIZE = 14 _MERGE_SIZE = 2 _MAX_PATCHES = 65536 @@ -459,11 +458,6 @@ def main() -> None: ) tt_config = cast(KimiK3Model.Config, model_registry(args.model_flavor).model) - # The released tokenizer emits ids across the full vocab, while debugmodel - # is sized for the test tokenizer. - tt_config.vocab_size = _VOCAB_SIZE - tt_config.tok_embeddings.num_embeddings = _VOCAB_SIZE - tt_config.lm_head.out_features = _VOCAB_SIZE torch.manual_seed(args.seed) tt_model = _build_tt_model(tt_config, dtype) hf_state_dict = KimiK3StateDictAdapter(tt_config, hf_assets_path=None).to_hf( diff --git a/tests/unit_tests/test_kimi_k3.py b/tests/unit_tests/test_kimi_k3.py index 425a63f8d9..e0fbf39be2 100644 --- a/tests/unit_tests/test_kimi_k3.py +++ b/tests/unit_tests/test_kimi_k3.py @@ -21,20 +21,13 @@ from torchtitan.config import CompileConfig, ParallelismConfig, TrainingConfig from torchtitan.distributed import ParallelDims -# FLA is a per-model dependency (kimi_k3/requirements.txt) imported at module -# scope for the KDA kernel, so skip rather than fail collection without it. -try: - from torchtitan.models.kimi_k3 import ( - _kimi_k3_config, - _vision_encoder_config, - parallelize_kimi_k3, - ) - from torchtitan.models.kimi_k3.model import KimiK3Model, KimiKDAKernel - from torchtitan.models.kimi_k3.state_dict_adapter import KimiK3StateDictAdapter -except ModuleNotFoundError as exc: - raise unittest.SkipTest( - f"Kimi K3 optional dependency unavailable: {exc.name}" - ) from exc +from torchtitan.models.kimi_k3 import ( + _kimi_k3_config, + _vision_encoder_config, + parallelize_kimi_k3, +) +from torchtitan.models.kimi_k3.model import KimiK3Model, KimiKDAKernel +from torchtitan.models.kimi_k3.state_dict_adapter import KimiK3StateDictAdapter def _small_model_config( @@ -59,7 +52,7 @@ def _small_model_config( qk_nope_head_dim=16, qk_rope_head_dim=16, v_head_dim=16, - kda_head_dim=32, + kda_head_dim=16, conv_kernel_size=3, dense_hidden_dim=128, latent_dim=32, @@ -146,25 +139,12 @@ def _kda_recurrent_reference( class TestKimiK3(unittest.TestCase): - def test_multimodal_forward(self): - # All-MLA so the forward runs without the CUDA-only KDA kernel; the - # KDA path is covered by the FSDP parity test below. - config = _small_model_config(full_attention_layers={0, 1}) + def test_flex_attention_mask(self): + config = _small_model_config() model = config.build() - model.init_states() - positions = torch.arange(6, dtype=torch.int32).unsqueeze(0) + positions = torch.arange(4, dtype=torch.int32).unsqueeze(0) attention_masks = model.get_attention_masks(positions) self.assertIsInstance(attention_masks, BlockMask) - with torch.no_grad(): - logits = model( - torch.tensor([[1, 7, 2, 3, 4, 5]]), - pixel_values=torch.randn(1, 4, 3 * 2 * 2), - grid_thw=torch.tensor([[1, 2, 2]]), - special_tokens={"image_id": 7}, - positions=positions, - attention_masks=attention_masks, - ) - self.assertEqual(logits.shape, (1, 6, config.vocab_size)) def test_update_from_config_propagates_moe_force_load_balance(self): from torchtitan.config import DebugConfig @@ -285,8 +265,10 @@ def world_size(self): @with_comms def test_fsdp_matches_non_distributed_forward_backward(self): torch.manual_seed(3) - # Layer 0 is KDA, layer 1 is MLA, so one run covers both attentions. - config = _small_model_config(attn_res_block_size=2) + config = _small_model_config( + attn_res_block_size=2, + full_attention_layers={0, 1}, + ) with torch.device("meta"): model = config.build() model.to_empty(device=self.device_type) diff --git a/torchtitan/models/common/vision_encoder.py b/torchtitan/models/common/vision_encoder.py index 832fa63422..4576a4bfd9 100644 --- a/torchtitan/models/common/vision_encoder.py +++ b/torchtitan/models/common/vision_encoder.py @@ -22,17 +22,13 @@ from collections.abc import Callable from dataclasses import dataclass, field -from typing import cast -import spmd_types as spmd import torch from torch.nn.attention.flex_attention import BlockMask, create_block_mask -from torchtitan.distributed.utils import get_spmd_backend from torchtitan.models.common import Linear from torchtitan.models.common.attention import FlexAttention, local_head_split -from torchtitan.models.common.nn_modules import GELU, LayerNorm, RMSNorm -from torchtitan.models.common.rope import _maybe_wrap_positions +from torchtitan.models.common.nn_modules import GELU, LayerNorm from torchtitan.protocols.module import Module compiled_create_block_mask = torch.compile(create_block_mask) @@ -70,148 +66,6 @@ def mask_mod(b, h, q_idx, kv_idx): ) -def get_temporal_pos_embed( - num_frames: int, - embed_dim: int, - *, - base: float = 10000.0, - device: torch.device | None = None, -) -> torch.Tensor: - """Fixed 1D sinusoidal embeddings for the temporal axis (video frames). - - Returns ``(num_frames, embed_dim)`` float32; the standard 1D sincos formula - over frame indices. - - Args: - num_frames: Number of video frames (temporal positions). - embed_dim: Embedding width per frame. - base: Sinusoid base (longest wavelength); the conventional PE constant. - device: Device for the returned tensor. - """ - grid = torch.arange(num_frames, dtype=torch.float32, device=device) - omega = torch.arange(embed_dim // 2, dtype=torch.float32, device=device) / ( - embed_dim / 2.0 - ) - omega = 1.0 / base**omega - out = torch.outer(grid, omega) - return torch.cat([out.sin(), out.cos()], dim=1) - - -def compute_2d_rope_cache( - freq_table: torch.Tensor, - grids: list[list[int]], - max_num_patch: int, - head_dim: int, -) -> torch.Tensor: - """Compute the padded 2D-RoPE complex ``freqs_cis`` cache in raster order. - - For head-dim pair index ``k`` (``k`` in ``[0, head_dim/4)``), even output - pairs are rotated by the *column* (x) position and odd pairs by the *row* - (y) position. The per-axis angle for a position ``p`` is ``p * inv_freq[k]``; - this looks it up by gathering row ``p`` of ``freq_table`` (built once by - ``VisionRotaryEmbedding2D`` and cached by the encoder) rather than - recomputing ``p * inv_freq`` each call. Frames repeat the spatial pattern. - - Returns a complex cache consumed by ``ComplexRoPE.apply_rotary_emb``; only - the cache is 2D/per-grid, which is why it is built here rather than by the - 1D ``ComplexRoPE`` cache machinery. - - Args: - freq_table: ``(max_hw, head_dim/4)`` position-to-frequency table, where - ``freq_table[p, k] = p * inv_freq[k]``. - grids: per-item ``[t, h, w]`` patch counts as host ints (``grid_thw`` - read to CPU once by the caller, so the per-item loop adds no syncs). - max_num_patch: Padded sequence length. - head_dim: Attention head dim (must be divisible by 4). - - Returns: - ``(N, max_num_patch, 1, head_dim/2)`` complex64 (head axis = 1 to - broadcast over the heads). - """ - device = freq_table.device - - angles = freq_table.new_zeros(len(grids), max_num_patch, head_dim // 2) - if get_spmd_backend() == "spmd_types" and spmd.is_type_checking(): - angles = spmd.mutate_type(angles, src=spmd.R, dst={"dp": spmd.V, "tp": spmd.I}) - - # Group by (h, w) so the per-resolution angle grid is built once. - hw_to_indices: dict[tuple[int, int], list[int]] = {} - for i, (_, h, w) in enumerate(grids): - hw_to_indices.setdefault((h, w), []).append(i) - - for (h, w), indices in hw_to_indices.items(): - # Raster order: position p -> (row = p // w, col = p % w). Gather each - # axis's angles from the precomputed table (freq_table[pos] = pos*inv_freq). - flat = torch.arange(h * w, device=device) - flat = cast(torch.Tensor, _maybe_wrap_positions(flat, freq_table)) - if get_spmd_backend() == "spmd_types" and spmd.is_type_checking(): - flat = spmd.mutate_type(flat, "tp", src=spmd.R, dst=spmd.I) - x_ang = freq_table[flat % w] # (h*w, head_dim/4) column - y_ang = freq_table[flat // w] # (h*w, head_dim/4) row - # Interleave x/y so pair 2k uses x-position, pair 2k+1 uses y-position. - ang = torch.stack([x_ang, y_ang], dim=-1).reshape(h * w, head_dim // 2) - for i in indices: - t = grids[i][0] - seq_len = t * h * w - angles[i, :seq_len] = ang.repeat(t, 1) - - # Complex unit-modulus cache; unsqueeze the head axis for broadcast. - return torch.polar(torch.ones_like(angles), angles).unsqueeze(2) - - -class VisionRotaryEmbedding2D(Module): - """2D rotary position embedding for the vision tower. - - Holds the per-axis frequencies ``inv_freq`` (``head_dim/4`` of them, shared - by the row and column axes). ``forward(seqlen)`` returns the - position-to-frequency table ``freq_table[p, k] = p * inv_freq[k]`` for - positions up to ``seqlen``; ``compute_2d_rope_cache`` gathers per-patch - row/col angles from it, and ``ComplexRoPE.apply_rotary_emb`` applies them. - ``head_dim`` must be divisible by 4. - """ - - @dataclass(kw_only=True, slots=True) - class Config(Module.Config): - head_dim: int - theta: float = 10000.0 - - def __init__(self, config: Config): - super().__init__() - if config.head_dim % 4 != 0: - raise ValueError( - f"2D RoPE requires head_dim divisible by 4, got {config.head_dim}." - ) - self.head_dim = config.head_dim - self.theta = config.theta - self.register_buffer("inv_freq", self._compute_inv_freq(), persistent=False) - - def _compute_inv_freq(self, *, device: torch.device | None = None) -> torch.Tensor: - # inv_freq[k] = theta**(-4k/head_dim) for k in [0, head_dim/4); the - # step of 4 leaves room for the row/col split of the 2D rotation. - return 1.0 / ( - self.theta - ** ( - torch.arange(0, self.head_dim, 4, dtype=torch.float32, device=device) - / self.head_dim - ) - ) - - def _init_self_buffers(self, *, buffer_device: torch.device | None = None) -> None: - """Re-compute inv_freq on the target device after to_empty().""" - device = buffer_device or self.inv_freq.device - self.inv_freq = self._compute_inv_freq(device=device) - - def forward(self, seqlen: int) -> torch.Tensor: - """Frequency table ``(seqlen, head_dim/4)`` for positions ``[0, seqlen)``.""" - seq = torch.arange( - seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype - ) - seq = cast(torch.Tensor, _maybe_wrap_positions(seq, self.inv_freq)) - if get_spmd_backend() == "spmd_types" and spmd.is_type_checking(): - seq = spmd.mutate_type(seq, "tp", src=spmd.R, dst=spmd.I) - return torch.outer(seq, self.inv_freq) - - class VisionMLP(Module): """Feed-forward network with GELU activation (fc1 -> act -> fc2).""" @@ -296,8 +150,8 @@ class VisionTransformerBlock(Module): @dataclass(kw_only=True, slots=True) class Config(Module.Config): - norm1: LayerNorm.Config | RMSNorm.Config - norm2: LayerNorm.Config | RMSNorm.Config + norm1: LayerNorm.Config + norm2: LayerNorm.Config attn: VisionAttention.Config mlp: VisionMLP.Config diff --git a/torchtitan/models/kimi_k3/README.md b/torchtitan/models/kimi_k3/README.md index f440c6697f..8d61bb00cf 100644 --- a/torchtitan/models/kimi_k3/README.md +++ b/torchtitan/models/kimi_k3/README.md @@ -44,20 +44,14 @@ describe the released model. ## Numerical Parity -`scripts/checkpoint_conversion/numerical_tests_kimi_k3.py` reduces the released -Hugging Face config to the debug model and compares both implementations on the -same text+image prompt, each side doing its own preprocessing. Float32 results: - -| Stage | Result | -|-------|--------| -| Pixel preprocessing | max difference `1.192e-7` | -| Projected vision features | cosine `1.000000`, max difference `3.152e-3` | -| MoE routing | `3936 / 3936` expert choices match | -| Last-token logits | KL `1.8215e-8`, top-1 match, top-5 `5 / 5` | - -Routed experts always run through the bf16 grouped GEMM, so the float32 logit -difference is bounded by bf16 rather than by float32. - -`tests/unit_tests/test_kimi_k3.py` covers the KDA kernel against a recurrent -reference, the HuggingFace state-dict round trip, a multimodal forward, and -FSDP2 forward/backward parity against non-distributed execution. +End-to-end KL divergence against the Hugging Face implementation (multimodal +inputs): **1.5370e-6**, with **100% top-1 and top-5 match**. + +Vision parity: pixel preprocessing max difference **1.192e-7**; projected vision +features cosine similarity **1.000000** and max difference **2.669e-3**. + +Test scripts: + +- `scripts/checkpoint_conversion/numerical_tests_kimi_k3.py` -- Hugging Face vs. + TorchTitan comparison +- `tests/unit_tests/test_kimi_k3.py` -- KDA and FSDP2 correctness diff --git a/torchtitan/models/kimi_k3/__init__.py b/torchtitan/models/kimi_k3/__init__.py index 51b1721a04..63e886e587 100644 --- a/torchtitan/models/kimi_k3/__init__.py +++ b/torchtitan/models/kimi_k3/__init__.py @@ -16,12 +16,7 @@ from torchtitan.models.common.moe import RoutedExperts, TokenChoiceTopKRouter from torchtitan.models.common.nn_modules import GELU, RMSNorm from torchtitan.models.common.token_dispatcher import LocalTokenDispatcher -from torchtitan.models.common.vision_encoder import ( - VisionAttention, - VisionMLP, - VisionRotaryEmbedding2D, - VisionTransformerBlock, -) +from torchtitan.models.common.vision_encoder import VisionAttention, VisionMLP from torchtitan.models.utils import validate_converter_order from torchtitan.protocols.model import ModelConfigConverter from torchtitan.protocols.model_spec import ModelSpec @@ -39,7 +34,12 @@ ) from .parallelize import parallelize_kimi_k3 from .state_dict_adapter import KimiK3StateDictAdapter -from .vision_encoder import KimiK3VisionEncoder, KimiK3VisionProjector +from .vision_encoder import ( + KimiK3VisionBlock, + KimiK3VisionEncoder, + KimiK3VisionProjector, + VisionRotaryEmbedding2D, +) __all__ = [ "KIMI_K3_SPECIAL_TOKENS", @@ -294,7 +294,7 @@ def _vision_encoder_config( eps=1e-5, param_init=_NORM_INIT, ) - block = VisionTransformerBlock.Config( + block = KimiK3VisionBlock.Config( norm1=vision_norm, norm2=vision_norm, attn=VisionAttention.Config( @@ -469,9 +469,7 @@ def _debugmodel(attn_backend: str) -> KimiK3Model.Config: dim = 256 return _kimi_k3_config( dim=dim, - # Sized for tests/assets/tokenizer, not the released 163840-token vocab, - # which would put 84% of the parameters in the embedding and lm_head. - vocab_size=2048, + vocab_size=163840, num_layers=13, full_attention_layers={3, 7, 11, 12}, attn_res_block_size=12, diff --git a/torchtitan/models/kimi_k3/config_registry.py b/torchtitan/models/kimi_k3/config_registry.py index 4737e7b905..afda14fd50 100644 --- a/torchtitan/models/kimi_k3/config_registry.py +++ b/torchtitan/models/kimi_k3/config_registry.py @@ -19,7 +19,7 @@ def kimi_k3_debugmodel() -> Trainer.Config: - model_spec = model_registry("debugmodel", attn_backend="flex") + model_spec = model_registry("debugmodel") return Trainer.Config( loss=ChunkedLossWrapper.Config( loss_fn=CrossEntropyLoss.Config( diff --git a/torchtitan/models/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py index d2a1e47f92..cbf27e7c64 100644 --- a/torchtitan/models/kimi_k3/model.py +++ b/torchtitan/models/kimi_k3/model.py @@ -453,8 +453,9 @@ def forward(self, x_BLD: torch.Tensor) -> torch.Tensor: -1, expert_ids_BLS, True ) num_tokens_per_expert_E = routing_map_BLE.sum(dim=(0, 1)) - with torch.no_grad(): - self.tokens_per_expert_E.add_(num_tokens_per_expert_E) + if self.training: + with torch.no_grad(): + self.tokens_per_expert_E.add_(num_tokens_per_expert_E) routed_BLD = self.routed_experts( self.routed_down(x_BLD), diff --git a/torchtitan/models/kimi_k3/parallelize.py b/torchtitan/models/kimi_k3/parallelize.py index ce0d65fa8f..9101a5af03 100644 --- a/torchtitan/models/kimi_k3/parallelize.py +++ b/torchtitan/models/kimi_k3/parallelize.py @@ -56,10 +56,8 @@ def parallelize_kimi_k3( if compile_config.enable and "model" in compile_config.components: raise NotImplementedError("Kimi K3 does not support model compilation.") if ac_config is not None: - # TODO: untested against the block's attention-residual tuple signature. raise NotImplementedError( - "Kimi K3 FSDP2 does not support activation checkpointing yet; " - "pass activation-checkpoint:none." + "Kimi K3 FSDP2 does not support activation checkpointing yet." ) if training.enable_cpu_offload: raise NotImplementedError( @@ -76,9 +74,7 @@ def parallelize_kimi_k3( if vision_encoder is not None: # TODO: An image batch on one DP rank and a text-only batch on another # execute different FSDP collectives, deadlock, and hit a 90-second - # timeout. Under CP the same deadlock is reachable even when every rank - # gets images, since a rank's sequence shard can hold no vision - # placeholders. A general solution is needed. + # timeout. A general solution is needed. apply_fsdp_to_vision_encoder( vision_encoder, dp_mesh, diff --git a/torchtitan/models/kimi_k3/vision_encoder.py b/torchtitan/models/kimi_k3/vision_encoder.py index 3f07042a4b..9a6960294e 100644 --- a/torchtitan/models/kimi_k3/vision_encoder.py +++ b/torchtitan/models/kimi_k3/vision_encoder.py @@ -4,12 +4,15 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -"""MoonViT-V2 vision encoder used by Kimi K3. +"""MoonViT3d vision encoder used by Kimi K3. Shape suffixes: - N = number of visual items - P = maximum patches per item (padded) - D = vision hidden dimension +- H = number of attention heads +- K = attention head dimension +- C = number of complex-valued head-dimension pairs - M = maximum merged tokens per item (padded) - F = merged feature dimension - O = projected text dimension @@ -19,21 +22,35 @@ import torch import torch.nn.functional as F +from torch.nn.attention.flex_attention import BlockMask from torchtitan.models.common import Linear from torchtitan.models.common.nn_modules import GELU, RMSNorm -from torchtitan.models.common.rope import ComplexRoPE from torchtitan.models.common.vision_encoder import ( compiled_create_block_mask, - compute_2d_rope_cache, - get_temporal_pos_embed, get_vision_block_mask_mod, - VisionRotaryEmbedding2D, - VisionTransformerBlock, + VisionAttention, + VisionMLP, ) from torchtitan.protocols.module import Module, ModuleDict +def _get_temporal_pos_embed( + num_frames: int, + embed_dim: int, + *, + device: torch.device, +) -> torch.Tensor: + """Return fixed 1D sinusoidal embeddings for video frame positions.""" + grid = torch.arange(num_frames, dtype=torch.float32, device=device) + omega = torch.arange(embed_dim // 2, dtype=torch.float32, device=device) / ( + embed_dim / 2.0 + ) + omega = 1.0 / 10000.0**omega + angles = torch.outer(grid, omega) + return torch.cat((angles.sin(), angles.cos()), dim=-1) + + def _pad_sequence(x: torch.Tensor, target_length: int) -> torch.Tensor: """Pad the leading sequence dimension without modifying ``x`` in place.""" padding_length = target_length - x.shape[0] @@ -87,7 +104,7 @@ def _compute_learned_pos_embeds( if num_frames == 1: item_pos = spatial else: - temporal = get_temporal_pos_embed(num_frames, dim, device=pos_embed.device) + temporal = _get_temporal_pos_embed(num_frames, dim, device=pos_embed.device) item_pos = spatial.unsqueeze(0) + temporal.unsqueeze(1).to(spatial.dtype) item_pos = item_pos.reshape(num_frames * grid_h * grid_w, dim) padded_positions.append(_pad_sequence(item_pos, max_num_patches)) @@ -95,6 +112,51 @@ def _compute_learned_pos_embeds( return torch.stack(padded_positions) +def _compute_2d_rope_cache( + freq_table: torch.Tensor, + grids: list[list[int]], + max_num_patches: int, + head_dim: int, +) -> torch.Tensor: + """Build the real-valued 2D RoPE cache in raster patch order.""" + cached_spatial: dict[tuple[int, int], torch.Tensor] = {} + padded_angles = [] + for num_frames, grid_h, grid_w in grids: + spatial = cached_spatial.get((grid_h, grid_w)) + if spatial is None: + flat = torch.arange(grid_h * grid_w, device=freq_table.device) + x_angles = freq_table[flat % grid_w] + y_angles = freq_table[flat // grid_w] + spatial = torch.stack((x_angles, y_angles), dim=-1).reshape( + grid_h * grid_w, head_dim // 2 + ) + cached_spatial[(grid_h, grid_w)] = spatial + item_angles = spatial.repeat(num_frames, 1) + padded_angles.append(_pad_sequence(item_angles, max_num_patches)) + + angles = torch.stack(padded_angles) + return torch.stack((angles.cos(), angles.sin()), dim=-1).unsqueeze(2) + + +def _apply_2d_rope( + q_NPHK: torch.Tensor, + k_NPHK: torch.Tensor, + rope_cache_NP1C2: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Apply 2D RoPE using the real form of complex multiplication.""" + + cos_NP1C = rope_cache_NP1C2[..., 0] + sin_NP1C = rope_cache_NP1C2[..., 1] + + def rotate(x_NPHK: torch.Tensor) -> torch.Tensor: + x_NPHC2 = x_NPHK.float().reshape(*x_NPHK.shape[:-1], -1, 2) + real_NPHC = x_NPHC2[..., 0] * cos_NP1C - x_NPHC2[..., 1] * sin_NP1C + imag_NPHC = x_NPHC2[..., 0] * sin_NP1C + x_NPHC2[..., 1] * cos_NP1C + return torch.stack((real_NPHC, imag_NPHC), dim=-1).flatten(-2) + + return rotate(q_NPHK).to(q_NPHK.dtype), rotate(k_NPHK).to(k_NPHK.dtype) + + def _temporal_pool_and_merge( hidden_NPD: torch.Tensor, grids: list[list[int]], @@ -127,6 +189,84 @@ def _temporal_pool_and_merge( return torch.stack(padded_items) +class VisionRotaryEmbedding2D(Module): + """Per-axis frequency table for MoonViT's interleaved 2D RoPE.""" + + @dataclass(kw_only=True, slots=True) + class Config(Module.Config): + head_dim: int + theta: float = 10000.0 + + def __init__(self, config: Config): + super().__init__() + if config.head_dim % 4 != 0: + raise ValueError( + "Vision 2D RoPE head_dim must be divisible by 4, " + f"got {config.head_dim}." + ) + self.head_dim = config.head_dim + self.theta = config.theta + self.register_buffer("inv_freq", self._compute_inv_freq(), persistent=False) + + def _compute_inv_freq(self, *, device: torch.device | None = None) -> torch.Tensor: + return 1.0 / ( + self.theta + ** ( + torch.arange( + 0, + self.head_dim, + 4, + dtype=torch.float32, + device=device, + ) + / self.head_dim + ) + ) + + def _init_self_buffers(self, *, buffer_device: torch.device | None = None) -> None: + device = buffer_device or self.inv_freq.device + self.inv_freq = self._compute_inv_freq(device=device) + + def forward(self, seqlen: int) -> torch.Tensor: + positions = torch.arange( + seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype + ) + return torch.outer(positions, self.inv_freq) + + +class KimiK3VisionBlock(Module): + """MoonViT pre-norm attention and MLP block.""" + + @dataclass(kw_only=True, slots=True) + class Config(Module.Config): + norm1: RMSNorm.Config + norm2: RMSNorm.Config + attn: VisionAttention.Config + mlp: VisionMLP.Config + + def __init__(self, config: Config): + super().__init__() + self.norm1 = config.norm1.build() + self.norm2 = config.norm2.build() + self.attn = config.attn.build() + self.mlp = config.mlp.build() + + def forward( + self, + x_NPD: torch.Tensor, + *, + rope_cache: torch.Tensor, + attention_mask: BlockMask, + ) -> torch.Tensor: + x_NPD = x_NPD + self.attn( + self.norm1(x_NPD), + rope_cache=rope_cache, + rope_apply=_apply_2d_rope, + attention_mask=attention_mask, + ) + return x_NPD + self.mlp(self.norm2(x_NPD)) + + class KimiK3VisionProjector(Module): """PatchMergerMLPV2 projector from merged vision features to text width.""" @@ -170,7 +310,7 @@ class Config(Module.Config): interpolation_mode: str patch_embed_proj: Linear.Config rotary_pos_emb: VisionRotaryEmbedding2D.Config - block: VisionTransformerBlock.Config + block: KimiK3VisionBlock.Config final_norm: RMSNorm.Config projector: KimiK3VisionProjector.Config @@ -217,7 +357,7 @@ def _compute_position_embeddings( self.interpolation_mode, self.max_num_frames, ) - rope_cache = compute_2d_rope_cache( + rope_cache = _compute_2d_rope_cache( self._cached_freq_table, grids, max_num_patches, @@ -275,7 +415,6 @@ def forward( hidden_NPD = block( hidden_NPD, rope_cache=rope_cache, - rope_apply=ComplexRoPE.apply_rotary_emb, attention_mask=attention_mask, ) hidden_NPD = self.final_norm(hidden_NPD) From 905b36674c00c69c5f85d2d341c5cb4d24fd5232 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Wed, 19 Aug 2026 04:05:24 +0000 Subject: [PATCH 42/67] reuse ComplexRoPE --- torchtitan/models/kimi_k3/README.md | 4 +-- torchtitan/models/kimi_k3/vision_encoder.py | 27 +++++---------------- 2 files changed, 8 insertions(+), 23 deletions(-) diff --git a/torchtitan/models/kimi_k3/README.md b/torchtitan/models/kimi_k3/README.md index 8d61bb00cf..49137970e8 100644 --- a/torchtitan/models/kimi_k3/README.md +++ b/torchtitan/models/kimi_k3/README.md @@ -45,10 +45,10 @@ describe the released model. ## Numerical Parity End-to-end KL divergence against the Hugging Face implementation (multimodal -inputs): **1.5370e-6**, with **100% top-1 and top-5 match**. +inputs): **9.6895e-8**, with **100% top-1 and top-5 match**. Vision parity: pixel preprocessing max difference **1.192e-7**; projected vision -features cosine similarity **1.000000** and max difference **2.669e-3**. +features cosine similarity **1.000000** and max difference **2.654e-3**. Test scripts: diff --git a/torchtitan/models/kimi_k3/vision_encoder.py b/torchtitan/models/kimi_k3/vision_encoder.py index 9a6960294e..c83b0d698c 100644 --- a/torchtitan/models/kimi_k3/vision_encoder.py +++ b/torchtitan/models/kimi_k3/vision_encoder.py @@ -26,6 +26,7 @@ from torchtitan.models.common import Linear from torchtitan.models.common.nn_modules import GELU, RMSNorm +from torchtitan.models.common.rope import ComplexRoPE from torchtitan.models.common.vision_encoder import ( compiled_create_block_mask, get_vision_block_mask_mod, @@ -135,26 +136,10 @@ def _compute_2d_rope_cache( padded_angles.append(_pad_sequence(item_angles, max_num_patches)) angles = torch.stack(padded_angles) - return torch.stack((angles.cos(), angles.sin()), dim=-1).unsqueeze(2) - - -def _apply_2d_rope( - q_NPHK: torch.Tensor, - k_NPHK: torch.Tensor, - rope_cache_NP1C2: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor]: - """Apply 2D RoPE using the real form of complex multiplication.""" - - cos_NP1C = rope_cache_NP1C2[..., 0] - sin_NP1C = rope_cache_NP1C2[..., 1] - - def rotate(x_NPHK: torch.Tensor) -> torch.Tensor: - x_NPHC2 = x_NPHK.float().reshape(*x_NPHK.shape[:-1], -1, 2) - real_NPHC = x_NPHC2[..., 0] * cos_NP1C - x_NPHC2[..., 1] * sin_NP1C - imag_NPHC = x_NPHC2[..., 0] * sin_NP1C + x_NPHC2[..., 1] * cos_NP1C - return torch.stack((real_NPHC, imag_NPHC), dim=-1).flatten(-2) - - return rotate(q_NPHK).to(q_NPHK.dtype), rotate(k_NPHK).to(k_NPHK.dtype) + # ComplexRoPE.apply_rotary_emb multiplies in complex64; float() only widens + # the container, so cos/sin keep whatever precision angles were computed in. + cos_sin = torch.stack((angles.cos(), angles.sin()), dim=-1).float() + return torch.view_as_complex(cos_sin).unsqueeze(2) def _temporal_pool_and_merge( @@ -261,7 +246,7 @@ def forward( x_NPD = x_NPD + self.attn( self.norm1(x_NPD), rope_cache=rope_cache, - rope_apply=_apply_2d_rope, + rope_apply=ComplexRoPE.apply_rotary_emb, attention_mask=attention_mask, ) return x_NPD + self.mlp(self.norm2(x_NPD)) From 0f45ff8da18b72079852ecb6af186b89beb9d915 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Wed, 19 Aug 2026 06:34:08 +0000 Subject: [PATCH 43/67] scale up k3 model --- torchtitan/models/kimi_k3/README.md | 4 ++-- torchtitan/models/kimi_k3/__init__.py | 12 ++++++------ 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/torchtitan/models/kimi_k3/README.md b/torchtitan/models/kimi_k3/README.md index 49137970e8..897c387354 100644 --- a/torchtitan/models/kimi_k3/README.md +++ b/torchtitan/models/kimi_k3/README.md @@ -45,10 +45,10 @@ describe the released model. ## Numerical Parity End-to-end KL divergence against the Hugging Face implementation (multimodal -inputs): **9.6895e-8**, with **100% top-1 and top-5 match**. +inputs): **6.7634e-7**, with **100% top-1 and top-5 match**. Vision parity: pixel preprocessing max difference **1.192e-7**; projected vision -features cosine similarity **1.000000** and max difference **2.654e-3**. +features cosine similarity **1.000000** and max difference **2.730e-3**. Test scripts: diff --git a/torchtitan/models/kimi_k3/__init__.py b/torchtitan/models/kimi_k3/__init__.py index 63e886e587..62e061248e 100644 --- a/torchtitan/models/kimi_k3/__init__.py +++ b/torchtitan/models/kimi_k3/__init__.py @@ -473,9 +473,9 @@ def _debugmodel(attn_backend: str) -> KimiK3Model.Config: num_layers=13, full_attention_layers={3, 7, 11, 12}, attn_res_block_size=12, - num_heads=4, - q_lora_rank=128, - kv_lora_rank=64, + num_heads=8, + q_lora_rank=256, + kv_lora_rank=128, qk_nope_head_dim=32, qk_rope_head_dim=16, v_head_dim=32, @@ -483,9 +483,9 @@ def _debugmodel(attn_backend: str) -> KimiK3Model.Config: conv_kernel_size=4, dense_hidden_dim=1024, latent_dim=128, - expert_hidden_dim=128, - num_experts=8, - top_k=2, + expert_hidden_dim=256, + num_experts=16, + top_k=4, num_shared_experts=2, vision_encoder=_vision_encoder_config( text_dim=dim, From 62732742af2641e8b621025a575e3ac298e0c394 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Wed, 19 Aug 2026 07:56:55 +0000 Subject: [PATCH 44/67] remove nit --- .../numerical_tests_kimi.py | 26 +------------------ 1 file changed, 1 insertion(+), 25 deletions(-) diff --git a/scripts/checkpoint_conversion/numerical_tests_kimi.py b/scripts/checkpoint_conversion/numerical_tests_kimi.py index 515e8ff98d..2ba33e8176 100644 --- a/scripts/checkpoint_conversion/numerical_tests_kimi.py +++ b/scripts/checkpoint_conversion/numerical_tests_kimi.py @@ -4,31 +4,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -"""Full text+image e2e logit parity: torchtitan Kimi-VL vs the released HF model. - -Runs the released HF Kimi-VL (``trust_remote_code``) and torchtitan in ONE -process on the same text+image prompt and compares last-token logits. torchtitan -does its OWN Kimi image processing (``process_image`` + ``vision_to_patches``, -raster order), so the full pipeline is exercised (preprocessing + vision + -projector + scatter + DeepSeek-V3 text tower), not just the forward. - -The released remote code targets transformers ~4.50.x and does NOT import on 5.x, -so run this in an env with ``transformers==4.50.3`` + ``tiktoken`` + ``blobfile``. - -Precision: the HF reference runs at ``--hf_dtype`` (default float32, the model's -"true" output); torchtitan runs at ``--dtype`` (text) with ``--vision_dtype`` -overriding only its vision encoder. ``--dtype float32`` is the correctness gate; -``--dtype bfloat16 --vision_dtype float16`` is the realistic config (~1e-2 KL). - -Usage: - CUDA_VISIBLE_DEVICES=0 python -m \\ - scripts.checkpoint_conversion.numerical_tests_kimi \\ - --hf_model_path ~/hf_assets/moonshotai/Kimi-VL-A3B-Instruct \\ - --tt_checkpoint_path outputs/kimi/kimi_vl_a3b_dcp --dtype float32 - -Add ``--force-hf-routing`` to make titan use HF's exact per-token expert -selections (diagnostic: removes the MoE routing-flip divergence). -""" +"""Full text+image e2e logit parity: torchtitan Kimi-VL vs the released HF model.""" import argparse import os From d56c69e49d3652661019d340744ad89768d2d80f Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Wed, 19 Aug 2026 09:03:58 +0000 Subject: [PATCH 45/67] remove unsed head_dim in KimiDeltaAttention --- tests/unit_tests/test_kimi_k3.py | 5 +---- torchtitan/models/kimi_k3/__init__.py | 6 +----- torchtitan/models/kimi_k3/model.py | 28 ++------------------------- 3 files changed, 4 insertions(+), 35 deletions(-) diff --git a/tests/unit_tests/test_kimi_k3.py b/tests/unit_tests/test_kimi_k3.py index e0fbf39be2..bdf055556d 100644 --- a/tests/unit_tests/test_kimi_k3.py +++ b/tests/unit_tests/test_kimi_k3.py @@ -191,10 +191,7 @@ def parameter(*shape: int) -> torch.Tensor: tensor.detach().clone().requires_grad_() for tensor in actual_inputs ) - kernel = KimiKDAKernel.Config( - head_dim=head_dim, - lower_bound=lower_bound, - ).build() + kernel = KimiKDAKernel.Config(lower_bound=lower_bound).build() actual_BLHV = kernel(*actual_inputs) expected_BLHV = _kda_recurrent_reference( *expected_inputs, diff --git a/torchtitan/models/kimi_k3/__init__.py b/torchtitan/models/kimi_k3/__init__.py index 62e061248e..dd8756ab2e 100644 --- a/torchtitan/models/kimi_k3/__init__.py +++ b/torchtitan/models/kimi_k3/__init__.py @@ -206,10 +206,7 @@ def conv() -> Conv1d.Config: forget_b=_linear(head_dim, projection_dim), beta=_linear(dim, num_heads), output_gate=_linear(dim, projection_dim), - kernel=KimiKDAKernel.Config( - head_dim=head_dim, - lower_bound=-5.0, - ), + kernel=KimiKDAKernel.Config(lower_bound=-5.0), output_norm=KimiRMSNormGated.Config( dim=head_dim, eps=1e-5, @@ -334,7 +331,6 @@ def _vision_encoder_config( block=block, final_norm=vision_norm, projector=KimiK3VisionProjector.Config( - merged_dim=merged_dim, linear_1=_linear( merged_dim, merged_dim, diff --git a/torchtitan/models/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py index cbf27e7c64..19257427dc 100644 --- a/torchtitan/models/kimi_k3/model.py +++ b/torchtitan/models/kimi_k3/model.py @@ -228,20 +228,15 @@ class KimiKDAKernel(Module): The gate activation, the beta sigmoid, and the query/key L2 norm are all fused into the kernel rather than materialized here, so the decay never - exists as a full ``(B, L, H, K)`` tensor. ``ReferenceKimiKDAKernel`` in - ``tests/unit_tests/test_kimi_k3.py`` implements the same interface as an - explicit recurrence and is the numerical baseline for this kernel; it also - lets the CPU test suite exercise the surrounding model without Triton. + exists as a full ``(B, L, H, K)`` tensor. """ @dataclass(kw_only=True, slots=True) class Config(Module.Config): - head_dim: int lower_bound: float | None = -5.0 def __init__(self, config: Config): super().__init__() - self.head_dim = config.head_dim self.lower_bound = config.lower_bound if self.lower_bound is not None and not (-5.0 <= self.lower_bound < 0.0): raise ValueError("KDA lower_bound must be in the safe range [-5, 0).") @@ -293,9 +288,6 @@ class Config(Module.Config): forget_b: Linear.Config beta: Linear.Config output_gate: Linear.Config - # Typed as the base config so the KDA kernel stays swappable, matching - # how VisionAttention types its inner_attention. Anything assigned here - # must accept KimiKDAKernel.forward's arguments. kernel: Module.Config output_norm: KimiRMSNormGated.Config output_proj: Linear.Config @@ -607,19 +599,7 @@ class Config(Decoder.Config): spatial_merge_size: int = 2 def update_from_config(self, *, config, **kwargs) -> None: - parallelism = config.parallelism - unsupported = { - "tensor parallel": parallelism.tensor_parallel_degree, - "pipeline parallel": parallelism.pipeline_parallel_degree, - "context parallel": parallelism.context_parallel_degree, - "expert parallel": parallelism.expert_parallel_degree, - } - enabled = [name for name, degree in unsupported.items() if degree > 1] - if enabled: - raise NotImplementedError( - "Kimi K3 supports FSDP2 data parallelism only; " - f"disable {', '.join(enabled)}." - ) + # Unsupported parallelisms are rejected in parallelize_kimi_k3. dataloader = getattr(config, "dataloader", None) if getattr(dataloader, "packing_buffer_size", 0) > 0: raise NotImplementedError( @@ -705,10 +685,6 @@ def _prepare_multimodal_embeds( num_tokens_per_item, special_tokens["image_id"], ) - if not vision_positions: - raise ValueError( - "pixel_values were provided but no image placeholder tokens were found." - ) return scatter_vision_embeds( embeddings, vision_embeds=vision_embeds, From 5e006ba3ec07f7389f1791f9a56cb674f6a28404 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Wed, 19 Aug 2026 09:04:02 +0000 Subject: [PATCH 46/67] remove unsed riase Error --- torchtitan/models/kimi_k3/vision_encoder.py | 14 -------------- 1 file changed, 14 deletions(-) diff --git a/torchtitan/models/kimi_k3/vision_encoder.py b/torchtitan/models/kimi_k3/vision_encoder.py index c83b0d698c..688db3eb2c 100644 --- a/torchtitan/models/kimi_k3/vision_encoder.py +++ b/torchtitan/models/kimi_k3/vision_encoder.py @@ -55,10 +55,6 @@ def _get_temporal_pos_embed( def _pad_sequence(x: torch.Tensor, target_length: int) -> torch.Tensor: """Pad the leading sequence dimension without modifying ``x`` in place.""" padding_length = target_length - x.shape[0] - if padding_length < 0: - raise ValueError( - f"Cannot pad a sequence of length {x.shape[0]} to {target_length}." - ) if padding_length == 0: return x padding = x.new_zeros(padding_length, *x.shape[1:]) @@ -257,7 +253,6 @@ class KimiK3VisionProjector(Module): @dataclass(kw_only=True, slots=True) class Config(Module.Config): - merged_dim: int linear_1: Linear.Config linear_2: Linear.Config post_norm: RMSNorm.Config @@ -265,18 +260,12 @@ class Config(Module.Config): def __init__(self, config: Config): super().__init__() - self.merged_dim = config.merged_dim self.linear_1 = config.linear_1.build() self.linear_2 = config.linear_2.build() self.post_norm = config.post_norm.build() self.activation = config.activation.build() def forward(self, merged_NMF: torch.Tensor) -> torch.Tensor: - if merged_NMF.shape[-1] != self.merged_dim: - raise ValueError( - f"Expected merged vision dim {self.merged_dim}, got " - f"{merged_NMF.shape[-1]}." - ) projected_NMO = self.linear_2(self.activation(self.linear_1(merged_NMF))) return self.post_norm(projected_NMO) @@ -301,9 +290,6 @@ class Config(Module.Config): def __init__(self, config: Config): super().__init__() - self.dim = config.dim - self.patch_size = config.patch_size - self.in_channels = config.in_channels self.merge_kernel_size = config.merge_kernel_size self.max_num_frames = config.max_num_frames self.interpolation_mode = config.interpolation_mode From 99c80c098fd9e61f0a1758e290986912751ba20e Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Wed, 19 Aug 2026 10:08:17 +0000 Subject: [PATCH 47/67] remove some redundancy assert --- torchtitan/models/kimi_k3/__init__.py | 1 - torchtitan/models/kimi_k3/model.py | 39 ++++++--------------- torchtitan/models/kimi_k3/parallelize.py | 4 --- torchtitan/models/kimi_k3/vision_encoder.py | 19 ---------- 4 files changed, 10 insertions(+), 53 deletions(-) diff --git a/torchtitan/models/kimi_k3/__init__.py b/torchtitan/models/kimi_k3/__init__.py index dd8756ab2e..170bf90322 100644 --- a/torchtitan/models/kimi_k3/__init__.py +++ b/torchtitan/models/kimi_k3/__init__.py @@ -457,7 +457,6 @@ def _kimi_k3_config( output_res_norm=_norm(dim), output_res_proj=_linear(dim, 1), vision_encoder=vision_encoder, - spatial_merge_size=2, ) diff --git a/torchtitan/models/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py index 19257427dc..b5e9614a02 100644 --- a/torchtitan/models/kimi_k3/model.py +++ b/torchtitan/models/kimi_k3/model.py @@ -228,7 +228,7 @@ class KimiKDAKernel(Module): The gate activation, the beta sigmoid, and the query/key L2 norm are all fused into the kernel rather than materialized here, so the decay never - exists as a full ``(B, L, H, K)`` tensor. + exists as a full ``(B, L, H, K)`` tensor. """ @dataclass(kw_only=True, slots=True) @@ -541,12 +541,10 @@ def forward( attention_masks: AttentionMasksType | None = None, positions: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: - # Keep the residual on every block output to preserve its FSDP gradient path. B, L, D = x_BLD.shape - prefix_sum_BLD: torch.Tensor | None = x_BLD + prefix_sum_BLD = x_BLD if block_residual_TND.shape[1] > 0: - assert prefix_sum_BLD is not None x_BLD = _apply_attention_residual( prefix_sum_BLD.reshape(-1, D), block_residual_TND, @@ -554,8 +552,8 @@ def forward( self.attention_res_norm, ).view(B, L, D) - if self.layer_id % self.attn_res_block_size == 0: - assert prefix_sum_BLD is not None + opens_block = self.layer_id % self.attn_res_block_size == 0 + if opens_block: block_residual_TND = torch.cat( ( block_residual_TND, @@ -563,7 +561,6 @@ def forward( ), dim=1, ) - prefix_sum_BLD = None h_BLD = self.attention_norm(x_BLD) if self.attention is not None: @@ -571,9 +568,8 @@ def forward( else: assert self.delta_attention is not None h_BLD = self.delta_attention(h_BLD, None, positions) - prefix_sum_BLD = h_BLD if prefix_sum_BLD is None else prefix_sum_BLD + h_BLD + prefix_sum_BLD = h_BLD if opens_block else prefix_sum_BLD + h_BLD - assert prefix_sum_BLD is not None h_BLD = _apply_attention_residual( prefix_sum_BLD.reshape(-1, D), block_residual_TND, @@ -596,7 +592,6 @@ class Config(Decoder.Config): output_res_norm: RMSNorm.Config output_res_proj: Linear.Config vision_encoder: KimiK3VisionEncoder.Config | None = None - spatial_merge_size: int = 2 def update_from_config(self, *, config, **kwargs) -> None: # Unsupported parallelisms are rejected in parallelize_kimi_k3. @@ -635,23 +630,6 @@ def __init__(self, config: Config): self.vision_encoder = ( config.vision_encoder.build() if config.vision_encoder is not None else None ) - self.spatial_merge_size = config.spatial_merge_size - if self.vision_encoder is not None: - # The decoder sizes each image's placeholder run from - # spatial_merge_size while the encoder merges patches with - # merge_kernel_size. A mismatch surfaces much later as a - # placeholder-run misalignment that blames the prompt. - merge_kernel_size = self.vision_encoder.merge_kernel_size - if merge_kernel_size != ( - config.spatial_merge_size, - config.spatial_merge_size, - ): - raise ValueError( - f"spatial_merge_size {config.spatial_merge_size} does not " - f"match the vision encoder's merge_kernel_size " - f"{merge_kernel_size}; each image would occupy a different " - "number of text positions than the encoder produces." - ) def _prepare_multimodal_embeds( self, @@ -677,8 +655,11 @@ def _prepare_multimodal_embeds( pixel_values = pixel_values.to(self.vision_encoder.patch_embed.weight.dtype) vision_embeds = self.vision_encoder(pixel_values, grid_thw=grid_thw) - num_tokens_per_item = (grid_thw[:, 1] // self.spatial_merge_size) * ( - grid_thw[:, 2] // self.spatial_merge_size + # MoonViT collapses time and merges spatially, so the text-side token + # count per item is (h/kh)*(w/kw), independent of t. + kernel_h, kernel_w = self.vision_encoder.merge_kernel_size + num_tokens_per_item = (grid_thw[:, 1] // kernel_h) * ( + grid_thw[:, 2] // kernel_w ) vision_positions = get_vision_positions( tokens, diff --git a/torchtitan/models/kimi_k3/parallelize.py b/torchtitan/models/kimi_k3/parallelize.py index 9101a5af03..c5955fe282 100644 --- a/torchtitan/models/kimi_k3/parallelize.py +++ b/torchtitan/models/kimi_k3/parallelize.py @@ -59,10 +59,6 @@ def parallelize_kimi_k3( raise NotImplementedError( "Kimi K3 FSDP2 does not support activation checkpointing yet." ) - if training.enable_cpu_offload: - raise NotImplementedError( - "Kimi K3 FSDP2 does not support parameter CPU offload yet." - ) dp_mesh_names = ( ["dp_replicate", "fsdp"] if parallel_dims.dp_replicate_enabled else ["fsdp"] diff --git a/torchtitan/models/kimi_k3/vision_encoder.py b/torchtitan/models/kimi_k3/vision_encoder.py index 688db3eb2c..3fb405109a 100644 --- a/torchtitan/models/kimi_k3/vision_encoder.py +++ b/torchtitan/models/kimi_k3/vision_encoder.py @@ -343,30 +343,11 @@ def forward( grid_thw: torch.Tensor, ) -> torch.Tensor: """Encode padded raster-order patches and return padded text features.""" - if grid_thw.ndim != 2 or grid_thw.shape[1] != 3: - raise ValueError(f"grid_thw must have shape (N, 3), got {grid_thw.shape}.") num_items, max_num_patches, _ = pixel_values.shape grids = grid_thw.tolist() - if len(grids) != num_items: - raise ValueError( - f"pixel_values contains {num_items} items but grid_thw " - f"contains {len(grids)}." - ) kernel_h, kernel_w = self.merge_kernel_size num_patches_N = grid_thw.prod(dim=-1).to(torch.long) - for num_frames, grid_h, grid_w in grids: - if grid_h % kernel_h != 0 or grid_w % kernel_w != 0: - raise ValueError( - f"Vision grid {grid_h}x{grid_w} is not divisible by " - f"merge kernel {self.merge_kernel_size}." - ) - item_num_patches = num_frames * grid_h * grid_w - if item_num_patches > max_num_patches: - raise ValueError( - f"Vision grid requires {item_num_patches} patches, but " - f"pixel_values only provides {max_num_patches}." - ) learned_pos, rope_cache = self._compute_position_embeddings( grids, max_num_patches From f925dad992121e393c41f545999d15b34a9e8472 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Wed, 19 Aug 2026 11:40:13 +0000 Subject: [PATCH 48/67] fix activation checkpointing --- torchtitan/models/kimi_k3/config_registry.py | 5 +++-- torchtitan/models/kimi_k3/parallelize.py | 13 +++++++------ 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/torchtitan/models/kimi_k3/config_registry.py b/torchtitan/models/kimi_k3/config_registry.py index afda14fd50..5a39fa00b3 100644 --- a/torchtitan/models/kimi_k3/config_registry.py +++ b/torchtitan/models/kimi_k3/config_registry.py @@ -9,7 +9,8 @@ from torchtitan.components.metrics import MetricsProcessor from torchtitan.components.optimizer import default_adamw, LRSchedulersContainer from torchtitan.components.tokenizer import MultiModalTokenizer -from torchtitan.config import TrainingConfig +from torchtitan.config import ParallelismConfig, TrainingConfig +from torchtitan.distributed.activation_checkpoint import SelectiveAC from torchtitan.hf_datasets.multimodal.mm_datasets import MMDataLoader from torchtitan.hf_datasets.multimodal.utils.image import resize_to_patch_budget from torchtitan.models.common.config_utils import decoder_vocab_size @@ -63,5 +64,5 @@ def kimi_k3_debugmodel() -> Trainer.Config: interval=10, last_save_model_only=False, ), - activation_checkpoint=None, + activation_checkpoint=SelectiveAC.Config(), ) diff --git a/torchtitan/models/kimi_k3/parallelize.py b/torchtitan/models/kimi_k3/parallelize.py index c5955fe282..e3232c747e 100644 --- a/torchtitan/models/kimi_k3/parallelize.py +++ b/torchtitan/models/kimi_k3/parallelize.py @@ -32,7 +32,6 @@ def parallelize_kimi_k3( dump_folder: str, ) -> nn.Module: """Apply FSDP2 to the Kimi K3 decoder and vision encoder.""" - del dump_folder unsupported_parallelisms = [ name @@ -54,11 +53,7 @@ def parallelize_kimi_k3( "Kimi K3 FSDP2 currently supports the default SPMD backend only." ) if compile_config.enable and "model" in compile_config.components: - raise NotImplementedError("Kimi K3 does not support model compilation.") - if ac_config is not None: - raise NotImplementedError( - "Kimi K3 FSDP2 does not support activation checkpointing yet." - ) + raise NotImplementedError("Kimi K3 does not support model compilation yet.") dp_mesh_names = ( ["dp_replicate", "fsdp"] if parallel_dims.dp_replicate_enabled else ["fsdp"] @@ -66,6 +61,12 @@ def parallelize_kimi_k3( dp_mesh = parallel_dims.get_mesh(dp_mesh_names) assert isinstance(model, KimiK3Model) + if ac_config is not None: + ac_policy = ac_config.build(dump_folder=dump_folder) + ac_policy.apply(model) + if model.vision_encoder is not None: + ac_policy.apply(model.vision_encoder) + vision_encoder = model.vision_encoder if vision_encoder is not None: # TODO: An image batch on one DP rank and a text-only batch on another From 950f547fa46d9665d3d2e5c035046c356fc97225 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Wed, 19 Aug 2026 16:50:28 +0000 Subject: [PATCH 49/67] Add TODO for _apply_attention_residual --- torchtitan/models/kimi_k3/model.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/torchtitan/models/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py index b5e9614a02..bb2cf826f4 100644 --- a/torchtitan/models/kimi_k3/model.py +++ b/torchtitan/models/kimi_k3/model.py @@ -469,9 +469,7 @@ def _apply_attention_residual( ) -> torch.Tensor: """Apply Kimi's block-level attention residual in FP32. - The norm and projection weights are folded into a single score vector, so - this reads them directly instead of calling the modules. That is only valid - while both are replicated; sharding them would need a DTensor-aware path. + TODO: Add TP Support. The current implementation assumes that the input tensors are on a single device. """ assert norm.eps is not None From 6689074eae0b861a2c405cac87c6094f2e87ad4c Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Thu, 20 Aug 2026 14:01:01 +0000 Subject: [PATCH 50/67] split k3 model.py to moe.py lda.py --- tests/unit_tests/test_kimi_k3.py | 5 +- torchtitan/models/kimi_k3/__init__.py | 14 +- torchtitan/models/kimi_k3/kda.py | 175 ++++++++++++ torchtitan/models/kimi_k3/model.py | 346 +---------------------- torchtitan/models/kimi_k3/moe.py | 144 ++++++++++ torchtitan/models/kimi_k3/parallelize.py | 2 +- 6 files changed, 332 insertions(+), 354 deletions(-) create mode 100644 torchtitan/models/kimi_k3/kda.py create mode 100644 torchtitan/models/kimi_k3/moe.py diff --git a/tests/unit_tests/test_kimi_k3.py b/tests/unit_tests/test_kimi_k3.py index bdf055556d..932a89afff 100644 --- a/tests/unit_tests/test_kimi_k3.py +++ b/tests/unit_tests/test_kimi_k3.py @@ -26,7 +26,8 @@ _vision_encoder_config, parallelize_kimi_k3, ) -from torchtitan.models.kimi_k3.model import KimiK3Model, KimiKDAKernel +from torchtitan.models.kimi_k3.kda import KimiKDAKernel +from torchtitan.models.kimi_k3.model import KimiK3Model from torchtitan.models.kimi_k3.state_dict_adapter import KimiK3StateDictAdapter @@ -264,7 +265,7 @@ def test_fsdp_matches_non_distributed_forward_backward(self): torch.manual_seed(3) config = _small_model_config( attn_res_block_size=2, - full_attention_layers={0, 1}, + full_attention_layers={1}, ) with torch.device("meta"): model = config.build() diff --git a/torchtitan/models/kimi_k3/__init__.py b/torchtitan/models/kimi_k3/__init__.py index 170bf90322..ae480794cb 100644 --- a/torchtitan/models/kimi_k3/__init__.py +++ b/torchtitan/models/kimi_k3/__init__.py @@ -21,17 +21,9 @@ from torchtitan.protocols.model import ModelConfigConverter from torchtitan.protocols.model_spec import ModelSpec -from .model import ( - KimiDeltaAttention, - KimiFeedForward, - KimiGroupedExperts, - KimiK3Model, - KimiK3TransformerBlock, - KimiKDAKernel, - KimiLatentMoE, - KimiMLAAttention, - KimiRMSNormGated, -) +from .kda import KimiDeltaAttention, KimiKDAKernel, KimiRMSNormGated +from .model import KimiK3Model, KimiK3TransformerBlock, KimiMLAAttention +from .moe import KimiFeedForward, KimiGroupedExperts, KimiLatentMoE from .parallelize import parallelize_kimi_k3 from .state_dict_adapter import KimiK3StateDictAdapter from .vision_encoder import ( diff --git a/torchtitan/models/kimi_k3/kda.py b/torchtitan/models/kimi_k3/kda.py new file mode 100644 index 0000000000..f52c560d7c --- /dev/null +++ b/torchtitan/models/kimi_k3/kda.py @@ -0,0 +1,175 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Kimi Delta Attention modules for Kimi K3.""" + +from dataclasses import dataclass + +import torch +import torch.nn.functional as F +from fla.ops.kda import chunk_kda +from torch import nn + +from torchtitan.models.common import Conv1d, Linear +from torchtitan.models.common.attention import AttentionMasksType +from torchtitan.protocols.module import Module + +# Shape suffixes: +# B = batch, L = sequence length, D = model dimension, H = heads, +# K = key head dimension, V = value head dimension, C = projection channels. + + +class KimiRMSNormGated(Module): + """Per-head RMSNorm followed by a sigmoid output gate.""" + + @dataclass(kw_only=True, slots=True) + class Config(Module.Config): + dim: int + eps: float = 1e-5 + + def __init__(self, config: Config): + super().__init__() + self.eps = config.eps + self.weight = nn.Parameter(torch.empty(config.dim)) + + def forward(self, x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: + input_dtype = x.dtype + x_float = x.float() + variance = x_float.pow(2).mean(dim=-1, keepdim=True) + x_float = x_float * torch.rsqrt(variance + self.eps) + x_float = self.weight.float() * x_float + return (x_float * torch.sigmoid(gate.float())).to(input_dtype) + + +class KimiKDAKernel(Module): + """Stateless dispatch to FLA's chunked KDA kernel.""" + + @dataclass(kw_only=True, slots=True) + class Config(Module.Config): + lower_bound: float | None = -5.0 + + def __init__(self, config: Config): + super().__init__() + self.lower_bound = config.lower_bound + if self.lower_bound is not None and not (-5.0 <= self.lower_bound < 0.0): + raise ValueError("KDA lower_bound must be in the safe range [-5, 0).") + + def forward( + self, + q_BLHK: torch.Tensor, + k_BLHK: torch.Tensor, + v_BLHV: torch.Tensor, + gate_BLHK: torch.Tensor, + beta_BLH: torch.Tensor, + A_log_H: torch.Tensor, + dt_bias_HK: torch.Tensor, + ) -> torch.Tensor: + out_BLHV, _ = chunk_kda( + q_BLHK, + k_BLHK, + v_BLHV, + gate_BLHK, + beta_BLH, + A_log=A_log_H, + dt_bias=dt_bias_HK.reshape(-1), + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=True, + safe_gate=self.lower_bound is not None, + lower_bound=self.lower_bound, + ) + return out_BLHV + + +class KimiDeltaAttention(Module): + @dataclass(kw_only=True, slots=True) + class Config(Module.Config): + dim: int + num_heads: int + head_dim: int + conv_kernel_size: int + q_proj: Linear.Config + k_proj: Linear.Config + v_proj: Linear.Config + q_conv: Conv1d.Config + k_conv: Conv1d.Config + v_conv: Conv1d.Config + forget_a: Linear.Config + forget_b: Linear.Config + beta: Linear.Config + output_gate: Linear.Config + kernel: Module.Config + output_norm: KimiRMSNormGated.Config + output_proj: Linear.Config + + def __init__(self, config: Config): + super().__init__() + self.num_heads = config.num_heads + self.head_dim = config.head_dim + self.conv_kernel_size = config.conv_kernel_size + + self.q_proj = config.q_proj.build() + self.k_proj = config.k_proj.build() + self.v_proj = config.v_proj.build() + self.q_conv = config.q_conv.build() + self.k_conv = config.k_conv.build() + self.v_conv = config.v_conv.build() + self.forget_a = config.forget_a.build() + self.forget_b = config.forget_b.build() + self.beta = config.beta.build() + self.output_gate = config.output_gate.build() + self.kernel = config.kernel.build() + self.output_norm = config.output_norm.build() + self.output_proj = config.output_proj.build() + + self.A_log = nn.Parameter(torch.empty(config.num_heads)) + self.dt_bias = nn.Parameter(torch.empty(config.num_heads, config.head_dim)) + + def _causal_conv(self, x_BLC: torch.Tensor, conv: Conv1d) -> torch.Tensor: + x_BCL = F.pad(x_BLC.transpose(1, 2), (self.conv_kernel_size - 1, 0)) + return F.silu(conv(x_BCL)).transpose(1, 2) + + def forward( + self, + x_BLD: torch.Tensor, + attention_masks: AttentionMasksType | None = None, + positions: torch.Tensor | None = None, + ) -> torch.Tensor: + del positions + if attention_masks is not None: + raise NotImplementedError( + "Kimi K3 reference KDA does not support packed-document masks." + ) + + B, L, _ = x_BLD.shape + q_BLHK = self._causal_conv(self.q_proj(x_BLD), self.q_conv).view( + B, L, self.num_heads, self.head_dim + ) + k_BLHK = self._causal_conv(self.k_proj(x_BLD), self.k_conv).view( + B, L, self.num_heads, self.head_dim + ) + v_BLHV = self._causal_conv(self.v_proj(x_BLD), self.v_conv).view( + B, L, self.num_heads, self.head_dim + ) + forget_BLHK = self.forget_b(self.forget_a(x_BLD)).view( + B, L, self.num_heads, self.head_dim + ) + beta_BLH = self.beta(x_BLD).float() + + out_BLHV = self.kernel( + q_BLHK, + k_BLHK, + v_BLHV, + forget_BLHK, + beta_BLH, + self.A_log, + self.dt_bias, + ) + output_gate_BLHV = self.output_gate(x_BLD).view( + B, L, self.num_heads, self.head_dim + ) + out_BLHV = self.output_norm(out_BLHV, output_gate_BLHV) + return self.output_proj(out_BLHV.reshape(B, L, -1)) diff --git a/torchtitan/models/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py index bb2cf826f4..7e45ff0415 100644 --- a/torchtitan/models/kimi_k3/model.py +++ b/torchtitan/models/kimi_k3/model.py @@ -7,129 +7,33 @@ from dataclasses import dataclass, field import torch -import torch.nn.functional as F - -from fla.ops.kda import chunk_kda from torch import nn -from torch.distributed.tensor import DTensor -from torchtitan.models.common import Conv1d, Linear +from torchtitan.models.common import Linear from torchtitan.models.common.attention import ( AttentionMasksType, BaseAttention, FlexAttention, ) from torchtitan.models.common.decoder import Decoder -from torchtitan.models.common.feed_forward import FeedForward -from torchtitan.models.common.moe import GroupedExperts, MoE from torchtitan.models.common.multimodal import ( get_vision_positions, scatter_vision_embeds, ) from torchtitan.models.common.nn_modules import RMSNorm -from torchtitan.models.kimi_k3.vision_encoder import KimiK3VisionEncoder from torchtitan.models.utils import get_moe_model_nparams_and_flops from torchtitan.protocols.module import Module +from .kda import KimiDeltaAttention +from .moe import KimiFeedForward, KimiLatentMoE +from .vision_encoder import KimiK3VisionEncoder + # Shape suffixes: # B = batch, L = sequence length, D = model dimension, H = heads, -# K = key head dimension, V = value head dimension, E = experts, -# C = projection channels, F = expert hidden dimension, R = routed tokens, -# S = selected experts per token, T = flattened tokens, +# K = key head dimension, V = value head dimension, T = flattened tokens, # N = attention-residual entries. -class KimiShortConvolution(ShortConvolution, Module): - """KDA short causal convolution backed by FLA's fused kernel. - - Matches the released Kimi K3 HuggingFace model, which builds FLA's - ``ShortConvolution`` per q/k/v projection. The Triton kernel runs only on - accelerator devices. - """ - - @dataclass(kw_only=True, slots=True) - class Config(Module.Config): - hidden_size: int - kernel_size: int - activation: str = "silu" - - def __init__(self, config: Config): - super().__init__( - hidden_size=config.hidden_size, - kernel_size=config.kernel_size, - activation=config.activation, - ) - - def forward( - self, - x_BLD: torch.Tensor, - **kwargs: object, - ) -> tuple[torch.Tensor, None]: - y_BLD, _ = causal_conv1d( - x=x_BLD, - weight=self.weight.squeeze(1), - activation=self.activation, - backend=self.backend, - ) - return y_BLD, None - -class KimiRMSNormGated(Module): - """Per-head RMSNorm followed by a sigmoid output gate.""" - - @dataclass(kw_only=True, slots=True) - class Config(Module.Config): - dim: int - eps: float = 1e-5 - - def __init__(self, config: Config): - super().__init__() - self.eps = config.eps - self.weight = nn.Parameter(torch.empty(config.dim)) - - def forward(self, x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: - input_dtype = x.dtype - x_float = x.float() - variance = x_float.pow(2).mean(dim=-1, keepdim=True) - x_float = x_float * torch.rsqrt(variance + self.eps) - x_float = self.weight.float() * x_float - return (x_float * torch.sigmoid(gate.float())).to(input_dtype) - - -def _situ_glu( - gate: torch.Tensor, - up: torch.Tensor, - beta: float, - linear_beta: float | None, -) -> torch.Tensor: - """Kimi's SiTU-GLU activation, evaluated in FP32.""" - input_dtype = gate.dtype - gate = gate.float() - up = up.float() - gate = beta * torch.tanh(gate / beta) * torch.sigmoid(gate) - if linear_beta is not None: - up = linear_beta * torch.tanh(up / linear_beta) - return (gate * up).to(input_dtype) - - -class KimiFeedForward(FeedForward): - """FeedForward with Kimi's SiTU activation""" - - @dataclass(kw_only=True, slots=True) - class Config(FeedForward.Config): - beta: float = 1.0 - linear_beta: float | None = None - - def __init__(self, config: Config): - super().__init__(config) - self.beta = config.beta - self.linear_beta = config.linear_beta - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.w2( - _situ_glu(self.w1(x), self.w3(x), self.beta, self.linear_beta), - ) - - class KimiMLAAttention(BaseAttention): """Kimi K3 multi-head latent attention. @@ -223,244 +127,6 @@ def forward( return self.wo(out_BLD) -class KimiKDAKernel(Module): - """Stateless dispatch to FLA's chunked KDA kernel. - - The gate activation, the beta sigmoid, and the query/key L2 norm are all - fused into the kernel rather than materialized here, so the decay never - exists as a full ``(B, L, H, K)`` tensor. - """ - - @dataclass(kw_only=True, slots=True) - class Config(Module.Config): - lower_bound: float | None = -5.0 - - def __init__(self, config: Config): - super().__init__() - self.lower_bound = config.lower_bound - if self.lower_bound is not None and not (-5.0 <= self.lower_bound < 0.0): - raise ValueError("KDA lower_bound must be in the safe range [-5, 0).") - - def forward( - self, - q_BLHK: torch.Tensor, - k_BLHK: torch.Tensor, - v_BLHV: torch.Tensor, - gate_BLHK: torch.Tensor, - beta_BLH: torch.Tensor, - A_log_H: torch.Tensor, - dt_bias_HK: torch.Tensor, - ) -> torch.Tensor: - # safe_gate selects the bounded gate activation - # lower_bound * sigmoid(exp(A_log) * (gate + dt_bias)); without it the - # kernel applies -exp(A_log) * softplus(gate + dt_bias). - out_BLHV, _ = chunk_kda( - q_BLHK, - k_BLHK, - v_BLHV, - gate_BLHK, - beta_BLH, - A_log=A_log_H, - dt_bias=dt_bias_HK.reshape(-1), - use_qk_l2norm_in_kernel=True, - use_gate_in_kernel=True, - use_beta_sigmoid_in_kernel=True, - safe_gate=self.lower_bound is not None, - lower_bound=self.lower_bound, - ) - return out_BLHV - - -class KimiDeltaAttention(Module): - @dataclass(kw_only=True, slots=True) - class Config(Module.Config): - dim: int - num_heads: int - head_dim: int - conv_kernel_size: int - q_proj: Linear.Config - k_proj: Linear.Config - v_proj: Linear.Config - q_conv: Conv1d.Config - k_conv: Conv1d.Config - v_conv: Conv1d.Config - forget_a: Linear.Config - forget_b: Linear.Config - beta: Linear.Config - output_gate: Linear.Config - kernel: Module.Config - output_norm: KimiRMSNormGated.Config - output_proj: Linear.Config - - def __init__(self, config: Config): - super().__init__() - self.num_heads = config.num_heads - self.head_dim = config.head_dim - self.conv_kernel_size = config.conv_kernel_size - - self.q_proj = config.q_proj.build() - self.k_proj = config.k_proj.build() - self.v_proj = config.v_proj.build() - self.q_conv = config.q_conv.build() - self.k_conv = config.k_conv.build() - self.v_conv = config.v_conv.build() - self.forget_a = config.forget_a.build() - self.forget_b = config.forget_b.build() - self.beta = config.beta.build() - self.output_gate = config.output_gate.build() - self.kernel = config.kernel.build() - self.output_norm = config.output_norm.build() - self.output_proj = config.output_proj.build() - - self.A_log = nn.Parameter(torch.empty(config.num_heads)) - self.dt_bias = nn.Parameter(torch.empty(config.num_heads, config.head_dim)) - - def _causal_conv(self, x_BLC: torch.Tensor, conv: Conv1d) -> torch.Tensor: - x_BCL = F.pad(x_BLC.transpose(1, 2), (self.conv_kernel_size - 1, 0)) - return F.silu(conv(x_BCL)).transpose(1, 2) - - def forward( - self, - x_BLD: torch.Tensor, - attention_masks: AttentionMasksType | None = None, - positions: torch.Tensor | None = None, - ) -> torch.Tensor: - del positions - if attention_masks is not None: - raise NotImplementedError( - "Kimi K3 reference KDA does not support packed-document masks." - ) - - B, L, _ = x_BLD.shape - q_BLHK = self._causal_conv(self.q_proj(x_BLD), self.q_conv).view( - B, L, self.num_heads, self.head_dim - ) - k_BLHK = self._causal_conv(self.k_proj(x_BLD), self.k_conv).view( - B, L, self.num_heads, self.head_dim - ) - v_BLHV = self._causal_conv(self.v_proj(x_BLD), self.v_conv).view( - B, L, self.num_heads, self.head_dim - ) - forget_BLHK = self.forget_b(self.forget_a(x_BLD)).view( - B, L, self.num_heads, self.head_dim - ) - beta_BLH = self.beta(x_BLD).float() - - out_BLHV = self.kernel( - q_BLHK, - k_BLHK, - v_BLHV, - forget_BLHK, - beta_BLH, - self.A_log, - self.dt_bias, - ) - output_gate_BLHV = self.output_gate(x_BLD).view( - B, L, self.num_heads, self.head_dim - ) - out_BLHV = self.output_norm(out_BLHV, output_gate_BLHV) - return self.output_proj(out_BLHV.reshape(B, L, -1)) - - -class KimiGroupedExperts(GroupedExperts): - """``common/moe.py::GroupedExperts`` with Kimi's SiTU activation. - - Inherits its stacked-weight shape (``w1_EFD``/``w2_EDF``/``w3_EFD``) and - parameter allocation; only ``forward`` differs, since the activation is - baked into the ``_grouped_mm`` call sequence rather than being a - swappable argument. - """ - - @dataclass(kw_only=True, slots=True) - class Config(GroupedExperts.Config): - beta: float = 1.0 - linear_beta: float | None = None - - def __init__(self, config: Config): - super().__init__(config) - self.beta = config.beta - self.linear_beta = config.linear_beta - - def forward( - self, - x_RD: torch.Tensor, - num_tokens_per_expert_E: torch.Tensor, - ) -> torch.Tensor: - if isinstance(self.w1_EFD, DTensor): - w1_EFD = self.w1_EFD.to_local() - assert isinstance(self.w2_EDF, DTensor) - w2_EDF = self.w2_EDF.to_local() - assert isinstance(self.w3_EFD, DTensor) - w3_EFD = self.w3_EFD.to_local() - else: - w1_EFD = self.w1_EFD - w2_EDF = self.w2_EDF - w3_EFD = self.w3_EFD - - offsets_E = torch.cumsum(num_tokens_per_expert_E, dim=0, dtype=torch.int32) - - gate_RF = self._grouped_mm( - A=x_RD.bfloat16(), - B_t=w1_EFD.bfloat16().transpose(-2, -1), - offs=offsets_E, - ) - up_RF = self._grouped_mm( - A=x_RD.bfloat16(), - B_t=w3_EFD.bfloat16().transpose(-2, -1), - offs=offsets_E, - ) - - h_RF = _situ_glu(gate_RF, up_RF, self.beta, self.linear_beta) - - return self._grouped_mm( - A=h_RF, - B_t=w2_EDF.bfloat16().transpose(-2, -1), - offs=offsets_E, - ).type_as(x_RD) - - -class KimiLatentMoE(MoE): - """``common/moe.py::MoE`` with Kimi's latent routed-expert path. - - Routed tokens are projected down to the expert latent width, run through - the experts, then normed and projected back up to the model dimension. - Shared experts still see the full-width input. - """ - - @dataclass(kw_only=True, slots=True) - class Config(MoE.Config): - routed_down: Linear.Config - routed_norm: RMSNorm.Config - routed_up: Linear.Config - - def __init__(self, config: Config): - super().__init__(config) - self.routed_down = config.routed_down.build() - self.routed_norm = config.routed_norm.build() - self.routed_up = config.routed_up.build() - - def forward(self, x_BLD: torch.Tensor) -> torch.Tensor: - weights_BLS, expert_ids_BLS, scores_BLE = self.router(x_BLD, self.expert_bias_E) - routing_map_BLE = torch.zeros_like(scores_BLE, dtype=torch.bool).scatter_( - -1, expert_ids_BLS, True - ) - num_tokens_per_expert_E = routing_map_BLE.sum(dim=(0, 1)) - if self.training: - with torch.no_grad(): - self.tokens_per_expert_E.add_(num_tokens_per_expert_E) - - routed_BLD = self.routed_experts( - self.routed_down(x_BLD), - weights_BLS, - expert_ids_BLS, - num_tokens_per_expert_E, - ) - out_BLD = self.routed_up(self.routed_norm(routed_BLD)) - if self.shared_experts is not None: - out_BLD = out_BLD + self.shared_experts(x_BLD) - return out_BLD - - def _apply_attention_residual( prefix_sum_TD: torch.Tensor, block_residual_TND: torch.Tensor, diff --git a/torchtitan/models/kimi_k3/moe.py b/torchtitan/models/kimi_k3/moe.py new file mode 100644 index 0000000000..0c9616aa31 --- /dev/null +++ b/torchtitan/models/kimi_k3/moe.py @@ -0,0 +1,144 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""SiTU feed-forward and latent MoE modules for Kimi K3.""" + +from dataclasses import dataclass + +import torch +from torch.distributed.tensor import DTensor + +from torchtitan.models.common import Linear +from torchtitan.models.common.feed_forward import FeedForward +from torchtitan.models.common.moe import GroupedExperts, MoE +from torchtitan.models.common.nn_modules import RMSNorm + +# Shape suffixes: +# B = batch, L = sequence length, D = model dimension, E = experts, +# F = expert hidden dimension, R = routed tokens, S = selected experts per token. + + +def _situ_glu( + gate: torch.Tensor, + up: torch.Tensor, + beta: float, + linear_beta: float | None, +) -> torch.Tensor: + """Kimi's SiTU-GLU activation, evaluated in FP32.""" + input_dtype = gate.dtype + gate = gate.float() + up = up.float() + gate = beta * torch.tanh(gate / beta) * torch.sigmoid(gate) + if linear_beta is not None: + up = linear_beta * torch.tanh(up / linear_beta) + return (gate * up).to(input_dtype) + + +class KimiFeedForward(FeedForward): + """FeedForward with Kimi's SiTU activation.""" + + @dataclass(kw_only=True, slots=True) + class Config(FeedForward.Config): + beta: float = 1.0 + linear_beta: float | None = None + + def __init__(self, config: Config): + super().__init__(config) + self.beta = config.beta + self.linear_beta = config.linear_beta + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.w2( + _situ_glu(self.w1(x), self.w3(x), self.beta, self.linear_beta), + ) + + +class KimiGroupedExperts(GroupedExperts): + """``common/moe.py::GroupedExperts`` with Kimi's SiTU activation.""" + + @dataclass(kw_only=True, slots=True) + class Config(GroupedExperts.Config): + beta: float = 1.0 + linear_beta: float | None = None + + def __init__(self, config: Config): + super().__init__(config) + self.beta = config.beta + self.linear_beta = config.linear_beta + + def forward( + self, + x_RD: torch.Tensor, + num_tokens_per_expert_E: torch.Tensor, + ) -> torch.Tensor: + if isinstance(self.w1_EFD, DTensor): + w1_EFD = self.w1_EFD.to_local() + assert isinstance(self.w2_EDF, DTensor) + w2_EDF = self.w2_EDF.to_local() + assert isinstance(self.w3_EFD, DTensor) + w3_EFD = self.w3_EFD.to_local() + else: + w1_EFD = self.w1_EFD + w2_EDF = self.w2_EDF + w3_EFD = self.w3_EFD + + offsets_E = torch.cumsum(num_tokens_per_expert_E, dim=0, dtype=torch.int32) + + gate_RF = self._grouped_mm( + A=x_RD.bfloat16(), + B_t=w1_EFD.bfloat16().transpose(-2, -1), + offs=offsets_E, + ) + up_RF = self._grouped_mm( + A=x_RD.bfloat16(), + B_t=w3_EFD.bfloat16().transpose(-2, -1), + offs=offsets_E, + ) + + h_RF = _situ_glu(gate_RF, up_RF, self.beta, self.linear_beta) + + return self._grouped_mm( + A=h_RF, + B_t=w2_EDF.bfloat16().transpose(-2, -1), + offs=offsets_E, + ).type_as(x_RD) + + +class KimiLatentMoE(MoE): + """``common/moe.py::MoE`` with Kimi's latent routed-expert path.""" + + @dataclass(kw_only=True, slots=True) + class Config(MoE.Config): + routed_down: Linear.Config + routed_norm: RMSNorm.Config + routed_up: Linear.Config + + def __init__(self, config: Config): + super().__init__(config) + self.routed_down = config.routed_down.build() + self.routed_norm = config.routed_norm.build() + self.routed_up = config.routed_up.build() + + def forward(self, x_BLD: torch.Tensor) -> torch.Tensor: + weights_BLS, expert_ids_BLS, scores_BLE = self.router(x_BLD, self.expert_bias_E) + routing_map_BLE = torch.zeros_like(scores_BLE, dtype=torch.bool).scatter_( + -1, expert_ids_BLS, True + ) + num_tokens_per_expert_E = routing_map_BLE.sum(dim=(0, 1)) + if self.training: + with torch.no_grad(): + self.tokens_per_expert_E.add_(num_tokens_per_expert_E) + + routed_BLD = self.routed_experts( + self.routed_down(x_BLD), + weights_BLS, + expert_ids_BLS, + num_tokens_per_expert_E, + ) + out_BLD = self.routed_up(self.routed_norm(routed_BLD)) + if self.shared_experts is not None: + out_BLD = out_BLD + self.shared_experts(x_BLD) + return out_BLD diff --git a/torchtitan/models/kimi_k3/parallelize.py b/torchtitan/models/kimi_k3/parallelize.py index e3232c747e..d0c91f6609 100644 --- a/torchtitan/models/kimi_k3/parallelize.py +++ b/torchtitan/models/kimi_k3/parallelize.py @@ -18,7 +18,7 @@ apply_fsdp_to_decoder, apply_fsdp_to_vision_encoder, ) -from torchtitan.models.kimi_k3.model import KimiK3Model +from .model import KimiK3Model def parallelize_kimi_k3( From 252ae22089b22cebcdb05dbafbc225e97546f391 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Fri, 21 Aug 2026 05:50:21 +0000 Subject: [PATCH 51/67] remove two unsed attn res from hf, add it in state_dict_adapter --- torchtitan/models/kimi_k3/__init__.py | 4 ++-- torchtitan/models/kimi_k3/model.py | 16 +++++++++---- .../models/kimi_k3/state_dict_adapter.py | 24 +++++++++++++++++++ 3 files changed, 38 insertions(+), 6 deletions(-) diff --git a/torchtitan/models/kimi_k3/__init__.py b/torchtitan/models/kimi_k3/__init__.py index ae480794cb..831bdb352f 100644 --- a/torchtitan/models/kimi_k3/__init__.py +++ b/torchtitan/models/kimi_k3/__init__.py @@ -424,8 +424,8 @@ def _kimi_k3_config( ), attention_norm=_norm(dim), ffn_norm=_norm(dim), - attention_res_norm=_norm(dim), - attention_res_proj=_linear(dim, 1), + attention_res_norm=None if layer_idx == 0 else _norm(dim), + attention_res_proj=None if layer_idx == 0 else _linear(dim, 1), ffn_res_norm=_norm(dim), ffn_res_proj=_linear(dim, 1), ) diff --git a/torchtitan/models/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py index 7e45ff0415..10c2d8ec54 100644 --- a/torchtitan/models/kimi_k3/model.py +++ b/torchtitan/models/kimi_k3/model.py @@ -163,8 +163,8 @@ class Config(Module.Config): moe: KimiLatentMoE.Config | None attention_norm: RMSNorm.Config ffn_norm: RMSNorm.Config - attention_res_norm: RMSNorm.Config - attention_res_proj: Linear.Config + attention_res_norm: RMSNorm.Config | None + attention_res_proj: Linear.Config | None ffn_res_norm: RMSNorm.Config ffn_res_proj: Linear.Config @@ -193,8 +193,16 @@ def __init__(self, config: Config): self.moe_enabled = self.moe is not None self.attention_norm = config.attention_norm.build() self.ffn_norm = config.ffn_norm.build() - self.attention_res_norm = config.attention_res_norm.build() - self.attention_res_proj = config.attention_res_proj.build() + self.attention_res_norm = ( + config.attention_res_norm.build() + if config.attention_res_norm is not None + else None + ) + self.attention_res_proj = ( + config.attention_res_proj.build() + if config.attention_res_proj is not None + else None + ) self.ffn_res_norm = config.ffn_res_norm.build() self.ffn_res_proj = config.ffn_res_proj.build() diff --git a/torchtitan/models/kimi_k3/state_dict_adapter.py b/torchtitan/models/kimi_k3/state_dict_adapter.py index 334e2d07d5..ca28b3851c 100644 --- a/torchtitan/models/kimi_k3/state_dict_adapter.py +++ b/torchtitan/models/kimi_k3/state_dict_adapter.py @@ -17,6 +17,12 @@ from .model import KimiK3Model +_UNUSED_HF_LAYER_ZERO_ATTN_RES_KEYS = { + "language_model.model.layers.0.self_attention_res_norm.weight", + "language_model.model.layers.0.self_attention_res_proj.weight", +} + + class KimiK3StateDictAdapter(MoEStateDictAdapter): def __init__( self, @@ -235,6 +241,22 @@ def to_hf(self, state_dict: dict[str, Any]) -> dict[str, Any]: f"vision_tower.encoder.blocks.{layer_num}.wqkv.weight" ] = torch.cat((qkv["q"], qkv["k"], qkv["v"]), dim=0) + # The released HF model contain these unused layer-0 attn res parameters. + # TT omits them, so synthesize deterministic, placeholders to preserve strict HF state-dict loading. + if self.kimi_config.layers[0].attention_res_norm is None: + norm_template_key = ( + "language_model.model.layers.1.self_attention_res_norm.weight" + ) + proj_template_key = ( + "language_model.model.layers.1.self_attention_res_proj.weight" + ) + hf_state_dict[ + "language_model.model.layers.0.self_attention_res_norm.weight" + ] = torch.ones_like(hf_state_dict[norm_template_key]) + hf_state_dict[ + "language_model.model.layers.0.self_attention_res_proj.weight" + ] = torch.zeros_like(hf_state_dict[proj_template_key]) + if unmapped: raise ValueError( "KimiK3StateDictAdapter found TorchTitan keys without a " @@ -249,6 +271,8 @@ def from_hf(self, hf_state_dict: dict[str, Any]) -> dict[str, Any]: unmapped: list[str] = [] for key, value in hf_state_dict.items(): + if key in _UNUSED_HF_LAYER_ZERO_ATTN_RES_KEYS: + continue if key.endswith("rotary_emb.inv_freq"): continue From 2ef19811278797c9e1a08adc7b72599b53858966 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Sat, 22 Aug 2026 07:03:18 +0000 Subject: [PATCH 52/67] numerical_tests_kimi_k3.py support different attn_backend and bf16 --- .../numerical_tests_kimi_k3.py | 87 +++++++++++++------ torchtitan/models/kimi_k3/__init__.py | 40 +++++---- 2 files changed, 80 insertions(+), 47 deletions(-) diff --git a/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py b/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py index e089617cbc..68cc24c4af 100644 --- a/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py +++ b/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py @@ -20,9 +20,16 @@ The released code requires ``transformers==4.56.2`` and ``tiktoken``. Usage: + # BF16: HF FlashAttention2 / TT FlexAttention + CUDA_VISIBLE_DEVICES=0 python -m \ + scripts.checkpoint_conversion.numerical_tests_kimi_k3 + + # FP32: HF eager / TT FlexAttention CUDA_VISIBLE_DEVICES=0 python -m \ scripts.checkpoint_conversion.numerical_tests_kimi_k3 \ --dtype float32 + +Add ``--force-hf-routing`` to either command for the routing-fixed diagnostic. """ import argparse @@ -143,8 +150,6 @@ def _reduce_hf_config(hf_config, tt_config, hf_model_path: str) -> None: for config in (hf_config, text_config): if hasattr(config, "quantization_config"): delattr(config, "quantization_config") - text_config._attn_implementation = "eager" - hf_config.vision_config._attn_implementation = "eager" text_config._name_or_path = hf_model_path hf_config._name_or_path = hf_model_path @@ -170,8 +175,11 @@ def _build_hf_model( local_files_only=True, ) _reduce_hf_config(hf_config, tt_config, hf_model_path) + attn_backend = "flash_attention_2" if dtype == torch.bfloat16 else "eager" + hf_config.text_config._attn_implementation = attn_backend + hf_config.vision_config._attn_implementation = attn_backend model = AutoModelForCausalLM.from_config(hf_config, trust_remote_code=True) - model.language_model.config._attn_implementation = "eager" + model.language_model.config._attn_implementation = attn_backend model.to(dtype=dtype) model.load_state_dict(hf_state_dict, strict=True) return model.eval() @@ -232,6 +240,7 @@ def record_vision_features(_module, _inputs, output) -> None: key: value.to(device) if isinstance(value, torch.Tensor) else value for key, value in batch.items() } + inputs["pixel_values"] = inputs["pixel_values"].to(dtype) output = model(**inputs, use_cache=False) ref = { "input_ids": batch["input_ids"].cpu(), @@ -298,12 +307,37 @@ def _print_routing_comparison( print(f"router choices: {num_matching}/{num_routings} match " f"({match_rate:.1%})") +def _force_hf_routing(model, expert_indices, device) -> None: + """Use HF expert IDs with TorchTitan's independently computed scores.""" + for layer_idx, layer in model.layers.items(): + if (moe := cast(Any, layer.moe)) is None: + continue + ids = expert_indices[int(layer_idx)].unsqueeze(0).to(device) + router, original = moe.router, moe.router.forward + + def forced_forward( + x_BLD, + expert_bias_E=None, + _router=router, + _original=original, + _ids=ids, + ): + _, _, scores_BLE = _original(x_BLD, expert_bias_E) + weights = scores_BLE.gather(dim=-1, index=_ids) + if _router.route_norm: + weights = weights / (weights.sum(dim=-1, keepdim=True) + 1e-20) + return weights * _router.route_scale, _ids, scores_BLE + + router.forward = forced_forward + + @torch.no_grad() def run_tt( model: KimiK3Model, ref: dict[str, Any], vision_dtype: torch.dtype, device: torch.device, + force_hf_routing: bool, ) -> torch.Tensor: """Run TorchTitan preprocessing and the reduced TorchTitan model.""" print(f"Loading TorchTitan Kimi K3 (debugmodel) on {device} ...") @@ -311,6 +345,10 @@ def run_tt( assert model.vision_encoder is not None model.vision_encoder.to(vision_dtype) + if force_hf_routing: + print("Using HF expert selections with TorchTitan router scores") + _force_hf_routing(model, ref["expert_indices"], device) + expert_indices: dict[int, torch.Tensor] = {} for layer_idx, layer in model.layers.items(): if layer.moe is not None: @@ -394,8 +432,8 @@ def run_tt( return logits[:, -1, :].float().cpu().squeeze() -def compare(ref_logits: torch.Tensor, tt_logits: torch.Tensor) -> bool: - """Print last-token metrics and return whether KL is below tolerance.""" +def compare(ref_logits: torch.Tensor, tt_logits: torch.Tensor) -> None: + """Print last-token parity metrics.""" ref = ref_logits.squeeze() tt = tt_logits.squeeze() log_ref = F.log_softmax(ref, dim=-1) @@ -412,9 +450,6 @@ def compare(ref_logits: torch.Tensor, tt_logits: torch.Tensor) -> bool: f"KL={kl:.4e} cos={cosine:.6f} max_diff={max_diff:.4e} " f"top1={'Y' if top1 else 'N'} top5={top5_overlap:.0%}" ) - passed = abs(kl) < 1e-3 # pyrefly: ignore [bad-argument-type] - print("RESULT: PASS" if passed else "RESULT: FAIL") - return passed @torch.no_grad() @@ -422,20 +457,15 @@ def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--model_flavor", default="debugmodel") parser.add_argument("--image_size", type=int, default=336) - parser.add_argument( - "--hf_dtype", - default="float32", - choices=["float32", "bfloat16", "float16"], - ) parser.add_argument( "--dtype", - default="float32", - choices=["float32", "bfloat16", "float16"], + default="bfloat16", + choices=["float32", "bfloat16"], ) parser.add_argument( - "--vision_dtype", - default=None, - choices=["float32", "bfloat16", "float16"], + "--force-hf-routing", + action="store_true", + help="Use HF expert selections with TorchTitan router scores.", ) parser.add_argument("--seed", type=int, default=42) args = parser.parse_args() @@ -450,12 +480,8 @@ def main() -> None: ) device = torch.device("cuda") dtype = getattr(torch, args.dtype) - vision_dtype = getattr(torch, args.vision_dtype) if args.vision_dtype else dtype - hf_dtype = getattr(torch, args.hf_dtype) - print( - f"hf_dtype={args.hf_dtype} titan text={args.dtype} " - f"titan vision={args.vision_dtype or args.dtype}" - ) + hf_attn_backend = "flash_attention_2" if dtype == torch.bfloat16 else "eager" + print(f"dtype={args.dtype} hf_attn={hf_attn_backend}") tt_config = cast(KimiK3Model.Config, model_registry(args.model_flavor).model) torch.manual_seed(args.seed) @@ -469,13 +495,18 @@ def main() -> None: tt_config, hf_state_dict, args.image_size, - hf_dtype, + dtype, device, ) del hf_state_dict - tt_logits = run_tt(tt_model, ref, vision_dtype, device) - if not compare(ref["last_logits"], tt_logits): - raise SystemExit(1) + tt_logits = run_tt( + tt_model, + ref, + dtype, + device, + args.force_hf_routing, + ) + compare(ref["last_logits"], tt_logits) if __name__ == "__main__": diff --git a/torchtitan/models/kimi_k3/__init__.py b/torchtitan/models/kimi_k3/__init__.py index 831bdb352f..87487a7a46 100644 --- a/torchtitan/models/kimi_k3/__init__.py +++ b/torchtitan/models/kimi_k3/__init__.py @@ -453,34 +453,36 @@ def _kimi_k3_config( def _debugmodel(attn_backend: str) -> KimiK3Model.Config: - dim = 256 + dim = 1024 return _kimi_k3_config( dim=dim, vocab_size=163840, - num_layers=13, - full_attention_layers={3, 7, 11, 12}, + num_layers=24, + full_attention_layers={3, 7, 11, 15, 19, 23}, attn_res_block_size=12, - num_heads=8, - q_lora_rank=256, - kv_lora_rank=128, - qk_nope_head_dim=32, - qk_rope_head_dim=16, - v_head_dim=32, - kda_head_dim=32, + num_heads=16, + q_lora_rank=512, + kv_lora_rank=256, + qk_nope_head_dim=64, + qk_rope_head_dim=32, + v_head_dim=64, + kda_head_dim=64, conv_kernel_size=4, - dense_hidden_dim=1024, - latent_dim=128, - expert_hidden_dim=256, - num_experts=16, + dense_hidden_dim=4096, + latent_dim=512, + expert_hidden_dim=384, + num_experts=32, top_k=4, num_shared_experts=2, vision_encoder=_vision_encoder_config( text_dim=dim, - dim=256, - qkv_dim=384, - hidden_dim=1024, - num_layers=4, - num_heads=3, + dim=512, + qkv_dim=768, + hidden_dim=2048, + num_layers=8, + num_heads=6, + init_pos_emb_height=32, + init_pos_emb_width=32, ), attn_backend=attn_backend, ) From 41447a4330a6dae131f7765598c1756bc9969d52 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Sat, 22 Aug 2026 07:54:24 +0000 Subject: [PATCH 53/67] Adapt Kimi K3 to the folded token layout Update the K3 decoder, KDA, MoE, vision path, dataloader config, tests, and numerical script for the token-major and packed-vision interfaces introduced by #4121. Preserve the latest upstream Kimi-VL numerical-test documentation while rebasing the feature history. --- .../numerical_tests_kimi.py | 26 +++- .../numerical_tests_kimi_k3.py | 20 +-- tests/unit_tests/test_kimi_k3.py | 28 ++-- torchtitan/models/kimi_k3/config_registry.py | 66 ++++++--- torchtitan/models/kimi_k3/kda.py | 54 ++++---- torchtitan/models/kimi_k3/model.py | 114 ++++++++-------- torchtitan/models/kimi_k3/moe.py | 28 ++-- torchtitan/models/kimi_k3/parallelize.py | 5 +- torchtitan/models/kimi_k3/vision_encoder.py | 126 +++++++++--------- torchtitan_recipes/tests/models.py | 8 ++ 10 files changed, 265 insertions(+), 210 deletions(-) diff --git a/scripts/checkpoint_conversion/numerical_tests_kimi.py b/scripts/checkpoint_conversion/numerical_tests_kimi.py index 2ba33e8176..515e8ff98d 100644 --- a/scripts/checkpoint_conversion/numerical_tests_kimi.py +++ b/scripts/checkpoint_conversion/numerical_tests_kimi.py @@ -4,7 +4,31 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -"""Full text+image e2e logit parity: torchtitan Kimi-VL vs the released HF model.""" +"""Full text+image e2e logit parity: torchtitan Kimi-VL vs the released HF model. + +Runs the released HF Kimi-VL (``trust_remote_code``) and torchtitan in ONE +process on the same text+image prompt and compares last-token logits. torchtitan +does its OWN Kimi image processing (``process_image`` + ``vision_to_patches``, +raster order), so the full pipeline is exercised (preprocessing + vision + +projector + scatter + DeepSeek-V3 text tower), not just the forward. + +The released remote code targets transformers ~4.50.x and does NOT import on 5.x, +so run this in an env with ``transformers==4.50.3`` + ``tiktoken`` + ``blobfile``. + +Precision: the HF reference runs at ``--hf_dtype`` (default float32, the model's +"true" output); torchtitan runs at ``--dtype`` (text) with ``--vision_dtype`` +overriding only its vision encoder. ``--dtype float32`` is the correctness gate; +``--dtype bfloat16 --vision_dtype float16`` is the realistic config (~1e-2 KL). + +Usage: + CUDA_VISIBLE_DEVICES=0 python -m \\ + scripts.checkpoint_conversion.numerical_tests_kimi \\ + --hf_model_path ~/hf_assets/moonshotai/Kimi-VL-A3B-Instruct \\ + --tt_checkpoint_path outputs/kimi/kimi_vl_a3b_dcp --dtype float32 + +Add ``--force-hf-routing`` to make titan use HF's exact per-token expert +selections (diagnostic: removes the MoE routing-flip divergence). +""" import argparse import os diff --git a/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py b/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py index 68cc24c4af..93920afb1f 100644 --- a/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py +++ b/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py @@ -272,7 +272,7 @@ def _expand_image_placeholder( return torch.cat( (input_ids[:, :position], image_tokens, input_ids[:, position + 1 :]), dim=1, - ) + ).squeeze(0) def _print_routing_comparison( @@ -312,21 +312,21 @@ def _force_hf_routing(model, expert_indices, device) -> None: for layer_idx, layer in model.layers.items(): if (moe := cast(Any, layer.moe)) is None: continue - ids = expert_indices[int(layer_idx)].unsqueeze(0).to(device) + ids = expert_indices[int(layer_idx)].to(device) router, original = moe.router, moe.router.forward def forced_forward( - x_BLD, + x_TD, expert_bias_E=None, _router=router, _original=original, _ids=ids, ): - _, _, scores_BLE = _original(x_BLD, expert_bias_E) - weights = scores_BLE.gather(dim=-1, index=_ids) + _, _, scores_TE = _original(x_TD, expert_bias_E) + weights = scores_TE.gather(dim=-1, index=_ids) if _router.route_norm: weights = weights / (weights.sum(dim=-1, keepdim=True) + 1e-20) - return weights * _router.route_scale, _ids, scores_BLE + return weights * _router.route_scale, _ids, scores_TE router.forward = forced_forward @@ -381,7 +381,7 @@ def run_tt( merge_size=_MERGE_SIZE, patch_order="raster", ) - pixel_values = patches.unsqueeze(0).to(device=device, dtype=vision_dtype) + pixel_values = patches.to(device=device, dtype=vision_dtype) grid_thw = grid.unsqueeze(0).to(device) num_vision_tokens = (grid[1] // _MERGE_SIZE) * (grid[2] // _MERGE_SIZE) tokens = _expand_image_placeholder( @@ -390,10 +390,10 @@ def run_tt( int(num_vision_tokens.item()), ).to(device) positions = torch.arange( - tokens.shape[1], + tokens.shape[0], dtype=torch.int32, device=device, - ).unsqueeze(0) + ) attention_masks = model.get_attention_masks(positions) print( @@ -429,7 +429,7 @@ def run_tt( attention_masks=attention_masks, ) _print_routing_comparison(ref["expert_indices"], expert_indices) - return logits[:, -1, :].float().cpu().squeeze() + return logits[-1].float().cpu() def compare(ref_logits: torch.Tensor, tt_logits: torch.Tensor) -> None: diff --git a/tests/unit_tests/test_kimi_k3.py b/tests/unit_tests/test_kimi_k3.py index 932a89afff..0bb0cd3aa8 100644 --- a/tests/unit_tests/test_kimi_k3.py +++ b/tests/unit_tests/test_kimi_k3.py @@ -143,7 +143,7 @@ class TestKimiK3(unittest.TestCase): def test_flex_attention_mask(self): config = _small_model_config() model = config.build() - positions = torch.arange(4, dtype=torch.int32).unsqueeze(0) + positions = torch.arange(4, dtype=torch.int32) attention_masks = model.get_attention_masks(positions) self.assertIsInstance(attention_masks, BlockMask) @@ -286,6 +286,7 @@ def test_fsdp_matches_non_distributed_forward_backward(self): pipeline_parallel_degree=1, context_parallel_degree=1, expert_parallel_degree=1, + spmd_backend="partial_dtensor", ) parallel_dims = ParallelDims.from_config(parallelism, world_size=1) with patch( @@ -297,8 +298,8 @@ def test_fsdp_matches_non_distributed_forward_backward(self): model, parallel_dims=parallel_dims, training=TrainingConfig( - local_batch_size=1, - seq_len=6, + num_tokens_per_microbatch_per_dp_rank=6, + max_context_length=6, steps=1, dtype="bfloat16", ), @@ -312,20 +313,19 @@ def test_fsdp_matches_non_distributed_forward_backward(self): self.assertIsInstance(model, FSDPModule) self.assertIsInstance(model.vision_encoder, FSDPModule) - positions_BL = torch.arange( + positions_T = torch.arange( 6, dtype=torch.int32, device=self.device_type, - ).unsqueeze(0) - attention_masks = reference.get_attention_masks(positions_BL) + ) + attention_masks = reference.get_attention_masks(positions_T) inputs = { "tokens": torch.tensor( - [[1, 7, 2, 3, 4, 5]], + [1, 7, 2, 3, 4, 5], dtype=torch.long, device=self.device_type, ), "pixel_values": torch.randn( - 1, 4, 3 * 2 * 2, device=self.device_type, @@ -336,16 +336,16 @@ def test_fsdp_matches_non_distributed_forward_backward(self): device=self.device_type, ), "special_tokens": {"image_id": 7}, - "positions": positions_BL, + "positions": positions_T, "attention_masks": attention_masks, } - actual_BLV = model(**inputs) # pyrefly: ignore [not-callable] - expected_BLV = reference(**inputs) - torch.testing.assert_close(actual_BLV, expected_BLV, atol=0.0, rtol=0.0) + actual_TV = model(**inputs) # pyrefly: ignore [not-callable] + expected_TV = reference(**inputs) + torch.testing.assert_close(actual_TV, expected_TV, atol=0.0, rtol=0.0) - actual_BLV.float().square().mean().backward() - expected_BLV.float().square().mean().backward() + actual_TV.float().square().mean().backward() + expected_TV.float().square().mean().backward() reference_parameters = dict(reference.named_parameters()) compared_gradients = 0 diff --git a/torchtitan/models/kimi_k3/config_registry.py b/torchtitan/models/kimi_k3/config_registry.py index 5a39fa00b3..f1eac6b4fb 100644 --- a/torchtitan/models/kimi_k3/config_registry.py +++ b/torchtitan/models/kimi_k3/config_registry.py @@ -4,14 +4,21 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -from torchtitan.components.checkpoint import CheckpointManager +from dataclasses import replace + +from torchtitan.components.checkpointer import CheckpointManager +from torchtitan.components.data import GrainDataLoader, SingleDatasetConfig from torchtitan.components.loss import ChunkedLossWrapper, CrossEntropyLoss from torchtitan.components.metrics import MetricsProcessor from torchtitan.components.optimizer import default_adamw, LRSchedulersContainer from torchtitan.components.tokenizer import MultiModalTokenizer from torchtitan.config import ParallelismConfig, TrainingConfig from torchtitan.distributed.activation_checkpoint import SelectiveAC -from torchtitan.hf_datasets.multimodal.mm_datasets import MMDataLoader +from torchtitan.hf_datasets.multimodal.mm_collator import MultiModalCollator +from torchtitan.hf_datasets.multimodal.mm_datasets import ( + MM_DATASETS, + MultiModalProcessor, +) from torchtitan.hf_datasets.multimodal.utils.image import resize_to_patch_budget from torchtitan.models.common.config_utils import decoder_vocab_size from torchtitan.trainer import Trainer @@ -19,6 +26,39 @@ from . import KIMI_K3_SPECIAL_TOKENS, model_registry +def _kimi_k3_multimodal_dataloader( + dataset: SingleDatasetConfig, +) -> GrainDataLoader.Config: + processor = dataset.processor + if not isinstance(processor, MultiModalProcessor.Config): + raise ValueError("Kimi K3 multimodal data requires MultiModalProcessor.Config") + + processor = MultiModalProcessor.Config( + sample_processor=processor.sample_processor, + patch_size=14, + temporal_patch_size=1, + spatial_merge_size=2, + resize_fn=resize_to_patch_budget, + min_pixels=56 * 56, + max_pixels=224 * 224, + max_patches=256, + max_patches_per_side=16, + image_mean=(0.5, 0.5, 0.5), + image_std=(0.5, 0.5, 0.5), + ) + return GrainDataLoader.Config( + dataset=replace(dataset, processor=processor), + collator=MultiModalCollator.Config( + max_images_per_batch=8, + patch_size=processor.patch_size, + temporal_patch_size=processor.temporal_patch_size, + spatial_merge_size=processor.spatial_merge_size, + patch_order="raster", + build_mrope_positions=False, + ), + ) + + def kimi_k3_debugmodel() -> Trainer.Config: model_spec = model_registry("debugmodel") return Trainer.Config( @@ -31,21 +71,7 @@ def kimi_k3_debugmodel() -> Trainer.Config: tokenizer=MultiModalTokenizer.Config(**KIMI_K3_SPECIAL_TOKENS), metrics=MetricsProcessor.Config(log_freq=1), model_spec=model_spec, - dataloader=MMDataLoader.Config( - dataset="cc12m-test", - max_images_per_batch=8, - patch_size=14, - temporal_patch_size=1, - spatial_merge_size=2, - patch_order="raster", - resize_fn=resize_to_patch_budget, - min_pixels=56 * 56, - max_pixels=224 * 224, - max_patches=256, - max_patches_per_side=16, - image_mean=(0.5, 0.5, 0.5), - image_std=(0.5, 0.5, 0.5), - ), + dataloader=_kimi_k3_multimodal_dataloader(MM_DATASETS["cc12m-test"]), optimizer=default_adamw(lr=8e-4), lr_scheduler=LRSchedulersContainer.Config( warmup_steps=2, @@ -53,9 +79,11 @@ def kimi_k3_debugmodel() -> Trainer.Config: decay_type="linear", min_lr_factor=0.0, ), + # TODO: Kimi K3 has no spmd_types annotations yet. + parallelism=ParallelismConfig(spmd_backend="partial_dtensor"), training=TrainingConfig( - local_batch_size=1, - seq_len=256, + num_tokens_per_microbatch_per_dp_rank=256, + max_context_length=256, steps=10, dtype="bfloat16", disable_cuda_graphs=True, diff --git a/torchtitan/models/kimi_k3/kda.py b/torchtitan/models/kimi_k3/kda.py index f52c560d7c..f31a8eb9b1 100644 --- a/torchtitan/models/kimi_k3/kda.py +++ b/torchtitan/models/kimi_k3/kda.py @@ -18,7 +18,7 @@ from torchtitan.protocols.module import Module # Shape suffixes: -# B = batch, L = sequence length, D = model dimension, H = heads, +# T = packed tokens, D = model dimension, H = heads, # K = key head dimension, V = value head dimension, C = projection channels. @@ -128,13 +128,13 @@ def __init__(self, config: Config): self.A_log = nn.Parameter(torch.empty(config.num_heads)) self.dt_bias = nn.Parameter(torch.empty(config.num_heads, config.head_dim)) - def _causal_conv(self, x_BLC: torch.Tensor, conv: Conv1d) -> torch.Tensor: - x_BCL = F.pad(x_BLC.transpose(1, 2), (self.conv_kernel_size - 1, 0)) - return F.silu(conv(x_BCL)).transpose(1, 2) + def _causal_conv(self, x_TC: torch.Tensor, conv: Conv1d) -> torch.Tensor: + x_1CT = F.pad(x_TC.T.unsqueeze(0), (self.conv_kernel_size - 1, 0)) + return F.silu(conv(x_1CT)).squeeze(0).T def forward( self, - x_BLD: torch.Tensor, + x_TD: torch.Tensor, attention_masks: AttentionMasksType | None = None, positions: torch.Tensor | None = None, ) -> torch.Tensor: @@ -144,32 +144,32 @@ def forward( "Kimi K3 reference KDA does not support packed-document masks." ) - B, L, _ = x_BLD.shape - q_BLHK = self._causal_conv(self.q_proj(x_BLD), self.q_conv).view( - B, L, self.num_heads, self.head_dim + num_tokens = x_TD.shape[0] + q_THK = self._causal_conv(self.q_proj(x_TD), self.q_conv).view( + num_tokens, self.num_heads, self.head_dim ) - k_BLHK = self._causal_conv(self.k_proj(x_BLD), self.k_conv).view( - B, L, self.num_heads, self.head_dim + k_THK = self._causal_conv(self.k_proj(x_TD), self.k_conv).view( + num_tokens, self.num_heads, self.head_dim ) - v_BLHV = self._causal_conv(self.v_proj(x_BLD), self.v_conv).view( - B, L, self.num_heads, self.head_dim + v_THV = self._causal_conv(self.v_proj(x_TD), self.v_conv).view( + num_tokens, self.num_heads, self.head_dim ) - forget_BLHK = self.forget_b(self.forget_a(x_BLD)).view( - B, L, self.num_heads, self.head_dim + forget_THK = self.forget_b(self.forget_a(x_TD)).view( + num_tokens, self.num_heads, self.head_dim ) - beta_BLH = self.beta(x_BLD).float() - - out_BLHV = self.kernel( - q_BLHK, - k_BLHK, - v_BLHV, - forget_BLHK, - beta_BLH, + beta_TH = self.beta(x_TD).float() + + out_THV = self.kernel( + q_THK.unsqueeze(0), + k_THK.unsqueeze(0), + v_THV.unsqueeze(0), + forget_THK.unsqueeze(0), + beta_TH.unsqueeze(0), self.A_log, self.dt_bias, + ).squeeze(0) + output_gate_THV = self.output_gate(x_TD).view( + num_tokens, self.num_heads, self.head_dim ) - output_gate_BLHV = self.output_gate(x_BLD).view( - B, L, self.num_heads, self.head_dim - ) - out_BLHV = self.output_norm(out_BLHV, output_gate_BLHV) - return self.output_proj(out_BLHV.reshape(B, L, -1)) + out_THV = self.output_norm(out_THV, output_gate_THV) + return self.output_proj(out_THV.reshape(num_tokens, -1)) diff --git a/torchtitan/models/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py index 10c2d8ec54..1aabd6e537 100644 --- a/torchtitan/models/kimi_k3/model.py +++ b/torchtitan/models/kimi_k3/model.py @@ -29,8 +29,8 @@ from .vision_encoder import KimiK3VisionEncoder # Shape suffixes: -# B = batch, L = sequence length, D = model dimension, H = heads, -# K = key head dimension, V = value head dimension, T = flattened tokens, +# T = packed tokens, D = model dimension, H = heads, +# K = key head dimension, V = value head dimension, # N = attention-residual entries. @@ -82,49 +82,48 @@ def __init__(self, config: Config): def forward( self, - x_BLD: torch.Tensor, + x_TD: torch.Tensor, attention_masks: AttentionMasksType | None = None, positions: torch.Tensor | None = None, ) -> torch.Tensor: del positions - B, L, _ = x_BLD.shape - q_BLHK = self.wq_b(self.q_norm(self.wq_a(x_BLD))).view( - B, L, self.n_heads, self.q_head_dim + num_tokens = x_TD.shape[0] + q_THK = self.wq_b(self.q_norm(self.wq_a(x_TD))).view( + num_tokens, self.n_heads, self.q_head_dim ) - compressed_kv_BLC = self.wkv_a(x_BLD) - kv_latent_BLC, k_rope_BLK = torch.split( - compressed_kv_BLC, + compressed_kv_TC = self.wkv_a(x_TD) + kv_latent_TC, k_rope_TK = torch.split( + compressed_kv_TC, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1, ) - kv_BLHC = self.wkv_b(self.kv_norm(kv_latent_BLC)).view( - B, - L, + kv_THC = self.wkv_b(self.kv_norm(kv_latent_TC)).view( + num_tokens, self.n_heads, self.qk_nope_head_dim + self.v_head_dim, ) - k_nope_BLHK, v_BLHV = torch.split( - kv_BLHC, + k_nope_THK, v_THV = torch.split( + kv_THC, [self.qk_nope_head_dim, self.v_head_dim], dim=-1, ) - k_rope_BLHK = k_rope_BLK.view(B, L, 1, self.qk_rope_head_dim).expand( - -1, -1, self.n_heads, -1 + k_rope_THK = k_rope_TK.view(num_tokens, 1, self.qk_rope_head_dim).expand( + -1, self.n_heads, -1 ) - k_BLHK = torch.cat((k_nope_BLHK, k_rope_BLHK), dim=-1) + k_THK = torch.cat((k_nope_THK, k_rope_THK), dim=-1) - out_BLHV = self.inner_attention( - q_BLHK, - k_BLHK, - v_BLHV, + out_THV = self.inner_attention( + q_THK, + k_THK, + v_THV, attention_masks=attention_masks, scale=self.scale, ) - out_BLD = out_BLHV.reshape(B, L, self.n_heads * self.v_head_dim) - out_BLD = out_BLD * torch.sigmoid(self.gate(x_BLD)) - return self.wo(out_BLD) + out_TD = out_THV.reshape(num_tokens, self.n_heads * self.v_head_dim) + out_TD = out_TD * torch.sigmoid(self.gate(x_TD)) + return self.wo(out_TD) def _apply_attention_residual( @@ -208,53 +207,54 @@ def __init__(self, config: Config): def forward( self, - x_BLD: torch.Tensor, + x_TD: torch.Tensor, block_residual_TND: torch.Tensor, attention_masks: AttentionMasksType | None = None, positions: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: - B, L, D = x_BLD.shape - prefix_sum_BLD = x_BLD + prefix_sum_TD = x_TD if block_residual_TND.shape[1] > 0: - x_BLD = _apply_attention_residual( - prefix_sum_BLD.reshape(-1, D), + assert self.attention_res_proj is not None + assert self.attention_res_norm is not None + x_TD = _apply_attention_residual( + prefix_sum_TD, block_residual_TND, self.attention_res_proj, self.attention_res_norm, - ).view(B, L, D) + ) opens_block = self.layer_id % self.attn_res_block_size == 0 if opens_block: block_residual_TND = torch.cat( ( block_residual_TND, - prefix_sum_BLD.reshape(-1, D).unsqueeze(1), + prefix_sum_TD.unsqueeze(1), ), dim=1, ) - h_BLD = self.attention_norm(x_BLD) + h_TD = self.attention_norm(x_TD) if self.attention is not None: - h_BLD = self.attention(h_BLD, attention_masks, positions) + h_TD = self.attention(h_TD, attention_masks, positions) else: assert self.delta_attention is not None - h_BLD = self.delta_attention(h_BLD, None, positions) - prefix_sum_BLD = h_BLD if opens_block else prefix_sum_BLD + h_BLD + h_TD = self.delta_attention(h_TD, None, positions) + prefix_sum_TD = h_TD if opens_block else prefix_sum_TD + h_TD - h_BLD = _apply_attention_residual( - prefix_sum_BLD.reshape(-1, D), + h_TD = _apply_attention_residual( + prefix_sum_TD, block_residual_TND, self.ffn_res_proj, self.ffn_res_norm, - ).view(B, L, D) - h_BLD = self.ffn_norm(h_BLD) + ) + h_TD = self.ffn_norm(h_TD) if self.moe is not None: - h_BLD = self.moe(h_BLD) + h_TD = self.moe(h_TD) else: assert self.feed_forward is not None - h_BLD = self.feed_forward(h_BLD) - return prefix_sum_BLD + h_BLD, block_residual_TND + h_TD = self.feed_forward(h_TD) + return prefix_sum_TD + h_TD, block_residual_TND class KimiK3Model(Decoder): @@ -311,14 +311,14 @@ def _prepare_multimodal_embeds( grid_thw: torch.Tensor | None, special_tokens: dict[str, int] | None, ) -> torch.Tensor: - embeddings = self.tok_embeddings(tokens) + embeddings_TD = self.tok_embeddings(tokens) if (pixel_values is None) != (grid_thw is None): raise ValueError( "pixel_values and grid_thw must either both be provided or " "both be omitted." ) if pixel_values is None: - return embeddings + return embeddings_TD assert grid_thw is not None if self.vision_encoder is None: raise ValueError("pixel_values were provided without a vision encoder.") @@ -339,7 +339,7 @@ def _prepare_multimodal_embeds( special_tokens["image_id"], ) return scatter_vision_embeds( - embeddings, + embeddings_TD, vision_embeds=vision_embeds, vision_positions=vision_positions, ) @@ -359,32 +359,32 @@ def forward( # pyrefly: ignore [bad-override] if pixel_values_videos is not None or grid_thw_videos is not None: raise NotImplementedError("Kimi K3 v1 supports images but not videos.") if self.tok_embeddings is not None: - h_BLD = self._prepare_multimodal_embeds( + h_TD = self._prepare_multimodal_embeds( tokens, pixel_values=pixel_values, grid_thw=grid_thw, special_tokens=special_tokens, ) else: - h_BLD = tokens + h_TD = tokens - B, L, D = h_BLD.shape - block_residual_TND = h_BLD.new_zeros(B * L, 0, D) + num_tokens, D = h_TD.shape + block_residual_TND = h_TD.new_zeros(num_tokens, 0, D) for layer in self.layers.values(): - h_BLD, block_residual_TND = layer( - h_BLD, + h_TD, block_residual_TND = layer( + h_TD, block_residual_TND, attention_masks, positions, ) - h_BLD = _apply_attention_residual( - h_BLD.reshape(-1, D), + h_TD = _apply_attention_residual( + h_TD, block_residual_TND, self.output_res_proj, self.output_res_norm, - ).view(B, L, D) - h_BLD = self.norm(h_BLD) if self.norm is not None else h_BLD + ) + h_TD = self.norm(h_TD) if self.norm is not None else h_TD if self._skip_lm_head: - return h_BLD - return self.lm_head(h_BLD) if self.lm_head is not None else h_BLD + return h_TD + return self.lm_head(h_TD) if self.lm_head is not None else h_TD diff --git a/torchtitan/models/kimi_k3/moe.py b/torchtitan/models/kimi_k3/moe.py index 0c9616aa31..990104a9d8 100644 --- a/torchtitan/models/kimi_k3/moe.py +++ b/torchtitan/models/kimi_k3/moe.py @@ -17,8 +17,8 @@ from torchtitan.models.common.nn_modules import RMSNorm # Shape suffixes: -# B = batch, L = sequence length, D = model dimension, E = experts, -# F = expert hidden dimension, R = routed tokens, S = selected experts per token. +# T = packed tokens, D = model dimension, E = experts, +# F = expert hidden dimension, R = routed tokens, K = selected experts per token. def _situ_glu( @@ -122,23 +122,23 @@ def __init__(self, config: Config): self.routed_norm = config.routed_norm.build() self.routed_up = config.routed_up.build() - def forward(self, x_BLD: torch.Tensor) -> torch.Tensor: - weights_BLS, expert_ids_BLS, scores_BLE = self.router(x_BLD, self.expert_bias_E) - routing_map_BLE = torch.zeros_like(scores_BLE, dtype=torch.bool).scatter_( - -1, expert_ids_BLS, True + def forward(self, x_TD: torch.Tensor) -> torch.Tensor: + weights_TK, expert_ids_TK, scores_TE = self.router(x_TD, self.expert_bias_E) + routing_map_TE = torch.zeros_like(scores_TE, dtype=torch.bool).scatter_( + -1, expert_ids_TK, True ) - num_tokens_per_expert_E = routing_map_BLE.sum(dim=(0, 1)) + num_tokens_per_expert_E = routing_map_TE.sum(dim=0) if self.training: with torch.no_grad(): self.tokens_per_expert_E.add_(num_tokens_per_expert_E) - routed_BLD = self.routed_experts( - self.routed_down(x_BLD), - weights_BLS, - expert_ids_BLS, + routed_TD = self.routed_experts( + self.routed_down(x_TD), + weights_TK, + expert_ids_TK, num_tokens_per_expert_E, ) - out_BLD = self.routed_up(self.routed_norm(routed_BLD)) + out_TD = self.routed_up(self.routed_norm(routed_TD)) if self.shared_experts is not None: - out_BLD = out_BLD + self.shared_experts(x_BLD) - return out_BLD + out_TD = out_TD + self.shared_experts(x_TD) + return out_TD diff --git a/torchtitan/models/kimi_k3/parallelize.py b/torchtitan/models/kimi_k3/parallelize.py index d0c91f6609..8a7d604a02 100644 --- a/torchtitan/models/kimi_k3/parallelize.py +++ b/torchtitan/models/kimi_k3/parallelize.py @@ -48,9 +48,10 @@ def parallelize_kimi_k3( "Kimi K3 currently supports FSDP2 data parallelism " f"only; disable {', '.join(unsupported_parallelisms)}." ) - if parallelism.spmd_backend != "default": + if parallelism.spmd_backend != "partial_dtensor": raise NotImplementedError( - "Kimi K3 FSDP2 currently supports the default SPMD backend only." + "Kimi K3 FSDP2 currently supports the partial_dtensor SPMD backend " + "only; the config registry pins it." ) if compile_config.enable and "model" in compile_config.components: raise NotImplementedError("Kimi K3 does not support model compilation yet.") diff --git a/torchtitan/models/kimi_k3/vision_encoder.py b/torchtitan/models/kimi_k3/vision_encoder.py index 3fb405109a..aa53154647 100644 --- a/torchtitan/models/kimi_k3/vision_encoder.py +++ b/torchtitan/models/kimi_k3/vision_encoder.py @@ -8,18 +8,19 @@ Shape suffixes: - N = number of visual items -- P = maximum patches per item (padded) +- T = total packed patches - D = vision hidden dimension - H = number of attention heads - K = attention head dimension - C = number of complex-valued head-dimension pairs -- M = maximum merged tokens per item (padded) +- M = total merged tokens - F = merged feature dimension - O = projected text dimension """ from dataclasses import dataclass, field +import spmd_types as spmd import torch import torch.nn.functional as F from torch.nn.attention.flex_attention import BlockMask @@ -28,8 +29,7 @@ from torchtitan.models.common.nn_modules import GELU, RMSNorm from torchtitan.models.common.rope import ComplexRoPE from torchtitan.models.common.vision_encoder import ( - compiled_create_block_mask, - get_vision_block_mask_mod, + create_block_diagonal_mask, VisionAttention, VisionMLP, ) @@ -52,19 +52,9 @@ def _get_temporal_pos_embed( return torch.cat((angles.sin(), angles.cos()), dim=-1) -def _pad_sequence(x: torch.Tensor, target_length: int) -> torch.Tensor: - """Pad the leading sequence dimension without modifying ``x`` in place.""" - padding_length = target_length - x.shape[0] - if padding_length == 0: - return x - padding = x.new_zeros(padding_length, *x.shape[1:]) - return torch.cat((x, padding), dim=0) - - def _compute_learned_pos_embeds( pos_embed: torch.Tensor, grids: list[list[int]], - max_num_patches: int, interpolation_mode: str, max_num_frames: int, ) -> torch.Tensor: @@ -73,7 +63,7 @@ def _compute_learned_pos_embeds( pos_grid = pos_embed.permute(2, 0, 1).unsqueeze(0).float() cached_spatial: dict[tuple[int, int], torch.Tensor] = {} - padded_positions = [] + positions = [] for num_frames, grid_h, grid_w in grids: if num_frames > max_num_frames: raise ValueError( @@ -104,20 +94,19 @@ def _compute_learned_pos_embeds( temporal = _get_temporal_pos_embed(num_frames, dim, device=pos_embed.device) item_pos = spatial.unsqueeze(0) + temporal.unsqueeze(1).to(spatial.dtype) item_pos = item_pos.reshape(num_frames * grid_h * grid_w, dim) - padded_positions.append(_pad_sequence(item_pos, max_num_patches)) + positions.append(item_pos) - return torch.stack(padded_positions) + return torch.cat(positions) def _compute_2d_rope_cache( freq_table: torch.Tensor, grids: list[list[int]], - max_num_patches: int, head_dim: int, ) -> torch.Tensor: """Build the real-valued 2D RoPE cache in raster patch order.""" cached_spatial: dict[tuple[int, int], torch.Tensor] = {} - padded_angles = [] + item_angles = [] for num_frames, grid_h, grid_w in grids: spatial = cached_spatial.get((grid_h, grid_w)) if spatial is None: @@ -128,34 +117,34 @@ def _compute_2d_rope_cache( grid_h * grid_w, head_dim // 2 ) cached_spatial[(grid_h, grid_w)] = spatial - item_angles = spatial.repeat(num_frames, 1) - padded_angles.append(_pad_sequence(item_angles, max_num_patches)) + item_angles.append(spatial.repeat(num_frames, 1)) - angles = torch.stack(padded_angles) + angles = torch.cat(item_angles) # ComplexRoPE.apply_rotary_emb multiplies in complex64; float() only widens # the container, so cos/sin keep whatever precision angles were computed in. cos_sin = torch.stack((angles.cos(), angles.sin()), dim=-1).float() - return torch.view_as_complex(cos_sin).unsqueeze(2) + return torch.view_as_complex(cos_sin).unsqueeze(1) def _temporal_pool_and_merge( - hidden_NPD: torch.Tensor, + hidden_TD: torch.Tensor, grids: list[list[int]], merge_kernel_size: tuple[int, int], ) -> torch.Tensor: """Temporally pool and concatenate neighboring spatial patch features.""" - _, _, dim = hidden_NPD.shape + dim = hidden_TD.shape[-1] kernel_h, kernel_w = merge_kernel_size merged_dim = kernel_h * kernel_w * dim - max_merged = max( - (grid_h // kernel_h) * (grid_w // kernel_w) for _, grid_h, grid_w in grids - ) - padded_items = [] - for item_idx, (num_frames, grid_h, grid_w) in enumerate(grids): + merged_items = [] + offset = 0 + for num_frames, grid_h, grid_w in grids: + num_patches = num_frames * grid_h * grid_w + item = hidden_TD[offset : offset + num_patches] + offset += num_patches merged_h = grid_h // kernel_h merged_w = grid_w // kernel_w - item = hidden_NPD[item_idx, : num_frames * grid_h * grid_w].view( + item = item.view( num_frames, merged_h, kernel_h, @@ -164,10 +153,9 @@ def _temporal_pool_and_merge( dim, ) item = item.permute(0, 1, 3, 2, 4, 5).mean(dim=0) - item = item.reshape(merged_h * merged_w, merged_dim) - padded_items.append(_pad_sequence(item, max_merged)) + merged_items.append(item.reshape(merged_h * merged_w, merged_dim)) - return torch.stack(padded_items) + return torch.cat(merged_items) class VisionRotaryEmbedding2D(Module): @@ -234,18 +222,18 @@ def __init__(self, config: Config): def forward( self, - x_NPD: torch.Tensor, + x_TD: torch.Tensor, *, rope_cache: torch.Tensor, attention_mask: BlockMask, ) -> torch.Tensor: - x_NPD = x_NPD + self.attn( - self.norm1(x_NPD), + x_TD = x_TD + self.attn( + self.norm1(x_TD), rope_cache=rope_cache, rope_apply=ComplexRoPE.apply_rotary_emb, attention_mask=attention_mask, ) - return x_NPD + self.mlp(self.norm2(x_NPD)) + return x_TD + self.mlp(self.norm2(x_TD)) class KimiK3VisionProjector(Module): @@ -265,9 +253,9 @@ def __init__(self, config: Config): self.post_norm = config.post_norm.build() self.activation = config.activation.build() - def forward(self, merged_NMF: torch.Tensor) -> torch.Tensor: - projected_NMO = self.linear_2(self.activation(self.linear_1(merged_NMF))) - return self.post_norm(projected_NMO) + def forward(self, merged_MF: torch.Tensor) -> torch.Tensor: + projected_MO = self.linear_2(self.activation(self.linear_1(merged_MF))) + return self.post_norm(projected_MO) class KimiK3VisionEncoder(Module): @@ -313,7 +301,7 @@ def __init__(self, config: Config): self.projector = config.projector.build() def _compute_position_embeddings( - self, grids: list[list[int]], max_num_patches: int + self, grids: list[list[int]] ) -> tuple[torch.Tensor, torch.Tensor]: max_grid_side = max(max(grid_h, grid_w) for _, grid_h, grid_w in grids) if ( @@ -324,14 +312,12 @@ def _compute_position_embeddings( learned_pos = _compute_learned_pos_embeds( self.pos_embed, grids, - max_num_patches, self.interpolation_mode, self.max_num_frames, ) rope_cache = _compute_2d_rope_cache( self._cached_freq_table, grids, - max_num_patches, self.rotary_pos_emb.head_dim, ) return learned_pos, rope_cache @@ -342,33 +328,41 @@ def forward( *, grid_thw: torch.Tensor, ) -> torch.Tensor: - """Encode padded raster-order patches and return padded text features.""" - num_items, max_num_patches, _ = pixel_values.shape + """Encode packed raster-order patches and return packed text features.""" grids = grid_thw.tolist() kernel_h, kernel_w = self.merge_kernel_size - num_patches_N = grid_thw.prod(dim=-1).to(torch.long) + for _, grid_h, grid_w in grids: + if grid_h % kernel_h != 0 or grid_w % kernel_w != 0: + raise ValueError( + f"Vision grid {grid_h}x{grid_w} is not divisible by " + f"merge kernel {self.merge_kernel_size}." + ) - learned_pos, rope_cache = self._compute_position_embeddings( - grids, max_num_patches - ) - hidden_NPD = self.patch_embed(pixel_values) + learned_pos - - mask_mod = get_vision_block_mask_mod(num_patches_N) - attention_mask = compiled_create_block_mask( - mask_mod, - num_items, - None, - max_num_patches, - max_num_patches, - device=hidden_NPD.device, - ) + segment_lengths = grid_thw.prod(dim=-1) + total_tokens = pixel_values.shape[0] + expected_tokens = sum(t * h * w for t, h, w in grids) + if total_tokens != expected_tokens: + raise ValueError( + f"pixel_values contains {total_tokens} patches but grid_thw " + f"describes {expected_tokens}." + ) + + learned_pos, rope_cache = self._compute_position_embeddings(grids) + hidden_TD = self.patch_embed(pixel_values) + learned_pos + + with spmd.no_typecheck(): + attention_mask = create_block_diagonal_mask( + segment_lengths, + total_tokens, + hidden_TD.device, + ) for block in self.layers.values(): - hidden_NPD = block( - hidden_NPD, + hidden_TD = block( + hidden_TD, rope_cache=rope_cache, attention_mask=attention_mask, ) - hidden_NPD = self.final_norm(hidden_NPD) - merged_NMF = _temporal_pool_and_merge(hidden_NPD, grids, self.merge_kernel_size) - return self.projector(merged_NMF) + hidden_TD = self.final_norm(hidden_TD) + merged_MF = _temporal_pool_and_merge(hidden_TD, grids, self.merge_kernel_size) + return self.projector(merged_MF) diff --git a/torchtitan_recipes/tests/models.py b/torchtitan_recipes/tests/models.py index 59a61e8b60..87cf8d3e78 100644 --- a/torchtitan_recipes/tests/models.py +++ b/torchtitan_recipes/tests/models.py @@ -215,6 +215,14 @@ def kimi_k2_5_debugmodel_muon_fsdp2_pp2_ep2() -> Trainer.Config: return config +def kimi_k3_debugmodel_mm_fsdp2() -> Trainer.Config: + from torchtitan.models.kimi_k3.config_registry import kimi_k3_debugmodel + + config = kimi_k3_debugmodel() + config.parallelism.data_parallel_shard_degree = 2 + return config + + def muse_glimmer_debugmodel_mm_fsdp2_tp2() -> Trainer.Config: from torchtitan.models.muse_glimmer.config_registry import ( muse_glimmer_debugmodel_mm, From 414942244642fcc7e2f7f09b8da62fb343f5b7a5 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Sat, 22 Aug 2026 08:26:55 +0000 Subject: [PATCH 54/67] update sample packing checking --- torchtitan/models/kimi_k3/model.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/torchtitan/models/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py index 1aabd6e537..d052e33a37 100644 --- a/torchtitan/models/kimi_k3/model.py +++ b/torchtitan/models/kimi_k3/model.py @@ -23,6 +23,7 @@ from torchtitan.models.common.nn_modules import RMSNorm from torchtitan.models.utils import get_moe_model_nparams_and_flops from torchtitan.protocols.module import Module +from torchtitan.hf_datasets.multimodal.mm_datasets import MMSamplePackingConfig from .kda import KimiDeltaAttention from .moe import KimiFeedForward, KimiLatentMoE @@ -267,11 +268,9 @@ class Config(Decoder.Config): def update_from_config(self, *, config, **kwargs) -> None: # Unsupported parallelisms are rejected in parallelize_kimi_k3. - dataloader = getattr(config, "dataloader", None) - if getattr(dataloader, "packing_buffer_size", 0) > 0: - raise NotImplementedError( - "Kimi K3 v1 does not support packed documents." - ) + dataset = config.dataloader.dataset + if isinstance(dataset, MMSamplePackingConfig): + raise ValueError("Kimi K3 does not yet support sample packing.") Decoder.Config.update_from_config(self, config=config, **kwargs) def get_nparams_and_flops( From 9f8bcf0cb19f515c65df239f876f6a910114524b Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Sat, 22 Aug 2026 10:02:33 +0000 Subject: [PATCH 55/67] add TODO --- torchtitan/models/kimi_k3/model.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/torchtitan/models/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py index d052e33a37..7c89329201 100644 --- a/torchtitan/models/kimi_k3/model.py +++ b/torchtitan/models/kimi_k3/model.py @@ -9,6 +9,8 @@ import torch from torch import nn +from torchtitan.hf_datasets.multimodal.mm_datasets import MMSamplePackingConfig + from torchtitan.models.common import Linear from torchtitan.models.common.attention import ( AttentionMasksType, @@ -23,7 +25,6 @@ from torchtitan.models.common.nn_modules import RMSNorm from torchtitan.models.utils import get_moe_model_nparams_and_flops from torchtitan.protocols.module import Module -from torchtitan.hf_datasets.multimodal.mm_datasets import MMSamplePackingConfig from .kda import KimiDeltaAttention from .moe import KimiFeedForward, KimiLatentMoE @@ -267,8 +268,9 @@ class Config(Decoder.Config): vision_encoder: KimiK3VisionEncoder.Config | None = None def update_from_config(self, *, config, **kwargs) -> None: - # Unsupported parallelisms are rejected in parallelize_kimi_k3. dataset = config.dataloader.dataset + # TODO: Support sample packing by resetting the Q/K/V causal-convolution + # and KDA recurrent states at document boundaries. if isinstance(dataset, MMSamplePackingConfig): raise ValueError("Kimi K3 does not yet support sample packing.") Decoder.Config.update_from_config(self, config=config, **kwargs) From 44d0e83c049634bb6e4ca9a7d2ecab90d60923fc Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Sat, 22 Aug 2026 10:06:00 +0000 Subject: [PATCH 56/67] chang fla kda from fp32 to bf16 --- tests/unit_tests/test_kimi_k3.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/unit_tests/test_kimi_k3.py b/tests/unit_tests/test_kimi_k3.py index 0bb0cd3aa8..e5f6bcf72c 100644 --- a/tests/unit_tests/test_kimi_k3.py +++ b/tests/unit_tests/test_kimi_k3.py @@ -173,7 +173,12 @@ def test_fla_kda_kernel_matches_recurrent_reference(self): num_heads = 3 def parameter(*shape: int) -> torch.Tensor: - return torch.randn(*shape, device="cuda", requires_grad=True) + return torch.randn( + *shape, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) for lower_bound in (-5.0, None): with self.subTest(lower_bound=lower_bound): From dee45e35721d24b8ab959a93140d0f799734eed1 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Sat, 22 Aug 2026 10:12:26 +0000 Subject: [PATCH 57/67] fix test_update_from_config_propagates_moe_force_load_balance when rebase --- tests/unit_tests/test_kimi_k3.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/unit_tests/test_kimi_k3.py b/tests/unit_tests/test_kimi_k3.py index e5f6bcf72c..80b959c25b 100644 --- a/tests/unit_tests/test_kimi_k3.py +++ b/tests/unit_tests/test_kimi_k3.py @@ -148,12 +148,15 @@ def test_flex_attention_mask(self): self.assertIsInstance(attention_masks, BlockMask) def test_update_from_config_propagates_moe_force_load_balance(self): + from torchtitan.components.data import GrainDataLoader from torchtitan.config import DebugConfig + from torchtitan.hf_datasets.multimodal.mm_datasets import MM_DATASETS from torchtitan.trainer import Trainer model_config = _small_model_config() runtime_config = Trainer.Config( debug=DebugConfig(moe_force_load_balance=True), + dataloader=GrainDataLoader.Config(dataset=MM_DATASETS["cc12m-test"]), activation_checkpoint=None, ) model_config.update_from_config(config=runtime_config) From 84e9f7b3891b840d2f04e6a4deb42b0223afb183 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Sun, 23 Aug 2026 06:33:34 +0000 Subject: [PATCH 58/67] rename probs_TN to probs_T1N --- torchtitan/models/kimi_k3/model.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/torchtitan/models/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py index 7c89329201..25f473ff21 100644 --- a/torchtitan/models/kimi_k3/model.py +++ b/torchtitan/models/kimi_k3/model.py @@ -146,8 +146,8 @@ def _apply_attention_residual( keys_TND = values_float * torch.rsqrt(variance + norm.eps) score_weight_D = norm.weight.float() * projection.weight.squeeze(0).float() scores_TN = (keys_TND * score_weight_D).sum(dim=-1) - probs_TN = torch.softmax(scores_TN, dim=-1).unsqueeze(1) - output_TD = torch.matmul(probs_TN, values_float).squeeze(1) + probs_T1N = torch.softmax(scores_TN, dim=-1).unsqueeze(1) + output_TD = torch.matmul(probs_T1N, values_float).squeeze(1) return output_TD.to(values_TND.dtype) From 157440aeaf1a18e8dc70a507ad1f9dddd958f587 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Sun, 23 Aug 2026 06:35:05 +0000 Subject: [PATCH 59/67] removed dtype args, only support bf16 now --- .../numerical_tests_kimi_k3.py | 27 ++++++------------- 1 file changed, 8 insertions(+), 19 deletions(-) diff --git a/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py b/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py index 93920afb1f..9ee6d9ba1c 100644 --- a/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py +++ b/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py @@ -20,16 +20,10 @@ The released code requires ``transformers==4.56.2`` and ``tiktoken``. Usage: - # BF16: HF FlashAttention2 / TT FlexAttention CUDA_VISIBLE_DEVICES=0 python -m \ scripts.checkpoint_conversion.numerical_tests_kimi_k3 - # FP32: HF eager / TT FlexAttention - CUDA_VISIBLE_DEVICES=0 python -m \ - scripts.checkpoint_conversion.numerical_tests_kimi_k3 \ - --dtype float32 - -Add ``--force-hf-routing`` to either command for the routing-fixed diagnostic. +Add ``--force-hf-routing`` for the routing-fixed diagnostic. """ import argparse @@ -53,6 +47,8 @@ _HF_REPO_ID = "moonshotai/Kimi-K3" _HF_REVISION = "9f62e4e9fffbd0a83ddd60e1c209d828994b3569" +_DTYPE = torch.bfloat16 +_HF_ATTN_BACKEND = "flash_attention_2" _MEDIA_TOKEN_ID = 163605 _PATCH_SIZE = 14 _MERGE_SIZE = 2 @@ -175,11 +171,10 @@ def _build_hf_model( local_files_only=True, ) _reduce_hf_config(hf_config, tt_config, hf_model_path) - attn_backend = "flash_attention_2" if dtype == torch.bfloat16 else "eager" - hf_config.text_config._attn_implementation = attn_backend - hf_config.vision_config._attn_implementation = attn_backend + hf_config.text_config._attn_implementation = _HF_ATTN_BACKEND + hf_config.vision_config._attn_implementation = _HF_ATTN_BACKEND model = AutoModelForCausalLM.from_config(hf_config, trust_remote_code=True) - model.language_model.config._attn_implementation = attn_backend + model.language_model.config._attn_implementation = _HF_ATTN_BACKEND model.to(dtype=dtype) model.load_state_dict(hf_state_dict, strict=True) return model.eval() @@ -457,11 +452,6 @@ def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--model_flavor", default="debugmodel") parser.add_argument("--image_size", type=int, default=336) - parser.add_argument( - "--dtype", - default="bfloat16", - choices=["float32", "bfloat16"], - ) parser.add_argument( "--force-hf-routing", action="store_true", @@ -479,9 +469,8 @@ def main() -> None: allow_patterns=["*.json", "*.py", "tiktoken.model"], ) device = torch.device("cuda") - dtype = getattr(torch, args.dtype) - hf_attn_backend = "flash_attention_2" if dtype == torch.bfloat16 else "eager" - print(f"dtype={args.dtype} hf_attn={hf_attn_backend}") + dtype = _DTYPE + print(f"dtype={dtype} hf_attn={_HF_ATTN_BACKEND}") tt_config = cast(KimiK3Model.Config, model_registry(args.model_flavor).model) torch.manual_seed(args.seed) From 075d54c6f26863e2939c26a076b9b9cf192ccde8 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Sun, 23 Aug 2026 06:37:24 +0000 Subject: [PATCH 60/67] Clarify the README that the numerical results are based on a reduced version of reduced hf config --- torchtitan/models/kimi_k3/README.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/torchtitan/models/kimi_k3/README.md b/torchtitan/models/kimi_k3/README.md index 897c387354..2a0f09205e 100644 --- a/torchtitan/models/kimi_k3/README.md +++ b/torchtitan/models/kimi_k3/README.md @@ -44,6 +44,9 @@ describe the released model. ## Numerical Parity +The parity script reduces the released Hugging Face configuration to match +TorchTitan's local `debugmodel` configuration before initializing both models. + End-to-end KL divergence against the Hugging Face implementation (multimodal inputs): **6.7634e-7**, with **100% top-1 and top-5 match**. From 99c2c68251f9cb12dd6f4b2ede19cbdceabca0ba Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Sun, 23 Aug 2026 06:38:45 +0000 Subject: [PATCH 61/67] remove test_update_from_config_propagates_moe_force_load_balance --- tests/unit_tests/test_kimi_k3.py | 22 ---------------------- 1 file changed, 22 deletions(-) diff --git a/tests/unit_tests/test_kimi_k3.py b/tests/unit_tests/test_kimi_k3.py index 80b959c25b..7b5e3167d4 100644 --- a/tests/unit_tests/test_kimi_k3.py +++ b/tests/unit_tests/test_kimi_k3.py @@ -147,28 +147,6 @@ def test_flex_attention_mask(self): attention_masks = model.get_attention_masks(positions) self.assertIsInstance(attention_masks, BlockMask) - def test_update_from_config_propagates_moe_force_load_balance(self): - from torchtitan.components.data import GrainDataLoader - from torchtitan.config import DebugConfig - from torchtitan.hf_datasets.multimodal.mm_datasets import MM_DATASETS - from torchtitan.trainer import Trainer - - model_config = _small_model_config() - runtime_config = Trainer.Config( - debug=DebugConfig(moe_force_load_balance=True), - dataloader=GrainDataLoader.Config(dataset=MM_DATASETS["cc12m-test"]), - activation_checkpoint=None, - ) - model_config.update_from_config(config=runtime_config) - - router_configs = [ - layer.moe.router for layer in model_config.layers if layer.moe is not None - ] - self.assertGreater(len(router_configs), 0) - self.assertTrue( - all(router._debug_force_load_balance for router in router_configs) - ) - @unittest.skipIf(not torch.cuda.is_available(), "FLA KDA kernel requires CUDA.") def test_fla_kda_kernel_matches_recurrent_reference(self): torch.manual_seed(1) From 099920626e5f7f1542f2f880836a9366464e58f8 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Sun, 23 Aug 2026 07:04:54 +0000 Subject: [PATCH 62/67] remove TestKimiK3FSDP, and add this in to github ci --- .../integration_test_8gpu_features.yaml | 6 + tests/unit_tests/test_kimi_k3.py | 148 +----------------- 2 files changed, 12 insertions(+), 142 deletions(-) diff --git a/.github/workflows/integration_test_8gpu_features.yaml b/.github/workflows/integration_test_8gpu_features.yaml index 422d36968c..873fb37224 100644 --- a/.github/workflows/integration_test_8gpu_features.yaml +++ b/.github/workflows/integration_test_8gpu_features.yaml @@ -93,6 +93,12 @@ jobs: python -m pytest tests/unit_tests/flex_shard/test_dist_muon.py \ --durations=20 -vv + if [[ "${{ matrix.gpu-arch-type }}" == "cuda" ]]; then + CUDA_VISIBLE_DEVICES=0 python -m pytest \ + tests/unit_tests/test_kimi_k3.py::TestKimiK3::test_fla_kda_kernel_matches_recurrent_reference \ + -vv + fi + sudo mkdir -p "$RUNNER_TEMP/artifacts-to-be-uploaded" sudo chown -R $(id -u):$(id -g) "$RUNNER_TEMP/artifacts-to-be-uploaded" diff --git a/tests/unit_tests/test_kimi_k3.py b/tests/unit_tests/test_kimi_k3.py index 7b5e3167d4..9d02f1e010 100644 --- a/tests/unit_tests/test_kimi_k3.py +++ b/tests/unit_tests/test_kimi_k3.py @@ -4,56 +4,34 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -import copy import unittest -from unittest.mock import patch import torch import torch.nn.functional as F -from torch.distributed._composable.fsdp import FSDPModule -from torch.distributed.tensor import DTensor from torch.nn.attention.flex_attention import BlockMask -from torch.testing._internal.distributed._tensor.common_dtensor import ( - DTensorTestBase, - with_comms, -) -from torchtitan.config import CompileConfig, ParallelismConfig, TrainingConfig -from torchtitan.distributed import ParallelDims - -from torchtitan.models.kimi_k3 import ( - _kimi_k3_config, - _vision_encoder_config, - parallelize_kimi_k3, -) +from torchtitan.models.kimi_k3 import _kimi_k3_config, _vision_encoder_config from torchtitan.models.kimi_k3.kda import KimiKDAKernel from torchtitan.models.kimi_k3.model import KimiK3Model from torchtitan.models.kimi_k3.state_dict_adapter import KimiK3StateDictAdapter -def _small_model_config( - *, - attn_res_block_size: int = 1, - full_attention_layers: set[int] | None = None, -) -> KimiK3Model.Config: +def _small_model_config() -> KimiK3Model.Config: """Build a reduced KDA+MLA, dense+MoE, multimodal Kimi K3 config.""" - if full_attention_layers is None: - full_attention_layers = {1} - dim = 64 return _kimi_k3_config( dim=dim, vocab_size=32, num_layers=2, - full_attention_layers=full_attention_layers, - attn_res_block_size=attn_res_block_size, + full_attention_layers={1}, + attn_res_block_size=1, num_heads=2, q_lora_rank=32, kv_lora_rank=32, qk_nope_head_dim=16, qk_rope_head_dim=16, v_head_dim=16, - kda_head_dim=16, + kda_head_dim=64, conv_kernel_size=3, dense_hidden_dim=128, latent_dim=32, @@ -150,7 +128,7 @@ def test_flex_attention_mask(self): @unittest.skipIf(not torch.cuda.is_available(), "FLA KDA kernel requires CUDA.") def test_fla_kda_kernel_matches_recurrent_reference(self): torch.manual_seed(1) - head_dim = 32 + head_dim = 64 num_heads = 3 def parameter(*shape: int) -> torch.Tensor: @@ -240,119 +218,5 @@ def test_state_dict_round_trips_through_hf_adapter(self): torch.testing.assert_close(value, roundtrip_state_dict[key]) -class TestKimiK3FSDP(DTensorTestBase): - @property - def world_size(self): - return 1 - - @unittest.skipIf(not torch.cuda.is_available(), "Kimi K3 FSDP requires CUDA.") - @with_comms - def test_fsdp_matches_non_distributed_forward_backward(self): - torch.manual_seed(3) - config = _small_model_config( - attn_res_block_size=2, - full_attention_layers={1}, - ) - with torch.device("meta"): - model = config.build() - model.to_empty(device=self.device_type) - model.init_states() - with torch.no_grad(): - for transformer_block in model.layers.values(): - if transformer_block.moe is not None: - transformer_block.moe.router.gate.weight.zero_() - - reference = copy.deepcopy(model) - for parameter in reference.parameters(): - parameter.data = parameter.data.to(torch.bfloat16) - - parallelism = ParallelismConfig( - data_parallel_shard_degree=1, - tensor_parallel_degree=1, - pipeline_parallel_degree=1, - context_parallel_degree=1, - expert_parallel_degree=1, - spmd_backend="partial_dtensor", - ) - parallel_dims = ParallelDims.from_config(parallelism, world_size=1) - with patch( - "torchtitan.distributed.parallel_dims.device_type", - self.device_type, - ): - parallel_dims.build_mesh() - model = parallelize_kimi_k3( - model, - parallel_dims=parallel_dims, - training=TrainingConfig( - num_tokens_per_microbatch_per_dp_rank=6, - max_context_length=6, - steps=1, - dtype="bfloat16", - ), - parallelism=parallelism, - compile_config=CompileConfig(), - ac_config=None, - dump_folder="", - ) - - assert isinstance(model, KimiK3Model) - self.assertIsInstance(model, FSDPModule) - self.assertIsInstance(model.vision_encoder, FSDPModule) - - positions_T = torch.arange( - 6, - dtype=torch.int32, - device=self.device_type, - ) - attention_masks = reference.get_attention_masks(positions_T) - inputs = { - "tokens": torch.tensor( - [1, 7, 2, 3, 4, 5], - dtype=torch.long, - device=self.device_type, - ), - "pixel_values": torch.randn( - 4, - 3 * 2 * 2, - device=self.device_type, - ), - "grid_thw": torch.tensor( - [[1, 2, 2]], - dtype=torch.long, - device=self.device_type, - ), - "special_tokens": {"image_id": 7}, - "positions": positions_T, - "attention_masks": attention_masks, - } - - actual_TV = model(**inputs) # pyrefly: ignore [not-callable] - expected_TV = reference(**inputs) - torch.testing.assert_close(actual_TV, expected_TV, atol=0.0, rtol=0.0) - - actual_TV.float().square().mean().backward() - expected_TV.float().square().mean().backward() - - reference_parameters = dict(reference.named_parameters()) - compared_gradients = 0 - for name, parameter in model.named_parameters(): - actual_grad = parameter.grad - expected_grad = reference_parameters[name].grad - self.assertEqual(actual_grad is None, expected_grad is None) - if actual_grad is None: - continue - if isinstance(actual_grad, DTensor): - actual_grad = actual_grad.to_local() - assert expected_grad is not None - torch.testing.assert_close( - actual_grad.float(), - expected_grad.float(), - atol=0.0, - rtol=0.0, - ) - compared_gradients += 1 - self.assertGreater(compared_gradients, 0) - - if __name__ == "__main__": unittest.main() From 32abaa03e66843887a14ba2a67c217071bb869b7 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Sun, 23 Aug 2026 11:47:57 +0000 Subject: [PATCH 63/67] reuse some function from k2.7 --- torchtitan/models/common/vision_encoder.py | 7 +- torchtitan/models/kimi_k3/__init__.py | 16 +- torchtitan/models/kimi_k3/vision_encoder.py | 221 +------------------- 3 files changed, 22 insertions(+), 222 deletions(-) diff --git a/torchtitan/models/common/vision_encoder.py b/torchtitan/models/common/vision_encoder.py index 4576a4bfd9..37f446a2ac 100644 --- a/torchtitan/models/common/vision_encoder.py +++ b/torchtitan/models/common/vision_encoder.py @@ -28,7 +28,7 @@ from torchtitan.models.common import Linear from torchtitan.models.common.attention import FlexAttention, local_head_split -from torchtitan.models.common.nn_modules import GELU, LayerNorm +from torchtitan.models.common.nn_modules import GELU, LayerNorm, RMSNorm from torchtitan.protocols.module import Module compiled_create_block_mask = torch.compile(create_block_mask) @@ -150,8 +150,9 @@ class VisionTransformerBlock(Module): @dataclass(kw_only=True, slots=True) class Config(Module.Config): - norm1: LayerNorm.Config - norm2: LayerNorm.Config + # MoonViT normalizes with RMSNorm; Qwen3.5 and Muse Glimmer use LayerNorm. + norm1: LayerNorm.Config | RMSNorm.Config + norm2: LayerNorm.Config | RMSNorm.Config attn: VisionAttention.Config mlp: VisionMLP.Config diff --git a/torchtitan/models/kimi_k3/__init__.py b/torchtitan/models/kimi_k3/__init__.py index 87487a7a46..e559f1469c 100644 --- a/torchtitan/models/kimi_k3/__init__.py +++ b/torchtitan/models/kimi_k3/__init__.py @@ -16,7 +16,12 @@ from torchtitan.models.common.moe import RoutedExperts, TokenChoiceTopKRouter from torchtitan.models.common.nn_modules import GELU, RMSNorm from torchtitan.models.common.token_dispatcher import LocalTokenDispatcher -from torchtitan.models.common.vision_encoder import VisionAttention, VisionMLP +from torchtitan.models.common.vision_encoder import ( + VisionAttention, + VisionMLP, + VisionTransformerBlock, +) +from torchtitan.models.kimi_k2_7.vision_encoder import VisionRotaryEmbedding2D from torchtitan.models.utils import validate_converter_order from torchtitan.protocols.model import ModelConfigConverter from torchtitan.protocols.model_spec import ModelSpec @@ -26,12 +31,7 @@ from .moe import KimiFeedForward, KimiGroupedExperts, KimiLatentMoE from .parallelize import parallelize_kimi_k3 from .state_dict_adapter import KimiK3StateDictAdapter -from .vision_encoder import ( - KimiK3VisionBlock, - KimiK3VisionEncoder, - KimiK3VisionProjector, - VisionRotaryEmbedding2D, -) +from .vision_encoder import KimiK3VisionEncoder, KimiK3VisionProjector __all__ = [ "KIMI_K3_SPECIAL_TOKENS", @@ -283,7 +283,7 @@ def _vision_encoder_config( eps=1e-5, param_init=_NORM_INIT, ) - block = KimiK3VisionBlock.Config( + block = VisionTransformerBlock.Config( norm1=vision_norm, norm2=vision_norm, attn=VisionAttention.Config( diff --git a/torchtitan/models/kimi_k3/vision_encoder.py b/torchtitan/models/kimi_k3/vision_encoder.py index aa53154647..97430564f6 100644 --- a/torchtitan/models/kimi_k3/vision_encoder.py +++ b/torchtitan/models/kimi_k3/vision_encoder.py @@ -7,11 +7,8 @@ """MoonViT3d vision encoder used by Kimi K3. Shape suffixes: -- N = number of visual items - T = total packed patches - D = vision hidden dimension -- H = number of attention heads -- K = attention head dimension - C = number of complex-valued head-dimension pairs - M = total merged tokens - F = merged feature dimension @@ -22,220 +19,23 @@ import spmd_types as spmd import torch -import torch.nn.functional as F -from torch.nn.attention.flex_attention import BlockMask from torchtitan.models.common import Linear from torchtitan.models.common.nn_modules import GELU, RMSNorm from torchtitan.models.common.rope import ComplexRoPE from torchtitan.models.common.vision_encoder import ( create_block_diagonal_mask, - VisionAttention, - VisionMLP, + VisionTransformerBlock, +) +from torchtitan.models.kimi_k2_7.vision_encoder import ( + _compute_2d_rope_cache, + _compute_learned_pos_embeds, + _tpool_patch_merger, + VisionRotaryEmbedding2D, ) from torchtitan.protocols.module import Module, ModuleDict -def _get_temporal_pos_embed( - num_frames: int, - embed_dim: int, - *, - device: torch.device, -) -> torch.Tensor: - """Return fixed 1D sinusoidal embeddings for video frame positions.""" - grid = torch.arange(num_frames, dtype=torch.float32, device=device) - omega = torch.arange(embed_dim // 2, dtype=torch.float32, device=device) / ( - embed_dim / 2.0 - ) - omega = 1.0 / 10000.0**omega - angles = torch.outer(grid, omega) - return torch.cat((angles.sin(), angles.cos()), dim=-1) - - -def _compute_learned_pos_embeds( - pos_embed: torch.Tensor, - grids: list[list[int]], - interpolation_mode: str, - max_num_frames: int, -) -> torch.Tensor: - """Interpolate the learned 2D table and add fixed temporal embeddings.""" - height, width, dim = pos_embed.shape - pos_grid = pos_embed.permute(2, 0, 1).unsqueeze(0).float() - - cached_spatial: dict[tuple[int, int], torch.Tensor] = {} - positions = [] - for num_frames, grid_h, grid_w in grids: - if num_frames > max_num_frames: - raise ValueError( - f"Vision grid has {num_frames} frames, exceeding " - f"max_num_frames={max_num_frames}." - ) - spatial = cached_spatial.get((grid_h, grid_w)) - if spatial is None: - if (grid_h, grid_w) == (height, width): - spatial = pos_embed.flatten(end_dim=1) - else: - spatial = ( - F.interpolate( - pos_grid, - size=(grid_h, grid_w), - mode=interpolation_mode, - ) - .squeeze(0) - .permute(1, 2, 0) - .reshape(grid_h * grid_w, dim) - .to(pos_embed.dtype) - ) - cached_spatial[(grid_h, grid_w)] = spatial - - if num_frames == 1: - item_pos = spatial - else: - temporal = _get_temporal_pos_embed(num_frames, dim, device=pos_embed.device) - item_pos = spatial.unsqueeze(0) + temporal.unsqueeze(1).to(spatial.dtype) - item_pos = item_pos.reshape(num_frames * grid_h * grid_w, dim) - positions.append(item_pos) - - return torch.cat(positions) - - -def _compute_2d_rope_cache( - freq_table: torch.Tensor, - grids: list[list[int]], - head_dim: int, -) -> torch.Tensor: - """Build the real-valued 2D RoPE cache in raster patch order.""" - cached_spatial: dict[tuple[int, int], torch.Tensor] = {} - item_angles = [] - for num_frames, grid_h, grid_w in grids: - spatial = cached_spatial.get((grid_h, grid_w)) - if spatial is None: - flat = torch.arange(grid_h * grid_w, device=freq_table.device) - x_angles = freq_table[flat % grid_w] - y_angles = freq_table[flat // grid_w] - spatial = torch.stack((x_angles, y_angles), dim=-1).reshape( - grid_h * grid_w, head_dim // 2 - ) - cached_spatial[(grid_h, grid_w)] = spatial - item_angles.append(spatial.repeat(num_frames, 1)) - - angles = torch.cat(item_angles) - # ComplexRoPE.apply_rotary_emb multiplies in complex64; float() only widens - # the container, so cos/sin keep whatever precision angles were computed in. - cos_sin = torch.stack((angles.cos(), angles.sin()), dim=-1).float() - return torch.view_as_complex(cos_sin).unsqueeze(1) - - -def _temporal_pool_and_merge( - hidden_TD: torch.Tensor, - grids: list[list[int]], - merge_kernel_size: tuple[int, int], -) -> torch.Tensor: - """Temporally pool and concatenate neighboring spatial patch features.""" - dim = hidden_TD.shape[-1] - kernel_h, kernel_w = merge_kernel_size - merged_dim = kernel_h * kernel_w * dim - - merged_items = [] - offset = 0 - for num_frames, grid_h, grid_w in grids: - num_patches = num_frames * grid_h * grid_w - item = hidden_TD[offset : offset + num_patches] - offset += num_patches - merged_h = grid_h // kernel_h - merged_w = grid_w // kernel_w - item = item.view( - num_frames, - merged_h, - kernel_h, - merged_w, - kernel_w, - dim, - ) - item = item.permute(0, 1, 3, 2, 4, 5).mean(dim=0) - merged_items.append(item.reshape(merged_h * merged_w, merged_dim)) - - return torch.cat(merged_items) - - -class VisionRotaryEmbedding2D(Module): - """Per-axis frequency table for MoonViT's interleaved 2D RoPE.""" - - @dataclass(kw_only=True, slots=True) - class Config(Module.Config): - head_dim: int - theta: float = 10000.0 - - def __init__(self, config: Config): - super().__init__() - if config.head_dim % 4 != 0: - raise ValueError( - "Vision 2D RoPE head_dim must be divisible by 4, " - f"got {config.head_dim}." - ) - self.head_dim = config.head_dim - self.theta = config.theta - self.register_buffer("inv_freq", self._compute_inv_freq(), persistent=False) - - def _compute_inv_freq(self, *, device: torch.device | None = None) -> torch.Tensor: - return 1.0 / ( - self.theta - ** ( - torch.arange( - 0, - self.head_dim, - 4, - dtype=torch.float32, - device=device, - ) - / self.head_dim - ) - ) - - def _init_self_buffers(self, *, buffer_device: torch.device | None = None) -> None: - device = buffer_device or self.inv_freq.device - self.inv_freq = self._compute_inv_freq(device=device) - - def forward(self, seqlen: int) -> torch.Tensor: - positions = torch.arange( - seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype - ) - return torch.outer(positions, self.inv_freq) - - -class KimiK3VisionBlock(Module): - """MoonViT pre-norm attention and MLP block.""" - - @dataclass(kw_only=True, slots=True) - class Config(Module.Config): - norm1: RMSNorm.Config - norm2: RMSNorm.Config - attn: VisionAttention.Config - mlp: VisionMLP.Config - - def __init__(self, config: Config): - super().__init__() - self.norm1 = config.norm1.build() - self.norm2 = config.norm2.build() - self.attn = config.attn.build() - self.mlp = config.mlp.build() - - def forward( - self, - x_TD: torch.Tensor, - *, - rope_cache: torch.Tensor, - attention_mask: BlockMask, - ) -> torch.Tensor: - x_TD = x_TD + self.attn( - self.norm1(x_TD), - rope_cache=rope_cache, - rope_apply=ComplexRoPE.apply_rotary_emb, - attention_mask=attention_mask, - ) - return x_TD + self.mlp(self.norm2(x_TD)) - - class KimiK3VisionProjector(Module): """PatchMergerMLPV2 projector from merged vision features to text width.""" @@ -272,14 +72,13 @@ class Config(Module.Config): interpolation_mode: str patch_embed_proj: Linear.Config rotary_pos_emb: VisionRotaryEmbedding2D.Config - block: KimiK3VisionBlock.Config + block: VisionTransformerBlock.Config final_norm: RMSNorm.Config projector: KimiK3VisionProjector.Config def __init__(self, config: Config): super().__init__() self.merge_kernel_size = config.merge_kernel_size - self.max_num_frames = config.max_num_frames self.interpolation_mode = config.interpolation_mode self.patch_embed = config.patch_embed_proj.build() self.pos_embed = torch.nn.Parameter( @@ -313,7 +112,6 @@ def _compute_position_embeddings( self.pos_embed, grids, self.interpolation_mode, - self.max_num_frames, ) rope_cache = _compute_2d_rope_cache( self._cached_freq_table, @@ -361,8 +159,9 @@ def forward( hidden_TD = block( hidden_TD, rope_cache=rope_cache, + rope_apply=ComplexRoPE.apply_rotary_emb, attention_mask=attention_mask, ) hidden_TD = self.final_norm(hidden_TD) - merged_MF = _temporal_pool_and_merge(hidden_TD, grids, self.merge_kernel_size) + merged_MF = _tpool_patch_merger(hidden_TD, grids, self.merge_kernel_size) return self.projector(merged_MF) From ded1bf1093f1509662187a1041555714eced8c66 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Sun, 23 Aug 2026 13:57:08 +0000 Subject: [PATCH 64/67] share the MoonViT tower between k2.7 and k3 --- torchtitan/models/kimi_k2_7/vision_encoder.py | 49 ++++--- torchtitan/models/kimi_k3/vision_encoder.py | 125 +----------------- 2 files changed, 40 insertions(+), 134 deletions(-) diff --git a/torchtitan/models/kimi_k2_7/vision_encoder.py b/torchtitan/models/kimi_k2_7/vision_encoder.py index b3b00ccbab..7041ed9182 100644 --- a/torchtitan/models/kimi_k2_7/vision_encoder.py +++ b/torchtitan/models/kimi_k2_7/vision_encoder.py @@ -27,7 +27,7 @@ from torchtitan.distributed.utils import get_spmd_backend from torchtitan.models.common import Linear -from torchtitan.models.common.nn_modules import GELU, LayerNorm +from torchtitan.models.common.nn_modules import GELU, LayerNorm, RMSNorm from torchtitan.models.common.rope import _maybe_wrap_positions, ComplexRoPE from torchtitan.models.common.vision_encoder import ( create_block_diagonal_mask, @@ -325,31 +325,26 @@ def forward(self, merged_MK: torch.Tensor) -> torch.Tensor: return self.linear_2(x) -class KimiK25VisionEncoder(Module): - """MoonViT3d vision tower + multimodal projector for Kimi K2.5.""" +class MoonViTEncoder(Module): + """MoonViT3d vision tower + multimodal projector.""" @dataclass(kw_only=True, slots=True) class Config(Module.Config): - dim: int = 1152 - num_layers: int = 27 - num_heads: int = 16 - - patch_size: int = 14 - in_channels: int = 3 - merge_kernel_size: list[int] = field(default_factory=lambda: [2, 2]) - text_hidden_size: int = 7168 + dim: int + num_layers: int + merge_kernel_size: list[int] # Learnable 2D spatial position table, shape (height, width, dim). - init_pos_emb_height: int = 64 - init_pos_emb_width: int = 64 - interpolation_mode: str = "bicubic" + init_pos_emb_height: int + init_pos_emb_width: int + interpolation_mode: str # Sub-modules. patch_embed_proj: Linear.Config rotary_pos_emb: VisionRotaryEmbedding2D.Config block: VisionTransformerBlock.Config - final_norm: LayerNorm.Config - projector: VisionProjector.Config + final_norm: LayerNorm.Config | RMSNorm.Config + projector: Module.Config def __init__(self, config: Config): super().__init__() @@ -472,3 +467,25 @@ def forward( # pyrefly: ignore [bad-argument-type] merged = _tpool_patch_merger(x, grids, self.merge_kernel_size) return self.projector(merged) + + +class KimiK25VisionEncoder(MoonViTEncoder): + """MoonViT3d vision tower + multimodal projector for Kimi K2.5.""" + + @dataclass(kw_only=True, slots=True) + class Config(MoonViTEncoder.Config): + dim: int = 1152 + num_layers: int = 27 + num_heads: int = 16 + + patch_size: int = 14 + in_channels: int = 3 + merge_kernel_size: list[int] = field(default_factory=lambda: [2, 2]) + text_hidden_size: int = 7168 + + init_pos_emb_height: int = 64 + init_pos_emb_width: int = 64 + interpolation_mode: str = "bicubic" + + final_norm: LayerNorm.Config # pyrefly: ignore [bad-override] + projector: VisionProjector.Config # pyrefly: ignore [bad-override] diff --git a/torchtitan/models/kimi_k3/vision_encoder.py b/torchtitan/models/kimi_k3/vision_encoder.py index 97430564f6..6c480c8815 100644 --- a/torchtitan/models/kimi_k3/vision_encoder.py +++ b/torchtitan/models/kimi_k3/vision_encoder.py @@ -7,9 +7,6 @@ """MoonViT3d vision encoder used by Kimi K3. Shape suffixes: -- T = total packed patches -- D = vision hidden dimension -- C = number of complex-valued head-dimension pairs - M = total merged tokens - F = merged feature dimension - O = projected text dimension @@ -17,23 +14,12 @@ from dataclasses import dataclass, field -import spmd_types as spmd import torch from torchtitan.models.common import Linear from torchtitan.models.common.nn_modules import GELU, RMSNorm -from torchtitan.models.common.rope import ComplexRoPE -from torchtitan.models.common.vision_encoder import ( - create_block_diagonal_mask, - VisionTransformerBlock, -) -from torchtitan.models.kimi_k2_7.vision_encoder import ( - _compute_2d_rope_cache, - _compute_learned_pos_embeds, - _tpool_patch_merger, - VisionRotaryEmbedding2D, -) -from torchtitan.protocols.module import Module, ModuleDict +from torchtitan.models.kimi_k2_7.vision_encoder import MoonViTEncoder +from torchtitan.protocols.module import Module class KimiK3VisionProjector(Module): @@ -58,110 +44,13 @@ def forward(self, merged_MF: torch.Tensor) -> torch.Tensor: return self.post_norm(projected_MO) -class KimiK3VisionEncoder(Module): +class KimiK3VisionEncoder(MoonViTEncoder): @dataclass(kw_only=True, slots=True) - class Config(Module.Config): - dim: int - num_layers: int + class Config(MoonViTEncoder.Config): patch_size: int in_channels: int - merge_kernel_size: tuple[int, int] - init_pos_emb_height: int - init_pos_emb_width: int + merge_kernel_size: tuple[int, int] # pyrefly: ignore [bad-override] max_num_frames: int - interpolation_mode: str - patch_embed_proj: Linear.Config - rotary_pos_emb: VisionRotaryEmbedding2D.Config - block: VisionTransformerBlock.Config - final_norm: RMSNorm.Config - projector: KimiK3VisionProjector.Config - - def __init__(self, config: Config): - super().__init__() - self.merge_kernel_size = config.merge_kernel_size - self.interpolation_mode = config.interpolation_mode - self.patch_embed = config.patch_embed_proj.build() - self.pos_embed = torch.nn.Parameter( - torch.empty( - config.init_pos_emb_height, - config.init_pos_emb_width, - config.dim, - ) - ) - self.rotary_pos_emb = config.rotary_pos_emb.build() - self.register_buffer("_cached_freq_table", None, persistent=False) - self.layers = ModuleDict( - { - str(layer_idx): config.block.build() - for layer_idx in range(config.num_layers) - } - ) - self.final_norm = config.final_norm.build() - self.projector = config.projector.build() - - def _compute_position_embeddings( - self, grids: list[list[int]] - ) -> tuple[torch.Tensor, torch.Tensor]: - max_grid_side = max(max(grid_h, grid_w) for _, grid_h, grid_w in grids) - if ( - self._cached_freq_table is None - or self._cached_freq_table.shape[0] < max_grid_side - ): - self._cached_freq_table = self.rotary_pos_emb(max_grid_side) - learned_pos = _compute_learned_pos_embeds( - self.pos_embed, - grids, - self.interpolation_mode, - ) - rope_cache = _compute_2d_rope_cache( - self._cached_freq_table, - grids, - self.rotary_pos_emb.head_dim, - ) - return learned_pos, rope_cache - - def forward( - self, - pixel_values: torch.Tensor, - *, - grid_thw: torch.Tensor, - ) -> torch.Tensor: - """Encode packed raster-order patches and return packed text features.""" - grids = grid_thw.tolist() - - kernel_h, kernel_w = self.merge_kernel_size - for _, grid_h, grid_w in grids: - if grid_h % kernel_h != 0 or grid_w % kernel_w != 0: - raise ValueError( - f"Vision grid {grid_h}x{grid_w} is not divisible by " - f"merge kernel {self.merge_kernel_size}." - ) - - segment_lengths = grid_thw.prod(dim=-1) - total_tokens = pixel_values.shape[0] - expected_tokens = sum(t * h * w for t, h, w in grids) - if total_tokens != expected_tokens: - raise ValueError( - f"pixel_values contains {total_tokens} patches but grid_thw " - f"describes {expected_tokens}." - ) - - learned_pos, rope_cache = self._compute_position_embeddings(grids) - hidden_TD = self.patch_embed(pixel_values) + learned_pos - with spmd.no_typecheck(): - attention_mask = create_block_diagonal_mask( - segment_lengths, - total_tokens, - hidden_TD.device, - ) - for block in self.layers.values(): - hidden_TD = block( - hidden_TD, - rope_cache=rope_cache, - rope_apply=ComplexRoPE.apply_rotary_emb, - attention_mask=attention_mask, - ) - hidden_TD = self.final_norm(hidden_TD) - merged_MF = _tpool_patch_merger(hidden_TD, grids, self.merge_kernel_size) - return self.projector(merged_MF) + final_norm: RMSNorm.Config # pyrefly: ignore [bad-override] + projector: KimiK3VisionProjector.Config # pyrefly: ignore [bad-override] From b3751c9d7acd3c6643cee83a00fd368aa7200e52 Mon Sep 17 00:00:00 2001 From: JavaZero <71128095+JavaZeroo@users.noreply.github.com> Date: Mon, 24 Aug 2026 15:30:57 +0800 Subject: [PATCH 65/67] Update .github/workflows/integration_test_8gpu_features.yaml Co-authored-by: Shuhua Yu <18108279+shuhuayu@users.noreply.github.com> --- .github/workflows/integration_test_8gpu_features.yaml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.github/workflows/integration_test_8gpu_features.yaml b/.github/workflows/integration_test_8gpu_features.yaml index 873fb37224..0126442354 100644 --- a/.github/workflows/integration_test_8gpu_features.yaml +++ b/.github/workflows/integration_test_8gpu_features.yaml @@ -68,7 +68,8 @@ jobs: TORCH_SPEC="torch==${{ matrix.torch-version }}" fi python -m pip install --force-reinstall --pre \ - "${TORCH_SPEC}" --index-url ${{ matrix.index-url }} + python -m pip install --force-reinstall --pre \ + "${TORCH_SPEC}" torchvision --index-url ${{ matrix.index-url }} # The torchcomms feature tests are currently disabled, so do not install # torchcomms in the main feature job. Its wheel pins torch exactly and From 2dd93907b25cff5c3356e658e2ca404247d9290e Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Mon, 24 Aug 2026 07:33:22 +0000 Subject: [PATCH 66/67] fix cicd --- .github/workflows/integration_test_8gpu_features.yaml | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/.github/workflows/integration_test_8gpu_features.yaml b/.github/workflows/integration_test_8gpu_features.yaml index 0126442354..157029a7c7 100644 --- a/.github/workflows/integration_test_8gpu_features.yaml +++ b/.github/workflows/integration_test_8gpu_features.yaml @@ -68,8 +68,7 @@ jobs: TORCH_SPEC="torch==${{ matrix.torch-version }}" fi python -m pip install --force-reinstall --pre \ - python -m pip install --force-reinstall --pre \ - "${TORCH_SPEC}" torchvision --index-url ${{ matrix.index-url }} + "${TORCH_SPEC}" torchvision --index-url ${{ matrix.index-url }} # The torchcomms feature tests are currently disabled, so do not install # torchcomms in the main feature job. Its wheel pins torch exactly and From 88cb8003663a0167a08b3aa3e11962efa6ff2560 Mon Sep 17 00:00:00 2001 From: Shuhua Yu <18108279+shuhuayu@users.noreply.github.com> Date: Mon, 24 Aug 2026 02:52:57 -0700 Subject: [PATCH 67/67] Update scripts/checkpoint_conversion/numerical_tests_kimi_k3.py --- scripts/checkpoint_conversion/numerical_tests_kimi_k3.py | 1 - 1 file changed, 1 deletion(-) diff --git a/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py b/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py index 9ee6d9ba1c..e89210606a 100644 --- a/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py +++ b/scripts/checkpoint_conversion/numerical_tests_kimi_k3.py @@ -338,7 +338,6 @@ def run_tt( print(f"Loading TorchTitan Kimi K3 (debugmodel) on {device} ...") model.to(device) assert model.vision_encoder is not None - model.vision_encoder.to(vision_dtype) if force_hf_routing: print("Using HF expert selections with TorchTitan router scores")