第9章：AI 加速技术栈

传统 CUDA Core 矩阵乘法：
┌────────┐     ┌────────┐
│ Thread │ --> │ 1 FMA  │ --> 单个元素
└────────┘     └────────┘

Tensor Core 矩阵乘法：
┌────────┐     ┌─────────────┐
│ Warp   │ --> │ 4x4x4 MMA   │ --> 整个子矩阵
└────────┘     └─────────────┘

4x4x4 矩阵乘法累加操作：
D = A × B + C

      [4x4]     [4x4]     [4x4]     [4x4]
        A    ×    B    +    C    =    D
    (FP16)   (FP16)   (FP32)   (FP32)

每个Tensor Core每周期完成：

- 64次乘法（4×4×4）
- 64次加法
- 总计128个浮点运算

// Warp Matrix Multiply-Accumulate
wmma::fragment<wmma::matrix_a, 16, 16, 16, half, wmma::row_major> a_frag;
wmma::fragment<wmma::matrix_b, 16, 16, 16, half, wmma::col_major> b_frag;
wmma::fragment<wmma::accumulator, 16, 16, 16, float> c_frag;

wmma::mma_sync(c_frag, a_frag, b_frag, c_frag);

┌─────────────────────────────────────┐
│      Turing Tensor Core 数据类型      │
├─────────────────────────────────────┤
│ FP16 × FP16 → FP16/FP32            │
│ INT8 × INT8 → INT32                │
│ INT4 × INT4 → INT32                │
│ INT1 (二值) × INT1 → INT32         │
└─────────────────────────────────────┘

FP32:  [1位符号][8位指数][23位尾数]
TF32:  [1位符号][8位指数][10位尾数]
FP16:  [1位符号][5位指数][10位尾数]
BF16:  [1位符号][8位指数][7位尾数]

动态范围对比：
FP32/TF32/BF16: ±3.4×10^38
FP16: ±65504

FP8 格式：
E4M3: [1符号][4指数][3尾数] - 更高精度
E5M2: [1符号][5指数][2尾数] - 更大范围

┌──────────────────────────────────┐
│    Transformer Engine 工作流程     │
├──────────────────────────────────┤
│ 1. 统计张量数值分布               │
│ 2. 选择最优精度（FP8/FP16）       │
│ 3. 自动插入量化/反量化            │
│ 4. 动态调整损失缩放因子           │
└──────────────────────────────────┘

精度谱系：
FP64 -> FP32 -> TF32 -> BF16 -> FP16 -> FP8 -> FP6 -> FP4
  ↓       ↓       ↓        ↓       ↓      ↓     ↓     ↓
科学   图形   AI训练   大模型  标准DL  高效  推理  极限

Microscaling (MX) 格式：

- MX6: 6位尾数 + 共享8位标度
- MX4: 4位尾数 + 共享8位标度
- 块大小：32个元素共享一个标度因子

Volta (V100):          每SM 8个Tensor Core
Ampere (A100):         每SM 4个第三代Tensor Core
Hopper (H100):         每SM 4个第四代Tensor Core
Blackwell (B200):      每SM 4个第五代Tensor Core

性能密度提升：
V100:  0.125 TFLOPS/Tensor Core
A100:  0.49 TFLOPS/Tensor Core
H100:  1.23 PFLOPS/Tensor Core
B200:  2.25 PFLOPS/Tensor Core

混合精度训练数据流：
┌──────────────────────────────────────────┐
│            FP32 Master Weights            │
└────────────┬───────────────┬──────────────┘
             ↓ 转换          ↑ 更新
┌──────────────────────────────────────────┐
│            FP16 Weights                   │
├──────────────────────────────────────────┤
│     前向传播 (FP16) → 损失计算            │
│             ↓                             │
│     反向传播 (FP16) ← 梯度计算            │
└──────────────────────────────────────────┘

白名单（FP16）- 计算密集型：
├── 矩阵乘法 (GEMM)
├── 卷积操作 (Conv)
└── RNN/LSTM/GRU

黑名单（FP32）- 精度敏感：
├── 损失函数计算
├── Softmax（大模型）
└── 批归一化统计量

灰名单（动态决定）：
├── 激活函数
├── 池化操作
└── 元素级运算

# 基础AMP训练循环
from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for data, target in dataloader:
    optimizer.zero_grad()

    # 自动混合精度区域
    with autocast():
        output = model(data)
        loss = loss_fn(output, target)

    # 损失缩放反向传播
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

标准反向传播：
Loss → Gradients → Weight Update

损失缩放反向传播：
Loss × Scale → Scaled Gradients → Gradients/Scale → Weight Update

典型缩放因子：2^8 到 2^24

初始化：scale = 2^16, growth_interval = 2000

每次迭代：

1. 计算：scaled_loss = loss × scale
2. 反向传播：compute gradients
3. 检查梯度：
   如果存在 inf/nan：
      scale = scale / 2
      跳过本次更新
   否则：
      更新权重：gradients / scale
      如果连续growth_interval次成功：
         scale = scale × 2

# 模拟大批量训练
accumulation_steps = 4
scale = scaler.get_scale()

for i, (data, target) in enumerate(dataloader):
    with autocast():
        output = model(data)
        loss = loss_fn(output, target)
        loss = loss / accumulation_steps

    scaler.scale(loss).backward()

    if (i + 1) % accumulation_steps == 0:
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad()

        符号  指数  尾数   范围          精度
FP32:    1    8    23    ±3.4×10^38    ~7位小数
FP16:    1    5    10    ±65,504       ~3位小数  
BF16:    1    8     7    ±3.4×10^38    ~2位小数

关键差异：

- BF16牺牲精度换取动态范围
- FP16精度更高但容易溢出

FP16训练：

- 需要精心调整损失缩放
- 可能出现训练不稳定
- 某些层需要FP32

BF16训练：

- 几乎无需调整即可工作
- 训练曲线平滑
- 所有层均可使用BF16

E4M3 (指数4位，尾数3位)：

- 范围：±448
- 精度：0.125
- 用途：前向传播

E5M2 (指数5位，尾数2位)：

- 范围：±57,344  
- 精度：0.25
- 用途：反向传播梯度

混合使用策略：
前向：E4M3 (需要精度)
反向：E5M2 (需要范围)
主权重：FP32 (保证收敛)

┌─────────────────────────────────────┐
│         FP32 Master Weights         │
└──────┬──────────────────┬───────────┘
       ↓ 量化            ↑ 更新
┌─────────────────────────────────────┐
│      E4M3 Weights & Activations     │
├─────────────────────────────────────┤
│    前向传播 → FP32损失计算          │
│         ↓                           │
│    E5M2梯度 ← 反向传播              │
└─────────────────────────────────────┘

import transformer_engine as te

# 自动FP8层替换
model = te.pytorch.TransformerLayer(
    hidden_size=4096,
    ffn_hidden_size=16384,
    num_attention_heads=32,
    fp8=True,  # 启用FP8
    fp8_format="hybrid"  # E4M3 + E5M2
)

# 训练时自动处理量化
with te.pytorch.fp8_autocast(enabled=True):
    output = model(input)
    loss = loss_fn(output, target)

挑战1：激活值分布变化大
解决：Per-tensor动态量化

挑战2：梯度消失/爆炸
解决：延迟量化 + 自适应缩放

挑战3：批归一化精度损失
解决：统计量保持FP32

挑战4：注意力机制数值不稳定
解决：Softmax前后使用FP16/FP32

密集矩阵：                2:4稀疏矩阵：
[1.2, 0.8, 0.3, 0.9]  →  [1.2, 0.0, 0.0, 0.9]
[0.5, 0.1, 0.7, 0.4]  →  [0.5, 0.0, 0.7, 0.0]
[0.2, 0.6, 0.8, 0.3]  →  [0.0, 0.6, 0.8, 0.0]

硬件存储格式：
值数组：   [1.2, 0.9, 0.5, 0.7, 0.6, 0.8]
索引数组： [0,   3,   0,   2,   1,   2]

┌─────────────────────────────────────────┐
│         1. 密集模型预训练                 │
└────────────────┬────────────────────────┘
                 ↓
┌─────────────────────────────────────────┐
│    2. 计算重要性分数（幅值/梯度）         │
└────────────────┬────────────────────────┘
                 ↓
┌─────────────────────────────────────────┐
│    3. 应用2:4掩码（保留最大2个）          │
└────────────────┬────────────────────────┘
                 ↓
┌─────────────────────────────────────────┐
│    4. 稀疏微调（恢复精度）                │
└────────────────┬────────────────────────┘
                 ↓
┌─────────────────────────────────────────┐
│    5. 部署稀疏模型（2×加速）              │
└─────────────────────────────────────────┘

import apex
from apex.contrib.sparsity import ASP

# 自动稀疏化
model = create_model()
optimizer = torch.optim.Adam(model.parameters())

# 配置ASP
ASP.prune_trained_model(model, optimizer)

# 训练循环保持不变
for epoch in range(num_epochs):
    for data, target in dataloader:
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()

量化公式：
q = round(x / scale) + zero_point

反量化公式：
x̂ = (q - zero_point) × scale

其中：

- x: 原始浮点值
- q: 量化整数值
- scale: 缩放因子
- zero_point: 零点偏移

对称量化（常用于权重）：
范围：[-128, 127]
scale = max(|x_max|, |x_min|) / 127
zero_point = 0

非对称量化（常用于激活）：
范围：[0, 255]
scale = (x_max - x_min) / 255
zero_point = round(-x_min / scale)

import tensorrt as trt

# INT8校准
class Int8Calibrator(trt.IInt8EntropyCalibrator2):
    def __init__(self, dataloader):
        self.dataloader = dataloader

    def get_batch(self):
        # 返回校准数据
        return next(self.dataloader)

# 构建INT8引擎
builder = trt.Builder(logger)
config = builder.create_builder_config()
config.set_flag(trt.BuilderFlag.INT8)
config.int8_calibrator = Int8Calibrator(calib_dataloader)

engine = builder.build_engine(network, config)

权重：4位量化
激活：16位保持

优势：

- 模型大小减少75%
- 内存带宽减少4×
- 计算密度提升2-4×

实现策略：

1. 分组量化（每128个权重共享scale）
2. 非均匀量化（更多比特给重要值）
3. 混合精度（关键层保持高精度）

标准前向：
y = Wx + b

QAT前向：
W_q = fake_quantize(W)
x_q = fake_quantize(x)
y = W_q × x_q + b

fake_quantize操作：

1. 量化：q = round(x / scale)
2. 裁剪：q = clip(q, q_min, q_max)
3. 反量化：x̂ = q × scale

import torch.quantization as quant

# 1. 准备QAT模型
model = create_model()
model.qconfig = quant.get_default_qat_qconfig('cuda')
model_prepared = quant.prepare_qat(model)

# 2. QAT训练
for epoch in range(num_epochs):
    for data, target in dataloader:
        output = model_prepared(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()

    # 3. 调整量化参数
    if epoch > 3:
        model_prepared.apply(torch.quantization.enable_observer)
    if epoch > 5:
        model_prepared.apply(torch.quantization.enable_fake_quant)

# 4. 转换为量化模型
model_quantized = quant.convert(model_prepared)

┌────────────────────────────┐
│   1. 收集激活值统计信息      │
│   (最小值、最大值、分布)     │
└─────────────┬──────────────┘
              ↓
┌────────────────────────────┐
│   2. 计算量化参数           │
│   (scale, zero_point)      │
└─────────────┬──────────────┘
              ↓
┌────────────────────────────┐
│   3. 量化权重和偏置         │
│   (离线转换)               │
└─────────────┬──────────────┘
              ↓
┌────────────────────────────┐
│   4. 插入量化/反量化节点     │
│   (运行时转换激活)          │
└────────────────────────────┘

# MinMax校准（快速但精度较低）
def minmax_calibration(model, dataloader):
    min_vals, max_vals = [], []
    for data in dataloader:
        output = model(data)
        min_vals.append(output.min())
        max_vals.append(output.max())
    return min(min_vals), max(max_vals)

# 百分位校准（平衡速度和精度）
def percentile_calibration(model, dataloader, percentile=99.9):
    values = []
    for data in dataloader:
        output = model(data)
        values.extend(output.flatten())
    return np.percentile(values, [100-percentile, percentile])

# KL散度校准（TensorRT默认，精度最高）
def kl_divergence_calibration(model, dataloader):
    # 最小化量化前后分布的KL散度
    # 复杂实现，参考TensorRT源码
    pass

特点：

- 量化参数预先计算固定
- 推理时无需统计计算
- 速度快但精度略低

实现：
输入 → [预计算scale/zp] → INT8计算 → [固定反量化] → 输出

适用场景：

- 输入分布稳定
- 边缘设备部署
- 实时性要求高

特点：

- 运行时计算量化参数
- 适应输入分布变化
- 精度高但有额外开销

实现：
输入 → [动态统计] → [计算scale/zp] → INT8计算 → [动态反量化] → 输出

适用场景：

- 输入分布变化大
- 精度要求高
- 服务器端推理

class HybridQuantizedModel(nn.Module):
    def __init__(self, model):
        super().__init__()
        # 权重静态量化
        self.conv_layers = quantize_static(model.conv_layers)

        # 激活动态量化
        self.fc_layers = quantize_dynamic(model.fc_layers)

        # 敏感层保持FP16
        self.attention = model.attention

    def forward(self, x):
        # 卷积层：INT8静态
        x = self.conv_layers(x)

        # 注意力：FP16
        x = self.attention(x.half())

        # 全连接：INT8动态
        x = self.fc_layers(x)
        return x

模型大小限制严格？
├─是→ INT4量化
│    └─精度要求高？
│        ├─是→ INT4 QAT
│        └─否→ INT4 PTQ (GPTQ/AWQ)
└─否→ INT8量化
     └─可以重训练？
         ├─是→ QAT
         └─否→ PTQ
             └─输入分布稳定？
                 ├─是→ 静态量化
                 └─否→ 动态量化