Gradscaler flags

[Adel-Moumen]

What does this PR do?

This PR add GradScaler flags to ease the customisation of fp16 flags. Furthermore, it adds back the check_loss_isfinite so that the training can be automatically interrupted when the model is locked in NaN/infs losses.

Before submitting

• Did you read the contributor guideline?
• Did you make sure your PR does only one thing, instead of bundling different changes together?
• Did you make sure to update the documentation with your changes? (if necessary)
• Did you write any new necessary tests? (not for typos and docs)
• Did you verify new and existing tests pass locally with your changes?
• Did you list all the breaking changes introduced by this pull request?
• Does your code adhere to project-specific code style and conventions?

PR review

Reviewer checklist

• Is this pull request ready for review? (if not, please submit in draft mode)
• Check that all items from Before submitting are resolved
• Make sure the title is self-explanatory and the description concisely explains the PR
• Add labels and milestones (and optionally projects) to the PR so it can be classified
• Confirm that the changes adhere to compatibility requirements (e.g., Python version, platform)
• Review the self-review checklist to ensure the code is ready for review

[Adel-Moumen]

Hey @asumagic, could you please have a look at this PR? Do you feel it helps to customise more the fp16 training ?

[TParcollet]

I am not in favor in adding these parameters as basic Brain parameters. They will never be used by 99.99% of our users.

[TParcollet]

Checking for all parameters to be finite is redundant with the GradScaler. It already does this. Also it can be crazy expensive for very large large models. Checking if the loss is not finite makes sense, not the parameters.

[asumagic]

The grad scaler does this but it does not care for how long the parameters have gone non-finite. An idea I suggested yesterday was to occasionally check the gradscaler scale for insane values with a patience mechanism, as I've sometimes seen the scale vanish or explode when issues occurred.
That does induce a CPU-GPU sync though, but I'm not sure how often we have one in the first place.

[Adel-Moumen]

I agree with @TParcollet. I will remove the part where I'm checking the NaNs/inf in the weights and will only cover the loss part. BTW, the function was also intended for other non-GradScaler use cases like fp32.

[mravanelli]

Personally, I'm not against adding this hparams. @Adel-Moumen what is the typical use case for that?

[asumagic]

I agree it might be a bit weird to have those flags exposed directly as a CLI parameter in the core especially since it is not obvious that those relate to the gradient scaler from their name alone.

Maybe it could be specified as an optional brain parameter of some sort? Not sure how feasible that is and how ergonomic it would be to use in the CLI.

[Adel-Moumen]

it is not obvious that those relate to the gradient scaler from their name alone.

They are all starting with gradscaler as the root of their name suggest that they are related to gradscaler :p

TBH, if you think no one will use them, I can remove them. However, I'd like to understand a bit better how to deal with these situations in the future. For instance, if I have some parameters like gradscaler scale_init, how should I provide to the user a way to modify them - without having to modify core.py manually AND without adding new flags in Brain class ?

[asumagic]

They are all starting with gradscaler as the root of their name suggest that they are related to gradscaler :p

Whoops, I didn't double check the code this morning. It wasn't the case when I checked yesterday :P

[mravanelli]

I think @Adel-Moumen is raising a good point here. Even if we think that most of the users won't change these options, what users are supposed to do if they want to change them? Changing core.py directly?

[asumagic]

I think @Adel-Moumen is raising a good point here. Even if we think that most of the users won't change these options, what users are supposed to do if they want to change them? Changing core.py directly?

You can hijack the Brain constructor in your training script to override the grad scaler defined by the core, but it's not very pretty...

[mravanelli]

As discussed privately, we remove all the gradscaler options but we keep the checks on the loss. Thank you @Adel-Moumen