📦 WP-193e — perf(ptl): replay CUDA segmentation-head predictions¶
Engineering log · Phase 18 — CUDA graph replay and pipeline throughput
WP-193e — preserve prediction gradient boundaries while replaying their work¶
A CUDA training trace on an L4 attributes 197.112 ms of four instrumented steps' host dispatch to the detection head's forward calls and 177.742 ms to its matching backward nodes, identified by their forward autograd sequence numbers. Together they occupy 28.40% of that 1.320 s host interval, before the separate prototype and semantic predictions. This diagnostic budgets submission work; complete-fit performance is measured separately below.
The opt-in capture_segment_heads records existing prediction calls through R43's graphed-callable API. Box/class stems retain independent input-gradient contributions. Each branch's coefficient routine, ProtoNet and the semantic predictor keep separate backward nodes, so an unused mask prediction leaves its parameters' gradients absent. The native per-level order, flattening, concatenations, tanh and prediction expressions remain unchanged. Backbone, neck, prototype fusion, assignment, losses, optimizer, schedules and EMA retain their existing execution.
Combining each branch's box/class predictions failed an exact state comparison, so the implementation retains separate stem contributions. FP32 replay also failed the semantic-bias gradient byte oracle. Replay is verified only for BF16 mixed precision on PyTorch 2.9 or newer; FP32 and FP16 retain native calls.
One fresh topology traversal feeds the guard checks. An attribute certificate retains bindings and checks mutable registry names and sequence elements, while current tensor storage, hooks, modes, precision and class dispatch remain independently checked. Changed or custom prediction stages, frozen prediction parameters, aliasing, logical negative/conjugate tensor views and unsupported runtime conditions use the native route. The native fusion is included in this guard without being recorded. Each prediction group retains one recorded input signature; a partial batch stays native. Transient holders keep original checkpoint keys and discard recordings on copying and at fit completion or failure.
Five alternating fresh-process pairs on an isolated NVIDIA L4 g2-standard-8, against main e4db46d after detection-loss and pose-feature replay landed, reduced complete nano segmentation fit time by 16.60%. Native/replay medians (ranges), seconds:
| Host | Native | Replay | Time reduction |
|---|---|---|---|
| g2-standard-8 | 202.794 (201.536–204.709) | 169.130 (168.167–169.786) | 16.60% |
Each fit uses 384 training and 32 disjoint validation COCO val2017 images with the original 80-class annotations, 10 epochs, batch 4, 640 px and BF16 mixed precision. Both arms enable prepared caching, device HSV, epoch prefetch, orthogonalization replay and update replay. The timer includes construction, cache preparation, graph recording, training, EMA, validation and checkpoints. Driver 580.178.04, torch 2.9.1+cu129, Lightning 2.6.6. Whole-process median time, including imports and verification, falls 16.16%.
Per-step-loss, model, optimizer, scheduler and EMA byte hashes and validation metrics match in the five pairs. A separate untimed ten-epoch pair matches eight dense prediction tensors at steps 0, 480 and 959 and matches the timed final state hashes. Runtime source SHA256: 1bc17340e8ede25574a4ad390e75790989dc58d9725f09d961b7cd115ab216f5. This page names neither the hashed artifact nor the hashing method, so a reader cannot reproduce the digest from it; until both are published, the digest records the author's measurement rather than checkable evidence. The two prediction-recorder test files report 134 passed on the L4 with torch 2.9.1 and 126 passed, 8 skipped on torch 2.4. Full pre-commit and the offline suite pass. Full make gate, accelerator goldens and long-run accuracy remain unverified.
Earlier cohorts are historical rather than measurements of this source: main 6eb6abb gave 16.42% on g2-standard-8 and 11.01% on g2-standard-4; the main 28549ba refresh gave 16.33% on g2-standard-8. Their original records remain in the earlier work-package revisions and detailed evidence; they are not a controlled comparison of CPU count or a current-source large-batch claim.
The speed claim is limited to the measured batch-4 workload. An earlier exploratory batch-16 L4 pair (640 px, BF16, three epochs / 72 steps) changed complete fit time from 40.413 s to 39.600 s, only 2.01%. Model, optimizer, scheduler, EMA and per-step-loss bytes and validation metrics matched. That screen used earlier source and a different harness revision; it is neither a repeated large-batch result nor a measurement of this current source. Larger-batch performance on the current source remains unverified.
The private prediction-runner seam in models/heads/detect.py changes dispatch around the existing expressions in their original order. Keeping independent box/class input-gradient additions requires recording the separate prediction calls: the earlier combined-branch screen failed the byte oracle. The default runner still invokes the native calls, and the recorder composes their adapters outside the head. This is a scheduling seam on a restricted surface, not a new method expression or prediction format.
The real automatic Trainer.fit regression uses MuSGD, EMA, valid instance masks and a validation epoch between two training epochs. It observes actual CUDA graph launches after recording and again after validation, instead of accepting populated entries as proof of replay. Forcing permanent native execution while retaining those entries makes this regression fail. The torch 2.4 route keeps its unsupported BF16 recorder native.
A _RecordedGroup's warmup shares the sibling network recorder's own hazard: a failed make_graphed_callables call can leave BatchNorm running-stat writes queued on its warmup stream, so the restore on the current stream must wait for a real synchronize first (test_network_capture_guards.py::TestFailedRecording::test_synchronizes_before_buffer_restore covers this for WP-193c). That case had no counterpart here; a new tests/ptl/test_segment_head_capture.py::test_recorded_group_synchronizes_before_buffer_restore injects the same failure into a real box stem's warmup on a local RTX 4060, asserts the device is synchronized before the BatchNorm buffer is restored, and confirms the failure still raises. Reverting the synchronize call in _RecordedGroup.__init__ makes it fail, so it is not vacuous. This is a local compatibility check, not a timed claim.
_SegmentHeadCapture.retire() shares _NetworkCapture.retire()'s disabled-guard contract too: only the first incompatibility of a fit warns, so a second cause in the same fit must stay silent (test_network_capture_guards.py::TestRetireWarning::test_later_cause_stays_silent covers this for WP-193c). A new tests/ptl/test_segment_head_capture.py::test_later_retire_cause_stays_silent calls retire() twice on a bare holder and confirms only the first call warns; reverting the if self.disabled: return guard makes it fail.
Two CPU-only adapter-call tests also distinguish active replay from a silent native fallback while retaining recordings. With an unchanged signature and no pending backward, the graphed callable must run; while a backward is pending, the native callable must run. The replay case fails under an always-native guard mutation.
Peak CUDA reserved memory medians are 6084→6474 MiB. Graph pools cost memory, and the result is specific to the measured small-batch workload. The implementation keeps first-order automatic optimization on one CUDA device without accumulation; unsupported modes remain native. Gradient accumulation stays native because this package validated replay only at accumulate_grad_batches=1, not because of a known defect in how the recorded backward protects live gradients.
Provenance: R43, this repository's native prediction calls and its own CUDA measurements; R12 supplies the public COCO data. No method arithmetic, coefficient, iteration count or method default changes. A15, A17 and A38 remain unchanged.