NeMo Lens

View as Markdown

NeMo Lens (nemo.lens) is a shared OpenTelemetry instrumentation library for the NVIDIA NeMo ecosystem. It provides unified tracing, metrics, and optional log bridging across distributed training, inference, and RL workloads. NeMo Lens keeps hot-path overhead low with span-group gating, where disabled groups use a fast frozenset lookup and create no span objects. Consumer projects can integrate it as an optional dependency through canonical no-op fallbacks.

Key Features

  • Low overhead when disabled: span group gating is a frozenset lookup; no span objects are allocated when a group is off
  • Distributed-training aware: rank-based export strategies, cross-rank trace correlation
  • Framework-agnostic primitives: managed_span, trace_fn, span_cm work everywhere
  • OTel idiomatic: real TracerProvider / MeterProvider, W3C propagation, standard semconv
  • Optional dependency: consumers import using try/except ImportError; NeMo Lens ships canonical no-op fallbacks
  • Pluggable: custom exporters, custom samplers, custom span groups
  • Resource auto-detection: SLURM, Kubernetes, local GPU count out of the box

Consumers

nemo-lens is consumed by three NeMo-ecosystem projects, each extending it with domain-specific span groups:

ConsumerDomainSpan group extension
Megatron-LMTransformer pre-trainingMegatronSpanGroup: microbatch, layer, communication, activation_offload, data_loading, inference
NeMo RLRLHF, GRPO, DPORLSpanGroup: rollout, generation, logprob, reward, advantage, policy_update, reference_policy, data_processing
NeMo GymRL environmentsGymSpanGroup: server, rollout_collection, verify, aggregate_metrics

NeMo Lens emits OTLP; choosing and running an observability stack is up to you. Quick links for sending telemetry to common destinations:

For trying things out locally, see the demo docker-compose stack (proof of concept, not for production).