[kimi k2_7] add kimi k2_7 by shuhuayu · Pull Request #3532 · pytorch/torchtitan

shuhuayu · 2026-06-04T20:35:00Z

Summary:

Adds the Kimi K2.7 model family — a DeepSeek-V3 (MLA + sigmoid-routed MoE) decoder paired with a MoonViT3d vision encoder. It maximizes reuse of existing torchtitan code: KimiK25Model extends DeepSeekV3Model, so the decoder inherits DeepSeek-V3's model, sharding, and state-dict adapter unchanged; the vision encoder, the vision to text scatter, and the vision-side adapter keys are the new parts.

This is a long pr, hope the following the partition can help reviewers.

Part 1 — Shared primitives (foundational, small)

torchtitan/models/common/multimodal.py — the scatter (get_vision_positions / scatter_vision_embeds). Tiny, used by both models.
torchtitan/models/common/vision_encoder.py — shared ViT blocks (VisionAttention/VisionMLP/VisionTransformerBlock).
torchtitan/distributed/fsdp.py — apply_fsdp_to_vision_encoder (new function).
torchtitan/distributed/pipeline_parallel.py — pipeline_vlm (new function).
torchtitan/models/deepseek_v3/init.py — generalized build_mla_moe_layers / make_mla_attention_config to take init callables + RoPE, so deepseek_v3 and kimi share the layer builder.
torchtitan/models/utils.py — get_moe_model_nparams_and_flops excludes the vision encoder from the per-token FLOP term (counts its params, not FLOPs).

Part 2 — Shared data pipeline

torchtitan/hf_datasets/multimodal/utils/image.py — resize_fn + vision_to_patches(patch_order=…).
torchtitan/hf_datasets/multimodal/mm_collator.py — patch_order threading + the MRoPE guard.
torchtitan/hf_datasets/multimodal/mm_datasets.py — Config.patch_order; removed resize_fn/budget defaults.

Part 3 — Kimi model

torchtitan/models/kimi_k2_7/vision_encoder.py — MoonViT3d (linear patch embed, 2D RoPE, temporal pool + 2×2 merge, projector).
torchtitan/models/kimi_k2_7/model.py — forward + scatter wiring + mixed-media guard.
torchtitan/models/kimi_k2_7/sharding.py — TP/SP/EP plan (incl. the embedding-Replicate-for-scatter SP trick).
torchtitan/models/kimi_k2_7/parallelize.py — applies parallelisms; delegates to pipeline_vlm.
torchtitan/models/kimi_k2_7/state_dict_adapter.py — checkpoint mapping (both projector spellings, fused-QKV split, Conv2d→Linear patch embed; the one piece not unit-tested in this PR).
torchtitan/models/kimi_k2_7/config_registry.py — flavors + _mm_dataloader (patch_order="raster").
torchtitan/models/kimi_k2_7/init.py — flavors, special tokens, pipeline_vlm wiring.
torchtitan/models/init.py — register kimi_k2_7 in the top-level model registry.

Part 4 — Qwen (shared-refactor blast radius; confirm no regression)

torchtitan/models/qwen3_5/model.py — scatter rewired to the shared helper.
torchtitan/models/qwen3_5/parallelize.py — FSDP + pipeline_vlm consolidation.
torchtitan/models/qwen3_5/init.py — pipeline_vlm registration.
torchtitan/models/qwen3_5/vision_encoder.py — minor (only the grid_thw host-int doc/imports).

Part 5 — Tests

tests/unit_tests/test_vision_encoder.py — shared scatter.
tests/unit_tests/test_mm_dataset_preprocessing.py — patch-budget resize + block/raster order.
tests/integration_tests/models.py — the kimi case (FSDP+TP+EP+PP, exercises SP).
tests/assets/tokenizer/{tokenizer,tokenizer_config}.json — adds Kimi media special tokens to the shared test tokenizer.
scripts/checkpoint_conversion/numerical_tests_kimi.py — the PR'd e2e numerics (HF Kimi-VL parity).

Part 6 — Docs

torchtitan/models/kimi_k2_7/README.md

Verification

HF Kimi-VL parity (float32, last-token logits): vision cosine 0.999977; KL 4.3e-2 (top-1 ✓, top-5 4/5), 5.3e-4 (top-1 ✓, top-5 5/5) with routing pinned to HF's selections.
Parallelism: bit-identical logits (max diff 0.0) for no-parallel / FSDP / FSDP+EP; within bf16 tolerance for FSDP+EP+TP (with SP).

weifengpy · 2026-06-04T20:50:28Z

+# 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 for Kimi K2.5.


this one has vision encoder - is it literally kimi 2.6?

yes, this is k2.5 which share the same architecture of k2.6, you can use this to load kimi k2.6 weights.

tianyu-l

quality seems very high, left some minor comments

tianyu-l · 2026-07-23T21:48:21Z

+      "<|media_pad|>": 2016,
+      "<|media_begin|>": 2017,
+      "<|media_end|>": 2018,
+      "[PAD]": 2019


add a comment on for which model they are introduced

Would be good to backfill the info for previous tokens (e.g. Qwen 3.5, RL, etc.)

agreed, but met a mechanical issue that this json file is loaded by a strict parser and does not allow inline comments starting with // and i tested it locally.

tianyu-l · 2026-07-23T21:49:07Z

            "gpt_oss_pp+fsdp+ep+sacop",
            ngpu=8,
        ),
+        # Integration Test Cases for Kimi K2.5


If K2.7 is sharing the same code, we should use the latest naming.

kimi 2.5, 2.6, 2.7 all share the same code, updated to kimi 2.7 to be latest. module wise, we still stick to kimi 2_5, matching the convention for kimi models on huggingface.

tianyu-l · 2026-07-23T21:51:47Z

        spatial_merge_size: int
        """Spatially merge visual tokens after encoder. e.g. 2 means 2x2=4 patches merged."""

+        patch_order: str = "block"


type should be a Literal?

indeed, otherwise will accept invalid strings not in "block" or "raster".

tianyu-l · 2026-07-23T21:53:09Z

+        ``spatial_merge_size**2`` group contiguous, or ``"raster"`` (row-major).
+        Must be ``"block"`` when ``build_mrope_positions`` is set."""
+
+        resize_fn: Callable[..., tuple[int, int, int, int]] = resize_to_pixel_budget


this may not be tyro-friendly, do we need to do some extra annotation to suppress warning / failure?

nice catch. used Annotated to suppress warning.

tianyu-l · 2026-07-23T21:58:26Z

+    "moonlight-16B-A3B": _moonlight_16b_a3b,
+    "Kimi-VL-A3B": _kimi_vl_a3b,
+    "1T-A32B": _1t_a32b,


names are confusing, is there a pattern?

i was following the naming in https://huggingface.co/moonshotai/collections. The reasoning is that i found we can reuse the code for kimi 2.5 to config the other two kimi models, moonlight-16b-a3b, and kimi-vl-a3b. i keep them because it provides users with more model choices if the 1t kimi 2.5-2.7 is too big for them.

where is 32 coming from in "A32B" ? I see "A3B" from "Kimi-VL-A3B" and "Moonlight-A3B"

I was thinking about 3 models from https://huggingface.co/collections/moonshotai/kimi-k25

moonshotai/Kimi-K2.5

moonshotai/Kimi-K2.6

moonshotai/Kimi-K2.7-Code

it's from the model summary, section 2 table: https://huggingface.co/moonshotai/Kimi-K2.7-Code

maybe changing to Kimi-K2_5 for the last one? more direct, and K2.6, K2.7-code share the same code.

it's from the model summary, section 2 table: https://huggingface.co/moonshotai/Kimi-K2.7-Code

do you mean 32B activated parameters? is there a convention to use activated parameters for model name?

yeah, like for qwen moe series: https://huggingface.co/collections/Qwen/qwen35, and Kimi-VL-A3B, meaning activated parameter number is 3b. now i am inclined to using Kimi-K2_5 to reduce confusion, since this is the name for the model.

TIL what A3B means

tianyu-l · 2026-07-23T22:00:05Z

+    NOTE: the passed-in model should preferably be on meta device; otherwise it
+    must fit in GPU or CPU memory.
+    """
+    if parallelism.spmd_backend == "full_dtensor":


Is spmd_types backend supported? If possible, that should be the only backend we support onwards. Also fine to wait for #3895

no. i would suggest doing this in a batch with qwen 3_5 and other models.

I published a PR for review to unblock myself from benchmarking kimi k2.7 #4063

tianyu-l · 2026-07-23T22:00:25Z

+    if parallel_dims.cp_enabled:
+        raise NotImplementedError(
+            "Context Parallel is not yet supported for Kimi K2.5: vision scatter "
+            "needs the full sequence before CP would shard it."
+        )


cc @jinsooihm

do we have any plans to support cp for vlms soon?

tianyu-l · 2026-07-23T22:06:42Z

    nparams_shared_experts = 0
    nparams_experts = 0
    nparams_dense = 0
+    # TODO: add a per-batch vision encoder FLOP term for accurate VLM MFU.


oh good point, didn't realize we were computing it wrong

right, opened an issue to track: #4009

weifengpy · 2026-07-29T06:28:55Z

+        attn_backend="flex",
+        converters=[
+            Float8LinearConverter.Config(
+                filter_fqns=["output", "router.gate"],


do you mean "lm_head" as the fqn ?

nice catch. i copied from dsv3, it looks like dsv3, llama3 both had the same bug.

made a pr to fix this for llama3 and dsv3 first: #4008. this is from a shared decoder refactor.

i removed this quantization for kimi_k2_5 after checking the technical report, where FP8-E4M3 is only used to compress selected saved activations, not for GEMMs.

weifengpy · 2026-07-29T06:45:30Z

+        if self.vision_encoder is None:
+            return super().to_hf(state_dict)
+
+        to_hf_map = {v: k for k, v in self.vision_from_hf_map.items()}


from codex

Kimi-VL checkpoint loading still emits mm_projector.* before DCP load; multi_modal_projector.* normalization happens afterward.

but honestly I don't know why we import models from HF "Kimi-VL-A3B" collection
I was assuming K2.5 - 2.7 referring to https://huggingface.co/collections/moonshotai/kimi-k25 by default. but correct me if i am wrong

these are smaller models that can reuse the kimi k2.5 code, we used a smaller model to validate numerics, and i think these smaller models are also useful for users.

from codex

Kimi-VL checkpoint loading still emits mm_projector.* before DCP load; multi_modal_projector.* normalization happens afterward.

this is indeed a problem, updated.

tianyu-l · 2026-07-29T09:13:29Z

kimi 2.5, 2.6, 2.7 all share the same code, updated to kimi 2.7 to be latest. module wise, we still stick to kimi 2_5, matching the convention for kimi models on huggingface.

Module name sounds fine, but can we modify the folder name to 2.7? I feel it is more eye-catching.

tianyu-l · 2026-07-29T09:14:36Z

@@ -0,0 +1,61 @@
+# Kimi K2.5


Suggested change

# Kimi K2.5

# Kimi K2.7

sounds good.

shuhuayu · 2026-07-29T20:10:18Z

made a test run of the kimi-vl-a3b variant on cc12m text + image datasets on 64 h100s, with dp=64 and ep=8, loss decreases normally.

i'll merge this now and iterate on the todos later.

shuhuayu requested review from fegin, tianyu-l, wconstab and wwwjn as code owners June 4, 2026 20:35

pytorch-bot Bot added the ciflow/8gpu label Jun 4, 2026

meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Jun 4, 2026

shuhuayu marked this pull request as draft June 4, 2026 20:35

weifengpy reviewed Jun 4, 2026

View reviewed changes

shuhuayu changed the title ~~[kimi 2_5] add kimi 2_5~~ [kimi k2_5] add kimi k2_5 Jun 25, 2026

shuhuayu force-pushed the atom branch 3 times, most recently from 601d1b6 to 0ac24c5 Compare June 28, 2026 06:16

shuhuayu marked this pull request as ready for review June 28, 2026 06:27

tianyu-l reviewed Jul 23, 2026

View reviewed changes

shuhuayu added 2 commits July 28, 2026 20:47

add kimi 2_5

998e18c

refactor and add tests

cee8dd5

shuhuayu force-pushed the atom branch 3 times, most recently from 4d71750 to e79df70 Compare July 29, 2026 06:11

weifengpy reviewed Jul 29, 2026

View reviewed changes

shuhuayu force-pushed the atom branch from e79df70 to 0142dae Compare July 29, 2026 07:22

polish

c7c7922

shuhuayu mentioned this pull request Jul 29, 2026

VLM MFU is miscomputed: vision-encoder FLOPs billed per decoder token, not per image patch #4009

Open

tianyu-l approved these changes Jul 29, 2026

View reviewed changes

weifengpy approved these changes Jul 29, 2026

View reviewed changes

shuhuayu changed the title ~~[kimi k2_5] add kimi k2_5~~ [kimi k2_7] add kimi k2_7 Jul 29, 2026

shuhuayu force-pushed the atom branch 2 times, most recently from b217cb5 to aa4ac72 Compare July 29, 2026 09:57

address comments and bug fixes

8c432cc

shuhuayu force-pushed the atom branch from aa4ac72 to 8c432cc Compare July 29, 2026 10:00

update readme

acf0df7

shuhuayu merged commit 85c549b into pytorch:main Jul 29, 2026
16 of 19 checks passed

Uh oh!

Conversation

shuhuayu commented Jun 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary:

Part 1 — Shared primitives (foundational, small)

Part 2 — Shared data pipeline

Part 3 — Kimi model

Part 4 — Qwen (shared-refactor blast radius; confirm no regression)

Part 5 — Tests

Part 6 — Docs

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tianyu-l left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

shuhuayu Jul 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

shuhuayu Jul 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

shuhuayu commented Jun 4, 2026 •

edited

Loading

shuhuayu Jul 29, 2026 •

edited

Loading

shuhuayu Jul 29, 2026 •

edited

Loading