[model] fix: Finalize providers before distributed model construction

[yaoyu-33]

What changed

Finalize legacy model providers in the training setup caller immediately before entering provide_distributed_model(), and add a focused regression test for heterogeneous Llama-Nemotron providers.

Bug and impact

Heterogeneous providers defer Megatron-Core post-initialization, so derived fields such as per_block_parameters do not exist until finalize() runs. The training construction path previously entered provide_distributed_model() before establishing that lifecycle boundary, causing an AttributeError when model construction consumed the field.

_build_distributed_model() now finalizes a legacy provider after distributed initialization and immediately before invoking provide_distributed_model(). Finalization remains an explicit caller responsibility rather than an implicit side effect of the provider method.

Validation

The regression test replaces the distributed construction method with an entry callback and verifies that the heterogeneous block configuration has already been parsed before that callback is entered.

python -m pytest tests/unit_tests/training/test_setup.py::TestBuildDistributedModel tests/unit_tests/models/test_model_provider_mixin.py
10 passed

uv run --active --no-sync pre-commit run --all-files
passed

Scope

This does not change conversion mappings, provider configuration fields, public model-builder APIs, or Megatron-Core. It only makes the training setup lifecycle ordering explicit for deprecated ModelProviderMixin providers.

[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.

[yaoyu-33]

/ok to test 626eb60

[claude]

LGTM — the fix is well-scoped and correctly tested.

Summary: Adding self.finalize() before get_model(...) in provide_distributed_model ensures deferred heterogeneous fields (e.g. per_block_parameters) are computed before the model is built. This closes a gap where a provider constructed and passed directly to provide_distributed_model (bypassing apply_overrides_and_finalize) would build with un-parsed config.

Verified, no concerns:

• finalize() is documented and implemented as idempotent (safe to call multiple times), so paths that already finalize elsewhere are unaffected.
• The new call runs before the torch.distributed init block. MegatronMIMOProvider.finalize() raises if distributed is not initialized, but MIMO overrides provide_distributed_model entirely and does not reach the modified base method, so no regression.
• The new unit test drives the real LlamaNemotronHeterogeneousProvider and asserts per_block_parameters is populated inside the model-build callback, which is exactly the behavior being fixed.

Suggested test cases:

• No perf tests impacted.
• test_provide_distributed_model_finalizes_heterogeneous_config_before_model_build (added, covers the fix)
• Consider also asserting finalize() runs before distributed init to lock in that deferred-field computation does not depend on torch.distributed being initialized for the base path.

[yaoyu-33]

/ok to test 26d2fb1

[copy-pr-bot]

/ok to test 26d2fb1084b8cbb8a20ca9634a1c8d74b1a5c44e

@yaoyu-33, there was an error processing your request: E2

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/

[yaoyu-33]

/ok to test 26d2fb1