Skip to content

Commit b5b7b0b

Browse files
fix(collectivex): capture rank alignment before graph timing (#3727)
1 parent f4a42a2 commit b5b7b0b

4 files changed

Lines changed: 41 additions & 78 deletions

File tree

‎collectivex/bench/ep_harness.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -196,9 +196,10 @@ def _write_json_atomic(path: str, value) -> None:
196196

197197

198198
def kernel_generation(backend) -> str:
199-
"""The adapter's kernel family; `-cudagraph` keeps replayed rows a separate series."""
199+
"""Identify the kernel family and captured-all-reduce measurement regime."""
200200
family = backend.kernel_generation or "n-a"
201-
return f"{family}-cudagraph" if backend.cuda_graph_enabled else family
201+
# Historical -cudagraph rows used an external barrier and alignment spin.
202+
return f"{family}-cudagraph-in-graph-ar" if backend.cuda_graph_enabled else family
202203

203204

204205
class _Collectives:

‎collectivex/bench/ep_timing.py‎

Lines changed: 21 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -32,18 +32,16 @@ def sample():
3232
return [sample() for _ in range(iters)]
3333

3434

35-
def time_cuda_graph_phase_us(torch, fn, warmup: int, iters: int, interval, align) -> list[float]:
35+
def time_cuda_graph_phase_us(torch, fn, warmup: int, iters: int, interval) -> list[float]:
3636
"""Time one event-record interval captured inside graph replay.
3737
38-
`align()` enqueues a device-side rank barrier before each replay, so replays start together
39-
rather than ~75us apart (b200 EP16), which the cross-rank MAX would report as latency.
38+
The graph contains an all-reduce before the timing events to reduce rank entry skew.
4039
"""
4140
for _ in range(max(0, warmup)):
4241
fn()
4342
torch.cuda.synchronize()
4443
samples = []
4544
for _ in range(iters):
46-
align()
4745
fn()
4846
torch.cuda.synchronize()
4947
samples.append(interval[0].elapsed_time(interval[1]) * 1000.0)
@@ -223,53 +221,24 @@ def poison(tensor):
223221

224222

225223
class GraphTiming:
226-
"""Captured dispatch -> combine pairs timed by event nodes inside the graph, each timed replay
227-
started behind a device-side rank barrier so the cross-rank MAX is the operation, not launch
228-
skew."""
224+
"""Measure dispatch/combine latency with CUDA events inside a replayed CUDA graph.
225+
226+
Each replay first runs an all-reduce to bring ranks closer together before timed work
227+
begins. The all-reduce is part of the graph but precedes the timing events, so its cost
228+
is excluded from the reported latency.
229+
"""
229230

230231
graph = True
231232
# Handle attributes dispatch writes; the replay value check poisons them.
232233
DISPATCH_OUTPUT_FIELDS = ("recv_x", "recv_scales", "dispatch_output")
233234

234235
def __init__(self, backend):
235236
self.backend = backend
236-
self.align_cycles = None
237237
self._align_token = None
238238

239239
def components(self):
240240
return ["roundtrip", "dispatch", "combine"]
241241

242-
def calibrate_align_spin(self, spin_us=100.0):
243-
"""Size the post-barrier spin to `spin_us` of wall time on this GPU.
244-
245-
The spin lets every host enqueue its replay before the stream reaches it. `_sleep` counts
246-
SM cycles, so a fixed count left ~15us of skew between differently clocked gb200 ranks.
247-
"""
248-
import torch
249-
250-
probe = 200_000
251-
start = torch.cuda.Event(enable_timing=True)
252-
end = torch.cuda.Event(enable_timing=True)
253-
torch.cuda._sleep(probe // 10) # ramp clocks before the measured spin
254-
start.record()
255-
torch.cuda._sleep(probe)
256-
end.record()
257-
torch.cuda.synchronize()
258-
elapsed_us = max(start.elapsed_time(end) * 1000.0, 1e-3)
259-
self.align_cycles = max(1, int(probe * spin_us / elapsed_us))
260-
261-
def align(self):
262-
"""Enqueue a device-side rank barrier on the current stream, without a host sync."""
263-
import torch
264-
import torch.distributed as dist
265-
266-
if self._align_token is None:
267-
self._align_token = torch.zeros(1, device=self.backend.device)
268-
if self.align_cycles is None:
269-
self.calibrate_align_spin()
270-
dist.all_reduce(self._align_token)
271-
torch.cuda._sleep(self.align_cycles)
272-
273242
def capture_pairs(self, problem, staged, pairs, marks):
274243
"""Capture `pairs` back-to-back pairs into one graph; `marks` picks the windows that get
275244
event nodes ("pair", "dispatch", "combine"). Returns (graph, {mark: (starts, ends)},
@@ -278,6 +247,9 @@ def capture_pairs(self, problem, staged, pairs, marks):
278247
import torch.distributed as dist
279248

280249
b = self.backend
250+
if self._align_token is None:
251+
self._align_token = torch.zeros(1, device=b.device)
252+
dist.all_reduce(self._align_token)
281253
stamps = {mark: ([graph_event() for _ in range(pairs)], [graph_event() for _ in range(pairs)])
282254
for mark in marks}
283255

@@ -290,6 +262,14 @@ def record(mark, edge, i):
290262
graph = torch.cuda.CUDAGraph()
291263
combined = handle = None
292264
with torch.cuda.graph(graph, capture_error_mode="relaxed"):
265+
# A separate all-reduce followed by a host-launched graph can let CPU launch
266+
# delays spread the ranks apart again, inflating the measured communication time.
267+
# Capture the all-reduce and its stream dependencies in this graph so each replay
268+
# reaches the timed work without another host launch after alignment. Place it
269+
# before the first timing event to exclude alignment cost, and only once per graph
270+
# to preserve back-to-back dispatch/combine pairs. This reduces entry skew; it
271+
# does not guarantee that every rank starts the timed work simultaneously.
272+
dist.all_reduce(self._align_token)
293273
for i in range(pairs):
294274
record("pair", 0, i)
295275
record("dispatch", 0, i)
@@ -313,10 +293,9 @@ def component(self, name, problem, warmup, iters):
313293
staged = self.backend.warm_and_hoist_stage(problem, warmup)
314294
mark = "pair" if name == "roundtrip" else name
315295
graph, stamps, combined, _ = self.capture_pairs(problem, staged, 1, (mark,))
316-
self.calibrate_align_spin() # clocks move with load and temperature: per timed series
317296
starts, ends = stamps[mark]
318297
samples = time_cuda_graph_phase_us(
319-
torch, graph.replay, warmup, iters, (starts[0], ends[0]), align=self.align
298+
torch, graph.replay, warmup, iters, (starts[0], ends[0])
320299
)
321300
combined.fill_(float("nan"))
322301
torch.cuda.synchronize()
@@ -329,16 +308,14 @@ def component(self, name, problem, warmup, iters):
329308

330309
def chain(self, problem, staged, iters, drop):
331310
"""Each sibling is ONE graph of `iters` unrolled pairs, the shape of a decode graph, replayed
332-
once untimed (upload) and once aligned and timed."""
311+
once untimed (upload) and once timed, with alignment captured before the first event."""
333312
import torch
334313

335314
floors, floor_stamps, _, _ = self.capture_pairs(problem, staged, iters, ("dispatch", "combine"))
336315
period, period_stamps, combined, _ = self.capture_pairs(problem, staged, iters, ("pair",))
337-
self.calibrate_align_spin()
338316
for graph in (floors, period):
339317
graph.replay()
340318
torch.cuda.synchronize()
341-
self.align()
342319
graph.replay()
343320
torch.cuda.synchronize()
344321
return _chain_result(period_stamps["pair"], floor_stamps, drop, combined)

‎collectivex/docs/methodology.md‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -335,9 +335,10 @@ every check, without changing its contract.
335335
Only the launch mechanism changes; every family keeps its eager meaning:
336336

337337
- **Fresh-entry components** (`roundtrip`, `dispatch`, `combine`) capture one pair with event
338-
nodes around the timed window. Each timed replay starts behind a device-side rank barrier (an
339-
all-reduce, then a spin of fixed wall time, calibrated per GPU because `torch.cuda._sleep`
340-
counts SM cycles); without it b200 EP16 ranks entered ~75 µs apart. A rank whose host launches
338+
nodes around the timed window. An all-reduce is captured before the first timing event,
339+
including its stream dependencies, so alignment stays outside the measured interval without
340+
a host-side handoff or spin kernel. This reduces skew but does not guarantee simultaneous
341+
release across ranks. A rank whose host launches
341342
late still stalls the others, so these series publish **only p50** (p90/p95/p99 and the matching
342343
token rate are null). Component origin is `cuda-graph-replay`.
343344
- **The chained family** (`pair_period`, floors, health) captures each sibling chain as one graph

‎collectivex/tests/test_chain.py‎

Lines changed: 13 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -241,19 +241,18 @@ def _graph_backend(self, stage_device_work=False, precision="bf16"):
241241
backend.mode, backend.CUDA_GRAPH_MODES = "normal", ("normal",)
242242
return backend
243243

244-
def _aligned_replays(self, calls):
245-
"""Replays preceded by the device-side rank barrier (all-reduce, then the spin)."""
246-
return [i for i, call in enumerate(calls)
247-
if call == "graph_replay" and calls[i - 2:i] == ["all_reduce", "align_spin"]]
248-
249-
def test_every_timed_replay_starts_behind_a_device_side_rank_barrier(self):
250-
# Without it each replay restarts from the preceding sync and ranks enter ~75us apart.
244+
def test_alignment_is_captured_before_the_timed_window_without_spin(self):
251245
backend = self._graph_backend()
252246
with mock.patch.dict(os.environ, {}, clear=True), \
253247
trace_torch(backend.clock, backend.calls):
254248
backend.benchmark_component("roundtrip", new_problem(), warmup=2, iters=3)
255-
replays = [i for i, call in enumerate(backend.calls) if call == "graph_replay"]
256-
self.assertEqual(self._aligned_replays(backend.calls), replays[2:5])
249+
calls = backend.calls
250+
begin, end = calls.index("capture_begin"), calls.index("capture_end")
251+
first_event = calls.index("record", begin)
252+
self.assertEqual(calls[begin:first_event].count("all_reduce"), 1)
253+
self.assertNotIn("all_reduce", calls[first_event:end])
254+
self.assertNotIn("all_reduce", calls[end:])
255+
self.assertNotIn("align_spin", calls)
257256

258257
def test_the_graph_chain_is_one_capture_of_unrolled_pairs_per_sibling(self):
259258
iters, drop = 5, 1
@@ -269,7 +268,11 @@ def test_the_graph_chain_is_one_capture_of_unrolled_pairs_per_sibling(self):
269268
for lo, hi in zip(begin, end):
270269
self.assertEqual(ops_only(calls[lo:hi]), ["dispatch", "stage", "combine"] * iters)
271270
self.assertEqual(calls.count("graph_replay"), 4)
272-
self.assertEqual(len(self._aligned_replays(calls)), 2)
271+
self.assertNotIn("align_spin", calls)
272+
for lo, hi in zip(begin, end):
273+
first_event = calls.index("record", lo)
274+
self.assertEqual(calls[lo:first_event].count("all_reduce"), 1)
275+
self.assertNotIn("all_reduce", calls[first_event:hi])
273276
self.assertEqual({len(series[k]) for k in ("pair", "dispatch", "combine")}, {iters - drop})
274277
self.assertEqual(len(series["start_to_start"]), iters - drop - 1)
275278
self.assertTrue(series["combined"].cloned)
@@ -290,25 +293,6 @@ def test_external_switch_restores_the_eager_component_pipeline(self):
290293

291294

292295
class GraphAlignmentAndValueCheck(unittest.TestCase):
293-
def test_the_alignment_spin_is_sized_to_wall_time_not_a_cycle_count(self):
294-
# A GPU spinning twice as fast (the probe spin takes half the time) must be handed twice
295-
# the cycles, so every rank releases after the same wall time whatever its SM clock.
296-
spins = []
297-
fake = types.SimpleNamespace(cuda=types.SimpleNamespace(
298-
_sleep=spins.append, synchronize=lambda: None,
299-
Event=lambda **kwargs: types.SimpleNamespace(
300-
record=lambda: None, elapsed_time=lambda other: probe_ms,
301-
),
302-
))
303-
timing = ep_timing.GraphTiming(_ChainBackend())
304-
results = {}
305-
for probe_ms in (0.1, 0.05): # the 200k-cycle probe took 100us, then 50us
306-
with mock.patch.dict(sys.modules, {"torch": fake}):
307-
timing.calibrate_align_spin()
308-
results[probe_ms] = timing.align_cycles
309-
self.assertEqual(results[0.1], 200_000) # the default 100us spin at 2 cycles/ns
310-
self.assertEqual(results[0.05], 2 * results[0.1])
311-
312296
def test_the_replay_value_check_poisons_what_dispatch_wrote_before_replaying(self):
313297
backend = _ChainBackend(stage_device_work=True, fp8_consume="native", precision="fp8")
314298
order, combined = [], _Combined(1.0)

0 commit comments

Comments
 (0)