Compile failed to finish within 1 hour.

[kirpal-dev]

Hi All,

Im trying to run inference on a 6B parameter model with google colab TPU environment.

I tried to compile my model, but I a RuntimeError that the compile failed to finish within 1 hour. Is there a way to extend the max compile time, or how can I compile my model in less than 1 hour?

[/usr/local/lib/python3.9/dist-packages/jax/interpreters/xla.py](https://localhost:8080/#) in backend_compile(backend, built_c, options)
--> 343   return backend.compile(built_c, compile_options=options)
RuntimeError: Internal: Compile failed to finish within 1 hour.

jax 0.2.12
jaxlib 0.1.68

• can't upgrade due to other dependencies *

[jakevdp]

It's hard to give any anwer without seeing the code that cause the issue.

One guess: are you using a lot of Python control flow in your computation? If so, it can cause long compile times, and you should try re-expressing your function in terms of operations like lax.fori_loop. See https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#python-control-flow-jit for some discussion of this.