From 97703ed61c45bd69377385ed32c1600fd09a386d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jes=C3=BAs=20Royeth?= Date: Tue, 20 Oct 2026 14:04:58 -0300 Subject: [PATCH 1/6] perf(ptl): replay native CUDA oriented networks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replay the native oriented network without changing arithmetic. Preserve six output fields and retain native execution for unsupported consumers and accumulation. WP: 200 Provenance: R43; native repository stages and measured CUDA profiles Assumptions: none Gate: tests/ptl/test_oriented_network_capture.py Signed-off-by: Jesús Royeth --- CHANGELOG.md | 2 + docs/PROVENANCE.md | 1 + docs/ROADMAP.md | 3 +- docs/TRAINING.md | 4 +- .../phase-17-pose-objective-parity.md | 2 + docs/work_packages/wp-193d.md | 2 +- docs/work_packages/wp-199.md | 4 +- docs/work_packages/wp-200.md | 43 ++ scripts/lint/audit_docs_present.py | 2 +- src/lucid_yolo/ptl/module.py | 84 +++- src/lucid_yolo/ptl/network_capture.py | 58 ++- tests/ptl/test_network_capture.py | 4 +- tests/ptl/test_oriented_network_capture.py | 369 ++++++++++++++++++ 13 files changed, 555 insertions(+), 23 deletions(-) create mode 100644 docs/work_packages/wp-200.md create mode 100644 tests/ptl/test_oriented_network_capture.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 94acb6f5..f2dbb0f9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,8 @@ All notable changes to lucid-yolo are documented here, following the Keep a Chan ### Added +- **Opt-in CUDA network replay for oriented training.** `--model.capture_network true` records the existing backbone, neck and both oriented head branches, including their angle outputs. The native objective and state updates remain outside capture. Accumulation and overridden objectives keep the native route; the default stays off. See [WP-200](docs/work_packages/wp-200.md) for complete-fit measurements and limits. + - **Opt-in CUDA replay of detection assignment and loss.** `--model.capture_loss true` records the existing loss forward/backward with a masked target workspace that grows from observed counts. Current label checks and schedules stay live; returned values own their storage and fit exit releases the recording. Default off, compatible with network and orthogonalization replay. See [WP-199](docs/work_packages/wp-199.md) for measurements and limits. - **Opt-in shared-feature CUDA replay for pose training.** `--model.capture_network true` can record the existing backbone and neck for `task: keypoints`; the head stays native so a batch with no assigned positives keeps absent point-stem gradients. Accumulated pose runs, runs with a fully frozen backbone and neck, and heads with changed topology, hooks or forwards stay native. The default remains off. The speed claim is limited to a small-batch workload (batch 4, one L4): a single batch-16 pair changed fit time by 2.50%, and replay lowered peak allocated memory while raising peak reserved memory, so check throughput and both memory figures on your own batch size. See [WP-193d](docs/work_packages/wp-193d.md) for complete-fit measurements and validation limits. diff --git a/docs/PROVENANCE.md b/docs/PROVENANCE.md index b24ca173..f6097102 100644 --- a/docs/PROVENANCE.md +++ b/docs/PROVENANCE.md @@ -146,3 +146,4 @@ The `R1–R20` row is a genuine collapse and stays one. Those twenty were transc | 2026-10-17 | R44 | Optional compiler for the existing colour expressions in [WP-198](work_packages/wp-198.md); own CUDA profiles, launch-shape measurements and public-pixel comparisons select the fusion. | | 2026-10-17 | R40 | The existing runtime's [persistent DataLoader iterator](https://github.com/pytorch/pytorch/blob/v2.9.1/torch/utils/data/dataloader.py#L481-L494) is read diagnostically for [WP-198](work_packages/wp-198.md); the same reset is scheduled earlier, with local RNG/worker and end-to-end evidence. No framework implementation copied. | | 2026-10-18 | R42 | Update-preparation scheduling for [WP-197c](work_packages/wp-197c.md), based on this repository's optimizer and our L4 traces; stable buffers, warmup rollback and non-nested capture use the existing API contract. No external optimizer or detector implementation consulted. | +| 2026-10-20 | R43 | The PyTorch 2.9 graphed-callable contract is read for [WP-200](work_packages/wp-200.md). Its six-output adapter and unchanged oriented network follow this repository and the measured L4 dispatch cost; no external model implementation consulted. | diff --git a/docs/ROADMAP.md b/docs/ROADMAP.md index ef06acfb..67ee1cf0 100644 --- a/docs/ROADMAP.md +++ b/docs/ROADMAP.md @@ -1,6 +1,6 @@ # 🗺️ Work-Package Roadmap -The agent's work queue: 199 numbered work packages, one commit each, executed in dependency order per the AGENTS.md loop, plus lettered rows recording follow-up work that landed after its package had closed. Transcribed from the governing blueprint (sec. 15) with a live status column. +The agent's work queue: 200 numbered work packages, one commit each, executed in dependency order per the AGENTS.md loop, plus lettered rows recording follow-up work that landed after its package had closed. Transcribed from the governing blueprint (sec. 15) with a live status column. **Scope says what a package does, not what executing it taught.** Measurements, rejected approaches and negative results live in RESEARCH_LOG.md, linked per row as `· log`; choices the papers left open live in ASSUMPTIONS.md and are cited by id; release results live in REPRODUCTION_REPORT.md. A row that restates any of those is duplicating a record that has an owner. @@ -363,6 +363,7 @@ Opened by issue 2 on the tracker: R1 sec. 3.4.2 publishes the pose objective thi | 197c | `perf(optim): replay CUDA update preparation` [DATA][GPU] | Opt-in group preparation replay retains native arithmetic, records momentum with rollback-safe setup and applies changing learning rates outside capture; layout alignment enables final foreach writes (D29) · [log](work_packages/wp-197c.md#wp-197c) | `tests/optim/test_update_replay.py`; `tests/ptl/test_cli.py::test_capture_updates_flag_reaches_optimizer`; focused optimizer/replay/CLI CPU/CUDA checks, mutation proof, five paired public-COCO fits with byte parity and 10.01% median complete-fit reduction (nano, batch 4, BF16, L4; batch-32 screen showed no benefit; batch 16 unmeasured); mypy and full precommit; full `make gate` not run | 197b | 🔄 | | 198 | `perf(data): overlap epoch preparation and fuse colour restore` [DATA][GPU] | Persistent workers prepare up to `num_workers x prefetch_factor` next-epoch batches during validation; an optional Triton pass restores the existing colour boundary without intermediate batches. Fixed sampler ownership, finite epochs and exact pixel codes remain the contract · [log](work_packages/wp-198.md#wp-198) | `tests/ptl/test_epoch_prefetch.py`; `tests/ptl/test_fused_device_hsv.py`; isolated-L4 complete-fit comparisons; full `make gate` not run | 197 | 🔄 | | 199 | `perf(ptl): replay detection assignment and loss` [DATA][GPU] | Opt-in loss forward/backward replay calls the existing assignment and scoring expressions with masked target capacity. One recording grows from observed counts; checks and live alpha stay current, and optimizer work is independent · [log](work_packages/wp-199.md#wp-199) | `tests/ptl/test_loss_capture.py`; focused CPU/CUDA, mutation checks and repeated complete public-COCO fits; full `make gate` and accelerator goldens not run | 198 | 🔄 | +| 200 | `perf(ptl): replay native CUDA oriented networks` [DATA][GPU] | Opt-in network replay retains all six native oriented outputs, while classification, rotated objectives and state updates stay native. Accumulation and changed objectives retain the original route · [log](work_packages/wp-200.md#wp-200) | `tests/ptl/test_oriented_network_capture.py`; actual graph execution, empty-angle gradients and validation transitions; five paired public-DOTA complete fits above 8%; focused CPU/CUDA and full pre-commit; full `make gate`, accelerator goldens and long-run accuracy not run | 193d, 199 | 🔄 | ## ⛓️ Critical path and parallelism diff --git a/docs/TRAINING.md b/docs/TRAINING.md index 65aa2aea..070210e0 100644 --- a/docs/TRAINING.md +++ b/docs/TRAINING.md @@ -135,7 +135,9 @@ Budget about `12 * sum(height * width)` bytes for the retained RGB pixels, plus For `task: keypoints`, the same flag records only backbone and neck; the complete dual head and the current RLE/OKS or Laplace objectives stay native. That preserves absent point-stem gradients on batches with no assigned positives. Gradient accumulation, a fully frozen backbone and neck, changed head topology or forwards, and head hooks/compilation keep pose native; a partial freeze that leaves any backbone or neck parameter trainable still replays. The trainable set is fixed per fit at the first recording, so a fully frozen start, or a later freeze or unfreeze of any backbone or neck parameter, keeps the rest of the fit native after one warning: a recipe that starts with backbone and neck fully frozen trains natively throughout, even after it unfreezes them. The head's ordinary dense classification path supplies gradients to all three shared features. The pose speed-up is likewise limited to a small-batch workload (batch 4): a single batch-16 pair changed fit time by 2.50%, and replay lowered peak allocated memory while raising peak reserved memory (see the memory note below), so check throughput and memory on your own batch size. The pose measurements and remaining validation limits are recorded in [WP-193d](work_packages/wp-193d.md). -This path requires single-device CUDA automatic optimisation and `task: detect` or `task: keypoints`. BF16 replay requires a final PyTorch 2.9 or newer release; 2.9 pre-release and NGC `a0` builds are excluded by design. 2.9 is the first verified runtime: older BF16 runs warn and stay eager because the 2.4 comparison did not meet the exact eager oracle. FP32 and FP16 replay are checked on 2.4 as well. CPU, MPS and other strategies warn and stay eager; other tasks or `compile_step` raise at construction. One image signature is recorded per fit. A partial batch uses eager execution without growing a graph cache. Changes to the recorded stages’ parameter/buffer storage, module modes, hooks or precision retire replay for the rest of that fit. The supplied network must keep its ordinary stage implementations and static Python configuration. Root forward overrides, stage instance-forward overrides and registered hooks use the eager path. Anomaly detection (`--trainer.detect_anomaly true`) also stays eager so backward NaN checks run normally. A failure while recording raises instead of falling back to eager (RNG state and buffers are restored first), so disable `--model.capture_network` if recording fails on your model. Higher-order differentiation is not supported. +For `task: obb`, the flag records backbone, neck and the complete dual head, including both angle outputs. The native classification and oriented losses stay outside capture. Their empty-target expressions retain zero angle gradients. Overridden training consumers, classification losses or oriented helpers keep the complete native call; gradient accumulation also stays native. The measurements and remaining limits are in [WP-200](work_packages/wp-200.md). + +This path requires single-device CUDA automatic optimisation and `task: detect`, `task: keypoints` or `task: obb`. BF16 replay requires a final PyTorch 2.9 or newer release; 2.9 pre-release and NGC `a0` builds are excluded by design. 2.9 is the first verified runtime: older BF16 runs warn and stay eager because the 2.4 comparison did not meet the exact eager oracle. FP32 and FP16 replay are checked on 2.4 as well. CPU, MPS and other strategies warn and stay eager; other tasks or `compile_step` raise at construction. One image signature is recorded per fit. A partial batch uses eager execution without growing a graph cache. Changes to the recorded stages’ parameter/buffer storage, module modes, hooks or precision retire replay for the rest of that fit. The supplied network must keep its ordinary stage implementations and static Python configuration. Root forward overrides, stage instance-forward overrides and registered hooks use the eager path. Anomaly detection (`--trainer.detect_anomaly true`) also stays eager so backward NaN checks run normally. A failure while recording raises instead of falling back to eager (RNG state and buffers are restored first), so disable `--model.capture_network` if recording fails on your model. Higher-order differentiation is not supported. Recording warmup restores RNG and BatchNorm buffers before the first real batch. Accumulation or custom optimiser/zero-grad hooks preserve live gradients outside the graph-owned backward buffers before they can be overwritten. AMP uses an uncached autocast scope as the PyTorch API requires. Validation, prediction and public `forward` calls stay eager. Checkpoints retain the ordinary parameter names and the flag, with no graph state. Each fit records afresh and drops its recording on completion or an exception. CUDA graph pools can increase reserved memory even when peak allocated memory falls; measure both on the intended workload. The measured scope and remaining validation gaps are in [WP-193c](work_packages/wp-193c.md). diff --git a/docs/engineering/phase-17-pose-objective-parity.md b/docs/engineering/phase-17-pose-objective-parity.md index 6016a2a4..21f418d3 100644 --- a/docs/engineering/phase-17-pose-objective-parity.md +++ b/docs/engineering/phase-17-pose-objective-parity.md @@ -35,3 +35,5 @@ --8<-- "work_packages/wp-198.md:engineering" --8<-- "work_packages/wp-199.md:engineering" + +--8<-- "work_packages/wp-200.md:engineering" diff --git a/docs/work_packages/wp-193d.md b/docs/work_packages/wp-193d.md index a3652217..25b522d9 100644 --- a/docs/work_packages/wp-193d.md +++ b/docs/work_packages/wp-193d.md @@ -2,7 +2,7 @@ > Engineering log · [Phase 17 — R1 pose-objective parity](../engineering/phase-17-pose-objective-parity.md) > -> 🔗 Depends on [WP-193c](../work_packages/wp-193c.md) · Needed by none +> 🔗 Depends on [WP-193c](../work_packages/wp-193c.md) · Needed by [WP-200](../work_packages/wp-200.md) > > ⬅️ [WP-193c](../work_packages/wp-193c.md) · [Roadmap](../ROADMAP.md) · [WP-194](../work_packages/wp-194.md) ➡️ diff --git a/docs/work_packages/wp-199.md b/docs/work_packages/wp-199.md index ad684d4e..910bc203 100644 --- a/docs/work_packages/wp-199.md +++ b/docs/work_packages/wp-199.md @@ -2,9 +2,9 @@ > Engineering log · [Phase 17 — R1 pose-objective parity](../engineering/phase-17-pose-objective-parity.md) > -> 🔗 Depends on [WP-198](../work_packages/wp-198.md) · Needed by none +> 🔗 Depends on [WP-198](../work_packages/wp-198.md) · Needed by [WP-200](../work_packages/wp-200.md) > -> ⬅️ [WP-198](../work_packages/wp-198.md) · [Roadmap](../ROADMAP.md) +> ⬅️ [WP-198](../work_packages/wp-198.md) · [Roadmap](../ROADMAP.md) · [WP-200](../work_packages/wp-200.md) ➡️ diff --git a/docs/work_packages/wp-200.md b/docs/work_packages/wp-200.md new file mode 100644 index 00000000..52f8d175 --- /dev/null +++ b/docs/work_packages/wp-200.md @@ -0,0 +1,43 @@ +# WP-200 — replay native CUDA oriented networks + +> Engineering log · [Phase 17 — R1 pose-objective parity](../engineering/phase-17-pose-objective-parity.md) +> +> 🔗 Depends on [WP-193d](../work_packages/wp-193d.md), [WP-199](../work_packages/wp-199.md) · Needed by none +> +> ⬅️ [WP-199](../work_packages/wp-199.md) · [Roadmap](../ROADMAP.md) + + + + + +## WP-200 — six native outputs, one recorded network + +**The oriented head is part of the recording.** The existing `capture_network` option now accepts `task: obb`. Backbone, neck and both head branches run their existing expressions through R43's graphed-callable API. The adapter retains the four dense detection fields and both angle tensors, rebuilding them by name. Assignment, rotated decoding, classification, ProbIoU, rotated L1, angle scoring, schedules, clipping, optimizer, EMA and validation stay in their current loop. No method arithmetic, gain, convention or frozen golden changes. + +**The complete native objective owns the gradient contract.** Its empty-target path retains zero gradients to both angle stems. An overridden objective can instead leave those gradients absent, so a changed training consumer, classification loss or oriented helper keeps the complete native call. Accumulation also stays native, as it does for pose: the existing shared recorder has no established accumulation contract. The ordinary network storage, hooks, precision, single-device CUDA and first-order constraints remain unchanged. One image signature is retained per fit; a temporary partial batch stays native and can return to the existing recording. The default stays off. + +**The stage was priced before implementation.** Four native CUDA steps occupied a 1.466 s instrumented window, with about 24.1% covered by recorded device activity. Network forward dispatch took 378.113 ms, including 174.956 ms in backbone and neck and 202.463 ms in the head. A separate complete fit with lightweight host timers took 68.263 s: network forwards took 15.943 s, or 23.35%, across 208 training and 32 validation calls. That includes validation, which remains native, and is an upper bound rather than a speed claim. The 21.261 s backward timer includes loss work as well as network work and is not wholly removable. Profiler runs are separate from the final timing comparisons. + +**The smaller recording did not clear the investigation's target.** A shared-backbone/neck-only screen preserved model, optimizer, scheduler, EMA and loss bytes, but changed complete-fit time from 68.277 to 63.800 s, a single-pair 6.56% reduction. It was set aside before the final cohort. The complete six-output screen preserved those same states and validation values and changed 68.742 to 59.853 s, a preliminary 12.93%; the repeated cohort, rather than that screen, determines the final claim. + +**Workload.** The registered nano oriented model trains on 104 public DOTA-v1.0 tiles with 2,075 annotations and validates on 13 tiles with 616 annotations. Training and validation use disjoint original source images, 16 and 4 respectively. Eight epochs, batch 4, 1024 px, BF16, seed 123, deterministic kernels, four persistent loader workers and one NVIDIA L4 on a g2-standard-8. Both arms enable prepared source caching, device HSV, epoch prefetch and the existing orthogonalization/update-preparation replay. The comparison uses separate upstream and candidate source trees in fresh processes. Timing brackets include model/trainer construction, cache and recording setup, augmentation, training, EMA, validation and checkpoint writes, with device synchronization at both boundaries. + +**[MEASURED] Five alternating fresh-process pairs.** One excluded warmup pair precedes the five measured pairs. Native and replay ranges do not overlap; every paired reduction is positive (11.669–13.572%). The exact model, optimizer, scheduler, EMA and per-step loss byte hashes agree in each pair and across native repeats, as do validation values. The package import path and hashes bind each arm to its separate source tree. The complete-fit median drops **12.011%**. + +| Metric, median (range across five runs) | Native | Replay | +| -- | -- | -- | +| Complete fit, seconds | 68.642 (68.213–69.333) | 60.397 (59.924–60.784) | +| Whole process, seconds | 73.795 (73.379–74.561) | 65.608 (65.169–65.987) | +| Warmed training step, ms | 183.275 (179.156–184.821) | 139.500 (138.453–141.266) | + +Warmed step diagnostics discard the first 64 steps and use CUDA events, synchronized before reading. Peak allocated memory is 2,572,018,688 bytes native and 1,889,886,208 bytes replay; peak reserved memory is 4,167,041,024 and 4,508,876,800 bytes, respectively, constant across the five runs. Extra graph-pool reservation is a tradeoff, not a capacity improvement. Dependency/compiler and OS caches are warm; each fit reconstructs its prepared image cache and CUDA recording. Package import and post-fit byte hashing are outside the fit timer and included in the separately reported process time. Runtime: Python 3.11.17, torch 2.9.1+cu129, Lightning 2.6.6, L4 driver 580.178.04. + +**The removed work is observed at its owner.** Separate four-step CUDA traces show 31,193 native `cudaLaunchKernel` calls and 12,609 with replay, while graph launches rise from 4 to 12, adding the network forward/backward recordings beside the existing optimizer replay. Device-event union is 353.374 and 358.045 ms; these instrumented windows support the dispatch diagnosis, not a kernel-speed claim. + +**A larger-batch probe limits the claim.** One batch-8 pair, with the same eight epochs and 1024 px inputs, changes complete-fit time from 60.494 to 58.014 s, only 4.099%. State and validation parity still hold, but that one pair is below the investigation's 8% target and is not a repeated performance claim. The five-pair headline remains scoped to batch 4. + +**Validation.** Thirty oriented/constructor cases passed on both torch 2.9.1 and 2.4.0; the floor retains its existing BF16 eager fallback, while FP32 and FP16 exercise real replay. A silent eager reversion and an unsafe objective-guard bypass each make the committed Trainer tests fail. Five additional helper-replacement cases pass on CPU after review. Local CPU checks and source doctests, with those cases included, report 17 passed and 19 CUDA skips. Strict mypy and the full pre-commit hooks pass. No frozen golden is changed. + +**Coverage and limits.** Focused tests compare all six fields, actual CUDA launches, empty-target gradient presence, temporary shape fallback, clipped accumulation, FP32/FP16/BF16, scaler overflow and recovery, native-objective fallback, EMA validation and fit cleanup. FP16 fixtures use a unit initial scale plus one injected overflow so at least one real optimizer update must occur. A silent native reversion must fail the launch assertions; bypassing the objective guard must be rejected by the absent-angle-gradient case. Full `make gate`, accelerator goldens and long-run acceptance accuracy are outside this validation run. + + diff --git a/scripts/lint/audit_docs_present.py b/scripts/lint/audit_docs_present.py index b950d463..8301793d 100644 --- a/scripts/lint/audit_docs_present.py +++ b/scripts/lint/audit_docs_present.py @@ -31,7 +31,7 @@ #: Lowest work-package count ROADMAP.md is allowed to hold. A ratchet, not a target: #: contiguity alone would not notice the last row being deleted. Raise it when adding #: a work package; never lower it. -_WP_FLOOR = 199 +_WP_FLOOR = 200 #: Lowest decision count DECISIONS.md is allowed to hold. A ratchet, not a target: #: contiguity alone would not notice the last row being deleted, since what remains diff --git a/src/lucid_yolo/ptl/module.py b/src/lucid_yolo/ptl/module.py index 6b81175d..bcf6e280 100644 --- a/src/lucid_yolo/ptl/module.py +++ b/src/lucid_yolo/ptl/module.py @@ -189,6 +189,7 @@ ) from lucid_yolo.eval.dota_eval import MAX_DETECTIONS, evaluate_rotated_map, rotated_detections_to_predictions from lucid_yolo.eval.segment_decode import decode_instance_masks +from lucid_yolo.losses import oriented_loss from lucid_yolo.losses.dual_loss import DualBranchLoss, DualLossOutput from lucid_yolo.losses.keypoint_nll_loss import LaplaceNLLLoss from lucid_yolo.losses.mask_loss import instance_mask_loss @@ -216,7 +217,7 @@ from lucid_yolo.optim.musgd import MuSGD from lucid_yolo.optim.schedule import warmup_decay_factor from lucid_yolo.ptl.coco_backend import build_mean_average_precision -from lucid_yolo.ptl.loss_capture import _LossCapture, _LossCaptureLifecycle +from lucid_yolo.ptl.loss_capture import _LossCapture, _LossCaptureLifecycle, _source_signature from lucid_yolo.ptl.network_capture import ( _detection_stages, _module_breach, @@ -788,19 +789,46 @@ def _pose_head_breach(head: torch.nn.Module, topology: tuple[tuple[object, ...], def _validate_replay_options(task: str, compile_step: bool, capture_network: bool, capture_loss: bool) -> None: """Reject replay options outside the objectives they record. - Network replay records the detection or pose stages; loss replay records the detection - objective only. + Network replay records the detection, pose or oriented stages; loss replay records + the detection objective only. Raises: ValueError: If network replay is combined with a task other than detection or - keypoints, loss replay with a task other than detection, or either with the compiler. + keypoints or oriented detection, loss replay with another task, or either with the compiler. """ - if capture_network and (task not in ("detect", "keypoints") or compile_step): - raise ValueError("capture_network=True requires task='detect' or task='keypoints', and compile_step=False") + if capture_network and (task not in ("detect", "keypoints", "obb") or compile_step): + raise ValueError("capture_network=True requires task='detect', 'keypoints' or 'obb', and compile_step=False") if capture_loss and (task != "detect" or compile_step): raise ValueError("capture_loss=True requires task='detect' and compile_step=False") +def _oriented_replay_breach(module: DetectionLitModule) -> str | None: + """Keep full-network replay on the native objective that connects all six outputs. + + The native empty-target path retains zero gradients for the angle stems. An + overridden objective may leave them absent instead, which a recorded backward + would materialize. Changed consumers therefore keep their complete native call. + """ + for name, expected in _ORIENTED_METHODS: + method = getattr(module, name) + if getattr(method, "__func__", method) is not expected: + return "the native oriented objective is overridden" + helpers = ( + decode_ltrb, + decode_rboxes, + oriented_branch_terms, + vars(oriented_loss)["square_angle_loss"], + oriented_loss._no_ground_truth, + ) + if helpers != _ORIENTED_HELPERS: + return "a native oriented objective helper changed" + try: + native_loss = _source_signature(module.loss) is not None + except AttributeError: + native_loss = False + return None if native_loss else "the native oriented classification objective changed" + + class DetectionLitModule(LightningModule): """Lightning detector: model stack, dual-branch loss, and MuSGD (WP-034). @@ -959,10 +987,10 @@ class DetectionLitModule(LightningModule): strides, dtype, device, step count and precision select each recording; this flag alone leaves learning rates, momentum and parameters outside the graphs. Extra workspace and setup can outweigh the saving on short runs. - capture_network: Replay fixed detection or keypoint network stages' CUDA forward and backward + capture_network: Replay fixed detection, keypoint or oriented network stages' CUDA forward and backward with :func:`torch.cuda.make_graphed_callables`. Defaults to ``False`` because capture pays setup and retains a graph pool. Requires - ``task="detect"`` or ``task="keypoints"``, ``compile_step=False`` and + ``task="detect"``, ``task="keypoints"`` or ``task="obb"``, ``compile_step=False`` and single-device CUDA automatic optimization. Pose records backbone and neck; its head stays native to preserve absent gradients in batches with no assigned positives. Pose training stays eager @@ -971,7 +999,10 @@ class DetectionLitModule(LightningModule): module is hooked, overridden or compiled. The trainable set is fixed per fit at the first recording: a fully frozen start, or a later freeze or unfreeze of a recorded parameter, keeps the rest of that fit eager after one warning. - Either task stays eager when the LightningModule carries a hook, overrides + Oriented training records the complete network, including both angle stems, + and retains the native classification and oriented objective. Accumulation + or an overridden objective keeps it eager. Every task stays eager when the + LightningModule carries a hook, overrides ``forward`` or is compiled. Assignment, losses, schedules, optimizer, EMA and evaluation stay eager. One image signature is recorded per fit; partial batches stay eager and @@ -995,7 +1026,7 @@ class DetectionLitModule(LightningModule): Raises: ValueError: If ``task`` is not one of ``"detect"``, ``"segment"``, ``"obb"``, ``"keypoints"``, if ``rotated_iou_form`` is not a known rotated-IoU form, - if ``capture_network`` is combined with a task other than detection or keypoints, + if ``capture_network`` is combined with a task other than detection, keypoints or oriented detection, ``capture_loss`` with a non-detection task, or either with ``compile_step=True``, if ``keypoint_loss`` is not a known keypoint objective, if ``task="keypoints"`` was asked for without a ``num_keypoints``, or if a @@ -1464,10 +1495,10 @@ def _detect_objective( def _training_network(self, images: Tensor) -> DualHeadOutput: """Run the dual head through the training recording when its contract holds. - Detection records all three stages; pose records only backbone and neck. + Detection and oriented detection record all three stages; pose records only backbone and neck. Its optional point stems stay native so batches with no assigned positives keep their absent gradients, rather than zero-valued captured gradients. - Both tasks keep changing instance counts and loss arithmetic outside capture. + All tasks keep changing instance counts and loss arithmetic outside capture. """ head_out: DualHeadOutput if ( @@ -1481,8 +1512,9 @@ def _training_network(self, images: Tensor) -> DualHeadOutput: ): # WP-193d reports the shared recorder's CUDA accumulation checks failing on the unmodified # baseline; whether detection must stay eager here too is open in docs/ESCALATION.md (2026-10-20). - if self._task == "keypoints" and self.trainer.accumulate_grad_batches != 1: - self._network_capture.retire("keypoint replay requires accumulate_grad_batches=1") + if self._task in ("keypoints", "obb") and self.trainer.accumulate_grad_batches != 1: + label = "keypoint" if self._task == "keypoints" else "oriented" + self._network_capture.retire(f"{label} replay requires accumulate_grad_batches=1") if self._task == "keypoints": breach = _pose_head_breach(self.head, self._pose_head_topology) if breach is not None: @@ -1503,6 +1535,8 @@ def _training_network(self, images: Tensor) -> DualHeadOutput: elif self._task == "keypoints": features = self._network_capture.run_features(images, (self.backbone, self.neck)) head_out = self(images) if features is None else cast("DualHeadOutput", self.head(features)) + elif self._task == "obb": + head_out = self._training_oriented_network(images) else: eager = cast("Callable[[Tensor], DualHeadOutput]", self) head_out = self._network_capture.run(images, (self.backbone, self.neck, self.head), eager) @@ -1510,6 +1544,15 @@ def _training_network(self, images: Tensor) -> DualHeadOutput: head_out = self(images) return head_out + def _training_oriented_network(self, images: Tensor) -> DualHeadOutput: + """Record the complete oriented network only while its native consumers stay intact.""" + reason = _oriented_replay_breach(self) + if reason is not None: + self._network_capture.retire(reason) + return cast("DualHeadOutput", self(images)) + eager = cast("Callable[[Tensor], DualHeadOutput]", self) + return self._network_capture.run_oriented(images, (self.backbone, self.neck, self.head), eager) + def _alpha_tensor(self, device: torch.device) -> Tensor: """Lift the current :attr:`alpha` to a 0-d ``float64`` tensor on ``device`` for the compiled region. @@ -2760,3 +2803,16 @@ def _anchor_grid(self, height: int, width: int, device: torch.device) -> tuple[T #: The unmodified public forward; instance/class overrides retain their eager dispatch. _EAGER_NETWORK_FORWARD = DetectionLitModule.forward + +#: Consumers whose native calls retain gradients to every oriented output, even on empty batches. +_ORIENTED_METHODS = tuple( + (name, getattr(DetectionLitModule, name)) + for name in ("training_step", "_shared_step", "_task_extra_loss", "_oriented_extra_loss", "_branch_oriented_terms") +) +_ORIENTED_HELPERS = ( + decode_ltrb, + decode_rboxes, + oriented_branch_terms, + vars(oriented_loss)["square_angle_loss"], + oriented_loss._no_ground_truth, +) diff --git a/src/lucid_yolo/ptl/network_capture.py b/src/lucid_yolo/ptl/network_capture.py index f5fbe725..eceef1fb 100644 --- a/src/lucid_yolo/ptl/network_capture.py +++ b/src/lucid_yolo/ptl/network_capture.py @@ -2,7 +2,7 @@ """Opt-in replay of fixed network stages' CUDA forward and backward work. The network has a fixed image signature while its ground-truth count changes per -batch. Detection records backbone, neck and head; pose records backbone and neck. +batch. Detection and oriented detection record backbone, neck and head; pose records backbone and neck. Both leave assignment, losses, schedules, optimizer and EMA in their existing loop. PyTorch's graphed-callable API owns the autograd bridge; no network arithmetic is rewritten. The recording is transient, @@ -72,6 +72,53 @@ def rebuild(tensors: tuple[Tensor, ...]) -> DualHeadOutput: return DualHeadOutput(o2m_cls=o2m_cls, o2m_box=o2m_box, o2o_cls=o2o_cls, o2o_box=o2o_box) +class _OrientedNetwork(nn.Module): + """Expose both native angle tensors beside the four dense detection fields. + + Examples: + ```pycon + >>> from lucid_yolo.ptl import DetectionLitModule + >>> model = DetectionLitModule(0.34, 0.25, 64, 2, task="obb").eval() + >>> network = _OrientedNetwork(model.backbone, model.neck, model.head) + >>> with torch.no_grad(): + ... tensors = network(torch.zeros(2, 3, 64, 64)) + >>> len(tensors), network.rebuild(tensors).o2m_angle.shape + (6, torch.Size([2, 84, 1])) + + ``` + """ + + def __init__(self, backbone: nn.Module, neck: nn.Module, head: nn.Module) -> None: + super().__init__() + self.backbone, self.neck, self.head = backbone, neck, head + + def forward(self, images: Tensor) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]: + """Run the existing stages and retain the six tensors the native oriented objective reads.""" + output = _detection_stages(self.backbone, self.neck, self.head, images) + assert output.o2m_angle is not None and output.o2o_angle is not None + return ( + output.o2m_cls, + output.o2m_box, + output.o2o_cls, + output.o2o_box, + output.o2m_angle, + output.o2o_angle, + ) + + @staticmethod + def rebuild(tensors: tuple[Tensor, ...]) -> DualHeadOutput: + """Rebuild every recorded field by name, leaving absent modalities absent.""" + o2m_cls, o2m_box, o2o_cls, o2o_box, o2m_angle, o2o_angle = tensors + return DualHeadOutput( + o2m_cls=o2m_cls, + o2m_box=o2m_box, + o2o_cls=o2o_cls, + o2o_box=o2o_box, + o2m_angle=o2m_angle, + o2o_angle=o2o_angle, + ) + + class _FeatureNetwork(nn.Module): """Expose backbone and neck without fixing which optional head gradients exist. @@ -363,6 +410,15 @@ def eager(_: Tensor) -> _Features | None: rebuild: Callable[[tuple[Tensor, ...]], _Features | None] = _FeatureNetwork.rebuild return self._run(images, stages, eager, lambda: _FeatureNetwork(*stages), rebuild) + def run_oriented( + self, + images: Tensor, + stages: tuple[nn.Module, nn.Module, nn.Module], + eager: Callable[[Tensor], DualHeadOutput], + ) -> DualHeadOutput: + """Replay all six network outputs consumed by the unchanged native oriented objective.""" + return self._run(images, stages, eager, lambda: _OrientedNetwork(*stages), _OrientedNetwork.rebuild) + def _run( self, images: Tensor, diff --git a/tests/ptl/test_network_capture.py b/tests/ptl/test_network_capture.py index 18c4194c..c06f96ce 100644 --- a/tests/ptl/test_network_capture.py +++ b/tests/ptl/test_network_capture.py @@ -90,14 +90,14 @@ def test_default_has_no_callback_or_registered_state() -> None: "kwargs", [ pytest.param({"task": "segment"}, id="segment"), - pytest.param({"task": "obb"}, id="obb"), + pytest.param({"task": "obb", "compile_step": True}, id="obb-compile-step"), pytest.param({"compile_step": True}, id="detect-compile-step"), pytest.param({"task": "keypoints", "num_keypoints": 5, "compile_step": True}, id="keypoints-compile-step"), ], ) def test_incompatible_options_are_explicit(kwargs: dict[str, object]) -> None: """Unsupported tasks and the separate compiler path fail before training begins.""" - with pytest.raises(ValueError, match="requires task='detect' or task='keypoints', and compile_step=False"): + with pytest.raises(ValueError, match=r"requires task=.*compile_step=False"): _module(capture_network=True, **kwargs) diff --git a/tests/ptl/test_oriented_network_capture.py b/tests/ptl/test_oriented_network_capture.py new file mode 100644 index 00000000..aacd7130 --- /dev/null +++ b/tests/ptl/test_oriented_network_capture.py @@ -0,0 +1,369 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Oriented replay retains angle fields, native gradients and validation transitions.""" + +from __future__ import annotations + +import copy +import warnings +from collections.abc import Iterator +from dataclasses import fields +from typing import Any +from unittest.mock import MagicMock + +import pytest +import torch +from pytorch_lightning import Callback, Trainer +from pytorch_lightning.plugins.precision import MixedPrecision +from torch.utils.data import DataLoader + +from lucid_yolo.cli.train import DetectionCLI, packaged_config +from lucid_yolo.data.rotated_geom import rboxes_to_polygons +from lucid_yolo.data.targets import Targets +from lucid_yolo.models.heads.detect import DualHeadOutput +from lucid_yolo.ptl import DetectionLitModule, EMACallback +from lucid_yolo.ptl.datamodule import DetectionDataModule +from lucid_yolo.ptl.module import _oriented_replay_breach +from lucid_yolo.ptl.network_capture import _OrientedNetwork + + +@pytest.fixture(autouse=True) +def _deterministic_kernels() -> Iterator[None]: + """Keep native comparisons deterministic and restore the caller's policy afterwards.""" + enabled = torch.are_deterministic_algorithms_enabled() + warn_only = torch.is_deterministic_algorithms_warn_only_enabled() + torch.use_deterministic_algorithms(True) + try: + yield + finally: + torch.use_deterministic_algorithms(enabled, warn_only=warn_only) + + +def _oriented(**kwargs: Any) -> DetectionLitModule: + """Build a small ordinary oriented model whose optimizer fits in a focused test. + + Examples: + >>> _oriented().task + 'obb' + """ + options = {"lr": 1e-4, "warmup_epochs": 0.0, "lrf": 1.0, "w_muon": 0.0, "w_sgd": 1.0} + options.update(kwargs) + return DetectionLitModule(0.34, 0.25, 64, 2, task="obb", **options) + + +def _same(left: Any, right: Any) -> None: + """Compare tensor bytes and recursive training state, including signed zero. + + Examples: + >>> _same({"counter": torch.tensor(2)}, {"counter": torch.tensor(2)}) + """ + if isinstance(left, torch.Tensor): + assert left.shape == right.shape and left.dtype == right.dtype + assert torch.equal( + left.detach().contiguous().reshape(-1).view(torch.uint8), + right.detach().contiguous().reshape(-1).view(torch.uint8), + ) + elif isinstance(left, dict): + assert left.keys() == right.keys() + for key, value in left.items(): + _same(value, right[key]) + elif isinstance(left, (list, tuple)): + assert len(left) == len(right) + for first, second in zip(left, right, strict=True): + _same(first, second) + else: + assert left == right + + +def _batches(layout: str) -> list[tuple[torch.Tensor, list[Targets]]]: + """Keep full, empty and temporary partial batches in one native loader contract. + + Examples: + >>> [len(targets) for _, targets in _batches("partial")] + [2, 1, 2] + """ + image = torch.rand(2, 3, 64, 64) + rboxes = torch.tensor([[32.0, 32.0, 32.0, 12.0, 0.3]]) + corners = rboxes_to_polygons(rboxes) + target = Targets( + boxes=torch.cat((corners.amin(dim=1), corners.amax(dim=1)), dim=1), + labels=torch.tensor([1]), + rboxes=rboxes, + ) + empty = Targets(boxes=torch.empty(0, 4), labels=torch.empty(0, dtype=torch.int64), rboxes=torch.empty(0, 5)) + batches = [(image.clone(), [copy.deepcopy(target), copy.deepcopy(target)]) for _ in range(3)] + if layout == "empty": + batches[1] = (image.clone(), [copy.deepcopy(empty), copy.deepcopy(empty)]) + elif layout == "partial": + batches[1] = (image[:1].clone(), [copy.deepcopy(target)]) + return batches + + +class _ReplayProbe(Callback): + """Observe real graph launches and angle-gradient presence after every training batch.""" + + def __init__(self, launches: list[int]) -> None: + self.launches = launches + self.deltas: list[int] = [] + self.angle_gradients: list[bool] = [] + self.losses: list[torch.Tensor] = [] + self.start = 0 + + def on_train_batch_start(self, *args: Any, **kwargs: Any) -> None: + """Mark the real replay count immediately before this batch's forward.""" + self.start = len(self.launches) + + def on_train_batch_end(self, trainer: Trainer, module: DetectionLitModule, outputs: Any, *args: Any) -> None: + """Retain launched work, live gradients and loss bytes after backward and the optimizer.""" + self.deltas.append(len(self.launches) - self.start) + angles = [parameter for name, parameter in module.named_parameters() if "angle_stems" in name] + assert angles + self.angle_gradients.append(all(parameter.grad is not None for parameter in angles)) + self.losses.append(outputs["loss"].detach().clone()) + + +class _OverflowFirst(Callback): + """Make one backward non-finite, then require the scaler to resume real updates.""" + + def __init__(self) -> None: + self.pending = True + + def on_before_backward(self, trainer: Trainer, module: DetectionLitModule, loss: torch.Tensor) -> None: + """Inject through the public callback without changing which network outputs are consumed.""" + if self.pending: + loss.mul_(float("inf")) + self.pending = False + + +def test_oriented_adapter_preserves_fields() -> None: + """The six-field adapter cannot swap the two angle branches or lose optional fields.""" + model = _oriented().eval() + with torch.no_grad(): + assert model.head.o2o.angle_stems is not None + for parameter in model.head.o2o.angle_stems.parameters(): + parameter.add_(0.1) + image = torch.rand(2, 3, 64, 64) + expected = model(image) + network = _OrientedNetwork(model.backbone, model.neck, model.head) + actual = network.rebuild(network(image)) + assert not torch.equal(expected.o2m_angle, expected.o2o_angle) + for field in fields(DualHeadOutput): + _same(getattr(expected, field.name), getattr(actual, field.name)) + + +def test_oriented_default_and_cpu_fallback() -> None: + """Enabling the existing flag changes no checkpoint keys and stays native without CUDA.""" + native, replayed = _oriented(), _oriented(capture_network=True) + assert native.state_dict().keys() == replayed.state_dict().keys() + assert not native.hparams.capture_network and replayed.hparams.capture_network + assert _oriented_replay_breach(replayed) is None + replayed.trainer = Trainer(accelerator="cpu", logger=False, enable_checkpointing=False) + with pytest.warns(UserWarning, match="single-device CUDA"): + replayed.setup("fit") + assert not replayed._capture_network_active + + +def test_oriented_cli_accepts_network_replay() -> None: + """The packaged oriented recipe builds the flag and keeps rotated-target transport enabled.""" + cli = DetectionCLI( + DetectionLitModule, + DetectionDataModule, + args=[ + "--config", + str(packaged_config("obb_nano_smoke")), + "--model.capture_network", + "true", + "--trainer.accelerator", + "cpu", + "--trainer.logger", + "false", + ], + run=False, + ) + assert cli.model.task == "obb" and cli.model.hparams.capture_network + assert cli.datamodule._rotated_targets + + +@pytest.mark.gpu +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA graph replay") +@pytest.mark.parametrize("precision", ["32-true", "16-mixed", "bf16-mixed"]) +@pytest.mark.parametrize("layout", ["uniform", "empty", "partial"]) +@pytest.mark.parametrize("accumulation", [1, 2]) +def test_oriented_trainer_preserves_native_state( + precision: str, + layout: str, + accumulation: int, + graph_recordings: MagicMock, + bf16_replay_expected: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Actual replay survives EMA validation and preserves empty-target gradients and temporary fallback.""" + torch.manual_seed(42) + batches = _batches(layout) + initial = _oriented().state_dict() + real_replay = torch.cuda.CUDAGraph.replay + launches: list[int] = [] + + def observe(graph: torch.cuda.CUDAGraph) -> None: + """Count actual graph execution while retaining the real runtime call.""" + launches.append(id(graph)) + real_replay(graph) + + monkeypatch.setattr(torch.cuda.CUDAGraph, "replay", observe) + results = [] + enabled = torch.are_deterministic_algorithms_enabled() + warn_only = torch.is_deterministic_algorithms_warn_only_enabled() + try: + for capture in (False, True): + torch.manual_seed(42) + model = _oriented(capture_network=capture) + model.load_state_dict(initial) + probe, ema = _ReplayProbe(launches), EMACallback() + callbacks: list[Callback] = [probe, ema] + precision_options: dict[str, Any] = {"precision": precision} + if precision == "16-mixed": + # The native short fixture skips every update at the default scale. + # Unit scale plus one explicit overflow proves both skipping and recovery. + scaler = torch.amp.GradScaler("cuda", init_scale=1.0) + precision_options = {"plugins": [MixedPrecision("16-mixed", "cuda", scaler=scaler)]} + callbacks.append(_OverflowFirst()) + trainer = Trainer( + accelerator="cuda", + devices=1, + max_epochs=2, + gradient_clip_val=1.0, + logger=False, + enable_checkpointing=False, + enable_progress_bar=False, + enable_model_summary=False, + num_sanity_val_steps=0, + deterministic=True, + callbacks=callbacks, + accumulate_grad_batches=accumulation, + **precision_options, + ) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + trainer.fit( + model, + train_dataloaders=DataLoader(copy.deepcopy(batches), batch_size=None), + val_dataloaders=DataLoader(copy.deepcopy(batches[:1]), batch_size=None), + ) + recording = capture and accumulation == 1 and (precision != "bf16-mixed" or bf16_replay_expected) + retired = [ + str(item.message) for item in caught if str(item.message).startswith("capture_network=True stays") + ] + assert len(retired) == int(capture and not recording), retired + assert graph_recordings.call_count == int(recording) + assert len(probe.deltas) == 6 and all(probe.angle_gradients) + assert trainer.global_step == (3 + accumulation - 1) // accumulation * 2 + for index, delta in enumerate(probe.deltas): + active = recording and not (layout == "partial" and index % 3 == 1) + assert delta >= 2 if active else delta == 0 + assert trainer.optimizers[0].state + assert model._network_capture._network is None + scaler = getattr(trainer.precision_plugin, "scaler", None) + if scaler is not None: + assert scaler.get_scale() < 1.0 + results.append( + copy.deepcopy( + ( + model.state_dict(), + trainer.optimizers[0].state_dict(), + ema.state_dict(), + probe.losses, + None if scaler is None else scaler.state_dict(), + trainer.callback_metrics, + ) + ) + ) + _same(*results) + finally: + torch.use_deterministic_algorithms(enabled, warn_only=warn_only) + + +@pytest.mark.parametrize("method", ["training_step", "_task_extra_loss", "_branch_oriented_terms"]) +def test_oriented_objective_override_stays_native(method: str, monkeypatch: pytest.MonkeyPatch) -> None: + """A changed consumer cannot silently replace absent gradients with recorded zeros.""" + model = _oriented(capture_network=True) + monkeypatch.setattr(model, method, lambda *args: torch.tensor(0.0)) + assert _oriented_replay_breach(model) == "the native oriented objective is overridden" + + +def test_oriented_classification_override_stays_native(monkeypatch: pytest.MonkeyPatch) -> None: + """An unchanged model does not certify a replaced classification consumer.""" + model = _oriented(capture_network=True) + monkeypatch.setattr(model.loss._o2m_loss, "_classification_loss", lambda *args: torch.tensor(0.0)) + assert _oriented_replay_breach(model) == "the native oriented classification objective changed" + + +@pytest.mark.parametrize( + "target", + [ + pytest.param("lucid_yolo.ptl.module.decode_ltrb", id="decode-ltrb"), + pytest.param("lucid_yolo.ptl.module.decode_rboxes", id="decode-rboxes"), + pytest.param("lucid_yolo.ptl.module.oriented_branch_terms", id="oriented-branch-terms"), + pytest.param("lucid_yolo.losses.oriented_loss.square_angle_loss", id="square-angle-loss"), + pytest.param("lucid_yolo.losses.oriented_loss._no_ground_truth", id="no-ground-truth"), + ], +) +def test_oriented_helper_replacement_stays_native(target: str, monkeypatch: pytest.MonkeyPatch) -> None: + """A swapped decode or loss helper retires replay, the oriented twin of WP-199's delegated-stage check.""" + model = _oriented(capture_network=True) + monkeypatch.setattr(target, lambda *args, **kwargs: None) + assert _oriented_replay_breach(model) == "a native oriented objective helper changed" + + +class _ClassificationOnly(DetectionLitModule): + """A supported native customization whose objective leaves angle gradients absent.""" + + def _task_extra_loss(self, context: Any, out: Any, stage: str) -> torch.Tensor: + """Keep classification connected while declining oriented supervision.""" + return torch.zeros_like(out.total) + + +@pytest.mark.gpu +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA training") +def test_oriented_override_keeps_absent_angle_gradients(graph_recordings: MagicMock) -> None: + """The public trainer fallback skips angle state instead of materializing zero gradients.""" + batches = _batches("uniform") + initial = _oriented().state_dict() + results = [] + for capture in (False, True): + model = _ClassificationOnly( + 0.34, + 0.25, + 64, + 2, + task="obb", + capture_network=capture, + w_muon=0.0, + w_sgd=1.0, + lr=1e-4, + warmup_epochs=0.0, + lrf=1.0, + ) + model.load_state_dict(initial) + trainer = Trainer( + accelerator="cuda", + devices=1, + max_epochs=1, + logger=False, + enable_checkpointing=False, + enable_progress_bar=False, + enable_model_summary=False, + num_sanity_val_steps=0, + limit_val_batches=0, + gradient_clip_val=1.0, + deterministic=True, + ) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + trainer.fit(model, train_dataloaders=DataLoader(copy.deepcopy(batches), batch_size=None)) + retired = [str(item.message) for item in caught if str(item.message).startswith("capture_network=True stays")] + assert len(retired) == int(capture) + assert not graph_recordings.call_count + angles = [parameter for name, parameter in model.named_parameters() if "angle_stems" in name] + assert angles and all(parameter.grad is None for parameter in angles) + assert all(parameter not in trainer.optimizers[0].state for parameter in angles) + results.append(copy.deepcopy((model.state_dict(), trainer.optimizers[0].state_dict()))) + _same(*results) From 5bca8a02638b81a810fb4202fcd8894f11d651c3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jes=C3=BAs=20Royeth?= Date: Tue, 20 Oct 2026 14:35:21 -0300 Subject: [PATCH 2/6] test(docs): keep roadmap fixtures aligned with package floor MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Derive the valid roadmap and header expectations from the configured floor. Check rejection exactly one package below the floor. WP: 200 Provenance: none (test fixture maintenance; no new sources) Assumptions: none Gate: scripts/_tests/test_audit_docs_present.py Signed-off-by: Jesús Royeth --- scripts/_tests/test_audit_docs_present.py | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/scripts/_tests/test_audit_docs_present.py b/scripts/_tests/test_audit_docs_present.py index a9754c62..a895b110 100644 --- a/scripts/_tests/test_audit_docs_present.py +++ b/scripts/_tests/test_audit_docs_present.py @@ -83,7 +83,8 @@ def _write_required_docs(docs_dir: Path, repo_root: Path) -> None: _write(docs_dir / "ESCALATION.md", "escalation\n") _write( docs_dir / "ROADMAP.md", - "199 numbered work packages.\n" + "".join(f"| {i:03d} | a | b | c | d | ✅ |\n" for i in range(1, 200)), + f"{audit._WP_FLOOR} numbered work packages.\n" + + "".join(f"| {i:03d} | a | b | c | d | ✅ |\n" for i in range(1, audit._WP_FLOOR + 1)), ) _write(docs_dir / "DATASETS.md", "datasets\n") _write(docs_dir / "TRAINING.md", "training\n") @@ -305,11 +306,16 @@ class TestCheckRoadmapWpIdsUniqueAndComplete: def test_flags_floor(self, tmp_path: Path) -> None: """A roadmap below the work-package floor is reported even when contiguous.""" _write_required_docs(tmp_path / "docs", tmp_path) - _write(tmp_path / "docs" / "ROADMAP.md", "1 numbered work packages.\n| 001 | ... |\n") + floor = audit._WP_FLOOR + _write( + tmp_path / "docs" / "ROADMAP.md", + f"{floor - 1} numbered work packages.\n" + + "".join(f"| {i:03d} | a | b | c | d | ✅ |\n" for i in range(1, floor)), + ) violations = audit.check_roadmap_wp_ids_unique_and_complete(tmp_path / "docs", tmp_path) - assert any("shrank below 199" in violation for violation in violations) + assert violations == [f"roadmap shrank below {floor} work packages: {floor - 1}"] def test_clean_at_floor(self, tmp_path: Path) -> None: """A contiguous roadmap at the floor count reports no violation.""" @@ -365,12 +371,13 @@ def test_flags_mismatched_count(self, tmp_path: Path) -> None: _write_required_docs(tmp_path / "docs", tmp_path) text = (tmp_path / "docs" / "ROADMAP.md").read_text(encoding="utf-8") _write( - tmp_path / "docs" / "ROADMAP.md", text.replace("199 numbered work packages.", "5 numbered work packages.") + tmp_path / "docs" / "ROADMAP.md", + text.replace(f"{audit._WP_FLOOR} numbered work packages.", "5 numbered work packages."), ) violations = audit.check_roadmap_header_count(tmp_path / "docs", tmp_path) - assert violations == ["header states 5 numbered work packages but the table carries 199"] + assert violations == [f"header states 5 numbered work packages but the table carries {audit._WP_FLOOR}"] def test_clean_when_counts_agree(self, tmp_path: Path) -> None: """A stated package count matching the numbered rows reports no violation.""" From 3780b3ec807432245176ec3313573ff2bcc27aa3 Mon Sep 17 00:00:00 2001 From: Jirka Borovec <6035284+Borda@users.noreply.github.com> Date: Wed, 21 Oct 2026 00:07:19 +0200 Subject: [PATCH 3/6] fix(ptl): judge oriented replay on loss identity MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Oriented network replay now asks only whether the loss objects and the callables they use are the native ones. A new _native_loss helper in loss_capture.py holds that identity check, and _source_signature calls it before its scalar setting checks, so its result is unchanged. The oriented guard calls _native_loss instead of wrapping _source_signature in a try/except, so a changed scalar loss setting such as an assigner epsilon no longer retires replay, and the retirement reason reads "the native loss or an assignment implementation changed". The five oriented helper bindings are read by one _oriented_helpers function, used for both the import-time _ORIENTED_HELPERS snapshot and the per-call comparison. [resolve group] PR https://github.com/Borda/lucid-YOLO/pull/23 — items 4 5 WP: 200 Provenance: none (review remediation; no new source) Assumptions: none Gate: tests/ptl/test_oriented_network_capture.py::test_oriented_loss_setting_change_keeps_replay; tests/ptl/test_oriented_network_capture.py::test_oriented_classification_override_stays_native; tests/ptl/test_oriented_network_capture.py::test_oriented_helper_replacement_stays_native --- Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com> Signed-off-by: Jirka Borovec <6035284+Borda@users.noreply.github.com> --- src/lucid_yolo/ptl/loss_capture.py | 25 +++++++++---- src/lucid_yolo/ptl/module.py | 41 +++++++++++----------- tests/ptl/test_oriented_network_capture.py | 17 ++++++++- 3 files changed, 54 insertions(+), 29 deletions(-) diff --git a/src/lucid_yolo/ptl/loss_capture.py b/src/lucid_yolo/ptl/loss_capture.py index 3e5019f3..794cbf27 100644 --- a/src/lucid_yolo/ptl/loss_capture.py +++ b/src/lucid_yolo/ptl/loss_capture.py @@ -91,17 +91,28 @@ def _functions(source: DualBranchLoss) -> tuple[object, ...]: _EXPECTED = _functions(DualBranchLoss()) -def _source_signature(source: DualBranchLoss) -> tuple[object, ...] | None: - """Certify the built-in objects and scalar settings, excluding the live alpha input.""" +def _native_loss(source: DualBranchLoss) -> bool: + """Tell whether the built-in loss objects and the implementations they call are unchanged. + + Only identity is judged: exact component types and the callables :func:`_functions` + reads. Scalar settings are not, so a caller whose loss runs outside any recording can + rely on this without inheriting the workspace limits :func:`_source_signature` adds. + """ if type(source) is not DualBranchLoss: - return None + return False objects = (source, source.o2m_assigner, source.o2o_assigner, source._o2m_loss, source._o2o_loss) - if ( + return ( tuple(type(obj) for obj in objects) - != (DualBranchLoss, SmallTargetAssigner, UniqueAssigner, DetectionBranchLoss, DetectionBranchLoss) - or _functions(source) != _EXPECTED - ): + == (DualBranchLoss, SmallTargetAssigner, UniqueAssigner, DetectionBranchLoss, DetectionBranchLoss) + and _functions(source) == _EXPECTED + ) + + +def _source_signature(source: DualBranchLoss) -> tuple[object, ...] | None: + """Certify the built-in objects and scalar settings, excluding the live alpha input.""" + if not _native_loss(source): return None + objects = (source, source.o2m_assigner, source.o2o_assigner, source._o2m_loss, source._o2o_loss) fields = ("topk", "alpha", "beta", "eps", "s_min", "s_ref") settings = tuple(getattr(assigner, name) for assigner in objects[1:3] for name in fields) settings += tuple(getattr(branch, name) for branch in objects[3:] for name in ("box_gain", "cls_gain", "l1_gain")) diff --git a/src/lucid_yolo/ptl/module.py b/src/lucid_yolo/ptl/module.py index b838700c..b7f88c07 100644 --- a/src/lucid_yolo/ptl/module.py +++ b/src/lucid_yolo/ptl/module.py @@ -217,7 +217,7 @@ from lucid_yolo.optim.musgd import MuSGD from lucid_yolo.optim.schedule import warmup_decay_factor from lucid_yolo.ptl.coco_backend import build_mean_average_precision -from lucid_yolo.ptl.loss_capture import _LossCapture, _LossCaptureLifecycle, _source_signature +from lucid_yolo.ptl.loss_capture import _LossCapture, _LossCaptureLifecycle, _native_loss from lucid_yolo.ptl.network_capture import ( _detection_stages, _module_breach, @@ -816,6 +816,21 @@ def _validate_replay_options( ) +def _oriented_helpers() -> tuple[object, ...]: + """Read the decode and oriented-loss helpers the native objective calls by global name. + + Resolved at call time, so the import-time call that builds :data:`_ORIENTED_HELPERS` + keeps the originals and every later call sees the current bindings. + """ + return ( + decode_ltrb, + decode_rboxes, + oriented_branch_terms, + vars(oriented_loss)["square_angle_loss"], + oriented_loss._no_ground_truth, + ) + + def _oriented_replay_breach(module: DetectionLitModule) -> str | None: """Keep full-network replay on the native objective that connects all six outputs. @@ -827,20 +842,9 @@ def _oriented_replay_breach(module: DetectionLitModule) -> str | None: method = getattr(module, name) if getattr(method, "__func__", method) is not expected: return "the native oriented objective is overridden" - helpers = ( - decode_ltrb, - decode_rboxes, - oriented_branch_terms, - vars(oriented_loss)["square_angle_loss"], - oriented_loss._no_ground_truth, - ) - if helpers != _ORIENTED_HELPERS: + if _oriented_helpers() != _ORIENTED_HELPERS: return "a native oriented objective helper changed" - try: - native_loss = _source_signature(module.loss) is not None - except AttributeError: - native_loss = False - return None if native_loss else "the native oriented classification objective changed" + return None if _native_loss(module.loss) else "the native loss or an assignment implementation changed" class DetectionLitModule(LightningModule): @@ -2844,10 +2848,5 @@ def _anchor_grid(self, height: int, width: int, device: torch.device) -> tuple[T (name, getattr(DetectionLitModule, name)) for name in ("training_step", "_shared_step", "_task_extra_loss", "_oriented_extra_loss", "_branch_oriented_terms") ) -_ORIENTED_HELPERS = ( - decode_ltrb, - decode_rboxes, - oriented_branch_terms, - vars(oriented_loss)["square_angle_loss"], - oriented_loss._no_ground_truth, -) +#: The oriented objective helpers as imported, against which :func:`_oriented_replay_breach` compares. +_ORIENTED_HELPERS = _oriented_helpers() diff --git a/tests/ptl/test_oriented_network_capture.py b/tests/ptl/test_oriented_network_capture.py index aacd7130..81a4b63f 100644 --- a/tests/ptl/test_oriented_network_capture.py +++ b/tests/ptl/test_oriented_network_capture.py @@ -22,6 +22,7 @@ from lucid_yolo.models.heads.detect import DualHeadOutput from lucid_yolo.ptl import DetectionLitModule, EMACallback from lucid_yolo.ptl.datamodule import DetectionDataModule +from lucid_yolo.ptl.loss_capture import _source_signature from lucid_yolo.ptl.module import _oriented_replay_breach from lucid_yolo.ptl.network_capture import _OrientedNetwork @@ -293,7 +294,21 @@ def test_oriented_classification_override_stays_native(monkeypatch: pytest.Monke """An unchanged model does not certify a replaced classification consumer.""" model = _oriented(capture_network=True) monkeypatch.setattr(model.loss._o2m_loss, "_classification_loss", lambda *args: torch.tensor(0.0)) - assert _oriented_replay_breach(model) == "the native oriented classification objective changed" + assert _oriented_replay_breach(model) == "the native loss or an assignment implementation changed" + + +def test_oriented_loss_setting_change_keeps_replay(monkeypatch: pytest.MonkeyPatch) -> None: + """A scalar loss setting outside the network recording leaves oriented replay eligible. + + The oriented loss always runs eagerly after the recorded network, so an assigner epsilon + the detection loss recording refuses cannot change which outputs carry gradients. Both + halves are asserted: the loss certificate still refuses the value, and the network + guard does not inherit that refusal. + """ + model = _oriented(capture_network=True) + monkeypatch.setattr(model.loss.o2m_assigner, "eps", 1e-40) + assert _source_signature(model.loss) is None + assert _oriented_replay_breach(model) is None @pytest.mark.parametrize( From d45c3f2aa12fc30042140b65befdd3ccc2a7f3ec Mon Sep 17 00:00:00 2001 From: Jirka Borovec <6035284+Borda@users.noreply.github.com> Date: Wed, 21 Oct 2026 00:07:19 +0200 Subject: [PATCH 4/6] test(ptl): exercise oriented replay routing on CPU MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CPU tests now reach the CUDA-gated oriented route through CPU images that report is_cuda: an unbroken module hands all three stages to run_oriented, accumulation and an overridden objective each retire once and return the native forward, and plain CPU images never reach the recorder. A subclass override retires replay, the objective-override parametrization covers all five guarded methods, the packaged OBB recipe must pass the replay guard, and setup warns about the ignored flag only for the fit stage while make_graphed_callables is patched to fail. The GPU trainer tests pin each retirement's reason, and the oriented parity oracle also compares LR-scheduler state, with lrf=0.1 so both arms run a decaying schedule. The empty-target OBB training step now runs backward and asserts that every angle_stems parameter of both branches receives a gradient, and the network-replay constructor error is matched in full. [resolve group] PR https://github.com/Borda/lucid-YOLO/pull/23 — items 3 9 6 7 8 10 11 WP: 200 Provenance: none (review remediation; no new source) Assumptions: none Gate: tests/ptl/test_oriented_network_capture.py; tests/ptl/test_obb_training.py::test_image_without_instances_gives_finite_loss; tests/ptl/test_network_capture.py::test_incompatible_options_are_explicit --- Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com> Signed-off-by: Jirka Borovec <6035284+Borda@users.noreply.github.com> --- tests/ptl/test_network_capture.py | 10 +- tests/ptl/test_obb_training.py | 8 +- tests/ptl/test_oriented_network_capture.py | 210 +++++++++++++++++++-- 3 files changed, 214 insertions(+), 14 deletions(-) diff --git a/tests/ptl/test_network_capture.py b/tests/ptl/test_network_capture.py index c06f96ce..fd998041 100644 --- a/tests/ptl/test_network_capture.py +++ b/tests/ptl/test_network_capture.py @@ -11,6 +11,7 @@ import copy import gc import pickle +import re import warnings import weakref from collections.abc import Callable, Iterator, Sequence @@ -96,8 +97,13 @@ def test_default_has_no_callback_or_registered_state() -> None: ], ) def test_incompatible_options_are_explicit(kwargs: dict[str, object]) -> None: - """Unsupported tasks and the separate compiler path fail before training begins.""" - with pytest.raises(ValueError, match=r"requires task=.*compile_step=False"): + """Unsupported tasks and the separate compiler path fail before training begins. + + The whole message is matched, so it names the network-replay option and every task + it accepts: a dropped or added task, or the loss-replay message, would fail here. + """ + message = "capture_network=True requires task='detect', 'keypoints' or 'obb', and compile_step=False" + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): _module(capture_network=True, **kwargs) diff --git a/tests/ptl/test_obb_training.py b/tests/ptl/test_obb_training.py index 01ff7d07..dedb0a45 100644 --- a/tests/ptl/test_obb_training.py +++ b/tests/ptl/test_obb_training.py @@ -399,17 +399,23 @@ def test_image_without_instances_gives_finite_loss() -> None: With no ground truth the padded rotated axis has length zero, which no ``gather`` can index; the term has to answer with an exact zero that is still connected to - the predictions. + the predictions. Connected means the backward reaches every angle stem of both + branches with a gradient, the native premise full-network replay rests on: an empty + batch that left an angle gradient absent would make a recorded backward differ. """ module = _tiny_module() recorder = _LogRecorder() module.log = recorder # type: ignore[method-assign] loss = module.training_step((torch.randn(1, 3, _IMG_SIZE, _IMG_SIZE), [Targets.empty()]), 0) + loss.backward() assert bool(torch.isfinite(loss)) assert recorder.values["train/rbox"] == 0.0 assert recorder.values["train/angle"] == 0.0 + angles = {name: parameter for name, parameter in module.named_parameters() if "angle_stems" in name} + assert any(name.startswith("head.o2o.") for name in angles) and any(name.startswith("head.o2m.") for name in angles) + assert all(parameter.grad is not None for parameter in angles.values()), angles.keys() def test_rotated_candidacy_reaches_assigner() -> None: diff --git a/tests/ptl/test_oriented_network_capture.py b/tests/ptl/test_oriented_network_capture.py index 81a4b63f..43c8fcc6 100644 --- a/tests/ptl/test_oriented_network_capture.py +++ b/tests/ptl/test_oriented_network_capture.py @@ -5,10 +5,11 @@ import copy import warnings -from collections.abc import Iterator +from collections.abc import Callable, Iterator from dataclasses import fields +from types import SimpleNamespace from typing import Any -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest import torch @@ -151,20 +152,37 @@ def test_oriented_adapter_preserves_fields() -> None: _same(getattr(expected, field.name), getattr(actual, field.name)) -def test_oriented_default_and_cpu_fallback() -> None: - """Enabling the existing flag changes no checkpoint keys and stays native without CUDA.""" +@pytest.mark.parametrize("stage", ["fit", "validate", "test", "predict"]) +def test_oriented_default_and_cpu_fallback(stage: str) -> None: + """Enabling the existing flag changes no checkpoint keys and stays native without CUDA. + + Every Lightning stage sets the module up, and only ``fit`` trains, so only ``fit`` + names the ignored flag: an evaluation run reporting a replay it never would use is + noise, and a CPU setup that turned recording on would be a silent regression. The + recorder entry point raises if anything reaches it, so the CPU setup is proven inert. + """ native, replayed = _oriented(), _oriented(capture_network=True) assert native.state_dict().keys() == replayed.state_dict().keys() assert not native.hparams.capture_network and replayed.hparams.capture_network assert _oriented_replay_breach(replayed) is None replayed.trainer = Trainer(accelerator="cpu", logger=False, enable_checkpointing=False) - with pytest.warns(UserWarning, match="single-device CUDA"): - replayed.setup("fit") + with ( + patch("torch.cuda.make_graphed_callables", side_effect=AssertionError("CPU must stay eager")), + warnings.catch_warnings(record=True) as caught, + ): + warnings.simplefilter("always") + replayed.setup(stage) + ignored = [str(item.message) for item in caught if str(item.message).startswith("capture_network=True is ignored")] + assert len(ignored) == int(stage == "fit"), ignored assert not replayed._capture_network_active def test_oriented_cli_accepts_network_replay() -> None: - """The packaged oriented recipe builds the flag and keeps rotated-target transport enabled.""" + """The packaged oriented recipe builds the flag, keeps rotated-target transport and passes the replay guard. + + A shipped setting that tripped the guard would leave every user of the recipe on the + eager path with only a retirement warning to show for it. + """ cli = DetectionCLI( DetectionLitModule, DetectionDataModule, @@ -182,6 +200,7 @@ def test_oriented_cli_accepts_network_replay() -> None: ) assert cli.model.task == "obb" and cli.model.hparams.capture_network assert cli.datamodule._rotated_targets + assert _oriented_replay_breach(cli.model) is None @pytest.mark.gpu @@ -197,9 +216,22 @@ def test_oriented_trainer_preserves_native_state( bf16_replay_expected: bool, monkeypatch: pytest.MonkeyPatch, ) -> None: - """Actual replay survives EMA validation and preserves empty-target gradients and temporary fallback.""" + """Actual replay survives EMA validation and preserves empty-target gradients and temporary fallback. + + A replaying arm that never records names exactly one cause, and the cause is pinned: + accumulation is checked before the recorder reads the precision policy, so with both + breached the accumulation reason is the one reported. A retirement for an unrelated + cause would otherwise satisfy the count while hiding the regression behind it. + """ torch.manual_seed(42) batches = _batches(layout) + retirements = ( + ["accumulate_grad_batches=1"] + if accumulation == 2 + else ["BF16 replay requires"] + if precision == "bf16-mixed" and not bf16_replay_expected + else [] + ) initial = _oriented().state_dict() real_replay = torch.cuda.CUDAGraph.replay launches: list[int] = [] @@ -216,7 +248,8 @@ def observe(graph: torch.cuda.CUDAGraph) -> None: try: for capture in (False, True): torch.manual_seed(42) - model = _oriented(capture_network=capture) + # A decaying LR keeps a real scheduler on both arms, so its state joins the oracle. + model = _oriented(capture_network=capture, lrf=0.1) model.load_state_dict(initial) probe, ema = _ReplayProbe(launches), EMACallback() callbacks: list[Callback] = [probe, ema] @@ -254,6 +287,8 @@ def observe(graph: torch.cuda.CUDAGraph) -> None: str(item.message) for item in caught if str(item.message).startswith("capture_network=True stays") ] assert len(retired) == int(capture and not recording), retired + expected = retirements if capture else [] + assert all(fragment in message for fragment, message in zip(expected, retired, strict=True)), retired assert graph_recordings.call_count == int(recording) assert len(probe.deltas) == 6 and all(probe.angle_gradients) assert trainer.global_step == (3 + accumulation - 1) // accumulation * 2 @@ -270,6 +305,7 @@ def observe(graph: torch.cuda.CUDAGraph) -> None: ( model.state_dict(), trainer.optimizers[0].state_dict(), + trainer.lr_scheduler_configs[0].scheduler.state_dict(), ema.state_dict(), probe.losses, None if scaler is None else scaler.state_dict(), @@ -282,7 +318,9 @@ def observe(graph: torch.cuda.CUDAGraph) -> None: torch.use_deterministic_algorithms(enabled, warn_only=warn_only) -@pytest.mark.parametrize("method", ["training_step", "_task_extra_loss", "_branch_oriented_terms"]) +@pytest.mark.parametrize( + "method", ["training_step", "_shared_step", "_task_extra_loss", "_oriented_extra_loss", "_branch_oriented_terms"] +) def test_oriented_objective_override_stays_native(method: str, monkeypatch: pytest.MonkeyPatch) -> None: """A changed consumer cannot silently replace absent gradients with recorded zeros.""" model = _oriented(capture_network=True) @@ -328,6 +366,140 @@ def test_oriented_helper_replacement_stays_native(target: str, monkeypatch: pyte assert _oriented_replay_breach(model) == "a native oriented objective helper changed" +class _CudaFlagged(torch.Tensor): + """CPU tensor that reports CUDA residency so the CUDA-gated oriented routing is reachable offline. + + Operations unwrap to plain tensors, so the real model still runs on the CPU. A copy of + the pose suite's class: ``tests/`` is not a package, so one test module cannot import + another's helpers. + """ + + is_cuda = True + + @classmethod + def __torch_function__(cls, func, types, args=(), kwargs=None): # type: ignore[no-untyped-def] + """Run the operation on the underlying data and return an ordinary tensor.""" + with torch._C.DisableTorchFunctionSubclass(): + return func(*args, **(kwargs or {})) + + +def _routed_oriented(accumulation: int = 1) -> DetectionLitModule: + """Build an oriented module that satisfies every training-time gate of the full-network route. + + Examples: + >>> _routed_oriented(2).trainer.accumulate_grad_batches + 2 + """ + model = _oriented(capture_network=True) + model._trainer = SimpleNamespace(training=True, accumulate_grad_batches=accumulation) + model._capture_network_active = True + return model + + +def _cuda_flagged_images() -> torch.Tensor: + """Images that pass the ``is_cuda`` gate while living on the CPU. + + Examples: + >>> _cuda_flagged_images().is_cuda + True + """ + return torch.rand(2, 3, 64, 64).as_subclass(_CudaFlagged) + + +def _override_objective(model: DetectionLitModule) -> None: + """Replace one native oriented consumer on this instance, as a user customization would. + + Examples: + >>> model = _oriented() + >>> _override_objective(model) + >>> "_task_extra_loss" in vars(model) + True + """ + model._task_extra_loss = lambda *args: torch.tensor(0.0) # type: ignore[method-assign] + + +def test_oriented_routing_dispatches_full_network_to_recorder() -> None: + """An unbroken oriented module hands all three stages to the oriented recorder and returns its output. + + The CUDA-only route cannot execute on CPU runners, so images spoof ``is_cuda`` and the + recorder is replaced: what is pinned is the dispatch itself, the oriented entry point + with the complete stage tuple and the module as its eager fallback, never the + detection route, and no retirement warning. + """ + model = _routed_oriented() + images = _cuda_flagged_images() + replayed = object() + with ( + patch.object(model._network_capture, "run_oriented", return_value=replayed) as oriented, + patch.object(model._network_capture, "run") as detect, + warnings.catch_warnings(), + ): + warnings.simplefilter("error") + output = model._training_network(images) + oriented.assert_called_once_with(images, (model.backbone, model.neck, model.head), model) + detect.assert_not_called() + assert output is replayed + + +def test_oriented_routing_leaves_cpu_images_native() -> None: + """Images that are not on CUDA never reach the oriented recorder, however many gates pass. + + Every other training-time gate is satisfied, so the image residency check alone + keeps the route native: the recorder is never called, no retirement is warned, and + all six fields equal the module's own eager forward byte for byte. + """ + model = _routed_oriented() + images = torch.rand(2, 3, 64, 64) + expected = model(images) + with ( + patch.object(model._network_capture, "run_oriented") as oriented, + warnings.catch_warnings(), + ): + warnings.simplefilter("error") + output = model._training_network(images) + oriented.assert_not_called() + _same( + [getattr(expected, field.name) for field in fields(DualHeadOutput)], + [getattr(output, field.name) for field in fields(DualHeadOutput)], + ) + + +@pytest.mark.parametrize( + ("accumulation", "mutate", "recorder_calls", "reason"), + [ + pytest.param(2, lambda model: None, 1, "oriented replay requires accumulate_grad_batches=1", id="accumulation"), + pytest.param(1, _override_objective, 0, "the native oriented objective is overridden", id="objective-override"), + ], +) +def test_oriented_routing_retires_then_runs_native( + accumulation: int, mutate: Callable[[DetectionLitModule], None], recorder_calls: int, reason: str +) -> None: + """A broken oriented contract retires the recording once and returns the full native call. + + Accumulation retires before dispatch, so the retired recorder is still reached and hands + back its eager fallback; an overridden consumer retires before the recorder is reached + at all. Either way the output is the module's own forward, run exactly once. + """ + model = _routed_oriented(accumulation) + mutate(model) + images = _cuda_flagged_images() + expected = model(images) + calls: list[None] = [] + model.backbone.register_forward_hook(lambda *_: calls.append(None)) + with ( + patch.object(model._network_capture, "run_oriented", wraps=model._network_capture.run_oriented) as spy, + pytest.warns(UserWarning, match=f"stays eager: {reason}"), + ): + output = model._training_network(images) + assert spy.call_count == recorder_calls + assert model._network_capture._disabled and model._network_capture._network is None + assert len(calls) == 1 + _same( + [getattr(expected, field.name) for field in fields(DualHeadOutput)], + [getattr(output, field.name) for field in fields(DualHeadOutput)], + ) + + class _ClassificationOnly(DetectionLitModule): """A supported native customization whose objective leaves angle gradients absent.""" @@ -336,10 +508,25 @@ def _task_extra_loss(self, context: Any, out: Any, stage: str) -> torch.Tensor: return torch.zeros_like(out.total) +def test_oriented_subclass_override_stays_native() -> None: + """A consumer overridden in a subclass retires replay, as an instance attribute does. + + The guard unwraps bound methods through ``__func__`` before comparing them. A subclass + method is a bound method whose function is not the native one, the path the + instance-level lambdas above never reach; this case runs it without a GPU. + """ + model = _ClassificationOnly(0.34, 0.25, 64, 2, task="obb", capture_network=True) + assert _oriented_replay_breach(model) == "the native oriented objective is overridden" + + @pytest.mark.gpu @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA training") def test_oriented_override_keeps_absent_angle_gradients(graph_recordings: MagicMock) -> None: - """The public trainer fallback skips angle state instead of materializing zero gradients.""" + """The public trainer fallback skips angle state instead of materializing zero gradients. + + The replaying arm retires once, and for the overridden objective specifically: a + retirement for any other cause would leave the same absent gradients behind it. + """ batches = _batches("uniform") initial = _oriented().state_dict() results = [] @@ -376,6 +563,7 @@ def test_oriented_override_keeps_absent_angle_gradients(graph_recordings: MagicM trainer.fit(model, train_dataloaders=DataLoader(copy.deepcopy(batches), batch_size=None)) retired = [str(item.message) for item in caught if str(item.message).startswith("capture_network=True stays")] assert len(retired) == int(capture) + assert all("the native oriented objective is overridden" in message for message in retired), retired assert not graph_recordings.call_count angles = [parameter for name, parameter in model.named_parameters() if "angle_stems" in name] assert angles and all(parameter.grad is None for parameter in angles) From ec54fee7c8e497029cab1587c271430698d32a73 Mon Sep 17 00:00:00 2001 From: Jirka Borovec <6035284+Borda@users.noreply.github.com> Date: Wed, 21 Oct 2026 00:07:19 +0200 Subject: [PATCH 5/6] docs(ptl): state oriented replay scope and limits MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit TRAINING.md and the CHANGELOG entry limit the oriented speed claim to Nano OBB at batch 4, BF16 on one L4, and cite the 4.10% batch-8 fit-time change and the rise in peak reserved memory from 3974 to 4300 MiB. The network-capture docstrings cover the oriented path alongside detection and pose, and the WP-200 coverage list names gradient clipping and accumulation retirement instead of clipped accumulation. [resolve group] PR https://github.com/Borda/lucid-YOLO/pull/23 — items 13 14 16 WP: 200 Provenance: none (wording of existing WP-200 measurements; no new source) Assumptions: none Gate: make gate (docs site and docstring doctests) --- Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com> Signed-off-by: Jirka Borovec <6035284+Borda@users.noreply.github.com> --- CHANGELOG.md | 2 +- docs/TRAINING.md | 2 +- docs/work_packages/wp-200.md | 2 +- src/lucid_yolo/ptl/module.py | 2 +- src/lucid_yolo/ptl/network_capture.py | 8 ++++---- 5 files changed, 8 insertions(+), 8 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index bf9949f0..8ea9e460 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,7 +6,7 @@ All notable changes to lucid-yolo are documented here, following the Keep a Chan ### Added -- **Opt-in CUDA network replay for oriented training.** `--model.capture_network true` records the existing backbone, neck and both oriented head branches, including their angle outputs. The native objective and state updates remain outside capture. Accumulation and overridden objectives keep the native route; the default stays off. See [WP-200](docs/work_packages/wp-200.md) for complete-fit measurements and limits. +- **Opt-in CUDA network replay for oriented training.** `--model.capture_network true` records the existing backbone, neck and both oriented head branches, including their angle outputs. The native objective and state updates remain outside capture. Accumulation and overridden objectives keep the native route; the default stays off. The speed claim is limited to a small-batch workload (Nano OBB at batch 4, BF16, one L4): a single batch-8 pair changed fit time by only 4.10%, and replay raised peak reserved memory from 3974 to 4300 MiB, so check throughput and memory on your own batch size. See [WP-200](docs/work_packages/wp-200.md) for complete-fit measurements and limits. - **Opt-in CUDA replay of detection assignment and loss.** `--model.capture_loss true` records the existing loss forward/backward with a masked target workspace that grows from observed counts. Current label checks and schedules stay live; returned values own their storage and fit exit releases the recording. Default off, compatible with network and orthogonalization replay. See [WP-199](docs/work_packages/wp-199.md) for measurements and limits. diff --git a/docs/TRAINING.md b/docs/TRAINING.md index 085e77e2..4e341e7c 100644 --- a/docs/TRAINING.md +++ b/docs/TRAINING.md @@ -143,7 +143,7 @@ Budget about `12 * sum(height * width)` bytes for the retained RGB pixels, plus For `task: keypoints`, the same flag records only backbone and neck; the complete dual head and the current RLE/OKS or Laplace objectives stay native. That preserves absent point-stem gradients on batches with no assigned positives. Gradient accumulation, a fully frozen backbone and neck, changed head topology or forwards, and head hooks/compilation keep pose native; a partial freeze that leaves any backbone or neck parameter trainable still replays. The trainable set is fixed per fit at the first recording, so a fully frozen start, or a later freeze or unfreeze of any backbone or neck parameter, keeps the rest of the fit native after one warning: a recipe that starts with backbone and neck fully frozen trains natively throughout, even after it unfreezes them. The head's ordinary dense classification path supplies gradients to all three shared features. The pose speed-up is likewise limited to a small-batch workload (batch 4): a single batch-16 pair changed fit time by 2.50%, and replay lowered peak allocated memory while raising peak reserved memory (see the memory note below), so check throughput and memory on your own batch size. The pose measurements and remaining validation limits are recorded in [WP-193d](work_packages/wp-193d.md). -For `task: obb`, the flag records backbone, neck and the complete dual head, including both angle outputs. The native classification and oriented losses stay outside capture. Their empty-target expressions retain zero angle gradients. Overridden training consumers, classification losses or oriented helpers keep the complete native call; gradient accumulation also stays native. The measurements and remaining limits are in [WP-200](work_packages/wp-200.md). +For `task: obb`, the flag records backbone, neck and the complete dual head, including both angle outputs. The native classification and oriented losses stay outside capture. Their empty-target expressions retain zero angle gradients. Overridden training consumers, classification losses or oriented helpers keep the complete native call; gradient accumulation also stays native. The oriented speed-up is likewise limited to a small-batch workload (Nano OBB at batch 4, BF16, one L4): a single batch-8 pair changed fit time by only 4.10%, and replay raised peak reserved memory from 3974 to 4300 MiB, so check throughput and memory on your own batch size. The measurements and remaining limits are in [WP-200](work_packages/wp-200.md). This path requires single-device CUDA automatic optimisation and `task: detect`, `task: keypoints` or `task: obb`. BF16 replay requires a final PyTorch 2.9 or newer release; 2.9 pre-release and NGC `a0` builds are excluded by design. 2.9 is the first verified runtime: older BF16 runs warn and stay eager because the 2.4 comparison did not meet the exact eager oracle. FP32 and FP16 replay are checked on 2.4 as well. CPU, MPS and other strategies warn and stay eager; other tasks or `compile_step` raise at construction. One image signature is recorded per fit. A partial batch uses eager execution without growing a graph cache. Changes to the recorded stages’ parameter/buffer storage, module modes, hooks or precision retire replay for the rest of that fit. The supplied network must keep its ordinary stage implementations and static Python configuration. Root forward overrides, stage instance-forward overrides and registered hooks use the eager path. Anomaly detection (`--trainer.detect_anomaly true`) also stays eager so backward NaN checks run normally. A failure while recording raises instead of falling back to eager (RNG state and buffers are restored first), so disable `--model.capture_network` if recording fails on your model. Higher-order differentiation is not supported. diff --git a/docs/work_packages/wp-200.md b/docs/work_packages/wp-200.md index 52f8d175..054c5f72 100644 --- a/docs/work_packages/wp-200.md +++ b/docs/work_packages/wp-200.md @@ -38,6 +38,6 @@ Warmed step diagnostics discard the first 64 steps and use CUDA events, synchron **Validation.** Thirty oriented/constructor cases passed on both torch 2.9.1 and 2.4.0; the floor retains its existing BF16 eager fallback, while FP32 and FP16 exercise real replay. A silent eager reversion and an unsafe objective-guard bypass each make the committed Trainer tests fail. Five additional helper-replacement cases pass on CPU after review. Local CPU checks and source doctests, with those cases included, report 17 passed and 19 CUDA skips. Strict mypy and the full pre-commit hooks pass. No frozen golden is changed. -**Coverage and limits.** Focused tests compare all six fields, actual CUDA launches, empty-target gradient presence, temporary shape fallback, clipped accumulation, FP32/FP16/BF16, scaler overflow and recovery, native-objective fallback, EMA validation and fit cleanup. FP16 fixtures use a unit initial scale plus one injected overflow so at least one real optimizer update must occur. A silent native reversion must fail the launch assertions; bypassing the objective guard must be rejected by the absent-angle-gradient case. Full `make gate`, accelerator goldens and long-run acceptance accuracy are outside this validation run. +**Coverage and limits.** Focused tests compare all six fields, actual CUDA launches, empty-target gradient presence, temporary shape fallback, gradient clipping, accumulation retirement (native fallback), FP32/FP16/BF16, scaler overflow and recovery, native-objective fallback, EMA validation and fit cleanup. FP16 fixtures use a unit initial scale plus one injected overflow so at least one real optimizer update must occur. A silent native reversion must fail the launch assertions; bypassing the objective guard must be rejected by the absent-angle-gradient case. Full `make gate`, accelerator goldens and long-run acceptance accuracy are outside this validation run. diff --git a/src/lucid_yolo/ptl/module.py b/src/lucid_yolo/ptl/module.py index b7f88c07..0720a9c7 100644 --- a/src/lucid_yolo/ptl/module.py +++ b/src/lucid_yolo/ptl/module.py @@ -1530,7 +1530,7 @@ def _training_network(self, images: Tensor) -> DualHeadOutput: """Run the dual head through the training recording when its contract holds. Detection and oriented detection record all three stages; pose records only backbone and neck. - Its optional point stems stay native so batches with no assigned positives + Pose's optional point stems stay native so batches with no assigned positives keep their absent gradients, rather than zero-valued captured gradients. All tasks keep changing instance counts and loss arithmetic outside capture. """ diff --git a/src/lucid_yolo/ptl/network_capture.py b/src/lucid_yolo/ptl/network_capture.py index eceef1fb..ce3bf524 100644 --- a/src/lucid_yolo/ptl/network_capture.py +++ b/src/lucid_yolo/ptl/network_capture.py @@ -3,7 +3,7 @@ The network has a fixed image signature while its ground-truth count changes per batch. Detection and oriented detection record backbone, neck and head; pose records backbone and neck. -Both leave assignment, losses, schedules, +All three leave assignment, losses, schedules, optimizer and EMA in their existing loop. PyTorch's graphed-callable API owns the autograd bridge; no network arithmetic is rewritten. The recording is transient, bounded to one input signature and discarded at the fit boundary. @@ -368,7 +368,7 @@ def retire(self, reason: str, *, stacklevel: int = 3) -> None: stacklevel: The frame the warning names, counted as :func:`warnings.warn` counts from this method. The default names the caller of whoever called ``retire``; :meth:`_run` passes 4 so a retirement it decides still names - the code that called :meth:`run` or :meth:`run_features`. + the code that called :meth:`run`, :meth:`run_features` or :meth:`run_oriented`. """ if self._disabled: return @@ -434,8 +434,8 @@ def _run( stages: The caller's stage modules, whose identities the recording keeps. eager: Produces the value returned whenever this call does not replay. It is a fallback value, not necessarily the native result: :meth:`run` passes the - native forward, while :meth:`run_features` returns ``None`` so its caller - makes the full native call itself. + native forward, and so does :meth:`run_oriented`, while :meth:`run_features` + returns ``None`` so its caller makes the full native call itself. build: Composes the stages into the module checked for eligibility until a recording exists, and captured by the first one. rebuild: Turns the replayed output tensors back into the caller's value. From 40fb443a391e9e04209e11416eed4b31c1d3360f Mon Sep 17 00:00:00 2001 From: Jirka Borovec <6035284+Borda@users.noreply.github.com> Date: Wed, 21 Oct 2026 00:54:34 +0200 Subject: [PATCH 6/6] docs(provenance): drop repeat R43 usage-log row MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The 2026-10-20 R43 row cited the PyTorch 2.9 graphed-callable contract, which R43 does not register (only the 2.4 API documentation and the v2.4.0 LICENSE), and R43 already has its first-use row on 2026-10-16 in a log defined as first-use only. The row is removed; WP-200's use of R43 stays traced through its commit trailer. [resolve group] PR https://github.com/Borda/lucid-YOLO/pull/23 — items 15 WP: 200 Provenance: none (usage-log correction; R43 register row unchanged) Assumptions: none Gate: pre-commit run --files docs/PROVENANCE.md; scripts/_tests/test_audit_docs_present.py --- Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com> Signed-off-by: Jirka Borovec <6035284+Borda@users.noreply.github.com> --- docs/PROVENANCE.md | 1 - 1 file changed, 1 deletion(-) diff --git a/docs/PROVENANCE.md b/docs/PROVENANCE.md index f6097102..b24ca173 100644 --- a/docs/PROVENANCE.md +++ b/docs/PROVENANCE.md @@ -146,4 +146,3 @@ The `R1–R20` row is a genuine collapse and stays one. Those twenty were transc | 2026-10-17 | R44 | Optional compiler for the existing colour expressions in [WP-198](work_packages/wp-198.md); own CUDA profiles, launch-shape measurements and public-pixel comparisons select the fusion. | | 2026-10-17 | R40 | The existing runtime's [persistent DataLoader iterator](https://github.com/pytorch/pytorch/blob/v2.9.1/torch/utils/data/dataloader.py#L481-L494) is read diagnostically for [WP-198](work_packages/wp-198.md); the same reset is scheduled earlier, with local RNG/worker and end-to-end evidence. No framework implementation copied. | | 2026-10-18 | R42 | Update-preparation scheduling for [WP-197c](work_packages/wp-197c.md), based on this repository's optimizer and our L4 traces; stable buffers, warmup rollback and non-nested capture use the existing API contract. No external optimizer or detector implementation consulted. | -| 2026-10-20 | R43 | The PyTorch 2.9 graphed-callable contract is read for [WP-200](work_packages/wp-200.md). Its six-output adapter and unchanged oriented network follow this repository and the measured L4 dispatch cost; no external model implementation consulted. |