企业官网建设流程全解析

Megatron-LM 可观测性扩展指南：为训练框架自定义 Span、Span Group 与 Metric

【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM

Megatron-LM 通过 nemo-lens（下文简称 lens）接入 OpenTelemetry（OTel），在训练循环、流水线调度、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.py（MegatronSpanGroup常量与预设）、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，可能为None（telemetry 未初始化时），因此调用方必须判空：

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", tracer=telemetry.tracer, param_count=1e9): ... # your code

span_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", iteration=iteration) 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_span=True, **{'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_enabled

megatron/core/telemetry/fallbacks.py是这一机制的落点：当 lens 已安装时，它直接再导出nemo.lens.fallbacks中的同名实现（is_span_group_enabled、managed_span、safe_set_span_attributes、span_cm、trace_fn），保证行为一致；当 lens 未安装时，则提供一套内联的 no-op：managed_spanyieldNone、span_cmyieldNone、is_span_group_enabled恒返回False、trace_fn原样返回函数。这样上层代码可以无差别调用，代价仅在未安装时为零。

training.py 的模块级导入正是这一模式的完整示范，它把回退后的函数统一重命名为_otel_*前缀，并在文件内所有插桩点复用。

2.4 仓库内真实调用示例

以下摘自 megatron/training/training.py，可作为撰写自定义 span 时的参照模板：

注意这些真实调用统一把 group 名作为字符串常量传入（'step'、'checkpoint'等），与MegatronSpanGroup中的常量值一一对应，二者可以混用。

三、命名规范：让 span 可检索、可归并

跨进程、跨团队的 telemetry 数据最终都会汇入同一后端（Jaeger、Tempo、Honeycomb 等），命名是否规范直接决定查询效率。Megatron 的约定如下：

具体规则：

• 当某个名字被多个消费者共享（如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_region）；
• layer：每层 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与_PRESETS：

class 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( name="megatron.training.my_new_metric_ms", unit="ms", description="...", ), } _INSTRUMENTS[meter] = instruments return instruments def record_my_metrics(meter, *, my_new_value_ms=None, ...): 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_ms（histogram）、megatron.training.loss（gauge）、megatron.training.throughput_tflops（gauge）、megatron.training.grad_norm（gauge）、megatron.training.skipped_iters（counter）、megatron.training.learning_rate（gauge）、megatron.training.tokens_per_sec（gauge）、megatron.training.memory_allocated_gb（gauge）。设计要点：

• 按 meter 缓存 instruments：用weakref.WeakKeyDictionary以 meter 为键做惰性缓存，Meter被 GC 后缓存自动释放，避免 instrumentation 重复创建导致的内存泄漏与指标重复注册；
• None值静默跳过：record_*的所有参数均可选，传None就不记录，调用方无需分支判断；
• 仅在导出 rank 上调用：通过 telemetry handle 的is_exporting判断，只让负责导出的 rank 调record_*(meter=handle.meter, ...)，避免多 rank 重复上报；
• opentelemetry 未安装时整体 no-op：training_metrics.py顶部对opentelemetry的导入失败置metrics = None，record_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 编写测试时，遵循原文档给出的三层清单：

1. 在tests/unit_tests/telemetry/下添加测试，断言 span 在所属 group 启用时被发射、禁用时不发射（可用is_span_group_enabled的返回值直接验证门控逻辑）；
2. 使用 lens 测试工具提供的InMemorySpanExporter捕获 span（该 exporter 通过 lens 的conftest.py以sys.path或测试工具方式共享，无需起后端服务）；
3. 断言 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_span（group 禁用时 yieldNone），所有 lens 导入都必须配 try/except 回退到 fallbacks.py；
• span 名遵循megatron.<subsystem>.<op>，属性遵循megatron.<attr>/dl.<attr>二分法，共享名字用nemo.lens.semconv常量；
• 先复用现有 group，确需新 group 时同步修改 span_groups.py、文档与 dashboard；
• metric 照抄 training_metrics.py 的 weakref 缓存 +None跳过 + 导出 rank 限定模式；
• 插桩行为用 tests/unit_tests/telemetry/ 下的无 OTel 依赖测试固化下来。

更宏观的插桩点位与指标清单可继续阅读同目录下的 configuration.md、metrics.md 与 span-groups.md；lens 本身的配置模型、插桩原语与测试 fixture 细节以 lens 官方文档为准。

【免费下载链接】Megatron-LMOngoing research training transformer models at scale项目地址: https://gitcode.com/GitHub_Trending/me/Megatron-LM