Skip to content

Tensor Metrics 1: Add framework for training metrics - #6647

Merged
jaredcasper merged 2 commits into
NVIDIA:mainfrom
philipcmonk:pmonk/tensor-metrics-framework
Sep 17, 2026
Merged

jaredcasper merged 2 commits into
NVIDIA:mainfrom
philipcmonk:pmonk/tensor-metrics-framework

Conversation

@philipcmonk

@philipcmonk philipcmonk commented Aug 19, 2026 •

Copy link
Copy Markdown
Contributor

This PR introduces a general training metrics framework. This is a reimplementation and generalization of #5492, and it's split out of #6552.

There are dozens of metrics that we want to keep track of during training, and currently we have to thread each one separately through the model. This is bug-prone and delicate, and it multiplies the complexity of making other changes that affect the sharding of tensors.

Here, we add a system that decouples obtaining a tensor from the metric computation itself. This lets users easily add metrics without changes to MCore.

Check the committed README at megatron/training/tensor_metrics/README.md for an overview of the system and how to add metrics.

This is part 1 of a four-PR stack:

  1. This PR: TensorMetric core framework
  2. Tensor Metrics 2: Expose MCore tensors through observation hooks #6648: 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

This PR contains only the authoring API, executor, documentation, and unit tests. It does not discover tensors, integrate with the training loop, register concrete metrics, or write logging results. Those pieces are added by later PRs in the stack.

@copy-pr-bot

copy-pr-bot Bot commented Aug 19, 2026

Copy link
Copy Markdown

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

Copy link
Copy Markdown
Contributor Author

/ok to test

@philipcmonk

Copy link
Copy Markdown
Contributor Author

/ok to test

@philipcmonk

Copy link
Copy Markdown
Contributor Author

/ok to test

Signed-off-by: Philip Monk <pmonk@nvidia.com>
@chtruong814
chtruong814 requested a review from a team as a code owner September 16, 2026 13:38
@chtruong814

Copy link
Copy Markdown
Contributor

/ok to test 3c13e61

@svcnvidia-nemo-ci svcnvidia-nemo-ci added the Approved All necessary approvals have been made label Sep 17, 2026
@jaredcasper
jaredcasper added this pull request to the merge queue Sep 17, 2026
@nemo-automation-bot

Copy link
Copy Markdown

🔄 Merge queue validation started!

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

Merged via the queue into NVIDIA:main with commit 3fb6298 Sep 17, 2026
106 of 108 checks passed

This branch was previously deployed

2 inactive deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Approved All necessary approvals have been made complexity: high nemotron

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants