How to enable BF16 on tpus?
With notebook_launcher(main, use_fp16=True) my data still fp32
2021-10-06T12:19:17Z
YEs that would be expected, this does not control bfloat16. Support for bfloat16 on TPUs is not in Accelerate yet.
2021-10-06T15:33:07Z
Any way I can do so? Torch xla docs say to use env var, but it doesn’t work too.
2022-08-11T16:58:50Z
I am trying to figure out about using bf16 on TPU. What I read fromthis GH issueis that setting the environment variableXLA_USE_BF16=1should automagically work?Can anyone confirm if this does in fact let one train on TPU with Accelerate in bf16? Or am I mis-reading this.
2022-08-11T17:04:22Z
sam-writer:Can anyone confirm if this does in fact let one train on TPU with Accelerate in bf16That’s one way, yes. But accelerate also now supportsmixed_precision="bf16". You can also specify bf16 duringaccleerate configRead more here on our docs:Training on TPUs with 🤗 Accelerate