kda in torchtitan (training)

.ci/docker/requirements-vlm.txt

@@ -3,3 +3,4 @@ einops
 pillow
 torchvision
 flash-linear-attention
+attention-gym

pyproject.toml

@@ -67,5 +67,5 @@ testpaths = ["tests"]
 
 [tool.pyrefly]
 project-excludes = ["torchtitan/experiments", "**/tests/**"]
-replace-imports-with-any = ["torchao.*", "torchft", "torchvision.*", "deep_ep.*", "jinja2.*", "fla.*", "helion", "helion.*", "batch_invariant_ops", "torchcomms"]  # optional dependencies
+replace-imports-with-any = ["torchao.*", "torchft", "torchvision.*", "deep_ep.*", "jinja2.*", "fla.*", "attn_gym.*", "helion", "helion.*", "batch_invariant_ops", "torchcomms"]  # optional dependencies
 search-path = ["../pytorch"]  # local built pytorch

tests/unit_tests/test_kda_attention.py

@@ -0,0 +1,119 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Unit tests for the KDA linear-attention layer."""
+
+import importlib.util
+import unittest
+
+import torch
+
+from torchtitan.models.common.attention import create_varlen_metadata_for_document
+
+_HAS_BLACKWELL = (
+    importlib.util.find_spec("attn_gym") is not None
+    and torch.cuda.is_available()
+    and torch.cuda.get_device_capability() >= (10, 0)
+)
+
+
+@unittest.skipUnless(
+    _HAS_BLACKWELL, "KDA requires attention-gym and CUDA capability 10.0 or newer"
+)
+class TestKDA(unittest.TestCase):
+    def _make_kda(self, *, backend: str = "chunked"):
+        """Build a KDA layer with deterministic weights."""
+        from torchtitan.models.common import Conv1d, Linear, RMSNorm
+        from torchtitan.models.common.attention import (
+            KDA,
+            KDAAttention,
+            KDAInnerAttention,
+        )
+
+        def linear(in_features: int, out_features: int) -> Linear.Config:
+            return Linear.Config(
+                in_features=in_features, out_features=out_features, bias=False
+            )
+
+        model = KDA.Config(
+            num_heads=2,
+            head_dim=128,
+            in_proj_qkv=linear(32, 768),
+            conv_qkv=Conv1d.Config(
+                in_channels=768,
+                out_channels=768,
+                kernel_size=4,
+                groups=768,
+                bias=False,
+            ),
+            gate_proj_a=linear(32, 128),
+            gate_proj_b=linear(128, 256),
+            beta_proj=linear(32, 2),
+            out_gate_proj_a=linear(32, 128),
+            out_gate_proj_b=linear(128, 256),
+            out_norm=RMSNorm.Config(normalized_shape=128),
+            out_proj=linear(256, 32),
+            attention=KDAAttention.Config(
+                head_dim=128,
+                inner_attention=KDAInnerAttention.Config(backend=backend),
+            ),
+        ).build()
+        model = model.to(device="cuda", dtype=torch.bfloat16)
+        with torch.no_grad():
+            for param in model.parameters():
+                values = torch.linspace(
+                    -0.2, 0.2, param.numel(), dtype=param.dtype, device=param.device
+                )
+                param.copy_(values.reshape_as(param))
+            model.A_log.fill_(0.0)
+            model.dt_bias.zero_()
+            model.out_norm.weight.fill_(1.0)
+        return model
+
+    def _inputs(self, seed: int, tokens: int = 128) -> torch.Tensor:
+        torch.manual_seed(seed)
+        return torch.randn(1, tokens, 32, device="cuda", dtype=torch.bfloat16)
+
+    def test_chunked_and_recurrent_backends_agree(self):
+        chunked = self._make_kda(backend="chunked")
+        recurrent = self._make_kda(backend="recurrent")
+        recurrent.load_state_dict(chunked.state_dict())
+
+        x_BLD = self._inputs(seed=0)
+        torch.testing.assert_close(
+            chunked(x_BLD).float(), recurrent(x_BLD).float(), rtol=5e-2, atol=5e-2
+        )
+
+    def test_varlen_matches_independent_document_forwards(self):
+        lengths = (37, 64, 91)
+        x_BLD = self._inputs(seed=2, tokens=sum(lengths))
+        positions = torch.tensor(
+            [[index for length in lengths for index in range(length)]],
+            device="cuda",
+            dtype=torch.int32,
+        )
+        masks = create_varlen_metadata_for_document(
+            positions, include_host_offsets=True
+        )
+
+        for backend in ("chunked", "recurrent"):
+            model = self._make_kda(backend=backend)
+            packed = model(x_BLD, masks)
+            start = 0
+            for document, length in enumerate(lengths):
+                with self.subTest(backend=backend, document=document):
+                    document_slice = slice(start, start + length)
+                    torch.testing.assert_close(
+                        packed[:, document_slice].float(),
+                        model(x_BLD[:, document_slice], None).float(),
+                        rtol=2e-2,
+                        atol=2e-2,
+                    )
+                start += length
+
+
+if __name__ == "__main__":
+    unittest.main()

torchtitan/models/common/__init__.py

@@ -16,12 +16,17 @@
     get_fixed_block_mask_mod,
     get_sliding_window_mask_mod,
     GQAttention,
+    KDA,
+    KDAAttention,
+    KDABackend,
+    KDAInnerAttention,
     QKVLinear,
     ScaledDotProductAttention,
     VarlenAttention,
     VarlenMetadata,
 )
 from .decoder import Decoder, TransformerBlock
+from .decoder_sharding import set_kda_sharding
 from .embedding import Embedding
 from .feed_forward import compute_ffn_hidden_dim, FeedForward, SigmoidGatedFeedForward
 from .linear import Linear, ScaledBiasRowwiseLinear
@@ -61,6 +66,10 @@
     "GQAttention",
     "GroupNorm",
     "Identity",
+    "KDA",
+    "KDAAttention",
+    "KDABackend",
+    "KDAInnerAttention",
     "LayerNorm",
     "Linear",
     "MoE",
@@ -69,6 +78,7 @@
     "RoPE",
     "ScaledBiasRowwiseLinear",
     "ScaledDotProductAttention",
+    "set_kda_sharding",
     "SiLU",
     "TransformerBlock",
     "VarlenAttention",

torchtitan/models/common/attention/__init__.py

@@ -0,0 +1,52 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from .attention import (
+    AttentionMasksType,
+    BaseAttention,
+    BaseQKVLinear,
+    create_attention_mask,
+    create_varlen_metadata_for_document,
+    FlexAttention,
+    FusedQKVLinear,
+    get_causal_mask_mod,
+    get_document_mask_mod,
+    get_efficient_causal_mask_mod_for_packed_document,
+    get_fixed_block_mask_mod,
+    get_sliding_window_mask_mod,
+    GQAttention,
+    local_head_split,
+    QKVLinear,
+    ScaledDotProductAttention,
+    VarlenAttention,
+    VarlenMetadata,
+)
+from .kda import KDA, KDAAttention, KDABackend, KDAInnerAttention
+
+__all__ = [
+    "AttentionMasksType",
+    "BaseAttention",
+    "BaseQKVLinear",
+    "create_attention_mask",
+    "create_varlen_metadata_for_document",
+    "FlexAttention",
+    "FusedQKVLinear",
+    "get_causal_mask_mod",
+    "get_document_mask_mod",
+    "get_efficient_causal_mask_mod_for_packed_document",
+    "get_fixed_block_mask_mod",
+    "get_sliding_window_mask_mod",
+    "GQAttention",
+    "KDA",
+    "KDAAttention",
+    "KDABackend",
+    "KDAInnerAttention",
+    "local_head_split",
+    "QKVLinear",
+    "ScaledDotProductAttention",
+    "VarlenAttention",
+    "VarlenMetadata",
+]