[jinyouzhi]

Stacked on #8 (affine-window-prepared).

Why

A TransferWindow (#7/#8) only maps peers that share a host, so every operation between hosts is staged through send/recv arenas. On torch builds whose symmetric memory has NCCL one-sided puts, the window can register with the NCCL backend and deliver packed regions to a remote window in place, with no staging copy on either side.

What

• deepspeed.comm: symm_mem_rendezvous(..., backend=), symm_mem_handle_at (a handle at a byte offset into a rendezvoused buffer) and symm_mem_put_signal. A torch without put_signal gets a clear error.
• TransferWindow: learns which processes share a host; a window spanning hosts registers with NCCL (CUDA only) and warms up NCCL's one-sided state collectively. NCCL counts signals per peer on one channel, so READY/DONE share it and waits take a count.
• prepare_transfers: an operation between hosts is one put when both regions are packed window buffers; strided ones keep staging. Mapped (same host) copies are unchanged.
• run: puts are issued first on the window stream. The target acknowledges after all puts from a peer have landed, and the putting process waits for that before run returns, because a put can still be reading its source after its stream has completed. Without the acknowledgement the 2×8 benchmark below saw the next step's source values in the target.
• Docs: one paragraph on windows across hosts, including when staging is faster (below).

Not duplicate

Follow-up to this fork's #7/#8; no upstream PR adds cross-host window transfers. deepspeedai#8966 calibrates link costs for the planner and does not move data; its link numbers are used below as a reference.

Hardware

2 hosts × 8× H20, RoCE 8×400G (mlx5_bond_0..7), driver 580.178.04, torch 2.13.0+cu129, NCCL 2.29.7.

Tests

Single host (8× H20), pytest tests/unit/checkpoint/test_affine_transfer_executor.py tests/unit/checkpoint/test_affine_transfer_executor_external_group.py tests/unit/checkpoint/test_affine_transfer_plan.py tests/unit/checkpoint/test_affine_shard_map.py:

Branch	Result
#8	263 passed, 1 skipped
this PR	265 passed, 1 skipped (new TestWindowAcrossHosts, world sizes 2 and 4, packed and strided)
deepspeedai#8966 (test_affine_transfer_calibration.py + executor)	59 passed

TestWindowAcrossHosts fakes hosts by patching socket.gethostname, so puts go over NVLink. It checks targets and the mapped/staged byte split. It does not catch the source-reuse race fixed above: on one host the put finishes reading before its signal, even at 8 MiB per process. That fix is covered only by the 2×8 benchmark, which checks every target exactly on every repetition.

pre-commit run --files <changed files> passes.

Benchmarks

Aggregate GB/s, median of 5 runs after 2 warm-up runs. Every target is checked exactly on every run; ! would mark a mismatch, and there are none.

Methods:

• stage = execute_transfers.
• stage_p = prepare_transfers(...).run() without a window.
• pull_p/push_p = prepared with a 4 GiB window.

One host (8 GPUs), #8's fan/peer scenarios — unchanged by this PR

pattern	w	MiB	upstream stage	#8 stage_p	#8 pull_p	#8 push_p	this stage_p	this pull_p	this push_p	best / upstream
peer (median of 56)	1	64	65.8	106.0	244.2	253.3	105.9	243.7	252.6	3.84x
fan_out	7	4	46.2	122.6	172.6	118.7	122.7	174.9	123.3	3.79x
fan_out	7	64	199.7	311.0	366.9	312.2	311.1	366.8	310.8	1.84x
fan_out	7	256	236.6	329.8	390.1	340.6	334.2	390.9	340.3	1.65x
fan_in	7	4	50.5	148.5	108.9	171.7	152.6	111.9	170.1	3.37x
fan_in	7	64	213.3	329.1	298.9	368.2	334.2	297.6	366.4	1.72x
fan_in	7	256	247.2	362.6	327.0	390.2	361.1	326.2	389.8	1.58x
pairs	4	4	36.7	130.7	194.1	197.1	131.2	190.9	196.4	5.35x
pairs	4	64	232.9	421.3	957.8	950.2	424.5	955.8	999.6	4.29x
pairs	4	256	320.1	468.0	1226.3	1272.6	468.1	1183.4	1273.0	3.98x
bipartite 2x2	2	4	33.8	117.0	135.4	141.7	113.1	138.9	139.3	4.12x
bipartite 2x2	2	64	220.9	409.5	540.6	561.7	409.8	542.5	564.8	2.56x
bipartite 2x2	2	256	307.3	188.8*	633.5	660.9	461.7	632.7	659.3	2.15x
bipartite 4x4	4	4	107.5	352.6	355.8	371.3	333.5	351.7	383.0	3.56x
bipartite 4x4	4	64	481.0	907.2	1138.4	1185.4	992.2	1144.5	1185.4	2.46x
bipartite 4x4	4	256	558.3	944.9	1283.5	1339.1	940.7	1276.0	1334.4	2.39x

best / upstream uses this PR's best method. #8 and this PR agree within noise in every cell (*one #8 outlier).

Two hosts (16 GPUs), the same scenarios across hosts

Upstream and #8 have no window across hosts, so only their staged numbers are shown. "this" = this PR.

pattern	w	MiB	upstream stage	#8 stage_p	this stage_p	this pull_p (put)	this push_p (put)
xpeer	1	4	5.4	13.0	13.0	13.4	13.5
xpeer	1	64	17.6	19.9	19.3	21.9	21.9
xpeer	1	256	19.8	21.1	21.2	22.5	22.5
xpairs	8	4	43.5	102.3	102.9	72.5	72.4
xpairs	8	64	126.7	149.2	149.0	104.4	121.5
xpairs	8	256	143.9	152.1	152.0	137.4	179.2
xring	16	4	62.1	186.5	178.2	138.4	139.3
xring	16	64	204.6	273.9	271.2	205.7	236.4
xring	16	256	230.8	280.6	278.4	270.1	351.5
xfan_out	8	4	26.8	54.5	56.2	19.6	19.6
xfan_out	8	64	66.5	75.6	75.7	22.6	22.6
xfan_out	8	256	69.9	76.4	77.0	22.7	22.7
xfan_in	8	4	14.3	30.7	19.6	26.2	25.4
xfan_in	8	64	20.4	36.9	22.1	31.6	31.7
xfan_in	8	256	20.8	37.2	24.0	33.6	33.6
xbipartite	8	4	84.4	140.7	75.2	74.3	62.5
xbipartite	8	64	151.5	173.1	172.8	72.5	73.9
mixed_fan_out (7 same host + 8 remote)	15	4	48.5	94.7	94.8	24.1	23.6
mixed_fan_out	15	64	116.2	141.2	141.2	30.3	32.3
mixed_fan_out	15	256	123.5	143.6	144.0	39.6	39.9
strided row→col (total)	16	512	141.9	236.3	234.8	238.1	239.5
strided row→col (total)	16	2048	220.9	303.2	298.0	288.2	295.2

• Staged cross-host fan-in and small bipartite vary between runs on the same code path (xfan_in 20–37 GB/s, xbipartite 4 MiB 75–145). The calibration below shows the same spread.
• Strided row→col: 896 MiB went by put and 1024 MiB was staged, as designed.

deepspeedai#8966 link calibration, same machines

calibrate_links + fit_links at 17aa897e6, NCCL send/recv, 2×8 ranks. Intra-node fans use fan_width=(1, 7). Cross-node fans use fan_width=(1, 4, 8) with fan_peers choosing peers on the other host. Run times: 63 s intra-node, 101 s cross-node.

fit	alpha	GB/s	max rel. error
local	9.4 us	1871	0.25
peer, same host	60.3 us	53.6	0.06
peer, across hosts	91.0 us	20.4	0.12

peer_pairs	pairs	min	median	max (GB/s)
same host	112	52.6 (12→9)	53.6	54.8 (1→7)
across hosts	128	19.7 (15→2)	20.2	28.5 (5→11)

fan	width	root aggregate median (min–max), GB/s	per-link median, GB/s
fan_out, same host	1	54.0 (53.6–54.7)	53.5
fan_out, same host	7	79.5 (79.1–79.9)	11.3
fan_in, same host	1	53.2 (52.5–55.0)	54.1
fan_in, same host	7	94.4 (94.3–94.8)	17.3
fan_out, across hosts	1	23.7 (23.5–30.9)	20.3
fan_out, across hosts	4	69.2 (67.8–70.9)	16.9
fan_out, across hosts	8	78.3 (78.0–79.1)	9.7
fan_in, across hosts	1	20.3 (19.8–29.6)	23.8
fan_in, across hosts	4	22.2 (21.6–39.5)	10.4
fan_in, across hosts	8	22.6 (21.5–39.2)	8.6

How the calibration lines up with the executor:

• One remote link: cross-host peer 20.4 GB/s ≈ xpeer staged 21.2 GB/s. A put reaches 22.5 GB/s with no per-message alpha.
• Cross-host fan_out: 78.3 GB/s ≈ xfan_out staged 77.0 GB/s. One GPU's sends are relayed over the host's other adapters. A put cannot use them and stays at one adapter (22.7 GB/s).
• Cross-host fan_in: flat at ~22 GB/s with outliers up to 39, the same spread as staged xfan_in. Puts reach 33–45 GB/s because each remote sender's adapter writes directly into the receiver.
• Same-host fans: the NCCL fan_out root total of 79.5 GB/s is below the executor's prepared window pull (390.9 GB/s at 256 MiB). The window's direct NVLink copies are what pays inside a host.

Weight sync, as in the #8 / deepspeedai#8966 threads

These scenarios follow the GH200/Slingshot measurements discussed in deepspeedai#8966 (direct / relay / sharded / sharded + intra-node gather, incast, staging size).

Setup: 8 GiB of bf16 weights (32 parameters of 256 MiB) go from the trainer node (ranks 0–7) to 8 engines on the other node (ranks 8–15). Each cell is the median of 3 runs after 1 warm-up run, and every engine buffer is checked exactly. Upstream and #8 cannot map a window across hosts, so they only have staged numbers.

Time to sync 8 GiB to 8 engines, ms (lower is better):

scenario	upstream stage	#8 stage_p	this stage_p	this pull_p	this push_p
direct: each engine pulls everything from trainer GPU 0	1130.5	975.4	973.7	4089.6	4090.5
relay: 0 → 8 across, then 8 → 9..15 inside	1394.8	1235.3	1098.5	534.3	557.3
sharded direct: trainer i sends its shard to every engine	530.3	394.2	393.9	585.2	574.0
sharded i → 8+i, then an intra-node executor plan	376.7	232.9	233.1	84.4	71.5
sharded i → 8+i, then NCCL all_gather_into_tensor	88.7	80.1	80.1	74.1	72.6

• Same shape as GH200. Splitting the source across ranks is what pays, and the model should cross the network once.
• The best plan: a sharded cross-node plan followed by an intra-node plan, run through windows, takes 71.5 ms. That matches NCCL's all-gather (72.6 ms) and is 5.3× faster than upstream's best executor-only plan (376.7 ms).
• With staging only, the intra-node stage costs ~150 ms more than the all-gather.
• Direct with puts is 4.1 s. One GPU's adapter carries all 64 GiB (~16 GB/s), whereas staged send/recv is relayed over the host's other adapters. This is the one-to-many case described above.

Incast across nodes, 1 GiB per sender into rank 0, aggregate GB/s:

senders	upstream stage	#8 stage_p	this stage_p	this pull_p	this push_p
1 → 1	20.1	26.2	18.8	21.9	22.3
2 → 1	19.5	32.4	19.8	43.5	44.5
4 → 1	21.5	36.8	21.3	45.5	45.5
8 → 1	23.5	37.4	21.5	30.5	30.5

• Staged incast is roughly flat, as on Slingshot, but varies between runs (18–37 GB/s) on the same code path. The upstream, Prepare window transfers once and run them per step #8 and this-PR staging paths are identical here, so this is fabric or NCCL routing, not code.
• Puts into one GPU reach ~45 GB/s from 2–4 remote senders. This goes beyond what one adapter gives staging: the senders' adapters write directly into the receiver.
• So on this fabric, splitting across remote holders does help when puts are used.

Staging budget per peer (StagingBudget.bytes_per_peer), 8 GiB sync, ms:

scenario	MiB	upstream stage	#8 stage_p	this stage_p
direct	8	1129.5	973.9	971.5
direct	64	964.8	950.3	948.5
direct	256	953.7	946.8	944.5
sharded direct	8	523.0	394.9	377.1
sharded direct	64	405.3	393.4	380.2
sharded direct	256	407.6	395.4	371.7

From 32 MiB up, all budgets are within ±3%. 8 MiB costs up to 17% per call on upstream, and less once prepared. So 64 MiB is a safe default on RoCE, as on IB and Slingshot.

Reading the numbers

• One host: nothing changes.
• Across hosts, puts win when:
  • each process sends to few remote peers: xpeer, push_p on xpairs/xring at 256 MiB (+18% / +26% over stage_p);
  • many processes send into one: xfan_in.
• Across hosts, staging wins when one process sends to many remote peers (xfan_out, mixed_fan_out, xbipartite). The calibration shows why: a sender's NCCL send/recv reaches ~78 GB/s at width 8, because NCCL relays through the other adapters of the host. A put leaves only through the sender's own 400G adapter (≈22.7 GB/s, the same as one cross-host link).
• Small messages (4 MiB): staging also wins on xpairs/xring.
• No automatic selection yet. This PR does not choose between put and staging per pattern. The docs say a window that excludes such buffers is faster for one-to-many traffic. A cost-based choice could reuse [AffineIR] Calibrate affine transfer links per path (Plan analysis & calibration Part-2) deepspeedai/DeepSpeed#8966's link model.

Matrix

Validated:

• CUDA, bf16/fp32 buffers, pull and push windows, 8 GPUs on one host and 2×8 across hosts over RoCE.
• External group (unit tests).

Not validated:

• XPU/other accelerators across hosts (rejected with an error by design).
• InfiniBand or other NICs.
• More than 2 hosts.
• torch < 2.13 (gives a clear error).
• End-to-end training or RL weight sync.

ZeRO stage, offload and pipeline/tensor parallelism are not touched.

AI assistance

This change, its tests and the benchmarks were developed with AI assistance (GitHub Copilot). A human must review every changed line before this leaves draft.