Lower jit times
Open
enhancement
- Dominant language
- Jupyter Notebook
- Stars
- 3.2k
- Forks
- 349
- PR merge metrics
- No merged PRs in 30d
Description
Jit times for training environments can be quite high, up to 6-7 minutes for free colabs communicating with TPUs. Some options here are:
- Code reduction and consolidation lowers tracing / XLA overhead
- Try to store HLO for different environment batch sizes to lower jit times: https://github.com/google/jax/pull/7997
- Bug JAX team to use Brax as a test case for lowering jit time
Contributor guide
Assessment
This issue has not been assessed yet.