Xla: LSTM model compiles very slowly

[BadrYoubiIdrissi]

I'm trying to train a language model on TPU. I've written some code that took way too long to compile the model for XLA. Here's some dummy code that reproduces my issue. I haven't come across any examples implementing RNNs for NLP in the test cases so maybe I'm doing something wrong?

import torch
import torch_xla
import torch.nn as nn
import time
import model

import torch_xla_py.xla_model as xm
import torch_xla_py.data_parallel as dp

class Loader:
    def __iter__(self):
        return self
    def __next__(self):
        return torch.rand(100, 100, 100)

m = nn.LSTM(100,100)
devices = xm.get_xla_supported_devices()
m = dp.DataParallel(m, device_ids=devices, batchdim=1)

def train_loop(model, loader, device, context):
    h = torch.rand(1, 100, 100).to(device), torch.rand(1, 100, 100).to(device)
    t = 0
    for _, e in loader:
        tic = time.time()
        out = model(e, h)
        toc = time.time()
        print(toc-tic)

m(train_loop, Loader())

I've also dumped XLA logs using export XLA_FLAGS=--xla_dump_to=/tmp/xlalogs and it gave the following logs :

module_0000.after_optimizations.txt
module_0000.after_optimizations-buffer-assignment.txt
module_0000.before_optimizations.txt
module_0000.ir-no-opt.txt
module_0000.ir-no-opt-noconst.txt
module_0000.txt

[dlibenzi]

Seems like you are using local XLA CPU, not TPU?
With pytorch/xla you cannot measure compile time like that.
Things are done asynchronously and compilation might even not happen within your timed section.
Did you export the environment variables . required to setup the TPU run?

https://github.com/pytorch/xla/blob/master/README.md

export XRT_TPU_CONFIG="tpu_worker;0;<IP of the TPU node>:8470"

You can get the graphs locally by using the XLA_SAVE_TENSORS_FILE and XLA_SAVE_TENSORS_FMT environment variables, or you can use (using out like in your example):

print(torch_xla._XLAC._get_xla_tensors_text([out]))

Or:

print(torch_xla._XLAC._get_xla_tensors_hlo([out]))

[dlibenzi]

While debugging, I also suggest you export:

export XLA_IR_DEBUG=1 XLA_HLO_DEBUG=1

This gives Python location information in the graphs.

[BadrYoubiIdrissi]

I have indeed explicitely used CPU to test the code out and it struck me as peculiar that the test networks compiled quite fast while this one takes around 10 to 20 minutes. The tic toc was to test the native pytorch cpu! But yeah I did notice that the execution was done asynchronously.

So does this mean that it is normal for an LSTM network to compile in such a long time (even on CPU) compared to other modules? I am going to try it out on TPU to see if it is much faster. I'll also check if the graph is correct. Thank you for the debugging suggestions!

[dlibenzi]

What is likely happening there is that you have a model that makes us switch to pytorch CPU device as we do not implement a given operation, so we end up compiling (which is pretty fast on XLA CPU), and then running it.
The slow you are experiencing is the running it part.
If you can:

print(torch_xla._XLAC._xla_metrics_report())

After each step, that will give us a better idea of the operators that we do not lower.

[BadrYoubiIdrissi]

The code seems to run at a normal speed after the first iteration though. I'm going to run with the metrics asap

[BadrYoubiIdrissi]

Metric: OutboundData
  TotalSamples: 3
  Counter: 30.90MB
  ValueRate: 628.23MB / second
  Rate: 60.9892 / second
  Percentiles: 1%=393.75KB; 5%=393.75KB; 10%=393.75KB; 20%=393.75KB; 50%=15.26MB; 80%=15.26MB; 90%=15.26MB; 95%=15.26MB; 99%=15.26MB
Metric: TransferToServerTime
  TotalSamples: 3
  Counter: 043ms23.650us
  ValueRate: 01s268ms36.562us / second
  Rate: 88.419 / second
  Percentiles: 1%=006ms552.906us; 5%=006ms552.906us; 10%=006ms552.906us; 20%=006ms552.906us; 50%=018ms932.658us; 80%=020ms538.086us; 90%=020ms538.086us; 95%=020ms538.086us; 99%=020ms538.086us
Counter: CreateDataHandles
  Value: 14
Counter: CreateXlaTensor
  Value: 1421
Counter: DestroyXlaTensor
  Value: 304
Counter: SyncTensorsToData
  Value: 6
Counter: XRTAllocateFromTensor_Empty
  Value: 10
Counter: XrtSessionCount
  Value: 2

This is what I get when printing metrics in the first iteration.

EDIT : I don't understand how or why but when i change this line out = model(e, h) to model(e,h) the model compiles and runs fast. Do we have to manually delete the variables in the train loop?

I've also tried to run the model by transposing it to the xla device (CPU in this case) and the loop is done without any issues.

This code runs without any issue :

import torch
import torch_xla
import torch.nn as nn
import model

device = 'xla:0'

m = nn.LSTM(100,100).to(device)
a = torch.rand(100, 100, 100).to(device)
h = torch.rand(1, 100, 100).to(device), torch.rand(1, 100, 100).to(device)
for i in range(100):
    out = m(a, h)

[dlibenzi]

From the metrics report you posted, that loop does not do anything.
I mean, we build an IR graph mapping your tensor operations, but you never materialize it (look at the result).

[BadrYoubiIdrissi]

I'm sorry I don't quite get how I can materialize it? Could you explain a little further?

[dlibenzi]

In "normal" pytorch, when you do:

c = a + b
d = c + 3

The tensors c and d are actually calculated at the time that Python line executes.
In pytorch/xla we use an approach called Lazy Tensor which instead of immediately computing the pytorch operations, it builds an IR graph.
It is only when you "look" (ie, print()) at the value, that the whole graph rooted to the materialized value, gets executed.

[BadrYoubiIdrissi]

Okay I think I understand ! So this code

import torch
import torch_xla
import torch.nn as nn
import model

device = 'xla:0'

m = nn.LSTM(100,100).to(device)
a = torch.rand(100, 100, 100).to(device)
h = torch.rand(1, 100, 100).to(device), torch.rand(1, 100, 100).to(device)
for i in range(100):
    out = m(a, h)

only seems to run fast because nothing is actually evaluated ? That would also explain why deleting the variable makes it faster. However it doesn't explain why the code goes faster after the first iteration in the original code (If the graph is built fast and it is only the evaluation that is costly)

[BadrYoubiIdrissi]

If the problem lies in an unsupported operation in the LSTM how would you suggest I go on about training a language model with pytorch/XLA? Has it ever been attempted in your tests?

[ailzhang]

@BadrYoubiIdrissi The code is expected to run much faster after the first iteration as it's hitting cache and it's designed in this way to speed up. We mainly look at the speed after no new graph compilation.
We are currently working on the model coverage but haven't done a language model ourselves yet. If you see the speed is still not good after the second/third iteration, and there're missing lowered ops, we should be able to prioritize those. Thanks!

[BadrYoubiIdrissi]

It does definitely speed up after the first iteration (nearly instant compared to the first) I'll update this issue once I have measured results on TPU with the language model. Thank you very much for your help!

[BadrYoubiIdrissi]

I finally (successfully) ran some tests on a real TPU and indeed compilation time is good. The same model ran on TPU is approximately 3 times faster (when normalizing with batch size*sequence length) than a GTX 1080. However the TPU's usage is very low (5% to 10%). How could I identify operations that are running slowly? (Maybe those processed on CPU). CPU usage is not at 100% either and Network usage isn't exceptionally large so I'm guessing it's not the data loading that is bottlenecking(?)

[dlibenzi]

The metrics report can give some insights:

print(torch_xla._XLAC._xla_metrics_report())

[dlibenzi]

I am closing this, but if there is further investigation feel free to re-open.