Fix bandwidth modeling for sparse two-node communicators

[AlbedoWang]

Summary

Use aggregated inter-node bandwidth for Ring and Tree whenever a communicator spans more than one node. Intra-node bandwidth is now selected only for a communicator contained within one node. The change is intentionally limited to the two bulk-bandwidth branches and adds focused coverage for sparse and dense two-node topologies.

Why this is required

For the outer dimension of a (2, 8) mesh on two eight-GPU hosts, derive_mesh_dim_topo() produces n_nodes=2 and ppn=1. Each two-rank communicator therefore contains one rank from each host. NVLink/NVSwitch terminates within a host, so a Ring AllReduce must carry both its reduce-scatter and all-gather phases through a NIC and the inter-host fabric, while a Tree AllReduce has a single remote tree edge. GPUDirect RDMA can avoid host-memory staging, but it cannot remove the network hop. Pricing this communicator with raw intra-node bandwidth violates that physical cut regardless of measured performance.

The previous condition resembles NCCL 2.25.1 tuning.cc, but NCCL's graphs[a]->bwIntra is not the same quantity as AutoParallel's static topo.bw_intra. NCCL marks every multi-node communicator as NCCL_HW_NET, explicitly accounts for inter-node steps, and constructs a topology-aware graph whose channels start at network endpoints. When NICs are present, NCCL initializes graph bandwidth candidates from its inter-node speed table, checks GPU-to-NIC paths, and reserves link and NIC capacity before tuning.cc consumes graphs[a]->bwIntra; see search.cc topology construction, path-capacity accounting, and network-aware bandwidth search. AutoParallel does not run that graph search for each candidate mesh dimension: its topo.bw_intra is a configured local-fabric value. Reusing NCCL's nNodes <= 2 condition with this different input changes the semantics and makes a sparse cross-node communicator look like an NVSwitch-local communicator.

Selecting bw_inter_agg for every n_nodes > 1 restores the physical invariant while preserving dense two-node behavior. For ppn=1, bw_inter_agg is the model's unaggregated inter-node bandwidth. For a dense two-node communicator with ppn=8, the existing min(bw_intra, bw_inter * ppn) expression still models parallel network injection and caps it at the intra-node ceiling. Other n_nodes <= 2 conditions are left unchanged because they belong to different algorithm- or latency-specific formulas.

Paired CUDA-trace validation

The cost-model change was evaluated before opening this PR in paired profile arms on the same physical two-host allocation and rank-to-GPU mapping: 16 H100 80GB GPUs on two Grand Teton RoCE hosts, a (dp=2, tp=8) mesh with cross-node DP and intra-node TP, LLaMA-3 8B with 32 layers, sequence length 8192, local/global batch 2/4, bf16 parameters with fp32 reductions, selective activation checkpointing with memory budget 0.5, Inductor compilation for model and loss, fused AdamW, and real pinned-revision C4 data through the legacy position loader. Serialized configs and every rank/step input hash matched across the before/after arms. Each arm captured one complete profiler-active training step on all 16 ranks; the analysis below uses raw NCCL kernel duration per rank and does not infer performance from E2E timing.

The incorrect DP AllReduce estimate

The before-fix solution emitted 65 identical 32 MiB DP AllReduce events per rank on the sparse (n_nodes=2, ppn=1) communicator. The old branch priced the cross-node payload using bw_intra; rescoring the identical event with the fixed branch uses bw_inter_agg.

Quantity	Old model	Fixed model	Before-fix CUDA trace
32 MiB DP AllReduce, per event	162.954 us	785.721 us	982.509 us
65 events per rank	10.592 ms	51.072 ms	63.863 ms
Trace / model	6.03x	1.25x	—

The old model underestimates this primitive by 6.03x because it assigns local-fabric bandwidth to a communicator whose only edge crosses the network. The fixed branch reduces the mismatch to 1.25x. This trace validates the magnitude of the correction; the physical cross-node route is the correctness argument.

Effect on the selected solution

The corrected objective changes the solver result rather than only changing a displayed estimate. It changes 256 graph-node placements, all from RS(1) to S(0)S(1), with 224 changes in forward and 32 in backward. The principal per-rank collective-count changes are:

Group / primitive	Before fix	After fix	Delta
DP mesh_fsdp AllReduce	65	0	-65
DP coalesced AllGather	164	100	-64
DP coalesced ReduceScatter	33	98	+65
TP base AllGather	42	32	-10
TP coalesced AllGather	171	191	+20
TP AllToAll	322	322	0

The 65 incorrectly cheap DP AllReduce events disappear from the after-fix solution. They are replaced by a different DP sharding/ReduceScatter pattern, showing that this bandwidth error was large enough to change the optimizer's placement decision.

Estimated versus traced communication delta

To compare the plans under one objective, every emitted mesh_fsdp and mesh_tp event from both traces was rescored with the fixed cost model using its traced message size. Values below are raw NCCL kernel duration per rank; overlap is intentionally not removed.

Dimension	Before estimated	After estimated	Model delta	Before trace	After trace	Trace delta
DP primitives	101.841 ms	74.180 ms	-27.661 ms	141.313 ms	118.268 ms	-23.045 ms
TP primitives	114.908 ms	115.208 ms	+0.300 ms	252.704 ms	299.111 ms	+46.408 ms
DP + TP	216.749 ms	189.388 ms	-27.361 ms	394.017 ms	417.379 ms	+23.362 ms

For the dimension affected by this PR, the model predicts 27.661 ms less DP communication and the trace observes 23.045 ms less. The corrected model therefore gets both the direction and approximate magnitude of the DP plan change right.

The overall DP+TP delta reverses because of a separate TP modeling error, not because the DP fix is wrong. TP AllGather is the largest remaining mismatch: the fixed model predicts only +0.300 ms while the trace records +47.879 ms. The after-fix solution replaces ten 117.44 MiB base AllGathers and adds twenty 58.72 MiB coalesced AllGathers; a size-only estimate sees nearly the same logical footprint but does not represent collective form, scheduling position, or concurrent resource contention. That TP coalescing/scheduling issue is outside this PR's two-line scope.

The conclusion is intentionally narrow: the old condition makes a sparse cross-node DP AllReduce appear 6.03x cheaper than its trace, materially changes the solver's placement choice, and leaves 65 expensive DP AllReduces in the selected plan. This PR restores the physical bandwidth choice, removes those collectives from the new solution, and brings the modeled DP delta close to the traced DP delta. It does not claim an E2E training speedup.

Testing

• python -m pytest tests/test_nccl_cost_model.py -q (119 passed)
• black --check autoparallel/cost_models/nccl_cost_model.py tests/test_nccl_cost_model.py
• flake8 autoparallel/cost_models/nccl_cost_model.py tests/test_nccl_cost_model.py