diff --git a/test/modules/test_rnn.py b/test/modules/test_rnn.py index 5a1c4a90e39..27d9967938b 100644 --- a/test/modules/test_rnn.py +++ b/test/modules/test_rnn.py @@ -662,6 +662,10 @@ def test_lstm_vmap_complex_model(self): in_key="embed", out_key="features", python_based=True, + # vmap cannot trace through ``torch._higher_order_ops.scan``; the + # 'pad' backend keeps the time loop as a plain Python call into + # the Python-based LSTM, which is fully vmap-compatible. + recurrent_backend="pad", ) mlp = TensorDictModule( MLP( @@ -2004,6 +2008,10 @@ def test_gru_vmap_complex_model(self): in_key="embed", out_key="features", python_based=True, + # vmap cannot trace through ``torch._higher_order_ops.scan``; the + # 'pad' backend keeps the time loop as a plain Python call into + # the Python-based GRU, which is fully vmap-compatible. + recurrent_backend="pad", ) mlp = TensorDictModule( MLP( diff --git a/torchrl/modules/tensordict_module/rnn.py b/torchrl/modules/tensordict_module/rnn.py index dbe4a1bf96a..70aef7dd12f 100644 --- a/torchrl/modules/tensordict_module/rnn.py +++ b/torchrl/modules/tensordict_module/rnn.py @@ -15,7 +15,7 @@ from tensordict import TensorDict, TensorDictBase, unravel_key_list from tensordict.base import NO_DEFAULT from tensordict.nn import dispatch, TensorDictModuleBase as ModuleBase -from tensordict.utils import expand_as_right, prod, set_lazy_legacy +from tensordict.utils import expand_as_right, set_lazy_legacy from torch import nn, Tensor from torch.nn.modules.rnn import RNNCellBase @@ -915,21 +915,24 @@ def forward(self, tensordict: TensorDictBase): # we want to get an error if the value input is missing, but not the hidden states defaults = [NO_DEFAULT, None, None] shape = tensordict.shape - tensordict_shaped = tensordict if self.recurrent_mode: - # if less than 2 dims, unsqueeze - ndim = tensordict_shaped.get(self.in_keys[0]).ndim - while ndim < 3: - tensordict_shaped = tensordict_shaped.unsqueeze(0) - ndim += 1 - if ndim > 3: - dims_to_flatten = ndim - 3 - # we assume that the tensordict can be flattened like this - nelts = prod(tensordict_shaped.shape[: dims_to_flatten + 1]) - tensordict_shaped = tensordict_shaped.apply( - lambda value: value.flatten(0, dims_to_flatten), - batch_size=[nelts, tensordict_shaped.shape[-1]], + # Straight-line shape normalization. Time is the last batch dim; + # all earlier batch dims are folded into a single leading B. + # Cheaper and simpler than the historical ``while ndim < 3`` loop + # plus ``prod(...)`` + ``apply(..., batch_size=[...])``. + td_ndim = tensordict.ndim + if td_ndim == 0: + raise ValueError( + "LSTMModule(recurrent_mode=True) requires the input " + "tensordict to have at least one batch dim (time). Got a " + "0-d tensordict." ) + elif td_ndim == 1: + tensordict_shaped = tensordict.unsqueeze(0) + elif td_ndim == 2: + tensordict_shaped = tensordict + else: + tensordict_shaped = tensordict.flatten(0, -2) else: tensordict_shaped = tensordict.reshape(-1).unsqueeze(-1) @@ -937,6 +940,9 @@ def forward(self, tensordict: TensorDictBase): splits = None backend = self.recurrent_backend if backend == "auto": + # In eager, CuDNN-backed pad is the fastest path; under torch.compile + # the data-dependent ``_split_and_pad_sequence`` branch is unfriendly, + # so prefer scan there. backend = "scan" if is_compiling() else "pad" use_scan = self.recurrent_mode and backend == "scan" use_triton = self.recurrent_mode and backend == "triton" @@ -948,12 +954,10 @@ def forward(self, tensordict: TensorDictBase): ): from torchrl.objectives.value.utils import _get_num_per_traj_init - # if we have consecutive trajectories, things get a little more complicated - # we have a tensordict of shape [B, T] - # we will split / pad things such that we get a tensordict of shape - # [N, T'] where T' <= T and N >= B is the new batch size, such that - # each index of N is an independent trajectory. We'll need to keep - # track of the indices though, as we want to put things back together in the end. + # Multi-trajectory rollouts under the pad backend: split each row + # into per-trajectory windows of shape [N, T'], run the LSTM on + # the padded result, then stitch them back. Required for correctness + # whenever ``is_init`` fires mid-row. splits = _get_num_per_traj_init(is_init) tensordict_shaped_shape = tensordict_shaped.shape tensordict_shaped = _split_and_pad_sequence( @@ -999,7 +1003,6 @@ def forward(self, tensordict: TensorDictBase): tensordict_shaped.set(self.out_keys[1], hidden0) tensordict_shaped.set(self.out_keys[2], hidden1) if splits is not None: - # let's recover our original shape tensordict_shaped = _inv_pad_sequence(tensordict_shaped, splits).reshape( tensordict_shaped_shape ) @@ -2120,21 +2123,21 @@ def forward(self, tensordict: TensorDictBase): # we want to get an error if the value input is missing, but not the hidden states defaults = [NO_DEFAULT, None] shape = tensordict.shape - tensordict_shaped = tensordict if self.recurrent_mode: - # if less than 2 dims, unsqueeze - ndim = tensordict_shaped.get(self.in_keys[0]).ndim - while ndim < 3: - tensordict_shaped = tensordict_shaped.unsqueeze(0) - ndim += 1 - if ndim > 3: - dims_to_flatten = ndim - 3 - # we assume that the tensordict can be flattened like this - nelts = prod(tensordict_shaped.shape[: dims_to_flatten + 1]) - tensordict_shaped = tensordict_shaped.apply( - lambda value: value.flatten(0, dims_to_flatten), - batch_size=[nelts, tensordict_shaped.shape[-1]], + # Straight-line shape normalization (see LSTMModule.forward). + td_ndim = tensordict.ndim + if td_ndim == 0: + raise ValueError( + "GRUModule(recurrent_mode=True) requires the input " + "tensordict to have at least one batch dim (time). Got a " + "0-d tensordict." ) + elif td_ndim == 1: + tensordict_shaped = tensordict.unsqueeze(0) + elif td_ndim == 2: + tensordict_shaped = tensordict + else: + tensordict_shaped = tensordict.flatten(0, -2) else: tensordict_shaped = tensordict.reshape(-1).unsqueeze(-1) @@ -2153,12 +2156,6 @@ def forward(self, tensordict: TensorDictBase): ): from torchrl.objectives.value.utils import _get_num_per_traj_init - # if we have consecutive trajectories, things get a little more complicated - # we have a tensordict of shape [B, T] - # we will split / pad things such that we get a tensordict of shape - # [N, T'] where T' <= T and N >= B is the new batch size, such that - # each index of N is an independent trajectory. We'll need to keep - # track of the indices though, as we want to put things back together in the end. splits = _get_num_per_traj_init(is_init) tensordict_shaped_shape = tensordict_shaped.shape tensordict_shaped = _split_and_pad_sequence( @@ -2190,7 +2187,6 @@ def forward(self, tensordict: TensorDictBase): tensordict_shaped.set(self.out_keys[0], val) tensordict_shaped.set(self.out_keys[1], hidden) if splits is not None: - # let's recover our original shape tensordict_shaped = _inv_pad_sequence(tensordict_shaped, splits).reshape( tensordict_shaped_shape )