Megatron-LM 可观测性扩展指南为训练框架自定义 Span、Span Group 与 Metric【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LMMegatron-LM 通过 nemo-lens下文简称 lens接入 OpenTelemetryOTel在训练循环、流水线调度、P2P 通信、梯度同步、检查点与评估等框架边界处发射 trace 与 metric。本文是面向 Megatron 开发者的插桩扩展手册你将学会使用span_cm/managed_span为自有代码添加 span、遵循megatron.subsystem.op命名规范、选择合适的 span group甚至定义新 group、按training_metrics.py的模式注册自定义 metric并通过仓库自带的 telemetry 单元测试验证插桩行为——同时严格保证在未安装 lens 的环境下代码依然可用。一、背景基于 nemo.lens 的 OpenTelemetry 插桩体系Megatron-LM 的可观测性设计遵循一条明确的分工原则通用机制span group 概念、插桩原语、配置模型、自定义 exporter、资源探测全部由 lens 提供Megatron 只负责定义自己的 span 名、metric 名、span group 扩展与命名约定。因此往 Megatron 代码里添加新的 span 或 metric本质上是使用 lens 的原语、遵守 Megatron 的约定。仓库内的相关代码集中在两个位置megatron/core/telemetry/telemetry 帮助模块包含 span_groups.pyMegatronSpanGroup常量与预设、fallbacks.py无 lens 时的 no-op 回退、training_metrics.py训练指标记录逻辑。模块的职责说明见init.py。megatron/training/training.py训练循环中数百处真实插桩点是学习如何正确使用这些原语的最佳范例。telemetry handle 通过 megatron/training/global_vars.py 中的get_telemetry()获取它直接返回全局_GLOBAL_TELEMETRY_HANDLE可能为Nonetelemetry 未初始化时因此调用方必须判空from megatron.training.global_vars import get_telemetry telemetry get_telemetry() if telemetry is not None: ... # 仅在 telemetry 激活时才访问 telemetry.tracer / telemetry.meter这一判空习惯贯穿全文所有代码示例。二、添加自定义 Span插桩原语有两条主线span_cm无条件创建与managed_span组门控分别面向冷路径与热路径。2.1 简单块span_cm对于低频、不处于性能关键路径的代码冷路径直接使用span_cm即可。只要 telemetry 处于激活状态它就始终创建 span没有任何 group 开关from megatron.training.global_vars import get_telemetry from nemo.lens.helpers import span_cm telemetry get_telemetry() if telemetry is not None: with span_cm(megatron.my_custom_op, tracertelemetry.tracer, param_count1e9): ... # your codespan_cm支持以关键字参数传入任意属性如上例的param_count会自动附加到创建的 span 上。仓库中这类用法很常见例如 training.py 中的span_cm(megatron.train.iteration.forward_backward, tracer_otel_step_tracer, num_microbatches...)与 training.py 中的span_cm(megatron.train.iteration.optimizer, tracer_otel_step_tracer)——注意它们都会先从nemo.lens.helpers导入span_cm。2.2 组门控块managed_span对于每轮迭代都会执行、甚至每 micro-batch 都会执行的热路径需要把插桩开销压到最低当组被禁用时不应分配任何 span 对象。managed_span正是为此设计——组禁用时它 yield 出None函数体照常执行from megatron.core.telemetry.span_groups import MegatronSpanGroup from nemo.lens.helpers import managed_span with managed_span(MegatronSpanGroup.STEP, megatron.my_custom_step, iterationiteration) as span: result do_work() if span is not None: span.set_attribute(megatron.my_custom.result, result)关键点span为None并不代表插桩失败而是该 group 当前被关闭的正常表现。在设置属性前必须检查if span is not None否则热路径下会引入不必要的对象访问与属性写入。managed_span除了第一个位置参数是 group 名字符串如step外其余签名与span_cm一致。仓库训练循环中的典型调用是with _otel_managed_span(step, megatron.train.iteration, is_goodput_spanTrue, **{megatron.iteration: iteration}) as _step_span:见 training.py其中is_goodput_span是 Megatron 用于区分有效训练时间goodput与停顿时间的附加标记。2.3 必备的 lens 导入回退模式lens 是 Megatron 的可选依赖未安装 lens 时Megatron 仍须完整可用。因此 Megatron 代码中对 lens 的每一次导入都必须使用 try/except 回退惯用法try: from nemo.lens.helpers import managed_span as _otel_managed_span from nemo.lens.state import is_span_group_enabled as _otel_sg_enabled except ImportError: from megatron.core.telemetry.fallbacks import managed_span as _otel_managed_span from megatron.core.telemetry.fallbacks import is_span_group_enabled as _otel_sg_enabledmegatron/core/telemetry/fallbacks.py是这一机制的落点当 lens 已安装时它直接再导出nemo.lens.fallbacks中的同名实现is_span_group_enabled、managed_span、safe_set_span_attributes、span_cm、trace_fn保证行为一致当 lens 未安装时则提供一套内联的 no-opmanaged_spanyieldNone、span_cmyieldNone、is_span_group_enabled恒返回False、trace_fn原样返回函数。这样上层代码可以无差别调用代价仅在未安装时为零。training.py 的模块级导入正是这一模式的完整示范它把回退后的函数统一重命名为_otel_*前缀并在文件内所有插桩点复用。2.4 仓库内真实调用示例以下摘自 megatron/training/training.py可作为撰写自定义 span 时的参照模板场景原语与 group代码位置模型初始化_otel_managed_span(model_init, megatron.startup.model_init, ...)training.pyDataLoader 构建_otel_managed_span(data_loading, megatron.startup.dataloader, ...)training.py加载检查点_otel_managed_span(load_checkpoint, megatron.checkpoint.load, ...)training.py保存检查点_otel_managed_span(checkpoint, megatron.checkpoint.save, ..., **{megatron.iteration: iteration})training.py单步训练_otel_managed_span(step, megatron.train.iteration, ...)training.py首次迭代预热_otel_managed_span(first_iteration, megatron.train.forward_pre_hook, ...)training.py参数范数计算_otel_managed_span(step, megatron.train.params_norm, ...)training.py注意这些真实调用统一把 group 名作为字符串常量传入step、checkpoint等与MegatronSpanGroup中的常量值一一对应二者可以混用。三、命名规范让 span 可检索、可归并跨进程、跨团队的 telemetry 数据最终都会汇入同一后端Jaeger、Tempo、Honeycomb 等命名是否规范直接决定查询效率。Megatron 的约定如下类型约定示例Span 名megatron.subsystem.opmegatron.train_step、megatron.microbatch.forwardSpan 属性Megatron 专属用megatron.attr跨消费者共享用dl.attrmegatron.iteration、dl.rankResource 属性Megatron 专属用megatron.attr共享用dl.attr主机/SLURM/K8s 属性用标准命名megatron.num_layers、dl.tensor_parallel.size、host.nameMetric 名megatron.subsystem.metricmegatron.training.loss具体规则当某个名字被多个消费者共享如DL_RANK、NEMO_RUN_ID时优先使用nemo.lens.semconv中的常量避免各模块手写字符串导致漂移仅属于 Megatron 的名字允许直接硬编码字符串——它们短小且可 grep硬编码反而便于全局搜索与代码审查。完整的 span 属性清单可参考 docs/user-guide/observability/span-groups.md 中的Span attributes小节例如megatron.model_type、megatron.global_batch_size、megatron.num_microbatches、megatron.microbatch_id、dl.pipeline_parallel.rank等。四、选择 Span Group粒度决策每个 span 都必须归属一个 group运行时由MEGATRON_OTEL_SPAN_GROUPS环境变量或--otel-span-groupsCLI 参数决定哪些 group 被发射。新增 span 时按下列决策树选择生产环境永远需要→job除 setup 类 span 外极少使用每轮迭代一次→step位于 forward/backward 内部→forward_backward或microbatch位于优化器内部→optimizer与检查点相关→checkpoint与评估相关→evaluate跨 rank 通信→communication推理请求路径→inference除非现有 group 确实都不合适否则不要自创 group——新增 group 需要改动MegatronSpanGroup并同步更新预设preset维护成本由全仓库共同承担。仓库中MegatronSpanGroup的真实定义见 megatron/core/telemetry/span_groups.py。它继承 lens 的共享基础组job、checkpoint、evaluate、model_init、load_checkpoint、step、forward_backward、optimizer并追加 Megatron 特有的细粒度组microbatch、layer、communication、activation_offload、data_loading、first_iteration、trace_region、inference。其中几个特殊组的语义值得注意first_iteration捕获本进程实际执行的首个训练迭代可能是恢复检查点或跳过迭代之后的迭代不一定是第 1 轮用于归集编译、CUDA graph 捕获、预取等一次性 warmup 开销与稳态stepspan 明确区分trace_region为megatron.core.perfetto_trace中约 85 个 perfetto 原生trace_region(...)标记点自动生成对应 lens span刻意不放进per_step预设需要显式开启如--otel-span-groups per_step,trace_regionlayer每层 forward含 attention 与 MLP 细分开销最高仅all预设包含。真实预设与文档示例略有出入以源码为准为_PRESETS { default: frozenset([JOB, CHECKPOINT, EVALUATE, FIRST_ITERATION, INFERENCE]), per_step: frozenset([JOB, CHECKPOINT, EVALUATE, MODEL_INIT, LOAD_CHECKPOINT, STEP, FORWARD_BACKWARD, OPTIMIZER, COMMUNICATION, DATA_LOADING, FIRST_ITERATION, INFERENCE]), profiling: ALL_GROUPS, all: ALL_GROUPS, }关于各预设的相对开销与 span 层级树从megatron.pretrain到megatron.layer.self_attention的完整父子关系见 docs/user-guide/observability/span-groups.md 的Span hierarchy一节它是验证新 span 挂接层级是否合理的权威参考。五、新增一个 Span Group若现有 group 确实无法承载你的场景按以下三步操作1. 编辑 megatron/core/telemetry/span_groups.py在MegatronSpanGroup中追加常量并更新ALL_GROUPS与_PRESETSclass MegatronSpanGroup(SpanGroup): # ... 现有 group 保持不变 ... MY_NEW_GROUP my_new_group # 新 group 常量 ALL_GROUPS SpanGroup.ALL_GROUPS | frozenset( [..., MY_NEW_GROUP] # 并入全集 ) _PRESETS { default: frozenset([...]), # 通常不把新 group 加进 default per_step: frozenset([...]), # 若按迭代发射则加入 per_step profiling: ALL_GROUPS, all: ALL_GROUPS, # 永远包含在 all 中 }注意ALL_GROUPS必须用frozenset且与基础组取并集_PRESETS必须作为子类自己的属性覆盖不能沿用基类预设。2. 在 docs/user-guide/observability/span-groups.md 中登记新 group说明它控制的 span 及典型发射频率保持文档与代码同步。3. 若新 group 引入了新的 metric 标签同步更新后端 dashboard 查询。新增 group 的可验证性由仓库自带的单元测试保障tests/unit_tests/telemetry/test_span_groups.py 会断言 group 名唯一、不与基础组冲突、ALL_GROUPS恰为基础组并 Megatron 组、预设必须满足default ⊆ per_step ⊆ all的嵌套关系且microbatch/layer/activation_offload/trace_region这些高开销组只能出现在all中。这些测试在未安装 lens 时同样可运行使用span_groups.py内建的SpanGroupstub。六、添加自定义 Metric领域特定 metric 的推荐做法是在megatron/core/telemetry/下新建模块完全模仿 training_metrics.py 的结构。核心模式如下# megatron/core/telemetry/my_domain_metrics.py import weakref from opentelemetry import metrics _INSTRUMENTS: weakref.WeakKeyDictionary weakref.WeakKeyDictionary() def _get_instruments(meter: metrics.Meter) - dict: instruments _INSTRUMENTS.get(meter) if instruments is None: instruments { my_new_metric: meter.create_histogram( namemegatron.training.my_new_metric_ms, unitms, description..., ), } _INSTRUMENTS[meter] instruments return instruments def record_my_metrics(meter, *, my_new_value_msNone, ...): i _get_instruments(meter) if my_new_value_ms is not None: i[my_new_metric].record(my_new_value_ms)该模式在 training_metrics.py 中有完整实现仓库实际注册了 8 个训练指标megatron.training.step_duration_mshistogram、megatron.training.lossgauge、megatron.training.throughput_tflopsgauge、megatron.training.grad_normgauge、megatron.training.skipped_iterscounter、megatron.training.learning_rategauge、megatron.training.tokens_per_secgauge、megatron.training.memory_allocated_gbgauge。设计要点按 meter 缓存 instruments用weakref.WeakKeyDictionary以 meter 为键做惰性缓存Meter被 GC 后缓存自动释放避免 instrumentation 重复创建导致的内存泄漏与指标重复注册None值静默跳过record_*的所有参数均可选传None就不记录调用方无需分支判断仅在导出 rank 上调用通过 telemetry handle 的is_exporting判断只让负责导出的 rank 调record_*(meterhandle.meter, ...)避免多 rank 重复上报opentelemetry 未安装时整体 no-optraining_metrics.py顶部对opentelemetry的导入失败置metrics Nonerecord_training_metrics直接返回因此该模块本身不引入硬依赖。七、测试新的插桩代码Megatron 的 telemetry 单元测试位于 tests/unit_tests/telemetry/共三个文件test_span_groups.py、test_training_metrics.py、test_fallbacks.py。它们采用 lens 测试体系中的全局 OTel 状态重置模式每个测试前后清理全局状态并刻意设计为不依赖真实 OTel SDK 也能运行test_training_metrics.py使用FakeMeter/FakeInstrument驱动record_training_metrics精确断言每个参数写入了哪个 instrument、什么类型histogram/gauge/counter、什么值同时用gc.collect()验证弱引用缓存的释放行为。为新增 span 编写测试时遵循原文档给出的三层清单在tests/unit_tests/telemetry/下添加测试断言 span 在所属 group 启用时被发射、禁用时不发射可用is_span_group_enabled的返回值直接验证门控逻辑使用 lens 测试工具提供的InMemorySpanExporter捕获 span该 exporter 通过 lens 的conftest.py以sys.path或测试工具方式共享无需起后端服务断言 span 的名称、属性以及父子关系新 span 应挂在既有层级树的正确位置可对照 span-groups.md 的 span hierarchy。针对 span group 新增仓库测试给出了一套可直接复用的断言范式见 test_span_groups.py预设必须嵌套、default必须保持粗粒度不能包含step/forward_backward/microbatch、verbose 组只能进all。八、何时不应添加插桩插桩不是越多越好。原文档明确划出三条红线不要在紧密内层循环中插桩如逐 token、逐参数级别。即使managed_span只是一次 frozenset 查找在数十亿次调用累计下也会成为不可忽略的开销不要对全 rank 运行且基数无界的代码插桩。若 span 属性包含类似 tensor shape 这类高方差值会在后端引发基数爆炸cardinality explosion拖垮 trace 存储与查询不要用 span 替代日志。结构化日志属于日志通道可通过 OTel log bridge 与活动 span 的 trace ID 关联span 描述的是有界操作不是每个有趣的事件。当拿不准粒度时先在该子系统的边界处加一个粗粒度 span而不是在每个内部调用点埋细粒度 span——粒度可以后续根据 profiling 数据逐步细化。九、快速回顾冷路径用span_cm热路径用managed_spangroup 禁用时 yieldNone所有 lens 导入都必须配 try/except 回退到 fallbacks.pyspan 名遵循megatron.subsystem.op属性遵循megatron.attr/dl.attr二分法共享名字用nemo.lens.semconv常量先复用现有 group确需新 group 时同步修改 span_groups.py、文档与 dashboardmetric 照抄 training_metrics.py 的 weakref 缓存 None跳过 导出 rank 限定模式插桩行为用 tests/unit_tests/telemetry/ 下的无 OTel 依赖测试固化下来。更宏观的插桩点位与指标清单可继续阅读同目录下的 configuration.md、metrics.md 与 span-groups.mdlens 本身的配置模型、插桩原语与测试 fixture 细节以 lens 官方文档为准。【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 Megatron-LM 可观测性扩展指南：为训练框架自定义 Span、Span Group 与 Metric
Megatron-LM 可观测性扩展指南为训练框架自定义 Span、Span Group 与 Metric【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LMMegatron-LM 通过 nemo-lens下文简称 lens接入 OpenTelemetryOTel在训练循环、流水线调度、P2P 通信、梯度同步、检查点与评估等框架边界处发射 trace 与 metric。本文是面向 Megatron 开发者的插桩扩展手册你将学会使用span_cm/managed_span为自有代码添加 span、遵循megatron.subsystem.op命名规范、选择合适的 span group甚至定义新 group、按training_metrics.py的模式注册自定义 metric并通过仓库自带的 telemetry 单元测试验证插桩行为——同时严格保证在未安装 lens 的环境下代码依然可用。一、背景基于 nemo.lens 的 OpenTelemetry 插桩体系Megatron-LM 的可观测性设计遵循一条明确的分工原则通用机制span group 概念、插桩原语、配置模型、自定义 exporter、资源探测全部由 lens 提供Megatron 只负责定义自己的 span 名、metric 名、span group 扩展与命名约定。因此往 Megatron 代码里添加新的 span 或 metric本质上是使用 lens 的原语、遵守 Megatron 的约定。仓库内的相关代码集中在两个位置megatron/core/telemetry/telemetry 帮助模块包含 span_groups.pyMegatronSpanGroup常量与预设、fallbacks.py无 lens 时的 no-op 回退、training_metrics.py训练指标记录逻辑。模块的职责说明见init.py。megatron/training/training.py训练循环中数百处真实插桩点是学习如何正确使用这些原语的最佳范例。telemetry handle 通过 megatron/training/global_vars.py 中的get_telemetry()获取它直接返回全局_GLOBAL_TELEMETRY_HANDLE可能为Nonetelemetry 未初始化时因此调用方必须判空from megatron.training.global_vars import get_telemetry telemetry get_telemetry() if telemetry is not None: ... # 仅在 telemetry 激活时才访问 telemetry.tracer / telemetry.meter这一判空习惯贯穿全文所有代码示例。二、添加自定义 Span插桩原语有两条主线span_cm无条件创建与managed_span组门控分别面向冷路径与热路径。2.1 简单块span_cm对于低频、不处于性能关键路径的代码冷路径直接使用span_cm即可。只要 telemetry 处于激活状态它就始终创建 span没有任何 group 开关from megatron.training.global_vars import get_telemetry from nemo.lens.helpers import span_cm telemetry get_telemetry() if telemetry is not None: with span_cm(megatron.my_custom_op, tracertelemetry.tracer, param_count1e9): ... # your codespan_cm支持以关键字参数传入任意属性如上例的param_count会自动附加到创建的 span 上。仓库中这类用法很常见例如 training.py 中的span_cm(megatron.train.iteration.forward_backward, tracer_otel_step_tracer, num_microbatches...)与 training.py 中的span_cm(megatron.train.iteration.optimizer, tracer_otel_step_tracer)——注意它们都会先从nemo.lens.helpers导入span_cm。2.2 组门控块managed_span对于每轮迭代都会执行、甚至每 micro-batch 都会执行的热路径需要把插桩开销压到最低当组被禁用时不应分配任何 span 对象。managed_span正是为此设计——组禁用时它 yield 出None函数体照常执行from megatron.core.telemetry.span_groups import MegatronSpanGroup from nemo.lens.helpers import managed_span with managed_span(MegatronSpanGroup.STEP, megatron.my_custom_step, iterationiteration) as span: result do_work() if span is not None: span.set_attribute(megatron.my_custom.result, result)关键点span为None并不代表插桩失败而是该 group 当前被关闭的正常表现。在设置属性前必须检查if span is not None否则热路径下会引入不必要的对象访问与属性写入。managed_span除了第一个位置参数是 group 名字符串如step外其余签名与span_cm一致。仓库训练循环中的典型调用是with _otel_managed_span(step, megatron.train.iteration, is_goodput_spanTrue, **{megatron.iteration: iteration}) as _step_span:见 training.py其中is_goodput_span是 Megatron 用于区分有效训练时间goodput与停顿时间的附加标记。2.3 必备的 lens 导入回退模式lens 是 Megatron 的可选依赖未安装 lens 时Megatron 仍须完整可用。因此 Megatron 代码中对 lens 的每一次导入都必须使用 try/except 回退惯用法try: from nemo.lens.helpers import managed_span as _otel_managed_span from nemo.lens.state import is_span_group_enabled as _otel_sg_enabled except ImportError: from megatron.core.telemetry.fallbacks import managed_span as _otel_managed_span from megatron.core.telemetry.fallbacks import is_span_group_enabled as _otel_sg_enabledmegatron/core/telemetry/fallbacks.py是这一机制的落点当 lens 已安装时它直接再导出nemo.lens.fallbacks中的同名实现is_span_group_enabled、managed_span、safe_set_span_attributes、span_cm、trace_fn保证行为一致当 lens 未安装时则提供一套内联的 no-opmanaged_spanyieldNone、span_cmyieldNone、is_span_group_enabled恒返回False、trace_fn原样返回函数。这样上层代码可以无差别调用代价仅在未安装时为零。training.py 的模块级导入正是这一模式的完整示范它把回退后的函数统一重命名为_otel_*前缀并在文件内所有插桩点复用。2.4 仓库内真实调用示例以下摘自 megatron/training/training.py可作为撰写自定义 span 时的参照模板场景原语与 group代码位置模型初始化_otel_managed_span(model_init, megatron.startup.model_init, ...)training.pyDataLoader 构建_otel_managed_span(data_loading, megatron.startup.dataloader, ...)training.py加载检查点_otel_managed_span(load_checkpoint, megatron.checkpoint.load, ...)training.py保存检查点_otel_managed_span(checkpoint, megatron.checkpoint.save, ..., **{megatron.iteration: iteration})training.py单步训练_otel_managed_span(step, megatron.train.iteration, ...)training.py首次迭代预热_otel_managed_span(first_iteration, megatron.train.forward_pre_hook, ...)training.py参数范数计算_otel_managed_span(step, megatron.train.params_norm, ...)training.py注意这些真实调用统一把 group 名作为字符串常量传入step、checkpoint等与MegatronSpanGroup中的常量值一一对应二者可以混用。三、命名规范让 span 可检索、可归并跨进程、跨团队的 telemetry 数据最终都会汇入同一后端Jaeger、Tempo、Honeycomb 等命名是否规范直接决定查询效率。Megatron 的约定如下类型约定示例Span 名megatron.subsystem.opmegatron.train_step、megatron.microbatch.forwardSpan 属性Megatron 专属用megatron.attr跨消费者共享用dl.attrmegatron.iteration、dl.rankResource 属性Megatron 专属用megatron.attr共享用dl.attr主机/SLURM/K8s 属性用标准命名megatron.num_layers、dl.tensor_parallel.size、host.nameMetric 名megatron.subsystem.metricmegatron.training.loss具体规则当某个名字被多个消费者共享如DL_RANK、NEMO_RUN_ID时优先使用nemo.lens.semconv中的常量避免各模块手写字符串导致漂移仅属于 Megatron 的名字允许直接硬编码字符串——它们短小且可 grep硬编码反而便于全局搜索与代码审查。完整的 span 属性清单可参考 docs/user-guide/observability/span-groups.md 中的Span attributes小节例如megatron.model_type、megatron.global_batch_size、megatron.num_microbatches、megatron.microbatch_id、dl.pipeline_parallel.rank等。四、选择 Span Group粒度决策每个 span 都必须归属一个 group运行时由MEGATRON_OTEL_SPAN_GROUPS环境变量或--otel-span-groupsCLI 参数决定哪些 group 被发射。新增 span 时按下列决策树选择生产环境永远需要→job除 setup 类 span 外极少使用每轮迭代一次→step位于 forward/backward 内部→forward_backward或microbatch位于优化器内部→optimizer与检查点相关→checkpoint与评估相关→evaluate跨 rank 通信→communication推理请求路径→inference除非现有 group 确实都不合适否则不要自创 group——新增 group 需要改动MegatronSpanGroup并同步更新预设preset维护成本由全仓库共同承担。仓库中MegatronSpanGroup的真实定义见 megatron/core/telemetry/span_groups.py。它继承 lens 的共享基础组job、checkpoint、evaluate、model_init、load_checkpoint、step、forward_backward、optimizer并追加 Megatron 特有的细粒度组microbatch、layer、communication、activation_offload、data_loading、first_iteration、trace_region、inference。其中几个特殊组的语义值得注意first_iteration捕获本进程实际执行的首个训练迭代可能是恢复检查点或跳过迭代之后的迭代不一定是第 1 轮用于归集编译、CUDA graph 捕获、预取等一次性 warmup 开销与稳态stepspan 明确区分trace_region为megatron.core.perfetto_trace中约 85 个 perfetto 原生trace_region(...)标记点自动生成对应 lens span刻意不放进per_step预设需要显式开启如--otel-span-groups per_step,trace_regionlayer每层 forward含 attention 与 MLP 细分开销最高仅all预设包含。真实预设与文档示例略有出入以源码为准为_PRESETS { default: frozenset([JOB, CHECKPOINT, EVALUATE, FIRST_ITERATION, INFERENCE]), per_step: frozenset([JOB, CHECKPOINT, EVALUATE, MODEL_INIT, LOAD_CHECKPOINT, STEP, FORWARD_BACKWARD, OPTIMIZER, COMMUNICATION, DATA_LOADING, FIRST_ITERATION, INFERENCE]), profiling: ALL_GROUPS, all: ALL_GROUPS, }关于各预设的相对开销与 span 层级树从megatron.pretrain到megatron.layer.self_attention的完整父子关系见 docs/user-guide/observability/span-groups.md 的Span hierarchy一节它是验证新 span 挂接层级是否合理的权威参考。五、新增一个 Span Group若现有 group 确实无法承载你的场景按以下三步操作1. 编辑 megatron/core/telemetry/span_groups.py在MegatronSpanGroup中追加常量并更新ALL_GROUPS与_PRESETSclass MegatronSpanGroup(SpanGroup): # ... 现有 group 保持不变 ... MY_NEW_GROUP my_new_group # 新 group 常量 ALL_GROUPS SpanGroup.ALL_GROUPS | frozenset( [..., MY_NEW_GROUP] # 并入全集 ) _PRESETS { default: frozenset([...]), # 通常不把新 group 加进 default per_step: frozenset([...]), # 若按迭代发射则加入 per_step profiling: ALL_GROUPS, all: ALL_GROUPS, # 永远包含在 all 中 }注意ALL_GROUPS必须用frozenset且与基础组取并集_PRESETS必须作为子类自己的属性覆盖不能沿用基类预设。2. 在 docs/user-guide/observability/span-groups.md 中登记新 group说明它控制的 span 及典型发射频率保持文档与代码同步。3. 若新 group 引入了新的 metric 标签同步更新后端 dashboard 查询。新增 group 的可验证性由仓库自带的单元测试保障tests/unit_tests/telemetry/test_span_groups.py 会断言 group 名唯一、不与基础组冲突、ALL_GROUPS恰为基础组并 Megatron 组、预设必须满足default ⊆ per_step ⊆ all的嵌套关系且microbatch/layer/activation_offload/trace_region这些高开销组只能出现在all中。这些测试在未安装 lens 时同样可运行使用span_groups.py内建的SpanGroupstub。六、添加自定义 Metric领域特定 metric 的推荐做法是在megatron/core/telemetry/下新建模块完全模仿 training_metrics.py 的结构。核心模式如下# megatron/core/telemetry/my_domain_metrics.py import weakref from opentelemetry import metrics _INSTRUMENTS: weakref.WeakKeyDictionary weakref.WeakKeyDictionary() def _get_instruments(meter: metrics.Meter) - dict: instruments _INSTRUMENTS.get(meter) if instruments is None: instruments { my_new_metric: meter.create_histogram( namemegatron.training.my_new_metric_ms, unitms, description..., ), } _INSTRUMENTS[meter] instruments return instruments def record_my_metrics(meter, *, my_new_value_msNone, ...): i _get_instruments(meter) if my_new_value_ms is not None: i[my_new_metric].record(my_new_value_ms)该模式在 training_metrics.py 中有完整实现仓库实际注册了 8 个训练指标megatron.training.step_duration_mshistogram、megatron.training.lossgauge、megatron.training.throughput_tflopsgauge、megatron.training.grad_normgauge、megatron.training.skipped_iterscounter、megatron.training.learning_rategauge、megatron.training.tokens_per_secgauge、megatron.training.memory_allocated_gbgauge。设计要点按 meter 缓存 instruments用weakref.WeakKeyDictionary以 meter 为键做惰性缓存Meter被 GC 后缓存自动释放避免 instrumentation 重复创建导致的内存泄漏与指标重复注册None值静默跳过record_*的所有参数均可选传None就不记录调用方无需分支判断仅在导出 rank 上调用通过 telemetry handle 的is_exporting判断只让负责导出的 rank 调record_*(meterhandle.meter, ...)避免多 rank 重复上报opentelemetry 未安装时整体 no-optraining_metrics.py顶部对opentelemetry的导入失败置metrics Nonerecord_training_metrics直接返回因此该模块本身不引入硬依赖。七、测试新的插桩代码Megatron 的 telemetry 单元测试位于 tests/unit_tests/telemetry/共三个文件test_span_groups.py、test_training_metrics.py、test_fallbacks.py。它们采用 lens 测试体系中的全局 OTel 状态重置模式每个测试前后清理全局状态并刻意设计为不依赖真实 OTel SDK 也能运行test_training_metrics.py使用FakeMeter/FakeInstrument驱动record_training_metrics精确断言每个参数写入了哪个 instrument、什么类型histogram/gauge/counter、什么值同时用gc.collect()验证弱引用缓存的释放行为。为新增 span 编写测试时遵循原文档给出的三层清单在tests/unit_tests/telemetry/下添加测试断言 span 在所属 group 启用时被发射、禁用时不发射可用is_span_group_enabled的返回值直接验证门控逻辑使用 lens 测试工具提供的InMemorySpanExporter捕获 span该 exporter 通过 lens 的conftest.py以sys.path或测试工具方式共享无需起后端服务断言 span 的名称、属性以及父子关系新 span 应挂在既有层级树的正确位置可对照 span-groups.md 的 span hierarchy。针对 span group 新增仓库测试给出了一套可直接复用的断言范式见 test_span_groups.py预设必须嵌套、default必须保持粗粒度不能包含step/forward_backward/microbatch、verbose 组只能进all。八、何时不应添加插桩插桩不是越多越好。原文档明确划出三条红线不要在紧密内层循环中插桩如逐 token、逐参数级别。即使managed_span只是一次 frozenset 查找在数十亿次调用累计下也会成为不可忽略的开销不要对全 rank 运行且基数无界的代码插桩。若 span 属性包含类似 tensor shape 这类高方差值会在后端引发基数爆炸cardinality explosion拖垮 trace 存储与查询不要用 span 替代日志。结构化日志属于日志通道可通过 OTel log bridge 与活动 span 的 trace ID 关联span 描述的是有界操作不是每个有趣的事件。当拿不准粒度时先在该子系统的边界处加一个粗粒度 span而不是在每个内部调用点埋细粒度 span——粒度可以后续根据 profiling 数据逐步细化。九、快速回顾冷路径用span_cm热路径用managed_spangroup 禁用时 yieldNone所有 lens 导入都必须配 try/except 回退到 fallbacks.pyspan 名遵循megatron.subsystem.op属性遵循megatron.attr/dl.attr二分法共享名字用nemo.lens.semconv常量先复用现有 group确需新 group 时同步修改 span_groups.py、文档与 dashboardmetric 照抄 training_metrics.py 的 weakref 缓存 None跳过 导出 rank 限定模式插桩行为用 tests/unit_tests/telemetry/ 下的无 OTel 依赖测试固化下来。更宏观的插桩点位与指标清单可继续阅读同目录下的 configuration.md、metrics.md 与 span-groups.mdlens 本身的配置模型、插桩原语与测试 fixture 细节以 lens 官方文档为准。【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考