train_gpt2.c / train_gpt2.cu 程式碼導讀

檔案架構

train_gpt2.c — CPU 參考實作

約 1182 行的純 C 程式碼，包含完整訓練管線：

區段	行數	說明
層實作	35-435	各層的 forward/backward 函式
模型定義	523-705	GPT2Config、ParameterTensors、ActivationTensors、GPT2 struct
模型操作	707-1033	build_from_checkpoint、forward、backward、zero_grad、update
採樣器	1049-1073	token 採樣、RNG
main()	1077-1182	訓練迴圈

train_gpt2.cu — CUDA 加速實作

約 1904 行，結構與 CPU 版本對應但更複雜：

• 混合精度 (floatX 抽象為 BF16/FP16/FP32)
• cuBLASLt 矩陣乘法
• 自訂 CUDA kernel（LayerNorm、Attention、GeLU、AdamW）
• 多 GPU 支援（NCCL、ZeRO）
• 中斷續跑功能

各層 forward 函式

encoder_forward     →   (B,T,C) token+position embeddings
layernorm_forward   →   normalize + scale + shift
matmul_forward      →   {Q,K,V} projections, attention projection, FC layers
attention_forward   →   scaled dot-product attention with causal mask
gelu_forward        →   GELU non-linearity
residual_forward    →   residual connection (element-wise add)
softmax_forward     →   softmax for output probabilities
crossentropy_forward →  cross-entropy loss

記憶體配置（CUDA 版本）

所有參數連續儲存在 GPU 記憶體中，透過 malloc_and_point_parameters 將各 tensor 指標指向正確偏移位置：

[wte][wpe][ln1w][ln1b][qkvw][qkvb]...[lnfw][lnfb]

激活值同樣連續配置，支援 recompute 標誌節省記憶體。