Xla: Incorrect gradient on TPU after in-place operation

[myleott]

🐛 Bug

Some in-place operations produce incorrect gradients on TPU. They work fine on CPU/GPU, however.

To Reproduce

import torch
import torch.nn as nn

def run_test(device):
    torch.manual_seed(0)
    linear = nn.Linear(512, 1024).to(device=device)
    batch = torch.rand(16, 512).to(device=device)
    x = linear(batch)
    x[:, :768] = 0
    loss = x.sum()
    loss.backward()
    print('\tloss', loss)
    print('\tgnorm', linear.weight.grad.norm())

if __name__ == "__main__":
    import torch_xla.core.xla_model as xm
    tpu = xm.xla_device()

    print('TPU:')
    run_test(tpu)
    print('CPU:')
    run_test('cpu')

# output:
# TPU:
#         loss tensor(42.1703, device='xla:1', grad_fn=<SumBackward0>)
#         gnorm tensor(1519.3051, device='xla:1')
# CPU:
#         loss tensor(42.3902, grad_fn=<SumBackward0>)
#         gnorm tensor(2936.6663)

Expected behavior

While the losses above are similar, the grad norms are very different.

Environment

• torch_xla version: master (0.8+766bfef)

[dlibenzi]

Must be the norm() op.
I changed that with sum() and results are close:

CPU:
    loss tensor(42.3902, grad_fn=<SumBackward0>)
    gnorm tensor(1052793.7500)
TPU:
    loss tensor(42.3902, device='xla:0', grad_fn=<SumBackward0>)
    gnorm tensor(1054687., device='xla:0')

[dlibenzi]

Never mind ... I was running on XLA_CPU :man_facepalming:

[myleott]

Yep, if I change the norm to:

print('\tgrad[0]', linear.weight.grad[0])

I get:

TPU:
        loss tensor(42.1703, device='xla:1', grad_fn=<SumBackward0>)
        grad[0] tensor([2.0337, 1.3750, 2.5273, 3.1680, 2.6660, 2.0762, 3.0352, 2.3018, 1.4619, ...
CPU:
        loss tensor(42.3902, grad_fn=<SumBackward0>)
        grad[0] tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., ...

[dlibenzi]

I simplified and reduced sizes (and got the IR graph).
But yeah, something is off with view+inplace+backward in that particular case.

https://gist.github.com/dlibenzi/6f50850f5aa843b274f0806882fbfae8

[ailzhang]

Hi I'm taking a look at this issue. Will reply here once I've some updates.

[ultrons]

An observation, just playing with this example:

import torch
import torch.nn as nn

def run_test(device):
    torch.manual_seed(0)
    linear = nn.Linear(8,2048 ).to(device=device)
    batch = torch.rand(1, 8).to(device=device)
    x = linear(batch)


    # For linear layer the gradient of weights should be equal to input vector transpose repeated 
     g_ref = batch.t().repeat(1, 2048).to(device=device)
    loss = x.sum()

    loss.backward()
    print('\tloss', loss)
    print('\tgnorm', linear.weight.grad.norm())
    print('\tgrefnorm', g_ref.norm())

if __name__ == "__main__":
    import torch_xla.core.xla_model as xm
    tpu = xm.xla_device()

    print('TPU:')
    run_test(tpu)
    print('CPU:')
    run_test('cpu')



# TPU:
#       loss tensor(24.1235, device='xla:1', grad_fn=<SumBackward0>)
#       gnorm tensor(76.4076, device='xla:1')
#       grefnorm tensor(76.4077, device='xla:1')
# CPU:
#       loss tensor(24.1235, grad_fn=<SumBackward0>)
#       gnorm tensor(76.4054)
#       grefnorm tensor(76.4030)

For some reason even reference gradient (just repeated input tensor) seems to be different in both the TPU and CPU cases, should we investigate that ?

[dlibenzi]

We know the issue, which is related to the way PyTorch assumes view semantics during the backward call generation.
We are looking for the best fitting solution.

[myleott]

Looks like this is fixed after https://github.com/pytorch/pytorch/pull/36073, right?

[dlibenzi]

Looks like this is fixed after pytorch/pytorch#36073, right?

Should be. Thanks @ailzhang !

[ailzhang]

Yup let's close this via https://github.com/pytorch/pytorch/pull/36073 :D Thanks for the patience!