Select a result to preview

CUDA4_使用CUDA实现Transformer结构

CUDA（四）：使用 CUDA 实现 Transformer 结构

作者: 紫气东来
原文链接: https://zhuanlan.zhihu.com/p/694416583
发布日期: 2024年

文章概述

本文是 CUDA 系列的第四篇，聚焦于使用 CUDA 从零实现 Transformer 架构。作者基于 OpenAI 科学家 Andrej Karpathy 开源的 llm.c 项目，详细解析了如何脱离深度学习框架，直接使用 CUDA 构建 LLM 的算子模块、训练过程和计算优化。这是一次深入底层的技术探索，旨在帮助开发者理解 Transformer 的实现细节和性能优化技巧。

核心价值：

一、关键算子及模块的实现

作者将算子和模块的实现分为基础算子和核心模块两大类，每个部分都有独立的详细文章。

1.1 基础算子的实现

1. LayerNorm 算子

2. Softmax 算子

3. Cross Entropy 损失函数

4. AdamW 优化器

5. GELU 激活函数与残差连接

6. 矩阵乘法（GEMM）

1.2 核心模块的实现

1. Embedding 层与 LM Head 层

2. Self-Attention 机制

二、训练过程的实现

2.1 训练前的准备工作

2.1.1 训练数据下载并分词

使用 Python 脚本进行数据准备：

python prepro_tinyshakespeare.py

执行操作:

  1. 数据下载: 从 GitHub 下载 Tiny Shakespeare 数据集
  2. 分词处理: 使用 GPT-2 tokenizer 进行分词
  3. 数据集划分: 分为训练集和验证集

输出结果:

Saved 32768 tokens to data/tiny_shakespeare_val.bin
Saved 305260 tokens to data/tiny_shakespeare_train.bin

2.1.2 标准模型转换

主要操作:

  1. 加载预训练模型: 加载标准的 GPT-2 模型
  2. 格式转换:
    • 保存为 FP32 和 BF16 两个版本的二进制文件
    • 词表大小从 50257 padding 到 50304（提高矩阵乘法效率）
  3. Tokenizer 转换: 将分词器保存为二进制格式
  4. 调试状态保存: 保存输入和中间状态用于验证

输出文件:

性能基准（PyTorch 版本）:

iteration 1, loss: 5.2700, time: 69.203ms, tok/s: 3699.27
iteration 2, loss: 4.0597, time: 48.847ms, tok/s: 5240.88
...
iteration 10, loss: 0.3765, time: 62.727ms, tok/s: 4081.15
final 9 iters avg: 58.621ms
peak memory consumption: 2392 MiB

2.2 训练过程

2.2.1 训练过程的实现

核心代码结构:

// 1. 构建 GPT-2 模型
GPT2 model;
gpt2_build_from_checkpoint(&model, load_filename);

// 2. 初始化分词器
Tokenizer tokenizer;
tokenizer_init(&tokenizer, "gpt2_tokenizer.bin");

// 3. 初始化数据加载器
DataLoader train_loader, val_loader;
dataloader_init(&train_loader, &multi_gpu_config, train_tokens_filename, B, T);
dataloader_init(&val_loader, &multi_gpu_config, val_tokens_filename, B, T);

// 4. 训练循环
for (int step = 0; step < max_steps; step++) {
    dataloader_next_batch(&train_loader);
    gpt2_forward(&model, train_loader.inputs, train_loader.targets, B, T, false);
    gpt2_zero_grad(&model);
    gpt2_backward(&model);
    gpt2_update(&model, learning_rate, 0.9f, 0.999f, 1e-8f, 0.0f, step+1);
}

编译与运行:

make train_gpt2cu
./train_gpt2cu

训练配置:

| Parameter             | Value                                              |
+-----------------------+----------------------------------------------------+
| batch size B          | 4                                                  |
| sequence length T     | 1024                                               |
| learning rate         | 3.000000e-04                                       |
| device                | NVIDIA GeForce RTX 4090                            |
| precision             | BF16                                               |
| num_layers L          | 12                                                 |
| num_heads NH          | 12                                                 |
| channels C            | 768                                                |
| num_parameters        | 124475904                                          |

内存分配:

训练性能（纯 CUDA 实现）:

step 1/74: train loss 4.364526 (43.898880 ms, 93305.343750 tok/s)
step 2/74: train loss 4.501646 (41.996288 ms, 97532.406250 tok/s)
...
step 20/74: train loss 3.743915 (42.046463 ms, 97677.789062 tok/s)
val loss 3.702019

生成示例:

Nay, thou amongst I: he never came to hear this, that I shall not entertain 
the will of myself. Call to play, Gregg! Take W.L<|endoftext|>Your pleas 
about my ignorance, Oh, move your way, charge, plunder: plan for This bloody 
Staff expedition:

2.2.2 使用 cuDNN 模块加速

编译与运行:

make train_gpt2cu USE_CUDNN=1
./train_gpt2cu

性能提升:

cuDNN 版本训练日志:

step 1/74: train loss 4.370480 (362.220551 ms, 11308.027344 tok/s)  # 首次编译
step 2/74: train loss 4.505543 (32.261120 ms, 126963.945312 tok/s)
step 3/74: train loss 4.421638 (32.068607 ms, 127354.804688 tok/s)
...
step 20/74: train loss 3.746799 (32.408577 ms, 127089.820312 tok/s)
val loss 3.702071

内存优化:

2.2.3 更大规模训练示例

使用 TinyStories 数据集进行训练：

# 数据准备
python prepro_tinystories.py

# 训练（更大的 batch size）
./train_gpt2cu -i data/TinyStories -v 250 -s 250 -g 144 -o stories.log -b 32

训练配置:

性能表现:

step 1/28248: train loss 2.386457 (780.998657 ms, 41956.539062 tok/s)
step 2/28248: train loss 3.283106 (215.723007 ms, 151898.468750 tok/s)
step 3/28248: train loss 2.365078 (219.082748 ms, 150703.875000 tok/s)
...
step 20/28248: train loss 1.869874 (220.104706 ms, 149131.281250 tok/s)

损失下降趋势:

三、环境配置要求

该项目对软件版本有较高要求，推荐配置：

OS: Ubuntu 22.04
Driver: 550.54.15
CUDA: 12.4
PyTorch: 2.4.0.dev20240513+cu121
cuDNN: 8.9.7.29
cudnn-frontend: 1.4.0

关键依赖:

四、性能分析与优化总结

4.1 性能对比

实现方式 吞吐量 (tok/s) 相对提升 内存占用 (激活)
PyTorch ~4,000 基准 未知
纯 CUDA ~97,000 24x 2,853 MiB
CUDA + cuDNN ~127,000 32x 1,703 MiB

4.2 关键优化技术

1. 算子层面优化

2. 内存优化

3. 计算优化

4. 系统层面优化

4.3 性能瓶颈分析

内存带宽瓶颈:

计算瓶颈:

同步开销:

五、实践建议与最佳实践

5.1 开发建议

  1. 从简单开始: 先实现正确的版本，再进行优化
  2. 逐步优化: 每次优化后验证正确性和性能提升
  3. 使用 profiler: 使用 Nsight Systems/Compute 分析性能
  4. 参考标准实现: 对比 PyTorch 的输出验证正确性
  5. 保存中间状态: 便于调试和验证

5.2 调试技巧

  1. 数值验证: 与 PyTorch 实现对比中间结果
  2. 梯度检查: 使用数值梯度验证反向传播
  3. 单元测试: 为每个算子编写测试
  4. 可视化: 使用 TensorBoard 监控训练过程
  5. 日志记录: 详细记录损失、梯度范数等指标

5.3 性能调优

  1. Profile 驱动: 先分析瓶颈再优化
  2. 内存优化优先: GPU 通常受内存带宽限制
  3. 批量大小调整: 找到最佳的 batch size
  4. 混合精度: 使用 BF16/FP16 加速训练
  5. 库函数优先: 优先使用 cuBLAS、cuDNN 等优化库

六、技术亮点与创新

6.1 项目特色

  1. 完全脱离框架: 不依赖 PyTorch/TensorFlow
  2. 教育价值高: 代码清晰，注释详细
  3. 性能优秀: 接近或超过框架性能
  4. 模块化设计: 易于理解和扩展
  5. 渐进式优化: 展示从基础到高级的优化过程

6.2 学习价值

七、总结与展望

7.1 核心要点

  1. 算子实现是基础: LayerNorm、Softmax、Attention 等算子的高效实现是关键
  2. 优化是系统工程: 需要从算法、实现、硬件多个层面综合考虑
  3. cuDNN 加速显著: 使用优化库可以获得 30% 以上的性能提升
  4. 内存管理重要: 合理的内存分配和复用可以支持更大的模型和 batch size

7.2 未来方向

  1. 更多优化技术: Flash Attention 2、PagedAttention 等
  2. 分布式训练: 多 GPU、多节点训练支持
  3. 更大模型: 支持 GPT-3、LLaMA 等更大规模模型
  4. 推理优化: 量化、剪枝、知识蒸馏等技术
  5. 新硬件支持: 适配 H100、A100 等新一代 GPU

7.3 参考资源

附录：关键代码片段

A.1 前向传播核心流程

void gpt2_forward(GPT2 *model, int* inputs, int* targets, int B, int T, bool training) {
    // 1. Token embedding + Position embedding
    encoder_forward(model->acts.encoded, inputs, model->params.wte, 
                    model->params.wpe, B, T, C);
    
    // 2. Transformer blocks
    for (int l = 0; l < L; l++) {
        // LayerNorm 1
        layernorm_forward(ln1, encoded, ln1w, ln1b, B, T, C);
        
        // Multi-head attention
        attention_forward(atty, att, ln1, l0, l1, B, T, C, NH);
        
        // Residual connection 1
        residual_forward(residual, encoded, atty, B*T*C);
        
        // LayerNorm 2
        layernorm_forward(ln2, residual, ln2w, ln2b, B, T, C);
        
        // MLP
        matmul_forward(fch, ln2, fcw, B, T, C, 4*C);
        gelu_forward(fch_gelu, fch, B*T*4*C);
        matmul_forward(fcproj, fch_gelu, fcprojw, B, T, 4*C, C);
        
        // Residual connection 2
        residual_forward(encoded, residual, fcproj, B*T*C);
    }
    
    // 3. Final LayerNorm
    layernorm_forward(model->acts.lnf, encoded, lnfw, lnfb, B, T, C);
    
    // 4. LM head
    matmul_forward(model->acts.logits, model->acts.lnf, 
                   model->params.wte, B, T, C, V);
    
    // 5. Loss calculation
    if (targets != NULL) {
        softmax_crossentropy_forward(model->acts.losses, model->acts.logits, 
                                      targets, B, T, V);
    }
}

A.2 反向传播核心流程

void gpt2_backward(GPT2 *model) {
    // 1. Loss gradient
    softmax_crossentropy_backward(dlogits, model->acts.losses, 
                                   model->acts.logits, targets, B, T, V);
    
    // 2. LM head backward
    matmul_backward(dlnf, dwte, dlogits, model->acts.lnf, 
                    model->params.wte, B, T, C, V);
    
    // 3. Final LayerNorm backward
    layernorm_backward(dresidual, dlnfw, dlnfb, dlnf, residual, lnfw, B, T, C);
    
    // 4. Transformer blocks backward (reverse order)
    for (int l = L-1; l >= 0; l--) {
        // Residual 2 backward
        residual_backward(dresidual2, dfcproj, dresidual, B*T*C);
        
        // MLP backward
        matmul_backward(dfch_gelu, dfcprojw, dfcproj, fch_gelu, 
                        fcprojw, B, T, 4*C, C);
        gelu_backward(dfch, fch, dfch_gelu, B*T*4*C);
        matmul_backward(dln2, dfcw, dfch, ln2, fcw, B, T, C, 4*C);
        
        // LayerNorm 2 backward
        layernorm_backward(dresidual, dln2w, dln2b, dln2, residual, ln2w, B, T, C);
        
        // Residual 1 backward
        residual_backward(dencoded, datty, dresidual, B*T*C);
        
        // Attention backward
        attention_backward(dln1, dl0, dl1, datty, ln1, l0, l1, B, T, C, NH);
        
        // LayerNorm 1 backward
        layernorm_backward(dencoded, dln1w, dln1b, dln1, encoded, ln1w, B, T, C);
    }
    
    // 5. Embedding backward
    encoder_backward(dwte, dwpe, dencoded, inputs, B, T, C);
}

A.3 AdamW 优化器更新

void adamw_update(float* params, float* grads, float* m, float* v, 
                  int n, float lr, float beta1, float beta2, 
                  float eps, float weight_decay, int t) {
    for (int i = 0; i < n; i++) {
        // Update biased first moment estimate
        m[i] = beta1 * m[i] + (1.0f - beta1) * grads[i];
        
        // Update biased second raw moment estimate
        v[i] = beta2 * v[i] + (1.0f - beta2) * grads[i] * grads[i];
        
        // Compute bias-corrected first moment estimate
        float m_hat = m[i] / (1.0f - powf(beta1, t));
        
        // Compute bias-corrected second raw moment estimate
        float v_hat = v[i] / (1.0f - powf(beta2, t));
        
        // Update parameters with weight decay
        params[i] = params[i] * (1.0f - lr * weight_decay) 
                    - lr * m_hat / (sqrtf(v_hat) + eps);
    }
}

诗句引用:

浅水池塘莲叶香，红尘道路柳阴长。 —— 彭汝砺 《和初夏》

本文档基于 llm.c 项目整理，详细代码请参考原项目仓库。