第九章:高级主题与前沿技术
章节大纲
9.1 分布式推理与模型并行
- 模型并行的编译优化
- Pipeline 并行与微批处理
- 张量并行的自动分片
- 分布式推理的通信优化
9.2 动态批处理与序列并行
- 动态批处理的编译挑战
- Continuous Batching 技术
- 序列并行的内存优化
- PagedAttention 与 vLLM 集成
9.3 硬件加速器适配
- NPU/TPU 的编译后端
- 自定义 ASIC 的图映射
- 异构计算的调度策略
- 硬件感知的图优化
9.4 编译器技术的未来发展
- AI 编译器的发展趋势
- 自适应编译与在线学习
- 跨框架的统一 IR
- 量子计算的编译挑战
开篇
在前面的章节中,我们深入探讨了 PyTorch 编译技术的核心概念和实践方法。本章将聚焦于高级主题和前沿技术,这些内容代表了当前深度学习编译器领域的最新进展和未来方向。在自动驾驶和具身智能系统中,这些技术正在推动着更大规模、更高效率的模型部署成为可能。
随着模型规模的指数级增长和硬件架构的多样化发展,传统的单机单卡推理已经无法满足需求。分布式推理、动态批处理、异构硬件适配等技术成为了必然选择。同时,编译器技术本身也在经历深刻的变革,从静态优化走向动态自适应,从单一框架走向跨平台统一。
本章的学习目标:
- 掌握大规模模型的分布式推理技术和编译优化方法
- 理解动态批处理和序列并行的实现原理与优化策略
- 学习不同硬件加速器的适配方法和编译后端设计
- 了解编译器技术的最新研究进展和未来发展趋势
这些高级技术虽然复杂,但在实际生产环境中却至关重要。让我们开始这段探索之旅。
9.1 分布式推理与模型并行
9.1.1 模型并行的编译优化
在自动驾驶的多模态大模型和具身智能的决策网络中,模型参数量往往超过单个 GPU 的内存容量。以 Tesla FSD 的视觉 Transformer 为例,其参数量可达数十亿,单卡 A100 的 80GB 显存难以容纳完整模型及其激活值。模型并行成为必然选择,但如何在编译阶段优化分布式执行图是一个关键挑战。
模型并行的核心思想是将神经网络的不同部分分配到不同的计算设备上。与数据并行不同,模型并行不是复制整个模型,而是将模型的参数和计算进行切分。这种方式特别适合参数量巨大但批次较小的推理场景,正如自动驾驶系统中的实时感知任务。
张量模型并行(Tensor Model Parallelism)
张量并行是一种细粒度的并行方式,它将单个算子的计算分布到多个设备上。这种方法的关键在于找到算子内部的可并行维度,并设计高效的通信模式。
对于最常见的线性层 Y = XW + b,我们有两种主要的切分策略:
- 列并行(Column Parallel):按输出维度切分
设备 0: W_0 = W[:, :k/2] → Y_0 = XW_0
设备 1: W_1 = W[:, k/2:] → Y_1 = XW_1
最终结果: Y = concat([Y_0, Y_1], dim=1)
这种切分方式的优势在于前向传播不需要通信,每个设备独立计算自己的输出部分。但在反向传播时需要 AllReduce 来同步梯度。
- 行并行(Row Parallel):按输入维度切分
设备 0: W_0 = W[:d/2, :] → 需要 X[:, :d/2]
设备 1: W_1 = W[d/2:, :] → 需要 X[:, d/2:]
最终结果: Y = Y_0 + Y_1 (需要 AllReduce)
行并行需要在前向传播时进行 AllReduce,但反向传播时各设备可以独立计算梯度。
编译器在选择并行策略时需要考虑多个因素:通信带宽、计算密度、内存占用等。现代编译器通过成本模型自动选择最优策略:
# 编译时的图转换示例
@torch.compile(options={"tensor_parallel": 2})
def parallel_linear(x, weight):
# 编译器自动分析并选择最优切分策略
# 考虑因素:
# 1. 通信/计算比
# 2. 内存带宽利用率
# 3. 设备间通信拓扑
return F.linear(x, weight)
对于注意力机制,张量并行更加复杂。自注意力的计算涉及 Q、K、V 三个投影和最终的输出投影,编译器需要协调这些操作的并行策略:
注意力头并行切分:
设备 0: 处理前 H/2 个注意力头
设备 1: 处理后 H/2 个注意力头
Q_0 = X @ W_q_0 (heads 0 to H/2-1)
K_0 = X @ W_k_0
V_0 = X @ W_v_0
Attn_0 = softmax(Q_0 @ K_0.T) @ V_0
最终需要 AllGather 来合并所有头的输出
流水线并行(Pipeline Parallelism)
流水线并行将模型按层切分到不同设备,形成计算流水线。这种方法特别适合深度网络,如自动驾驶中的多层感知网络。编译器需要解决的核心问题包括:
- 自动阶段划分:基于计算量和内存占用平衡各阶段
理想的划分应该使每个阶段的计算时间大致相等,同时考虑内存约束。编译器通过分析每层的 FLOPs 和参数量来决定划分点:
阶段划分算法:
total_flops = sum(layer_flops)
target_flops_per_stage = total_flops / num_devices
current_stage_flops = 0
for layer in model.layers:
current_stage_flops += layer.flops
if current_stage_flops >= target_flops_per_stage:
# 在此处划分阶段
create_stage_boundary()
current_stage_flops = 0
- 微批调度:优化 1F1B(一次前向一次反向)调度策略
微批处理将一个批次分成多个小批次,通过流水线并行处理。1F1B 调度策略在稳定状态下保持一个前向传播立即跟随一个反向传播,最小化激活值的生命周期:
阶段 0 (GPU0): Layers[0:4] ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
↓
阶段 1 (GPU1): Layers[4:8] ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
↓
阶段 2 (GPU2): Layers[8:12] ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
微批调度时序图:
GPU0: F0 F1 F2 F3 B3 B2 B1 B0
GPU1: F0 F1 F2 F3 B3 B2 B1 B0
GPU2: F0 F1 F2 F3 B3 B2 B1 B0
其中 F 表示前向,B 表示反向,数字表示微批次编号
这种调度方式的优势在于减少了"气泡"(idle time),提高了设备利用率。编译器需要计算最优的微批次数量:
最优微批次数 = min(
内存允许的最大值,
4 * 流水线阶段数 # 经验值
)
- 激活重计算:在内存和计算之间权衡
流水线并行中,每个阶段需要保存激活值直到反向传播。对于深度网络,这会消耗大量内存。选择性激活重计算可以减少内存占用:
class PipelineStage(nn.Module):
def forward(self, x):
# 编译器自动插入检查点
if self.training and self.use_checkpoint:
# 只保存输入,中间激活值在反向时重算
return checkpoint(self.layers, x)
else:
return self.layers(x)
9.1.2 编译时的通信优化
分布式推理的性能瓶颈往往在于设备间通信。在自动驾驶场景中,多个摄像头的数据需要在不同 GPU 间同步和融合,通信延迟直接影响系统的实时性。现代编译器通过多种策略来优化通信开销。
通信与计算重叠
编译器的核心优化思想是将通信操作与计算操作并行执行,隐藏通信延迟。这需要仔细分析数据依赖关系,找出可以提前启动的通信操作:
# 编译器生成的优化代码模式
def optimized_forward(x, layer1, layer2):
# 分析数据流,识别可以异步的通信
# 异步启动通信
handle = dist.all_reduce(x, async_op=True)
# 在通信进行时执行与通信结果无关的本地计算
local_out = layer1(x)
# 在真正需要通信结果前等待
handle.wait()
# 使用通信结果继续计算
return layer2(local_out + x) # x 已经完成 all_reduce
更复杂的场景下,编译器会构建通信-计算依赖图,通过拓扑排序找出最优的执行顺序:
依赖图分析:
节点 A (计算) ─→ 节点 B (通信) ─→ 节点 D (计算)
↓ ↑
节点 C (计算) ─────────────────────┘
优化后的执行顺序:
1. 启动 B (异步通信)
2. 执行 A
3. 执行 C
4. 等待 B 完成
5. 执行 D
通信融合与批处理
小消息的通信开销主要来自于延迟而非带宽。编译器通过识别和合并多个小规模通信操作来摊销延迟开销:
优化前(3次通信,3倍延迟):
AllReduce(tensor1, size=1MB) # 延迟 100us
AllReduce(tensor2, size=1MB) # 延迟 100us
AllReduce(tensor3, size=1MB) # 延迟 100us
总时间 = 300us + 传输时间
优化后(1次通信,1倍延迟):
combined = concat([tensor1, tensor2, tensor3])
AllReduce(combined, size=3MB) # 延迟 100us
tensor1, tensor2, tensor3 = split(combined)
总时间 = 100us + 传输时间
编译器还会考虑通信拓扑来优化通信模式。例如,在 NVLink 连接的 GPU 间使用 Ring AllReduce,在 InfiniBand 网络上使用 Tree AllReduce:
class CommunicationOptimizer:
def optimize_allreduce(self, tensors, topology):
if topology == "nvlink" and len(tensors) <= 8:
# NVLink 适合 Ring AllReduce
return self.ring_allreduce(tensors)
elif topology == "infiniband":
# InfiniBand 适合 Tree AllReduce
return self.tree_allreduce(tensors)
else:
# 默认使用 NCCL 的自动选择
return self.nccl_allreduce(tensors)
梯度累积与延迟通信
在推理场景中虽然没有梯度,但在在线学习或少样本适应的场景下,编译器仍需要优化梯度通信:
# 编译器生成的梯度累积代码
def gradient_accumulation_step(model, data_batch, accumulation_steps=4):
accumulated_grads = {}
for step in range(accumulation_steps):
# 前向和反向传播
loss = model(data_batch[step])
loss.backward()
# 累积梯度而不立即通信
for name, param in model.named_parameters():
if name not in accumulated_grads:
accumulated_grads[name] = param.grad.clone()
else:
accumulated_grads[name] += param.grad
# 批量通信所有梯度
handles = []
for name, grad in accumulated_grads.items():
handle = dist.all_reduce(grad, async_op=True)
handles.append(handle)
# 等待所有通信完成
for handle in handles:
handle.wait()
9.1.3 分布式推理的内存优化
在分布式环境下,内存优化变得更加复杂。每个设备不仅要管理自己的参数和激活值,还要考虑通信缓冲区和同步开销。
激活检查点(Activation Checkpointing)的分布式策略
激活检查点技术通过丢弃中间激活值并在反向传播时重新计算来节省内存。在分布式环境下,检查点的设置需要考虑通信边界:
@torch.compile(options={
"distributed": True,
"activation_checkpoint": "selective"
})
def distributed_model(x):
# 策略1:在通信密集点设置检查点
# 减少需要保存的分布式激活值
x = checkpoint(layer1, x) # 通信前检查点
x = all_reduce(x) # 通信操作
x = layer2(x) # 不设置检查点,避免重复通信
# 策略2:在内存峰值点设置检查点
if x.numel() > memory_threshold:
x = checkpoint(heavy_layer, x)
return x
编译器通过分析激活值的生命周期和通信模式,自动确定最优的检查点位置:
激活值生命周期分析:
Layer1 output: [创建 t=0] ──────→ [使用 t=10] → [释放]
Layer2 output: [创建 t=2] ──────────────→ [使用 t=15] → [释放]
Layer3 output: [创建 t=5] ───→ [使用 t=8] → [释放]
内存峰值出现在 t=5-8,应在 Layer2 后设置检查点
零冗余优化器(ZeRO)集成
ZeRO 技术通过将模型状态(参数、梯度、优化器状态)分片到多个设备来减少每个设备的内存占用。编译器需要生成相应的通信和计算代码:
ZeRO-1: 优化器状态分片
- 每个 rank 只保存 1/N 的优化器状态
- 内存节省: (optimizer_states) * (N-1)/N
ZeRO-2: 优化器状态 + 梯度分片
- 每个 rank 只保存 1/N 的梯度
- 梯度在 AllReduce 后立即分片
- 内存节省: (optimizer_states + gradients) * (N-1)/N
ZeRO-3: 优化器状态 + 梯度 + 参数分片
- 每个 rank 只保存 1/N 的模型参数
- 前向/反向时按需获取参数
- 内存节省: (params + gradients + optimizer_states) * (N-1)/N
编译器生成的 ZeRO-3 执行流程:
class ZeRO3Executor:
def forward(self, layer_id, input):
# 1. 收集该层参数(如果不在本地)
params = self.gather_parameters(layer_id)
# 2. 执行前向计算
output = functional_forward(input, params)
# 3. 立即释放非本地参数
if not self.owns_parameters(layer_id):
del params # 释放内存
return output
def backward(self, layer_id, grad_output):
# 1. 重新收集参数
params = self.gather_parameters(layer_id)
# 2. 计算梯度
grad_input, grad_params = functional_backward(
grad_output, params
)
# 3. 规约并分片梯度
self.reduce_scatter_gradients(grad_params)
return grad_input
内存池与预分配策略
分布式环境下的内存碎片化问题更加严重。编译器通过内存池管理来优化:
class DistributedMemoryPool:
def __init__(self, device_id, pool_size):
# 预分配大块连续内存
self.pool = torch.cuda.memory_pool(device_id)
self.pool.set_size(pool_size)
# 为不同用途划分内存区域
self.regions = {
'parameters': (0, pool_size * 0.4),
'activations': (pool_size * 0.4, pool_size * 0.7),
'comm_buffer': (pool_size * 0.7, pool_size * 0.9),
'workspace': (pool_size * 0.9, pool_size)
}
def allocate(self, size, purpose='activations'):
# 从对应区域分配内存
region_start, region_end = self.regions[purpose]
return self.pool.allocate(size, region_start)
9.2 动态批处理与序列并行
9.2.1 动态批处理的编译挑战
在自动驾驶的实时感知系统中,输入批次大小和序列长度经常变化。传统的静态编译无法有效处理这种动态性。
Continuous Batching 技术
连续批处理允许不同长度的序列在同一批次中处理:
传统批处理(填充到最大长度):
Batch = [
[token1, token2, token3, PAD, PAD],
[token1, token2, PAD, PAD, PAD],
[token1, token2, token3, token4, token5]
]
连续批处理(动态打包):
Active_tokens = [
(seq0, pos0), (seq0, pos1), (seq0, pos2),
(seq1, pos0), (seq1, pos1),
(seq2, pos0), (seq2, pos1), (seq2, pos2), (seq2, pos3), (seq2, pos4)
]
编译器需要生成支持动态索引的核函数:
@torch.compile(dynamic=True)
def continuous_attention(queries, keys, values, seq_lens):
# 编译器生成变长序列的注意力计算
# 使用 CUDA 图捕获不同长度的执行路径
outputs = []
for i, seq_len in enumerate(seq_lens):
q = queries[i, :seq_len]
k = keys[i, :seq_len]
v = values[i, :seq_len]
out = scaled_dot_product_attention(q, k, v)
outputs.append(out)
return outputs
9.2.2 PagedAttention 与内存管理
PagedAttention 将 KV 缓存组织成固定大小的页面,实现细粒度的内存管理:
KV 缓存的页表结构:
┌─────────────┬─────────────┬─────────────┐
│ Page 0 │ Page 1 │ Page 2 │
│ (32 tokens)│ (32 tokens)│ (32 tokens)│
└─────────────┴─────────────┴─────────────┘
↑ ↑ ↑
Seq 0 Seq 0/1 Seq 1
页表映射:
Seq 0: [Page0, Page1_partial]
Seq 1: [Page1_partial, Page2]
编译器优化策略:
class PagedAttentionCompiler:
def compile_attention(self, q, k_cache, v_cache, page_table):
# 生成分页访问的 CUDA 核函数
kernel = """
__global__ void paged_attention_kernel(
float* q, float* k_pages, float* v_pages,
int* page_table, int* page_sizes
) {
int seq_id = blockIdx.x;
int head_id = blockIdx.y;
// 通过页表访问 KV 缓存
for (int page_idx : page_table[seq_id]) {
float* k_page = k_pages + page_idx * PAGE_SIZE;
float* v_page = v_pages + page_idx * PAGE_SIZE;
// 计算注意力...
}
}
"""
return compile_cuda(kernel)
9.2.3 序列并行的优化策略
对于超长序列(如自动驾驶的长时间轨迹预测),序列并行将序列维度分布到多个设备:
序列并行的注意力计算:
设备 0: seq[0:L/2] 设备 1: seq[L/2:L]
Q₀ ─────┐ Q₁ ─────┐
K₀ ─────┼─ AllGather → K ─────┤
V₀ ─────┘ V₁ ─────┘
↓ ↓
Attn₀ Attn₁
↓ ↓
Output₀ Output₁
Ring Attention 优化:
def ring_attention(q_local, k_local, v_local, rank, world_size):
"""
环形注意力:每个设备只持有部分 KV,通过环形传递计算完整注意力
"""
output = torch.zeros_like(q_local)
for step in range(world_size):
# 计算当前 KV 块的注意力
attn_weights = torch.matmul(q_local, k_local.transpose(-2, -1))
attn_output = torch.matmul(attn_weights, v_local)
output += attn_output
# 环形传递 KV 到下一个 rank
k_local = ring_send_recv(k_local, rank, world_size)
v_local = ring_send_recv(v_local, rank, world_size)
return output
9.3 硬件加速器适配
9.3.1 NPU/TPU 的编译后端
不同的硬件加速器有其独特的架构特点,编译器需要生成相应的优化代码。
华为昇腾 NPU 适配
昇腾 NPU 使用 CANN 计算架构,编译器需要:
# PyTorch 到 CANN 图的转换
class AscendCompiler:
def compile(self, torch_model):
# 1. 图捕获与分析
graph = torch.fx.symbolic_trace(torch_model)
# 2. 算子映射
cann_graph = self.map_to_cann_ops(graph)
# 3. Cube 单元优化(矩阵计算单元)
self.optimize_cube_ops(cann_graph)
# 4. Vector 单元优化(向量计算单元)
self.optimize_vector_ops(cann_graph)
# 5. 数据排布优化(NCHW -> NC1HWC0)
self.optimize_data_layout(cann_graph)
return cann_graph
Google TPU 适配
TPU 的 systolic array 架构需要特殊的编译优化:
TPU v4 架构特点:
- 128x128 systolic array
- HBM2e 内存
- 高带宽互联
编译优化要点:
1. 矩阵分块以匹配 systolic array 大小
2. XLA 编译器集成
3. bfloat16 自动混合精度
@torch_xla.compile()
def tpu_optimized_model(x):
# XLA 编译器自动优化:
# - 算子融合
# - 内存布局优化
# - systolic array 利用率最大化
with torch_xla.amp.autocast():
return model(x)
9.3.2 自定义 ASIC 的图映射
对于专用 AI 芯片(如 Tesla Dojo、Waymo TPU),需要自定义编译后端:
class CustomASICCompiler:
def __init__(self, hw_spec):
self.hw_spec = hw_spec # 硬件规格
self.op_library = self.load_op_library() # 算子库
def compile(self, model):
# 1. 图级优化
graph = self.capture_graph(model)
graph = self.fuse_operators(graph)
# 2. 算子选择
for node in graph.nodes():
if node.op in self.op_library:
# 使用硬件原生算子
node.impl = self.op_library[node.op]
else:
# 分解为基础算子
node.impl = self.decompose(node.op)
# 3. 内存分配
memory_plan = self.allocate_memory(graph)
# 4. 指令生成
instructions = self.generate_instructions(graph, memory_plan)
return ASICExecutable(instructions)
9.3.3 异构计算的调度策略
在自动驾驶系统中,经常需要协调 CPU、GPU、NPU 等多种计算资源:
异构执行示例:
┌─────────────┐ ┌─────────────┐ ┌─────────────┐
│ CPU │ │ GPU │ │ NPU │
│ 预处理 │ --> │ 主干网络 │ --> │ 后处理 │
│ 数据加载 │ │ 特征提取 │ │ NMS │
└─────────────┘ └─────────────┘ └─────────────┘
编译器的异构调度优化:
class HeterogeneousScheduler:
def schedule(self, graph, devices):
# 1. 设备能力建模
device_profiles = self.profile_devices(devices)
# 2. 算子放置决策
placement = {}
for node in graph.nodes():
# 基于算子特性和设备能力匹配
best_device = self.find_best_device(
node, device_profiles
)
placement[node] = best_device
# 3. 插入数据传输
self.insert_transfers(graph, placement)
# 4. 流水线优化
self.optimize_pipeline(graph, placement)
return placement
9.4 编译器技术的未来发展
9.4.1 AI 编译器的发展趋势
- 自适应编译
未来的编译器将具备运行时自适应能力:
class AdaptiveCompiler:
def __init__(self):
self.perf_history = {}
self.compilation_cache = {}
def compile_adaptive(self, model, input_shape):
# 基于历史性能数据选择编译策略
if input_shape in self.perf_history:
best_config = self.perf_history[input_shape]
else:
# 在线搜索最优配置
best_config = self.auto_tune(model, input_shape)
# 增量编译
if self.should_recompile(best_config):
compiled = self.recompile(model, best_config)
self.compilation_cache[input_shape] = compiled
return self.compilation_cache[input_shape]
- 神经网络辅助编译
使用机器学习模型预测最优编译策略:
编译决策网络:
输入特征:[图结构, 硬件特性, 数据分布]
↓
特征提取
↓
决策网络
↓
输出:[算子融合策略, 内存布局, 并行策略]
9.4.2 跨框架的统一 IR
MLIR(Multi-Level Intermediate Representation)
MLIR 提供了多层次的中间表示,支持渐进式降低:
PyTorch Model
↓
Torch Dialect ←── 高层抽象
↓
Linalg Dialect ←── 线性代数操作
↓
Affine Dialect ←── 循环优化
↓
LLVM Dialect ←── 底层代码生成
↓
Target Code
统一 IR 的优势:
- 复用优化 passes
- 跨框架模型迁移
- 硬件厂商统一适配
9.4.3 量子-经典混合计算
随着量子计算的发展,编译器需要支持量子-经典混合编程:
class QuantumHybridCompiler:
def compile_hybrid(self, classical_model, quantum_circuit):
# 1. 识别量子加速机会
quantum_ops = self.identify_quantum_advantage(
classical_model
)
# 2. 量子电路优化
optimized_circuit = self.optimize_quantum_circuit(
quantum_circuit
)
# 3. 经典-量子接口生成
interface = self.generate_interface(
classical_model, optimized_circuit
)
# 4. 错误缓解策略
error_mitigation = self.add_error_mitigation(
optimized_circuit
)
return HybridExecutable(
classical_model, optimized_circuit, interface
)
9.4.4 编译器的可解释性
未来的编译器将提供更好的可解释性和调试能力:
class ExplainableCompiler:
def compile_with_explanation(self, model):
decisions = []
# 记录每个优化决策
with self.decision_tracker() as tracker:
# 算子融合决策
fusion_plan = self.plan_fusion(model)
tracker.record("fusion", fusion_plan, self.explain_fusion)
# 内存规划决策
memory_plan = self.plan_memory(model)
tracker.record("memory", memory_plan, self.explain_memory)
# 并行策略决策
parallel_plan = self.plan_parallelism(model)
tracker.record("parallel", parallel_plan, self.explain_parallel)
# 生成可视化报告
report = self.generate_report(decisions)
return compiled_model, report
本章小结
本章深入探讨了 PyTorch 编译技术的高级主题和前沿发展方向。我们学习了四个核心领域:
关键概念回顾
-
分布式推理与模型并行 - 张量并行通过切分算子计算实现设备间并行 - 流水线并行通过阶段划分和微批调度优化吞吐量 - 编译器优化通信与计算重叠、通信融合等关键技术 - 激活检查点和 ZeRO 技术降低内存占用
-
动态批处理与序列并行 - Continuous Batching 消除填充开销,提高 GPU 利用率 - PagedAttention 通过页表管理实现细粒度内存控制 - Ring Attention 支持超长序列的分布式处理 - 编译器需要生成支持动态形状的高效核函数
-
硬件加速器适配 - NPU/TPU 等专用硬件需要定制编译后端 - 自定义 ASIC 要求完整的图映射和指令生成 - 异构计算需要智能的设备调度和数据传输优化 - 硬件感知的编译优化是性能关键
-
编译器技术的未来 - 自适应编译根据运行时反馈动态调整策略 - 神经网络辅助编译使用 ML 预测最优配置 - MLIR 等统一 IR 实现跨框架优化复用 - 量子-经典混合计算带来新的编译挑战
核心公式总结
- 模型并行的通信复杂度:
通信量 = O(N × H × P)
其中 N 是序列长度,H 是隐藏维度,P 是并行度
- 动态批处理的内存节省:
内存节省率 = 1 - (实际 token 数 / (批大小 × 最大长度))
- Ring Attention 的通信步数:
通信轮数 = P (设备数)
每轮通信量 = 2 × N/P × H (KV 的大小)
实践要点
在应用这些高级技术时,需要注意:
- 性能建模:准确评估不同并行策略的收益和开销
- 动态适应:根据实际负载动态调整批处理和并行策略
- 硬件特性:深入理解目标硬件的架构特点
- 端到端优化:综合考虑计算、内存、通信的整体优化