Integrate flex attention

recipes/full_finetune_distributed.py

@@ -21,7 +21,7 @@
 from torch.optim import Optimizer
 from torch.utils.data import DataLoader, DistributedSampler
 from torchtune import config, modules, training, utils
-from torchtune.data import padded_collate_sft
+from torchtune.data import padded_collate_packed, padded_collate_sft
 from torchtune.datasets import ConcatDataset
 from torchtune.recipe_interfaces import FTRecipeInterface
 from torchtune.training import DummyProfiler, PROFILER_KEY
@@ -500,14 +500,15 @@ def _setup_data(
             dataset=ds,
             batch_size=batch_size,
             sampler=sampler,
-            collate_fn=(
-                partial(
-                    padded_collate_sft,
-                    padding_idx=self._tokenizer.pad_id,
-                    ignore_idx=self._loss_fn.ignore_index,
-                )
-                if not packed
-                else None
+            collate_fn=partial(
+                padded_collate_sft,
+                padding_idx=self._tokenizer.pad_id,
+                ignore_idx=self._loss_fn.ignore_index,
+            )
+            if not packed
+            else partial(
+                padded_collate_packed,
+                device=self._device,
             ),
         )
 

recipes/full_finetune_single_device.py

@@ -19,7 +19,7 @@
 from torch.utils.data import DataLoader, DistributedSampler
 
 from torchtune import config, modules, training, utils
-from torchtune.data import padded_collate_sft
+from torchtune.data import padded_collate_packed, padded_collate_sft
 from torchtune.datasets import ConcatDataset
 from torchtune.recipe_interfaces import FTRecipeInterface
 from torchtune.training import DummyProfiler, PROFILER_KEY
@@ -464,14 +464,15 @@ def _setup_data(
             dataset=ds,
             batch_size=batch_size,
             sampler=sampler,
-            collate_fn=(
-                partial(
-                    padded_collate_sft,
-                    padding_idx=self._tokenizer.pad_id,
-                    ignore_idx=self._loss_fn.ignore_index,
-                )
-                if not packed
-                else None
+            collate_fn=partial(
+                padded_collate_sft,
+                padding_idx=self._tokenizer.pad_id,
+                ignore_idx=self._loss_fn.ignore_index,
+            )
+            if not packed
+            else partial(
+                padded_collate_packed,
+                device=self._device,
             ),
         )
 

recipes/lora_finetune_distributed.py

@@ -26,7 +26,7 @@
 from torch.optim import Optimizer
 from torch.utils.data import DataLoader, DistributedSampler
 from torchtune import config, modules, training, utils
-from torchtune.data import padded_collate_sft
+from torchtune.data import padded_collate_packed, padded_collate_sft
 from torchtune.datasets import ConcatDataset
 from torchtune.modules.peft import (
     get_adapter_params,
@@ -551,14 +551,15 @@ def _setup_data(
             dataset=ds,
             batch_size=batch_size,
             sampler=sampler,
-            collate_fn=(
-                partial(
-                    padded_collate_sft,
-                    padding_idx=self._tokenizer.pad_id,
-                    ignore_idx=self._loss_fn.ignore_index,
-                )
-                if not packed
-                else None
+            collate_fn=partial(
+                padded_collate_sft,
+                padding_idx=self._tokenizer.pad_id,
+                ignore_idx=self._loss_fn.ignore_index,
+            )
+            if not packed
+            else partial(
+                padded_collate_packed,
+                device=self._device,
             ),
         )
 

recipes/lora_finetune_single_device.py

@@ -19,7 +19,7 @@
 from torch.optim import Optimizer
 from torch.utils.data import DataLoader, DistributedSampler
 from torchtune import config, modules, training, utils
-from torchtune.data import padded_collate_sft
+from torchtune.data import padded_collate_packed, padded_collate_sft
 from torchtune.datasets import ConcatDataset
 from torchtune.modules.peft import (
     get_adapter_params,
@@ -495,14 +495,15 @@ def _setup_data(
             dataset=ds,
             sampler=sampler,
             batch_size=batch_size,
-            collate_fn=(
-                partial(
-                    padded_collate_sft,
-                    padding_idx=self._tokenizer.pad_id,
-                    ignore_idx=self._loss_fn.ignore_index,
-                )
-                if not packed
-                else None
+            collate_fn=partial(
+                padded_collate_sft,
+                padding_idx=self._tokenizer.pad_id,
+                ignore_idx=self._loss_fn.ignore_index,
+            )
+            if not packed
+            else partial(
+                padded_collate_packed,
+                device=self._device,
             ),
         )
 

recipes/qat_distributed.py

@@ -21,7 +21,7 @@
 from torch.optim import Optimizer
 from torch.utils.data import DataLoader, DistributedSampler
 from torchtune import config, modules, training, utils
-from torchtune.data import padded_collate_sft
+from torchtune.data import padded_collate_packed, padded_collate_sft
 from torchtune.datasets import ConcatDataset
 from torchtune.recipe_interfaces import FTRecipeInterface
 from torchtune.training import DummyProfiler, PROFILER_KEY
@@ -523,14 +523,15 @@ def _setup_data(
             dataset=ds,
             batch_size=batch_size,
             sampler=sampler,
-            collate_fn=(
-                partial(
-                    padded_collate_sft,
-                    padding_idx=self._tokenizer.pad_id,
-                    ignore_idx=self._loss_fn.ignore_index,
-                )
-                if not packed
-                else None
+            collate_fn=partial(
+                padded_collate_sft,
+                padding_idx=self._tokenizer.pad_id,
+                ignore_idx=self._loss_fn.ignore_index,
+            )
+            if not packed
+            else partial(
+                padded_collate_packed,
+                device=self._device,
             ),
         )
 

tests/torchtune/data/test_collate.py

@@ -6,14 +6,19 @@
 
 # (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary.
 
+from unittest import mock
+
 import pytest
 import torch
+from tests.test_utils import gpu_test
 from torchtune.data import (
     left_pad_sequence,
     padded_collate,
     padded_collate_dpo,
+    padded_collate_packed,
     padded_collate_sft,
 )
+from torchtune.modules.attention_utils import _SUPPORTS_FLEX_ATTENTION
 
 
 class TestPaddedCollateSFT:
@@ -47,6 +52,110 @@ def test_batch_pad_sequence(self):
             padded_label, torch.tensor([10, ignore_idx, ignore_idx])
         )
 
+    @mock.patch("torchtune.modules.attention_utils._SUPPORTS_FLEX_ATTENTION", False)
+    def test_padded_collate_packed_sdpa(self):
+        token_pairs = [
+            {
+                "tokens": torch.tensor([1, 2, 3, 4, 5, 6]),
+                "labels": torch.tensor([7, 8, 9, 10, 11, 12]),
+                "input_pos": torch.tensor([0, 1, 2, 0, 1, 0]),
+                "seq_lens": torch.tensor([3, 2, 1]),
+            },
+            {
+                "tokens": torch.tensor([13, 14, 15, 16, 17, 18]),
+                "labels": torch.tensor([19, 20, 21, 22, 23, 24]),
+                "input_pos": torch.tensor([0, 1, 0, 1, 0, 1]),
+                "seq_lens": torch.tensor([2, 2, 2]),
+            },
+        ]
+        collated = padded_collate_packed(
+            batch=token_pairs,
+        )
+        torch.testing.assert_close(
+            collated["tokens"],
+            torch.tensor([[1, 2, 3, 4, 5, 6], [13, 14, 15, 16, 17, 18]]),
+        )
+        torch.testing.assert_close(
+            collated["labels"],
+            torch.tensor([[7, 8, 9, 10, 11, 12], [19, 20, 21, 22, 23, 24]]),
+        )
+        torch.testing.assert_close(
+            collated["input_pos"],
+            torch.tensor([[0, 1, 2, 0, 1, 0], [0, 1, 0, 1, 0, 1]]),
+        )
+        torch.testing.assert_close(
+            collated["mask"],
+            torch.tensor(
+                [
+                    [
+                        [1, 0, 0, 0, 0, 0],
+                        [1, 1, 0, 0, 0, 0],
+                        [1, 1, 1, 0, 0, 0],
+                        [0, 0, 0, 1, 0, 0],
+                        [0, 0, 0, 1, 1, 0],
+                        [0, 0, 0, 0, 0, 1],
+                    ],
+                    [
+                        [1, 0, 0, 0, 0, 0],
+                        [1, 1, 0, 0, 0, 0],
+                        [0, 0, 1, 0, 0, 0],
+                        [0, 0, 1, 1, 0, 0],
+                        [0, 0, 0, 0, 1, 0],
+                        [0, 0, 0, 0, 1, 1],
+                    ],
+                ],
+                dtype=torch.bool,
+            ),
+        )
+
+    @pytest.mark.skipif(
+        not _SUPPORTS_FLEX_ATTENTION,
+        reason="Please install a nightly build of torch to run this test.",
+    )
+    @gpu_test(gpu_count=1)
+    def test_padded_collate_packed_flex(self):
+        # create_block_mask requires that seq_len be divisible by 128, the default block size.
+        # see https://github.com/pytorch/pytorch/blob/main/torch/nn/attention/flex_attention.py#L636
+        batch = [
+            {
+                "tokens": torch.ones(128, dtype=torch.long),
+                "labels": torch.ones(128, dtype=torch.long),
+                "input_pos": torch.zeros(128, dtype=torch.long),
+                "seq_lens": torch.ones(64, dtype=torch.long) * 2,
+            },
+            {
+                "tokens": torch.ones(128, dtype=torch.long),
+                "labels": torch.ones(128, dtype=torch.long),
+                "input_pos": torch.zeros(128, dtype=torch.long),
+                "seq_lens": torch.ones(32, dtype=torch.long) * 4,
+            },
+        ]
+        collated = padded_collate_packed(
+            batch=batch,
+        )
+        torch.testing.assert_close(
+            collated["tokens"],
+            torch.stack(
+                [torch.ones(128, dtype=torch.long), torch.ones(128, dtype=torch.long)]
+            ),
+        )
+        torch.testing.assert_close(
+            collated["labels"],
+            torch.stack(
+                [torch.ones(128, dtype=torch.long), torch.ones(128, dtype=torch.long)]
+            ),
+        )
+        torch.testing.assert_close(
+            collated["input_pos"],
+            torch.stack(
+                [torch.zeros(128, dtype=torch.long), torch.zeros(128, dtype=torch.long)]
+            ),
+        )
+        torch.testing.assert_close(
+            collated["mask"].to_dense(),
+            torch.tensor([[[[1]]], [[[1]]]], dtype=torch.int32, device="cuda"),
+        )
+
 
 class TestLeftPadSequence:
     def test_left_pad_sequence(self):