vrr/unsloth
2
0
Fork 0
mirror of https://github.com/unslothai/unsloth.git synced 2026-10-19 08:54:21 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 39

Compute RMSNorm weight gradient when the norm weight is trainable (#13594)

Browse source 
* Compute RMSNorm weight gradient when the norm weight is trainable

* Chunk the eager RMSNorm weight-gradient reduction to bound fp32 temporaries

* Compute the RMSNorm weight gradient inside the Triton backward kernels

A TRAIN_W constexpr switches both backward kernels (one-row and multi-row)
to persistent programs that accumulate dL/dW in fp32 registers over their
rows and store one partial each, like flash-linear-attention's
layer_norm_bwd_kernel_row. A small Triton kernel sums the partials in a
fixed order, so eager and torch.compile give the same bytes. With a frozen
weight TRAIN_W is False and the kernels compile to the previous code.
Replaces the torch-side reduction from the earlier commits.

* Launch one RMSNorm dW program for an empty input instead of dividing by zero

* Sum RMSNorm dW partials in 64x64 tiles instead of one row at a time

The serial row loop cost about 73 us per 148 partials on a B200; tiles cut
the sum 4-8x and keep a fixed order, so eager and compiled bytes still match.

* Trim comments in the RMSNorm weight-gradient change

* Allow one bf16 ulp between trainable and frozen RMSNorm dX, as on Blackwell

* Raise RMSNorm dW backward occupancy and pipeline the partials sum

- Accumulate the per-row dW term before the row sum, so the fp32 dY row is
  not held across it (fewer registers).
- On 64-warp SMs (A100 / H100 / B200, CUDA only), rows of at most 16 elements
  per lane launch 32 warps of programs per SM with maxnreg = 64.
- Pipeline the 64x64-tile dW partials reduction (num_stages = 3).
- Cache the row-path launch choice per (device, n_cols): at small shapes the
  backward is host bound, and recomputing it per call cost about 12 us of
  Python per forward + backward on a Colab A100 host.

* Use the forward's rounded normed for the non-Gemma RMSNorm dW

* Trim comments in the RMSNorm dW occupancy and rounding changes
This commit is contained in:
Daniel Han 2026-10-14 09:41:32 -07:00 • committed by GitHub
parent 01a130beeb
commit fc4995a4b4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 457 additions and 39 deletions
Download patch file Download diff file Expand all files Collapse all files

15
tests/test_rms_rope_compile.py
View file

@ -87,6 +87,12 @@ def _rms_case(fn, shape, dtype, gemma, non_contiguous):
return Y.detach(), X.grad, norm.weight.grad
def _assert_rms_matches(eager, compiled):
assert eager[2] is not None and compiled[2] is not None, "trainable weight got no grad"
for e, c, name in zip(eager, compiled, ("Y", "dX", "dW")):
assert _bytes_equal(e, c), name
@pytest.mark.parametrize("dtype", DTYPES, ids = str)
@pytest.mark.parametrize("gemma", [False, True], ids = ["llama", "gemma"])
@pytest.mark.parametrize(
@ -100,8 +106,7 @@ def test_rms_layernorm_compiled_matches_eager(shape, non_contiguous, gemma, dtyp
compiled = _rms_case(
_compile(lambda n, x, gm: fast_rms_layernorm(n, x, gm)), shape, dtype, gemma, non_contiguous
)
for e, c, name in zip(eager, compiled, ("Y", "dX", "dW")):
assert _bytes_equal(e, c), name
_assert_rms_matches(eager, compiled)
if TRACEABLE:
assert _graph_breaks() == 0
@ -246,7 +251,8 @@ def test_bf16_without_native_support_stays_untraced(monkeypatch):
False,
False,
)
assert all(_bytes_equal(e, c) for e, c in zip(eager, compiled)) and _graph_breaks() == 0
_assert_rms_matches(eager, compiled)
assert _graph_breaks() == 0
torch._dynamo.reset()
dynamo_utils.counters.clear()
eager = _rms_case(fast_rms_layernorm, (64, 256), torch.bfloat16, False, False)
@ -257,7 +263,8 @@ def test_bf16_without_native_support_stays_untraced(monkeypatch):
False,
False,
)
assert all(_bytes_equal(e, c) for e, c in zip(eager, compiled)) and _graph_breaks() > 0
_assert_rms_matches(eager, compiled)
assert _graph_breaks() > 0
def test_rope_dynamic_shapes():

142
tests/test_rmsnorm_weight_grad.py Normal file
View file

@ -0,0 +1,142 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
import pytest
from real_accelerator import (
has_real_cuda,
)
import torch
pytestmark = [
pytest.mark.gpu,
pytest.mark.skipif(not has_real_cuda(), reason = "CUDA Triton kernels required"),
]
def _reference(
inputs,
weights,
grad,
gemma,
eps = 1e-6,
):
x = inputs.clone().requires_grad_()
w = weights.clone().requires_grad_()
normed = x.float() * torch.rsqrt(x.float().square().mean(-1, keepdim = True) + eps)
if gemma:
y = (normed * (w.float() + 1.0)).to(x.dtype)
else:
y = w * normed.to(x.dtype)
y.backward(grad)
return x.grad, w.grad
def _fast(
path,
inputs,
weights,
grad,
gemma,
eps = 1e-6,
):
from unsloth.kernels import rms_layernorm
x = inputs.clone().requires_grad_()
w = weights.clone().requires_grad_()
if path == "eager":
y = rms_layernorm.Fast_RMS_Layernorm.apply(x, w, eps, gemma)
else:
if not rms_layernorm._TRACEABLE:
pytest.skip("torch.library.triton_op unavailable")
y, _ = torch.ops.unsloth.rms_layernorm(x.reshape(-1, x.shape[-1]), w, eps, gemma)
y.view(x.shape).backward(grad.clone())
return x.grad, w.grad
# Issue #3014: a trainable norm weight (full finetuning without full_finetuning=True) got no grad.
@pytest.mark.parametrize("path", ["eager", "traced"])
@pytest.mark.parametrize("gemma", [False, True])
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
@pytest.mark.parametrize("rows,cols", [(3, 128), (9000, 128), (3, 384), (9000, 384)])
def test_rmsnorm_trainable_weight_gets_grad(path, gemma, dtype, rows, cols):
torch.manual_seed(3407)
inputs = torch.randn(rows, cols, device = "cuda", dtype = dtype)
weights = torch.rand(cols, device = "cuda", dtype = dtype)
grad = torch.randn(rows, cols, device = "cuda", dtype = dtype)
ref_dx, ref_dw = _reference(inputs, weights, grad, gemma)
dx, dw = _fast(path, inputs, weights, grad, gemma)
assert dw is not None and dw.dtype == weights.dtype and dw.shape == weights.shape
torch.testing.assert_close(dx, ref_dx, rtol = 2e-2, atol = 1e-2)
exact = (
ref_dw.double()
if dtype == torch.float32
else _reference(inputs.float(), weights.float(), grad.float(), gemma)[1].double()
)
cos = torch.nn.functional.cosine_similarity(dw.double(), exact, dim = 0)
assert cos > 0.9999, cos
torch.testing.assert_close(dw.double(), exact, rtol = 1e-2, atol = 1e-2 * exact.abs().max().item())
# The non-Gemma forward computes W * normed.to(W.dtype), so dW must use the rounded normed.
@pytest.mark.parametrize("path", ["eager", "traced"])
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
@pytest.mark.parametrize("cols", [128, 384])
def test_rmsnorm_dw_uses_the_forward_rounding(path, dtype, cols):
torch.manual_seed(3407)
inputs = torch.randn(9000, cols, device = "cuda", dtype = dtype)
weights = torch.rand(cols, device = "cuda", dtype = dtype)
grad = torch.randn(9000, cols, device = "cuda", dtype = dtype)
dw = _fast(path, inputs, weights, grad, False)[1].double()
x = inputs.float()
normed = x * torch.rsqrt(x.square().mean(-1, keepdim = True) + 1e-6)
rounded = (grad.double() * normed.to(dtype).double()).sum(0)
unrounded = (grad.double() * normed.double()).sum(0)
assert (dw - rounded).norm() < (dw - unrounded).norm()
@pytest.mark.parametrize("gemma", [False, True])
@pytest.mark.parametrize("cols", [128, 384])
def test_rmsnorm_trainable_weight_dx_matches_frozen(gemma, cols):
"""The trainable row loop may contract differently (Gemma's (W + 1), and plain rows on Blackwell): at most
a bf16 ulp apart."""
torch.manual_seed(3407)
inputs = torch.randn(777, cols, device = "cuda", dtype = torch.bfloat16)
weights = torch.rand(cols, device = "cuda", dtype = torch.bfloat16)
grad = torch.randn(777, cols, device = "cuda", dtype = torch.bfloat16)
from unsloth.kernels.rms_layernorm import Fast_RMS_Layernorm
frozen = inputs.clone().requires_grad_()
Fast_RMS_Layernorm.apply(frozen, weights, 1e-6, gemma).backward(grad.clone())
dx, _ = _fast("eager", inputs, weights, grad, gemma)
torch.testing.assert_close(dx.float(), frozen.grad.float(), rtol = 2**-7, atol = 1e-6)
@pytest.mark.parametrize("gemma", [False, True])
def test_rmsnorm_weight_grad_eager_matches_traced_bytes(gemma):
torch.manual_seed(3407)
inputs = torch.randn(5000, 256, device = "cuda", dtype = torch.bfloat16)
weights = torch.rand(256, device = "cuda", dtype = torch.bfloat16)
grad = torch.randn(5000, 256, device = "cuda", dtype = torch.bfloat16)
_, eager = _fast("eager", inputs, weights, grad, gemma)
_, traced = _fast("traced", inputs, weights, grad, gemma)
assert torch.equal(eager.view(torch.int16), traced.view(torch.int16))
def test_rmsnorm_frozen_weight_stays_gradless():
from unsloth.kernels.rms_layernorm import Fast_RMS_Layernorm
x = torch.randn(4, 64, device = "cuda", requires_grad = True)
w = torch.rand(64, device = "cuda")
Fast_RMS_Layernorm.apply(x, w, 1e-6, False).sum().backward()
assert x.grad is not None and w.grad is None
@pytest.mark.parametrize("cols", [128, 384])
def test_rmsnorm_trainable_weight_empty_input(cols):
from unsloth.kernels.rms_layernorm import Fast_RMS_Layernorm
x = torch.randn(0, cols, device = "cuda", dtype = torch.bfloat16, requires_grad = True)
w = torch.rand(cols, device = "cuda", dtype = torch.bfloat16, requires_grad = True)
Fast_RMS_Layernorm.apply(x, w, 1e-6, False).sum().backward()
assert x.grad.shape == (0, cols)
assert torch.equal(w.grad, torch.zeros_like(w))

339
unsloth/kernels/rms_layernorm.py
View file

@ -9,6 +9,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import functools
import hashlib
import os
import warnings
@ -63,7 +64,8 @@ def _rms_layernorm_forward(
tl.store(Y + col_offsets, output, mask = mask)
def _rms_layernorm_backward(
@triton.jit
def _rms_backward_row(
dY,
dY_row_stride: tl.constexpr,
dX,
@ -71,26 +73,16 @@ def _rms_layernorm_backward(
X,
X_row_stride: tl.constexpr,
W,
W_row_stride: tl.constexpr,
r,
r_row_stride: tl.constexpr,
row_idx,
col_offsets,
mask,
n_cols: tl.constexpr,
eps: tl.constexpr,
GEMMA: tl.constexpr,
LONG_INDEXING: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
dW_acc,
TRAIN_W: tl.constexpr,
):
"""
Fast RMS Layernorm kernel for the backward pass
Inspiration from a Triton tutorial:
https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html
"""
row_idx = tl.program_id(0)
if LONG_INDEXING:
row_idx = row_idx.to(tl.int64)
col_offsets = tl.arange(0, BLOCK_SIZE)
mask = col_offsets < n_cols
dY += row_idx * dY_row_stride
X += row_idx * X_row_stride
r += row_idx * r_row_stride
@ -107,6 +99,13 @@ def _rms_layernorm_backward(
# Get saved row variance
inv_var = tl.load(r).to(tl.float32)
normed = X_row * inv_var
if TRAIN_W:
# dW before the row sum so the fp32 dY row is not held across it (fewer registers, all programs resident).
if GEMMA:
dW_acc += dY_row * normed
else:
# The forward multiplies W by normed rounded to W's dtype.
dW_acc += dY_row * normed.to(W.dtype.element_ty).to(tl.float32)
if GEMMA:
dY_W = dY_row * (W_row + 1.0)
@ -116,6 +115,84 @@ def _rms_layernorm_backward(
rowsum_dY_normed = tl.sum(dY_W * normed, axis = 0)
output = inv_var / n_cols * (n_cols * dY_W - normed * rowsum_dY_normed)
tl.store(dX + col_offsets, output, mask = mask)
return dW_acc
def _rms_layernorm_backward(
dY,
dY_row_stride: tl.constexpr,
dX,
dX_row_stride: tl.constexpr,
X,
X_row_stride: tl.constexpr,
W,
W_row_stride: tl.constexpr,
r,
r_row_stride: tl.constexpr,
dW,
n_rows,
rows_per_program,
n_cols: tl.constexpr,
eps: tl.constexpr,
GEMMA: tl.constexpr,
LONG_INDEXING: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
TRAIN_W: tl.constexpr,
):
"""
Fast RMS Layernorm kernel for the backward pass
Inspiration from a Triton tutorial:
https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html
TRAIN_W: per-program fp32 dW partials (fla-org/flash-linear-attention layer_norm_bwd_kernel_row).
"""
pid = tl.program_id(0)
if LONG_INDEXING:
pid = pid.to(tl.int64)
col_offsets = tl.arange(0, BLOCK_SIZE)
mask = col_offsets < n_cols
if TRAIN_W:
dW_acc = tl.zeros((BLOCK_SIZE,), dtype = tl.float32)
row_start = pid * rows_per_program
for row_idx in range(row_start, tl.minimum(row_start + rows_per_program, n_rows)):
dW_acc = _rms_backward_row(
dY,
dY_row_stride,
dX,
dX_row_stride,
X,
X_row_stride,
W,
r,
r_row_stride,
row_idx,
col_offsets,
mask,
n_cols,
GEMMA,
dW_acc,
True,
)
tl.store(dW + pid * n_cols + col_offsets, dW_acc, mask = mask)
else:
_rms_backward_row(
dY,
dY_row_stride,
dX,
dX_row_stride,
X,
X_row_stride,
W,
r,
r_row_stride,
pid,
col_offsets,
mask,
n_cols,
GEMMA,
0.0,
False,
)
_rms_layernorm_backward = triton.jit(_rms_layernorm_backward)
@ -123,6 +200,7 @@ _rms_layernorm_backward = triton.heuristics(
{
"GEMMA": lambda args: bool(args["GEMMA"]),
"LONG_INDEXING": lambda args: bool(args["LONG_INDEXING"]),
"TRAIN_W": lambda args: bool(args["TRAIN_W"]),
}
)(_rms_layernorm_backward)
@ -240,7 +318,7 @@ def _rms_layernorm_forward_rows(
@triton.jit
def _rms_layernorm_backward_rows(
def _rms_backward_rows_block(
dY,
dY_row_stride: tl.constexpr,
dX,
@ -249,19 +327,14 @@ def _rms_layernorm_backward_rows(
X_row_stride: tl.constexpr,
W,
r,
rows,
n_rows,
n_cols: tl.constexpr,
eps: tl.constexpr,
GEMMA: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
ROWS: tl.constexpr,
WARPS: tl.constexpr,
LONG_INDEXING: tl.constexpr,
):
pid = tl.program_id(0)
if LONG_INDEXING:
pid = pid.to(tl.int64)
rows = pid * ROWS + tl.arange(0, ROWS)
col_offsets = tl.arange(0, BLOCK_SIZE)
row_mask = rows < n_rows
col_mask = col_offsets < n_cols
@ -290,6 +363,95 @@ def _rms_layernorm_backward_rows(
tl.store(dX + rows[:, None] * dX_row_stride + col_offsets[None, :], output, mask = mask)
else:
tl.store(dY + rows[:, None] * dY_row_stride + col_offsets[None, :], output, mask = mask)
# Masked rows load zeros, so they add nothing; dead code unless TRAIN_W.
if not GEMMA:
# The forward multiplies W by normed rounded to W's dtype.
normed = normed.to(W.dtype.element_ty).to(tl.float32)
return tl.sum(dY_rows * normed, axis = 0)
@triton.jit
def _rms_layernorm_backward_rows(
dY,
dY_row_stride: tl.constexpr,
dX,
dX_row_stride: tl.constexpr,
X,
X_row_stride: tl.constexpr,
W,
r,
dW,
n_rows,
blocks_per_program,
n_cols: tl.constexpr,
eps: tl.constexpr,
GEMMA: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
ROWS: tl.constexpr,
WARPS: tl.constexpr,
LONG_INDEXING: tl.constexpr,
TRAIN_W: tl.constexpr,
):
pid = tl.program_id(0)
if LONG_INDEXING:
pid = pid.to(tl.int64)
if TRAIN_W:
dW_acc = tl.zeros((BLOCK_SIZE,), dtype = tl.float32)
block_start = pid * blocks_per_program
block_end = tl.minimum(block_start + blocks_per_program, tl.cdiv(n_rows, ROWS))
for block in range(block_start, block_end):
dW_acc += _rms_backward_rows_block(
dY,
dY_row_stride,
dX,
dX_row_stride,
X,
X_row_stride,
W,
r,
block * ROWS + tl.arange(0, ROWS),
n_rows,
n_cols,
GEMMA,
BLOCK_SIZE,
ROWS,
WARPS,
)
col_offsets = tl.arange(0, BLOCK_SIZE)
tl.store(dW + pid * n_cols + col_offsets, dW_acc, mask = col_offsets < n_cols)
else:
_rms_backward_rows_block(
dY,
dY_row_stride,
dX,
dX_row_stride,
X,
X_row_stride,
W,
r,
pid * ROWS + tl.arange(0, ROWS),
n_rows,
n_cols,
GEMMA,
BLOCK_SIZE,
ROWS,
WARPS,
)
@triton.jit
def _rms_sum_dw_partials(
partials, dW, n_programs, n_cols, BLOCK_ROWS: tl.constexpr, BLOCK_COLS: tl.constexpr
):
# Fixed tile order, so eager and torch.compile give the same bytes.
cols = tl.program_id(0) * BLOCK_COLS + tl.arange(0, BLOCK_COLS)
col_mask = cols < n_cols
acc = tl.zeros((BLOCK_ROWS, BLOCK_COLS), dtype = tl.float32)
for start in tl.range(0, n_programs, BLOCK_ROWS, num_stages = 3):
rows = start + tl.arange(0, BLOCK_ROWS)
mask = (rows < n_programs)[:, None] & col_mask[None, :]
acc += tl.load(partials + rows[:, None] * n_cols + cols[None, :], mask = mask, other = 0)
tl.store(dW + cols, tl.sum(acc, axis = 0).to(dW.dtype.element_ty), mask = col_mask)
# Off on ROCm (64-lane wavefronts reduce in another order) and before Triton 3.6 (slower backward).
@ -454,6 +616,50 @@ def _rms_forward(
return Y, r
_DW_PROGRAMS_PER_SM = 4
@functools.lru_cache(maxsize = None)
def _sm_count(device):
try:
return int(torch.cuda.get_device_properties(device).multi_processor_count)
except Exception:
return 128
@functools.lru_cache(maxsize = None)
def _max_sm_threads(device):
try:
return int(torch.cuda.get_device_properties(device).max_threads_per_multi_processor)
except Exception:
return 0
def _dw_split(units, max_programs):
per_program = max(1, triton.cdiv(units, max(1, min(units, max_programs))))
return max(1, triton.cdiv(units, per_program)), per_program
def _dw_launch(units, device):
return _dw_split(units, _sm_count(device) * _DW_PROGRAMS_PER_SM)
@functools.lru_cache(maxsize = None)
def _dw_row_config(device, n_cols):
# Cached (backward is host bound at small shapes). 64-warp SMs, <= 16 elems/lane: cap registers at 64
# so all programs stay resident; uncapped, wide rows kept one program per SM. T4: no measured gain.
BLOCK_SIZE, num_warps = calculate_settings(n_cols)
per_sm, launch = _DW_PROGRAMS_PER_SM, {}
if (
device.type == "cuda"
and not torch.version.hip
and BLOCK_SIZE <= 512 * num_warps
and _max_sm_threads(device) >= 2048
):
per_sm, launch = max(1, 32 // num_warps), {"maxnreg": 64}
return BLOCK_SIZE, num_warps, _sm_count(device) * per_sm, launch
def _rms_backward(
dY,
dX,
@ -464,16 +670,22 @@ def _rms_backward(
gemma,
wrap,
multirow = None,
train_w = False,
):
# Non-Gemma writes dX over dY (the kernel ignores dX); Gemma writes into dX.
n_rows, n_cols = dY.shape
if multirow is None:
multirow = _multirow_checked(X, W, eps, gemma)
long = long_indexing(dY, dX, X, block = _INDEX_MARGIN)
partials = None
if multirow:
BLOCK_SIZE, ROWS, WARPS, num_warps = multirow
programs, per_program = triton.cdiv(n_rows, ROWS), 1
if train_w:
programs, per_program = _dw_launch(programs, dY.device)
partials = torch.empty((programs, n_cols), dtype = torch.float32, device = dY.device)
try:
wrap(_rms_layernorm_backward_rows)[((n_rows + ROWS - 1) // ROWS,)](
wrap(_rms_layernorm_backward_rows)[(programs,)](
dY,
dY.stride(0),
dX,
@ -482,7 +694,9 @@ def _rms_backward(
X.stride(0),
W,
r,
r if partials is None else partials,
n_rows,
per_program,
n_cols,
eps,
GEMMA = gemma,
@ -490,15 +704,23 @@ def _rms_backward(
ROWS = ROWS,
WARPS = WARPS,
LONG_INDEXING = long,
TRAIN_W = train_w,
num_warps = num_warps,
)
return
except Exception as error:
if _multirow_must_raise(error, wrap):
raise
_multirow_disable(f"launch failed: {error!r}")
BLOCK_SIZE, num_warps = calculate_settings(n_cols)
wrap(_rms_layernorm_backward)[(n_rows,)](
else:
return _rms_sum_partials(partials, W, wrap)
programs, per_program, launch = n_rows, 1, {}
if train_w:
BLOCK_SIZE, num_warps, max_programs, launch = _dw_row_config(dY.device, int(n_cols))
programs, per_program = _dw_split(n_rows, max_programs)
partials = torch.empty((programs, n_cols), dtype = torch.float32, device = dY.device)
else:
BLOCK_SIZE, num_warps = calculate_settings(n_cols)
wrap(_rms_layernorm_backward)[(programs,)](
dY,
dY.stride(0),
dX,
@ -509,13 +731,30 @@ def _rms_backward(
W.stride(0),
r,
r.stride(0),
r if partials is None else partials,
n_rows,
per_program,
n_cols,
eps,
GEMMA = gemma,
LONG_INDEXING = long,
BLOCK_SIZE = BLOCK_SIZE,
TRAIN_W = train_w,
num_warps = num_warps,
**launch,
)
return _rms_sum_partials(partials, W, wrap)
def _rms_sum_partials(partials, W, wrap):
if partials is None:
return None
programs, n_cols = partials.shape
dW = torch.empty(n_cols, dtype = W.dtype, device = partials.device)
wrap(_rms_sum_dw_partials)[(triton.cdiv(n_cols, 64),)](
partials, dW, programs, n_cols, BLOCK_ROWS = 64, BLOCK_COLS = 64
)
return dW
_eager_kernel = lambda kernel: kernel
@ -550,9 +789,19 @@ class Fast_RMS_Layernorm(torch.autograd.Function):
X, W, r = ctx.saved_tensors
dX = torch.empty_like(dY) if ctx.GEMMA else dY
with torch_gpu_device(dY.device):
_rms_backward(dY, dX, X, W, r, ctx.eps, ctx.GEMMA, _eager_kernel)
dW = _rms_backward(
dY,
dX,
X,
W,
r,
ctx.eps,
ctx.GEMMA,
_eager_kernel,
train_w = ctx.needs_input_grad[1],
)
dX = dX.view(*shape)
return dX, None, None, None
return dX, dW, None, None
# Compiled graphs take these ops: same kernels, fresh outputs (the Function writes dX over dY).
@ -604,13 +853,30 @@ if _TRACEABLE:
@torch.library.triton_op("unsloth::rms_layernorm_backward", mutates_args = ())
def _rms_layernorm_backward_op(
dY: torch.Tensor, X: torch.Tensor, W: torch.Tensor, r: torch.Tensor, eps: float, gemma: bool
) -> torch.Tensor:
dY: torch.Tensor,
X: torch.Tensor,
W: torch.Tensor,
r: torch.Tensor,
eps: float,
gemma: bool,
train_w: bool,
) -> Tuple[torch.Tensor, torch.Tensor]:
dX = torch.empty_like(dY)
# Non-Gemma rewrites its dY argument, so it gets a copy.
with torch_gpu_device(dY.device):
_rms_backward(dY if gemma else dX.copy_(dY), dX, X, W, r, eps, gemma, _traced_kernel)
return dX
dW = _rms_backward(
dY if gemma else dX.copy_(dY),
dX,
X,
W,
r,
eps,
gemma,
_traced_kernel,
train_w = train_w,
)
# Empty dW = frozen (ops return tensors only).
return dX, dW if train_w else W.new_empty(0)
def _rms_setup_context(ctx, inputs, output):
X, W, eps, gemma = inputs
@ -619,8 +885,11 @@ if _TRACEABLE:
def _rms_layernorm_op_backward(ctx, dY, dr):
X, W, r = ctx.saved_tensors
dX = torch.ops.unsloth.rms_layernorm_backward(dY.contiguous(), X, W, r, ctx.eps, ctx.gemma)
return dX, None, None, None
train_w = ctx.needs_input_grad[1]
dX, dW = torch.ops.unsloth.rms_layernorm_backward(
dY.contiguous(), X, W, r, ctx.eps, ctx.gemma, train_w
)
return dX, dW if train_w else None, None, None
_rms_layernorm_op.register_autograd(
_rms_layernorm_op_backward, setup_context = _rms_setup_context