Repository: pytorch/torchtune

Issue:
Investigate possible memory leak with sample packing
Following the gist below, a dataset of size ~3 GB with 2M entries takes up more than 900 GB of CPU memory while doing offline sample packing. 

https://gist.github.com/joecummings/642eaa9ce539ad93360ee3f999dbcfa3

Need to get to the bottom of this.


Discussion:
To process 50k samples (~ 2.2% of the entire dataset) consumes 41gb of memory.

To process 100k samples (~ 4.3% of the entire dataset) consumes 81 gb of memory. 

This seems like the memory is scaling linearly, based on this looks like the entire process would consume ~1886 gb of memory.
Using the following code, I'm calculating the size of the packs directly:

```
packs = [
	{"tensor": ... , "mask": ... , "labels": ... , "input_pos": ... },
	{"tensor": ... , "mask": ... , "labels": ... , "input_pos": ... },
	...
]


def mem_packs(packs):
    total_size = 0
    for pack in packs:
        for elem in pack.values():
			# this adds python object memory (same for every object) + size of tensor elem * num of tensors
            total_size += sys.getsizeof(elem) + elem.element_size() * elem.nelement()
    return f"{round(total_size / (1024**3), 2)}G"
```


Here's how it plays out for every pack:

```
tokens: 88 (fixed size of Python object) + 8 (size of torch.int64) * 3072 (max seq len)
labels: 88 + 8  * 3072
mask: 88 + 1 (size of torch.bool) * 3072 * 3072 (max seq len x max seq len)
input_pos: 88 + 8 * 3072
------------------------------
9511264 bytes ~= 9.5 MB
```

The mask is HUGGGEEE - we need to change how we calculate this.
See comment soln seems to be to calculate mask in each batch: https://github.com/pytorch/torchtune/pull/1083/files#diff-d069dcb7c967a8f3aaaca697349854143ff7b0631f885a6e7bd62dda19f6eee1R169

We could also technically do this for `input_pos`. This would slow down the processing but save memory:

```
tokens: 88 (fixed size of Python object) + 8 (size of torch.int64) * 3072 (max seq len)
labels: 88 + 8  * 3072
----------------------
49328 bytes ~= 0.05 MB
```
Thought I might chip in if it's helpful. I have a little helper function for doing this for each batch in my PPO fork:

```python
def get_causal_mask(
    tokens: torch.Tensor,
    padding_mask: torch.Tensor,
    *,
    dtype: torch.dtype = torch.float32,
) -> torch.Tensor:
    """
    Generates a causal attention mask for the given tokens and padding mask suitable for
    consumtion by :func:`~torch.nn.functional.scaled_dot_product_attention~` where the
    mask is added to the attention score.

    HF uses a similar implementation internally, see
    https://github.com/huggingface/transformers/blob/a564d10afe1a78c31934f0492422700f61a0ffc0/src/transformers/models/mistral/modeling_mistral.py#L1096

    Args:
        tokens (torch.Tensor): tensor of token IDs with shape [bsz x seq_length]
        padding_mask (torch.Tensor): tensor of padding ID masks with shape [bsz x seq_length]
        dtype (torch.dtype): dtype to infer fill value for masking
    Returns:
        torch.Tensor: Casual mask with shape [bsz x seq_length x seq_length]
    """
    fill_value = torch.finfo(dtype).min
    mask = torch.triu(
        torch.full((tokens.shape[-1], tokens.shape[-1]), fill_value), diagonal=1
    ).to(tokens.device, dtype=dtype)
    mask = mask.masked_fill(
        padding_mask.unsqueeze(1).expand(-1, tokens.shape[1], tokens.shape[1]),
        fill_value,
    return mask
    )
 ```
This is inspired by how HF does things to keep memory usage low: only creating `padding_mask` during dataset sampling, which is `[bsz, seq_len]`, and reshaping this on-the-fly to a `[bsz, seq_len, seq_len]` before passing to SDPA. Here's an example in context which also generates position IDs. 

```python
                    # create position IDs and causal masks for the current trajectory
                    padding_masks = query_responses == self._tokenizer.pad_id
                    # we only need custom causal masks for sequences with left-padding
                    if padding_masks.any():
                        masks = ppo_utils.get_causal_mask(
                            query_responses,
                            padding_mask=padding_masks,
                            dtype=self._dtype,
                        )
                        position_ids = (~padding_masks).cumsum(-1) - (~padding_masks).long()
                        position_ids = position_ids.to(device=self._device, dtype=torch.int)
                    else:
                        # defer SDPA to handle causal masks
                        masks, position_ids = None, None
```

after discussing on disc this needs to be updated so atnn masks are bool, not floats, but the idea is the same.



Patch:
diff --git a/torchtune/datasets/_chat.py b/torchtune/datasets/_chat.py
index 320d3b4e6d..82729ed03a 100644
--- a/torchtune/datasets/_chat.py
+++ b/torchtune/datasets/_chat.py
@@ -185,4 +185,8 @@ def chat_dataset(
         train_on_input=train_on_input,
         **load_dataset_kwargs,
     )
-    return PackedDataset(ds, max_seq_len=max_seq_len) if packed else ds
+    return (
+        PackedDataset(ds, max_seq_len=max_seq_len, padding_idx=tokenizer.pad_id)
+        if packed
+        else ds
+    )
diff --git a/torchtune/datasets/_instruct.py b/torchtune/datasets/_instruct.py
index 31fcfa4086..c06e7397e6 100644
--- a/torchtune/datasets/_instruct.py
+++ b/torchtune/datasets/_instruct.py
@@ -177,4 +177,8 @@ def instruct_dataset(
         max_seq_len=max_seq_len,
         **load_dataset_kwargs,
     )
-    return PackedDataset(ds, max_seq_len=max_seq_len) if packed else ds
+    return (
+        PackedDataset(ds, max_seq_len=max_seq_len, padding_idx=tokenizer.pad_id)
+        if packed
+        else ds
+    )
diff --git a/torchtune/datasets/_packed.py b/torchtune/datasets/_packed.py
index 6d951ed50c..fb41b983ba 100644
--- a/torchtune/datasets/_packed.py
+++ b/torchtune/datasets/_packed.py
@@ -4,7 +4,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from typing import Any, Dict, List, Optional
+from typing import Dict, List, Optional, Union
 
 import torch
 from torch.nn import functional as F
@@ -14,6 +14,8 @@
 from torchtune.utils import get_world_size_and_rank
 from tqdm import tqdm
 
+PACK_TYPE = Dict[str, Union[torch.Tensor, List[int]]]
+
 
 class PackedDataset(Dataset):
     """
@@ -71,7 +73,8 @@ class on packed samples with a ``Sampler`` as part of the dataloader. Currently,
         ds (Dataset): dataset to sample pack. This should return a dictionary with field
             "tokens" and "labels" containing the tokenized and label samples.
         max_seq_len (int): Maximum number of tokens to pack
-        max_packs (Optional[int]): maximum number of packs. Default is None, which will create as many
+        padding_idx (int): padding index for the tokenizer. Default is 0.
+        max_packs (Optional[int]): Maximum number of packs. Default is None, which will create as many
             packs as possible.
         split_across_pack (bool): if the last sample in a pack does not fit in ``max_seq_len``,
             split the sample into the next pack, or move it entirely to the beginning of the next pack.
@@ -83,34 +86,33 @@ class on packed samples with a ``Sampler`` as part of the dataloader. Currently,
     def __init__(
         self,
         ds: Dataset,
+        *,
         max_seq_len: int,
+        padding_idx: int = 0,
         max_packs: Optional[int] = None,
         split_across_pack: bool = False,
     ) -> None:
         self.ds = ds
         self.max_seq_len = max_seq_len
+        self.padding_idx = padding_idx
         self.max_packs = max_packs
         self.split_across_pack = split_across_pack
-        # where final samples will be held
-        self.packs: List[Dict[str, List[int]]] = []
+        # Where final samples will be held
+        self.packs: List[PACK_TYPE] = []
+        self.previous_sample_boundary: int = 0
         self._pack()
 
     def _pack(self) -> None:
-        """
-        Iterate through the dataset. Use a buffer to hold samples until max_seq_len,
+        """Iterate through the dataset. Use a buffer to hold samples until max_seq_len,
         then append the buffer to self.packs as a single "packed" sample. Continue
-        until max_packs or end of dataset.
-        """
-        # buffer to hold samples until they are long enough to be added to self.packs
+        until max_packs or end of dataset."""
+        # Buffer to hold samples until they are long enough to be added to self.packs
         current_pack = {
             "tokens": [],
             "labels": [],
-            "mask": [],
             "input_pos": [],
+            "seq_lens": [],
         }
-        # Keep track of what index the previous sample ends in case we need
-        # to end a pack early
-        previous_sample_boundary = 0
 
         # Only show progress bar on rank 0
         _, rank = get_world_size_and_rank()
@@ -119,156 +121,170 @@ def _pack(self) -> None:
 
         for sample in self.ds:
             tokens, labels = sample["tokens"], sample["labels"]
+
             # If the dataset outputs samples that are larger than the specified
             # max_seq_len and we're unable to split it, user needs to modify
             # one of the two parameters
             seq_len = len(tokens)
             if seq_len > self.max_seq_len and not self.split_across_pack:
                 raise ValueError(
-                    f"Dataset sample is too long ({len(tokens)} > {self.max_seq_len}). "
+                    f"Dataset sample is too long ({seq_len} > {self.max_seq_len}). "
                     "Please set `split_across_pack=True` or increase `max_seq_len`."
                 )
 
-            # Create integer mask and position ids for current sample and extend
-            # current pack
-            current_sample = {
-                "tokens": tokens,
-                "labels": labels,
-                # Mask is simply a causal mask within this sample length
-                "mask": [torch.tril(torch.ones(seq_len, seq_len, dtype=torch.bool))],
-                "input_pos": list(range(seq_len)),
-            }
-            current_pack = {k: v + current_sample[k] for k, v in current_pack.items()}
-
-            # If the current pack is long enough, add it to self.packs and retain
-            # any truncated samples for next pack, if splitting samples
+            # Update the current pack
+            current_pack["tokens"] += tokens
+            current_pack["labels"] += labels
+            current_pack["input_pos"] += list(range(seq_len))
+            current_pack["seq_lens"] += [seq_len]
+
+            # If the current pack is over the max_seq_len, add it to self.packs and
+            # retain any truncated or bumped samples for next pack
             if len(current_pack["tokens"]) > self.max_seq_len:
-                current_pack = self._add_pack(
-                    current_pack=current_pack,
-                    boundary=self.max_seq_len
-                    if self.split_across_pack
-                    else previous_sample_boundary,
-                )
+                current_pack = self._split_and_add_pack(current_pack)
 
             if rank == 0:
                 pbar.update()
-            previous_sample_boundary = len(current_pack["tokens"])
-            if self.max_packs is not None and len(self.packs) >= self.max_packs:
+
+            # Keep track of previous sample boundary
+            self.previous_sample_boundary = len(current_pack["tokens"])
+
+            # If max packs is set, stop packing when we reach that number
+            if self.max_packs is not None and len(self.packs) == self.max_packs:
                 break
 
-        # Add the last pack with remaining samples that did not fit in previous
+        # Handle the last pack if there's leftover and we haven't filled up the max packs
         if len(current_pack["tokens"]) > 0 and (
             self.max_packs is None or len(self.packs) < self.max_packs
         ):
-            current_pack = self._add_pack(
-                current_pack=current_pack, boundary=len(current_pack["tokens"])
-            )
-            assert len(current_pack["tokens"]) == 0
+            # No need to handle splitting at this point so we can just add the current pack
+            self._add_pack(current_pack)
+
+    def _split_and_add_pack(self, current_pack: PACK_TYPE) -> PACK_TYPE:
+        """Splits the current pack at the boundary, processes it, adds it to ``self.packs`` and
+        returns the start of the next pack."""
+
+        if self.split_across_pack:
+            boundary = self.max_seq_len
+            # The last elem in ``seq_lens`` ensures that ``sum(seq_lens) == self.max_seq_len``
+            seq_len_padding = [self.max_seq_len - sum(current_pack["seq_lens"][:-1])]
+        else:
+            boundary = self.previous_sample_boundary
+            # If we aren't splitting across packs, we leave out the last sample b/c
+            # it will go into the next pack
+            seq_len_padding = []
 
-    def _add_pack(
-        self, current_pack: Dict[str, List[int]], boundary: int
-    ) -> Dict[str, List[int]]:
-        """
-        Pad and add the current pack to self.packs and return what's remaining.
-        """
-        # So far we've kept a list of causal masks, one for each sample in the pack.
-        # Now we need to combine them into a single mask for the entire pack.
-        packing_mask = torch.block_diag(*current_pack["mask"])
         pack = {
             "tokens": current_pack["tokens"][:boundary],
             "labels": current_pack["labels"][:boundary],
-            "mask": packing_mask[:boundary, :boundary],
             "input_pos": current_pack["input_pos"][:boundary],
+            "seq_lens": current_pack["seq_lens"][:-1] + seq_len_padding,
         }
 
-        # Resultant shapes after padding
-        # tokens: [max_seq_len, ]
-        # labels: [max_seq_len, ]
-        # mask: [max_seq_len, max_seq_len]
-        # input_pos: [max_seq_len, ]
-        padded_pack = self._pad_pack(pack)
-        self.packs.append(padded_pack)
+        # Process and add the pack
+        self._add_pack(pack)
+
+        # Return the length of the first sample in next pack if we are splitting across packs,
+        # otherwise return the length of the last sample in the current pack
+        next_seq_len = (
+            len(current_pack["tokens"][boundary:])
+            if self.split_across_pack
+            else current_pack["seq_lens"][-1]
+        )
 
-        # Keep sample that did not fit or got truncated for the next pack
-        updated_pack = {
+        return {
             "tokens": current_pack["tokens"][boundary:],
             "labels": current_pack["labels"][boundary:],
-            "mask": [packing_mask[boundary:, boundary:]],
             "input_pos": current_pack["input_pos"][boundary:],
+            "seq_lens": [next_seq_len],
         }
 
-        return updated_pack
+    def _add_pack(self, pack: PACK_TYPE) -> None:
+        """Processes, pads and adds a pack to ``self.packs``."""
+        pack = self._convert_to_tensors(pack)
+        pack = self._pad_pack(pack, padding_idx=self.padding_idx)
+        self.packs.append(pack)
 
-    def _pad_pack(
-        self,
-        sample: Dict[str, Any],
-        padding_idx: int = 0,
-        ignore_idx: int = CROSS_ENTROPY_IGNORE_IDX,
-    ) -> Dict[str, torch.Tensor]:
-        """Pad a group of packed sequences to max sequence length in the group, and
-        convert integer lists to tensors. Account for attention mask and position
-        ids.
-
-        This is a sample-wise collator and should not be used with the dataloader.
-
-        Sample should look like::
-            {
-                "tokens": List[int],
-                "labels": List[int],
-                "mask": Tensor,
-                "input_pos": List[int],
-            }
-
-        Args:
-            sample (Dict[str, Any]): A dictionary containing tokens, labels, mask,
-                and position ids
-            padding_idx (int): Padding index for input ids. Defaults to 0.
-            ignore_idx (int): Padding index for labels. Defaults to -100.
-
-        Returns:
-            Collated tokens, labels, mask, and position ids.
+    def _convert_to_tensors(self, pack: PACK_TYPE) -> PACK_TYPE:
+        """Converts a pack into tensors. Pack comes in as a dict of lists and is converted to tensors.
+        The only key that does not get converted is ``seq_lens``.
         """
-        tokens = sample["tokens"]
-        labels = sample["labels"]
-        mask = sample["mask"]
-        input_pos = sample["input_pos"]
-
-        # Pad to max sequence length
-        tokens = F.pad(
-            torch.tensor(tokens), (0, self.max_seq_len - len(tokens)), value=padding_idx
-        )
-        labels = F.pad(
-            torch.tensor(labels), (0, self.max_seq_len - len(labels)), value=ignore_idx
-        )
-        # For attention mask, simply use identity matrix for the pad tokens
-        mask_pad = torch.eye(self.max_seq_len - mask.shape[0], dtype=torch.bool)
-        mask = torch.block_diag(mask, mask_pad)
-        # For position ids, continue to increment for pad tokens
-        next_pos = input_pos[-1] + 1
-        input_pos_pad = torch.arange(
-            next_pos, next_pos + self.max_seq_len - len(input_pos)
+        return {
+            "tokens": torch.tensor(pack["tokens"]),
+            "labels": torch.tensor(pack["labels"]),
+            "input_pos": torch.tensor(pack["input_pos"]),
+            "seq_lens": pack["seq_lens"],
+        }
+
+    def _pad_pack(self, pack: PACK_TYPE, padding_idx: int) -> PACK_TYPE:
+        """Pads a pack to ``self.max_seq_len``."""
+        # Pad tokens
+        padded_tokens = F.pad(
+            pack["tokens"],
+            (0, self.max_seq_len - len(pack["tokens"])),
+            value=padding_idx,
         )
-        # Do not go beyond max_seq_len - 1
-        input_pos_pad = input_pos_pad.clamp(max=self.max_seq_len - 1)
-        input_pos = torch.cat(
-            [
-                torch.tensor(input_pos),
-                input_pos_pad,
-            ]
+
+        # Pad labels
+        padded_labels = F.pad(
+            pack["labels"],
+            (0, self.max_seq_len - len(pack["labels"])),
+            value=CROSS_ENTROPY_IGNORE_IDX,
         )
 
-        assert tokens.shape == labels.shape == input_pos.shape
-        assert tokens.shape[0] == mask.shape[0]
+        # Pad input_pos continuing the sequence from last value
+        # in input_pos
+        # e.g. [0 1 2] -> [0 1 2 3 4 5] for self.max_seq_len = 6
+        num_range = torch.arange(
+            pack["input_pos"][-1] + 1,
+            pack["input_pos"][-1] + self.max_seq_len - len(pack["input_pos"]) + 1,
+        )
+        # Clamp to max_seq_len - 1 to avoid out of bounds error
+        clamped_num_range = torch.clamp(num_range, 0, self.max_seq_len - 1)
+        padded_input_pos = torch.cat([pack["input_pos"], clamped_num_range])
 
         return {
-            "tokens": tokens,
-            "labels": labels,
-            "mask": mask,
-            "input_pos": input_pos,
+            "tokens": padded_tokens,
+            "labels": padded_labels,
+            "input_pos": padded_input_pos,
+            "seq_lens": pack["seq_lens"],  # seq_len is untouched
         }
 
-    def __len__(self):
+    def __len__(self) -> int:
         return len(self.packs)
 
-    def __getitem__(self, index: int) -> Dict[str, List[int]]:
-        return self.packs[index]
+    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
+        """Constructs the attention mask on-the-fly and returns whole sample."""
+        current_pack = self.packs[idx]
+
+        num_samples_in_pack = len(current_pack["seq_lens"])
+        total_seq_len = 0
+
+        block_attn_masks = []
+
+        for i, seq_len in enumerate(current_pack["seq_lens"]):
+            total_seq_len += seq_len
+
+            # Append lower triangular matrix for causal mask
+            block_attn_masks.append(
+                torch.tril(torch.ones(seq_len, seq_len, dtype=torch.bool))
+            )
+
+            # If we're at the last sample and the total seq len is less than the max seq len,
+            # we need to pad with identity matrix for the remainder
+            if i == num_samples_in_pack - 1 and total_seq_len < self.max_seq_len:
+                block_attn_masks.append(
+                    torch.eye(
+                        self.max_seq_len - total_seq_len,
+                        self.max_seq_len - total_seq_len,
+                        dtype=torch.bool,
+                    )
+                )
+
+        return {
+            "tokens": current_pack["tokens"],
+            "labels": current_pack["labels"],
+            "input_pos": current_pack["input_pos"],
+            # Assemble the mask into a block causal matrix
+            "mask": torch.block_diag(*block_attn_masks),
+        }
diff --git a/torchtune/datasets/_text_completion.py b/torchtune/datasets/_text_completion.py
index 92ab83da19..f9b0ed0c48 100644
--- a/torchtune/datasets/_text_completion.py
+++ b/torchtune/datasets/_text_completion.py
@@ -122,4 +122,8 @@ def text_completion_dataset(
         max_seq_len=max_seq_len,
         **load_dataset_kwargs,
     )
-    return PackedDataset(ds, max_seq_len=max_seq_len) if packed else ds
+    return (
+        PackedDataset(ds, max_seq_len=max_seq_len, padding_idx=tokenizer.pad_id)
+        if packed
+        else ds
+    )


Test patch:
diff --git a/tests/torchtune/datasets/test_packed_dataset.py b/tests/torchtune/datasets/test_packed_dataset.py
index 57f1b0dfaa..8afb532c0b 100644
--- a/tests/torchtune/datasets/test_packed_dataset.py
+++ b/tests/torchtune/datasets/test_packed_dataset.py
@@ -77,9 +77,29 @@ def _get_expected_mask_and_input_pos(
 
         return mask[:max_seq_len, :max_seq_len], torch.tensor(input_pos[:max_seq_len])
 
+    def _calculate_num_packs(
+        self, dataset_size, max_seq_len, sample_size, split_across_pack, max_packs
+    ):
+        # First see how many samples we can fit in a single pack
+        num_samples_per_pack, remainder = divmod(max_seq_len, sample_size)
+
+        # If we split across pack (and the samples don't fit perfectly in max_seq_len), we can fit more
+        if split_across_pack and remainder > 0:
+            # Now we need the fractional to see how many we can partially fit in each pack
+            num_samples_per_pack = max_seq_len / sample_size
+
+        # If we don't split across pack, we will need more packs
+        num_packs, remainder = divmod(dataset_size, num_samples_per_pack)
+
+        # If there's leftover, we need to add one more pack
+        if remainder > 0:
+            num_packs += 1
+
+        return num_packs if num_packs < max_packs else max_packs
+
     @pytest.mark.parametrize("max_seq_len", [25])
     @pytest.mark.parametrize("sample_size", [2, 5])
-    @pytest.mark.parametrize("max_packs", [5])
+    @pytest.mark.parametrize("max_packs", [5, 200])
     @pytest.mark.parametrize("split_across_pack", [True, False])
     def test_packed_dataset(
         self, max_seq_len, sample_size, max_packs, split_across_pack
@@ -91,8 +111,13 @@ def test_packed_dataset(
             max_packs=max_packs,
             split_across_pack=split_across_pack,
         )
+
         # Check we get right number of packs
-        assert len(packed) == max_packs
+        correct_num_packs = self._calculate_num_packs(
+            len(dataset), max_seq_len, sample_size, split_across_pack, max_packs
+        )
+        assert len(packed) == correct_num_packs
+
         # Check all fields are same length
         assert (
             len(packed[0]["tokens"])
@@ -105,7 +130,7 @@ def test_packed_dataset(
         if split_across_pack:
             # If we split samples, we'll know how many samples by taking the
             # full length and dividing by sample size
-            last_index, remainder = divmod(max_packs * max_seq_len, sample_size)
+            last_index, remainder = divmod(len(packed) * max_seq_len, sample_size)
             # Account for remaining sample that didn't fit in window
             last_index = last_index if remainder > 0 else last_index - 1
         else:
@@ -113,7 +138,7 @@ def test_packed_dataset(
             # how much fits in a single window and multiplying by max rows.
             # If there is a remainder, this will end up being a pad token.
             last_index = (
-                (max_seq_len // sample_size) * max_packs - 1
+                (max_seq_len // sample_size) * len(packed) - 1
                 if max_seq_len % sample_size == 0
                 else 0
             )
@@ -207,11 +232,11 @@ def test_packed_dataset_real_data(self):
 
     def test_pad_pack(self):
         padding_idx = -8
-        ignore_idx = -9
+        ignore_idx = -100  # Same as CROSS_ENTROPY_IGNORE_IDX
         pack = {
             "tokens": [2, 5],
             "labels": [3, 7],
-            "mask": torch.tensor([[True, False], [True, True]]),
+            "seq_lens": [1, 1],
             # Let the first token be the end of the previous sample (pos 8),
             # and the second token the start of the next sample (pos 0). Collate
             # should continue from 0 -> 1, 2, ...
@@ -224,11 +249,11 @@ def test_pad_pack(self):
             max_seq_len=4,
         )
 
-        padded = packed._pad_pack(pack, padding_idx=padding_idx, ignore_idx=ignore_idx)
+        pack = packed._convert_to_tensors(pack)
+        padded = packed._pad_pack(pack, padding_idx=padding_idx)
 
         padded_input = padded["tokens"]
         padded_label = padded["labels"]
-        padded_mask = padded["mask"]
         padded_input_pos = padded["input_pos"]
 
         torch.testing.assert_close(
@@ -237,15 +262,12 @@ def test_pad_pack(self):
         torch.testing.assert_close(
             padded_label, torch.tensor([3, 7, ignore_idx, ignore_idx])
         )
-        assert torch.equal(
-            padded_mask,
-            torch.tensor(
-                [
-                    [True, False, False, False],
-                    [True, True, False, False],
-                    [False, False, True, False],
-                    [False, False, False, True],
-                ]
-            ),
-        )
         torch.testing.assert_close(padded_input_pos, torch.tensor([8, 0, 1, 2]))
+
+    def test_pack_errors_if_sample_too_long(self):
+        dataset = DummyDataset(8)
+        with pytest.raises(ValueError, match="Dataset sample is too long"):
+            PackedDataset(
+                dataset,
+                max_seq_len=4,
+            )