Issue with unit tests on NVIdia A100 (GPU)

Hi everyone.

I see the issue when run unit tests on NVidia A100 (GPU). Here is the [link](https://github.com/NVIDIA/JAX-Toolbox/actions/runs/9099672951/job/25013673011?pr=590#step:7:593) for more details.

Briefly:
```
=========================== short test summary info ============================
FAILED opt/gemma/gemma/layers_test.py::EinsumTest::test_rmsnorm0 - AssertionE...
FAILED opt/gemma/gemma/modules_test.py::FeedForwardTest::test_ffw0 - Assertio...
FAILED opt/gemma/gemma/positional_embeddings_test.py::PositionalEmbeddingsTest::test_adds_positional_embeddings0
FAILED opt/gemma/gemma/sampler_test.py::SamplerTest::test_forward_equivalence
================== 4 failed, 12 passed, 2 warnings in 26.55s ===================
```

The first 3 is similar to issues on V100 (#32), but the last one:
4. test_forward_equivalence [link](https://github.com/NVIDIA/JAX-Toolbox/actions/runs/9099672951/job/25013673011?pr=590#step:7:574). Can you relax the tolerance when run on GPUs?