Equivalent of GradScaler in the C++ API

MaxC2: <p>I currently implement a training of a MSS (Music Source Separation) model in C++ and I need the equivalent of GradScaler (from <code>torch.cuda.amp.grad_scaler</code> in python).</p>
<p>As explain <a href="https://discuss.pytorch.org/t/deploy-mixed-precision-model-in-libtorch/89046/5">here</a> two years ago by Michael Carilli:</p>
<aside class="quote no-group quote-modified" data-username="mcarilli" data-post="5" data-topic="89046">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/mcarilli/48/28517_2.png" class="avatar"><a href="https://discuss.pytorch.org/t/deploy-mixed-precision-model-in-libtorch/89046/5">Deploy mixed precision model in libtorch</a></div>
<blockquote>
<p>Imitating <code>torch.cuda.amp.GradScaler</code> on the C++ side is not as easy as imitating <code>torch.cuda.amp.autocast</code>. <code>autocast</code>’s Python source is (mostly) a thin wrapper around the calls above, but <code>GradScaler</code>’s Python source contains a fair amount of important logic. I do plan to create an equivalent C++ <code>GradScaler</code>, but that’s a substantial effort.</p>
</blockquote>
</aside>
<p>I definitively not found anything about this point in the different places of information.</p>
<p>I do use the <code>torch.cuda.amp.autocast</code> imitating trick, with a revelant 6x performance boost  (as expected in the NVIDIA documentation, thank’s to Tensor Cores!), but because of the topology of my model, I really need a gradient scaling of my loss function.</p>
<p>Has anyone ever faced this problem? Is there a shareable solution or should I think about a custom implementation?</p>
<p>Thank you in advance for your help!</p>

MaxC2: <p>Self-Reply</p>
<p>I’ve done the C++ implementation of the <strong>GradScaler</strong> class following the python implementation.</p>
<p>You can find the <code>gradscaler.hpp</code> and is <code>gradscaler_test.hpp</code> files in my github gist:</p>
<aside class="onebox githubgist" data-onebox-src="https://gist.github.com/dorpxam/67ad2bc222b2cf567d4a6fc298375e13">
  <header class="source">

      <a href="https://gist.github.com/dorpxam/67ad2bc222b2cf567d4a6fc298375e13" target="_blank" rel="noopener nofollow ugc">gist.github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://gist.github.com/dorpxam/67ad2bc222b2cf567d4a6fc298375e13" target="_blank" rel="noopener nofollow ugc">https://gist.github.com/dorpxam/67ad2bc222b2cf567d4a6fc298375e13</a></h4>

  <h5>gradscaler.hpp</h5>
  <pre><code class="C++">#pragma once

#include &lt;iterator&gt;
#include &lt;type_traits&gt;

#include &lt;torch/torch.h&gt;
#include &lt;torch/serialize/archive.h&gt;

namespace torch {
namespace amp {</code></pre>
    This file has been truncated. <a href="https://gist.github.com/dorpxam/67ad2bc222b2cf567d4a6fc298375e13" target="_blank" rel="noopener nofollow ugc">show original</a>
  <h5>gradscaler_test.hpp</h5>
  <pre><code class="C++">#pragma once

/*
	Testing 'GradScaler' using adapted tests provided in the python api

		&gt; pytorch/test/test_cuda.py

	The dependencies for running the tests are:

		&gt; https://github.com/catchorg/Catch2</code></pre>
    This file has been truncated. <a href="https://gist.github.com/dorpxam/67ad2bc222b2cf567d4a6fc298375e13" target="_blank" rel="noopener nofollow ugc">show original</a>

<p>
</p>

  </article>

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

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

<p>Unfortunatly, this is not tested on multi-GPU systems, but pass the test on single GPU without problem.</p>
<p>Just be carefull with the <code>scale()</code> method that support both Tensor or a iterable of Tensors in a generic way. The method support all the <code>std-like</code> containers that support the back insertion (<code>std::back_inserter</code>) for build N-dimensional output recursively.</p>
<p>For a better support, I use all the <code>c10 namespace</code> tools when possible (optional, variant, …).</p>
<p>Hope that help the Libtorch C++ community.</p>