Support for Training with BF16
Author: JamesDeAntonis
Pull request: yes
Labels: name: WIP

# What does this PR do?

As seen in [this pr](https://github.com/huggingface/transformers/pull/10956), there is demand for bf16 compatibility in training of transformers models. The pytorch folks just added [this feature](https://github.com/pytorch/pytorch/pull/61002) to their master branch, so we are now able to work on adding it to this repo.  This pr follows from [this issue](https://github.com/huggingface/transformers/issues/13170).

Fixes https://github.com/huggingface/transformers/issues/13170

------------------
(OP edited by @stas00)

Also merged here and adapted changes proposed by @manuelciosici at https://github.com/huggingface/transformers/pull/14448

This PR:

- adds helper utils: `require_torch_bf16` and `is_torch_bf16_available`
- modifies `invert_attention_mask` and one `forward` in t5 to include bf16 mode switches

HF Trainer: 
- adds `--bf16` and `--bf16_full_eval` modes - same as fp16 equivalents
- renames and deprecates `--fp16_backend` and replaces it with `--half_precision_backend` - since we now have more than one half precision mode

Tests:
- adds `--bf16` and `--bf16_full_eval` tests

@sgugger, @LysandreJik, 

Also tagging @patrickvonplaten, @patil-suraj since once this is merged you can start sending users that have problems with bf16 pre-trained models and have Amphere hardware to use this `--bf16` mode. Deepspeed bf16 support will follow soon.

Comments:
ok, pt-nightly installed so that I could test the new functionality.

so I added:
```
- `require_torch_bf16`
- `is_torch_bf16_available`
- super basic test that validates that `--bf16` doesn't fail
- placeholder for the future bf16 full eval test
```

So now the tests should be expanded to actually validate that bf16 is happening, with some number checking - e.g. we could check that the numbers are indeed bf16.


I'm not observing as much of a memory improvement as I expected.

Memory improvements I'm seeing are 0-15%, whereas I expected around 40% (per derrick's observation [here](https://github.com/huggingface/transformers/pull/10956#issuecomment-841431396)).  Is there anywhere in master where autocast is disabled for some reason?  For example, that was going to be the case [here](https://github.com/huggingface/transformers/pull/10956/files#diff-ebaff1224cad711fd3cefb771ce17d1392ae2dfc7f74dc7da9dd014d7642a344R308), but that change is not currently in master.

The two questions I just commented are part of my digging into whether there is a bug somewhere.

EDIT: I found it interesting that `fp16` was giving similarly lackluster gains as `bf16`.  That suggests it's not a `bf16`-specific issue

I saw some recent comments on the torch slack that suggestion that bf16 hasn't quite been figured out performance-wise and can actually be slower depending on the hardware. One issue being cuDNN has no `bfloat16` at the moment, the other issue is that many bf16 kernels are simply not there, so it falls back to some slower functionality.

May I suggest to ask this question on https://discuss.pytorch.org/ and hopefully some knowledgeable devs with experience could give us a definitive answer. Please share the link if you do.

I think on our side the main priority for providing bf16 support is to overcome the overflow issue in models pretrained in mixed bf16, performance being secondary. But of course, it'd be great to actually benefit from the new Ampere cards which have a lot of power but almost a year passed and we still can't quite harness that power.

BTW, which card are you testing it with?

>BTW, which card are you testing it with?

RTX A6000.  I'm pretty sure it's not related to this pr, per [this](https://discuss.pytorch.org/t/amp-autocast-not-faster-than-fp32/111757/11?u=jadeantonis)

Thank you for posting there, @JamesDeAntonis. Let's see if Piotr has some encouraging feedback. 

Otherwise the whole thing is very veiled at the moment as nobody wrote any definitive answers.

Also fyi, I updated `s/fast_dtype/dtype/` as it changed in nightly. But the current nightly has a broken `is_bf16_supported()` function, so it won't work - I reported it - hope should be fixed in a day or two.

So best don't update your nightly just yet.

Quoting from: https://github.com/pytorch/pytorch/issues/57806#issuecomment-834697571
> Concerning the lousy "speedup" with amp on A100, first of all I'd expect less of a relative difference because A100 should use TF32 tensor cores by default for FP32 (non-amp) runs, which is 2X less throughput than FP16 for matmuls on paper, but much faster than not using tensor cores at all. This closes the performance gap to amp, so I do expect the FP32 vs amp difference to be more modest on Ampere. It's possible that with FP32 backed by TF32 library math, ops that benefit most from tensor cores (ie matmuls, convs) have been accelerated enough that the network is mainly bound by CPU overhead or by ops Amp doesn't affect as much, so turning on Amp doesn't squeeze much more blood from the stone.

So if that is so, then you're not seeing an improvement from amp/bf16 because behind the scenes it already uses tf32.

We definitely need some more definitive guides as currently we can only collect such comments shared here and there and no proper document that covers all the grounds.

>So if that is so, then you're not seeing an improvement from amp/bf16 because behind the scenes it already uses tf32.

Interesting, but what about the below snippet from [here](https://moocaholic.medium.com/fp64-fp32-fp16-bfloat16-tf32-and-other-members-of-the-zoo-a1ca7897d407)

>For comparison, A100’s peak performances are:
>FP32 without tensor core: 19.5 TFLOPS
>TF32 tensor core: 156 TFLOPS (so, using TF32 in place of FP32 can give you an easy speed improvement).
>FP16/BF16 tensor core: 312 TFLOPS (so, thoughtfully designed switch to FP16/BF16 can give you more speed improvements, but the costs are higher).

It looks like the gains are still "immodest" in the presence of an fp16/bf16 tensor core.  Is the point that 156 TFLOPS is already so fast that further improvements are not worth the costs of making the switch from fp32 to fp/bf16?  Because otherwise, it should still be at least somewhat faster, not slower.

I don't yet have the understanding of this domain to comment intelligently. My feeling is that the devil is in the detail and will heavily depend on each model's components. And it's best to discuss this subject matter on the pytorch side where the experts who understand this domain are.

Using my limited understanding my answer to your question would be:
Most likely if you were to take a single tensor and run it through an OP that natively supports TF32 and BF16 you should see the numbers you quoted. But since there is a lot of back and forth casting happening around amp and not all ops support these native functions, the overall results with hundreds of different ops combined in the single fwd/bwd pass of a model the results are quite different.

In lieu of having an expert advice the other approach is to run your code through a native torch profiler, watch which ops get invoked on what dtypes, look them up whether they support the new tensor cores, etc.


Updates:
- pt-nightly from 09.01 can be used with this PR (dates before that had a bug)
- merged the 2 bf16 util functions as suggested by Sylvain

@JamesDeAntonis, you now have a green light to address the proposed changes after updating your install of pt-nightly, then when it's done we will update/complete the tests and then we can merge this.

If something is unclear please let us know. If you're not sure how to deal with deprecation, then you can just complete the new API and will add the deprecation afterwards. But you can look up at other cli args deprecations done in `training_args.py`.

Thank you!

Update on this: my teammate is investigating the slowdown and doing some tests on both inference and training.  He should have some results pretty soon that we can work with in a resumed discussion, including some new commits.  Thanks for all your help so far!

Awesome, thank you for the update, @JamesDeAntonis!

Hi @stas00, what do you think of these results and justifications?  All numbers from `t5-3b` on A100 cards

```
bf16 train:
INFO     - Finished in 140.42053532600403s
INFO     - Peak memory usage: 68.577 GB
fp32 train:
INFO     - Finished in 131.2668480873108s
INFO     - Peak memory usage: 71.861 GB
bf16 generate 32 tokens:
INFO     - Finished in 1.271615743637085s
INFO     - Peak memory usage: 6.927 GB
fp32 generate 32 tokens:
INFO     - Finished in 1.2117650508880615s
INFO     - Peak memory usage: 13.057 GB
```

## Some justifications:

Memory:
* 32-token generation: 47% improvement because, in the fp32 case, even though computations are done in 19-bit, all the 32-bit weights are stored in memory. in the bf16 case, memory is only allocated for 16-bit weights
* Training: 5% improvement because the only difference between fp32 and bf16 is that fp32 does default auto-casting to tf32/bf19 while bf16 does auto-casting to bf16.  so, the gains come from the 16% bit reduction during computation

Time:
* 5-6% time increase both times when using bf16.  I don't understand why this is happening

@stas00 one other detail is that we're having trouble training to the same loss as regular precision (1.5 for bf16 amp vs 1.2 for full precision).  Furthermore, when we generate with the 1.5-loss model, we get gibberish regardless of whether generating at full precision or half.  This leads me to question whether our branch is completely correct.

With this in mind, I don't understand why we wouldn't scale when training with bf16.  Rationale: if fp32 is -1000.0 to 1000.0 (precise to tenth's place), fp16 is like the integers -500 to 500 and bf16 is like the even integers -1000 to 1000.  To avoid underflow, fp16 amp convention in this analogy is to scale by a factor of 10 to make fp16's most precise unit (integer) analogous to fp32's most precise unit (tenth's place).  By this logic, bf16 should be scaled by 20 to have the same effect.

Do you agree with that logic, or do you understand where I go wrong?

> Hi @stas00, what do you think of these results and justifications? All numbers from `t5-3b` on A100 cards
> [...]
> Memory:
> 
> * 32-token generation: 47% improvement because, in the fp32 case, even though computations are done in 19-bit, all the 32-bit weights are stored in memory. in the bf16 case, memory is only allocated for 16-bit weights
> 
> * Training: 5% improvement because the only difference between fp32 and bf16 is that fp32 does default auto-casting to tf32/bf19 while bf16 does auto-casting to bf16.  so, the gains come from the 16% bit reduction during computation

Shouldn't bf16 be 2x faster than tf32 according to nvidia GPU specs? At least for some ops?

I don't suppose we have a way to tell pytorch to tell cuda not to cast to tf32 - so that we could compare bf16 to the actual fp32.

I'd say post all these benchmarks in that pytorch thread, since that's the bf16 experts are. And ask whether what you got makes sense and if it doesn't why and how can we/they fix that.

> With this in mind, I don't understand why we wouldn't scale when training with bf16.

Excellent catch! It's because we forgot to do that! It's currently done for fp16 only:

https://github.com/huggingface/transformers/blob/5e3b4a70d3d17f2482d50aea230f7ed42b3a8fd0/src/transformers/trainer.py#L436-L444

while at it perhaps check if there are other `if args.fp16` checks that need to have ` or args.bf16` added.

>Excellent catch! It's because we forgot to do that! It's currently done for fp16 only:

Ok, and I think the default of `2 ** 16` would work for bf16, because precision is pooled up by 16 bits (by the way, I think `2 ** 16` is overkill for fp16 because precision for fp16 is only pooled up by only 13 bits [the other three bits are saved by decreasing range, ie the source of the original issue], but it doesn't really matter)

Let's see if your latest code leads to a better outcome in your quality and speed benchmarks.

🎉🪄🥇🏆🚀

Unfortunately, it doesn't seem like scaling fixed the issue.  The loss went down to the fp32 level (actually eclipsed it), but inference still gave gibberish

Sorry to hear it didn't help, James.

Here are some ideas to potentially proceed with:

Is this something we could ask for the pytorch team to reproduce? 

i.e. ideally writing a few lines of code that they could reproduce the issue with?

Do you know if the inference works ok, if you were to train in amp/bf16 but then doing inference in fp32? and if it has to do with amp/bf16 or full bf16?

Perhaps something is wrong only in the inference stage?

The other or an additional approach could be to take a normally trained public model and to try to run inference on it in (1) amp/bf16 (2) full bf16 and comparing the outcome with the fp32 mode?

Hi @JamesDeAntonis 
I am trying to train mt5-xxl-13B model with 8x40GB A100.
I was wondering what is the condition for this PR. Is this ok to use it for training or any red flags?

Hi James, I'm trying to fine-tune T5-3B on a single A100 GPU (40gb memory) and I tried this PR out of desperate search. It seems like a promising direction to use `bf16` as it's natively supported by pytorch. However, while `fp16` with `amp` didn't, this version of the code seems to give `CUDA_MEMORY_ERROR` even with batch size 1

OK, since we have 2 half-baked PRs, 

https://github.com/huggingface/transformers/pull/13207
https://github.com/huggingface/transformers/pull/14448
I'm going to try to merge the 2 to keep the credits and start a new PR.

If you have something to push now is the time.

@manuelciosici, FYI: we will deal with deepspeed in a separate PR https://github.com/huggingface/transformers/pull/14569 - in particular since the ZeRO3 support hasn't been merged yet and we always need a new release from deepspeed to be able to update our integration side.

@sgugger, please kindly have a look. I merged 2 PRs and cleaned things up and added a single deprecation.

I also reverted the earlier attempt to use a shared `--half_precision_full_eval` since it didn't make sense -  `--fp16_full_eval` and `--bf16_full_eval` are similar but 2 different modes requiring different code. If we want a shared one then we have to additionally require either `--fp16` or `--bf16` and then adjust the logic accordingly. If you prefer that let me know.

Since bf16 has a much larger dynamic range most of the fp16 workarounds of that type aren't needed. So I grep'ed for `if torch.float16` checks and I didn't see anything other 2 places. I'm sure I may have missed some, but it'll surely let itself known when we start using it.

Note, I've updated the OP with the up-to-date list of changes, so please refer to it for an overview.

So I think we just need a couple of tests and if everybody is happy this is good to go. (tests added)

The CI failure is unrelated.

OK, a few tests added.

@JamesDeAntonis and @manuelciosici - please have a look - let me know if anything else is needed in your opinion. Thanks.

@sgugger, would it help to document the bf16 API as experimental and a subject to change at a moment's notice? 

Yes please!

Thanks a lot for the review and the suggestions, @manuelciosici - all integrated, plus added a warning that this API is experimental, so if once we start using it we find that we could improve it we can.

I have been working on a guide to all these new modes including tf32, @manuelciosici, et al - if you get a chance to proofread, please have a look at https://github.com/huggingface/transformers/pull/14579

Thank you!