refactor(training): initialize runtime services after config validation by cuichenx · Pull Request #8272 · NVIDIA/Megatron-LM

cuichenx · 2026-10-15T21:30:46Z

I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Move runtime-service initialization into initialize_megatron(), after run-config validation and before distributed/RNG setup.

Training entrypoints construct/register the config and call pretrain(); inference entrypoints call initialize_megatron() without separately constructing runtime services.
Resolve tokenizer-derived vocabulary centrally after tokenizer construction, before writers and other services consume it. Preserve explicit and checkpoint-restored vocabulary.
Forward build_tokenizer=False for heterogeneous MiMo, which supplies its own tokenization.
Keep the lightweight checkpoint-converter bootstrap, validating its config before creating services.

Most entrypoint changes remove the old explicit runtime/vocabulary calls and their imports. This does not change config ownership, training/inference selection, or distributed/RNG behavior.

This is the separate startup-order follow-up discussed in #7778; #7778 does not depend on it. Pending #7778 and #8149 changes will need to be merged here once they land.

Issue tracking

Related to #3554 and #7778.

Contribution process

Pre-checks

I have added relevant unit tests
I have added relevant functional tests
If this PR adds or changes a GPU kernel (Triton, jit_fuser/torch.compile, CUDA extension, TE or external-library dispatch, or a scatter/index accumulation), I have added or updated its bit-exact determinism test and registered it in tests/unit_tests/determinism/kernels/manifest.py
I have added proper typing to my code
I have added relevant documentation
I have run the autoformatter on my PR

Validation at 77521cccbe1451984d949873ccc647ec34e9bc45: 405 targeted tests passed on each of two GPU ranks across argument adapters, initialization lifecycle, inference, RNG, profiling, kernel warmup, autotuning, run config, and container validation. Existing lifecycle tests cover validation before side effects, tokenizer-derived/explicit/checkpoint vocabulary, tokenizer suppression, and the RL entrypoint's single initialization boundary. Independent full-diff review completed.

Changed-test formatting and diff checks pass. The all-files pre-commit run reports existing lint/format issues in untouched files; those are left unchanged to keep this diff focused.

Code review

Draft for startup-order review. No GPU kernels or numerical operations changed.

Signed-off-by: Chen Cui <chcui@nvidia.com>

copy-pr-bot · 2026-10-15T21:30:52Z

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

cuichenx

the important two refactors are

cuichenx · 2026-10-15T22:00:18Z

        )

+    # Runtime services consume the validated config, before distributed/RNG setup.
+    initialize_runtime_services(args, build_tokenizer=build_tokenizer, training=training)


cuichenx · 2026-10-15T22:00:25Z

+
    if build_tokenizer:
        _build_tokenizer(args)
+    resolve_tokenizer_vocab_size(get_run_config(), getattr(args, "padded_vocab_size", None))


before, this wasn't always running every time initialize_runtime_services() was called. is it fine to do that now? should it be on some condition, eg if build_tokenizer or if build_tokenizer and training ?

maanug-nv · 2026-10-16T00:37:37Z

        skip_model_parallel_init=True,
        p2p_communicator=communicator,
        pg_collection=topology.schedule_pg_collection,
+        build_tokenizer=False,


note for later, this should be handled by a cfg.tokenizer=None

maanug-nv

just 1 minor question.

maanug-nv · 2026-10-16T00:55:06Z

+
    if build_tokenizer:
        _build_tokenizer(args)
+    resolve_tokenizer_vocab_size(get_run_config(), getattr(args, "padded_vocab_size", None))


before, this wasn't always running every time initialize_runtime_services() was called. is it fine to do that now? should it be on some condition, eg if build_tokenizer or if build_tokenizer and training ?

maanug-nv

addressed offline. good to merge now.

Signed-off-by: Chen Cui <chcui@nvidia.com>

tdene · 2026-10-16T09:54:19Z

    assert cfg_container is not None

    _ensure_var_is_not_initialized(_GLOBAL_ARGS, 'args')
+    cfg_container.validate()


There are a couple config in the inference container that are environment-dependent: --inference-kv-cache-management-mode and --inference-static-kv-memory-pointers. They both require torch_memory_saver to be installed under some of the argument values.

This validation check leads to an odd case, where a checkpoint that was RL-trained on a machine that has torch_memory_saver can fail validation if the checkpoint is loaded on a machine without torch_memory_saver.

I think that before this PR, most codepaths already did validation and would encounter this issue. But I see that tools/checkpoint/convert.py specifically did not validate before this, and this is new behavior.

I don't think this is a problem. But I'm noting it down here, because it's a strange interaction.

Signed-off-by: Chen Cui <chcui@nvidia.com>

nemo-automation-bot · 2026-10-18T08:09:17Z

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/39110341563

nemo-automation-bot · 2026-10-18T19:08:46Z

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/39157400748

…tion config migration Merge main including NVIDIA#8272. Resolve batch defaults during one-time finalization, keep repeated validation read-only, and consume config-owned batch policy throughout training. Signed-off-by: Chen Cui <chcui@nvidia.com>

Resolve conflicts with the runtime-services refactor (NVIDIA#8272), the config validate/finalize split, and disaggregated inference process groups. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Asha Anoosheh <aanoosheh@nvidia.com>

Repair main-to-dev sync compatibility for the 2026-10-20 slice (main 13a9d46 merged into dev ff4fbc9): - moe/router: InferenceTopKRouter._compiled_topk_routing accepts and forwards precomputed_indices, matching the base router call site. - tensor_parallel/random: MHCCheckpointManager subclasses CheckpointWithoutOutputManager so main's checkpoint consumers accept it. - moe/token_dispatcher: _HybridEPManager.setup_metadata clears the stale num_permuted_tokens budget when no capacity factor is configured. - tests: align dev-only tests with main API changes (test_training TrainingConfig and TrainState, test_mhc_schedule _model_chunk_state, test_hybrid_hash_routing tensor_parallel import, test_route_metadata_cpu token_dispatcher fake, MTP argmax acceptance-count test). - tests: test_train_step_schedule_plumbing provides the cfg.train, cfg.validation, cfg.ddp and TrainState inputs main's train_step and setup_model_and_optimizer now read. - tests: test_virtual_experts quantile cases select the micro_batch estimation scope that main's quantile estimator implements; the merged default keeps dev's global_batch scope. - tests/test_utils/recipes/gb200/unit-tests.yaml: drop the duplicated tests/unit_tests/ssm/test_gated_delta_net.py bucket produced by the merge, restoring the one-bucket-per-file unit-test scope contract. - tests: test_hybrid_distillation_entry expects the multimodal_dev entry point to leave runtime services and tokenizer vocabulary resolution to initialize_megatron, which main NVIDIA#8272 made their owner. - examples/mimo_bagel/train.py: read the training iteration from get_train_state() for the PP>1 data seed and the loss-alignment audit log, because main NVIDIA#7932 no longer sets args.curr_iteration or args.iteration (review 5609274485). - pyproject.toml: carry main NVIDIA#7494's [tool.setuptools.package-data] entry so wheel and sdist installs ship megatron/core/allocator/csrc/vmm_symm_allocator.cpp (review 5609274485). Signed-off-by: svcnemo-autobot <svcnemo-autobot@nvidia.com>

Repair main-to-dev sync compatibility for the 2026-10-20 slice (main 0156c21 merged into dev ae76c60): - moe/router: InferenceTopKRouter._compiled_topk_routing accepts and forwards precomputed_indices, matching the base router call site. - tensor_parallel/random: MHCCheckpointManager subclasses CheckpointWithoutOutputManager so main's checkpoint consumers accept it. - moe/token_dispatcher: _HybridEPManager.setup_metadata clears the stale num_permuted_tokens budget when no capacity factor is configured. - tests: align dev-only tests with main API changes (test_training TrainingConfig and TrainState, test_mhc_schedule _model_chunk_state, test_hybrid_hash_routing tensor_parallel import, test_route_metadata_cpu token_dispatcher fake, MTP argmax acceptance-count test). - tests: test_train_step_schedule_plumbing provides the cfg.train, cfg.validation, cfg.ddp and TrainState inputs main's train_step and setup_model_and_optimizer now read. - tests: test_virtual_experts quantile cases select the micro_batch estimation scope that main's quantile estimator implements; the merged default keeps dev's global_batch scope. - tests/test_utils/recipes/gb200/unit-tests.yaml: drop the duplicated tests/unit_tests/ssm/test_gated_delta_net.py bucket produced by the merge, restoring the one-bucket-per-file unit-test scope contract. - tests: test_hybrid_distillation_entry expects the multimodal_dev entry point to leave runtime services and tokenizer vocabulary resolution to initialize_megatron, which main NVIDIA#8272 made their owner. - examples/mimo_bagel/train.py: read the training iteration from get_train_state() for the PP>1 data seed and the loss-alignment audit log, because main NVIDIA#7932 no longer sets args.curr_iteration or args.iteration (review 5609274485). - pyproject.toml: carry main NVIDIA#7494's [tool.setuptools.package-data] entry so wheel and sdist installs ship megatron/core/allocator/csrc/vmm_symm_allocator.cpp (review 5609274485). Signed-off-by: svcnemo-autobot <svcnemo-autobot@nvidia.com>

cuichenx added 8 commits September 30, 2026 22:41

refactor(training): migrate DistributedInitConfig

374faff

Signed-off-by: Chen Cui <chcui@nvidia.com>

Merge main and refresh distributed config ownership

1199bf4

Signed-off-by: Chen Cui <chcui@nvidia.com>

chore: preserve original copyright year

1cefa93

Signed-off-by: Chen Cui <chcui@nvidia.com>

test(training): stabilize checkpoint setup and consolidate dist coverage

ac091bd

Signed-off-by: Chen Cui <chcui@nvidia.com>

Merge main and fix distributed config integration

f994dea

Signed-off-by: Chen Cui <chcui@nvidia.com>

refactor(training): initialize TP user buffers from model config

db5da7f

Signed-off-by: Chen Cui <chcui@nvidia.com>

refactor(training): generate fake process group flag from config

4b55bec

Signed-off-by: Chen Cui <chcui@nvidia.com>

refactor(training): initialize runtime services after config validation

77521cc

Signed-off-by: Chen Cui <chcui@nvidia.com>

cuichenx commented Oct 15, 2026

View reviewed changes

cuichenx marked this pull request as ready for review October 15, 2026 22:00

cuichenx requested review from a team as code owners October 15, 2026 22:00

github-actions Bot added the complexity: medium label Oct 15, 2026

copy-pr-bot Bot deployed to public October 15, 2026 22:04 Active

cuichenx mentioned this pull request Oct 15, 2026

refactor(training): migrate DistributedInitConfig #7778

Merged

5 of 7 tasks

maanug-nv reviewed Oct 16, 2026

View reviewed changes

maanug-nv approved these changes Oct 16, 2026

View reviewed changes

style(tests): format startup initialization import

ff221ba

Signed-off-by: Chen Cui <chcui@nvidia.com>

copy-pr-bot Bot deployed to test October 16, 2026 04:19 Active

copy-pr-bot Bot deployed to public October 16, 2026 04:21 Active

Phlip79 approved these changes Oct 16, 2026

View reviewed changes

tdene approved these changes Oct 16, 2026

View reviewed changes

Merge main into distributed config migration

47a1b35

Signed-off-by: Chen Cui <chcui@nvidia.com>

Merge branch 'main' into chcui/omar/validated-runtime-startup

e4f71c6

cuichenx enabled auto-merge October 18, 2026 07:09

copy-pr-bot Bot had a problem deploying to test October 18, 2026 07:10 Error

test(training): use real RL config in startup lifecycle test

b28f727

Signed-off-by: Chen Cui <chcui@nvidia.com>

copy-pr-bot Bot deployed to public October 18, 2026 07:13 Active

copy-pr-bot Bot deployed to test October 18, 2026 07:16 Active

copy-pr-bot Bot deployed to public October 18, 2026 07:20 Active

cuichenx added this pull request to the merge queue Oct 18, 2026

github-merge-queue Bot removed this pull request from the merge queue due to failed status checks Oct 18, 2026

Merge branch 'main' into chcui/omar/validated-runtime-startup

5b3bd64

tdene enabled auto-merge October 18, 2026 16:27

copy-pr-bot Bot deployed to test October 18, 2026 16:28 Active

copy-pr-bot Bot deployed to public October 18, 2026 16:30 Active

tdene added this pull request to the merge queue Oct 18, 2026

Merged via the queue into NVIDIA:main with commit ce852bf Oct 18, 2026
110 checks passed

svcnemo-autobot mentioned this pull request Oct 21, 2026

chore: AUT-2598 nightly sync main into dev (2026-10-20) #8402

Merged

svcnvidia-nemo-ci mentioned this pull request Oct 22, 2026

chore: nightly sync main into dev (22_10_2026) #8434

Open

Conversation

cuichenx commented Oct 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Issue tracking

Contribution process

Pre-checks

Code review

Uh oh!

copy-pr-bot Bot commented Oct 15, 2026

Uh oh!

cuichenx left a comment

Choose a reason for hiding this comment

Uh oh!

cuichenx Oct 15, 2026

Choose a reason for hiding this comment

Uh oh!

cuichenx Oct 15, 2026

Choose a reason for hiding this comment

Uh oh!

maanug-nv Oct 16, 2026

Choose a reason for hiding this comment

Uh oh!

maanug-nv Oct 16, 2026

Choose a reason for hiding this comment

Uh oh!

maanug-nv left a comment

Choose a reason for hiding this comment

Uh oh!

maanug-nv Oct 16, 2026

Choose a reason for hiding this comment

Uh oh!

maanug-nv left a comment

Choose a reason for hiding this comment

Uh oh!

tdene Oct 16, 2026

Choose a reason for hiding this comment

Uh oh!

nemo-automation-bot Bot commented Oct 18, 2026

Uh oh!

Uh oh!

nemo-automation-bot Bot commented Oct 18, 2026

Uh oh!

Uh oh!

This branch was successfully deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

cuichenx commented Oct 15, 2026 •

edited

Loading