kda in torchtitan (training)

.ci/docker/requirements.txt

@@ -8,3 +8,4 @@ safetensors
 einops
 pillow
 spmd_types==0.2.5
+attn-gym[linear]==0.0.5

pyproject.toml

@@ -28,6 +28,7 @@ dependencies = [
     "einops",
     "pillow",
     "spmd_types==0.2.5",
+    "attn-gym[linear]==0.0.5",
 ]
 dynamic = ["version"]
 

tests/integration_tests/__init__.py

@@ -35,6 +35,8 @@ class OverrideDefinitions:
     ngpu: int = 4
     disabled: bool = False
     skip_rocm_test: bool = False
+    required_cuda_capabilities: Sequence[tuple[int, int]] = ()
+    """CUDA compute capabilities on which the test can run."""
     timeout: int | None = None
     golden_numerics_path: str | None = None
     """Run through loss_compare.py using this mode-specific golden path."""

tests/integration_tests/models.py

@@ -206,5 +206,6 @@ def build_model_tests_list() -> list[OverrideDefinitions]:
             test_descr="Kimi K3 multimodal FSDP",
             test_name="kimi_k3_mm_fsdp",
             ngpu=2,
+            required_cuda_capabilities=((10, 0), (10, 3)),
         ),
     ]

tests/integration_tests/run_tests.py

@@ -14,6 +14,8 @@
 from concurrent.futures import Future, ThreadPoolExecutor
 from pathlib import Path
 
+import torch
+
 from torchtitan.tools.logging import logger
 from torchtitan.trainer import Trainer
 
@@ -332,17 +334,28 @@ def run_single_test(
 
 def _filter_tests(
     args, test_list: list[OverrideDefinitions]
-) -> tuple[list[OverrideDefinitions], list[OverrideDefinitions]]:
+) -> tuple[
+    list[OverrideDefinitions],
+    list[OverrideDefinitions],
+    list[OverrideDefinitions],
+]:
     """Filter tests by name, scope, disabled state, architecture, and GPU count.
 
-    Returns (runnable, skipped_due_to_ngpu).
+    Returns (runnable, skipped_due_to_ngpu, skipped_due_to_cuda_capability).
     """
     exclude_set = set()
     if hasattr(args, "exclude") and args.exclude:
         exclude_set = {name.strip() for name in args.exclude.split(",")}
 
     runnable: list[OverrideDefinitions] = []
     skipped_ngpu: list[OverrideDefinitions] = []
+    skipped_cuda_capability: list[OverrideDefinitions] = []
+    cuda_capability = (
+        torch.cuda.get_device_capability()
+        if getattr(args, "gpu_arch_type", "cuda") == "cuda"
+        and torch.cuda.is_available()
+        else None
+    )
     for test_flavor in test_list:
         if args.test_name != "all" and test_flavor.test_name != args.test_name:
             continue
@@ -361,11 +374,17 @@ def _filter_tests(
             and test_flavor.skip_rocm_test
         ):
             continue
+        if (
+            test_flavor.required_cuda_capabilities
+            and cuda_capability not in test_flavor.required_cuda_capabilities
+        ):
+            skipped_cuda_capability.append(test_flavor)
+            continue
         if execution_mode != "fake_pg" and args.ngpu < test_flavor.ngpu:
             skipped_ngpu.append(test_flavor)
             continue
         runnable.append(test_flavor)
-    return runnable, skipped_ngpu
+    return runnable, skipped_ngpu, skipped_cuda_capability
 
 
 def run_tests(
@@ -374,12 +393,18 @@ def run_tests(
     parallel: bool = True,
 ):
     """Run all integration tests to test the core features of TorchTitan."""
-    runnable, skipped_ngpu = _filter_tests(args, test_list)
+    runnable, skipped_ngpu, skipped_cuda_capability = _filter_tests(args, test_list)
     for test_flavor in skipped_ngpu:
         logger.info(
             f"Skipping test {test_flavor.test_name} that requires {test_flavor.ngpu} gpus,"
             f" because --ngpu arg is {args.ngpu}"
         )
+    for test_flavor in skipped_cuda_capability:
+        logger.info(
+            f"Skipping test {test_flavor.test_name} because its required CUDA "
+            "capability is unavailable; supported capabilities are "
+            f"{tuple(test_flavor.required_cuda_capabilities)}"
+        )
 
     failed_tests: list[tuple[str, str]] = []
     execution_mode = getattr(args, "execution_mode", "real_pg")

tests/unit_tests/cpu/test_integration_test_definitions.py

@@ -18,7 +18,11 @@
 from tests.integration_tests.flux import build_flux_test_list
 from tests.integration_tests.h100 import build_h100_tests_list
 from tests.integration_tests.models import build_model_tests_list
-from tests.integration_tests.run_tests import _parse_test_suites, run_single_test
+from tests.integration_tests.run_tests import (
+    _filter_tests,
+    _parse_test_suites,
+    run_single_test,
+)
 
 
 def test_hf_checkpoint_load_path_comes_from_test_config(monkeypatch) -> None:
@@ -72,6 +76,41 @@ def test_parse_multiple_integration_test_suites() -> None:
     )
 
 
+def test_filter_tests_skips_unsupported_cuda_capability(monkeypatch) -> None:
+    monkeypatch.setattr(
+        "tests.integration_tests.run_tests.torch.cuda.is_available", lambda: True
+    )
+    monkeypatch.setattr(
+        "tests.integration_tests.run_tests.torch.cuda.get_device_capability",
+        lambda: (8, 6),
+    )
+    supported = OverrideDefinitions(test_name="supported")
+    blackwell_only = OverrideDefinitions(
+        test_name="blackwell_only",
+        required_cuda_capabilities=((10, 0), (10, 3)),
+    )
+    args = type(
+        "Args",
+        (),
+        {
+            "test_name": "all",
+            "execution_mode": "real_pg",
+            "test_scope": "all",
+            "gpu_arch_type": "cuda",
+            "ngpu": 8,
+            "exclude": None,
+        },
+    )()
+
+    runnable, skipped_ngpu, skipped_cuda_capability = _filter_tests(
+        args, [supported, blackwell_only]
+    )
+
+    assert runnable == [supported]
+    assert not skipped_ngpu
+    assert skipped_cuda_capability == [blackwell_only]
+
+
 def test_h100_tests_are_registered_in_separate_suite() -> None:
     assert {test.test_name for test in build_h100_tests_list()} == {
         "2d_asynctp_compile",

tests/unit_tests/test_kda_attention.py

@@ -0,0 +1,133 @@
+# 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 import Conv1d, Linear
+from torchtitan.models.common.attention import create_varlen_metadata_for_document
+from torchtitan.models.kimi_k3.kda import InnerKDA, KDA, KDAKernel, KimiRMSNormGated
+
+_HAS_BLACKWELL = (
+    importlib.util.find_spec("attn_gym") is not None
+    and torch.cuda.is_available()
+    and torch.cuda.get_device_capability() in {(10, 0), (10, 3)}
+)
+
+
+def _kda_config() -> KDA.Config:
+    def linear(in_features: int, out_features: int) -> Linear.Config:
+        return Linear.Config(
+            in_features=in_features,
+            out_features=out_features,
+            bias=False,
+        )
+
+    projection_dim = 256
+
+    def conv() -> Conv1d.Config:
+        return Conv1d.Config(
+            in_channels=projection_dim,
+            out_channels=projection_dim,
+            kernel_size=4,
+            groups=projection_dim,
+            bias=False,
+        )
+
+    return KDA.Config(
+        num_heads=2,
+        head_dim=128,
+        conv_kernel_size=4,
+        q_proj=linear(32, projection_dim),
+        k_proj=linear(32, projection_dim),
+        v_proj=linear(32, projection_dim),
+        q_conv=conv(),
+        k_conv=conv(),
+        v_conv=conv(),
+        forget_a=linear(32, 128),
+        forget_b=linear(128, projection_dim),
+        beta=linear(32, 2),
+        output_gate=linear(32, projection_dim),
+        inner_kda=InnerKDA.Config(
+            head_dim=128,
+            kernel=KDAKernel.Config(),
+        ),
+        output_norm=KimiRMSNormGated.Config(dim=128),
+        output_proj=linear(projection_dim, 32),
+    )
+
+
+@unittest.skipUnless(
+    _HAS_BLACKWELL, "KDA requires Attention Gym on CUDA capability 10.0 or 10.3"
+)
+class TestKDA(unittest.TestCase):
+    def _make_kda(self):
+        model = _kda_config().build()
+        model = model.to(device="cuda", dtype=torch.bfloat16)
+        torch.manual_seed(1)
+        with torch.no_grad():
+            for param in model.parameters():
+                param.normal_(mean=0.0, std=0.02)
+            model.A_log.uniform_(1.0, 16.0).log_()
+            model.dt_bias.zero_()
+            model.output_norm.weight.fill_(1.0)
+        return model
+
+    def _inputs(self, seed: int, tokens: int = 128) -> torch.Tensor:
+        torch.manual_seed(seed)
+        return torch.randn(tokens, 32, device="cuda", dtype=torch.bfloat16)
+
+    def test_varlen_matches_independent_documents(self):
+        lengths = (37, 64, 91)
+        x_TD = self._inputs(seed=2, tokens=sum(lengths)).requires_grad_()
+        positions_T = torch.tensor(
+            [index for length in lengths for index in range(length)],
+            device="cuda",
+            dtype=torch.int32,
+        )
+        masks = create_varlen_metadata_for_document(
+            positions_T,
+            include_host_offsets=True,
+        )
+        self.assertEqual(masks.cu_seq_q_host, (0, 37, 101, 192))
+
+        model = self._make_kda()
+        packed_TD = model(x_TD, masks)
+        independent_TD = torch.cat(
+            [model(document_TD, None) for document_TD in x_TD.split(lengths)]
+        )
+        torch.testing.assert_close(
+            packed_TD.float(),
+            independent_TD.float(),
+            rtol=2e-2,
+            atol=2e-2,
+        )
+        output_grad_TD = torch.randn_like(packed_TD)
+        parameters = tuple(model.parameters())
+        packed_grads = torch.autograd.grad(
+            packed_TD,
+            (x_TD, *parameters),
+            output_grad_TD,
+        )
+        independent_grads = torch.autograd.grad(
+            independent_TD,
+            (x_TD, *parameters),
+            output_grad_TD,
+        )
+        torch.testing.assert_close(
+            packed_grads,
+            independent_grads,
+            rtol=2e-2,
+            atol=2e-2,
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()