Fix 2D tensor communication for asymmetric DP in Bridge Communicator

[yashaswikarnati]

Summary

Adds native 2D tensor support to BridgeCommunicator for modules that produce [B*S, H] outputs (e.g., vision encoders), fixing a crash when DP sizes differ across modules.

The problem

The bridge communicator assumed all tensors are 3D [S, B, H] and used unsqueeze(-1) / squeeze(-1) to adapt 2D tensors. This breaks with asymmetric DP because the bridge splits/concatenates along dim_mapping['b']=1 — but after unsqueeze(-1), dim 1 is the hidden dimension, not batch:

[577, 4096] → unsqueeze(-1) → [577, 4096, 1]
                                 s=0  b=1   h=2
                                      ^^^
                          bridge thinks this is batch
                          but it's actually hidden

Fan-in (DP=3 → DP=1): cat(dim=1) on 3 tensors concatenates along hidden → [577, 12288] instead of [1731, 4096].
Fan-out (DP=1 → DP=3): split(dim=1, 3) splits hidden into thirds → [1731, 1365] instead of [577, 4096].

Silent when DP ratio is 1:1 (cat/split of 1 tensor are no-ops).

The fix

Add tensor_ndim parameter to BridgeCommunicator. For 2D tensors, batch is folded into dim 0, so fan-in/fan-out uses cat/split at dim 0 directly — no unsqueeze needed:

Fan-in:   3 × [577, H]  → cat(dim=0)    → [1731, H]
Fan-out:  [1731, H]      → split(3, dim=0) → 3 × [577, H]

MultiModulePipelineCommunicator accepts module_output_ndim dict to configure each bridge per source module.

Files changed

• bridge_communicator.py — tensor_ndim param, _batch_dim property
• multimodule_communicator.py — module_output_ndim param, passes to bridges
• test_mimo_1f1b_schedule.py — 4 new 8-GPU asymmetric DP test configs
• test_bridge_communicator.py — 2 new 2D fan-in/fan-out tests

Test plan

• 8-GPU: Enc DP=4 → LLM TP=2 PP=2 DP=1 (fan-in)
• 8-GPU: Enc TP=2 PP=2 DP=1 → LLM DP=4 (fan-out)
• 8-GPU: Enc DP=2 → LLM TP=2 PP=3 DP=1 (fan-in)
• 8-GPU: Enc TP=2 PP=2, LLM TP=2 PP=2 (symmetric, no regression)
• 8-GPU: Bridge-level 2D fan-in/fan-out fwd+bwd

🤖 Generated with Claude Code

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

[github-actions]

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

1. Add the oncall reviewer (optional reviewer)
2. Add required review teams based on your changes

See the contribution guide for more details.

[yashaswikarnati]

/claude review

[claude]

LGTM

[yashaswikarnati]

/ok to test 32422aa

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

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