RL Infra 系统设计:GPU硬件×网络拓扑×模型结构→并行策略×资源配比×流水线
从硬件约束出发,推导RL训练的最优并行策略、资源分配和流水线设计
RL Infra 系统设计:最大化RL训练效率
一、核心问题框架
1.1 RL训练效率的决定因素
RL训练效率 = 有效token吞吐量 = f(GPU硬件, 网络拓扑, 模型结构, 并行策略, 资源配比, 流水线)
目标函数:
其中:
- = 并行策略 (TP/PP/DP/EP/SP)
- = 资源配比 (生成GPU数 vs 训练GPU数)
- = 流水线调度 (同步/异步/重叠)
约束条件:
1.2 RL训练的特殊性
与预训练/SFT不同,RL训练有两个截然不同的计算阶段:
┌─────────────────────────────────────────────────────────────────┐
│ RL Training Loop │
│ │
│ ┌──────────────────┐ ┌──────────────────┐ │
│ │ 阶段1: 生成 │ │ 阶段2: 训练 │ │
│ │ (Rollout) │ → │ (Update) │ │
│ │ │ │ │ │
│ │ 特点: │ │ 特点: │ │
│ │ - 推理密集 │ │ - 训练密集 │ │
│ │ - 自回归生成 │ │ - forward+backward │ │
│ │ - 高带宽需求 │ │ - 高计算需求 │ │
│ │ - 内存用于KV Cache │ │ - 内存用于激活+梯度 │ │
│ │ - batch size大 │ │ - micro batch小 │ │
│ └──────────────────┘ └──────────────────┘ │
│ │
│ 关键问题: 两个阶段对硬件资源的需求完全不同! │
└─────────────────────────────────────────────────────────────────┘
二、硬件约束分析
2.1 GPU硬件参数
主流GPU规格对比
| 参数 | H100 SXM | H200 SXM | A100 SXM | L40S |
|---|---|---|---|---|
| FP16/BF16算力 | 989 TFLOPS | 989 TFLOPS | 312 TFLOPS | 362 TFLOPS |
| FP8算力 | 1979 TFLOPS | 1979 TFLOPS | - | - |
| INT8算力 | 1979 TOPS | 1979 TOPS | 624 TOPS | 724 TOPS |
| 显存容量 | 80 GB HBM3 | 141 GB HBM3e | 80 GB HBM2e | 48 GB GDDR6X |
| 显存带宽 | 3.35 TB/s | 4.8 TB/s | 2.0 TB/s | 864 GB/s |
| NVLink带宽 | 900 GB/s | 900 GB/s | 600 GB/s | - |
| NVLink拓扑 | 8 GPU全互联 | 8 GPU全互联 | 8 GPU全互联 | - |
| PCIe | Gen5 x16 | Gen5 x16 | Gen4 x16 | Gen4 x16 |
| TDP | 700W | 700W | 400W | 350W |
关键硬件约束
1. 显存容量约束
模型参数内存:
BF16: params × 2 bytes
FP16: params × 2 bytes
INT8: params × 1 byte
INT4: params × 0.5 bytes
优化器状态 (Adam):
FP32 master weights: params × 4 bytes
FP32 momentum: params × 4 bytes
FP32 variance: params × 4 bytes
总计: params × 12 bytes
梯度:
BF16/FP16: params × 2 bytes
单模型总内存 (BF16 + Adam):
params × (2 + 12 + 2) = params × 16 bytes
示例 (70B模型):
70B × 16 bytes = 1,120 GB ≈ 14张H100 (仅1个模型)
RL需要3-4个模型 → 需要42-56张H100 (仅参数)
2. 计算约束
训练FLOPs (近似):
C ≈ 6 × params × tokens (forward + backward)
单GPU吞吐 (理论峰值50%利用率):
H100: 989 TFLOPS × 50% = 494.5 TFLOPS
70B模型: 6 × 70B = 420 TFLOPs/token
单GPU吞吐: 494.5T / 420T ≈ 1.18 tokens/GPU/s
实际吞吐 (考虑通信、流水线气泡):
典型: 0.3-0.8 tokens/GPU/s
3. 通信约束
AllReduce通信量 (数据并行):
V = 2 × params × bytes_per_param (Ring AllReduce)
70B BF16: 2 × 70B × 2 = 280 GB
NVLink AllReduce时间 (8 GPU):
T = V / NVLink_bw = 280 GB / 900 GB/s ≈ 0.31 s
跨节点AllReduce时间 (InfiniBand 400Gb/s):
T = V / IB_bw = 280 GB / 50 GB/s ≈ 5.6 s
2.2 网络拓扑
常见集群拓扑
┌─────────────────────────────────────────────────────────────────┐
│ 典型GPU集群拓扑 │
│ │
│ ┌─────────────────────────────────────────────────────┐ │
│ │ InfiniBand / RoCE │ │
│ │ (400-800 Gb/s) │ │
│ └──────────┬──────────┬──────────┬──────────┬─────────┘ │
│ │ │ │ │ │
│ ┌──────┴──────┐ ┌─┴────────┐ ┌────────┴──────┐ │
│ │ Node 1 │ │ Node 2 │ │ Node N │ │
│ │ ┌────────┐ │ │ ┌──────┐ │ │ ┌────────┐ │ │
│ │ │NVSwitch│ │ │ │NVSwitch│ │ │ │NVSwitch│ │ │
│ │ └───┬────┘ │ │ └──┬───┘ │ │ └───┬────┘ │ │
│ │ ┌───┼───┐ │ │ ┌──┼──┐ │ │ ┌───┼───┐ │ │
│ │ │GPU│GPU│ │ │ │GPU│GPU│ │ │ │GPU│GPU│ │ │
│ │ │0 │7 │ │ │ │0 │7 │ │ │ │0 │7 │ │ │
│ │ └───┴───┘ │ │ └──┴──┘ │ │ └───┴───┘ │ │
│ │ 8×H100 │ │ 8×H100 │ │ 8×H100 │ │
│ │ NVLink全互联│ │ NVLink │ │ NVLink │ │
│ └─────────────┘ └──────────┘ └──────────────┘ │
│ │
│ 带宽层级: │
│ - 节点内NVLink: 900 GB/s (双向) │
│ - 节点间IB/RoCE: 50-100 GB/s │
│ - 比率: ~10-18x 差距 │
└─────────────────────────────────────────────────────────────────┘
网络带宽层级
| 层级 | 带宽 | 延迟 | 适用通信 |
|---|---|---|---|
| GPU内SRAM | ~19 TB/s | ~1ns | Tensor Core计算 |
| GPU内HBM | 2-4.8 TB/s | ~10ns | 参数/激活读写 |
| 节点内NVLink | 900 GB/s | ~1μs | TP/PP通信 |
| 节点内PCIe | 64 GB/s | ~2μs | CPU-GPU offload |
| 节点间IB 400G | 50 GB/s | ~2μs | DP AllReduce |
| 节点间RoCE 400G | 50 GB/s | ~5μs | DP AllReduce |
| 节点间IB 800G | 100 GB/s | ~2μs | 未来集群 |
三、并行策略推导
3.1 并行策略选择矩阵
核心原则:通信量大的并行维度放在高带宽层级。
并行策略选择:
┌─────────────────────────────────────────────────────────────┐
│ │
│ 通信带宽: NVLink >> IB/RoCE >> PCIe │
│ (900GB/s) (50-100GB/s) (64GB/s) │
│ │
│ 并行策略: TP/PP → DP/EP → Offload │
│ 通信模式: AllReduce/ AllGather/ PCIe/CPU │
│ P2P ReduceScatter │
│ │
│ 规则: │
│ 1. TP维度 ≤ 节点内GPU数 (8) │
│ 2. PP跨节点时,通信量需可控 │
│ 3. DP跨节点,通信量大但可异步 │
│ 4. EP跨节点,通信量取决于专家数量 │
│ │
└─────────────────────────────────────────────────────────────┘
3.2 各并行策略的通信分析
张量并行 (TP)
通信模式: 每层AllReduce (2次)
通信量: 2 × batch × seq_len × hidden_dim × bytes
通信位置: 节点内NVLink
适用条件:
- hidden_dim大 → 通信量大 → 必须节点内
- 注意力头数 ≥ TP_degree
- 节点内GPU数限制 (通常TP ≤ 8)
示例 (70B, TP=8, 节点内):
每层通信: 2 × 2048 × 8192 × 2 bytes = 64 MB
NVLink时间: 64 MB / 900 GB/s ≈ 0.07 ms
可忽略
流水线并行 (PP)
通信模式: P2P传输激活
通信量: batch × seq_len × hidden_dim × bytes (每micro-batch)
通信位置: 可跨节点
适用条件:
- 层数多,单GPU放不下
- 跨节点时通信量可控
- micro-batch数 ≥ PP_degree × 2 (减少气泡)
气泡率:
bubble_ratio = (PP_degree - 1) / (PP_degree - 1 + micro_batches)
示例 (PP=4, micro_batches=16):
bubble_ratio = 3 / (3 + 16) = 15.8%
数据并行 (DP)
通信模式: AllReduce梯度
通信量: 2 × params × bytes (Ring AllReduce)
通信位置: 可跨节点
适用条件:
- 模型能放进单GPU (或ZeRO分片后)
- batch size需要放大
- 通信可异步 (gradient accumulation)
ZeRO分片:
ZeRO-1: 分片优化器状态 → 内存减少4x
ZeRO-2: +分片梯度 → 内存减少8x
ZeRO-3: +分片参数 → 内存减少N_gpu x
专家并行 (EP)
通信模式: AlltoAll
通信量: token数 × hidden_dim × bytes × 2 (dispatch + combine)
通信位置: 可跨节点
适用条件:
- MoE模型
- 专家数 > GPU数
- 路由稀疏 → 通信量可控
示例 (DeepSeek-V2, 160专家, EP=64):
每token通信: 2 × 5120 × 2 bytes = 20 KB
批量通信: 20 KB × 2048 × 2048 ≈ 80 GB
序列并行 (SP)
通信模式: AllGather/ReduceScatter
通信量: batch × seq_len/SP_degree × hidden_dim × bytes
通信位置: 节点内NVLink
适用条件:
- 长序列 (seq_len > 32K)
- 与TP结合使用
- 减少激活内存
3.3 并行策略推导算法
def derive_parallelism(
model_params, # 模型参数量 (e.g., 70B)
num_layers, # 层数
hidden_dim, # 隐藏维度
num_heads, # 注意力头数
num_experts, # 专家数 (MoE)
seq_len, # 序列长度
gpus_per_node, # 节点内GPU数
gpu_memory, # GPU显存 (GB)
gpu_compute, # GPU算力 (TFLOPS)
nvlink_bw, # NVLink带宽 (GB/s)
ib_bw, # InfiniBand带宽 (GB/s)
total_gpus, # 总GPU数
):
# Step 1: 确定TP (通信密集,必须节点内)
tp_max = min(gpus_per_node, num_heads)
# 选择TP使得单层参数能放进单GPU
params_per_layer = model_params / num_layers
bytes_per_param = 2 # BF16
tp_candidates = [1, 2, 4, 8]
tp = 1
for t in tp_candidates:
if t <= tp_max:
layer_mem = params_per_layer / t * bytes_per_param
if layer_mem < gpu_memory * 0.3: # 留30%给激活
tp = t
# Step 2: 确定PP
gpus_per_tp_group = tp
gpus_for_pp = total_gpus // gpus_per_tp_group
# PP尽量在节点内,跨节点通信量小
pp_max_per_node = gpus_per_node // tp
if gpus_for_pp <= pp_max_per_node:
pp = gpus_for_pp
else:
pp = pp_max_per_node # 节点内PP
# 跨节点用DP
pp = min(pp, num_layers // 4) # 每PP stage至少4层
# Step 3: 确定DP
gpus_per_pp_pipeline = tp * pp
dp = total_gpus // gpus_per_pp_pipeline
# Step 4: 确定EP (如果是MoE)
ep = 1
if num_experts > 1:
# EP通常等于节点数或GPU数
ep = min(dp, num_experts)
# Step 5: 确定SP (长序列)
sp = 1
if seq_len > 32768:
sp = min(tp, seq_len // 8192)
return {
'tp': tp,
'pp': pp,
'dp': dp,
'ep': ep,
'sp': sp,
'total': tp * pp * dp
}
3.4 典型配置示例
70B Dense模型 + 64×H100
模型: 70B Dense, 80层, hidden=8192, 64头
硬件: 8节点 × 8×H100 = 64 GPU
网络: NVLink (节点内) + IB 400G (节点间)
并行策略:
TP = 8 (节点内全互联,每层参数70B/80/8=109M,可接受)
PP = 2 (跨2个节点,每stage 40层)
DP = 4 (4个数据并行组)
总计: 8 × 2 × 4 = 64 ✓
通信分析:
TP通信 (NVLink): 64 MB/层 × 80层 = 5.1 GB → 5.1/900 = 5.7 ms
PP通信 (IB): 2048 × 8192 × 2 = 32 MB/micro-batch → 32/50 = 0.64 ms
DP通信 (IB): 2 × 70B × 2 = 280 GB → 280/50 = 5.6 s (可异步)
内存分析:
参数: 70B × 2 / (8×2) = 8.75 GB/GPU (TP=8, PP=2)
优化器: 70B × 12 / (8×2×4) = 2.63 GB/GPU (ZeRO-1)
激活: ~20-30 GB/GPU (gradient checkpointing)
总计: ~35-45 GB/GPU (80GB可用)
236B MoE模型 (21B激活) + 128×H100
模型: DeepSeek-V2, 60层, hidden=5120, 160专家, 激活21B
硬件: 16节点 × 8×H100 = 128 GPU
网络: NVLink + IB 400G
并行策略:
TP = 4 (节点内,减少通信)
PP = 4 (跨节点,每stage 15层)
EP = 8 (8路专家并行,每GPU 20专家)
DP = 1 (128 / (4×4×8) = 1)
总计: 4 × 4 × 8 × 1 = 128 ✓
通信分析:
TP通信 (NVLink): 每层32 MB × 60层 = 1.9 GB → 2.1 ms
PP通信 (IB): 32 MB/micro-batch → 0.64 ms
EP通信 (IB): AlltoAll, ~10-50 ms/层 (取决于路由)
DP通信: 无 (DP=1)
内存分析:
Dense参数: 21B × 2 / (4×4) = 2.6 GB/GPU
Expert参数: 160 × (expert_size) / 8 = ~20 GB/GPU
优化器: ~15 GB/GPU
总计: ~40-50 GB/GPU
1.5B小模型 + 8×H100 (纯RL训练)
模型: 1.5B Dense, 24层, hidden=2048, 16头
硬件: 1节点 × 8×H100
场景: PPO训练 (Policy + Reference + Value + Reward)
并行策略:
TP = 1 (模型小,不需要)
PP = 1 (单GPU放得下)
DP = 8 (8路数据并行)
总计: 1 × 1 × 8 = 8 ✓
资源分配:
Policy模型: 1.5B × 16 = 24 GB (参数+优化器+梯度)
Reference模型: 1.5B × 2 = 3 GB (仅参数,offload到CPU)
Value模型: 1.5B × 16 = 24 GB (与Policy共享优化器)
Reward模型: 1.5B × 2 = 3 GB (仅参数,可量化)
总计: ~54 GB/GPU (可用80GB)
优化:
- Reference模型offload到CPU
- Reward模型INT8量化
- 生成和训练重叠
四、RL流水线资源配比
4.1 RL流水线阶段分析
┌─────────────────────────────────────────────────────────────────┐
│ RL训练完整流水线 │
│ │
│ 时间 → │
│ ┌──────┐ ┌──────┐ ┌──────┐ ┌──────┐ ┌──────┐ ┌──────┐ │
│ │Prefill│→│Decode │→│Reward│→│Advant│→│Train │→│Sync │ │
│ │(生成) │ │(生成) │ │计算 │ │计算 │ │(更新)│ │(同步)│ │
│ └──────┘ └──────┘ └──────┘ └──────┘ └──────┘ └──────┘ │
│ │
│ 阶段 计算特点 GPU利用率 瓶颈 │
│ ───────────────────────────────────────────────────────────── │
│ Prefill 计算密集 高 无 │
│ Decode 内存密集(逐token) 低 KV Cache带宽 │
│ Reward 推理密集 中 批量大小 │
│ Advantage CPU/GPU混合 低 Value模型推理 │
│ Train 计算密集 高 梯度通信 │
│ Sync 通信密集 - 网络带宽 │
│ │
│ 关键瓶颈: Decode阶段GPU利用率最低 (~10-30%) │
└─────────────────────────────────────────────────────────────────┘
4.2 生成-训练资源配比
核心问题:多少GPU用于生成,多少GPU用于训练?
配比推导
设:
N_gen = 生成GPU数
N_train = 训练GPU数
N_total = N_gen + N_train
T_gen = 生成一批数据的时间
T_train = 训练这批数据的时间
目标: 最大化训练吞吐量
throughput = batch_tokens / max(T_gen, T_train)
最优配比: 使 T_gen ≈ T_train (负载均衡)
N_gen / N_train ≈ T_gen_per_gpu / T_train_per_gpu
典型配比分析
| 场景 | 模型规模 | 生成特点 | 推荐配比 | 原因 |
|---|---|---|---|---|
| 小模型+短序列 | 1-7B | 快速生成 | 2:8 | 训练是瓶颈 |
| 中模型+中序列 | 7-70B | 中等生成 | 4:6 | 均衡 |
| 大模型+长序列 | 70B+ | 慢速生成 | 6:4 | 生成是瓶颈 |
| 推理密集型 | 70B+ | 长CoT | 7:3 | 生成占主导 |
动态配比优化
def optimize_resource_ratio(
model_params, # 模型参数量
seq_len, # 序列长度
num_gen_per_prompt, # 每prompt生成数
gpu_compute, # GPU算力
gpu_memory, # GPU显存
total_gpus, # 总GPU数
):
# 估算生成时间
# Prefill: 1次forward (计算密集)
prefill_flops = 6 * model_params * seq_len
prefill_time = prefill_flops / (gpu_compute * 0.5)
# Decode: seq_len次forward (内存密集)
decode_flops_per_token = 2 * model_params # 仅forward
decode_time_per_token = decode_flops_per_token / (gpu_compute * 0.1) # 利用率低
decode_time = decode_time_per_token * seq_len * num_gen_per_prompt
total_gen_time = prefill_time + decode_time
# 估算训练时间
train_flops = 6 * model_params * seq_len * num_gen_per_prompt
train_time = train_flops / (gpu_compute * 0.4) # 训练利用率40%
# 计算最优配比
gen_ratio = total_gen_time / (total_gen_time + train_time)
N_gen = int(total_gpus * gen_ratio)
N_train = total_gpus - N_gen
# 约束: N_gen 和 N_train 必须是合法的并行度
N_gen = round_to_valid_parallelism(N_gen)
N_train = round_to_valid_parallelism(N_train)
return N_gen, N_train
4.3 流水线重叠策略
策略1: 串行执行(简单但低效)
GPU: [=====生成=====][=====训练=====][=====生成=====]...
利用率: 50% (一半时间GPU空闲)
策略2: 异步重叠(推荐)
生成GPU: [=====生成=====][=====生成=====][=====生成=====]...
训练GPU: [=====训练=====][=====训练=====][=====训练=====]...
↑ ↑
数据传输 数据传输
利用率: ~90% (少量数据传输开销)
策略3: 流水线重叠(高级)
GPU组1: [生成batch1][训练batch1][生成batch3]...
GPU组2: [生成batch2][训练batch2]...
GPU组3: [生成batch3][训练batch3]...
利用率: ~95% (完美重叠)
策略4: 生成-训练混合(最新)
同一GPU: [prefill][decode][train][prefill][decode][train]...
←─── 生成 ───→←─ 训练─→
利用率: ~80% (切换开销)
优势: 无需专用生成GPU,资源灵活
五、RL训练完整流水线设计
5.1 端到端流水线架构
┌─────────────────────────────────────────────────────────────────┐
│ RL训练完整流水线架构 │
│ │
│ ┌───────────────────────────────────────────────────────────┐ │
│ │ 调度器 (Scheduler) │ │
│ │ - 动态分配GPU给生成/训练 │ │
│ │ - 监控GPU利用率 │ │
│ │ - 负载均衡 │ │
│ └───────────────────────┬───────────────────────────────────┘ │
│ │ │
│ ┌────────────────┼────────────────┐ │
│ ▼ ▼ ▼ │
│ ┌───────────────┐ ┌───────────────┐ ┌───────────────┐ │
│ │ 生成集群 │ │ 训练集群 │ │ 评估集群 │ │
│ │ │ │ │ │ │ │
│ │ ┌───────────┐ │ │ ┌───────────┐ │ │ ┌───────────┐ │ │
│ │ │vLLM引擎 │ │ │ │训练框架 │ │ │ │评估框架 │ │ │
│ │ │- TP并行 │ │ │ │- 3D并行 │ │ │ │- 批量推理 │ │ │
│ │ │- PagedAttn │ │ │ │- ZeRO-3 │ │ │ │- 奖励计算 │ │ │
│ │ │- 连续批处理│ │ │ │- 梯度累积 │ │ │ │ │ │ │
│ │ └───────────┘ │ │ └───────────┘ │ │ └───────────┘ │ │
│ └───────┬───────┘ └───────┬───────┘ └───────┬───────┘ │
│ │ │ │ │
│ └─────────────────┼─────────────────┘ │
│ │ │
│ ┌───────┴───────┐ │
│ │ 经验缓冲 │ │
│ │ (Replay) │ │
│ └───────────────┘ │
└─────────────────────────────────────────────────────────────────┘
5.2 单步RL训练详细流程
class RLTrainingPipeline:
def __init__(self, config):
# 模型
self.policy = load_model(config.policy_path) # 生成+训练
self.reference = load_model(config.ref_path) # 仅推理
self.reward = load_model(config.reward_path) # 仅推理
self.value = load_model(config.value_path) # 仅推理 (PPO)
# 推理引擎
self.gen_engine = vLLMEngine(self.policy) # 生成引擎
self.reward_engine = vLLMEngine(self.reward) # 奖励引擎
# 训练框架
self.trainer = DeepSpeedTrainer(self.policy) # 训练器
def step(self, prompts):
# ============ 阶段1: 生成 ============
# 使用vLLM生成response
with self.gen_engine:
responses = self.gen_engine.generate(
prompts,
n=self.num_gen_per_prompt, # 每prompt生成K个
max_tokens=self.max_seq_len,
temperature=self.temperature,
)
# ============ 阶段2: 奖励计算 ============
# 使用Reward模型计算奖励
with self.reward_engine:
rewards = self.reward_engine.evaluate(responses)
# ============ 阶段3: 优势计算 ============
# 计算log probabilities
old_logprobs = self.policy.compute_logprobs(responses)
ref_logprobs = self.reference.compute_logprobs(responses)
# 计算KL惩罚
kl_penalty = old_logprobs - ref_logprobs
# 计算优势 (GRPO: group relative)
if self.algorithm == 'GRPO':
advantages = compute_grpo_advantages(
rewards, kl_penalty, self.kl_coef
)
elif self.algorithm == 'PPO':
values = self.value.compute_values(responses)
advantages = compute_gae_advantages(
rewards, values, kl_penalty, self.gamma, self.lam
)
# ============ 阶段4: 训练 ============
# 模型同步 (vLLM → 训练框架)
sync_model_weights(self.gen_engine, self.trainer)
# 梯度更新
for epoch in range(self.num_epochs):
for batch in create_batches(responses, advantages):
loss = self.trainer.train_step(batch)
self.trainer.step()
# ============ 阶段5: 同步 ============
# 训练框架 → vLLM 同步模型
sync_model_weights(self.trainer, self.gen_engine)
return {
'reward': rewards.mean(),
'kl': kl_penalty.mean(),
'loss': loss,
}
5.3 关键优化点
优化1: 生成-训练重叠
# 异步生成-训练重叠
class AsyncRLPipeline:
def __init__(self):
self.gen_queue = Queue() # 生成队列
self.train_queue = Queue() # 训练队列
async def run(self):
# 并行运行生成和训练
await asyncio.gather(
self.generation_loop(),
self.training_loop(),
)
async def generation_loop(self):
while True:
prompts = await self.get_prompts()
# 异步生成
responses = await self.gen_engine.async_generate(prompts)
await self.train_queue.put(responses)
async def training_loop(self):
while True:
responses = await self.train_queue.get()
# 异步训练
await self.trainer.async_train(responses)
优化2: 动态批处理
# 生成阶段: 大batch,充分利用GPU
gen_batch_size = 256 # 同时生成256个response
# 训练阶段: 小batch,梯度累积
train_micro_batch = 4
gradient_accumulation = 64
effective_batch = 4 * 64 = 256 # 与生成batch对齐
优化3: 模型权重同步优化
# 问题: 生成引擎和训练框架的模型格式不同
# vLLM: 连续内存,PagedAttention
# DeepSpeed: ZeRO分片
# 方案1: 直接同步参数 (简单但慢)
def sync_model_weights(src, dst):
for name, param in src.named_parameters():
dst.get_parameter(name).data.copy_(param.data)
# 方案2: 使用共享内存 (快但复杂)
# 两个引擎共享同一块GPU内存
# 方案3: 增量同步 (仅同步变化的参数)
# 适用于LoRA训练
def sync_lora_weights(src, dst):
for name, param in src.named_parameters():
if 'lora' in name:
dst.get_parameter(name).data.copy_(param.data)
六、RL Infra 关键指标与监控
6.1 核心性能指标
| 指标 | 定义 | 目标值 | 影响因素 |
|---|---|---|---|
| Token吞吐量 | tokens/sec/GPU | >1000 | 并行策略、batch size |
| GPU利用率 | 计算时间/总时间 | >80% | 流水线重叠、批处理 |
| 生成效率 | 生成tokens/sec | >5000 | 推理引擎、batch size |
| 训练效率 | 训练tokens/sec | >2000 | 并行策略、梯度累积 |
| 生成-训练比 | 生成时间/训练时间 | ~1.0 | 资源配比 |
| 通信开销 | 通信时间/总时间 | <10% | 网络带宽、并行策略 |
| 内存利用率 | 已用显存/总显存 | 70-85% | 模型大小、batch size |
| 经验效率 | 有效tokens/生成tokens | >0.8 | 数据过滤、复用率 |
6.2 监控仪表板
┌─────────────────────────────────────────────────────────────────┐
│ RL训练监控仪表板 │
│ │
│ ┌─────────────────┐ ┌─────────────────┐ ┌─────────────────┐ │
│ │ Token吞吐量 │ │ GPU利用率 │ │ 内存使用 │ │
│ │ 1,234 tok/s/gpu │ │ 87.3% │ │ 72.1 GB/80 GB │ │
│ │ ▲ 5.2% │ │ ▲ 2.1% │ │ 90.1% │ │
│ └─────────────────┘ └─────────────────┘ └─────────────────┘ │
│ │
│ ┌─────────────────┐ ┌─────────────────┐ ┌─────────────────┐ │
│ │ 生成-训练比 │ │ 通信开销 │ │ KL散度 │ │
│ │ 0.95 │ │ 6.8% │ │ 0.0234 │ │
│ │ 目标: ~1.0 │ │ 目标: <10% │ │ 目标: <0.1 │ │
│ └─────────────────┘ └─────────────────┘ └─────────────────┘ │
│ │
│ ┌───────────────────────────────────────────────────────────┐ │
│ │ 时间分解 │ │
│ │ Prefill: 12% │ Decode: 45% │ Reward: 8% │ Train: 25% │ Sync: 10%│
│ └───────────────────────────────────────────────────────────┘ │
└─────────────────────────────────────────────────────────────────┘
七、实战配置模板
7.1 小规模: 7B模型 + 8×H100
# 场景: PPO训练,7B模型
cluster:
nodes: 1
gpus_per_node: 8
gpu_type: H100-SXM-80GB
network: NVLink (节点内)
model:
policy: 7B
reference: 7B
value: 7B
reward: 7B
parallelism:
tp: 1
pp: 1
dp: 8
ep: 1
resource_allocation:
generation:
gpus: 4 # 4 GPU用于生成
engine: vLLM
batch_size: 256
max_seq_len: 4096
training:
gpus: 4 # 4 GPU用于训练
framework: DeepSpeed
micro_batch: 4
gradient_accumulation: 16
effective_batch: 512
optimization:
reference_model: offload_to_cpu
reward_model: int8_quantize
value_model: share_with_policy
gradient_checkpointing: true
mixed_precision: bf16
pipeline:
mode: async_overlap # 生成-训练异步重叠
buffer_size: 2 # 2个batch缓冲
7.2 中规模: 70B模型 + 64×H100
# 场景: GRPO训练,70B模型
cluster:
nodes: 8
gpus_per_node: 8
gpu_type: H100-SXM-80GB
network: NVLink (节点内) + IB 400G (节点间)
model:
policy: 70B
reference: 70B
reward: 70B
# GRPO无Value模型
parallelism:
# 生成并行策略
generation:
tp: 8 # 节点内TP
pp: 1
dp: 8 # 8个数据并行组
# 训练并行策略
training:
tp: 8
pp: 1
dp: 8
zero_stage: 3 # ZeRO-3分片
resource_allocation:
generation:
gpus: 32 # 4节点用于生成
engine: vLLM
batch_size: 512
max_seq_len: 8192
num_gen_per_prompt: 16
training:
gpus: 32 # 4节点用于训练
framework: DeepSpeed
micro_batch: 2
gradient_accumulation: 32
effective_batch: 512
optimization:
reference_model: offload_to_cpu # 70B offload
reward_model: int8_quantize # 70B INT8
gradient_checkpointing: true
mixed_precision: bf16
activation_checkpointing: true
pipeline:
mode: async_overlap
buffer_size: 4
sync_frequency: 1 # 每步同步模型
7.3 大规模: 236B MoE + 256×H100
# 场景: GRPO训练,DeepSeek-V2
cluster:
nodes: 32
gpus_per_node: 8
gpu_type: H100-SXM-80GB
network: NVLink + IB 400G
model:
policy: 236B (21B active)
reference: 236B
reward: 236B
parallelism:
# 生成并行策略
generation:
tp: 4
pp: 2
ep: 8
dp: 4
# 训练并行策略
training:
tp: 4
pp: 2
ep: 8
dp: 4
zero_stage: 1 # MoE不适合ZeRO-3
resource_allocation:
generation:
gpus: 128 # 16节点
engine: vLLM
batch_size: 1024
max_seq_len: 16384
training:
gpus: 128 # 16节点
framework: Megatron-DeepSpeed
micro_batch: 1
gradient_accumulation: 64
optimization:
reference_model: offload_to_cpu
reward_model: int4_quantize # 236B INT4
expert_parallelism: 8
communication: async_allreduce
gradient_compression: true
pipeline:
mode: pipeline_parallel
buffer_size: 8
sync_frequency: 4 # 每4步同步
八、总结
8.1 并行策略选择决策树
模型能放进单GPU?
├── 是 → TP=1, PP=1, DP=GPU数
│ (适用于 <10B模型)
└── 否 → 需要模型并行
│
├── 单层参数能放进单GPU?
│ ├── 是 → 选择TP使得通信最小
│ │ TP = min(节点内GPU数, 注意力头数)
│ └── 否 → 需要更大的TP或PP
│
└── 选择PP
├── 节点内能放下? → PP在节点内
└── 否 → PP跨节点,注意通信量
MoE模型?
├── 是 → 添加EP
│ EP = min(节点数, 专家数)
└── 否 → 无EP
长序列 (>32K)?
├── 是 → 添加SP
│ SP = min(TP, seq_len/8K)
└── 否 → 无SP
8.2 资源配比决策树
生成时间 vs 训练时间?
├── 生成 > 训练 (长序列/多response)
│ └── 增加生成GPU: N_gen:N_train = 7:3
├── 训练 > 生成 (短序列/少response)
│ └── 增加训练GPU: N_gen:N_train = 3:7
└── 均衡
└── N_gen:N_train = 5:5 或 4:6
8.3 关键优化优先级
优先级1 (效果最大):
1. 选择合适的推理引擎 (vLLM/SGLang)
2. 生成-训练异步重叠
3. 动态批处理和连续批处理
优先级2 (效果显著):
4. Reference模型offload到CPU
5. Reward模型量化 (INT8/INT4)
6. 梯度检查点
优先级3 (效果中等):
7. 通信优化 (异步AllReduce、梯度压缩)
8. 混合精度训练 (BF16)
9. LoRA训练 (减少可训练参数)
优先级4 (效果较小):
10. 数据过滤和课程学习
11. 经验回放优化
12. 容错和弹性训练