From 99323e78dfde184d90773f28e99da1e4436a54cd Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 7 Aug 2026 18:07:57 -0700 Subject: [PATCH 1/2] Let FastGroupNorm emit its input's dtype `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. --- ScaFFold/unet/group_norm.py | 150 ++++++++++++++++++-- ScaFFold/unet/triton_group_norm.py | 93 +++++++++++- tests/test_groupnorm.py | 210 ++++++++++++++++++++++++++++ tests/test_groupnorm_wiring_edge.py | 10 +- tests/test_triton_group_norm.py | 79 +++++++++++ 5 files changed, 518 insertions(+), 24 deletions(-) diff --git a/ScaFFold/unet/group_norm.py b/ScaFFold/unet/group_norm.py index 6e0729f..2975f5d 100644 --- a/ScaFFold/unet/group_norm.py +++ b/ScaFFold/unet/group_norm.py @@ -39,10 +39,14 @@ same shapes, same numerics -- only the kernel differs, so checkpoints are interchangeable in both directions with any other GroupNorm-based build. The one addition is the optional fused ``activation`` (see below), which adds no -state either. +state either. There is one deliberate departure from ``nn.GroupNorm``, in the +output *dtype* under autocast, and it has its own section in +``FastGroupNorm``'s docstring; read it before consuming a GroupNorm output as +fp32. -All three return the input's memory format, so the rungs are interchangeable in -everything a caller can observe (see :func:`_match_memory_format`). +All three rungs return the input's memory format *and* the input's dtype, so +they are interchangeable in everything a caller can observe (see +:func:`_match_memory_format` and :func:`_match_input_dtype`). Routing rejections, in the order they are tested: @@ -581,6 +585,32 @@ def _match_memory_format(out, reference): return out.contiguous(memory_format=torch.channels_last_3d) +def _match_input_dtype(out, reference): + """Give ``out`` ``reference``'s dtype, casting only if it differs. + + The compiled and eager rungs call ``F.group_norm``, which carries + autocast's **fp32** cast policy and therefore returns fp32 for a bf16 input + inside an autocast region. ``FastGroupNorm`` emits the input's dtype + instead (see its docstring for why, and for the warning that goes with it), + and the Triton rung does it in the store; these two have to do it here. + + Free outside autocast and free for an fp32 input, where ``F.group_norm`` + already returns the input's dtype -- the identity check, not the cast, is + what runs. Inside autocast on a narrower input it is one cast, which is + exactly the cast the *consumer* was going to do a moment later, so nothing + is spent that was not being spent already; what it buys is that a rung + change cannot change the dtype of an activation. + + ``.to(dtype)`` preserves the memory format, so this composes with + :func:`_match_memory_format` in either order; it is applied last because + the fused-activation rung applies its activation before the store, and the + other two must round after the ReLU for the same reason. + """ + if out.dtype is reference.dtype: + return out + return out.to(reference.dtype) + + class FastGroupNorm(nn.GroupNorm): """``nn.GroupNorm`` with a Triton GPU kernel and an optional fused ReLU. @@ -595,6 +625,75 @@ class FastGroupNorm(nn.GroupNorm): correctness of the model therefore does not depend on which kernel runs; only the number of memory passes does. + Output dtype -- read this before consuming a GroupNorm output + ============================================================= + **This module returns its input's dtype, which is not what + ``nn.GroupNorm``/``F.group_norm`` returns under autocast.** ``at::group_norm`` + carries autocast's ``fp32`` cast policy, so stock GroupNorm returns fp32 + for *any* input dtype inside an enabled autocast region; this module + returns bf16 for a bf16 input, fp16 for an fp16 one, and fp32 for an fp32 + one. Concretely, in the shipped configuration (``torch_amp: 1``, bf16) the + activations a ``FastGroupNorm`` hands out are **bf16**; under + ``torch_amp: 0``, or in ``eval``/``inference_mode`` outside an autocast + region, they are **fp32** and nothing here differs from stock; under an + fp16 autocast they are fp16. "The input's dtype" is the rule; "bf16" is + only what that rule evaluates to in production. + + The statistics are *not* narrowed: mean and variance are accumulated in + fp32 on every rung, exactly as before, and the normalized value is computed + in fp32. Only the store changes, so the output is the fp32 answer rounded + once -- not a narrower computation. What this is buying is the round trip + that rounding used to cost twice over: every consumer of a GroupNorm in + this network immediately narrowed the fp32 result back to bf16 (see below), + and writing fp32 to memory to read it back and write bf16 is a full-volume + read plus two full-volume writes that do no arithmetic. Priced end to end + before this landed, paired and alternating arms: **-36.39 +/- 1.55 ms of a + 449.56 ms step at scale 8 on one MI300A** (peak memory 42.78 -> 38.49 GiB) + and -3.93 +/- 1.19 ms of 66.49 at scale 7, of which -22.13 ms at scale 8 is + the cast traffic itself and the rest is the GroupNorm's own halved store + and a cheaper max-pool. + + **Why it is safe here, and what would invalidate that.** Every consumer of + a GroupNorm output in this UNet rounds it to autocast's dtype before doing + any arithmetic: the following convolution (``aten::convolution`` carries the + ``lower_precision_fp`` policy), the skip concatenation (``_skip_concat`` + casts to :func:`~ScaFFold.unet.unet_parts._consumer_dtype` for exactly this + reason), and ``max_pool3d``, which is a selection and commutes with + rounding. So no consumer in *this* model ever sees the fp32 bits it used + to be handed. A future consumer that does want them -- an fp32 residual + add, a loss term, a normalization whose statistics are taken over the + GroupNorm output, anything reading it outside an autocast region -- will + get bf16 **silently**, with no error and no warning, only ~3 decimal digits + where it expected ~7. Such a consumer must upcast explicitly, or take its + input from somewhere else. This is the one place ``FastGroupNorm`` is not + a drop-in replacement, and it is deliberate. + + The departure is *not* free numerically, in one respect worth naming: it + does not change the forward's value (which is the same fp32 number, rounded + once), but it changes the backward's reduction order, so a run is not + bitwise comparable with one taken before it. Measured at 4.4e-06 relative + on the loss by step 23 at scale 8 -- the same class and size as the other + fp32-noise changes on this branch. Each configuration remains bitwise + reproducible with itself. + + All three rungs do this, not just the Triton one. The Triton kernel is + told to store the input's dtype (:func:`_triton_forward` passes + ``out_dtype``); the compiled and eager rungs cast afterwards + (:func:`_match_input_dtype`). Divergence would have been cheaper -- the + fallback rungs pay one cast -- and it was rejected: the rungs are chosen + per module and per *process*, a latch can demote an unproven module + mid-run, and a proven module still falls back when its kernel raises + outside a backward replay. A dtype that depended on that choice would be + an activation width that changes mid-step, differs between DDP ranks with + different latch histories, and -- the sharp one -- breaks + ``torch.utils.checkpoint``, whose non-reentrant recompute compares the + *dtype* of every saved tensor against the original and raises + ``CheckpointError`` on a mismatch. The rungs already go to some length to + be indistinguishable in everything a caller can observe (see + ``_match_memory_format`` and the module docstring's "Latches"); dtype is + now on that list, and the cast that keeps it there is one the consumer + would have paid anyway. + DistConv's ``DCTensor`` gets the fast kernels too, by unwrapping to the local shard in front of them rather than by letting the op dispatch through the wrapper. Both would work -- the Triton kernel is a real dispatcher op, @@ -687,26 +786,47 @@ def _activate(self, out): ) def _triton_forward(self, local): - """The native channels-last kernel, with the activation fused in.""" + """The native channels-last kernel, with the activation fused in. + + ``out_dtype=local.dtype`` is this module's departure from + ``F.group_norm``'s autocast contract, spelled at the one call site that + wants it rather than in the kernel module's default -- the standalone + ``triton_group_norm`` is documented as reproducing ``F.group_norm``'s + dtype and still does. The kernel stores that dtype directly, so unlike + the two rungs below there is no fp32 intermediate to cast; the + statistics are fp32 either way. See the class docstring's "Output + dtype". + """ return _get_triton_module().triton_group_norm( - local, self.num_groups, self.weight, self.bias, self.eps, self.activation + local, + self.num_groups, + self.weight, + self.bias, + self.eps, + self.activation, + out_dtype=local.dtype, ) def _compiled_forward(self, local): - return self._activate( - _match_memory_format( - _get_compiled_group_norm()( - local, self.num_groups, self.weight, self.bias, self.eps - ), - local, - ) + return _match_input_dtype( + self._activate( + _match_memory_format( + _get_compiled_group_norm()( + local, self.num_groups, self.weight, self.bias, self.eps + ), + local, + ) + ), + local, ) def _eager_forward(self, input): # super().forward() is the stock kernel; deferring to it keeps the eager - # path identical to nn.GroupNorm's (plus the ReLU and the relayout) by - # construction. - return self._activate(_match_memory_format(super().forward(input), input)) + # path identical to nn.GroupNorm's (plus the ReLU, the relayout and the + # dtype narrowing) by construction. + return _match_input_dtype( + self._activate(_match_memory_format(super().forward(input), input)), input + ) def forward(self, input): global _compile_failed, _triton_failed diff --git a/ScaFFold/unet/triton_group_norm.py b/ScaFFold/unet/triton_group_norm.py index 2dd28e9..41003c3 100644 --- a/ScaFFold/unet/triton_group_norm.py +++ b/ScaFFold/unet/triton_group_norm.py @@ -61,11 +61,12 @@ Public API ========== ``triton_group_norm(input, num_groups, weight=None, bias=None, eps=1e-5, -activation=None)`` +activation=None, *, out_dtype=MATCH_STOCK_DTYPE)`` Drop-in for ``F.group_norm`` (plus an optionally fused ReLU) with first-order autograd support. Accepts *anything* ``F.group_norm`` accepts; inputs the Triton kernel cannot serve fall back to ``F.group_norm`` - internally (see "Layouts" below). + internally (see "Layouts" below). ``out_dtype`` is an opt-*in* that + overrides the dtype rule below; its default keeps the rule exactly. ``is_supported(input, num_groups, weight=None, bias=None, activation=None)`` Cheap, side-effect-free predicate: ``True`` exactly when the native Triton @@ -87,7 +88,13 @@ without materializing the fp32 copy of the input that autocast's cast would create: the kernels read the input at its native width and accumulate in fp32, which is bit-for-bit the same computation as upcasting first, but reads - half the bytes. Gradients follow the same rule: ``d_input`` has the input's + half the bytes. A caller who wants a *different* output width has to ask + for it: ``out_dtype=`` overrides this rule (``None`` spells "the input's + dtype"), the statistics stay fp32 either way, and the default -- + ``MATCH_STOCK_DTYPE`` -- *is* the rule above, so nothing that does not ask + can be surprised. ``FastGroupNorm`` is the one caller that asks; see its + docstring for why, and for what makes it safe there. + Gradients follow the same rule: ``d_input`` has the input's dtype, ``d_weight``/``d_bias`` have the parameter's dtype. Forward time at ``[1,64,256^3]`` / ``[1,128,128^3]`` / ``[1,256,64^3]`` in ms: fp32 4.28/1.20/0.35, bf16 2.42/0.68/0.22, fp16 2.43/0.66/0.22, and bf16-in with @@ -1708,6 +1715,55 @@ def _autocast_out_dtype(input: torch.Tensor) -> Optional[torch.dtype]: return None +class _MatchStockDtype: + """The type of :data:`MATCH_STOCK_DTYPE`; see it.""" + + __slots__ = () + + def __repr__(self): # pragma: no cover - diagnostics only + return "MATCH_STOCK_DTYPE" + + +#: :func:`triton_group_norm`'s default ``out_dtype``: "whatever +#: ``F.group_norm`` would have returned for this call", i.e. fp32 inside an +#: autocast region and the input's dtype outside one. +#: +#: A distinct sentinel rather than ``None`` because ``None`` already means +#: something else here -- it is the *op*-level spelling of "the input's dtype", +#: and that is precisely one of the things a caller may ask for explicitly. The +#: default has to be a third value, and a named one makes the two askable +#: answers visible at the call site instead of hiding one of them behind the +#: absence of an argument. +MATCH_STOCK_DTYPE = _MatchStockDtype() + + +def _checked_out_dtype(out_dtype): + """Validate :func:`triton_group_norm`'s ``out_dtype`` and return it. + + Argument-only, so it can run before the input has been shown to be a tensor + at all: a bad ``out_dtype`` is the caller's error either way, and raising it + here keeps the stock ``F.group_norm`` error for a bad *input* intact. + """ + if out_dtype is MATCH_STOCK_DTYPE or out_dtype is None: + return out_dtype + if out_dtype in SUPPORTED_DTYPES: + return out_dtype + raise ValueError( + f"out_dtype must be MATCH_STOCK_DTYPE (the default), None (the input's " + f"dtype) or one of {SUPPORTED_DTYPES}, got {out_dtype!r}" + ) + + +def _out_dtype_for(input: torch.Tensor, out_dtype) -> Optional[torch.dtype]: + """Resolve a validated ``out_dtype`` against ``input``. + + Returns the op's spelling: a concrete dtype, or ``None`` for "the input's". + """ + if out_dtype is MATCH_STOCK_DTYPE: + return _autocast_out_dtype(input) + return out_dtype + + def is_supported( input, num_groups: int, @@ -1784,6 +1840,8 @@ def triton_group_norm( bias=None, eps: float = 1e-5, activation: Optional[str] = None, + *, + out_dtype=MATCH_STOCK_DTYPE, ): """GroupNorm with an optionally fused activation, channels-last native. @@ -1794,16 +1852,37 @@ def triton_group_norm( callers with a faster fallback should branch on :func:`is_supported` themselves. - The output has the input's memory format and ``F.group_norm``'s dtype; see - the module docstring for the full contract. + The output has the input's memory format and, by default, + ``F.group_norm``'s dtype; see the module docstring for the full contract. + + ``out_dtype`` is the one way to depart from that dtype rule, and it is + keyword-only and opt-in: :data:`MATCH_STOCK_DTYPE` (the default) reproduces + ``F.group_norm`` exactly, ``None`` asks for the input's dtype -- which is + the same thing outside autocast and *narrower* inside it -- and a + :data:`SUPPORTED_DTYPES` member asks for that dtype. It is honoured on both + routes, the kernel's and the ``F.group_norm`` fallback's, so the answer + never depends on which one served the call: the kernel stores the requested + dtype directly (no round trip through fp32), the fallback casts afterwards. + Only the *store* changes -- the statistics are accumulated in fp32 + regardless, so a narrowed output is the fp32 answer rounded once, not a + narrower computation. """ if activation not in SUPPORTED_ACTIVATIONS: raise ValueError( f"activation must be one of {SUPPORTED_ACTIVATIONS}, got {activation!r}" ) + out_dtype = _checked_out_dtype(out_dtype) if not is_supported(input, num_groups, weight, bias, activation): out = F.group_norm(input, num_groups, weight, bias, eps) - return F.relu(out) if activation == "relu" else out + if activation == "relu": + out = F.relu(out) + # Resolved only now: `input` has been through F.group_norm, so it is a + # tensor, and MATCH_STOCK_DTYPE on this route is by construction a + # no-op -- F.group_norm just produced exactly that dtype. + target = _out_dtype_for(input, out_dtype) + if target is None: + target = input.dtype + return out if out.dtype is target else out.to(target) out, _mean, _rstd = torch.ops.scaffold_gn.group_norm( input, num_groups, @@ -1811,6 +1890,6 @@ def triton_group_norm( bias, float(eps), activation, - _autocast_out_dtype(input), + _out_dtype_for(input, out_dtype), ) return out diff --git a/tests/test_groupnorm.py b/tests/test_groupnorm.py index 4b1cb05..cf338a4 100644 --- a/tests/test_groupnorm.py +++ b/tests/test_groupnorm.py @@ -1581,6 +1581,216 @@ def run(triton): _assert_close(triton[index], eager[index], 1e-4, what) +# --------------------------------------------------------------------------- +# Output dtype -- the deliberate departure from F.group_norm's autocast rule +# --------------------------------------------------------------------------- +# +# `FastGroupNorm` returns its *input's* dtype where `F.group_norm` carries +# autocast's fp32 cast policy and returns fp32. Three claims are pinned below, +# and they are not the same claim: that the departure happens at all (and is +# visible in the model, not only in a unit shape); that it happens on *every* +# rung, so a mid-run fallback cannot change an activation's width; and that +# outside autocast -- `torch_amp: 0`, eval, inference -- nothing moves. +# +# The standalone `triton_group_norm()` keeps the stock rule; that contract is +# pinned in tests/test_triton_group_norm.py and must not move with these. + +#: How far two GroupNorm results may differ once both have been rounded to a +#: narrow dtype. Just over one ulp at the output's own magnitude (2^-8 for +#: bf16, 2^-11 for fp16), because the two rungs' fp32 answers differ in the +#: last places and each is then rounded independently -- a difference the +#: fp32-output comparisons elsewhere in this file never see. +_NARROW_TOL = {torch.bfloat16: 8e-3, torch.float16: 1e-3} + + +def _pin_rung(rung): + """Route the next call to exactly one rung of the ladder.""" + gn_mod.set_triton_enabled(rung == "triton") + gn_mod.set_compile_enabled(rung == "compiled") + + +@pytest.mark.gpu +@pytest.mark.parametrize("rung", ["triton", "compiled", "eager"]) +@pytest.mark.parametrize("autocast_dtype", [torch.bfloat16, torch.float16]) +def test_gpu_output_dtype_is_the_inputs_under_autocast(rung, autocast_dtype): + """Under autocast the output is the *input's* dtype, on all three rungs. + + The control in the same region is stock ``F.group_norm``, which returns + fp32 here -- that is the contract being departed from, so it is measured + rather than asserted from memory. The eager rung is additionally held to + the sharp form of the claim: its result is bitwise the stock fp32 result + rounded once, so the narrowing is a store and not a narrower computation. + """ + device = torch.device("cuda") + generator = torch.Generator(device=device).manual_seed(71) + x = torch.randn(1, 64, 16, 16, 16, device=device, generator=generator).to( + dtype=autocast_dtype, memory_format=torch.channels_last_3d + ) + + fast = FastGroupNorm(_GROUPS, 64).to(device) + with torch.no_grad(): + fast.weight.normal_(1.0, 0.1, generator=generator) + fast.bias.normal_(0.0, 0.1, generator=generator) + + calls = [] + original = FastGroupNorm._triton_forward + + def spy(self, local): + calls.append(local.dtype) + return original(self, local) + + _pin_rung(rung) + FastGroupNorm._triton_forward = spy + try: + with torch.autocast("cuda", dtype=autocast_dtype): + stock = nn.functional.group_norm(x, _GROUPS, fast.weight, fast.bias) + out = fast(x) + finally: + FastGroupNorm._triton_forward = original + + assert len(calls) == (1 if rung == "triton" else 0), "wrong rung answered" + assert stock.dtype is torch.float32, "assumption about stock GroupNorm broke" + assert out.dtype is x.dtype, "FastGroupNorm did not emit its input's dtype" + assert _channels_last(out) + _assert_close(out.float(), stock.float(), _NARROW_TOL[autocast_dtype], "output") + if rung == "eager": + assert torch.equal(out, stock.to(x.dtype)), "not a plain rounding of stock" + + +@pytest.mark.gpu +@pytest.mark.parametrize("rung", ["triton", "compiled", "eager"]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_gpu_output_dtype_outside_autocast_is_unchanged(rung, dtype): + """``torch_amp: 0``, ``eval`` and ``inference_mode`` see no change at all. + + Outside an autocast region ``F.group_norm`` already returns the input's + dtype, so "the input's dtype" and "``F.group_norm``'s dtype" are the same + rule and this module's departure is invisible. The fp32 parametrization is + the shipped ``torch_amp: 0`` configuration; the bf16 one (a bf16 module, + which is how a hand-cast model arrives) checks that the answer follows the + input rather than being pinned to fp32 by accident. + """ + device = torch.device("cuda") + generator = torch.Generator(device=device).manual_seed(73) + x = torch.randn(1, 64, 16, 16, 16, device=device, generator=generator).to( + dtype=dtype, memory_format=torch.channels_last_3d + ) + fast = FastGroupNorm(_GROUPS, 64).to(device=device, dtype=dtype) + + calls = [] + original = FastGroupNorm._triton_forward + + def spy(self, local): + calls.append(local.dtype) + return original(self, local) + + _pin_rung(rung) + FastGroupNorm._triton_forward = spy + try: + stock = nn.functional.group_norm(x, _GROUPS, fast.weight, fast.bias) + out = fast(x) + finally: + FastGroupNorm._triton_forward = original + + assert len(calls) == (1 if rung == "triton" else 0), "wrong rung answered" + assert stock.dtype is dtype, "assumption about stock GroupNorm broke" + assert out.dtype is stock.dtype + _assert_close(out.float(), stock.float(), _NARROW_TOL.get(dtype, 1e-5), "output") + + +@pytest.mark.gpu +def test_gpu_a_rung_fallback_does_not_change_the_output_dtype(caplog): + """The reason all three rungs narrow rather than only the fast one. + + A rung is chosen per module and per process, and a kernel failure demotes a + module mid-run (see the module docstring's "Latches"). If only the Triton + rung emitted the input's dtype, that demotion would silently change the + width of an activation between one step and the next -- and between DDP + ranks with different latch histories -- and would break + ``torch.utils.checkpoint``, which compares the dtype of every recomputed + saved tensor. So the same call is answered twice, once by the kernel and + once by the fallback it lands on when the kernel raises, and the two + answers must agree on dtype. + """ + from ScaFFold.unet.triton_group_norm import TritonKernelError + + device = torch.device("cuda") + generator = torch.Generator(device=device).manual_seed(77) + x = torch.randn(1, 64, 16, 16, 16, device=device, generator=generator).to( + dtype=torch.bfloat16, memory_format=torch.channels_last_3d + ) + fast = FastGroupNorm(_GROUPS, 64).to(device) + gn_mod.set_triton_enabled(None) + gn_mod.set_compile_enabled(True) + + with torch.autocast("cuda", dtype=torch.bfloat16): + served = fast(x) + assert fast._triton_ok, "the Triton rung did not serve the first call" + + original = FastGroupNorm._triton_forward + + def _raises(self, local): + raise TritonKernelError("simulated Triton failure") + + FastGroupNorm._triton_forward = _raises + try: + with caplog.at_level(logging.WARNING, logger=gn_mod.__name__): + with torch.autocast("cuda", dtype=torch.bfloat16): + fell_back = fast(x) + finally: + FastGroupNorm._triton_forward = original + + assert gn_mod._triton_failed is True, "the failure did not latch" + assert any("falling back" in r.message for r in caplog.records) + assert fell_back.dtype is served.dtype is x.dtype + assert _channels_last(fell_back) + _assert_close( + fell_back.float(), served.float(), _NARROW_TOL[torch.bfloat16], "output" + ) + + +@pytest.mark.gpu +@pytest.mark.parametrize("autocast", [True, False]) +def test_gpu_unet_group_norm_outputs_follow_the_activation_dtype(autocast): + """The claim as the model sees it: every site, not one unit shape. + + Under bf16 autocast every ``FastGroupNorm`` in the UNet consumes bf16 (its + producer is a convolution, which autocast casts) and must now hand back + bf16 rather than the fp32 stock GroupNorm would return -- that round trip + is what the change removes. With autocast off the same census must read + fp32 everywhere, which is the ``torch_amp: 0`` path. + """ + device = torch.device("cuda") + model = _make_unet(seed=0).to(device, memory_format=torch.channels_last_3d) + x = _make_input(seed=9).to(device).contiguous(memory_format=torch.channels_last_3d) + + census = [] + + def hook(module, inputs, output): + census.append((inputs[0].dtype, output.dtype)) + + for module in model.modules(): + if isinstance(module, FastGroupNorm): + module.register_forward_hook(hook) + + gn_mod.set_triton_enabled(None) + with ( + torch.autocast("cuda", dtype=torch.bfloat16, enabled=autocast), + torch.no_grad(), + ): + model(x) + + assert census, "no GroupNorm ran" + expected = torch.bfloat16 if autocast else torch.float32 + assert all(seen_in is expected for seen_in, _ in census), ( + f"a GroupNorm did not see {expected}: {census}" + ) + wrong = [ + i for i, (seen_in, seen_out) in enumerate(census) if seen_in is not seen_out + ] + assert not wrong, f"GroupNorm changed the activation dtype at sites {wrong}" + + @pytest.mark.gpu def test_gpu_unet_keeps_the_channels_last_chain(monkeypatch): """The whole point: GroupNorm stops breaking the layout chain in the model. diff --git a/tests/test_groupnorm_wiring_edge.py b/tests/test_groupnorm_wiring_edge.py index 86a6597..7602b37 100644 --- a/tests/test_groupnorm_wiring_edge.py +++ b/tests/test_groupnorm_wiring_edge.py @@ -520,10 +520,16 @@ def test_a_proven_rung_does_not_degrade_while_a_backward_replays_it(monkeypatch, proven_flag = "_triton_ok" if rung == "triton" else "_compiled_ok" def make_kernel(fail_always): - def _kernel(input, num_groups, weight, bias, eps, *activation): + # ``out_dtype`` is keyword-only on the real ``triton_group_norm`` and + # ``_triton_forward`` passes it on every call; the double honours it + # rather than swallowing it, so that a double never hides a dtype the + # kernel would have produced. The compiled rung's kernel takes no such + # argument, which is why this signature has to accept it optionally. + def _kernel(input, num_groups, weight, bias, eps, *activation, out_dtype=None): if fail_always or gn_mod._replaying_a_forward(): raise failure("simulated kernel failure") - return F.group_norm(input, num_groups, weight, bias, eps) + out = F.group_norm(input, num_groups, weight, bias, eps) + return out if out_dtype in (None, out.dtype) else out.to(out_dtype) return _kernel diff --git a/tests/test_triton_group_norm.py b/tests/test_triton_group_norm.py index ff4d0ee..a2ddf11 100644 --- a/tests/test_triton_group_norm.py +++ b/tests/test_triton_group_norm.py @@ -603,6 +603,85 @@ def test_output_dtype_matches_stock(dtype, autocast_dtype): assert _rel(got.float(), stock.float()) <= _TOL[stock.dtype] +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("autocast_dtype", [None, torch.bfloat16]) +def test_out_dtype_opt_in_overrides_the_stock_rule(dtype, autocast_dtype): + """``out_dtype=`` is the only way past the rule above, and it is opt-in. + + Three spellings, one call each: the default (:data:`tgn.MATCH_STOCK_DTYPE`) + reproduces ``F.group_norm``; ``None`` asks for the input's dtype, which + differs from stock exactly inside an autocast region; and an explicit dtype + asks for that one. ``FastGroupNorm`` is the caller that uses the middle + spelling -- see ``ScaFFold.unet.group_norm``. + """ + device = torch.device("cuda") + x, weight, bias, _ = _tensors((1, 64, 5, 5, 5), dtype, device, seed=61) + if autocast_dtype is not None: + weight = weight.float() + bias = bias.float() + ctx = ( + torch.autocast("cuda", dtype=autocast_dtype) + if autocast_dtype is not None + else torch.autocast("cuda", enabled=False) + ) + with ctx: + default = triton_group_norm(x, GROUPS, weight, bias, EPS) + narrow = triton_group_norm(x, GROUPS, weight, bias, EPS, out_dtype=None) + wide = triton_group_norm(x, GROUPS, weight, bias, EPS, out_dtype=torch.float32) + stock = torch.float32 if autocast_dtype is not None else dtype + assert default.dtype == stock, "the default moved; it is the standalone contract" + assert narrow.dtype == dtype + assert wide.dtype == torch.float32 + for got in (default, narrow, wide): + assert got.is_contiguous(memory_format=CL) + # Only the *store* changes: the statistics and the normalized value are + # fp32 on every one of these calls, and the tiling plan is a function of + # the shape alone (`_plan` takes no dtype), so the narrow answer is the + # wide one rounded once -- bitwise, not approximately. + assert torch.equal(narrow, wide.to(dtype)) + assert torch.equal( + default, wide if default.dtype is torch.float32 else wide.to(dtype) + ) + + +@pytest.mark.gpu +def test_out_dtype_is_honoured_on_the_fallback_route_too(): + """A contiguous input takes ``F.group_norm``, and still answers in the + requested dtype -- otherwise the dtype of a result would depend on which + kernel served it, which is the thing every other part of this contract + exists to prevent.""" + device = torch.device("cuda") + gen = torch.Generator(device=device).manual_seed(63) + x = torch.randn(2, 64, 5, 6, 7, device=device, dtype=torch.bfloat16, generator=gen) + weight = torch.randn(64, device=device, generator=gen) + bias = torch.randn(64, device=device, generator=gen) + with torch.autocast("cuda", dtype=torch.bfloat16): + assert not is_supported(x, GROUPS, weight, bias), "meant to be a fallback" + stock = triton_group_norm(x, GROUPS, weight, bias, EPS) + narrow = triton_group_norm(x, GROUPS, weight, bias, EPS, out_dtype=None) + relu = triton_group_norm( + x, GROUPS, weight, bias, EPS, "relu", out_dtype=torch.bfloat16 + ) + assert stock.dtype is torch.float32, "assumption about stock GroupNorm broke" + assert narrow.dtype is torch.bfloat16 + assert torch.equal(narrow, stock.to(torch.bfloat16)) + # ... and the activation is applied before the narrowing, as it is in the + # kernel's store, so a fused and an unfused ReLU still agree bitwise. + assert relu.dtype is torch.bfloat16 + assert torch.equal(relu, F.relu(stock).to(torch.bfloat16)) + + +@pytest.mark.parametrize("bad", [torch.float64, torch.int32, "bfloat16", 16]) +def test_out_dtype_rejects_what_the_kernel_cannot_store(bad): + """Validated as an argument, before anything looks at the input, so the + message names the argument rather than surfacing as a Triton compile error + deep inside a launch.""" + x = torch.randn(1, 64, 2, 2, 2).to(memory_format=CL) + with pytest.raises(ValueError, match="out_dtype"): + triton_group_norm(x, GROUPS, out_dtype=bad) + + @pytest.mark.gpu def test_autocast_gradient_dtypes_match_stock(): """Under autocast, ``d_input`` keeps the input's dtype and the parameter From bc2b59ed95c22719d8c728f6de189959ac6bc435 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Fri, 7 Aug 2026 19:13:30 -0700 Subject: [PATCH 2/2] Record what the GroupNorm dtype change actually measures 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. --- ScaFFold/unet/group_norm.py | 54 ++++++++++++++++++++++++++----------- 1 file changed, 39 insertions(+), 15 deletions(-) diff --git a/ScaFFold/unet/group_norm.py b/ScaFFold/unet/group_norm.py index 2975f5d..af0a8f2 100644 --- a/ScaFFold/unet/group_norm.py +++ b/ScaFFold/unet/group_norm.py @@ -646,20 +646,24 @@ class FastGroupNorm(nn.GroupNorm): that rounding used to cost twice over: every consumer of a GroupNorm in this network immediately narrowed the fp32 result back to bf16 (see below), and writing fp32 to memory to read it back and write bf16 is a full-volume - read plus two full-volume writes that do no arithmetic. Priced end to end - before this landed, paired and alternating arms: **-36.39 +/- 1.55 ms of a - 449.56 ms step at scale 8 on one MI300A** (peak memory 42.78 -> 38.49 GiB) - and -3.93 +/- 1.19 ms of 66.49 at scale 7, of which -22.13 ms at scale 8 is - the cast traffic itself and the rest is the GroupNorm's own halved store - and a cheaper max-pool. + read plus two full-volume writes that do no arithmetic. Measured end to + end on the landed code, paired, arms alternating within each of 6 + replicates: **-35.73 +/- 0.93 ms of a 440.28 ms step at scale 8 on one + MI300A** (0.9189x; peak memory 42.776 -> 38.486 GiB) and **-4.03 +/- 0.57 ms + of 66.00 at scale 7** (0.9390x; 6.102 -> 5.572 GiB). Of the 35.4 ms of + device time that buys at scale 8, 22.2 is the cast traffic itself, 9.1 the + GroupNorm's own halved store and load, and 1.4 a max-pool that reads and + writes half the bytes; the remaining 2.8 sits in rows this change cannot + reach and is not claimed. **Why it is safe here, and what would invalidate that.** Every consumer of a GroupNorm output in this UNet rounds it to autocast's dtype before doing any arithmetic: the following convolution (``aten::convolution`` carries the ``lower_precision_fp`` policy), the skip concatenation (``_skip_concat`` casts to :func:`~ScaFFold.unet.unet_parts._consumer_dtype` for exactly this - reason), and ``max_pool3d``, which is a selection and commutes with - rounding. So no consumer in *this* model ever sees the fp32 bits it used + reason), and ``max_pool3d``, whose *forward* is a selection and commutes + with rounding (its backward does not; see below). So no consumer in *this* + model ever sees the fp32 bits it used to be handed. A future consumer that does want them -- an fp32 residual add, a loss term, a normalization whose statistics are taken over the GroupNorm output, anything reading it outside an autocast region -- will @@ -668,13 +672,33 @@ class FastGroupNorm(nn.GroupNorm): input from somewhere else. This is the one place ``FastGroupNorm`` is not a drop-in replacement, and it is deliberate. - The departure is *not* free numerically, in one respect worth naming: it - does not change the forward's value (which is the same fp32 number, rounded - once), but it changes the backward's reduction order, so a run is not - bitwise comparable with one taken before it. Measured at 4.4e-06 relative - on the loss by step 23 at scale 8 -- the same class and size as the other - fp32-noise changes on this branch. Each configuration remains bitwise - reproducible with itself. + **The departure is not free numerically, and the reason is not the + kernel.** The forward is exact -- every site's output is bitwise the wide + answer rounded once, and the whole model's output is bitwise unchanged -- + but the backward moves, so a run is not bitwise comparable with one taken + before this. Measured on the loss at production geometry: identical at + step 1, first differing at step 2, worst **3.3e-06 relative by step 13 at + scale 8** and 5.4e-06 by step 18 at scale 7 over 24 steps, which is the same + class and size as the other fp32-noise changes on this branch. Two + mechanisms, both *downstream* of this module and both measured + (``work/profile/PROFILE_GN_DTYPE.md`` SS6; the kernel itself is bitwise + invariant to its output dtype, and its tiling plan is a pure function of the + shape, so an earlier guess that the plan changed is refuted): + + 1. **Gradient accumulation at fan-out.** Every encoder block's output is + consumed twice -- the max-pool into the next block and the skip + concatenation -- so autograd sums two cotangents at it. When this + module returned fp32 both arrived upcast from bf16 and the sum was + exact; now it is rounded to bf16. A GroupNorm with one consumer is + bitwise unchanged; the same one with two is not. + 2. **``max_pool3d``'s tie-breaking.** Rounding to bf16 creates ties inside + a pooling window that fp32 broke strictly, and the backward scatters + through *indices*, so the gradient lands on a different element. 0.76 % + of windows have a tied bf16 maximum on a ReLU'd GroupNorm output. + + Each configuration remains bitwise reproducible **with itself**: 6 + independent processes per arm per scale produce identical 24-step loss + vectors. All three rungs do this, not just the Triton one. The Triton kernel is told to store the input's dtype (:func:`_triton_forward` passes