Xla: Einsum slow and consume large memory

Created on 20 Mar 2020 · 9Comments · Source: pytorch/xla

❓ Questions and Help

Hi guys,

I have calculated vectors using einsum and found that it is much slower on TPU than on CPU,
For example the toy code below:

import torch
import torch_xla
import torch_xla.core.xla_model as xm


device = xm.xla_device()
# device='cpu'

eijk = torch.zeros((3, 3, 3), device=device, dtype=torch.float32)
eijk[0, 1, 2] = eijk[1, 2, 0] = eijk[2, 0, 1] = 1
eijk[0, 2, 1] = eijk[2, 1, 0] = eijk[1, 0, 2] = -1

v1 = torch.rand(1000, 3, device=device)
d1 = torch.rand(1000, 3, device=device)
t_val = torch.rand(1000, 3, device=device)

p_val = torch.einsum('uj,vk,ijk->uvi', d1, v1, eijk)    
d_val = torch.einsum('ijk,jk->ij', p_val, v1)
u_val = torch.einsum('ijk,jk->ij',p_val, t_val) * d_val

u_val.to('cpu')

on TPU it takes half a minutes, but on CPU it takes just 1 secone.

And on my program I found that the codes consumes 29 G memory to store a tensor with size (10000, 5896, 3) .

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

Total hbm usage >= 29.32G:
    reserved        529.00M 
    program          28.80G 
    arguments       unknown size 

Output size unknown.

Program hbm requirement 28.80G:
    reserved           4.0K
    global             4.0K
    HLO temp         28.80G (4.5% utilization: Unpadded (1.31G) Padded (28.80G), 0.0% fragmentation (100.0K))

  Largest program allocations in hbm:

  1. Size: 27.92G
     Shape: f32[10000,5856,3,1,1]{2,0,1,4,3:T(8,128)}
     Unpadded size: 670.17M
     Extra memory due to padding: 27.27G (42.7x expansion)
     XLA label: %copy.92 = f32[10000,5856,3,1,1]{2,0,1,4,3:T(8,128)} copy(f32[10000,5856,3,1,1]{1,2,0,3,4:T(4,128)} %bitcast.40)
     Allocation type: HLO temp
     ==========================

  2. Size: 898.44M
     Shape: f32[10000,3,1,5856,1]{3,1,0,4,2:T(4,128)}
     Unpadded size: 670.17M
     Extra memory due to padding: 228.27M (1.3x expansion)
     XLA label: %reshape.291 = f32[10000,3,1,5856,1]{3,1,0,4,2:T(4,128)} reshape(f32[30000,5856]{1,0:T(8,128)} %fusion.47)
     Allocation type: HLO temp
     ==========================

  3. Size: 92.0K
     Shape: f32[5856,3]{0,1:T(4,128)}
     Unpadded size: 68.6K
     Extra memory due to padding: 23.4K (1.3x expansion)
     XLA label: %reshape.292 = f32[5856,3]{0,1:T(4,128)} reshape(f32[5856,1,3]{1,0,2:T(8,128)} %get-tuple-element.8)
     Allocation type: HLO temp
     ==========================

Could anyone help me on this problem? Thank you in advance!

stale

Source

magicknight

All 9 comments

Thanks for reporting. We will take a look!

dlibenzi on 20 Mar 2020

👍1

HLO graph:

https://gist.github.com/dlibenzi/0e183890f27f541d72626fc1043f07de

dlibenzi on 20 Mar 2020

👍1

I dont understand the HLO graph but it seems like something on memroy allocation?

Extra memory due to padding: 27.27G (42.7x expansion)

Is this line telling us that the padding cusumes a lot more memory than the original tensor?

magicknight on 17 Apr 2020

Let me file an XLA bug internally.
This is likely layout assignment taking an odd decision.

dlibenzi on 17 Apr 2020

The graph that I generated from your example does not seem to match the the tensor size which blows up memory.
Do you mind changing the example to be 100% matching what created the error?

dlibenzi on 17 Apr 2020

Hi dlibenzi,
Thank you for reply,
Here is the code which created the error:

import torch
import torch_xla
import torch_xla.core.xla_model as xm

device = xm.xla_device()
# device = 'cpu'
print(device)

tensor_1 = torch.rand([5856, 3, 3], device=device)
tensor_2 = torch.rand(3, device=device)
tensor_3 = torch.rand([40000, 3], device=device)

eijk = torch.zeros((3, 3, 3), device=device, dtype=torch.float32)
eijk[0, 1, 2] = eijk[1, 2, 0] = eijk[2, 0, 1] = 1
eijk[0, 2, 1] = eijk[2, 1, 0] = eijk[1, 0, 2] = -1
eps = torch.tensor(0.000001, device=device, dtype=torch.float32)

tensor_14 = tensor_1[:, 0, :] - tensor_1[:, 1, :]
tensor_15 = tensor_1[:, 0, :] - tensor_1[:, 2, :]
tensor_4 = torch.einsum('ijk,uj,vk->uvi', eijk, tensor_3, tensor_15)

tensor_5 = torch.einsum('ijk,jk->ij', tensor_4, tensor_14)
tensor_6 = torch.abs(tensor_5) > eps
tensor_5 = 1.0 / tensor_5
tensor_5 = tensor_6 * tensor_5
tensor_7 = tensor_2 - tensor_1[:, 0, :]

tensor_8 = torch.einsum('jk,ijk->ij', tensor_7, tensor_4) * tensor_5

tensor_9 = (tensor_8 > 0) & (tensor_8 < 1)
tensor_5 = tensor_9 * tensor_5

tensor_10 = torch.cross(tensor_7, tensor_14)

tensor_11 = torch.einsum('ik,jk->ij', tensor_3,
                     tensor_10) * tensor_5
tensor_12 = (tensor_11 > 0) & (tensor_8+tensor_11 < 1)

tensor_5 = tensor_12 * tensor_5

tensor_13 = torch.einsum(
    'j ,ij->ij', torch.einsum('ik,ik->i', tensor_15, tensor_10), tensor_5)

tensor_13.to('cpu')

Running this code results in the output below:

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

Total hbm usage >= 115.23G:
    reserved         18.00M 
    program         115.21G 
    arguments       unknown size 

Output size unknown.

Program hbm requirement 115.21G:
    reserved           4.0K
    global            36.0K
    HLO temp        115.21G (4.5% utilization: Unpadded (5.24G) Padded (115.21G), 0.0% fragmentation (52.0K))

  Largest program allocations in hbm:

  1. Size: 111.69G
     Shape: f32[40000,5856,3,1,1]{2,0,1,4,3:T(8,128)}
     Unpadded size: 2.62G
     Extra memory due to padding: 109.08G (42.7x expansion)
     XLA label: %copy.53 = f32[40000,5856,3,1,1]{2,0,1,4,3:T(8,128)} copy(f32[40000,5856,3,1,1]{1,2,0,3,4:T(4,128)} %bitcast.16)
     Allocation type: HLO temp
     ==========================

  2. Size: 3.51G
     Shape: f32[40000,3,1,5856,1]{3,1,0,4,2:T(4,128)}
     Unpadded size: 2.62G
     Extra memory due to padding: 913.09M (1.3x expansion)
     XLA label: %reshape.2358 = f32[40000,3,1,5856,1]{3,1,0,4,2:T(4,128)} reshape(f32[120000,5856]{1,0:T(8,128)} %fusion.2524)
     Allocation type: HLO temp
     ==========================

  3. Size: 5.72M
     Shape: f32[5856,1,3]{2,1,0:T(2,128)}
     Unpadded size: 68.6K
     Extra memory due to padding: 5.65M (85.3x expansion)
     XLA label: %fusion.2520 = (f32[5856,1,3]{2,1,0:T(2,128)}, f32[5856,1,3]{2,1,0:T(2,128)}, f32[5856,1,3]{2,1,0:T(2,128)}, f32[5856,1,3]{2,1,0:T(2,128)}) fusion(f32[3]{0:T(256)} %fusion.2537, f32[5856,3,3]{2,1,0:T(4,128)} %reshape.42632), kind=kLoop, calls=%fused_comput...
     Allocation type: HLO temp
     ==========================

  4. Size: 470.0K
     Shape: bf16[30000,4]{0,1}
     Unpadded size: 234.4K
     Extra memory due to padding: 235.6K (2.0x expansion)
     XLA label: %copy.60 = bf16[30000,4]{0,1} copy(bf16[30000,4]{1,0:T(8,128)(2,1)} %fusion.2147)
     Allocation type: HLO temp
     ==========================

  5. Size: 92.0K
     Shape: f32[5856,3]{0,1:T(4,128)}
     Unpadded size: 68.6K
     Extra memory due to padding: 23.4K (1.3x expansion)
     XLA label: %fusion.2527 = (f32[5856,3]{0,1:T(4,128)}, f32[5856,3]{0,1:T(4,128)}) fusion(f32[5856,1,3]{0,2,1:T(4,128)} %copy.47, f32[5856,1,3]{0,2,1:T(4,128)} %copy.46, f32[3]{0:T(256)} %fusion.2537), kind=kLoop, calls=%fused_computation.2525
     Allocation type: HLO temp
     ==========================

magicknight on 18 Apr 2020

Filed a bug internally.
I ended up using a smaller size version.

import torch
import torch_xla
import torch_xla.core.xla_model as xm

device = xm.xla_device()

tensor_1 = torch.rand([5856, 3, 3], device=device)
tensor_2 = torch.rand(3, device=device)
tensor_3 = torch.rand([15000, 3], device=device)

eijk = torch.zeros((3, 3, 3), device=device, dtype=torch.float32)
eijk[0, 1, 2] = eijk[1, 2, 0] = eijk[2, 0, 1] = 1
eijk[0, 2, 1] = eijk[2, 1, 0] = eijk[1, 0, 2] = -1
eps = torch.tensor(0.000001, device=device, dtype=torch.float32)

tensor_14 = tensor_1[:, 0, :] - tensor_1[:, 1, :]
tensor_15 = tensor_1[:, 0, :] - tensor_1[:, 2, :]
tensor_4 = torch.einsum('ijk,uj,vk->uvi', eijk, tensor_3, tensor_15)

tensor_5 = torch.einsum('ijk,jk->ij', tensor_4, tensor_14)
tensor_6 = torch.abs(tensor_5) > eps
tensor_5 = 1.0 / tensor_5
tensor_5 = tensor_6 * tensor_5
tensor_7 = tensor_2 - tensor_1[:, 0, :]
tensor_8 = torch.einsum('jk,ijk->ij', tensor_7, tensor_4) * tensor_5

tensor_9 = (tensor_8 > 0) & (tensor_8 < 1)
tensor_5 = tensor_9 * tensor_5
tensor_10 = torch.cross(tensor_7, tensor_14)
tensor_11 = torch.einsum('ik,jk->ij', tensor_3,
                     tensor_10) * tensor_5
tensor_12 = (tensor_11 > 0) & (tensor_8+tensor_11 < 1)
tensor_5 = tensor_12 * tensor_5
tensor_13 = torch.einsum(
    'j ,ij->ij', torch.einsum('ik,ik->i', tensor_15, tensor_10), tensor_5)

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

print(tensor_13.cpu()[0][0])

dlibenzi on 19 Apr 2020

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.