Kimi K3 for All

The Missing Training Pipeline

Moonshot AI released Kimi K3 in July 2026, boasting 2.8 trillion parameters, with 104 billion active, native vision, a million-token context, and an interesting architecture (Kimi Delta Attention, Attention Residuals, Stable LatentMoE, SiTU-GLU, Per-Head Muon). What they released alongside it was a HuggingFace modeling_*.py file. Weights and an inference forward pass. The training code was left out.

That is pretty normal. Almost nobody ships a training pipeline with a model release. But this left an interesting question open “how do you build the K3 architecture”. The paper describes the idea behind it but the code runs a pretrained checkpoint. Neither mentions what breaks when you actually try to build this architecture yourself and train a model from a random init.

The project started as a simple “implement the K3 architecture and train a small model.” Somewhere in the first week it evolved further. The released code cannot train, in four separate and non-obvious ways, and getting from there to a real training loop is most of the work. So the 1.27B-parameter model below is not really the point, it is a test case (one which could've been made a lot better). If the pipeline can take K3’s exact architecture from a random init to coherent English on a laptop-sized budget, it works. Is the checkpoint any good? Not especially, it only saw one epoch of Wikipedia, which was a questionable decision on my part.

The K3 Architecture

Here's a quick rundown of five things to know for the later parts to make sense.

Kimi Delta Attention (KDA): Most layers are not standard attention. KDA is a linear attention variant with a gated delta rule, maintaining a recurrent state instead of an all-pairs matrix, so cost grows linearly with sequence length. K3 uses a lower-bounded decay gate so the state can’t forget too aggressively.

Gated MLA: Every fourth layer is full attention, compressed: Multi-head Latent Attention projects keys and values through a low-rank latent instead of storing them at full width, plus an output gate. A total 93-layer K3 comprising of 69 KDA + 24 Gated MLA, 3:1, final layer is MLA.

Attention Residuals (AttnRes): Instead of uniform residual accumulation, K3 keeps a stack of block-level residual snapshots and lets each layer select among them with a softmax over learned scores, so information can skip depth.

Stable LatentMoE: The FFN blocks are a mixture of experts operating in a lower-dimensional latent space, with down/up-projections and a norm in between. 896 routed experts, 16 active per token, plus 2 shared experts that always run.

SiTU-GLU: β·tanh(gate/β)·sigmoid(gate) · up, with both branches bounded. It is an activation built for training stability at scale, which turns out to matter (see “Known Issues” below).

If you want to dive deeper into the architecture developed by moonshot, consider reading this detailed K3 architecture breakdown: x.com/waterloo_intern/status/2081762065392541951

Making K3 Trainable

This is where the exciting stuff happens.

Moonshot’s release is inference-only, and not in one obvious documented place. It is in four separate places that a training attempt hits one after another:

1. dt_bias and e_score_correction_bias are created with torch.empty(...) and never initialised. _init_weights only handles Linear and Embedding. Invisible when loading pretrained weights (the checkpoint overwrites them). You’d get uninitialised memory in your decay gate and your router bias, and the model diverges without an obvious cause.
2. KimiMoEGate.forward contains assert not self.training. The router refuses to run in training mode causing a full stop.
3. KimiSparseMoeBlock.forward raises NotImplementedError when self.training is true, and the dispatch it does have (moe_infer) is wrapped in @torch.no_grad(). This leads to absolutely no gradient path through the experts.
4. gradient_checkpointing_enable() is a silent no-op. KimiLinearModel.forward loops over its layers without ever calling a checkpoint function. This was the trickiest one to fix. Everything seems normal but memory usage is higher than expected with no clear indications.

None of this is a criticism towards Moonshot. Inference-only releases are the norm. But it means the actual engineering content of “train a K3” is these four fixes and what comes after them, and unfortuantely it wasn't documented as well as I hoped.

Here’s what patch 2 and 3 look like in the released code, straight from kimi_k3/kimi_original/modeling_kimi_linear.py:

# KimiMoEGate.forward
assert not self.training
scores = scores.view(bsz * seq_len, -1)
scores_for_choice = scores + self.e_score_correction_bias.unsqueeze(0)

and the fix, in kimi_k3/train/patches.py, same math, only the bias is detached so it only steers which experts get picked without picking up a gradient of its own:

def _gate_forward(self: KimiMoEGate, hidden_states):
    ...
    scores_for_choice = scores + self.e_score_correction_bias.detach().unsqueeze(0)

And patch 4, the gradient checkpointing that was a documented no-op:

_ORIG_LAYER_FORWARD = KimiDecoderLayer.forward

def _layer_forward(self: KimiDecoderLayer, hidden_states, **kwargs):
    if getattr(self, "_grad_ckpt", False) and self.training and torch.is_grad_enabled():
        return checkpoint(_ORIG_LAYER_FORWARD, self, hidden_states,
                           use_reentrant=False, **kwargs)
    return _ORIG_LAYER_FORWARD(self, hidden_states, **kwargs)

KimiDecoderLayer.forward is monkey-patched with this wrapper instead of being edited in place, which is the pattern for all four fixes: the original method stays reachable, the patch just decides when to call through to it.

All four are patched at runtime, in kimi_k3/train/patches.py. It is deliberately kept separate from the architecture code so kimi_k3/kimi_original/ stays byte-identical to the HuggingFace release. This allows the pipeline to change again if needed. The MoE training dispatch was verified bit-identical in forward and gradient-matched against the reference per-expert loop (other/test_moe_equiv.py) before it went anywhere near a real run.

The Training Code

The four patches make K3 trainable in the first palce. The rest of training/ is what makes a multi-day, interruption-heavy run actually finish, and this is the part that’s reusable for any K3 based model, not just this one.

The Epoch: The obvious way to feed a training loop is to sample random offsets into the token stream. It’s also wrong as with replacement, a nominal epoch touches only 1 − 1/e ≈ 63% of unique tokens, quietly wasting a third of the corpus. TokenLoader instead partitions the stream into non-overlapping windows and shuffles their order, so one pass really is one pass, and position is checkpointed so a resume continues the sweep instead of restarting it:

class TokenLoader:
    def _build_order(self):
        rng = np.random.default_rng(self.seed + self.epoch)
        self.order = (rng.permutation(self.n_windows) if self.shuffle
                      else np.arange(self.n_windows))

    def batch(self):
        if self.pos + self.batch_size > self.n_windows:
            self.epoch += 1
            self.pos = 0
            self._build_order()
        idx = self.order[self.pos:self.pos + self.batch_size]
        self.pos += self.batch_size
        ...

Per-Head Muon: One of K3’s stated design choices, orthogonalizes each attention head’s update block independently rather than the whole projection matrix at once, which needed its own optimizer:

if nh > 1:  # per-head: split rows into nh blocks, orthogonalize each
    out = upd.size(0)
    hd = out // nh
    ortho = torch.empty_like(upd)
    for h in range(nh):
        sl = slice(h * hd, (h + 1) * hd)
        ortho[sl] = _newton_schulz5(upd[sl], ns)
else:
    ortho = _newton_schulz5(upd, ns)

Checkpointing: A checkpoint that’s missing the data loader’s position, the RNG state, or half the optimizer state doesn’t resume, but rather restarts with extra steps. Every checkpoint here is model, MTP head, both optimizers, step, token count, the loader’s sweep position, and both the CPU and CUDA RNG states, written atomically so a crash mid-write can’t corrupt the last good one:

def save_ckpt(path, model, mtp, muon, adamw, step, tokens, cfg, loader, args):
    tmp = str(path) + ".tmp"
    torch.save({
        "model": model.state_dict(), "mtp": mtp.state_dict(),
        "muon": muon.state_dict(), "adamw": adamw.state_dict(),
        "step": step, "tokens": tokens, "loader_rng": loader.state(),
        "torch_rng": torch.get_rng_state(),
        "cuda_rng": torch.cuda.get_rng_state_all(),
    }, tmp)
    os.replace(tmp, path)   # atomic: no partial checkpoint on crash

Supervisor Script: supervisor.sh restarts the trainer from the last checkpoint on failure, but only counts a restart as progress if the step counter actually moved since the last one. Repeated identical failures trip a retry budget instead of burning the GPU forever:

cur_step=$(grep -a "^step " "$LOG" | tail -1 | sed -E 's/^step +([0-9]+).*/\1/')
if [[ "$cur_step" -gt "$last_step" ]]; then
  fails=1                      # made progress since last failure; reset budget
else
  fails=$((fails + 1))
fi
last_step=$cur_step
if [[ $fails -ge $MAX_FAILS ]]; then
  echo "GIVING UP: $fails failures with no progress past step $cur_step" >> "$LOG"
  break
fi

The combination of the above components let the actual run survive 18 pause/resume cycles with zero discontinuity in the loss curve. None of it is specific to Wikipedia or to this parameter count. It’s the part of the pipeline meant to outlive this particular checkpoint which may prove helpful to others aiming to replicate this.

Making the Pipeline Faster

A pipeline that runs is necessary but not sufficient. The first working version trained at 4,400 tokens/second on an H100, a three-week epoch. Four changes took it to 17,500 tok/s:

• Batched MoE dispatch: The reference implementation loops over experts one at a time: 40 experts × 23 layers = 920 tiny GEMMs per forward, entirely launch-bound. Stacking expert weights and padding each bucket to a common capacity turns that into 3 bmm calls.
• Chunked cross-entropy: A 163,840-wide vocabulary at batch 12 × seq 2048 costs ~5 GiB of fp32 logits, twice over (main head + MTP head). Splitting the token axis and recomputing each chunk in the backward pass keeps one chunk resident.
• Removing hidden device→host syncs: torch.bincount and torch.nonzero size their output from the data, forcing a GPU→CPU sync. Profiling showed cudaMemcpyAsync at 52% of CPU time. scatter_add into fixed-size buffers removed it:

  # torch.bincount(flat_ids, minlength=E) would sync here; this doesn't
  counts = torch.zeros(E, dtype=torch.long, device=x.device).scatter_add_(
      0, flat_ids, torch.ones_like(flat_ids)
  )

• torch.compile on the elementwise hot spots. Worth ~1.45×.

Chunked cross-entropy looks like this, splitting the token axis so only one chunk of vocab-163,840 logits is ever resident at once, and recomputing each chunk under checkpoint in the backward pass instead of keeping all of them around:

def chunked_ce(hidden, weight, targets, n_chunks=8):
    h = hidden.reshape(-1, hidden.size(-1))
    t = targets.reshape(-1)
    total = 0.0
    for hc, tc in zip(h.chunk(n_chunks), t.chunk(n_chunks)):
        total = total + checkpoint(_chunk_ce_sum, hc, weight, tc, use_reentrant=False)
    return total / t.numel()

Then it stopped improving, at about 13% MFU. Turning gradient checkpointing off, removing 19% of the FLOPs, made it slower. This led to the conclusion that the pipeline is bandwidth-bound, not FLOP-bound. At hidden size 1024 every GEMM is far below what an H100 wants, and K3’s design is deliberately elementwise-heavy (fp32 SiTU, fp32 RMSNorm, an fp32 AttnRes softmax every layer). Those burn bandwidth and contribute nothing to the FLOP numerator. Longer sequences made it worse, not better (−4% at 4096, −31% at 8192). MFU is simply the wrong yardstick for this architecture at this size which is a lesson about the pipeline, independent of the model it happens to be training.

The Resulting Model

One epoch, 15,600 steps, 4.601 billion tokens of English Wikipedia, every token seen exactly once, is the pipeline’s end-to-end test. If it produces a coherent model with no NaNs, no divergence, and a resumable/checkpointable run across interruptions, the pipeline passed.

Measurements

HellaSwag at 33.7% against a 25% floor means the model did learn real commonsense structure. It is far short of what a compute-optimal model this size would reach; that gap is tokens (4.6B against a Chinchilla-optimal ~25B).

Inference tests on Macbook Air M1: inference/generate.py on an M-series Mac, MPS backend, the pure-PyTorch fla_shim KDA fallback (no Triton), float32: 1.7 tok/s, greedy, batch 1. That is 8.5× the ~0.2 tok/s previously measured on a server CPU, and confirms the whole point of this exercise: the pipeline’s output, a 2.4 GiB checkpoint, runs on a no GPU laptop. (bf16 currently fails on MPS with a dtype mismatch inside fla_shim’s depthwise conv, a real, filed limitation and not swept under the rug; float32 works.)

Greedily decoded samples:

“World War II began in” → September 1939, and the first aircraft arrived in the

“The capital of France is” → in the 15th arrondissement of Paris. The

“Photosynthesis is the process by which” → photosynthesis is converted to a form of energy. The photosynthesis process is the process by which photos

“The Pacific Ocean is” → a large, shallow, shallow water body of the Pacific Ocean, located in the Pacific Ocean.

The first is correct, including the month. The rest are fluent and confidently wrong or circular, which matches the profile of a pipeline that works correctly on an undertrained model. The syntax is solid (the pipeline learned to produce English), while semantics are not (18% of a Chinchilla-optimal budget isn’t enough to fix facts). It’s what was expected from a correct pipeline given too few tokens, and quite different from what you’d see if something in the training loop were subtly broken.

MoE at This Scale

After all this I had another important question. Is MOE even relevant at this scale?

The monitoring metric I originally trusted, moe_imb (max expert load / mean), spent the whole run pinned near its ceiling, which I read repeatedly as “the router has collapsed.” It had not. That metric reports the worst single layer and cannot distinguish “4 experts absorb everything” (fatal) from “1 expert is popular” (harmless): both read exactly num_experts / top_k. This turned out to be a pipeline-instrumentation bug, not a training bug, and it’s discussed properly in the next section.

Measured properly, with other/probe_moe_load.py, which hooks every gate, accumulates real routing counts over a few validation batches, and computes entropy alongside the old max/mean number instead of trusting it alone:

imb = (c.max() / c.mean()).item()
dead = int((c == 0).sum())
p = frac[frac > 0]
ent = float(-(p * p.log()).sum() / np.log(E))   # normalised 0..1

Nothing was starved. But is the routing useful, or just alive? Measuring Jensen-Shannon divergence of expert usage across four text domains (math, biography, geography, code) against a within-domain noise floor:

Specialisation emerges top-down: Deep layers ended up routing by content such as biography and mathematics select disjoint expert sets, math and code share them. The first two or three MoE layers show essentially no specialisation (ratio ≈ 1×).

But the router’s central advantage doesn’t fully survive the scale-down. K3 activates 16 of 896 experts, 1.8% sparsity. With 40 experts the closest workable setting is 4 of 40, 10%. Fine-grained MoE works because 896 experts can each occupy a narrow slice of the input distribution; at 40, every expert ends up more of a generalist. A dense ~0.4B model would plausibly match this checkpoint today. The batched dispatch and load balancing are doing what they’re supposed to. There’s just less for them to work with at this size.

Known Issues

Here’s some things that broke along the way (interested people can definitely make PRs to fix these or anything else they find).

The bias-drift bug: moe_imb crept from ~1.5 to ~5.9 over 4,000 steps. Inspecting e_score_correction_bias in a checkpoint found the cause. Every expert’s bias had drifted to ~+4.4 with a spread of only 0.09. Since top-k is invariant to a constant added to all scores, almost all of that magnitude was doing nothing, while the tiny differential signal that actually balances load stayed buried. Two separate defects, both fixed in the pipeline:

• Counts were overwritten: The gate stored per-expert counts on every forward. With grad_accum=6 the balancing step only ever saw the last micro-batch: 1/6 of the tokens. The sign update was chasing sampling noise.
• The update wasn’t zero-centred: A plain ±gamma sign update drifts in common mode whenever over- and under-loaded expert counts differ. With 40 experts at top-4, ~36 sit below the mean each step, pushing the whole bias vector up by ~0.8·gamma/step. Fixed by mean-subtracting the update, which is routing-neutral:

  delta = -gamma * torch.sign(counts - target)
  delta = delta - delta.mean()          # zero-centre: no common-mode drift
  m.e_score_correction_bias.add_(delta.to(m.e_score_correction_bias.dtype))

Raising the balancing step size didn’t fix imbalence: Tried at 2× the value causing the router to collapse to the ceiling within one step. The sign update never converges, it random-walks with step size gamma, so a bigger step means a bigger oscillation amplitude, not better balance. Sustained imbalance is a router problem, not a dispatcher problem. A related attempt to fix it by lowering the dispatch capacity threshold also made things worse (higher VRAM, no balance improvement) and was reverted.

Collapse during warmup: With a short LR ramp the router collapses onto a handful of experts within ~20 steps and pins at the ceiling, leaving ~90% of the model dead. Warmup is specified as a fraction of total steps, so shortening a run for a smoke test silently shortens warmup too. Dropping max_steps from 100k to 2.2k took warmup from ~1000 steps to 22 and collapsed a real run, not just the smoke test. Fixed by making warmup fraction and not step count.

Misleading monitoring metric: As covered above, max/mean of the worst layer conflates two states that look identical (9.9–10.0) but mean opposite things. This one doesn’t have a training fix. It has a pipeline fix, other/probe_moe_load.py and other/probe_moe_specialisation.py now exist because the scalar dashboard number wasn’t trustworthy on its own, and a pipeline that can’t tell you when it’s actually broken isn’t finished yet.

Limitations

Altough the entire pipeline works there’s some limitations I faced with respect to the training run I conducted to make the test ~1.3B model:

• Undertrained: 4.601B tokens against a Chinchilla-optimal ~25B for this size, about 18%. Facts are unreliable. This is a token-budget fact (Wikipedia caps at 4.6B, one epoch, no repetition), not a data-quality one.
• Wikipedia only: No code, conversation, or diverse web text, which is narrow for a base model.
• Expert sparsity 10% vs. K3’s 1.8%: This was forced by the smaller expert count, not chosen.
• Quantile Balancing replaced: K3’s actual load-balancing method is absent from the released inference code and described only qualitatively. The DeepSeek-style sign update substituted here random-walks and never converges.
• Text-only: MoonViT-V2 (401M) is dropped.
• bf16, not MXFP4/MXFP8: SDPA, not FlashAttention-2 (numerically equivalent, different kernel). MTP enabled where the released config ships it off.
• No instruction tuning: This is a base model.
• Mac MPS bf16 is currently broken: (dtype mismatch in fla_shim’s conv1d); float32 works at 1.7 tok/s.

Future Work

The pipeline is at a usable state. Future work has been proposed as steps to train a better test model.

Finer experts: The clearest lesson from the specialisation analysis is that granularity is nearly free at this parameter count.

Four times as many experts, each a quarter the size resulting in same total parameters. Sparsity moves from 10% toward K3’s 1.8%, directly targeting the knob that governs whether specialisation can sharpen in the shallow layers. first_k_dense_replace moves from 1 to 3, since the first three MoE layers measurably do no routing in v1, so they shouldn’t pay for experts.

More and better data: Wikipedia caps at 4.6B tokens. v2 of the model can be trained on FineWeb-Edu, targeting 25B tokens ≈ Chinchilla-optimal.

Pipeline instrumentation, not just model tuning: moe_imb is retired in favor of routing entropy and dead-expert count as the primary dashboard signals. The v1 run generated several false alarms and one actively harmful “fix” because the wrong number was being watched. That is a permanent pipeline improvement, independent of which model gets trained next.

Quantile Balancing implementation: It would remove the largest remaining deviation from K3, if the training-side procedure can be reconstructed faithfully. It’s currently described only qualitatively in the paper/blog.

If the granularity hypothesis is right, v2 of the trained model should show specialisation reaching the early layers and materially lower perplexity. If it’s wrong, if 3.8% sparsity at 160 experts still doesn’t specialise in layers 0–2. That would let us infer that something about the minimum viable expert count for fine-grained MoE, learned from a pipeline built to measure it properly rather than assume it.

I would’ve loved to train v2 myself but I’m out of credits on runpod.io.

Try It

git clone https://github.com/ArneshBanerjee/Kimi-K3-for-All
cd Kimi-K3-for-All
pip install torch transformers==4.57.1 safetensors tiktoken blobfile einops numpy
# download the weights from the Hugging Face repo into checkpoints/hf/

python inference/generate.py -p "The history of computing began"

It picks its own backend: real fla Triton kernels on CUDA, and a pure-PyTorch KDA implementation on Mac/CPU where Triton is unavailable. Both were verified to produce identical text from the same checkpoint.

To run the pipeline yourself, from scratch:

python training/prepare_data.py     # ~13 min: download + tokenize Wikipedia
./training/launch_training.sh       # the full run

Paths live in paths.conf, so point DATA_DIR and CKPT_DIR wherever you like.

Credits

The architecture is Moonshot AI’s. kimi_k3/kimi_original/ is their released code, redistributed unmodified under the Kimi K3 License (MLA and MoE gating adapted from DeepSeek-V3 under Apache 2.0). All credit for the design belongs to them; this project only built the training pipeline they didn’t release, at a size that fits on a laptop.

• Model: huggingface.co/moonshotai/Kimi-K3
• Paper: arxiv.org/abs/2607.24653
• Blog: kimi.com/blog/kimi-k3
• Article (this helped me a lot): x.com/waterloo_intern/status/2081762065392541951

Training data: wikimedia/wikipedia 20231101.en, CC BY-SA 3.0 / GFDL.

Code: github.com/ArneshBanerjee/Kimi-K3-for-All · Weights: huggingface.co/ArneshBanerjee/Kimi-K3-for-All
Author: arneshbanerjee.dev · Contact: arneshbanerjee24 [at] gmail [dot] com