Xla: `tensor.gather` slows down `loss.backward()`

Created on 3 Aug 2019  路  50Comments  路  Source: pytorch/xla

Using tensor.gather in the network slows down loss.backward(). Here's a code snippet to reproduce the issue:

vocab_size = 30522
import torch
import torch_xla_py.xla_model as xm
xla_device = xm.xla_device()

ids = torch.LongTensor(32, 20, vocab_size).random_(0, 127)
labels = torch.LongTensor(32, 20).random_(0, vocab_size - 1)
preds = torch.rand(32, 128, vocab_size, requires_grad=True)

ids = ids.to(xla_device)
labels = labels.to(xla_device)
preds = preds.to(xla_device)

gathered_preds = preds.gather(1, ids)
loss = torch.nn.functional.cross_entropy(gathered_preds.view(-1, vocab_size), labels.view(-1))

loss.backward()  # takes > 7 minutes

Notes:
1- Even though the second batch should be processed faster than the first (compiled graph is cached), it remains slow for all batches
2- slight variation on the vocab_size (e.g, 30000 vs 30522) has a huge impact on the processing time

All 50 comments

Below is the metric report after processing 5 batches.

1- It takes 7 minutes to compile. Is it normal for compilation time to be that long?
2- Execution time is 7 seconds/batch. This seems very slow. Is that expected?

Metric: CompileTime
  TotalSamples: 1     <<<<<<<<<<<   one graph
  Counter: 07m10s593ms51.293us    <<<<<<<<<<<    7 minutes to compile
  Percentiles: 1%=07m10s593ms51.293us; 5%=07m10s593ms51.293us; 10%=07m10s593ms51.293us; 20%=07m10s593ms51.293us; 50%$
07m10s593ms51.293us; 80%=07m10s593ms51.293us; 90%=07m10s593ms51.293us; 95%=07m10s593ms51.293us; 99%=07m10s593ms51.29$
us
Metric: ExecuteTime
  TotalSamples: 5   <<<<<<<<<<<   5 batches
  Counter: 35s656ms984.922us    <<<<<<<<<<<   7 seconds/batch 
  ValueRate: 069ms497.156us / second
  Rate: 0.0100267 / second
  Percentiles: 1%=04s107ms782.846us; 5%=04s107ms782.846us; 10%=04s107ms782.846us; 20%=04s107ms404.151us; 50%=04s107m$
409.359us; 80%=18s110ms13.488us; 90%=18s110ms13.488us; 95%=18s110ms13.488us; 99%=18s110ms13.488us
Metric: InboundData
  TotalSamples: 5
  Counter: 2.32GB
  ValueRate: 4.77MB / second
  Rate: 0.0100343 / second
  Percentiles: 1%=475.34MB; 5%=475.34MB; 10%=475.34MB; 20%=475.34MB; 50%=475.34MB; 80%=475.34MB; 90%=475.34MB; 95%=4$
5.34MB; 99%=475.34MB
Metric: OutboundData
  TotalSamples: 15
  Counter: 3.05GB
  ValueRate: 3.29MB / second
  Rate: 0.0158021 / second
  Percentiles: 1%=5.00KB; 5%=5.00KB; 10%=5.00KB; 20%=5.00KB; 50%=148.54MB; 80%=475.34MB; 90%=475.34MB; 95%=475.34MB; 99%=475.34MB
Metric: ReleaseDataHandlesTime
  TotalSamples: 9
  Counter: 099ms2.896us
  ValueRate: 198.695us / second
  Rate: 0.0180627 / second
  Percentiles: 1%=001ms268.288us; 5%=001ms268.288us; 10%=001ms268.288us; 20%=001ms282.276us; 50%=002ms629.583us; 80%$032ms232.240us; 90%=035ms434.710us; 95%=035ms434.710us; 99%=035ms434.710us
Metric: TransferFromServerTime
  TotalSamples: 5
  Counter: 14s388ms477.041us
  ValueRate: 029ms875.516us / second
  Rate: 0.0100343 / second
  Percentiles: 1%=03s640ms522.230us; 5%=03s640ms522.230us; 10%=03s640ms522.230us; 20%=03s640ms635.544us; 50%=03s919m$777.648us; 80%=03s177ms649.724us; 90%=03s177ms649.724us; 95%=03s177ms649.724us; 99%=03s177ms649.724us
Metric: TransferToServerTime
  TotalSamples: 15
  Counter: 20s941ms645.680us
  ValueRate: 021ms45.316us / second
  Rate: 0.015831 / second
  Percentiles: 1%=006ms652.113us; 5%=006ms652.113us; 10%=007ms90.362us; 20%=008ms972.505us; 50%=02s968ms405.952us; 8$%=02s038ms605.435us; 90%=02s087ms118.331us; 95%=02s132ms626.147us; 99%=02s132ms626.147us
Counter: CachedSyncTensors
  Value: 4
Counter: CreateCompileHandles
  Value: 1
Counter: CreateDataHandles
  Value: 20
Counter: CreateXlaTensor
  Value: 65
Counter: DestroyDataHandles
  Value: 17
Counter: DestroyXlaTensor
  Value: 60
Counter: ReleaseDataHandles
  Value: 17
Counter: UncachedSyncTensors
  Value: 1
Counter: XRTAllocateFromTensor_Empty
  Value: 3
Counter: XrtCompile_Empty
  Value: 256
Counter: XrtExecuteChained_Empty
  Value: 256
Counter: XrtExecute_Empty
  Value: 256
Counter: XrtRead_Empty
  Value: 256
Counter: XrtReleaseAllocationHandle_Empty
  Value: 256
Counter: XrtReleaseCompileHandle_Empty
  Value: 256
Counter: XrtSessionCount
  Value: 3
Counter: XrtSubTuple_Empty
  Value: 256

Thanks for reporting. Transfer to/from server also seem to take long.
What are the differences in ExecuteTime between the two vocab sizes you mentioned?
It is also suggested to let run for 10..20 steps, in order to collect more info.

Certain values for vocab_size have long CompileTime (e.g. 30522) while other vocab_size values have short CompileTime (e.g. 30000), but ExecuteTime is the same for all.

Also, some values of vocab_size consistently result into Segmentation fault (e.g. 29999)

Below are the details for the 3 configurations after 15 steps.

vocab_size = 30522

Metric: CompileTime
  TotalSamples: 1
  Counter: 07m02s535ms633.740us
  Percentiles: 1%=07m02s535ms633.740us; 5%=07m02s535ms633.740us; 10%=07m02s535ms633.740us; 20%=07m02s535ms633.740us; 50%=07m02s535ms633.740us; 80%=07m02s535ms633.740us; 90%=07m02s535ms633.740us; 95%=07m02s535ms633.740us; 99%=07m02s535ms633.740us
Metric: ExecuteTime
  TotalSamples: 15
  Counter: 01m16s327ms657.379us
  ValueRate: 408ms415.405us / second
  Rate: 0.0802633 / second
  Percentiles: 1%=04s121ms990.205us; 5%=04s121ms990.205us; 10%=04s121ms38.381us; 20%=04s121ms73.344us; 50%=04s121ms147.301us; 80%=04s122ms243.487us; 90%=04s256ms950.065us; 95%=18s494ms141.018us; 99%=18s494ms141.018us
Metric: InboundData
  TotalSamples: 15
  Counter: 6.99GB
  ValueRate: 38.49MB / second
  Rate: 0.0807127 / second
  Percentiles: 1%=476.91MB; 5%=476.91MB; 10%=476.91MB; 20%=476.91MB; 50%=476.91MB; 80%=476.91MB; 90%=476.91MB; 95%=476.91MB; 99%=476.91MB
Metric: OutboundData
  TotalSamples: 45
  Counter: 9.17GB
  ValueRate: 14.98MB / second
  Rate: 0.0717884 / second
  Percentiles: 1%=5.00KB; 5%=5.00KB; 10%=5.00KB; 20%=5.00KB; 50%=149.03MB; 80%=476.91MB; 90%=476.91MB; 95%=476.91MB; 99%=476.91MB
Metric: ReleaseDataHandlesTime
  TotalSamples: 29
  Counter: 153ms502.748us
  ValueRate: 821.501us / second
  Rate: 0.156217 / second
  Percentiles: 1%=001ms258.042us; 5%=001ms281.702us; 10%=001ms300.678us; 20%=001ms356.017us; 50%=001ms466.428us; 80%=002ms706.401us; 90%=024ms95.045us; 95%=044ms193.398us; 99%=045ms408.000us
Metric: TransferFromServerTime
  TotalSamples: 15
  Counter: 46s087ms991.782us
  ValueRate: 248ms987.135us / second
  Rate: 0.0807127 / second
  Percentiles: 1%=03s807ms664.792us; 5%=03s807ms664.792us; 10%=03s814ms340.640us; 20%=03s847ms532.730us; 50%=03s975ms932.595us; 80%=03s240ms545.570us; 90%=04s513ms119.326us; 95%=04s847ms265.730us; 99%=04s847ms265.730us
Metric: TransferToServerTime
  TotalSamples: 45
  Counter: 01m01s937ms957.225us
  ValueRate: 098ms538.582us / second
  Rate: 0.0720291 / second
  Percentiles: 1%=003ms691.167us; 5%=003ms975.521us; 10%=004ms110.244us; 20%=006ms563.218us; 50%=02s936ms546.091us; 80%=02s061ms873.077us; 90%=02s142ms777.376us; 95%=02s191ms674.152us; 99%=03s573ms115.636us
Counter: CachedSyncTensors
  Value: 14
Counter: CreateCompileHandles
  Value: 1
Counter: CreateDataHandles
  Value: 60
Counter: CreateXlaTensor
  Value: 195
Counter: DestroyDataHandles
  Value: 57
Counter: DestroyXlaTensor
  Value: 190
Counter: ReleaseDataHandles
  Value: 57
Counter: UncachedSyncTensors
  Value: 1
Counter: XRTAllocateFromTensor_Empty
  Value: 3
Counter: XrtCompile_Empty
  Value: 256
Counter: XrtExecuteChained_Empty
  Value: 256
Counter: XrtExecute_Empty
  Value: 256
Counter: XrtRead_Empty
  Value: 256
Counter: XrtReleaseAllocationHandle_Empty
  Value: 256
Counter: XrtReleaseCompileHandle_Empty
  Value: 256
Counter: XrtSessionCount
  Value: 3
Counter: XrtSubTuple_Empty
  Value: 256

vocab_size = 30000

Metric: CompileTime
  TotalSamples: 1
  Counter: 042ms70.149us
  Percentiles: 1%=042ms70.149us; 5%=042ms70.149us; 10%=042ms70.149us; 20%=042ms70.149us; 50%=042ms70.149us; 80%=042ms70.149us; 90%=042ms70.149us; 95%=042ms70.149us; 99%=042ms70.149us
Metric: ExecuteTime
  TotalSamples: 15
  Counter: 01m01s169ms461.549us
  ValueRate: 330ms372.360us / second
  Rate: 0.081014 / second
  Percentiles: 1%=04s056ms53.196us; 5%=04s056ms53.196us; 10%=04s056ms191.283us; 20%=04s057ms529.497us; 50%=04s057ms737.660us; 80%=04s057ms280.381us; 90%=04s187ms433.975us; 95%=04s245ms13.492us; 99%=04s245ms13.492us
Metric: InboundData
  TotalSamples: 15
  Counter: 6.87GB
  ValueRate: 38.04MB / second
  Rate: 0.0811518 / second
  Percentiles: 1%=468.75MB; 5%=468.75MB; 10%=468.75MB; 20%=468.75MB; 50%=468.75MB; 80%=468.75MB; 90%=468.75MB; 95%=468.75MB; 99%=468.75MB
Metric: OutboundData
  TotalSamples: 45
  Counter: 9.01GB
  ValueRate: 48.66MB / second
  Rate: 0.237282 / second
  Percentiles: 1%=5.00KB; 5%=5.00KB; 10%=5.00KB; 20%=5.00KB; 50%=146.48MB; 80%=468.75MB; 90%=468.75MB; 95%=468.75MB; 99%=468.75MB
Metric: ReleaseDataHandlesTime
  TotalSamples: 29
  Counter: 143ms243.878us
  ValueRate: 775.794us / second
  Rate: 0.157061 / second
  Percentiles: 1%=001ms187.762us; 5%=001ms228.535us; 10%=001ms431.039us; 20%=001ms470.628us; 50%=002ms551.110us; 80%=002ms841.982us; 90%=023ms426.264us; 95%=037ms269.406us; 99%=041ms271.595us
Metric: TransferFromServerTime
  TotalSamples: 15
  Counter: 45s944ms687.421us
  ValueRate: 243ms150.699us / second
  Rate: 0.0811518 / second
  Percentiles: 1%=03s709ms852.691us; 5%=03s709ms852.691us; 10%=03s799ms723.038us; 20%=03s819ms237.376us; 50%=03s917ms988.237us; 80%=03s185ms437.898us; 90%=03s212ms24.858us; 95%=04s534ms159.243us; 99%=04s534ms159.243us
Metric: TransferToServerTime
  TotalSamples: 45
  Counter: 01m02s748ms867.091us
  ValueRate: 329ms482.670us / second
  Rate: 0.240117 / second
  Percentiles: 1%=002ms312.820us; 5%=003ms210.478us; 10%=004ms78.145us; 20%=005ms231.926us; 50%=02s925ms549.395us; 80%=02s123ms446.850us; 90%=02s195ms537.121us; 95%=02s281ms298.584us; 99%=03s614ms330.203us
Counter: CachedSyncTensors
  Value: 14
Counter: CreateCompileHandles
  Value: 1
Counter: CreateDataHandles
  Value: 60
Counter: CreateXlaTensor
  Value: 195
Counter: DestroyDataHandles
  Value: 57
Counter: DestroyXlaTensor
  Value: 190
Counter: ReleaseDataHandles
  Value: 57
Counter: UncachedSyncTensors
  Value: 1
Counter: XRTAllocateFromTensor_Empty
  Value: 3
Counter: XrtCompile_Empty
  Value: 256
Counter: XrtExecuteChained_Empty
  Value: 256
Counter: XrtExecute_Empty
  Value: 256
Counter: XrtRead_Empty
  Value: 256
Counter: XrtReleaseAllocationHandle_Empty
  Value: 256
Counter: XrtReleaseCompileHandle_Empty
  Value: 256
Counter: XrtSessionCount
  Value: 3
Counter: XrtSubTuple_Empty
  Value: 256

vocab_size = 29999 , crashes at loss.backward()

2019-08-05 15:40:05.173133: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] >>> Dumping Computation 0                                                              [46/125712]
2019-08-05 15:40:05.173239: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] HloModule SyncTensorsGraph.79
2019-08-05 15:40:05.173248: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-05 15:40:05.173260: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] %MaxComputation.10 (x.11: f32[], y.12: f32[]) -> f32[] {
2019-08-05 15:40:05.173267: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %x.11 = f32[] parameter(0)
2019-08-05 15:40:05.173273: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %y.12 = f32[] parameter(1)
2019-08-05 15:40:05.173280: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   ROOT %maximum.13 = f32[] maximum(f32[] %x.11, f32[] %y.12)
2019-08-05 15:40:05.173286: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] }
2019-08-05 15:40:05.173292: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-05 15:40:05.173298: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] %AddComputation.19 (x.20: f32[], y.21: f32[]) -> f32[] {
2019-08-05 15:40:05.173304: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %x.20 = f32[] parameter(0)
2019-08-05 15:40:05.173311: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %y.21 = f32[] parameter(1)
2019-08-05 15:40:05.173317: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   ROOT %add.22 = f32[] add(f32[] %x.20, f32[] %y.21)
2019-08-05 15:40:05.173323: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] }
2019-08-05 15:40:05.173341: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-05 15:40:05.173348: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] %AddComputation.44 (x.45: s32[], y.46: s32[]) -> s32[] {
2019-08-05 15:40:05.173357: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %x.45 = s32[] parameter(0)
2019-08-05 15:40:05.173365: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %y.46 = s32[] parameter(1)
2019-08-05 15:40:05.173387: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   ROOT %add.47 = s32[] add(s32[] %x.45, s32[] %y.46)
2019-08-05 15:40:05.173394: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] }
2019-08-05 15:40:05.173400: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-05 15:40:05.173406: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] %AddComputation.54 (x.55: f32[], y.56: f32[]) -> f32[] {
2019-08-05 15:40:05.173412: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %x.55 = f32[] parameter(0)
2019-08-05 15:40:05.173418: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %y.56 = f32[] parameter(1)
2019-08-05 15:40:05.173426: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   ROOT %add.57 = f32[] add(f32[] %x.55, f32[] %y.56)
2019-08-05 15:40:05.173436: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] }
2019-08-05 15:40:05.173442: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-05 15:40:05.173448: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] %ScatterCombiner.73 (p0.74: f32[], p1.75: f32[]) -> f32[] {
2019-08-05 15:40:05.173459: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %p0.74 = f32[] parameter(0)
2019-08-05 15:40:05.173468: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %p1.75 = f32[] parameter(1)
2019-08-05 15:40:05.173491: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   ROOT %add.76 = f32[] add(f32[] %p0.74, f32[] %p1.75)
2019-08-05 15:40:05.173507: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] }
2019-08-05 15:40:05.173514: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-05 15:40:05.173520: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] ENTRY %SyncTensorsGraph.79 (param_0.1: s64[32,20,29999], param_1.2: f32[32,128,29999], param_2.27$
 s64[32,20]) -> (f32[32,128,29999]) {
2019-08-05 15:40:05.173527: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %constant.64 = f32[] constant(0)
2019-08-05 15:40:05.173533: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reshape.65 = f32[1,1,1]{2,1,0} reshape(f32[] %constant.64)
2019-08-05 15:40:05.173540: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.66 = f32[1,1,1]{2,1,0} broadcast(f32[1,1,1]{2,1,0} %reshape.65), dimensions={0,1,2}
2019-08-05 15:40:05.173547: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reshape.67 = f32[] reshape(f32[1,1,1]{2,1,0} %broadcast.66)
2019-08-05 15:40:05.173555: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.68 = f32[32,128,29999]{2,1,0} broadcast(f32[] %reshape.67), dimensions={}
2019-08-05 15:40:05.173566: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %iota.69 = s64[32,20,29999,1]{2,0,1,3} iota(), iota_dimension=0
2019-08-05 15:40:05.173577: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %param_0.1 = s64[32,20,29999]{2,0,1} parameter(0)
2019-08-05 15:40:05.173586: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reshape.70 = s64[32,20,29999,1]{3,2,1,0} reshape(s64[32,20,29999]{2,0,1} %param_0.1)
2019-08-05 15:40:05.173596: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %iota.71 = s64[32,20,29999,1]{2,0,1,3} iota(), iota_dimension=2
2019-08-05 15:40:05.173607: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %concatenate.72 = s64[32,20,29999,3]{3,2,1,0} concatenate(s64[32,20,29999,1]{2,0,1,3} %iota.69, 
s64[32,20,29999,1]{3,2,1,0} %reshape.70, s64[32,20,29999,1]{2,0,1,3} %iota.71), dimensions={3}
2019-08-05 15:40:05.173620: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %param_2.27 = s64[32,20]{0,1} parameter(2)
2019-08-05 15:40:05.173635: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reshape.28 = s64[640]{0} reshape(s64[32,20]{0,1} %param_2.27)
2019-08-05 15:40:05.173646: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.32 = s64[640,29999]{1,0} broadcast(s64[640]{0} %reshape.28), dimensions={0}
2019-08-05 15:40:05.173657: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %iota.31 = s64[1,29999]{1,0} iota(), iota_dimension=1
2019-08-05 15:40:05.173673: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reshape.33 = s64[29999]{0} reshape(s64[1,29999]{1,0} %iota.31)
2019-08-05 15:40:05.173684: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.34 = s64[640,29999]{1,0} broadcast(s64[29999]{0} %reshape.33), dimensions={1}
2019-08-05 15:40:05.173695: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %compare.35 = pred[640,29999]{1,0} compare(s64[640,29999]{1,0} %broadcast.32, s64[640,29999]{1,$} %broadcast.34), direction=EQ
2019-08-05 15:40:05.173719: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %constant.29 = f32[] constant(1)
2019-08-05 15:40:05.173727: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.36 = f32[640,29999]{1,0} broadcast(f32[] %constant.29), dimensions={}
2019-08-05 15:40:05.173738: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %constant.30 = f32[] constant(0)
2019-08-05 15:40:05.173748: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.37 = f32[640,29999]{1,0} broadcast(f32[] %constant.30), dimensions={}
2019-08-05 15:40:05.173756: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %select.38 = f32[640,29999]{1,0} select(pred[640,29999]{1,0} %compare.35, f32[640,29999]{1,0} %$roadcast.36, f32[640,29999]{1,0} %broadcast.37)
2019-08-05 15:40:05.173763: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %negate.50 = f32[640,29999]{1,0} negate(f32[640,29999]{1,0} %select.38)
2019-08-05 15:40:05.173774: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %constant.39 = s64[] constant(-100)
2019-08-05 15:40:05.173783: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.40 = s64[640]{0} broadcast(s64[] %constant.39), dimensions={}
2019-08-05 15:40:05.173794: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %compare.41 = pred[640]{0} compare(s64[640]{0} %reshape.28, s64[640]{0} %broadcast.40), directi$n=NE
2019-08-05 15:40:05.173802: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %convert.42 = s32[640]{0} convert(pred[640]{0} %compare.41)
2019-08-05 15:40:05.173809: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %constant.43 = s32[] constant(0)
2019-08-05 15:40:05.173823: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reduce.48 = s32[] reduce(s32[640]{0} %convert.42, s32[] %constant.43), dimensions={0}, to_appl$=%AddComputation.44
2019-08-05 15:40:05.173830: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %convert.49 = f32[] convert(s32[] %reduce.48)
2019-08-05 15:40:05.173836: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.51 = f32[640,29999]{1,0} broadcast(f32[] %convert.49), dimensions={}
2019-08-05 15:40:05.173842: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %divide.52 = f32[640,29999]{1,0} divide(f32[640,29999]{1,0} %negate.50, f32[640,29999]{1,0} %br$adcast.51)
2019-08-05 15:40:05.173849: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %param_1.2 = f32[32,128,29999]{2,1,0} parameter(1)
2019-08-05 15:40:05.173856: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %iota.3 = s64[32,20,29999,1]{2,0,1,3} iota(), iota_dimension=0
2019-08-05 15:40:05.173875: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reshape.4 = s64[32,20,29999,1]{3,2,1,0} reshape(s64[32,20,29999]{2,0,1} %param_0.1)
2019-08-05 15:40:05.173881: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %iota.5 = s64[32,20,29999,1]{2,0,1,3} iota(), iota_dimension=2
2019-08-05 15:40:05.173891: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %concatenate.6 = s64[32,20,29999,3]{3,2,1,0} concatenate(s64[32,20,29999,1]{2,0,1,3} %iota.3, s$4[32,20,29999,1]{3,2,1,0} %reshape.4, s64[32,20,29999,1]{2,0,1,3} %iota.5), dimensions={3}
2019-08-05 15:40:05.173899: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %gather.7 = f32[32,20,29999]{2,1,0} gather(f32[32,128,29999]{2,1,0} %param_1.2, s64[32,20,29999,
3]{3,2,1,0} %concatenate.6), offset_dims={}, collapsed_slice_dims={0,1,2}, start_index_map={0,1,2}, index_vector_dim=3, slice_sizes={1,1,1}
2019-08-05 15:40:05.173906: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reshape.8 = f32[640,29999]{1,0} reshape(f32[32,20,29999]{2,1,0} %gather.7)
2019-08-05 15:40:05.173912: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %constant.9 = f32[] constant(-inf)
2019-08-05 15:40:05.173936: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reduce.14 = f32[640]{0} reduce(f32[640,29999]{1,0} %reshape.8, f32[] %constant.9), dimensions={
1}, to_apply=%MaxComputation.10
2019-08-05 15:40:05.173943: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.15 = f32[640,29999]{1,0} broadcast(f32[640]{0} %reduce.14), dimensions={0}
2019-08-05 15:40:05.173949: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %subtract.16 = f32[640,29999]{1,0} subtract(f32[640,29999]{1,0} %reshape.8, f32[640,29999]{1,0} 
%broadcast.15)
2019-08-05 15:40:05.173956: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %exponential.17 = f32[640,29999]{1,0} exponential(f32[640,29999]{1,0} %subtract.16)
2019-08-05 15:40:05.173964: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %constant.18 = f32[] constant(0)
2019-08-05 15:40:05.173974: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reduce.23 = f32[640]{0} reduce(f32[640,29999]{1,0} %exponential.17, f32[] %constant.18), dimens
ions={1}, to_apply=%AddComputation.19
2019-08-05 15:40:05.173983: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %log.24 = f32[640]{0} log(f32[640]{0} %reduce.23)
2019-08-05 15:40:05.173994: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.25 = f32[640,29999]{1,0} broadcast(f32[640]{0} %log.24), dimensions={0}
2019-08-05 15:40:05.174004: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %subtract.26 = f32[640,29999]{1,0} subtract(f32[640,29999]{1,0} %subtract.16, f32[640,29999]{1,0
} %broadcast.25)
2019-08-05 15:40:05.174014: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %exponential.59 = f32[640,29999]{1,0} exponential(f32[640,29999]{1,0} %subtract.26)
2019-08-05 15:40:05.174024: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %constant.53 = f32[] constant(0)
2019-08-05 15:40:05.174034: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reduce.58 = f32[640]{0} reduce(f32[640,29999]{1,0} %divide.52, f32[] %constant.53), dimensions=
{1}, to_apply=%AddComputation.54
2019-08-05 15:40:05.174044: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.60 = f32[640,29999]{1,0} broadcast(f32[640]{0} %reduce.58), dimensions={0}
2019-08-05 15:40:05.174055: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %multiply.61 = f32[640,29999]{1,0} multiply(f32[640,29999]{1,0} %exponential.59, f32[640,29999]{
1,0} %broadcast.60)
2019-08-05 15:40:05.174067: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %subtract.62 = f32[640,29999]{1,0} subtract(f32[640,29999]{1,0} %divide.52, f32[640,29999]{1,0} 
%multiply.61)
2019-08-05 15:40:05.174074: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reshape.63 = f32[32,20,29999]{2,1,0} reshape(f32[640,29999]{1,0} %subtract.62)
2019-08-05 15:40:05.174081: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %scatter.77 = f32[32,128,29999]{2,1,0} scatter(f32[32,128,29999]{2,1,0} %broadcast.68, s64[32,20
,29999,3]{3,2,1,0} %concatenate.72, f32[32,20,29999]{2,1,0} %reshape.63), update_window_dims={}, inserted_window_dims={0,1,2}, scatter_dims_to_operand_dims={0,1,2}, index_vector_d
im=3, to_apply=%ScatterCombiner.73
2019-08-05 15:40:05.174090: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   ROOT %tuple.78 = (f32[32,128,29999]{2,1,0}) tuple(f32[32,128,29999]{2,1,0} %scatter.77)
2019-08-05 15:40:05.174100: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] }
2019-08-05 15:40:05.174110: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-05 15:40:05.174122: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-05 15:40:05.174130: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] StackTrace:
2019-08-05 15:40:05.174137: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] *** Begin stack trace ***
2019-08-05 15:40:05.174143: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]        tensorflow::CurrentStackTrace[abi:cxx11]()
2019-08-05 15:40:05.174150: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]        xla::util::ReportComputationError(tensorflow::Status const&, absl::Span<xla::XlaComputation
 const* const>)
2019-08-05 15:40:05.174157: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]        xla::util::ShapeHash(xla::Shape const&)
2019-08-05 15:40:05.174165: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]        xla::XrtComputationClient::ExecuteComputation(xla::ComputationClient::Computation const&, a
bsl::Span<std::shared_ptr<xla::ComputationClient::Data> const>, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&, xla::ComputationClient::Exe
cuteComputationOptions const&)
2019-08-05 15:40:05.174172: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-05 15:40:05.174179: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-05 15:40:05.174186: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-05 15:40:05.174193: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-05 15:40:05.174200: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-05 15:40:05.174211: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]        clone
2019-08-05 15:40:05.174220: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] *** End stack trace ***
2019-08-05 15:40:05.174227: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-05 15:40:05.174237: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] Status: Internal: From /job:tpu_worker/replica:0/task:0:
2019-08-05 15:40:05.174245: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] Accelerator device halted prematurely, perhaps due to an on-device check-failure. Node 0 halted un
expectedly at tag:pc 1:0xbdf (from 1:0xc87); want 998:0x6: bounds check 0 [dereference of %s24] for %26 = dma_hbm_to_vmem /*hbm=*/%s24, /*size_in_granules=*/4096, /*vmem=*/%s21, /
*dst_syncflagno=*/%s18
2019-08-05 15:40:05.174262: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] hlo: %reshape.14 = reshape(%copy.46)
2019-08-05 15:40:05.174269: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-05 15:40:05.174278: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-05 15:40:05.174287: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]         [[{{node XRTExecute}}]]
Segmentation fault (core dumped)

Here's gdb output after Segmentation Fault

(gdb) bt
#0  0x00007fffe89c1f36 in __cxxabiv1::__dynamic_cast (src_ptr=0x0, src_type=0x7fffcebc6998 <typeinfo for xla::ComputationClient::Data>, 
    dst_type=0x7fffcebc69b0 <typeinfo for xla::XrtComputationClient::XrtData>, src2dst=0)
    at /home/nwani/m3/conda-bld/compilers_linux-64_1560109574129/work/.build/x86_64-conda_cos6-linux-gnu/src/gcc/libstdc++-v3/libsupc++/dyncast.cc:50
#1  0x00007fffc5ee0ebe in xla::XrtComputationClient::TransferFromServer(absl::Span<std::shared_ptr<xla::ComputationClient::Data> const>) ()
   from /root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch_xla/lib/libxla_computation_client.so
#2  0x00007fffcf294bed in torch_xla::XLATensor::GetTensorsFused (tensors=0x7fff9a7f2a00, writeable=0x0) at torch_xla/csrc/tensor.cpp:867
#3  0x00007fffcf294b18 in torch_xla::XLATensor::GetTensors (tensors=0x7fff9a7f2a00, writeable=0x0) at torch_xla/csrc/tensor.cpp:847
#4  0x00007fffcf2ca9f2 in torch_xla::bridge::XlaCreateTensorList (tensors=..., writeable=0x0) at torch_xla/csrc/aten_xla_bridge.cpp:112
#5  0x00007fffcf1fbf39 in torch_xla::AtenXlaType::_copy_from (self=..., dst=..., non_blocking=<optimized out>) at torch_xla/csrc/aten_xla_type.cpp:265
#6  0x00007fffe49b4791 in at::native::copy_(at::Tensor&, at::Tensor const&, bool) () from /root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/lib/libtorch.so
#7  0x00007fffe4eabda6 in at::TypeDefault::copy_(at::Tensor&, at::Tensor const&, bool) () from /root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/lib/libtorch.so
#8  0x00007fffe711bd3d in torch::autograd::VariableType::copy_(at::Tensor&, at::Tensor const&, bool) ()
   from /root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/lib/libtorch.so
#9  0x00007fffe4bd9334 in at::native::to_impl(at::Tensor const&, c10::TensorOptions const&, bool) ()
   from /root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/lib/libtorch.so
#10 0x00007fffe4bd8db6 in at::native::to(at::Tensor const&, c10::TensorOptions const&, bool, bool) ()
   from /root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/lib/libtorch.so
#11 0x00007fffe4eb96a9 in at::TypeDefault::to(at::Tensor const&, c10::TensorOptions const&, bool, bool) ()
   from /root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/lib/libtorch.so
#12 0x00007fffe6cda318 in torch::autograd::VariableType::to(at::Tensor const&, c10::TensorOptions const&, bool, bool) ()
   from /root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/lib/libtorch.so
#13 0x00007fffe489dca2 in at::Tensor::to(c10::TensorOptions const&, bool, bool) const () from /root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/lib/libtorch.so
#14 0x00007fffe4abe2a9 in at::Tensor::toType(at::DeprecatedTypeProperties const&, bool) const ()
   from /root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/lib/libtorch.so
#15 0x00007fffe7104fa9 in torch::autograd::CopyBackwards::apply(std::vector<torch::autograd::Variable, std::allocator<torch::autograd::Variable> >&&) ()
   from /root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/lib/libtorch.so
#16 0x00007fffe70ff18b in torch::autograd::Node::operator()(std::vector<torch::autograd::Variable, std::allocator<torch::autograd::Variable> >&&) ()
   from /root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/lib/libtorch.so
#17 0x00007fffe70f6496 in torch::autograd::Engine::evaluate_function(torch::autograd::NodeTask&) ()
   from /root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/lib/libtorch.so
#18 0x00007fffe70f557b in torch::autograd::Engine::thread_main(torch::autograd::GraphTask*) () from /root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/lib/libtorch.so
#19 0x00007fffe70f5481 in torch::autograd::Engine::thread_init(int) () from /root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/lib/libtorch.so
#20 0x00007fffe8fb8a47 in torch::autograd::python::PythonEngine::thread_init(int) () from /root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/lib/libtorch_python.so
#21 0x00007fffe89df19d in std::execute_native_thread_routine (__p=0x555556dcdc50)
    at /home/nwani/m3/conda-bld/compilers_linux-64_1560109574129/work/.build/x86_64-conda_cos6-linux-gnu/src/gcc/libstdc++-v3/src/c++11/thread.cc:80
#22 0x00007ffff7bc34a4 in start_thread (arg=0x7fff9a7f4700) at pthread_create.c:456
#23 0x00007ffff7905d0f in clone () at ../sysdeps/unix/sysv/linux/x86_64/clone.S:97

The compile time difference. We do have a compilation cache service side, so a 40ms compilation is definitely a cache hit.
WRT the crash, it seems a machine check. Can you try exporting the following env and rerun?

XLA_USE_32BIT_LONG=1

The compile time difference. We do have a compilation cache service side, so a 40ms compilation is definitely a cache hit.

That means I can run it with vocab_size = 30522, wait for a few minutes, kill the process, then run again, and it will pull the graph from the cache. But this is not happening. CompileTime for the second run is still as slow as the first.

Will try the env and let you know.

Should be.

Should be.

You are right. So right now the problem is the long ExecuteTime, TransferFromServerTime and TransferToServerTime

I tried XLA_USE_32BIT_LONG=1, still got the same segmentation fault. The logs are below

2019-08-05 19:26:48.375519: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] >>> Dumping Computation 0
2019-08-05 19:26:48.375644: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] HloModule SyncTensorsGraph.79
2019-08-05 19:26:48.375652: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-05 19:26:48.375665: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] %MaxComputation.10 (x.11: f32[], y.12: f32[]) -> f32[] {
2019-08-05 19:26:48.375672: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %x.11 = f32[] parameter(0)
2019-08-05 19:26:48.375700: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %y.12 = f32[] parameter(1)
2019-08-05 19:26:48.375707: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   ROOT %maximum.13 = f32[] maximum(f32[] %x.11, f32[] %y.12)
2019-08-05 19:26:48.375715: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] }
2019-08-05 19:26:48.375724: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-05 19:26:48.375730: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] %AddComputation.19 (x.20: f32[], y.21: f32[]) -> f32[] {
2019-08-05 19:26:48.375737: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %x.20 = f32[] parameter(0)
2019-08-05 19:26:48.375755: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %y.21 = f32[] parameter(1)
2019-08-05 19:26:48.375762: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   ROOT %add.22 = f32[] add(f32[] %x.20, f32[] %y.21)
2019-08-05 19:26:48.375768: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] }
2019-08-05 19:26:48.375773: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-05 19:26:48.375779: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] %AddComputation.44 (x.45: s32[], y.46: s32[]) -> s32[] {
2019-08-05 19:26:48.375786: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %x.45 = s32[] parameter(0)
2019-08-05 19:26:48.375791: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %y.46 = s32[] parameter(1)
2019-08-05 19:26:48.375815: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   ROOT %add.47 = s32[] add(s32[] %x.45, s32[] %y.46)
2019-08-05 19:26:48.375821: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] }
2019-08-05 19:26:48.375827: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-05 19:26:48.375833: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] %AddComputation.54 (x.55: f32[], y.56: f32[]) -> f32[] {
2019-08-05 19:26:48.375840: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %x.55 = f32[] parameter(0)
2019-08-05 19:26:48.375846: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %y.56 = f32[] parameter(1)
2019-08-05 19:26:48.375853: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   ROOT %add.57 = f32[] add(f32[] %x.55, f32[] %y.56)
2019-08-05 19:26:48.375859: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] }
2019-08-05 19:26:48.375864: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-05 19:26:48.375870: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] %ScatterCombiner.73 (p0.74: f32[], p1.75: f32[]) -> f32[] {
2019-08-05 19:26:48.375888: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %p0.74 = f32[] parameter(0)
2019-08-05 19:26:48.375895: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %p1.75 = f32[] parameter(1)
2019-08-05 19:26:48.375905: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   ROOT %add.76 = f32[] add(f32[] %p0.74, f32[] %p1.75)
2019-08-05 19:26:48.375912: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] }
2019-08-05 19:26:48.375921: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-05 19:26:48.375931: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] ENTRY %SyncTensorsGraph.79 (param_0.1: s32[32,20,29999], param_1.2: f32$
32,128,29999], param_2.27: s32[32,20]) -> (f32[32,128,29999]) {
2019-08-05 19:26:48.375952: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %constant.64 = f32[] constant(0)
2019-08-05 19:26:48.375962: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reshape.65 = f32[1,1,1]{2,1,0} reshape(f32[] %constant.64)
2019-08-05 19:26:48.375969: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.66 = f32[1,1,1]{2,1,0} broadcast(f32[1,1,1]{2,1,0} %reshap$
.65), dimensions={0,1,2}
2019-08-05 19:26:48.375976: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reshape.67 = f32[] reshape(f32[1,1,1]{2,1,0} %broadcast.66)
2019-08-05 19:26:48.375982: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.68 = f32[32,128,29999]{2,1,0} broadcast(f32[] %reshape.67)$
 dimensions={}
2019-08-05 19:26:48.375989: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %iota.69 = s32[32,20,29999,1]{2,0,1,3} iota(), iota_dimension=0
2019-08-05 19:26:48.375995: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %param_0.1 = s32[32,20,29999]{2,0,1} parameter(0)
2019-08-05 19:26:48.376001: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reshape.70 = s32[32,20,29999,1]{3,2,1,0} reshape(s32[32,20,29999]{2,$
,1} %param_0.1)
2019-08-05 19:26:48.376009: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %iota.71 = s32[32,20,29999,1]{2,0,1,3} iota(), iota_dimension=2
2019-08-05 19:26:48.376017: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %concatenate.72 = s32[32,20,29999,3]{3,2,1,0} concatenate(s32[32,20,2$
999,1]{2,0,1,3} %iota.69, s32[32,20,29999,1]{3,2,1,0} %reshape.70, s32[32,20,29999,1]{2,0,1,3} %iota.71), dimensions={3}
2019-08-05 19:26:48.376026: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %param_2.27 = s32[32,20]{0,1} parameter(2)
2019-08-05 19:26:48.376041: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reshape.28 = s32[640]{0} reshape(s32[32,20]{0,1} %param_2.27)
2019-08-05 19:26:48.376049: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.32 = s32[640,29999]{1,0} broadcast(s32[640]{0} %reshape.28$
, dimensions={0}
2019-08-05 19:26:48.376062: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %iota.31 = s32[1,29999]{1,0} iota(), iota_dimension=1
2019-08-05 19:26:48.376068: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reshape.33 = s32[29999]{0} reshape(s32[1,29999]{1,0} %iota.31)
2019-08-05 19:26:48.376078: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.34 = s32[640,29999]{1,0} broadcast(s32[29999]{0} %reshape.$
3), dimensions={1}
2019-08-05 19:26:48.376089: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %compare.35 = pred[640,29999]{1,0} compare(s32[640,29999]{1,0} %broad$
ast.32, s32[640,29999]{1,0} %broadcast.34), direction=EQ
2019-08-05 19:26:48.376100: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %constant.29 = f32[] constant(1)
2019-08-05 19:26:48.376110: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.36 = f32[640,29999]{1,0} broadcast(f32[] %constant.29), dim
ensions={}
2019-08-05 19:26:48.376118: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %constant.30 = f32[] constant(0)
2019-08-05 19:26:48.376126: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.37 = f32[640,29999]{1,0} broadcast(f32[] %constant.30), dim
ensions={}
2019-08-05 19:26:48.376138: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %select.38 = f32[640,29999]{1,0} select(pred[640,29999]{1,0} %compare.
35, f32[640,29999]{1,0} %broadcast.36, f32[640,29999]{1,0} %broadcast.37)
2019-08-05 19:26:48.376149: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %negate.50 = f32[640,29999]{1,0} negate(f32[640,29999]{1,0} %select.38
)
2019-08-05 19:26:48.376156: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %constant.39 = s32[] constant(-100)
2019-08-05 19:26:48.376165: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.40 = s32[640]{0} broadcast(s32[] %constant.39), dimensions=
{}
2019-08-05 19:26:48.376172: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %compare.41 = pred[640]{0} compare(s32[640]{0} %reshape.28, s32[640]{0
} %broadcast.40), direction=NE
2019-08-05 19:26:48.376181: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %convert.42 = s32[640]{0} convert(pred[640]{0} %compare.41)
2019-08-05 19:26:48.376189: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %constant.43 = s32[] constant(0)
2019-08-05 19:26:48.376196: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reduce.48 = s32[] reduce(s32[640]{0} %convert.42, s32[] %constant.43)
, dimensions={0}, to_apply=%AddComputation.44
2019-08-05 19:26:48.376204: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %convert.49 = f32[] convert(s32[] %reduce.48)
2019-08-05 19:26:48.376211: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.51 = f32[640,29999]{1,0} broadcast(f32[] %convert.49), dime
nsions={}
2019-08-05 19:26:48.376219: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %divide.52 = f32[640,29999]{1,0} divide(f32[640,29999]{1,0} %negate.50
, f32[640,29999]{1,0} %broadcast.51)
2019-08-05 19:26:48.376227: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %param_1.2 = f32[32,128,29999]{2,1,0} parameter(1)
2019-08-05 19:26:48.376238: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %iota.3 = s32[32,20,29999,1]{2,0,1,3} iota(), iota_dimension=0
2019-08-05 19:26:48.376247: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reshape.4 = s32[32,20,29999,1]{3,2,1,0} reshape(s32[32,20,29999]{2,0,
1} %param_0.1)
2019-08-05 19:26:48.376256: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %iota.5 = s32[32,20,29999,1]{2,0,1,3} iota(), iota_dimension=2
2019-08-05 19:26:48.376266: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %concatenate.6 = s32[32,20,29999,3]{3,2,1,0} concatenate(s32[32,20,299
99,1]{2,0,1,3} %iota.3, s32[32,20,29999,1]{3,2,1,0} %reshape.4, s32[32,20,29999,1]{2,0,1,3} %iota.5), dimensions={3}
2019-08-05 19:26:48.376276: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %gather.7 = f32[32,20,29999]{2,1,0} gather(f32[32,128,29999]{2,1,0} %p
aram_1.2, s32[32,20,29999,3]{3,2,1,0} %concatenate.6), offset_dims={}, collapsed_slice_dims={0,1,2}, start_index_map={0,1,2}, index_vector_dim=3, slice_s
izes={1,1,1}
2019-08-05 19:26:48.376285: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reshape.8 = f32[640,29999]{1,0} reshape(f32[32,20,29999]{2,1,0} %gath
er.7)
2019-08-05 19:26:48.376295: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %constant.9 = f32[] constant(-inf)
2019-08-05 19:26:48.376302: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reduce.14 = f32[640]{0} reduce(f32[640,29999]{1,0} %reshape.8, f32[] 
%constant.9), dimensions={1}, to_apply=%MaxComputation.10
2019-08-05 19:26:48.376313: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.15 = f32[640,29999]{1,0} broadcast(f32[640]{0} %reduce.14),
 dimensions={0}
2019-08-05 19:26:48.376321: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %subtract.16 = f32[640,29999]{1,0} subtract(f32[640,29999]{1,0} %resha
pe.8, f32[640,29999]{1,0} %broadcast.15)
2019-08-05 19:26:48.376328: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %exponential.17 = f32[640,29999]{1,0} exponential(f32[640,29999]{1,0} 
%subtract.16)
2019-08-05 19:26:48.376334: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %constant.18 = f32[] constant(0)
2019-08-05 19:26:48.376343: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reduce.23 = f32[640]{0} reduce(f32[640,29999]{1,0} %exponential.17, f
32[] %constant.18), dimensions={1}, to_apply=%AddComputation.19
2019-08-05 19:26:48.376350: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %log.24 = f32[640]{0} log(f32[640]{0} %reduce.23)
2019-08-05 19:26:48.376357: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.25 = f32[640,29999]{1,0} broadcast(f32[640]{0} %log.24), di
mensions={0}
2019-08-05 19:26:48.376363: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %subtract.26 = f32[640,29999]{1,0} subtract(f32[640,29999]{1,0} %subtr
act.16, f32[640,29999]{1,0} %broadcast.25)
2019-08-05 19:26:48.376369: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %exponential.59 = f32[640,29999]{1,0} exponential(f32[640,29999]{1,0} 
%subtract.26)
2019-08-05 19:26:48.376376: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %constant.53 = f32[] constant(0)
2019-08-05 19:26:48.376382: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reduce.58 = f32[640]{0} reduce(f32[640,29999]{1,0} %divide.52, f32[] 
%constant.53), dimensions={1}, to_apply=%AddComputation.54
2019-08-05 19:26:48.376388: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.60 = f32[640,29999]{1,0} broadcast(f32[640]{0} %reduce.58),
 dimensions={0}
2019-08-05 19:26:48.376395: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %multiply.61 = f32[640,29999]{1,0} multiply(f32[640,29999]{1,0} %expon
ential.59, f32[640,29999]{1,0} %broadcast.60)
2019-08-05 19:26:48.376401: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %subtract.62 = f32[640,29999]{1,0} subtract(f32[640,29999]{1,0} %divid
e.52, f32[640,29999]{1,0} %multiply.61)
2019-08-05 19:26:48.376408: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reshape.63 = f32[32,20,29999]{2,1,0} reshape(f32[640,29999]{1,0} %sub
tract.62)
2019-08-05 19:26:48.376414: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %scatter.77 = f32[32,128,29999]{2,1,0} scatter(f32[32,128,29999]{2,1,0
} %broadcast.68, s32[32,20,29999,3]{3,2,1,0} %concatenate.72, f32[32,20,29999]{2,1,0} %reshape.63), update_window_dims={}, inserted_window_dims={0,1,2}, 
scatter_dims_to_operand_dims={0,1,2}, index_vector_dim=3, to_apply=%ScatterCombiner.73
2019-08-05 19:26:48.376422: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   ROOT %tuple.78 = (f32[32,128,29999]{2,1,0}) tuple(f32[32,128,29999]{2,
1,0} %scatter.77)
2019-08-05 19:26:48.376428: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] }
2019-08-05 19:26:48.376434: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-05 19:26:48.376440: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-05 19:26:48.376445: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] StackTrace:
2019-08-05 19:26:48.376451: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] *** Begin stack trace ***
2019-08-05 19:26:48.376458: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]        tensorflow::CurrentStackTrace[abi:cxx11]()
2019-08-05 19:26:48.376464: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]        xla::util::ReportComputationError(tensorflow::Status const&, absl
::Span<xla::XlaComputation const* const>)
2019-08-05 19:26:48.376471: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]        xla::util::ShapeHash(xla::Shape const&)
2019-08-05 19:26:48.376477: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]        xla::XrtComputationClient::ExecuteComputation(xla::ComputationCli
ent::Computation const&, absl::Span<std::shared_ptr<xla::ComputationClient::Data> const>, std::__cxx11::basic_string<char, std::char_traits<char>, std::a
llocator<char> > const&, xla::ComputationClient::ExecuteComputationOptions const&)
2019-08-05 19:26:48.376485: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-05 19:26:48.376491: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-05 19:26:48.376497: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-05 19:26:48.376503: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-05 19:26:48.376509: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-05 19:26:48.376515: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]        clone
2019-08-05 19:26:48.376521: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] *** End stack trace ***
2019-08-05 19:26:48.376527: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-05 19:26:48.376533: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] Status: Internal: From /job:tpu_worker/replica:0/task:0:
2019-08-05 19:26:48.376539: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] Accelerator device halted prematurely, perhaps due to an on-device check
-failure. Node 0 halted unexpectedly at tag:pc 1:0x937 (from 1:0x9df); want 998:0x6: bounds check 0 [dereference of %s24] for %26 = dma_hbm_to_vmem /*hbm
=*/%s24, /*size_in_granules=*/4096, /*vmem=*/%s21, /*dst_syncflagno=*/%s18
2019-08-05 19:26:48.376547: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] hlo: %reshape.2 = reshape(%copy.38)
2019-08-05 19:26:48.376553: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-05 19:26:48.376563: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-05 19:26:48.376569: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]         [[{{node XRTExecute}}]]
#0  0x00007fffe89c1f36 in __cxxabiv1::__dynamic_cast (src_ptr=0x0, src_type=0x7fffcebc6998 <typeinfo for xla::ComputationClient::Data>, 
    dst_type=0x7fffcebc69b0 <typeinfo for xla::XrtComputationClient::XrtData>, src2dst=0)
    at /home/nwani/m3/conda-bld/compilers_linux-64_1560109574129/work/.build/x86_64-conda_cos6-linux-gnu/src/gcc/libstdc++-v3/libsupc++/dyncast.cc:50
#1  0x00007fffc5ee0ebe in xla::XrtComputationClient::TransferFromServer(absl::Span<std::shared_ptr<xla::ComputationClient::Data> const>) ()
   from /root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch_xla/lib/libxla_computation_client.so
#2  0x00007fffcf294bed in torch_xla::XLATensor::GetTensorsFused (tensors=0x7fff6de42a00, writeable=0x0) at torch_xla/csrc/tensor.cpp:867
#3  0x00007fffcf294b18 in torch_xla::XLATensor::GetTensors (tensors=0x7fff6de42a00, writeable=0x0) at torch_xla/csrc/tensor.cpp:847
#4  0x00007fffcf2ca9f2 in torch_xla::bridge::XlaCreateTensorList (tensors=..., writeable=0x0) at torch_xla/csrc/aten_xla_bridge.cpp:112
#5  0x00007fffcf1fbf39 in torch_xla::AtenXlaType::_copy_from (self=..., dst=..., non_blocking=<optimized out>) at torch_xla/csrc/aten_xla_type.cpp:265
#6  0x00007fffe49b4791 in at::native::copy_(at::Tensor&, at::Tensor const&, bool) ()
   from /root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/lib/libtorch.so
#7  0x00007fffe4eabda6 in at::TypeDefault::copy_(at::Tensor&, at::Tensor const&, bool) ()
   from /root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/lib/libtorch.so
#8  0x00007fffe711bd3d in torch::autograd::VariableType::copy_(at::Tensor&, at::Tensor const&, bool) ()
   from /root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/lib/libtorch.so
#9  0x00007fffe4bd9334 in at::native::to_impl(at::Tensor const&, c10::TensorOptions const&, bool) ()
   from /root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/lib/libtorch.so
#10 0x00007fffe4bd8db6 in at::native::to(at::Tensor const&, c10::TensorOptions const&, bool, bool) ()
   from /root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/lib/libtorch.so
#11 0x00007fffe4eb96a9 in at::TypeDefault::to(at::Tensor const&, c10::TensorOptions const&, bool, bool) ()
   from /root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/lib/libtorch.so
#12 0x00007fffe6cda318 in torch::autograd::VariableType::to(at::Tensor const&, c10::TensorOptions const&, bool, bool) ()
   from /root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/lib/libtorch.so
#13 0x00007fffe489dca2 in at::Tensor::to(c10::TensorOptions const&, bool, bool) const ()
   from /root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/lib/libtorch.so
#14 0x00007fffe4abe2a9 in at::Tensor::toType(at::DeprecatedTypeProperties const&, bool) const ()
   from /root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/lib/libtorch.so
#15 0x00007fffe7104fa9 in torch::autograd::CopyBackwards::apply(std::vector<torch::autograd::Variable, std::allocator<torch::autograd::Variable> >&&) ()
   from /root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/lib/libtorch.so
#16 0x00007fffe70ff18b in torch::autograd::Node::operator()(std::vector<torch::autograd::Variable, std::allocator<torch::autograd::Variable> >&&) ()
   from /root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/lib/libtorch.so
#17 0x00007fffe70f6496 in torch::autograd::Engine::evaluate_function(torch::autograd::NodeTask&) ()
   from /root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/lib/libtorch.so
#18 0x00007fffe70f557b in torch::autograd::Engine::thread_main(torch::autograd::GraphTask*) ()
   from /root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/lib/libtorch.so
#19 0x00007fffe70f5481 in torch::autograd::Engine::thread_init(int) ()
   from /root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/lib/libtorch.so
#20 0x00007fffe8fb8a47 in torch::autograd::python::PythonEngine::thread_init(int) ()
   from /root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/lib/libtorch_python.so
#21 0x00007fffe89df19d in std::execute_native_thread_routine (__p=0x555556e0e420)
    at /home/nwani/m3/conda-bld/compilers_linux-64_1560109574129/work/.build/x86_64-conda_cos6-linux-gnu/src/gcc/libstdc++-v3/src/c++11/thread.cc:80
#22 0x00007ffff7bc34a4 in start_thread (arg=0x7fff6de44700) at pthread_create.c:456
#23 0x00007ffff7905d0f in clone () at ../sysdeps/unix/sysv/linux/x86_64/clone.S:97

I hate to sound like the usual support guy, but since that is a machine check, can you try restarting that TPU node?
Are you able to try with a new node?

does it matter if it is a restart or a new node? I can try either.

does it matter if it is a restart or a new node? I can try either.

The intent is to exclude we are not dealing with a bad TPU node.

Restarting the same node gave the same segmentation fault.

Using a new node, it hangs on loss.backward. Here's the gdb backtrace where the hanging happens,

#0  pthread_cond_wait@@GLIBC_2.3.2 () at ../sysdeps/unix/sysv/linux/x86_64/pthread_cond_wait.S:185
#1  0x00007fffe89db4cb in __gthread_cond_wait (__mutex=<optimized out>, __cond=<optimized out>)
    at /home/nwani/m3/conda-bld/compilers_linux-64_1560109574129/work/.build/x86_64-conda_cos6-linux-gnu/build/build-cc-gcc-final/x86_64-conda_cos6-linux
-gnu/libstdc++-v3/include/x86_64-conda_cos6-linux-gnu/bits/gthr-default.h:878
#2  std::condition_variable::wait (this=<optimized out>, __lock=...)
    at /home/nwani/m3/conda-bld/compilers_linux-64_1560109574129/work/.build/x86_64-conda_cos6-linux-gnu/src/gcc/libstdc++-v3/src/c++11/condition_variabl
e.cc:53
#3  0x00007fffc5ed816b in xla::util::MultiWait::Wait() ()
   from /root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch_xla/lib/libxla_computation_client.so
#4  0x00007fffc5edeffc in xla::XrtComputationClient::InitializeDevices(std::unique_ptr<tensorflow::tpu::TopologyProto, std::default_delete<tensorflow::tp
u::TopologyProto> >) () from /root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch_xla/lib/libxla_computation_client.so
#5  0x00007fffc5eddc8f in xla::XrtComputationClient::XrtComputationClient(xla::XrtComputationClient::Options, std::unique_ptr<tensorflow::tpu::TopologyPr
oto, std::default_delete<tensorflow::tpu::TopologyProto> >) ()
   from /root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch_xla/lib/libxla_computation_client.so
#6  0x00007fffc5ec8c1f in xla::ComputationClient::Create() ()
   from /root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch_xla/lib/libxla_computation_client.so
#7  0x00007fffc5ecb797 in xla::ComputationClient::Get() ()
   from /root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch_xla/lib/libxla_computation_client.so
#8  0x00007fffcf1e13e7 in torch_xla::(anonymous namespace)::InitXlaModuleBindings(pybind11::module)::$_6::operator()[abi:cxx11]() const (
    this=<optimized out>) at torch_xla/csrc/init_python_bindings.cpp:226
#9  pybind11::detail::argument_loader<>::call_impl<std::vector<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >, std::allo
cator<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > > >, torch_xla::(anonymous namespace)::InitXlaModuleBindings(pybind
11::module)::$_6&, , pybind11::detail::void_type>(torch_xla::(anonymous namespace)::InitXlaModuleBindings(pybind11::module)::$_6&, std::integer_sequence<
unsigned long>, pybind11::detail::void_type&&) (this=<optimized out>, f=...)
    at /root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/include/pybind11/cast.h:1931
#10 pybind11::detail::argument_loader<>::call<std::vector<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >, std::allocator
<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > > >, pybind11::detail::void_type, torch_xla::(anonymous namespace)::Init
XlaModuleBindings(pybind11::module)::$_6&> (this=<optimized out>, f=...)
    at /root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/include/pybind11/cast.h:1908
#11 void pybind11::cpp_function::initialize<torch_xla::(anonymous namespace)::InitXlaModuleBindings(pybind11::module)::$_6, std::vector<std::__cxx11::bas
ic_string<char, std::char_traits<char>, std::allocator<char> >, std::allocator<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<ch
ar> > > >, , pybind11::name, pybind11::scope, pybind11::sibling>(torch_xla::(anonymous namespace)::InitXlaModuleBindings(pybind11::module)::$_6&&, std::v
ector<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >, std::allocator<std::__cxx11::basic_string<char, std::char_traits<c
har>, std::allocator<char> > > > (*)(), pybind11::name const&, pybind11::scope const&, pybind11::sibling const&)::{lambda(pybind11::detail::function_call
&)#1}::operator()(pybind11::detail::function_call&) const (call=..., this=<optimized out>)
    at /root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/include/pybind11/pybind11.h:156
#12 void pybind11::cpp_function::initialize<torch_xla::(anonymous namespace)::InitXlaModuleBindings(pybind11::module)::$_6, std::vector<std::__cxx11::bas
ic_string<char, std::char_traits<char>, std::allocator<char> >, std::allocator<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<ch
ar> > > >, , pybind11::name, pybind11::scope, pybind11::sibling>(torch_xla::(anonymous namespace)::InitXlaModuleBindings(pybind11::module)::$_6&&, std::v
ector<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >, std::allocator<std::__cxx11::basic_string<char, std::char_traits<c
har>, std::allocator<char> > > > (*)(), pybind11::name const&, pybind11::scope const&, pybind11::sibling const&)::{lambda(pybind11::detail::function_call
&)#1}::__invoke(pybind11::detail::function_call&) (call=...)
    at /root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/include/pybind11/pybind11.h:133
#13 0x00007fffcf1e6e9c in pybind11::cpp_function::dispatcher (self=<optimized out>, args_in=<optimized out>, kwargs_in=0x0)
    at /root/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/include/pybind11/pybind11.h:620
#14 0x0000555555665744 in _PyCFunction_FastCallDict () at /tmp/build/80754af9/python_1546130271559/work/Objects/methodobject.c:231
#15 0x00005555556ec42c in call_function () at /tmp/build/80754af9/python_1546130271559/work/Python/ceval.c:4851
#16 0x000055555571138a in _PyEval_EvalFrameDefault () at /tmp/build/80754af9/python_1546130271559/work/Python/ceval.c:3335
#17 0x00005555556e7289 in _PyEval_EvalCodeWithName (qualname=0x0, name=<optimized out>, closure=0x0, kwdefs=0x0, defcount=0, defs=0x0, kwstep=2, 
    kwcount=<optimized out>, kwargs=0x0, kwnames=0x0, argcount=0, args=0x0, locals=0x7fffcfe58438, globals=0x7fffcfe58438, _co=0x7fffcf7eef60)
    at /tmp/build/80754af9/python_1546130271559/work/Python/ceval.c:4166
#18 PyEval_EvalCodeEx () at /tmp/build/80754af9/python_1546130271559/work/Python/ceval.c:4187
#19 0x00005555556e801c in PyEval_EvalCode (co=co@entry=0x7fffcf7eef60, globals=globals@entry=0x7fffcfe58438, locals=locals@entry=0x7fffcfe58438)
    at /tmp/build/80754af9/python_1546130271559/work/Python/ceval.c:731
#20 0x000055555570ec8b in builtin_exec_impl.isra.11 (locals=0x7fffcfe58438, globals=0x7fffcfe58438, source=0x7fffcf7eef60)
    at /tmp/build/80754af9/python_1546130271559/work/Python/bltinmodule.c:983
#21 builtin_exec () at /tmp/build/80754af9/python_1546130271559/work/Python/clinic/bltinmodule.c.h:283
#22 0x0000555555668711 in PyCFunction_Call () at /tmp/build/80754af9/python_1546130271559/work/Objects/methodobject.c:114
#23 0x00005555557164ad in do_call_core (kwdict=0x7ffff6e4c678, callargs=0x7fffcf867208, func=0x7ffff7fe4990)
    at /tmp/build/80754af9/python_1546130271559/work/Python/ceval.c:5116
#24 _PyEval_EvalFrameDefault () at /tmp/build/80754af9/python_1546130271559/work/Python/ceval.c:3404
---Type <return> to continue, or q <return> to quit---
#25 0x00005555556e58e4 in _PyEval_EvalCodeWithName () at /tmp/build/80754af9/python_1546130271559/work/Python/ceval.c:4166
#26 0x00005555556e6771 in fast_function () at /tmp/build/80754af9/python_1546130271559/work/Python/ceval.c:4992
#27 0x00005555556ec505 in call_function () at /tmp/build/80754af9/python_1546130271559/work/Python/ceval.c:4872
#28 0x000055555571138a in _PyEval_EvalFrameDefault () at /tmp/build/80754af9/python_1546130271559/work/Python/ceval.c:3335
#29 0x00005555556e653b in _PyFunction_FastCall (globals=<optimized out>, nargs=2, args=<optimized out>, co=<optimized out>)
    at /tmp/build/80754af9/python_1546130271559/work/Python/ceval.c:4933
#30 fast_function () at /tmp/build/80754af9/python_1546130271559/work/Python/ceval.c:4968
#31 0x00005555556ec505 in call_function () at /tmp/build/80754af9/python_1546130271559/work/Python/ceval.c:4872
#32 0x000055555571138a in _PyEval_EvalFrameDefault () at /tmp/build/80754af9/python_1546130271559/work/Python/ceval.c:3335
#33 0x00005555556e653b in _PyFunction_FastCall (globals=<optimized out>, nargs=1, args=<optimized out>, co=<optimized out>)
    at /tmp/build/80754af9/python_1546130271559/work/Python/ceval.c:4933
#34 fast_function () at /tmp/build/80754af9/python_1546130271559/work/Python/ceval.c:4968
#35 0x00005555556ec505 in call_function () at /tmp/build/80754af9/python_1546130271559/work/Python/ceval.c:4872
#36 0x000055555571138a in _PyEval_EvalFrameDefault () at /tmp/build/80754af9/python_1546130271559/work/Python/ceval.c:3335
#37 0x00005555556e653b in _PyFunction_FastCall (globals=<optimized out>, nargs=2, args=<optimized out>, co=<optimized out>)
    at /tmp/build/80754af9/python_1546130271559/work/Python/ceval.c:4933
#38 fast_function () at /tmp/build/80754af9/python_1546130271559/work/Python/ceval.c:4968
#39 0x00005555556ec505 in call_function () at /tmp/build/80754af9/python_1546130271559/work/Python/ceval.c:4872
#40 0x000055555571138a in _PyEval_EvalFrameDefault () at /tmp/build/80754af9/python_1546130271559/work/Python/ceval.c:3335
#41 0x00005555556e6bab in _PyFunction_FastCall (globals=<optimized out>, nargs=2, args=<optimized out>, co=<optimized out>)
    at /tmp/build/80754af9/python_1546130271559/work/Python/ceval.c:4933
#42 _PyFunction_FastCallDict () at /tmp/build/80754af9/python_1546130271559/work/Python/ceval.c:5035
#43 0x0000555555665b0f in _PyObject_FastCallDict () at /tmp/build/80754af9/python_1546130271559/work/Objects/abstract.c:2310
#44 0x00005555556a7810 in _PyObject_CallMethodIdObjArgs () at /tmp/build/80754af9/python_1546130271559/work/Objects/abstract.c:2796
#45 0x000055555565cb10 in PyImport_ImportModuleLevelObject () at /tmp/build/80754af9/python_1546130271559/work/Python/import.c:1578
#46 0x0000555555713a8b in import_name (level=0x555555892aa0 <small_ints+160>, fromlist=0x55555584bb30 <_Py_NoneStruct>, name=0x7ffff6ea5fa8, 
    f=0x5555559017a8) at /tmp/build/80754af9/python_1546130271559/work/Python/ceval.c:5245
#47 _PyEval_EvalFrameDefault () at /tmp/build/80754af9/python_1546130271559/work/Python/ceval.c:2899
#48 0x00005555556e7289 in _PyEval_EvalCodeWithName (qualname=0x0, name=<optimized out>, closure=0x0, kwdefs=0x0, defcount=0, defs=0x0, kwstep=2, 
    kwcount=<optimized out>, kwargs=0x0, kwnames=0x0, argcount=0, args=0x0, locals=0x7ffff7f64090, globals=0x7ffff7f64090, _co=0x7ffff7f1fae0)
    at /tmp/build/80754af9/python_1546130271559/work/Python/ceval.c:4166
#49 PyEval_EvalCodeEx () at /tmp/build/80754af9/python_1546130271559/work/Python/ceval.c:4187
#50 0x00005555556e801c in PyEval_EvalCode (co=co@entry=0x7ffff7f1fae0, globals=globals@entry=0x7ffff7f64090, locals=locals@entry=0x7ffff7f64090)
    at /tmp/build/80754af9/python_1546130271559/work/Python/ceval.c:731
#51 0x000055555576a3c4 in run_mod () at /tmp/build/80754af9/python_1546130271559/work/Python/pythonrun.c:1025
#52 0x000055555576a7c1 in PyRun_FileExFlags () at /tmp/build/80754af9/python_1546130271559/work/Python/pythonrun.c:978
#53 0x000055555576a9c3 in PyRun_SimpleFileExFlags () at /tmp/build/80754af9/python_1546130271559/work/Python/pythonrun.c:419
#54 0x000055555576aacd in PyRun_AnyFileExFlags () at /tmp/build/80754af9/python_1546130271559/work/Python/pythonrun.c:81
#55 0x000055555576e4b3 in run_file (p_cf=0x7fffffffe96c, filename=0x5555558a8120 L"tmp/test_segf.py", fp=0x555555937090)
    at /tmp/build/80754af9/python_1546130271559/work/Modules/main.c:340
#56 Py_Main () at /tmp/build/80754af9/python_1546130271559/work/Modules/main.c:811
#57 0x000055555563702e in main () at /tmp/build/80754af9/python_1546130271559/work/Programs/python.c:69
#58 0x00007ffff783d2e1 in __libc_start_main (main=0x555555636f40 <main>, argc=2, argv=0x7fffffffeb78, init=<optimized out>, fini=<optimized out>, 
    rtld_fini=<optimized out>, stack_end=0x7fffffffeb68) at ../csu/libc-start.c:291
#59 0x0000555555717e0e in _start () at ../sysdeps/x86_64/elf/start.S:103

The hang is in InitializeDevices() ... likely you changed the TPU node but did not update the XRT_TPU_CONFIG environment variable with the new IP.

The long time in torch.gather() seems like a corner case happening in the TPU code generation, triggered by your tensor's shape.
We are looking into it internally.

I tried this again with the nightly build but still getting the same problem,

  • with vocab_size = 29999, it crashes
  • with vocab_size = 30522, takes 7 minutes to compile and 5-7 seconds to execute one step

Can you show me how the metrics look like with the new code, and the full crash dump?

code:

vocab_size = 30522  #  or 29999
import torch
import torch_xla_py.xla_model as xm
xla_device = xm.xla_device()

for i in range(10):
    ids = torch.LongTensor(32, 20, vocab_size).random_(0, 127)
    labels = torch.LongTensor(32, 20).random_(0, vocab_size - 1)
    preds = torch.rand(32, 128, vocab_size, requires_grad=True)

    ids = ids.to(xla_device)
    labels = labels.to(xla_device)
    preds = preds.to(xla_device)

    gathered_preds = preds.gather(1, ids)
    loss = torch.nn.functional.cross_entropy(gathered_preds.view(-1, vocab_size), labels.view(-1))

    loss.backward()
import torch_xla
print(torch_xla._XLAC._xla_metrics_report())

with vocab_size = 30522

Metric: CompileTime                                                                                                          
  TotalSamples: 1                                                                                                            
  Counter: 07m13s014ms374.496us                                                                                              
  Percentiles: 1%=07m13s014ms374.496us; 5%=07m13s014ms374.496us; 10%=07m13s014ms374.496us; 20%=07m13s014ms374.496us; 50%=07m1
3s014ms374.496us; 80%=07m13s014ms374.496us; 90%=07m13s014ms374.496us; 95%=07m13s014ms374.496us; 99%=07m13s014ms374.496us     
Metric: ExecuteTime                                                                                                          
  TotalSamples: 10                                                                                                           
  Counter: 01m07s595ms283.077us                                                                                              
  ValueRate: 488ms899.298us / second                                                                                         
  Rate: 0.0732633 / second                                                                                                   
  Percentiles: 1%=05s363ms632.271us; 5%=05s363ms632.271us; 10%=05s363ms151.577us; 20%=05s363ms304.431us; 50%=05s364ms502.492u
s; 80%=05s487ms231.716us; 90%=18s201ms82.580us; 95%=18s201ms82.580us; 99%=18s201ms82.580us                                   
Metric: InboundData                                                                                                          
  TotalSamples: 10                                                                                                           
  Counter: 4.66GB                                                                                                            
  ValueRate: 35.34MB / second                                                                                                
  Rate: 0.0740955 / second                                                                                                   
  Percentiles: 1%=476.91MB; 5%=476.91MB; 10%=476.91MB; 20%=476.91MB; 50%=476.91MB; 80%=476.91MB; 90%=476.91MB; 95%=476.91MB; 
99%=476.91MB                                                                                                                 
Metric: OutboundData                                                                                                         
  TotalSamples: 30                                                                                                           
  Counter: 6.11GB                                                                                                            
  ValueRate: 10.66MB / second                                                                                                
  Rate: 0.0510877 / second                                                                                                   
  Percentiles: 1%=5.00KB; 5%=5.00KB; 10%=5.00KB; 20%=5.00KB; 50%=149.03MB; 80%=476.91MB; 90%=476.91MB; 95%=476.91MB; 99%=476$91MB
Metric: ReleaseDataHandlesTime
  TotalSamples: 19
  Counter: 111ms873.440us
  ValueRate: 822.056us / second
  Rate: 0.140873 / second
  Percentiles: 1%=001ms406.439us; 5%=001ms406.439us; 10%=002ms515.820us; 20%=002ms521.075us; 50%=002ms602.293us; 80%=002ms59.753us; 90%=028ms48.942us; 95%=030ms334.786us; 99%=030ms334.786us
Metric: TransferFromServerTime
  TotalSamples: 10
  Counter: 35s536ms119.141us
  ValueRate: 256ms897.211us / second
  Rate: 0.0740955 / second
  Percentiles: 1%=03s127ms598.948us; 5%=03s127ms598.948us; 10%=03s257ms249.781us; 20%=03s260ms135.089us; 50%=03s292ms772.131us; 80%=03s471ms695.892us; 90%=05s807ms571.249us; 95%=05s807ms571.249us; 99%=05s807ms571.249us
Metric: TransferToServerTime
  TotalSamples: 30
  Counter: 41s834ms809.831us
  ValueRate: 070ms884.415us / second
  Rate: 0.0513431 / second
  Percentiles: 1%=003ms560.909us; 5%=003ms586.445us; 10%=003ms655.259us; 20%=003ms932.400us; 50%=02s810ms270.810us; 80%=02s015ms837.280us; 90%=02s364ms226.410us; 95%=02s408ms724.464us; 99%=04s591ms4.674us
Counter: CachedSyncTensors
  Value: 9
Counter: CreateCompileHandles
  Value: 1
Counter: CreateDataHandles
  Value: 40
Counter: CreateXlaTensor
  Value: 130
Counter: DestroyDataHandles
  Value: 37
Counter: DestroyXlaTensor
  Value: 125
Counter: ReleaseDataHandles
  Value: 37
Counter: UncachedSyncTensors
  Value: 1
Counter: XRTAllocateFromTensor_Empty
  Value: 3
Counter: XrtCompile_Empty
  Value: 256
Counter: XrtExecuteChained_Empty
  Value: 256
Counter: XrtExecute_Empty
  Value: 256
Counter: XrtRead_Empty
  Value: 256
Counter: XrtReleaseAllocationHandle_Empty
  Value: 256
Counter: XrtReleaseCompileHandle_Empty
  Value: 256
Counter: XrtSessionCount
  Value: 3
Counter: XrtSubTuple_Empty
  Value: 256

with vocab_size = 29999

2019-08-19 16:38:34.800472: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] >>> Dumping Computation 0                   
2019-08-19 16:38:34.800609: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] HloModule SyncTensorsGraph.84               
2019-08-19 16:38:34.800619: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]                                             
2019-08-19 16:38:34.800640: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] %MaxComputation.11 (x.12: f32[], y.13: f32[$
) -> f32[] {                                                                                                                 
2019-08-19 16:38:34.800648: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %x.12 = f32[] parameter(0)                
2019-08-19 16:38:34.800655: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %y.13 = f32[] parameter(1)                
2019-08-19 16:38:34.800662: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   ROOT %maximum.14 = f32[] maximum(f32[] %x$
12, f32[] %y.13)                                                                                                             
2019-08-19 16:38:34.800685: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] }                                           
2019-08-19 16:38:34.800693: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]                                             
2019-08-19 16:38:34.800700: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] %AddComputation.20 (x.21: f32[], y.22: f32[$
) -> f32[] {                                                                                                                 
2019-08-19 16:38:34.800731: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %x.21 = f32[] parameter(0)                
2019-08-19 16:38:34.800750: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %y.22 = f32[] parameter(1)                
2019-08-19 16:38:34.800772: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   ROOT %add.23 = f32[] add(f32[] %x.21, f32$
] %y.22)                                                                                                                     
2019-08-19 16:38:34.800786: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] }                                           
2019-08-19 16:38:34.800797: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]                                             
2019-08-19 16:38:34.800810: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] %AddComputation.45 (x.46: s32[], y.47: s32[$
) -> s32[] {                                                                                                                 
2019-08-19 16:38:34.800822: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %x.46 = s32[] parameter(0)                
2019-08-19 16:38:34.800829: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %y.47 = s32[] parameter(1)                
2019-08-19 16:38:34.800836: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   ROOT %add.48 = s32[] add(s32[] %x.46, s32$
] %y.47)                                                                                                                     
2019-08-19 16:38:34.800848: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] }                                           
2019-08-19 16:38:34.800860: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]                                             
2019-08-19 16:38:34.800872: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] %AddComputation.59 (x.60: f32[], y.61: f32[$
) -> f32[] {                                                                                                                 
2019-08-19 16:38:34.800885: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %x.60 = f32[] parameter(0)                
2019-08-19 16:38:34.800897: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %y.61 = f32[] parameter(1)                
2019-08-19 16:38:34.800910: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   ROOT %add.62 = f32[] add(f32[] %x.60, f32$
] %y.61)                                                                                                                     
2019-08-19 16:38:34.800922: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] }                                           
2019-08-19 16:38:34.800929: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]                                             
2019-08-19 16:38:34.800940: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] %ScatterCombiner.78 (p0.79: f32[], p1.80: f$
2[]) -> f32[] {                                                                                                              
2019-08-19 16:38:34.800953: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %p0.79 = f32[] parameter(0)               
2019-08-19 16:38:34.800965: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %p1.80 = f32[] parameter(1)               
2019-08-19 16:38:34.800978: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   ROOT %add.81 = f32[] add(f32[] %p0.79, f3$
[] %p1.80)                                                                                                                   
2019-08-19 16:38:34.800990: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] }                                           
2019-08-19 16:38:34.801002: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]                                             
2019-08-19 16:38:34.801014: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] ENTRY %SyncTensorsGraph.84 (param_0.1: s64[$
2,20,29999], param_1.2: f32[32,128,29999], param_2.28: s64[32,20]) -> (f32[32,128,29999]) {                                  
2019-08-19 16:38:34.801027: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %constant.69 = f32[] constant(0)          
2019-08-19 16:38:34.801039: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reshape.70 = f32[1,1,1]{2,1,0} reshape(f3
2[] %constant.69)
2019-08-19 16:38:34.801052: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.71 = f32[1,1,1]{2,1,0} broadcas
t(f32[1,1,1]{2,1,0} %reshape.70), dimensions={0,1,2}
2019-08-19 16:38:34.801065: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reshape.72 = f32[] reshape(f32[1,1,1]{2,1
,0} %broadcast.71)
2019-08-19 16:38:34.801078: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.73 = f32[32,128,29999]{2,1,0} b
roadcast(f32[] %reshape.72), dimensions={}
2019-08-19 16:38:34.801091: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %iota.74 = s64[32,20,29999,1]{2,0,1,3} iot
a(), iota_dimension=0
2019-08-19 16:38:34.801104: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %param_0.1 = s64[32,20,29999]{2,0,1} param
eter(0)
2019-08-19 16:38:34.801116: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reshape.75 = s64[32,20,29999,1]{3,2,1,0}
reshape(s64[32,20,29999]{2,0,1} %param_0.1)
2019-08-19 16:38:34.801129: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %iota.76 = s64[32,20,29999,1]{2,0,1,3} iot
a(), iota_dimension=2
2019-08-19 16:38:34.801142: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %concatenate.77 = s64[32,20,29999,3]{3,2,1
,0} concatenate(s64[32,20,29999,1]{2,0,1,3} %iota.74, s64[32,20,29999,1]{3,2,1,0} %reshape.75, s64[32,20,29999,1]{2,0,1,3} %i
ota.76), dimensions={3}
2019-08-19 16:38:34.801156: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %param_2.28 = s64[32,20]{0,1} parameter(2)
2019-08-19 16:38:34.801169: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reshape.29 = s64[640]{0} reshape(s64[32,2
0]{0,1} %param_2.28)
2019-08-19 16:38:34.801182: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.33 = s64[640,29999]{1,0} broadc
ast(s64[640]{0} %reshape.29), dimensions={0}
2019-08-19 16:38:34.801194: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %iota.32 = s64[1,29999]{1,0} iota(), iota_
dimension=1
2019-08-19 16:38:34.801208: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reshape.34 = s64[29999]{0} reshape(s64[1,
29999]{1,0} %iota.32)
2019-08-19 16:38:34.801221: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.35 = s64[640,29999]{1,0} broadc
ast(s64[29999]{0} %reshape.34), dimensions={1}
2019-08-19 16:38:34.801234: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %compare.36 = pred[640,29999]{1,0} compare
(s64[640,29999]{1,0} %broadcast.33, s64[640,29999]{1,0} %broadcast.35), direction=EQ
2019-08-19 16:38:34.801246: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %constant.30 = f32[] constant(1)
2019-08-19 16:38:34.801259: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.37 = f32[640,29999]{1,0} broadc
ast(f32[] %constant.30), dimensions={}
2019-08-19 16:38:34.801271: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %constant.31 = f32[] constant(0)
2019-08-19 16:38:34.801283: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.38 = f32[640,29999]{1,0} broadc
ast(f32[] %constant.31), dimensions={}
2019-08-19 16:38:34.801296: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %select.39 = f32[640,29999]{1,0} select(pr
ed[640,29999]{1,0} %compare.36, f32[640,29999]{1,0} %broadcast.37, f32[640,29999]{1,0} %broadcast.38)
2019-08-19 16:38:34.801309: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %negate.55 = f32[640,29999]{1,0} negate(f3
2[640,29999]{1,0} %select.39)
2019-08-19 16:38:34.801321: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %constant.40 = s64[] constant(-100)
2019-08-19 16:38:34.801333: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.41 = s64[640]{0} broadcast(s64[
] %constant.40), dimensions={}
2019-08-19 16:38:34.801346: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %compare.42 = pred[640]{0} compare(s64[640
]{0} %reshape.29, s64[640]{0} %broadcast.41), direction=NE
2019-08-19 16:38:34.801359: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %convert.43 = s32[640]{0} convert(pred[640
]{0} %compare.42)
2019-08-19 16:38:34.801371: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %constant.44 = s32[] constant(0)
2019-08-19 16:38:34.801383: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reduce.49 = s32[] reduce(s32[640]{0} %con
vert.43, s32[] %constant.44), dimensions={0}, to_apply=%AddComputation.45
2019-08-19 16:38:34.801396: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %convert.50 = f32[] convert(s32[] %reduce.
49)
2019-08-19 16:38:34.801408: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %constant.51 = f32[] constant(0)
2019-08-19 16:38:34.801421: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %compare.53 = pred[] compare(f32[] %conver
t.50, f32[] %constant.51), direction=NE
2019-08-19 16:38:34.801433: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %constant.52 = f32[] constant(1)
2019-08-19 16:38:34.801446: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %select.54 = f32[] select(pred[] %compare.
53, f32[] %convert.50, f32[] %constant.52)
2019-08-19 16:38:34.801458: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.56 = f32[640,29999]{1,0} broadc
ast(f32[] %select.54), dimensions={}
2019-08-19 16:38:34.801471: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %divide.57 = f32[640,29999]{1,0} divide(f3
2[640,29999]{1,0} %negate.55, f32[640,29999]{1,0} %broadcast.56)
2019-08-19 16:38:34.801483: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %param_1.2 = f32[32,128,29999]{2,1,0} para
meter(1)
2019-08-19 16:38:34.801495: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %iota.4 = u32[32,20,29999,1]{2,0,1,3} iota
(), iota_dimension=0
2019-08-19 16:38:34.801508: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %convert.3 = u32[32,20,29999]{2,0,1} conve
rt(s64[32,20,29999]{2,0,1} %param_0.1)
2019-08-19 16:38:34.801520: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reshape.5 = u32[32,20,29999,1]{3,2,1,0} r
eshape(u32[32,20,29999]{2,0,1} %convert.3)
2019-08-19 16:38:34.801533: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %iota.6 = u32[32,20,29999,1]{2,0,1[22/734]
(), iota_dimension=2
2019-08-19 16:38:34.801545: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %concatenate.7 = u32[32,20,29999,3]{3,2,1,
0} concatenate(u32[32,20,29999,1]{2,0,1,3} %iota.4, u32[32,20,29999,1]{3,2,1,0} %reshape.5, u32[32,20,29999,1]{2,0,1,3} %iota
.6), dimensions={3}
2019-08-19 16:38:34.801559: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %gather.8 = f32[32,20,29999]{2,1,0} gather
(f32[32,128,29999]{2,1,0} %param_1.2, u32[32,20,29999,3]{3,2,1,0} %concatenate.7), offset_dims={}, collapsed_slice_dims={0,1,
2}, start_index_map={0,1,2}, index_vector_dim=3, slice_sizes={1,1,1}
2019-08-19 16:38:34.801572: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reshape.9 = f32[640,29999]{1,0} reshape(f
32[32,20,29999]{2,1,0} %gather.8)
2019-08-19 16:38:34.801585: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %constant.10 = f32[] constant(-inf)
2019-08-19 16:38:34.801597: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reduce.15 = f32[640]{0} reduce(f32[640,29
999]{1,0} %reshape.9, f32[] %constant.10), dimensions={1}, to_apply=%MaxComputation.11
2019-08-19 16:38:34.801611: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.16 = f32[640,29999]{1,0} broadc
ast(f32[640]{0} %reduce.15), dimensions={0}
2019-08-19 16:38:34.801623: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %subtract.17 = f32[640,29999]{1,0} subtrac
t(f32[640,29999]{1,0} %reshape.9, f32[640,29999]{1,0} %broadcast.16)
2019-08-19 16:38:34.801636: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %exponential.18 = f32[640,29999]{1,0} expo
nential(f32[640,29999]{1,0} %subtract.17)
2019-08-19 16:38:34.801649: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %constant.19 = f32[] constant(0)
2019-08-19 16:38:34.801661: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reduce.24 = f32[640]{0} reduce(f32[640,29
999]{1,0} %exponential.18, f32[] %constant.19), dimensions={1}, to_apply=%AddComputation.20
2019-08-19 16:38:34.801674: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %log.25 = f32[640]{0} log(f32[640]{0} %red
uce.24)
2019-08-19 16:38:34.801686: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.26 = f32[640,29999]{1,0} broadc
ast(f32[640]{0} %log.25), dimensions={0}
2019-08-19 16:38:34.801699: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %subtract.27 = f32[640,29999]{1,0} subtrac
t(f32[640,29999]{1,0} %subtract.17, f32[640,29999]{1,0} %broadcast.26)
2019-08-19 16:38:34.801716: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %exponential.64 = f32[640,29999]{1,0} expo
nential(f32[640,29999]{1,0} %subtract.27)
2019-08-19 16:38:34.801729: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %constant.58 = f32[] constant(0)
2019-08-19 16:38:34.801737: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reduce.63 = f32[640]{0} reduce(f32[640,29
999]{1,0} %divide.57, f32[] %constant.58), dimensions={1}, to_apply=%AddComputation.59
2019-08-19 16:38:34.801757: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.65 = f32[640,29999]{1,0} broadc
ast(f32[640]{0} %reduce.63), dimensions={0}
2019-08-19 16:38:34.801767: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %multiply.66 = f32[640,29999]{1,0} multipl
y(f32[640,29999]{1,0} %exponential.64, f32[640,29999]{1,0} %broadcast.65)
2019-08-19 16:38:34.801774: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %subtract.67 = f32[640,29999]{1,0} subtrac
t(f32[640,29999]{1,0} %divide.57, f32[640,29999]{1,0} %multiply.66)
2019-08-19 16:38:34.801782: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reshape.68 = f32[32,20,29999]{2,1,0} resh
ape(f32[640,29999]{1,0} %subtract.67)
2019-08-19 16:38:34.801790: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %scatter.82 = f32[32,128,29999]{2,1,0} sca
tter(f32[32,128,29999]{2,1,0} %broadcast.73, s64[32,20,29999,3]{3,2,1,0} %concatenate.77, f32[32,20,29999]{2,1,0} %reshape.68
), update_window_dims={}, inserted_window_dims={0,1,2}, scatter_dims_to_operand_dims={0,1,2}, index_vector_dim=3, to_apply=%S
catterCombiner.78
2019-08-19 16:38:34.801799: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   ROOT %tuple.83 = (f32[32,128,29999]{2,1,0}
) tuple(f32[32,128,29999]{2,1,0} %scatter.82)
2019-08-19 16:38:34.801816: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] }
2019-08-19 16:38:34.801824: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]
2019-08-19 16:38:34.801830: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]
2019-08-19 16:38:34.801830: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]
2019-08-19 16:38:34.801837: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] StackTrace:
2019-08-19 16:38:34.801845: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] *** Begin stack trace ***
2019-08-19 16:38:34.801854: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]        tensorflow::CurrentStackTrace[abi:cxx
11]()
2019-08-19 16:38:34.801863: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]        xla::util::ReportComputationError(ten
sorflow::Status const&, absl::Span<xla::XlaComputation const* const>)
2019-08-19 16:38:34.801882: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]        xla::util::ShapeHash(xla::Shape const
&)
2019-08-19 16:38:34.801891: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]        xla::XrtComputationClient::ExecuteCom
putation(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::ExecuteComputationOp
tions const&)
2019-08-19 16:38:34.801899: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]
2019-08-19 16:38:34.801906: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]
2019-08-19 16:38:34.801924: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]
2019-08-19 16:38:34.801935: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]
2019-08-19 16:38:34.801947: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]
2019-08-19 16:38:34.801958: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]        clone
2019-08-19 16:38:34.801971: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] *** End stack trace ***
2019-08-19 16:38:34.801983: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]
2019-08-19 16:38:34.801994: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] Status: Internal: From /job:tpu_worker/repli
ca:0/task:0:
2019-08-19 16:38:34.802007: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] Accelerator device halted prematurely, perha
ps due to an on-device check-failure. Node 0 halted unexpectedly at tag:pc 1:0x1e86 (from 1:0x1f2e); want 993:0x13: bounds ch
eck 0 [dereference of %s25] for %27 = dma_hbm_to_vmem /*hbm=*/%s25, /*size_in_granules=*/4096, /*vmem=*/%s22, /*dst_syncflagn
o=*/%s19
2019-08-19 16:38:34.802021: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] hlo: %reshape.13 = reshape(%copy.46)
2019-08-19 16:38:34.802033: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]
2019-08-19 16:38:34.802044: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]
2019-08-19 16:38:34.802056: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]         [[{{node XRTExecute}}]]

We might have used a too conservative test to switch among Sparse and Dense gather lowering.
Let me drop the threshold ...

944 is in. LMK at the next nightly (tomorrow, since I just merged) if you still see issues.

That should address the performance issue, though the high compile time will need to be checked.

Also, would it be possible to get the XLA HLO dump (the one that comes with the crash) in a non word-wrapped fashion?
Maybe redirect to file?

With vocab_size=30522, it is slightly faster but still executes in seconds

Metric: CompileTime
  TotalSamples: 1
  Counter: 06m19s261ms801.922us
  Percentiles: 1%=06m19s261ms801.922us; 5%=06m19s261ms801.922us; 10%=06m19s261ms801.922us; 20%=06m19s261ms801.922us; 50%=06m19s261ms801.922us; 80%=06m19s261ms801.922us; 90%=06m19s261ms801.922us; 95%=06m19s261ms801.922us; 99%=06m19s261ms801.922us
Metric: ExecuteTime
  TotalSamples: 10
  Counter: 45s373ms834.898us
  ValueRate: 380ms436.463us / second
  Rate: 0.0838467 / second
  Percentiles: 1%=04s868ms679.384us; 5%=04s868ms679.384us; 10%=04s868ms772.534us; 20%=04s868ms812.180us; 50%=04s868ms4.417us; 80%=04s989ms871.644us; 90%=10s439ms547.747us; 95%=10s439ms547.747us; 99%=10s439ms547.747us
Metric: InboundData
  TotalSamples: 10
  Counter: 4.66GB
  ValueRate: 40.14MB / second
  Rate: 0.0841615 / second
  Percentiles: 1%=476.91MB; 5%=476.91MB; 10%=476.91MB; 20%=476.91MB; 50%=476.91MB; 80%=476.91MB; 90%=476.91MB; 95%=476.91MB; 99%=476.91MB
Metric: OutboundData
  TotalSamples: 30
  Counter: 6.11GB
  ValueRate: 12.30MB / second
  Rate: 0.0589481 / second
  Percentiles: 1%=5.00KB; 5%=5.00KB; 10%=5.00KB; 20%=5.00KB; 50%=149.03MB; 80%=476.91MB; 90%=476.91MB; 95%=476.91MB; 99%=476.91MB
Metric: ReleaseDataHandlesTime
  TotalSamples: 19
  Counter: 115ms61.878us
  ValueRate: 968.690us / second
  Rate: 0.159958 / second
  Percentiles: 1%=001ms196.089us; 5%=001ms196.089us; 10%=001ms453.069us; 20%=002ms542.581us; 50%=002ms619.556us; 80%=002ms860.966us; 90%=030ms576.969us; 95%=032ms238.926us; 99%=032ms238.926us
Metric: TransferFromServerTime
  TotalSamples: 10
  Counter: 32s444ms219.187us
  ValueRate: 273ms55.293us / second
  Rate: 0.0841615 / second
  Percentiles: 1%=03s959ms589.445us; 5%=03s959ms589.445us; 10%=03s991ms374.189us; 20%=03s019ms569.185us; 50%=03s314ms241.892us; 80%=04s526ms304.989us; 90%=04s555ms642.843us; 95%=04s555ms642.843us; 99%=04s555ms642.843us
Metric: TransferToServerTime
  TotalSamples: 30
  Counter: 39s707ms400.699us
  ValueRate: 076ms331.426us / second
  Rate: 0.0591603 / second
  Percentiles: 1%=002ms361.096us; 5%=002ms430.936us; 10%=003ms549.033us; 20%=003ms657.558us; 50%=02s774ms733.970us; 80%=02s049ms232.018us; 90%=02s090ms897.206us; 95%=02s104ms981.997us; 99%=02s387ms822.662us
Counter: CachedSyncTensors

with vocab_size=29999, it still crashes
https://gist.github.com/ibeltagy/2fd2d7a6f31c44f0e7a5f02a4ff10813

The good news is that there is no more XLA gather in there.
We will have to investigate the crash issue (which does not seem an OOM anymore) and why it is so slow.

Can you try to:

export XLA_SAVE_TENSORS_FILE=/PATH/TO/FILE
export XLA_SAVE_TENSORS_FMT=hlo

Then run with vocab_size=30522 for a few steps?
The /PATH/TO/FILE is always appended into, so clear after each run.

Also, the logs being posted are truncated (not sure whether the github UI does it), so it is a bit of a pain to rebuild a valid text graph.
For reporting, maybe it is better to redirect to file and avoid the copy&paste from the console to truncate lines.

ok, running the stuff you mentioned. About the log, I think it is github UI. Can you try to scroll horizontally and tell me if you can see the whole thing?

It is not the UI. Or, an horizontal scroll still shows truncated lines.
But we have noticed the embedding lookup still uses a sparse gather.
We will fix that ...

Here's the tensors file: https://gist.github.com/ibeltagy/56d2062d2fe02019b8d6124d96ebfb26
Curious what you are looking at in this file?

Hmm, the 30522 vocab size does not have gather.
Can you try 29999 with:

export XLA_DENSE_GATHER_FACTOR=1000

It should not crash but be as slow as 30522.

Are you using a batch size 128? Mind trying 64 and report exec and compile time metrics?

batch size is 32. The 128 dimension is the one I am doing the gather over. Gather selects 20 values out of the 128.

We do end up with tensors like this, which is 10GB:

%broadcast.9 = f32[32,128,20,30522]{3,2,1,0} broadcast(f32[] %constant.8), dimensions={}

That does not get materialized, but I think the slow compile time might be due to the time spent within the rematerialization code.

I am not sure if 10GB is a lot of not, but the same operation is performed by tensorflow on tpus with no problem.

I think we need to revise our embedding lowering (forward and backward).
Today we use pytorch's native to lower the forward in terms of supported XLA ops, but that generates sub-optimal code.

It still crashes with vocab_size = 29999, export XLA_DENSE_GATHER_FACTOR=1000 and reducing size to 64

Here's the graph:

[ScheduleSyncTensorsGraph]
TensorsGraphInfo:

HloModule IrToHlo.92

%add_F32.12 (lhs.13: f32[], rhs.14: f32[]) -> f32[] {
  %lhs.13 = f32[] parameter(0)
  %rhs.14 = f32[] parameter(1)
  ROOT %add.15 = f32[] add(f32[] %lhs.13, f32[] %rhs.14)
}

%MaxComputation.19 (x.20: f32[], y.21: f32[]) -> f32[] {
  %x.20 = f32[] parameter(0)
  %y.21 = f32[] parameter(1)
  ROOT %maximum.22 = f32[] maximum(f32[] %x.20, f32[] %y.21)
}

%AddComputation.28 (x.29: f32[], y.30: f32[]) -> f32[] {
  %x.29 = f32[] parameter(0)
  %y.30 = f32[] parameter(1)
  ROOT %add.31 = f32[] add(f32[] %x.29, f32[] %y.30)
}

%AddComputation.53 (x.54: s32[], y.55: s32[]) -> s32[] {
  %x.54 = s32[] parameter(0)
  %y.55 = s32[] parameter(1)
  ROOT %add.56 = s32[] add(s32[] %x.54, s32[] %y.55)
}

%AddComputation.67 (x.68: f32[], y.69: f32[]) -> f32[] {
  %x.68 = f32[] parameter(0)
  %y.69 = f32[] parameter(1)
  ROOT %add.70 = f32[] add(f32[] %x.68, f32[] %y.69)
}

%ScatterCombiner.86 (p0.87: f32[], p1.88: f32[]) -> f32[] {
  %p0.87 = f32[] parameter(0)
  %p1.88 = f32[] parameter(1)
  ROOT %add.89 = f32[] add(f32[] %p0.87, f32[] %p1.88)
}

ENTRY %IrToHlo.92 (param_0.1: s64[32,20,29999], param_1.2: f32[32,64,29999], param_2.36: s64[32,20]) -> (f32[32,64,29999]) {
  %constant.77 = f32[] constant(0)
  %reshape.78 = f32[1,1,1]{2,1,0} reshape(f32[] %constant.77)
  %broadcast.79 = f32[1,1,1]{2,1,0} broadcast(f32[1,1,1]{2,1,0} %reshape.78), dimensions={0,1,2}
  %reshape.80 = f32[] reshape(f32[1,1,1]{2,1,0} %broadcast.79)
  %broadcast.81 = f32[32,64,29999]{2,1,0} broadcast(f32[] %reshape.80), dimensions={}
  %iota.82 = s64[32,20,29999,1]{2,0,1,3} iota(), iota_dimension=0
  %param_0.1 = s64[32,20,29999]{2,0,1} parameter(0)
  %reshape.83 = s64[32,20,29999,1]{3,2,1,0} reshape(s64[32,20,29999]{2,0,1} %param_0.1)
  %iota.84 = s64[32,20,29999,1]{2,0,1,3} iota(), iota_dimension=2
  %concatenate.85 = s64[32,20,29999,3]{3,2,1,0} concatenate(s64[32,20,29999,1]{2,0,1,3} %iota.82, s64[32,20,29999,1]{3,2,1,0} %reshape.83, s64[32,20,29999,1]{2,0,1,3} %iota.84), dimensions={3}
  %param_2.36 = s64[32,20]{0,1} parameter(2)
  %reshape.37 = s64[640]{0} reshape(s64[32,20]{0,1} %param_2.36)
  %broadcast.41 = s64[640,29999]{1,0} broadcast(s64[640]{0} %reshape.37), dimensions={0}
  %iota.40 = s64[1,29999]{1,0} iota(), iota_dimension=1
  %reshape.42 = s64[29999]{0} reshape(s64[1,29999]{1,0} %iota.40)
  %broadcast.43 = s64[640,29999]{1,0} broadcast(s64[29999]{0} %reshape.42), dimensions={1}
  %compare.44 = pred[640,29999]{1,0} compare(s64[640,29999]{1,0} %broadcast.41, s64[640,29999]{1,0} %broadcast.43), direction=EQ
  %constant.38 = f32[] constant(1)
  %broadcast.45 = f32[640,29999]{1,0} broadcast(f32[] %constant.38), dimensions={}
  %constant.39 = f32[] constant(0)
  %broadcast.46 = f32[640,29999]{1,0} broadcast(f32[] %constant.39), dimensions={}
  %select.47 = f32[640,29999]{1,0} select(pred[640,29999]{1,0} %compare.44, f32[640,29999]{1,0} %broadcast.45, f32[640,29999]{1,0} %broadcast.46)
  %negate.63 = f32[640,29999]{1,0} negate(f32[640,29999]{1,0} %select.47)
  %constant.48 = s64[] constant(-100)
  %broadcast.49 = s64[640]{0} broadcast(s64[] %constant.48), dimensions={}
  %compare.50 = pred[640]{0} compare(s64[640]{0} %reshape.37, s64[640]{0} %broadcast.49), direction=NE
  %convert.51 = s32[640]{0} convert(pred[640]{0} %compare.50)
  %constant.52 = s32[] constant(0)
  %reduce.57 = s32[] reduce(s32[640]{0} %convert.51, s32[] %constant.52), dimensions={0}, to_apply=%AddComputation.53
  %convert.58 = f32[] convert(s32[] %reduce.57)
  %constant.59 = f32[] constant(0)
  %compare.61 = pred[] compare(f32[] %convert.58, f32[] %constant.59), direction=NE
  %constant.60 = f32[] constant(1)
  %select.62 = f32[] select(pred[] %compare.61, f32[] %convert.58, f32[] %constant.60)
  %broadcast.64 = f32[640,29999]{1,0} broadcast(f32[] %select.62), dimensions={}
  %divide.65 = f32[640,29999]{1,0} divide(f32[640,29999]{1,0} %negate.63, f32[640,29999]{1,0} %broadcast.64)
  %convert.3 = u32[32,20,29999]{2,0,1} convert(s64[32,20,29999]{2,0,1} %param_0.1)
  %broadcast.4 = u32[32,64,20,29999]{3,2,1,0} broadcast(u32[32,20,29999]{2,0,1} %convert.3), dimensions={0,2,3}
  %iota.5 = u32[32,64,20,29999]{3,2,1,0} iota(), iota_dimension=1
  %compare.6 = pred[32,64,20,29999]{3,2,1,0} compare(u32[32,64,20,29999]{3,2,1,0} %broadcast.4, u32[32,64,20,29999]{3,2,1,0} %iota.5), direction=EQ
  %param_1.2 = f32[32,64,29999]{2,1,0} parameter(1)
  %broadcast.7 = f32[32,64,20,29999]{3,2,1,0} broadcast(f32[32,64,29999]{2,1,0} %param_1.2), dimensions={0,1,3}
  %constant.8 = f32[] constant(0)
  %broadcast.9 = f32[32,64,20,29999]{3,2,1,0} broadcast(f32[] %constant.8), dimensions={}
  %select.10 = f32[32,64,20,29999]{3,2,1,0} select(pred[32,64,20,29999]{3,2,1,0} %compare.6, f32[32,64,20,29999]{3,2,1,0} %broadcast.7, f32[32,64,20,29999]{3,2,1,0} %broadcast.9)
  %constant.11 = f32[] constant(0)
  %reduce.16 = f32[32,20,29999]{2,1,0} reduce(f32[32,64,20,29999]{3,2,1,0} %select.10, f32[] %constant.11), dimensions={1}, to_apply=%add_F32.12
  %reshape.17 = f32[640,29999]{1,0} reshape(f32[32,20,29999]{2,1,0} %reduce.16)
  %constant.18 = f32[] constant(-inf)
  %reduce.23 = f32[640]{0} reduce(f32[640,29999]{1,0} %reshape.17, f32[] %constant.18), dimensions={1}, to_apply=%MaxComputation.19
  %broadcast.24 = f32[640,29999]{1,0} broadcast(f32[640]{0} %reduce.23), dimensions={0}
  %subtract.25 = f32[640,29999]{1,0} subtract(f32[640,29999]{1,0} %reshape.17, f32[640,29999]{1,0} %broadcast.24)
  %exponential.26 = f32[640,29999]{1,0} exponential(f32[640,29999]{1,0} %subtract.25)
  %constant.27 = f32[] constant(0)
  %reduce.32 = f32[640]{0} reduce(f32[640,29999]{1,0} %exponential.26, f32[] %constant.27), dimensions={1}, to_apply=%AddComputation.28
  %log.33 = f32[640]{0} log(f32[640]{0} %reduce.32)
  %broadcast.34 = f32[640,29999]{1,0} broadcast(f32[640]{0} %log.33), dimensions={0}
  %subtract.35 = f32[640,29999]{1,0} subtract(f32[640,29999]{1,0} %subtract.25, f32[640,29999]{1,0} %broadcast.34)
  %exponential.72 = f32[640,29999]{1,0} exponential(f32[640,29999]{1,0} %subtract.35)
  %constant.66 = f32[] constant(0)
  %reduce.71 = f32[640]{0} reduce(f32[640,29999]{1,0} %divide.65, f32[] %constant.66), dimensions={1}, to_apply=%AddComputation.67
  %broadcast.73 = f32[640,29999]{1,0} broadcast(f32[640]{0} %reduce.71), dimensions={0}
  %multiply.74 = f32[640,29999]{1,0} multiply(f32[640,29999]{1,0} %exponential.72, f32[640,29999]{1,0} %broadcast.73)
  %subtract.75 = f32[640,29999]{1,0} subtract(f32[640,29999]{1,0} %divide.65, f32[640,29999]{1,0} %multiply.74)
  %reshape.76 = f32[32,20,29999]{2,1,0} reshape(f32[640,29999]{1,0} %subtract.75)
  %scatter.90 = f32[32,64,29999]{2,1,0} scatter(f32[32,64,29999]{2,1,0} %broadcast.81, s64[32,20,29999,3]{3,2,1,0} %concatenate.85, f32[32,20,29999]{2,1,0} %reshape.76), update_window_dims={}, inserted_window_dims={0,1,2}, scatter_dims_to_operand_dims={0,1,2}, index_vector_dim=3, to_apply=%ScatterCombiner.86
  ROOT %tuple.91 = (f32[32,64,29999]{2,1,0}) tuple(f32[32,64,29999]{2,1,0} %scatter.90)
}

and here's the core dump (sorry it is copied from the console again, but I guess you don't need it as the log above is enough)

2019-08-28 01:01:39.824633: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] >>> Dumping Computation 0
2019-08-28 01:01:39.824793: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] HloModule SyncTensorsGraph.92
2019-08-28 01:01:39.824804: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-28 01:01:39.824829: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] %add_F32.12 (lhs.13: f32[], rhs.14: f32[]) -> f32[] {
2019-08-28 01:01:39.824837: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %lhs.13 = f32[] parameter(0)
2019-08-28 01:01:39.824845: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %rhs.14 = f32[] parameter(1)
2019-08-28 01:01:39.824853: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   ROOT %add.15 = f32[] add(f32[] %lhs.13, f32[] %rhs.14)
2019-08-28 01:01:39.824861: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] }
2019-08-28 01:01:39.824868: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-28 01:01:39.824875: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] %MaxComputation.19 (x.20: f32[], y.21: f32[]) -> f32[] {
2019-08-28 01:01:39.824883: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %x.20 = f32[] parameter(0)
2019-08-28 01:01:39.824891: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %y.21 = f32[] parameter(1)
2019-08-28 01:01:39.824898: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   ROOT %maximum.22 = f32[] maximum(f32[] %x.20, f32[] %y.21)
2019-08-28 01:01:39.824906: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] }
2019-08-28 01:01:39.824913: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-28 01:01:39.824920: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] %AddComputation.28 (x.29: f32[], y.30: f32[]) -> f32[] {
2019-08-28 01:01:39.824927: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %x.29 = f32[] parameter(0)
2019-08-28 01:01:39.824936: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %y.30 = f32[] parameter(1)
2019-08-28 01:01:39.824944: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   ROOT %add.31 = f32[] add(f32[] %x.29, f32[] %y.30)
2019-08-28 01:01:39.824951: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] }
2019-08-28 01:01:39.824958: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-28 01:01:39.824979: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] %AddComputation.53 (x.54: s32[], y.55: s32[]) -> s32[] {
2019-08-28 01:01:39.824987: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %x.54 = s32[] parameter(0)
2019-08-28 01:01:39.824995: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %y.55 = s32[] parameter(1)
2019-08-28 01:01:39.825004: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   ROOT %add.56 = s32[] add(s32[] %x.54, s32[] %y.55)
2019-08-28 01:01:39.825014: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] }
2019-08-28 01:01:39.825031: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-28 01:01:39.825048: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] %AddComputation.67 (x.68: f32[], y.69: f32[]) -> f32[] {
2019-08-28 01:01:39.825059: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %x.68 = f32[] parameter(0)
2019-08-28 01:01:39.825103: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %y.69 = f32[] parameter(1)
2019-08-28 01:01:39.825111: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   ROOT %add.70 = f32[] add(f32[] %x.68, f32[] %y.69)
2019-08-28 01:01:39.825118: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] }
2019-08-28 01:01:39.825125: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-28 01:01:39.825133: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] %ScatterCombiner.86 (p0.87: f32[], p1.88: f32[]) -> f32[] {
2019-08-28 01:01:39.825140: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %p0.87 = f32[] parameter(0)
2019-08-28 01:01:39.825147: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %p1.88 = f32[] parameter(1)
2019-08-28 01:01:39.825155: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   ROOT %add.89 = f32[] add(f32[] %p0.87, f32[] %p1.88)
2019-08-28 01:01:39.825164: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] }
2019-08-28 01:01:39.825174: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-28 01:01:39.825185: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] ENTRY %SyncTensorsGraph.92 (param_0.1: s64[32,20,29999], param_1.2: f32[32,64,29999],param_2.36: s64[32,20]) -> (f32[32,64,29999]) {
2019-08-28 01:01:39.825198: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %constant.77 = f32[] constant(0)
2019-08-28 01:01:39.825216: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reshape.78 = f32[1,1,1]{2,1,0} reshape(f32[] %constant.77)
2019-08-28 01:01:39.825228: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.79 = f32[1,1,1]{2,1,0} broadcast(f32[1,1,1]{2,1,0} %reshape.78), dimensions={0,1,2}
2019-08-28 01:01:39.825238: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reshape.80 = f32[] reshape(f32[1,1,1]{2,1,0} %broadcast.79)
2019-08-28 01:01:39.825249: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.81 = f32[32,64,29999]{2,1,0} broadcast(f32[] %reshape.80), dimensions={}
2019-08-28 01:01:39.825261: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %iota.82 = s64[32,20,29999,1]{2,0,1,3} iota(), iota_dimension=0
2019-08-28 01:01:39.825269: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %param_0.1 = s64[32,20,29999]{2,0,1} parameter(0)
2019-08-28 01:01:39.825276: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reshape.83 = s64[32,20,29999,1]{3,2,1,0} reshape(s64[32,20,29999]{2,0,1} %param_0.1)
2019-08-28 01:01:39.825284: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %iota.84 = s64[32,20,29999,1]{2,0,1,3} iota(), iota_dimension=2
2019-08-28 01:01:39.825292: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %concatenate.85 = s64[32,20,29999,3]{3,2,1,0} concatenate(s64[32,20,29999,1]{2,0,1,3} %iota.82, s64[32,20,29999,1]{3,2,1,0} %reshape.83, s64[32,20,29999,1]{2,0,1,3} %iota.84), dimensions={3}
2019-08-28 01:01:39.825300: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %param_2.36 = s64[32,20]{0,1} parameter(2)
2019-08-28 01:01:39.825308: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reshape.37 = s64[640]{0} reshape(s64[32,20]{0,1} %param_2.36)
2019-08-28 01:01:39.825319: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.41 = s64[640,29999]{1,0} broadcast(s64[640]{0} %reshape.37), dimensions={0}
2019-08-28 01:01:39.825329: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %iota.40 = s64[1,29999]{1,0} iota(), iota_dimension=1
2019-08-28 01:01:39.825338: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reshape.42 = s64[29999]{0} reshape(s64[1,29999]{1,0} %iota.40)
2019-08-28 01:01:39.825358: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.43 = s64[640,29999]{1,0} broadcast(s64[29999]{0} %reshape.42), dimensions={1}
2019-08-28 01:01:39.825367: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %compare.44 = pred[640,29999]{1,0} compare(s64[640,29999]{1,0} %broadcast.41, s64[640,29999]{1,0} %broadcast.43), direction=EQ
2019-08-28 01:01:39.825375: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %constant.38 = f32[] constant(1)
2019-08-28 01:01:39.825382: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.45 = f32[640,29999]{1,0} broadcast(f32[] %constant.38), dimensions={}
2019-08-28 01:01:39.825390: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %constant.39 = f32[] constant(0)
2019-08-28 01:01:39.825410: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.46 = f32[640,29999]{1,0} broadcast(f32[] %constant.39), dimensions={}
2019-08-28 01:01:39.825418: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %select.47 = f32[640,29999]{1,0} select(pred[640,29999]{1,0} %compare.44, f32[640,29999]{1,0} %broadcast.45, f32[640,29999]{1,0} %broadcast.46)
2019-08-28 01:01:39.825426: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %negate.63 = f32[640,29999]{1,0} negate(f32[640,29999]{1,0} %select.47)
2019-08-28 01:01:39.825435: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %constant.48 = s64[] constant(-100)
2019-08-28 01:01:39.825442: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.49 = s64[640]{0} broadcast(s64[] %constant.48), dimensions={}
2019-08-28 01:01:39.825450: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %compare.50 = pred[640]{0} compare(s64[640]{0} %reshape.37, s64[640]{0} %broadcast.49), direction=NE
2019-08-28 01:01:39.825459: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %convert.51 = s32[640]{0} convert(pred[640]{0} %compare.50)
2019-08-28 01:01:39.825466: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %constant.52 = s32[] constant(0)
2019-08-28 01:01:39.825474: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reduce.57 = s32[] reduce(s32[640]{0} %convert.51, s32[] %constant.52), dimensions={0}, to_apply=%AddComputation.53
2019-08-28 01:01:39.825482: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %convert.58 = f32[] convert(s32[] %reduce.57)
2019-08-28 01:01:39.825490: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %constant.59 = f32[] constant(0)
2019-08-28 01:01:39.825502: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %compare.61 = pred[] compare(f32[] %convert.58, f32[] %constant.59), direction=NE
2019-08-28 01:01:39.825512: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %constant.60 = f32[] constant(1)
2019-08-28 01:01:39.825523: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %select.62 = f32[] select(pred[] %compare.61, f32[] %convert.58, f32[] %constant.60)
2019-08-28 01:01:39.825548: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.64 = f32[640,29999]{1,0} broadcast(f32[] %select.62), dimensions={}
2019-08-28 01:01:39.825556: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %divide.65 = f32[640,29999]{1,0} divide(f32[640,29999]{1,0} %negate.63, f32[640,29999]{1,0} %broadcast.64)
2019-08-28 01:01:39.825564: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %convert.3 = u32[32,20,29999]{2,0,1} convert(s64[32,20,29999]{2,0,1} %param_0.1)
2019-08-28 01:01:39.825571: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.4 = u32[32,64,20,29999]{3,2,1,0} broadcast(u32[32,20,29999]{2,0,1} %convert.3), dimensions={0,2,3}
2019-08-28 01:01:39.825579: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %iota.5 = u32[32,64,20,29999]{3,2,1,0} iota(), iota_dimension=1
2019-08-28 01:01:39.825587: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %compare.6 = pred[32,64,20,29999]{3,2,1,0} compare(u32[32,64,20,29999]{3,2,1,0} %broadcast.4, u32[32,64,20,29999]{3,2,1,0} %iota.5), direction=EQ
2019-08-28 01:01:39.825595: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %param_1.2 = f32[32,64,29999]{2,1,0} parameter(1)
2019-08-28 01:01:39.825603: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.7 = f32[32,64,20,29999]{3,2,1,0} broadcast(f32[32,64,29999]{2,1,0} %param_1.2), dimensions={0,1,3}
2019-08-28 01:01:39.825611: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %constant.8 = f32[] constant(0)
2019-08-28 01:01:39.825635: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.9 = f32[32,64,20,29999]{3,2,1,0} broadcast(f32[] %constant.8), dimensions={}
2019-08-28 01:01:39.825645: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %select.10 = f32[32,64,20,29999]{3,2,1,0} select(pred[32,64,20,29999]{3,2,1,0} %compare.6, f32[32,64,20,29999]{3,2,1,0} %broadcast.7, f32[32,64,20,29999]{3,2,1,0} %broadcast.9)
2019-08-28 01:01:39.825658: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %constant.11 = f32[] constant(0)
2019-08-28 01:01:39.825667: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reduce.16 = f32[32,20,29999]{2,1,0} reduce(f32[32,64,20,29999]{3,2,1,0} %select.10, f32[] %constant.11), dimensions={1}, to_apply=%add_F32.12
2019-08-28 01:01:39.825675: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reshape.17 = f32[640,29999]{1,0} reshape(f32[32,20,29999]{2,1,0} %reduce.16)
2019-08-28 01:01:39.825688: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %constant.18 = f32[] constant(-inf)
2019-08-28 01:01:39.825708: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reduce.23 = f32[640]{0} reduce(f32[640,29999]{1,0} %reshape.17, f32[] %constant.18), dimensions={1}, to_apply=%MaxComputation.19
2019-08-28 01:01:39.825719: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.24 = f32[640,29999]{1,0} broadcast(f32[640]{0} %reduce.23), dimensions={0}
2019-08-28 01:01:39.825727: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %subtract.25 = f32[640,29999]{1,0} subtract(f32[640,29999]{1,0} %reshape.17, f32[640,29999]{1,0} %broadcast.24)
2019-08-28 01:01:39.825738: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %exponential.26 = f32[640,29999]{1,0} exponential(f32[640,29999]{1,0} %subtract.25)
2019-08-28 01:01:39.825747: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %constant.27 = f32[] constant(0)
2019-08-28 01:01:39.825755: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reduce.32 = f32[640]{0} reduce(f32[640,29999]{1,0} %exponential.26, f32[] %constant.27), dimensions={1}, to_apply=%AddComputation.28
2019-08-28 01:01:39.825769: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %log.33 = f32[640]{0} log(f32[640]{0} %reduce.32)
2019-08-28 01:01:39.825786: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.34 = f32[640,29999]{1,0} broadcast(f32[640]{0} %log.33), dimensions={0}
2019-08-28 01:01:39.825800: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %subtract.35 = f32[640,29999]{1,0} subtract(f32[640,29999]{1,0} %subtract.25, f32[640,29999]{1,0} %broadcast.34)
2019-08-28 01:01:39.825821: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %exponential.72 = f32[640,29999]{1,0} exponential(f32[640,29999]{1,0} %subtract.35)
2019-08-28 01:01:39.825832: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %constant.66 = f32[] constant(0)
2019-08-28 01:01:39.825840: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reduce.71 = f32[640]{0} reduce(f32[640,29999]{1,0} %divide.65, f32[] %constant.66), dimensions={1}, to_apply=%AddComputation.67
2019-08-28 01:01:39.825851: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %broadcast.73 = f32[640,29999]{1,0} broadcast(f32[640]{0} %reduce.71), dimensions={0}
2019-08-28 01:01:39.825861: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %multiply.74 = f32[640,29999]{1,0} multiply(f32[640,29999]{1,0} %exponential.72, f32[640,29999]{1,0} %broadcast.73)
2019-08-28 01:01:39.825871: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %subtract.75 = f32[640,29999]{1,0} subtract(f32[640,29999]{1,0} %divide.65, f32[640,29999]{1,0} %multiply.74)
2019-08-28 01:01:39.825886: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %reshape.76 = f32[32,20,29999]{2,1,0} reshape(f32[640,29999]{1,0} %subtract.75)
2019-08-28 01:01:39.825899: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   %scatter.90 = f32[32,64,29999]{2,1,0} scatter(f32[32,64,29999]{2,1,0} %broadcast.81, s64[32,20,29999,3]{3,2,1,0} %concatenate.85, f32[32,20,29999]{2,1,0} %reshape.76), update_window_dims={}, inse
rted_window_dims={0,1,2}, scatter_dims_to_operand_dims={0,1,2}, index_vector_dim=3, to_apply=%ScatterCombiner.86
2019-08-28 01:01:39.825917: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]   ROOT %tuple.91 = (f32[32,64,29999]{2,1,0}) tuple(f32[32,64,29999]{2,1,0} %scatter.90)
2019-08-28 01:01:39.825928: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] }
2019-08-28 01:01:39.825938: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-28 01:01:39.825945: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-28 01:01:39.825954: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] StackTrace:
2019-08-28 01:01:39.825966: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] *** Begin stack trace ***
2019-08-28 01:01:39.825974: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]        tensorflow::CurrentStackTrace[abi:cxx11]()
2019-08-28 01:01:39.825983: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]        xla::util::ReportComputationError(tensorflow::Status const&, absl::Span<xla::XlaComputation const* const>)
2019-08-28 01:01:39.825997: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]        xla::util::ShapeHash(xla::Shape const&)
2019-08-28 01:01:39.826012: E 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, st
d::char_traits<char>, std::allocator<char> > const&, xla::ComputationClient::ExecuteComputationOptions const&)
2019-08-28 01:01:39.826026: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-28 01:01:39.826041: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-28 01:01:39.826051: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-28 01:01:39.826091: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-28 01:01:39.826099: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-28 01:01:39.826111: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]        clone
2019-08-28 01:01:39.826122: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] *** End stack trace ***
2019-08-28 01:01:39.826131: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-28 01:01:39.826141: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] Status: Internal: From /job:tpu_worker/replica:0/task:0:
2019-08-28 01:01:39.826151: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] Accelerator device halted prematurely, perhaps due to an on-device check-failure. Node 0 halted unexpectedly at tag:pc 16:0x18d9 (from 16:0x1981); want 498:0x13: bounds check 0 [dereference of %s25
] for %27 = dma_hbm_to_vmem /*hbm=*/%s25, /*size_in_granules=*/4096, /*vmem=*/%s22, /*dst_syncflagno=*/%s19
2019-08-28 01:01:39.826165: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] hlo: %reshape.12 = reshape(%copy.19)
2019-08-28 01:01:39.826175: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-28 01:01:39.826182: E tensorflow/compiler/xla/xla_client/xla_util.cc:72] 
2019-08-28 01:01:39.826191: E tensorflow/compiler/xla/xla_client/xla_util.cc:72]         [[{{node XRTExecute}}]]
Segmentation fault (core dumped)

We have two tasks for us here.
First, we need to look at the TPU crash and compile slow performance (filed an internal bug).
Second, we need to improve the embedding forward/backward lowering.

I am not sure how much faster we can make that gather.
On XLA CPU takes about 3.3s while on PyTorch CPU takes about 3s.
On TPU we are at 4.8s.
@ibeltagy Do you happen to have GPU times?

On Titan V,

  • gather, cross_entropy loss, loss.backward take 0.3s
  • only gather takes 0.11ms

So, our gather is fast.
Running this:

import torch
import torch_xla
import torch_xla_py.xla_model as xm

vocab_size = 30522

xla_device = xm.xla_device()

ids = torch.LongTensor(32, 20, vocab_size).random_(0, 127)
preds = torch.rand(32, 128, vocab_size, requires_grad=True)
ids = ids.to(xla_device)
preds = preds.to(xla_device)
gathered_preds = preds.gather(1, ids)
print(torch_xla._XLAC._get_xla_tensors_hlo([gathered_preds]))
torch_xla._XLAC._xla_sync_multi([gathered_preds], [], True, True)
print(torch_xla._XLAC._xla_metrics_report())

Gives this:

HloModule IrToHlo.18

%add_F32.12 (lhs.13: f32[], rhs.14: f32[]) -> f32[] {
  %lhs.13 = f32[] parameter(0)
  %rhs.14 = f32[] parameter(1)
  ROOT %add.15 = f32[] add(f32[] %lhs.13, f32[] %rhs.14)
}

ENTRY %IrToHlo.18 (param_0.1: s64[32,20,30522], param_1.2: f32[32,128,30522]) -> (f32[32,20,30522]) {
  %param_0.1 = s64[32,20,30522]{2,0,1} parameter(0)
  %convert.3 = u32[32,20,30522]{2,0,1} convert(s64[32,20,30522]{2,0,1} %param_0.1)
  %broadcast.4 = u32[32,128,20,30522]{3,2,1,0} broadcast(u32[32,20,30522]{2,0,1} %convert.3), dimensions={0,2,3}
  %iota.5 = u32[32,128,20,30522]{3,2,1,0} iota(), iota_dimension=1
  %compare.6 = pred[32,128,20,30522]{3,2,1,0} compare(u32[32,128,20,30522]{3,2,1,0} %broadcast.4, u32[32,128,20,30522]{3,2,1,0} %iota.5), direction=EQ
  %param_1.2 = f32[32,128,30522]{2,1,0} parameter(1)
  %broadcast.7 = f32[32,128,20,30522]{3,2,1,0} broadcast(f32[32,128,30522]{2,1,0} %param_1.2), dimensions={0,1,3}
  %constant.8 = f32[] constant(0)
  %broadcast.9 = f32[32,128,20,30522]{3,2,1,0} broadcast(f32[] %constant.8), dimensions={}
  %select.10 = f32[32,128,20,30522]{3,2,1,0} select(pred[32,128,20,30522]{3,2,1,0} %compare.6, f32[32,128,20,30522]{3,2,1,0} %broadcast.7, f32[32,128,20,30522]{3,2,1,0} %broadcast.9)
  %constant.11 = f32[] constant(0)
  %reduce.16 = f32[32,20,30522]{2,1,0} reduce(f32[32,128,20,30522]{3,2,1,0} %select.10, f32[] %constant.11), dimensions={1}, to_apply=%add_F32.12
  ROOT %tuple.17 = (f32[32,20,30522]{2,1,0}) tuple(f32[32,20,30522]{2,1,0} %reduce.16)
}


Metric: CompileTime
  TotalSamples: 1
  Counter: 031ms881.421us
  Percentiles: 1%=031ms881.421us; 5%=031ms881.421us; 10%=031ms881.421us; 20%=031ms881.421us; 50%=031ms881.421us; 80%=031ms881.421us; 90%=031ms881.421us; 95%=031ms881.421us; 99%=031ms881.421us
Metric: ExecuteTime
  TotalSamples: 1
  Counter: 043ms209.812us
  Percentiles: 1%=043ms209.812us; 5%=043ms209.812us; 10%=043ms209.812us; 20%=043ms209.812us; 50%=043ms209.812us; 80%=043ms209.812us; 90%=043ms209.812us; 95%=043ms209.812us; 99%=043ms209.812us

It is the scatter we do on the backward pass which blows up compilation and runtime.
We are looking into that.

On Titan V,

  • gather, cross_entropy loss, loss.backward take 0.3s
  • only gather takes 0.11ms

@ibeltagy ...

0.11ms or 0.11s?
0.11ms seems a bit too fast given the amount of memory it needs to read/write.

sorry, it is 11ms not 0.11ms

@dlibenzi, any updates about this?

@dlibenzi, any updates about this?

Sorry about that, but not yet. We are working on it internally.
There are two issues to be tackled. Compilation speed and runtime performance.
Will update the issue when we have something ...

We have an implementation of dense scatter which passes our tests, though there are a few more things which needs to be checked:

https://github.com/pytorch/xla/commits/wire_dense_scatter

That PR went in.
We are still looking into the XLA scatter issues, but this PR might be able to get around that.
Let us know if it did ...

Thanks. It is a lot better; Compile and and Execute are fast, and it doesn't crash with vocab_size=29999.

However, it seems that all the time is being spent in TransferFromServer and TransferToServer, both are still slow (in the order of seconds).

Metric: CompileTime
  TotalSamples: 1
  Counter: 129ms687.525us
  Percentiles: 1%=129ms687.525us; 5%=129ms687.525us; 10%=129ms687.525us; 20%=129ms687.525us; 50%=129ms687.525us; 80%=129ms687.525us; 90%=129ms687.525us; 95%=129ms687.525us; 99%=129ms687.525us
Metric: ExecuteTime
  TotalSamples: 10
  Counter: 01s188ms859.320us
  ValueRate: 014ms453.299us / second
  Rate: 0.121675 / second
  Percentiles: 1%=103ms974.501us; 5%=103ms974.501us; 10%=103ms1.919us; 20%=103ms50.919us; 50%=103ms323.326us; 80%=147ms406.419us; 90%=215ms984.119us; 95%=215ms984.119us; 99%=215ms984.119us
Metric: InboundData
  TotalSamples: 10
  Counter: 4.66GB
  ValueRate: 58.32MB / second
  Rate: 0.122291 / second
  Percentiles: 1%=476.91MB; 5%=476.91MB; 10%=476.91MB; 20%=476.91MB; 50%=476.91MB; 80%=476.91MB; 90%=476.91MB; 95%=476.91MB; 99%=476.91MB
Metric: OutboundData
  TotalSamples: 30
  Counter: 6.11GB
  ValueRate: 72.51MB / second
  Rate: 0.347519 / second
  Percentiles: 1%=5.00KB; 5%=5.00KB; 10%=5.00KB; 20%=5.00KB; 50%=149.03MB; 80%=476.91MB; 90%=476.91MB; 95%=476.91MB; 99%=476.91MB
Metric: ReleaseDataHandlesTime
  TotalSamples: 19
  Counter: 115ms753.795us
  ValueRate: 001ms403.640us / second
  Rate: 0.232403 / second
  Percentiles: 1%=002ms727.645us; 5%=002ms727.645us; 10%=002ms740.455us; 20%=002ms845.226us; 50%=002ms989.087us; 80%=002ms161.697us; 90%=029ms687.395us; 95%=029ms433.816us; 99%=029ms433.816us
Metric: TransferFromServerTime
  TotalSamples: 10
  Counter: 33s281ms297.003us
  ValueRate: 407ms999.701us / second
  Rate: 0.122291 / second
  Percentiles: 1%=03s109ms347.112us; 5%=03s109ms347.112us; 10%=03s151ms601.461us; 20%=03s192ms159.595us; 50%=03s384ms68.899us; 80%=04s564ms315.624us; 90%=04s583ms896.839us; 95%=04s583ms896.839us; 99%=04s583ms896.839us
Metric: TransferToServerTime
  TotalSamples: 30
  Counter: 38s866ms368.913us
  ValueRate: 449ms432.873us / second
  Rate: 0.356068 / second
  Percentiles: 1%=005ms489.553us; 5%=006ms959.086us; 10%=006ms120.860us; 20%=007ms765.345us; 50%=02s819ms329.882us; 80%=02s886ms108.281us; 90%=02s921ms511.044us; 95%=02s938ms604.071us; 99%=03s523ms42.903us
Counter: CachedSyncTensors
  Value: 9
Counter: CreateCompileHandles
  Value: 1
Counter: CreateDataHandles
  Value: 40
Counter: CreateXlaTensor
  Value: 130
Counter: DestroyDataHandles
  Value: 37
Counter: DestroyXlaTensor
  Value: 125
Counter: ReleaseDataHandles
  Value: 37
Counter: UncachedSyncTensors
  Value: 1
Counter: XRTAllocateFromTensor_Empty
  Value: 3
Counter: XrtCompile_Empty
  Value: 256
Counter: XrtExecuteChained_Empty
  Value: 256
Counter: XrtExecute_Empty
  Value: 256
Counter: XrtRead_Empty
  Value: 256
Counter: XrtReleaseAllocationHandle_Empty
  Value: 256
Counter: XrtReleaseCompileHandle_Empty
  Value: 256
Counter: XrtSessionCount
  Value: 3
Counter: XrtSubTuple_Empty
  Value: 256

The model seems fully fused (no aten::* counters), and your TPU step time seems 100..150ms.
Your outbound data suggests you are sending to TPU different tensor sizes. Up to 476MB.

Metric: OutboundData
  TotalSamples: 30
  Counter: 6.11GB
  ValueRate: 72.51MB / second
  Rate: 0.347519 / second
  Percentiles: 1%=5.00KB; 5%=5.00KB; 10%=5.00KB; 20%=5.00KB; 50%=149.03MB; 80%=476.91MB; 90%=476.91MB; 95%=476.91MB; 99%=476.91MB

The transfer time is very high, agreed. We do hide TPU device data load behind TPU computation, but now the TPU computation is too small to hide 4s uploads.

Metric: TransferFromServerTime
  TotalSamples: 10
  Counter: 33s281ms297.003us
  ValueRate: 407ms999.701us / second
  Rate: 0.122291 / second
  Percentiles: 1%=03s109ms347.112us; 5%=03s109ms347.112us; 10%=03s151ms601.461us; 20%=03s192ms159.595us; 50%=03s384ms68.899us; 80%=04s564ms315.624us; 90%=04s583ms896.839us; 95%=04s583ms896.839us; 99%=04s583ms896.839us

Is your user VM in the same zone as you TPU VM?

Your outbound data suggests you are sending to TPU different tensor sizes. Up to 476MB.

Yes, three tensors are sent on each step, 2 very small, and 1 a lot larger (code below). Is 4 seconds a reasonable time for 476MB or is that slow?

Is your user VM in the same zone as you TPU VM?

No, both are in us-central1-a

vocab_size = 30522
import torch
import torch_xla_py.xla_model as xm
xla_device = xm.xla_device()

for i in range(10):
    ids = torch.LongTensor(32, 20, vocab_size).random_(0, 127)
    labels = torch.LongTensor(32, 20).random_(0, vocab_size - 1)
    preds = torch.rand(32, 128, vocab_size, requires_grad=True)

    ids = ids.to(xla_device)
    labels = labels.to(xla_device)
    preds = preds.to(xla_device)

    gathered_preds = preds.gather(1, ids)
    loss = torch.nn.functional.cross_entropy(gathered_preds.view(-1, vocab_size), labels.view(-1))

    loss.backward()
import torch_xla
print(torch_xla._XLAC._xla_metrics_report())

wait, in practice, we don't send that much data to the TPU. Such large tensors are computed as the output of another layer. In that case, this code snippet is not representative of the actual work load.

One problem is that you are not using our DataParallel interface.
That one sends data in background, and all these three tensor uploads would go in parallel, and overlapping with TPU computation.

But yes, a 476MB upload taking 4s is WAY too much.

@ibeltagy I am closing this as it is no more about gather.

Was this page helpful?
0 / 5 - 0 ratings