Preserve GroupNorm dtype - #105
Open
ndryden wants to merge 2 commits into
Open
Conversation
`at::group_norm` carries autocast's fp32 cast policy, so under `autocast(bf16)` every GroupNorm in the UNet returned fp32 and every consumer of it -- the next convolution (`lower_precision_fp` policy), the skip concatenation, the max-pool feeding the next block -- immediately cast it back to bf16. That round trip is a full-volume read and two full-volume writes per site doing no arithmetic; profiling put it at 24.31 ms/step at scale 8 (86% of peak HBM for the traffic it moves) and 3.59 at scale 7. `FastGroupNorm` now emits its input's dtype. The statistics are still accumulated in fp32 on every rung and the normalized value is still computed in fp32 -- only the store narrows, so the output is the fp32 answer rounded once rather than a narrower computation. Priced by monkeypatch before landing: -36.39 +/- 1.55 ms of a 449.56 ms step at scale 8 and -3.93 +/- 1.19 of 66.49 at scale 7, with peak memory 42.78 -> 38.49 GiB and 6.10 -> 5.57. The standalone `triton_group_norm()` is unchanged: its documented guarantee is that the output dtype is exactly `F.group_norm`'s, tests pin it, and it still holds. The departure is an opt-in the kernel module now offers -- a keyword-only `out_dtype=`, defaulting to a `MATCH_STOCK_DTYPE` sentinel -- which `FastGroupNorm._triton_forward` is the one caller to use. It is honoured on the `F.group_norm` fallback route as well as the kernel's, so a result's dtype never depends on which one served it. All three rungs narrow, not just the Triton one. A rung is chosen per module and per process; a latch demotes an unproven module mid-run and a proven module still falls back when its kernel raises outside a backward replay. A dtype that followed the rung would be an activation width that changes mid-run, differs between DDP ranks with different latch histories, and breaks `torch.utils.checkpoint`, which compares the dtype of every recomputed saved tensor. The compiled and eager rungs therefore cast after `F.group_norm` -- one cast, which is exactly the cast their consumer was about to do. Not bitwise free, and the class docstring says so: the forward is the same fp32 value rounded once, but the backward's accumulation moves and the loss diverges ~4.4e-06 relative by step 23 at scale 8. Each configuration stays bitwise reproducible with itself. 24 tests: the dtype under autocast on all three rungs and both autocast dtypes, the no-autocast and `torch_amp: 0` paths keeping stock behaviour, a whole-UNet census of every site, a simulated mid-run rung fallback answering in the same dtype, and the kernel module's new keyword on both routes including its rejections.
Two corrections to the docstring, both from measuring the landed code rather than the monkeypatch that priced it. The numbers. Paired, arms alternating within each of 6 replicates, on one MI300A: -35.73 +/- 0.93 ms of a 440.28 ms step at scale 8 and -4.03 +/- 0.57 of 66.00 at scale 7, with peak memory 42.776 -> 38.486 and 6.102 -> 5.572 GiB. Both reproduce the prediction (-36.39 +/- 1.55 and -3.93 +/- 1.19) and the memory figures land on it to three digits. The device-time breakdown of the 35.4 ms is now itemized instead of summarized, and the 2.8 ms that does not belong to any row this change can reach is named as unclaimed. The reason it is not bitwise free. The docstring said "it changes the backward's reduction order", inheriting a guess that the kernel's tiling plan depends on the output dtype. That guess is wrong: `_plan` takes no dtype, both directions build it from the shape alone, and running the backward op on the same values as fp32 and as bf16 gives bitwise identical d_input/d_weight/d_bias. The kernel is not involved. What moves is downstream of it -- gradient accumulation at the encoder's fan-out sites, which used to sum two cotangents in fp32 and now sums them in bf16, and `max_pool3d`'s argmax tie-breaking, which changes when rounding makes a window's maximum non-unique. The second one also falsifies the "max_pool commutes with rounding, so it is free" argument that this docstring was repeating: the forward commutes, the backward scatters through indices and does not. Both mechanisms are demonstrated in isolation, with everything else held fixed, and located in the model: the first site the backward reaches whose gradient differs is exactly the deepest encoder output that has two consumers, and every purely serial site before it is bitwise identical. The forward is exact at every site and the model output is bitwise unchanged. Each configuration is bitwise reproducible with itself across 6 independent processes. Measurement and probes: work/profile/PROFILE_GN_DTYPE.md.
michaelmckinsey1
approved these changes
Aug 13, 2026
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
2 participants
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.
This should reduce some memory and conversion/copy overheads.
Code by Claude.