Slow aten::fill_ and aten::add_

username: klye
cooked: <p>Given some m-by-m covariance matrix C with missing entries, I would like to find a low rank m-by-k matrix L where k &lt;&lt; m so that LL^T is a good approximation to C. To do so, I’ve coded a distributed version of stochastic gradient descent using Pytorch’s distributed framework.</p>
<p>I noticed that for a fixed number of observed points in C there is a dramatic increase in the time reqired complete a backward pass after m exceeds some threshold. I profiled my code and realized that after m reaches this threshold, the <code>CPU time avg</code> for calls to <code>aten::fill_</code> and <code>aten::add_</code> during the backward pass spike then plateau. Below is a minimum working example demonstrating this behavior.</p>
<pre data-code-wrap="python"><code class="lang-python"># mwe.py
import argparse
import os
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn as nn
from torch.nn.functional import mse_loss
from typing import Callable


class LowRankCovariance(nn.Module):

    def __init__(
            self, 
            num_vars: int, 
            num_facs: int
        ):
        super().__init__()
        self.num_facs = num_facs
        self.loads = torch.randn(num_vars, num_facs, dtype=torch.float64)
        self.loads.requires_grad_()
        self.loads = nn.Parameter(self.loads)

    def forward(self, idx0, idx1):
        return (self.loads[idx0] * self.loads[idx1]).sum(dim=1)

    
def init_process(
        rank: int, 
        world_size: int, 
        fcn: Callable, 
        path_shared: str,
        num_vars: int, 
        num_facs: int, 
        batch_size: int,
        lr: float, 
        seed: int
    ):
    dist.init_process_group(
        'gloo', init_method=f'file://{path_shared}',
        rank=rank, world_size=world_size
    )
    fcn(rank, num_vars, num_facs, batch_size, lr, seed)


def train(
        rank: int, 
        num_vars: int, 
        num_facs: int, 
        batch_size: int,
        lr: float, 
        seed: int
    ) -&gt; None:

    gen = torch.Generator().manual_seed(seed)

    # Generate fake data
    sz = 10000
    points1 = torch.randint(0, num_vars, (sz,), dtype=torch.int32, generator=gen)
    points2 = torch.randint(0, num_vars, (sz,), dtype=torch.int32, generator=gen)
    points = torch.column_stack([points1, points2])
    cov = torch.randn(sz, dtype=torch.float64, generator=gen)

    model = LowRankCovariance(num_vars, num_facs)
    optimizer = torch.optim.SGD(model.parameters(), lr=lr)

    with torch.autograd.profiler.profile(enabled=True) as prof:

        cnt = 0
        while cnt &lt; sz: 

            # Get batch of data
            last_idx = min(cnt, cnt + batch_size)
            points_batch = points[cnt:last_idx]
            cov_batch = cov[cnt:last_idx]
            cnt += batch_size

            # Forward pass
            preds = model(points_batch[:,0], points_batch[:,1])
            loss = mse_loss(preds, cov_batch)

            # Backward pass
            optimizer.zero_grad()
            with torch.autograd.profiler.record_function('backward_pass'):
                loss.backward()
            optimizer.step()

    if rank == 0:
        script_dir = os.path.dirname(os.path.realpath(__file__))
        prof.export_chrome_trace(f'{script_dir}/profile_{num_vars}.json')
        print(prof.key_averages().table(sort_by="self_cpu_time_total"))

    dist.destroy_process_group()



if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument('--num_vars', type=int)
    parser.add_argument('--world_size', type=int)
    parser.add_argument('--num_facs', type=int)
    parser.add_argument('--batch_size', type=int)
    parser.add_argument('--lr', type=float)
    parser.add_argument('--seed', type=int, default=12345)
    args = parser.parse_args()


    # ---------- DISTRIBUTED TRAINING ---------- #

    # Refresh file for shared file-system initialization
    script_dir = os.path.dirname(os.path.realpath(__file__))
    path_shared = os.path.join(script_dir, 'shared')
    if os.path.exists(path_shared):
        os.remove(path_shared)

    mp.set_start_method('spawn')
    processes = []
    for rank in range(args.world_size):
        p = mp.Process(
            target=init_process, 
            args=(
                rank, args.world_size, train, path_shared,
                args.num_vars, args.num_facs, args.batch_size,
                args.lr, args.seed + rank
            )
        )
        p.start()
        processes.append(p)

    for p in processes:
        p.join()

</code></pre>
<p>If you run the script five times (see commands below) with m = 8190, 8191, 8192, 8193, and 8194, you’ll notice that the average CPU times for <code>aten::fill_</code> and <code>aten:add_</code> start low, increase slightly at m = 8192 then increase dramatically at m = 8193 and remain high.</p>
<pre data-code-wrap="bash"><code class="lang-bash">$ python mwe.py --num_vars 8190 --world_size 5 --num_facs 4 --batch_size 512 --lr 0.1
$ python mwe.py --num_vars 8191 --world_size 5 --num_facs 4 --batch_size 512 --lr 0.1
$ python mwe.py --num_vars 8192 --world_size 5 --num_facs 4 --batch_size 512 --lr 0.1
$ python mwe.py --num_vars 8193 --world_size 5 --num_facs 4 --batch_size 512 --lr 0.1
$ python mwe.py --num_vars 8194 --world_size 5 --num_facs 4 --batch_size 512 --lr 0.1
</code></pre>
<p>Below are two screenshots of profile traces for m = 8191 and m = 8193.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/5/d5b8854135060aa8afa26252e185f39297ba836f.png" data-download-href="https://discuss.pytorch.org/uploads/default/d5b8854135060aa8afa26252e185f39297ba836f" title="Screenshot 2024-02-16 at 6.32.22 PM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/d/5/d5b8854135060aa8afa26252e185f39297ba836f_2_690x105.png" alt="Screenshot 2024-02-16 at 6.32.22 PM" data-base62-sha1="uuF292uEgxPCGP0LEmWOgR2PEkf" width="690" height="105" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/d/5/d5b8854135060aa8afa26252e185f39297ba836f_2_690x105.png, https://discuss.pytorch.org/uploads/default/optimized/3X/d/5/d5b8854135060aa8afa26252e185f39297ba836f_2_1035x157.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/d/5/d5b8854135060aa8afa26252e185f39297ba836f_2_1380x210.png 2x" data-dominant-color="DBD9D9"><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">Screenshot 2024-02-16 at 6.32.22 PM</span><span class="informations">2980×456 95.2 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><br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/3/c/3cb642fcc16366f1f2559fcc3145bb3e8f84fc6d.png" data-download-href="https://discuss.pytorch.org/uploads/default/3cb642fcc16366f1f2559fcc3145bb3e8f84fc6d" title="Screenshot 2024-02-16 at 6.32.34 PM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/3/c/3cb642fcc16366f1f2559fcc3145bb3e8f84fc6d_2_690x105.png" alt="Screenshot 2024-02-16 at 6.32.34 PM" data-base62-sha1="8F58O33kxoUIImL8sbF7tm06ZI1" width="690" height="105" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/3/c/3cb642fcc16366f1f2559fcc3145bb3e8f84fc6d_2_690x105.png, https://discuss.pytorch.org/uploads/default/optimized/3X/3/c/3cb642fcc16366f1f2559fcc3145bb3e8f84fc6d_2_1035x157.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/3/c/3cb642fcc16366f1f2559fcc3145bb3e8f84fc6d_2_1380x210.png 2x" data-dominant-color="D8D1D2"><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">Screenshot 2024-02-16 at 6.32.34 PM</span><span class="informations">2984×456 87.6 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>Notice that in the m = 8193 trace, there are extended calls to <code>aten::fill_</code> and <code>aten::add_</code> that are not present in the m = 8191 trace.</p>
<p>Environment details:</p>
<ul>
<li>OS: Red Hat Enterprise Linux 7</li>
<li>Python version: 3.10.12</li>
<li>Pytorch version: 2.1.2</li>
</ul>
<p>Other comments:</p>
<ul>
<li>I was also able to replicate this slowdown locally (MacOS Ventura 13.2) using <code>mwe.py</code>, although it was not as pronounced.</li>
<li>The slowdown worsens as the world size increases.</li>
</ul>
<p>Does anyone know what causes this slowdown and how I can avoid it?</p>