TraceML：用三行代码为训练循环加入 step 级诊断

一直以来，你都能盯着 loss 曲线看；TraceML 让你看见训练循环内部的效率。

在每个训练步骤内部，时间究竟是如何在数据加载、前向、反向和优化器之间分配的，你其实并不清楚。在查看训练运行时，工程师常用的工具链大致是这样：

nvidia-smi和集群仪表盘，用于查看 GPU 利用率

W&B、MLflow 或 TensorBoard，用于查看 loss 曲线和运行历史

PyTorch Profiler 或 Nsight Systems，用于需要深入检查的场景

这些工具都有用，但在"任务正在运行"和"打开一个重量级 profiler"之间还缺了一层：正常训练期间、step 粒度的轻量级可见性，并且只需要极少的配置。

系统监控看到的是机器；实验追踪器看到的是结果；深度 profiler 看到的是 kernel 和 timeline，前提是你已经怀疑了某个具体问题，并且愿意承担相应的开销。

但如果只是想知道，在任意一次运行中，时间是如何在训练 step 内部分配的，又该怎么办？

这正是 TraceML 提供的能力。

TraceML 集成只需少量代码改动。把训练 step 标记一次：

import traceml

traceml.init(mode="auto")

for batch in dataloader:

  with traceml.trace_step(model):

      optimizer.zero_grad(set_to_none=True)

      outputs = model(batch["x"])

      loss = criterion(outputs, batch["y"])

      loss.backward()

      optimizer.step()

然后运行：

traceml run train.py

把训练 step 标记一次，TraceML 把这个边界变成结构化的诊断信息：TraceML 利用 step 边界，把时序、内存、rank、进程和系统信号组织起来，整合成一份关于"运行把时间浪费在了哪里"的诊断。

任务运行时，TraceML 会在日志旁边打开一个实时终端视图。

在一次 PyTorch 训练运行上的实时仪表盘。这次运行被判定为 compute-bound（计算受限），反向传播占据了 step 时间的主要部分，内存面板则提示在观察窗口内 reserved memory 持续增长。

一眼就能回答这些问题：

dataloader fetch 是否占用了过多的 step 时间？

计算是不是主导部分？

在支持的分布式运行中，各个 rank 的表现是否存在差异？

内存是稳定的，还是出现了压力？

这次运行是否足够均衡，到不需要再做更深入 profiling 的程度？

运行结束时，TraceML 会写出final_summary.json——结构化的 JSON，而不是一份庞大的 trace 文件。它足够小，便于存储，容易记录到日志，也方便在不同运行之间做对比。下面是一个简化后的 summary 示例：

{

"step_time": {

  "diagnosis": "INPUT_BOUND",

  "dataloader_pct": 47.0,

  "forward_pct": 31.0,

  "backward_pct": 18.0,

  "optimizer_pct": 4.0

},

"step_memory": {

  "diagnosis": "BALANCED"

}

}

每一项诊断都直接对应到下一步该看哪里：

很多工程时间都花在了去调查那些其实并不需要调查的运行上。知道一次运行是均衡的，就意味着可以放心去做别的事情。

TraceML 没有打算取代 PyTorch Profiler 或 Nsight Systems。需要 kernel 级别的 trace、CUDA timeline 或 NCCL 行为信息时那些才是合适的工具。

他的定位是轻量化的profiler，先跑 TraceML给这次运行做分类，再决定更深入的 profiling 是否值得做、应该指向哪里。

TraceML 是 PyTorch 训练诊断中的 tool zero（零号工具）。

TraceML 最有用的部分不只是实时视图，而是final_summary.json。

这个文件很小、结构化、容易在每个任务上收集。它可以被记录到 W&B 或 MLflow，作为 CI 工件存储，也可以与之前的运行做差异对比。TraceML 还提供了一个 compare 工作流：

traceml compare run_a.json run_b.json

当一次训练性能回退出现时，问题不应该只是吞吐量是否下降了？，而应该是时间转移到哪里去了？。这正是final_summary.json设计要在不同运行之间收集的那种信号。

pip install traceml-ai

traceml run train.py

仓库地址：https://github.com/traceopt-ai/traceml

TraceML 是开源的目前支持单 GPU 以及单节点 DDP/FSDP；多节点支持很快会推出。

如果你从事 PyTorch 训练、分布式任务或 ML 基础设施相关的工作，可以事实这个包，他能让你在打开 profiler 之前，先弄清楚自己面对的是哪一种"慢"。

作者：TraceOpt

点个在看你最好看！