Span Groups | NVIDIA NeMo Lens

Span groups are the NeMo Lens mechanism for controlling trace granularity at runtime without code changes. Every instrumentation site tags itself with a group name; at startup, only the enabled groups actually emit spans.

Why Use Span Groups

A training job can use different levels of tracing detail for different environments:

Minimize production overhead. Enable the default preset to emit only coarse spans (such as job, checkpoint, and evaluate) for the lowest performance cost.
Add iteration boundaries. Enable the per_step preset in staging to add per-iteration boundaries for a moderate performance cost.
Diagnose performance hangs. Enable the all preset during active debugging to instrument every site, including per-microbatch and per-layer spans, which carries the highest performance cost.

The primary goal is to toggle granularity through a single env var without code changes, relying on a fast gating check when disabled.

Base Span Groups

The SpanGroup class ships with eight groups covering typical training workflows:

Coarse-Grained Groups in the `default` Preset

Group	Typical Spans	Frequency
`job`	Pretrain/train root spans	Once per job
`checkpoint`	Checkpoint save	Every N iterations
`evaluate`	Evaluation pass	Every N iterations

Medium-Grained Groups in the `per_step` Preset

Group	Typical Spans	Frequency
`model_init`	Model construction	Once at startup
`load_checkpoint`	Checkpoint load	Once at startup
`step`	Training step boundary	Every iteration
`forward_backward`	Forward and backward pass	Every iteration
`optimizer`	Optimizer step	Every iteration

Presets

Presets bundle groups by use case:

Preset	Groups Included	Relative Cost	Use Case
`default`	`job`, `checkpoint`, and `evaluate`	Lowest	Safe for production
`per_step`	`default` plus `model_init`, `load_checkpoint`, `step`, `forward_backward`, and `optimizer`	Moderate	Staging and profiling
`all`	Every group in the class	Highest	Active debugging

Presets are specific to each subclass; for example, MegatronSpanGroup.ALL_GROUPS contains more groups than SpanGroup.ALL_GROUPS.

Specification String

The config.span_groups field accepts a comma-separated specification string that mixes preset keywords and individual group names:

$ NEMO_LENS_SPAN_GROUPS=default                  # just default preset
$ NEMO_LENS_SPAN_GROUPS=per_step                 # per_step preset
$ NEMO_LENS_SPAN_GROUPS=default,step             # default + one extra group
$ NEMO_LENS_SPAN_GROUPS=step,optimizer,checkpoint # individual groups only
$ NEMO_LENS_SPAN_GROUPS=all                      # everything

Resolution occurs once at setup_telemetry through config.resolved_span_groups. The resulting frozenset is registered with the state module and consulted at every instrumentation site.

Unknown keywords raise ValueError at resolution time with a list of valid options.

State Machinery

Enabled groups live in a module-level frozenset in nemo.lens.state:

1 from nemo.lens.state import is_span_group_enabled, set_enabled_span_groups
2 
3 is_span_group_enabled('step')        # returns bool; a frozenset lookup
4 set_enabled_span_groups(frozenset(['job', 'step']))

The read path (is_span_group_enabled) is lock-free and safe to call from any thread. The write path (set_enabled_span_groups) is lock-protected and typically called once by setup_telemetry.

set_enabled_span_groups is also a top-level public export (from nemo.lens import set_enabled_span_groups), allowing you to override the active groups at runtime without reaching into nemo.lens.state.

Extend with Library-Specific Subclasses

Subclass SpanGroup to add domain-specific groups. For example, Megatron uses the following extension:

1 from nemo.lens.groups import SpanGroup
2 
3 class MegatronSpanGroup(SpanGroup):
4     MICROBATCH = "microbatch"
5     COMMUNICATION = "communication"
6     ACTIVATION_OFFLOAD = "activation_offload"
7     DATA_LOADING = "data_loading"
8     INFERENCE = "inference"
9     LAYER = "layer"
10 
11     ALL_GROUPS = frozenset([
12         *SpanGroup.ALL_GROUPS,
13         MICROBATCH, COMMUNICATION, ACTIVATION_OFFLOAD,
14         DATA_LOADING, INFERENCE, LAYER,
15     ])
16 
17     _PRESETS = {
18         "default": frozenset([SpanGroup.JOB, SpanGroup.CHECKPOINT, SpanGroup.EVALUATE, INFERENCE]),
19         "per_step": frozenset([
20             *SpanGroup.ALL_GROUPS, COMMUNICATION, DATA_LOADING, INFERENCE,
21         ]),
22         "all": ALL_GROUPS,
23     }

Pass this subclass to from_env:

1 cfg = NemoLensConfig.from_env(
2     prefix='MEGATRON_OTEL',
3     fallback_prefix='NEMO_LENS',
4     span_group_cls=MegatronSpanGroup,
5 )

The config.resolved_span_groups field now resolves against MegatronSpanGroup.ALL_GROUPS and its custom _PRESETS.

Design Notes

Configure groups at runtime. Groups are runtime knobs rather than compile-time settings. Toggling an env var and restarting your application constitutes the entire configuration workflow, requiring no code changes.
Ensure orthogonal controls. Groups operate independently of rank sampling (NEMO_LENS_SAMPLER_ENABLED and NEMO_LENS_EXPORT_SAMPLE_RATE) and export strategy (NEMO_LENS_EXPORT_STRATEGY). You can combine these settings to enable per_step groups, sample 10% of ranks, and export from only a single rank.
Filter at multiple granularities. Groups act as coarse-grained filters. For fine-grained control, such as tracing only iterations where loss exceeds a specific threshold, add a runtime check inside your instrumented code, as is_span_group_enabled is only one signal.