[BUG] `LossModule.convert_to_functional` can NOT correctly deal with a `Linear` module in the PyG lib

[Sefank]

Describe the bug

LossModule.convert_to_functional can NOT correctly deal with a Linear module in the PyG lib.

As some modules in PyG lib depend on the Linear module of the PyG version to work, if LossModule.convert_to_functional fails to work with it, it would cause severe incompatibility with the PyG lib.

To Reproduce

This example is mostly copied from the TorchRL's documentation, with only a substitution of torch.nn.Linear with torch_geometric.nn.dense.Linear.

Code

import torch
from torch import nn
from torch_geometric.nn.dense import Linear
from torchrl.data.tensor_specs import BoundedTensorSpec
from torchrl.modules.distributions.continuous import NormalParamWrapper, TanhNormal
from torchrl.modules.tensordict_module.actors import ProbabilisticActor, ValueOperator
from torchrl.modules.tensordict_module.common import SafeModule
from torchrl.objectives.ppo import PPOLoss
from tensordict.tensordict import TensorDict
n_act, n_obs = 4, 3
spec = BoundedTensorSpec(-torch.ones(n_act), torch.ones(n_act), (n_act,))
base_layer = Linear(n_obs, 5)
net = NormalParamWrapper(nn.Sequential(base_layer, Linear(5, 2 * n_act)))
module = SafeModule(net, in_keys=["observation"], out_keys=["loc", "scale"])
actor = ProbabilisticActor(
    module=module,
    distribution_class=TanhNormal,
    in_keys=["loc", "scale"],
    spec=spec)
module = nn.Sequential(base_layer, Linear(5, 1))
value = ValueOperator(
    module=module,
    in_keys=["observation"])
loss = PPOLoss(actor, value)
batch = [2, ]
action = spec.rand(batch)
data = TensorDict({"observation": torch.randn(*batch, n_obs),
        "action": action,
        "sample_log_prob": torch.randn_like(action[..., 1]),
        ("next", "done"): torch.zeros(*batch, 1, dtype=torch.bool),
        ("next", "terminated"): torch.zeros(*batch, 1, dtype=torch.bool),
        ("next", "reward"): torch.randn(*batch, 1),
        ("next", "observation"): torch.randn(*batch, n_obs),
    }, batch)
loss(data)

Traceback Info

Traceback (most recent call last):
  File "<project_root>/issuse.py", line 35, in <module>
    loss(data)
  File "<project_root>/venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<project_root>/venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1568, in _call_impl
    result = forward_call(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<project_root>/venv/lib/python3.11/site-packages/tensordict/_contextlib.py", line 126, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "<project_root>/venv/lib/python3.11/site-packages/tensordict/nn/common.py", line 282, in wrapper
    return func(_self, tensordict, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<project_root>/venv/lib/python3.11/site-packages/torchrl/objectives/ppo.py", line 433, in forward
    self.value_estimator(
  File "<project_root>/venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<project_root>/venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<project_root>/venv/lib/python3.11/site-packages/torchrl/objectives/value/advantages.py", line 63, in new_func
    return fun(self, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<project_root>/venv/lib/python3.11/site-packages/torchrl/objectives/value/advantages.py", line 52, in new_fun
    return fun(self, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<project_root>/venv/lib/python3.11/site-packages/tensordict/nn/common.py", line 282, in wrapper
    return func(_self, tensordict, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<project_root>/venv/lib/python3.11/site-packages/torchrl/objectives/value/advantages.py", line 1224, in forward
    value, next_value = _call_value_nets(
                        ^^^^^^^^^^^^^^^^^
  File "<project_root>/venv/lib/python3.11/site-packages/torchrl/objectives/value/advantages.py", line 138, in _call_value_nets
    data_out = vmap(value_net, (0, 0))(data_in, params_stack)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<project_root>/venv/lib/python3.11/site-packages/torch/_functorch/apis.py", line 188, in wrapped
    return vmap_impl(func, in_dims, out_dims, randomness, chunk_size, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<project_root>/venv/lib/python3.11/site-packages/torch/_functorch/vmap.py", line 266, in vmap_impl
    return _flat_vmap(
           ^^^^^^^^^^^
  File "<project_root>/venv/lib/python3.11/site-packages/torch/_functorch/vmap.py", line 38, in fn
    return f(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^
  File "<project_root>/venv/lib/python3.11/site-packages/torch/_functorch/vmap.py", line 379, in _flat_vmap
    batched_outputs = func(*batched_inputs, **kwargs)
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<project_root>/venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<project_root>/venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<project_root>/venv/lib/python3.11/site-packages/tensordict/nn/functional_modules.py", line 572, in new_fun
    old_params = _assign_params(
                 ^^^^^^^^^^^^^^^
  File "<project_root>/venv/lib/python3.11/site-packages/tensordict/nn/functional_modules.py", line 649, in _assign_params
    return _swap_state(module, params, make_stateless, return_old_tensordict)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<project_root>/venv/lib/python3.11/site-packages/tensordict/nn/functional_modules.py", line 389, in _swap_state
    _old_value = _swap_state(
                 ^^^^^^^^^^^^
  File "<project_root>/venv/lib/python3.11/site-packages/tensordict/nn/functional_modules.py", line 389, in _swap_state
    _old_value = _swap_state(
                 ^^^^^^^^^^^^
  File "<project_root>/venv/lib/python3.11/site-packages/tensordict/nn/functional_modules.py", line 378, in _swap_state
    raise Exception(f"{model}\nhas no stateless attribute.")
Exception: Linear(3, 5, bias=False)
has no stateless attribute.

Looking into the traceback, it's found that the state attribute _is_stateless is missing, which should be correctly copied in LossModule.convert_to_functional:

rl/torchrl/objectives/common.py

Line 247 in bf264e0

functional_module = deepcopy(module)

However, Linear of the PyG version overrides the default deepcopy behavior:

This code is copied from PyG repo.

    def __deepcopy__(self, memo):
        out = Linear(self.in_channels, self.out_channels, self.bias
                     is not None, self.weight_initializer,
                     self.bias_initializer)
        if self.in_channels > 0:
            out.weight = copy.deepcopy(self.weight, memo)
        if self.bias is not None:
            out.bias = copy.deepcopy(self.bias, memo)
        return out

which will cause the extra attribute added to the module and not aware by the Linear module will be ignored.

Expected behavior

LossModule.convert_to_functional can correctly deal with a Linear module in the PyG lib, just as the standard Linear module in Pytorch.

System info

• TorchRL 0.2.0
• Torch 2.1.0
• Python 3.11.6 | packaged by conda-forge
• OS: WSL2 - Ubuntu 22.04 LTS

Checklist

• I have checked that there is no similar issue in the repo (required)
• I have read the documentation (required)
• I have provided a minimal working example to reproduce the bug (required)

[Sefank]

I don't know if it's common to require modules to be deep copied correctly along with externally added attributes. Maybe it would be better if we could find other ways to relax this requirement.

[vmoens]

Thanks for raising this.
We noticed the same thing with @matteobettini following #1576.

I think I can come up with a solution. This fits in a more general discussion about functional calls with tensordict and functional calls in general.

Context

We want to have modules that are parameter free (such that serializing them is cheap) and can be called with tensordict parameter as input. Long term the best thing will be to use torch.func.functional_call and I think we can make it work for v0.3.
Right now we use a trick from TensorDict where we patch the modules with a _is_functional and _is_stateless attributes, and call module(input, params). The forward method is wrapped in a function that scans the input for parameters and if found, places them within the module and executes it before taking them back.

Issues

This has several advantages. For instance, we can decorate any function in the module. However, it is a bit of a hack, and we want functional calls to be uniform within the PyTorch ecosystem. We're working on making TensorDict a possible input to functional_call, but that would require trnsordict to be integrated into core to be stable and reliable (without a ton of monkey patching). So until that happens we're stuck with our current method.

The functional call API was designed to take any TensorDict type as input. The decorated forward does these checks:

• if a params kwarg is passed, we pop it from the kwargs and populate the module
• if the module is stateless (=> check attribute) and the last input is a tensordict, we use it to repopulate the module. This is super clunky, as the last input could be something else (if you forgot to pass the tensordict) or the module could be stateful (and you'd still want to swap the old for the new params).

Solution (to this and other problems)

We have introduced a special parameter carrying container, TensorDictParams. What we could do is drop the private attributes in functional modules and just check the type of the last arg (if the params kwarg is not present). If found, we use it to repopulate. Only caveat is that I'm not sure of how to make this change in a non-bc-breaking manner (ie, still accept non TensorDictParams instances with a warning).

This will not completely solve your issue as the forward won't be declared anymore (because of how PyG deepcopy works). But we can patch __deepcopy__ during the functionalization too!

I will give it a shot next week and keep you posted.

Additional context

The move of trnsordict core features to PyTorch core is a proposal I want to make soon. I will ping you once I have an initial PR, any support or input would be appreciated!

[Sefank]

@vmoens Thanks so much for your help with this issue. Looking forward to the solution and if there's anything I can do to help, just let me know.

[janblumenkamp]

I'm interested in a solution to this problem too!

[vmoens]

i'll be working on this. The first step is to implement the to_module decorator in tensordict, but before that there are a couple of important PRs to merge. This will be part of 0.3, I promise!
ETA of this feature on nightlies in 2 weeks top!

#1612 (comment)

[vmoens]

I think #1707 and #1711 should solve this:

#1707 gets rid of the previous functional calls that relied on the clunky module(input, params=params) in favour of

with params.to_module(module):
    module(input)

I'm a bit puzzled about the fixes in #1711.
Here's the context:
we want the modules to be functionally called within the loss to allow for a bunch of things like ensemble learning (ie, a single forward pass for all q-value networks etc).
To do that we save the parameters on one side, and the module separately. We want to avoid for loss_module.parameters() to return a duplicated list of parameters. There are 2 solutions:

1. Save the module without nn.Module.__setattr__, directly in loss_module.__dict__:

   self.__dict__[module_name] = module

   That avoids allocating the module in loss_module._modules which exposes it to parameters(), modules() etc. The problem I have with that is that we loose track of the module content, ie. if we call loss_module.cuda(2) hoping that the memory on cuda(0) is freed, we'll still have some parameters sitting there. We should then patch loss_module to have some wacky behaviour that does dispatching but not lookups, which isn't really doing us a favour.
2. Use deepcopy to store a copy of the module with parameters empty of any content. The idea is to use the "meta" device + functionalization to swap the module parameters with contentless parameters:

   def _make_meta_params(param):
       is_param = isinstance(param, nn.Parameter)
   
       pd = param.detach().to("meta")
   
       if is_param:
           pd = nn.Parameter(pd, requires_grad=False)
       return pd
   with params.apply(_make_meta_params, device=torch.device("meta")).to_module(
       module
   ):
       # avoid buffers and params being exposed
       self.__dict__[module_name] = deepcopy(module)

   We would still store it in self.__dict__ directly (because otherwise we'll have some contentless parameters in the parameter list)
   Now I see that there is one main problem in this issue with deepcopy which is that we lose module attributes like _is_stateless but this isn't used anymore so that should work.
   If anyone's interested in providing feedback, I'd love to hear concerns about using deepcopy is the only thing we do is (a) populate module with "meta"-parameters and (2) deepcopy such module within the loss module. Am I right to assume that if we don't have any weird patch of the module occurring like it used to be the case, this will work ok?

[matteobettini]

Thanks for this discussion.

One thing that I am not fully grasping yet is how solution 2 mititgates the issues of solution 1.
If we are still setting the modules with self.__dict__[module_name] = deepcopy(module), these will not be present in calls like modules().

That said, to me it seems like we should avoid solution 1 as I would not introduce wacky behaviors in losses to customly treat only some operations.
The use of deepcopy in general should be fine and solution 2 would solve the issues we have with PyG. Although I do not know if in general there are torch libraries that do not admit deepcopies.

I have some more general questions:

1. is with params.to_module(module); module(input) the same as torch.func.functional_call(module, params, input)?
2. I assume there is no way to store a module structure without params in pytorch and that is why we are populating it with meta params as place holders? It would be cool to have some more flexibility in separating modules from their params from torch core as this may also allow to avoid deepcopying just to have a functional version

[vmoens]

One thing that I am not fully grasping yet is how solution 2 mititgates the issues of solution 1.

since the parameters have 0 content, we don't mind if calling loss_module.cuda() ignores them. Actually they must be ignored because if they were not calling param.to("meta").cuda() would fail (bc content has been erased).

is with params.to_module(module); module(input) the same as torch.func.functional_call(module, params, input)?

In this case yes but our solution allows you to call other methods than forward (like get_dist). We monkey-patch functional_call such that it does the same thing so the behaviour for forward is 100% identical, just slower.

I assume there is no way to store a module structure without params in pytorch and that is why we are populating it with meta params as place holders?

We could use None but we would loose the "parameter" type I believe, which isn't great.

It would be cool to have some more flexibility in separating modules from their params from torch core as this may also allow to avoid deepcopying just to have a functional version

Not sure how you can do that without deepcopying? What would the API look like?
There is a way of having a stateless module, which is exactly what we do. But if you don't deepcopy it you simply lose the content, which I don't think is very wise in general.
If I was writing a code for myself, I would not deepcopy, I would simply do

params = TensorDict.from_module(module)
# make module stateless
params.to('meta').to_module(module)
# carry on with code

that way params and module are completely separated, but module still knows what is a param and what is a buffer (which I don't think is the case if you replace them with None)

[matteobettini]

params = TensorDict.from_module(module)
# make module stateless
params.to('meta').to_module(module)
# carry on with code

Right and we cannot do this because otherwise you would make the collection policy functional as well, thus the added deepcopy.

I think this all makes sense for me.

since the parameters have 0 content, we don't mind if calling loss_module.cuda() ignores them. Actually they must be ignored because if they were not calling param.to("meta").cuda() would fail (bc content has been erased).

Here I was referring to loss_module.modules() which will not list any module if i understand correctly

[vmoens]

The goal is for them not to appear in modules() so that's not a bug, it's a feature :D