Xla: Check failed: program_shape.parameters_size() == po_data->parameters_data.size()
In my code, I have the following error:
RuntimeError: torch_xla/csrc/tensor.cpp:1534 : Check failed: program_shape.parameters_size() == po_data->parameters_data.size() (1276 vs. 1274)
I am curious what this error refers to? The stack trace is not very helpful.
I have attached the full stack trace for reference. Please ignore the beginning part which is part of my custom code.
paramsizeerror.txt
tmabraham
All 32 comments
Hi @tmabraham This means there is a parameter size mismatch between the xlaComputation we built(program_shape) and the data we passed to the Compile function(po_data means post order data). Does this fail consistently? Could you share a small repo so we can better debug this issue?
JackCaoG
on 3 Aug 2020
@JackCaoG It fails consistently. Sometimes I get just the SIGSEGV fault error‚ and sometimes I get the full stack trace. I tried with single core (n_procs=1), which different batch sizes, and with bfloat16. Each time the error arose a couple of batches into training.
Does this error mean there is something wrong with the data passed to the model?
tmabraham
on 3 Aug 2020
@JackCaoG BTW the code is here:
https://github.com/tmabraham/fastai_tpu/blob/master/fastai_v2/tpu_distributed_fastai2.py
It's for the fastai library, so let me know if you have any questions about how this code works.
tmabraham
on 3 Aug 2020
@tmabraham What happened was when the code performs a SyncLiveTensorsGraph, we picked the tensor that is alive and compile a graph out of it. In high level, you can think number of the input(parameter) of the graph should match the number of tensors, but in your case parameters_data(number of tensors treated as input to graph) is 2 less than the compiled graph required input size. We do have some special handling of the tensor->parameter, so I am wondering if that's why the mismatch happens.
Thanks for sharing the code, do you have some instruction of how to setup the environment(I assume I need to install some fastai libaray)? I will try to look it soon.
JackCaoG
on 3 Aug 2020
@JackCaoG As far as I can tell, just doing pip install fastai2 in the torch-xla-nightly conda environment will work. fastai2 is compatible with torch nightly.
tmabraham
on 3 Aug 2020
@JackCaoG Any update on this? Or is there any confusion about how the code works (given that it uses fastai2)?
tmabraham
on 6 Aug 2020
@tmabraham Hi I was able to reproduce the failure locally, will update this thread when I have more clues.
JackCaoG
on 6 Aug 2020
@JackCaoG Okay sounds good. Please keep me updated and let me know if there's anything I can do to help you debug this faster.
tmabraham
on 6 Aug 2020
The problem is that number of device data is less than the number of the parameter required by the graph.. which is weird. I dump the device data and graph parameter graph and run a simple parser. I saw
f32[64] match
f32[64,3,7,7] match
mismatch found
f32[128,3,224,224] 3 4
f32[64,64,1,1] match
f32[64,64,3,3] match
f32[256] match
f32[256,64,1,1] match
f32[64,256,1,1] match
f32[128] match
f32[128,256,1,1] match
f32[128,128,3,3] match
f32[512] match
f32[512,128,1,1] match
f32[512,256,1,1] match
f32[128,512,1,1] match
f32[256,512,1,1] match
f32[256,256,3,3] match
f32[1024] match
f32[1024,256,1,1] match
f32[1024,512,1,1] match
f32[256,1024,1,1] match
f32[512,1024,1,1] match
f32[512,512,3,3] match
f32[2048] match
f32[2048,512,1,1] match
f32[2048,1024,1,1] match
f32[512,2048,1,1] match
f32[4096] match
f32[128,2048,1,1] match
f32[512,4096] match
f32[] match
s64[] match
f32[101,512] match
s64[128] match
f32[128,2048,7,7] match
This means that graph requires 4 f32[128,3,224,224] but there is only 3 is presented in the device data tensor. I noticed that f32[128,3,224,224] is the last parameter(I run the code couple times and I saw it failed at different places but this pattern remains, graph requires one more 4d tensor at the end). I will look at xla_builder into detail to see where is this last parameter coming from.
device data ops
....
s64[] xla::device_data, [email protected]:113, device=TPU:0
s64[] xla::device_data, [email protected]:113, device=TPU:0
s64[] xla::device_data, [email protected]:113, device=TPU:0
f32[128,3,224,224]{3,2,1,0} xla::device_data, location=_rebuild_xla_tensor@_utils.py:177, device=TPU:0
s64[128]{0} xla::device_data, location=_rebuild_xla_tensor@_utils.py:177, device=TPU:0
f32[128,3,224,224]{3,2,1,0} xla::device_data, location=_rebuild_xla_tensor@_utils.py:177, device=TPU:0
s64[128]{0} xla::device_data, location=_rebuild_xla_tensor@_utils.py:177, device=TPU:0
graph parameters
(p0: f32[64], p1: f32[64], p2: f32[64], p3: f32[64], p4: f32[64,3,7,7], p5: f32[128,3,224,224], p6: f32[64],......
p1264: s64[], p1265: s64[], p1266: f32[128,3,224,224], p1267: s64[128], p1268: f32[128,3,224,224], p1269: s64[128],
p1270: f32[128,3,224,224])
@JackCaoG So f32[128,3,224,224] would be the batch being passed into the model. When you say that 4 f32[128,3,224,224] is required, this is probably because its 4 iterations in, right? So somehow the batches are being lost somewhere? Also, if a batch from each iteration is being passed into the model then why does the model need 4 of them at one time?
Sorry if these questions seem a little naive, I admittedly don't know much about what's going on under the hood of the TPU 🤐 .
tmabraham
on 7 Aug 2020
I can't comment much about the real problem as I am still debugging it. I don't think batches is lost during the computation. XLA is a graph based compiler, every pytorch op will be transform to one or more ops defined in here. You can think of this as a tree structure, node add would take two parameter and have one output, this output might be the parameter(input) or another node. We will have a graph that takes multiple inputs.
The current problem is that, given a set of pytorch tensors, we form a graph. However the number of device_data(parameter) in the pytorch tensors is mismatch with the number of parameter required for the graph.
JackCaoG
on 8 Aug 2020
@JackCaoG OK i think that makes more sense. I guess the question is where in the model it requires 4 f32[128,3,224,224], right? I find it odd though, because f32[128,3,224,224] is actually the shape of the batch...
tmabraham
on 8 Aug 2020
It turned out that the last op (the ops that requires additional parameter that was missing) was a as_strided node.
node is f32[128,3,224,224]{3,2,1,0} aten::as_strided, location=mark_step@xla_model.py:719, size=(128, 3, 224, 224), stride=(150528, 50176, 224, 1), storage_offset=0
GetOutputOp, output is not lowered yet number of nodes needs to be lowered 3
now lowering f32[128,3,224,224]{3,2,1,0} xla::device_data, location=_rebuild_xla_tensor@_utils.py:177, device=TPU:0
Device data with handle 4438875643412039 is not in the map, adding it
GetParameter after push_back parameter size= 1271
GetParameter after push_back parameter size= 1271
now lowering f32[128,3,224,224]{3,2,1,0} xla::as_strided_view_update, location=mark_step@xla_model.py:719, size=(128, 3, 224, 224), stride=(150528, 50176, 224, 1), storage_offset=0
output f32[128,3,224,224]{3,2,1,0} aten::expand, location=_rebuild_xla_tensor@_utils.py:177, size=(128, 3, 224, 224)already lowered
output f32[128,3,224,224]{3,2,1,0} xla::device_data, location=_rebuild_xla_tensor@_utils.py:177, device=TPU:0already lowered
now lowering f32[128,3,224,224]{3,2,1,0} aten::as_strided, location=mark_step@xla_model.py:719, size=(128, 3, 224, 224), stride=(150528, 50176, 224, 1), storage_offset=0
output f32[128,3,224,224]{3,2,1,0} xla::as_strided_view_update, location=mark_step@xla_model.py:719, size=(128, 3, 224, 224), stride=(150528, 50176, 224, 1), storage_offset=0already lowered
One of its operand f32[128,3,224,224]{3,2,1,0} xla::device_data, location=_rebuild_xla_tensor@_utils.py:177 is missing from the parameter.. This _rebuild_xla_tensor seems suspicious. It is from https://github.com/pytorch/pytorch/blob/master/torch/_utils.py#L176. There might be something special about this tensor. @ailzhang Do you have any ideas?
JackCaoG
on 10 Aug 2020
I added a pdb in def _rebuild_xla_tensor(data, dtype, device, requires_grad): and found that _rebuild_xla_tensor is called after we collect all of the parameter tensors and build graph. There is either a (indirect) misuse of this function or we have a bug in this interaction. I will check with Ailing.
JackCaoG
on 10 Aug 2020
@tmabraham It seems like fastaiv2 patch pytorch's _reduce_ex_ function(which will call _rebuild_xla_tensor) and we have xla specific logic in that function. After deleting fastai patch of _reduce_ex_ in fastai2/torch_core.py and __reduce_ex__ patch in tpu_distributed_fastai2.py, code runs fine.
JackCaoG
on 10 Aug 2020
@JackCaoG Fastai2 has their own type of tensor (TensorBase) that subclasses the PyTorch Tensor. It implements almost the same __reduce_ex__ behavior as PyTorch. I simply added the same XLA-specific functionality to __reduce_ex__ from the original PyTorch Tensor code to TensorBase. Is there something wrong with this implementation? Because it should be essentially the same as PyTorch's __reduce_ex__.
tmabraham
on 11 Aug 2020
Maybe it is
if type(self) is not Tensor and has_torch_function(relevant_args):
return handle_torch_function(Tensor.__reduce_ex__, relevant_args, self, proto)
I am not sure if that's the cause. If you can try to use pytorch native _reduce_ex_ and play with it for a while you might have a better idea.
JackCaoG
on 11 Aug 2020
@JackCaoG But this shouldn't be called because we patch TensorBase with our __reduce_ex__ method that is essentially the same as PyTorch's method. Is this not recommended? I wouldn't think it would be a problem.
Basically, I am not seeing how the situation is much different in regular PyTorch. It would still call _rebuild_xla_tensor after the all of the parameter tensors are collected and the graph is built, right?
tmabraham
on 11 Aug 2020
I removed all patches and run it for a bit longer and saw the same problem, so that patch is not the problem..sigh.... I will take another look tomorrow, maybe it is related to our copy(calls _rebuild_xla_tensor) performs a deep copy while pytorch copy is a shallow copy.
JackCaoG
on 11 Aug 2020
@JackCaoG Oh wow okay sorry it's such a difficult bug 😅 I hope the bug isn't something dumb on my end lol. Let me know if you discover anything new!
tmabraham
on 11 Aug 2020
Oh, don't worry 😄 , this kind of bug usually takes couple back and forth. Thanks for your patience through!
JackCaoG
on 11 Aug 2020
After adding more debugging statement I saw events happens in the following order
- a lot of
xm.send_cpu_data_to_device->....-> _rebuild_xla_tensor is called, mixed with asyncSyncLiveTensorsGraph. - SyncLiveTensorsGraph ->.. SyncTensorsGraphInternal -> RunPostOrder(*tensors, coll.indices) to collect the device data(parameter) and live tensor nodes in post_order.
- While we computing the post order data,
GetIrValueForTensoris called and sending data to the device. Handle get back is3143821633530146 - SyncTensorsGraphInternal -> Compile() is called, and last node
f32[128,3,224,224]{3,2,1,0} aten::as_stridedtakes 3 operands. One of them is device_data (with metadatalocation=_rebuild_xla_tensor@_utils.py:177) with handle3143821633530146which is missing from thepo_data->parameters_data. checkXLA_CHECK_EQ(program_shape.parameters_size(),po_data->parameters_data.size());failed.
I think what happened in step 4 make sense as while we collecting the device_data in live tensors, device_data with 3143821633530146 has not been sent to the device yet. The part that I am not sure about is why send_to_device is not done before we start syncing the tensors. I would image there is a barrier that we need to wait for all operand's data being sent to the device before we start syncing the tensor.
@davidel Am I missing something here? This feel like a race condition but it is a such general case. If we don't do this properly we should see this error way more frequent. I removed fastaiv2's patch for __reduce_ex__ but there might be other patches/overwrites I am not aware of. If you could point me to where to perform the barrier(I imaged we should handle this case somewhere) I could check if that part of code is being run properly.
The stack trace for _rebuild_xla_tensor is
File "/home/jackcao/anaconda3/envs/pytorch/lib/python3.7/threading.py", line 890, in _bootstrap
self._bootstrap_inner()
File "/home/jackcao/anaconda3/envs/pytorch/lib/python3.7/threading.py", line 926, in _bootstrap_inner
self.run()
File "/home/jackcao/anaconda3/envs/pytorch/lib/python3.7/threading.py", line 870, in run
self._target(*self._args, **self._kwargs)
File "/home/jackcao/anaconda3/envs/pytorch/lib/python3.7/site-packages/torch_xla-1.6-py3.7-linux-x86_64.egg/torch_xla/distributed/parallel_loader.py", line 172, in _worker
batch = xm.send_cpu_data_to_device(batch, device)
File "/home/jackcao/anaconda3/envs/pytorch/lib/python3.7/site-packages/torch_xla-1.6-py3.7-linux-x86_64.egg/torch_xla/core/xla_model.py", line 847, in send_cpu_data_to_device
return ToXlaTensorArena(convert_fn, select_fn).transform(data)
File "/home/jackcao/anaconda3/envs/pytorch/lib/python3.7/site-packages/torch_xla-1.6-py3.7-linux-x86_64.egg/torch_xla/core/xla_model.py", line 411, in transform
return self._replace_tensors(inputs)
File "/home/jackcao/anaconda3/envs/pytorch/lib/python3.7/site-packages/torch_xla-1.6-py3.7-linux-x86_64.egg/torch_xla/core/xla_model.py", line 405, in _replace_tensors
convert_fn)
File "/home/jackcao/anaconda3/envs/pytorch/lib/python3.7/site-packages/torch_xla-1.6-py3.7-linux-x86_64.egg/torch_xla/utils/utils.py", line 214, in for_each_instance_rewrite
return _for_each_instance_rewrite(value, select_fn, fn, rwmap)
File "/home/jackcao/anaconda3/envs/pytorch/lib/python3.7/site-packages/torch_xla-1.6-py3.7-linux-x86_64.egg/torch_xla/utils/utils.py", line 194, in _for_each_instance_rewrite
result.append(_for_each_instance_rewrite(x, select_fn, fn, rwmap))
File "/home/jackcao/anaconda3/envs/pytorch/lib/python3.7/site-packages/torch_xla-1.6-py3.7-linux-x86_64.egg/torch_xla/utils/utils.py", line 194, in _for_each_instance_rewrite
result.append(_for_each_instance_rewrite(x, select_fn, fn, rwmap))
File "/home/jackcao/anaconda3/envs/pytorch/lib/python3.7/site-packages/torch_xla-1.6-py3.7-linux-x86_64.egg/torch_xla/utils/utils.py", line 202, in _for_each_instance_rewrite
result = copy.copy(value)
File "/home/jackcao/anaconda3/envs/pytorch/lib/python3.7/copy.py", line 106, in copy
return _reconstruct(x, None, *rv)
File "/home/jackcao/anaconda3/envs/pytorch/lib/python3.7/copy.py", line 274, in _reconstruct
y = func(*args)
File "/home/jackcao/anaconda3/envs/pytorch/lib/python3.7/site-packages/torch/_utils.py", line 178, in _rebuild_xla_tensor
From a quick look (I'm in Rest Mode 😄 ) it seems the loader wrapped by the parallel loader, is feeding XLA tensors, not PT/CPU ones, which is not allowed.
The error is likely that the XLA tensors being fed, are actually live tensors for the current step. So pretty much it's a Kaboom right there.
davidel
on 12 Aug 2020
so Davide's theory is right, the loader is trying to fed the XLA tensor instead of CPU tensor. The problem is that there is a dl = dl.to(self.device) in https://github.com/tmabraham/fastai_tpu/blob/master/fastai_v2/tpu_distributed_fastai2.py#L65 which should be removed. @tmabraham I found https://github.com/pytorch/xla/blob/master/test/test_train_mp_mnist.py to be a great example of how to do a simple mp.
Removing that line gave me a new error of RuntimeError: Input type (torch.FloatTensor) and weight type (XLAFloatType) should be the same through
JackCaoG
on 13 Aug 2020
@tmabraham I am not sure how distributed_dl works and why you passed that to the ParallelLoader. ParallelLoader already performs multi-processing data parallelism. Could you try with a basic data loader and removed the dl.to(self.device) I mentioned above?
JackCaoG
on 13 Aug 2020
@JackCaoG DistributedDL is meant to replace PyTorch's DistributedSampler+DataLoader pattern for multi-GPU training. Given that this pattern is also applicable here, I assume it will work here as well. I can try without it but then the DataLoader would go over every single batch in the dataset and not properly divide the sampling over the cores.
tmabraham
on 13 Aug 2020
I think our parallel loader already does that? https://github.com/pytorch/xla/blob/54f3e16a72be3f5e706aad0022df37e3a69e91c4/torch_xla/distributed/parallel_loader.py#L50
JackCaoG
on 13 Aug 2020
@JackCaoG I don't think ParallelLoader implements distributed sampling, right? You still have to define a DistributedSampler, create a regular PyTorch DataLoader with your data and DistributedSampler (before passing to ParallelLoader) right? Basically, DistributedDL does all of that. My understanding is that it would replace this:
https://github.com/pytorch/xla/blob/54f3e16a72be3f5e706aad0022df37e3a69e91c4/test/test_train_mp_mnist.py#L88-L99
Also, this is an embarrassingly simple question that I should know the answer to, but where in the code does ParallelLoader put the batch on the device?
tmabraham
on 13 Aug 2020
I see, you are right, sorry for the misinformation. 😄 I have to admit that I am more familiar with the compiler side of things.
For your question, I think it is
https://github.com/pytorch/xla/blob/54f3e16a72be3f5e706aad0022df37e3a69e91c4/torch_xla/distributed/parallel_loader.py#L172
The batch(to send to device) should be CPU tensors, but if you do dl = dl.to(self.device) then batch(before sending) actually become XLATensor, which will cause the parameter mismatch error we saw above.
JackCaoG
on 13 Aug 2020
@JackCaoG No worries!
So I fixed the error (it at least runs a single epoch) by removing dl.to(self.device) but adding to the callback self.learn.xb = [xb_item.to(self.device) for xb_item in self.xb] and self.learn.xb = [yb_item.to(self.device) for yb_item in self.yb] in the step right before the batch gets passed into the model.
This leads me to another question, most code I see using PyTorch XLA (and code I have written myself as well) put the batch returned by ParallelLoader onto the TPU device before passing into the XLA model. The example code you linked or the API_GUIDE.md (and PyTorch XLA documentation) does not do this. However, it was only with this fix (added to the callback) that my code would run. So which is the correct way of doing this?
tmabraham
on 13 Aug 2020
The ParallelLoader (and its PerDeviceLoader) should have CPU tensors fed into, and return XLA tensors.
You should not need to perform any to(device) on the data returned by the ParallelLoader.
Unless either we have a bug in our data structure walk-and-convert code, or the model does something "strange" in its logic.
davidel
on 13 Aug 2020
@tmabraham I will close this issue for now, feel free to reopen if you have any other questions.
JackCaoG
on 20 Aug 2020
Related issues
ibeltagy
·
5Comments
nosound2
·
6Comments
butchland
·
7Comments
CalebEverett
·
4Comments
Arjuna197
·
5Comments