feat: upgrade to Python 3.14 and PyTorch 2.14; add torch.compile and bfloat16

- Python 3.14, torch 2.14.1 (cu130 wheels), NumPy 2; frozen parity
references pass unchanged
- trainer.compile (dynamic-shape torch.compile) and trainer.precision
(bfloat16 autocast; graph sums, loss, optimizer in float32)
- CUDA graph steps use device-tensor learning rates, so one capture per
bucket follows the schedule; checkpoints store portable optimizer state
- trainer=paper_fast: compile + CUDA graphs + bfloat16, 5.9 vs 19.2 ms per
GamGam batch for float32 eager on an A100
- loader workers fork explicitly (Python 3.14 defaults to forkserver);
ONNX export keeps the TorchScript exporter

Correction to 5b73bcb: the pre-change float32 eager step was 21.3 ms per
batch (not ~29 ms; that profile used unshuffled ttH-only batches), so the
default path improved by about 1.2x and the large gains are opt-in.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

.python-version

@@ -1 +1 @@
-3.12
+3.14

CHANGELOG.md

@@ -11,8 +11,10 @@
   selection (`checkpoint=best_auc`), and an ensemble Slurm script.
 - `data.derived_branches` (object selection and tag thresholds) and
   `data.fold_modulus` for ntuples outside the Delphes layout.
-- Opt-in CUDA graph training steps (`trainer.cuda_graphs`) and TF32 matmuls
-  (`trainer.allow_tf32`); `trainer=paper_fast` enables both.
+- Opt-in `torch.compile` (`trainer.compile`), bfloat16 autocast
+  (`trainer.precision`), CUDA graph training steps (`trainer.cuda_graphs`), and
+  TF32 matmuls (`trainer.allow_tf32`); `trainer=paper_fast` enables all four,
+  about 3x faster per GamGam epoch than float32 eager on an A100.
 - `data.loader.num_workers`/`pin_memory` for background batch collation.
 - Per-epoch `history.jsonl`, `inference.metrics_output`, early stopping on
   `train_loss`, and macro precision/recall/F1 for multiclass tasks.
@@ -21,6 +23,12 @@
 
 ### Changed
 
+- The environment moved to Python 3.14 and PyTorch 2.14.1 (CUDA 13.0 wheels;
+  NumPy 2), which is required for `torch.compile`. Frozen ROOT-GNN parity
+  references are reproduced unchanged. ONNX export keeps the TorchScript
+  exporter (`dynamo=False`), and loader workers are explicitly forked because
+  Python 3.14 no longer forks by default.
+
 - DGL was removed. Preparation is vectorized and writes a model-neutral packed
   `EventCache` (cache schema 2) instead of pickled DGL graphs; graph batches
   (`TensorGraph`, edge-index layout) and transformer token batches are built

README_PROJECT.md

@@ -21,7 +21,8 @@ backend.
 
 ## Installation
 
-The supported development environment is Python 3.12 (`>=3.12,<3.13`). Core
+The supported development environment is Python 3.14 (`>=3.14,<3.15`) with
+PyTorch 2.14 (CUDA 13.0 wheels on Linux x86-64). Core
 development is supported on macOS and Linux:
 
 ```bash

docs/configuration.md

@@ -150,14 +150,21 @@ distributed preparation and should be benchmarked on the target machine.
 
 ## Training speed options
 
-`trainer.allow_tf32: true` uses TensorFloat-32 CUDA matrix multiplication and
-`trainer.cuda_graphs: true` captures complete training steps (forward,
-backward, Adam step) as CUDA graphs and replays them. Graph batches are padded
-into node/edge capacity buckets with discarded padding graphs, and steps are
-captured again when the learning rate changes. CUDA graphs require a CUDA
-device, a single process, a graph model, and Adam or AdamW. Both options are
-off by default: results then agree with float32 eager training to
-floating-point rounding rather than bitwise. See [performance.md](performance.md).
+| Option | Effect |
+| --- | --- |
+| `trainer.compile: true` | `torch.compile` the model forward with dynamic shapes (Inductor); the first epoch includes about a minute of compilation |
+| `trainer.precision: bfloat16` | bfloat16 autocast for the forward pass; parameters, optimizer, loss, and graph sums stay float32 |
+| `trainer.allow_tf32: true` | TensorFloat-32 for float32 CUDA matrix multiplication |
+| `trainer.cuda_graphs: true` | capture complete training steps (forward, backward, Adam step) as CUDA graphs and replay them |
+
+For CUDA graphs, graph batches are padded into node/edge capacity buckets with
+discarded padding graphs, and learning rates become device tensors that
+schedulers update in place, so one capture per bucket serves the whole run.
+They require a CUDA device, a single process, a graph model, and Adam or
+AdamW; checkpoints store plain learning rates and resume in any mode. All four
+options are off by default, and with any of them results agree with float32
+eager training to floating-point rounding rather than bitwise.
+`trainer=paper_fast` enables all four. See [performance.md](performance.md).
 
 ## Checkpoints and resolved configuration
 

docs/migration.md

@@ -114,6 +114,14 @@ Intentional differences from the previous main:
 - the scheduler is now stepped before the epoch hook rather than after the
   early-stopping check; the stopping epoch still does not step it.
 
+Environment upgrade (Python 3.12 / PyTorch 2.2.2 to Python 3.14 / PyTorch
+2.14.1, CUDA 13.0, NumPy 2): every frozen parity reference, including the
+bitwise DGL forward outputs and prepared graphs, passes unchanged. ONNX export
+passes `dynamo=False` to keep the TorchScript exporter and its `dynamic_axes`
+contract (the default dynamo exporter needs `onnxscript`); it emits deprecation
+warnings. CUDA graph steps hold learning rates as device tensors; checkpoints
+store them as floats with `capturable` disabled.
+
 ## Validation requirements
 
 Every new model family must provide unit tests for its representation and

docs/performance.md

@@ -28,9 +28,44 @@ threads) against the previous DGL graph cache:
 | Batch construction (1024 events) | 83 ms | 3.3 ms |
 | GPU training epoch | 166.5 ms/batch | 25.8 ms/batch |
 
-Training is now bound by the GPU step itself. Mixed precision,
-`torch.compile`, and size-aware batching are the next levers and have not been
-adopted without measurements; size-aware batching and streaming prediction
+At that point eager training became bound by host-side kernel launches (about
+a thousand small kernels per step for this model), not by GPU arithmetic.
+
+## GamGam training step
+
+Measured on the ATLAS Open Data GamGam task (1024-event batches of about 5,600
+nodes and 30,000 edges; paper model: hidden size 64, 4 MLP layers, 4 graph
+steps), one A100-SXM4-80GB, `OMP_NUM_THREADS=8`, loader workers with pinned
+memory, steady-state `Trainer.train_epoch` over 300 shuffled batches:
+
+| Options | torch 2.2.2, Python 3.12 | torch 2.14.1, Python 3.14 |
+| --- | ---: | ---: |
+| float32 eager (default) | 18.4 ms | 19.2 ms |
+| bfloat16 eager | | 22.4 ms |
+| CUDA graphs + TF32 | 8.1 ms | 8.5 ms |
+| `torch.compile` | n/a | 13.0 ms |
+| compile + CUDA graphs + TF32 | n/a | 6.1 ms |
+| compile + CUDA graphs + bfloat16 (`paper_fast`) | n/a | 5.9 ms |
+
+Before this work the float32 eager step (collation, transfer, and step on the
+training thread) took 21.3 ms per batch against 17.7 ms now with the same
+torch 2.2.2 stack, and epoch metrics used a pairwise ROC AUC that could not
+evaluate the 2M-event training split at all. Training-set AUC after three epochs agreed
+across all rows to within GPU run-to-run variation (0.862-0.867). A full
+`trainer=paper_fast` epoch, including validation, takes about 12.5 s (37 s for
+the first, which compiles) against about 38 s for float32 eager.
+
+Observations:
+
+- bfloat16 alone is slower: the model is small (64 features), so it is not
+  matrix-multiplication bound, and eager autocast adds cast kernels. Under
+  `torch.compile` the casts fuse and bfloat16 is on par with TF32.
+- `torch.compile` alone removes about a third of the launch overhead; CUDA
+  graphs remove the rest, and the two compose.
+- Graph sums accumulate in float32 under autocast (bfloat16 represents
+  integers exactly only to 256, and means divide by edge counts).
+
+Size-aware batching and streaming prediction remain unexplored because they
 affect ordering or output semantics.
 
 ## CPU threads

docs/quickstart.md

@@ -12,7 +12,8 @@ ROOT -> features -> EventCache (packed objects) -> graph or sequence batch -> mo
 
 ## Install
 
-The supported environment is Python 3.12. From the repository root:
+The supported environment is Python 3.14 with PyTorch 2.14. From the
+repository root:
 
 ```bash
 uv sync --dev
@@ -77,8 +78,8 @@ uv run gnn4colliders evaluate --config-name config_atlas_opendata_gamgam \
 
 Preparation takes about 15 s with 8 workers. Training follows the paper
 recipe and keeps `best.pt`, the epoch with the best validation AUC; add
-`trainer=paper_fast` for CUDA graph replay and TF32 (about 2x faster per
-epoch on an A100). The task, its configuration choices, and the comparison
+`trainer=paper_fast` for `torch.compile`, CUDA graph replay, and a bfloat16
+forward (about 3x faster per epoch on an A100). The task, its configuration choices, and the comparison
 with the paper are described in
 [atlas_opendata_gamgam.md](atlas_opendata_gamgam.md).
 

pyproject.toml

@@ -7,14 +7,14 @@ name = "gnn4colliders"
 dynamic = ["version"]
 description = "Machine-learning tools for collider event analysis."
 readme = "README.md"
-requires-python = ">=3.12,<3.13"
+requires-python = ">=3.14,<3.15"
 dependencies = [
     "awkward>=2.0",
     "hydra-core>=1.3",
     "huggingface-hub>=0.36,<1",
-    "numpy>=1.24,<2",
+    "numpy>=2.1",
     "scikit-learn>=1.3",
-    "torch==2.2.2",
+    "torch==2.14.1",
     "uproot>=5.0",
 ]
 
@@ -49,12 +49,12 @@ version = { attr = "gnn4colliders.__version__" }
 
 [tool.uv.sources]
 torch = [
-    { index = "pytorch-cu121", marker = "sys_platform == 'linux' and platform_machine == 'x86_64'" },
+    { index = "pytorch-cu130", marker = "sys_platform == 'linux' and platform_machine == 'x86_64'" },
 ]
 
 [[tool.uv.index]]
-name = "pytorch-cu121"
-url = "https://download.pytorch.org/whl/cu121"
+name = "pytorch-cu130"
+url = "https://download.pytorch.org/whl/cu130"
 explicit = true
 
 [tool.pytest.ini_options]

src/gnn4colliders/config/application.py

@@ -287,7 +287,7 @@ class _CheckpointPolicy:
             "epoch": epoch,
             "seconds": now - self.started,
             "learning_rate": [
-                group["lr"] for group in self.trainer.optimizer.param_groups
+                float(group["lr"]) for group in self.trainer.optimizer.param_groups
             ],
             "train": {"loss": train_result.loss, **train_result.metrics},
         }

src/gnn4colliders/config/factories.py

@@ -179,4 +179,6 @@ def build_trainer(
         scheduler_step=scheduler_step,
         distributed_context=distributed_context,
         cuda_graphs=bool(_get(config, "cuda_graphs", False)),
+        compile_model=bool(_get(config, "compile", False)),
+        precision=str(_get(config, "precision", "float32")),
     )

src/gnn4colliders/configs/trainer/default.yaml

@@ -18,3 +18,7 @@ early_stopping:
 # a single process, graph batches, and an Adam optimizer.
 allow_tf32: false
 cuda_graphs: false
+# torch.compile the model forward (dynamic shapes), and the forward precision:
+# float32 or bfloat16 (autocast; loss, graph sums, and optimizer in float32).
+compile: false
+precision: float32

src/gnn4colliders/configs/trainer/paper.yaml

@@ -19,3 +19,5 @@ early_stopping:
   min_delta: 0.0
 allow_tf32: false
 cuda_graphs: false
+compile: false
+precision: float32

src/gnn4colliders/configs/trainer/paper_fast.yaml

@@ -2,7 +2,11 @@ defaults:
   - paper
   - _self_
 
-# The paper recipe with captured CUDA graph steps and TF32 matmuls: about 2x
-# faster per epoch on an A100 for the GamGam task (docs/performance.md).
-allow_tf32: true
+# The paper recipe with a compiled model, captured CUDA graph steps, and a
+# bfloat16 forward: about 3x faster per epoch than float32 eager on an A100
+# for the GamGam task (docs/performance.md). Results agree with float32 to
+# rounding, not bitwise.
+compile: true
 cuda_graphs: true
+precision: bfloat16
+allow_tf32: true

src/gnn4colliders/data/loader.py

@@ -114,6 +114,10 @@ class EventLoader:
             num_workers=self.num_workers,
             pin_memory=self.pin_memory,
             prefetch_factor=self.prefetch_factor,
+            # Forked workers share the events copy-on-write; Python 3.14's
+            # default (forkserver) would pickle them to every worker. Workers
+            # only collate on the CPU, so forking a CUDA process is safe.
+            multiprocessing_context="fork",
         )
 
 

src/gnn4colliders/export/onnx.py

@@ -304,6 +304,9 @@ def export_root_gnn_onnx(
                 "logits": {0: "graphs"},
             },
             do_constant_folding=True,
+            # The TorchScript exporter keeps the dynamic_axes contract and
+            # needs no onnxscript dependency.
+            dynamo=False,
         )
         validate_onnx(temporary)
         with torch.inference_mode():

src/gnn4colliders/models/root_gnn/edge_network.py

@@ -15,8 +15,10 @@ from .blocks import make_mlp
 def _sum_by_group(
     values: torch.Tensor, groups: torch.Tensor, count: int
 ) -> torch.Tensor:
-    result = values.new_zeros((count, values.shape[1]))
-    return result.index_add_(0, groups, values)
+    # Sums accumulate in at least float32, also under bfloat16 autocast.
+    dtype = torch.promote_types(values.dtype, torch.float32)
+    result = values.new_zeros((count, values.shape[1]), dtype=dtype)
+    return result.index_add_(0, groups, values.to(dtype))
 
 
 def _gather(values: torch.Tensor, index: torch.Tensor) -> torch.Tensor:
@@ -30,8 +32,9 @@ def _mean_by_group(
     values: torch.Tensor, groups: torch.Tensor, sizes: torch.Tensor
 ) -> torch.Tensor:
     """Per-group means; ``sizes`` holds the (exact, integer) group sizes."""
-    sums = _sum_by_group(values, groups, int(sizes.numel()))
-    return sums / sizes.to(values.dtype)[:, None].clamp_min(1)
+    # ``shape[0]`` stays symbolic under torch.compile (no recompile per size).
+    sums = _sum_by_group(values, groups, sizes.shape[0])
+    return sums / sizes.to(sums.dtype)[:, None].clamp_min(1)
 
 
 class EdgeNetwork(nn.Module):

src/gnn4colliders/training/checkpoint.py

@@ -57,6 +57,21 @@ def _state_dict(model: nn.Module) -> dict[str, Any]:
     }
 
 
+def _portable_optimizer_state(optimizer: torch.optim.Optimizer) -> dict[str, Any]:
+    """Optimizer state without CUDA-graph settings (tensor learning rates and
+    ``capturable``), so any training mode can resume from it."""
+    state = optimizer.state_dict()
+    groups = []
+    for group in state["param_groups"]:
+        group = dict(group)
+        if isinstance(group.get("lr"), torch.Tensor):
+            group["lr"] = float(group["lr"])
+        if "capturable" in group:
+            group["capturable"] = False
+        groups.append(group)
+    return {**state, "param_groups": groups}
+
+
 def _rng_state() -> dict[str, Any]:
     state: dict[str, Any] = {
         "python": random.getstate(),
@@ -151,7 +166,7 @@ class CheckpointManager:
             "epoch": epoch,
             "global_step": int(trainer_state.global_step),
             "model_state_dict": _state_dict(model),
-            "optimizer_state_dict": optimizer.state_dict()
+            "optimizer_state_dict": _portable_optimizer_state(optimizer)
             if optimizer is not None
             else None,
             "scheduler_state_dict": scheduler.state_dict()

src/gnn4colliders/training/cuda_graphs.py

@@ -9,7 +9,7 @@ captures serve a whole epoch.
 
 Results agree with eager training to floating-point rounding, not bitwise:
 matrix multiplications see padded row counts and the optimizer runs in its
-``capturable`` form.
+``capturable`` form with device-tensor learning rates.
 """
 
 from __future__ import annotations
@@ -30,15 +30,31 @@ class _Capture:
     buffers: StaticGraphBatch
     loss: torch.Tensor
     logits: torch.Tensor
-    learning_rates: tuple[float, ...]
+
+
+def make_capturable(optimizer: torch.optim.Optimizer) -> None:
+    """Prepare an Adam optimizer for capture (idempotent).
+
+    Learning rates become device tensors, which schedulers update in place,
+    so captured steps follow the schedule without being captured again.
+    Optimizer step counters move to the parameters' device.
+    """
+    for group in optimizer.param_groups:
+        device = group["params"][0].device
+        group["capturable"] = True
+        if not isinstance(group["lr"], torch.Tensor):
+            group["lr"] = torch.tensor(float(group["lr"]), device=device)
+        for parameter in group["params"]:
+            state = optimizer.state.get(parameter, {})
+            if isinstance(state.get("step"), torch.Tensor):
+                state["step"] = state["step"].to(device)
 
 
 class CudaGraphTrainStep:
     """Replay captured training steps for graph batches on a CUDA device.
 
-    The model must not synchronize with the host. A step is captured for each
-    ``(graphs, nodes, edges)`` bucket and captured again whenever a learning
-    rate changes, because captured kernels hold the rate as a constant.
+    The model must not synchronize with the host. A step is captured once per
+    ``(graphs, nodes, edges)`` capacity bucket.
     """
 
     def __init__(
@@ -53,8 +69,7 @@ class CudaGraphTrainStep:
     ) -> None:
         if not isinstance(optimizer, (torch.optim.Adam, torch.optim.AdamW)):
             raise ValueError("CUDA graph training supports the Adam optimizers")
-        for group in optimizer.param_groups:
-            group["capturable"] = True
+        make_capturable(optimizer)
         self.model, self.task, self.optimizer = model, task, optimizer
         self.forward = forward
         self.node_multiple, self.edge_multiple = node_multiple, edge_multiple
@@ -64,13 +79,9 @@ class CudaGraphTrainStep:
     def supports(batch: Any) -> bool:
         return isinstance(batch, GraphBatch) and batch.graph.node_features.is_cuda
 
-    def _learning_rates(self) -> tuple[float, ...]:
-        return tuple(float(group["lr"]) for group in self.optimizer.param_groups)
-
-    def _capture(
-        self, key: tuple[int, int, int], batch: GraphBatch, rates: tuple[float, ...]
-    ) -> _Capture:
-        self._captures.pop(key, None)
+    def _capture(self, key: tuple[int, int, int], batch: GraphBatch) -> _Capture:
+        # A restored checkpoint may have replaced the optimizer settings.
+        make_capturable(self.optimizer)
         buffers = StaticGraphBatch(batch, key)
         buffers.load(batch)
         # Gradients allocated during capture belong to this graph's pool.
@@ -83,7 +94,7 @@ class CudaGraphTrainStep:
             loss = self.task.loss(logits, buffers.batch)
             loss.backward()
             self.optimizer.step()
-        capture = _Capture(graph, buffers, loss, logits, rates)
+        capture = _Capture(graph, buffers, loss, logits)
         self._captures[key] = capture
         return capture
 
@@ -92,10 +103,9 @@ class CudaGraphTrainStep:
         key = graph_batch_capacity(
             batch, node_multiple=self.node_multiple, edge_multiple=self.edge_multiple
         )
-        rates = self._learning_rates()
         capture = self._captures.get(key)
-        if capture is None or capture.learning_rates != rates:
-            capture = self._capture(key, batch, rates)
+        if capture is None:
+            capture = self._capture(key, batch)
         capture.buffers.load(batch)
         capture.graph.replay()
         return capture.loss.detach().clone(), capture.logits.detach().clone()

src/gnn4colliders/training/trainer.py

@@ -33,8 +33,22 @@ CUDA_GRAPH_WARMUP_STEPS = 3
 _POSITIONAL_ARGUMENTS: dict[type, int] = {}
 
 
+#: ``trainer.precision`` values and their autocast dtype (``None``: float32).
+PRECISIONS = {"float32": None, "bfloat16": torch.bfloat16}
+
+
+def _unwrapped(model: nn.Module) -> nn.Module:
+    """The user module inside ``torch.compile`` and DDP wrappers."""
+    while True:
+        inner = getattr(model, "_orig_mod", None) or getattr(model, "module", None)
+        if not isinstance(inner, nn.Module):
+            return model
+        model = inner
+
+
 def _positional_arguments(model: nn.Module) -> int:
     """Number of positional ``forward`` parameters, cached per model class."""
+    model = _unwrapped(model)
     cached = _POSITIONAL_ARGUMENTS.get(type(model))
     if cached is not None:
         return cached
@@ -99,9 +113,13 @@ class Trainer:
         scheduler_step: str = "epoch",
         distributed_context: DistributedContext | None = None,
         cuda_graphs: bool = False,
+        compile_model: bool = False,
+        precision: str = "float32",
     ) -> None:
         if scheduler_step not in {"epoch", "validation"}:
             raise ValueError("scheduler_step must be 'epoch' or 'validation'")
+        if precision not in PRECISIONS:
+            raise ValueError(f"precision must be one of {sorted(PRECISIONS)}")
         self.distributed_context = (
             distributed_context or DistributedContext.single_process(device)
         )
@@ -113,6 +131,13 @@ class Trainer:
         self.early_stopping = early_stopping
         self.scheduler_step = scheduler_step
         self.state = TrainerState()
+        self.precision = precision
+        self._autocast_dtype = PRECISIONS[precision]
+        # ``self.model`` stays the eager module for checkpoints and callers;
+        # forward passes go through the compiled wrapper when requested.
+        self._forward_module = (
+            torch.compile(self.model, dynamic=True) if compile_model else self.model
+        )
         self._graph_step = None
         if cuda_graphs:
             if self.distributed_context.enabled:
@@ -122,9 +147,19 @@ class Trainer:
             from .cuda_graphs import CudaGraphTrainStep
 
             self._graph_step = CudaGraphTrainStep(
-                self.model, task, optimizer, _model_forward
+                self.model, task, optimizer, lambda _model, batch: self._forward(batch)
             )
 
+    def _forward(self, batch: Any) -> torch.Tensor:
+        """Model logits, in float32 even when computed under autocast."""
+        if self._autocast_dtype is None:
+            return _model_forward(self._forward_module, batch)
+        with torch.autocast(
+            device_type=torch.device(self.device).type, dtype=self._autocast_dtype
+        ):
+            logits = _model_forward(self._forward_module, batch)
+        return logits.float()
+
     def _step(self, batch: Any) -> tuple[torch.Tensor, torch.Tensor]:
         """One optimizer step; returns the detached loss and logits.
 
@@ -132,7 +167,7 @@ class Trainer:
         batch while this one runs.
         """
         self.optimizer.zero_grad(set_to_none=True)
-        logits = _model_forward(self.model, batch)
+        logits = self._forward(batch)
         loss = self.task.loss(logits, batch)
         loss.backward()
         self.optimizer.step()
@@ -206,7 +241,7 @@ class Trainer:
         with torch.inference_mode():
             for batch in loader:
                 moved = batch.to(self.device, non_blocking=True)
-                outputs.append(_model_forward(self.model, moved).detach())
+                outputs.append(self._forward(moved).detach())
                 batches.append(_summary(moved))
         if not outputs:
             raise ValueError("cannot evaluate an empty loader")

tests/fixtures/root_gnn_dgl_reference/generate.py

@@ -14,10 +14,10 @@ from pathlib import Path
 
 import numpy as np
 import torch
+from gnn4colliders.graphs.preparation import build_graph_samples, open_event_source
 
 from gnn4colliders.compat import load_legacy_checkpoint
 from gnn4colliders.data import GraphDataLoader, GraphDataset
-from gnn4colliders.graphs.preparation import build_graph_samples, open_event_source
 from gnn4colliders.models.root_gnn import (
     EdgeNetwork,
     FineTunedEdgeNetwork,

tests/unit/training/test_lifecycle_policies.py

@@ -123,3 +123,62 @@ def test_cuda_graph_steps_match_eager_gradients():
         assert logits.shape == (16, 3)
         for left, right in zip(eager.parameters(), graphed.parameters()):
             torch.testing.assert_close(right.grad, left.grad, rtol=1e-4, atol=1e-6)
+
+
+def test_bfloat16_autocast_trains_with_float32_logits_and_weights():
+    events = _events(seed=6)
+    trainer = _trainer(events, precision="bfloat16")
+    batch = next(iter(EventLoader(events, collate_graph_batch, 8)))
+    loss, logits = trainer._step(batch)
+    assert logits.dtype == torch.float32 and torch.isfinite(loss)
+    assert all(p.dtype == torch.float32 for p in trainer.model.parameters())
+    reference = _trainer(events)._forward(batch)
+    torch.testing.assert_close(
+        _trainer(events, precision="bfloat16")._forward(batch),
+        reference,
+        rtol=5e-2,
+        atol=5e-2,
+    )
+
+
+def test_unknown_precision_is_rejected():
+    with pytest.raises(ValueError, match="precision"):
+        _trainer(_events(), precision="float16")
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA unavailable")
+def test_compiled_graph_steps_with_bfloat16_track_eager_training():
+    events = _events(96, seed=7)
+    loader = EventLoader(events, collate_graph_batch, 16)
+
+    def run(**options):
+        trainer = _trainer(events, device="cuda", **options)
+        return [trainer.train_epoch(loader).loss for _ in range(3)]
+
+    eager = run()
+    fast = run(compile_model=True, cuda_graphs=True, precision="bfloat16")
+    assert fast == pytest.approx(eager, rel=5e-2)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA unavailable")
+def test_cuda_graph_steps_follow_the_learning_rate_schedule(tmp_path):
+    from gnn4colliders.training import CheckpointManager
+
+    events = _events(96, seed=8)
+    loader = EventLoader(events, collate_graph_batch, 16)
+    trainer = _trainer(events, device="cuda", cuda_graphs=True)
+    trainer.train_epoch(loader)  # warmup steps, then captured replays
+    # A zero rate must reach the already captured steps.
+    torch.optim.lr_scheduler.ExponentialLR(trainer.optimizer, gamma=0.0).step()
+    before = [p.detach().clone() for p in trainer.model.parameters()]
+    trainer.train_epoch(loader)
+    for old, new in zip(before, trainer.model.parameters()):
+        assert torch.equal(old, new)
+    # Checkpoints stay loadable by eager training on any device.
+    path = CheckpointManager(tmp_path).save(
+        model=trainer.model, trainer_state=trainer.state, optimizer=trainer.optimizer
+    )
+    groups = CheckpointManager.load(path)["optimizer_state_dict"]["param_groups"]
+    assert all(type(g["lr"]) is float and not g["capturable"] for g in groups)