Make loss calculation deterministic and sync-free - #103
Open
ndryden wants to merge 3 commits into
Open
Conversation
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.
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.
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
reduction=sumis nondeterministic. Switching to no reduction and following it with a.sum()is deterministic and costs little.bincountrequired a host/device synchronization. This has been replaced.Code by Claude.