llm.c 文件重點

技術核心

1. 純 C 實作 (train_gpt2.c)

• 無框架依賴: 不使用 PyTorch/TensorFlow，僅用標準 C 函式庫
• FP32 精度: 所有計算使用 float 型別
• OpenMP 並行: 在矩陣乘法、LayerNorm 等層使用 #pragma omp parallel for
• 啟發式優化: matmul 使用 loop unrolling (LOOP_UNROLL=8) 提升快取命中

2. CUDA 實作 (train_gpt2.cu)

• 混合精度訓練: 預設 BF16，支援 FP32 和 FP16
• cuBLASLt 矩陣乘法: 使用 cublasLtMatmul 取代手寫 kernel（速度接近硬體上限）
• cuDNN Flash Attention: 選用 cuDNN 的 fused attention 實作
• 手寫 Kernel (llmc/ 目錄):
  • LayerNorm: layernorm_forward_kernel3 (無 shared memory) 和 kernel6 (有 shared memory)
  • Attention: 自訂 kernel 支援 causal mask
  • AdamW: adamw_kernel3 融合更新
  • GeLU: 自訂 forward/backward kernel

3. 記憶體管理

• 參數佈局: 所有權重連續配置在單一 GPU 記憶體區塊，用迭代器分配各 tensor
• 激活記憶體: 21 個 activation tensor，支援 recompute 選項以節省 VRAM
• ZeRO 優化: Stage 0-3，支援分散式梯度降

4. 訓練管線

資料載入 → 編碼 → 12層Transformer → LayerNorm → Linear → Softmax → Loss
                                                               ↓
AdamW 更新 ← 梯度裁剪 ← 梯度同步(NCCL) ← 反向傳播 ← CrossEntropy

5. 多 GPU 支援

• MPI 初始化分散式環境
• NCCL AllReduce 梯度同步
• ZeRO-1 對優化器狀態進行分片
• 支援 Slurm 排程器

學習資源

• doc/layernorm/layernorm.md: 從 PyTorch 到 C 的逐步教學，涵蓋 forward/backward
• dev/cuda/: 手寫 kernel 的 playground
• test_gpt2.c / test_gpt2.cu: 與 PyTorch 輸出比對的一致性測試