先给采样侧加计时,再决定动哪个参数。训练推进慢不等于采样慢:只有把单次采样拆成环境步进、前向计算、数据传输三段,并且同时记录训练侧等待采样的时长,才能判断是该查环境调用、并行度没生效,还是批大小设置不合适。在没有分段数据之前调批大小,通常只会把两种不同的瓶颈混在一起,测出来的结论也无法复现。
排查顺序建议固定:先看环境步进是不是主要耗时,再看训练是否真的在等待采样,然后才用单变量法调批大小,最后确认并行、工作进程类配置是否真的生效。每轮测试只改一个参数,记录每秒采样条数、单次耗时、进程数和内存占用。如果环境步进占比大,调大批大小一般收益有限;如果前向与数据传输占主要时间,再考虑批大小和并行度。
给采样侧加计时,拆出各段耗时
这一步的目标是确定耗时落在哪一段,而不是立刻优化。直接在采样循环里打三处时间戳,累计后打印单次均值,骨架如下,函数名和调用方式按你自己的采样器替换:
import time
t = {'env': 0.0, 'fwd': 0.0, 'xfer': 0.0}
n = 0
def sample_once(env, policy, transfer_fn, action):
global n
t0 = time.perf_counter()
obs = env.step(action) # 环境步进
t1 = time.perf_counter()
out = policy(obs) # 前向计算
t2 = time.perf_counter()
transfer_fn(out) # 数据传输/落盘/入队
t3 = time.perf_counter()
t['env'] += t1 - t0
t['fwd'] += t2 - t1
t['xfer'] += t3 - t2
n += 1
if n % 100 == 0:
print('step=%d env=%.1fms fwd=%.1fms xfer=%.1fms total=%.1fms'
% (n, t['env']/n*1e3, t['fwd']/n*1e3,
t['xfer']/n*1e3,
(t['env']+t['fwd']+t['xfer'])/n*1e3))
区分预热与稳定阶段有两种简单做法:一是打印前若干步的逐次耗时,先看首次调用是否明显偏高(常见于环境初始化、模型首次前向、显存分配);二是把累计计数在预热步之后清零,再开始统计均值。报告结论时要说明用的是哪种口径,否则两次测试的数字不可比。如果环境步进项的均值明显高于前向和数据传输两项,下一步优先查环境调用方式与并行度,而不是批大小。
确认训练侧是否在等待采样
采样慢不一定是问题,只有训练在等采样时,它才是训练推进速度的瓶颈。在取样本的调用外层再打一对时间戳,同时记录本步的计算耗时:
tw0 = time.perf_counter()
batch = sampler.get(batch_size) # 等待采样返回
wait = time.perf_counter() - tw0
wait_ratio = wait / (wait + compute_time)
log_every(k, wait, compute_time, wait_ratio)
打印时同时给出等待绝对时长和等待占比,按固定步数汇总一次。如果等待占比高,说明采样侧确实是瓶颈,继续走后面几节的排查;如果等待占比低,训练时间主要花在计算或优化器上,调批大小对整体推进速度影响有限,应先把注意力放回训练侧,避免在采样配置上反复试错。需要注意:单次采样存在尖峰,建议看一段窗口的分布而不只看最大值,并结合日志确认是否伴随数据加载或落盘抖动。
只改批大小并对比采样吞吐
批大小的影响必须用单变量法确认。控制变量做法是:固定并行度配置、固定环境数量、固定前向模型与精度、固定数据传输方式,只改采样批量,其余保持不变。每个取值跑同一段步数,记录每秒采样条数(样本数除以采样总耗时,不含训练计算),并把第一节的三段耗时一起打印出来,方便看出耗时转移到哪一段。
内存占用要同步观察:进程常驻内存可用 ps、top 或 /proc 下的状态查看,显存可用框架自带的内存统计接口或显卡查询命令,重点看峰值而非瞬时值。批大小调大后如果吞吐没有上升而内存峰值接近上限,说明继续加批量可能换来换入换出或分配失败,通常应停在该取值附近,不再硬顶。结论要写成“在该并行度配置下,批量从 A 到 B 的每秒采样条数变化”,不要泛化成通用最优值。
检查并行度相关配置是否生效
配置写了但没生效是采样慢的常见原因。先在配置文件和启动入口里定位并行、工作进程类配置项,通常名字包含 worker、num_envs、num_workers、parallel、n_proc、threads 这类词,用搜索命令把它们连同默认值和生效位置一起列出来:
grep -rn -E 'worker|num_envs|parallel|n_proc|threads' configs/ src/ | head -50
改完之后必须验证生效:看进程或线程数(如 ps -eLf 后按进程过滤统计行数),或在工作进程启动处打印进程号与读取到的配置值,确认数量与配置一致、没有回退到默认值。增减工作进程时,如果日志里的进程号数量不变,或者环境步进耗时几乎不动,就说明这个配置没有进入实际执行路径,可能是被上层覆盖、参数名拼错、或该后端不支持多进程。确认生效后,再按第三节的单变量法对比不同并行度下的每秒采样条数,同时观察内存与显存峰值的变化。
把多轮测试写成一张对比表
每轮测试只动一个参数,测完立刻填表,避免事后回忆。字段建议如下,一行一次配置:
| 配置项 | 取值 | 耗时 | 吞吐 | 备注 |
|---|---|---|---|---|
| 采样批大小 | 基准值 | 环境/前向/传输分段均值 | 每秒采样条数 | 其余参数固定,记录预热口径 |
| 采样批大小 | 上调一档 | 同上 | 同上 | 记录内存与显存峰值 |
| 工作进程数 | 当前值 | 同上 | 同上 | 附进程数验证结果 |
| 工作进程数 | 增减一档 | 同上 | 同上 | 确认配置已被读取 |
| 训练侧等待占比 | — | 等待时长/步 | — | 判断瓶颈是否真在采样侧 |
表里要让别人能看出:改了哪个参数、其他参数是什么、耗时用哪一段数据、吞吐怎么算。同一张表里混用不同预热口径或不同步数的数据,结论就不可核对。最终建议按“先环境调用与并行度、后批大小”的顺序收敛,并保留原始日志,方便复测时对齐。