[layernorm] Add backward pass for training (PR 3/3, #769)#801
Open
jhinpan wants to merge 4 commits into
Open
Conversation
Makes RMSNorm training-capable (PR 1 of 3 for ROCm#769): - build_rmsnorm_module gains store_rstd=False: when enabled, writes per-row rstd (1/RMS, fp32, shape (M,)) needed by backward. Default off keeps the existing launcher signature and all callers byte-for-byte unchanged; covers fast, generic, and small-N (N<=2048) paths. - build_rmsnorm_bwd_module: fused single kernel, one block per row. Pass 1 computes c1 = mean_N(x_hat*wdy); pass 2 stores dx = (wdy - x_hat*c1)*rstd and atomicAdds dw = dy*x_hat into an fp32 DWeight[N]. All reduction and weight-grad accumulation in fp32; only dx cast back to I/O dtype. - fp32 atomicAdd chosen for the cross-row dweight reduction after benchmarking it against a two-pass scratch+finalizer variant on MI355X: atomic never blows up and wins the large-N (real LLM hidden size) regime. - RMSNormFunction (torch.autograd.Function) + public rmsnorm(x, weight, eps), quack-aligned, in the kernel layer. Math matches quack rmsnorm_bwd_ref. - Tests: gradient checks vs torch.autograd across f32/f16/bf16, both N paths (incl. unaligned and N<=2048), plus an end-to-end rmsnorm() autograd test covering batched (>2D) reshape and grads on x + weight. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Fixes from the PR ROCm#795 code review: - eps is now baked into the forward kernel (build_rmsnorm_module / small-N builder gain an `eps` param) instead of being silently ignored in favor of the module EPS=1e-5. rmsnorm(x, w, eps=1e-6) now produces correct numerics; eps is part of the fwd compile-cache key. - Multi-GPU correctness: compiled-fn cache keys now include x.device, and compile+launch run under `with torch.cuda.device(x.device)` so a kernel built on cuda:0 is never launched on cuda:1 (previously faulted with hipErrorInvalidDevice + memory access fault). - Deduplicate the LLVM-ptr helper: hoist get_llvm_ptr into kernels_common.py (was a byte-for-byte copy of hgemm_splitk's helper, which CLAUDE.md says belongs in kernels_common) and use it. - Public rmsnorm() now asserts x.shape[-1] == weight length; rmsnorm_fwd/ rmsnorm_bwd assert contiguous inputs (make_buffer_tensor assumes row-major contiguous). - Document the backward kernel's deferred perf follow-ups (vectorized fast path + pass-1/pass-2 caching) instead of leaving them implicit. - Tests: add test_rmsnorm_eps_honored (eps must change output / match a torch ref at 1e-5/1e-6/1e-2) and test_rmsnorm_multi_gpu (marked multi_gpu; runs rmsnorm fwd+bwd on cuda:0 and cuda:1). Deferred (perf/altitude, not correctness; noted in code): vectorized backward fast path, pass-2 load caching, and the store_rstd launcher duplication (kept to preserve the byte-for-byte store_rstd=False path). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…m#769) Adds the fused-add (prenorm) RMSNorm backward, extends the fused-add forward to save rstd, and wires a quack-aligned autograd wrapper. Follows all the patterns established/fixed in PR 1 (eps baked into the kernel + in the compile-cache key, device in the cache key with a torch.cuda.device guard, shared get_llvm_ptr, contiguity/shape asserts). - Forward: build_fused_add_rmsnorm_module gains store_rstd + eps. When store_rstd, per-row rstd of `added = x + residual` is written (tid 0, fast + generic paths). Default off keeps the existing launcher signature byte-for-byte; the quant fused-add path is untouched. - Backward: build_fused_add_rmsnorm_bwd_module — one block/row. d_added = (wdy - a_hat*c1)*rstd; total = d_added + dresidual_out. Since added = x + residual, dx == dresidual unconditionally, so the kernel writes dx once and the wrapper returns it aliased as dresidual (no second buffer / no double store). dweight via fp32 atomicAdd. - Wrappers + FusedAddRMSNormFunction (prenorm: returns (out, residual_out); backward consumes the residual_out grad) + public fused_add_rmsnorm(x, residual, weight, eps, prenorm=True). - Guards (from review): assert x/residual/weight (and added/weight/dout) dtypes match — the kernel reads all operands with one elem dtype, so a mismatch would silently bit-reinterpret bytes. Tests: fused-add backward vs torch.autograd across f32/f16/bf16 (both dresidual_out=None and non-None), end-to-end fused_add_rmsnorm() autograd with 3D reshape, and a dtype-mismatch guard test. Full non-large suite: 13/13 pass. Deferred (perf/altitude, noted in code): vectorized backward fast path, pass-2 load caching, reduction-helper dedup, store_rstd launcher fork. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Adds plain LayerNorm backward (the stretch item of ROCm#769), mirroring the reviewed rmsnorm backward (PR 1) template and its fixes. - Forward: build_layernorm_module gains store_stats + eps. When store_stats, per-row mean AND rstd (both fp32, (M,)) are written (tid 0, fast + generic paths). Default off keeps the existing launcher byte-for-byte; the fused-add and quant layernorm paths are untouched. - Backward: build_layernorm_bwd_module — one block/row. With x_hat=(x-mean)*rstd, wdy=dy*gamma: c1=mean(wdy), c2=mean(wdy*x_hat) via a two-value block reduction; dx=(wdy-c1-x_hat*c2)*rstd; dgamma and dbias each accumulated over rows via fp32 atomicAdd. - Wrappers + LayerNormFunction + public layernorm(x, weight, bias, eps), following all PR 1/2 fixes: eps baked into kernel + in cache key, device in cache key with torch.cuda.device guard, shared get_llvm_ptr, contiguity/shape/dtype-equality asserts, dgamma+dbias zeroed after compile. Hardening (from review): fwd/bwd assert weight/bias length == N; bwd asserts mean/rstd are fp32. Tests: backward vs torch.autograd (F.layer_norm) across f32/f16/bf16 (dx, dgamma, dbias, mean, rstd), end-to-end layernorm() autograd with 3D reshape, eps-honored, multi-GPU (marked), and dtype-mismatch guard. Full non-large layernorm suite: 11/11 pass. Deferred (perf/altitude, noted in code): vectorized backward fast path, pass-2 load caching, reduction-helper dedup. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Contributor
Author
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
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.
Summary
PR 3 of 3 for #769 — adds plain LayerNorm backward (the "stretch" item). Stacked on #800 (PR 2) → #795 (PR 1) — review those first; this PR's own change is the single layernorm commit.
What's included
mean+rstd:build_layernorm_modulegainsstore_stats+eps. LayerNorm backward needs both stats (unlike rmsnorm's single rstd). Default off ⇒ existing launcher byte-for-byte unchanged; fused-add + quant layernorm paths untouched.build_layernorm_bwd_module— one block/row. Withx_hat=(x−mean)·rstd,wdy=dy·γ:c1=mean_N(wdy),c2=mean_N(wdy·x_hat)(two-value block reduction);dx=(wdy−c1−x_hat·c2)·rstd;dγ=Σ_rows(dy·x_hat)anddβ=Σ_rows(dy)each via fp32 atomicAdd.LayerNormFunction+ publiclayernorm(x, weight, bias, eps). Follows all PR 1/2 fixes (eps baked + in cache key, device in cache key +torch.cuda.deviceguard, sharedget_llvm_ptr, contiguity/shape/dtype-equality asserts, dgamma+dbias zeroed after compile).Self-review (multi-angle review + on-device verification)
No correctness or cross-file issues found. Applied two cheap hardening guards it surfaced:
fwd/bwdassertweight/biaslength == N;bwdassertsmean/rstdare fp32 (direct-caller footguns; the public path was already safe).Deferred (perf/altitude, noted in code, consistent with PR 1/2): vectorized backward fast path, pass-2 load caching, reduction-helper dedup.
Testing (MI355X, gfx950)
Backward vs
torch.autograd(F.layer_norm) across f32/f16/bf16 (dx, dgamma, dbias, mean, rstd); end-to-endlayernorm()autograd with 3D reshape; eps-honored; multi-GPU (marked, exercised); dtype-mismatch guard. Full non-large layernorm suite: 11/11 pass.Related
Part of #769. Completes the 3-PR series (#795, #800, this).
🤖 Generated with Claude Code