NVIDIA / NVIDIA/Megatron-LM

Separate configuration preparation from runtime initialization

Open
#7,324 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 6h
Merged PRs (30d)
271

Description

Parent initiative: #3554

## Goal

Make configuration preparation and runtime initialization explicit so the training entrypoints can construct and register `PretrainConfigContainer` before its downstream consumers run. Preserve current Megatron-LM behavior during this structural change.

## Context

`parse_and_validate_args()` currently applies checkpoint overrides and validation, then calls `set_global_variables(args)`. That registers args and constructs the tokenizer, microbatch calculator, loggers, timers, and other services. Entrypoints construct the model config and container afterward.

Tokenizer construction is a real prerequisite: it can derive `padded_vocab_size`, which GPT/hybrid model-config construction consumes. Moving all initialization after container construction would break that dependency.

## Scope

- Separate input parsing/checkpoint overrides/validation, tokenizer-dependent vocabulary resolution, container construction, and runtime-service initialization into explicit phases.
- Reuse `pretrain_cfg_container_from_args()` for container construction.
- Register the active container and expose `get_cfg()` with corresponding cleanup.
- Make tokenizer bootstrap explicit, preserve checkpoint-provided vocabulary precedence, and avoid constructing the tokenizer twice.
- Wire the in-scope training entrypoints through the explicit phases. Preserve the existing parse helper's behavior for callers outside the migration boundary through a compatibility wrapper where needed.
- Preserve required startup order, including early fault-tolerance monitoring, logging/timing, distributed setup, and restart behavior. Do not mechanically reorder services whose position affects behavior.
- Add focused startup tests and reusable baseline comparisons in the existing test structure.

This is the first initialization PR. The extracted helpers may still consume args; migrating their configuration ownership belongs to the follow-up subissue.

## Completion criteria

- Startup phases and their dependencies are explicit; the container is registered before migrated downstream consumers access it.
- Representative CLI configurations produce the same resolved model/config inputs and startup behavior as original MLM.
- Tests cover tokenizer-derived and explicitly/checkpoint-provided vocabulary sizes, single construction of runtime services, required initialization ordering, and cleanup/reinitialization.
- Existing training entrypoints and out-of-scope callers retain their supported behavior.

## Out of scope

- Making all initialization helpers or the full training loop config-authoritative.
- Porting all of `validate_args()` or migrating mutable progress to `TrainState` (#3555).
- New checkpoint metadata, callback APIs, or PEFT behavior.

Validation compares original MLM with refactored MLM.

Contributor guide

Open the contributing guide

Research direction

Start by tracing parse_and_validate_args() and pretrain_cfg_container_from_args() to map parsing, tokenizer-dependent vocabulary resolution, container construction, and runtime-service initialization. Review get_cfg() registration and cleanup, then use the existing test structure to compare original and refactored MLM startup behavior. Done means explicit phases, preserved ordering and compatibility, single service construction, and coverage for vocabulary, reinitialization, and cleanup.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Refactor
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.