Integrate flex attention

recipes/full_finetune_distributed.py

@@ -469,7 +469,10 @@ def _setup_data(
                 ignore_idx=self._loss_fn.ignore_index,
             )
             if not packed
-            else None,
+            else partial(
+                utils.padded_collate_packed,
+                device=self._device,
+            ),
         )
 
         if self._is_rank_zero:

recipes/full_finetune_single_device.py

@@ -432,7 +432,10 @@ def _setup_data(
                 ignore_idx=self._loss_fn.ignore_index,
             )
             if not packed
-            else None,
+            else partial(
+                utils.padded_collate_packed,
+                device=self._device,
+            ),
         )
 
         log.info("Dataset and Sampler are initialized.")

recipes/lora_finetune_distributed.py

@@ -520,14 +520,15 @@ def _setup_data(
             dataset=ds,
             batch_size=batch_size,
             sampler=sampler,
-            collate_fn=(
-                partial(
-                    utils.padded_collate,
-                    padding_idx=self._tokenizer.pad_id,
-                    ignore_idx=self._loss_fn.ignore_index,
-                )
-                if not packed
-                else None
+            collate_fn=partial(
+                utils.padded_collate,
+                padding_idx=self._tokenizer.pad_id,
+                ignore_idx=self._loss_fn.ignore_index,
+            )
+            if not packed
+            else partial(
+                utils.padded_collate_packed,
+                device=self._device,
             ),
         )
 

recipes/lora_finetune_single_device.py

@@ -451,14 +451,15 @@ def _setup_data(
             dataset=ds,
             sampler=sampler,
             batch_size=batch_size,
-            collate_fn=(
-                partial(
-                    utils.padded_collate,
-                    padding_idx=self._tokenizer.pad_id,
-                    ignore_idx=self._loss_fn.ignore_index,
-                )
-                if not packed
-                else None
+            collate_fn=partial(
+                utils.padded_collate,
+                padding_idx=self._tokenizer.pad_id,
+                ignore_idx=self._loss_fn.ignore_index,
+            )
+            if not packed
+            else partial(
+                utils.padded_collate_packed,
+                device=self._device,
             ),
         )
 

recipes/qat_distributed.py

@@ -445,7 +445,10 @@ def _setup_data(
                 ignore_idx=self._loss_fn.ignore_index,
             )
             if not packed
-            else None,
+            else partial(
+                utils.padded_collate_packed,
+                device=self._device,
+            ),
         )
 
         if self._is_rank_zero:

tests/torchtune/datasets/test_packed_dataset.py

@@ -48,34 +48,27 @@ def __len__(self):
 
 
 class TestPackedDataset:
-    def _get_expected_mask_and_input_pos(
+    def _get_expected_seq_lens_and_input_pos(
         self, max_seq_len, sample_size, split_across_pack
     ):
         """
-        Generate expected integer mask and position ids for given max sequence
+        Generate expected seq lens and position ids for given max sequence
         length and sample length
         """
         num_samples, remainder = divmod(max_seq_len, sample_size)
+        seq_lens = [sample_size] * num_samples
         if split_across_pack and remainder > 0:
             num_samples += 1
-        mask = torch.block_diag(
-            *[
-                torch.tril(torch.ones(sample_size, sample_size, dtype=torch.bool))
-                for i in range(1, num_samples + 1)
-            ]
-        )
         input_pos = [list(range(sample_size)) for i in range(1, num_samples + 1)]
         input_pos = list(itertools.chain(*input_pos))
 
-        # Emulate mask and position id padding
-        if not split_across_pack and remainder > 0:
-            mask = torch.block_diag(
-                mask,
-                torch.eye(remainder, dtype=torch.bool),
-            )
-            input_pos.extend(list(range(sample_size, sample_size + remainder)))
+        # Emulate seq len and position id padding
+        if remainder > 0:
+            if not split_across_pack:
+                input_pos.extend(list(range(sample_size, sample_size + remainder)))
+            seq_lens.extend([remainder])
 
-        return mask[:max_seq_len, :max_seq_len], torch.tensor(input_pos[:max_seq_len])
+        return torch.tensor(seq_lens), torch.tensor(input_pos[:max_seq_len])
 
     def _calculate_num_packs(
         self, dataset_size, max_seq_len, sample_size, split_across_pack, max_packs
@@ -122,7 +115,6 @@ def test_packed_dataset(
         assert (
             len(packed[0]["tokens"])
             == len(packed[0]["labels"])
-            == len(packed[0]["mask"])
             == len(packed[0]["input_pos"])
         )
         # Check that samples are packed correctly - very last individual sample
@@ -145,10 +137,13 @@ def test_packed_dataset(
 
         assert packed[-1]["tokens"][-1].item() == last_index
 
-        expected_mask, expected_input_pos = self._get_expected_mask_and_input_pos(
+        (
+            expected_seq_lens,
+            expected_input_pos,
+        ) = self._get_expected_seq_lens_and_input_pos(
             max_seq_len, sample_size, split_across_pack
         )
-        torch.testing.assert_close(packed[0]["mask"], expected_mask)
+        torch.testing.assert_close(packed[0]["seq_lens"], expected_seq_lens)
         torch.testing.assert_close(packed[0]["input_pos"], expected_input_pos)
 
     def test_packed_dataset_real_data(self):
@@ -162,48 +157,15 @@ def test_packed_dataset_real_data(self):
             torch.tensor([5, 2, 6, 4, 3, 8, -1, 0, 4, 3]),
             torch.tensor([4, 3, 2, 5, 7, -1, -100, -100, -100, -100]),
         ]
-        expected_mask = [
+        expected_seq_lens = [
             torch.tensor(
-                [
-                    [1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
-                    [1, 1, 0, 0, 0, 0, 0, 0, 0, 0],
-                    [1, 1, 1, 0, 0, 0, 0, 0, 0, 0],
-                    [1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
-                    [1, 1, 1, 1, 1, 0, 0, 0, 0, 0],
-                    [1, 1, 1, 1, 1, 1, 0, 0, 0, 0],
-                    [1, 1, 1, 1, 1, 1, 1, 0, 0, 0],
-                    [0, 0, 0, 0, 0, 0, 0, 1, 0, 0],
-                    [0, 0, 0, 0, 0, 0, 0, 1, 1, 0],
-                    [0, 0, 0, 0, 0, 0, 0, 1, 1, 1],
-                ]
+                [7, 3],
             ),
             torch.tensor(
-                [
-                    [1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
-                    [1, 1, 0, 0, 0, 0, 0, 0, 0, 0],
-                    [1, 1, 1, 0, 0, 0, 0, 0, 0, 0],
-                    [1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
-                    [1, 1, 1, 1, 1, 0, 0, 0, 0, 0],
-                    [1, 1, 1, 1, 1, 1, 0, 0, 0, 0],
-                    [1, 1, 1, 1, 1, 1, 1, 0, 0, 0],
-                    [0, 0, 0, 0, 0, 0, 0, 1, 0, 0],
-                    [0, 0, 0, 0, 0, 0, 0, 1, 1, 0],
-                    [0, 0, 0, 0, 0, 0, 0, 1, 1, 1],
-                ]
+                [7, 3],
             ),
             torch.tensor(
-                [
-                    [1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
-                    [1, 1, 0, 0, 0, 0, 0, 0, 0, 0],
-                    [1, 1, 1, 0, 0, 0, 0, 0, 0, 0],
-                    [1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
-                    [1, 1, 1, 1, 1, 0, 0, 0, 0, 0],
-                    [1, 1, 1, 1, 1, 1, 0, 0, 0, 0],
-                    [0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
-                    [0, 0, 0, 0, 0, 0, 0, 1, 0, 0],
-                    [0, 0, 0, 0, 0, 0, 0, 0, 1, 0],
-                    [0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
-                ]
+                [6, 4],
             ),
         ]
         expected_input_pos = [
@@ -219,16 +181,16 @@ def test_packed_dataset_real_data(self):
         )
 
         for i in range(len(packed)):
-            prompt, label, mask, input_pos = (
+            prompt, label, seq_lens, input_pos = (
                 packed[i]["tokens"],
                 packed[i]["labels"],
-                packed[i]["mask"],
+                packed[i]["seq_lens"],
                 packed[i]["input_pos"],
             )
             torch.testing.assert_close(prompt, expected_tokenized_prompts[i])
             torch.testing.assert_close(label, expected_tokenized_labels[i])
             torch.testing.assert_close(input_pos, expected_input_pos[i])
-            torch.testing.assert_close(mask, expected_mask[i].to(dtype=torch.bool))
+            torch.testing.assert_close(seq_lens, expected_seq_lens[i])
 
     def test_pad_pack(self):
         padding_idx = -8
@@ -255,6 +217,7 @@ def test_pad_pack(self):
         padded_input = padded["tokens"]
         padded_label = padded["labels"]
         padded_input_pos = padded["input_pos"]
+        padded_seq_lens = padded["seq_lens"]
 
         torch.testing.assert_close(
             padded_input, torch.tensor([2, 5, padding_idx, padding_idx])
@@ -263,6 +226,7 @@ def test_pad_pack(self):
             padded_label, torch.tensor([3, 7, ignore_idx, ignore_idx])
         )
         torch.testing.assert_close(padded_input_pos, torch.tensor([8, 0, 1, 2]))
+        torch.testing.assert_close(padded_seq_lens, torch.tensor([1, 1, 2]))
 
     def test_pack_errors_if_sample_too_long(self):
         dataset = DummyDataset(8)

tests/torchtune/utils/test_attention_bias.py

@@ -0,0 +1,81 @@
+# 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.
+
+# (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary.
+
+from unittest import mock
+
+import pytest
+import torch
+
+from torchtune.utils.attention_bias import (
+    _get_document_ids_from_seq_lens,
+    create_block_causal_mask,
+    packed_block_causal_mask,
+)
+
+
+class TestBlockCausalMask:
+    @pytest.fixture
+    def seq_lens(self):
+        return torch.tensor([[2, 3, 1, 0], [2, 2, 2, 0]])
+
+    def test_get_document_ids_from_seq_lens(self, seq_lens):
+        actual = _get_document_ids_from_seq_lens(seq_lens)
+        expected = torch.tensor([[0, 0, 1, 1, 1, 2], [0, 0, 1, 1, 2, 2]])
+        torch.testing.assert_close(actual, expected)
+
+    def test_create_block_causal_mask(self, seq_lens):
+        actual = create_block_causal_mask(seq_lens)
+        expected = torch.tensor(
+            [
+                [
+                    [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, 1, 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,
+        )
+        torch.testing.assert_close(actual, expected)
+
+    @mock.patch("torchtune.utils.attention_bias.torch_version_ge")
+    def test_packed_block_causal_mask_sdpa(self, mock_version, seq_lens):
+        mock_version.return_value = False
+        actual = packed_block_causal_mask(seq_lens, device="cpu")
+        expected = torch.tensor(
+            [
+                [
+                    [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, 1, 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,
+        )
+        torch.testing.assert_close(actual, expected)