Half Precision based training adaptations

nath-11743: <p>Hello,<br>
I Wanted to try a few training recipes in fixed precision</p>
<ol>
<li>looking for documentation on how different modules in a model should be converted correctly<br>
my current setup just uses <code>model.half()</code> or <code>model.to(torch.float16)</code> it hasnot thrown any blocking errors, but i would like to setup my code correctly for future use</li>
<li>loss scaling
<ol>
<li>current LRs for different highlevel modules are set to a known and stabilised config that works in single precision and mixed precision with <code>float16</code>/<code>bfloat16</code>
<ol>
<li>[edit] ofc i dont expect to keep these as final lrs in half precision, only using them as a starting point</li>
</ol>
</li>
<li>when running half precision step 2 features outputs jump to <code>nan</code> from early layers</li>
<li>looking to use a scaler for the lossfunction at this stage but the <code>amp.GradScaler</code> returns</li>
<li>
<pre data-code-wrap="bash"><code class="lang-bash">  File ".../trainers.py", line 92, in run_step
    self.scaler.unscale_(self.optimizer)
  File ".../python3.12/site-packages/torch/amp/grad_scaler.py", line 342, in unscale_
    optimizer_state["found_inf_per_device"] = self._unscale_grads_(
                                              ^^^^^^^^^^^^^^^^^^^^^
  File ".../python3.12/site-packages/torch/amp/grad_scaler.py", line 264, in _unscale_grads_
    raise ValueError("Attempting to unscale FP16 gradients.")
ValueError: Attempting to unscale FP16 gradients.

</code></pre>
</li>
<li>
<pre data-code-wrap="python"><code class="lang-python">#scaling code
            self.scaler.scale(loss).backward()
            self.scaler.unscale_(self.optimizer)
            if self.cfg.clip_grad is not None:
                torch.nn.utils.clip_grad_norm_(
                    self.model.parameters(), self.cfg.clip_grad
                )
            self.scaler.step(self.optimizer)

            # When enable amp, optimizer.step call are skipped if the loss scaling factor is too large.
            # Fix torch warning scheduler step before optimizer step.
            scaler = self.scaler.get_scale()
            self.scaler.update()
            if scaler &lt;= self.scaler.get_scale():
                self.scheduler.step()
</code></pre>
</li>
<li>What would be the correct way to implement loss scaling in such a configuration with any pytorch based modules?</li>
</ol>
</li>
</ol>

ptrblck: <p>Did you check the <a href="https://docs.pytorch.org/docs/stable/notes/amp_examples.html">mixed precision training examples</a> which demonstrate how to use <code>autocast</code> instead of explicitly transforming the parameters to a lower <code>dtype</code> which can cause numerical stability issues?</p>

nath-11743: <p>hi <a class="mention" href="/u/ptrblck">@ptrblck</a></p>
<p>thanks for reverting.</p>
<p>ofcourse ive already checked amp based training and as mentioned in 2.1, ive already got the training code working with autocast with desired model outcomes.</p>
<p>I intend now to train a model in half precision only.</p>
<p>Im getting an error using the gradscaler and likely not doing it correctly.</p>
<p>for a minimal exampe im adapting the example code from <a href="https://docs.pytorch.org/docs/stable/notes/amp_examples.html" class="inline-onebox" rel="noopener nofollow ugc">Automatic Mixed Precision examples — PyTorch 2.7 documentation</a></p>
<pre data-code-wrap="python"><code class="lang-python">scaler = GradScaler()
model.half() # model.to(torch.float16)
for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()
        output_logits = model(input)
        #output_logits.dtype = torch.float16
        #target is torch.int64
        loss = loss_fn(output_logits.float(), target)
        scaler.scale(loss).backward()

        # Unscales the gradients of optimizer's assigned params in-place

        ##error!! ValueError: Attempting to unscale FP16 gradients.
        scaler.unscale_(optimizer)

        # Since the gradients of optimizer's assigned params are unscaled, clips as usual:
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)

</code></pre>
<p>scaler.unscale_(optimizer) gives the <code>ValueError: Attempting to unscale FP16 gradients.</code></p>
<p>In hindsight i couldve phrased my initial question as just point 2.3 rather than the full set. apologies for any confusion</p>

Dhia-naouali: <p>Hii</p>
<p><code>GradScaler</code> is designed to work in mixed precision training,</p>
<p><code>scaler.unscale_(optimizer)</code> expects model params and grads to be in <code>float32</code> to unscale the grads after back propagating a scaled loss</p>
<p>you’re facing the error cuz your model weights are in <code>float16</code> which isn’t supported by <code>GradScaler</code></p>
<p>when working fully in half precision you’ll have to handle loss scale / unscale manually (and carefully) to avoid over / under flows,</p>
<p>(no need to unscale grads if your model is fully in half precision tho since everything is already in <code>float16</code>)</p>