File size: 7,067 Bytes
e1acc15 | 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 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 | """CPU stream model: two distinct CUDA devices, never solo-emulation evidence."""
from __future__ import annotations
import ast
import inspect
import textwrap
from collections.abc import Iterator
from contextlib import contextmanager
from pathlib import Path
from types import SimpleNamespace
from typing import Any, cast
import pytest
import torch
from offload.dspark.probabilistic.pp4 import pair_probe
from offload.dspark.worker import sampled_fused, tp2
from offload.dspark.worker.tests.test_sampled_dispatch import paired_fixture
class Streams:
def __init__(self) -> None:
self.device = torch.device('cuda:0')
self.captured: set[torch.device] = set()
self.launches: list[tuple[torch.device, bool]] = []
self.streams: list[torch.device] = []
@contextmanager
def on_device(self, dev: torch.device) -> Iterator[None]:
before, self.device = self.device, dev
try:
yield
finally:
self.device = before
def stream(self, *, device: torch.device) -> torch.device:
self.streams.append(device)
return device
@contextmanager
def graph(self, graph: object, *, stream: torch.device | None = None) -> Iterator[None]:
# Reproduce PyTorch's shared default: it was first created on rank0.
captured = torch.device('cuda:0') if stream is None else stream
self.captured.add(captured)
try:
yield
finally:
self.captured.remove(captured)
def __getitem__(self, grid: object) -> Any:
def launch(*args: Any, **kwargs: Any) -> None:
self.launches.append((self.device, self.device in self.captured))
return launch
def context(monkeypatch: pytest.MonkeyPatch) -> tuple[Streams, sampled_fused.FusedContext, dict[int, torch.Tensor], dict[int, torch.Tensor], dict[int, torch.Tensor]]:
streams = Streams()
ctx = sampled_fused.FusedContext(paired_fixture(), 1, 4)
ctx.model.devs = [torch.device('cuda:0'), torch.device('cuda:1')]
logits = {r: torch.zeros(1, 4, 64, dtype=torch.bfloat16) for r in (0, 1)}
previous = {r: torch.zeros(1, dtype=torch.int64) for r in (0, 1)}
outputs = {r: torch.zeros(1, 4, dtype=torch.int64) for r in (0, 1)}
ctx.primed = {(r, pos, 256, 1) for r in (0, 1) for pos in range(4)}
# Previous-token column views after position0 have stride4.
ctx.primed.update((r, pos, 256, 4) for r in (0, 1) for pos in range(4))
monkeypatch.setattr(torch.cuda, 'device', streams.on_device)
monkeypatch.setattr(torch.cuda, 'Stream', streams.stream)
monkeypatch.setattr(torch.cuda, 'CUDAGraph', object)
monkeypatch.setattr(torch.cuda, 'graph', streams.graph)
monkeypatch.setattr(torch.cuda, 'is_current_stream_capturing', lambda: streams.device in streams.captured)
monkeypatch.setattr(sampled_fused, '_kernel', streams)
return streams, ctx, logits, previous, outputs
def test_actual_pair_capture_uses_distinct_device_streams(monkeypatch: pytest.MonkeyPatch) -> None:
streams, ctx, logits, previous, outputs = context(monkeypatch)
graphs = pair_probe.capture_pair(ctx, logits, previous, outputs, 4)
assert set(graphs) == {0, 1}
assert streams.streams == [torch.device('cuda:0'), torch.device('cuda:1')]
assert streams.launches == [(torch.device(f'cuda:{r}'), True) for r in (0, 1) for _ in range(4)]
assert not ctx.capture_expected
def test_shared_default_stream_refuses_before_eager_peer_launch(monkeypatch: pytest.MonkeyPatch) -> None:
streams, ctx, logits, previous, outputs = context(monkeypatch)
with pytest.raises(RuntimeError, match='rank 1: expected capture'), ctx.capturing((1,)), streams.graph(object()):
ctx.select((1,), logits, previous, outputs, 0)
assert streams.launches == []
assert not ctx.capture_expected
# A failed capture cannot leave intent armed during ordinary eager work.
ctx.select((0, 1), logits, previous, outputs, 0)
assert streams.launches == [(torch.device('cuda:0'), False), (torch.device('cuda:1'), False)]
def test_capture_rank_ownership_cannot_change(monkeypatch: pytest.MonkeyPatch) -> None:
streams, ctx, logits, previous, outputs = context(monkeypatch)
with ctx.capturing((0,)), pytest.raises(RuntimeError, match='rank ownership'):
ctx.select((1,), logits, previous, outputs, 0)
assert not streams.launches
@pytest.mark.parametrize('wrong_rank1_stream', [False, True])
def test_real_worker_capture_path(monkeypatch: pytest.MonkeyPatch, wrong_rank1_stream: bool) -> None:
streams, ctx, logits, previous, outputs = context(monkeypatch)
class Stream:
def __init__(self, device: torch.device) -> None:
self.device = device
def wait_stream(self, other: Stream) -> None:
pass
def graph(g: object, *, pool: object, stream: Stream) -> Any:
return streams.graph(g, stream=stream.device)
monkeypatch.setattr(torch.cuda, 'Stream', Stream)
monkeypatch.setattr(torch.cuda, 'stream', lambda st: streams.on_device(st.device))
monkeypatch.setattr(torch.cuda, 'current_stream', lambda dev: Stream(dev))
monkeypatch.setattr(torch.cuda, 'synchronize', lambda dev: None)
monkeypatch.setattr(torch.cuda, 'graph', graph)
ranks = [SimpleNamespace(r=r, dev=dev, pool=None,
cap_stream=Stream(ctx.model.devs[0] if wrong_rank1_stream else dev)) for r, dev in enumerate(ctx.model.devs)]
model = SimpleNamespace(ALL=(0, 1), devs=ctx.model.devs, emulated=False, ranks=ranks, _pp4_context=ctx)
def fn(active: tuple[int, ...]) -> dict[str, object]:
ctx.select(active, logits, previous, outputs, 0)
return {}
if wrong_rank1_stream:
with pytest.raises(RuntimeError, match='rank 1: expected capture'):
tp2.TPDrafter._capture_one(cast(tp2.TPDrafter, model), fn)
assert not ctx.capture_expected
assert streams.launches[-1] == (torch.device('cuda:0'), True)
else:
result = tp2.TPDrafter._capture_one(cast(tp2.TPDrafter, model), fn)
assert set(result['graphs']) == {0, 1}
assert streams.launches[-2:] == [(torch.device('cuda:0'), True), (torch.device('cuda:1'), True)]
assert streams.launches[:4] == [(torch.device(f'cuda:{r}'), False) for _ in range(2) for r in (0, 1)]
def test_all_paired_capture_sites_have_explicit_stream_and_intent() -> None:
for fn in (pair_probe.capture_pair, tp2.TPDrafter._capture_one):
tree = ast.parse(textwrap.dedent(inspect.getsource(fn)))
calls = [n for n in ast.walk(tree) if isinstance(n, ast.Call)]
graphs = [n for n in calls if isinstance(n.func, ast.Attribute) and n.func.attr == 'graph']
assert graphs
assert all(any(k.arg == 'stream' and not isinstance(k.value, ast.Constant) for k in n.keywords) for n in graphs)
assert any(isinstance(n.func, ast.Attribute) and n.func.attr == 'capturing' for n in calls)
source = Path(inspect.getfile(tp2)).read_text()
assert 'self.cap_stream = torch.cuda.Stream(dev)' in source
|