ChatGPT中的Transformer架构实战：从原理到高效实现

1次阅读

共计 2172 个字符，预计需要花费 6 分钟才能阅读完成。

Transformer 架构在处理长序列时面临三个主要挑战：

计算复杂度高：自注意力机制的计算复杂度与序列长度的平方成正比，当处理长文本时（如超过 2048 个 token），显存占用和计算时间会急剧增加。
内存瓶颈：KV（Key-Value）缓存随着序列长度线性增长，在批量推理时容易触发 OOM（内存不足）错误。
位置信息丢失：原始 Transformer 的位置编码在长序列场景下可能无法有效捕捉远距离依赖关系。

多头注意力机制
原始：固定维度的多头投影
ChatGPT：采用分组查询注意力(GQA)，减少 KV 头的数量
位置编码
原始：绝对位置编码
ChatGPT：旋转位置编码(RoPE)，更好地建模相对位置关系
归一化层
原始：后置层归一化
ChatGPT：前置层归一化（Pre-LN），训练更稳定

import torch
import torch.nn as nn
import math

class EfficientAttention(nn.Module):
    def __init__(self, embed_dim, num_heads, group_size=4):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.group_size = min(group_size, num_heads)  # GQA 分组大小

        # 投影矩阵初始化
        self.q_proj = nn.Linear(embed_dim, embed_dim)
        self.kv_proj = nn.Linear(embed_dim, 2 * self.head_dim * self.group_size)
        self.out_proj = nn.Linear(embed_dim, embed_dim)

    def forward(self, x, attention_mask=None):
        batch_size, seq_len, _ = x.shape

        # 查询向量投影
        q = self.q_proj(x)
        q = q.view(batch_size, seq_len, self.num_heads, self.head_dim)

        # 键值向量分组投影
        kv = self.kv_proj(x)
        kv = kv.view(batch_size, seq_len, 2, self.group_size, self.head_dim)
        k, v = kv.unbind(2)  # [B, L, G, D]

        # 注意力得分计算
        q = q.transpose(1, 2)  # [B, H, L, D]
        k = k.transpose(1, 2).transpose(2, 3)  # [B, G, D, L]
        attn_weights = torch.matmul(q, k) / math.sqrt(self.head_dim)

        if attention_mask is not None:
            attn_weights += attention_mask

        attn_probs = torch.softmax(attn_weights, dim=-1)

        # 价值向量加权
        v = v.transpose(1, 2)  # [B, G, L, D]
        output = torch.matmul(attn_probs, v)
        output = output.transpose(1, 2).contiguous()
        output = output.view(batch_size, seq_len, -1)

        return self.out_proj(output)

def apply_rotary_pos_emb(q, k, sin, cos):
    """应用旋转位置编码到查询和键向量"""
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed

def rotate_half(x):
    """将输入张量的后半部分旋转 180 度"""
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat((-x2, x1), dim=-1)

梯度裁剪：在反向传播前添加torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
学习率预热：使用线性或余弦预热调度器，前 5% 的训练步骤逐步提高学习率
权重初始化 ：对注意力层的投影矩阵使用nn.init.xavier_uniform_() 初始化