[aroshanghias-nvd]

What does this PR do?

Install training callbacks on the configuration consumed by Megatron's scheduler. This fixes a silent gradient-synchronization failure when a provider, such as canonical Nemotron Omni, copies its configuration while constructing the model.

Why the existing binding is incorrect

setup() currently passes the provider configuration (or the GPT/Hybrid builder's transformer configuration) to _update_model_config_funcs(). This assumes that the selected object is also the runtime model configuration.

Canonical Omni breaks that identity assumption:

1. The provider copies its configuration into language_cfg and passes it to the canonical model.
2. The model constructor stores that copy as model.config.
3. After building the model and optimizer, Bridge installs finalize_model_grads_func, grad_scale_func, and applicable synchronization callbacks on the original provider.
4. Megatron reads get_model_config(model). Its finalization callback is still None, so the conditional scheduler call is skipped.

Conceptually, the existing path does this:

runtime_config = deepcopy(provider_config)
model.config = runtime_config
install_callbacks(provider_config)
# provider_config.finalize_model_grads_func is set
# model.config.finalize_model_grads_func is still None

This PR obtains the callback target from get_model_config(model[0]), matching the scheduler's lookup. The existing installer continues to construct the callbacks and per-chunk callback lists. The reduction algorithms and normalization settings are unchanged.

Training consequence—and why a run can look healthy

In the reproduced PP=1 configuration with overlap_grad_reduce=False, gradient finalization completes the required DP/CP reductions and applicable TP reductions for replicated parameters. The distributed optimizer assumes those operations have completed before it consumes its owned gradient slices. Without them, a shard can update from its owner's local contribution instead of the intended combined gradient.

A small example uses 512 zero-initialized parameters and fresh distributed Adam state. Every parameter receives a local contribution of +1 on DP rank 0 and -3 on DP rank 1. Rank 0 owns the first half of the optimizer update; rank 1 owns the second half.

Case	Existing binding	Corrected binding
DP=2, TP=1; contributions +1, -3	Owned gradients remain +1 and -3; 256 parameters decrease and 256 increase	The combined gradient is negative; all 512 parameters increase
TP=2, DP=1; replicated parameter contributions +1, -1	TP replicas update in opposite directions	Contributions cancel; all parameters remain zero

Parameter all-gather can still give DP replicas identical final weights after the incorrect updates. It assembles the already-updated parameter shards; it does not combine the gradients or recompute the optimizer updates. Every replica therefore receives the same mixture of shard updates based on different ranks' local gradients. Backward also continues to produce nonzero gradients with useful learning signal. Consequently, matching DP weights, stable losses, or successful training completion do not establish that synchronization was correct.

This is a training-correctness finding. It does not establish that existing trained models are unusable, quantify model-quality degradation, or demonstrate a full-model TMPE improvement. Affected historical runs require their exact provider, entrypoint, runtime configuration, and overlap settings to be checked. This patch does not repair prior optimizer history or prescribe a checkpoint reset.

Changelog

• Resolve the runtime configuration from the constructed model when installing training callbacks.
• Add regression coverage through the real setup() call and callback installer, with model construction and external services mocked.
• Cover copied Omni configurations and ordinary GPT/Hybrid providers, retain the GPT/Hybrid ModelConfig tests that traverse the actual builder dispatch, and include one representative two-chunk case for overlapped callback installation.
• Exercise the actual PP=1 scheduler with three microbatches: assert exactly one finalizer call with the correct model and process groups during training, and none during forward-only inference.

Validation

Validation base: Bridge 2eae3c971dd6971beb403e6536987427b7e8a3b2, with its pinned MCore c9b53d0a87cb926f47115259593ecfeb351ca29f.

• Before/after regression: running all eight new cases against unpatched setup.py produces three expected failures and five passes. Both copied-Omni binding cases fail because config.finalize_model_grads_func is None; the training scheduler case fails because the finalizer is called zero times. All eight pass with the fix.
• Targeted unit tests: 46 passed, 0 failed, 0 skipped. This includes the eight new callback cases, existing setup and setup-restart-hook tests, and the four state-reset tests.
• Two-GPU numerical validation: all eight cases passed for the unchanged production fix. Tested existing/corrected bindings, DP=2 or TP=2, and both two-item/three-item loss contracts with actual Megatron DDP, scheduler, gradient finalization, and precision-aware distributed Adam (PyTorch 2.11.0+cu130).
• Lint/format: all repository pre-commit hooks passed; git diff --check passed.
• Static typing limitation: strict mypy was attempted in the lightweight local check environment, but did not pass because runtime imports were unavailable and the existing source reports typing errors. This is not a clean mypy result.

The eight new unit-test cases check callback binding and scheduler invocation. The GPT/Hybrid ModelConfig cases traverse the real _build_distributed_model() dispatch with a lightweight recording builder; other model construction is mocked. The scheduler test also mocks forward/backward computation and the finalizer itself. Numerical gradient correctness is covered by the separate two-GPU experiment described above, not by these unit tests.

The targeted test commands, in an environment with Bridge and its pinned MCore installed, are:

uv run python -m pytest tests/unit_tests/training/test_runtime_callbacks.py \
  tests/unit_tests/training/test_setup.py \
  tests/unit_tests/training/test_setup_restart_hooks.py -q
uv run python -m pytest tests/unit_tests/training/test_state.py -k reset_for_restart -q

To reproduce the original failure, run the new test file against unpatched setup.py:

uv run python -m pytest tests/unit_tests/training/test_runtime_callbacks.py \
  -q

The numerical check distinguishes both loss contracts: the two-item, already-normalized loss used by GRPO requires the sum -2; the three-item loss with two global tokens produces -1. The corrected binding preserves both contracts. Numerical topology coverage is PP=1, CP=1, with DP and TP tested separately. Callback-list tests cover multiple chunks and overlap wiring; they are not a claim of full distributed PP>1, CP>1, or FSDP qualification.

GitHub Actions CI

The new regression tests live under tests/unit_tests/training/. Repository CI remains the merge gate; the standalone numerical experiment is additional validation, not a new CI job in this PR.

Before your PR is "Ready for review"

• Read and followed the contributor guidelines.
• Added necessary regression tests.
• Explained the behavioral change, impact, and validation limits above; no user-facing configuration change requires separate documentation.
• No new dependency or optional-library requirement.

Additional information

Related history: #4974 (July 23) explicitly preserved the provider configuration in the legacy Omni wrapper. #4885 (July 29 Pacific / July 30 UTC) introduced the canonical path using a copied runtime configuration and also removed the legacy wrapper's explicit config=self. This identifies a recurrence in that rewrite relative to its parent; it is not a claim that July was the first-ever occurrence or that every subsequent training run was affected.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[claude]

LGTM — clean, well-targeted correctness fix.

The change resolves the callback-target config from the constructed model via get_model_config(model[0]), exactly the lookup the Megatron scheduler uses. This fixes the silent gradient-sync failure when a provider (e.g. canonical Nemotron Omni) deepcopies its config into the model, so callbacks installed on the provider never reached model.config.

Verified:

• get_model_config unwraps DDP/FSDP and returns the runtime model.config the scheduler reads — matches intent.
• The dropped isinstance(..., GPTModelConfig/HybridModelConfig) branch is subsumed: for builders the model stores transformer, which is what get_model_config returns. Those imports remain used elsewhere in setup.py, so no dead-import lint issue.
• Regression coverage is strong: exercises the real setup() + installer across copied-Omni, plain providers/builders, 1/2 chunks, overlap on/off, restart rebinding, and drives the actual PP=1 scheduler to assert exactly one finalizer call per step and none in forward-only.

No functional changes needed.

Suggested test cases

• tests/unit_tests/training/test_runtime_callbacks.py::test_setup_installs_callbacks_on_scheduler_config (kind in copied_omni, gpt_provider, hybrid_provider, gpt_builder, hybrid_builder; chunk_count 1 and 2; overlap False and True)
• tests/unit_tests/training/test_runtime_callbacks.py::test_restart_binds_callbacks_to_rebuilt_model (copied False and True)
• tests/unit_tests/training/test_runtime_callbacks.py::test_scheduler_finalizes_once_after_setup (forward_only False and True; microbatches 1 and 3)
• tests/unit_tests/training/test_setup.py and tests/unit_tests/training/test_setup_restart_hooks.py (regression guard on the setup path)

No perf tests impacted.

[yaoyu-33]

Checked the new ModelConfig path. The production fix already targets the correct runtime config for both paths: GPTModelBuilder and HybridModelBuilder pass their nested transformer config into the constructed model, and the scheduler reads that same object through get_model_config(model[0]). I pushed 65bece8 to strengthen the regression coverage: the test now uses BridgeGPTModelConfig and actually traverses _build_distributed_model for both GPT and Hybrid instead of mocking the builder dispatch. @aroshanghias-nvd please review the added coverage.

[yaoyu-33]

/ok to test 65bece8

[aroshanghias-nvd]

looks good to me. I'm removing some of the tests added, since they were not that helpful. I will make a commit soon.

[aroshanghias-nvd]

/ok to test 33c47e4

[cuichenx]

/ok to test 234789c

[cuichenx]

/ok to test fe56c9e

[cuichenx]

/ok to test fc71484

[cuichenx]

/ok to test 3e6c19a