Xla: model_parallel doesn't cover the whole dataset

[nosound2]

In the example

test/test_train_mnist.py

there are 10000 test images, and each of the 8 workers processes 1152 of them. The rest 784 images are never accessed. I want to implement

• all test images are processed
• the returned results order can be manipulated to match the dataset order

Is this possible with model_parallel? Question 2 is interesting even if each worker has to process the same number of batches, - are there any return order guarantees, how can we know what indices have been processed?

And a related question, as a workaround, how can I run a single TPU worker but with the same simplicity (no model_parallel) as a CUDA worker? Just replacing device = 'cuda' with device = 'xla:1' or similar didn't work for me.

Thank you!

[dlibenzi]

All the 8 cores batches must be the same.
If you set --batch_size=125 you should be good.

You can run on single core like this:

import torch_xla_py.xla_model as xm

device = xm.xla_device()
model = MNIST().to(device=device)
…
for batch_idx, (data, target) in enumerate(train_loader):
  data = data.to(device=device)
  target = target.to(device=device)
  optimizer.zero_grad()
  o = model(data)
  l = loss(o, target)
  l.backward()
  xm.optimizer_step(optimizer)

[nosound2]

Regarding the single core run, it is very fast for me on inference but it is extremely slow on loss calculation. Specifically the following line takes a lot of time:

class NLLLoss(_WeightedLoss):
    ...
    def forward(self, input, target):
        ...
        input_cpu = input.cpu()

it also take increasingly more time from batch to batch. 1 second, then 3 seconds, 6, 9 and so on. Any ideas?

[dlibenzi]

So the DataParallel issues the tensor's graph barrier automatically:

https://github.com/pytorch/xla/blob/0cfb18127cdd23601f4ca8dc5b73ac594dcb8f6b/torch_xla_py/data_parallel.py#L59

If you are using a by-hand approach, you will need to issue it.

Keep in mind though that not using DataParallel you will not be overlapping sending data to TPU with TPU computations, so throughput could be considerably lower.

[nosound2]

Thanks a lot, with xm.mark_step() it is flying for me now!

[dlibenzi]

Great!
Though I suggest you using DataParallel anyway, and eventually pad the latest batch with other samples, if performance means something to you.
Unless you have a very heavy TPU step time, and very light data sent to TPU, overlapping the TPU data send with TPU computation can boost performance quite a bit.

[nosound2]

1. It is still good to have a very simple change that one can do to his code as a sanity check.
2. Regarding your proposal with padding, can be nice to have it automatically. And also it returns me to the question I asked before, what is the order of returned values? How do I aggregate 8 returned results into original dataset order?

[dlibenzi]

1. Absolutely. But when it comes to performance, there might be a few mods which can maximize the overall training costs in $$ terms.

2. It is hard to do automatically, as it depends on the dataset and model architecture. For single core inference it is easy. Just use one core and set drop_last=False in the DataParallel creation. This will cause an extra compile at the last batch, but due to service side compile caching, the cost will be much smaller at the next run. But when replication (multi core) is in place, you really need to have the same computation (hence shapes) running on all cores.

[mruberry]

Closing because it appears the initial question was resolved.