Instrument Megatron-LM with OpenTelemetry traces and metrics

[deepakn94]

What does this PR do?

Wires the span groups and metric instruments added in #6415 into the code that produces them, and adds the configuration surface to turn it all on.

This replaces #6385, rebased onto current main now that #6415 has merged. Original changes by @rhewett-nv — please close #6385 in favour of this one.

What's instrumented

Subsystem	File
Training loop, model init, evaluation	megatron/training/training.py
Checkpoint save/load, async checkpointing	megatron/training/checkpointing.py, async_utils.py
Pipeline schedules, microbatch forward/backward	megatron/core/pipeline_parallel/schedules.py
P2P send/recv	megatron/core/pipeline_parallel/p2p_communication.py
Gradient sync	megatron/core/distributed/distributed_data_parallel.py
Activation offload/reload	megatron/core/pipeline_parallel/fine_grained_activation_offload.py
Transformer and Mamba layers	megatron/core/transformer/transformer_layer.py, megatron/core/ssm/mamba_layer.py
Dataset builder	megatron/core/datasets/blended_megatron_dataset_builder.py

Adds --otel-enabled, --otel-service-name, --otel-span-groups, the MEGATRON_OTEL_* environment variables, and telemetry setup/teardown in megatron/training/global_vars.py. Publishes the Observability guide by adding its docs/index.md toctree entry.

Cost when disabled

Spans are gated per group: a disabled group costs a frozenset lookup and no span object is allocated. Non-exporting ranks resolve to an empty group set. Every nemo.lens import falls back to the no-ops in megatron.core.telemetry.fallbacks, so nothing here requires nemo-lens to be installed.

Changes since the first review

Thanks @ahmadki and @santhnm2 — all four review threads are addressed.

• Inference-server instrumentation reverted. megatron/core/inference/text_generation_server/text_generation_server.py is byte-identical to main again; this PR no longer touches that module. Per @santhnm2 it is effectively deprecated, so the GenAI-semconv spans and metrics will be ported to dynamic_text_gen_server/text_generation_server.py as a follow-up rather than reworked here.
• The nemo.lens test suite is gone. Removed all seven files (conftest.py, test_config, test_handle, test_helpers, test_instrumentation, test_instruments, test_propagation). Per @ahmadki they test the Lens library rather than Megatron and belong upstream. They were inherited verbatim from Telemetry/otel lens upstream #6385. The four telemetry test files already on main from Add megatron.core.telemetry base layer and observability guide #6415 are untouched and still pass.
• RandomIdGenerator monkeypatch removed from _set_telemetry. It reached into opentelemetry.sdk process-globally to work around identically-seeded span IDs across data-parallel ranks; nemo-lens handles that in its own provider setup.
• nemo-lens pin bumped to b85578fc in both the dev group and the otel extra.
• Rebased onto main across 68 commits. Two conflicts: distributed_data_parallel.py (kept both main's PARAM_READY_CALLBACK_ATTR import and the lens trace_fn fallback) and training.py, where main had removed torch.cuda.empty_cache() from save_checkpoint_and_time in Remove torch cuda empty cache #5678 — adopted that removal rather than resurrecting the call.
• Ran tools/autoformat.sh, which Telemetry/otel lens upstream #6385 never had; black reformatted 9 files and isort fixed 6.

⚠️ opentelemetry-api pin widened — needs a look

uv.lock had to be regenerated for the new dependency, and doing so surfaced an unsatisfiable conflict that exists on main today:

• dev pinned opentelemetry-api~=1.33.1, capping OpenTelemetry at 1.33.x.
• That forces opentelemetry-exporter-otlp-proto-grpc <=1.34.1, which needs protobuf>=3.19,<6.0.
• nvidia-resiliency-ext==0.6.0 needs grpcio-tools>=1.76.0, which needs protobuf>=6.31.1.
• Irreconcilable, so megatron-core[dev] and megatron-core[otel] could not resolve together.

nemo-lens is not the constraint; it asks only for opentelemetry-api>=1.20.0. The cap was. The conflict is invisible on main only because nothing currently pulls the otel extra.

Fix: opentelemetry-api~=1.33.1 -> >=1.33.1,<2. This moves every dev environment from OpenTelemetry 1.33.1 to 1.44.0, not just telemetry runs, which is the main thing worth reviewing here.

Two notes on the lockfile:

• Generated with uv 0.7.2, the version CI pins in both setup-uv (cicd-main.yml:396) and Dockerfile.ci.dev. Verified in the CI container with uv lock --check (exit 0) and uv sync --locked --only-group linting (exit 0), which is what the linting job runs.
• The diff shows revision = 3 -> revision = 2. Main's lock was generated by a newer uv than CI pins, so 0.7.2 rewrites it in its own format. Pre-existing inconsistency, not introduced here.

Issue tracking

Linked issue: replaces #6385

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

Verified locally

• Sphinx builds clean with --fail-on-warning (the flag the docs job uses), including the newly published Observability toctree.
• black, isort, ruff, and pylint (10.00/10) pass on every changed file in megatron/core.
• All changed Python files compile.
• The 83 telemetry unit tests from Add megatron.core.telemetry base layer and observability guide #6415 still pass against this branch.

Known follow-ups

Extracting the span-lifecycle helpers out of training.py — immediate follow-up, already planned.

This PR adds 1,281 lines to megatron/training/training.py (now 5,603 lines), and a large share of that is telemetry bookkeeping rather than training logic:

• 11 span-lifecycle functions in a contiguous block (_start_otel_job_spans, _end_otel_startup_span, _start_otel_train_span, _reroot_otel_interval, _maybe_reroot_otel_interval, _end_otel_interval_span, _end_otel_train_span, _end_otel_job_spans, _otel_telemetry_active, _otel_mark_goodput, _backdated_otel_span)
• 9 module-level _otel_* globals holding span handles, context tokens and idempotence flags
• the SIGTERM handler, force-flush, graceful-drain and exit-hook installer

Roughly 400-500 lines will move to a dedicated module, leaving training.py with imports plus the call-site usage that should have been the only footprint all along: 21 with _otel_managed_span(...), 6 _otel_sg_enabled(...), and a handful of attribute stamps.

Deliberately not done in this PR: it would rewrite the exact training.py diff reviewers have already read, and the span lifecycle has no unit-test coverage here (the four telemetry tests on main cover megatron/core/telemetry, not this), so a move touching global mutation ordering and atexit/SIGTERM idempotence would be unguarded. As a standalone follow-up the diff is almost entirely a move, reviewable with git diff -M, and can land with tests for the lifecycle.

Also queued: port the GenAI-semconv inference instrumentation to dynamic_text_gen_server/, and restore the inference sections of the Observability guide alongside it (they were removed here in f007fd2, since this PR no longer ships that instrumentation).

🤖 Generated with Claude Code

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[santhnm2]

This file is pretty much deprecated, changes should go in megatron/core/inference/text_generation_server/dynamic_text_gen_server/text_generation_server.py

[deepakn94]

Good catch — reverted text_generation_server.py to main in f521a8e, so this PR no longer touches the deprecated module at all.

Porting the GenAI-semconv spans and metrics to dynamic_text_gen_server/text_generation_server.py will come as a follow-up rather than being reworked here, to keep this PR to the training-loop instrumentation.

[ahmadki]

The entire test suit should be dropped, these tests belong to Lens - some of them are even direct copies

[asolergi-nv]

/claude review

[claude]

"full" isn't a valid preset — MegatronSpanGroup._PRESETS defines default, per_step, profiling, all.

Suggested change

	'("default", "per_step", "full", "all") or individual group names '
	'("default", "per_step", "profiling", "all") or individual group names '

[claude]

Span is named megatron.startup.launch_script (training.py:497), not launch_script_setup — the comment 30 lines below (line 492) already uses the correct name.

Suggested change

	# through env vars (see megatron.startup.launch_script_setup/.container_load
	# through env vars (see megatron.startup.launch_script/.container_load

[claude]

Light review — mostly docs/tests. Two single-line typos left as inline suggestions; the rest below.

1. The observability docs published by this PR describe the pre-refactor span taxonomy

docs/index.md adds the Observability toctree, so these pages go from orphaned to user-facing — but they still document span names the code does not emit:

documented	actually emitted
megatron.pretrain	megatron.startup + megatron.train
megatron.train_step	megatron.train.iteration
megatron.forward_backward	megatron.train.iteration.forward_backward
megatron.optimizer_step	megatron.train.iteration.optimizer
megatron.save_checkpoint* / megatron.load_checkpoint*	megatron.checkpoint.save* / megatron.checkpoint.load*
megatron.model_init	megatron.startup.model_init

Affected: docs/user-guide/observability/index.md:50-55 and span-groups.md:32-44, 68-92, 103-114. Two more gaps in the same set:

• span-groups.md:22-25 preset table omits the first_iteration / trace_region groups and the profiling preset that span_groups.py defines (and lists full, matching the arguments.py typo).
• pipeline-parallel.md documents broadcast_trace_context and create_linked_span, and span-groups.md documents megatron.pp.recv_forward.linked — grep finds none of these anywhere under megatron/. Since the trace is now segmented via Links in training.py instead, that page should be rewritten or left out of the toctree.

2. megatron.training.loss exports one arbitrary key

megatron/training/training.py:3572-3577:

_loss_keys = [k for k in total_loss_dict if k not in _meta_keys]
if _loss_keys:
    _otel_loss_snapshot = total_loss_dict[_loss_keys[0]].item() / float(...)

_loss_keys[0] depends on total_loss_dict insertion order, so for models returning more than one loss key (e.g. MoE aux loss) the exported metric silently tracks whichever key happened to land first, and can change between configs. The log path a few lines down averages every non-meta key. Suggest pinning it to a canonical key and skipping the metric when absent:

_loss_key = "lm loss" if "lm loss" in total_loss_dict else None
if _loss_key is not None:
    _otel_loss_snapshot = total_loss_dict[_loss_key].item() / float(
        max(1, total_loss_dict.get(advanced_iters_key, 1))
    )

3. _otel_shutdown_done survives destroy_global_state()

_otel_shutdown_done (training.py:364) is module-level and never reset, but destroy_global_state() -> destroy_global_vars() does clear _GLOBAL_TELEMETRY_HANDLE. A process that calls pretrain() twice (unit tests, notebooks, multi-stage drivers) gets a fresh telemetry handle on the second run with _otel_shutdown_done = True, so _end_otel_job_spans() returns early and the second run never flushes its spans. Suggest clearing it alongside:

def destroy_global_state():
    global _otel_shutdown_done
    _otel_shutdown_done = False
    destroy_global_vars()
    ...

(_otel_exit_hooks_installed should stay True — the atexit/SIGTERM handlers are not uninstalled.)

4. No tests added

The PR body checks "I have added relevant unit tests," but tests/unit_tests/telemetry/ still contains only the four files from #6415 and the diff adds no tests elsewhere. The pure-Python helpers here are unit-testable without a GPU or a live collector:

• build_telemetry_resource_attrs(args) (global_vars.py) — SLURM/NVRx env var mapping, uuid5 determinism for job_uuid/run_uuid, _detect_gpu_identity returning {} when pynvml is unavailable.
• build_otel_worker_bootstrap(args) (async_utils.py) — the returned dict has to stay picklable across the spawn boundary, and inert when --otel-enabled is off.
• _maybe_reroot_otel_interval / _backdated_otel_span (training.py) — the re-root cadence vs. save_interval (including the save_interval <= 0 never-re-root path) is pure arithmetic and easy to get wrong on resume.

The no-op fallbacks in megatron/core/telemetry/fallbacks.py mean these run fine without nemo-lens installed.

Clean

No new direct parallel_state.get_*_group() / get_*_group() reads in megatron/core — the new instrumentation in distributed_data_parallel.py, p2p_communication.py, schedules.py, transformer_layer.py, and mamba_layer.py is decorator-only. The async-checkpoint call_idx tagging matches the upstream return types (schedule_async_request -> int, maybe_finalize_async_calls -> List[int]), and the record_training_metrics kwargs match its signature.

[asolergi-nv]

There are no critical issues to address, but I wonder if we should start regulating which comments belong in Python files and which ones should be in the documentation.

[asolergi-nv]

@maanug-nv This won't be directly exposed in MBridge, right? Should we include OTel in LoggerConfig?

[asolergi-nv]

Do we need ALL these comments & details here?

[rhewett-nv]

There are no critical issues to address, but I wonder if we should start regulating which comments belong in Python files and which ones should be in the documentation.

Comments used to be for people but now they are for agents. Having the documentation closer to the code seems to work better for them...