Add ColocatedBridgeCommunicator for heterogeneous TP/DP MIMO training (NMFW-17) - #4368

Merged

yashaswikarnati merged 29 commits into

NVIDIA:mainfrom

yashaswikarnati:ykarnati/nmfw-17-colocated-bridge-pp1

Apr 28, 2026

Commits on Apr 17, 2026

Add ColocatedBridgeCommunicator for heterogeneous TP/DP MIMO training (NMFW-17)

yashaswikarnati
and
claude
committed
9cc2d37
Fix fan-out backward: all-gather across sibling dest ranks (NMFW-17)

yashaswikarnati
and
claude
committed
eb8191e
Apply review fixes to ColocatedBridgeCommunicator (NMFW-17)

yashaswikarnati
and
claude
committed
e238bcb

Commits on Apr 20, 2026

Fix review bugs in ColocatedBridgeCommunicator (NMFW-17)

yashaswikarnati
and
claude
committed
16756b5
Add bridge gradient correctness tests and edge-case hardening (NMFW-17)

yashaswikarnati
and
claude
committed
1b8c568
Address PR #10 review comments (NMFW-17)

yashaswikarnati
and
claude
committed
137adac
Fix encoder gradient scaling for heterogeneous DP and add mean-CE correctness test (NMFW-17)

yashaswikarnati
and
claude
committed
e04d75b
Prune redundant bridge tests per SPARK review (NMFW-17)

yashaswikarnati
and
claude
committed
9548fa0
Rewrite correctness test: vocab head, global-mean CE, full ref comparison (NMFW-17)

yashaswikarnati
and
claude
committed
236cb22
Fix e2e loss formulation and grad scaling per SPARK review (NMFW-17)

yashaswikarnati
and
claude
committed
9cc7e7d

Commits on Apr 21, 2026

Reuse 1F1B test infra in colocated tests (NMFW-17)

yashaswikarnati
and
claude
committed
7454c40
Unify fan-in/fan-out attributes behind BridgeDirection (NMFW-17)

yashaswikarnati
and
claude
committed
17227a1
Unify RankRole factories via RankRole.build (NMFW-17)

yashaswikarnati
and
claude
committed
8a144f2
Test cleanup per PR review comments 14 and 19 (NMFW-17)

yashaswikarnati
and
claude
committed
8d42043
Validate gradient_reduce_div_factor inputs (SENTINEL NMFW-17)

yashaswikarnati
and
claude
committed
552d8c2
Replace correctness test with DP=1-reference oracle (NMFW-17)

yashaswikarnati
and
claude
committed
9b469eb
Add gradient_reduce_div_factor validation tests (NMFW-17)

yashaswikarnati
and
claude
committed
8a0d421
Add colocated fan-out grad accumulation test (NMFW-17)

yashaswikarnati
and
claude
committed
5bc7e6c
Fix correctness oracle: use equal-DP reference instead of TP=1 (NMFW-17)
yashaswikarnati
committed
238f08c
Verify bridge numeric correctness with fp32 + no bias/dropout (NMFW-17)

yashaswikarnati
and
claude
committed
a1125b1
Drop gradient_reduce_div_factor; use calculate_per_token_loss in MIMO

yashaswikarnati
and
claude
committed
6cb093f

Commits on Apr 22, 2026

Release colocated PGs and forward force_all_reduce (NMFW-17)

yashaswikarnati
and
claude
committed
1993fff
Address PR #10 review comments (NMFW-17)

yashaswikarnati
and
claude
committed
34fd7ef

Commits on Apr 23, 2026

Address PR review comments: fix copyrights to 2026, remove docstring comment

yashaswikarnati
and
claude
committed
5073ad0

Commits on Apr 27, 2026

Fix CI lint and copyright header (NMFW-17)

yashaswikarnati
and
claude
committed
189a9eb
Add docstring to _ColocatedCommunicate.forward to satisfy pylint (NMFW-17)

yashaswikarnati
and
claude
committed
58a7687
Skip colocated communicator setup when grids lack TP/DP (NMFW-17)

yashaswikarnati
and
claude
committed
63bfc5e
Use distinct rank_offsets for the PP role-determination test (NMFW-17)

yashaswikarnati
and
claude
committed
1d4a117
Fix CI lint: collapse generator onto one line (NMFW-17)

yashaswikarnati
and
claude
committed
641a861