> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/lens/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/lens/_mcp/server.

# Span Groups

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:

```bash
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`:

```python
from nemo.lens.state import is_span_group_enabled, set_enabled_span_groups

is_span_group_enabled('step')        # returns bool; a frozenset lookup
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:

```python
from nemo.lens.groups import SpanGroup

class MegatronSpanGroup(SpanGroup):
    MICROBATCH = "microbatch"
    COMMUNICATION = "communication"
    ACTIVATION_OFFLOAD = "activation_offload"
    DATA_LOADING = "data_loading"
    INFERENCE = "inference"
    LAYER = "layer"

    ALL_GROUPS = frozenset([
        *SpanGroup.ALL_GROUPS,
        MICROBATCH, COMMUNICATION, ACTIVATION_OFFLOAD,
        DATA_LOADING, INFERENCE, LAYER,
    ])

    _PRESETS = {
        "default": frozenset([SpanGroup.JOB, SpanGroup.CHECKPOINT, SpanGroup.EVALUATE, INFERENCE]),
        "per_step": frozenset([
            *SpanGroup.ALL_GROUPS, COMMUNICATION, DATA_LOADING, INFERENCE,
        ]),
        "all": ALL_GROUPS,
    }
```

Pass this subclass to `from_env`:

```python
cfg = NemoLensConfig.from_env(
    prefix='MEGATRON_OTEL',
    fallback_prefix='NEMO_LENS',
    span_group_cls=MegatronSpanGroup,
)
```

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.