Backward Hooks with AMP

andrea_bragagnolo: <p>HI, I tried looking in the documentation but I couldn’t find what I’m looking for.<br>
My question is: how are intermediate gradient collected by the backward hooks affected by amp?</p>
<p>I understand that during training we should use <code>GradScaler</code> on the loss as <code>GradScaler.scale(loss).backward()</code>, but what about the gradient collected by a backward hook? Should that also be scaled?</p>
<p>EDIT:<br>
I tried doing some testing with a simple setup and found out that if amp is NOT used I obtain different gradients from when amp IS used e.g.</p>
<pre><code class="lang-auto">tensor([[-0.6644,  0.6644]], device='cuda:0') # normal
tensor([[-43552.,  43552.]], device='cuda:0', dtype=torch.float16) # amp
</code></pre>
<p>but using GradScaler.scale on the gradient did not return the non-amp one. Is there a way to use amp and hooks to gather intermediate gradients?</p>
<p>Here the main parts of my testing setup<br>
Model:</p>
<pre><code class="lang-auto">class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(10, 5, bias=True)
        self.r1 = nn.ReLU()
        self.fc2 = nn.Linear(5, 2, bias=True)
    
    def forward(self, x):
        x = self.fc1(x)
        x = self.r1(x)
        x = self.fc2(x)
        return x
</code></pre>
<p>Hook:</p>
<pre><code class="lang-auto">def backward_hook(module, grad_input, grad_output):
    print(grad_output[0])
</code></pre>
<p>Training:</p>
<pre><code class="lang-auto">dataset = TensorDataset(x, y)
dataloader = DataLoader(dataset, shuffle=False, batch_size=1, num_workers=0)

for data, target in dataloader:
    data, target = data.to(device), target.to(device)
    optimizer.zero_grad()
    if amp:
        with autocast():
            out = model(data)
            loss = loss_fn(out, target)
        scaled_loss = scaler.scale(loss)
        scaled_loss.backward()
            
        scaler.step(optimizer)
        scaler.update()
    else:
        out = model(data)
        loss = loss_fn(out, target)
        loss.backward()
        optimizer.step()
</code></pre>

ptrblck: <p>If you want to work with unscaled gradients, you could follow the <a href="https://pytorch.org/docs/stable/notes/amp_examples.html#working-with-unscaled-gradients">AMP - Working with unscaled gradients tutorial</a>, which explains how to unscale the gradients before using them in e.g. gradient clipping.</p>

andrea_bragagnolo: <p>Hi, thanks for the reply, but unless I misunderstood something this isn’t what I’m looking for: in the link you provided it is suggested to unscale the optimizer so that we can unscales gradients held by optimizer’s assigned parameters, but those are not the gradient I’m interested in. I’ll try to explaing it better.</p>
<p>In my project I need to use the intermediate gradients of the neurons (NOT the gradient of the loss w.r.t. the weights) and to fetch those I employ some backward hooks.<br>
When I run the <code>backward()</code> pass WITHOUT amp I obtain the following grad tensor <code>tensor([[-0.6644,  0.6644]], device='cuda:0')</code>.<br>
When I run the <code>backward()</code> pass WITH amp, instead, I get <code>tensor([[-43552.,  43552.]], device='cuda:0', dtype=torch.float16)</code>. Now I’d like to unscale this float16 gradients to the values I get without amp, so that I can still use them in the application. Is this possible?</p>

ptrblck: <p>Yes, the next section of the tutorials shows how to <a href="https://pytorch.org/docs/stable/notes/amp_examples.html#gradient-penalty">unscale gradients manually</a> using <code>inv_scale = 1./scaler.get_scale()</code>, which you could also use in your hooks.</p>

andrea_bragagnolo: <p>Oh wow, I really missed that. Thank you very much for helping me.</p>

mcarilli: <p>If performance is a concern, prefer <code>1./scaler._get_scale_async()</code>, which will be faster because it retrieves the scale factor tensor directly. <code>get_scale()</code> calls .item() on the scale factor tensor which incurs a cpu-&gt;gpu memcopy and sync. I should probably give <code>_get_scale_async()</code> a documented, public exposure.</p>