AI手撕代码笔记
Attn相关
基础attn
细节：
- 注意K的转秩
- 注意mask写法是置为-1e9(masked_fill需要转bool取反)
- dropout在softmax之后
- softmax只做最后一个维度
python
import torch
import torch.nn as nn
from typing import Optional, Tuple
import torch.nn.functional as F
import math
class ScaledDotProductAttention(nn.Module):
def __init__(self, dropout_p: float = 0.0):
super().__init__()
self.dropout = nn.Dropout(dropout_p)
def forward(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
mask: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
# q/k/v: [B, H, S, D]；返回 (output, attn_weights)
# print(q, k, v, mask)
d = q.shape[-1]
scores = q @ k.transpose(-1, -2) / math.sqrt(d)
if mask is not None:
scores = scores.masked_fill_(~mask.bool(), -1e9)
attn_weights = F.softmax(scores, dim=-1)
attn_weights = self.dropout(attn_weights)
output = attn_weights @ v
return (output, attn_weights)
causal mask
常规版本(多头注意力)
python
def create_batch_causal_mask(batch_size: int, n_head: int, seq_len: int, device=None):
"""输出 [B, H, L, L]"""
causal_2d = torch.tril(torch.ones(seq_len, seq_len, dtype=torch.bool, device=device))
mask = causal_2d[None, None, :, :].expand(batch_size, n_head, seq_len, seq_len)
return mask对于decoder部分，如果以及有了一部分生成，带滑窗便宜的casual mask：
python
q_idx = torch.arange(seq_len, device=x.device).unsqueeze(1) # [seq_len, 1]
k_idx = torch.arange(total_len, device=x.device).unsqueeze(0) # [1, total_len]
causal_mask = k_idx > (start_pos + q_idx) # [seq_len, total_len]多头注意力
细节：
- 分头：先reshape然后permute
- 记得要转回来
python
class MultiHeadAttention(nn.Module):
def __init__(self, embed_dim: int, num_heads: int):
super().__init__()
assert embed_dim % num_heads == 0, "embed_dim必须可以被num_heads整除"
self.embed_dim = embed_dim
self.num_heads = num_heads
self.head_dim = embed_dim // num_heads
# Q K V 投影
self.w_q = nn.Linear(embed_dim, embed_dim)
self.w_k = nn.Linear(embed_dim, embed_dim)
self.w_v = nn.Linear(embed_dim, embed_dim)
# 输出投影
self.w_o = nn.Linear(embed_dim, embed_dim)
def forward(self, x, mask=None):
"""
Args:
x: [B, L, C] 输入
mask: [B, 1, Lq, Lk] bool, True=允许访问, False=mask掉
Returns:
out: [B, L, C]
"""
B, Lq, _ = x.shape
# 1. 投影
q = self.w_q(x) # [B, Lq, C]
k = self.w_k(x) # [B, Lk, C]
v = self.w_v(x)
Lk = k.size(1)
# 2. split heads: [B, H, L, head_dim]
def split_head(t):
B, L, C = t.shape
return t.reshape(B, L, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
q = split_head(q) # [B, H, Lq, hd]
k = split_head(k) # [B, H, Lk, hd]
v = split_head(v) # [B, H, Lk, hd]
# 3. scaled dot‑product
scale = self.head_dim ** (-0.5)
attn_score = torch.matmul(q, k.transpose(-1, -2)) * scale # [B, H, Lq, Lk]
# apply mask: False的位置填‑inf，softmax后权重为0
if mask is not None:
attn_score = attn_score.masked_fill(~mask, float("-inf"))
attn_weight = F.softmax(attn_score, dim=-1) # [B, H, Lq, Lk]
out = torch.matmul(attn_weight, v) # [B, H, Lq, hd]
# concat heads
out = out.permute(0, 2, 1, 3).reshape(B, Lq, self.embed_dim) # [B, Lq, C]
out = self.w_o(out)
return out, attn_weightcross attn
注意：Q和K/V来源不一致，其他与self-attn一样
辅助函数
Softmax
python
def softmax(x, dim=-1):
e_x = np.exp(x - np.max(x, axis=-1, keepdims=True))
return e_x / np.sum(e_x, axis=-1, keepdims=True)RL
advantage计算
python
import torch
def compute_mc_advantage(rewards, dones, values, gamma=0.99):
T = len(rewards)
returns = torch.zeros_like(rewards)
running_ret = 0.0
for t in reversed(range(T)):
if dones[t]:
running_ret = 0.0
running_ret = rewards[t] + gamma * running_ret
returns[t] = running_ret
adv = returns - values
adv = (adv - adv.mean())/(adv.std()+1e-8)
return adv, returnsGAE
python
import torch
def compute_gae(rewards: torch.Tensor,
dones: torch.Tensor,
values: torch.Tensor,
gamma: float = 0.99,
lam: float = 0.95):
"""
Args:
rewards: [T] 单段轨迹奖励
dones: [T] bool/float，1代表episode结束
values: [T] critic预测的state value
gamma: 折扣因子
lam: GAE λ参数
Returns:
advantages: [T]
returns: [T] = advantages + values
"""
T = len(rewards)
advantages = torch.zeros_like(rewards)
last_advantage = 0.0
for t in reversed(range(T)):
# 下一个时刻value，最后一步t=T‑1没有next state → 0
next_val = values[t+1].item() if (t + 1 < T) else 0.0
# TD‑error δ_t
delta = rewards[t] + gamma * next_val * (1.0 - dones[t]) - values[t].item()
# GAE递推
advantages[t] = delta + gamma * lam * (1.0 - dones[t]) * last_advantage
last_advantage = advantages[t].item()
returns = advantages + values
# 可选：advantage标准化，ppo训练必用
advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
return advantages, returns
python
import torch
def compute_group_advantages(rewards):
# GRPO组内优势，无critic，用组均值做基线
group_mean = rewards.mean()
advantages = rewards - group_mean
advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
return advantages
def grpo_calc_loss(old_log_probs, new_log_probs, advantages, ref_log_probs, clip_epsilon=0.2, kl_beta=0.04):
ratio = torch.exp(new_log_probs - old_log_probs)
surr1 = ratio * advantages
surr2 = torch.clamp(ratio, 1-clip_epsilon, 1+clip_epsilon) * advantages
policy_loss = -torch.min(surr1, surr2).mean()
# KL惩罚 ref || current
kl = torch.exp(new_log_probs) * (new_log_probs - ref_log_probs)
kl_div = kl.mean()
total_loss = policy_loss + kl_beta * kl_div
return total_loss
def main():
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 超参
group_size = 4
clip_epsilon = 0.2
kl_beta = 0.04
update_epochs = 2
total_iters = 50
for iter_idx in range(total_iters):
# -------- rollout 采样得到这批数据（不展开模型、采样细节） --------
old_log_probs = torch.randn(group_size, device=device)
ref_log_probs = torch.randn(group_size, device=device)
rewards = torch.randn(group_size, device=device)
# -------- 计算advantage --------
advantages = compute_group_advantages(rewards)
# -------- PPO‑clip 更新循环 --------
for _ in range(update_epochs):
new_log_probs = torch.randn(group_size, device=device) # 模型前向，不展开
# -------- 计算loss --------
loss = grpo_calc_loss(old_log_probs, new_log_probs, advantages, ref_log_probs, clip_epsilon, kl_beta)
# -------- 优化器步骤（不展开模型与optimizer定义） --------
loss.backward()
# optimizer.step()
# optimizer.zero_grad()
if iter_idx %10 ==0:
print(f"iter {iter_idx}, loss: {loss.item():.3f}")
if __name__ == "__main__":
main()