假设 Q 和 K 的每个元素是独立的均值为 0、方差为 1 的随机变量，则 Q·K^T 的每个元素的方差为 d_k（d_k 个方差为 1 的乘积之和）。
如果不除以 √d_k：
除以 √d_k 后，点积的方差被归一化为 1，softmax 工作在合理区间，梯度正常流动。
scores = Q @ K.T / math.sqrt(d_k) # 方差归一化 attn = softmax(scores) @ V
Multi-Head Attention 的核心优势：
# Multi-Head: 8 heads, d_model=512, d_k=64 # 每个 head 学习独立的 W_Q, W_K, W_V (512 -> 64) head_i = Attention(Q @ W_Q_i, K @ W_K_i, V @ W_V_i) output = Concat(head_1, ..., head_8) @ W_O
MHA (Multi-Head Attention): 每个 head 有独立的 K, V。KV Cache 大小 = 2 * n_heads * d_k * seq_len * batch_size。
MQA (Multi-Query Attention): 所有 Q heads 共享同一组 K, V。KV Cache 缩小为原来的 1/n_heads。例如 32 heads 时省 32x 内存。
GQA (Grouped-Query Attention): 将 Q heads 分成 g 组，每组共享一组 KV。KV Cache 缩小为原来的 g/n_heads。
具体内存节省（以 LLaMA-2 70B 为例，n_heads=64, GQA g=8）：
GQA 在内存节省和模型质量之间取得了较好的平衡，因此被 LLaMA-2 70B、Mistral 等模型采用。
核心思想：将 Q, K, V 分成小 block (tile)，在 SRAM 中逐块计算，不在 HBM 中物化 N*N 的 attention matrix。
在线 Softmax 推导：
标准 softmax: softmax(x_i) = exp(x_i) / sum(exp(x_j))
分块处理时，处理完第一块得到局部 max m1 和局部 sum l1。处理第二块时，新的全局 max m_new = max(m1, m2)，需要校正：
// 处理 block 2 后更新：
m_new = max(m_old, m_block2)
l_new = l_old * exp(m_old - m_new) + l_block2 * exp(m_block2 - m_new)
// 校正之前的输出：
O_new = O_old * (l_old * exp(m_old - m_new) / l_new)
+ softmax_block2(S_block2) @ V_block2 * (l_block2 * exp(m_block2 - m_new) / l_new)
内存分析：
Prefill 阶段：
Decode 阶段：
为什么不缓存 Q：
挑战分析（以 70B 模型, d=8192, 80 layers, GQA g=8 为例）：
推荐方案：多层次组合设计
1. Ring Attention：将序列切分到多个设备，每个设备处理局部 block，KV 通过 ring 传递。内存从 O(N) 降为 O(N/P)，P 为设备数。
2. 稀疏 + 分层注意力：
3. KV Cache 压缩：
4. Position Encoding：使用 RoPE + YaRN/NTK-aware scaling 扩展位置编码到 1M
最终效果：通过 Ring Attention (16 devices) + sparse (effective 8x reduction) + KV quant (4x)，实际内存需求降至约 3.2TB / (16*8*4) = ~6.4 GB/device，可行。
作用：确保位置 i 的 token 只能 attend 到位置 <= i 的 tokens，实现自回归属性。没有 causal mask 的话，模型能"偷看"未来信息，训练时的 teacher forcing 会泄漏答案。
实现方式：
# 方法1: 上三角 mask, 将未来位置设为 -inf mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1) * (-1e9) scores = Q @ K.T / sqrt(d_k) + mask # 加 mask attn = softmax(scores) # -inf 位置 softmax 后为 0 # 方法2: Flash Attention 中通过 block 计算范围隐式实现 # 只计算 j <= i 的 (i, j) 对，跳过上三角区域
注意：Causal mask 是下三角矩阵（包含对角线），上三角为 -inf。在 Flash Attention 中，通过只遍历有效 block 来避免无谓计算，同时节省约 50% 的计算量。
核心思想：对 Q 和 K 向量施加位置相关的旋转变换，使得 q_m^T * k_n 的值仅取决于相对位置 (m-n) 和内容。
数学原理：
m * theta_i 旋转（m 是位置, theta_i = 10000^(-2i/d)）# RoPE 实现核心
def apply_rope(x, freqs):
# x: [batch, seq, heads, d]
# 将 x 拆成实部虚部对
x_complex = torch.view_as_complex(x.reshape(*x.shape[:-1], -1, 2))
# freqs: [seq, d/2], 每个位置的旋转角
freqs_complex = torch.polar(torch.ones_like(freqs), freqs)
# 复数乘法 = 2D 旋转
return torch.view_as_real(x_complex * freqs_complex).flatten(-2)
优势：无需额外参数；天然支持相对位置；可通过 NTK-aware 或 YaRN 方法扩展到更长序列。
Flash Attention 2 的三大优化：
1. 减少非 matmul FLOPs：
2. 并行化维度改变（sequence 维度并行）：
3. Warp 间更好的工作分配：
效果：在 A100 上达到理论 FLOPs 的 ~72%（FA1 约 50%），速度提升约 2x。
H100 参考数据：FP16 峰值 989 TFLOPS, HBM 带宽 3.35 TB/s, 计算/带宽比 = 989/3.35 ~= 295 FLOPs/byte (machine balance point)。
Prefill 阶段 (batch=1, seq_len=L=2048, d=4096)：
Decode 阶段 (batch=1, 新增 1 token, cached L=2048)：
关键区别：Prefill 是矩阵乘矩阵 (GEMM)，AI ~= seq_len；Decode 是矩阵乘向量 (GEMV)，AI ~= 1。这解释了为什么 batching 对 decode 极为重要——增大 batch 将 GEMV 变为 GEMM，提高 AI。
定义：Warp divergence 发生在同一个 warp 的 32 个线程遇到条件分支时，部分线程走 if 路径，部分走 else 路径。
影响：
// Bad: warp divergence
if (threadIdx.x % 2 == 0) { do_A(); }
else { do_B(); }
// Better: 让相邻线程走同一分支
if (threadIdx.x < 16) { do_A(); } // 前半 warp
else { do_B(); } // 后半 warp (仍有 divergence 但更少 warps 受影响)
// Best: 避免分支，用算术替代
val = condition * A + (1-condition) * B;
物理结构：从 Volta 架构开始，每个 SM 的 Shared Memory 和 L1 Cache 共享同一块 SRAM（如 A100 每 SM 192KB, H100 每 SM 256KB）。
关键区别：
配置方式：
// 可通过 API 配置分配比例
cudaFuncSetAttribute(kernel,
cudaFuncAttributePreferredSharedMemoryCarveout,
cudaSharedmemCarveoutMaxShared); // 最大化 shared memory
// H100: 最多 228KB shared memory per SM
// A100: 最多 164KB shared memory per SM
选择策略：需要精确控制数据布局和复用模式时用 Shared Memory；访问模式不规则或难以手动管理时依赖 L1 Cache。Flash Attention 等高性能 kernel 大量使用 Shared Memory 做 tiling。
Roofline Model 基本原理：
Performance (FLOPS/s) = min(Peak Compute, Bandwidth * Arithmetic Intensity)
判断方法：
H100 SXM: Peak FP16 Compute: 989 TFLOPS HBM Bandwidth: 3.35 TB/s Ridge Point = 989 / 3.35 = 295 FLOPs/Byte 如果 kernel 的 AI < 295: Memory-bound (在斜线段) 如果 kernel 的 AI > 295: Compute-bound (在水平线段)
实际例子：
优化方向：memory-bound 优化带宽利用（fusion, coalescing）；compute-bound 优化算力利用（Tensor Core, tiling）。
三个阶段：
// PyTorch 用法
g = torch.cuda.CUDAGraph()
with torch.cuda.graph(g): # Capture
output = model(static_input)
# Instantiate 在 context exit 时自动完成
g.replay() # Launch
限制条件：
优势：消除 kernel launch overhead（从 ~5-10us/kernel 降至整个 graph 一次 ~15us），对 decode 阶段的小 kernel 串效果显著。
Tensor Core MMA 基本操作： D = A * B + C，其中 A, B, C, D 是小矩阵块。
条件要求：
// WMMA API 使用 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::load_matrix_sync(a_frag, a_ptr, lda); wmma::load_matrix_sync(b_frag, b_ptr, ldb); wmma::mma_sync(c_frag, a_frag, b_frag, c_frag);
上层要求：PyTorch 中要利用 Tensor Core，矩阵维度需为 8 (FP16) 或 16 (INT8) 的倍数，否则会 padding 或 fallback 到 CUDA Core。
Step 1: 计算理论 FLOPs
GEMM [M,K] x [K,N] 的 FLOPs = 2 * M * K * N
= 2 * 4096 * 4096 * 4096 = 2 * 4096^3 = 137.4 GFLOPs
Step 2: 计算实际 throughput
实际 Performance = 137.4 GFLOPs / 0.5ms = 137.4 / 0.0005 = 274.9 TFLOPS
Step 3: 计算 MFU
H100 SXM FP16 Tensor Core 峰值 = 989 TFLOPS
MFU = 274.9 / 989 = 27.8%
Step 4: 判断瓶颈
可能原因：
定义：Occupancy = 每个 SM 上活跃 warp 数 / SM 支持的最大 warp 数。例如 H100 每 SM 最多 64 warps (2048 threads)。
影响因素：
# 查看 kernel occupancy
# 编译时: nvcc --ptxas-options=-v 查看寄存器和 shared memory 使用
# 运行时:
cudaOccupancyMaxActiveBlocksPerMultiprocessor(
&numBlocks, kernel, blockSize, sharedMemSize);
注意：高 occupancy 不一定意味着高性能（有时低 occupancy 但高数据复用反而更快），但过低的 occupancy 通常意味着无法有效隐藏内存延迟。
memory_snapshot 能看到：
# 启用 memory 记录
torch.cuda.memory._record_memory_history(max_entries=100000)
# ... 运行模型 ...
# 导出 snapshot
snapshot = torch.cuda.memory._snapshot()
torch.cuda.memory._dump_snapshot("mem_snapshot.pickle")
# 也可用更简单的统计
print(torch.cuda.memory_summary())
# 关键指标:
# Allocated: 当前模型实际使用的
# Reserved: PyTorch caching allocator 持有的（含已释放但未归还 OS 的）
OOM 排查步骤：
torch.cuda.memory_summary() 看 peak memorytorch.cuda.empty_cache()CUDA Stream 定义：一个有序的 GPU 操作（kernel launch, memcpy 等）队列。同一 stream 内操作严格按序执行，不同 stream 的操作可以并行（硬件资源允许时）。
用途：
同步机制：
// 1. cudaStreamSynchronize - 阻塞 CPU 直到 stream 完成
cudaStreamSynchronize(stream1);
// 2. cudaEvent - 跨 stream 同步（非阻塞 GPU 端）
cudaEvent_t event;
cudaEventCreate(&event);
cudaEventRecord(event, stream1); // stream1 中记录 event
cudaStreamWaitEvent(stream2, event); // stream2 等待该 event
// stream2 后续操作会在 stream1 到达 event 后才执行
// 3. cudaDeviceSynchronize - 等待所有 stream 完成
// PyTorch 中:
s1 = torch.cuda.Stream()
s2 = torch.cuda.Stream()
with torch.cuda.stream(s1):
a = torch.matmul(x, y)
event = s1.record_event()
s2.wait_event(event)
with torch.cuda.stream(s2):
b = a + 1 # 安全地使用 s1 的结果
三级 Tiling 层次：
Level 1: Thread Block Tile (Block-level)
Level 2: Warp Tile (Warp-level)
Level 3: Thread Tile (Thread-level / MMA-level)
// 典型配置 (CUTLASS style): // Block tile: 128x128x32 (M x N x K) // Warp tile: 64x64x32 // MMA tile: 16x8x16 // // 每个 Block 有 4 个 warps (2x2 warp grid) // 每个 warp 做 (64/16)*(64/8) = 4*8 = 32 次 MMA per K-step // // 数据流: // HBM --async_copy--> SMEM --ldmatrix--> RF --mma--> RF --store--> SMEM/HBM
优化关键：Software pipelining (多级 buffer 重叠 load 和 compute)、Shared Memory 的 swizzle/padding 避免 bank conflict、异步拷贝 (cp.async) 隐藏 HBM 延迟。
Ring AllReduce 过程：
通信复杂度分析：
为什么与 GPU 数量 N 无关：
虽然步数是 2*(N-1)，但每步发送的数据量是 D/N。总通信量 = 2*(N-1)*D/N = 2D*(1-1/N)，渐近为 2D，与 N 无关。这是因为增加 GPU 时，步数增加但每步数据量等比减少，二者抵消。
总时间：2*(N-1) * (latency + D/(N*bandwidth)) ≈ 2D/bandwidth（带宽主导时）
各 Stage 分片内容和内存节省：
通信量对比：设模型参数量为 M
Trade-off: Stage 3 以 1.5x 通信量换取 N 倍内存节省，适合训练超大模型。
Column Parallel Linear (按列切分 W)：
Row Parallel Linear (按行切分 W)：
Megatron-LM 的经典组合：
// MLP: Column Parallel -> GeLU -> Row Parallel -> AllReduce // 输入 X 广播给所有 GPU Y1 = X @ W1_col_i // Column Parallel, 无通信 Z1 = GeLU(Y1) // 各 GPU 独立计算 Y2 = Z1 @ W2_row_i // Row Parallel, 得到 partial sum Output = AllReduce(Y2) // 唯一的通信点 // Attention 类似: QKV projection (Column) -> Output projection (Row) -> AllReduce
总结：Row Parallel 需要 AllReduce，因为每个 GPU 只计算了最终结果的 partial sum。Column Parallel 的输出本身就是完整结果的一个分片。
Bubble Ratio 公式：
GPipe 方式下: bubble ratio = (p - 1) / (m + p - 1)
其中 p = pipeline stages 数, m = microbatches 数。
1F1B (One Forward One Backward) 调度：
1F1B 的优势：
进一步优化：Interleaved 1F1B (如 virtual stages) 将每个 GPU 分配 v 个 virtual stages，bubble 降为 (p-1)/(m*v+p-1)。
集群拓扑分析：
推荐方案: TP=8, PP=2, DP=4
内存验证 (每 GPU):
替代方案: TP=4, PP=4, DP=4 — 如果节点内带宽不够 TP=8，或需要更小的 bubble。
MoE 的通信模式：
每个 MoE 层需要两次 All-to-All:
All-to-All 通信量计算：
设: N = Expert Parallel degree (EP, expert 分布的 GPU 数) B = batch size (tokens per GPU) d = hidden dimension top_k = 每个 token 选择的 expert 数 E = 总 expert 数 每 GPU 发送: B * top_k * d * (N-1)/N * element_size (每个 token 的 top_k 选择大概率分布在其他 GPU 上) 总通信量 per GPU ≈ B * top_k * d * 2 bytes (FP16) (发送 + 接收各一份, 两次 All-to-All 共 4x) 例: B=4096, top_k=2, d=4096, EP=8 单次 All-to-All per GPU: 4096 * 2 * 4096 * 2 = 64 MB 每层两次: 128 MB 如果模型有 64 个 MoE 层: 总通信 8.2 GB
不均衡性问题：
优化：将 EP 放在节点内 NVLink 上；用 hierarchical All-to-All 减少跨节点通信；DeepSeek-V3 使用 fine-grained expert segmentation 来改善负载均衡。
三种 Backend：
# 典型用法 import torch.distributed as dist # GPU 训练 (推荐) dist.init_process_group(backend="nccl") # CPU 训练 dist.init_process_group(backend="gloo") # 混合: GPU 用 nccl, CPU 用 gloo dist.init_process_group(backend="nccl") # PyTorch 会自动为 CPU tensor 操作使用 gloo
最佳实践：多 GPU 训练统一使用 nccl。如果有 CPU 端的通信需求（如参数服务器），可同时初始化 gloo process group。
原理：
将一个大 batch 分成 K 个 micro-batch，逐个计算梯度并累加，最后一起更新参数：
optimizer.zero_grad()
for i in range(accumulation_steps):
output = model(micro_batch[i])
loss = criterion(output, target[i])
loss = loss / accumulation_steps # 归一化
loss.backward() # 梯度累加到 .grad
# 等价于 effective_batch = micro_batch_size * accumulation_steps
optimizer.step()
数学等价性：
大 batch 梯度: g = (1/B) * sum_{i=1}^{B} grad_i
累积 K 个 micro-batch (each size b, B=K*b):
g = (1/K) * sum_{k=1}^{K} [(1/b) * sum_{i in batch_k} grad_i] = (1/B) * sum_{i=1}^{B} grad_i
数学上完全等价（除了 BatchNorm 的统计量可能不同）。
优势：
# DDP + Gradient Accumulation 优化
for i, batch in enumerate(dataloader):
# 前 K-1 步不通信
ctx = model.no_sync() if (i+1) % K != 0 else nullcontext()
with ctx:
loss = model(batch) / K
loss.backward()
if (i+1) % K == 0:
optimizer.step() # AllReduce 在这里自动发生
optimizer.zero_grad()
Overlap 核心思想：反向传播是从后向前逐层计算梯度的。当后面几层的梯度算好后，可以立即开始 AllReduce，同时前面几层继续计算梯度。通信和计算在时间上重叠。
DDP Gradient Bucketing：
# DDP bucketing 示意 (从后向前)
# 时间线:
# 计算: [grad_layer_N] [grad_layer_N-1] ... [grad_layer_1]
# 通信: [AllReduce bucket_1] [AllReduce bucket_2] ...
# (overlap!)
# 配置:
model = DDP(model,
bucket_cap_mb=25, # 每个 bucket 大小
find_unused_parameters=False, # 静态图优化
gradient_as_bucket_view=True # 减少内存拷贝
)
实现细节：
效果：理想情况下通信几乎完全被计算隐藏，training throughput 接近无通信的理论值。
模型假设：
并行策略设计: TP=8, EP=16, PP=4, DP=4
验证: 8 * 16 * 4 * 4 = 不对。需要重新设计，因为 EP 和 TP/DP 有重叠：
修正方案 (分层设计)：
实际分组: 2048 = TP(8) * PP(4) * DP(64) = 2048，EP(64) 在每个 PP stage 的 512 GPU 中从 DP 维度选 64。
简化方案 (参考 DeepSeek-V3 风格): TP=8, PP=16, DP=16, EP=16
总 GPU = 2048 TP = 8 (节点内 NVLink, dense attention) PP = 16 (128 layers / 16 = 8 layers per stage) DP = 16 (2048 / 8 / 16 = 16) EP = 16 (expert parallel, 与 DP 同维度, 64 experts / 16 = 4 experts per GPU) 通信分析: - TP: AllReduce 2x per layer, 节点内 NVLink (~900GB/s) - PP: 点对点 activation, 跨节点 IB (~50 GB/s), 量小 - DP: AllReduce gradients (dense params only ~50B/TP8/PP16 = 0.4B per GPU) - EP: All-to-All per MoE layer, 跨节点, 主要开销 内存 per GPU: - Dense params: 50B / 8(TP) / 16(PP) * 2 bytes = 0.78 GB - Expert params: (1000B-50B) / 8(TP) / 16(EP) * 2 bytes = 14.8 GB (每 GPU 4 experts) - Optimizer: ~3x params in FP32 / ZeRO-1(DP16) ≈ 3 GB - Activation: gradient checkpointing, ~10-20 GB - Total: ~35-45 GB ✓
关键优化：Expert Parallelism 的 All-to-All 是瓶颈，需要（1）节点内 EP 优先、（2）通信计算 overlap、（3）auxiliary loss 保证负载均衡。
nvidia-smi 和 torch.cuda.memory_summary() 定位内存占用分布。常见原因及解决方案：(1) KV Cache 增长超限（Inference）：长 prompt + 大 batch 导致 KV Cache 超出预分配空间。解决：设置 max_num_seqs 限制最大并发；启用 PagedAttention 避免预分配浪费；配置 mem_fraction_static（如 SGLang 的 --mem-fraction-static 0.85）预留安全余量。(2) Activation Memory 爆炸（Training）：batch size 过大或 sequence length 过长时中间激活占用巨大（正比于 batch×seq_len×hidden×layers）。解决：启用 gradient checkpointing（用计算换内存，减少约 60% 激活内存）；减小 micro-batch size 并增大 gradient accumulation steps；使用 mixed precision training 减半激活内存。(3) Memory Fragmentation：频繁分配释放不同大小的 tensor 导致碎片，虽然总 free memory 够但没有足够大的连续块。解决：torch.cuda.empty_cache() 释放缓存；使用 memory pool（如 CUDA Memory Pool）；设置 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True 允许非连续分配。额外工具：torch.cuda.memory_snapshot() 生成内存时间线，可视化分析碎片。nvidia-smi -q 检查 GPU clock frequency（是否降频，thermal throttling）、ECC error count（单 bit 错误会触发重试降速）、PCIe bandwidth（是否降级到 x8/Gen3）。检查 dmesg 是否有 Xid error。如果某卡 SM clock 从 1980MHz 降到 1500MHz，说明 thermal throttle。(2) 通信层检查：nccl-tests 跑 all_reduce 带宽测试，确认 NVLink/IB 是否正常。如果某条 NVLink 故障，NCCL 会 fallback 到 PCIe，带宽从 600GB/s 降到 64GB/s。检查 IB port 状态 ibstat。(3) 数据层检查：data loading 是否成为瓶颈——检查 DataLoader worker 的 CPU/IO utilization。如果共享存储（NFS/Lustre）性能下降或某个 epoch 数据变大（如 curriculum learning 切换到长序列），会拖慢整体。Profiler 看 GPU 是否有 idle gap 在等数据。(4) 框架层检查：是否 dynamic shape 导致 CUDA Graph 失效，每个 iteration 重新 compile kernel？某些 batch 的 sequence length 特别长触发更多计算？gradient accumulation 配置是否被意外修改？(5) 系统层：检查是否有其他进程抢占 GPU（nvidia-smi pmon），或 CPU 资源被其他 job 占用（cgroup 限制）。用 Nsight Systems/PyTorch Profiler 做 trace 对比正常时和异常时的 kernel timeline 差异。torch.distributed.all_gather 收集各 rank 的 grad norm 做对比。(2) 数据问题排查：检查异常 rank 的 data loader 是否读到 corrupted 数据（如 NaN in input、truncated tokenization）。验证方法：在异常 rank 上打印几个 batch 的 input 和 label，确认合理性。特别注意 padding/masking 的实现在不同 rank 是否一致。(3) 通信正确性：NCCL allreduce 后各 rank 的参数应完全一致。检查：训练 N 步后比较各 rank 的 model checksum（sum(p.data.sum() for p in model.parameters())）。如果不一致，说明某次 allreduce 出错或有 rank 跳过了同步步骤。设置 NCCL_DEBUG=WARN 检查通信错误。(4) 数值稳定性：Mixed precision training 中 loss scale 是否在某些 rank 上频繁 overflow 导致 gradient 被跳过？检查 grad_scaler._scale 和 skip 次数。如果某些 rank 的数据分布极端，可能反复触发 overflow。(5) 硬件 ECC 错误：GPU 的 memory bit flip 可能导致计算结果随机错误。检查 nvidia-smi --query-gpu=ecc.errors.uncorrected.aggregate.total -i GPU_ID。(6) 修复策略：确认问题后——数据问题则修复数据 pipeline；通信问题则更换硬件或强制 NCCL barrier + checksum 验证；数值问题则调整 loss scale 策略或对异常 gradient 做 clipping。gc.disable() 在 critical path，预分配 memory pool。(6) 网络抖动：多节点 TP 时某次 allreduce 被网络延迟拖慢。排查：监控 NCCL op 的 per-call latency。