diff --git a/src/liger_kernel/transformers/__init__.py b/src/liger_kernel/transformers/__init__.py index 26bdef91b..9235ae1da 100644 --- a/src/liger_kernel/transformers/__init__.py +++ b/src/liger_kernel/transformers/__init__.py @@ -3,6 +3,7 @@ from typing import TYPE_CHECKING # Always-safe imports (independent of 'transformers') +from liger_kernel.transformers.attn_res import LigerAttnRes # noqa: F401 from liger_kernel.transformers.cross_entropy import LigerCrossEntropyLoss # noqa: F401 from liger_kernel.transformers.dyt import LigerDyT # noqa: F401 from liger_kernel.transformers.fused_add_rms_norm import LigerFusedAddRMSNorm # noqa: F401 @@ -172,6 +173,7 @@ def __getattr__(name: str): # Shared symbols in all environments __all__ = [ "is_transformers_available", + "LigerAttnRes", "LigerCrossEntropyLoss", "LigerDyT", "LigerFusedLinearCrossEntropyLoss", diff --git a/src/liger_kernel/transformers/attn_res.py b/src/liger_kernel/transformers/attn_res.py new file mode 100644 index 000000000..ebff58a35 --- /dev/null +++ b/src/liger_kernel/transformers/attn_res.py @@ -0,0 +1,43 @@ +import torch +import torch.nn as nn + +from liger_kernel.ops import LigerAttnResFunction + + +class LigerAttnRes(nn.Module): + """Attention Residuals (AttnRes) from Kimi/Moonshot AI (arXiv:2603.15031). + + Replaces the standard residual connection ``h = h_prev + f(RMSNorm(h_prev))`` + with a softmax attention over the depth (block) dimension: the stacked outputs + of ``N`` blocks are each RMSNorm'd, scored against a learned pseudo-query, and + the resulting per-block softmax weights produce a weighted sum. This is the + module wrapper around :func:`~liger_kernel.transformers.functional.liger_attn_res`. + + Args: + hidden_size: hidden dimension ``D`` of each block output. + eps: epsilon for the per-block RMSNorm (default: 1e-6). + query_init_std: standard deviation of the normal initialization for the + learned pseudo-query ``w_query`` (default: 0.02). ``w_norm`` (the + per-block RMSNorm weight) is initialized to ones. + + Shape: + - Input: ``[N, B, T, D]`` stacked block outputs, or a list of ``N`` + tensors each of shape ``[B, T, D]``. + - Output: ``[B, T, D]``. + """ + + def __init__(self, hidden_size: int, eps: float = 1e-6, query_init_std: float = 0.02): + super().__init__() + self.hidden_size = hidden_size + self.eps = eps + self.query_init_std = query_init_std + self.w_query = nn.Parameter(torch.randn(hidden_size) * query_init_std) + self.w_norm = nn.Parameter(torch.ones(hidden_size)) + + def forward(self, V): + if isinstance(V, (list, tuple)): + V = torch.stack(V) + return LigerAttnResFunction.apply(V, self.w_query, self.w_norm, self.eps) + + def extra_repr(self): + return f"hidden_size={self.hidden_size}, eps={self.eps}" diff --git a/test/transformers/test_attn_res.py b/test/transformers/test_attn_res.py index 064ccdd9e..0cf5ab440 100644 --- a/test/transformers/test_attn_res.py +++ b/test/transformers/test_attn_res.py @@ -8,6 +8,7 @@ from test.utils import supports_bfloat16 from liger_kernel.ops import LigerAttnResFunction +from liger_kernel.transformers.attn_res import LigerAttnRes from liger_kernel.transformers.functional import liger_attn_res from liger_kernel.utils import infer_device @@ -185,3 +186,99 @@ def test_correctness_functional(N, B, T, D, dtype, atol, rtol): y2.backward(grad) assert_verbose_allclose(V1.grad, V2.grad, atol=atol, rtol=rtol) + + +@pytest.mark.flaky(reruns=3, reruns_delay=2) +@pytest.mark.parametrize( + "N, B, T, D", + [ + (4, 2, 64, 512), + (8, 2, 32, 256), + # weird shapes + (3, 5, 37, 123), + ], +) +@pytest.mark.parametrize( + "dtype, atol, rtol", + [ + (torch.float32, 1e-4, 1e-5), + (torch.float16, 1e-2, 1e-3), + pytest.param( + torch.bfloat16, + 1e-1, + 1e-2, + marks=pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU"), + ), + ], +) +def test_module_matches_reference(N, B, T, D, dtype, atol, rtol): + """LigerAttnRes module matches the PyTorch reference (fwd) and trains its params (bwd).""" + set_seed(0) + model = LigerAttnRes(hidden_size=D, eps=1e-6).to(device).to(dtype) + + V = torch.randn(N, B, T, D, device=device, dtype=dtype) + V_in = V.clone().requires_grad_(True) + out = model(V_in) + ref = pytorch_attn_res(V, model.w_query.detach(), model.w_norm.detach(), eps=1e-6) + assert out.shape == (B, T, D) + assert_verbose_allclose(out, ref, atol=atol, rtol=rtol) + + out.backward(torch.randn_like(out)) + assert V_in.grad is not None and torch.isfinite(V_in.grad).all() + for name, p in model.named_parameters(): + assert p.grad is not None and torch.isfinite(p.grad).all(), f"no/invalid grad for {name}" + + +@pytest.mark.flaky(reruns=3, reruns_delay=2) +@pytest.mark.parametrize( + "N, B, T, D", + [ + (4, 2, 64, 512), + (3, 5, 37, 123), + ], +) +def test_module_param_gradients_match_reference(N, B, T, D): + """The learned params (w_query, w_norm) must receive gradients that match the + PyTorch reference — the whole point of the module is to train them. + + Checked in fp32 only: w_query/w_norm gradients are a sum-reduction over all + ``N*B*T`` tokens, and the kernel accumulates that reduction in fp32 while a + same-dtype PyTorch reference does not, so in fp16/bf16 the reference is the + less-accurate baseline (verified against an fp64 ground truth) and an + element-wise comparison would test reduction noise, not correctness. The + low-precision forward and input-grad paths are covered by the tests above. + """ + set_seed(0) + atol, rtol = 1e-3, 1e-4 + model = LigerAttnRes(hidden_size=D, eps=1e-6).to(device) + + V = torch.randn(N, B, T, D, device=device) + do = torch.randn(B, T, D, device=device) + + # Reference: same parameter values, plain PyTorch autograd. + V_ref = V.clone().requires_grad_(True) + wq_ref = model.w_query.detach().clone().requires_grad_(True) + wn_ref = model.w_norm.detach().clone().requires_grad_(True) + pytorch_attn_res(V_ref, wq_ref, wn_ref, eps=1e-6).backward(do) + + # Module (Triton kernel) path. + V_mod = V.clone().requires_grad_(True) + model(V_mod).backward(do) + + assert_verbose_allclose(V_mod.grad, V_ref.grad, atol=atol, rtol=rtol) + assert_verbose_allclose(model.w_query.grad, wq_ref.grad, atol=atol, rtol=rtol) + assert_verbose_allclose(model.w_norm.grad, wn_ref.grad, atol=atol, rtol=rtol) + + +def test_module_list_input_and_repr(): + """Module accepts a list of blocks (equivalent to the stacked tensor) and reprs its config.""" + set_seed(0) + N, B, T, D = 4, 2, 16, 64 + model = LigerAttnRes(hidden_size=D).to(device) + blocks = [torch.randn(B, T, D, device=device) for _ in range(N)] + + out_list = model(blocks) + out_stacked = model(torch.stack(blocks)) + assert out_list.shape == (B, T, D) + assert_verbose_allclose(out_list, out_stacked, atol=1e-6, rtol=1e-6) + assert f"hidden_size={D}" in model.extra_repr()