Large CPU overhead on small model

oliverhamilton1: <p>I am trying to deploy a relatively small (~50M params) open source model, and I notice that this block in the decoder portion of the model has lots of overhead. Specifically, in the profile graph, I see lots of gaps in the GPU utilization and lots of CPU overhead (images attached).</p>
<p>Attempts to <code>torch.compile</code> and optimize the module inflate latency by quite a large amount, so im not sure that is the solution.</p>
<p>Any tips on how to optimize to reduce CPU overhead and latency? And to reduce the gaps in the GPU utilization? This module block is present quite a number of times in the model and it is responsible for able 70% of the model’s latency currently. Thanks!</p>
<p>Here is the code:</p>
<pre><code class="lang-auto">import torch
import torch.nn.functional as F
import torch.nn as nn
from torch.nn import Conv1d, ConvTranspose1d, AvgPool1d, Conv2d
from torch.nn.utils import weight_norm, remove_weight_norm, spectral_norm

import math
import random
import numpy as np 

def get_padding(kernel_size, dilation=1):
    return int((kernel_size*dilation - dilation)/2)

def init_weights(m, mean=0.0, std=0.01):
    classname = m.__class__.__name__
    if classname.find("Conv") != -1:
        m.weight.data.normal_(mean, std)

class AdaIN1d(nn.Module):
    def __init__(self, style_dim, num_features):
        super().__init__()
        self.norm = nn.InstanceNorm1d(num_features, affine=False)
        self.fc = nn.Linear(style_dim, num_features*2)

    def forward(self, x, s):
        h = self.fc(s)
        h = h.view(h.size(0), h.size(1), 1)
        gamma, beta = torch.chunk(h, chunks=2, dim=1)
        return (1 + gamma) * self.norm(x) + beta

class AdaINResBlock1(torch.nn.Module):
    def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5), style_dim=64):
        super(AdaINResBlock1, self).__init__()
        self.convs1 = nn.ModuleList([
            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],
                               padding=get_padding(kernel_size, dilation[0]))),
            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],
                               padding=get_padding(kernel_size, dilation[1]))),
            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[2],
                               padding=get_padding(kernel_size, dilation[2])))
        ])
        self.convs1.apply(init_weights)

        self.convs2 = nn.ModuleList([
            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
                               padding=get_padding(kernel_size, 1))),
            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
                               padding=get_padding(kernel_size, 1))),
            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
                               padding=get_padding(kernel_size, 1)))
        ])
        self.convs2.apply(init_weights)
        
        self.adain1 = nn.ModuleList([
            AdaIN1d(style_dim, channels),
            AdaIN1d(style_dim, channels),
            AdaIN1d(style_dim, channels),
        ])
        
        self.adain2 = nn.ModuleList([
            AdaIN1d(style_dim, channels),
            AdaIN1d(style_dim, channels),
            AdaIN1d(style_dim, channels),
        ])
        
        self.alpha1 = nn.ParameterList([nn.Parameter(torch.ones(1, channels, 1)) for i in range(len(self.convs1))])
        self.alpha2 = nn.ParameterList([nn.Parameter(torch.ones(1, channels, 1)) for i in range(len(self.convs2))])


    def forward(self, x, s):
        for c1, c2, n1, n2, a1, a2 in zip(self.convs1, self.convs2, self.adain1, self.adain2, self.alpha1, self.alpha2):
            xt = n1(x, s)
            xt = xt + (1 / a1) * (torch.sin(a1 * xt) ** 2)  # Snake1D
            xt = c1(xt)
            xt = n2(xt, s)
            xt = xt + (1 / a2) * (torch.sin(a2 * xt) ** 2)  # Snake1D
            xt = c2(xt)
            x = xt + x
        return x

    def remove_weight_norm(self):
        for l in self.convs1:
            remove_weight_norm(l)
        for l in self.convs2:
            remove_weight_norm(l)
</code></pre>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/5/b5ec0c27a0279766c536afdaed8abaeff35af48e.png" data-download-href="https://discuss.pytorch.org/uploads/default/b5ec0c27a0279766c536afdaed8abaeff35af48e" title="Screenshot 2024-10-19 at 9.45.11 AM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/b/5/b5ec0c27a0279766c536afdaed8abaeff35af48e_2_290x500.png" alt="Screenshot 2024-10-19 at 9.45.11 AM" data-base62-sha1="pXm9oxb931OTCGskUrjHhhywVVs" width="290" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/b/5/b5ec0c27a0279766c536afdaed8abaeff35af48e_2_290x500.png, https://discuss.pytorch.org/uploads/default/optimized/3X/b/5/b5ec0c27a0279766c536afdaed8abaeff35af48e_2_435x750.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/b/5/b5ec0c27a0279766c536afdaed8abaeff35af48e_2_580x1000.png 2x" data-dominant-color="E4DFEB"><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-10-19 at 9.45.11 AM</span><span class="informations">1192×2052 153 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>

ptrblck: <p>Try to use <code>CUDAGraphs</code> as described in <a href="https://pytorch.org/blog/accelerating-pytorch-with-cuda-graphs/">this blog post</a> to reduce the CPU overhead or use <code>torch.compile</code> with <code>mode="reduce-overhead"</code> as it will graph capture the workload internally.</p>

oliverhamilton1: <p>For this particular open source audio model, we were seeing <code>mode="reduce-overhead"</code> not provide adequate optimization because of variably sized input tensors. Even if the input tensors to the whole model were constant in shape, the model has some non-determinism such that tensors are not of constant size when they reach the “bottleneck” <code>torch.nn.Module</code> in the model. Any tips on how to optimize a case like this?</p>