diff --git a/ScaFFold/utils/losses.py b/ScaFFold/utils/losses.py index 0a1636d..5d929c2 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,59 @@ 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" plus an explicit .sum(), not reduction="sum": the + # fused CUDA reduction accumulates with atomicAdd, so its summation + # order follows whichever order the blocks retire in and the loss + # value varies run to run. It has no deterministic implementation, and + # `more_determinism` passes warn_only=True, so there it would only + # warn and stay nondeterministic. The separate .sum() is a + # fixed-order reduction and is bitwise reproducible. 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: that + # drains the pipeline mid-step, since the host stops submitting until + # the read returns and the queue runs dry behind it. + # torch.cuda.set_sync_debug_mode("error") catches a regression. 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, where new_tensor would stage that float + # through a pageable host-to-device copy. + 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. Not torch.bincount: it sizes its output + # from the largest label, so it reads that label back to the host + # even when minlength already fixes the width. + 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..57765c6 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,108 @@ 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(): + # Pins that the CE numerator is summed outside the loss kernel: + # reduction="sum" on CUDA accumulates with atomicAdd, so its 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 rather than + # allclose. + # + # The volume is load-bearing, not incidental: the atomics only collide + # once the shape puts many blocks in flight, so shrinking this test for + # speed would quietly turn it into one that passes either way. 128**3 with + # 7 classes is scale 7 with the shipped n_categories. + 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(): + # Pins that the CE term issues no host-device sync. It runs once per + # training step, so a 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. A sync is invisible in the loss value, so + # only a test catches one. set_sync_debug_mode is documented as a + # prototype that does not catch every synchronizing op, so this is a + # floor, not a proof. + # + # The shape matters only in that it covers both normalizer branches and + # both entry points; the debug mode does the checking, so it stays small. + 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: keep the results on device and inspect them + # after the debug mode is back off. + outs = [] + torch.cuda.synchronize() + torch.cuda.set_sync_debug_mode("error") + try: + 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(): + # Pins that the CE path is legal under strict determinism. + # more_determinism sets use_deterministic_algorithms(warn_only=True), so a + # nondeterministic kernel only warns there and the run stays + # irreproducible -- a regression in this path would be invisible under the + # config that is supposed to catch it. Strict mode raises instead, since + # the fused loss reduction has no deterministic implementation. The shape + # only has to be a valid volume, so it stays small. + 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