core.tensor_observation

Scoped notifications for observing short-lived tensors in Megatron Core.

Module Contents

Functions

capture_tensor_observations	Notify observer about the selected tensor source kinds within this scope.
suspend_tensor_observations	Temporarily suppress notifications to the active tensor observer.
is_observing_tensor	Return whether source_kind has an active observer.
observe_tensor	Notify the active observer about one tensor produced during the logical forward.
observe_layer_residuals	Observe a layer’s incoming residual accumulator and net contribution.
_notify_observer

Data

TensorObservationCallback
_ACTIVE_OBSERVER
_OBSERVATION_SUSPENDED

API

core.tensor_observation.TensorObservationCallback

None

core.tensor_observation._ACTIVE_OBSERVER: contextvars.ContextVar[tuple[core.tensor_observation.TensorObservationCallback, frozenset[str]] | None]

‘ContextVar(…)’

core.tensor_observation._OBSERVATION_SUSPENDED: contextvars.ContextVar[bool]

‘ContextVar(…)’

core.tensor_observation.capture_tensor_observations(

observer: core.tensor_observation.TensorObservationCallback,

source_kinds: frozenset[str],

) → collections.abc.Iterator[None]

Notify observer about the selected tensor source kinds within this scope.

Parameters:

• observer – Callback receiving the owner, local name, source kind, tensor, and its optional tensor-parallel, sequence, and batch dimensions.

• source_kinds – Source kinds to notify. Other observation sites remain no-ops.

core.tensor_observation.suspend_tensor_observations() → collections.abc.Iterator[None]

Temporarily suppress notifications to the active tensor observer.

Activation checkpoint implementations use this around their backward recomputation so the original forward is observed exactly once.

core.tensor_observation.is_observing_tensor(source_kind: str) → bool

Return whether source_kind has an active observer.

core.tensor_observation.observe_tensor(

owner: object,

name: str,

source_kind: str,

tensor: torch.Tensor,

*,

tp_shard_dim: int | None = None,

sequence_dim: int | None = None,

batch_dim: int | None = None,

) → None

Notify the active observer about one tensor produced during the logical forward.

Activation checkpoint implementations suppress notifications during their backward recomputation, so the original no-grad forward remains the single observed execution.

Parameters:

• owner – Model object used to resolve a canonical site name.

• name – Name local to owner.

• source_kind – Stable source category used by metric selection.

• tensor – Short-lived observed tensor.

• tp_shard_dim – Tensor dimension sharded over tensor parallel ranks, or None when the tensor is replicated over tensor parallel ranks.

• sequence_dim – Tensor dimension identifying the context-parallel sequence population, or None when that population is not represented by one tensor dimension.

• batch_dim – Tensor dimension identifying the data- and gradient-tensor-parallel batch population, or None when it is not represented by one tensor dimension.

core.tensor_observation.observe_layer_residuals(

layer: object,

accumulator: torch.Tensor,

output: torch.Tensor,

) → None

Observe a layer’s incoming residual accumulator and net contribution.

With per-layer CUDA graphs, the contribution is reconstructed as output - accumulator outside the captured layer. Full-iteration CUDA graphs include this call and do not replay Python observation notifications. The contribution is computed without autograd only when a contribution metric is active.

Parameters:

• layer – Layer owning the residual sites.

• accumulator – Residual stream entering the layer.

• output – Residual stream leaving the layer.

core.tensor_observation._notify_observer(

owner: object,

name: str,

source_kind: str,

tensor: torch.Tensor,

tp_shard_dim: int | None,

sequence_dim: int | None,

batch_dim: int | None,

) → None