[Feature] Packing

test/llm/test_wrapper.py

@@ -12,7 +12,7 @@
 
 import pytest
 import torch
-from tensordict import lazy_stack, set_list_to_stack, TensorDict
+from tensordict import assert_close, lazy_stack, set_list_to_stack, TensorDict
 
 from tensordict.utils import _zip_strict
 from torchrl.data.llm import History
@@ -163,6 +163,22 @@ def sample_tokens(vllm_instance):
     return tokenized["input_ids"], tokenized["attention_mask"]
 
 
+@pytest.fixture
+def sample_tokens_unpadded(vllm_instance):
+    """Create sample tokens for testing."""
+    model, tokenizer = vllm_instance
+    text = [
+        "Are you happy? Say yes or no.",
+        "Explain the difference between a cat and a dog. Be very detailed.",
+    ]
+    tokenized = tokenizer(text, padding=False)
+    return torch.nested.nested_tensor(
+        [torch.tensor(t) for t in tokenized["input_ids"]], layout=torch.jagged
+    ), torch.nested.nested_tensor(
+        [torch.tensor(t) for t in tokenized["attention_mask"]], layout=torch.jagged
+    )
+
+
 def check_output_shapes(out, pad_output, requested_log_probs=False):
     if pad_output:
         # We can get all tensors or they are none
@@ -1538,8 +1554,6 @@ def test_log_probs_consistency(
         vllm_lp_result = vllm_lp_wrapper(new_data.copy())
         tf_lp_result = tf_lp_wrapper(new_data.copy())
 
-        from tensordict import assert_close
-
         assert_close(
             vllm_lp_result, tf_lp_result, atol=1e-1, rtol=1e-1, intersection=True
         )
@@ -1707,6 +1721,100 @@ def test_transformers_custom_masking(
         assert hasattr(dist, "log_prob")
 
 
+@pytest.mark.skipif(not _has_transformers, reason="transformers not available")
+@pytest.mark.parametrize("pad_output", [False, True])
+class TestPacking:
+    def test_packing_history(
+        self, transformers_instance, sample_history_assistant, pad_output
+    ):
+        model, tokenizer = transformers_instance
+
+        wrapper_packed = TransformersWrapper(
+            model,
+            tokenizer=tokenizer,
+            input_mode="history",
+            generate=False,
+            return_log_probs=True,
+            pad_output=pad_output,
+            pad_model_input=False,
+        )
+        wrapped_padded = TransformersWrapper(
+            model,
+            tokenizer=tokenizer,
+            input_mode="history",
+            generate=False,
+            return_log_probs=True,
+            pad_output=pad_output,
+            pad_model_input=True,
+        )
+
+        td = TensorDict(
+            {"history": ChatHistory(full=sample_history_assistant)}, batch_size=(2,)
+        ).to_lazystack(0)
+
+        result_padded = wrapped_padded(td)
+        result_packed = wrapper_packed(td)
+        assert_close(result_packed["log_probs"], result_padded["log_probs"])
+
+    def test_packing_text(self, transformers_instance, sample_text, pad_output):
+        model, tokenizer = transformers_instance
+        wrapper_packed = TransformersWrapper(
+            model,
+            tokenizer=tokenizer,
+            input_mode="text",
+            generate=False,
+            return_log_probs=True,
+            pad_output=pad_output,
+            pad_model_input=False,
+        )
+        wrapped_padded = TransformersWrapper(
+            model,
+            tokenizer=tokenizer,
+            input_mode="text",
+            generate=False,
+            return_log_probs=True,
+            pad_output=pad_output,
+            pad_model_input=True,
+        )
+        td = TensorDict({"text": Text(full=sample_text)}, batch_size=(2,))
+        result_packed = wrapper_packed(td)
+        result_padded = wrapped_padded(td)
+        assert_close(result_packed["log_probs"], result_padded["log_probs"])
+
+    def test_packing_tokens(
+        self, transformers_instance, sample_tokens_unpadded, pad_output
+    ):
+        model, tokenizer = transformers_instance
+        wrapper_packed = TransformersWrapper(
+            model,
+            tokenizer=tokenizer,
+            input_mode="tokens",
+            generate=False,
+            return_log_probs=True,
+            pad_output=pad_output,
+            pad_model_input=False,
+        )
+        wrapped_padded = TransformersWrapper(
+            model,
+            tokenizer=tokenizer,
+            input_mode="tokens",
+            generate=False,
+            return_log_probs=True,
+            pad_output=pad_output,
+            pad_model_input=True,
+        )
+        td = TensorDict(
+            {
+                "tokens": Tokens(full=sample_tokens_unpadded[0]),
+                "masks": Masks(all_attention_mask=sample_tokens_unpadded[1]),
+            },
+            batch_size=(2,),
+        ).to_lazystack(0)
+        result_padded = wrapped_padded(td)
+        result_packed = wrapper_packed(td)
+        assert_close(result_packed["log_probs"], result_padded["log_probs"])
+
+
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()
     pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown)

torchrl/modules/llm/policies/common.py

@@ -328,6 +328,8 @@ class LLMWrapperBase(TensorDictModuleBase):
         generate_kwargs: Additional arguments to pass to the model's generate method.
         tokenizer_kwargs: Additional arguments to pass to the tokenizer.
         pad_output: Whether to pad the output sequences to a uniform length.
+        pad_model_input: Whether to pad the model input sequences to a uniform length.
+            May not be supported by all models.
         inplace: Determines how the module should handle in-place operations.
         device: The device to use for computation.
         layout: The layout to use for the output tensors when pad_output=False.