test(train): add GPU bitwise-reproducibility test and document determinism

train/README.md

@@ -28,3 +28,19 @@ make install
 uv run dvc repro            # full pipeline (uses GPU for training when available)
 uv run dvc repro train      # just the training stage (data-prep cached)
 ```
+
+## Determinism
+
+Training is fully deterministic: `train.py` seeds Python/NumPy/torch and the
+DataLoader workers (`L.seed_everything(seed, workers=True)`) and runs with
+`Trainer(deterministic=True)`, which also enables strict
+`torch.use_deterministic_algorithms` and sets `CUBLAS_WORKSPACE_CONFIG` on
+CUDA. Same seed + same hardware (same CPU or same GPU model) + same
+torch/CUDA versions produce a bitwise-identical checkpoint (verified
+end-to-end on real data, including optimizer state and best-epoch selection).
+
+Scope: a CPU run and a GPU run with the same seed do **not** match each
+other, and neither do different GPU models — different kernels round
+floating-point sums differently. That is inherent to floating point, not a
+bug. `tests/test_reproducibility.py` guards same-seed weight reproducibility
+on CPU and GPU (the GPU variant skips when CUDA is unavailable, e.g. in CI).

train/tests/test_reproducibility.py

@@ -3,7 +3,8 @@
 Runs two short Lightning fits with the same seed on a tiny fake dataset and
 asserts that every weight in the final ``state_dict`` is bitwise identical.
 A third run with a different seed acts as a negative control so the test
-cannot silently pass if nothing is actually random.
+cannot silently pass if nothing is actually random. Runs on CPU always and
+on GPU when CUDA is available (skipped otherwise, e.g. in CI).
 
 Exercises the full seeding path used by ``scripts/train.py``:
 ``L.seed_everything(seed, workers=True)`` + ``Trainer(deterministic=True)``
@@ -18,6 +19,7 @@
 
 import lightning as L
 import numpy as np
+import pytest
 import torch
 from PIL import Image
 from torch.utils.data import DataLoader
@@ -70,7 +72,7 @@ def _make_split(
 
 
 def _fit_once_transformer(
-    seed: int, train_dir: Path, val_dir: Path, log_dir: Path
+    seed: int, train_dir: Path, val_dir: Path, log_dir: Path, accelerator: str
 ) -> dict:
     L.seed_everything(seed, workers=True)
 
@@ -106,7 +108,7 @@ def _fit_once_transformer(
 
     trainer = L.Trainer(
         max_epochs=2,
-        accelerator="cpu",
+        accelerator=accelerator,
         devices=1,
         deterministic=True,
         logger=False,
@@ -120,8 +122,8 @@ def _fit_once_transformer(
     return {k: v.detach().clone() for k, v in lit.state_dict().items()}
 
 
-def test_transformer_training_is_bitwise_reproducible_with_fixed_seed(
-    tmp_path: Path,
+def _assert_bitwise_reproducible(
+    tmp_path: Path, accelerator: str, *, negative_control: bool = True
 ) -> None:
     train_dir = _make_split(
         tmp_path,
@@ -130,16 +132,46 @@ def test_transformer_training_is_bitwise_reproducible_with_fixed_seed(
     )
     val_dir = _make_split(tmp_path, "val", [("e", 1, 4), ("f", 0, 3)])
 
-    run1 = _fit_once_transformer(SEED, train_dir, val_dir, tmp_path / "run1")
-    run2 = _fit_once_transformer(SEED, train_dir, val_dir, tmp_path / "run2")
-    run_other = _fit_once_transformer(
-        OTHER_SEED, train_dir, val_dir, tmp_path / "run_other"
+    run1 = _fit_once_transformer(
+        SEED, train_dir, val_dir, tmp_path / "run1", accelerator
+    )
+    run2 = _fit_once_transformer(
+        SEED, train_dir, val_dir, tmp_path / "run2", accelerator
     )
 
-    assert run1.keys() == run2.keys() == run_other.keys()
+    assert run1.keys() == run2.keys()
     for key in run1:
         assert torch.equal(run1[key], run2[key]), (
             f"Same-seed transformer runs diverged at {key!r}"
         )
-    differing = [key for key in run1 if not torch.equal(run1[key], run_other[key])]
-    assert differing, "Different-seed run produced identical transformer weights"
+
+    if negative_control:
+        run_other = _fit_once_transformer(
+            OTHER_SEED, train_dir, val_dir, tmp_path / "run_other", accelerator
+        )
+        assert run_other.keys() == run1.keys()
+        differing = [key for key in run1 if not torch.equal(run1[key], run_other[key])]
+        assert differing, "Different-seed run produced identical transformer weights"
+
+
+def test_transformer_training_is_bitwise_reproducible_with_fixed_seed(
+    tmp_path: Path,
+) -> None:
+    _assert_bitwise_reproducible(tmp_path, accelerator="cpu")
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a CUDA GPU")
+def test_transformer_training_is_bitwise_reproducible_on_gpu(tmp_path: Path) -> None:
+    """GPU twin of the CPU test.
+
+    ``Trainer(deterministic=True)`` enables strict
+    ``torch.use_deterministic_algorithms`` and sets
+    ``CUBLAS_WORKSPACE_CONFIG``, so CUDA kernels must be deterministic too.
+    Guards against changes (e.g. mixed precision, attention backends) that
+    would silently break GPU run-to-run reproducibility.
+
+    Skips the different-seed negative control: seeds diverge the model at
+    init on the CPU before the GPU is involved, so the control proves
+    nothing GPU-specific — and the CPU test always runs it.
+    """
+    _assert_bitwise_reproducible(tmp_path, accelerator="gpu", negative_control=False)