【PYTORCH】迷你DECODER-ONLY模型
前置知识
torch.triu
设置下三角为0，从第一行开始，其中第一行从左到有有diagonal个0
mask = torch.triu(torch.ones(5, 5),diagonal=0)
print(mask)
shell
tensor([[1., 1., 1., 1., 1.],
[0., 1., 1., 1., 1.],
[0., 0., 1., 1., 1.],
[0., 0., 0., 1., 1.],
[0., 0., 0., 0., 1.]])mask = torch.triu(torch.ones(5, 5),diagonal=1)
print(mask)
shell
tensor([[0., 1., 1., 1., 1.],
[0., 0., 1., 1., 1.],
[0., 0., 0., 1., 1.],
[0., 0., 0., 0., 1.],
[0., 0., 0., 0., 0.]])masked_fill
将矩阵里面等于mask的值设值为第二个参数
将矩阵值为0的设置为-inf
mask = torch.triu(torch.ones(5, 5),diagonal=0)
mask.masked_fill(mask == 0, float("-inf"))
print(mask)
shell
tensor([[1., 1., 1., 1., 1.],
[-inf, 1., 1., 1., 1.],
[-inf, -inf, 1., 1., 1.],
[-inf, -inf, -inf, 1., 1.],
[-inf, -inf, -inf, -inf, 1.]])将矩阵值为1的设置为-inf
mask = torch.triu(torch.ones(5, 5),diagonal=0)
mask.masked_fill(mask == 1, float("-inf"))
print(mask)
shell
tensor([[-inf, -inf, -inf, -inf, -inf],
[0., -inf, -inf, -inf, -inf],
[0., 0., -inf, -inf, -inf],
[0., 0., 0., -inf, -inf],
[0., 0., 0., 0., -inf]])register_buffer
self.register_buffer("mask", mask)可以理解为self.mask = mask，但好处是
model.to(device)→ buffer 自动跟着移到对应设备torch.save(model.state_dict())会把 mask 存进去，load 的时候自动恢复
完整代码
python
import torch
import torch.nn as nn
import torch.nn.functional as F
# ========== 超参 ==========
vocab_size = 100
hidden_dim = 64
n_heads = 2
n_layers = 2
max_seq_len = 35
head_dim = hidden_dim // n_heads
device = "cuda" if torch.cuda.is_available() else "cpu"
# ========== 对话特殊token定义 ==========
IM_START = 90
IM_END = 91
PAD = 0
# 简易词表映射（演示用，真实场景用tokenizer）
token2text = {
0: "[PAD]",
1: "你", 2: "好", 3: "是", 4: "谁",
5: "我", 6: "豆", 7: "包",
14: "谢", # 新加
16: "不", # 新加
17: "客", # 新加
18: "气", # 新加
IM_START: "<|im_start|>",
IM_END: "<|im_end|>"
}
text2token = {v:k for k,v in token2text.items() if isinstance(k,int)}
# ========== 带KV Cache支持的因果自注意力 ==========
class CausalSelfAttention(nn.Module):
def __init__(self):
super().__init__()
self.q_proj = nn.Linear(hidden_dim, hidden_dim)
self.k_proj = nn.Linear(hidden_dim, hidden_dim)
self.v_proj = nn.Linear(hidden_dim, hidden_dim)
self.out_proj = nn.Linear(hidden_dim, hidden_dim)
# 因果掩码：上三角mask，看不到未来token
mask = torch.triu(torch.ones(max_seq_len, max_seq_len), diagonal=1)
mask = mask.masked_fill(mask == 1, float("-inf"))
self.register_buffer("mask", mask)
def forward(self, x, past_kv=None):
"""
x: [B, T, D]
past_kv: tuple(K_cache, V_cache) 推理时传入，训练为None
return: out, new_kv
"""
B, T, D = x.shape
q = self.q_proj(x) # [B,T,D]
k = self.k_proj(x)
v = self.v_proj(x)
# 多头切分
q = q.view(B, T, n_heads, head_dim).transpose(1,2) # [B, nh, T, hd]
k = k.view(B, T, n_heads, head_dim).transpose(1,2)
v = v.view(B, T, n_heads, head_dim).transpose(1,2)
# ===== KV Cache 拼接逻辑 =====
if past_kv is not None:
past_k, past_v = past_kv
k = torch.cat([past_k, k], dim=2)
v = torch.cat([past_v, v], dim=2)
new_kv = (k, v)
kv_len = k.shape[2] # 历史总长度
# QK点积注意力
attn_score = q @ k.transpose(-1,-2) / (head_dim ** 0.5) # [B,nh,T,kv_len]
# 取对应长度的mask
attn_score = attn_score + self.mask[:T, :kv_len]
attn_weight = F.softmax(attn_score, dim=-1)
out = attn_weight @ v # [B,nh,T,hd]
# 多头合并
out = out.transpose(1,2).contiguous().view(B, T, D)
out = self.out_proj(out)
return out, new_kv
# ========== Transformer Block（残差+LN+FFN） ==========
class TransformerBlock(nn.Module):
def __init__(self):
super().__init__()
self.attn = CausalSelfAttention()
self.ln1 = nn.LayerNorm(hidden_dim)
self.ln2 = nn.LayerNorm(hidden_dim)
self.ffn = nn.Sequential(
nn.Linear(hidden_dim, hidden_dim*4),
nn.GELU(),
nn.Linear(hidden_dim*4, hidden_dim)
)
def forward(self, x, past_kv_list=None):
new_kv_list = []
if past_kv_list is None:
past_kv_list = [None]*n_layers
# 注意力 + 残差
attn_out, new_kv = self.attn(self.ln1(x), past_kv_list[0])
x = x + attn_out
new_kv_list.append(new_kv)
# FFN + 残差
x = x + self.ffn(self.ln2(x))
return x, new_kv_list
# ========== Decoder-only 大模型主体 ==========
class MiniChatModel(nn.Module):
def __init__(self):
super().__init__()
self.embedding = nn.Embedding(vocab_size, hidden_dim)
self.blocks = nn.ModuleList([TransformerBlock() for _ in range(n_layers)])
self.ln_f = nn.LayerNorm(hidden_dim)
self.lm_head = nn.Linear(hidden_dim, vocab_size)
def forward(self, tokens, past_kv_list=None):
"""
tokens: [B, T]
past_kv_list: 每层的(K,V)，推理增量解码用
return: logits, new_kv_list
"""
B, T = tokens.shape
x = self.embedding(tokens) # [B,T,D]
new_kv_list = []
if past_kv_list is None:
past_kv_list = [None]*n_layers
for i, block in enumerate(self.blocks):
x, kv = block(x, [past_kv_list[i]])
new_kv_list.append(kv[0])
x = self.ln_f(x)
logits = self.lm_head(x) # [B,T,vocab_size]
return logits, new_kv_list
# ========== SFT Loss：只计算assistant部分loss ==========
def sft_cross_entropy(logits, tokens, loss_mask):
"""
logits: [B, T, V]
tokens: [B, T] label
loss_mask: [B, T] bool, True=要算loss(assistant), False=忽略(user/prompt)
"""
B, T, V = logits.shape
logits_flat = logits[:,:-1,:].reshape(-1, V)
labels_flat = tokens[:,1:].reshape(-1)
mask_flat = loss_mask[:,1:].reshape(-1)
ce = F.cross_entropy(logits_flat, labels_flat, reduction="none")
ce = ce * mask_flat
return ce.sum() / (mask_flat.sum() + 1e-8)
# ========== 对话生成函数（KV Cache增量解码） ==========
@torch.no_grad()
def generate(model, prompt_tokens, max_new_tokens=10):
model.eval()
tokens = torch.tensor(prompt_tokens, dtype=torch.long, device=device).unsqueeze(0)
past_kv = None
for _ in range(max_new_tokens):
logits, past_kv = model(tokens, past_kv)
next_logit = logits[:,-1,:]
next_id = torch.argmax(next_logit, dim=-1, keepdim=True) # greedy
if next_id.item() == IM_END:
break
tokens = torch.cat([tokens, next_id], dim=-1)
return tokens[0].tolist()
# ========== 【演示：一条对话样本训练】 ==========
if __name__ == "__main__":
model = MiniChatModel().to(device)
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
# 构造一条对话样本
# <|im_start|>user 你是谁 <|im_end|> <|im_start|>assistant 我是豆包 <|im_end|>
dialog_tokens = [
IM_START, # <|im_start|>
14, 14, # 谢 谢
IM_END, # <|im_end|>
IM_START, # <|im_start|>
16, 17, 18, # 不 客 气
IM_END, # <|im_end|>
IM_START,
1,3,4,
IM_END,
IM_START,
5,3,6,7,
IM_END
]
# loss_mask：False=user部分不算loss；True=assistant部分算loss
loss_mask = [False, False, False, False, True, True, True, True, True,False, False, False, False,False,True, True, True, True, True,True]
tokens = torch.tensor(dialog_tokens, dtype=torch.long, device=device).unsqueeze(0)
mask_tensor = torch.tensor(loss_mask, dtype=torch.float, device=device).unsqueeze(0)
print("==== 开始SFT训练一轮 ====")
for step in range(30):
logits, _ = model(tokens)
loss = sft_cross_entropy(logits, tokens, mask_tensor)
opt.zero_grad()
loss.backward()
opt.step()
if step % 5 == 0:
print(f"step {step}, loss: {loss.item():.4f}")
print("\n==== 对话推理测试 ====")
# prompt：<|im_start|>user 你是谁 <|im_end|> <|im_start|>assistant
prompt = [IM_START, 14,14, IM_END, IM_START]
out_tokens = generate(model, prompt, max_new_tokens=8)
print("输出token序列：", out_tokens)