性能调优

Relax 训练吞吐量最大化实践指南。本文提到的所有参数均可在配置参考手册中查阅。

性能分析

调优前先定位瓶颈。Relax 内置三套互补的 profiling 工具，覆盖 推理引擎、训练后端 和 GPU 内存 三个维度。所有 trace 文件默认保存在 traces/<tb_experiment_name>/ 目录下，按子目录区分：

工具	目标	默认输出目录	查看方式
SGLang Profiling	Rollout 推理的 CUDA kernel / 算子分析	`traces/<tb_experiment_name>/sglang_trace/`	TensorBoard or `https://ui.perfetto.dev/`
Training Profiling	Actor 训练 / log-probs 计算的算子分析	`traces/<tb_experiment_name>/train_trace/`	TensorBoard or `https://ui.perfetto.dev/`
Memory Profiling	GPU 显存分配历史，OOM 诊断	`traces/<tb_experiment_name>/memory_snapshot/`	PyTorch Memory Viz

Trace 文件命名规则

训练 trace 文件名包含 rank{global}_dp{dp}_tp{tp}_pp{pp} 标识，例如 train_overall_rank0_dp0_tp0_pp0.1713780123.pt.trace.json.gz
内存快照 文件名同样包含 rank 标识，例如 memory_snapshot_time1713780123_rank0_dp0_tp0_pp0_snapshot.pickle
SGLang trace 文件以 engine{i} 为前缀区分不同引擎实例，例如 engine0-1713780123-TP-0.trace.json.gz

1. SGLang 推理 Profiling

对 Rollout 阶段所有 SGLang 引擎进行 torch.profiler 采集。通过 HTTP API /start_profile 和 /stop_profile 触发，不影响训练侧的 profiler。

示例用法 — 每个 rollout step 都采集：

bash

python3 relax/entrypoints/train.py \
    --sglang-profile \
    --tb-experiment-name my-experiment \
    # ... 其他参数

指定 rollout step 范围 — 仅在 step 2、3、4 采集（start/end 均 inclusive，推荐用法，避免大量 trace 文件）：

bash

python3 relax/entrypoints/train.py \
    --sglang-profile \
    --sglang-profile-step-start 2 \
    --sglang-profile-step-end 4 \
    --tb-experiment-name my-experiment \
    # ... 其他参数

也可以用 --sglang-profile-steps 指定不连续的 step 列表（优先级高于 start/end）：

bash

--sglang-profile-steps 2 5 10

所有 step 参数均使用 绝对 rollout ID（0-indexed），即第 0 轮、第 1 轮 ... 与 --start-rollout-id 无关。

进阶参数：

参数	默认值	说明
`--sglang-profile-step-start`	None	profiling 起始 rollout step（inclusive，0-indexed）
`--sglang-profile-step-end`	None	profiling 结束 rollout step（inclusive，0-indexed）。例如 start=2, end=4 采集 step 2, 3, 4
`--sglang-profile-steps`	None	不连续 step 列表，优先级高于 start/end
`--sglang-profile-num-steps`	3	每轮 Rollout 中采集的 SGLang 前向步数。-1 表示整轮采集
`--sglang-profile-activities`	CPU GPU	要采集的活动类型
`--sglang-profile-by-stage`	False	按 prefill / decode 阶段分别采集
`--sglang-profile-with-stack`	False	记录 Python 调用栈
`--sglang-profile-record-shapes`	False	记录张量形状信息
`--sglang-profile-output-dir`	None	自定义输出目录。默认 `traces/<tb_experiment_name>/sglang_trace`

2. 训练 Profiling（PyTorch Profiler）

对 Actor 训练步骤进行 torch.profiler 采集，生成兼容 TensorBoard 的 trace 文件。

示例用法 — 采集第 2、3、4 步（start/end 均 inclusive）：

bash

python3 relax/entrypoints/train.py \
    --use-pytorch-profiler \
    --profile-step-start 2 \
    --profile-step-end 4 \
    --tb-experiment-name my-experiment \
    # ... 其他参数

TIP

--profile-step-start 和 --profile-step-end 均为 inclusive，是从本次训练启动后的 步数偏移，不是绝对 rollout ID。断点续训时计数从 0 重新开始。例如 start=2, end=4 采集 step 2, 3, 4（共 3 步）。

语义与 --sglang-profile-step-start/end 相同（两端均 inclusive）。

详细信息标志：

标志	作用
`--profile-with-stack`	在每个 trace 事件中记录 Python 调用栈。用于定位触发高开销操作的代码路径
`--profile-with-memory`	在 trace 中跟踪 CUDA 显存分配/释放。用于发现显存尖峰
`--profile-with-flops`	估算每个算子的 FLOPs。用于计算硬件利用率 (MFU)

完整示例：

bash

python3 relax/entrypoints/train.py \
    --use-pytorch-profiler \
    --profile-target train_overall \
    --profile-step-start 2 \
    --profile-step-end 4 \
    --profile-with-stack \
    --profile-with-memory \
    --profile-with-flops \
    --tb-experiment-name my-experiment \
    # ... 其他参数

WARNING

启用 --profile-with-stack 和 --profile-with-memory 会增加额外开销。建议仅在诊断时使用，不用于生产训练。

3. GPU 内存 Profiling

记录 CUDA 显存分配/释放历史，用于诊断显存泄漏和 OOM 问题。在发生 OOM 时会自动 dump 内存快照。

最小用法 — 开启记录 + 在第 2 步后主动 dump：

bash

python3 relax/entrypoints/train.py \
    --record-memory-history \
    --memory-snapshot-num-steps 2 \
    --tb-experiment-name my-experiment \
    # ... 其他参数

进阶参数：

参数	默认值	说明
`--memory-snapshot-path`	snapshot.pickle	快照文件名后缀
`--memory-snapshot-dir`	None	自定义输出目录。默认 `traces/<tb_experiment_name>/memory_snapshot`
`--memory-snapshot-num-steps`	None	在指定步数后主动 dump 快照（0-indexed，设 3 表示第 2 步后 dump）
`--memory-recorder`	torch	后端选择：`torch`（PyTorch 内置）、`memray`（需 `pip install memray`）

查看快照：访问 PyTorch Memory Viz，拖入生成的 .pickle 文件。

三种 Profiling 联合使用

实际诊断中，可同时开启三种 profiling 以获得全面视图。以下是一个完整的联合使用示例：

bash

python3 relax/entrypoints/train.py \
    # --- SGLang 推理 Profiling ---
    --sglang-profile \
    --sglang-profile-step-start 2 \
    --sglang-profile-step-end 4 \
    # --- 训练 Profiling ---
    --use-pytorch-profiler \
    --profile-step-start 2 \
    --profile-step-end 4 \
    # --- 内存 Profiling ---
    --record-memory-history \
    --memory-snapshot-num-steps 2 \
    # --- 实验名（决定 trace 输出目录）---
    --tb-experiment-name my-profiling-run \
    # ... 其他训练参数

上述配置会产出如下目录结构：

traces/my-profiling-run/
├── sglang_trace/                          # SGLang 引擎 trace（按 rollout step 分目录）
│   ├── rollout_2/
│   │   ├── engine0-...-TP-0.trace.json.gz
│   │   ├── engine0-...-TP-1.trace.json.gz
│   │   ├── engine1-...-TP-0.trace.json.gz
│   │   └── ...
│   ├── rollout_3/
│   │   └── ...
│   └── rollout_4/
│       ├── engine0-...-TP-0.trace.json.gz
│       └── ...
├── train_trace/                           # 训练 trace
│   ├── train_overall_rank0_dp0_tp0_pp0.....pt.trace.json.gz
│   ├── train_overall_rank1_dp0_tp1_pp0.....pt.trace.json.gz
│   └── ...
└── memory_snapshot/                       # 内存快照
    ├── memory_snapshot_time..._rank0_dp0_tp0_pp0_snapshot.pickle
    ├── memory_snapshot_time..._rank1_dp0_tp1_pp0_snapshot.pickle
    └── ...

动态批处理

动态批处理根据样本长度动态打包，使每个 micro-batch 的总 Token 数接近目标值，相比固定大小的 micro-batch 能更好地利用 GPU。它同时也是一种有效的 OOM 防护机制——使用固定的 --micro-batch-size 时，一批异常长的序列可能超出 GPU 显存，而动态批处理通过 --max-tokens-per-gpu 限制每个 micro-batch 的总 Token 数，使显存使用可预测。

bash

--use-dynamic-batch-size \
--max-tokens-per-gpu 9216

使用 Context Parallelism (CP) 时，--max-tokens-per-gpu 应设为约 max_response_len / cp_size。

如果计算 log probs 是单独的瓶颈，可以为该阶段设置不同的 Token 预算：

bash

--log-probs-max-tokens-per-gpu 12288

TIP

如果训练中遇到 OOM，将固定的 --micro-batch-size 切换为 --use-dynamic-batch-size 并设置保守的 --max-tokens-per-gpu 通常是第一步。更多详情请参阅 OOM 排查指南。

并行配置

张量并行与序列并行

对于可放入单节点的模型，张量并行 (TP) + 序列并行 (SP) 是最常见的配置：

bash

--tensor-model-parallel-size 2 \
--sequence-parallel

较大模型（30B+）通常使用 TP=2 或 TP=4 并启用 SP。

MoE 专家并行

对于 MoE 模型（如 Qwen3-30B-A3B），将专家分布到多个 GPU：

bash

--expert-model-parallel-size 2 \
--expert-tensor-parallel-size 1

上下文并行

长上下文训练时，将序列拆分到多个 GPU：

bash

--context-parallel-size 2

激活重计算

重计算以计算换内存。对大多数 RL 训练场景，建议启用重计算：

bash

--recompute-granularity full \
--recompute-method uniform \
--recompute-num-layers 1

该配置在所有层上均匀重计算激活。根据显存/计算的权衡调整 --recompute-num-layers。关于 selective 粒度和 block 方法的更多细节，请参考 Megatron-LM 文档。

多模态处理并行化

训练多模态模型（如 Qwen3-VL）时，HuggingFace 的图像/视频数据处理器可能因 Python GIL 成为 CPU 瓶颈。--mm-processor-pool-size 参数创建 ProcessPoolExecutor 来绕过 GIL 竞争：

bash

--mm-processor-pool-size 4

值	行为
`0`（默认）	使用 `ThreadPoolExecutor` — 受 GIL 竞争影响
`> 0`	创建指定数量 worker 的 `ProcessPoolExecutor`，实现真正的 CPU 并行

TIP

建议从每 GPU 可用的 CPU 核数开始设置。例如，在 64 CPU、8 GPU 的节点上，可尝试 --mm-processor-pool-size 8。

SGLang 推理引擎调优

内存分配

控制 SGLang 为 KV Cache 预留的 GPU 显存比例：

bash

--sglang-mem-fraction-static 0.8

通常取值在 0.6 到 0.85 之间。Colocate 模式下，较高的值（0.8）会减少训练可用的显存但提高推理吞吐。全异步模式下使用独立 GPU 时，可以设得更高。

推理 TP 大小

设置每个推理引擎实例使用的 GPU 数：

bash

--rollout-num-gpus-per-engine 1

对于大模型，需要增大此值以满足模型最低 TP 需求。在可能的情况下使用 TP=1 可获得最佳的单请求吞吐。

Partial Rollout（部分生成）

在长响应场景（如代码生成、思维链推理）中，等待所有样本完全生成完毕会导致训练 GPU 长时间空闲。Partial Rollout 通过允许未完成的样本被中断并回收到数据缓冲区来避免这个问题，使训练可以先使用已完成的样本继续进行：

bash

--partial-rollout

控制中断频率

默认情况下，一个样本可以被无限次中断。设置 --partial-rollout-max-aborted-count 来保证样本在被中断一定次数后最终完成生成：

bash

--partial-rollout \
--partial-rollout-max-aborted-count 3

On-Policy 掩码

当一个样本被回收并在后续 rollout 中继续生成时，其早期 Token 是在之前的策略版本下生成的。使用 --mask-offpolicy-in-partial-rollout 来掩码这些 off-policy Token，使只有 on-policy 生成的 Token 参与训练：

bash

--partial-rollout \
--mask-offpolicy-in-partial-rollout

TIP

Partial Rollout 最适合 max_response_len 较大（如 8K+）且响应长度差异较大的工作负载。对于短且长度均匀的响应，回收的开销可能不值得。

数据加载优化

流式数据集

对于内存放不下的超大数据集，使用流式模式：

bash

--use-streaming-dataset \
--streaming-buffer-size 10000

数据均衡

在数据并行 rank 间均匀分配 Token 数量以减少空闲等待：

bash

--balance-data

WARNING

--balance-data 仅在 colocate 模式下可用。全异步模式（--fully-async）不支持该功能，因为 TransferQueue 数据消费路径与数据均衡不兼容。同时启用两者会在启动时报错。

WARNING

使用 --balance-data 时，同一 Prompt 的不同响应可能被分到不同训练步。

权重更新流水线

对于参数量大的 MoE 模型，分块进行权重更新以避免显存尖峰：

bash

--update-weight-buffer-size 536870912  # 512 MB

如果在权重同步期间观察到显存压力，可减小该值。

全异步训练

如需最大化吞吐，使用全异步训练流水线，为训练和推理分配独立的 GPU 集群：

bash

--fully-async \
--max-staleness 1 \
--num-data-storage-units 1

陈旧度调优

--max-staleness 参数控制 rollout 数据相对于当前训练模型可以落后多少个版本。它直接影响吞吐量与数据新鲜度之间的权衡：

值	行为
`1`（默认）	训练仅消费当前或前一个版本的数据。吞吐较低但数据更新鲜
`2-3`	允许适度陈旧的数据。对大多数工作负载是较好的平衡
更高	减少训练等待时间从而提高吞吐，但数据可能由较旧的策略生成

TIP

建议从 --max-staleness 1 开始，如果观察到训练进程频繁等待新鲜的 rollout 数据，再逐步增大。关注训练 loss 的稳定性——如果更高的陈旧度导致 loss 不稳定，则降低该值。

完整配置请参阅全异步训练流水线。

下一步

OOM 排查指南 — 调优导致内存问题时的处理方法
配置参考手册 — 完整参数列表
调试指南 — 隔离排查训练和推理问题

性能调优

性能分析

Trace 文件命名规则

1. SGLang 推理 Profiling

2. 训练 Profiling（PyTorch Profiler）

3. GPU 内存 Profiling

三种 Profiling 联合使用

动态批处理

并行配置

张量并行与序列并行

MoE 专家并行

上下文并行

激活重计算

多模态处理并行化

SGLang 推理引擎调优

内存分配

推理 TP 大小

Partial Rollout（部分生成）

控制中断频率

On-Policy 掩码

数据加载优化

流式数据集

数据均衡

权重更新流水线

全异步训练

陈旧度调优

推荐配置

Qwen3-4B 8 GPU（Colocate）

Qwen3-30B-A3B 16 GPU（全异步）

下一步

性能调优 ​

性能分析 ​

Trace 文件命名规则 ​

1. SGLang 推理 Profiling ​

2. 训练 Profiling（PyTorch Profiler） ​

3. GPU 内存 Profiling ​

三种 Profiling 联合使用 ​

动态批处理 ​

并行配置 ​

张量并行与序列并行 ​

MoE 专家并行 ​

上下文并行 ​

激活重计算 ​

多模态处理并行化 ​

SGLang 推理引擎调优 ​

内存分配 ​

推理 TP 大小 ​

Partial Rollout（部分生成） ​

控制中断频率 ​

On-Policy 掩码 ​

数据加载优化 ​

流式数据集 ​

数据均衡 ​

权重更新流水线 ​

全异步训练 ​

陈旧度调优 ​

推荐配置 ​

Qwen3-4B 8 GPU（Colocate） ​

Qwen3-30B-A3B 16 GPU（全异步） ​

下一步 ​

性能调优

性能分析

Trace 文件命名规则

1. SGLang 推理 Profiling

2. 训练 Profiling（PyTorch Profiler）

3. GPU 内存 Profiling

三种 Profiling 联合使用

动态批处理

并行配置

张量并行与序列并行

MoE 专家并行

上下文并行

激活重计算

多模态处理并行化

SGLang 推理引擎调优

内存分配

推理 TP 大小

Partial Rollout（部分生成）

控制中断频率

On-Policy 掩码

数据加载优化

流式数据集

数据均衡

权重更新流水线

全异步训练

陈旧度调优

推荐配置

Qwen3-4B 8 GPU（Colocate）

Qwen3-30B-A3B 16 GPU（全异步）

下一步