Make PretrainConfigContainer and TrainState authoritative
- 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
Assessment
This issue has not been assessed yet.