Introduce a callback system for the Megatron-LM training loop
- 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
Assessment
This issue has not been assessed yet.