From b4027d74ea42fe749a09e9c08327deac36bf7042 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Thu, 6 Aug 2026 16:09:06 -0700 Subject: [PATCH 1/3] Sum the cross-entropy numerator outside the loss kernel reduction="sum" is not just a convenience on CUDA: nll_loss and cross_entropy fuse the reduction into the loss kernel and accumulate with atomicAdd, so the summation order is whatever order the blocks retire in and the value changes between identical calls. Compute per-voxel losses with reduction="none" and .sum() them, which is an ordinary tree reduction over a fixed shape. This was the last thing in the model that was not bitwise reproducible. The Triton kernels closed everything else (0 of 64 gradient tensors vary against 15 with them off), which left the loss *value* alone drifting -- 3-4 distinct values per 100 calls at 128**3, rel 2.2e-7. That is small, but it is not confined to a log line: it lands in train_stats.csv, and is_best is keyed on val_loss_avg, so it can decide a near-tie between two epochs. more_determinism did not catch it and could not: worker.py passes warn_only=True, so torch printed "nll_loss2d_forward_out_cuda_template does not have a deterministic implementation" and carried on. The config that exists to make runs reproducible was reporting the defect rather than fixing it. Measured, this tree, 4 runs at scale 7 / 128**3, plus 3 more with more_determinism=1: parameters, per-batch losses, per-batch dice, first-batch forward activations, CE class weights and every train_stats.csv column (excluding wall-clock) are now identical bit for bit, and the default config and more_determinism=1 produce the *same bits* -- more_determinism is now redundant for reproducibility rather than partial. Same result at 2 ranks and at 4 ranks with spatial sharding, DDP bucket fingerprints identical on both sides of the all-reduce. The op probe at 128**3 goes from 1 of 66 tensors varying to 0 of 66, while an in-process control calling the old F.nll_loss(reduction="sum") still gives 3-4 distinct values -- so the node is colliding today and the result is not a vacuous pass. Not a cost at the scale that matters. fwd+bwd, 7 classes, paired alternating arms: 241 -> 147 us at 128**3 (0.61x -- the split form wins outright, the atomics were serializing) and 654 -> 715 us at 256**3 (1.09x, +61 us), against 74 ms and 458 ms steps. Peak memory is unchanged at both sizes: the per-voxel fp32 tensor is freed before the backward allocates a gradient the size of log_probs, which is 7x larger and sets the peak. The two tests are pinned to 128**3 deliberately. Sweeping the pre-fix code, the atomics do not collide at all below ~96**3 -- 48**3 and 64**3 give one distinct value no matter which reduction is used -- so a smaller, faster test would pass with the bug present. The strict-mode test asserts against use_deterministic_algorithms(True) rather than more_determinism's warn_only form, for the reason above: under warn_only a regression here is a log line, not a failure. --- ScaFFold/utils/losses.py | 44 +++++++++++++++++++------ tests/test_perf_hotpath.py | 67 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 101 insertions(+), 10 deletions(-) diff --git a/ScaFFold/utils/losses.py b/ScaFFold/utils/losses.py index 0a1636d..cfabb00 100644 --- a/ScaFFold/utils/losses.py +++ b/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,25 +146,49 @@ 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() if class_weights is None: # Sum the actual local voxel counts across spatial shards. We use diff --git a/tests/test_perf_hotpath.py b/tests/test_perf_hotpath.py index eac6a5f..5470e57 100644 --- a/tests/test_perf_hotpath.py +++ b/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,72 @@ 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_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 From 61b3918086b52c5d521bd447f6155bc0c0a74a28 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Thu, 6 Aug 2026 16:48:19 -0700 Subject: [PATCH 2/3] Build the CE normalizer without reading device memory The loss head did a blocking device-to-host read every training step. torch.bincount sizes its output from the largest label, so it copies that value back even though minlength already fixes the width; the unweighted branch's new_tensor() staged a Python float through a pageable H2D copy. Both are invisible in the result, which is why neither was noticed. The cost is not the copy, it is the drain. The host blocks until the read returns, which means it stops submitting, and the queue empties out behind it -- a stall in the middle of a step whose kernels are otherwise long enough to keep the device saturated. Replace the histogram with the definition it was computing: the weighted denominator is sum(weight[target_i]) over voxels, so gather the per-voxel weights and sum them. That is one pass instead of a bincount plus a dot, and it is also cheaper in its own right: 83.6 -> 26.1 us at 128**3 and 373.6 -> 90.8 us at 256**3. The unweighted branch uses new_full, which fills on the device with the value as a kernel argument. Measured with minibatch_bench, arms alternating within each rep, 22 steady steps a run, 3 reps, 6 classes: config A (scale 7, 1 GPU) 73.27 -> 66.67 ms/step -6.60 ms 9.0% config B (scale 8, 1 GPU) 456.72 -> 450.06 ms/step -6.66 ms 1.5% Per-rep ranges are disjoint at both (A: 73.13-74.82 vs 66.09-67.55; B: 455.07-456.75 vs 449.68-450.15). The saving is the same *absolute* size at both scales while bincount's own kernel cost differs between them by 290 us, which is the evidence that what was removed is a fixed-latency pipeline drain and not a cheaper kernel. Under torch.cuda.set_sync_debug_mode("error") the whole compute step -- forward, CE, dice, backward, Adam -- now completes with no synchronizing op; before, the CE call alone raised. This changes the numbers. Summing gathered weights is a different summation order from counts-dot-weights, so the normalizer moves by up to ~1.3e-7 relative (7 of 20 seeds at 128**3). End to end that is one fp32 ULP at step 0 (6.1e-8) and then training amplifies it: by step 8, 1.2e-5. Runs before and after this commit are therefore not bitwise comparable to each other. Each is still bitwise reproducible with itself, which is the property that matters and which was re-verified after this change -- 4 default runs and 3 with more_determinism=1 at scale 7, all identical including every train_stats.csv column, and default still produces the same bits as more_determinism=1. The test asserts on set_sync_debug_mode, which the docs call a prototype that does not catch every synchronizing op. It is a floor rather than a proof, but it does catch both of the calls removed here -- checked by running it against the previous code. --- ScaFFold/utils/losses.py | 29 +++++++++++++++++------- tests/test_perf_hotpath.py | 45 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 66 insertions(+), 8 deletions(-) diff --git a/ScaFFold/utils/losses.py b/ScaFFold/utils/losses.py index cfabb00..eaa75fd 100644 --- a/ScaFFold/utils/losses.py +++ b/ScaFFold/utils/losses.py @@ -190,20 +190,33 @@ def compute_sharded_cross_entropy_loss( ) 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 diff --git a/tests/test_perf_hotpath.py b/tests/test_perf_hotpath.py index 5470e57..cc6f71d 100644 --- a/tests/test_perf_hotpath.py +++ b/tests/test_perf_hotpath.py @@ -104,6 +104,51 @@ def test_gpu_ce_numerator_is_bitwise_reproducible(): 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 From 785dcf4c79ac3f5602d42586f4c8dce390aaa79b Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Thu, 6 Aug 2026 17:11:55 -0700 Subject: [PATCH 3/3] Correct two comments the block-list emptying invalidated Both said the segmentation head is kept on MIOpen. _policy_declines has returned False unconditionally since 2026-08-04, so the head runs on the Triton rung like every other convolution -- the rung census confirms it (FastConv3d.triton 228 over a 2-epoch scale-7 run: 19 convolutions a forward, including outc, times 12 forwards). The consequence is not cosmetic in the first case. _TritonConv3dFn.backward's grad_bias line was documented as unreachable in ScaFFold, and outc is the only biased convolution in the model, so it is now on the hot path -- a full-volume reduction once a step that nothing was accounting for. The second is in _TritonConvTranspose3dFn's docstring, where the claim was the contrast that made grad_bias a difference between the two ladders; there is no difference now. --- ScaFFold/unet/conv3d.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/ScaFFold/unet/conv3d.py b/ScaFFold/unet/conv3d.py index c5d74b6..e21352a 100644 --- a/ScaFFold/unet/conv3d.py +++ b/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``