企业官网建设流程全解析

Megatron-LM Span Groups 完全指南：用 OpenTelemetry 按需控制 Trace 粒度

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

Megatron-LM 通过nemo-lens接入 OpenTelemetry，在训练循环、检查点、评估、流水线并行通信等框架边界发出 trace。Span Groups（跨度分组）是控制 trace 粒度的核心机制：你只需设置一个环境变量或 CLI 参数，就能在"生产环境低开销的粗粒度追踪"与"开发调试用的逐层细粒度剖析"之间自由切换。读完本文，你将掌握 Megatron 的全部 span 分组预设、完整 span 层级树、关键 span 属性，以及如何结合 rank 策略与采样器在真实训练任务中落地这套可观测性方案。

Span Groups 机制概述

在 Megatron-LM 中，每个 span 都被打上一个分组标签（span group），运行时根据当前启用的分组集合决定是否真正创建并发出该 span。分组粒度由MEGATRON_OTEL_SPAN_GROUPS环境变量（或--otel-span-groupsCLI 标志）控制，取值可以是预设关键词、单独的分组名、或二者的混合。

该机制本身由nemo-lens库提供通用能力，Megatron 在此之上定义了MegatronSpanGroup（位于 megatron/core/telemetry/span_groups.py），扩展出 Megatron 专属的分组，并维护了完整的 span 层级。Megatron 侧只需要：

• 启用遥测总开关：MEGATRON_OTEL_ENABLED=1（默认关闭）；
• 通过MEGATRON_OTEL_SPAN_GROUPS指定粒度规格（默认default）。

典型的最小可用配置如下（详见 Observability 快速开始）：

export MEGATRON_OTEL_ENABLED=1 export OTEL_EXPORTER_OTLP_ENDPOINT=http://localhost:4317 export MEGATRON_OTEL_SPAN_GROUPS=default # 粗粒度，生产安全 torchrun --nproc_per_node=8 pretrain_gpt.py ...

预设关键词（Preset Keywords）

MEGATRON_OTEL_SPAN_GROUPS接受三类取值：预设关键词、单个分组名、或混合（逗号分隔）。三个预设关键词覆盖了从生产到调试的全部粒度需求：

从源码看，MegatronSpanGroup._PRESETS中per_step预设还额外包含first_iteration分组（用于捕获恢复后首个迭代的一次性预热开销），并定义了profiling作为all的别名，两者都等于ALL_GROUPS。default预设实际由{job, checkpoint, evaluate, first_iteration, inference}组成（见 span_groups.py）。

MegatronSpanGroup：Megatron 专属分组

MegatronSpanGroup定义于 megatron/core/telemetry/span_groups.py，继承自 lens 的基类SpanGroup，在通用分组之上增加了 Megatron 特有的细粒度分组：

源码中还定义了文档表格之外的两个分组，均只应显式选择而非纳入常用预设：

• first_iteration：进程内实际执行的首个训练迭代（注意：不一定是第 1 次迭代，例如恢复检查点或跳过迭代之后），用于捕获稳态迭代中不存在的单次预热开销（编译、CUDA graph 捕获、预取）；
• trace_region：为megatron.core.perfetto_trace中约 85 个原生trace_region(...)标记自动生成对应的 lens span，覆盖检查点/数据集/加载等子阶段。它刻意不在per_step预设中，只属于all；需要时显式启用，例如--otel-span-groups per_step,trace_region。

依赖缺失时的降级

span_groups.py采用"可选导入"策略：当nemo-lens未安装时，模块内提供一个最小的SpanGroup桩类，保证MegatronSpanGroup常量始终可用；但此时SpanGroup.resolve()会抛出RuntimeError，提示需pip install nemo-lens才能解析 span-group 规格。同样地，fallbacks.py 在 lens 缺失时提供trace_fn、managed_span、is_span_group_enabled等无操作替身——这意味着所有调用点可以无条件 import 该模块，遥测天然是可选的。

如何配置 Span Groups

环境变量与 CLI 标志

两个入口完全等价，CLI 标志优先级更高：

# 环境变量方式 export MEGATRON_OTEL_SPAN_GROUPS=per_step # CLI 标志方式（覆盖同名环境变量） python pretrain_gpt.py --otel-span-groups per_step ...

在 megatron/training/global_vars.py 的_set_telemetry()中，配置通过NemoLensConfig.from_env(prefix='MEGATRON_OTEL', fallback_prefix='NEMO_LENS', span_group_cls=MegatronSpanGroup)从环境变量构建，随后用三个 CLI 参数覆盖对应字段：

• --otel-enabled：覆盖MEGATRON_OTEL_ENABLED，置位后config.enabled = True；
• --otel-service-name NAME：覆盖OTEL_SERVICE_NAME（未显式设置时默认megatron-lm）；
• --otel-span-groups SPEC：覆盖MEGATRON_OTEL_SPAN_GROUPS，写入config.span_groups。

注意MEGATRON_OTEL_*变量都是NemoLensConfig字段的别名，以NEMO_LENS_*作为回退前缀，并非独立配置项——设置MEGATRON_OTEL_ENABLED=1等价于NEMO_LENS_ENABLED=1（详见 Observability 配置文档）。

常用配置组合示例

# 仅粗粒度 span —— 默认 MEGATRON_OTEL_SPAN_GROUPS=default # 包含逐 step span MEGATRON_OTEL_SPAN_GROUPS=per_step # default + 仅 microbatch（跳过 step/optimizer 分组） MEGATRON_OTEL_SPAN_GROUPS=default,microbatch # 全部 span MEGATRON_OTEL_SPAN_GROUPS=all

生产环境典型组合（配合 10% trace 采样）：

export MEGATRON_OTEL_ENABLED=1 export MEGATRON_OTEL_SPAN_GROUPS=per_step export OTEL_TRACES_SAMPLER=parentbased_traceidratio export OTEL_TRACES_SAMPLER_ARG=0.1 # 保留 10% 的 trace

完整 Span 层级树

以下是 Megatron 可能发出的全部 span 树，每个 span 右侧标注了控制它的分组。理解这棵树是定位性能问题的基础：例如发现某个megatron.layer.mlp耗时异常，你至少需要启用microbatch与layer分组才能看到它。

megatron.pretrain # job ├── megatron.model_init # model_init ├── megatron.load_checkpoint # load_checkpoint │ └── megatron.load_checkpoint.io_read # load_checkpoint └── megatron.train # job ├── megatron.train_step # step │ ├── megatron.forward_backward # forward_backward │ │ ├── megatron.microbatch.forward # microbatch (×N) │ │ │ └── megatron.layer.forward # layer (×L per microbatch) │ │ │ ├── megatron.layer.self_attention │ │ │ └── megatron.layer.mlp │ │ ├── megatron.microbatch.backward # microbatch (×N) │ │ ├── megatron.pp.recv_forward.linked # communication — link to sender's context (PP > 1) │ │ ├── megatron.p2p.recv_forward # communication │ │ ├── megatron.p2p.send_forward # communication │ │ ├── megatron.p2p.recv_backward # communication │ │ ├── megatron.p2p.send_backward # communication │ │ ├── megatron.activation.offload # activation_offload │ │ └── megatron.activation.reload # activation_offload │ ├── megatron.grad_sync.start # communication │ ├── megatron.grad_sync.finish # communication │ └── megatron.optimizer_step # optimizer ├── megatron.save_checkpoint # checkpoint │ ├── megatron.save_checkpoint.state_dict # checkpoint │ └── megatron.save_checkpoint.io_write # checkpoint └── megatron.evaluate # evaluate └── megatron.evaluate.step # evaluate (×N)

从调用点看，这棵树在源码中确实逐层落地：训练循环中的megatron.pretrain、megatron.train、megatron.train_step、megatron.forward_backward、megatron.optimizer_step、megatron.model_init、megatron.evaluate等 span 位于 megatron/training/training.py；megatron.microbatch.forward/backward与megatron.pp.recv_forward.linked位于 megatron/core/pipeline_parallel/schedules.py；P2P 四件套位于 megatron/core/pipeline_parallel/p2p_communication.py；梯度同步 span 位于 megatron/core/distributed/distributed_data_parallel.py；检查点 span 位于 megatron/training/checkpointing.py。

需要说明的是：仓库实现中实际发出的 span 名在个别位置与上树存在简写差异（例如训练循环内部实际使用megatron.startup.model_init、megatron.checkpoint.save、megatron.train.iteration等更细的名称，并带有is_goodput_span=True标记），上表是文档给出的规范层级骨架；排查具体 trace 时以实际导出的 span 名为准。

Span 属性（Span Attributes）

除了分层级，Megatron 还会在关键 span 上附加语义属性，供查询与告警使用。Megatron 专属的 span 属性如下：

此外，_set_telemetry()还会把训练配置写入 OTelResource属性（在 Jaeger 中显示为整个运行每个 span 共有的 Process 标签），包括dl.tensor_parallel.size、dl.pipeline_parallel.size、dl.data_parallel.size、megatron.num_layers、megatron.hidden_size、megatron.precision等，并叠加 lens 自动探测的 hostname、PID、GPU 数量、SLURM/Kubernetes 元数据（详见 configuration.md）。span 属性用于刻画"某次前向/某次保存"的个体特征，resource 属性用于给整个运行打上配置标签，两者互补。

粒度指南与开销分析

不同粒度组合的开销差异显著，选择建议如下：

一个值得强调的实现细节是禁用路径零分配：非导出 rank 的 span 分组是frozenset()，is_span_group_enabled()在所有分组上都返回False，因此根本不会创建 span 对象。禁用路径只是一次 frozenset 查找然后立即返回，而不是一个仍会分配内存的空操作 span（fallback 实现在 fallbacks.py 中同样以无操作替身保证该路径的零成本）。默认情况下整个分布式任务只由一个 rank（最后一个 rank）导出遥测；需要多 rank 数据时，可通过MEGATRON_OTEL_RANK_STRATEGY切换为all_ranks、sampled或first_rank_per_node。

进阶：理解调用链与扩展新 Span

如果你要在 Megatron 中新增一个 span，遵循的模式是：在 span_groups.py 中声明分组常量并加入ALL_GROUPS，然后在目标调用点使用_otel_managed_span(group, name, ...)（上下文管理器，对应 lens 的managed_span）或_otel_trace_fn(group, name)（装饰器）包裹，例如 training.py 中megatron.evaluate的装饰器用法。由于is_span_group_enabled()的提前返回语义，关闭的分组在热路径上只有一次 frozenset 查找的开销——这正是per_step、all等细粒度分组可以安全地内嵌在每次迭代/每个 micro-batch 循环中的原因。

若需在分布式（特别是流水线并行）场景下关联各 rank 的 trace，megatron.pp.recv_forward.linked以及 pipeline-parallel.md 中介绍的 trace 关联机制是首选入口；训练指标（loss、吞吐、梯度范数）的导出细节见 metrics.md。

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