process memory issue when training on TPU Pod v5e-64

I'm running into an issue where my processes run out of memory (I'm reaching 100% system memory utilization) when I try to train on a v5e-64 TPU pod. My konfig works fine on a v4-8.

the following appears before taking any train steps.
```
71885 Killed PYTHONPATH=$PYTHONPATH:. python -m kauldron.main --cfg=finetuning/pod_config_minload.py --cfg.workdir=/tmp/kauldron_oss/workdir
```

How to reproduce:
1. create a v5e-64 pod with v2-alpha-tpuv5-lite base image
2. follow tpu-pod-commander setup (https://github.com/young-geng/tpu_pod_commander/blob/main/README.md)
3. install gemma, jax[tpu], opencv-python-headless

```
with konfig.imports():
import jax.numpy as jnp
from gemma import gm
from kauldron import kd
import optax

import jax
jax.distributed.initialize()
```

and
```
def get_config():
batch_size = 64
import jax
if not jax.distributed.is_initialized():
jax.distributed.initialize()
```

If relevant, I'm loading my data from a Json file.