NVIDIA / NVIDIA/Megatron-LM

Make PretrainConfigContainer and TrainState authoritative

Open
#3,554 0 comments 0 reactions 0 assignees View on GitHub
enhancement
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 6h
Merged PRs (30d)
271

Description

## Goal

Make `PretrainConfigContainer` the source of truth for configuration used by the Megatron-LM training loop. Make `TrainState` the source of truth for mutable training progress.

Argparse-based entry points must continue to work. They should construct the same configuration container rather than remain a second source of truth.

## Current state

`pretrain()` accepts a `PretrainConfigContainer`, but initialization and training-loop code still read extensively from the global `args` namespace.

The existing container also lacks some configuration needed by the training loop. Mutable progress such as iteration, consumed samples, phase flags, and completed floating-point operations is still stored across `args` and other checkpoint fields.

## Scope
- Define a mapping from `args` fields to config attributes. Ensure every `args` field has a destination or mark it explicitly out of scope. (currently tracked in a spreadsheet)
- Complete `PretrainConfigContainer` coverage for existing training configuration.
- Add the required dataset, fault-tolerance, distillation, RL, and VLM configuration.
- Expand `pretrain_cfg_container_from_args()` to translate every in-scope argparse option into the appropriate nested config.
- Register the active `PretrainConfigContainer` as a process-global object and expose it through `get_cfg()`.
- Migrate initialization and code downstream of `pretrain()` from `get_args()` and `args.*` to `get_cfg()` and nested config fields.
- Move mutable training progress into `TrainState` and pass it through training, evaluation, checkpointing, and resume paths.
- Replace configuration fields that represent runtime helpers, such as data-parallel size or world size, with helper functions.
- Move validation and dependent-field derivation into the appropriate config dataclass, container, or runtime object.
- Preserve the behavior of every supported argparse option.

## Completion criteria

- `PretrainConfigContainer` supplies configuration throughout the agreed training-loop boundary.
- `TrainState` supplies mutable training progress throughout that boundary.
- Remaining `get_args()` or `args.*` use inside the boundary is documented and explicitly approved.
- Argparse-based entry points construct configuration and training state equivalent to direct dataclass-based entry points.
- Representative argparse configurations produce equivalent config objects and derived inputs at the model, optimizer, scheduler, distributed, checkpoint, logging, tokenizer, and dataset seams.
- Each new top-level config has translation tests and coverage for its direct consumers.
- The full convergence and performance test suites remain within agreed tolerances.

## Out of scope

- Removing argparse-based entry points.
- Removing `get_args()` from callers outside the migrated training-loop boundary.
- Making configuration immutable.
- Changing checkpoint metadata serialization. That work is tracked in #2319.

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.