Make loss calculation deterministic and sync-free

ScaFFold/unet/conv3d.py

@@ -1205,10 +1205,11 @@ def backward(ctx, grad_output):
             return grad_x, grad_w, grad_b, None, None, None
 
         # d(bias) is the sum of grad_output over every axis but the channel one,
-        # whatever the forward kernel was.  Only the segmentation head has a
-        # bias and it is on the policy block-list, so this is unreachable in
-        # ScaFFold today -- but ``is_supported`` accepts a bias, and the class is
-        # a public drop-in, so it has to be right.
+        # whatever the forward kernel was.  The segmentation head is the only
+        # biased convolution in ScaFFold, and emptying the block-list on
+        # 2026-08-04 (:func:`_policy_declines`) put it on this rung -- so this
+        # line runs once a step on a full-volume ``grad_output``, where it used
+        # to be unreachable outside the tests.
         grad_b = grad_output.sum(dim=(0, 2, 3, 4)) if (has_bias and needs_b) else None
         return grad_x, grad_w, grad_b, None, None, None
 
@@ -1632,9 +1633,11 @@ class _TritonConvTranspose3dFn(torch.autograd.Function):
 
     * **``grad_bias`` is on the path.**  All four transposed sites have one
       (``nn.ConvTranspose3d`` defaults to ``bias=True`` and ``unet_parts`` does
-      not turn it off), where the only biased ordinary convolution is the
-      segmentation head, which the block-list keeps on MIOpen.  So the reduction
-      below is reached on every step of every run, not only by a test.
+      not turn it off), so the reduction below is reached on every step of every
+      run, not only by a test.  This used to be a *difference* between the two
+      ladders, because the ordinary operator's only biased site -- the
+      segmentation head -- was blocked; emptying the block-list on 2026-08-04
+      put that on the rung too, so both are now hot.
     * **The rung-flip hazard is wider.**  ``DCTensor`` mirrors its local shard's
       size, strides, dtype and device exactly, and this ladder never adds a halo
       -- so the tensor this node saves and the one ``nn.ConvTranspose3d.forward``

ScaFFold/utils/losses.py

@@ -128,7 +128,7 @@ def compute_sharded_cross_entropy_loss(
 
     Each rank only sees a local spatial shard, so we cannot use the local
     `reduction="mean"` result directly. Instead we:
-    1. compute the local CE numerator with `reduction="sum"`,
+    1. compute the local CE numerator by summing the per-voxel losses,
     2. build the correct global denominator,
     3. all-reduce numerator and denominator together across the spatial mesh
        in a single collective, and
@@ -146,40 +146,77 @@ def compute_sharded_cross_entropy_loss(
 
     autocast_device = device_type if device_type != "mps" else "cpu"
     with torch.autocast(autocast_device, enabled=False):
-        # Accumulate CE in full precision. Using reduction="sum" gives us the
-        # numerator of the final global mean; if class weights are present,
-        # PyTorch applies the target-class weight to each voxel here. When the
-        # caller already computed log-softmax, NLL over it is identical to
-        # cross-entropy over the raw logits but avoids a second full upcast.
+        # Accumulate CE in full precision. Summing the per-voxel losses gives
+        # us the numerator of the final global mean; if class weights are
+        # present, PyTorch applies the target-class weight to each voxel here.
+        # When the caller already computed log-softmax, NLL over it is
+        # identical to cross-entropy over the raw logits but avoids a second
+        # full upcast.
+        #
+        # reduction="none" followed by .sum() rather than reduction="sum",
+        # which is not the same computation on CUDA: the fused reduction
+        # accumulates with atomicAdd, so the summation order is whatever order
+        # the blocks happen to retire in and the loss value changes run to run
+        # (measured at 128**3: 3-4 distinct values per 100 calls, rel 2.2e-7 --
+        # enough to perturb train_stats.csv and to flip a best-checkpoint tie).
+        # It has no deterministic implementation at all: under
+        # torch.use_deterministic_algorithms(True) the fused form raises, and
+        # `more_determinism` passes warn_only=True, so it warns and stays
+        # nondeterministic. The separate .sum() is an ordinary tree reduction
+        # over a fixed shape and is bitwise reproducible.
+        #
+        # It is also not a cost at the scale that matters. Measured here,
+        # fwd+bwd, 7 classes, paired alternating arms:
+        #     128**3 (scale 7)  241 -> 147 us   0.61x -- the split form wins,
+        #                                       the atomics were serializing
+        #     256**3 (scale 8)  654 -> 715 us   1.09x, +61 us
+        # against 74 ms and 458 ms steps respectively, so under 0.15% either
+        # way. Peak memory is unchanged at both: the per-voxel fp32 tensor
+        # (8/64 MiB) is freed by the time the backward allocates a gradient
+        # the size of log_probs (56/448 MiB), which sets the peak.
         if log_probs is not None:
-            local_ce_sum = F.nll_loss(
+            local_ce = F.nll_loss(
                 log_probs,
                 local_labels,
                 weight=class_weights,
-                reduction="sum",
+                reduction="none",
             )
         else:
-            local_ce_sum = F.cross_entropy(
+            local_ce = F.cross_entropy(
                 local_preds.float(),
                 local_labels,
                 weight=class_weights,
-                reduction="sum",
+                reduction="none",
             )
-
+        local_ce_sum = local_ce.sum()
+
+        # Neither branch below may read device memory from the host. Both used
+        # to: torch.bincount sizes its output from the largest label, so it
+        # copies that value back even when minlength already fixes the width,
+        # and new_tensor() stages a Python float through a pageable H2D copy.
+        # Either one is a full pipeline drain in the middle of the step -- the
+        # host blocks, stops submitting, and the queue runs dry behind it.
+        # Measured with torch.cuda.set_sync_debug_mode("error"), which flags
+        # both and neither replacement.
         if class_weights is None:
             # Sum the actual local voxel counts across spatial shards. We use
             # an all-reduced count instead of numel()*num_shards because shard
             # sizes can differ at chunk boundaries.
-            local_normalizer = local_ce_sum.new_tensor(float(local_labels.numel()))
+            #
+            # new_full rather than new_tensor: numel() is shape metadata the
+            # host already has, and full() fills on the device with the value
+            # as a kernel argument instead of copying it across.
+            local_normalizer = local_ce_sum.new_full((), float(local_labels.numel()))
         else:
             # Weighted CE divides by sum(weight[target_i]) over all voxels.
-            # Build that denominator from the local label histogram.
-            local_class_counts = torch.bincount(
-                local_labels.reshape(-1), minlength=class_weights.numel()
-            ).to(dtype=local_ce_sum.dtype)
-            local_normalizer = torch.dot(
-                local_class_counts, class_weights.to(dtype=local_ce_sum.dtype)
-            )
+            # Gather the per-voxel weights and sum them, which is that
+            # definition transcribed -- the histogram this replaced was an
+            # indirect route to the same number, and the more expensive one:
+            # bincount's own kernel plus the dot cost 83.6 us at 128**3 and
+            # 373.6 us at 256**3 against 26.1 and 90.8 us for the gather.
+            local_normalizer = class_weights.to(dtype=local_ce_sum.dtype)[
+                local_labels
+            ].sum()
 
     # Reduce the CE numerator and its denominator across the spatial shards in
     # one collective (they share the same mesh) rather than two, halving the

tests/test_perf_hotpath.py

@@ -20,6 +20,7 @@
 numerically equivalent to the straightforward reference it replaced.
 """
 
+import pytest
 import torch
 import torch.nn.functional as F
 
@@ -67,6 +68,117 @@ def test_ce_log_probs_path_matches_cross_entropy():
         assert torch.allclose(plain, ref, atol=1e-6)
 
 
+@pytest.mark.gpu
+def test_gpu_ce_numerator_is_bitwise_reproducible():
+    # The CE numerator must be summed outside the loss kernel. reduction="sum"
+    # on CUDA accumulates with atomicAdd, so the value depends on block retire
+    # order and changes between otherwise identical calls. Both entry points
+    # (precomputed log_probs via NLL, and raw logits via CE) go through the
+    # same reduction, so both are checked.
+    #
+    # Bitwise, not allclose: the whole point is that repeated calls agree to
+    # the last bit. The pre-fix code yields 4 distinct values here.
+    #
+    # 128**3 with 7 classes is scale 7 with the shipped n_categories, and the
+    # size is load-bearing, not incidental: sweeping the pre-fix code, a shape
+    # has to be big enough to put many blocks in flight before the atomics
+    # collide at all. 48**3 and 64**3 give 1 distinct value, 96**3 gives 2-3,
+    # 128**3 gives 4. Shrinking this test for speed would quietly turn it into
+    # a test that passes either way.
+    torch.manual_seed(3)
+    device = torch.device("cuda")
+    b, c, n = 1, 7, 128
+    preds = torch.randn(b, c, n, n, n, device=device)
+    labels = torch.randint(0, c, (b, n, n, n), device=device)
+    weights = torch.rand(c, device=device) + 0.5
+    log_probs = F.log_softmax(preds.float(), dim=1)
+
+    for w in (None, weights):
+        for kwargs in ({"log_probs": log_probs}, {}):
+            values = {
+                compute_sharded_cross_entropy_loss(
+                    preds, labels, None, (1,), "cuda", w, **kwargs
+                ).item()
+                for _ in range(50)
+            }
+            assert len(values) == 1, f"{len(values)} distinct CE values: {values}"
+
+
+@pytest.mark.gpu
+def test_gpu_ce_does_not_synchronize():
+    # The CE term runs once per training step, so a host-device sync in it
+    # drains the whole launch queue mid-step: the host stops submitting until
+    # the read comes back, and the GPU runs out of queued work behind it.
+    #
+    # Both branches have had one. torch.bincount sizes its output from the
+    # largest label and copies that value back even when minlength already
+    # fixes the width; new_tensor() stages a Python float through a pageable
+    # H2D copy. Both are invisible in the result and in the loss value, which
+    # is why they need a test rather than a review.
+    #
+    # set_sync_debug_mode is documented as a prototype that does not catch
+    # every synchronizing op, so this is a floor, not a proof. It does catch
+    # both of the above -- verified by running it against the previous code.
+    torch.manual_seed(5)
+    device = torch.device("cuda")
+    b, c, n = 1, 7, 64
+    preds = torch.randn(b, c, n, n, n, device=device)
+    labels = torch.randint(0, c, (b, n, n, n), device=device)
+    weights = torch.rand(c, device=device) + 0.5
+    log_probs = F.log_softmax(preds.float(), dim=1)
+
+    # .item() is itself a sync, so results are kept on device and only
+    # inspected after the debug mode is back off.
+    outs = []
+    torch.cuda.synchronize()
+    torch.cuda.set_sync_debug_mode("error")
+    try:
+        # Both weighting branches (the normalizer is built differently in
+        # each) and both entry points.
+        for w in (weights, None):
+            for kwargs in ({"log_probs": log_probs}, {}):
+                outs.append(
+                    compute_sharded_cross_entropy_loss(
+                        preds, labels, None, (1,), "cuda", w, **kwargs
+                    )
+                )
+    finally:
+        torch.cuda.set_sync_debug_mode("default")
+
+    torch.cuda.synchronize()
+    assert all(torch.isfinite(o) for o in outs)
+
+
+@pytest.mark.gpu
+def test_gpu_ce_survives_strict_deterministic_algorithms():
+    # more_determinism sets use_deterministic_algorithms(warn_only=True), so a
+    # nondeterministic kernel only warns there and the run stays irreproducible
+    # -- a determinism regression in this path would be invisible under the
+    # config that is supposed to catch it. Assert against strict mode instead,
+    # where the fused reduction raises:
+    #   "nll_loss2d_forward_out_cuda_template does not have a deterministic
+    #    implementation"
+    torch.manual_seed(4)
+    device = torch.device("cuda")
+    b, c = 1, 5
+    preds = torch.randn(b, c, 16, 16, 16, device=device)
+    labels = torch.randint(0, c, (b, 16, 16, 16), device=device)
+    weights = torch.rand(c, device=device) + 0.5
+    log_probs = F.log_softmax(preds.float(), dim=1)
+
+    was_deterministic = torch.are_deterministic_algorithms_enabled()
+    was_warn_only = torch.is_deterministic_algorithms_warn_only_enabled()
+    torch.use_deterministic_algorithms(True)
+    try:
+        for w in (None, weights):
+            for kwargs in ({"log_probs": log_probs}, {}):
+                compute_sharded_cross_entropy_loss(
+                    preds, labels, None, (1,), "cuda", w, **kwargs
+                )
+    finally:
+        torch.use_deterministic_algorithms(was_deterministic, warn_only=was_warn_only)
+
+
 def test_ce_uses_single_spatial_collective(monkeypatch):
     # The CE numerator and its normalizer are reduced together in one
     # SpatialAllReduce, not two. Count applications; the packed path issues