Torch.cuda.amp equivalent of apex.amp.initialize?

RylanSchaeffer: <p>I’m working on a project forked from a HuggingFace code base that used NVIDIA’s <code>apex.amp.initialize</code>.</p>
<pre><code class="lang-auto">            try:
                from apex import amp
            except ImportError:
                raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use fp16 training.")
            self.student, self.optimizer = amp.initialize(
                self.student,
                self.optimizer,
                opt_level=self.run_config['devices']['fp16_opt_level']
            )
</code></pre>
<p>What’s the equivalent function call using <code>torch.cuda.amp</code>?</p>

RylanSchaeffer: <p><a class="mention" href="/u/mcarilli">@mcarilli</a>, can I ask you for your insight?</p>

RylanSchaeffer: <p>Also, follow up question: apex had levels e.g. ‘fp16_opt_level’: ‘O1’,  # For fp16: Apex AMP optimization level selected in [‘O0’, ‘O1’, ‘O2’, and ‘O3’]. Does torch.cuda.amp have anything similar I need to keep in mind?</p>

ptrblck: <p>You don’t need to initialize the model and optimizer and should use <code>torch.cuda.amp.autocast</code> as well as <code>torch.cuda.amp.GradScaler</code> as described <a href="https://pytorch.org/docs/stable/notes/amp_examples.html#amp-examples">here</a>.</p>
<aside class="quote no-group" data-username="RylanSchaeffer" data-post="3" data-topic="132598">
<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/rylanschaeffer/48/17130_2.png" class="avatar"> RylanSchaeffer:</div>
<blockquote>
<p>Does torch.cuda.amp have anything similar I need to keep in mind?</p>
</blockquote>
</aside>
<p>No, right now native amp is similar to <code>apex.amp</code> O1. We are experimenting with an O2-style mode, which is still WIP.</p>

mcarilli: <p>The two ingredients of native amp (<code>torch.cuda.amp.autocast</code> and <code>torch.cuda.amp.GradScaler</code>) do not affect the model or optimizer in a stateful way.</p>
<p><code>autocast</code> casts inputs to listed functions on the fly. The casted inputs are temporaries. They may end up stashed for backward, but that’s the longest they last. They never overwrite any model attributes.</p>
<p><code>GradScaler</code> scales the loss and (either in an explicit call to <code>scaler.unscale_(optimizer)</code> or implicitly inside <code>scaler.step(optimizer)</code> unscales .grad attributes of the optimizer’s params in place, but it doesn’t overwrite or replace any optimizer attributes.</p>
<p>This is by design, to avoid the black box monkey-patching spaghetti that can take place in <code>apex.amp.initialize</code>. <code>torch.cuda.amp</code> has no equivalent to <code>initialize</code>, so just delete the line.</p>
<p>To toggle amp on or off based on <code>run_config</code> without needing to write divergent code, use autocast and GradScaler’s <code>enabled=</code> argument as shown <a href="https://pytorch.org/tutorials/recipes/recipes/amp_recipe.html#all-together-automatic-mixed-precision" rel="noopener nofollow ugc">here</a>. There’s no equivalent of opt_levels with native amp, it’s either on or off, so you’ll have to decide how to interpret or change <code>fp16_opt_level</code>.</p>