bug/feature: Fixing mixed precision training

[isamu-isozaki]

This makes it so we can do lower precision training. @piercus might be interested

[limiteinductive]

@deltheil PTAL. @isamu-isozaki could you rebase?

[limiteinductive]

When trying with bfloat16 I've got this error

[limiteinductive]

even with float16 I have an error

[limiteinductive]

See runtime bugs found

[isamu-isozaki]

@limiteinductive Ok, a slight update. For bfloat16 it doesn't seem like grad scaling is necessary because the number of exponent bits is the exact same as float32. Though in fp16 it is necessary. The steps needed to make this work for float16 seem like, traditionally,

1. Keep the model in fp32 precision
2. Have a scaler
3. Autocast to fp16 during training
   For autocasting usually it's done like

with autocast(device_type='cuda', dtype=torch.float16):

but in accelerate they do

if self.mixed_precision == "fp16" and is_torch_version(">=", "1.10"):
    model.forward = torch.cuda.amp.autocast(dtype=torch.float16)(model.forward)

and then do

model.forward = convert_outputs_to_fp32(model.forward)

I think I'll try to find a simpler way since these need the model to be "wrapped".

[isamu-isozaki]

Now that I think about it, even in bfloat16 the gradients should be fp32

[isamu-isozaki]

Ok fixed. For this training, the models getting trained have to be in fp32 while all other models can be self.dtype.
The reason I modified the DTypes in the pipelines is that they automatically convert the training model's fp32 to the self.dtype.
For the training step and evaluation, I think we have to do autocast and my guess is it's the simplest way.
In my dino pr, these changes made the gradients fp32 while training on bfloat16

[limiteinductive]

@deltheil PTAL

[deltheil]

Some comments, PTAL cc @limiteinductive

[deltheil]

I guess this is because with AMP some layers/ops behave better (range) with float32, right? Doesn't this deserve a comment?

[deltheil]

Since AMP is on by default (opt-out), and given that this dtype config is used jointly given the autocast(dtype=self.dtype, enabled=self.config.training.automatic_mixed_precision) call, any reason not to pick a sane default here? Namely: either float16 or bfloat16. Might also deserve a comment.

(accelerate even offers to "Choose from ‘no’,‘fp16’,‘bf16 or ‘fp8’.")

[limiteinductive]

We could set it to False by default, and even though float32 + amp doesn't do much, I don't see this as a big issue; I would rather keep the API straightforward.

[deltheil]

Note: as discussed offline, a follow up PR will be needed to properly unscale the gradients before gradient clipping (which happens in between the backward step and optimizer step, as per https://pytorch.org/docs/stable/notes/amp_examples.html#working-with-unscaled-gradients)

But also, what about gradient accumulation? https://pytorch.org/docs/stable/notes/amp_examples.html#working-with-unscaled-gradients

[isamu-isozaki]

I found that if we do

# For all parameters we train in automatic mixed precision we want them to be in float32.
for learnable_parameter in learnable_parameters:
    learnable_parameter.to(dtype=float32)

while it's more clear it doesn't actually set the dtype of all the parameters to float32. So, the correct way to do it is to set the parameter dtype in the model before register_model(when requires_grad is set), and then if not amp, set grad to self.dtype I think.