Xla: Out of memory error?

[tmabraham]

Is the following an out-of-memory error from the TPU?:

The text just keeps scrolling with similar messages.

It's surprising I get this error, because all I wanted to do is have a batch of 512 for 224x224 images, which I thought the TPU could handle.

[tmabraham]

Experimenting a little more, it seems as if a similar error comes about when setting num_threads=2. But I have to look into this a little further.

[dlibenzi]

I can't see the error.
When there is a compilation error, we dump the XLA computation.
So that could be an OOM, but not necessarily so.
WRT image and batch size, it is not simply a matter of those dimensions, but the model code manipulating them.
Resnet50 for example cannot handle that batch size on TPU v3.

What is num_threads==2?
The number of worker threads for the DataLoader?
If that is the number of cores, we do not support replication with num_cores==2.

[tmabraham]

@dlibenzi Thanks for your response.
"Resnet50 for example cannot handle that batch size on TPU v3." --> this was a ResNet152

Yes num_threads=2 refers to the number of worker threads for the DataLoader.

[dlibenzi]

For resnet152 at minibatch 512 we probably need to wait for TPU v5 😉

[tmabraham]

@dlibenzi oh wow ok... This is a little disappointing because I can easily train with bs=128 on a Tesla T4 or P100 but only a max of 256 on a TPU. So I can only double the batch size with a TPUv3? That too, these aren't huge images, just 224x224.

Anyway, so is num_threads=2 not supported?

[dlibenzi]

You can try using BF16, and post the whole OOM log, to get an idea how far we are.

export XLA_USE_BF16=1

About num_threads, we don't really care. You need to tune the number depending on the vCPU cores you have.

[tmabraham]

@dlibenzi BF16 does not work either. Here is the log:

2019-12-03 23:49:11.662610: E   19849 tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.3442 = bf16[2048]{0} broadcast(bf16[] %constant.3438), dimensions={}
2019-12-03 23:49:11.662621: E   19849 tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %get-tuple-element.3436 = bf16[2048]{0} get-tuple-element((bf16[512,2048,7,7]{3,2,1,0}, bf16[2048]{0}, bf16[2048]{0}) %batch-norm-training.3433), index=2
2019-12-03 23:49:11.662632: E   19849 tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %constant.3437 = bf16[] constant(1.00136e-05)
2019-12-03 23:49:11.662643: E   19849 tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.3439 = bf16[2048]{0} broadcast(bf16[] %constant.3437), dimensions={}
2019-12-03 23:49:11.662676: E   19849 tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %add.3440 = bf16[2048]{0} add(bf16[2048]{0} %get-tuple-element.3436, bf16[2048]{0} %broadcast.3439)
2019-12-03 23:49:11.662711: E   19849 tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %sqrt.3441 = bf16[2048]{0} sqrt(bf16[2048]{0} %add.3440)
2019-12-03 23:49:11.662730: E   19849 tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %divide.3443 = bf16[2048]{0} divide(bf16[2048]{0} %broadcast.3442, bf16[2048]{0} %sqrt.3441)
2019-12-03 23:49:11.662744: E   19849 tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %get-tuple-element.3434 = bf16[512,2048,7,7]{3,2,1,0} get-tuple-element((bf16[512,2048,7,7]{3,2,1,0}, bf16[2048]{0}, bf16[2048]{0}) %batch-norm-training.3433), index=0
2019-12-03 23:49:11.662758: E   19849 tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %constant.1 = bf16[] constant(1)
2019-12-03 23:49:11.662774: E   19849 tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reshape.2 = bf16[1,1,1,1]{3,2,1,0} reshape(bf16[] %constant.1)
2019-12-03 23:49:11.662787: E   19849 tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.3 = bf16[1,1,1,1]{3,2,1,0} broadcast(bf16[1,1,1,1]{3,2,1,0} %reshape.2), dimensions={0,1,2,3}
2019-12-03 23:49:11.662799: E   19849 tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reshape.4 = bf16[] reshape(bf16[1,1,1,1]{3,2,1,0} %broadcast.3)
2019-12-03 23:49:11.662810: E   19849 tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.5 = bf16[512,2048,7,7]{3,2,1,0} broadcast(bf16[] %reshape.4), dimensions={}
2019-12-03 23:49:11.662821: E   19849 tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %multiply.3386 = bf16[512,2048,7,7]{3,2,1,0} multiply(bf16[512,2048,7,7]{3,2,1,0} %maximum.3385, bf16[512,2048,7,7]{3,2,1,0} %broadcast.5)
2019-12-03 23:49:11.662833: E   19849 tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %add.3444 = bf16[512,2048,7,7]{3,2,1,0} add(bf16[512,2048,7,7]{3,2,1,0} %get-tuple-element.3434, bf16[512,2048,7,7]{3,2,1,0} %multiply.3386)
2019-12-03 23:49:11.662844: E   19849 tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %constant.3445 = bf16[] constant(0)
2019-12-03 23:49:11.662855: E   19849 tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.3446 = bf16[512,2048,7,7]{3,2,1,0} broadcast(bf16[] %constant.3445), dimensions={}
2019-12-03 23:49:11.662866: E   19849 tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %maximum.3447 = bf16[512,2048,7,7]{3,2,1,0} maximum(bf16[512,2048,7,7]{3,2,1,0} %add.3444, bf16[512,2048,7,7]{3,2,1,0} %broadcast.3446)
2019-12-03 23:49:11.662877: E   19849 tensorflow/compiler/xla/xla_client/xla_util.cc:72]   ROOT %tuple.3448 = (bf16[512,2048,7,7]{3,2,1,0}) tuple(bf16[512,2048,7,7]{3,2,1,0} %maximum.3447)
2019-12-03 23:49:11.662888: E   19849 tensorflow/compiler/xla/xla_client/xla_util.cc:72] }
2019-12-03 23:49:11.662900: E   19849 tensorflow/compiler/xla/xla_client/xla_util.cc:72]
2019-12-03 23:49:11.662910: E   19849 tensorflow/compiler/xla/xla_client/xla_util.cc:72]
2019-12-03 23:49:11.662921: E   19849 tensorflow/compiler/xla/xla_client/xla_util.cc:72] StackTrace:
2019-12-03 23:49:11.662932: E   19849 tensorflow/compiler/xla/xla_client/xla_util.cc:72] *** Begin stack trace ***
2019-12-03 23:49:11.662942: E   19849 tensorflow/compiler/xla/xla_client/xla_util.cc:72]        tensorflow::CurrentStackTrace[abi:cxx11]()
2019-12-03 23:49:11.662952: E   19849 tensorflow/compiler/xla/xla_client/xla_util.cc:72]        xla::util::ReportComputationError(tensorflow::Status const&, absl::Span<xla::XlaComputation const* const>)
2019-12-03 23:49:11.662964: E   19849 tensorflow/compiler/xla/xla_client/xla_util.cc:72]        xla::util::ShapeHash(xla::Shape const&)
2019-12-03 23:49:11.662971: E   19849 tensorflow/compiler/xla/xla_client/xla_util.cc:72]        xla::XrtComputationClient::ExecuteComputation(xla::ComputationClient::Computation const&, absl::Span<std::shared_ptr<xla::ComputationClient::Data> const>, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&, xla::ComputationClient::ExecuteComputationOptions const&)
2019-12-03 23:49:11.662985: E   19849 tensorflow/compiler/xla/xla_client/xla_util.cc:72]
2019-12-03 23:49:11.662995: E   19849 tensorflow/compiler/xla/xla_client/xla_util.cc:72]
2019-12-03 23:49:11.663005: E   19849 tensorflow/compiler/xla/xla_client/xla_util.cc:72]
2019-12-03 23:49:11.663016: E   19849 tensorflow/compiler/xla/xla_client/xla_util.cc:72]
2019-12-03 23:49:11.663026: E   19849 tensorflow/compiler/xla/xla_client/xla_util.cc:72]
2019-12-03 23:49:11.663036: E   19849 tensorflow/compiler/xla/xla_client/xla_util.cc:72]        clone
2019-12-03 23:49:11.663047: E   19849 tensorflow/compiler/xla/xla_client/xla_util.cc:72] *** End stack trace ***
2019-12-03 23:49:11.663058: E   19849 tensorflow/compiler/xla/xla_client/xla_util.cc:72]
2019-12-03 23:49:11.663069: E   19849 tensorflow/compiler/xla/xla_client/xla_util.cc:72] Status: Unavailable: From /job:tpu_worker/replica:0/task:0:
2019-12-03 23:49:11.663079: E   19849 tensorflow/compiler/xla/xla_client/xla_util.cc:72] Driver not open.
2019-12-03 23:49:11.663090: E   19849 tensorflow/compiler/xla/xla_client/xla_util.cc:72]         [[{{node XRTExecute}}]]Traceback (most recent call last):
  File "tpu_distributed_fastai.py", line 107, in <module>
    xmp.spawn(train_loop,args=())
  File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/torch_xla/distributed/xla_multiprocessing.py", line 154, in spawn
    _start_fn, args=(fn, args), nprocs=nprocs, join=join, daemon=daemon)
  File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/torch/multiprocessing/spawn.py", line 171, in spawn
    while not spawn_context.join():
  File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/torch/multiprocessing/spawn.py", line 107, in join    (error_index, name)
Exception: process 0 terminated with signal SIGSEGV

This is just the end, it keeps printing what looks like something regarding the model parameters. I cannot get that in full, because it is so much text, that the terminal doesn't even store the whole thing.

Is there more information you will need from me?

[dlibenzi]

Maybe you can pipe the output to a file and post that?

[zcain117]

Maybe you can pipe the output to a file and post that?

e.g python mytrain.py 2>&1 | tee myfile.txt

The OOM message is earlier up in the logs

[tmabraham]

@dlibenzi @zcain117 I here's the file. Let me know if there are any issues:
OOMLog.txt

[zcain117]

Doesn't look to me like an obvious TPU OOM. As a sanity check, have you tried with batch_size=1?

[tmabraham]

@zcain117 I ran my code with batch size of 256 yesterday, with bfloat16 enabled and it worked.

[zcain117]

I see. An OOM error usually looks more like this paste below, which is different from what you're seeing. But if the problem is 100% fixed by smaller batch, maybe some other error is happening as the TPU is running out of memory. @dlibenzi might see more clues from your logs but at the moment it seems like 256 batch size is the way to go.

Traceback (most recent call last):
  File "/home/zcain/pytorch/xla/torch_xla_py/data_parallel.py", line 226, in _module_runner
    result.result = loop_fn(module, loader, torch.device(device), context)
  File "xla/test/test_train_imagenet.py", line 136, in train_loop_fn
    xm.optimizer_step(optimizer)
  File "/home/zcain/pytorch/xla/torch_xla_py/xla_model.py", line 336, in optimizer_step
    _mark_step(replication)
  File "/home/zcain/pytorch/xla/torch_xla_py/xla_model.py", line 326, in _mark_step
    torch_xla._XLAC._xla_get_default_device(), devices, wait=False)
RuntimeError: tensorflow/compiler/xla/xla_client/xrt_computation_client.cc:398 : Check failed: mwait.Wait() == ::tensorflow::Status::OK() (Internal: Resource exhausted: From /job:tpu_worker/replica:0/task:0:
Ran out of memory in memory space hbm. Used 10.04G of 8.00G hbm. Exceeded hbm capacity by 2.04G.

Total hbm usage >= 10.04G:
    reserved        528.00M
    program           9.52G
    arguments       unknown size

Output size unknown.

Program hbm requirement 9.52G:
    reserved           4.0K
    global             4.0K
    HLO temp          9.52G (27.9% utilization, 0.0% fragmentation (1.43M))

  Largest program allocations in hbm:

  1. Size: 588.00M
     Operator: op_type="aten::clamp" source_file="[email protected]" source_line=953
     Shape: f32[128,96,112,112]{0,1,3,2:T(8,128)}
     Unpadded size: 588.00M
     XLA label: %fusion.5036 = f32[128,96,112,112]{0,1,3,2:T(8,128)} fusion(f32[96]{0:T(256)} %param_23.226, f32[96]{0:T(256)} %param_22.225, f32[96]{0:T(256)} %fusion.81161, f32[]{:T(256)} %param_8.46, f32[128,96,112,112]{0,1,3,2:T(8,128)} %get-tuple-element.4233, f32[96...
     Allocation type: HLO temp
     ==========================

  2. Size: 220.50M
     Operator: op_type="aten::cat" source_file="[email protected]" source_line=340
     Shape: f32[128,144,56,56]{0,1,3,2:T(8,128)}
     Unpadded size: 220.50M
     XLA label: %concatenate.2 = f32[128,144,56,56]{0,1,3,2:T(8,128)} concatenate(f32[128,1,56,56]{0,1,3,2:T(2,128)} %fusion.45709, f32[128,1,56,56]{0,1,3,2:T(2,128)} %fusion.45706, f32[128,1,56,56]{0,1,3,2:T(2,128)} %fusion.45703, f32[128,1,56,56]{0,1,3,2:T(2,128)} %fusi...
     Allocation type: HLO temp
     ==========================
...

[dlibenzi]

Can you try to grab recent wheels?
I see there log messages that should not be there, in recent code.
Like these:

2019-12-04 00:08:12.172238: I    3128 torch_xla/csrc/aten_xla_type.cpp:93] XLA GIT revision: ed6172751cd919adbb2da9c5dc04f7ceaa741125
2019-12-04 00:08:12.565401: I    3127 torch_xla/csrc/aten_xla_type.cpp:92] PyTorch GIT revision: ca7d0803e9dc98f58c2c809bfef9016adf21d645
2019-12-04 00:08:17.781730: I    3121 tensorflow/compiler/xla/xla_client/xrt_computation_client.cc:217] XRT device (LOCAL) CPU:0 -> /job:tpu_worker/replica:0/task:0/device:XLA_CPU:0
2019-12-04 00:08:17.781776: I    3121 tensorflow/compiler/xla/xla_client/xrt_computation_client.cc:217] XRT device (LOCAL) TPU:0 -> /job:tpu_worker/replica:0/task:0/device:TPU:0
2019-12-04 00:08:17.781785: I    3121 tensorflow/compiler/xla/xla_client/xrt_computation_client.cc:217] XRT device (REMOTE) TPU:1 -> /job:tpu_worker/replica:0/task:0/device:TPU:1

I see no OOM message in your log.

[tmabraham]

@zcain117 Ok thanks. I have never seen this error ever during my time working with Pytorch XLA/TPUs (past 2 months)! But to be fair, it could also be an issue with PyTorch interacting with fastai as I have been using fastai with PyTorch XLA and TPU.

[dlibenzi]

Yeah, ran it myself.
Does not fit at 512:

RuntimeError: Resource exhausted: From /job:tpu_worker/replica:0/task:0:
Ran out of memory in memory space hbm. Used 18.51G of 16.00G hbm. Exceeded hbm capacity by 2.51G.

Total hbm usage >= 18.51G:
    reserved        529.00M 
    program          17.99G 
    arguments       unknown size 

Output size unknown.

Program hbm requirement 17.99G:
    reserved           4.0K
    global             4.0K
    scoped           16.06M
    HLO temp         17.98G (99.9% utilization: Unpadded (17.95G) Padded (17.96G), 0.1% fragmentation (14.02M))

[dlibenzi]

It did fit 384 though:

XLA_USE_BF16=1 python test/test_train_mp_imagenet.py --model=resnet152 --batch_size=384 --num_cores=1 --fake_data

[xla:1/0](0) Loss=6.90625 Rate=5.68 GlobalRate=5.68 Time=Wed Dec  4 01:39:56 2019
[xla:1/0](20) Loss=0.03076 Rate=51.78 GlobalRate=50.18 Time=Wed Dec  4 01:41:29 2019
[xla:1/0](40) Loss=0.00000 Rate=218.65 GlobalRate=85.57 Time=Wed Dec  4 01:41:52 2019
[xla:1/0](60) Loss=0.00000 Rate=285.46 GlobalRate=113.01 Time=Wed Dec  4 01:42:16 2019
[xla:1/0](80) Loss=0.00000 Rate=312.25 GlobalRate=134.92 Time=Wed Dec  4 01:42:39 2019
[xla:1/0](100) Loss=0.00000 Rate=322.97 GlobalRate=152.82 Time=Wed Dec  4 01:43:02 2019
[xla:1/0](120) Loss=0.00000 Rate=327.20 GlobalRate=167.70 Time=Wed Dec  4 01:43:25 2019

[tmabraham]

@dlibenzi To grab the recent wheels I use this?:

$ cd /usr/share/torch-xla-nightly/pytorch/xla
$ . ./scripts/update_nightly_torch_wheels.sh'

Because this didn't work.

[tmabraham]

@dlibenzi umm... I have the same issue.

My torch_xla version is '0.8+ec87d32'

[jysohn23]

@dlibenzi oh wow ok... This is a little disappointing because I can easily train with bs=128 on a Tesla T4 or P100 but only a max of 256 on a TPU. So I can only double the batch size with a TPUv3? That too, these aren't huge images, just 224x224.

Wondering whether 256 was mini or global batch size? If mini-batch is 256 per chip vs when you specify --batch_size=256 on our runners this would be per v3 core (so 256*8 = 2048 global batch size per v3-8). Roughly speaking in terms of $, 1 V100 Chip ~= 2 v3 TPU core.

Just in case it wasn't clear 😄

[tmabraham]

@jysohn23 Yes it's per core. But I am comparing it to 4x Tesla T4's, which can handle bs=128. I haven't tried bs=256 with the multi-GPU setup though. But even then, I don't think this isn't a fair comparison yet because the TPU has 8 cores and I am comparing to 4 GPUs.

[dlibenzi]

Comparing core to core does not make much sense.
If for the price of one Xeon core one could buy 10 ARM64 cores, one should compare those, especially when the workloads are highly parallel like in model training :wink:
Like @jysohn23 explained, on Cloud TPU, for the price of one V100 you can get 2 TPU v3 cores, so this is our compare metric. And that ration is likely to get even better in the future.

[tmabraham]

@dlibenzi By 2 TPU v3 cores, you mean half a TPUv3 chip, right? So 2 V100's is an acceptable comparison to a single TPUv3, right?

[tmabraham]

Also, did I install the correct version? Because I still didn't see any message like the one you showed.

[dlibenzi]

@dlibenzi By 2 TPU v3 cores, you mean half a TPUv3 chip, right? So 2 V100's is an acceptable comparison to a single TPUv3, right?

4 V100 costs like a TPU v3-8, approximately, on Google Cloud.

[zcain117]

A single Cloud TPU device consists of four chips, each of which has two TPU cores.

So one v3-8 TPU has 8 cores total and is comparable to 4 V100 GPUs. This is the comparison we use most frequently since you can't reserve less than 1 TPU (8 cores) in Cloud as far as I know.

More details here on the chip vs. core vs. device situation: https://cloud.google.com/tpu/docs/tpus

[tmabraham]

@zcain117 Thanks for the clarification.

@dlibenzi Do I have an up-to-date version of torch_xla installed?:

My torch_xla version is '0.8+ec87d32'

[dlibenzi]

Yes, looks like recent version.

[tmabraham]

@dlibenzi Hmmm it's interesting why the error messages are different. Maybe there's isn't just a OOM error but some other problems too? I will attach an error log of when I run with the more recent version.

[tmabraham]

@dlibenzi Here is the OOMLog file:

OOMLog.txt

Getting similar error, running batch size=512 again. I know this is too big for the TPU, but I want to know if there is some other problem due to which I am not getting the same OOM error you are getting.

[tmabraham]

Also, just tried batch size=384. It did not work for me.

[dlibenzi]

Can you try single core?
Also, make you you did not hit CTRL^C (or killed) and there are still live processes talking to the HW.

[zcain117]

Good point Davide - @tmabraham sometimes with multiprocessing you can get straggler processes from past runs which you can clean up with pkill -9 python

[tmabraham]

@dlibenzi I am trying it out right now. It is taking longer, since it's not distributed, but it seems to be running fine. I am currently at 12% after 10 minutes.

[tmabraham]

I accidentally used most of my GCP credits so I won't be working on this. I will close this issue for now.

It's unfortunate, but I cannot continue development on a fastai wrapper for PyTorch XLA.