NVIDIA / NVIDIA/Megatron-LM

Introduce a callback system for the Megatron-LM training loop

Open
#7,266 0 comments 0 reactions 1 assignee Claimed by @maanug-nv View on GitHub
enhancement
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 6h
Merged PRs (30d)
271

Description

## Context

Callbacks are part of the first stage of consolidating the Megatron-LM and Megatron-Bridge training loops. They let integrations extend the training loop without adding framework-specific branches to it.

The implementation already exists in #5610. This issue tracks its review and merge.

## Included in #5610

- Adds `CallbackContext`, `Callback`, and `CallbackManager` to `megatron.training`.
- Adds hooks for setup, data initialization, training, logging, evaluation, and testing. This includes step-level hooks needed by RL training.
- Connects those hooks to the existing Megatron-LM training loop.
- Preserves user state across callback events.
- Adds unit tests for callback registration, ordering, dispatch, error propagation, context fields, and user state.

Implementation: #5610

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.