CS336 LecTure 3 笔记补充
Attention 机制的核心参数与学习过程
在基础的 Attention 模块中，输入 Input 通常是经过 Embedding 处理后的向量 \(x\)。Query (\(q\))、Key (\(k\)) 和 Value (\(v\)) 均由该输入 \(x\) 乘以各自对应的可学习参数矩阵（如 \(\omega_q, \omega_k, \omega_v\)）生成。
- 相似度计算：通过 Query 和 Key 的点积（Query times Key）来计算序列中元素之间的相似度。
- 参数学习：需要特别明确的是，\(\omega_q, \omega_k, \omega_v\) 这些权重参数是模型在训练过程中需要不断学习和优化的，而非固定不变的 Embedding 权重。
多头注意力（Multi-Head Attention）
多头的核心作用
多头机制决定了模型能够从多少个不同的“子空间”或角度同时观察输入序列。例如，在一个多头注意力模块中，某一个注意力头可能专门负责关注句子的语法结构，而另一个头则负责捕捉语义信息。
关键维度定义的区别
在多头注意力的公式 MultiHead(Q, K, V) = Concat(head1, ..., headh) 中，涉及几个极易混淆的维度概念：
- \(d_{model}\)（全局维度/模型维度）：这是输入和输出向量的全局长度。例如，一个词（如 "I"）经过 Embedding 后可能是一个 768 维的向量，这里的 768 就是 \(d_{model}\)。
- \(h\)（注意力头数）：代表将全局维度切分成的独立头的数量。例如，将 768 维切分为 8 个头，则 \(h=8\)。
- \(d_k\)（分头维度/子空间维度）：这是每个独立注意力头实际处理的向量长度。其计算公式为 \(d_k = d_{model} / h\)。在上述例子中，每个头只需要处理 64 维（\(768 / 8 = 64\)）的数据。\(d_k\) 也是 Query 和 Key 进行点积计算相关性时使用的空间大小，同时决定了每个头输出的特征空间大小。
张量维度的变换过程
在标准的张量运算中，输入数据的维度通常表示为 [Batch Size, Sequence Length, Dimension]：
- Batch Size：批次大小。
- Sequence Length：序列长度（即句子中 Word 的数量）。
- Dimension：全局维度 \(d_{model}\)。
为了引入多头机制，我们对第三个参数（Dimension）进行扩展与重塑。具体操作是将 \(d_{model}\) 拆分为 \(h \times d_k\)。在代码实现和维度变换时，通常会将代表头数的 \(h\) 提前，使得张量维度变为 [Batch Size, Sequence Length, h, d_k]。这样，每个头就能在独立的 \(d_k\) 局部区域内并行计算注意力，最后再将所有头的结果拼接（Concat）起来，恢复为原始的全局维度。
Embedding、RoPE 位置编码、前馈网络（FFN）、交叉熵损失以及 AdamW 优化器是五个很重要的核心组件，和 CS336 笔记 中的 FLashAttention 结合我们就可以得到一个 Transformer
1. Embedding（词元嵌入层）
Embedding 的本质是一个形状为 (vocab_size, d_model) 的权重矩阵。前向传播时，不需要做矩阵乘法，而是直接将输入的 token id 作为索引，去权重矩阵中“查表”取出对应的行向量。
首先，为什么我们需要 Embedding？
如果没有 Embedding，我们直接给一个词汇表里的各个单词编号，第五号和第六号单词可能天差地别，但是模型会认为二者是接近的（因为五和六离得很近），为了解决这个问题，我们需要建立 Embedding
它本质上是什么？
一个巨大的、可学习的查找表。把它想象成一个巨大的 Excel 表格：
- 行 : 代表词汇表里的每一个词（或
token）。总共有vocab_size行。 - 列 : 代表每个词的“特征”或“语义”。总共有
embedding_dim列。
所以，Embedding 层就是一个形状为 (vocab_size, embedding_dim) 的矩阵。
工作流程非常简单，就是查表：
- 输入： 一个词的 ID，比如
cat的 ID 是 42。 - 查找： Embedding 层接收到 42，然后去它的“大表格”里找到第 42 行。
- 输出： 取出第 42 行的整行数据。这行数据就是一个长度为
embedding_dim的向量。
这个向量，就是 "cat" 这个词的 Embedding。
import torch
import torch.nn as nn
import triton
import triton.language as tl
import math
import torch
import torch.nn as nn
class Embedding(nn.Module):
# 定义一个名为 Embedding 的类，它继承自 PyTorch 的 nn.Module
# 这意味着它是一个标准的神经网络模块，可以被添加到模型中，并且它的参数可以被自动追踪和优化
def __init__(self, num_embeddings: int, embedding_dim: int):
# 初始化
# num_embeddings: 词汇表的大小 (Vocab Size)
# embedding_dim: 每个词向量的维度 (Embedding Dimension)
super().__init__()
# 调用父类 nn.Module 的初始化函数 继承时的标准写法
self.num_embeddings = num_embeddings
self.embedding_dim = embedding_dim
# 保存传入的参数，方便后续使用或调试
self.W = nn.Parameter(torch.empty(num_embeddings, embedding_dim))
# 创建权重矩阵 W
# torch.empty(...) 创建一个形状为 (词汇表大小, 向量维度) 的空矩阵
# nn.Parameter(...) 告诉 PyTorch：这个矩阵 W 是模型的可学习参数，在训练时，自动计算它的梯度并更新
std = 1.0 / (embedding_dim ** 0.5)
# 计算初始化的标准差
# 这是一个非常经典的初始化技巧（Xavier 初始化的一种变体）
# 目的是让初始化的数值不要太大也不要太小，保持方差稳定，防止梯度消失或爆炸
# 维度越高，每个数就要分得越小，所以除以根号下维度
nn.init.trunc_normal_(self.W, mean=0.0, std=std, a=-3.0 * std, b=3.0 * std)
# 对矩阵 W 进行“截断正态分布”初始化
# 意思是 从均值为 0，标准差为 std 的正态分布中随机取值填入 W
# 如果有数值超过了 [-3*std, 3*std] 的范围，就扔掉重取。
# 避免初始权重中出现极端的离群值，让训练更稳定
def forward(self, token_ids: torch.Tensor) -> torch.Tensor:
# 前向传播函数
return self.W[token_ids]
# 查表
# 在 PyTorch 中，直接用一个张量作为索引（self.W[...]），会自动进行“高级索引”操作
# 它会根据 token_ids 里的数字，去 W 矩阵里把对应的行抠出来
# 比如 token_ids 是 [10, 25]，它就会把 W 的第 10 行和第 25 行取出来，拼成一个新的张量返回
# 输入形状: (Batch, Seq_Len) -> 输出形状: (Batch, Seq_Len, Embedding_Dim)
2. RoPE（旋转位置编码）
RoPE 的核心思想是将 Q 和 K 的向量每两个维度看作二维平面上的点，根据 token 的位置施加不同角度的旋转。这样在计算注意力点积时，结果自然只与相对位置差有关。
在 Transformer 模型中，注意力机制（Self-Attention）本身是“无序”的，它不知道哪个词在前面，哪个词在后面。因此，我们需要给词加上“位置信息”。
传统的做法（绝对位置编码）是给每个词加上一个固定的位置向量，但这种方式在处理超长文本时效果不好。而 RoPE 巧妙地结合了绝对位置和相对位置的优点，它的核心目标是：通过绝对位置编码的方式，实现相对位置编码。
RoPE 的灵感来源于几何学中的二维旋转。
- 旋转机制：它不是把位置信息“加”到词向量上，而是把词向量在二维平面中“旋转”一个与位置相关的角度。比如，第一个位置的向量旋转 \(\theta\) 角度，第二个位置旋转 \(2\theta\) 角度，依此类推。
- 保留相对距离：这种旋转有一个非常优雅的数学性质——当两个词（Query 和 Key）进行点积计算注意力分数时，它们各自的旋转角度会相互抵消一部分，最终结果只与这两个词的相对距离（即位置差）有关。
- 不改变向量长度：旋转操作不会改变原始向量的长度（模长），这有助于保持模型训练的稳定性。
def apply_rope(x, cos, sin):
# 将特征维度对半拆分，相当于把高维向量看作多个二维平面的组合
x1, x2 = x.chunk(2, dim=-1)
# 二维旋转的数学公式：(x1, x2) -> (-x2, x1)
rotated_x = torch.cat([-x2, x1], dim=-1)
# 利用预计算好的正余弦值，完成旋转操作
return x * cos + rotated_x * sin
注：在实际 Transformer Block 中，通常预先生成一张包含所有位置的正余弦表（缓存机制），在前向传播时根据当前序列长度切片使用。
3. FFN（前馈网络 - SwiGLU 变体）
现代 LLM通常使用 SwiGLU 替代传统的 ReLU FFN。它包含三个线性层（参数量是传统 FFN 的 1.5 倍），通过门控机制控制信息流。
Attention 层只是把相关词汇的信息聚合在了一起，但它并没有对这些信息进行深度的非线性处理。FFN 的作用就是对 Attention 提取出的特征进行逐位置的独立加工。
Transformer 模型中绝大部分的参数和学到的知识（比如语法规则、世界常识）都储存在 FFN 层中。
class SwiGLU_FFN(nn.Module):
def __init__(self, D_MODEL: int, D_FF: int):
super().__init__()
# W1: 门控分支 (Gate) - 决定“开多大阀门”
self.w1 = nn.Linear(D_MODEL, D_FF, bias=False)
# W3: 值分支 (Value) - 提供“实际的水流（信息）”
self.w3 = nn.Linear(D_MODEL, D_FF, bias=False)
# W2: 输出投影 - 将处理后的信息映射回原来的维度
self.w2 = nn.Linear(D_FF, D_MODEL, bias=False)
def forward(self, x):
# 门控分支：通过 SiLU 激活函数（一种平滑的 Sigmoid 变体）生成门控信号
gate = torch.nn.functional.silu(self.w1(x))
# 值分支：提取原始信息
value = self.w3(x)
# 门控机制：将门控信号与值分支逐元素相乘
# 门控信号的值通常在 0 到 1 之间，相当于一个“过滤器”
# 乘以 gate 后，不重要的信息会被抑制（接近0），重要的信息会被保留
return self.w2(gate * value)
4. Cross Entropy（交叉熵损失）
在语言模型中，计算交叉熵时需要将 logits 转为 float 以避免精度和溢出问题，同时将 target 转为 int64 以使用 gather 高效提取正确类别的概率。
模型的任务是“预测下一个词”
- 假设当前上下文中，下一个词真实应该是“猫”。
- 模型经过一系列计算后，输出了一个概率分布，它认为下一个词是“猫”的概率只有 10%，是“狗”的概率是 90%。
这时候，模型“考砸了”。交叉熵损失函数的作用，就是把这种“考砸了”的程度量化成一个具体的数字（Loss 值）。预测得越离谱，Loss 值越大；预测得越准，Loss 值越小。模型在训练时，唯一的目标就是通过反向传播（loss.backward()）去不断降低这个 Loss 值。
交叉熵的本质来源于信息论，它的公式是：
Loss = - log(模型预测正确类别的概率)
具体可以看CS336 笔记
- 如果模型对正确答案的预测概率是 0.9，Loss = -log(0.9) ≈ 0.105 （惩罚很小，说明模型很棒）
- 如果模型对正确答案的预测概率是 0.1，Loss = -log(0.1) ≈ 2.302 （惩罚很大，说明模型很差）
- 如果模型对正确答案的预测概率是 0.01，Loss = -log(0.01) ≈ 4.605 （惩罚极大）
def cross_entropy_loss(logits, targets, ignore_index=-100):
"""
Args:
logits: 模型输出的原始分数 (Batch, Seq_Len, Vocab_Size)
targets: 真实的下一个 Token ID (Batch, Seq_Len)
ignore_index: 需要被忽略的 Token ID（通常 Padding 会被设为 -100）
"""
# 获取张量的形状
BATCH, N_CTX, VOCAB_SIZE = logits.shape
# 提升精度并展平
# 将 logits 转为 float32 防止后续 exp() 计算时发生半精度溢出
# 将形状从 (BATCH, N_CTX, VOCAB_SIZE) 展平为 (BATCH * N_CTX, VOCAB_SIZE)
logits = logits.float().view(-1, VOCAB_SIZE)
# 将 targets 转为 int64 (long) 类型，这是 PyTorch 索引操作的硬性要求
# 形状从 (BATCH, N_CTX) 展平为 (BATCH * N_CTX,)
targets = targets.long().view(-1)
# 过滤掉需要忽略的 Token
# 生成一个布尔掩码，标记出哪些位置是真实的 Token
valid_mask = targets != ignore_index
# 处理极端边界情况
# 如果当前 Batch 中全都是 Padding（即没有任何有效的 Token）
if not valid_mask.any():
# 直接返回一个值为 0.0 的标量张量
# 必须加上 requires_grad=True，否则后续 loss.backward() 会报错
return torch.tensor(0.0, device=logits.device, requires_grad=True)
# 利用掩码提取出所有有效的 logits 和 targets
# 这一步将原本包含 Padding 的二维矩阵，压缩成了一维的有效数据
logits = logits[valid_mask]
targets = targets[valid_mask]
# 减去当前行的最大值
# 保持维度 (N, 1) 以便后续能正确广播
max_logits, _ = torch.max(logits, dim=-1, keepdim=True)
# 减去最大值，这在数学上等价于原公式，但能完美防止 exp() 溢出
logits = logits - max_logits
# 计算 LogSoftmax
# 计算分母的对数形式：log(sum(exp(logits)))
log_sum_exp = torch.log(torch.sum(torch.exp(logits), dim=-1))
# 计算每个类别的对数概率：logits - log_sum_exp
log_probs = logits - log_sum_exp.unsqueeze(-1)
# 获取真实类别的对数概率
# targets.unsqueeze(-1) 将形状变为 (N, 1) 以配合 gather 的索引要求
# gather 沿着最后一个维度（Vocab_Size），根据 targets 提供的索引精准提取概率
# squeeze(-1) 去掉多余的维度，恢复为 (N,)
target_log_probs = log_probs.gather(dim=-1, index=targets.unsqueeze(-1)).squeeze(-1)
# 计算最终的 Loss
# 加上负号（因为我们要最大化对数概率，等价于最小化负对数概率）
# 对所有有效 Token 的 Loss 求平均，得到当前 Batch 的最终损失值
return -target_log_probs.mean()
5. AdamW 优化器
在传统的 Adam 优化器中，为了防止模型过拟合（比喻一下就是死记硬背），通常会加入 L2 正则化（即权重衰减）。但 Adam 在实现时，错误地把权重衰减加到了梯度计算里，导致当学习率发生变化时（比如使用 Warmup 或余弦退火），正则化的强度也会跟着剧烈波动，严重削弱了模型的泛化能力。
AdamW 的核心突破在于，将“权重衰减”与“梯度更新”彻底解耦。它让 Adam 专心负责高效地寻找最优解，而让权重衰减独立地在每次更新后对参数进行惩罚，从而更好地控制模型复杂度，防止过拟合。
AdamW 在每次更新参数时会进行如下操作运行：
- 自适应学习率与动量：
- 一阶矩（动量）：记录梯度的历史方向，让模型在平坦的误差曲面上加速，减少震荡。
- 二阶矩（自适应学习率）：记录梯度的历史平方，为每个参数提供个性化的学习率。更新频繁的参数学习率会变小，更新少的参数学习率会变大。
- 偏差修正：因为模型刚开始训练时，历史矩估计为 0，会产生偏差。AdamW 会进行数学修正，让初始阶段的更新更准确。
- 解耦的权重衰减：
在参数更新时，不经过梯度，而是直接用参数本身乘以一个衰减系数（1 - 学习率 * 权重衰减系数）。这确保了正则化的强度是恒定可控的，不受学习率调度策略的干扰。 - 参数更新：
将上述两步的结果结合起来，完成最终的参数更新。
from torch.optim import Optimizer
class AdamW(Optimizer):
def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.01):
# 将超参数打包成字典
defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
# 调用父类初始化，将模型参数和超参数绑定
super().__init__(params, defaults)
@torch.no_grad()
def step(self):
# 关闭自动求导 节省大量显存和计算资源
for group in self.param_groups:
# 从参数组中提取超参数
beta1, beta2 = group['betas']
# 遍历该组中的每一个可学习参数（如权重矩阵 W、偏置 b 等）
for p in group['params']:
# 如果该参数没有梯度（例如被冻结的层），直接跳过
if p.grad is None:
continue
# 获取该参数的内部状态（用于存储历史梯度的动量）
state = self.state[p]
# 状态初始化（仅在第一次遇到该参数时执行）
if len(state) == 0:
state['step'] = 0
# 一阶矩估计：记录梯度的历史均值
state['exp_avg'] = torch.zeros_like(p)
# 二阶矩估计：记录梯度的历史方差（平方均值）
state['exp_avg_sq'] = torch.zeros_like(p)
# 更新步数计数器
state['step'] += 1
# 获取当前参数的梯度
grad = p.grad
# 更新一阶矩和二阶矩（使用 In-place 原地操作，节省显存）
# 一阶矩 = 0.9 * 旧一阶矩 + 0.1 * 当前梯度
state['exp_avg'].mul_(beta1).add_(grad, alpha=1 - beta1)
# 二阶矩 = 0.999 * 旧二阶矩 + 0.001 * 当前梯度的平方
state['exp_avg_sq'].mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
# 偏差修正
# 因为初始时刻矩估计为 0，会导致初期更新步长偏小，这里进行数学修正
bias_correction1 = 1 - beta1 ** state['step']
bias_correction2 = 1 - beta2 ** state['step']
# 计算修正后的学习率
step_size = group['lr'] / bias_correction1
# 计算分母：sqrt(二阶矩 / 修正因子) + eps
# eps 是为了防止分母为 0 导致除零错误
denom = (state['exp_avg_sq'].sqrt() / (bias_correction2 ** 0.5)).add_(group['eps'])
# AdamW 的核心：解耦的权重衰减
# 直接在参数上乘以 (1 - 学习率 * 衰减系数)
p.mul_(1 - group['lr'] * group['weight_decay'])
# 参数更新：p = p - step_size * (一阶矩 / 分母)
# addcdiv_ 表示 add(当前值, 分子 / 分母 * value)
p.addcdiv_(state['exp_avg'], denom, value=-step_size)
Transformer
将所有代码结合在一起我们就得到了需要的 Transformer
Show Code
# transformer
import torch
import torch.nn as nn
import triton
import triton.language as tl
import math
# ================= Embedding =================
class Embedding(nn.Module):
def __init__(self, num_embeddings: int, embedding_dim: int):
super().__init__()
self.num_embeddings = num_embeddings
self.embedding_dim = embedding_dim
self.W = nn.Parameter(torch.empty(num_embeddings, embedding_dim))
std = 1.0 / (embedding_dim ** 0.5)
nn.init.trunc_normal_(self.W, mean=0.0, std=std, a=-3.0 * std, b=3.0 * std)
def forward(self, token_ids: torch.Tensor) -> torch.Tensor:
return self.W[token_ids]
# ================= RoPE =================
def apply_rope(x, cos, sin):
"""
x: (BATCH, N_HEADS, N_CTX, D_HEAD)
cos, sin: (N_CTX, 1, 1, D_HEAD)
"""
x1, x2 = x.chunk(2, dim=-1)
rotated_x = torch.cat([-x2, x1], dim=-1)
return x * cos + rotated_x * sin
# ================= SwiGLU FFN =================
class SwiGLU_FFN(nn.Module):
def __init__(self, D_MODEL: int, D_FF: int):
super().__init__()
self.w1 = nn.Linear(D_MODEL, D_FF, bias=False) # Gate
self.w3 = nn.Linear(D_MODEL, D_FF, bias=False) # Value
self.w2 = nn.Linear(D_FF, D_MODEL, bias=False) # Output
def forward(self, x):
gate = torch.nn.functional.silu(self.w1(x))
value = self.w3(x)
return self.w2(gate * value)
# ================= Cross Entropy Loss =================
def cross_entropy_loss(logits, targets, ignore_index=-100):
"""
logits: (BATCH, N_CTX, VOCAB_SIZE)
targets: (BATCH, N_CTX)
"""
BATCH, N_CTX, VOCAB_SIZE = logits.shape
logits = logits.float().view(-1, VOCAB_SIZE)
targets = targets.long().view(-1)
# 过滤掉 ignore_index 的 token
valid_mask = targets != ignore_index
if not valid_mask.any():
return torch.tensor(0.0, device=logits.device, requires_grad=True)
logits = logits[valid_mask]
targets = targets[valid_mask]
max_logits, _ = torch.max(logits, dim=-1, keepdim=True)
logits = logits - max_logits
log_sum_exp = torch.log(torch.sum(torch.exp(logits), dim=-1))
log_probs = logits - log_sum_exp.unsqueeze(-1)
target_log_probs = log_probs.gather(dim=-1, index=targets.unsqueeze(-1)).squeeze(-1)
return -target_log_probs.mean()
# ================= AdamW Optimizer =================
from torch.optim import Optimizer
class AdamW(Optimizer):
def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.01):
defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
super().__init__(params, defaults)
@torch.no_grad()
def step(self):
for group in self.param_groups:
beta1, beta2 = group['betas']
for p in group['params']:
if p.grad is None:
continue
state = self.state[p]
if len(state) == 0:
state['step'] = 0
state['exp_avg'] = torch.zeros_like(p)
state['exp_avg_sq'] = torch.zeros_like(p)
state['step'] += 1
grad = p.grad
state['exp_avg'].mul_(beta1).add_(grad, alpha=1 - beta1)
state['exp_avg_sq'].mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
bias_correction1 = 1 - beta1 ** state['step']
bias_correction2 = 1 - beta2 ** state['step']
step_size = group['lr'] / bias_correction1
denom = (state['exp_avg_sq'].sqrt() / (bias_correction2 ** 0.5)).add_(group['eps'])
# 解耦的权重衰减
p.mul_(1 - group['lr'] * group['weight_decay'])
p.addcdiv_(state['exp_avg'], denom, value=-step_size)
# ================= Triton Flash Attention =================
@triton.jit
def _fwd_kernel(
Q, K, V, sm_scale, Out,
stride_qz, stride_qh, stride_qm, stride_qk,
stride_kz, stride_kh, stride_kn, stride_kk,
stride_vz, stride_vh, stride_vn, stride_vk,
stride_oz, stride_oh, stride_om, stride_ok,
Z, H, N_CTX,
BLOCK_M: tl.constexpr, BLOCK_DMODEL: tl.constexpr, BLOCK_N: tl.constexpr,
):
start_m = tl.program_id(0)
off_hz = tl.program_id(1)
off_z = off_hz // H
off_h = off_hz % H
qvk_offset = off_z.to(tl.int64) * stride_qz + off_h.to(tl.int64) * stride_qh
v_offset = off_z.to(tl.int64) * stride_vz + off_h.to(tl.int64) * stride_vh
o_offset = off_z.to(tl.int64) * stride_oz + off_h.to(tl.int64) * stride_oh
Q_block_ptr = tl.make_block_ptr(base=Q + qvk_offset, shape=(N_CTX, BLOCK_DMODEL), strides=(stride_qm, stride_qk),
offsets=(start_m * BLOCK_M, 0), block_shape=(BLOCK_M, BLOCK_DMODEL), order=(1, 0))
K_block_ptr = tl.make_block_ptr(base=K + qvk_offset, shape=(BLOCK_DMODEL, N_CTX), strides=(stride_kk, stride_kn),
offsets=(0, 0), block_shape=(BLOCK_DMODEL, BLOCK_N), order=(0, 1))
V_block_ptr = tl.make_block_ptr(base=V + v_offset, shape=(N_CTX, BLOCK_DMODEL), strides=(stride_vn, stride_vk),
offsets=(0, 0), block_shape=(BLOCK_N, BLOCK_DMODEL), order=(1, 0))
O_block_ptr = tl.make_block_ptr(base=Out + o_offset, shape=(N_CTX, BLOCK_DMODEL), strides=(stride_om, stride_ok),
offsets=(start_m * BLOCK_M, 0), block_shape=(BLOCK_M, BLOCK_DMODEL), order=(1, 0))
m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
# 增加 boundary_check 避免 N_CTX 不整除 BLOCK_SIZE 时的越界
q = tl.load(Q_block_ptr, boundary_check=(0, 1))
for start_n in range(0, N_CTX, BLOCK_N):
k = tl.load(K_block_ptr, boundary_check=(0, 1))
v = tl.load(V_block_ptr, boundary_check=(0, 1))
qk = tl.dot(q, k) * sm_scale
col_offsets = start_n + tl.arange(0, BLOCK_N)
row_offsets = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
# Causal Mask 和 序列长度边界 Mask
causal_mask = col_offsets[None, :] > row_offsets[:, None]
ctx_mask = col_offsets[None, :] < N_CTX
qk = tl.where(causal_mask | ~ctx_mask, float("-inf"), qk)
m_ij = tl.maximum(m_i, tl.max(qk, 1))
alpha = tl.exp(m_i - m_ij)
l_ij = tl.exp(qk - m_ij[:, None])
l_i = l_i * alpha + tl.sum(l_ij, 1)
acc = acc * alpha[:, None] + tl.dot(l_ij.to(v.dtype), v)
m_i = m_ij
K_block_ptr = tl.advance(K_block_ptr, (0, BLOCK_N))
V_block_ptr = tl.advance(V_block_ptr, (BLOCK_N, 0))
acc = acc / l_i[:, None]
tl.store(O_block_ptr, acc.to(Out.type.element_ty), boundary_check=(0, 1))
def triton_flash_attention(Q, K, V, causal=True):
BATCH, N_HEADS, N_CTX, D_HEAD = Q.shape
sm_scale = 1.0 / math.sqrt(D_HEAD)
O = torch.empty_like(Q)
BLOCK_M, BLOCK_N = 64, 64
grid = (triton.cdiv(N_CTX, BLOCK_M), BATCH * N_HEADS)
_fwd_kernel[grid](
Q, K, V, sm_scale, O,
Q.stride(0), Q.stride(1), Q.stride(2), Q.stride(3),
K.stride(0), K.stride(1), K.stride(2), K.stride(3),
V.stride(0), V.stride(1), V.stride(2), V.stride(3),
O.stride(0), O.stride(1), O.stride(2), O.stride(3),
BATCH, N_HEADS, N_CTX,
BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_DMODEL=D_HEAD
)
return O
# ================ Transformer =================
class RMSNorm(nn.Module):
def __init__(self, dim: int, eps: float = 1e-6):
super().__init__()
self.eps = eps
self.weight = nn.Parameter(torch.ones(dim))
def forward(self, x):
return x / torch.sqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight
class TransformerBlock(nn.Module):
def __init__(self, D_MODEL: int, N_HEADS: int, D_FF: int):
super().__init__()
self.D_MODEL = D_MODEL
self.N_HEADS = N_HEADS
self.D_HEAD = D_MODEL // N_HEADS
self.attn_norm = RMSNorm(D_MODEL)
self.ffn_norm = RMSNorm(D_MODEL)
self.wq = nn.Linear(D_MODEL, D_MODEL, bias=False)
self.wk = nn.Linear(D_MODEL, D_MODEL, bias=False)
self.wv = nn.Linear(D_MODEL, D_MODEL, bias=False)
self.wo = nn.Linear(D_MODEL, D_MODEL, bias=False)
# 使用 SwiGLU FFN
self.ffn = SwiGLU_FFN(D_MODEL, D_FF)
def forward(self, x, cos, sin):
BATCH, N_CTX, _ = x.shape
norm_x = self.attn_norm(x)
Q = self.wq(norm_x).view(BATCH, N_CTX, self.N_HEADS, self.D_HEAD).transpose(1, 2)
K = self.wk(norm_x).view(BATCH, N_CTX, self.N_HEADS, self.D_HEAD).transpose(1, 2)
V = self.wv(norm_x).view(BATCH, N_CTX, self.N_HEADS, self.D_HEAD).transpose(1, 2)
# 应用 RoPE
Q = apply_rope(Q, cos, sin)
K = apply_rope(K, cos, sin)
# 调用 Triton Flash Attention
attn_output = triton_flash_attention(Q, K, V, causal=True)
attn_output = attn_output.transpose(1, 2).contiguous().view(BATCH, N_CTX, self.D_MODEL)
x = x + self.wo(attn_output)
# 使用 SwiGLU FFN
x = x + self.ffn(self.ffn_norm(x))
return x
class DecoderOnlyTransformer(nn.Module):
def __init__(self, VOCAB_SIZE: int, D_MODEL: int, N_HEADS: int, N_LAYERS: int, D_FF: int, MAX_SEQ_LEN: int):
super().__init__()
# 使用 Embedding
self.tok_embeddings = Embedding(VOCAB_SIZE, D_MODEL)
self.layers = nn.ModuleList([TransformerBlock(D_MODEL, N_HEADS, D_FF) for _ in range(N_LAYERS)])
self.norm = RMSNorm(D_MODEL)
self.output = nn.Linear(D_MODEL, VOCAB_SIZE, bias=False)
# 预计算 RoPE (修复维度拼接问题)
D_HEAD = D_MODEL // N_HEADS
freqs = 1.0 / (10000.0 ** (torch.arange(0, D_HEAD, 2).float() / D_HEAD))
t = torch.arange(MAX_SEQ_LEN, dtype=torch.float32)
freqs = torch.outer(t, freqs) # (MAX_SEQ_LEN, D_HEAD/2)
# 将 freqs 在最后一维拼接，使其大小为 D_HEAD，从而与 x 的最后一维匹配
freqs_cis = torch.cat([freqs, freqs], dim=-1) # (MAX_SEQ_LEN, D_HEAD)
self.register_buffer('cos_cached', torch.cos(freqs_cis).view(MAX_SEQ_LEN, 1, 1, D_HEAD))
self.register_buffer('sin_cached', torch.sin(freqs_cis).view(MAX_SEQ_LEN, 1, 1, D_HEAD))
def forward(self, token_ids):
BATCH, N_CTX = token_ids.shape
h = self.tok_embeddings(token_ids)
cos = self.cos_cached[:N_CTX].to(h.device)
sin = self.sin_cached[:N_CTX].to(h.device)
for layer in self.layers:
h = layer(h, cos, sin)
h = self.norm(h)
return self.output(h)
# ================= test =================
if __name__ == "__main__":
# 初始化模型、优化器
model = DecoderOnlyTransformer(VOCAB_SIZE=1000, D_MODEL=256, N_HEADS=4, N_LAYERS=4, D_FF=512, MAX_SEQ_LEN=128).to(
'cuda').half()
optimizer = AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)
# 模拟 Batch 的数据
input_ids = torch.randint(0, 1000, (2, 64), device='cuda')
targets = torch.randint(0, 1000, (2, 64), device='cuda')
# 前向传播
logits = model(input_ids)
# 计算 Loss
loss = cross_entropy_loss(logits, targets)
print(f"Initial Loss: {loss.item():.4f}")
# 反向传播与 AdamW 更新
optimizer.zero_grad()
loss.backward()
optimizer.step()
print("Forward & Backward pass with custom components completed successfully!")
浙公网安备 33010602011771号