Does PyTorch not support bfloat16 when running on xla backend?
I'm still confused when dtype is set to bfloat16:
When torch.cuda.is_available() returns False, torch_xla uses the XLA:GPU backend. However, in this patch it seems to fall back to use the GPU backend.
Similarly, when torch.cuda.is_available() returns True, torch_xla uses the GPU backend, but the dtype is still forced to float16 instead of bfloat16.

Hi @baoleai good catch -- I meant to do what the err msg says. We don't need to force torch.float16 on cuda backend.

Thanks, will update the PR.

Ok, let me go over -- let me know if I am missing anything.

1. "When torch.cuda.is_available() returns False, torch_xla uses the XLA:GPU backend. However, in this patch it seems to fall back to use the GPU backend."

When cuda is not available (False), we go into this code path

if xr.is_bf16_supported() and not torch.cuda.is_available():
          # XLA:GPU with bfloat16 should run on `xla` backend
          # unless torch.autocast is compiled with cuda.
          backend = 'xla'
          self._cuda_bfloat16 = True

We uses xla backend for autocast.

2. Similarly, when torch.cuda.is_available() returns True, torch_xla uses the GPU backend, but the dtype is still forced to float16 instead of bfloat16.

This is the part I need to address, instead of else we want elif {cuda is not available}, to implement the intention laid out in the err msg.

So something like this, commit 50c1d2f1c598c2ec4ada8f7177861cc05056ff36

if dtype is None:
        dtype = torch.float16
elif dtype == torch.bfloat16 and not torch.cuda.is_available():
        if xr.is_bf16_supported():
          # XLA:GPU with bfloat16 should run on `xla` backend
          # unless torch.autocast is compiled with cuda.
          backend = 'xla'
          self._cuda_bfloat16 = True
        else:
          # This has been the default behavior for unsupported bfloat16 dtype
          dtype = torch.float16
          error_message = "In XLA:GPU autocast, but bfloat16 is not supported on this HW.\n"
          error_message += ("Using the default cuda autocast dtype float16.")

I am very confuse what this _cuda_bfloat16 actually means. Do you mind adding a comment above line 28 to explain all possible combinations and the expected behaviors?

+1, realized that this variable should be called _xla_bfloat16. Let me add a brief comment, too.

Something like this,

self._xla_bfloat16 = False  # True if xla backend with bfloat16 dtype.
if dtype is None:
   dtype = torch.float16
elif dtype == torch.bfloat16 and not torch.cuda.is_available():
   if xr.is_bf16_supported():
       # XLA:GPU with bfloat16 should run on `xla` backend
       # unless torch.autocast is compiled with cuda.
      backend = 'xla'
      self._xla_bfloat16 = True

Another confusing question is why there is only a special treatment for torch.cuda.is_available() is False and dtype=bfloat16 here, is it because XLA:GPU itself doesn't support bfloat16? But when torch.cuda.is_available() is True, there is no special treatment for bfloat16, then the following is still used

torch.set_autocast_xla_enabled(self.prev)
torch.set_autocast_xla_dtype(self.prev_dtype)

Hi @baoleai XLA:GPU supports bfloat16, but we were using the cuda backend for autocast, when torch.cuda.is_available() was false. Instead we want to use xla backend.

When torch.cuda.is_available() is true, we use cuda backend for autocast, and since we are entering/existing the autocast context via torch_xla.amp.autocast, we still need to set

torch.set_autocast_xla_enabled(self.prev)
torch.set_autocast_xla_dtype(self.prev_dtype)

if _xla_bfloat16 , shouldn't we use set_autocast_xla_enabled instead?

it will be set by torch.autocast, we need to set both -- since we are wrapping and calling torch.autocast. So the xla autocast is enabled if one is using torch_xla.amp.autocast, and as torch autocast is enabled we are also calling the torch.autocast with cuda or xla backend.

Can we leave a comment here, it is really confusing to see that if we are using xla_bf16, we need setup upstream autocast to enabled.

same question

ditto