原理精讲

Tokenization 基础 BPE 算法详解 GPT 系列 Tokenizer BPE 训练工程化

从 token ids 到 logits Embedding 与 LM Head

Attention 机制

Self-Attention 到 GQA Attention Sink

位置编码基础 RoPE 数学推导 RoPE 代码实现长度外推

GPU 编程基础

GPU 架构基础张量布局 Triton 入门：向量加法

Flash Attention 原理详解从朴素实现到 Auto-Tuning Block Pointer 与多维支持 Causal Masking 优化 Grouped Query Attention 反向传播实现

分布式训练

数据并行 ZeRO 优化器全分片数据并行张量并行流水线并行多维混合并行

动手训练

预训练数据 Tokenizer 训练模型架构数据流水线训练循环监控与验证

系统工程FlashAttention

反向传播实现

会员专享

实现 Flash Attention 的梯度计算，通过 Recomputation 实现内存高效的训练。

登录以继续阅读

这是一篇付费内容，请登录您的账户以访问完整内容。

Grouped Query Attention

实现 GQA/MQA 支持,让多个 Query Head 共享 KV,优化 KV Cache 内存占用。

分布式训练

从数据并行到多维混合并行，理解大模型训练的核心并行策略

目录

为什么需要自定义反向传播？

PyTorch 自动微分的局限性

Recomputation 策略

Attention 反向传播的数学原理

前向传播回顾

\frac{\partial \mathcal{L}}{\partial \mathbf{V}}

\frac{\partial \mathcal{L}}{\partial \mathbf{P}}

\frac{\partial \mathcal{L}}{\partial \mathbf{S}}

(Softmax 反向传播)

\frac{\partial \mathcal{L}}{\partial \mathbf{Q}}

\frac{\partial \mathcal{L}}{\partial \mathbf{K}}

完整的梯度计算流程

代码实现解析

Forward Kernel 的修改

Backward Kernel 实现

关键实现细节

1. 循环顺序的反转

2. Atomic Add 处理 dQ

3. 重新计算 P 而非保存

torch.autograd.Function 封装

数值正确性测试

内存占用对比

设计权衡与优化方向

Recomputation 的开销

进一步优化方向