ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

TorchTitan 数值调试实战:补丁 Profiler 以开启逐算子激活捕获(activation capture)

TorchTitan 数值调试实战:补丁 Profiler 以开启逐算子激活捕获(activation capture) 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.pyclass ActivationCaptureProfiler其语义是由Profiler的step()驱动、在指定 step 上捕获激活step()每完成一个训练 step 被调用一次第 N-1 个 step 结束后上膛armDebugModeTracer第 N 个 step 结束后落盘。但当前仓库的 profiler.py 只内置了 Kineto profiler 与 memory snapshot 两条通道build_torch_profiler/build_memory_profiler没有任何 activation capture 相关字段或钩子——这正是补丁文档第 1 节要补全的部分。2. 补丁 0Import 引导import bootstrapactivation_tracer.py位于.claude/skills/numerics_debugging/scripts/。由于.claude不是合法的 Python 标识符该目录无法作为包被点分路径导入所以每个补丁点都必须先把scripts/目录挂到sys.path上再导入。补丁文档给出的引导函数def _numerics_scripts_on_path() - None: Put the numerics_debugging skills 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. 补丁 1torchtitan/observability/profiler.py——新增配置字段与生命周期钩子这是补丁的主体需要补齐四处import、配置字段、构造参数与生命周期钩子、构建器方法。a文件顶部加入引导函数调用与导入# top of file _numerics_scripts_on_path() from activation_tracer import ActivationCaptureProfilerbProfiler.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上的定位——它是一个调试开关。cProfiler.__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_folderself._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( enabledTrue, modelself._model, dump_dirdump_dir, capture_stepcfg.profile_freq, ) profiler.__enter__() return profiler注意两个细节捕获输出目录是{base_folder}/numerics即config.dump_folder下的numerics/子目录capture_step复用现有的profile_freq配置项所以捕获步与 profiling 频率共享同一个 CLI 参数--profiler.profile_freq。4. 补丁 2torchtitan/trainer.py——把 model 传给 Profiler在 trainer.py 的训练主循环中config.profiler.build(...)的调用点当前未传model需要补一行with config.profiler.build( global_stepself.step, base_folderconfig.dump_folder, modelself.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. 补丁 3graph_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_modulebackward 算子。trace 路径下后两级全部失效必须靠第一级。修复思路trace 提交阶段已经把这些上下文暂存在node.meta里custom.module_fqn、stack_trace、autograd_backward。补丁是一个逐节点行走的 FX 解释器把这些元数据恢复成 ContextVar使捕获获得与 eager 相同的上下文。3a.torchtitan/experiments/graph_trainer/debug_utils.py追加FQNInterpreterclass 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_clsself._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 float32python .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_freq2且training.steps2时step 1 是热身step 2 是快照。这与第 3 节build_activation_capture_profiler里capture_stepcfg.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_numel1000的张量输出且排除_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 dont 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.pyDebugModeTracer/ActivationCaptureProfiler/ ContextVar /is_numerics_capture_activediff 工具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创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表