From 01bf0b3f4a5b9ae39930a2bb2d83773a708155c8 Mon Sep 17 00:00:00 2001 From: iback Date: Wed, 14 Oct 2026 07:29:33 +0000 Subject: [PATCH] fix: pin the region-moment matmul to fp32, which autocast overflows `_region_moments` computes per-region voxel count, mean and variance as `region_mask @ design`, where the design matrix's first column is all ones -- so the first output column is a plain voxel count. `autocast` demotes that matmul to fp16, which maxes at 65504, and the count saturates to `inf` for any region above about 40^3 voxels. At the 128^3 training patch size the background region always does. The failure is silent and total. `counts=inf` with `sums=inf` gives `means=NaN`, every per-region Gaussian contribution in `RandomRedistributeSegGPU` becomes NaN, the `tmax - tmin > 1e-8` guard is False on NaN, and `torch.where` returns the input untouched. `retain_stats` then maps the untouched min-max normalisation back to the original mean and std, so the transform is an exact identity. The `torch.isfinite(x).all()` check cannot see it, because the NaN only ever lived in `to_add`, and `record_applications` still reports `changed=True` off the 2.4e-7 of fp32 round-trip noise. The matmul form arrived in 20360d8; before it the same quantity was `mask_flat.sum(dim=1)`, a reduction, which autocast keeps in fp32. So redistribution has been dead in every run trained since, measured on a real preprocessed Dataset804 patch: max abs change 2.4e-7 on data of range 1.7, for every seed and at `in_seg` 0.25 and 1.0. Pinning the matmul to fp32 restores it -- 100% of voxels moved by up to 1.83, bit-identical to the same call with autocast off. The counts are represented exactly (128^3 is far below 2^24) and TF32 keeps the matmul as fast as the fp16 one: 0.33 -> 0.40 ms per call, about 0.04 s on a 53.7 s epoch. The moments are returned in fp32 deliberately, since casting back to `design.dtype` would reintroduce the overflow. Retraining the paper's "Ours" configuration with this patch and nothing else changed recovers the out-of-domain transfer the four-sequence retrains had lost: out-of-domain mean Dice on Dataset804 0.639 -> 0.777 against 0.760 before the regression, SpineGAN Dixon water 0.301 -> 0.769, its spinal-canal class 0.110 -> 0.863, with in-domain unchanged to within 0.006. The new tests pin the count against overflow and the transform against being an identity. They fail before this change and pass after, and they need no GPU because fp16 CPU autocast demotes the same operation. They do need a 64^3 patch: the identity requires two overflows, since with only the count at `inf` the mean comes out 0 rather than NaN, which is still wrong but still moves the data. Co-Authored-By: Claude Opus 5 --- smauglab/transforms/gpu/fromSeg.py | 14 +++- unit_tests/test_redistribute_autocast.py | 102 +++++++++++++++++++++++ 2 files changed, 115 insertions(+), 1 deletion(-) create mode 100644 unit_tests/test_redistribute_autocast.py diff --git a/smauglab/transforms/gpu/fromSeg.py b/smauglab/transforms/gpu/fromSeg.py index 023f8a5..a53155a 100644 --- a/smauglab/transforms/gpu/fromSeg.py +++ b/smauglab/transforms/gpu/fromSeg.py @@ -216,8 +216,20 @@ def _region_moments(region_mask: torch.Tensor, design: torch.Tensor) -> tuple[to `E[x^2] - E[x]^2`, which is the same quantity the explicit `((x - mean) * mask) ** 2` pass computed -- the mask is binary, so squaring it changes nothing -- without materialising the centred array. + + The matmul is pinned to fp32, and must stay that way. `autocast` demotes matmuls + to fp16, which maxes at 65504, so the `1` column -- a plain voxel count -- saturates + to `inf` for any region above about 40^3 voxels. At the 128^3 training patch size + the background region always does, and then `counts=inf` gives `means=NaN`, every + per-region Gaussian contribution is NaN, the `tmax - tmin > 1e-8` guard in + `RandomRedistributeSegGPU` is False on NaN, and the transform returns its input + untouched -- an exact identity that no NaN check sees, because the NaN never + reaches the output. fp32 represents these counts exactly (128^3 is far below 2^24) + and TF32 keeps the matmul as fast as the fp16 one. Returned in fp32 deliberately: + casting back to `design.dtype` would reintroduce the overflow under autocast. """ - sums = region_mask.to(design.dtype) @ design # (R, 3) + with torch.autocast(device_type=design.device.type, enabled=False): + sums = region_mask.to(torch.float32) @ design.to(torch.float32) # (R, 3) counts = sums[:, 0].clamp_min(1) means = sums[:, 1] / counts variances = (sums[:, 2] / counts - means * means).clamp_min(0) diff --git a/unit_tests/test_redistribute_autocast.py b/unit_tests/test_redistribute_autocast.py new file mode 100644 index 0000000..cef98bc --- /dev/null +++ b/unit_tests/test_redistribute_autocast.py @@ -0,0 +1,102 @@ +"""`RandomRedistributeSegGPU` must still redistribute under `autocast`. + +`_region_moments` computes each region's voxel count, mean and variance as one +`[R, S] x [S, 3]` matmul. `autocast` demotes matmuls to fp16, whose largest finite +value is 65504, so the `1` column -- a plain voxel count -- saturates to `inf` for +any region above about 40^3 voxels. At the 128^3 patch size nnU-Net trains on, the +background region always does. + +Nothing then raises. `counts=inf` makes `means=NaN`, every per-region Gaussian +contribution is NaN, `tmax - tmin > 1e-8` is False on NaN so the `torch.where` +returns the input, and `retain_stats` maps that back to the original mean and std -- +an exact identity. The transform's own `torch.isfinite(x).all()` guard passes, +because the NaN only ever lived in `to_add`. + +So the augmentation reported itself as applied and did nothing, for every real +training run between `20360d8` (which introduced the matmul; the reduction it +replaced was fp32 under autocast and did not overflow) and this test. It sits at +`p=0.4` in the paper's "Ours" config, the highest probability of any transfer +augmentation there. + +fp16 autocast on CPU demotes the same op, so these do not need a GPU. +""" + +from __future__ import annotations + +import torch + +from smauglab.transforms.gpu.fromSeg import RandomRedistributeSegGPU, _region_moments +from unit_tests.helpers import SmaugLabTestCase + +#: Big enough to reproduce the identity, small enough to stay a unit test. Both the +#: count column and the `sum(x)` column have to overflow: with only the count at `inf` +#: the mean comes out 0 rather than NaN, which is still wrong but still moves the patch. +#: `x` is the transform's [0, 1] working space, so `sum(x)` needs roughly 2 * 65504 +#: voxels. 48^3 is not enough; 64^3 is, and matches the real failure at 128^3. +SHAPE = (1, 1, 64, 64, 64) + + +def pair() -> tuple[torch.Tensor, torch.Tensor]: + """A z-scored patch and the single-channel integer label map the trainer passes.""" + seg = torch.zeros(*SHAPE) + seg[:, :, 16:32, 16:32, 16:32] = 1.0 + seg[:, :, 36:48, 36:48, 36:48] = 2.0 + return torch.randn(*SHAPE), seg + + +class TestRegionCountsSurviveAutocast(SmaugLabTestCase): + def test_the_count_does_not_overflow_fp16(self): + values = torch.rand(100_000) + masks = torch.ones(2, 100_000, dtype=torch.bool) + design = torch.stack([torch.ones_like(values), values, values * values], dim=1) + + with torch.autocast("cpu", torch.float16): + counts, means, stds = _region_moments(masks, design) + + self.assertTrue(torch.isfinite(counts).all(), f"counts overflowed: {counts}") + self.assertEqual(float(counts[0]), 100_000.0) + self.assertAlmostEqual(float(means[0]), float(values.mean()), places=4) + self.assertAlmostEqual(float(stds[0]), float(values.std(unbiased=False)), places=4) + + def test_it_agrees_with_the_same_call_outside_autocast(self): + values = torch.rand(100_000) + masks = torch.rand(3, 100_000) > 0.5 + design = torch.stack([torch.ones_like(values), values, values * values], dim=1) + + plain = _region_moments(masks, design) + with torch.autocast("cpu", torch.float16): + amp = _region_moments(masks, design) + + for got, want in zip(amp, plain): + self.assertTrue(torch.equal(got, want)) + + +class TestRedistributeIsNotAnIdentityUnderAutocast(SmaugLabTestCase): + def test_it_moves_the_patch(self): + for seed in range(5): + with self.subTest(seed=seed): + torch.manual_seed(seed) + img, seg = pair() + transform = RandomRedistributeSegGPU(in_seg=1.0, p=1.0, p_batch=1.0) + + with torch.autocast("cpu", torch.float16): + out = transform.apply_transform(img.clone(), {"seg": seg}, {}) + + moved = (out.float() - img).abs().max().item() + self.assertTrue(torch.isfinite(out).all()) + # The perturbation is up to 2.0 in the transform's [0, 1] working space, + # so anything near zero means the identity branch fired. + self.assertGreater(moved, 1e-2, "redistribute was an identity under autocast") + + def test_autocast_does_not_change_the_result(self): + """The moments are the only thing autocast could touch here, and they are pinned.""" + img, seg = pair() + transform = RandomRedistributeSegGPU(in_seg=1.0, p=1.0, p_batch=1.0) + + torch.manual_seed(0) + plain = transform.apply_transform(img.clone(), {"seg": seg}, {}) + torch.manual_seed(0) + with torch.autocast("cpu", torch.float16): + amp = transform.apply_transform(img.clone(), {"seg": seg}, {}) + + self.assertTrue(torch.allclose(plain, amp.float(), atol=1e-5))