[kimi k3] add eager reference model with FSDP2 by JavaZeroo · Pull Request #4025 · pytorch/torchtitan

JavaZeroo · 2026-07-30T07:49:26Z

Summary

Adds a PyTorch-native eager reference implementation of Kimi K3 together with
FSDP2/HSDP data parallelism.

Implements the hybrid KDA/MLA decoder, attention residuals, LatentMoE,
MoonViT-V2 vision encoder, multimodal projector, and image-feature scatter.
Reuses existing torchtitan components: Decoder, MoE,
TokenChoiceTopKRouter, RoutedExperts, GroupedExperts,
LocalTokenDispatcher, FeedForward, VisionAttention, VisionMLP,
ComplexRoPE.apply_rotary_emb, the shared vision block-mask helpers, and the
shared multimodal scatter. KimiFeedForward, KimiGroupedExperts, and
KimiLatentMoE each subclass the corresponding common class and override
only what Kimi actually changes -- the SiTU-GLU activation and the latent
expert projection.
Uses FLA's chunk_kda for the training KDA path, and keeps an explicit
recurrent implementation in the unit test as its numerical reference.

Why

Related to RFC #3029, which tracks the broader Kimi K3 pre-training,
post-training, and multi-dimensional parallelism effort.

Changes outside `models/kimi_k3/`

models/common/decoder.py: update_from_config assumed every attention
config carries a RoPE. Kimi K3's MLA sets mla_use_nope=True and has no RoPE
at all, so the sequence-length check and the cache resize are skipped when
there is none. No behavior change for models that do have RoPE.
tests/unit_tests/test_no_new_cli_options.py: registers kimi_k3 in
_GUARDED_CONFIGS. test_every_model_is_guarded requires an entry for every
model in _supported_models. The freeze snapshot passes, so this model
introduces no new command-line options.

Numerical validation

scripts/checkpoint_conversion/numerical_tests_kimi_k3.py compares the full
text+image path against the released HuggingFace implementation. The script
downloads the model from a pinned HuggingFace revision, reduces that config to
TorchTitan's debugmodel, and loads the randomly initialized TorchTitan state
dict into the HuggingFace model.

Float32:

pixel preprocessing max difference: 1.192e-7 (0 of 338688 above 1e-6)
vision feature cosine / max difference: 1.000000 / 2.730e-3
routing choices: 7869 / 7872 match
last-token logits: |KL| 6.76e-7, top-1 match, top-5 5 / 5.

The HF model, the TorchTitan text model, and the TorchTitan vision encoder can
also each use a different dtype:

HF / text / vision dtype	Routing match	\|Logits KL\|	Top-1	Top-5
float32 / float32 / float32	100.0%	`6.76e-7`	yes	5/5
bfloat16 / bfloat16 / bfloat16	98.5%	`1.20e-5`	yes	5/5
float16 / float16 / float16	99.7%	`8.21e-7`	yes	5/5
float32 / bfloat16 / float16	99.5%	`5.20e-6`	yes	5/5

Parallelism validation

debugmodel, same lr, 200 steps, comparing single-device / FSDP / HSDP losses:

Config	Loss at step 200	Per-step difference vs single
single	12.2693	--
FSDP	12.2744	mean `-0.0004`, std `0.0027`
HSDP	12.2745	mean `-0.0005`, std `0.0027`

The per-step differences are zero-mean. Bitwise logit comparison across
parallelism configurations is not done yet.

Current scope

Eager execution and FSDP2/HSDP data parallelism on the partial_dtensor
backend are supported.
TP, EP, PP, CP, full_dtensor, and spmd_types are not.
Image inputs are supported; video inputs are not.
torch.compile, packed documents, and generation cache are not supported yet.
The Kimi-K3 flavor carries the released 93-layer topology (2.78T
parameters) so parallelism work can build against it, but no trainer config
is registered for it yet.
MoonViT still overlaps with kimi_k2_7 (learned position-embedding
interpolation, the 2D RoPE frequency table, the temporal patch merger). The
two differ in in-place versus out-of-place construction, so unifying them in
models/common/ needs a numerics re-check on both models; follow-up.

pytorch-bot · 2026-07-30T07:49:34Z

The following ciflow label(s) have been added but CI has not been triggered yet because the workflows are awaiting approval:

ciflow/8gpu

Once a maintainer approves the workflows (scroll to the bottom of the PR page), the corresponding CI jobs will be triggered automatically. Please ping one of the reviewers if you do not have access to approve and run workflows.

shuhuayu · 2026-07-31T00:27:19Z

@JavaZeroo Thanks for the pr. It looks workable to me in general, have some general comments, do you have bandwidth to work on them timely?

Can you reuse more common modules in torchtitan today, like MoonViT in common/vision_encoder, common/multi_modal, and reuse MoE router.
this pr should include a reproducible numerical parity check test, likely using some reduced hf reference, since the original hf ckpt is too big.
The current code only supports fsdp2, is there bottleneck to support other parallelisms?
Currently fsdps wrap vision encoder independently, but vit is used conditionally, this will be problematic of one dp rank gets text only but another dp rank gets images.

shuhuayu · 2026-07-31T00:29:03Z

@claude please review.

claude · 2026-07-31T00:29:29Z

Claude finished @shuhuayu's task in 4m 3s —— View job

Code Review: Kimi K3 eager reference + FSDP2

I read through model.py, vision_encoder.py, parallelize.py, state_dict_adapter.py, config_registry.py, both unit tests, and the relevant shared infra (common/moe.py, common/vision_encoder.py, common/multimodal.py, distributed/fsdp.py, components/optimizer.py). Overall this is a clean, well-documented eager reference that follows the repo's shape-suffix and config conventions and already reuses Decoder, VisionMLP, get_vision_positions, and the shared FSDP helpers. Below I fold in @shuhuayu's four questions plus concrete findings.

1. Reuse of common modules (maintainer Q1) — largely agree, several concrete duplications

MoE router. KimiMoERouter (model.py:426-460) re-implements what common/moe.py::TokenChoiceTopKRouter already does: score_func="sigmoid", route_norm, route_scale, and expert_bias_E-based routing with topk(sorted=False) + gather. The only real difference is that Kimi computes the gate in fp32 via F.linear(x.float(), self.gate.weight.float()) instead of an autocast(float32) around self.gate(x). That is a very thin reason to fork the router — recommend reusing TokenChoiceTopKRouter (it already returns scores_BLE, which you need for the routing map) and, if the fp32 gate matters for HF parity, adding that as an option there rather than a parallel class.
Vision block. KimiK3VisionBlock (vision_encoder.py:320-349) is structurally identical to common/vision_encoder.py::VisionTransformerBlock (norm -> attn -> residual -> norm -> mlp). The reason for the fork is the eager, block-diagonal attention vs. the shared FlexAttention path. The shared VisionAttention already injects RoPE through a rope_apply callable and takes an inner_attention config — an eager block-diagonal attention module implementing that same interface would let you reuse VisionTransformerBlock/VisionAttention verbatim. Worth attempting; if the padded per-item Python loop genuinely can't fit the attention_mask: BlockMask contract, document why in the module docstring.
Vision embed scatter. _replace_vision_embeds (model.py:38-65) duplicates common/multimodal.py::scatter_vision_embeds. The only difference is out-of-place index_copy vs. in-place assignment (added in commit cb8d92c). This belongs in common/multimodal.py as the canonical out-of-place variant (or make the existing one autograd-safe), not a private copy in the model — other VLMs will want the same. Fix this →
KimiRMSNorm subclassing RMSNorm just to expose kimi_eps and force the fp32 reduction is reasonable, but check whether common/nn_modules.py::RMSNorm already reduces in fp32 — if so the subclass may be unnecessary.

2. Reproducible numerical parity test (maintainer Q2) — missing, should be added

The PR body reports strong parity numbers (text logits max abs err 1.7e-4, exact routed-expert IDs), but there is no committed test that reproduces them. test_kimi_k3.py covers topology, the KDA kernel vs. a local recurrence, GELU, and a to_hf/from_hf round-trip — all good — but nothing pins numerics against an HF reference. As @shuhuayu notes, the full checkpoint is too big, so the right move is a tiny randomly-initialized config loaded into both the TorchTitan module and a minimal HF-equivalent reference (or a saved golden-logits tensor) with an explicit atol/rtol. Without this, silent math drift in future refactors won't be caught, which is exactly the guarantee CLAUDE.md asks for.

3. Bottlenecks for other parallelisms (maintainer Q3)

The eager reference has two structural blockers to TP/CP/PP/EP beyond "not wired up yet":

Data-dependent, per-token Python loops. The KDA recurrence (KimiKDAKernel, model.py:317-330) loops over sequence length, and KimiLatentMoE.forward (model.py:548-558) loops over experts with torch.nonzero/index_select. These are inherently incompatible with torch.compile and CP, and the list-of-modules expert layout can't use the grouped-GEMM EP path in common/moe.py. EP specifically will require moving to GroupedExperts (or an equivalent grouped layout) — the list-backed KimiRoutedExperts is fine for a reference but is a dead end for EP.
KimiK3Model.forward overrides Decoder.forward with an incompatible signature (model.py:802, # pyrefly: ignore [bad-override]) and threads a block_residual_TND tuple through every layer. PP's pipelining stage-splitting assumes the standard decoder contract, so PP will need this reworked. Reasonable to defer, but worth stating explicitly in the README's "Initial scope" as the reason, not just "rejected."

4. Vision encoder FSDP + conditional execution (maintainer Q4) — real hang risk, agree

This is the most important correctness concern. parallelize.py:76-85 wraps vision_encoder as its own FSDP unit, but KimiK3Model.forward only calls it when pixel_values is not None (model.py:774-783). Under FSDP2 the unshard all-gather (forward) and reduce-scatter (backward) for the vision params only fire on ranks that actually run the encoder. If DP rank A gets an image batch and rank B gets a text-only batch, the collectives are mismatched and training deadlocks. The kimi_k3_mm_fsdp integration test never exercises this because the cc12m-test dataloader always yields images on both ranks, so CI won't catch it. Options: require every rank to run the encoder every step (e.g. a zero-token dummy forward when no images) and document the invariant, or gate on a per-rank flag with an explicit collective barrier. At minimum, add a test/assert for the mixed text-only/image case rather than leaving it latent.

Other findings

tokens_per_expert_E is reassigned out-of-place in KimiLatentMoE.forward (model.py:539-541: self.tokens_per_expert_E = self.tokens_per_expert_E + ...) whereas common/moe.py uses in-place .add_(). The load-balancing hook (optimizer.py:520-521) does moe.expert_bias_E.add_() and moe.tokens_per_expert_E.zero_(), and register_moe_load_balancing_hook is wired via post_optimizer_build_fn. Reassigning replaces the buffer object each step; confirm the hook still sees the live buffer (the unit test asserts identity after one forward, but the hook reads it after multiple forwards / grad-accum). In-place add_ as in common MoE would be safer and matches the documented AC double-count behavior.
KimiMoERouter reads self.gate.weight directly instead of calling self.gate(x). This works under pure FSDP (params are unsharded plain tensors inside forward) but silently breaks under TP, where weight is a sharded DTensor and F.linear(plain, dtensor) won't do what you want. TP is rejected today, so this is latent, but it's the kind of direct-.weight access that the distributed rules warn against.
MLA allocates a full L x L causal mask every forward (model.py:244-250) and vision attention runs a Python per-item loop. Both are fine as a reference; just flagging the O(L^2)/O(num_items) cost so it isn't mistaken for a training-ready path.
Param budget is dominated by the vocab. vocab_size=163840 with dim=256 puts ~84M of the ~100M params in tok_embeddings + lm_head, so the "100M debug model" exercises a very small transformer. That's a defensible choice (retains released vocab for exact tokenizer/state-dict parity), but consider noting it so reviewers don't read "100M" as transformer capacity.
Unused config field: KimiMLAAttention.Config.q_lora_rank is stored in the config but never used in __init__ (dims are derived from the projection configs). Either use it for an assertion or drop it.

Style / minor

ASCII-only in new comments/docstrings looks respected. Good.
parallelize.py:62 if ac_config is not None: rejects AC, but config_registry.py:76 sets activation_checkpoint=None, so the debug config is internally consistent — good.
README documents scope clearly; please add the parity-test reproduction steps once test Add FSDP grad scaler to the train loop #2 lands.

None of the above are blocking for an explicitly-scoped eager reference, but I'd prioritize (4) the conditional-vision-encoder FSDP hang and (2) a committed parity test before merge, and treat the router/vision-block/scatter reuse as the main follow-ups for the "reduce duplication" goal.
· agent/add-kimi-k3-reference-model

JavaZeroo · 2026-07-31T03:34:22Z

@JavaZeroo Thanks for the pr. It looks workable to me in general, have some general comments, do you have bandwidth to work on them timely?

Can you reuse more common modules in torchtitan today, like MoonViT in common/vision_encoder, common/multi_modal, and reuse MoE router.

this pr should include a reproducible numerical parity check test, likely using some reduced hf reference, since the original hf ckpt is too big.

The current code only supports fsdp2, is there bottleneck to support other parallelisms?

Currently fsdps wrap vision encoder independently, but vit is used conditionally, this will be problematic of one dp rank gets text only but another dp rank gets images.

Thanks @shuhuayu. Yes, I have bandwidth and am working on these items now.

I am reusing VisionTransformerBlock/VisionMLP, the common multimodal scatter helper, and TokenChoiceTopKRouter. I am currently running numerical validation for these changes.
Agreed. I am adding a reproducible reduced-model numerical parity test against the pinned Hugging Face Kimi K3 reference.
I would like to implement a version of fsdp with correct precision first. other parallel implementations on k3 may require more adaptation work, especially cp and pp. i think it is possible to quickly support a version of fsdp for k3 first, enough to support our experiments with small models. In the meantime I'm happy to be able to add ep tp support to this pr, do you see any need to add it in this pr?
I reproduced the mixed-modality FSDP issue and am working on the fix.

shuhuayu · 2026-07-31T06:05:35Z

I would like to implement a version of fsdp with correct precision first. other parallel implementations on k3 may require more adaptation work, especially cp and pp. i think it is possible to quickly support a version of fsdp for k3 first, enough to support our experiments with small models. In the meantime I'm happy to be able to add ep tp support to this pr, do you see any need to add it in this pr?

Support fsdp first sounds good to me. I think we should target for a training ready version using kernels for kda, one option is to use from fla.ops.kda import chunk_kda (similarly we used fla kernels for qwen 3.5, and put the pytorch native reference implementation into the numerical tests). We may have plan to use our own kernel for kda in the future.

cc: @tianyu-l

Addresses the review feedback on pytorch#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 <noreply@anthropic.com>

JavaZeroo · 2026-08-01T09:18:59Z

I would like to implement a version of fsdp with correct precision first. other parallel implementations on k3 may require more adaptation work, especially cp and pp. i think it is possible to quickly support a version of fsdp for k3 first, enough to support our experiments with small models. In the meantime I'm happy to be able to add ep tp support to this pr, do you see any need to add it in this pr?

Support fsdp first sounds good to me. I think we should target for a training ready version using kernels for kda, one option is to use from fla.ops.kda import chunk_kda (similarly we used fla kernels for qwen 3.5, and put the pytorch native reference implementation into the numerical tests). We may have plan to use our own kernel for kda in the future.

cc: @tianyu-l

Thanks @shuhuayu, I have already made the changes you requested.

Router is TokenChoiceTopKRouter now, and the vision scatter uses the shared
scatter_vision_embeds. The vision block is not folded into VisionTransformerBlock yet,
because that one hardcodes LayerNorm for the norms and BlockMask for the mask, and in eager
the padding rows would softmax to NaN. Reusing it would be a fairly big change. The MoE is
not reused either, I kept the for-loop form. Should that switch to grouped_mm, or stay eager?
Added test_kimi_k3_hf_parity.py, which freezes fp32 outputs from the released hf code, covering text logits, routed expert ids, vision features and multimodal logits.
Fixed.
kda uses chunk_kda now, with the gate, beta sigmoid and qk l2norm fused into the kernel,
same split as qwen3.5. The torch recurrence moved into the tests as the reference, and I
added a cuda test comparing fwd/bwd.

…ation pytorch/torchtitan#4025 adds Kimi K3 upstream, and checking it against the reasons this branch existed inverts the argument. It constructs nn.Linear positionally rather than through a config tree -- it has config dataclasses but does not declare child Linear.Config fields or build them -- so "return to the titan standard" was never true; upstream's own K3 does what ours does. And it supports FSDP2 only, explicitly rejecting HSDP, TP, PP, CP, EP, activation checkpointing, torch.compile and CPU offload, with the author noting TP/PP/CP would need significant adaptation because of data-dependent Python loops and incompatible forward signatures. So the parallelism work upstream declines to do is exactly what this fork has: 14/14 matrix legs producing loss, PP verified per-parameter at 0.00000 over 548 parameters, and two TP defects found and fixed, one of which also fixes upstream deepseek_v3. Refactoring toward a style upstream does not use, at the cost of breaking that, is the wrong trade. The cost was measured rather than estimated: converting three MLA linears to Linear.Config(...).build() failed 12 of 14 legs with silent exit=0 hangs. The branch keeps its two gated commits in case the LoRAConverter question returns. Work moves back to finishing veRL. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01WBy1d9YVu44nYCVqykRqL1

QIU023 · 2026-08-03T21:31:36Z

Thanks folks for connecting this PR with the broader RFC I have opened for the entire K3 support, I will review this asap in recent days and get it landed and aligned with my broad parallelism for all the reasonable interfaces needed, post-training and QAT support on top of this model backbone

torchtitan/distributed/fsdp.py already has apply_fsdp_to_vision_encoder. This folder carried its own apply_fsdp_vision, a 48-line duplicate of it that no caller ever reached, so the tower rode along inside the root wrap fully replicated on every DP rank. Invisible at the debug tower's 4 layers / hidden 256; not an option at MoonViT-V2's real 447.4M against k3mini's 80.9M text side, where the encoder is 5.5x the model it serves. Deleted the duplicate and called the core helper before the decoder, as its docstring asks. This also matches how pytorch#4025 wires the same thing, so the rebase is a deletion rather than a merge. Vendored add_zero_valued_dependency from that PR verbatim, with a note to drop it when the PR lands. It covers a hazard our own CP fix does not: FSDP2 issues the tower's all-gather from its pre-forward hook and its reduce-scatter from the output's autograd hooks, so once the tower is actually sharded, a rank that skips it desynchronizes the process group. Our fix only aligned our own all_reduce. One trap on the way: with the tower sharded its params are DTensors too, so encode_images' "is the weight a DTensor" test no longer distinguished TP's replication from FSDP's sharding. It lifted the input onto the FSDP mesh, where it met the plain all-gathered weight inside the conv. parallelize now records the tp mesh explicitly instead. 12/12 multimodal legs, 10 steps, seed 42 deterministic: bit-identical to the unsharded run on every leg (mm_fsdp2 7.73923 -> 5.32836 ... mm_ep2_fsdp2_pp2_cp2 7.71223 -> 5.26428). Vision confirmed live. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01WBy1d9YVu44nYCVqykRqL1

…multimodal Evidence index for the update to pytorch/torchtitan#3029, covering what is implemented and reproduced today rather than proposed: 13 text parallelism combinations and 12 multimodal ones, 10 steps each, seed 42 deterministic, all monotone; PP8xVP4 at |Dloss| 0.0018 against the no-PP reference; CP built on fla's merged KCP (fla-org/flash-linear-attention#691) rather than a private recurrence. Records the defects alongside, because each one passes every check that reads a loss curve: the Block AttnRes 1/tp over-reduction, the moe_sharding in_grad_placements drop that also reproduces on unmodified deepseek_v3, the non-autograd-aware conv halo that left ~60% gradient error on W-1 boundary tokens while the forward stayed bit-exact, and ten multimodal defects of which six silently reverted forward to its text-only branch. Also states what the matrices are NOT: bf16 with fp32 reduction, no QAT. K3's MXFP4 is post-training only -- the report puts QAT across SFT and RL, not pretraining -- so a pretraining-shaped matrix should not carry it. The kimi_k3_mini_qat_mxfp4 flavor implements the released scheme separately. Open gaps stated rather than omitted: LoRA's TP gradient defect (ratio up to 2.26 at tp4 on the rowwise lora_b, invisible to cold-seed checks because B is zero at init), and the report's sec 5.2.3 encoder optimizations. Refs: pytorch/torchtitan#3029, pytorch/torchtitan#4025 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01WBy1d9YVu44nYCVqykRqL1

Compiled 13-leg matrix: 11 pass with a worst eager-vs-compiled delta of 0.013, against the 0.10-0.40 spread the parallelism configurations show among themselves -- compile is numerically fine where it runs. Two fail, both EP with pipeline parallel, on _grouped_mm receiving a [224, 0] operand. First reading was that this is a core limitation: the call site, models/common/moe.py:95/101/106, is byte-identical in this fork and in #4025's tree, and has no empty-group guard. That reading is wrong. The rest of the file is not identical -- this fork rewrote the routing-map scatter under TP+EP (129e29de0), and that map determines the group boundaries _grouped_mm is handed. Control on #4025's tree, which carries upstream's unmodified moe.py: deepseek_v3_debugmodel at dp2 x ep2 x tp2 x pp2 with --compile.enable passes (loss 8.13452). Same call site, same parallelism, same compile flag. So the defect is in this fork, and the routing-map change is the prime suspect. Also records that #4025 declares torch.compile out of scope and defaults CompileConfig to enable=False, so the published comparison stays compile-off on both sides and needs no adjustment. Not fixed. Next step is to instrument num_tokens_per_expert_E under the failing configuration and find which expert goes empty, rather than adding a guard that hides the cause. Refs: pytorch/torchtitan#3029, pytorch/torchtitan#4025 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01WBy1d9YVu44nYCVqykRqL1

kimi_k3_debugmodel_pr_4025 mirrors pytorch#4025's debugmodel architecture exactly -- 13 layers at dim 256, 4 heads, q_lora 128 / kv_lora 64, qk_nope 32 / qk_rope 16 / v 32, full attention on {4, 8, 12} with KDA elsewhere, AttnRes block 12, LatentMoE latent 128 / 8 experts top-2 / 2 shared, vocab 163840, and a 4-layer 3-head MoonViT at dim 256 / qkv 384 / hidden 1024. Same model on both sides, so the comparison is our parallelism against theirs rather than two different debug models. The first version inherited k3mini's kda_layers, a 15-entry list, into a 13-layer model -- two descriptions of the same stack contradicting each other. Deriving it from full_attn_layers fixes ep2_fsdp2, which now runs all 5 steps. Verified: FSDP2 runs with vision live (20 encode_images calls, 30/30 tower parameters with gradients), starting loss 12.06 against pytorch#4025's own 12.48 on the same vocab. Refs: pytorch#3029, pytorch#4025 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01WBy1d9YVu44nYCVqykRqL1

Two desynchronizations on the PR-4025 twin flavor under CP, both surfacing as a 100-second NCCL watchdog timeout rather than an error. The sentinel-count all_reduce sat after forward's `pixel_values is None` early return, so a rank whose batch happened to carry no images returned without entering it while its CP peers waited there forever (NumelIn=2 on mesh_cp). Hoisted to the top of forward, gated on cp_world_size > 1 -- a property every rank agrees on before looking at any data. Second, now that the tower is FSDP-sharded, skipping it also skips the all-gather FSDP2 issues from its pre-forward hook (_ALLGATHER_BASE, NumelIn=10486144 on mesh_fsdp). An image-free batch now runs the tower on a minimal placeholder and keeps the graph edge through add_zero_valued_dependency, so every rank issues the same collectives and the tower's contribution to the data-parallel average is a correct zero. That is the hazard pytorch#4025 added that helper for, reached here by a second route. Both are real and both are fixed. They are NOT sufficient: fsdp2_tp2_cp2 and ep2_fsdp2_tp2_cp2 still hang at step 2 on the same NumelIn=2 all_reduce, so a third path leaves a rank out of it. Ruled out: it is not KCP (that is fla's KDA recurrence, not this collective) and not the sentinel-count assertion (which never fires in the logs). Next step is per-rank instrumentation of the entry to _exchange_sentinel_counts rather than more hypotheses. No regression: fsdp2 on the twin flavor is bit-identical (12.05716 12.04941 12.04791 11.98434 11.78795), and ep2_fsdp2 -- which the kda_layers fix repaired -- still runs all 5 steps. Refs: pytorch#3029, pytorch#4025 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01WBy1d9YVu44nYCVqykRqL1

Per-rank instrumentation of the sentinel-count exchange on the PR-4025 twin. The sharding premise holds: rank 0 and rank 2 are a CP pair reporting local counts 255 and 34, summing to 289 -- exactly 17x17, one 34x34-patch image after 2x2 merge. Each rank does hold a complementary slice. What does not hold is the number of times the exchange runs. forward executes several times per step over different microbatches (pixel_values of 1120, 1140, 1156 and 1092 patches were observed), and the entry counts differ between ranks within a step. A collective whose count differs across participants hangs the same way as one whose participants differ, which is why fixing the two data-dependent entry conditions was necessary but not sufficient. So the remaining defect is in how many times a per-forward collective runs relative to the microbatch loop, not in which slice a rank takes. Refs: pytorch/torchtitan#3029, pytorch/torchtitan#4025 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01WBy1d9YVu44nYCVqykRqL1

…aths State file so the diagnosis can resume without re-deriving it: the failing command, the two collectives that time out, the four hypotheses killed by measurement (KCP, the sentinel assertion, the shard arithmetic, the call counts), the two defects fixed on the way, and the exact next probe -- flush per collective rather than per step, so a partial step-2 trace survives. Also answers the question the twin's failure raised about the published 12/12 multimodal matrix: if the same code paths hang there, was that result luck? kimi_k3_mini_vl at dp2 x tp2 x cp2 runs 30 steps clean (7.73550 -> 2.78104), three times the published horizon, on the leg most likely to be fragile. So the difference between the two flavors is configuration, not chance, and the 12/12 holds. Which configuration difference triggers it is still open. max_patches and seq_len are identical in both, so the obvious candidate is out; what remains is vocab 2020 vs 163840, dim 512 vs 256, 21 vs 13 layers, 15 vs 10 KDA layers, and local_batch_size. Bisecting the twin one field at a time toward k3mini is ~2 minutes per run. Refs: pytorch/torchtitan#3029, pytorch/torchtitan#4025 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01WBy1d9YVu44nYCVqykRqL1

… k3mini too Bisected the twin flavor toward kimi_k3_mini_vl one field at a time. The difference was not in the model at all: the published multimodal matrix passed --training.local-batch-size 4, the twin matrix did not, and both flavors default to 1. At global batch 8 over dp2 that is one forward per step versus four gradient-accumulation microbatches -- exactly the forward=4 the per-rank probe recorded. kimi_k3_debugmodel_pr_4025 local_batch 4: 5 steps pass local_batch 1: hangs kimi_k3_mini_vl local_batch 4: 30 steps pass local_batch 1: hangs So the defect reproduces on kimi_k3_mini_vl as well. The published 12-leg multimodal matrix did not exercise it because that run's local batch was large enough to avoid accumulation entirely. That qualifies the published number: it is "passes without gradient accumulation", not "passes", and the qualification has to travel with it -- accumulation is standard at any real scale. Why accumulation breaks it is still open. The shard arithmetic is correct and step-1 call counts match across all eight ranks, so the suspect is state carried across microbatches within a step. Refs: pytorch/torchtitan#3029, pytorch/torchtitan#4025 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01WBy1d9YVu44nYCVqykRqL1

…tuple Refs: pytorch#3029 A non-last PP stage returns (hidden_state, block_residuals) -- the AttnRes adapter ships the block payload alongside the activation -- so handing that straight to add_zero_valued_dependency raised AttributeError: 'tuple' object has no attribute 'dtype'. Both tower-alive call sites did it: the image-free path (latent, never reached by a passing leg) and the zero-sentinel CP path added in the previous commit, which is what surfaced it. Route both through a local helper that puts the graph edge on the hidden state and rebuilds the tuple, the same thing the adapter's own _keepalive_touch does. Kept out of add_zero_valued_dependency so that helper stays byte-identical to pytorch#4025's and the rebase stays a clean delete. Twin-flavor multimodal matrix, 3 steps, seed 42, deterministic: the three PP+CP legs go from FAIL to passing, taking the matrix to 10/13. tp2_pp2_cp2 12.07205 12.02185 11.98569 fsdp2_pp2_cp2 12.05744 12.03855 11.97321 ep2_fsdp2_pp2_cp2 12.06617 12.02984 11.96501

Update the K3 decoder, KDA, MoE, vision path, dataloader config, tests, and numerical script for the token-major and packed-vision interfaces introduced by pytorch#4121. Preserve the latest upstream Kimi-VL numerical-test documentation while rebasing the feature history.

…base

shuhuayu

thanks for addressing my comments, looks good to me overall, please fix some final comments.

QIU023 · 2026-08-23T05:28:48Z

seems this PR is approved by maintainers and will be merged soon, already start rebasing now and will raise the parallelism supports first (EP, PP, CP, TP)

…version of reduced hf config

JavaZeroo · 2026-08-23T13:58:38Z

@shuhuayu @tianyu-l thx for the reviewing, I've implemented your final reviews.

In addition, I abstracted a MoonViTEncoder, so that k3 can also reuse k2.7's ViTEncoder.

both the K3 and K2.5 vision encoders produce bitwise-identical output compare with before

Co-authored-by: Shuhua Yu <18108279+shuhuayu@users.noreply.github.com>

shuhuayu · 2026-08-24T09:57:16Z

@JavaZeroo thanks for iterating on this pr, i merged it and will keep working on pushing more model features and parallelism support.

JavaZeroo · 2026-08-24T10:24:13Z

@JavaZeroo thanks for iterating on this pr, i merged it and will keep working on pushing more model features and parallelism support.

Thanks for helping me iterate on this pr, that help me learned a lot about titian. I'd like to implement support for document packing, if that is fine with you.

pytorch-bot Bot added the ciflow/8gpu label Jul 30, 2026

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

JavaZeroo changed the title ~~[kimi k3] add Kimi K3 reference model~~ [kimi k3] add eager reference model with FSDP2 Jul 30, 2026

JavaZeroo mentioned this pull request Jul 30, 2026

[RFC] Kimi K3 Architecture support for torchtitan based pre-training and post-training (with veRL) #3029

Open

JavaZeroo force-pushed the agent/add-kimi-k3-reference-model branch from 3c6de86 to eae7ea9 Compare July 30, 2026 14:45

tianyu-l requested a review from shuhuayu July 30, 2026 20:44

JavaZeroo marked this pull request as ready for review August 1, 2026 09:19

JavaZeroo requested review from fegin, tianyu-l, wconstab and wwwjn as code owners August 1, 2026 09:19

JavaZeroo added 5 commits August 22, 2026 07:46

split k3 model.py to moe.py lda.py

6689074

remove two unsed attn res from hf, add it in state_dict_adapter

252ae22

numerical_tests_kimi_k3.py support different attn_backend and bf16

2ef1981

update sample packing checking

4149422

JavaZeroo force-pushed the agent/add-kimi-k3-reference-model branch from fae4093 to 4149422 Compare August 22, 2026 08:29

JavaZeroo added 3 commits August 22, 2026 10:02

add TODO

9f8bcf0

chang fla kda from fp32 to bf16

44d0e83

fix test_update_from_config_propagates_moe_force_load_balance when re…

dee45e3

…base

JavaZeroo requested a review from shuhuayu August 23, 2026 03:07

shuhuayu approved these changes Aug 23, 2026

View reviewed changes

tianyu-l approved these changes Aug 23, 2026

View reviewed changes

JavaZeroo added 7 commits August 23, 2026 06:33

rename probs_TN to probs_T1N

84e9f7b

removed dtype args, only support bf16 now

157440a

Clarify the README that the numerical results are based on a reduced …

075d54c

…version of reduced hf config

remove test_update_from_config_propagates_moe_force_load_balance

99c2c68

remove TestKimiK3FSDP, and add this in to github ci

0999206

reuse some function from k2.7

32abaa0

share the MoonViT tower between k2.7 and k3

ded1bf1

shuhuayu reviewed Aug 24, 2026

View reviewed changes

Comment thread .github/workflows/integration_test_8gpu_features.yaml Outdated

JavaZeroo and others added 2 commits August 24, 2026 15:30

Update .github/workflows/integration_test_8gpu_features.yaml

b3751c9

Co-authored-by: Shuhua Yu <18108279+shuhuayu@users.noreply.github.com>

fix cicd

2dd9390

shuhuayu reviewed Aug 24, 2026

View reviewed changes

Comment thread scripts/checkpoint_conversion/numerical_tests_kimi_k3.py Outdated

Update scripts/checkpoint_conversion/numerical_tests_kimi_k3.py

88cb800

shuhuayu merged commit 5fecad9 into pytorch:main Aug 24, 2026
3 checks passed

shuhuayu mentioned this pull request Aug 24, 2026

kda in torchtitan (training) #4164

Open

Uh oh!

Conversation

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

Uh oh!

Summary

Why

Changes outside models/kimi_k3/

Numerical validation

Parallelism validation

Current scope

Uh oh!

pytorch-bot Bot commented Jul 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

shuhuayu commented Jul 31, 2026

Uh oh!

shuhuayu commented Jul 31, 2026

Uh oh!

claude Bot commented Jul 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Code Review: Kimi K3 eager reference + FSDP2

1. Reuse of common modules (maintainer Q1) — largely agree, several concrete duplications

2. Reproducible numerical parity test (maintainer Q2) — missing, should be added

3. Bottlenecks for other parallelisms (maintainer Q3)

4. Vision encoder FSDP + conditional execution (maintainer Q4) — real hang risk, agree

Other findings

Style / minor

Uh oh!

JavaZeroo commented Jul 31, 2026

Uh oh!

shuhuayu commented Jul 31, 2026

Uh oh!

JavaZeroo commented Aug 1, 2026

Uh oh!

QIU023 commented Aug 3, 2026

Uh oh!

shuhuayu left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

QIU023 commented Aug 23, 2026

Uh oh!

JavaZeroo commented Aug 23, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

shuhuayu commented Aug 24, 2026

Uh oh!

JavaZeroo commented Aug 24, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

JavaZeroo commented Jul 30, 2026 •

edited

Loading

Changes outside `models/kimi_k3/`

pytorch-bot Bot commented Jul 30, 2026 •

edited

Loading

claude Bot commented Jul 31, 2026 •

edited

Loading