Update

[ghstack-poisoned]

Author: Vincent Moens

test/test_actors.py

@@ -3,19 +3,27 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 import argparse
+import importlib.util
 import os
 
 import pytest
 import torch
-
 from tensordict import NonTensorStack, TensorDict
 from tensordict.nn import CompositeDistribution, TensorDictModule
 from tensordict.nn.distributions import NormalParamExtractor
 
 from torch import distributions as dist, nn
 from torchrl.data import Binary, Bounded, Categorical, Composite, MultiOneHot, OneHot
+from torchrl.data.llm import LLMData
 from torchrl.data.llm.dataset import _has_transformers
-from torchrl.modules import from_hf_transformers, MLP, SafeModule, TanhDelta, TanhNormal
+from torchrl.modules import (
+    from_hf_transformers,
+    from_vllm,
+    MLP,
+    SafeModule,
+    TanhDelta,
+    TanhNormal,
+)
 from torchrl.modules.tensordict_module.actors import (
     _process_action_space_spec,
     ActorValueOperator,
@@ -37,6 +45,8 @@
     from _utils_internal import get_default_devices
     from mocking_classes import NestedCountingEnv
 
+_has_vllm = importlib.util.find_spec("vllm") is not None
+
 
 @pytest.mark.parametrize(
     "log_prob_key",
@@ -908,6 +918,7 @@ def test_lmhead_actorvalueoperator(device):
 
 
 @pytest.mark.skipif(not _has_transformers, reason="missing transformers dependencies")
+@pytest.mark.skipif(not _has_vllm, reason="missing vllm dependencies")
 class TestTransformerActor:
     @pytest.mark.parametrize(
         "from_text, generate, tokens, attention_mask",
@@ -924,7 +935,6 @@ class TestTransformerActor:
         ],
     )
     def test_from_hf_transformers(self, from_text, generate, tokens, attention_mask):
-        from torchrl.data.llm import LLMData
         from transformers import AutoTokenizer, GPT2Config, GPT2LMHeadModel
 
         tokenizer = AutoTokenizer.from_pretrained("gpt2")
@@ -934,26 +944,157 @@ def test_from_hf_transformers(self, from_text, generate, tokens, attention_mask)
         m = from_hf_transformers(
             model, tokenizer=tokenizer, from_text=from_text, generate=generate
         )
+        self._run_check(m, tokens, attention_mask, generate, from_text, has_logits=True)
+
+    @pytest.mark.parametrize(
+        "from_text, generate, tokens, attention_mask",
+        [
+            (True, True, None, None),
+            (True, False, None, None),
+            (
+                False,
+                True,
+                torch.randint(1024, (1, 10)),
+                torch.ones(1, 10, dtype=torch.int64),
+            ),
+            (False, True, torch.randint(1024, (1, 10)), None),
+        ],
+    )
+    def test_from_vllm(self, from_text, generate, tokens, attention_mask):
+        from vllm import LLM
+
+        model = LLM(model="facebook/opt-125m")
+        m = from_vllm(model, from_text=from_text, generate=generate)
+        self._run_check(
+            m, tokens, attention_mask, generate, from_text, has_logits=False
+        )
+
+    def _make_data(
+        self,
+        m,
+        tokens,
+        attention_mask,
+        generate,
+        from_text,
+        has_logits,
+        text_response=None,
+        tokens_response=None,
+    ):
+        lp_kwargs = {}
         if from_text:
-            tdin = LLMData(text=NonTensorStack("a text"), batch_size=1)
+            if not generate:
+                text_response = (
+                    NonTensorStack(" and another text that follows")
+                    if text_response is None
+                    else text_response
+                )
+                if not isinstance(text_response, NonTensorStack):
+                    if isinstance(text_response, list):
+                        text_response = NonTensorStack(*text_response)
+                    else:
+                        text_response = NonTensorStack(text_response)
+                lp_kwargs.update({"text_response": text_response})
+            tdin = LLMData(text=NonTensorStack("a text"), **lp_kwargs, batch_size=1)
         else:
-            tdin = LLMData(tokens=tokens, attention_mask=attention_mask, batch_size=1)
+            if not generate:
+                if tokens_response is None:
+                    shape_response = tokens.shape
+                    shape_response = shape_response[:-1] + (shape_response[-1] * 2,)
+                    tokens_response = torch.randint(1024, shape_response)
+                lp_kwargs.update({"tokens_response": tokens_response})
+            tdin = LLMData(
+                tokens=tokens, attention_mask=attention_mask, **lp_kwargs, batch_size=1
+            )
+        return tdin
+
+    def _run_check(self, m, tokens, attention_mask, generate, from_text, has_logits):
+        tdin = self._make_data(
+            m, tokens, attention_mask, generate, from_text, has_logits
+        )
+        if from_text and generate:
+            assert tdin.text_response is None
+        elif from_text and not generate:
+            assert tdin.text_response is not None
+
         td = m(tdin)
         assert td is tdin
         assert isinstance(td, LLMData)
         if from_text and generate:
             assert td.text_response is not None
-        else:
-            assert td.text_response is None
-        if attention_mask is not None or from_text:
-            assert td.attention_mask is not None
+        if generate and (attention_mask is not None or from_text):
+            assert td.attention_mask is not None, (generate, generate, from_text)
         else:
             assert td.attention_mask is None
         if not generate:
-            assert td.text_response is None
-            assert td.tokens_response is None
+            # logprobs are computed on text response of tokens_response
+            assert td.text_response is not None or td.tokens_response is not None
             assert td.log_probs is not None
-            assert td.logits is not None
+            if has_logits:
+                assert td.logits is not None
+
+        # Test the shapes
+        assert td.tokens_response is not None, (generate, has_logits, from_text)
+
+        # If from text and not generating, the tokens are not returned for now
+        if not (from_text and not generate):
+            assert td.tokens_response.shape[:-1] == td.tokens.shape[:-1]
+            # The convention is that the response only has new tokens
+            assert (
+                td.tokens_response[..., : td.tokens.shape[-1]]
+                != td.tokens[..., : td.tokens_response.shape[-1]]
+            ).any()
+
+    @pytest.mark.parametrize(
+        "from_text, tokens, attention_mask",
+        [
+            (True, None, None),
+            (
+                False,
+                torch.randint(1024, (1, 10)),
+                torch.ones(1, 10, dtype=torch.int64),
+            ),
+            (False, torch.randint(1024, (1, 10)), None),
+        ],
+    )
+    def test_from_vllm_logprobs(self, from_text, tokens, attention_mask):
+        from vllm import LLM
+
+        model = LLM(model="facebook/opt-125m")
+        m_generate = from_vllm(model, from_text=from_text, generate=True)
+        m_logprobs = from_vllm(model, from_text=from_text, generate=False)
+        self._check_lps(
+            m_generate, m_logprobs, tokens, attention_mask, from_text, has_logits=False
+        )
+
+    def _check_lps(
+        self,
+        model_generate,
+        model_logprobs,
+        tokens,
+        attention_mask,
+        from_text,
+        has_logits,
+    ):
+        # Checks that the log-probs gathered with generate=False equate those with generate=True
+        tdin_genetate = self._make_data(
+            model_generate, tokens, attention_mask, True, from_text, has_logits
+        )
+        td_generate = model_generate(tdin_genetate)
+        tdin_logprobs = self._make_data(
+            model_logprobs,
+            tokens,
+            attention_mask,
+            False,
+            from_text,
+            has_logits,
+            tokens_response=td_generate.tokens_response,
+            text_response=td_generate.text_response,
+        )
+        td_logprobs = model_logprobs(tdin_logprobs)
+        print(td_generate.log_probs / td_logprobs.log_probs)
+        torch.testing.assert_close(
+            td_generate.log_probs, td_logprobs.log_probs, rtol=1e-2, atol=1e-2
+        )
 
 
 if __name__ == "__main__":

test/test_env.py

@@ -1692,6 +1692,34 @@ def test_parallel_env_device(
             env_serial.close(raise_if_closed=False)
             env0.close(raise_if_closed=False)
 
+    @pytest.mark.skipif(not _has_gym, reason="no gym")
+    @pytest.mark.parametrize("env_device", [None, "cpu"])
+    def test_parallel_env_device_vs_no_device(self, maybe_fork_ParallelEnv, env_device):
+        def make_env() -> GymEnv:
+            env = GymEnv(PENDULUM_VERSIONED(), device=env_device)
+            return env.append_transform(DoubleToFloat())
+
+        # Rollouts work with a regular env
+        parallel_env = maybe_fork_ParallelEnv(
+            num_workers=1, create_env_fn=make_env, device=None
+        )
+        parallel_env.reset()
+        parallel_env.set_seed(0)
+        torch.manual_seed(0)
+
+        parallel_rollout = parallel_env.rollout(max_steps=10)
+
+        # Rollout doesn't work with Parallelnv
+        parallel_env = maybe_fork_ParallelEnv(
+            num_workers=1, create_env_fn=make_env, device="cpu"
+        )
+        parallel_env.reset()
+        parallel_env.set_seed(0)
+        torch.manual_seed(0)
+
+        parallel_rollout_cpu = parallel_env.rollout(max_steps=10)
+        assert_allclose_td(parallel_rollout, parallel_rollout_cpu)
+
     @pytest.mark.skipif(not _has_gym, reason="no gym")
     @pytest.mark.flaky(reruns=3, reruns_delay=1)
     @pytest.mark.parametrize(

test/test_storage_map.py

@@ -350,6 +350,17 @@ def test_edges(self):
         edges_check = {(0, 1), (0, 2), (1, 3), (1, 4), (2, 5), (2, 6)}
         assert edges == edges_check
 
+    def test_make_node(self):
+        td = TensorDict({"obs": torch.tensor([0])})
+        tree = Tree(node_data=td)
+        assert tree.node_data is not None
+
+        tree = Tree.make_node(data=td)
+        assert tree.node_data is not None
+
+        tree = Tree.make_node(td)
+        assert tree.node_data is not None
+
 
 class TestMCTSForest:
     def dummy_rollouts(self) -> Tuple[TensorDict, ...]:

torchrl/_utils.py

@@ -18,7 +18,6 @@
 import warnings
 from contextlib import nullcontext
 from copy import copy
-from distutils.util import strtobool
 from functools import wraps
 from importlib import import_module
 from typing import Any, Callable, cast, TypeVar
@@ -35,6 +34,21 @@
 except ImportError:
     from torch._dynamo import is_compiling
 
+
+def strtobool(val: Any) -> bool:
+    """Convert a string representation of truth to a boolean.
+
+    True values are 'y', 'yes', 't', 'true', 'on', and '1'; false values are 'n', 'no', 'f', 'false', 'off', and '0'.
+    Raises ValueError if 'val' is anything else.
+    """
+    val = val.lower()
+    if val in ("y", "yes", "t", "true", "on", "1"):
+        return True
+    if val in ("n", "no", "f", "false", "off", "0"):
+        return False
+    raise ValueError(f"Invalid truth value {val!r}")
+
+
 LOGGING_LEVEL = os.environ.get("RL_LOGGING_LEVEL", "INFO")
 logger = logging.getLogger("torchrl")
 logger.setLevel(getattr(logging, LOGGING_LEVEL))

torchrl/data/map/tree.py

@@ -122,7 +122,7 @@ def make_node(
         return cls(
             count=torch.zeros(()),
             wins=torch.zeros(()),
-            node=data.exclude("action", "next"),
+            node_data=data.exclude("action", "next"),
             rollout=rollout,
             subtree=subtree,
             device=device,