File size: 1,241 Bytes
0693474 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 | """Eager peer-copy oracle: PyTorch copies/events, no raw peer pointers or spin protocol.
Debug only, outside graph capture. FP32 addition then optional BF16 rounding matches the push collective.
"""
from __future__ import annotations
import torch
def all_reduce(parts: tuple[torch.Tensor, torch.Tensor], dtype: torch.dtype = torch.bfloat16
) -> tuple[torch.Tensor, torch.Tensor]:
if any(x.dtype != torch.float32 or not x.is_cuda for x in parts) or parts[0].shape != parts[1].shape:
raise ValueError("peer-copy reference needs two same-shaped FP32 CUDA tensors")
ready = []
for x in parts:
with torch.cuda.device(x.device):
if torch.cuda.is_current_stream_capturing():
raise RuntimeError("peer-copy reference is eager only")
event = torch.cuda.Event()
event.record(torch.cuda.current_stream(x.device))
ready.append(event)
result = []
for r, x in enumerate(parts):
with torch.cuda.device(x.device):
torch.cuda.current_stream(x.device).wait_event(ready[1 - r])
peer = parts[1 - r].to(x.device, non_blocking=True)
result.append((x + peer).to(dtype))
return result[0], result[1]
|