NVIDIA / NVIDIA/Megatron-LM

Make TrainState authoritative for mutable training progress

Open
#3,555 1 comment 0 reactions 2 assignees Claimed by @cuichenx View on GitHub
enhancement
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

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.