dataset changes for dpo

Author: tobyzl2

fast_llm/data/data/gpt/data.py

@@ -32,20 +32,34 @@ class GPTBatch:
     token_ids: torch.Tensor
     loss_masking_spans: list[torch.Tensor] | None = None
     sequence_lengths: list[torch.Tensor] | None = None
+    chosen_loss_masking_spans: list[torch.Tensor] | None = None
+    rejected_loss_masking_spans: list[torch.Tensor] | None = None
 
 
 def gpt_data_collate_fn(
-    batch: list[GPTSample], use_loss_masking_spans: bool, cross_document_attention: bool
+    batch: list[GPTSample], 
+    use_loss_masking_spans: bool, 
+    cross_document_attention: bool,
+    use_preference_loss_masking_spans: bool
 ) -> GPTBatch:
     stacked_ids = np.stack([sample.token_ids for sample in batch])
     stacked_spans = None
     sequence_lengths = None
+    stacked_chosen_spans = None
+    stacked_rejected_spans = None
     if use_loss_masking_spans:
         stacked_spans = [torch.from_numpy(sample.loss_masking_spans) for sample in batch]
+    if use_preference_loss_masking_spans:
+        stacked_chosen_spans = [torch.from_numpy(sample.chosen_loss_masking_spans) for sample in batch]
+        stacked_rejected_spans= [torch.from_numpy(sample.rejected_loss_masking_spans) for sample in batch]
     if not cross_document_attention:
         sequence_lengths = [torch.tensor(sample.sequence_lengths) for sample in batch]
     return GPTBatch(
-        token_ids=torch.from_numpy(stacked_ids), loss_masking_spans=stacked_spans, sequence_lengths=sequence_lengths
+        token_ids=torch.from_numpy(stacked_ids), 
+        loss_masking_spans=stacked_spans, 
+        sequence_lengths=sequence_lengths,
+        chosen_loss_masking_spans=stacked_chosen_spans,
+        rejected_loss_masking_spans=stacked_rejected_spans
     )
 
 
@@ -169,6 +183,7 @@ def get_iterator(
                     gpt_data_collate_fn,
                     use_loss_masking_spans=self._config.sampling.use_loss_masking_spans,
                     cross_document_attention=self._cross_document_attention,
+                    use_preference_loss_masking_spans=self._config.sampling.use_preference_loss_masking_spans
                 ),
                 multiprocessing_context=self._config.multiprocessing_context.value if num_workers > 0 else None,
             )

fast_llm/data/dataset/gpt/memmap.py

@@ -106,7 +106,7 @@ def _init(self, name: str, prefix: pathlib.Path | str, num_documents: int | None
                         dtype=np.int32,
                         count=2,
                         offset=chosen_span_offset + idx * 2 * np.dtype(np.int32).itemsize,
-                    ).reshape(-1, 2)
+                    )
                 )
 
             rejected_span_offset = offset + self._document_sizes.nbytes + self._pointers.nbytes + np.array(self._chosen_spans).nbytes
@@ -117,7 +117,7 @@ def _init(self, name: str, prefix: pathlib.Path | str, num_documents: int | None
                         dtype=np.int32,
                         count=2,
                         offset=rejected_span_offset + idx * 2 * np.dtype(np.int32).itemsize,
-                    ).reshape(-1, 2)
+                    )
                 )
 
         self._bin_buffer_mmap = np.memmap(self._prefix.with_suffix(".bin"), mode="r", order="C")
@@ -169,30 +169,30 @@ def get(
             chosen_spans = self._chosen_spans[idx]
 
             # filter spans that are outside the range of the selected tokens in the document
-            chosen_sample_spans = chosen_spans[
-                (chosen_spans[:, 0] < offset + len(token_ids)) & (chosen_spans[:, 1] >= offset)
-            ]
+            chosen_spans = chosen_spans[
+                (chosen_spans[0] < offset + len(token_ids)) & (chosen_spans[1] >= offset)
+            ][0]
 
             # subtract by offset to normalize span boundaries
-            chosen_spans[:, 0] = np.maximum(chosen_spans[:, 0], offset) - offset # offset 
-            chosen_spans[:, 1] = np.minimum(chosen_spans[:, 1], offset + len(token_ids) - 1) - offset
+            chosen_spans[0] = np.maximum(chosen_spans[0], offset) - offset # offset 
+            chosen_spans[1] = np.minimum(chosen_spans[1], offset + len(token_ids) - 1) - offset
 
             rejected_spans = self._rejected_spans[idx]
 
             # filter spans that are outside the range of the selected tokens in the document
-            rejected_sample_spans = rejected_spans[
-                (rejected_spans[:, 0] < offset + len(token_ids)) & (rejected_spans[:, 1] >= offset)
-            ]
+            rejected_spans = rejected_spans[
+                (rejected_spans[0] < offset + len(token_ids)) & (rejected_spans[1] >= offset)
+            ][0]
 
             # subtract by offset to normalize span boundaries
-            rejected_spans[:, 0] = np.maximum(rejected_spans[:, 0], offset) - offset # offset 
-            rejected_spans[:, 1] = np.minimum(rejected_spans[:, 1], offset + len(token_ids) - 1) - offset
+            rejected_spans[0] = np.maximum(rejected_spans[0], offset) - offset # offset 
+            rejected_spans[1] = np.minimum(rejected_spans[1], offset + len(token_ids) - 1) - offset
 
         return GPTSample(
             token_ids=token_ids, 
             loss_masking_spans=sample_spans, 
-            chosen_loss_masking_spans=chosen_sample_spans, 
-            rejected_loss_masking_spans=rejected_sample_spans
+            chosen_loss_masking_spans=chosen_spans, 
+            rejected_loss_masking_spans=rejected_spans
         )
 
     @property

fast_llm/data/dataset/gpt/sampled.py

@@ -120,6 +120,9 @@ def __init__(
             # contains cumulative sum of document sizes grouped by TOKEN_CUMSUM_RATE in shuffled order
             self._token_cumsum_shuffled = MemmapArray(base_path.with_name(base_path.name + "_shuffled_cumsum.npy"))
             self._token_cumsum_unshuffled = MemmapArray(base_path.with_name(base_path.name + "_unshuffled_cumsum.npy"))
+
+            self._document_sizes = MemmapArray(base_path.with_name(base_path.name + "_shuffled_cumsum.npy"))
+
             self._yaml_path = base_path.with_suffix(".yaml")
             # Sample or validate the dataset of a given rank.
             if sampling.distributed.config.rank == sampling.get_next_rank():
@@ -132,11 +135,11 @@ def _sample(self) -> None:
         Create a `GPTSampledDataset` with the requested parameters.
         """
         # Get the document sizes, the main information needed for sampling.
-        self.document_sizes = torch.from_numpy(self._indexed_dataset.get_document_sizes()).to(self._device)
+        document_sizes = torch.from_numpy(self._indexed_dataset.get_document_sizes()).to(self._device)
 
         # Calculate basic stats.
-        documents_per_epoch = self.document_sizes.numel()
-        tokens_per_epoch = self.document_sizes.sum().item()
+        documents_per_epoch = document_sizes.numel()
+        tokens_per_epoch = document_sizes.sum().item()
         # We produce sequences of length `self._sequence_length + 1` so the last token has a label,
         # but we also include that last label in the following sample,
         # so we need `sequence_length * num_samples + 1` tokens in total.
@@ -160,7 +163,7 @@ def _sample(self) -> None:
             "dataset": {
                 "name": self._indexed_dataset.name,
                 "documents_per_epoch": documents_per_epoch,
-                "tokens_per_epoch": tokens_per_epoch,
+                "tokens_per_epoch": tokens_per_epoch
             },
             "num_samples": self._num_samples,
             "unshuffled_epochs": unshuffled_epochs,
@@ -247,7 +250,7 @@ def _sample(self) -> None:
         if self._config.enable_packing:
             if shuffled_epochs > 0:
                 token_cumsum_shuffled = self._get_token_cumsum(
-                    self.document_sizes[
+                    document_sizes[
                         # Torch indexing only works with int32 or int64
                         document_shuffling.to(
                             dtype=torch.int64 if document_shuffling.dtype == torch.int64 else torch.int32
@@ -268,15 +271,17 @@ def _sample(self) -> None:
 
             if unshuffled_epochs > 0:
                 token_cumsum_unshuffled = self._get_token_cumsum(
-                    self.document_sizes, offset=0, dtype=get_unsigned_integer_type(tokens_per_epoch * num_epochs).torch
+                    document_sizes, offset=0, dtype=get_unsigned_integer_type(tokens_per_epoch * num_epochs).torch
                 )
                 self._token_cumsum_unshuffled.save(token_cumsum_unshuffled.numpy(force=self._config.gpu))
         else:
-            self._document_shuffling.save(
-                document_shuffling[:self._num_samples].numpy(
-                    force=self._config.gpu
+            if shuffled_epochs > 0:
+                self._document_shuffling.save(
+                    document_shuffling[:self._num_samples].numpy(
+                        force=self._config.gpu
+                    )
                 )
-            )
+            self._document_sizes.save(document_sizes.numpy(force=self._config.gpu))
 
     def _get_token_cumsum(self, sizes: torch.Tensor, offset: int, dtype: torch.dtype) -> torch.Tensor:
         # Create the output tensor.
@@ -385,11 +390,15 @@ def __getitem__(self, index: int) -> typing.Any:
             sample = self._indexed_dataset.get(
                 document_index,
                 offset=0,
-                length=self.document_sizes[document_index],
+                length=self._document_sizes[document_index],
                 use_loss_masking_spans=self._config.use_loss_masking_spans,
                 use_preference_loss_masking_spans=self._config.use_preference_loss_masking_spans
             )
 
+            chosen_loss_masking_span_end = sample.chosen_loss_masking_spans[1] + 1
+            sequence_lengths = np.array([chosen_loss_masking_span_end, len(sample.token_ids) - chosen_loss_masking_span_end])
+            sample.sequence_lengths = sequence_lengths
+
             return sample
 
     @property

fast_llm/data/tokenizer.py

@@ -22,6 +22,8 @@ def __init__(self, config: TokenizerConfig):
             raise ValueError("Tokenizer does not have an BOS token.")
         self.eod_id = self.tokenizer.eos_token_id
         self.bod_id = self.tokenizer.bos_token_id
+        self.eod_token = self.tokenizer.eos_token
+        self.bod_token = self.tokenizer.bos_token
 
     @property
     def vocab_size(self) -> int:
@@ -52,6 +54,9 @@ def tokenize_with_spans(
         token_spans = []
         char_pos = 0
         beginning_of_text = True
+        if text.startswith(self.bod_token):
+            beginning_of_text = False
+        
         for start, end in char_spans:
             if char_pos < start:
                 curr_text = text[char_pos:start]
@@ -60,7 +65,11 @@ def tokenize_with_spans(
                 input_ids.extend(tokenized_text)
             curr_text = text[start : end + 1]
             if end >= len(text) - 1:
-                tokenized_text = self.tokenize(curr_text, begin=beginning_of_text, end=True)
+                tokenized_text = self.tokenize(
+                    curr_text, 
+                    begin=beginning_of_text, 
+                    end=True if not curr_text.endswith(self.eod_token) else False
+                )
             else:
                 tokenized_text = self.tokenize(curr_text, begin=beginning_of_text, end=False)
             beginning_of_text = False
@@ -69,7 +78,11 @@ def tokenize_with_spans(
             char_pos = end + 1
         if char_pos < len(text):
             curr_text = text[char_pos:]
-            tokenized_text = self.tokenize(curr_text, begin=beginning_of_text, end=True)
+            tokenized_text = self.tokenize(
+                curr_text, 
+                begin=beginning_of_text, 
+                end=True if not curr_text.endswith(self.eod_token) else False
+            )
             input_ids.extend(tokenized_text)
         return input_ids, token_spans