is it possible that from_module and to_module doesn't work well for LSTM modules?

[roger-creus]

Hello!

I am using this logic in my codebase, inspired from this

self.agent_inference = type(self.policy)(**inference_kwargs)
self.agent_inference_p = from_module(self.policy).data
self.agent_inference_p.to_module(self.agent_inference)

and while this works well for all policy classes in my project, if there is an LSTM inside the policy, then it stops working.

I have verified this is the only reason why the LSTM policies don't work, because when I comment out these 3 lines and use self.policy instead of self.agent_inference, the LSTM-equipped agents actually learn well.

What is going on?

[vmoens]

Hello
This works ok on my side:

import torch
from torch import nn
from tensordict import TensorDict

m1 = torch.nn.Sequential(
    nn.Linear(10, 10),
    nn.LSTM(10, 10, 2),
)

m2 = torch.nn.Sequential(
    nn.Linear(10, 10),
    nn.LSTM(10, 10, 2),
)

x = torch.randn(10, 10, 10)
y1= m1(x)
with TensorDict.from_module(m1).to_module(m2):
    y2 = m2(x)

torch.testing.assert_close(y1, y2)

but I know LSTM has this thing that they flatten params so maybe that's the issue?
Do you have a MRE?

[roger-creus]

thanks for your reply! I think I found it, it was the .data in the second line above. That seems to break things if the module has an LSTM in it

[vmoens]

ah yeah makes sense! And I think we to that somewhere in the data collector in torchrl. Maybe we should create a param that does not require a gradient instead!

[roger-creus]

:)

One additional question: can we expect full reproducibility when using compile and cudagraphs if we fix seeds? or there will always be some form of non-determinism?

Related question: should the exact same code, but compiling+cudagraphs vs not doing it yield the exact same results?

EDIT: See attached image, the change above fixed it, so now the lstm-policy trains, but when compile+cudagraphs is enabled it will hit NaNs in the loss and gradients at some point in training and collapse. Without compile and cudagraphs, the LSTM policies train nicely as you can see with the fix above.

[vmoens]

What do you call "compile + cudagraph"?
cudagraph wrapper from tensordict + compile, or compile and cudagraph (with reduce_overhead=True)?
The latter is safer but slower, the former will just assume that your graph can be captured as is the first time it is executed. If you have control flow or anything like that it could break badly. It's a dirty trick to get your code to run very fast without safety net.

[roger-creus]

I see! What I have been using is:

from tensordict.nn import CudaGraphModule
self.policy_fn = self.agent_inference.get_action_and_value
mode = "reduce-overhead" if not cudagraphs_enabled else None
self.policy_fn = torch.compile(self.policy_fn, mode=None)
self.policy_fn = CudaGraphModule(self.policy_fn, warmup=20)

but the results above were run with compile_enabled=True and cudagraphs_enabled=True so mode=None

So you suggest I should use mode = "reduce-overhead" for better safety? I don't recall doing any explicit control flow... what kind of operations would potentially break things in the backend?

[vmoens]

Yes I would try with reduce overhead and compare the results with the plot above!