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`` diff --git a/ScaFFold/utils/losses.py b/ScaFFold/utils/losses.py index 0a1636d..eaa75fd 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,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 diff --git a/tests/test_perf_hotpath.py b/tests/test_perf_hotpath.py index eac6a5f..cc6f71d 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,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