-
### Description
I am trying to fine-tune Gemma 2 on TPU and got the following error:
```
Traceback (most recent call last):
File "/usr/local/lib/python3.10/dist-packages/jax/_src/compiler.py", l…
-
### Description
```python
import functools
import jax
from jax.experimental import pallas as pl
import jax.numpy as jnp
@functools.partial(
pl.pallas_call,
out_shape=jax.ShapeDtypeSt…
-
### Description
Using [accelerators.md](https://github.com/ray-project/ray/blob/master/python/ray/util/accelerators/accelerators.md) as a reference point, we can graduate TPUs from custom resources (…
-
### Your current environment
```text
The output of `python collect_env.py`
Collecting environment information...
WARNING:root:libtpu.so and TPU device found. Setting PJRT_DEVICE=TPU.
INFO 11-04…
Bihan updated
18 hours ago
-
### Description
Hi JAX team,
In the past two days, I've been using GCP's queued-resources to create spot TPU v4-256/v4-64, and then running the following Python script.
```python
import jax
j…
-
How to solve the 'RESOURCE_EXHAUSTED' error when loading 'gemma2_instruct_2b_en' (the script is from kaggle and runs on colab with TPU)? Errors shown following:
: 2.5.1+c…
-
### Description
With `0.4.35` release, the flash attention kernel will hit compilation OOM for long sequence length inputs. It failed during compiling the reference attention implementation for cos…
-
GCP catalog collection [started to fail](https://github.com/dstackai/gpuhunt/actions/runs/10877908671/job/30180003847) due to errors like `google.api_core.exceptions.PermissionDenied: 403 Location eur…