Getting the "One of the variables needed for gradient computation has been modified by an inplace operation" Error while implementing PPO with a shared Module between actor and critic

Ebrahim.Pichka: <h1><a name="summary-of-the-problem-1" class="anchor" href="#summary-of-the-problem-1"></a>Summary of the problem</h1>
<p>I’m implementing PPO for a combinatorial optimization task. The Actor and Critic network share an Encoder network module at the beginning, and then each follows a different path of Modules for different Policy and Value outputs. Roughly speaking it would be like this:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/f/df8ca01a81d7c18bfc2071fccd1965451ac9cdec.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/df8ca01a81d7c18bfc2071fccd1965451ac9cdec" title="Page-1 (1)"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/d/f/df8ca01a81d7c18bfc2071fccd1965451ac9cdec_2_690x395.jpeg" alt="Page-1 (1)" data-base62-sha1="vTBLoxLn6Gut3FRFfuEgTB18My0" width="690" height="395" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/d/f/df8ca01a81d7c18bfc2071fccd1965451ac9cdec_2_690x395.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/d/f/df8ca01a81d7c18bfc2071fccd1965451ac9cdec_2_1035x592.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/d/f/df8ca01a81d7c18bfc2071fccd1965451ac9cdec_2_1380x790.jpeg 2x" data-dominant-color="F9F0E9"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Page-1 (1)</span><span class="informations">1640×941 139 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>While doing a round of updates on the Actor and Critic networks after an episode is finished, I get the following runtime error. After following</p>
<pre><code class="lang-auto">RuntimeError: one of the variables needed for gradient computation
has been modified by an inplace operation:
[torch.cuda.FloatTensor [128, 128]], which is output 0 of AsStridedBackward0,
is at version 2; expected version 1 instead.
Hint: the backtrace further above shows the operation that failed to compute its gradient.
The variable in question was changed in there or anywhere later. Good luck!
</code></pre>
<h1><a name="implementation-details-2" class="anchor" href="#implementation-details-2"></a>Implementation details</h1>
<p>The basic structure of the <strong>models</strong> is as follows:</p>
<pre data-code-wrap="py"><code class="lang-py">
#######################
#### ENCODER MODEL ####
#######################

class Encoder(nn.Module):
    def __init__(self, num_instances, embed_dim, device):
        super(Encoder, self).__init__()
        self.num_instances = num_instances      # n
        self.embed_dim = embed_dim              # d
        self.device = device

        self.lin1 = nn.Linear(num_instances, embed_dim)
        self.lin2 = nn.Linear(embed_dim, embed_dim)

    def forward(self, state):
        # state: [b, n, n]
        state_emb = F.relu(self.lin1(state))
        state_emb = F.relu(self.lin2(state_emb))
        # state_emb: [b, n, d]
        return state_emb 

#####################
#### ACTOR MODEL ####
#####################


class Actor(nn.Module):
    def __init__(self, shared_encoder: nn.Module, config):
        super(Actor, self).__init__()
        self.num_instances = config.num_instances      # n
        self.embed_dim = config.embed_dim              # d

        self.shared_encoder = shared_encoder
        self.device = config.device
            
        self.transformer_decoder = nn.TransformerEncoderLayer(
                d_model=config.embed_dim,
                nhead=config.num_heads,
                dim_feedforward=config.embed_dim*2,
                dropout=config.dropout,
                batch_first=True
                )

        self.linear = nn.Linear(config.embed_dim * config.num_instances, config.num_instances)

    def forward(self, state):
        # state: [b, n, n]

        batch_size = state.size(0)  # b
        n = state.size(1)           # n
        embed_dim = self.embed_dim  # d
     
        emb = self.shared_encoder(state) # -&gt; (b, n, d)

        # transform the embeddings using transformer decoder
        src_mask = torch.zeros(n, n).bool().to(self.device)
        x = self.transformer_decoder(emb, src_mask) # (b, n, d)
        x = x.view(batch_size, -1)  # (b, n*d)
        x = self.linear(x)          # (b, n)

        return x #  (b, n) =&gt; Policy logits to construct probabilities during training


######################
#### CRITIC MODEL ####
######################


class Critic(nn.Module):
    def __init__(self, shared_encoder: nn.Module, config):
        """ Critic network for the actor-critic algorithm. 
        Only uses Fully Connected layers after the shared encoder.
        """
        super(Critic, self).__init__()
        
        self.num_instances = config.num_instances
        self.embed_dim = config.embed_dim
        self.device = config.device

        self.shared_encoder = shared_encoder

        self.linear1 = nn.Linear(config.embed_dim * config.num_instances, config.embed_dim)
        self.linear2 = nn.Linear(config.embed_dim, config.embed_dim)
        self.linear3 = nn.Linear(config.embed_dim, 1)

    def forward(self, state):
        # state: [b, n, n]
        batch_size = state.size(0)
        state_emb = self.shared_encoder(state) # -&gt; (b, n, d)
        
        x = state_emb.view(batch_size, -1) # (b, n*d)
        x = self.linear1(x)
        x = F.relu(x)
        x = self.linear2(x)                 # (b, d)
        x = F.relu(x)
        x = self.linear3(x)                 # (b, 1)

        return x # (b, 1) =&gt; V(s)

</code></pre>
<p><br>
<br>
And the whole <strong>training procedure</strong> of the PPO algorithm is implemented this way:</p>
<pre data-code-wrap="py"><code class="lang-py">class PPO:
    def __init__(self, config):

        self.config = config
        self.n = config.num_instances
        self.batch_size = config.batch_size
        self.lr_actor = config.lr_actor
        self.lr_critic = config.lr_critic
        self.gamma = config.gamma
        self.clip_ratio = config.clip_ratio
        self.epochs = config.epochs

        self.shared_encoder = Encoder(self.n, config.embed_dim, config.dropout, config.device).to(config.device)
        self.actor = Actor(self.shared_encoder, config).to(config.device)
        self.critic = Critic(self.shared_encoder, config).to(config.device)
        

        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=self.lr_actor)
        self.critic_optimizer = torch.optim.Adam(self.critic.parameters(), lr=self.lr_critic)
        self.critic_loss_fn = nn.MSELoss()

        self.memory = []

    def get_action(self, state, action_mask):
        if isinstance(state, np.ndarray):
            state = torch.tensor(state, dtype=torch.float32).unsqueeze(0).to(self.config.device)
        logits = self.actor(state).squeeze()
        masked_logits = logits.detach().cpu() + (1 - np.squeeze(action_mask)) * -1e8
        dist = torch.distributions.Categorical(logits=masked_logits)
        action = dist.sample()
        log_prob = dist.log_prob(action)
        return action.item(), log_prob

    def store_transition(self, state, action, reward, next_state, done, log_prob):
        self.memory.append((state, action, reward, next_state, done, log_prob))

    def update(self):
        states, actions, rewards, next_states, dones, old_log_probs = zip(*self.memory)
        dev = self.config.device
        states = torch.tensor(states, dtype=torch.float32).to(dev)
        actions = torch.tensor(actions, dtype=torch.int64).view(-1, 1).to(dev)
        rewards = torch.tensor(rewards, dtype=torch.float32).view(-1, 1).to(dev)
        next_states = torch.tensor(next_states, dtype=torch.float32).to(dev)
        dones = torch.tensor(dones, dtype=torch.float32).view(-1, 1).to(dev)
        old_log_probs = torch.tensor(old_log_probs, dtype=torch.float32).view(-1, 1).to(dev)

        # Compute the advantages
        with torch.no_grad():
            target_values = rewards + self.gamma * (1 - dones) * self.critic(next_states)
            advantages = target_values - self.critic(states)

        # PPO update for the actor and critic
        for _ in range(self.epochs):
            logits = self.actor(states)
            dist = torch.distributions.Categorical(logits=logits)
            log_probs = dist.log_prob(actions)
            ratios = torch.exp(log_probs - old_log_probs)

            # Clipped surrogate objective for the policy fn
            actor_loss = -torch.min(ratios * advantages, torch.clamp(
                ratios, 1 - self.clip_ratio, 1 + self.clip_ratio) * advantages).mean()

                
            # Value function loss
            critic_loss = self.critic_loss_fn(self.critic(states), target_values)

            self.actor_optimizer.zero_grad()
            actor_loss.backward()
            self.actor_optimizer.step()

            self.critic_optimizer.zero_grad()
            critic_loss.backward()
            self.critic_optimizer.step()


    def train(self, env, num_episodes):
        self.actor.train()
        self.critic.train()
        for episode in range(num_episodes):
            state, _ = env.reset()
            done = False
            total_reward = 0

            while not done:
                action_mask = env.get_action_mask()
                action, log_prob = self.get_action(state, action_mask)
                next_state, reward, done, _, _ = env.step(action)
                self.store_transition(state, action, reward, next_state, done, log_prob)
                total_reward += reward
                state = next_state

            self.update()
            del self.memory[:]
            print(f"Episode: {episode + 1}, Total Reward: {total_reward}")
</code></pre>
<p><br>
<br>
Where the issue happens during the <code>update</code> method.</p>
<h2><a name="reproduction-3" class="anchor" href="#reproduction-3"></a>Reproduction</h2>
<p>For reproduction use such env and config:</p>
<pre data-code-wrap="py"><code class="lang-py">import numpy as np 
import torch 
import torch.nn as nn
import torch.nn.functional as F

# random env
class RndEnv:
    def __init__(self, n ):
        self.n = n
        self.step = 0
        
    def reset(self):
        return np.random.rand(n,n), {}

    def step(self, action):
        self.step += 1
        if self.step == self.n:
            return np.random.rand(n,n), np.random.rand(), True, False, {}
        return np.random.rand(n,n), {}, False, False, {}

    def get_action_mask(self):
        return np.ones(self.n, dtype=np.float32)

class TestConfig():
    """The config class."""
    num_instances = 7 # (n)
    train_episodes = 100
    test_episodes = 10
    embed_dim = 128
    num_heads = 8
    dropout = 0.1
    lr_actor = 1e-4
    lr_critic = 1e-4
    gamma = 1.0
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    batch_size = 64
    epochs = 10


config = TestConfig()
env = RndEnv(config.num_instances)

ppo = PPO(config)
ppo.train(env)


</code></pre>
<h1><a name="solutions-tried-so-far-4" class="anchor" href="#solutions-tried-so-far-4"></a>Solutions tried so far</h1>
<p>After using the <code>autograd.set_detect_anomaly(True)</code> it is shown that the error happened while computing the backward through the <code>self.critic</code> network for updating the critic optimizer. meaning at the <code>critic_loss.backward()</code> and it happened at the last linear layer of the <code>Encoder</code> meaning <code>state_emb = F.relu(self.lin2(state_emb))</code>.</p>
<p>Based on what said and reading similar questions to mine on this forum I have found that it is kind of similar to <a href="https://discuss.pytorch.org/t/one-of-the-variables-needed-for-gradient-computation-has-been-modified-by-an-inplace-operation-torch-cuda-floattensor-3-48-3-3-is-at-version-2-expected-version-1-instead/83241">this issue</a> , <a href="https://discuss.pytorch.org/t/runtimeerror-one-of-the-variables-needed-for-gradient-computation-has-been-modified-by-an-inplace-operation-torch-cuda-floattensor-is-at-version-1-expected-version-0-instead-hint-the-backtrace-further-above-shows-the-operation-that-failed-to-com/168485/3">this issue</a>, or maybe <a href="https://discuss.pytorch.org/t/runtimeerror-one-of-the-variables-needed-for-gradient-computation-has-been-modified-by-an-inplace-operation-torch-floattensor-3-1-which-is-output-0-of-tanhbackward-is-at-version-1-expected-version-0-instead/87630/8">this one.</a>. I tried cloning the states and advantages for each pass, but it didn’t work. I also checked for all in-place operations, but nothing was found.</p>
<p>I mainly suspect the issue is with the encoder being shared and the inputs being used twice. As I tried instantiating to separate encoders and use each for actor and critic, the training will work seamlessly. So the shared layer is causing the issue but I can’t find why.</p>
<p>I appreciate any help or suggestions.</p>
<h1><a name="full-error-message-5" class="anchor" href="#full-error-message-5"></a>Full error message</h1>
<pre data-code-wrap="py"><code class="lang-py">---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[3], line 11
      7 ppo_model = PPO(config)
      9 env = QAPEnv(config.num_instances)
---&gt; 11 ppo_model.train(env, config.train_episodes)

File c:\Users\...\trainer.py:117, in PPO.train(self, env, num_episodes)
    114     total_reward += reward
    115     state = next_state
--&gt; 117 self.update()
    118 del self.memory[:]
    119 print(f"Episode: {episode + 1}, Total Reward: {total_reward}")

File c:\Users\...\trainer.py:92, in PPO.update(self)
     89 self.actor_optimizer.step()
     91 self.critic_optimizer.zero_grad()
---&gt; 92 critic_loss.backward()
     93 self.critic_optimizer.step()
     95 # v_ =  self.critic(states)
     96 # critic_loss = self.critic_loss_fn(v_, target_values)

File c:\Users\...\.venv\Lib\site-packages\torch\_tensor.py:525, in Tensor.backward(self, gradient, retain_graph, create_graph, inputs)
    515 if has_torch_function_unary(self):
    516     return handle_torch_function(
    517         Tensor.backward,
    518         (self,),
   (...)
    523         inputs=inputs,
    524     )
--&gt; 525 torch.autograd.backward(
    526     self, gradient, retain_graph, create_graph, inputs=inputs
    527 )

File c:\Users\...\.venv\Lib\site-packages\torch\autograd\__init__.py:267, in backward(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)
    262     retain_graph = create_graph
    264 # The reason we repeat the same comment below is that
    265 # some Python versions print out the first line of a multi-line function
    266 # calls in the traceback and some print out the last line
--&gt; 267 _engine_run_backward(
    268     tensors,
    269     grad_tensors_,
    270     retain_graph,
    271     create_graph,
    272     inputs,
    273     allow_unreachable=True,
    274     accumulate_grad=True,
    275 )

File c:\Users\...\.venv\Lib\site-packages\torch\autograd\graph.py:744, in _engine_run_backward(t_outputs, *args, **kwargs)
    742     unregister_hooks = _register_logging_hooks_on_whole_graph(t_outputs)
    743 try:
--&gt; 744     return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
    745         t_outputs, *args, **kwargs
    746     )  # Calls into the C++ engine to run the backward pass
    747 finally:
    748     if attach_logging_hooks:


RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.cuda.FloatTensor [128, 128]], which is output 0 of AsStridedBackward0, is at version 2; expected version 1 instead. Hint: enable anomaly detection to find the operation that failed to compute its gradient, with torch.autograd.set_detect_anomaly(True).
</code></pre>

KFrank: <p>Hi Ebrahim!</p>
<aside class="quote no-group" data-username="Ebrahim.Pichka" data-post="1" data-topic="206633" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ebrahim.pichka/48/71568_2.png" class="avatar"> Ebrahim.Pichka:</div>
<blockquote>
<p>The Actor and Critic network share an Encoder network module at the beginning,</p>
</blockquote>
</aside>
<p>I believe that the <code>shared_encoder</code> is the source of your problem.</p>
<aside class="quote no-group">
<blockquote>
<pre><code class="lang-auto">RuntimeError: one of the variables needed for gradient computation
has been modified by an inplace operation:
[torch.cuda.FloatTensor [128, 128]], which is output 0 of AsStridedBackward0,
is at version 2; expected version 1 instead.
</code></pre>
<p>…</p>
<pre data-code-wrap="py"><code class="lang-py">class Encoder(nn.Module):
    def __init__(self, num_instances, embed_dim, device):
        ...
        self.lin1 = nn.Linear(num_instances, embed_dim)
        self.lin2 = nn.Linear(embed_dim, embed_dim)
</code></pre>
</blockquote>
</aside>
<p>Does <code>shared_encoder.lin2</code> have shape <code>[128, 128]</code>?</p>
<aside class="quote no-group">
<blockquote>
<pre data-code-wrap="py"><code class="lang-py">class PPO:
    ...
    def update(self):
            ...
            self.actor_optimizer.zero_grad()
            actor_loss.backward()
            self.actor_optimizer.step()

            self.critic_optimizer.zero_grad()
            critic_loss.backward()
            self.critic_optimizer.step()
</code></pre>
</blockquote>
</aside>
<p>Does <code>shared_encoder._version</code> change from <code>1</code> to <code>2</code> when you call<br>
<code>self.actor_optimizer.step()</code>?</p>
<pre data-code-wrap="py"><code class="lang-py">            self.actor_optimizer.zero_grad()
</code></pre>
<p>This zeros <code>PPO.actor</code>’s grads, including <code>shared_encoder</code>’s grads.</p>
<pre data-code-wrap="py"><code class="lang-py">            actor_loss.backward()
</code></pre>
<p>This backpropagates through <code>PPO.actor</code>, including through<br>
<code>shared_encoder</code>.</p>
<pre data-code-wrap="py"><code class="lang-py">            self.actor_optimizer.step()
</code></pre>
<p>This performs an optimization step on <code>PPO.actor</code>, including on<br>
<code>shared_encoder</code>.  An optimization step is a kind of inplace modification,<br>
so <code>shared_encoder.lin2</code> (and <code>lin1</code>) have been modified inplace.<br>
You can confirm this inplace modification by printing out<br>
<code>shared_encoder.lin2._version</code> before and after the optimizer step.</p>
<pre data-code-wrap="py"><code class="lang-py">            self.critic_optimizer.zero_grad()
</code></pre>
<p>This again zeros out <code>shared_encoder</code>’s grads.  This is logically sensible.</p>
<pre data-code-wrap="py"><code class="lang-py">            critic_loss.backward()
</code></pre>
<p>You now backpropagate through <code>PPO.critic</code>, <em>including through<br>
<code>shared_encoder</code> again.</em>  But <code>shared_encoder.lin2</code> has been modified<br>
inplace, hence the error.</p>
<pre data-code-wrap="py"><code class="lang-py">            self.critic_optimizer.step()
</code></pre>
<p>If you had reached this line without error, you would run a second<br>
optimization step on <code>shared_encoder</code>.  You do want to do something<br>
like this because (if I understand the overall logic of your training)<br>
you do want to optimize <code>shared_encoder</code> with respect to the gradients<br>
it got from <code>PPO.critic</code>.  This specific approach is not necessarily wrong,<br>
but I would suggest having a separate optimizer for <code>shared_encoder</code><br>
and calling it once when <code>shared_encoder</code> has its gradients from both<br>
<code>PPO.actor</code> and <code>PPO.critic</code> accumulated together.</p>
<aside class="quote no-group">
<blockquote>
<p>the error happened while computing the backward through the <code>self.critic</code> network for updating the critic optimizer. meaning at the <code>critic_loss.backward()</code> and it happened at the last linear layer of the <code>Encoder</code> meaning <code>state_emb = F.relu(self.lin2(state_emb))</code>.</p>
</blockquote>
</aside>
<p>Yes, I agree with this.</p>
<p>I don’t think that this would be the cleanest approach, but you should be<br>
able to fix your error with:</p>
<pre><code class="lang-auto">            self.actor_optimizer.zero_grad()
            self.critic_optimizer.zero_grad()

            actor_loss.backward()
            critic_loss.backward()

            self.actor_optimizer.step()

            self.actor_optimizer.zero_grad()
            self.critic_optimizer.step()
</code></pre>
<p>First you zero out all of the grads (redundantly zeroing out<br>
<code>shared_encoder</code>’s grads twice).  Then you backpropagate both<br>
losses. <code>shared_encoder</code> now has its gradients from both <code>PPO.actor</code><br>
and <code>PPO.critic</code> accumulated together.</p>
<p>Next you optimize <code>PPO.actor</code> <em>and</em> <code>shared_encoder</code>.  <code>shared_encoder</code><br>
is optimized with respect to both its <code>actor</code> and <code>critic</code> gradients.</p>
<p>Then you use <code>actor_optimizer</code> to zero out <code>shared_encoder</code>’s gradients<br>
so it (to first approximation) doesn’t get optimized twice.  Lastly, you<br>
optimize <code>PPO.critic</code>.  This does perform a second optimization step<br>
on <code>shared_encoder</code>, but with zero gradients.</p>
<p>Note that many optimizers, including <code>Adam</code>, do change parameter values<br>
even when those parameters have zero gradients.  This may or may not<br>
be a problem, but it seems a little messy to me and makes your training<br>
harder to reason about.</p>
<p>I would recommend setting up three optimizers:  an <code>actor_optimizer</code><br>
that contains <code>actor</code>’s parameters, <em>not including</em> the <code>shared_encoder</code><br>
parameters, a <code>critic_optimizer</code>, also without <code>shared_encoder</code>’s<br>
parameters, and a separate <code>encoder_optimizer</code> that contains just<br>
<code>shared_encoder</code>’s parameters.</p>
<p>Then:</p>
<pre data-code-wrap="py"><code class="lang-py">            self.actor_optimizer.zero_grad()
            self.critic_optimizer.zero_grad()
            self.encoder_optimizer.zero_grad()

            actor_loss.backward()
            critic_loss.backward()

            self.actor_optimizer.step()
            self.critic_optimizer.step()
            self.encoder_optimizer.step()
</code></pre>
<p>As noted above, when you call both <code>actor_loss.backward()</code> and<br>
<code>critic_loss.backward()</code>, you accumulate into <code>shared_encoder</code><br>
its gradients due to both <code>actor</code> and <code>critic</code>.  So when you call<br>
<code>self.encoder_optimizer.step()</code> (just once), you optimize<br>
<code>shared_encoder</code> with respect to both its <code>actor</code> and <code>critic</code> gradients.</p>
<p>Best.</p>
<p>K. Frank</p>