Skip to content

[model] fix: forward provider pipeline rank to MTP specs - #6406

Draft
yaoyu-33 wants to merge 1 commit into
mainfrom
codex/fix-mtp-spec-rank-20261010
Draft

yaoyu-33 wants to merge 1 commit into
mainfrom
codex/fix-mtp-spec-rank-20261010

Conversation

@yaoyu-33

Copy link
Copy Markdown
Contributor

With a GPTModelProvider, mtp_num_layers > 0, and dist.use_decentralized_pg=True, model construction reaches MCore's global pipeline-rank fallback even though Bridge has already supplied the provider's process-group collection. The global MPU pipeline group is unset in that mode, so MTP spec construction fails before training starts.

Forward the rank from the provider's pipeline group to the final get_gpt_mtp_block_spec call. This completes the Bridge caller migration started in #5752: its decoder callable receives pp_rank, but the final MTP helper still omitted it. Calls without a provider process-group collection retain their existing fallback. There are no public signature, dependency, or submodule changes.

The supported path is Bridge pretrain/setup → decentralized initialization → ModelProviderMixin.get_model → GPTModelProvider.provide → mtp_block_spec. This change is scoped to that provider path; it does not claim validation of all decentralized distributed training combinations.

Validation uses Bridge base 9166429859ed9a29860ca2e1f3fc317173df8c0e and its exact MCore pin 8949def81eb744b9219e2d794c80265928c40a5b. MCore discovery base was b86e75c184483a54689d17b1a9eb6c5bcc999425; a later default-branch recheck observed 6a6c8aa6e9ee258db08bed381e882c0f754d14bc, which does not change the pinned test dependency.

  • Red, with only the new regression applied to unchanged production code: uv run python -m pytest tests/unit_tests/models/test_gpt_provider.py -k test_mtp_spec_uses_provider_pipeline_group -q -p no:cacheprovider — exit 1, all three cases fail with AssertionError: MTP spec must not read global pipeline rank. The test calls the real pinned MCore spec builder; it covers PP=1 and both PP=2 ranks.
  • Green, identical command after the fix — exit 0, 3 passed.
  • Adjacent: uv run python -m pytest tests/unit_tests/models/test_gpt_provider.py tests/unit_tests/models/test_model_provider_mixin.py -q -p no:cacheprovider — exit 0, 45 passed, 1 existing GPU-dependent test skipped.
  • uv run pre-commit run --all-files — exit 0; no files modified.
  • Supplemental uv run mypy --strict src/megatron/bridge/models/gpt_provider.py — mypy 2.2.0 exits 2 with the same internal tool error on both the unmodified base and patched source. No passing type-check result is claimed.

Validation is CPU contract coverage, not a GPU training result. A separate read-only review of the final diff, source ownership, duplicate search, supported path, and red/green evidence found no blocking findings. All-state searches in both repositories found related upstream caller work in NVIDIA/Megatron-LM#8048 and fallback changes in NVIDIA/Megatron-LM#8057, but neither updates this Bridge call.

Signed-off-by: Yu Yao <yaoyu.094@gmail.com>
@copy-pr-bot

copy-pr-bot Bot commented Oct 10, 2026

Copy link
Copy Markdown

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

Copy link
Copy Markdown
Contributor Author

/ok to test 825d491

@cuichenx cuichenx added bug Something isn't working area:model Model implementations and HF bridge logic labels Oct 10, 2026

This branch was successfully deployed

1 active deployment
test — 825d491b Deployed Oct 10, 2026 by copy-pr-bot[bot] via cicd-wait-in-queue #22731
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:model Model implementations and HF bridge logic bug Something isn't working

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants