Xla: tensor has requires_grad=True on GPU, false on TPU

[hrbigelow]

🐛 Bug

To Reproduce

Steps to reproduce the behavior:
GPU colab
TPU colab

Both Colabs run the same code, but on GPU or TPU. Cell 3 (in either workbook), shown below, the second line should be commented out, and the second to last line, break at model.py:350

in the GPU version, batch.mel_input.requires_grad == True, while in the TPU version, it is False. Both notebooks use the same version of torch (while the TPU additionally uses torch_xla). The ae-wavenet codebase simply doesn't load torch_xla if it is run with the -hw GPU option. Also, this problem is the same whether run in single-process TPU or multi-process using xmp.spawn.

!cd ae-wavenet && git checkout master && git pull
# !cd ae-wavenet && git checkout da3141c

# delete all currently loaded ae-wavenet modules
import sys
d = list(map(lambda m: m[0], filter(lambda m: hasattr(m[1], '__file__') and 'ae-wavenet' in m[1].__file__, sys.modules.items())))
def dn(n): del sys.modules[n]
list(map(dn, d))

%env TF_CPP_VMODULE=tensor=5
%env TF_CPP_VMODULE=
%env XLA_IR_DEBUG=
VER=!git -C ae-wavenet rev-parse --short HEAD
XLA_SAVE_TENSORS_FILE='xla_save_tensors.{}.txt'.format(VER[0])
XLA_METRICS_FILE='xla_metrics.{}.txt'.format(VER[0])
!echo > {XLA_SAVE_TENSORS_FILE}
!echo > {XLA_METRICS_FILE}
%env XLA_METRICS_FILE={XLA_METRICS_FILE}
%env XLA_SAVE_TENSORS_FILE={XLA_SAVE_TENSORS_FILE}
%run -b ae-wavenet/model.py:350 -d ae-wavenet/train.py new -af ae-wavenet/par/arch.ae.json -tf ae-wavenet/par/train.basic.json -hw TPU-single -pi 1 -lrr 1e-6 -lrs 0 -nw 100 -nb 4 -si 5000 ./basic.full.%.ckpt ae-wavenet/dat/librispeech.some.dat
# !python ae-wavenet/train.py new -af ae-wavenet/par/arch.ae.json -tf ae-wavenet/par/train.basic.json -hw TPU-single -pi 1 -lrr 1e-6 -lrs 0 -nw 100 -nb 4 -si 5000 ./basic.full.%.ckpt ae-wavenet/dat/librispeech.some.dat

I think this is the relevant code but for some reason a breakpoint will not stop here.

#data.py:447
        vb = VirtualBatch(self.batch_size, self.max_wav_len, self.max_mel_len,
                self.max_embed_len, self.n_mel_chan)
        vb.mel_input.detach_()
        vb.mel_input.requires_grad_(False)
        for b in range(vb.batch_size):
            vb.set_one(b, self.calc_slice(), self)

        assert vb.valid()
        if self.target_device:
            vb.to(self.target_device)
        vb.mel_input.requires_grad_(True)

Expected behavior

Environment

• reproducible on XLA backend [CPU/TPU]:
• torch_xla version:
• Any other relevant information:

Additional context

[hrbigelow]

By the way, I'm happy to help with this - would a minimal example be helpful?

One problem I have unfortunately is that the Colab TPU that I'm using appears to be "sick" now - even calls like xm.xla_device() are taking several seconds, and the code that used to run now appears to hang indefinitely. It's been this way for a few days across several resets.

[dlibenzi]

A small repro would help, thanks!

A xm.xla_device() call the first time you call it in a binary, can take a few seconds as it has to initialize the TPU system.
WRT the hang, did you change wheels/dockers versions between the no-hang and the hang tests?

[hrbigelow]

WRT the hang, did you change wheels/dockers versions between the no-hang and the hang tests?

Probably - I'm not entirely sure. I occasionally had to restart the Colab runtime either because it timed out ('Runtime Disconnected' -> Reconnect), or because a cell was running and it was not responsive to the 'Stop' button or Ctrl-M + I. And, so I restarted it. And, sometimes (not always), the restart required rerunning the first notebook cell, which installs torch_xla. In my Colab, the wheels are defined (from Jin Young) as:

XRT_VERSION = "nightly"
DIST_BUCKET = "gs://tpu-pytorch/wheels"
TORCH_WHEEL = "torch-{}-cp36-cp36m-linux_x86_64.whl".format(XRT_VERSION)
TORCH_XLA_WHEEL = "torch_xla-{}-cp36-cp36m-linux_x86_64.whl".format(XRT_VERSION)
TORCHVISION_WHEEL = "torchvision-0.3.0-cp36-cp36m-linux_x86_64.whl"

I'm not sure what the background setup is in Colab - is it in Docker? How could I tell what version is being run?

[dlibenzi]

OK, you are installing nightly wheels.
Is it really hanging or compiling?
Because when you restart the TPU, you may wipe the compilation cache.
In case of apparent hanging, run on one core and:

export TF_CPP_VMODULE=tensor=5

This will log a bunch of stuff and should be ready to tell whether you are hanging on a compile.

[hrbigelow]

Hi @dlibenzi ,

Thanks for the suggestion. So, it seems like it may be doing a recompilation, but for the last 7 minutes. The last several lines of output are below, and here is the complete output:

test_output.txt

2019-11-23 19:24:05.629247: I torch_xla/csrc/tensor.cpp:1112] Executing IR graph hash 4512587929171635148 on device TPU:0 ...
2019-11-23 19:24:05.769657: I torch_xla/csrc/tensor.cpp:1117] Executing IR graph hash 4512587929171635148 on device TPU:0 done!
2019-11-23 19:24:05.772178: I torch_xla/csrc/tensor.cpp:981] Waiting on device barrier for device TPU:0 ...
2019-11-23 19:24:05.772227: I torch_xla/csrc/tensor.cpp:984] Waiting on device barrier for device TPU:0 done!
2019-11-23 19:24:05.772253: I torch_xla/csrc/tensor.cpp:1022] Tensors graph hash 8978071329308415931 on device TPU:0
2019-11-23 19:24:05.783208: I torch_xla/csrc/tensor.cpp:1280] SyncTensorsGraphSize=1482
2019-11-23 19:24:05.785649: I torch_xla/csrc/tensor.cpp:1293] Compiling IR graph hash 8978071329308415931 on device TPU:0 ...
2019-11-23 19:24:12.104844: I torch_xla/csrc/tensor.cpp:1298] Compiling IR graph hash 8978071329308415931 on device TPU:0 done!
2019-11-23 19:24:12.259199: I torch_xla/csrc/tensor.cpp:1112] Executing IR graph hash 8978071329308415931 on device TPU:0 ...
2019-11-23 19:24:12.400923: I torch_xla/csrc/tensor.cpp:1117] Executing IR graph hash 8978071329308415931 on device TPU:0 done!
2019-11-23 19:24:12.402783: I torch_xla/csrc/tensor.cpp:981] Waiting on device barrier for device TPU:0 ...
2019-11-23 19:24:12.402838: I torch_xla/csrc/tensor.cpp:984] Waiting on device barrier for device TPU:0 done!
2019-11-23 19:24:12.402856: I torch_xla/csrc/tensor.cpp:1022] Tensors graph hash 1362349540356948101 on device TPU:0
2019-11-23 19:24:12.410994: I torch_xla/csrc/tensor.cpp:1280] SyncTensorsGraphSize=1483
2019-11-23 19:24:12.413379: I torch_xla/csrc/tensor.cpp:1293] Compiling IR graph hash 1362349540356948101 on device TPU:0 ...
2019-11-23 19:24:18.674580: I torch_xla/csrc/tensor.cpp:1298] Compiling IR graph hash 1362349540356948101 on device TPU:0 done!
2019-11-23 19:24:18.828995: I torch_xla/csrc/tensor.cpp:1112] Executing IR graph hash 1362349540356948101 on device TPU:0 ...
2019-11-23 19:24:18.968631: I torch_xla/csrc/tensor.cpp:1117] Executing IR graph hash 1362349540356948101 on device TPU:0 done!
2019-11-23 19:24:18.970560: I torch_xla/csrc/tensor.cpp:981] Waiting on device barrier for device TPU:0 ...
2019-11-23 19:24:18.970620: I torch_xla/csrc/tensor.cpp:984] Waiting on device barrier for device TPU:0 done!
2019-11-23 19:24:18.970645: I torch_xla/csrc/tensor.cpp:1022] Tensors graph hash 12162494415976796445 on device TPU:0
2019-11-23 19:24:18.972593: I torch_xla/csrc/tensor.cpp:1280] SyncTensorsGraphSize=139
2019-11-23 19:24:18.973022: I torch_xla/csrc/tensor.cpp:1293] Compiling IR graph hash 12162494415976796445 on device TPU:0 ...
2019-11-23 19:24:19.278171: I torch_xla/csrc/tensor.cpp:1298] Compiling IR graph hash 12162494415976796445 on device TPU:0 done!
2019-11-23 19:24:19.292252: I torch_xla/csrc/tensor.cpp:1112] Executing IR graph hash 12162494415976796445 on device TPU:0 ...
2019-11-23 19:24:19.297617: I torch_xla/csrc/tensor.cpp:1117] Executing IR graph hash 12162494415976796445 on device TPU:0 done!
2019-11-23 19:24:19.299586: I torch_xla/csrc/tensor.cpp:981] Waiting on device barrier for device TPU:0 ...
2019-11-23 19:24:19.299636: I torch_xla/csrc/tensor.cpp:984] Waiting on device barrier for device TPU:0 done!
2019-11-23 19:24:19.299658: I torch_xla/csrc/tensor.cpp:1022] Tensors graph hash 5965884261943469193 on device TPU:0
2019-11-23 19:24:19.324906: I torch_xla/csrc/tensor.cpp:1280] SyncTensorsGraphSize=3713
2019-11-23 19:24:19.331429: I torch_xla/csrc/tensor.cpp:1293] Compiling IR graph hash 5965884261943469193 on device TPU:0 ...

So, that last message was at 19:24:19 and the time is now 19:31 and it is still spinning. I haven't seen it take anything like this long before though.

Thanks again,

Henry

[dlibenzi]

The graph does not seem huge (3713 nodes).
I'd let it run and see where it lands.
Eventually we can collect the HLO graphs and we can see whether there is something which triggers a compiler bad scalability issue.

[hrbigelow]

Okay, it is continuing running. I've attached the updated output. I see a few places in which there are ~20 minute delays, lines:

# line 477
2019-11-23 19:24:19.331429: I torch_xla/csrc/tensor.cpp:1293] Compiling IR graph hash 5965884261943469193 on device TPU:0 ...
2019-11-23 19:44:02.005397: I torch_xla/csrc/tensor.cpp:1298] Compiling IR graph hash 5965884261943469193 on device TPU:0 done!

# line 489
2019-11-23 19:44:19.702109: I torch_xla/csrc/tensor.cpp:1293] Compiling IR graph hash 988384775979973159 on device TPU:0 ...
2019-11-23 20:04:25.192849: I torch_xla/csrc/tensor.cpp:1298] Compiling IR graph hash 988384775979973159 on device TPU:0 done!

# line 926
2019-11-23 20:05:08.159797: I torch_xla/csrc/tensor.cpp:1293] Compiling IR graph hash 9409769670077423997 on device TPU:0 ...
2019-11-23 20:25:23.317723: I torch_xla/csrc/tensor.cpp:1298] Compiling IR graph hash 9409769670077423997 on device TPU:0 done!

There are also a few cases of one-minute delays:

#line 928
2019-11-23 20:25:24.305744: I torch_xla/csrc/tensor.cpp:981] Waiting on device barrier for device TPU:0 ...
2019-11-23 20:26:16.023602: I torch_xla/csrc/tensor.cpp:1117] Executing IR graph hash 9409769670077423997 on device TPU:0 done!

But, after the first 1000 lines, there are no more big delays. The full output is now 7400 lines long.

test_output2.txt

[hrbigelow]

Okay, I've created a minimal Colab example of the difference in behavior between GPU and TPU code here

And the test code is xla-tests and the revision I just tested is b319a1f

I will continue to play with it and let you know what I find.

Thanks again!

[dlibenzi]

Your code acts differently from GPU to TPU AFAICT.
On GPU, you set requires_grad after the tensor has been moved to GPU, while when you set back requires_grad here, the tensor will still be on CPU:

https://github.com/hrbigelow/xla-tests/blob/2ab3b236493452014f302ed65e2be40fb1e97ef9/requires_grad_test.py#L23

[hrbigelow]

=== please see next comment ===

Okay, so I corrected a few things in the code - apologies for that confusing test - in fact, the detach_() and requires_grad_(False) calls are completely unnecessary and I've removed them.

But, the issue is clear now. pl.ParallelLoader moves the tensor from CPU to TPU. However, it does not preserve the requires_grad flag when it does so.

See here and here

[hrbigelow]

Hi @dlibenzi Please ignore the previous example, I created a much simpler test, see here

Thanks again!

[dlibenzi]

This will likely fix that:

https://github.com/pytorch/xla/pull/1411