diff --git a/ScaFFold/unet/group_norm.py b/ScaFFold/unet/group_norm.py index 6e0729f..af0a8f2 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,99 @@ 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. 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``, 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 + 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, 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 + ``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 +810,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