Errors in backward when using FSDP -- RuntimeError: setStorage: sizes [32000, 4096], strides [4096, 1], storage offset 131076096, and itemsize 2 requiring a storage size of 524296192 are out of bounds for storage of size 0

Nardien: <p>Hi. I am fine-tuning llama-7b model on 4 A100 GPUs, utilizing FSDP and LoRA (peft) based on the <code>llama-recipes</code> library (<a href="https://github.com/facebookresearch/llama-recipes" class="inline-onebox" rel="noopener nofollow ugc">GitHub - facebookresearch/llama-recipes: Examples and recipes for Llama 2 model</a>).</p>
<p>Torch version: 2.2.0+cu118</p>
<p>My goal is to compute the “per-token-gradient” for each parameter for analysis purposes.</p>
<p>Here’s how I’ve modified the training code:</p>
<ol>
<li>
<p>I modified the model code to return non-averaged loss by setting the reduction in the CrossEntropyLoss class as ‘none’.</p>
</li>
<li>
<p>Instead of calling backward on the averaged loss, I now call backward on the loss for “each individual” token.</p>
</li>
</ol>
<p>In below code, I changed this section:</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/facebookresearch/llama-recipes/blob/98b122e57a8df44f5b88fa9fdab8818cc6e4969f/src/llama_recipes/utils/train_utils.py#L125">
  <header class="source">

      <a href="https://github.com/facebookresearch/llama-recipes/blob/98b122e57a8df44f5b88fa9fdab8818cc6e4969f/src/llama_recipes/utils/train_utils.py#L125" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/facebookresearch/llama-recipes/blob/98b122e57a8df44f5b88fa9fdab8818cc6e4969f/src/llama_recipes/utils/train_utils.py#L125" target="_blank" rel="noopener nofollow ugc">facebookresearch/llama-recipes/blob/98b122e57a8df44f5b88fa9fdab8818cc6e4969f/src/llama_recipes/utils/train_utils.py#L125</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="115" style="counter-reset: li-counter 114 ;">
          <li>            if train_config.enable_fsdp:</li>
          <li>                model.clip_grad_norm_(train_config.gradient_clipping_threshold)</li>
          <li>            else:</li>
          <li>                torch.nn.utils.clip_grad_norm_(model.parameters(), train_config.gradient_clipping_threshold)</li>
          <li>        scaler.step(optimizer)</li>
          <li>        scaler.update()</li>
          <li>        optimizer.zero_grad()</li>
          <li>        pbar.update(1)</li>
          <li>else:</li>
          <li>    # regular backpropagation when fp16 is not used</li>
          <li class="selected">    loss.backward()</li>
          <li>    if (step + 1) % gradient_accumulation_steps == 0 or step == len(train_dataloader) - 1:</li>
          <li>        if train_config.gradient_clipping and train_config.gradient_clipping_threshold &gt; 0.0:</li>
          <li>            if train_config.enable_fsdp:</li>
          <li>                model.clip_grad_norm_(train_config.gradient_clipping_threshold)</li>
          <li>            else:</li>
          <li>                torch.nn.utils.clip_grad_norm_(model.parameters(), train_config.gradient_clipping_threshold)</li>
          <li>        optimizer.step()</li>
          <li>        optimizer.zero_grad()</li>
          <li>        pbar.update(1)</li>
          <li></li>
      </ol>
    </code></pre>



  </article>

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

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

<p>to:</p>
<pre><code class="lang-auto">loss[0].backward(retrain_graph=True)
optimizer.zero_grad()
loss[1].backward()
</code></pre>
<p>However, this modification results in an error:</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File ".../lib/python3.8/site-packages/torch/autograd/__init__.py", line 411, in grad
    result = Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
RuntimeError: setStorage: sizes [32000, 4096], strides [4096, 1], storage offset 131076096, and itemsize 2 requiring a storage size of 524296192 are out of bounds for storage of size 0
</code></pre>
<p>I also attempted <code>torch.autograd.grad</code> as follows:</p>
<pre><code class="lang-auto">grad1 = torch.autograd.grad(loss[0], updated_parameters, retain_graph=True, allow_unused=True)
optimizer.zero_grad()
grad2 = torch.autograd.grad(loss[1], updated_parameters, allow_unused=True)
</code></pre>
<p>but it led to the same error.</p>
<p>Considering I’m using FSDP and LoRA, this seems to be a complex issue.<br>
One possible solution might be to perform forward computation for each token individually instead of the entire sequence, but this could be very time-consuming.</p>
<p>Additionally, I’m curious about the cause of this problem.<br>
Does anyone have any insights?</p>

ty_w: <p>Same problem. How did you fix it? Thanks.</p>