[layernorm] Add backward pass for training (PR 3/3, #769)#801
Conversation
b375c78 to
a87fa37
Compare
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> Co-authored-by: Cursor <cursoragent@cursor.com>
a87fa37 to
d58dd56
Compare
|
cc @coderfeli @zhiding512 Can u help review this final PR of this series? Thx so much! |
|
Change fx.Tensor to fx.Pointer, maybe much better? |
Use raw pointer arguments for every LayerNorm backward tensor, reconstructing the static per-row views in the kernel. This removes redundant dynamic layout operands while keeping M runtime-variable.\n\nUse the shared atomic_add helper for fp32 gradient accumulation and strengthen the wrapper contracts required by the pointer-only ABI. Update direct builder tests to compile and launch with typed pointers.
|
@zhiding512 Good suggestion — addressed in 45b786c. The backward kernel and launcher now use fx.Pointer for all eight tensor arguments, not only DGamma/DBias. Each block reconstructs static row/scalar views from the raw addresses, while M remains a runtime argument, so the compiled cache still works across different row counts. This reduces the call ABI from 16 dispatch slots to 10. I also switched the two accumulations to the shared atomic_add helper and added the contiguous/shape/device guards required by the pointer-only contract. Verified on MI355X: 11/11 non-large LayerNorm tests plus the pointer-argument unit test pass, including backward, autograd, and two-GPU coverage. |
|
cc @zhiding512 @coderfeli CI all passed. Could u pls have some final check? Thx! |
|
@zhiding512 Follow-up to the pointer suggestion is now in bc0cebe. The earlier 45b786c converted backward; this commit also converts the plain forward Input/Gamma/Beta/Output/Mean/Rstd kernel and launcher ABI, reconstructs static row/scalar views from raw addresses, and migrates every in-tree direct caller. The wrapper now explicitly enforces the contiguous, dtype, device, and 1D length-N contracts that Tensor metadata previously carried. MI355X store-stats A/B shows 1.23-1.41x better direct hot-call throughput on the representative small/medium shapes and 3-5% better allocation-included hot wrapper time; 4096x4096 is compute/bandwidth dominated and improves about 1%. Generic, Vec8, stats, backward/autograd, nonzero storage offsets, large M, non-default stream, and two-GPU/device-mismatch validation all pass. Full protocol and numbers are in the updated PR body; current-SHA CI is running. |
Summary
PR 3 of 3 for #769. #795 and #800 are already merged; this final PR adds plain LayerNorm training support and completes the series.
The change now covers the complete plain LayerNorm path:
Fused-add and quantized LayerNorm kernels remain unchanged.
Implementation
Forward stats and pointer ABI
build_layernorm_module gains store_stats and eps. When training needs gradients, each row writes FP32 mean and rstd for backward.
Input, Gamma, Beta, Output, Mean, and Rstd are now fx.Pointer arguments in both the kernel and JIT launcher. The wrapper owns the contiguous layout contract and reconstructs only the static views the current block needs:
Compile misses receive typed flyc.from_c_void_p values; hot calls pass raw data_ptr() addresses. M remains runtime, and the existing cache key still owns N, dtype, store_stats, eps, and device.
For forward, this reduces the launcher ABI including M and stream:
Direct users of build_layernorm_module must therefore compile with typed pointers and invoke the compiled function with raw addresses. All in-tree direct callers and the shared benchmark helper are migrated.
Because pointer arguments no longer carry tensor metadata, the Python wrapper explicitly owns the invariants: contiguous storage, matching dtype/device, and exactly one-dimensional length-N affine vectors. Nonzero storage offsets remain supported.
Backward
build_layernorm_bwd_module runs one block per row. With x_hat=(x-mean)rstd and wdy=dygamma:
dgamma and dbias accumulate in FP32. The backward kernel and launcher also use pointer arguments throughout; the public wrapper zeroes the accumulators after compile and before the real launch.
Autograd and cache ownership
LayerNormFunction provides layernorm(x, weight, bias, eps). eps is a compile-time constant in the forward cache key. Compiled entries remain device-specific, and compile/launch occur under the tensor device guard.
Performance
MI355X (gfx950), same physical GPU, parent 45b786c versus pointer-forward commit bc0cebe. Each number is the median of three independent processes; each process reports a five-run median after 30-50 warmups and 150-1000 iterations depending on size.
The direct column measures hot CompiledFunction call throughput with output/stats preallocated. The wrapper column measures hot layernorm_fwd(..., store_stats=True) wall time with output/stats allocation included.
The no-stats path shows the same pattern: 1.17-1.37x better direct hot-call throughput on the first three shapes and about 1.2-4.3% better wrapper time; 4096x4096 is bandwidth/compute dominated and improves about 1%.
This is an ABI/launch-throughput result, not a claim that register use universally falls. Generated gfx950 metadata is unchanged or slightly higher for some Pointer kernels; the measured benefit comes from eliminating dynamic layout operands and Tensor metadata packing.
Validation
MI355X local validation on the final implementation:
Two independent read-only reviews found no correctness or design blocker.
Scope
This PR does not add the deferred vectorized backward fast path, pass-2 load caching, or reduction-helper refactoring. Those are independent performance changes and should remain separately measurable.
Related
Part of #769. Completes the series after #795 and #800.