企业官网建设流程全解析

TorchTitan 数值调试实战：补丁 Profiler 以开启逐算子激活捕获（activation capture）

【免费下载链接】torchtitanA PyTorch native platform for training generative AI models项目地址: https://gitcode.com/GitHub_Trending/to/torchtitan

本文基于 TorchTitan 仓库中数值调试技能（numerics_debugging skill）的补丁参考文档 patching.md 展开，系统讲解如何临时修改torchtitan源码，把ActivationCaptureProfiler接入训练循环的Profiler生命周期，使训练在指定 step 上落盘逐算子（per-op）激活统计日志；当使用graph_trainer的aot_fx_trace编译路径时，还需补上一个 FX 解释器补丁，让 trace 图重放中的算子也能被正确归属到模块 FQN。读完本文，你将掌握完整的补丁点清单、每个补丁的源码级理由，以及捕获-对比工作流的参数要求。

1. 为什么需要补丁：捕获机制的设计边界

TorchTitan 的数值调试工具链由 SKILL.md 描述：核心是两个位于 scripts 目录 的脚本——activation_tracer.py（运行时捕获，基于torch.utils._debug_mode.DebugMode）和compare_numerics.py（对比两份日志并生成 HTML 报告）。捕获由Profiler里的ActivationCaptureProfiler驱动，输出到{dump_folder}/numerics/rank_{N}_activations.log。

关键在于设计边界：这两个脚本刻意放在torchtitan包之外。SKILL.md明确指出，核心torchtitan或graph_trainer中没有任何代码引用它们，Agent 必须在捕获运行前编辑torchtitan把 tracer 接进来，并在完成后还原——这些改动"不属于main"。

ActivationCaptureProfiler的实现确实存在于 activation_tracer.py（class ActivationCaptureProfiler），其语义是"由Profiler的step()驱动、在指定 step 上捕获激活"：step()每完成一个训练 step 被调用一次，第 N-1 个 step 结束后"上膛"（arm）DebugModeTracer，第 N 个 step 结束后落盘。但当前仓库的 profiler.py 只内置了 Kineto profiler 与 memory snapshot 两条通道（build_torch_profiler/build_memory_profiler），没有任何 activation capture 相关字段或钩子——这正是补丁文档第 1 节要补全的部分。

2. 补丁 0：Import 引导（import bootstrap）

activation_tracer.py位于.claude/skills/numerics_debugging/scripts/。由于.claude不是合法的 Python 标识符，该目录无法作为包被点分路径导入，所以每个补丁点都必须先把scripts/目录挂到sys.path上再导入。补丁文档给出的引导函数：

def _numerics_scripts_on_path() -> None: """Put the numerics_debugging skill's scripts/ on sys.path.""" import sys from pathlib import Path for parent in Path(__file__).resolve().parents: scripts = parent / ".claude" / "skills" / "numerics_debugging" / "scripts" if scripts.is_dir(): if str(scripts) not in sys.path: sys.path.insert(0, str(scripts)) return raise RuntimeError("numerics_debugging skill scripts/ not found")

文档特别强调：每个后续代码片段都应在 import 之前立即调用它——可以把该助手函数粘贴到需要它的文件里，也可以放在一个共享位置再 import。

为什么从__file__向上遍历父目录，而不能锚定torchtitan.__file__？这一点文档给出了精确解释：向上遍历锚定的是"包含被补丁文件的这份 checkout"。看起来等价的torchtitan.__file__锚定法在 editable 安装下并不等价——import torchtitan可能根据工作目录解析到另一份checkout，此时引导函数要么加载了别的代码树的 tracer，要么直接失败。这是多 checkout 环境（常见于开发机）下容易踩的坑。

3. 补丁 1：torchtitan/observability/profiler.py——新增配置字段与生命周期钩子

这是补丁的主体，需要补齐四处：import、配置字段、构造参数与生命周期钩子、构建器方法。

（a）文件顶部：加入引导函数调用与导入：

# top of file _numerics_scripts_on_path() from activation_tracer import ActivationCaptureProfiler

（b）Profiler.Config：在enable_memory_snapshot字段旁新增配置项。对照 profiler.py 现有结构，Config是kw_only的 dataclass，已有enable_profiling、profile_freq、enable_memory_snapshot等字段，新增字段为：

# inside Profiler.Config — add next to enable_memory_snapshot dump_numerics: bool = False """Dump per-op activation logs for numerics debugging. Writes ``{dump_folder}/numerics/rank_{rank}_activations.log`` (per-op stats + norm hashes of inputs / outputs)."""

dump_numerics默认False，因此补丁合入前对默认训练零影响；这也符合"补丁不留在main上"的定位——它是一个调试开关。

（c）Profiler.__init__：新增model关键字参数。现有签名（见 profiler.py 的__init__）接收config以及global_step/base_folder/leaf_folder，补丁后为：

# Profiler.__init__ — add model kwarg and slots def __init__( self, config: "Profiler.Config", *, global_step: int = 0, base_folder: str = "", leaf_folder: str = "", model: torch.nn.Module | None = None, # for activation capture profiler ) -> None: ... self.activation_capture_profiler = None # ActivationCaptureProfiler registers global module forward hooks on # the model so backward ops can recover their owning FQN. self._model = model

为什么ActivationCaptureProfiler需要model？从 activation_tracer.py 的DebugModeTracer.__enter__源码可以看到原因：它通过nn.modules.module.register_module_forward_pre_hook/register_module_forward_hook注册全局forward hooks，在 forward 后置钩子里从每个模块输出的 autograd 图出发做 DFS，把grad_fn -> FQN记入_grad_fn_to_module。这样 backward 阶段（由 C++ autograd engine 驱动，DebugMode的ModTracker模块栈为空）的算子才能恢复其所属模块的 FQN。没有 model 就没有这层归属信息。

（d）生命周期钩子：__enter__、__exit__、step三处，与现有 torch profiler / memory profiler 并列驱动：

# Profiler.__enter__ — build the activation capture profiler alongside the memory profiler self.activation_capture_profiler = self.build_activation_capture_profiler( base_folder=self._base_folder, ) # Profiler.__exit__ — teardown if self.activation_capture_profiler is not None: self.activation_capture_profiler.__exit__(exc_type, exc_val, exc_tb) self.activation_capture_profiler = None # Profiler.step — drive the capture-step cadence if self.activation_capture_profiler is not None: self.activation_capture_profiler.step()

Profiler.step()在 trainer.py 的训练主循环中每个 step 之后被调用（profiler.step()），这正好匹配ActivationCaptureProfiler.step()"每完成一个训练 step 调用一次"的设计节奏。

（e）新增构建方法：

# new method def build_activation_capture_profiler(self, *, base_folder: str): """Create and return an :class:`ActivationCaptureProfiler`, or ``None`` if disabled.""" cfg = self._config if not cfg.dump_numerics or self._model is None: return None dump_dir = os.path.join(base_folder, "numerics") profiler = ActivationCaptureProfiler( enabled=True, model=self._model, dump_dir=dump_dir, capture_step=cfg.profile_freq, ) profiler.__enter__() return profiler

注意两个细节：捕获输出目录是{base_folder}/numerics，即config.dump_folder下的numerics/子目录；capture_step复用现有的profile_freq配置项，所以捕获步与 profiling 频率共享同一个 CLI 参数（--profiler.profile_freq）。

4. 补丁 2：torchtitan/trainer.py——把 model 传给 Profiler

在 trainer.py 的训练主循环中，config.profiler.build(...)的调用点（当前未传model）需要补一行：

with config.profiler.build( global_step=self.step, base_folder=config.dump_folder, model=self.model_parts[0], # add this line ) as profiler: ...

补丁文档解释了为什么是model_parts[0]：在 pipeline 并行切分下model_parts是本 rank 拥有的模型片段列表，rank-0 拥有的 eager 模型就是model_parts[0]——ActivationCaptureProfiler正是在这个对象上安装 forward hooks，使得DebugMode的ModTracker与_grad_fn_to_module能把 backward 算子归属到正确的 FQN。

5. 补丁 3：graph_trainer专属——用 FQNInterpreter 重放 trace 图

这一节只在激活的训练路径是--compile.mode aot_fx_trace（即graph_trainer）时才需要。

问题本质：trace 图以gm(*flat_inputs)的方式整体调用，绕过了所有nn.Module.forward。于是DebugMode的ModTracker无法把算子归属到 FQN，日志退化为满屏<none>/op_N_*。从 activation_tracer.py 的record_hook可以看到 FQN 的三级回退优先级：_current_module_nameContextVar（由 FQNInterpreter 在 trace 重放时设置）→ModTracker模块栈（eager 模式）→_grad_fn_to_module（backward 算子）。trace 路径下后两级全部失效，必须靠第一级。

修复思路：trace 提交阶段已经把这些上下文暂存在node.meta里（custom.module_fqn、stack_trace、autograd_backward）。补丁是一个逐节点行走的 FX 解释器，把这些元数据恢复成 ContextVar，使捕获获得与 eager 相同的上下文。

3a.torchtitan/experiments/graph_trainer/debug_utils.py追加FQNInterpreter：

class FQNInterpreter(torch.fx.Interpreter): """Interpreter that sets activation tracer context vars from node metadata.""" def run_node(self, n: torch.fx.Node): from contextvars import Token _numerics_scripts_on_path() from activation_tracer import ( _current_module_name, _current_phase_override, _current_stack_frames, _parse_stack_trace, ) fqn = (n.meta.get("custom") or {}).get("module_fqn") stack_trace = n.meta.get("stack_trace") is_backward = n.meta.get("autograd_backward", False) phase = "backward" if is_backward else "forward" tokens: list[Token] = [] if fqn: tokens.append(_current_module_name.set(fqn)) if stack_trace: tokens.append(_current_stack_frames.set(_parse_stack_trace(stack_trace))) tokens.append(_current_phase_override.set(phase)) try: return super().run_node(n) finally: for token in reversed(tokens): token.var.reset(token)

逐节点执行前设置三个 ContextVar，执行完用Token逆序复位——finally块保证节点抛异常时上下文也不泄漏到下一个节点。这三个 ContextVar 与_parse_stack_trace在 activation_tracer.py 中均有定义，record_hook会优先读取它们（见第 3 节所述的 FQN 优先级）。

3b.torchtitan/experiments/graph_trainer/trainer.py——仅在捕获步注入解释器：

def _maybe_get_fqn_interpreter(self) -> type | None: _numerics_scripts_on_path() from activation_tracer import ( is_numerics_capture_active, ) if is_numerics_capture_active(): from torchtitan.experiments.graph_trainer.debug_utils import FQNInterpreter return FQNInterpreter return None # in forward_backward_step, where run_traced is invoked: outputs = run_traced( ..., interpreter_cls=self._maybe_get_fqn_interpreter(), )

这里有两点源码层面的印证：

• make_fx_tracer.py 中的run_traced已经接受interpreter_cls: type | None = None参数，并在非None时通过interpreter_cls(traced_result.gm).run(*flat_inputs)执行 trace 图——所以这一处无需任何补丁，只要把FQNInterpreter传进去即可。
• is_numerics_capture_active()检查的是 activation_tracer.py 中的模块级标志_numerics_capture_active，它在ActivationCaptureProfiler._setup()（上膛）时置True、_teardown()（落盘）时置回False。因此解释器只在捕获步生效，稳态训练路径（直接gm(*inputs)）完全不受影响。文档也明确了这一设计意图："the interpreter only kicks in on the capture step (and only under aot_fx_trace), so steady-state training is untouched."

6. 补丁之后：捕获-对比工作流与参数约束

补丁完成后的完整操作流（引自 SKILL.md 与补丁文档）：对每个要对比的运行各捕获一次，再用compare_numerics.py做 diff：

./run_train.sh \ --dump_folder ./outputs/run_A \ --training.steps 2 \ --profiler.dump_numerics \ --profiler.profile_freq 2 \ --debug.seed 42 \ --debug.deterministic \ --training.mixed_precision_param float32

python .claude/skills/numerics_debugging/scripts/compare_numerics.py \ outputs/run_A/numerics/rank_0_activations.log \ outputs/run_B/numerics/rank_0_activations.log \ --name1 run_A --name2 run_B \ -o diff.html

几个关键约束与语义，直接关系到补丁参数的正确取值：

• 捕获步 =profile_freq。profile_freq=2且training.steps=2时，step 1 是热身，step 2 是快照。这与第 3 节build_activation_capture_profiler里capture_step=cfg.profile_freq的取值一致。
• 内存开销只发生在捕获步。DebugModeTracer的统计量（L2 norm、Mean 等）在捕获时以内联方式以 float64 计算（见 activation_tracer.py 的_compute_stats），不克隆、不持有张量本身，捕获步额外内存约 10–40%。
• 两次运行必须使用相同 dtype 与 seed。compare_numerics.py的匹配器以 shape + float64 L1 norm 为键，若精度不同（如 bf16 对 fp32）每行都会发散，匹配器退化到仅结构化的stats通道。所以--debug.deterministic与统一--debug.seed是硬性前提。
• 默认只捕获float32/float16/bfloat16、元素数 ≥min_numel=1000的张量输出，且排除_EXCLUDED_OPS中的基础设施算子（view/reshape/indexing/cast 等）；通信类算子（all_gather_into_tensor、reduce_scatter_tensor等）不在默认排除列表内，因为它们常在 eager 与 traced 之间产生差异，需要可见。
• 捕获日志每行格式为module_fqn/op_N_opname（如layers.0.attention.qkv_linear.wq/op_0_mm），带 phase 标注；op_N是模块内计数器。这解释了补丁 3 中"日志退化为<none>/op_N_*"的含义——FQN 缺失后只剩计数器与算子名。

7. 补丁纪律与深入阅读路径

补丁文档开篇即给出纪律要求："Apply these patches before a capture run, then revert them when you are done (they don't belong onmain)."捕获机制对核心训练路径零侵入（dump_numerics默认False、FQNInterpreter 仅在捕获步注入），但sys.path引导代码和profiler.py/trainer.py的改动属于调试设施，不应回流到主分支。

若要在 diff 报告中进一步定制（排除算子、min_numel/ dtype 过滤、hash 函数、手工 override CSV 格式、HTML 外观），参见同目录的 customization.md；它与本文的补丁文档同属 numerics_debugging skill 的 references，分别回答"怎么接进来"与"接进来之后怎么调"两个问题。涉及的核心文件清单：

• 捕获实现：activation_tracer.py（DebugModeTracer/ActivationCaptureProfiler/ ContextVar /is_numerics_capture_active）
• diff 工具：compare_numerics.py（纯标准库，不依赖 torch）
• 补丁目标：profiler.py、trainer.py、debug_utils.py、graph_trainer/trainer.py、make_fx_tracer.py

【免费下载链接】torchtitanA PyTorch native platform for training generative AI models项目地址: https://gitcode.com/GitHub_Trending/to/torchtitan