HI, I tried looking in the documentation but I couldn’t find what I’m looking for.
My question is: how are intermediate gradient collected by the backward hooks affected by amp?
I understand that during training we should use GradScaler on the loss as GradScaler.scale(loss).backward(), but what about the gradient collected by a backward hook? Should that also be scaled?
EDIT:
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.
tensor([[-0.6644,  0.6644]], device='cuda:0') # normal
tensor([[-43552.,  43552.]], device='cuda:0', dtype=torch.float16) # amp

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?
Here the main parts of my testing setup
Model:
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

Hook:
def backward_hook(module, grad_input, grad_output):
    print(grad_output[0])

Training:
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()