Thank you for all this, it really helps alot!



 mcarilli:

manually unscaling the out-of-place grads before using them to set up the double-backward is expected


Good to know, thanks



 mcarilli:

why not also use an out-of-place torch.autograd.grad to create gradsH


I tested this today and couldn’t get it to work for some reason, collecting gradsH and params from the optimizer seemed to be the only thing that worked (i.e. it trained with no errors but results were poorer).



 mcarilli:

FYI GradScaler also has an enabled argument


Missed that, thanks



 mcarilli:

Your optimizer is vanilla SGD and you don’t appear to use hvs , so it’s hard for me to recommend a way to create it.


Oh yep, that was just some bare bones code to be able to run a single step. hvs is used to calculate the Hutchinson trace, which is subsequently used in a modification of AdamW
v = [torch.randint_like(p, high=2, device=device) for p in params]

for v_i in v: v_i[v_i == 0] = -1
        
hvs = torch.autograd.grad(gradsH, params, grad_outputs=v, only_inputs=True, create_graph=False, retain_graph=False)

hutchinson_trace = []
for hv, vi in zip(hvs, v):
    param_size = hv.size()
    if len(param_size) <= 1:  
        tmp_output = torch.abs(hv * vi)
        hutchinson_trace.append(tmp_output)

hutchinson_trace subsequently replaces the gradient in the denominator for the calculation of the AdamW step
BatchNorm had to be set to FP32
After implementing manual scaling, I got native AMP to work with my ResNet18, however I had to modify the input to my BatchNorm to explicitly be FP32 for it to work with larger input sizes (it worked for toy-size problems without modifying the input to BatchNorm).
Have you heard of others with this BN issue?
I see in the autocast docs its not explicitly defined whether BN should be FP16 or FP32. Maybe it just wasn’t working because of my double backward case…
Memory usage - autocast debugging
Despite successful training, when checking memory use via ‘nvidia-smi’ or ‘nvtop’ it doesn’t seem like the GPU memory usage has decreased. Its about the same as without native AMP and in fact after about 30 steps GPU memory usage increases by approx 800mb. I have identified that the increase happens after the second backward below is called for the 30th time.
hvs = torch.autograd.grad(gradsH, params, grad_outputs=v, only_inputs=True, retain_graph=False)

Note that 800mb also appears to be the increase in memory when the hvs is calculated in the very first step for the very first time, so maybe a copy is being made?
What might the best way to check what precision different operations are done in under autocast?
Checking dtype of the inputs/outputs of each layer?
Is there anything that you are aware of that would cause torch.autograd.grad memory usage to increase?
I checked the size of the tensors being input as below, but the inputs stay the exact same size from start to finish (as expected). The size of the output also remains the same.
gsize += g.element_size() * g.nelement()
Thanks again for your detailed reply above!