Xla: Cannot train complex model with attention

[EugKar]

❓ Questions and Help

When training a model with a for-loop, e.g. encoder-decoder model with attention (in particular, I'm training modified version of Tacotron2, training hangs. Is there a way to train such models using XLA?
Stack trace showed that it hangs inside XLATensor::TryLimitGraphSize after completing several decoder steps. This repeats several times at decoder's loop until crash.
I tried to assign high values to environment variables TRIM_GRAPH_CHECK_FREQUENCY and TRIM_GRAPH_SIZE. However, in this case it hung inside loss.item() call when printing loss value (as far as I understand, this is the place where actual computations on TPU start). This time stack trace shwed that it was inside xla::XrtComputationClient::Compile(). Here is the metrics output just before the item() call (notice huge value of tensors counter):

Metric: OutboundData
  TotalSamples: 100
  Accumulator: 115.25MB
  ValueRate: 9.80MB / second
  Rate: 8.504 / second
  Percentiles: 1%=4.00B; 5%=4.00B; 10%=4.00B; 20%=320.00B; 50%=2.00KB; 80%=490.00KB; 90%=5.00MB; 95%=6.12MB; 99%=24.50MB
Metric: TransferToServerTime
  TotalSamples: 100
  Accumulator: 867ms839.460us
  ValueRate: 074ms754.783us / second
  Rate: 8.50847 / second
  Percentiles: 1%=765.640us; 5%=823.710us; 10%=895.508us; 20%=001ms92.462us; 50%=002ms864.861us; 80%=008ms587.956us; 90%=013ms609.068us; 95%=035ms557.536us; 99%=161ms36.493us
Counter: CreateDataHandles
  Value: 100
Counter: CreateXlaTensor
  Value: 437618
Counter: DestroyXlaTensor
  Value: 434066
Counter: XRTAllocateFromTensor_Empty
  Value: 32
Counter: XrtSessionCount
  Value: 1

Without item() call and having TRIM_GRAPH_SIZE and TRIM_GRAPH_CHECK_FREQUENCY set to very high values, training steps are completed successfully, while there are no signs that there are any computations on TPU, like spikes on memory/CPU graphs. When I printed _get_xla_tensors_text with XLA_IR_DEBUG=1 XLA_HLO_DEBUG=1, I noticed that the majority of the tensors had been created in LSTM code.
So, it looks like a new branch of the graph is created at each step of the for-loop. Is there any way to fix that?
Decoder's loop has fixed steps count, and it looks like all shapes remain unchanged at every step.
P.S. I ran training on a single TPU core.

[dlibenzi]

The metrics shows it did not even finish one compile.
Any chance you can point me to the loop in question, instead of me fishing in the whole code? :smile:

One could place xm.mark_step() call every N loop iterations, to avoid us unrolling the whole thing.

Training on a single core is a good idea when debugging.

[EugKar]

I mean this loop. I've modified it in the following way, however, it does not affect the speed:

mel_outputs = torch.zeros([max_output_length, batch_size, self.n_mel_channels],
    dtype=decoder_inputs.dtype, device=decoder_inputs.device)
gate_outputs = torch.zeros([max_output_length, batch_size], dtype=mel_outputs.dtype,
    device=mel_outputs.device)
alignments = torch.zeros([max_output_length, mel_outputs.size(1), max_memory_length],
    dtype=mel_outputs.dtype, device=mel_outputs.device)
for i in range(max_output_length):
    decoder_input = decoder_inputs[i]
    mel_output, gate_output, attention_weights = self.decode(
        decoder_input, z_latent, z_observed)
    mel_outputs[i, ...] = mel_output.squeeze(1)
    gate_outputs[i, ...] = gate_output.squeeze()
    alignments[i, ...] = attention_weights

[taylanbil]

tpu branch in the pytorch-tpu/fairseq repo is an example of such models. We do train transformer without issues there. Specifically see this function.

[dlibenzi]

What is the value of max_output_length?

[EugKar]

It is 1600. The first XLATensor::TryLimitGraphSize occurs approximately at step 300.

[dlibenzi]

OK, so it is likely going to create a pretty monster graph w/out any cuts, and likely it is taking long time to compile.
The good news is that, if shapes are not changing, the compilation happen only at the beginning, and it is also cached on the server side (unless you reboot the TPU node), so the next runs should take far less time.
Can you try to let the compilation go and see what happens?

Also, can you enable debug logging on tensor so we can see the progress?

export TF_CPP_VMODULE=tensor=5

[EugKar]

Link to the output. After almost 2 hours of multiple iterations of compiling and executing IR graph, it finally crashed. Top memory usage on TPU was 58G at that point (I use v3-8), so it doesn't look like out of memory.

[dlibenzi]

Where have you seen to TPU HBM being 58GB? I don't see it in the logs...
Each core has 16GB, so 58GB is well above the max per core HBM (memory is not shared across core, so even if a TPU v3-8 has a total of 16GB*8 HBM memory, the limit is 16GB).

In any case, since the graph hash is constantly changing, it likely means there is some dynamism in your model.
So either the model graph itself changes per-step, or the shapes of the tensors.
Any chance you can print the debug metrics after each step?

import torch_xla.debug.metrics as met
print(met.metrics_report())

[EugKar]

The value 58GB was taken from TPU memory usage graph in GCP Console.
Here is a small piece of step-wise metrics output, just between steps when the first IR compilation happens (full log will be available later):

Step 307
Metric: OutboundData
  TotalSamples: 55
  Accumulator: 98.63MB
  ValueRate: 83.01MB / second
  Rate: 46.2891 / second
  Percentiles: 1%=4.00B; 5%=4.00B; 10%=128.00B; 20%=2.00KB; 50%=6.00KB; 80%=1.00MB; 90%=5.00MB; 95%=16.00MB; 99%=24.50MB
Metric: TransferToServerTime
  TotalSamples: 55
  Accumulator: 795ms528.322us
  ValueRate: 683ms972.019us / second
  Rate: 47.2777 / second
  Percentiles: 1%=845.356us; 5%=981.663us; 10%=001ms209.574us; 20%=001ms347.565us; 50%=005ms453.715us; 80%=011ms545.791us; 90%=036ms440.901us; 95%=066ms233.371us; 99%=166ms395.218us
Counter: CreateDataHandles
  Value: 55
Counter: CreateXlaTensor
  Value: 35594
Counter: DestroyXlaTensor
  Value: 14057
Counter: XRTAllocateFromTensor_Empty
  Value: 28
Counter: XrtSessionCount
  Value: 1
2019-10-13 16:54:20.932888: I torch_xla/csrc/tensor.cpp:973] Waiting on device barrier for device TPU:0 ...
2019-10-13 16:54:20.932965: I torch_xla/csrc/tensor.cpp:976] Waiting on device barrier for device TPU:0 done!
2019-10-13 16:54:20.933015: I torch_xla/csrc/tensor.cpp:1236] Syncing graph hash 5418318861087816881 on device 'TPU:0'
2019-10-13 16:54:21.732418: I torch_xla/csrc/tensor.cpp:1268] Compiling IR graph hash 5418318861087816881 on device TPU:0 ...
2019-10-13 16:55:19.867977: I torch_xla/csrc/tensor.cpp:1273] Compiling IR graph hash 5418318861087816881 on device TPU:0 done!
2019-10-13 16:55:19.868210: I torch_xla/csrc/tensor.cpp:1105] Executing IR graph hash 5418318861087816881 on device TPU:0 ...
2019-10-13 16:55:26.768482: I torch_xla/csrc/tensor.cpp:1110] Executing IR graph hash 5418318861087816881 on device TPU:0 done!

Step 308
Metric: OutboundData
  TotalSamples: 55
  Accumulator: 98.63MB
  ValueRate: 83.01MB / second
  Rate: 46.2891 / second
  Percentiles: 1%=4.00B; 5%=4.00B; 10%=128.00B; 20%=2.00KB; 50%=6.00KB; 80%=1.00MB; 90%=5.00MB; 95%=16.00MB; 99%=24.50MB
Metric: TransferToServerTime
  TotalSamples: 55
  Accumulator: 795ms528.322us
  ValueRate: 683ms972.019us / second
  Rate: 47.2777 / second
  Percentiles: 1%=845.356us; 5%=981.663us; 10%=001ms209.574us; 20%=001ms347.565us; 50%=005ms453.715us; 80%=011ms545.791us; 90%=036ms440.901us; 95%=066ms233.371us; 99%=166ms395.218us
Counter: CreateDataHandles
  Value: 55
Counter: CreateXlaTensor
  Value: 35675
Counter: DestroyXlaTensor
  Value: 14096
Counter: XRTAllocateFromTensor_Empty
  Value: 28
Counter: XrtSessionCount
  Value: 1

Step 309
Metric: OutboundData
  TotalSamples: 55
  Accumulator: 98.63MB
  ValueRate: 83.01MB / second
  Rate: 46.2891 / second
  Percentiles: 1%=4.00B; 5%=4.00B; 10%=128.00B; 20%=2.00KB; 50%=6.00KB; 80%=1.00MB; 90%=5.00MB; 95%=16.00MB; 99%=24.50MB
Metric: TransferToServerTime
  TotalSamples: 55
  Accumulator: 795ms528.322us
  ValueRate: 683ms972.019us / second
  Rate: 47.2777 / second
  Percentiles: 1%=845.356us; 5%=981.663us; 10%=001ms209.574us; 20%=001ms347.565us; 50%=005ms453.715us; 80%=011ms545.791us; 90%=036ms440.901us; 95%=066ms233.371us; 99%=166ms395.218us
Counter: CreateDataHandles
  Value: 55
Counter: CreateXlaTensor
  Value: 35756
Counter: DestroyXlaTensor
  Value: 14135
Counter: XRTAllocateFromTensor_Empty
  Value: 28
Counter: XrtSessionCount
  Value: 1

Step 310
Metric: OutboundData
  TotalSamples: 55
  Accumulator: 98.63MB
  ValueRate: 83.01MB / second
  Rate: 46.2891 / second
  Percentiles: 1%=4.00B; 5%=4.00B; 10%=128.00B; 20%=2.00KB; 50%=6.00KB; 80%=1.00MB; 90%=5.00MB; 95%=16.00MB; 99%=24.50MB
Metric: TransferToServerTime
  TotalSamples: 55
  Accumulator: 795ms528.322us
  ValueRate: 683ms972.019us / second
  Rate: 47.2777 / second
  Percentiles: 1%=845.356us; 5%=981.663us; 10%=001ms209.574us; 20%=001ms347.565us; 50%=005ms453.715us; 80%=011ms545.791us; 90%=036ms440.901us; 95%=066ms233.371us; 99%=166ms395.218us
Counter: CreateDataHandles
  Value: 55
Counter: CreateXlaTensor
  Value: 35837
Counter: DestroyXlaTensor
  Value: 14174
Counter: XRTAllocateFromTensor_Empty
  Value: 28
Counter: XrtSessionCount
  Value: 1

Step 311
Metric: CompileTime
  TotalSamples: 1
  Accumulator: 58s135ms422.889us
  Percentiles: 1%=58s135ms422.889us; 5%=58s135ms422.889us; 10%=58s135ms422.889us; 20%=58s135ms422.889us; 50%=58s135ms422.889us; 80%=58s135ms422.889us; 90%=58s135ms422.889us; 95%=58s135ms422.889us; 99%=58s135ms422.889us
Metric: ExecuteTime
  TotalSamples: 1
  Accumulator: 07s900ms159.089us
  Percentiles: 1%=07s900ms159.089us; 5%=07s900ms159.089us; 10%=07s900ms159.089us; 20%=07s900ms159.089us; 50%=07s900ms159.089us; 80%=07s900ms159.089us; 90%=07s900ms159.089us; 95%=07s900ms159.089us; 99%=07s900ms159.089us
Metric: OutboundData
  TotalSamples: 55
  Accumulator: 98.63MB
  ValueRate: 83.01MB / second
  Rate: 46.2891 / second
  Percentiles: 1%=4.00B; 5%=4.00B; 10%=128.00B; 20%=2.00KB; 50%=6.00KB; 80%=1.00MB; 90%=5.00MB; 95%=16.00MB; 99%=24.50MB
Metric: SyncTensorsGraphSize
  TotalSamples: 1
  Accumulator: 51967.00
  Percentiles: 1%=51967.00; 5%=51967.00; 10%=51967.00; 20%=51967.00; 50%=51967.00; 80%=51967.00; 90%=51967.00; 95%=51967.00; 99%=51967.00
Metric: TransferToServerTime
  TotalSamples: 55
  Accumulator: 795ms528.322us
  ValueRate: 683ms972.019us / second
  Rate: 47.2777 / second
  Percentiles: 1%=845.356us; 5%=981.663us; 10%=001ms209.574us; 20%=001ms347.565us; 50%=005ms453.715us; 80%=011ms545.791us; 90%=036ms440.901us; 95%=066ms233.371us; 99%=166ms395.218us
Counter: CreateCompileHandles
  Value: 1
Counter: CreateDataHandles
  Value: 56
Counter: CreateXlaTensor
  Value: 35918
Counter: DestroyXlaTensor
  Value: 14213
Counter: TrimIrGraph
  Value: 1
Counter: UncachedSyncTensors
  Value: 1
Counter: XRTAllocateFromTensor_Empty
  Value: 28
Counter: XrtCompile_Empty
  Value: 128
Counter: XrtExecuteChained_Empty
  Value: 128
Counter: XrtExecute_Empty
  Value: 128
Counter: XrtRead_Empty
  Value: 128
Counter: XrtReleaseAllocationHandle_Empty
  Value: 128
Counter: XrtReleaseCompileHandle_Empty
  Value: 128
Counter: XrtSessionCount
  Value: 2
Counter: XrtSubTuple_Empty
  Value: 128

[dlibenzi]

Unless there is a new feature I am not aware of, that memory is CPU host memory of the TPU VM, not the TPU on-board HBM memory.

Any longer data log to post?
The posted log data seems a bit strange to me.
If these are really step numbers 309,310,311 we should see much more metric data.

[EugKar]

Attached a full log with metrics output
It also seems strange for me that the first output telling about compiling IR appears after step 307, and CompileTime metric appears in 4 steps, at step 311. IR dump appears after step 1198, however, stack trace output only appears after step 1217, looks like they are unrelated.

[dlibenzi]

From the experience we have with model's behavior, that log looks strange.
Where do you print the step count?
There should be 1 Execute per step, while in your case there are many steps w/out Execute.
Are you using XLA's DataParallel, XLA multi-processing with ParallelLoader, or your own thing sending tensors with to() calls?
If the latter, do you have barrier=True as xm.optimizer_step() argument?

Your graph is recompiling because it is growing step over step (see SyncTensorsGraphSize metric growing), not because there are trips to PyTorch CPU ops (no aten::* counters).

[EugKar]

I print step count at every iteration of the for-loop from this comment. To avoid confusions, by "step" I mean iteration in that loop, not loss minimization step.
As soon as I'm testing on a single TPU core, I use my own to() call, and I tried with and without barrier=True. However, as you can see in the log I've sent you earlier, xm.optimizer_step() even haven't been reached.

[dlibenzi]

Oh, I see, that explains it! :smile:
In your loop, can you try something like:

for i in range(max_output_length):
  ...
  if i % 100 == 0:
    xm.mark_step()

You can play with that 100.

Would be interesting to see what happens on the TPU VM side and why you get a "driver not open" error.

BTW, are you running on Cloud TPU or on Colab?

[dlibenzi]

Also, maybe using torch.stack() might be a better way to write that loop?

[EugKar]

I'm running on Cloud TPU. I tried torch.stack() before, however, without debug output - it didn't work for me. Maybe, it's worth trying again with all that debug output, as well as mark_step().

[dlibenzi]

Maybe PyTorch is generating better code WRT your slice-update.
Worth a try IMHO.

[EugKar]

Tried torch.stack() version with xm.mark_step() using different intervals - 100, 10, and 1.
This is output with interval 1.. It did not crash, but hung on step 1596 for several hours - I couldn't even ssh to the instance...

[dlibenzi]

Hi!
We will have to look into this model, and understand what are the issues with it.
Clearly unrolling a 1600 long loop might be out of the picture.
Any way to try with shorter loop iterations?
AFAICT the decode graph is pretty thick, and we are repeating it 1600 times.
Might just be too much.

To you test, it seems the model is still trimming, while in theory it should not, especially with intervals as low as 1.
Do you still have the trim environment variables set to high values?

[EugKar]

No, I didn't set trim environment variables in the last experiment.

[EugKar]

Hi,
after updating to the latest nightly build and running the same code which I was able to run before, I get this error at LSTM.forward():

  File "/home/eugkar/tacotron2-pytorch/model.py", line 188, in forward
    outputs, _ = self.lstm(x)
  File "/anaconda3/envs/pytorch-nightly/lib/python3.6/site-packages/torch/nn/modules/module.py", line 541, in __call__
    result = self.forward(*input, **kwargs)
  File "/anaconda3/envs/pytorch-nightly/lib/python3.6/site-packages/torch/nn/modules/rnn.py", line 570, in forward
    return self.forward_tensor(input, hx)
  File "/anaconda3/envs/pytorch-nightly/lib/python3.6/site-packages/torch/nn/modules/rnn.py", line 549, in forward_tensor
    output, hidden = self.forward_impl(input, hx, batch_sizes, max_batch_size, sorted_indices)
  File "/anaconda3/envs/pytorch-nightly/lib/python3.6/site-packages/torch/nn/modules/rnn.py", line 532, in forward_impl
    self.dropout, self.training, self.bidirectional, self.batch_first)
RuntimeError: Input and parameter tensors are not at the same device, found input tensor at xla:1 and parameter tensor at cpu

to has been called for model, and there are successful forward() calls of convolutional layers before this LSTM call.

[dlibenzi]

Funny, I was trying on Saturday to look how a PyTorch LSTM network looked like, graph wise, and I got the same error, from here:

https://github.com/pytorch/pytorch/blob/6301d62e0b205c53a445ffb87ce7be1ac52d9cb9/torch/nn/modules/rnn.py#L532

What I think it is happening, the LSTM module does not seem to be clean WRT the to() operation, as it caches weights in a side list, which do not get converted to XLA device:

https://github.com/pytorch/pytorch/blob/6301d62e0b205c53a445ffb87ce7be1ac52d9cb9/torch/nn/modules/rnn.py#L88

The LSTMCell is clean from that POV, IIRC.

[EugKar]

Hi,
I've just updated to the latest nightly build, LSTM issue seems to be fixed, and not only this issue.
Now with xm.mark_step() each 100 step, all 1600 steps run without crash for just several minutes. However, I get OOM after the final step. at the stage of loss calculation.
The only thing I've changed in my code is using ParallelLoader instead of DataParallel.

[dlibenzi]

Yes, @mruberry fixed the LSTM issue.

Just curious. If you bump up the TRIM env variables really high, and let run w/out mark_step, how long does it take to compile, and what does the OOM message look like?

[EugKar]

Here is log with OOM error
When started with high TRIM variables (1e7) without mark_step, it hung up after step 1599 on compiling, and is still compiling...

[dlibenzi]

Argh! So close 😄

Total hbm usage >= 16.61G:

I assume that OOM error is related to the mark_step every 100 steps attempt, right?

Did it eventually finish compiling?

[EugKar]

I assume that OOM error is related to the mark_step every 100 steps attempt, right?

Yes

Did it eventually finish compiling?

Still waiting.

[dlibenzi]

Wonder which monster comes out of it.
Next time you run, can you export:

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

We need to take a look at that and see how we can do better.
My guess, the graph is big and OOM-ing, and the compiler is trying to find a path to not make it to.

[EugKar]

Hi, here is a tensors output

[dlibenzi]

Thanks!
That is 1/2 the graph (the forward only). That is likely because there is a loss.item() or print(loss) somewhere.
If you remove that we should be able to fetch the whole graph.
Did it ever finish compiling that?
You can check the tensors file, it should have one more entry (the forward + backward) if it did.

[EugKar]

Yes, there is print(loss.item()), however it didn't reach that point and hung even before backward step, according to the log.

[EugKar]

You are right, there was something wrong with logging. After I commented item() call, it successfully finished each step. However, the graph file has not been saved. At what stage is graph saved?

[dlibenzi]

If should be saved as soon as it reaches the barrier, next time you dequeue the next sample from the loader iterator ...

[dlibenzi]

So the forward (the graph above the loss) takes about 25m to compile. We are looking internally to speed that up.
Would be nice to get the whole graph, in order to debug that as well.

[dlibenzi]

We have an internal change coming into nightly which cuts the compilation time by half.
We will do more investigation also as to further reduce that.

[EugKar]

Would be nice to get the whole graph, in order to debug that as well.

Unfortunately, with huge Trim variables, and without item() call, tensors file is not saved. If I use default values for Trim variables, or add item(), only graph for the forward step is saved. I'll try it again tomorrow with nightly updates.

[dlibenzi]

Thanks!
They should be saved. That's what we do before even starting compiling.
Unless there is an infinite loop somewhere.

[EugKar]

Unfortunately, graph is still not saved. It is only saved when I call mark_step() or item(). Is there a way to explicitly invoke graph saving?

[dlibenzi]

Yes, the graph is saved from code inside the mark_step() path.
And mark_step() is called from within the PyTorch/XLA data iterators.
Aren't you using those?

[EugKar]

Yes, I'm using ParallelLoader.

[dlibenzi]

Then it should be called (mark_step()):

https://github.com/pytorch/xla/blob/2d6feaf221f7352440a731bbebf031887d05461b/torch_xla/distributed/parallel_loader.py#L34

[EugKar]

Actually, there was a bug in the way how I iterated over the loader, however, fix didn't help so far. I'll wait for a while and see whether the tensors file appears...

[EugKar]

Finally, I reduced the sequence length to 10, and got tensors file! Seems that several optimization steps have passed. Here it is: tensors.log

[dlibenzi]

Thanks, but 10 is far from 1600.
Does it run with 10?
We need to get a 1600 long graph to have an idea of the final size.

[EugKar]

Yes, it can run more than one step with 10, but with 1600 it does not even save tensors log file...

[dlibenzi]

With 10, can you let run 10..20 steps and print the metrics?
How far can you stretch that 10 to? Does 100 work?

[EugKar]

100 works, too. I manually stopped it after 41 iterations.
Here is the log

[dlibenzi]

Looks good to me.
You are only compiling 2 times:

Metric: CompileTime
  TotalSamples: 2
  Accumulator: 07m58s793ms768.238us
  ValueRate: 02s816ms572.895us / second
  Rate: 0.00869126 / second
  Percentiles: 1%=03m28s607ms1.760us; 5%=03m28s607ms1.760us; 10%=03m28s607ms1.760us; 20%=03m28s607ms1.760us; 50%=04m30s186ms766.478us; 80%=04m30s186ms766.478us; 90%=04m30s186ms766.478us; 95%=04m30s186ms766.478us; 99%=04m30s186ms766.478us

TPU step time ranges from 160ms to 300ms. Likely the 160ms is the stable one:

Metric: ExecuteTime
  TotalSamples: 41
  Accumulator: 35s537ms992.219us
  ValueRate: 055ms927.147us / second
  Rate: 0.0652058 / second
  Percentiles: 1%=158ms271.634us; 5%=159ms649.915us; 10%=159ms797.303us; 20%=159ms89.631us; 50%=160ms218.506us; 80%=275ms953.521us; 90%=281ms812.964us; 95%=306ms682.391us; 99%=13s459ms465.283us

Your graph is 136K IR nodes in size:

Metric: SyncTensorsGraphSize
  TotalSamples: 41
  Accumulator: 5610335.00
  ValueRate: 6585.74 / second
  Rate: 0.0481282 / second
  Percentiles: 1%=136615.00; 5%=136843.00; 10%=136843.00; 20%=136843.00; 50%=136843.00; 80%=136843.00; 90%=136843.00; 95%=136843.00; 99%=136843.00

And you do not seem (given the TPU step time size) to be bound to data upload:

Metric: TransferToServerTime
  TotalSamples: 196
  Accumulator: 02s753ms904.779us
  ValueRate: 002ms19.322us / second
  Rate: 0.225789 / second
  Percentiles: 1%=001ms125.114us; 5%=001ms392.330us; 10%=001ms457.428us; 20%=002ms552.796us; 50%=003ms646.754us; 80%=004ms868.126us; 90%=004ms393.619us; 95%=014ms677.927us; 99%=053ms224.169us

[dlibenzi]

And more importantly, no aten::* counters which means the graph is fully fused into a single XLA computation.

[EugKar]

I was able to run with 500, too. It was compiling for about 40 minutes. Do you need log and tensors file from that run (3.2G for 64 iterations)?

[dlibenzi]

No need to.
Just wanted to check whether the model was doing the right things.
It seems OK.
1600 could be just too much, though, I find it strange does not even get to the graph generation.

I would be interested in the HLO graph for the 500 case, which took 40m to compile.
I can run it internally and see whether there are visible bottlenecks in the compilation process.

[EugKar]

This is a huge hlo graph for 500 case (its uncompressed size is >3GB!). tensors_500.zip
One observation is that this file's size was increasing after each of 64 iterations, while the graph was compiled only twice.

[dlibenzi]

Yeah, that's a big one! 😄
Thanks. We will check if there is something we can do to speed that up.
Can you confirm that if you run the same model (ie, the 500 long) twice, the second time compilation takes much shorter?

[EugKar]

Yes, the compilation takes less than a minute if I run the model for the second time.

[stale[bot]]

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