s1lv3rj1nx/ch5

Chapter 5: DeepSeek - The Final Adventure

Overview

The goal of this chapter is to implement and train DeepSeek's innovative architecture, which introduces two groundbreaking techniques: Multi-Head Latent Attention (MLA) and Mixture of Experts (MoE). Building upon the concepts from previous chapters (KV Cache, MQA, GQA), we explore how DeepSeek achieves superior memory efficiency while maintaining model expressiveness.

Specifically, we implement key innovations from the DeepSeek research papers:

• Multi-Head Latent Attention (MLA): A novel attention mechanism that caches a single compact latent matrix instead of separate key-value matrices for each head, dramatically reducing memory requirements
• Mixture of Experts (MoE): Computational sparsity through expert routing, enabling larger model capacity with efficient inference
• Decoupled Rotary Position Embeddings (RoPE): Integration of positional information while preserving the absorption trick in MLA

Key Innovations

1. Multi-Head Latent Attention (MLA)

MLA addresses the fundamental trade-off between memory efficiency and model expressiveness. Instead of caching separate K and V matrices for each attention head, MLA:

• Projects inputs to a low-dimensional latent space: (C_{KV} = \text{LayerNorm}(X W_{DKV}))
• Caches only the compact latent matrix (C_{KV} \in \mathbb{R}^{s \times d_L}) where (d_L \ll n \times h)
• Uses the absorption trick to precompute combined projections: (W_{QK} = W_Q W_{UK}^T)
• Reduces KV cache from (2 \times n \times h) to (d_L), achieving 10-20× memory reduction

2. Decoupled Rotary Position Embeddings

To integrate RoPE with MLA while preserving the absorption trick:

• Splits queries into RoPE and non-RoPE components
• Applies RoPE only to a subset of dimensions, maintaining positional awareness
• Keeps the absorption trick intact for the non-RoPE components
• Balances positional encoding with memory efficiency

3. Mixture of Experts (MoE)

MoE enables computational sparsity through expert specialization:

• Routes each token to a subset of expert networks (top-k routing)
• Implements load balancing to prevent expert collapse
• Increases model capacity without proportional computational cost
• Each expert specializes in different aspects of language understanding

Codebase

For the entire code refer to: Github

Training

The model is trained on a diverse multi-source dataset combining:

• TinyStories: High-quality synthetic stories for coherence
• WikiText-2: Encyclopedia-style text for factual knowledge
• OpenWebText: Web content for diverse writing styles

This diverse dataset prevents overfitting and improves generalization.

Training from Scratch

To train the model from scratch, you have two options:

Option 1: FP8 Quantized Training (H100 GPUs)

For H100 GPUs with FP8 support, enable quantization in train.py (around line 159):

# In train.py, locate TRAINING_CONFIG and set:
TRAINING_CONFIG = {
    ...
    "enable_quantization": True,  # Enable FP8 quantization for H100
    ...
}

Then run:

cd ch5
python train.py

Benefits of FP8 training:

• 2× training speedup
• 50% memory savings
• Maintains model quality with minimal accuracy loss

Option 2: Standard Precision Training (All GPUs)

For GPUs without FP8 support (A100, V100, RTX series, etc.), keep quantization disabled (default):

# In train.py, TRAINING_CONFIG should have:
TRAINING_CONFIG = {
    ...
    "enable_quantization": False,  # Use standard mixed precision (FP16/FP32)
    ...
}

Then run:

cd ch5
python train.py

The training script includes:

• FP8 quantization for 2× speedup and 50% memory savings
• Gradient clipping to prevent exploding gradients
• Cosine annealing learning rate schedule
• Automatic checkpointing and WandB logging
• Mixed-precision training with automatic scaling
• Multi-dataset training (TinyStories, WikiText, OpenWebText)

Using Pre-trained Weights

Pre-trained weights are available on Hugging Face. Download the best_model.pth file and place it in the ch5/checkpoints directory.

Understanding the Training Dynamics

Training Curves Analysis

The training curves reveal several important characteristics of the DeepSeek model:

1. Rapid Initial Learning (Steps 0-2000):

   • Training loss drops sharply from ~4.6 to ~3.6
   • Perplexity decreases dramatically from ~100 to ~36
   • This indicates the model quickly learns basic language patterns and token distributions
   • The steep gradient shows efficient optimization with the AdamW optimizer and cosine annealing schedule

2. Convergence Phase (Steps 2000-3500):

   • Loss stabilizes around ~3.55-3.58
   • Perplexity plateaus at ~35-36
   • The model has learned the primary patterns in the training data
   • Diminishing returns suggest the model is approaching its capacity for this dataset size

3. Validation Performance:

   • Validation loss closely tracks training loss (5.86 → 6.01)
   • Validation perplexity remains stable (350 → 410)
   • The gap between training and validation metrics indicates some overfitting, which is expected given:
     • Small dataset size (~3,000 examples)
     • Complex architecture with MLA and MoE components
     • Limited training time (24-hour window)

Why the Validation Gap?

The difference between training and validation performance is natural and expected:

• Dataset Diversity: While we use three diverse sources (TinyStories, WikiText, OpenWebText), the relatively small sample size means the model memorizes training patterns
• Model Capacity: DeepSeek's architecture (with experts and latent attention) has high capacity, allowing it to fit the training data well
• Training Duration: With more data and longer training, this gap would typically decrease
• Perplexity Scale: A validation perplexity of ~400 is reasonable for a small model trained on limited data

Key Takeaways

Despite the validation gap, the model demonstrates:

• ✅ Stable Training: No gradient explosions or divergence
• ✅ Consistent Learning: Smooth loss curves without erratic jumps
• ✅ Functional Architecture: MLA and MoE components work correctly
• ✅ Reasonable Performance: The model learns language patterns and can generate coherent text

For production use, you would want:

• Larger and more diverse training dataset (100K+ examples)
• Longer training duration (multiple days)
• Regularization techniques (dropout, weight decay tuning)
• More sophisticated data augmentation

Inference

To generate text using the trained model:

cd ch5
python inference.py --model_path <model_path>/best_model.pt --prompt "Once upon a time" --max_new_tokens 100

The model can generate text in various styles depending on the input prompt, thanks to the diverse training data.

Sample Output

Initial text: Once upon a time

Generated Text:
--------------------------------------------------
Once upon a time, there was a little girl named Lucy. She loved to play 
in the garden with her toys. One sunny day, Lucy found a mysterious box 
hidden behind the old oak tree. Inside the box was a shiny golden key...
--------------------------------------------------

Model Configuration

The implementation includes two configurations:

Small Configuration (Training)

• Model Dimension: 512
• Hidden Dimension: 1376
• Attention Heads: 8
• KV Heads: 4 (Grouped Query Attention)
• Latent Dimension: 256 (for MLA)
• Layers: 8
• Experts: 4 (top-2 routing)
• Vocabulary Size: 50,257 (GPT-2 tokenizer)
• Context Length: 256 tokens

Full Configuration (Reference)

Based on DeepSeek-V2 specifications for scaling to production models.

Memory Efficiency Comparison

Comparing KV cache memory requirements for a single layer with 256 tokens:

Method	Cache Size	Formula
Multi-Head Attention (MHA)	32,768 bytes	(2 \times 8 \times 64 \times 256 \times 2)
Multi-Query Attention (MQA)	4,096 bytes	(2 \times 1 \times 64 \times 256 \times 2)
Grouped Query Attention (GQA)	8,192 bytes	(2 \times 2 \times 64 \times 256 \times 2)
Multi-Head Latent Attention (MLA)	2,048 bytes	(128 \times 256 \times 2)

MLA achieves 16× memory reduction compared to standard MHA!

Technical Highlights

Absorption Trick

The mathematical elegance of MLA lies in the absorption trick:

Standard: Q @ K^T = (X @ W_Q) @ (C_KV @ W_UK)^T
Absorbed: Q @ K^T = X @ (W_Q @ W_UK^T) @ C_KV^T = X @ W_QK @ C_KV^T

By precomputing (W_{QK} = W_Q W_{UK}^T), we eliminate the need to cache individual K matrices.

Load Balancing in MoE

To prevent expert collapse (all tokens routed to few experts), we implement:

• Auxiliary loss: Encourages uniform expert utilization
• Expert capacity: Limits tokens per expert
• Balanced routing: Monitors and adjusts expert load distribution

FP8 Quantization

Optional FP8 quantization provides:

• 2× training speedup on modern GPUs (H100, A100)
• 50% memory reduction for activations and gradients
• Minimal accuracy loss with proper scaling
• Automatic fallback to FP16 if FP8 is unavailable

Requirements

• Python 3.8+
• PyTorch 2.0+
• tiktoken
• datasets
• tqdm
• wandb (optional, for experiment tracking)

Citation

@misc{deepseek-v2,
      title={DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model}, 
      author={DeepSeek-AI and Aixin Liu and Bei Feng and Bin Wang and Bingxuan Wang and Bo Liu and Chenggang Zhao and Chengqi Dengr and Chong Ruan and Damai Dai and Daya Guo and Dejian Yang and Deli Chen and Dongjie Ji and Erhang Li and Fangyun Lin and Fuli Luo and Guangbo Hao and Guanting Chen and Guowei Li and H. Zhang and Hanwei Xu and Hao Yang and Haowei Zhang and Honghui Ding and Huajian Xin and Huazuo Gao and Hui Li and Hui Qu and J. L. Cai and Jian Liang and Jianzhong Guo and Jiaqi Ni and Jiashi Li and Jin Chen and Jingyang Yuan and Junjie Qiu and Junxiao Song and Kai Dong and Kaige Gao and Kang Guan and Lean Wang and Lecong Zhang and Lei Xu and Leyi Xia and Liang Zhao and Liyue Zhang and Meng Li and Miaojun Wang and Mingchuan Zhang and Minghua Zhang and Minghui Tang and Mingming Li and Ning Tian and Panpan Huang and Peiyi Wang and Peng Zhang and Qihao Zhu and Qinyu Chen and Qiushi Du and R. J. Chen and R. L. Jin and Ruiqi Ge and Ruizhe Pan and Runxin Xu and Ruyi Chen and S. S. Li and Shanghao Lu and Shangyan Zhou and Shanhuang Chen and Shaoqing Wu and Shengfeng Ye and Shirong Ma and Shiyu Wang and Shuang Zhou and Shuiping Yu and Shunfeng Zhou and Size Zheng and T. Wang and Tian Pei and Tian Yuan and Tianyu Sun and W. L. Xiao and Wangding Zeng and Wei An and Wen Liu and Wenfeng Liang and Wenjun Gao and X. Q. Li and Xiangyue Jin and Xianzu Wang and Xiao Bi and Xiaodong Liu and Xiaohan Wang and Xiaojin Shen and Xiaokang Chen and Xiaosha Chen and Xiaotao Nie and Xiaowen Sun and Xiaoxiang Wang and Xin Liu and Xin Xie and Xingkai Yu and Xinnan Song and Xinyi Zhou and Xinyu Yang and Xuan Lu and Xuecheng Su and Y. Wu and Y. K. Li and Y. X. Wei and Y. X. Zhu and Yanhong Xu and Yanping Huang and Yao Li and Yao Zhao and Yaofeng Sun and Yaohui Li and Yaohui Wang and Yi Zheng and Yichao Zhang and Yiliang Xiong and Yilong Zhao and Ying He and Ying Tang and Yishi Piao and Yixin Dong and Yixuan Tan and Yiyuan Liu and Yongji Wang and Yongqiang Guo and Yuchen Zhu and Yuduan Wang and Yuheng Zou and Yukun Zha and Yunxian Ma and Yuting Yan and Yuxiang You and Yuxuan Liu and Z. Z. Ren and Zehui Ren and Zhangli Sha and Zhe Fu and Zhen Huang and Zhen Zhang and Zhenda Xie and Zhewen Hao and Zhihong Shao and Zhiniu Wen and Zhipeng Xu and Zhongyu Zhang and Zhuoshu Li and Zihan Wang and Zihui Gu and Zilin Li and Ziwei Xie},
      year={2024},
      eprint={2405.04434},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2405.04434}, 
}

@misc{deepseek-v3,
      title={DeepSeek-V3 Technical Report}, 
      author={DeepSeek-AI and Aixin Liu and Bei Feng and Bing Xue and Bingxuan Wang and Bochao Wu and Chengda Lu and Chenggang Zhao and Chengqi Deng and Chenyu Zhang and Chong Ruan and Damai Dai and Daya Guo and Dejian Yang and Deli Chen and Dongjie Ji and Erhang Li and Fangyun Lin and Fucong Dai and Fuli Luo and Guangbo Hao and Guanting Chen and Guowei Li and H. Zhang and Han Bao and Hanwei Xu and Haocheng Wang and Haowei Zhang and Honghui Ding and Huajian Xin and Huazuo Gao and Hui Li and Hui Qu and J. L. Cai and Jian Liang and Jianzhong Guo and Jiaqi Ni and Jiashi Li and Jiawei Wang and Jin Chen and Jingchang Chen and Jingyang Yuan and Junjie Qiu and Junlong Li and Junxiao Song and Kai Dong and Kai Hu and Kaige Gao and Kang Guan and Kexin Huang and Kuai Yu and Lean Wang and Lecong Zhang and Lei Xu and Leyi Xia and Liang Zhao and Litong Wang and Liyue Zhang and Meng Li and Miaojun Wang and Mingchuan Zhang and Minghua Zhang and Minghui Tang and Mingming Li and Ning Tian and Panpan Huang and Peiyi Wang and Peng Zhang and Qiancheng Wang and Qihao Zhu and Qinyu Chen and Qiushi Du and R. J. Chen and R. L. Jin and Ruiqi Ge and Ruisong Zhang and Ruizhe Pan and Runji Wang and Runxin Xu and Ruoyu Zhang and Ruyi Chen and S. S. Li and Shanghao Lu and Shangyan Zhou and Shanhuang Chen and Shaoqing Wu and Shengfeng Ye and Shengfeng Ye and Shirong Ma and Shiyu Wang and Shuang Zhou and Shuiping Yu and Shunfeng Zhou and Shuting Pan and T. Wang and Tao Yun and Tian Pei and Tianyu Sun and W. L. Xiao and Wangding Zeng and Wanjia Zhao and Wei An and Wen Liu and Wenfeng Liang and Wenjun Gao and Wenqin Yu and Wentao Zhang and X. Q. Li and Xiangyue Jin and Xianzu Wang and Xiao Bi and Xiaodong Liu and Xiaohan Wang and Xiaojin Shen and Xiaokang Chen and Xiaokang Zhang and Xiaosha Chen and Xiaotao Nie and Xiaowen Sun and Xiaoxiang Wang and Xin Cheng and Xin Liu and Xin Xie and Xingchao Liu and Xingkai Yu and Xinnan Song and Xinxia Shan and Xinyi Zhou and Xinyu Yang and Xinyuan Li and Xuecheng Su and Xuheng Lin and Y. K. Li and Y. Q. Wang and Y. X. Wei and Y. X. Zhu and Yang Zhang and Yanhong Xu and Yanhong Xu and Yanping Huang and Yao Li and Yao Zhao and Yaofeng Sun and Yaohui Li and Yaohui Wang and Yi Yu and Yi Zheng and Yichao Zhang and Yifan Shi and Yiliang Xiong and Ying He and Ying Tang and Yishi Piao and Yisong Wang and Yixuan Tan and Yiyang Ma and Yiyuan Liu and Yongqiang Guo and Yu Wu and Yuan Ou and Yuchen Zhu and Yuduan Wang and Yue Gong and Yuheng Zou and Yujia He and Yukun Zha and Yunfan Xiong and Yunxian Ma and Yuting Yan and Yuxiang Luo and Yuxiang You and Yuxuan Liu and Yuyang Zhou and Z. F. Wu and Z. Z. Ren and Zehui Ren and Zhangli Sha and Zhe Fu and Zhean Xu and Zhen Huang and Zhen Zhang and Zhenda Xie and Zhengyan Zhang and Zhewen Hao and Zhibin Gou and Zhicheng Ma and Zhigang Yan and Zhihong Shao and Zhipeng Xu and Zhiyu Wu and Zhongyu Zhang and Zhuoshu Li and Zihui Gu and Zijia Zhu and Zijun Liu and Zilin Li and Ziwei Xie and Ziyang Song and Ziyi Gao and Zizheng Pan},
      year={2025},
      eprint={2412.19437},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2412.19437}, 
}

@misc{deepseek-moe,
      title={DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models}, 
      author={Damai Dai and Chengqi Deng and Chenggang Zhao and R. X. Xu and Huazuo Gao and Deli Chen and Jiashi Li and Wangding Zeng and Xingkai Yu and Y. Wu and Zhenda Xie and Y. K. Li and Panpan Huang and Fuli Luo and Chong Ruan and Zhifang Sui and aWenfeng Liang},
      year={2024},
      eprint={2401.06066},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2401.06066}, 
}

Acknowledgments

This implementation is inspired by the DeepSeek research papers and builds upon concepts from:

• "Attention Is All You Need" (Vaswani et al., 2017)
• "Language Models are Unsupervised Multitask Learners" (Radford et al., 2019)
• "The Llama 3 Herd of Models" (Meta AI, 2024)
• "GQA: Training Generalized Multi-Query Transformer Models" (Ainslie et al., 2023)

Special thanks to the DeepSeek team for their innovative contributions to efficient transformer architectures.