NVIDIA / NVIDIA/Megatron-LM

Load run_config.yaml into PretrainConfigContainer

Open
#6,003 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

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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.