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

[QIU023]

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

(Updated the main content to describle progress, completed items, and scoping for overall K3 support for every components needed in torchtitan stack with 5D parallelism and QLoRA and veRL post-training integration)

Author

Yiqiao Qiu — Member of Technical Staff, Cloud Infrastructure - Fireworks AI LinkedIn

Latest Status Summary

Since K3's release we have been building the parallelism side against the full architecture -- KDA, Gated MLA, Stable LatentMoE, Block AttnRes, and the MoonViT-V2 vision pathway. Posting now as early evidence rather than waiting for the eager reference (#4025) to land. Branches (fork QIU023/torchtitan), to be sliced per-topic and rebased onto torchtitan/models/kimi_k3/ once it does:

slot	scope	branch
A	TP: KDA NoParallel, Gated MLA, LatentMoE placements	k3_pr_a_tp_kda
B	EP + grouped GEMM	k3_pr_b_ep_grouped
C	PP + Block AttnRes cross-stage adapter	k3_pr_c_pp_attnres
D	CP: Ulysses for MLA, fla's merged KCP for KDA	k3_pr_d_cp_ulysses

(The dev branch also carries LoRA, Muon, quantile balancing, and MXFP4 QAT -- not proposed in this pass.)

Parallelism matrix

Multimodal (MoonViT + backbone), seed 42, deterministic, gradient accumulation on. Eighteen configurations: single-GPU, fsdp2, pp2, cp2, tp2, every 3-of-4 combination of FSDP/TP/PP/CP, all repeated with EP, plus the max-degree cells on 8 GPUs.

	100 steps
eager	18/18
compiled	15/18 (18/18 with the shim below)

max-degree cell	result
ep8 x fsdp8 (1 expert/rank)	pass, eager and compiled
pp4 (uneven 13-layer split, vision tower on stage 0)	pass
tp4, cp4 (1 head/rank)	pass
pp8	pass, after one fix (dense-grad on PP inputs)

Every cell starts from a shared step-0 checkpoint (--checkpoint.create_seed_checkpoint + --checkpoint.initial_load_path, as suggested on the eager PR; load verified -- cold start step-1 is 12.05342, seeded 12.07418). With init shared, the spread decomposes: the seven cells with no TP and no CP (fsdp2, pp2, pp4, pp8, ep2 x fsdp2, ep8 x fsdp8, single-GPU) agree bit-for-bit at step 1 (12.07418), and the only cells that differ are exactly those where TP or CP changes reduction order -- 12.07565 (CP), 12.07566 (TP), 12.07590 (both), a 0.009 spread fully accounted for.

Across the run the cells stay within a band rather than on top of each other, which is the documented criterion: CONTRIBUTING's numerics guidance asks for bit-wise identical loss only with the same parallelisms, and changing parallel degree changes reduction order, i.e. the computation. All eighteen decrease monotonically to the same order (0.343-0.386 at step 100), and two independent full runs of the same configuration reproduce bit-for-bit.

A dense control (the same cells with every layer a plain FFN) shows the same band -- 0.029 absolute spread at step 100 vs MoE's 0.043 -- so the band tracks reduction order, not routing.

Full per-cell tables, eager and compiled, are in the seed-checkpoint record.

• The compiled gaps are one upstream operator limitation, not ours: torch._grouped_mm rejects a zero contraction dimension -- the weight-gradient shape a rank sees when it ends up with zero routed tokens in total, which the ep2 x TP/PP/CP combinations produce (per-group empties alone do not: ep8 x fsdp8 passes compiled unshimmed). Five-line reproduction, no model/compile/distributed, being filed against pytorch/pytorch; with a local shim that re-strides exactly as the proposed patch would, every affected cell passes -- the shim is ours, so the unshimmed 15/18 stays the reported number.
• PP depth beyond the matrix: PP8xVP4 (32 virtual stages, Interleaved1F1B, multimodal) runs clean at seq_len 1024, 2048, 4096 and 8192 -- at 8192, 5 steps monotone 12.03512 -> 10.68649, peak 1.82 GiB of 15.5. The AttnRes cross-stage adapter is bit-identical on and off at every length, which is exactly what it claims: it changes what crosses a PP hop, not the arithmetic. Reaching 8192 took a code change rather than a bigger box -- the ceiling was the vocabulary-sized logits tensor (at V=163840 and L=8192 its fp32 upcast alone is 5.37 GiB), so the model grew a _skip_lm_head forward and the flavor uses chunked loss (verified bit-for-bit against plain CE at seq 1024). Details.
• (Fix merged) A shared-code gradient fix fell out of the matrix: TP-without-EP legs under-reduced gradients below the routed experts (3.10x whole-model grad_norm gap, invisible in loss -- Adam is scale-invariant); root cause in the shared MoE sharding config, reproduces on unmodified deepseek_v3, filed as [MoE] fix: pass the computed in_grad_placements without EP too — TP gradients below the experts lose their reduction #4054.
• The released 2.8T config was audited against report Table 1 on a meta device: 2.779T total vs 2.78T, 105.8B activated vs 104.2B, layer composition / latent / experts / vision tower matching field for field; the MTP layer is in the report but shipped disabled (num_nextn_predict_layers: 0), and enabling it costs ~30B that Table 1's total does not appear to include -- scale audit.

Caveats, briefly: the 100-step runs overfit a smoke dataset (stability evidence, not convergence -- real pretraining runs are queued); cross-stack numerical alignment is now blocked on our state-dict adapter round-tripping against theirs rather than on the rebase -- fixing init removed the RNG obstacle, but the same weights still have to land in both stacks, so these tables show eighteen configurations of one stack, not two stacks agreeing. Environment: torch 2.14.0.dev20260802+cu130, fla-core 0.5.1, 8x RTX 5060 Ti -- which cells can run is hardware-dependent (details in the matrix records), the losses are not. Real-data long runs are pending bigger hardware and will be posted then, not implied now.

Interface changes the PP and EP slices will touch -- all from my side

Nothing to merge yet and nothing we are asking anyone to build. PP (slot C): the eager decoder loop already threads (h_BLD, block_residual_TND) through every layer; the two missing seams -- an optional initial-stack argument and a layer-range slice, returning the pair -- land as our PP PR's first, separately reviewable commit. EP (slot B): the per-expert <-> grouped-GEMM conversion (bidirectional key map with sliced-tensor indexing, ...experts.3.w1.weight <-> w1_EFD[3], plus the grouped experts module) already exists on the branch and ships with that PR, composing with the eager per-expert adapter rather than replacing it. TP and CP need no seam and rebase cleanly.

Summary for K3 support initial scoping

Currently, my proposed plan is to add torchtitan/experiments/kimi_k3/ — the K3 model family (KDA + MLA + MoE + AttnRes) in the standard experiment layout (model.py / config_registry.py / parallelize.py / state_dict_adapter.py, following the qwen3_5 structure as the hybrid linear-attention precedent). K3 architecture model and its pre-training/post-training implementation will be aligned with the official Kimi technical report, which is tracked to release on 07-27. I will verify through a downscaled model given GPU resource constraint, but config-only scalable to official 2.8T version from Kimi released weight checkpoints.

Reference

• Full Kimi K3 report
• Attention Residual
  • Paper: arXiv:2603.15031
  • Reference impl (README + PDF only): MoonshotAI/Attention-Residuals
  • Kimi infra engineer's implementation notes: zhihu
• Branch: QIU023/torchtitan@attention_residual_dev
• Owner: Yiqiao Qiu, @QIU023

Archived old RFC status contents:

Previous Block Attn Res only RFC with comprehensive results for Block Attn Res experiments and implementations for PP adapter was archived in this link.

K3 pre-release status RFC with plans and scopes initially planned before alignments with K3 official info.

[tianyu-l]

Continued training of llama3 using attention residuals sounds interesting, but we don't have precedence in adding a model component without its actual being used in a model's original definition.

[QIU023]

Hi Tianyu, thanks a lot for the feedback and I totally agree that we shouldn’t bolt a foreign component onto Llama.

Would it work if I frame this as a new model under experiments/ instead? I’d implement a dense AttnRes transformer inspired by the paper, at the scale range studied in its scaling-law experiments (~200M–500M), following the reference pseudocode in Moonshot’s public repo. The PP integration of it (cross-stage caching adapter under Interleaved 1F1B) would be the core contribution. Therefore I prefer to keep it dense for now to stay focused on the pipeline parallelism and scaling to MoE + EP would be a natural follow-up.

Happy to adjust if you have a different direction in mind for AttnRes integration in torchtitan — just let me know what works best. Thanks!

[tianyu-l]

I personally would prefer waiting until it is used in and released from a "sota"-level model. IIUC it's not in Kimi K2.6?

Before that, having a working implementation in a fork of torchtitan that can be shared with the community sounds good enough to me.

[QIU023]

Sounds good! I’ll finish the implementation and maintain it as a fork following the existing experiments/ conventions, so the upstream path stays open if AttnRes lands in a SOTA release down the line (eg. Next Kimi model)

I’ll share the fork with the community once it’s in a clean state. Really appreciate the thoughtful guidance from you throughout this discussion!

[QIU023]

Hi Tianyu!

I would like to follow up on this and continue to
achieve closest Kimi K3 reproduction in torchtitan! I am excited that I just saw that the Kimi has released the K3: https://www.kimi.com/blog/kimi-k3 that is really using the Attention Residual in a overall 2.8T MoE model, given I have this previous work about Attention Residual pipeline parallelism distributed training support (and Kimi Linear I also built separately) so it should be very straightforward and convenient to keep moving forward!

I would keep following up to fully reproduce and support it with whatever GPU resources I can afford and start raising PRs once you and torchtitan team wants further formal support of Kimi K3.

Please let me know how you think about this following plans, thank you so much!

PS. In my own fork, the previous AttnRes PP adapter I have implemented and tested 8 PP mesh X 4 VP stages and seems working fine. Can present more details in this RFC ticket for you to review if needed

[tianyu-l]

@QIU023 we should definitely support Kimi K3 with attention residual. Looking forward to your contribution.

cc @shuhuayu

[QIU023]

Yeah thank you so much for your confirmation Tianyu! Sorry for delayed response.

I will add major updates in this RFC in above content about upcoming plans today. And I have already completed some parallelism support and 7.27 full alignment of model architecture details for K3 full reproduction and support in torchtitan, scoping mainly in K3 full architecture model with config-only needed scalable to official 2.8T model, including pre-training / SFT and post-training (given vLLM and sglang will have the official K3 inference support) and with both full-param and QLoRA MXFP4/FP8 (officially described by Moonshot team). Then sequentially requesting maintainers PR review for all individual changes, thanks again!