[Refactor] Simplify LLMEnv

ghstack-source-id: 9367bb1
Pull Request resolved: #2897

Author: Vincent Moens

test/mocking_classes.py

@@ -2478,6 +2478,8 @@ def _step(
 
 class DummyStrDataLoader:
     def __init__(self, batch_size=0):
+        if isinstance(batch_size, tuple):
+            batch_size = torch.Size(batch_size).numel()
         self.batch_size = batch_size
 
     def generate_random_string(self, length=10):
@@ -2489,21 +2491,25 @@ def __iter__(self):
 
     def __next__(self):
         if self.batch_size == 0:
-            return self.generate_random_string()
+            return {"text": self.generate_random_string()}
         else:
-            return [self.generate_random_string() for _ in range(self.batch_size)]
+            return {
+                "text": [self.generate_random_string() for _ in range(self.batch_size)]
+            }
 
 
 class DummyTensorDataLoader:
     def __init__(self, batch_size=0, max_length=10, padding=False):
+        if isinstance(batch_size, tuple):
+            batch_size = torch.Size(batch_size).numel()
         self.batch_size = batch_size
         self.max_length = max_length
         self.padding = padding
 
     def generate_random_tensor(self):
         """Generate a tensor of random int64 values."""
         length = random.randint(1, self.max_length)
-        rt = torch.randint(0, 100, (length,))
+        rt = torch.randint(1, 10000, (length,))
         return rt
 
     def pad_tensor(self, tensor):
@@ -2517,11 +2523,12 @@ def __iter__(self):
     def __next__(self):
         if self.batch_size == 0:
             tensor = self.generate_random_tensor()
-            return self.pad_tensor(tensor) if self.padding else tensor
+            tokens = self.pad_tensor(tensor) if self.padding else tensor
         else:
             tensors = [self.generate_random_tensor() for _ in range(self.batch_size)]
             if self.padding:
                 tensors = [self.pad_tensor(tensor) for tensor in tensors]
-                return torch.stack(tensors)
+                tokens = torch.stack(tensors)
             else:
-                return tensors
+                tokens = tensors
+        return {"tokens": tokens, "attention_mask": tokens != 0}

test/test_actors.py

@@ -1361,37 +1361,55 @@ def test_vllm_batch_run(self, pad, generate, use_tensorclass, vllm_instance):
         else:
             assert isinstance(tokens, list)
 
-    def test_vllm_collection(self, vllm_instance):
+    @pytest.mark.parametrize("from_text", [True])
+    def test_vllm_collection(self, vllm_instance, from_text):
         policy = vLLMWrapper(
             vllm_instance,
             return_log_probs=True,
-            generate_kwargs={"max_tokens": 10},
+            generate_kwargs={"max_tokens": 32},
+            from_text=from_text in (True, None),
         )
-        self._run_check_collector(policy)
+        tokenizer = vllm_instance.get_tokenizer()
+        self._run_check_collector(policy, from_text=from_text, tokenizer=tokenizer)
 
     def test_transformers_collection(self):
         ...
 
     @classmethod
-    def env_constructor(cls):
-        dl = DummyStrDataLoader(batch_size=32)
-        env = LLMEnv.from_dataloader(
-            dl,
-            batch_size=16,
-            repeats=4,
-            # str2str=True, group_repeats=True
-        )
-        assert env.batch_size == (64,)
-        return env
+    def env_constructor(cls, **kwargs):
+        def make():
+            # if kwargs.get("from_text", True):
+            dl = DummyStrDataLoader(batch_size=32)
+            # else:
+            #     dl = DummyTensorDataLoader(batch_size=32)
+            env = LLMEnv.from_dataloader(
+                dl,
+                batch_size=4,
+                repeats=4,
+                **kwargs,
+            )
+            assert env.batch_size == (16,)
+            return env
+
+        return make
 
-    def _run_check_collector(self, policy):
+    def _run_check_collector(self, policy, from_text, tokenizer):
+        if from_text is None:
+            kwargs = {"eos_token_id": tokenizer.eos_token_id}
+        else:
+            kwargs = {
+                "from_text": from_text,
+                "tokenizer": tokenizer,
+                "eos_token_id": tokenizer.eos_token_id,
+            }
         collector = SyncDataCollector(
-            self.env_constructor,
+            self.env_constructor(**kwargs),
             policy=policy,
-            frames_per_batch=128,
-            total_frames=512,
+            frames_per_batch=32,
+            total_frames=128,
             use_buffers=False,
         )
+        t = 0
         for data in collector:
             assert isinstance(data, LazyStackedTensorDict)
             assert isinstance(data.reshape(-1).get("text_response"), NonTensorStack)
@@ -1403,6 +1421,10 @@ def _run_check_collector(self, policy):
             assert ("next", "text") in data
             # tokens
             assert "tokens" in data
+
+            t += data.numel()
+            assert collector._frames == t
+            assert t < 512, t
             # assert ("next", "tokens") in data
 
     def test_vllm_generate_multiple_trajs(self, vllm_instance):

test/test_collector.py

@@ -3636,9 +3636,10 @@ def _run_collector_test(self, total_steps, rb, policy, tokenizer):
 
         env = LLMEnv.from_dataloader(
             dataloader=dataloader,
-            str2str=True,
+            from_text=True,
             batch_size=bsz,
             group_repeats=True,
+            eos_token_id=tokenizer.eos_token_id,
         )
         if rb:
             rb = ReplayBuffer(storage=LazyStackStorage(max_size=total_steps * 2))
@@ -3695,7 +3696,7 @@ async def test_llm_collector_start(self, vllm_instance):
 
         env = LLMEnv.from_dataloader(
             dataloader=dataloader,
-            str2str=True,
+            from_text=True,
             batch_size=bsz,
             group_repeats=True,
         )
@@ -3748,7 +3749,7 @@ def test_llm_collector_completed(
 
         env = LLMEnv.from_dataloader(
             dataloader=dataloader,
-            str2str=True,
+            from_text=True,
             batch_size=bsz,
             group_repeats=True,
             eos_token_id=tokenizer.eos_token_id,
@@ -3854,7 +3855,7 @@ def test_llm_collector_completed_async(
         def env_maker():
             env = LLMEnv.from_dataloader(
                 dataloader=dataloader,
-                str2str=True,
+                from_text=True,
                 batch_size=(),
                 group_repeats=True,
                 eos_token_id=tokenizer.eos_token_id,

test/test_cost.py

@@ -16708,7 +16708,8 @@ def test_hf(self, from_text):
             dl,
             tokenizer=tokenizer if not from_text else None,
             batch_size=(32,),
-            str2str=True,
+            from_text=True,
+            eos_token_id=tokenizer.eos_token_id,
         )
 
         class RewardTransform(Transform):