Make TrainState authoritative for mutable training progress
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
Parent initiative: #3554
## Context
Megatron-LM currently stores mutable training progress across `args`, local variables, and checkpoint return values. This includes the current iteration, consumed and skipped samples, validation samples, phase flags, and completed floating-point operations.
Megatron-Bridge collects these values in a `TrainState` object. Megatron-LM needs the same runtime contract so training progress is separate from configuration and has one source of truth.
## Scope
- Add a Megatron-LM `TrainState` with the following fields:
- `step`
- `consumed_train_samples`
- `skipped_train_samples`
- `consumed_valid_samples`
- `floating_point_operations_so_far`
- `do_train`
- `do_valid`
- `do_test`
- Initialize `TrainState` for new training runs.
- Pass the active state from its initialization point into the training, evaluation, and logging paths that read or update training progress. Alternatively, make the `TrainState` a global variable, as is planned for the config container.
- Expose the active state to the checkpoint boundary. Checkpoint serialization and reconstruction are tracked in #6004.
- Replace reads and mutations of the corresponding `args` fields and local progress variables with `TrainState`.
- Keep configuration in `PretrainConfigContainer`. Do not move mutable progress into config dataclasses.
## Completion criteria
- `TrainState` is the source of truth for all fields listed above during training.
- Training and skipped steps update the step, sample, and floating-point operation counters correctly.
- Evaluation updates consumed validation samples correctly.
- Data-loader initialization sets the train, validation, and test phase flags on `TrainState`.
- Logging, scheduling, and checkpoint callers receive progress from `TrainState` rather than `args`.
- Tests cover initialization and state changes during training, skipped steps, evaluation, and phase selection.
- Any remaining compatibility copies in `args` are documented and are not treated as authoritative.
## Out of scope
- Saving or loading `train_state.pt`. Checkpoint serialization and legacy checkpoint fallback are tracked in #6004.
- Migrating Megatron-Bridge `GlobalState`, `FaultToleranceState`, loggers, or other runtime singletons.
Contributor guide
Assessment
This issue has not been assessed yet.