[rmsnorm] Add fused-add / residual backward (PR 2/3, #769)#800
Open
jhinpan wants to merge 3 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>
Contributor
Author
|
Note: this is stacked on #795 (PR 1). Since the base branch lives in a fork, GitHub shows this against |
This was referenced Jul 3, 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
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 2 of 3 for #769 — adds the fused-add / residual (prenorm) RMSNorm backward. Stacked on #795 (PR 1) — review that first; this PR's own changes are the fused-add additions.
What's included
rstd:build_fused_add_rmsnorm_modulegainsstore_rstd+eps. Writes per-rowrstdofadded = x + residual. Default off ⇒ existing launcher byte-for-byte unchanged; quant fused-add path untouched.build_fused_add_rmsnorm_bwd_module— one block/row.d_added = (wdy − a_hat·c1)·rstd;total = d_added + dresidual_out. Sinceadded = x + residual,dx == dresidualunconditionally, so the kernel writesdxonce and the wrapper returns it aliased asdresidual(no second buffer, no double store — the quack pattern).dweightvia fp32 atomicAdd.FusedAddRMSNormFunction(prenorm: returns(out, residual_out), backward consumes the residual_out grad) + publicfused_add_rmsnorm(x, residual, weight, eps, prenorm=True). Follows all PR 1 fixes (eps baked + in cache key, device in cache key +torch.cuda.deviceguard, sharedget_llvm_ptr, contiguity/shape asserts).Self-review (ran the multi-angle review + verified on MI355X)
Applied two fixes it surfaced:
x/residual/weight(andadded/weight/dout) dtypes match — the kernel reads all operands with one elem dtype, so a mismatch silently bit-reinterpreted bytes (verified: mixed bf16/f16 produced err ~5.0 before the guard).Deferred (perf/altitude, noted in code): vectorized backward fast path, pass-2 load caching, reduction-helper dedup, store_rstd launcher fork.
Testing (MI355X, gfx950)
Fused-add backward vs
torch.autogradacross f32/f16/bf16 (bothdresidual_out=Noneand non-None), end-to-endfused_add_rmsnorm()autograd with 3D reshape, and a dtype-mismatch guard test. Full non-large suite: 13/13 pass.Related
Part of #769. Stacked on #795.
🤖 Generated with Claude Code