Repository navigation
Add backward passes to the fused rms_norm, layer_norm and softmax_last_dim ops - #4059
Open
DrJesseGlass wants to merge 3 commits into
Open
DrJesseGlass wants to merge 3 commits into
DrJesseGlass wants to merge 3 commits into
Conversation
The fused ops were registered with apply_op*_no_bwd, so a training graph through them failed at backward(). This adds: - a softmax_bwd kernel (CUDA) and CPU implementation, dx = y * (dy - sum(dy * y)), with the reduction in f32 for the half types and f64 for f64, and a bwd hook on SoftmaxLastDim that uses it on CPU/CUDA and a composed form elsewhere - composed backward rules for RmsNorm and LayerNorm with the statistics in f32 for the half types - the fused softmax forward now does its arithmetic in f32 for bf16/f16 on both CPU and CUDA and rounds once at the store; before, exp was stored in bf16 and then scaled, rounding twice. f32 and f64 are unchanged. Tests check gradient parity against the composed formulations on CPU and CUDA.
The f64 kernels cast inputs to float and accumulated the statistics in float, so for f64 inputs the fused forward used different normalization statistics than the backward, which recomputes them in f64. The kernels now take an accumulator type: float for bf16, f16 and f32 (unchanged), double for f64. layernorm keeps two scalar partial sums instead of a float2 so the same code serves both accumulators. A CUDA test checks the fused f64 forward and gradient against the composed formulation to 1e-10.
The comment described forward_diff as the faster kernel. It is the composed differentiable path; the fused kernel is the fast one and now has a backward pass as well.
This branch has not been deployed
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.
rms_norm, layer_norm, and softmax_last_dim ops return tensors with no backprop op, so the graph ends there. This was raised in #2168, #2977, and #3011 and has five open fix PRs, none merged: #3012 (layer_norm, conflicting, maintainer asked for a kernel-level fix), #3526 (rms_norm only, removes the fast path, never reviewed), #3613 (layer_norm only, conflicting with main), #3724 (softmax_last_dim only, mergeable, only reviewed by copilot), #3823 (softmax and layer_norm, mergeable).
This PR adds bwd for all three ops and keeps the fused forward on every backend. softmax_last_dim backward is a kernel on CPU and CUDA with a composed fallback on Metal. The rms_norm and layer_norm backward passes are composed with statistics in f32 for the half types. The CUDA norm kernels now accumulate statistics in f64. Tests cover f32, f64, bf16 forward and a non-contiguous gradient. The forward_diff doc comment is updated since the fused path now supports backward.
Repro on main:
Prints false false on main, true true with this branch.
With the fused forward kept, forward plus backward on CPU is roughly 6 to 7x faster than the composed path for softmax and 2 to 3x for the norms, and bf16 softmax forward error drops from about 4.5e-2 to 1e-3. We came to this adding fused backward kernels to close the training-time gap with PyTorch and found the ops stopped gradients.