Unexpected f64 div

[ymwangg]

🐛 Bug

We observed unexpected f64 div in the backward pass of nn.MultiheadAttention.

To Reproduce

import torch_xla
import torch_xla.core.xla_model as xm
from torch import nn

device = xm.xla_device()
mod = nn.MultiheadAttention(768, 12, batch_first=True)
mod.to(device)

x = torch.rand(24,512,768).to(device)
y,_ = mod(x,x,x,need_weights=False)
y.sum().backward()

print(torch_xla._XLAC._get_xla_tensors_text([y]))
print(torch_xla._XLAC._get_xla_tensors_text([p.grad for p in mod.parameters() if p.requires_grad]))

In the forward pass IR:

%31 = f32[] xla::device_data(), xla_shape=f32[], device=GPU:0
%35 = f32[24,12,512,64]{3,2,1,0} aten::div(%34, %31), xla_shape=f32[24,12,512,64]{3,2,1,0}

In the backward pass IR

%73 = f64[] xla::device_data(), xla_shape=f64[], device=GPU:0
%78 = f64[24,12,512,64]{3,2,1,0} aten::div(%77, %73), xla_shape=f64[24,12,512,64]{3,2,1,0}

One interesting I noticed is the aten::div in the forward pass is through the at::Scalar& other signature:

at::Tensor XLANativeFunctions::div(const at::Tensor& self,
                                   const at::Scalar& other) {

while the aten::div in the backward pass is through the at::Tensor& other signature:

at::Tensor XLANativeFunctions::div(
    const at::Tensor& self, const at::Tensor& other,
    c10::optional<c10::string_view> rounding_mode) {

and in both cases other.scalar_type() is f64.

Expected behavior

There should not be f64 div as it's very costly on some GPUs.

Environment

• Reproducible on XLA backend [CPU/TPU]: GPU
• torch_xla version: master

Additional context

The div is coming from here.

I'm not sure if this is an issue of pytorch or it should be fixed in torch_xla.

[JackCaoG]

Hmm, @bdhirsh I think this has to do in some case we wrap the scalar as a tensor? Something similar to

xla/torch_xla/csrc/aten_xla_type.cpp

Line 64 in 6d83ba7

if (tensor.unsafeGetTensorImpl()->is_wrapped_number()) {

.

Maybe the right way to handle it is to downcast it same type as self if we know it is just a wrapped number.

[ymwangg]

@JackCaoG I just checked and it appears tensor.unsafeGetTensorImpl()->is_wrapped_number() is evaluated to false in this case.
Another way to do it is to check the number of elements like

 at::Tensor UnwrapNumber(const at::Tensor& tensor, at::ScalarType dtype) {
+  if (tensor.numel() == 1) {
+    return tensor.to(dtype);
+  }
   return tensor.unsafeGetTensorImpl()->is_wrapped_number() ? tensor.to(dtype)
                                                            : tensor;
 }

Not sure if this is the best way to do it.

BTW, not sure if it's relevant but I observed the denominator is SymFloat in this case.

const auto embed_size = SymFloat(query_.sym_size(-1));
const auto scaling_factor = embed_size.sqrt().sqrt();
const auto query = query_ / scaling_factor;

[JackCaoG]

if (tensor.numel() == 1) seems to be a hack... I would wait for Brian or @ezyang to help regarding why tensor is passed to backward and whether that's expected before we hack it in pytorch/xla layer.

[ymwangg]

I see what's going wrong here. The backward div has a different signature.

Tensor div(const Tensor& self, const Scalar& other) {
  return self.div(wrapped_scalar_tensor(other)); // redispatch!
}

- name: div.Scalar(Tensor self, Scalar other) -> Tensor
  self: div_tensor_self_backward(grad, at::lift_fresh(at::scalar_to_tensor(other)), self.scalar_type())
  result: self_t / other

So use the same wrapped_scalar_tensor in the backward signature should solve this issue. I'll submit a PR in pytorch to fix it.

[ezyang]

It would make me feel better if we just created a Scalar variant of div_tensor_self_backward and relied on a Tensor / Scalar division inside its body. It's duplicate-y but avoids wrapped_scalar_tensor naughtiness (which is not guaranteed to do anything! beware!)