From cc913df8623da60653b4cdefcf09076a6e54b21b Mon Sep 17 00:00:00 2001 From: Chouffe Date: Thu, 11 Jun 2026 18:33:58 +0200 Subject: [PATCH 1/3] test(train): add GPU bitwise-reproducibility test and document determinism Same-seed training is bitwise reproducible on GPU as well as CPU: Trainer(deterministic=True) enables strict use_deterministic_algorithms and sets CUBLAS_WORKSPACE_CONFIG. The GPU twin of the reproducibility test guards this (skipped where CUDA is unavailable, e.g. CI). README documents the guarantee and its scope: same seed + same device type + same torch/CUDA versions; CPU vs GPU (or different GPU models) diverge by floating-point rounding, which is inherent. Closes #36 --- train/README.md | 16 +++++++++++ train/tests/test_reproducibility.py | 41 ++++++++++++++++++++++------- 2 files changed, 48 insertions(+), 9 deletions(-) diff --git a/train/README.md b/train/README.md index dfe2dd0..6dc6853 100644 --- a/train/README.md +++ b/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 device type + same torch/CUDA versions produce a +bitwise-identical checkpoint, 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 the guarantee (the GPU variant +skips when CUDA is unavailable, e.g. in CI). diff --git a/train/tests/test_reproducibility.py b/train/tests/test_reproducibility.py index d7400f4..5747546 100644 --- a/train/tests/test_reproducibility.py +++ b/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 = "cpu" ) -> 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,9 +122,7 @@ 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, -) -> None: +def _assert_bitwise_reproducible(tmp_path: Path, accelerator: str) -> None: train_dir = _make_split( tmp_path, "train", @@ -130,10 +130,14 @@ 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") + 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 + ) run_other = _fit_once_transformer( - OTHER_SEED, train_dir, val_dir, tmp_path / "run_other" + OTHER_SEED, train_dir, val_dir, tmp_path / "run_other", accelerator ) assert run1.keys() == run2.keys() == run_other.keys() @@ -143,3 +147,22 @@ def test_transformer_training_is_bitwise_reproducible_with_fixed_seed( ) 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. + """ + _assert_bitwise_reproducible(tmp_path, accelerator="gpu") From c74b536eb8e519ee41cf6514b0162bb780dabec5 Mon Sep 17 00:00:00 2001 From: Chouffe Date: Thu, 11 Jun 2026 18:45:03 +0200 Subject: [PATCH 2/3] fix(train): tighten determinism docs and trim redundant GPU control fit MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review follow-ups: the README guarantee now says same GPU model (not device type, which two different GPUs share) and scopes the test to what it actually asserts — same-seed weight reproducibility; optimizer state and best-epoch selection were verified end-to-end, not by the test. The GPU test drops the different-seed negative control (seeds diverge init on CPU, so it proves nothing GPU-specific) and _fit_once_transformer's accelerator parameter is now required. --- train/README.md | 10 +++++----- train/tests/test_reproducibility.py | 29 ++++++++++++++++++++--------- 2 files changed, 25 insertions(+), 14 deletions(-) diff --git a/train/README.md b/train/README.md index 6dc6853..4ed5716 100644 --- a/train/README.md +++ b/train/README.md @@ -35,12 +35,12 @@ 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 device type + same torch/CUDA versions produce a -bitwise-identical checkpoint, including optimizer state and best-epoch -selection. +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 the guarantee (the GPU variant -skips when CUDA is unavailable, e.g. in CI). +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). diff --git a/train/tests/test_reproducibility.py b/train/tests/test_reproducibility.py index 5747546..dbe5bad 100644 --- a/train/tests/test_reproducibility.py +++ b/train/tests/test_reproducibility.py @@ -72,7 +72,7 @@ def _make_split( def _fit_once_transformer( - seed: int, train_dir: Path, val_dir: Path, log_dir: Path, accelerator: str = "cpu" + seed: int, train_dir: Path, val_dir: Path, log_dir: Path, accelerator: str ) -> dict: L.seed_everything(seed, workers=True) @@ -122,7 +122,9 @@ def _fit_once_transformer( return {k: v.detach().clone() for k, v in lit.state_dict().items()} -def _assert_bitwise_reproducible(tmp_path: Path, accelerator: str) -> None: +def _assert_bitwise_reproducible( + tmp_path: Path, accelerator: str, *, negative_control: bool = True +) -> None: train_dir = _make_split( tmp_path, "train", @@ -136,17 +138,22 @@ def _assert_bitwise_reproducible(tmp_path: Path, accelerator: str) -> None: run2 = _fit_once_transformer( SEED, train_dir, val_dir, tmp_path / "run2", accelerator ) - run_other = _fit_once_transformer( - OTHER_SEED, train_dir, val_dir, tmp_path / "run_other", 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( @@ -164,5 +171,9 @@ def test_transformer_training_is_bitwise_reproducible_on_gpu(tmp_path: Path) -> ``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") + _assert_bitwise_reproducible(tmp_path, accelerator="gpu", negative_control=False) From 3d43cdba5ca8e6c565f8bb0415e36968c99a6b9b Mon Sep 17 00:00:00 2001 From: Chouffe Date: Fri, 12 Jun 2026 08:41:50 +0200 Subject: [PATCH 3/3] style(train): apply ruff format to reproducibility test --- train/tests/test_reproducibility.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/train/tests/test_reproducibility.py b/train/tests/test_reproducibility.py index dbe5bad..9617c46 100644 --- a/train/tests/test_reproducibility.py +++ b/train/tests/test_reproducibility.py @@ -150,9 +150,7 @@ def _assert_bitwise_reproducible( 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]) - ] + differing = [key for key in run1 if not torch.equal(run1[key], run_other[key])] assert differing, "Different-seed run produced identical transformer weights"