Vmap mlp ensemble zero grads after update

username: hal9000
cooked: <p>I’m quite sure I did something wrong, but I’m not really sure what went wrong.</p>
<p>Basically there are no grads on the parameters after the update (optimizer.step()) and the loss doesn’t minimize ( I test it using a trivial mapping from a x ~ N(0,1) to y ~ N(0,1) where both are fixed before training)<br>
the code I use is:</p>
<pre><code class="lang-auto">import torch 
import torch.nn as nn 
from torch import vmap 
from architectures import mlp 
from torch.func import stack_module_state,functional_call
import copy 
from tqdm import tqdm 
from torch import optim 


class network(nn.Module):
    def __init__(self,embedding_size:int,num_modules:int):
        super().__init__()    
        self.ensamble_modules = nn.ModuleList([mlp([embedding_size,embedding_size]) for _ in range(num_modules)])

    def forward(self,x:torch.Tensor) -&gt; torch.Tensor:
        base_model = copy.deepcopy(self.ensamble_modules[0]).to(device='meta')

        ensemble_params, ensemble_buffers = stack_module_state(self.ensamble_modules)
        ensemble_call = lambda ensemble_params, ensemble_buffers,x: functional_call(base_model,(ensemble_params, ensemble_buffers),x)

        return vmap(ensemble_call)(ensemble_params, ensemble_buffers,x) 


if __name__ == "__main__":
    BATCH_SIZE = 1
    EMBEDDING_SIZE = 10
    NUM_MODULES = 4
    x = torch.randn(NUM_MODULES,BATCH_SIZE,EMBEDDING_SIZE)
    y = torch.randn(NUM_MODULES,BATCH_SIZE,EMBEDDING_SIZE)
    net = network(EMBEDDING_SIZE,NUM_MODULES)

    optimizer = optim.Adam(net.parameters(),lr=0.001)
    
    for _ in tqdm(range(100000),ascii=True):
        optimizer.zero_grad()
        y_pred =  net(x)
        loss = nn.MSELoss()(y_pred,y)
        loss.backward()
        print(loss.item())
        optimizer.step()

</code></pre>

username: hal9000
cooked: <p>architectures.py  is</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import math

class mlp(nn.Module):
    def __init__(self,struct:list,act_f:nn.Module = nn.ReLU,add_layer_norm:bool=True) -&gt; None:
        super().__init__()
        self.struct = struct 
        self.act_f = act_f
        self.add_layer_norm = add_layer_norm
        self.layers = self._build()

    def _build(self) -&gt; nn.Sequential:
        layers = []
        for index,(i,j) in enumerate(zip(self.struct[:-1],self.struct[1:])):
            layers.append(nn.Linear(i,j))
            if(not (index == len(self.struct)-2)):#if not last layer
                if(self.add_layer_norm):
                    layers.append(nn.LayerNorm(j))
                layers.append(self.act_f())    
        return nn.Sequential(*layers)

    def forward(self,x:torch.Tensor) -&gt; torch.Tensor:
        return self.layers(x)
    
</code></pre>

username: hal9000
cooked: <aside class="quote no-group" data-username="hal9000" data-post="1" data-topic="220649">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/hal9000/48/27270_2.png" class="avatar"> hal9000:</div>
<blockquote>
<p><code>stack_module_state</code></p>
</blockquote>
</aside>
<p>I’m an idiot, I just see in the documentation this:</p>
<aside class="onebox allowlistedgeneric" data-onebox-src="https://docs.pytorch.org/docs/stable/generated/torch.func.stack_module_state.html">
  <header class="source">

      <a href="https://docs.pytorch.org/docs/stable/generated/torch.func.stack_module_state.html" target="_blank" rel="noopener nofollow ugc">docs.pytorch.org</a>
  </header>

  <article class="onebox-body">
    

<h3><a href="https://docs.pytorch.org/docs/stable/generated/torch.func.stack_module_state.html" target="_blank" rel="noopener nofollow ugc">torch.func.stack_module_state — PyTorch 2.7 documentation</a></h3>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>Given a list of <code>M</code> <code>nn.Modules</code> of the same class, returns two dictionaries that stack all of their parameters and buffers together, indexed by name. The stacked parameters are optimizable (i.e. they are new leaf nodes in the autograd history that are unrelated to the original parameters and can be passed directly to an optimizer).</p>