Transformer 与 PyTorch/CUDA 完全指南

Transformer 架构

Attention Is All You Need

整体架构

Encoder (编码器)

输入嵌入 + 位置编码

Input Embedding + Positional Encoding

Multi-Head Self-Attention

Add & Norm

Feed Forward Network

Add & Norm

× N 层堆叠

Decoder (解码器)

输出嵌入 + 位置编码

Output Embedding + Positional Encoding

Masked Multi-Head Self-Attention

Add & Norm

Multi-Head Cross-Attention

Add & Norm

Feed Forward Network

Add & Norm

× N 层堆叠

Self-Attention

计算序列中每个位置与其他所有位置的注意力权重，捕捉全局依赖关系。

Multi-Head Attention

并行使用多组 Q/K/V 投影，让模型在不同子空间中关注不同信息。

Feed Forward

每个位置独立应用的两层全连接网络，引入非线性变换能力。

Layer Normalization

对每个样本的所有特征进行归一化，稳定深层网络的训练。

Positional Encoding

为模型注入序列位置信息，因为 Attention 本身不具备位置感知能力。

Residual Connection

跳跃连接缓解梯度消失问题，让网络可以训练到数百层深度。

Attention 机制详解

理解现代 AI 的核心运算

Scaled Dot-Product Attention

Attention(Q, K, V) = softmax(QK^T / √d_k)V

QQuery 查询矩阵

KKey 键矩阵

VValue 值矩阵

d_kKey 的维度，用于缩放

步骤 1 - 计算相似度：通过矩阵乘法 QK^T 计算查询与所有键的相似度分数。

步骤 2 - 缩放：除以 √d_k 防止点积过大导致 softmax 梯度消失。

步骤 3 - Softmax：将分数转换为概率分布，每行之和为 1。

步骤 4 - 加权求和：用注意力权重对 Value 进行加权求和，得到输出。

attention.py PyTorch 实现

import torch
import torch.nn as nn
import math

class ScaledDotProductAttention(nn.Module):
    def __init__(self, dropout=0.1):
        super().__init__()
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, Q, K, V, mask=None):
        # Q, K, V: (batch, seq_len, d_k)
        d_k = Q.size(-1)
        
        # 1. 计算注意力分数
        scores = torch.matmul(Q, K.transpose(-2, -1)) \
                 / math.sqrt(d_k)
        # scores: (batch, seq_len, seq_len)
        
        # 2. 应用 Mask（可选）
        if mask is not None:
            scores = scores.masked_fill(
                mask == 0, -1e9
            )
        
        # 3. Softmax 归一化
        attn_weights = torch.softmax(scores, dim=-1)
        attn_weights = self.dropout(attn_weights)
        
        # 4. 加权求和
        output = torch.matmul(attn_weights, V)
        
        return output, attn_weights

Multi-Head Attention

将 Q、K、V 分别投影到 h 个不同的子空间中并行计算 Attention，最后将结果拼接并再次线性变换。

MultiHead(Q, K, V) = Concat(head₁, ..., head_h)W^O

where head_i = Attention(QW^Q_i, KW^K_i, VW^V_i)

• 每个 head 的维度 d_k = d_model / h

• 标准 Transformer 中 h = 8, d_model = 512

• 不同 head 可以学习不同的关注模式（语法、语义、指代等）

multi_head_attention.py

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model=512, num_heads=8):
        super().__init__()
        assert d_model % num_heads == 0
        
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads  # 64
        
        # 线性投影层
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)
        
        self.attention = ScaledDotProductAttention()
    
    def forward(self, Q, K, V, mask=None):
        batch_size = Q.size(0)
        
        # 1. 线性投影并分头
        # (batch, seq, d_model) -> 
        # (batch, heads, seq, d_k)
        Q = self.W_q(Q).view(
            batch_size, -1, self.num_heads, self.d_k
        ).transpose(1, 2)
        K = self.W_k(K).view(
            batch_size, -1, self.num_heads, self.d_k
        ).transpose(1, 2)
        V = self.W_v(V).view(
            batch_size, -1, self.num_heads, self.d_k
        ).transpose(1, 2)
        
        # 2. 计算 Attention
        out, attn = self.attention(Q, K, V, mask)
        
        # 3. 拼接多头并线性变换
        out = out.transpose(1, 2).contiguous() \
                   .view(batch_size, -1, self.d_model)
        
        return self.W_o(out)

Masked Attention & Cross Attention

Masked Self-Attention (Decoder)

防止解码器在生成第 i 个词时看到后续位置的信息。通过上三角掩码实现：

mask = torch.tril(torch.ones(seq_len, seq_len))
# 下三角矩阵，对角线及以下为 1

Cross Attention (Encoder-Decoder)

Decoder 使用自身的 Q，但使用 Encoder 输出的 K 和 V，实现源语言到目标语言的注意力对齐。

Q = decoder_hidden # 来自解码器
K, V = encoder_output # 来自编码器

Feed Forward Network (FFN)

FFN(x) = max(0, xW₁ + b₁)W₂ + b₂

ReLU 激活的两层全连接网络

• 中间维度 d_ff = 2048（通常是 d_model 的 4 倍）

• 对每个位置独立应用，不共享信息

• 引入非线性，增强模型表达能力

• 现代变体常用 GELU 替代 ReLU

Layer Normalization

LN(x) = γ ⊙ (x - μ) / √(σ² + ε) + β

• 与 BatchNorm 的区别：LN 对每个样本的特征维度归一化，BN 对 batch 维度归一化

• 在 NLP 中更稳定，因为序列长度变化大

• γ 和 β 是可学习的缩放和平移参数

• Pre-LN 结构（LayerNorm 在残差之前）训练更稳定

Residual Connection (残差连接)

output = LayerNorm(x + Sublayer(x))

原始 Transformer 使用 Post-LN（残差连接后接 LayerNorm）。现代模型（如 GPT-3、LLaMA）普遍改用 Pre-LN，将 LayerNorm 放在子层之前，训练更加稳定：

output = x + Sublayer(LayerNorm(x)) # Pre-LN

class TransformerBlock(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        super().__init__()
        self.attention = MultiHeadAttention(d_model, num_heads)
        self.ffn = FeedForward(d_model, d_ff)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, x, mask=None):
        # Pre-LN 结构
        # Attention 子层
        attn_out = self.attention(
            self.norm1(x), 
            self.norm1(x), 
            self.norm1(x), 
            mask
        )
        x = x + self.dropout(attn_out)
        
        # FFN 子层
        ffn_out = self.ffn(self.norm2(x))
        x = x + self.dropout(ffn_out)
        
        return x

PyTorch 核心概念

动态图深度学习框架

T

Tensor (张量)

PyTorch 的核心数据结构，类似于 NumPy 的 ndarray，但可以在 GPU 上运行并支持自动求导。

import torch

# 创建 Tensor
x = torch.tensor([[1, 2], [3, 4]])
x = torch.zeros(2, 3)
x = torch.randn(3, 3)  # 标准正态分布

# GPU 张量
x = x.cuda()           # 移到 GPU
x = x.to('cuda:0')     # 指定 GPU

# 关键属性
x.shape    # torch.Size([2, 2])
x.dtype    # torch.float32
x.device   # cuda:0

∇

Autograd (自动求导)

PyTorch 的自动微分引擎，通过计算图自动追踪所有操作，支持反向传播计算梯度。

x = torch.tensor(2.0, requires_grad=True)
y = torch.tensor(3.0, requires_grad=True)

z = x**2 + y**3      # z = 4 + 27 = 31
z.backward()         # 反向传播

print(x.grad)        # 4.0  (∂z/∂x = 2x)
print(y.grad)        # 27.0 (∂z/∂y = 3y²)

# 清零梯度（重要！）
optimizer.zero_grad()

M

nn.Module

所有神经网络模块的基类。通过继承 nn.Module 定义模型，PyTorch 会自动管理参数和子模块。

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 32, 3)
        self.fc = nn.Linear(32*26*26, 10)
    
    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = x.view(x.size(0), -1)
        return self.fc(x)

model = Net()
model = model.to('cuda')

# 查看参数
for name, param in model.named_parameters():
    print(name, param.shape)

O

Optimizer & Loss

优化器根据梯度更新参数，损失函数衡量模型预测与真实值的差距。

# 损失函数
criterion = nn.CrossEntropyLoss()
loss = criterion(outputs, labels)

# 优化器
optimizer = torch.optim.Adam(
    model.parameters(), 
    lr=1e-3,
    betas=(0.9, 0.999)
)

# 训练步骤（标准模板）
optimizer.zero_grad()   # 1. 清零梯度
loss.backward()         # 2. 反向传播
optimizer.step()        # 3. 更新参数

# 学习率调度
scheduler = torch.optim.lr_scheduler.\
    StepLR(optimizer, step_size=10, gamma=0.1)

标准训练流程模板

import torch
import torch.nn as nn
from torch.utils.data import DataLoader

def train_epoch(model, dataloader, criterion, optimizer, device):
    model.train()
    total_loss = 0
    
    for batch_idx, (data, target) in enumerate(dataloader):
        # 1. 数据移动到 GPU
        data, target = data.to(device), target.to(device)
        
        # 2. 前向传播
        output = model(data)
        loss = criterion(output, target)
        
        # 3. 反向传播
        optimizer.zero_grad()
        loss.backward()
        
        # 4. 梯度裁剪（防止梯度爆炸）
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        
        # 5. 更新参数
        optimizer.step()
        
        total_loss += loss.item()
    
    return total_loss / len(dataloader)

# 主训练循环
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = MyModel().to(device)

for epoch in range(num_epochs):
    train_loss = train_epoch(model, train_loader, criterion, optimizer, device)
    val_loss = evaluate(model, val_loader, criterion, device)
    print(f"Epoch {epoch}: Train Loss={train_loss:.4f}, Val Loss={val_loss:.4f}")
    
    # 保存最佳模型
    if val_loss < best_loss:
        torch.save(model.state_dict(), 'best_model.pt')
        best_loss = val_loss

CUDA 与 GPU 编程

释放 GPU 并行计算的力量

什么是 CUDA？

🎯

并行计算平台

NVIDIA 推出的并行计算平台和编程模型，让开发者能使用 GPU 进行通用计算（GPGPU）。

⚡

海量并行核心

现代 GPU 拥有数千个 CUDA Core（如 RTX 4090 有 16384 个），可同时执行大量线程。

🔧

编程接口

基于 C/C++ 的扩展，通过 __global__、__device__ 等关键字定义 GPU 核函数。

GPU 硬件架构基础

SM

Streaming Multiprocessor (SM)

GPU 的核心计算单元，每个 SM 包含多个 CUDA Core、Tensor Core、共享内存和寄存器。GPU 通常有 80-144 个 SM。

TC

Tensor Core

专用矩阵运算单元，专门加速混合精度（FP16/BF16/TF32）的矩阵乘加运算（GEMM），是 Transformer 训练加速的关键。

Mem

内存层次结构

从快到慢：寄存器 (Register) → 共享内存 (Shared Memory) → L1/L2 缓存 → 全局内存 (Global Memory/显存 HBM)。

内存带宽与计算能力对比

H100 (HBM3) 3.35 TB/s

A100 (HBM2e) 2.0 TB/s

RTX 4090 (GDDR6X) 1.0 TB/s

CPU DDR5 ~50 GB/s

注：H100 FP16 Tensor Core 算力达 989 TFLOPS，是 CPU 的数百倍。

PyTorch 中的 CUDA 使用

cuda_basics.py

import torch

# 检查 CUDA 是否可用
print(torch.cuda.is_available())   # True
print(torch.cuda.device_count())   # GPU 数量
print(torch.cuda.get_device_name(0))  # NVIDIA RTX 4090

# 创建 GPU 张量
x = torch.randn(3, 3).cuda()
x = torch.randn(3, 3, device='cuda:0')

# 将模型移到 GPU
model = MyModel().to('cuda')

# 指定数据类型（节省显存）
x = x.half()           # FP16
x = x.bfloat16()       # BF16 (Ampere+)

# 混合精度训练
from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler()

with autocast():
    output = model(input)
    loss = criterion(output, target)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

关键概念

to() 将张量或模型移动到指定设备（CPU/GPU），注意：这是浅拷贝操作

non_blocking 异步数据传输：tensor.to('cuda', non_blocking=True)

pin_memory 锁页内存加速 CPU→GPU 传输：DataLoader(..., pin_memory=True)

torch.no_grad() 推理时禁用梯度计算，减少显存占用并加速

显存管理技巧

• torch.cuda.empty_cache() - 释放未使用的缓存显存

• del tensor + torch.cuda.empty_cache() - 主动释放

• 使用 gradient checkpointing 时间换空间

• nvidia-smi 实时监控显存使用

CUDA C++ 编程入门

理解底层 CUDA 编程有助于优化 PyTorch 性能。以下是向量加法的 CUDA 核函数示例：

// CUDA 核函数：每个线程处理一个元素
__global__ void vectorAdd(
    const float *A, 
    const float *B, 
    float *C, 
    int numElements
) {
    // 计算全局线程 ID
    int i = blockDim.x * blockIdx.x + threadIdx.x;
    
    if (i < numElements) {
        C[i] = A[i] + B[i];
    }
}

// 主机代码调用
int numElements = 50000;
int threadsPerBlock = 256;
int blocksPerGrid = (numElements + threadsPerBlock - 1) 
                    / threadsPerBlock;

vectorAdd<<>>(
    d_A, d_B, d_C, numElements
);

CUDA 执行模型

Grid（网格）由多个 Block 组成

Block（块）由多个 Thread 组成

Thread（线程）最小执行单元

内存类型

__global__ 全局内存（显存，最大最慢）

__shared__ 共享内存（SM 内，块级可见，快）

__constant__ 常量内存（只读，缓存）

__device__ 设备函数，仅在 GPU 调用

核心要点总结

Transformer

Self-Attention 实现全局依赖建模，Multi-Head 增强表达能力，FFN 提供非线性，残差连接和 LayerNorm 保障深层训练稳定。

PyTorch

Tensor 是基础数据结构，Autograd 自动计算梯度，nn.Module 构建网络，配合 Optimizer 和 DataLoader 完成完整训练流程。

CUDA

GPU 通过海量并行核心加速矩阵运算，Tensor Core 专为 AI 优化。理解内存层次和异步传输是性能优化的关键。

Transformer 与 PyTorch / CUDA

Transformer 架构

整体架构

Self-Attention

Multi-Head Attention

Feed Forward

Layer Normalization

Positional Encoding

Residual Connection

Attention 机制详解

Scaled Dot-Product Attention

Multi-Head Attention

Masked Attention & Cross Attention

Masked Self-Attention (Decoder)

Cross Attention (Encoder-Decoder)

Feed Forward Network (FFN)

Layer Normalization

Residual Connection (残差连接)

PyTorch 核心概念

Tensor (张量)

Autograd (自动求导)

nn.Module

Optimizer & Loss

标准训练流程模板

CUDA 与 GPU 编程

什么是 CUDA？

并行计算平台

海量并行核心

编程接口

GPU 硬件架构基础

Streaming Multiprocessor (SM)

Tensor Core

内存层次结构

PyTorch 中的 CUDA 使用

关键概念

显存管理技巧

CUDA C++ 编程入门

CUDA 执行模型

内存类型

核心要点总结

Transformer

PyTorch

CUDA

Transformer
与 PyTorch / CUDA