Load run_config.yaml into PretrainConfigContainer
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
Parent initiative: #2319
## Context
New checkpoints will store configuration in `run_config.yaml`. Older Megatron-LM checkpoints contain only the legacy `state_dict["args"]` metadata.
Both paths must produce a `PretrainConfigContainer`, which is the configuration source of truth inside the training loop. Loading `run_config.yaml` should not recreate a flat `argparse.Namespace`.
## Scope
- When `run_config.yaml` is present, deserialize it into `PretrainConfigContainer`.
- Use the loaded container as the checkpoint configuration.
- When `run_config.yaml` is absent, load the legacy `args` metadata and convert it with `pretrain_cfg_container_from_args`.
- Keep any field-name mappings required by the legacy `args` conversion in that fallback path.
- Apply the same validation and dependent-field derivation to configurations loaded through either path.
## Completion criteria
- Megatron-LM loads a Bridge-style `run_config.yaml` into `PretrainConfigContainer`.
- Megatron-LM can reconstruct `PretrainConfigContainer` from older checkpoints that contain only `state_dict["args"]`.
- Equivalent YAML and legacy metadata produce equivalent configuration containers.
- Tests cover both load paths, including known field-name differences in the legacy conversion.
Contributor guide
Research direction
Start by locating the checkpoint-loading code that handles run_config.yaml and state_dict["args"], then read PretrainConfigContainer and pretrain_cfg_container_from_args. Run the existing checkpoint tests before adding coverage for both paths, including legacy field-name mappings. Done means both formats yield equivalent validated containers with dependent fields derived.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 55/100