F.conv2d runs x15 slower in mixed-precision/half-precision mode

helson73: <p>SYSTEM info:</p>
<ul>
<li>OS: Ubuntu 18.04</li>
<li>CUDA: 10.1</li>
<li>pytorch: 1.4.0 installed from conda</li>
<li>NVIDIA apex: latest installed from pip</li>
<li>GPU: Titan RTX 24GB</li>
<li>Driver: 430.64</li>
<li>cudnn version: 7603</li>
</ul>
<p>I tested my system with following code:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
from torch import optim
from apex import amp
from torch.nn import functional as F


def main(opt_level):

    assert torch.backends.cudnn.enabled, 'Cudnn is not enabled!'
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)

    padding = 3 // 2
    # input - torch.Size([1, 512, 256, 256])
    # weight - torch.Size([512, 128, 3, 3])
    class net(nn.Module):
        def __init__(self):
            super().__init__()
            self.weight = nn.Parameter(torch.randn(512, 128, 3, 3))
            self.padding = 3 // 2

        def forward(self, input):
            out = F.conv2d(input, self.weight, padding=self.padding, groups=4)
            return out


    input = torch.randn(1, 512, 256, 256).float().cuda()
    model = net().cuda()
    optimizer = optim.Adam(model.parameters(), lr=0.01, betas=(0, 0.999))
    model, _ = amp.initialize(model, optimizer, keep_batchnorm_fp32=None if opt_level == 'O1' else True,
                              loss_scale='dynamic', opt_level=opt_level, num_losses=1)

    for i in range(10):
        start.record()
        out = model(input)
        end.record()
        torch.cuda.synchronize()
        print('--- %.4f' % start.elapsed_time(end))

if __name__ == '__main__':
    import argparse
    p = argparse.ArgumentParser()
    p.add_argument('-opt_level', type=str, required=True,
                    choices=['O0', 'O1', 'O2', 'O3'])
    opt = vars(p.parse_args())
    main(opt['opt_level'])
</code></pre>
<p>and results are:</p>
<p>$ python test_conv2d.py -opt_level O0</p>
<pre><code class="lang-auto">Selected optimization level O0:  Pure FP32 training.

Defaults for this optimization level are:
enabled                : True
opt_level              : O0
cast_model_type        : torch.float32
patch_torch_functions  : False
keep_batchnorm_fp32    : None
master_weights         : False
loss_scale             : 1.0
Processing user overrides (additional kwargs that are not None)...
After processing overrides, optimization options are:
enabled                : True
opt_level              : O0
cast_model_type        : torch.float32
patch_torch_functions  : False
keep_batchnorm_fp32    : True
master_weights         : False
loss_scale             : dynamic
--- 378.2777
--- 5.2879
--- 4.7411
--- 4.6719
--- 4.6950
--- 4.6825
--- 4.7071
--- 4.6776
--- 4.6902
--- 4.6961
</code></pre>
<p>$ python test_conv2d.py -opt_level O1<br>
$ python test_conv2d.py -opt_level O2<br>
$ python test_conv2d.py -opt_level O3</p>
<pre><code class="lang-auto">Selected optimization level O2:  FP16 training with FP32 batchnorm and FP32 master weights.

Defaults for this optimization level are:
enabled                : True
opt_level              : O2
cast_model_type        : torch.float16
patch_torch_functions  : False
keep_batchnorm_fp32    : True
master_weights         : True
loss_scale             : dynamic
Processing user overrides (additional kwargs that are not None)...
After processing overrides, optimization options are:
enabled                : True
opt_level              : O2
cast_model_type        : torch.float16
patch_torch_functions  : False
keep_batchnorm_fp32    : True
master_weights         : True
loss_scale             : dynamic
--- 446.0503
--- 63.7256
--- 62.7091
--- 62.7143
--- 61.9786
--- 60.7601
--- 60.6553
--- 60.5181
--- 60.6509
--- 60.7128
</code></pre>
<p>Here is documentation of opt_levels in APEX.<br>
<a href="https://nvidia.github.io/apex/amp.html" rel="nofollow noopener">Opt_levels</a><br>
O0: float32<br>
O1/O2: mixed-precision<br>
O3: half-precision</p>
<p>Does anyone encounter similar issue?</p>

ptrblck: <p>Grouped convolutions might not trigger the FP16 path and thus might not use TensorCores.</p>

helson73: <p>Thanks, problem solved.</p>

helson73: <p>Hi, what should I do for triggering the fp16 with multiple groups?<br>
Should I modify pytorch source code?<br>
I found alternative implementation (apply each row of weight separately to each batch) is much slower than conv2d with groups (5.3 ms vs 4.6 ms). Thus speed boost with fp16 (3.2 ms) is not that significant (usually should be multiplied by 2~3, in this case less than 2.3 ms).<br>
I guess I’d better to modify the source code in pytorch to make it fully support fp16 in conv2d (with multiple groups).</p>

ptrblck: <p>Unfortunately, there is not much you can do at the moment besides using <code>torch.backends.cudnn.benchmark = True</code>.<br>
I profiled the calls and all options call into the cudnn path <a href="https://github.com/pytorch/pytorch/blob/c57f8984e6edc8410a324755aa0da193d99998df/aten/src/ATen/native/Convolution.cpp#L667">here</a>.<br>
Once cudnn provides faster kernels, you will see the potential speedup.</p>

helson73: <p>So it depends on cudnn’s update …<br>
Because my current model (pytorch version of stylegan2 [Karras]) uses a lot of custom conv layers(must process convolutions by groups), O1/O2 setting makes no big difference with O0 mode. Hope Nvidia updates faster kernels soon.</p>
<p>I also found some pytorch functions (e.g. torch.repeat, torch.cat) is slow in O1 (functions are patched as fp16) mode. In O2, these functions are faster as O0(fp32) but overall backpropagation is slower than O1.</p>

helson73: <p>I guess NHWC may matter …<br>
<a href="https://devtalk.nvidia.com/default/topic/1071156/jetson-nano/poor-group-convolution-performance-in-fp16/" class="onebox" target="_blank" rel="nofollow noopener">https://devtalk.nvidia.com/default/topic/1071156/jetson-nano/poor-group-convolution-performance-in-fp16/</a><br>
<aside class="onebox githubissue">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/issues/28619" target="_blank" rel="nofollow noopener">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 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/28619" target="_blank" rel="nofollow noopener">Channels Last (NHWC) support plan.</a>
    </h4>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2019-10-24" data-time="21:21:22" data-timezone="UTC">09:21PM - 24 Oct 19 UTC</span>
      </div>


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

<div class="github-row">
  <p class="github-content">Release goals:
Support ResNet50, AlexNet, ResNeXt and proof performance gain on mixed precision models + Volta GPUs.
Unlock developers to extend channels last...</p>
</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;">topic: memory format</span>
    <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">triaged</span>
</div>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>
</p>
<p>since it’s still on progress, I guess I could try it later when pytorch is update to date.<br>
With these updates for pytorch I can apply my models to NHWC kernels without aplex’s updates, right?</p>

ptrblck: <p>Some PRs landed already and you could try it out on some vision models. <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=9" title=":wink:" class="emoji" alt=":wink:"></p>

helson73: <p>Yeah, found an updated example in apex <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"><br>
Seems like latest apex is for pytorch-nightly and it requires a little more additional arguments related to memory format.<br>
For now this looks like the best option but still hope memory format bindings could be handled only once at env setting.</p>

helson73: <p>I’ve test it with pytorch-nightly (1.5.0) and latest apex but this time got x1000 slower in all modes (O0, O1, O2, O3) … (with channel last)<br>
They said it’s for volta and I am wordering whether it also works in turing …</p>