[philipcmonk]

This PR introduces a general training metrics subsystem. This is a reimplementation and generalization of #5492.

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.md for an overview of the system and how to add metrics.

This is a large PR, I recommend it be reviewed in pieces (in whatever order is natural to you):

• megatron/training/tensor_metrics/{core,observer}.py: The core system, with Observer and Executor implementations plus the TensorMetric interface.
• megatron/training/tensor_metrics/{definitions,examples,router_metrics}.py: Initial set of implemented metrics.
• megatron/training/training.py: How the observer is called from the training loop.
• megatron/core/parameter_names.py and megatron/training/tensor_metrics/optimizer_sources.py: Parameter/wgrad discovery across optimizers, canonical naming, manifests, and placement metadata.
• Remainder of megatron/core/*: How activation tensors are observed.
• tests/*: Unit tests.

This is a draft while the design is reviewed. The implementation is generally complete and I've started using it to debug issues, but details are not fully stable yet.

[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

[janEbert]

I think this needs to be split up, just based on the 8.9k LOC difference. :d
Is it possible in any way to abstract different features included here into different PRs?

[philipcmonk]

I think the natural boundaries are probably these:

• All MCore changes (new parameter names module, new tensor observation module, hooks into various places to observe tensors): ~500 lines + 500 lines tests
• TensorMetric framework: 900 lines + 500 lines tests
• Training integration: 1400 lines + 1500 lines tests
• Metric definitions: 2300 lines + 1200 lines tests

We could make those as four separate PRs or combine some of the training PRs since they'll probably have the same reviewers.

[philipcmonk]

This is now split up into

• Tensor Metrics 1: Add framework for training metrics #6647
• Tensor Metrics 2: Expose MCore tensors through observation hooks #6648
• Tensor Metrics 3: Add tensor metric definitions #6649
• Tensor Metrics 4: Integrate tensor observation into training #6650