%%{init: {'theme': 'base'}}%%
graph TB
A["训练性能"] --> B["计算效率<br/>FLOPs/s"]
A --> C["通信效率<br/>Bandwidth Util"]
A --> D["内存效率<br/>Peak Memory"]
A --> E["I/O 效率<br/>Samples/s"]
B --> B1["混合精度<br/>BF16/FP8"]
B --> B2["算子融合<br/>自定义 Kernel"]
B --> B3["编译优化<br/>TorchDynamo"]
C --> C1["通信隐藏<br/>计算-通信重叠"]
C --> C2["梯度压缩<br/>Top-K / 阈值"]
C --> C3["拓扑优化<br/>分层 AllReduce"]
D --> D1["激活重计算<br/>Gradient Checkpointing"]
D --> D2["ZeRO 分片<br/>FSDP"]
D --> D3["CPU Offload"]
E --> E1["预取流水线<br/>DataLoader 调优"]
E --> E2["内存映射<br/>mmap 数据格式"]
E --> E3["数据格式<br/>WebDataset / Parquet"]
训练性能优化
训练性能优化
性能优化是一个永无止境的过程。但前 80% 的收益往往来自 20% 的优化——关键在于找到瓶颈在哪。
章节导入：性能优化的第一原则
在开始任何优化之前，先问自己三个问题：
- 我的瓶颈是什么？（计算、通信、内存、I/O？）
- 我优化的代价是什么？（精度损失？代码复杂度？维护成本？）
- 这个优化能推广吗？（只对当前模型有效，还是通用？）
没有 profiler 的优化就是盲人摸象。本章的每个建议都基于一个假设：你已经用工具看到了瓶颈。
8.1 混合精度训练与 BF16/FP8
混合精度训练是最”性价比高”的优化：几乎不改动代码，就能获得 2-3 倍的加速和约 50% 的内存节省。
数值格式对比
| 格式 | 总位数 | 符号位 | 指数位 | 尾数位 | 动态范围 | 精度 | GPU 支持 |
|---|---|---|---|---|---|---|---|
| FP32 | 32 | 1 | 8 | 23 | ~10⁻³⁸ ~ 10³⁸ | 高 | 所有 |
| FP16 | 16 | 1 | 5 | 10 | ~10⁻⁵ ~ 65504 | 中 | Volta+ |
| BF16 | 16 | 1 | 8 | 7 | ~10⁻³⁸ ~ 10³⁸ | 低 | Ampere+ |
| FP8 (E4M3) | 8 | 1 | 4 | 3 | ~10⁻⁶ ~ 448 | 极低 | Hopper+ |
| FP8 (E5M2) | 8 | 1 | 5 | 2 | ~10⁻¹⁴ ~ 57344 | 极低 | Hopper+ |
%%{init: {'theme': 'base'}}%%
graph LR
subgraph "BF16 vs FP16"
FP16["FP16: 指数 5 位 + 尾数 10 位<br/>精度高，但范围小（容易溢出）"]
BF16["BF16: 指数 8 位 + 尾数 7 位<br/>精度低，但范围与 FP32 相同"]
end
Note["BF16 是大模型训练的默认选择<br/>不需要 Loss Scaling！"]
BF16 混合精度训练
import torch
# 方式 1：autocast（推荐，最简洁）
model = MyModel().cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
for batch in dataloader:
# 前向传播使用 BF16
with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
loss = model(batch)
# 反向传播：梯度自动转为 FP32 计算
loss.backward()
optimizer.step()
optimizer.zero_grad()# 方式 2：模型直接转为 BF16（省更多显存）
model = model.to(dtype=torch.bfloat16)
# 注意：某些层需要保持 FP32 以保稳定性
for module in model.modules():
if isinstance(module, (nn.LayerNorm, nn.BatchNorm1d)):
module.float() # LayerNorm 保持 FP32FP8 训练（Hopper 架构）
FP8 是 NVIDIA Hopper（H100/H200）引入的 8 位浮点格式，能将计算吞吐量再翻一倍。
# FP8 训练需要 Transformer Engine 库
import transformer_engine.pytorch as te
from transformer_engine.common.recipe import Format, DelayedScaling
# 使用 TE 的模块替代标准 nn.Module
model = te.TransformerLayer(
hidden_size=4096,
ffn_hidden_size=16384,
num_attention_heads=32,
layernorm_eps=1e-5
)
# FP8 recipe：使用 Delayed Scaling 策略
fp8_recipe = DelayedScaling(
fp8_format=Format.HYBRID, # 前向用 E4M3，反向用 E5M2
amax_history_len=16, # 用 16 步的 amax 历史来更新缩放因子
amax_compute_algo="max", # 取历史最大值
)
# 训练循环
for batch in dataloader:
with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
loss = model(batch)
loss.backward()
optimizer.step()
optimizer.zero_grad()
WarningFP8 的注意事项
- 不是所有层都适合 FP8：Embedding 层、Loss 计算层、概率计算层（Softmax 前）应保持高精度
- Delayed Scaling 需要预热：前几步用 FP16/BF16，等 amax 统计稳定后再切换到 FP8
- 小 batch 不友好：FP8 的精度极低，batch 太小时统计不稳定
- 目前仅限 Hopper+：A100 不支持 FP8 硬件加速
性能基准对比
模型: GPT-3 175B (类似配置)
硬件: 8 × H100 80GB SXM5
| 精度模式 | TFLOPS | 相对吞吐量 | 显存使用 |
|---------------|--------|----------|---------|
| FP32 | ~500 | 1.0× | 100% |
| FP16 (AMP) | ~1000 | 2.0× | ~60% |
| BF16 (AMP) | ~1000 | 2.0× | ~60% |
| BF16 (全模型) | ~1000 | 2.1× | ~50% |
| FP8 (TE) | ~1800 | 3.6× | ~35% |8.2 激活重计算与内存换计算
当显存不够时，激活重计算（Activation Recomputation / Gradient Checkpointing）是最有效的内存优化手段之一。
原理
%%{init: {'theme': 'base'}}%%
graph TB
subgraph "标准前向传播"
S1["保存所有中间激活值"] --> S2["反向传播直接使用"]
S3["内存：O(n_layers × batch × seq × hidden)"]
end
subgraph "激活重计算"
R1["只保存关键检查点"] --> R2["反向传播时重新计算"]
R1 --> R3["内存：O(sqrt(n_layers) × batch × seq × hidden)"]
R2 --> R4["计算：多约 33%"]
end
选择式激活重计算
不是所有层都需要重计算。研究表明，Transformer 中注意力矩阵占用的内存远大于 MLP 的中间激活值。只重计算 Attention 部分就能节省大量内存，而计算开销很小。
from torch.utils.checkpoint import checkpoint
class SelectiveCheckpointedLayer(nn.Module):
"""只重计算 Attention，保留 MLP 的激活值"""
def __init__(self, config):
super().__init__()
self.attention = MultiHeadAttention(config)
self.mlp = MLP(config)
self.norm1 = nn.LayerNorm(config.hidden_size)
self.norm2 = nn.LayerNorm(config.hidden_size)
def forward(self, x):
# Attention 使用 checkpoint（重计算）
x = x + checkpoint(
self._attention_block, x,
use_reentrant=False
)
# MLP 不使用 checkpoint（保留激活值）
x = x + self._mlp_block(x)
return x
def _attention_block(self, x):
return self.attention(self.norm1(x))
def _mlp_block(self, x):
return self.mlp(self.norm2(x))内存收益估算
70B 模型，seq_len=4096, batch_size=4
| 策略 | 激活值内存 | 计算开销 | 推荐场景 |
|-----|----------|---------|---------|
| 无重计算 | ~80 GB | 0% | 显存充足，追求最大吞吐 |
| 选择式（仅 Attention） | ~40 GB | ~5% | ✅ 推荐默认方案 |
| 全部重计算 | ~15 GB | ~33% | 极端内存不足 |
| CPU Offload 激活值 | ~10 GB GPU | 0% + PCIe | 最极端情况 |
Tip选择性重计算的公式
Transformer 层中，Attention 的中间激活值占整层的比例约为： \[\frac{\text{Attention memory}}{\text{Total}} \approx \frac{2 \times \text{seq\_len}}{2 \times \text{seq\_len} + 4 \times \text{hidden\_size}}\]
当 seq_len=4096, hidden_size=4096 时，Attention 约占 50%。只重计算这部分就能省一半内存，代价仅为 5-8% 的额外计算。这是最划算的内存优化之一。
8.3 通信优化
在大规模分布式训练中，通信往往是最大的瓶颈——特别是在跨节点场景下。
AllReduce 拓扑优化
Ring AllReduce 是最经典的算法，但当节点数增加时，环的延迟会累积。分层 AllReduce 将节点内和节点间通信分开：
%%{init: {'theme': 'base'}}%%
graph TB
subgraph "扁平 Ring AllReduce"
R["GPU 0 → GPU 1 → ... → GPU 63 → GPU 0"]
R_NOTE["需要 63 步才能完成一圈"]
end
subgraph "分层 AllReduce"
subgraph "Node 0 (NVLink 900GB/s)"
L00["GPU 0"] <--> L01["GPU 1"]
L01 <--> L02["GPU 2"]
L02 <--> L03["GPU 3"]
end
subgraph "Node 1"
L10["GPU 4"] <--> L11["GPU 5"]
L11 <--> L12["GPU 6"]
L12 <--> L13["GPU 7"]
end
L00 <-.->|IB 100GB/s| L10
L01 <-.-> L11
end
# 使用分层 AllReduce（节点内 NCCL，节点间 Ring）
os.environ["NCCL_INTER_NODE_CONNECTIVITY"] = "ring"
os.environ["NCCL_INTRA_NODE_CONNECTIVITY"] = "nvlink"
os.environ["NCCL_MIN_NCHANNELS"] = "4" # 多通道并行
os.environ["NCCL_NET_GDR_LEVEL"] = "2" # GPUDirect RDMA 级别
# 或者使用自定义的分层通信
def hierarchical_all_reduce(tensor):
world_size = dist.get_world_size()
local_size = 8 # 节点内 GPU 数
# 第一步：节点内 AllReduce（快，NVLink）
local_group = get_local_group()
dist.all_reduce(tensor, group=local_group, op=dist.ReduceOp.SUM)
# 第二步：跨节点 AllReduce（慢，IB）
# 只有每个节点的 rank 0 参与
if dist.get_rank() % local_size == 0:
inter_group = get_inter_node_group()
dist.all_reduce(tensor, group=inter_group, op=dist.ReduceOp.SUM)
# 第三步：节点内广播结果
dist.broadcast(tensor, src=0, group=local_group)
tensor /= world_size
return tensor通信-计算重叠
最理想的模式是在计算下一层的同时同步当前层的梯度：
class CommunicationOverlap:
"""梯度同步与反向传播重叠"""
def __init__(self, model):
self.model = model
self.comm_stream = torch.cuda.Stream()
def backward_with_overlap(self, loss):
"""手动控制梯度的通信时机"""
layers = list(self.model.parameters())
for i, layer in enumerate(layers):
# 在默认流上计算梯度
loss.backward(retain_graph=(i < len(layers) - 1))
if i > 0:
# 等待上一层通信完成
torch.cuda.current_stream().wait_stream(self.comm_stream)
# 上一层优化器更新
self.optimizer_step(layers[i-1])
# 在通信流上异步同步当前层梯度
with torch.cuda.stream(self.comm_stream):
dist.all_reduce(layer.grad, async_op=True)
# 处理最后一层
torch.cuda.current_stream().wait_stream(self.comm_stream)
self.optimizer_step(layers[-1])
Tip使用框架内置的通信重叠
现代框架（DDP、FSDP、DeepSpeed）已经内置了通信重叠优化。手动实现的复杂度极高，不推荐自行实现。理解原理有助于调参和诊断问题。
FSDP 的 forward_prefetch=True 就是自动做参数预取与计算重叠。
梯度压缩
当通信带宽受限时，可以通过压缩梯度来减少传输量：
# Top-K 稀疏化：只传输最大的 K 个梯度
def top_k_compress(tensor, ratio=0.01):
"""只保留梯度中 top-k% 的元素"""
k = max(1, int(tensor.numel() * ratio))
values, indices = torch.topk(tensor.abs(), k)
# 构造稀疏梯度
compressed = torch.zeros_like(tensor)
compressed[indices] = tensor[indices]
return compressed, values, indices
# 1-bit SGD：梯度量化为 ±1
def one_bit_compress(tensor):
"""将梯度量化为符号"""
sign = tensor.sign()
# 缩放因子：均值
scale = tensor.abs().mean()
return sign * scale
# 使用方式（在梯度同步前压缩）
for param in model.parameters():
if param.grad is not None:
compressed_grad = top_k_compress(param.grad, ratio=0.01)
# AllReduce 压缩后的梯度
dist.all_reduce(compressed_grad)
param.grad.copy_(compressed_grad)
Warning梯度压缩的代价
梯度压缩会引入噪声，可能影响收敛。特别是 1-bit 量化在高精度要求（如 fine-tune 阶段）时不推荐使用。预训练阶段一般可以容忍 1% 的 Top-K 稀疏化，但对最终精度的影响需要在你的具体模型上验证。
8.4 数据加载与预处理流水线
数据加载往往是训练中最容易被忽视的瓶颈。当 GPU 算力越来越强，数据供给的速度越来越容易跟不上。
诊断：数据加载是否是瓶颈
import time
import torch.profiler
# 方法 1：简单计时
class DataLoadTimer:
def __init__(self):
self.data_time = 0
self.compute_time = 0
def profile_step(self, loader_iter, train_fn):
torch.cuda.synchronize()
t0 = time.time()
batch = next(loader_iter)
torch.cuda.synchronize()
t1 = time.time()
train_fn(batch)
torch.cuda.synchronize()
t2 = time.time()
self.data_time += t1 - t0
self.compute_time += t2 - t1
# 数据加载时间 > 计算时间的 20% → 瓶颈
if self.data_time > 0.2 * self.compute_time:
logging.warning(
f"Data loading is slow! "
f"Data: {self.data_time:.2f}s, "
f"Compute: {self.compute_time:.2f}s"
)
# 方法 2：PyTorch Profiler
with torch.profiler.profile(
activities=[
torch.profiler.ProfilerActivity.CPU,
torch.profiler.ProfilerActivity.CUDA,
],
schedule=torch.profiler.schedule(wait=1, warmup=1, active=3),
on_trace_ready=torch.profiler.tensorboard_trace_handler('./profiler_logs'),
) as prof:
for step, batch in enumerate(loader):
loss = model(batch)
loss.backward()
optimizer.step()
prof.step() # 必须调用DataLoader 优化清单
# 一个优化过的 DataLoader 配置
dataloader = DataLoader(
dataset,
batch_size=32,
shuffle=False, # 使用 DistributedSampler
sampler=DistributedSampler(dataset),
num_workers=8, # ⭐ 每个 GPU 8 个 worker 进程
pin_memory=True, # ⭐ 锁页内存，加速 CPU→GPU 传输
persistent_workers=True, # ⭐ 跨 epoch 复用 worker
prefetch_factor=4, # ⭐ 每个 worker 预取 4 个 batch
drop_last=True, # 避免最后一个不完整 batch
collate_fn=custom_collate, # 高效的 collate 函数
worker_init_fn=worker_init_fn, # 确保每个 worker 有不同 seed
)
def worker_init_fn(worker_id):
"""确保每个 DataLoader worker 有不同的随机种子"""
worker_seed = torch.initial_seed() % 2**32 + worker_id
np.random.seed(worker_seed)
random.seed(worker_seed)高效数据格式
| 格式 | 读取方式 | 适用场景 | 随机访问 |
|---|---|---|---|
| CSV/JSON | 顺序解析 | 小数据集、原型 | ✅ |
| TFRecord | 顺序读 | TensorFlow 生态 | ❌ |
| HDF5 | mmap | 科学计算 | ✅ |
| Parquet | 列式读 | 结构化数据 | ✅ |
| WebDataset (TAR) | 顺序流式 | ⭐ 大规模训练 | ❌（但可 shard） |
| MMAP Binary | 直接内存映射 | ⭐ 低延迟随机访问 | ✅ |
# WebDataset：大规模训练的最佳实践
import webdataset as wds
# 将数据组织为 tar shard
# shard-000000.tar, shard-000001.tar, ...
# 每个 tar 包含: sample_0.json, sample_0.bin, sample_1.json, sample_1.bin, ...
dataset = (
wds.WebDataset("shards/shard-{000000..000999}.tar")
.shuffle(1000) # 缓冲 1000 个样本后打乱
.decode("torch") # 自动解码为 PyTorch tensor
.batched(32) # 打包 batch
)
dataloader = torch.utils.data.DataLoader(
dataset,
num_workers=4,
batch_size=None, # WebDataset 内部已 batch
)预处理流水线优化
# 将耗时的预处理操作移到 GPU 上
class GPUPreprocessor:
def __init__(self, device="cuda"):
self.device = device
def __call__(self, batch):
# 将原始数据直接搬到 GPU，然后在 GPU 上做预处理
batch = {k: v.to(self.device, non_blocking=True) for k, v in batch.items()}
# 在 GPU 上做 tokenization / augmentation
# 比在 CPU 上快 10-100x
batch["input_ids"] = self.tokenize_gpu(batch["text"])
batch["labels"] = self.create_labels_gpu(batch["input_ids"])
return batch
# 使用 NVIDIA DALI 实现高性能数据流水线
# DALI 在 GPU 上执行完整的预处理管线
from nvidia.dali import pipeline_def
import nvidia.dali.types as types
import nvidia.dali.fn as fn
@pipeline_def
def training_pipeline():
images, labels = fn.readers.file(
file_root=data_dir,
random_shuffle=True,
)
images = fn.decoders.image(images, device="mixed") # GPU 解码
images = fn.resize(images, resize_x=224, resize_y=224)
images = fn.crop_mirror_normalize(
images,
crop=(224, 224),
mean=[0.485 * 255, 0.456 * 255, 0.406 * 255],
std=[0.229 * 255, 0.224 * 255, 0.225 * 255],
)
return images, labels
pipe = training_pipeline(batch_size=256, num_threads=4, device_id=0)
pipe.build()8.5 算子级优化与自定义 Kernel
当框架的通用实现不够快时，自定义 kernel 可以带来最后的性能提升。
算子融合（Operator Fusion）
将多个小算子合并为一个大算子，减少 kernel launch 开销和中间结果的访存：
# 标准实现：3 次 kernel launch + 2 次中间结果写回
def standard_ffn(x, w1, b1, w2, b2):
h = torch.mm(x, w1.t()) + b1 # kernel 1
h = torch.gelu(h) # kernel 2
out = torch.mm(h, w2.t()) + b2 # kernel 3
return out
# 融合实现：1 次调用
# 使用 TorchInductor 自动融合
@torch.compile(mode="max-autotune")
def fused_ffn(x, w1, b1, w2, b2):
h = torch.mm(x, w1.t()) + b1
h = torch.gelu(h)
out = torch.mm(h, w2.t()) + b2
return out
# Inductor 会自动将 GELU 和 Bias addition 融合到矩阵乘法的 epilogueFlashAttention：最重要的算子优化
FlashAttention（Dao et al., 2022）通过分块计算和减少 HBM 访问，将 Attention 的内存复杂度从 \(O(n^2)\) 降到 \(O(n)\)，同时速度提升 2-4 倍。
# 使用 FlashAttention 2
from flash_attn import flash_attn_func
class FlashAttention(nn.Module):
def forward(self, q, k, v, causal=True):
# q, k, v: [batch, seq_len, num_heads, head_dim]
# 不需要实例化 [seq_len, seq_len] 的 attention matrix！
output = flash_attn_func(q, k, v, causal=causal)
return output
# 对比标准 Attention
class StandardAttention(nn.Module):
def forward(self, q, k, v, causal=True):
# 需要 O(seq_len²) 的中间矩阵
attn_weights = torch.matmul(q, k.transpose(-2, -1))
if causal:
mask = torch.triu(torch.ones_like(attn_weights), diagonal=1)
attn_weights = attn_weights.masked_fill(mask == 1, float('-inf'))
attn_weights = F.softmax(attn_weights, dim=-1)
output = torch.matmul(attn_weights, v)
return output
# 性能对比（seq_len=8192, batch=4, d_head=128, A100 80GB）：
# Standard: 12.4 GB 激活值, 42ms/step
# FlashAttn: 0.8 GB 激活值, 11ms/step → 3.8× 加速, 15.5× 省内存使用 Triton 编写自定义 Kernel
import triton
import triton.language as tl
@triton.jit
def fused_layer_norm_gelu_kernel(
X_PTR, W_PTR, B_PTR, Y_PTR,
N_ROWS, N_COLS,
stride_x_row, stride_y_row,
eps: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
):
"""融合 LayerNorm + GELU：一次访存完成两个操作"""
row = tl.program_id(0)
# 计算行的偏移
x_row_ptr = X_PTR + row * stride_x_row
y_row_ptr = Y_PTR + row * stride_y_row
# 加载整行数据
cols = tl.arange(0, BLOCK_SIZE)
mask = cols < N_COLS
x = tl.load(x_row_ptr + cols, mask=mask, other=0.0).to(tl.float32)
# LayerNorm
mean = tl.sum(x, axis=0) / N_COLS
x_centered = x - mean
variance = tl.sum(x_centered * x_centered, axis=0) / N_COLS
x_normed = x_centered / tl.sqrt(variance + eps)
# 应用 weight 和 bias
w = tl.load(W_PTR + cols, mask=mask)
b = tl.load(B_PTR + cols, mask=mask)
x_normed = x_normed * w + b
# GELU（融合到这里）
gelu = x_normed * 0.5 * (1.0 + tl.tanh(
0.7978845608 * (x_normed + 0.044715 * x_normed * x_normed * x_normed)
))
# 写回结果
tl.store(y_row_ptr + cols, gelu.to(tl.bfloat16), mask=mask)
# 在模型中使用
def fused_layernorm_gelu(x, weight, bias, eps=1e-5):
n_rows, n_cols = x.shape
block_size = triton.next_power_of_2(n_cols)
output = torch.empty_like(x)
fused_layer_norm_gelu_kernel[(n_rows,)](
x, weight, bias, output,
n_rows, n_cols,
x.stride(0), output.stride(0),
eps=eps,
BLOCK_SIZE=block_size,
)
return output
Tip什么时候需要自定义 Kernel
- 前 80% 的性能 来自框架配置（混合精度、FSDP、FlashAttention 等）
- 后 20% 才需要自定义 Kernel
- 优先使用现成的优化库：
flash_attn,xformers,transformer_engine,apex - 只有当 profiler 显示某个算子占据了显著时间，且没有现成优化时，才写 Triton Kernel
性能 Profiling 完整流程
#!/bin/bash
# 训练性能诊断完整流程
# 1. PyTorch Profiler
python -c "
import torch.profiler
with torch.profiler.profile(
activities=[
torch.profiler.ProfilerActivity.CPU,
torch.profiler.ProfilerActivity.CUDA,
],
schedule=torch.profiler.schedule(wait=1, warmup=1, active=5),
on_trace_ready=torch.profiler.tensorboard_trace_handler('./log'),
) as prof:
for i, batch in enumerate(loader):
train_step(batch)
prof.step()
"
# 2. NCCL 通信分析
os.environ["NCCL_PROFILE"] = "1"
os.environ["NCCL_DEBUG"] = "INFO"
# 3. GPU 利用率监控
nvidia-smi dmon -s pucvmet -d 1 # 每秒采样
# 4. 分析结果
# 查看 TensorBoard 的 PyTorch Profiler 标签页
# 关注：
# - GPU 利用率（目标 > 90%）
# - kernel 时间分布
# - 通信操作占比
# - CPU 等待时间
# - 内存使用峰值小结
性能优化的优先级——按照投入产出比排序：
| 优先级 | 优化手段 | 预期收益 | 实施难度 |
|---|---|---|---|
| 1 | 混合精度（BF16） | 2× 吞吐、50% 内存 | ⭐ |
| 2 | FlashAttention | 3-4× Attention 速度 | ⭐ |
| 3 | FSDP / ZeRO | 内存大幅减少 | ⭐⭐ |
| 4 | 激活重计算 | 50%+ 激活值内存 | ⭐ |
| 5 | DataLoader 优化 | 消除数据瓶颈 | ⭐⭐ |
| 6 | 通信重叠 | 10-20% 吞吐提升 | ⭐⭐ |
| 7 | torch.compile |
10-30% 计算 | ⭐ |
| 8 | FP8 训练 | 额外 2× 吞吐 | ⭐⭐⭐ |
| 9 | 自定义 Triton Kernel | 5-20% 计算 | ⭐⭐⭐⭐ |
记住：永远先 profile，再优化。不要凭直觉——GPU 内部发生的事情往往和你的直觉不同。
延伸阅读
- FlashAttention：Dao & Guo, “FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning” (2023)
- FP8 训练：NVIDIA, “Transformer Engine: FP8 Training” (2023)
- Triton 教程：https://triton-lang.org/main/getting-started/tutorials/
- PyTorch Profiler：官方文档 “PyTorch Profiler”
- DeepSpeed ZERO：Rajbhandari et al. (2020)
- 选择性激活重计算：Chen et al., “ActNN: Reducing Training Memory Footprint via 2-Bit Activations” (2021)
- 通信优化综述：Narayanan et al., “Efficient Large-Scale Language Model Training” (2021)