[jinyouzhi]

Calibrate affine transfer links per path

Stacked on #8809 only for its docs page (affine-transfer.rst); the module itself needs nothing beyond master. Until #8809 merges, please review only the commits after the #8809 merge.

Why

A replica policy or cost model choosing between plans (#8252, #8935/#8936) needs link costs measured on the machine it runs on. Executor measurements show that round count and receiver sharing dominate sync time, so per-pair numbers are not enough on their own; this PR also measures 1→N and N→1.

What

deepspeed/checkpoint/affine_transfer_calibration.py:

• calibrate_links(sizes, repeats, warmup, path_of, fan_width, fan_peers, group): collective over group (the default group unless given). Returns the same raw samples on every process, with no aggregation, for:

  • local copy;
  • peer (each ordered pair alone);
  • fan_out (1→N);
  • fan_in (N→1, receiver contention).

  Sizes run from 4 KiB to 256 MiB by default.

• Each message is labelled by path_of(from_process, to_process). default_path_of gives local/peer; node_path_of(n) gives intra_node/inter_node; any callable works.

• fit_links(calibration): per-path alpha + bytes / bandwidth, fitted to the per-size medians; peer_pairs per ordered pair; fans keyed by (root, peers), with send/receive giving each root's aggregate rate.

Weight transfers are one-way (a training process sends, a rollout process receives), so no pattern has a process send and receive at once.

It measures only. It does not plan, select replicas, or change the planner or executor.

Example (one node, 8 devices, default sizes)

pattern	B70 xccl: alpha (µs) / GB/s	L20 nccl (PCIe): alpha (µs) / GB/s	H20 nccl (NVLink): alpha (µs) / GB/s
local	73 / 359	16 / 1534*	9 / 1881
peer	301 / 28.2	72 / 25.3	54 / 275
fan_out (7 peers)	531 / 3.4	141 / 3.6	88 / 49.6
fan_in (7 peers)	670 / 2.6	56 / 3.8	46 / 54.5
run time	88 s	34 s	13 s

* poor linear fit (max relative error 0.77); small copies are cache-resident.

Multi-node (2 × 8 H20, RoCE 8 × 200 Gb/s per node, node_path_of(8), 16 processes)

pattern	alpha (µs) / GB/s	max rel. error
peer, intra_node	57 / 53.8	0.07
peer, inter_node	97 / 16.6 (pairs 14.0–24.5)	0.50
fan_in, 7 intra-node peers: root receive	94.6 GB/s
fan_out, 8 cross-node peers: root send	62.1 GB/s median (33–78 over roots)
run time	72 s (+111 s for cross-node fans of width 1/4/8)

• The same NVLink pair measured 275 GB/s in the 8-process group and 54 GB/s in the 16-process group over RoCE (105 GB/s with NCCL over TCP), so calibrate on the group and communication settings the transfer uses.

Tests

tests/unit/checkpoint/test_affine_transfer_calibration.py:

• fit on synthetic data with known link costs, per-pair and per-peer-set fits;
• schema and argument checks;
• 4-rank runs checking completeness, sample counts, and fan path direction with an asymmetric path_of.

10 passed on 8× B70 (xccl). pytest tests/unit/checkpoint/ -k affine: 344 passed, 1 skipped on 8× H20 with torch 2.13.0+cu129 / NCCL 2.29.7 and with torch 2.9.1+cu129 / NCCL 2.29.7. Mutation check: flipping the fan_in label direction fails a test. Cross-node runs used NCCL 2.29.7 on both nodes (torch 2.13.0 and 2.9.1).

AI assistance (GitHub Copilot) was used for the multi-node runs and this description; I reviewed the changes and the results.

Follow-ups

• Sampled pairs for large worlds.
• Joining the samples to a plan's messages (analysis PR, stacked on this one and [checkpoint] Execute affine transfer plans with bounded peer staging #8809).

[Achyuthan-S]

Thanks, this is the input the cost model needs. I'll build the cost-aware policy on #8936 from fit_links here and analyze_transfers in #8967, rather than measuring separately. Two things the receiver side of the model needs from the fan-in data:

1. Fan-in at several widths. The policy caps how many senders feed one receiver at once, so it needs the cost as a function of the number of concurrent senders, not only at world - 1. Could fan_width take a list, say (1, 2, 4, world - 1), with the fan fits keyed by width? The fans are cheap next to the pair sweep.
2. The root's time, fitted. For a fan-in, the finish time is when the root has all k messages: root_seconds against k * bytes. That is the receiver's ingest bandwidth, and it currently stays in the raw result. A fitted entry beside the per-peer one would save every consumer from refitting it.

Your B70 numbers already show why the cap matters. If I read the per-link fit right, seven senders give about 7 × 2.6 ≈ 18 GB/s into one receiver, which is less than one pair alone (28.2 GB/s).

[jinyouzhi]

@Achyuthan-S Yes, you're very observant—I agree. Based on these experiments, we can refine some of our policies.

[Achyuthan-S]

Thanks! If it's easier, I can send both as a small commit to your branch: fan_width accepting a list, with the fan fits keyed by width, and a fitted entry for the root's time. One refinement since: @0z5a measured flat incast on GH200 over Slingshot, so the policy should read the receive cap from the fit at each width rather than assume one.

[jinyouzhi]

@Achyuthan-S Sure, I have added you as my fork's collabration or I also can cherry-pick a commit provided by you. Thanks

[0z5a]

Read the last commit (3cfa4153), as asked. The measurement design is careful: the barrier and synchronize brackets keep collectives out of timed regions, the pair sweep isolates each link, and the fans separate sender and receiver sharing. On CPU (gloo, 4 ranks), the calibration tests pass together with #8967's and the executor's (62).

A few thoughts, mostly about feeding a cost model directly:

• Per-process send and receive capacity. A model of the form t = max over processes of max(sent / send_capacity, received / receive_capacity) needs each process's aggregate send and receive capacity. fit_links gives per-path link costs; the fan fits use the peers' times, and the roots' times, which carry the aggregate, stay only in the raw result. Computing bytes * len(peers) / root_seconds from the raw fan entries is a few lines. On CPU it gives a per-root figure for each direction. Having fit_links fit those per process (or per root) as send and receive would let #8967's ProcessReport.sent_bytes / received_bytes be priced directly.
• Fan peers. The peers are root + 1 … root + fan_width, so with fan_width below the processes per node every fan stays within a node. If receive capacity differs across nodes, a way to choose the peers (a fan_peers(root) hook, or a stride) would let the fans cover it.
• Defaults at scale. With the default fan_width = world - 1, each process can allocate (1 + world) * max(sizes), about 16 GiB at 64 ranks and 256 MiB, and the pair sweep is O(world²) barriers per size. The docstring says so, but a default capped at the processes per node, or a check against free memory, might save someone a surprise.
• Process group. Calibration always uses the default group, while the executor runs on whatever group TransferEndpoints names. An optional group argument would let it measure the group the executor will actually use, e.g. an external trainer–rollout group, with the same rank numbering the analysis report uses.

[jinyouzhi]

@Achyuthan-S @0z5a Thanks both. I agree with the direction: a list of fan widths and a fitted root time (per-process send / receive). Before we key the fits by width alone, here is fan-in data at four widths on 8x B70 (PCIe, xccl, 256 MiB, median of 10). It shows that on a PCIe tree, which senders feed a receiver matters more than how many.

Topology. Four PCIe switches, each with two cards behind one Gen5 x16 root port, and two switches per socket: {0,1} {2,3} on socket 0 and {4,5} {6,7} on socket 1 (assuming ranks follow PCI address order). GH200 over Slingshot is a Clos, so every sender set looks alike there, which fits your flat incast.

Pair bandwidth, GB/s (row = sender, column = receiver):

from\to   0     1     2     3  |  4     5     6     7
0         -    8.3   7.9   8.3 | 27.3  27.5  27.6  27.8
1        6.0    -    5.5   3.6 | 27.4  27.5  27.4  27.5
2       53.9  53.7    -   53.6 | 14.5  14.6  14.6  14.6
3       53.6  53.5  53.8    -  | 14.6  14.6  14.6  14.6
4       27.4  27.3  27.4  27.3 |  -   14.4  14.3  14.4
5       27.3  27.3  27.4  27.4 | 25.2   -   25.3  25.3
6       27.6  27.7  27.0  27.3 |  8.8   8.7   -    8.6
7       27.5  27.8  27.3  27.4 | 42.8  42.9  43.0   -

These are stable across runs. Cross-socket traffic sits near 27 GB/s. Within a socket, bandwidth depends on the sender and the direction (2→0 is 54, 0→2 is 7.9, on the same switches), which the tree alone does not predict.

Fan-in root aggregate (k * bytes / root_seconds), GB/s, by width:

width	min ~ max over roots	spread
1	7.0 ~ 53.9	7.7x
2	12.6 ~ 55.2	4.4x
4	18.7 ~ 37.3	2.0x
7	16.4 ~ 23.6	1.4x

Three cases explain most of it:

• 0 <- [1, 2] = 12.6: the root waits for the slowest sender (1→0 at 6.0), so the aggregate is about 2 × 6.
• 1 <- [2, 3] = 55.2: two 54 GB/s senders reach only about 55 together, so a receive-side limit binds.
• 3 <- [4, 5, 6, 7] = 27.9: four 27 GB/s senders together still get about 28, because the shared bottleneck is the cross-socket link, not the receiver.

So the finish time looks like max(slowest sender's bytes / pair bandwidth, bytes through a shared link / that link's capacity). Both terms depend on the sender set. A single fit per width would average 12.6 and 55.2 at width 2, and a single receive capacity per process would miss the shared cross-socket link.

Proposal for this PR:

1. Width list and root fit: keep both as you describe. They are right for Clos and other uniform fabrics.
2. Per-pair fit: add an optional per-pair entry in fit_links, keyed by (from, to) and fitted from the existing pair sweep, so no new measurement. It captures sender and direction asymmetry, which a path-class fit hides (the 28 GB/s peer fit here is the median of 3.6–54).
3. Choosing fan peers: add a fan_peers(root, width) hook (0z5a's suggestion). The fan fits are keyed by the peer set, or by the peers' path classes, not by width alone. That way a fan can target "same switch", "same socket" or "cross socket" senders and measure the shared-link capacity directly. The raw entries already record peers.
4. Docs: state that width-keyed fits assume a uniform fabric, and include this table as the counterexample.

On 0z5a's other points:

• I'll add the optional group argument, with ranks numbered as in the analysis report.
• I'll cap the default fan_width at the processes per node and document the memory needed.

@Achyuthan-S Please push your width-list and root-fit commit. I'll add the per-pair fit and the fan_peers hook on top so they don't conflict. Raw JSON for all four widths is available if you want to check the fits.

[Achyuthan-S]

Pushed the width-list and root-fit commit:

• fan_width takes one width or a list. Each fan sample records its width, and the result lists fan_widths.
• fit_links keys the fan fits by width, and adds send and receive: per width, a fit per process of the root's own time against all the bytes it moved. That is the aggregate capacity @0z5a described, so [AffineIR] Analyze affine transfer batches before execution (Plan analysis & calibration Part-1) #8967's sent_bytes / received_bytes can be priced against it directly.
• Docs and tests are updated: a synthetic test that recovers per-process capacities at two widths, and the 4-rank tests cover a list of widths. On CPU (gloo, 4 ranks), the calibration, executor and planner tests pass together (180).

Your B70 table is a great counterexample. It fits the same form once every shared resource is a term, not only the endpoints: a fan finishes when the busiest resource it crosses does, i.e. the max over each sender's link, the receiver, and any shared link such as your cross-socket hop, of the bytes through it over its capacity. Your three cases are one of each:

• 0 <- [1, 2] is bound by the 1→0 link;
• 1 <- [2, 3] by the receiver;
• 3 <- [4..7] by the cross-socket link.

So your per-pair fit and fan_peers are what calibrate those resources. I kept the width as the only key here so your peer-set keys can replace it without a conflict.

[jinyouzhi]

@Achyuthan-S @0z5a Pushed ba8b39f (on top of Achyuthan's commit, with #8809's 11fe945 merged in).

Updates

• Fan fits and send/receive are keyed by (root, tuple(sorted(peers))) instead of width, as Achyuthan anticipated; width is len(peers), and the width list and root fit stay.
• fits['peer_pairs'][(src, dst)]: one fit per direction.
• fan_peers(root, width): choose which peers form each fan.
• group: calibrate inside a process group; ranks match [AffineIR] Analyze affine transfer batches before execution (Plan analysis & calibration Part-1) #8967's report.
• Default width is capped at processes-per-node minus one, same-host peers first; docs note the buffer memory.
• Docs: uniform-fabric note, the B70 table, and the busiest-resource model.

On the B70 counterexample: send/receive are the completion rate of a fan with those peers, bound by the busiest resource, not the root's capacity. Per-resource capacities are left to the #8936 policy, built from peer_pairs plus targeted fan_peers fans.

Validation: 8× L20 (driver 580.173.02, torch 2.14.1+cu130): calibration tests on 4 GPUs, 10 passed; external-group test on CPU passed 20 of 20 runs; pre-commit passed. Multi-node not validated.

[Achyuthan-S]

Thanks, this is the right split. Keying by peer set captures exactly what the B70 table showed, and I agree on the semantics: send/receive are the completion rate of a given fan, and the per-resource capacities belong in the policy. I'll derive them in the #8936 policy from peer_pairs plus targeted fan_peers fans, and check the result against your B70 table and 0z5a's GH200 and B300 runs.

One small thing for consumers: the fits now have tuple keys (peer_pairs[(src, dst)], fans by (root, peers)), so json.dumps(fits) raises. The raw calibration still serializes, so caching that and refitting on load works. A line in the docstring saying so would save someone the surprise.

[jinyouzhi]

Good catch. I added a note to the fit_links() docstring: the returned fits contain tuple keys and are not directly JSON-serializable; callers should persist the raw calibrate_links() result and call fit_links() again after loading. The doc update is in my local worktree and will be included in the next signed-off PR update.

[Achyuthan-S]

Thanks for the JSON note and the H20 multi-node numbers. The same NVLink pair measuring 275 GB/s in one group and 54 GB/s in another is a useful warning, and I'll make "calibrate in the transfer's own group" explicit on the policy side.

One concern with dropping the concurrent pattern: some transfers do have processes sending and receiving at once.

• Colocated resharding: on the same GPUs, ZeRO-3 → AutoTP, each GPU sends its partition and receives its TP shard.
• Stage 2 of a relay: the engines on the target node all-gather, so each one sends and receives.
• Pipelined relay: the relay GPUs receive stage 1 while they send stage 2.
• Moves between trainers of different DP sizes.

@0z5a's B300 runs show what this costs: 618 GB/s per direction duplex, against 664 GB/s one-way. On GH200 the overlap held because the NIC's receive and the NVLink send are separate resources. So the policy needs a duplex number per path.

The shift pattern's labelling was the problem, not the measurement, and a simpler pattern would do: for example, each ordered pair exchanging both ways at once, labelled by the pair's path in each direction. If you'd rather keep this PR one-way, a line in the docs saying duplex isn't measured, plus a follow-up, works too.

[jinyouzhi]

@Achyuthan-S Thanks. I had assumed weight_sync in an RL loop is always one-way, so I missed the cases where a process sends and receives at the same time. Colocated resharding, relay stage 2 and the pipelined relay are all real ones, and the B300 duplex vs. one-way gap shows the policy needs this number. I'll bring the concurrent pattern back in this PR and add the duplex.

I agree the labelling was the problem. The old shift pattern made it unclear which participant each sample described.

[jinyouzhi]

@Achyuthan-S Pushed d00ba30, following your suggestion:

• peer both ways: peer samples now carry both_ways. A two-way sample has a pair send to each other at once while every other process is idle, and gives one sample per direction, labelled with that direction's path. peer / peer_pairs stay one-way. peer_both_ways / peer_both_ways_pairs hold the two-way fits.
• relay (the old shift pattern, relabelled): every process receives from p - shift while sending to p + shift. Each sample is one process's own time, labelled with its from_process / to_process and both receive_path and send_path. fit_links keys relay by (receive_path, send_path), so nothing is filed under a path it only half used.

H20, one node, 4 GPUs over NVLink, 256 MiB, median of 10, GB/s per direction: one-way 331–335, both ways 276–287 (0.83–0.86×), relay 282–287 (≈0.85×). Calibration tests pass on 4 GPUs (12 passed). Cross-node relay is not measured yet.