Update

[ghstack-poisoned]

Author: Vincent Moens

.github/unittest/linux_libs/scripts_llm/environment.yml

@@ -17,6 +17,6 @@ dependencies:
     - pyyaml
     - scipy
     - hydra-core
-    - transformers<4.42.0
+    - transformers
     - datasets
     - vllm

examples/rlhf/models/actor_critic.py

@@ -34,4 +34,4 @@ def init_actor_critic(model_cfg, sys_cfg):
     critic = model.get_value_operator()
     critic_head = model.get_value_head()
 
-    return actor, VmapModule(critic), critic_head, base_model
+    return actor, VmapModule(critic, mock=True), critic_head, base_model

torchrl/modules/tensordict_module/common.py

@@ -436,29 +436,38 @@ class VmapModule(TensorDictModuleBase):
         >>> assert (sample_in_td["x"][:, 0] == sample_in_td["y"]).all()
     """
 
-    def __init__(self, module: TensorDictModuleBase, vmap_dim=None):
+    def __init__(self, module: TensorDictModuleBase, vmap_dim=None, mock: bool = False):
         if not _has_functorch:
             raise ImportError("VmapModule requires torch>=2.0.")
         super().__init__()
         self.in_keys = module.in_keys
         self.out_keys = module.out_keys
         self.module = module
         self.vmap_dim = vmap_dim
+        self.mock = mock
         if torch.__version__ >= "2.0":
             self._vmap = torch.vmap
         else:
             import functorch
 
             self._vmap = functorch.vmap
 
+    def mock_(self, value: bool = True):
+        self.mock = value
+
     def forward(self, tensordict):
         # TODO: there is a risk of segfault if input is not a tensordict.
         # We should investigate (possibly prevent it c++ side?)
         vmap_dim = self.vmap_dim
         if vmap_dim is None:
             ndim = tensordict.ndim
             vmap_dim = ndim - 1
-        td = self._vmap(self.module, (vmap_dim,), (vmap_dim,))(tensordict)
+        if self.mock:
+            td = torch.stack(
+                [self.module(_td) for _td in tensordict.unbind(vmap_dim)], vmap_dim
+            )
+        else:
+            td = self._vmap(self.module, (vmap_dim,), (vmap_dim,))(tensordict)
         return tensordict.update(td)