Skip to content

Commit 78dd9bb

Browse files
committed
fix(gliformer): admit only qualified exact CUDA graphs
1 parent 824462e commit 78dd9bb

3 files changed

Lines changed: 64 additions & 42 deletions

File tree

‎packages/sie_server/src/sie_server/adapters/gliformer/adapter.py‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -674,17 +674,17 @@ def __init__(
674674
decoded in float32. CPU runs float32.
675675
revision: HuggingFace revision to pin when downloading artifacts.
676676
dense_dim: Catalog-declared embedding dimension, checked at load.
677-
cuda_graphs: Operator setting: "off" (default), "exact" or
678-
"bucketed". Capture the supported text encoder on CUDA;
679-
task heads and decoding stay eager. Bucketed mode adds
680-
masked padding and can change floating-point rounding.
677+
cuda_graphs: Operator setting: "off" (default) or "exact".
678+
Capture the supported text encoder on CUDA without adding
679+
padding; task heads and decoding stay eager. Bucketed mode
680+
is unsupported because it failed pretrained output parity.
681681
**kwargs: Additional arguments (ignored, for compatibility).
682682
"""
683683
_ = kwargs
684684
if cuda_graphs is False: # YAML's unquoted off
685685
cuda_graphs = "off"
686686
if not isinstance(cuda_graphs, str) or cuda_graphs not in GRAPH_MODES:
687-
raise ValueError("GLiFormer cuda_graphs must be 'off', 'exact' or 'bucketed'")
687+
raise ValueError("GLiFormer cuda_graphs must be 'off' or 'exact'; bucketed mode is unsupported")
688688
for name, value in (
689689
("inference_batch_tokens", inference_batch_tokens),
690690
("max_prompt_tokens", max_prompt_tokens),

‎packages/sie_server/src/sie_server/adapters/gliformer/cuda_graphs.py‎

Lines changed: 8 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -4,11 +4,10 @@
44
the packed word RNN, task heads, output checks and decoding stay eager. Layout,
55
image, packed and otherwise nonstandard encoder inputs also stay eager.
66
7-
``exact`` records a repeated input shape without padding; ``bucketed`` rounds
8-
lengths up to multiples of 32, preserving the batch size. Graphs hold at most
7+
``exact`` records a repeated input shape without padding. Graphs hold at most
98
2,048 tokens (1,024 for encoders wider than 768), and at most 64 shapes. These
109
are admission bounds, not a performance or numerical-equivalence claim.
11-
Bucketed execution can change floating-point rounding through padding.
10+
Bucketed execution is not supported: padding failed pretrained output parity.
1211
1312
Recording follows the process-wide lock, 10% free-memory headroom and 4%
1413
device-memory budget. A token bucket permits 16 recordings initially, then
@@ -49,7 +48,7 @@
4948

5049
logger = logging.getLogger(__name__)
5150

52-
GRAPH_MODES = ("off", "exact", "bucketed")
51+
GRAPH_MODES = ("off", "exact")
5352
_MAX_GRAPHS = 64
5453
_MAX_SIGHTINGS = 4096
5554
_RECORDING_BURST = 16
@@ -96,11 +95,10 @@ def _capture_device_bytes(graph: Any, device: torch.device, free_before: int) ->
9695

9796

9897
def graph_shape(rows: int, length: int, *, mode: str, window: int, max_tokens: int) -> Key | None:
99-
"""Choose a bounded shape without adding batch rows."""
100-
if mode not in ("exact", "bucketed") or rows < 1 or not 0 < length <= window:
98+
"""Admit a bounded exact shape without adding rows or token positions."""
99+
if mode != "exact" or rows < 1 or not 0 < length <= window:
101100
return None
102-
padded = length if mode == "exact" else min(window, -(-length // 32) * 32)
103-
return (rows, padded) if rows * padded <= max_tokens else None
101+
return (rows, length) if rows * length <= max_tokens else None
104102

105103

106104
def unsupported_reason(transformer: Any, device: str | torch.device) -> str | None:
@@ -153,8 +151,8 @@ def __init__(
153151
name: str = "",
154152
clock: Callable[[], float] = time.monotonic,
155153
) -> None:
156-
if mode not in ("exact", "bucketed") or window < 1:
157-
raise ValueError("GLiFormer graph runner needs exact/bucketed mode and a positive window")
154+
if mode != "exact" or window < 1:
155+
raise ValueError("GLiFormer graph runner needs exact mode and a positive window")
158156
self._transformer = transformer
159157
self._encoder = transformer.model
160158
self._layers = self._encoder.encoder

‎packages/sie_server/tests/adapters/test_gliformer_graphs.py‎

Lines changed: 51 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,7 @@ def __init__(self, **kwargs: Any) -> None:
5656
self.replays = 0
5757
super().__init__(
5858
_Transformer().eval(),
59-
mode=kwargs.pop("mode", "bucketed"),
59+
mode=kwargs.pop("mode", "exact"),
6060
window=256,
6161
pad_token_id=0,
6262
clock=lambda: self.now,
@@ -92,18 +92,28 @@ def room(monkeypatch: pytest.MonkeyPatch) -> dict[str, Any]:
9292
return state
9393

9494

95-
@pytest.mark.parametrize("mode", ["exact", "bucketed"])
96-
def test_shape_policy_is_bounded_and_preserves_batch_rows(mode: str) -> None:
95+
def test_shape_policy_is_bounded_and_preserves_exact_dimensions() -> None:
9796
for rows in (1, 2, 3, 8, 31, 64, 1024):
9897
for length in (1, 15, 31, 32, 33, 100, 255, 256, 257):
99-
key = graphs.graph_shape(rows, length, mode=mode, window=256, max_tokens=1024)
98+
key = graphs.graph_shape(rows, length, mode="exact", window=256, max_tokens=1024)
10099
if key is not None:
101-
assert key[0] == rows
102-
assert length <= key[1] <= 256
100+
assert key == (rows, length)
101+
assert length <= 256
103102
assert key[0] * key[1] <= 1024
104-
assert key[1] == (length if mode == "exact" else -(-length // 32) * 32)
105103
for shape in ((0, 4), (1, 0), (1, 257), (9, 256)):
106-
assert graphs.graph_shape(*shape, mode=mode, window=256, max_tokens=1024) is None
104+
assert graphs.graph_shape(*shape, mode="exact", window=256, max_tokens=1024) is None
105+
106+
107+
def test_bucketed_mode_is_rejected_before_model_or_forward_access() -> None:
108+
class UnavailableModel:
109+
def __getattr__(self, name: str) -> Any:
110+
pytest.fail(f"invalid mode accessed model attribute {name}")
111+
112+
assert graphs.GRAPH_MODES == ("off", "exact")
113+
for length in (31, 32, 33, 256):
114+
assert graphs.graph_shape(1, length, mode="bucketed", window=256, max_tokens=1024) is None
115+
with pytest.raises(ValueError, match="exact mode"):
116+
graphs.EncoderGraphRunner(UnavailableModel(), mode="bucketed", window=256, pad_token_id=0)
107117

108118

109119
def test_second_sighting_records_and_cache_replays(room: dict[str, Any]) -> None:
@@ -119,6 +129,7 @@ def test_second_sighting_records_and_cache_replays(room: dict[str, Any]) -> None
119129

120130
def test_no_headroom_global_recording_lock_and_credit_limits(room: dict[str, Any]) -> None:
121131
runner = _PolicyRunner()
132+
assert runner.offer() is None # exact mode's first sighting
122133
room["headroom"] = False
123134
assert runner.offer() is None
124135
room["headroom"] = True
@@ -128,12 +139,13 @@ def test_no_headroom_global_recording_lock_and_credit_limits(room: dict[str, Any
128139
assert runner.offer() is None
129140
runner.now = 2
130141
assert runner.offer() is not None
131-
assert set(runner.stats.eager) == {"no_headroom", "busy", "recording_paused"}
142+
assert set(runner.stats.eager) == {"first_shape", "no_headroom", "busy", "recording_paused"}
132143
runner.close()
133144

134145

135146
def test_over_budget_capture_is_not_retained(room: dict[str, Any]) -> None:
136147
runner = _PolicyRunner()
148+
assert runner.offer() is None # exact mode's first sighting
137149
room["budget"] = 1
138150
assert runner.offer() is None
139151
assert runner.held_bytes == runner.graph_count == 0
@@ -202,6 +214,7 @@ def accounted_capture(key: graphs.Key, device: torch.device) -> graphs._Graph:
202214
return entry
203215

204216
monkeypatch.setattr(runner, "_capture", accounted_capture)
217+
assert runner.offer() is None # exact mode's first sighting
205218
assert runner.offer() is None
206219
assert runner.graph_count == runner.held_bytes == runner.replays == 0
207220
assert runner.stats.recording_failures == 1
@@ -229,9 +242,10 @@ def accounted_capture(key: graphs.Key, device: torch.device) -> graphs._Graph:
229242

230243
monkeypatch.setattr(runner, "_capture", accounted_capture)
231244
room["budget"] = 4096
245+
assert runner.offer() is None # exact mode's first sighting
232246
assert runner.offer() is None
233247
assert runner.graph_count == runner.held_bytes == runner.replays == 0
234-
assert runner.stats.eager == {"budget_full": 1}
248+
assert runner.stats.eager == {"first_shape": 1, "budget_full": 1}
235249
runner.close()
236250

237251

@@ -240,13 +254,15 @@ def test_capture_failure_restores_eager_and_bounds_retries(room: dict[str, Any],
240254
runner = _PolicyRunner()
241255
runner.capture_error = torch.OutOfMemoryError("allocation failed") if oom else RuntimeError("not graphable")
242256
for length in (1, 33, 65, 97):
257+
assert runner.offer(length=length) is None # first sighting, or still paused/disabled
243258
assert runner.offer(length=length) is None
244259
assert runner.held_bytes == 0
245260
if oom:
246261
assert runner.stats.drops == 1
247262
assert not runner._failed
248263
runner.now = 61
249264
runner.capture_error = None
265+
assert runner.offer() is None # first sighting of this shape
250266
assert runner.offer() is not None
251267
else:
252268
assert len(runner._failed) == graphs._MAX_FAILURES
@@ -257,6 +273,7 @@ def test_capture_failure_restores_eager_and_bounds_retries(room: dict[str, Any],
257273

258274
def test_existing_replay_oom_releases_cache_and_propagates(room: dict[str, Any]) -> None:
259275
runner = _PolicyRunner()
276+
assert runner.offer() is None # exact mode's first sighting
260277
assert runner.offer() is not None
261278
runner.replay_error = torch.OutOfMemoryError("allocation failed")
262279
with pytest.raises(torch.OutOfMemoryError):
@@ -358,6 +375,7 @@ def test_replay_copies_and_slices_outputs_on_the_owned_stream(monkeypatch: pytes
358375

359376
def test_pressure_clear_waits_before_dropping_graph_storage(room: dict[str, Any]) -> None:
360377
runner = _PolicyRunner()
378+
assert runner.offer() is None # exact mode's first sighting
361379
assert runner.offer() is not None
362380
observed: list[int] = []
363381
runner._stream = SimpleNamespace(synchronize=lambda: observed.append(runner.graph_count))
@@ -368,7 +386,7 @@ def test_pressure_clear_waits_before_dropping_graph_storage(room: dict[str, Any]
368386
runner.close()
369387

370388

371-
@pytest.mark.parametrize("mode", ["off", False, "exact", "bucketed"])
389+
@pytest.mark.parametrize("mode", ["off", False, "exact"])
372390
def test_operator_option_and_default_eager_path(mode: str | bool, monkeypatch: pytest.MonkeyPatch) -> None:
373391
reject_unknown_loadtime_options(GLiFormerAdapter, {"cuda_graphs": mode}, model_name="tiny")
374392
adapter = GLiFormerAdapter("tiny", cuda_graphs=mode)
@@ -382,18 +400,19 @@ def test_operator_option_and_default_eager_path(mode: str | bool, monkeypatch: p
382400
assert adapter.release_optional_memory() == 0
383401

384402

385-
@pytest.mark.parametrize("mode", [True, None, 1, "on", "EXACT", []])
403+
@pytest.mark.parametrize("mode", [True, None, 1, "on", "EXACT", "bucketed", []])
386404
def test_invalid_modes_fail_before_loading(mode: Any) -> None:
387405
with pytest.raises(ValueError, match="cuda_graphs"):
388406
GLiFormerAdapter("tiny", cuda_graphs=mode)
389407

390408

391409
def test_pressure_release_keeps_weights_and_unload_restores_forward(room: dict[str, Any]) -> None:
392410
runner = _PolicyRunner()
393-
adapter = GLiFormerAdapter("tiny", cuda_graphs="bucketed")
411+
adapter = GLiFormerAdapter("tiny", cuda_graphs="exact")
394412
transformer = runner._transformer
395413
adapter._model = transformer
396414
adapter._graphs = runner
415+
assert runner.offer() is None # exact mode's first sighting
397416
assert runner.offer() is not None
398417
held = runner.held_bytes
399418
assert adapter.has_releasable_memory()
@@ -534,28 +553,37 @@ def observed_helper(*args: Any) -> torch.Tensor:
534553
ids = torch.arange(7).reshape(1, 7).long()
535554
mask = torch.ones_like(ids)
536555
outputs: list[torch.Tensor] = []
556+
baseline_ready, capture_ready = threading.Event(), threading.Event()
537557

538558
def eager_call() -> None:
539559
with torch.inference_mode():
540560
outputs.append(runner._forward(input_ids=ids, attention_mask=mask))
561+
baseline_ready.set()
562+
assert capture_ready.wait(timeout=5)
563+
outputs.append(runner._forward(input_ids=ids, attention_mask=mask))
541564

565+
thread = threading.Thread(target=eager_call)
542566
try:
543567
with torch.inference_mode():
544568
expected = runner._forward(input_ids=ids, attention_mask=mask)
569+
thread.start()
570+
assert baseline_ready.wait(timeout=5)
545571
table = runner._build_relative_pos(torch.empty((1, 7, 1)))
546572
with runner._layout_attention_scales(torch.device("cpu")), runner._recording(table):
547573
calls.clear()
548574
assert torch.equal(runner._forward(input_ids=ids, attention_mask=mask), expected)
549575
assert calls == []
550-
thread = threading.Thread(target=eager_call)
551-
thread.start()
576+
capture_ready.set()
552577
thread.join(timeout=5)
553578
assert not thread.is_alive()
554579
assert calls == [thread.ident] * len(runner._layers.layer)
555580
assert deberta_2d.scaled_size_sqrt is observed_helper
556-
assert len(outputs) == 1
557-
assert torch.equal(outputs[0], expected)
581+
assert len(outputs) == 2
582+
assert torch.equal(outputs[1], outputs[0])
558583
finally:
584+
capture_ready.set()
585+
if thread.ident is not None:
586+
thread.join(timeout=5)
559587
runner.close()
560588

561589

@@ -846,16 +874,14 @@ def _assert_output_close(actual: Any, expected: Any, tolerance: float) -> None:
846874
assert actual == expected
847875

848876

849-
@pytest.mark.parametrize("mode", ["exact", "bucketed"])
850-
def test_real_all_task_heads_and_usage_survive_encoder_only_padding(
877+
def test_real_all_task_heads_and_usage_survive_exact_encoder_replay(
851878
tiny_checkpoint: Path,
852-
mode: str,
853879
room: dict[str, Any],
854880
monkeypatch: pytest.MonkeyPatch,
855881
) -> None:
856882
"""Exercise the actual SIE load/inference/unload path on a generated checkpoint."""
857883
baseline = GLiFormerAdapter(tiny_checkpoint, max_seq_length=128, compute_precision="float32")
858-
candidate = GLiFormerAdapter(tiny_checkpoint, max_seq_length=128, compute_precision="float32", cuda_graphs=mode)
884+
candidate = GLiFormerAdapter(tiny_checkpoint, max_seq_length=128, compute_precision="float32", cuda_graphs="exact")
859885
monkeypatch.setattr(adapter_module, "unsupported_reason", lambda *_: None)
860886
monkeypatch.setattr(adapter_module, "EncoderGraphRunner", _EagerReplay)
861887
items = [Item(text="Alice works at Acme."), Item(text="Bob likes coffee.")]
@@ -869,19 +895,18 @@ def observe(_module: Any, _inputs: Any, output: Any) -> None:
869895

870896
hook = candidate._model.model.register_forward_hook(observe)
871897
try:
872-
tolerance = 0 if mode == "exact" else 1e-5 # synthetic fp32 padding, not a pretrained/GPU bound
873898
for task in _TASKS:
874899
expected = baseline.extract(items, **task)
875900
candidate.extract(items, **task) # exact mode's first sighting
876901
candidate.extract(items, **task) # capture, then test an ordinary replay
877902
actual = candidate.extract(items, **task)
878903
for field in ("entities", "classifications", "relations", "data", "errors", "input_token_counts"):
879-
_assert_output_close(getattr(actual, field), getattr(expected, field), tolerance)
904+
_assert_output_close(getattr(actual, field), getattr(expected, field), 0)
880905
eager_dense = baseline.encode(items, ["dense"]).dense
881906
candidate.encode(items, ["dense"])
882907
candidate.encode(items, ["dense"])
883908
dense = candidate.encode(items, ["dense"]).dense
884-
torch.testing.assert_close(torch.from_numpy(dense), torch.from_numpy(eager_dense), rtol=0, atol=tolerance)
909+
torch.testing.assert_close(torch.from_numpy(dense), torch.from_numpy(eager_dense), rtol=0, atol=0)
885910
assert observed
886911
assert {"ner_logits", "cat_logits", "joint_rel_logits", "structuring_logits"} <= set().union(
887912
*(set(row) for row in observed)
@@ -903,14 +928,13 @@ def observe(_module: Any, _inputs: Any, output: Any) -> None:
903928

904929
@pytest.mark.skipif(not torch.cuda.is_available(), reason="actual CUDA graph parity needs a GPU")
905930
@pytest.mark.parametrize("layout", [False, True], ids=["deberta", "layout-deberta"])
906-
@pytest.mark.parametrize("mode", ["exact", "bucketed"])
907-
def test_actual_cuda_capture_matches_eager_padded_shape_and_owns_output(layout: bool, mode: str) -> None:
931+
def test_actual_cuda_capture_matches_eager_shape_and_owns_output(layout: bool) -> None:
908932
transformer = _tiny_transformer(layout).to(device="cuda", dtype=torch.float16)
909933
original = transformer.forward
910934
ids = torch.arange(14, device="cuda").reshape(2, 7).long()
911935
mask = torch.ones_like(ids)
912936
mask[1, -2:] = 0
913-
runner = graphs.EncoderGraphRunner(transformer, mode=mode, window=64, pad_token_id=0)
937+
runner = graphs.EncoderGraphRunner(transformer, mode="exact", window=64, pad_token_id=0)
914938
try:
915939
with torch.inference_mode():
916940
key = runner.key(*ids.shape)

0 commit comments

Comments
 (0)