Add select float16 or bfloat16

[NeonBohdan]

autocast supports selecting dtype, float16 or bloat16(float16 is default)
And there is no way to switch to bfloat16

Trainer/trainer/trainer.py

Line 975 in 0b3c88a

with torch.cuda.amp.autocast(enabled=config.mixed_precision):

[NeonBohdan]

Tried on local version of Trainer with:
with torch.cuda.amp.autocast(enabled=config.mixed_precision, dtype=torch.bfloat16):

Returned RuntimeError: "reflection_pad1d_out_template" not implemented for 'BFloat16' oon VITS training
So no much reason for now

[erogol]

What is the benefit of bfloat16 training actually?

[NeonBohdan]

Float16 has reduced exponent 8->5 bits compared to float32. So during training gradients can overflow to "NaN"
To omit that used gradient scaling, but it can zeroing gradients, so that network will train slower in a meaning of number of epochs

Bfloat16 solves it. It keeps exponent the same 8 bits, and it makes it applicable to networks that previously wasn't enough stable with regular fp16

[erogol]

Awesome! Thanks for the explanation.