Incorporate synthetic data generation for token efficiency.
The first 10T tokens were trained with a constant learning rate of 2e-4 and a 500-step warmup.
This was followed by 5.5T tokens with a cosine decay from 2e-4 to 2e-5.
Weight decay was 0.1 throughout.
Towards the end, the learning rate decayed from 2e-5 to 7e-6.
Data
High-quality tokens are limited.
A naive approach to increasing token utility would be repeated exposure to the same tokens, increasing the epoch count, etc., but this leads to overfitting and poor model generalization.
They introduce synthetic data generation to increase token utility.
They use a synthetic rephrasing framework:
Style- and perspective-diverse prompting: prompt an LLM to generate rephrasings of the original text.
Chunkwise generation: texts are divided into segments, rephrased individually, and then stitched back together to form complete passages.
The overall pretraining data corpus has 15.5T tokens.
Data domains: web text, code, mathematics, and knowledge.
Model Architecture
1.04-trillion-parameter MoE, with 32 billion active parameters.
7,168 model hidden dimension.
2,048 MoE expert hidden dimension.
384 experts.
64 attention heads to balance computational overhead.
DeepSeek-V3 uses 256 experts and 128 attention heads, but Kimi wanted to increase sparsity according to their sparsity scaling laws.
YaRN is used to extend the context window to 128K.
Number of attention heads:
As context length increases, doubling attention heads, as DeepSeek-V3 does, leads to significant inference overhead.
Doubling attention heads did not give any significant improvement to validation loss either.
They therefore use 64 attention heads instead of doubling.
Ultra-sparse MoE.
MLA.
MuonClip:
The issue with scaling Muon is that it causes instability.
Existing solutions such as logit soft-capping clip the attention logits, but the QK dot products before that can still blow up. QK normalization cannot be used with MLA.
It uses QK-Clip with Muon instead: a weight-clipping mechanism that limits or rescales Q and K projections after the update to bound the growth of attention logits.
It does not alter the forward or backward pass in the current step.
Because only a few attention heads have exploding logits in practice, they use a per-head scaling factor and apply QK-Clip only to the exploding heads.
Here, Muon is integrated with weight decay, RMS matching, and QK-Clip into MuonClip.
Combines 16-way expert parallelism (EP), 8-way data parallelism (DP), 8-way pipeline parallelism (PP), and sequence parallelism (SP).
1F1B schedule: one forward pass and one backward pass.
Parameters are stored in BF16.
Gradient accumulation happens in FP32.
Activation reduction: remaining GPU memory is insufficient to hold MoE activations.
They store less important activations, such as the MoE up projection and SwiGLU, in FP8-E4M3.
They offload some activations to the CPU. During the 1F1B phase, they offload forward activations from the previous micro-batch and prefetch backward activations for the next.
