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
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,
vocab_size = 29999, it crashes vocab_size = 30522, takes 7 minutes to compile and 5-7 seconds to execute one stepCan 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 ...
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.3sgather takes 0.11msSo, 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.backwardtake 0.3s- only
gathertakes 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:
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.