File size: 4,950 Bytes
ed12f96 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 | # Trace 可视化
slime 可以为每条 rollout sample 挂上轻量级执行 trace。它会记录生成、奖励模型等 span 事件，并且可以在保存下来的 rollout debug dump 中离线查看。
## 保存 rollout trace 数据
如果想在运行结束后查看 trace，可以在训练时打开 rollout debug dump：
```bash
python train.py \
... \
--save-debug-rollout-data /path/to/debug/rollout_{rollout_id}.pt
```
每个保存出来的 `.pt` 文件都会包含 rollout samples，以及对应的 `trace` 数据。之后也可以通过 `--load-debug-rollout-data` 复用同一份 dump。
## 打开时间线查看器
对保存好的 rollout dump 运行：
```bash
python tools/trace_timeline_viewer.py /path/to/debug/rollout_0.pt
```
脚本会生成：
- `rollout_0.trace_timeline_cache.json`
- `rollout_0.trace_timeline_viewer.html`
默认情况下，它还会启动一个本地静态文件服务，方便直接在浏览器里打开。如果只想生成文件，可以加 `--no-serve`。
## 如何理解可视化结果
- 每一行对应一条 sample。
- 条形块表示 span，点表示瞬时事件。
- `trace_span(...)` 在开始和结束时记录的属性，都会显示在详情面板里。
- 当 SGLang 返回 PD 分离相关时延时，viewer 会自动补出 `[P]` 和 `[D]` 两条虚拟 lane，用来拆开展示 prefill/decode。
- 如果没有开启 PD，这两条虚拟 lane 不会出现，基础 trace 也仍然可以正常渲染。
## 给自定义代码打点
在自定义 rollout 或 reward 逻辑中，可以直接复用 `slime.utils.trace_utils` 里的工具：
- `trace_span(target, name, attrs=...)`：记录一段持续时间。
- `trace_event(target, name, attrs=...)`：记录一个瞬时事件。
- `trace_function(name, ...)`：用 decorator 把整个同步/异步函数包成一个 span。
- `bind_trace(sample)`：在 sample 被传递到其他 helper 或任务之前，确保它已经绑定好 trace carrier。
### `trace_span` 和 `trace_function` 的区别
如果你只想给函数内部的一小段逻辑打点，或者需要在代码块执行过程中补充 span 结束属性，就用 `trace_span(...)`。
如果整个函数调用都应该对应一个 span，就用 `trace_function(...)`。它本质上会先解析 trace target，然后在函数调用外层自动套一层 `trace_span(...)`，因此同步函数和异步函数都能直接用。
slime 主 rollout 流程里就是这样用的。例如 `generate_and_rm(...)` 按 sample 打点，而 `generate_and_rm_group(...)` 按 group 打点：
```python
from slime.utils.trace_utils import trace_function
@trace_function("generate_and_rm", target="sample")
async def generate_and_rm(args, sample, sampling_params, evaluation=False):
...
@trace_function(
"generate_and_rm_group",
target="group",
attrs_getter=lambda args, group, sampling_params, evaluation=False: {"group_size": len(group)},
)
async def generate_and_rm_group(args, group, sampling_params, evaluation=False):
...
```
### 如何选择 target
`trace_function(...)` 需要一个 trace target，通常是 `Sample`、`TraceHandle` 或它们的列表。
- 如果 target 本来就是函数参数，优先写成 `target="sample"` 或 `target="group"`。
- 如果 target 需要从参数里推导出来，就用 `target_getter=...`。
- 不建议默认依赖自动推断。实现里确实可以从参数或当前 trace 上下文里推断 target，但显式写出来更稳定，也更不容易出现歧义。
### 给 decorator span 增加属性
如果想在 span 开始时附带属性，可以用 `attrs_getter=...`：
```python
@trace_function(
"custom_rollout_batch",
target="samples",
attrs_getter=lambda samples, **_: {"batch_size": len(samples)},
)
async def custom_rollout_batch(samples, **kwargs):
...
```
如果你需要在函数执行到一半之后再补充属性，那么不要只靠 decorator，应该在函数内部再嵌一层 `trace_span(...)`。一个比较常见的模式是：
- 外层用 `trace_function(...)` 表示整个函数生命周期
- 内层用 `trace_span(...)` 标记 generation、RM、filter、post-process 等关键子步骤
如果想统一记录 SGLang 返回的 generation 元信息，可以复用 `build_sglang_meta_trace_attrs`：
```python
from slime.utils.trace_utils import build_sglang_meta_trace_attrs, trace_span
with trace_span(sample, "sglang_generate") as span:
output = await post(url, payload)
span.update(build_sglang_meta_trace_attrs(output["meta_info"]))
```
## 使用建议
- 先保存少量 rollout；单个 dump 的 sample 数量适中时，viewer 会更容易阅读。
- viewer 直接基于保存下来的 `.pt` dump 工作，因此可以把文件拷到别的机器离线分析。
- 如果你想看的是 SGLang 自身的 GPU / kernel 级 profiling trace，请参考 [性能分析](./profiling.md)。
|