core.telemetry.span_groups

Span group definitions for Megatron-LM telemetry.

Tries to import the real SpanGroup from nemo.lens. When nemo-lens is not installed, a minimal stub is provided so that MegatronSpanGroup constants are always available.

Module Contents

Classes

MegatronSpanGroup

Span groups for Megatron-LM instrumentation.

API

class core.telemetry.span_groups.MegatronSpanGroup

Bases: nemo.lens.groups.SpanGroup

Span groups for Megatron-LM instrumentation.

Extends the shared groups with Megatron-specific fine-grained groups.

MICROBATCH

‘microbatch’

Per-microbatch forward/backward spans.

LAYER

‘layer’

Per-transformer-layer forward (attention + MLP breakdown).

COMMUNICATION

‘communication’

P2P send/recv and gradient AllReduce/ReduceScatter.

ACTIVATION_OFFLOAD

‘activation_offload’

GPU<->CPU activation offload/reload spans.

DATA_LOADING

‘data_loading’

Data loading and batch preparation.

FIRST_ITERATION

‘first_iteration’

The first training iteration actually executed in this process (post checkpoint-resume, post iteration-skip) — not necessarily iteration 1, and distinct from the per-step STEP span since it captures one-off warmup costs (compilation, CUDA graph capture, prefetch) absent from steady-state iterations.

TRACE_REGION

‘trace_region’

Shadows every perfetto-native trace_region(...) marker with a lens span (see megatron.core.perfetto_trace) — ~85 checkpoint/dataset/load sub-phase markers, covered without per-site instrumentation. Verbose and fine-grained: deliberately NOT in the per_step preset (only all); opt in explicitly, e.g. --otel-span-groups per_step,trace_region.

INFERENCE

‘inference’

Inference server request spans.

ALL_GROUPS: Final[frozenset]

None

_PRESETS: ClassVar[dict]

None