[philipcmonk]

Adds dormant, scoped MCore hooks that expose selected short-lived tensors to an external observer without depending on the training tensor-metric implementation. These will be used in a later training-focused PR.

This is part 2 of a four-PR stack. However, this and PR 1 are completely independent and can be reviewed/merged in either order (just rebase this onto main if this is to be merged first).

1. Tensor Metrics 1: Add framework for training metrics #6647: TensorMetric core framework
2. This PR: MCore tensor observation sources
3. Tensor Metrics 3: Add tensor metric definitions #6649: metric definitions
4. Tensor Metrics 4: Integrate tensor observation into training #6650: training-loop integration

Some useful training diagnostics require tensors that do not survive until the optimizer step, including residual streams and router or output logits.

MCore should expose these values with minimal instrumentation while leaving metric selection, distributed reduction, scheduling, and logging outside megatron/core.

Design

megatron.core.tensor_observation provides a scoped callback mechanism based on a context variable. Outside an active observation scope, every hook is a cheap no-op.

The callback receives:

• the owning module;
• a name local to that owner;
• a stable "source kind";
• the tensor;
• its optional tensor-parallel shard dimension.

The source hook does not import or understand TensorMetric classes.

Observation sites

This PR exposes:

• residual accumulators;
• net per-layer residual contributions;
• GPT and HybridModel output logits;
• MTP logits;
• raw router logits before forced load-balancing modifications;
• normalized router decision scores;
• compact per-sequence router diagnostics.

We will later want to expose more observation sites, but this is the list of sites that we have so far used for metrics. More can be added when needed.

Canonical parameter names

This PR also adds CanonicalParameterNameMap, which creates stable logical names across:

• pipeline-local layer numbering;
• MTP layers;
• grouped and sequential experts;
• expert-parallel ownership;
• tied parameters and model chunks.

The logic included here is replicated ad hoc in several places in the codebase, so it seems worth centralizing it here. However, this could be moved into training/tensor_metrics if needed.

[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.

[philipcmonk]

/ok to test

[philipcmonk]

/ok to test

[philipcmonk]

/ok to test

[yashaswikarnati]

when we do THD / no inter document masking, i wonder if the metrics we capture here needs to take into account per-sequence

[philipcmonk]

Yes, right now all of these are per-sample not per-document. At some point it would be nice to support THD, but I think it's fine to defer for now?

[yashaswikarnati]

sgtm

[nemo-automation-bot]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/35283924913

[chtruong814]

Ran CI as part of an aggregated branch that included multiple PRs in the queue.
#7459

CI passed
https://github.com/NVIDIA/Megatron-LM/actions/runs/35291473099?pr=7459

Will fast merge this

[nemo-automation-bot]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/35302211621