Ensure Equivalence With DPO Padding and Padding Free

open_instruct/dpo_tune_cache.py

@@ -76,9 +76,6 @@
     simpo_loss,
     wpo_loss,
 )
-from open_instruct.padding_free_collator import (
-    TensorDataCollatorWithFlatteningDPO
-)
 from open_instruct.model_utils import push_folder_to_hub, save_with_accelerate
 from open_instruct.padding_free_collator import TensorDataCollatorWithFlatteningDPO
 from open_instruct.utils import (

open_instruct/dpo_utils.py

@@ -195,7 +195,10 @@ def process_batch(
     return processed
 
 
-def concatenated_inputs(batch: Dict[str, Union[List, torch.LongTensor]]) -> Dict[str, torch.LongTensor]:
+def concatenated_inputs(
+    batch: Dict[str, Union[List, torch.LongTensor]],
+    pad_token_id: int = 0,
+) -> Dict[str, torch.LongTensor]:
     """Concatenate the chosen and rejected inputs into a single tensor.
 
     Args:
@@ -209,12 +212,12 @@ def concatenated_inputs(batch: Dict[str, Union[List, torch.LongTensor]]) -> Dict
     concatenated_batch = {}
     for k in batch:
         if k.startswith("chosen") and isinstance(batch[k], torch.Tensor):
-            pad_value = -100 if "labels" in k else 0
+            pad_value = -100 if "labels" in k else pad_token_id
             concatenated_key = k.replace("chosen", "concatenated")
             concatenated_batch[concatenated_key] = pad_to_length(batch[k], max_length, pad_value=pad_value)
     for k in batch:
         if k.startswith("rejected") and isinstance(batch[k], torch.Tensor):
-            pad_value = -100 if "labels" in k else 0
+            pad_value = -100 if "labels" in k else pad_token_id
             concatenated_key = k.replace("rejected", "concatenated")
             concatenated_batch[concatenated_key] = torch.cat(
                 (concatenated_batch[concatenated_key], pad_to_length(batch[k], max_length, pad_value=pad_value)), dim=0
@@ -234,7 +237,11 @@ def concatenated_forward(
     We do this to avoid doing two forward passes, because it's faster for FSDP.
     """
     if not packing:
-        concatenated_batch = concatenated_inputs(batch)
+        try:
+            pad_token_id = model.tokenizer.pad_token_id
+        except:
+            pad_token_id = 0
+        concatenated_batch = concatenated_inputs(batch, pad_token_id=pad_token_id)
     else:
         concatenated_batch, bs = pf_concatenated_inputs(batch)
 

tests/test_padding_free.py

@@ -20,8 +20,19 @@
 open_instruct_dir = Path(__file__).parent.parent.absolute()
 sys.path.append(open_instruct_dir)
 from open_instruct.dataset_processor import CHAT_TEMPLATES
-from open_instruct.dataset_transformation import sft_tulu_tokenize_and_truncate_v1
-from open_instruct.padding_free_collator import TensorDataCollatorWithFlattening
+from open_instruct.dataset_transformation import (
+    sft_tulu_tokenize_and_truncate_v1, 
+    preference_span_search_mask_out
+)
+from open_instruct.padding_free_collator import (
+    TensorDataCollatorWithFlattening,
+    TensorDataCollatorWithFlatteningDPO
+)
+from open_instruct.dpo_utils import (
+    DataCollatorForSeq2SeqDPO,
+    concatenated_forward
+)
+
 
 try:
     import mamba_ssm  # noqa
@@ -79,8 +90,11 @@ class TestPaddingFree:
     seqlen = 128
     batch_size = 2
     dtype = torch.bfloat16
+    model = None
 
     def get_fa2_model_and_cfg(self, model_name: str, vocab_size: int) -> nn.Module:
+        if self.model is not None:
+            return self.model, self.model.config
         model_cls = MODEL_CLASSES[model_name]
         model_cfg = MODEL_CFGS[model_name]
         model_kwargs = MODEL_KWARGS[model_name]
@@ -93,6 +107,7 @@ def get_fa2_model_and_cfg(self, model_name: str, vocab_size: int) -> nn.Module:
             }
         )
         model = model_cls(cfg).to("cuda", dtype=self.dtype)
+        self.model = model
         return model, cfg
 
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="Padding free tests require CUDA")
@@ -113,8 +128,6 @@ def test_padding_free(self, model_name: str, loss_type: str) -> None:
         model.initialize_weights()
         pf_model = deepcopy(model)
 
-        inputs = torch.randint(cfg.vocab_size, size=(self.batch_size, self.seqlen), device="cpu")
-
         data = {
             0: {
                 "messages": [
@@ -195,3 +208,65 @@ def test_padding_free(self, model_name: str, loss_type: str) -> None:
             non_nan_grads.add(k)
         print(f"{non_nan_grads=}")
         print(f"{nan_grads=}")
+
+
+    def test_padding_free_dpo(self, model_name: str) -> None:
+        if model_name == "bamba" and not mamba_and_causal_conv_available:
+            pytest.skip("bamba padding-free tests require mamba_ssm and causal_conv1d")
+        torch.manual_seed(42)
+
+        tokenizer = AutoTokenizer.from_pretrained("ibm-ai-platform/Bamba-9B-v2")
+        tokenizer.add_special_tokens({"pad_token": "<pad>"})
+        tokenizer.chat_template = CHAT_TEMPLATES["tulu"]
+        vocab_size = len(tokenizer)
+
+        model, cfg = self.get_fa2_model_and_cfg(model_name, vocab_size)
+        model.initialize_weights()
+        pf_model = deepcopy(model)
+
+        data = {
+            0: {
+                "chosen": [
+                    {"role": "user", "content": "Why did the chicken cross the road?"},
+                    {"role": "assistant", "content": "To get to the other side"},
+                ],
+                "rejected": [
+                    {"role": "user", "content": "Why did the chicken cross the road?"},
+                    {"role": "assistant", "content": "To make friends with a cow. Trying to make a long response to simulate hallucination."},
+                ],
+            },
+            1: {
+                "chosen": [
+                    {"role": "user", "content": "What is one plus two?"},
+                    {"role": "assistant", "content": "The answer is 3"},
+                ],
+                "rejected": [
+                    {"role": "user", "content": "What is one plus two?"},
+                    {"role": "assistant", "content": "The answer is 12"},
+                ]
+            },
+        }
+
+        tok_data = {k: preference_span_search_mask_out(v, tokenizer, max_seq_length=2**30) for k, v in data.items()}
+
+        collate_fn = DataCollatorForSeq2SeqDPO(tokenizer=tokenizer, model=model, padding="longest")
+        dataloader = DataLoader(tok_data, shuffle=False, collate_fn=collate_fn, batch_size=self.batch_size)
+
+        pf_collate_fn = TensorDataCollatorWithFlatteningDPO()
+        pf_dataloader = DataLoader(tok_data, shuffle=False, collate_fn=pf_collate_fn, batch_size=self.batch_size)
+
+        batch = next(iter(dataloader))
+        pf_batch = next(iter(pf_dataloader))
+        for b in (batch, pf_batch):
+            for k in b:
+                if torch.is_tensor(b[k]):
+                    b[k] = b[k].cuda()
+
+        assert batch["chosen_input_ids"].shape[0] == 2
+        assert pf_batch["chosen_input_ids"].shape[0] == 1
+        chosen_logps, rejected_logps, _ = concatenated_forward(model, batch, average_log_prob=False)
+        pf_chosen_logos, pf_rejected_logps, _ = concatenated_forward(pf_model, pf_batch, average_log_prob=False, packing=True)
+
+        torch.testing.assert_close(chosen_logps, pf_chosen_logos, atol=1e-3, rtol=1e-3)
+        torch.testing.assert_close(rejected_logps, pf_rejected_logps, atol=1e-3, rtol=1e-3)
+