Many AsStridedBackward0 during the backward

username: pansn
cooked: <p>Our model needs to split the input data. Formatting input data within the forward function of our model causes extremely poor performance when training with torch.compile.</p>
<p>torch: v2.1.0<br>
GPU: V100</p>
<pre><code class="lang-auto">def CustomModel(torch.nn.Module):
    def __init__(self) -&gt; None:
        ...

    def format_data(self, data):
        data = torch.reshape(data, [-1, 5, 121024])
        data_list = torch.permute(data, [1, 0, 2])
        data_split_shape = [1,1,1,1,1,950,14,5,1,1,1,1,1,1,5,5,5,5,1,5,5,1,1,1,1,1,14,25,5,1,1,1,1,1,1,1,1,1,5,25,42,42,42,42,3,61,3,61]
        each_data_length = sum(data_split_shape)
        sequence_data_length = each_data_length * 16
        sequence_data_split_shape = [sequence_data_length, 4096, 256, 4096, 256]
        sequence_data, data1, data2, data3, data4 = torch.split(data_list, sequence_data_split_shape, dim=-1)
        sequence_data = torch.reshape(sequence_data, [5, -1, each_data_length])
        each_data_list = [torch.split(sequence_data[i], data_split_shape, dim=-1) for i in range(5)]
        d1 = data1.reshape([-1, 4096])
        d3 = data3.reshape([-1, 4096])
        d2 = data2.reshape([-1, 256])
        d4 = data4.reshape([-1, 256])
        return each_data_list, (d2, d1), (d4, d3)

    def forward(self, data):
        data_list = self.format_data(data)
        return ...

model = CustomModel().to("cuda")
model.train()

compiled_model = torch.compile(model)

total_step = 10000
for step in range(total_step):
    res = compiled_model(data)
    loss = compiled_model.loss_def(res)
    loss.backward()
</code></pre>
<p>The latency comparison of each step during training is as follows</p>
<div class="md-table">
<table>
<thead>
<tr>
<th>batch_size</th>
<th>format_in_forward</th>
<th>not_in_forward</th>
</tr>
</thead>
<tbody>
<tr>
<td>16</td>
<td>601.5 ms</td>
<td>781.4ms</td>
</tr>
</tbody>
</table>
</div><p>Through torch.profile, we found a large number of AsStridedBackward0 during backward.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/7/6/761cf7dc4bdf356c7159e926f7accf90c72db75b.png" data-download-href="https://discuss.pytorch.org/uploads/default/761cf7dc4bdf356c7159e926f7accf90c72db75b" title="16970141232922"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/7/6/761cf7dc4bdf356c7159e926f7accf90c72db75b_2_690x90.png" alt="16970141232922" data-base62-sha1="gQSp99mwPMplc9uXj63mnCsfrHZ" width="690" height="90" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/7/6/761cf7dc4bdf356c7159e926f7accf90c72db75b_2_690x90.png, https://discuss.pytorch.org/uploads/default/optimized/3X/7/6/761cf7dc4bdf356c7159e926f7accf90c72db75b_2_1035x135.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/7/6/761cf7dc4bdf356c7159e926f7accf90c72db75b_2_1380x180.png 2x" data-dominant-color="ECDBDC"><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">16970141232922</span><span class="informations">2389×312 52.8 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>

username: pansn
cooked: <aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/7bcf7da3a268b435777fe87c7794c382f444e86d/torch/_functorch/aot_autograd.py#L632">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/7bcf7da3a268b435777fe87c7794c382f444e86d/torch/_functorch/aot_autograd.py#L632" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/7bcf7da3a268b435777fe87c7794c382f444e86d/torch/_functorch/aot_autograd.py#L632" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/7bcf7da3a268b435777fe87c7794c382f444e86d/torch/_functorch/aot_autograd.py#L632</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="622" style="counter-reset: li-counter 621 ;">
          <li>    )</li>
          <li></li>
          <li></li>
          <li>def gen_alias_from_base(aliased_base_tensor, target_meta_tensor, target_requires_grad):</li>
          <li>    # Try to do view-replay if possible.</li>
          <li>    # fall back to .as_strided() if we can't.</li>
          <li>    if target_meta_tensor._base is not None:</li>
          <li>        # The base that we want to replay our view off of might have a different shape than the view's original base.</li>
          <li>        b = target_meta_tensor._base</li>
          <li>        abt = aliased_base_tensor</li>
          <li class="selected">        # Don't unnecessarily call as_strided if nothing changed; as_strided's</li>
          <li>        # backward is poorly implemented and slow</li>
          <li>        if abt is not b and (</li>
          <li>            abt.size() != b.size() or</li>
          <li>            abt.stride() != b.stride() or</li>
          <li>            abt.storage_offset() != b.storage_offset()</li>
          <li>        ):</li>
          <li>            reshaped_base_tensor = aliased_base_tensor.as_strided(</li>
          <li>                b.size(), b.stride(), b.storage_offset()</li>
          <li>            )</li>
          <li>        else:</li>
      </ol>
    </code></pre>



  </article>

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

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

<p>I found this comment while reading the source code. Is it related to the problem I am facing? How should I solve this problem? Should I only move format_data out of forward?</p>

username: pansn
cooked: <aside class="onebox githubissue" data-onebox-src="https://github.com/pytorch/pytorch/issues/111081">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/issues/111081" target="_blank" rel="noopener nofollow ugc">github.com/pytorch/pytorch</a>
  </header>

  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Issue">
	  <svg width="60" height="60" class="github-icon" viewBox="0 0 14 16" aria-hidden="true"><path fill-rule="evenodd" d="M7 2.3c3.14 0 5.7 2.56 5.7 5.7s-2.56 5.7-5.7 5.7A5.71 5.71 0 0 1 1.3 8c0-3.14 2.56-5.7 5.7-5.7zM7 1C3.14 1 0 4.14 0 8s3.14 7 7 7 7-3.14 7-7-3.14-7-7-7zm1 3H6v5h2V4zm0 6H6v2h2v-2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/pytorch/pytorch/issues/111081" target="_blank" rel="noopener nofollow ugc">AOTAutograd perf: avoid as_strided() calls when we have intermediate bases</a>
    </h4>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2023-10-11" data-time="20:23:00" data-timezone="UTC">08:23PM - 11 Oct 23 UTC</span>
      </div>


      <div class="user">
        <a href="https://github.com/bdhirsh" target="_blank" rel="noopener nofollow ugc">
          <img alt="bdhirsh" src="https://avatars.githubusercontent.com/u/16311747?v=4" class="onebox-avatar-inline" width="20" height="20">
          bdhirsh
        </a>
      </div>
    </div>

    <div class="labels">
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          triaged
        </span>
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          oncall: pt2
        </span>
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          module: aotdispatch
        </span>
    </div>
  </div>
</div>

  <div class="github-row">
    <p class="github-body-container">This is a more targeted version of an existing issue around as_strided calls in <span class="show-more-container"><a href="" rel="noopener" class="show-more">…</a></span><span class="excerpt hidden">AOTAutograd, https://github.com/pytorch/pytorch/issues/109237. Came from an internal issue

Simple repro:
```
import torch

@torch.compile
def f(x):
    out = x.mul(2)
    return out.view(out.shape), out.view(out.shape)

inps = (torch.randn(5), torch.tensor(0))
x = torch.randn(4, requires_grad=True)
out1, out2 = f(x)
print(out1.grad_fn)
```

prints:
```
&lt;AsStridedBackward0 object at 0x7f4c7d4d2260&gt;
```

We end up calling as_strided in the compiled forward, so an `AsStridedBackward` node shows up in the backward, which in general is not implemented to be particularly fast.

Why does this happen?

(1) AOTAutograd has logic for "intermediate bases". If we have two outputs of our graph that are aliases of each other (and of the same graph intermediate), today, AOTAutograd will just have the shared intermediate be an output to the graph. AOTAutograd will then replay the views off of the intermediate, so that autograd properly realizes that the outputs alias.

(2) AOTAutograd has a function to try to do the view replay, but it hits a slow path in that function that causes it to go to as_strided. We should figure out why and fix this: https://github.com/pytorch/pytorch/blob/main/torch/_functorch/aot_autograd.py#L807

cc @ezyang @msaroufim @wconstab @anijain2305</span></p>
  </div>

  </article>

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

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


username: bdhirsh
cooked: <p>Hey <a class="mention" href="/u/pansn">@pansn</a> do you have a self-contained code repro I can use that replicates the issue? That would be helpful for tracking down the as_strider calls.</p>

username: pansn
cooked: <p><a class="mention" href="/u/bdhirsh">@bdhirsh</a> Sorry, maybe my previous code was too abstract. Here is a simple reproducible sample code. Could you please take a look and help me out?</p>
<pre><code class="lang-auto">import torch
import numpy as np
import time

class CustomModel(torch.nn.Module):
    def format_data(self, data):
        data = torch.reshape(data, [-1, 5, 31024])
        data_list = torch.permute(data, [1, 0, 2])
        data_split_shape = [1,1,1,1,1,950,14,5,1,1,1,1,1,1,5,5,5,5,1,5,5,1,1,1,1,1,14,25,5,1,1,1,1,1,1,1,1,1,5,25,42,42,42,42,3,61,3,61]
        each_data_length = sum(data_split_shape)
        sequence_data_length = each_data_length * 16
        sequence_data_split_shape = [sequence_data_length, 4096, 256, 4096, 256]
        sequence_data, data1, data2, data3, data4 = torch.split(data_list, sequence_data_split_shape, dim=-1)
        sequence_data = torch.reshape(sequence_data, [5, -1, each_data_length])
        each_data_list = [torch.split(sequence_data[i], data_split_shape, dim=-1) for i in range(5)]
        d1 = data1.reshape([-1, 4096])
        d3 = data3.reshape([-1, 4096])
        d2 = data2.reshape([-1, 256])
        d4 = data4.reshape([-1, 256])
        return each_data_list, (d2, d1), (d4, d3)

    def forward(self, data):
        return self.format_data(data)
    
    def loss_def(self, res1, res2, res3):
        loss = 0
        for r in res1:
            for x in r:
                loss = loss + torch.sum(x)
        for r in res2:
            loss = loss + torch.sum(r)
        for r in res3:
            loss = loss + torch.sum(r)
        return loss

device = "cuda"

model = CustomModel().to(device)
model.train()

dtype = torch.float32

compiled_model = torch.compile(model)
warmup_step = 10
total_step = 100

for step in range(total_step):
    if step == warmup_step:
        torch.cuda.synchronize(device=device)
        start_time = time.perf_counter()
    data = torch.tensor(np.random.random([16, 5*31024]) + 0.1, requires_grad=True).type(dtype).to(device)
    res1, res2, res3 = compiled_model(data)
    loss = compiled_model.loss_def(res1, res2, res3 )
    loss.backward()

torch.cuda.synchronize(device=device)
end_time = time.perf_counter()

print(
    "torch.compile avg step time: {} ms".format(
        (end_time - start_time) * 1e3 / (total_step - warmup_step)
    )
)

for step in range(total_step):
    if step == warmup_step:
        torch.cuda.synchronize(device=device)
        start_time = time.perf_counter()
    data = torch.tensor(np.random.random([16, 5*31024]) + 0.1, requires_grad=True).type(dtype).to(device)
    res1, res2, res3 = model(data)
    loss = model.loss_def(res1, res2, res3 )
    loss.backward()

torch.cuda.synchronize(device=device)
end_time = time.perf_counter()

print(
    "torch avg step time: {} ms".format(
        (end_time - start_time) * 1e3 / (total_step - warmup_step)
    )
)
</code></pre>
<p>V100 result</p>
<pre><code class="lang-auto">torch.compile avg step time: 49.37022762993971 ms
torch avg step time: 35.64678467810154 ms
</code></pre>

username: bdhirsh
cooked: <p>Hey <a class="mention" href="/u/pansn">@pansn</a> - I believe the AsStridedBackward slowness should be fixed by <a href="https://github.com/pytorch/pytorch/pull/111411" rel="noopener nofollow ugc">https://github.com/pytorch/pytorch/pull/111411</a>. I tried running your repro on top of that fix on a nightly, and when I print out the <code>grad_fn</code> of every output of the forward, I no longer see any AsStridedBackward nodes.</p>