Issue: Inference Fails with Fine-Tuned Gemma3-4B (Kauldron + DPO), and OOM Despite 4×A100 40GB GPUs

Context:
I fine-tuned the Gemma 3-4B model using the official [examples/dpo.py](https://github.com/google-deepmind/gemma/blob/main/examples/dpo.py) script via Kauldron. The DPO training ran successfully and saved a checkpoint at:

/home/jupyter/models/my_dpo_gemma3_4b/ckpt_1000
The checkpoint folder contains:

ckpt_1000/
├── _CHECKPOINT_METADATA
├── default/
│ ├── _METADATA
│ ├── manifest.ocdbt
│ └── ocdbt.process_0/
└── ...
### **Problem 1: Inference using fine-tuned model fails with AnchoredPolicy**
I’m trying to perform inference using the fine-tuned checkpoint like this:

from gemma import gm

model = gm.nn.AnchoredPolicy(
policy=gm.nn.Gemma3_4B(tokens="batch.tokens", text_only=True),
)

params = gm.ckpts.load_params(
path="/home/jupyter/models/my_dpo_gemma3_4b/ckpt_1000",
text_only=True,
)

sampler = gm.text.ChatSampler(
model=model,
params=params,
multi_turn=False,
)

response = sampler.chat("I have been having this headache for the last year.")
This fails with the error:

_AttributeError: "AnchoredPolicy" object has no attribute "INFO". If "INFO" is defined in '.setup()', remember these fields are only accessible from inside 'init' or 'apply'._
❓Question:
**Is this the correct way to perform inferencing on a dpo tuned Gemma-3-4B model?**

### **Problem 2: OOM on A100 40GB ×4 (Kauldron + Gemma-3-4B)**
During training, I encountered:

RESOURCE_EXHAUSTED: Out of memory while trying to allocate 32480133448 bytes.
W ... bfc_allocator.cc:501] Allocator (GPU_1_bfc) ran out of memory trying to allocate 30.25GiB ...
Despite using:

**4× A100 40GB (160GB total GPU memory)**

The only way to prevent OOM was to set:
batch_size = 4
max_seq_length = 128
**These are very conservative settings given the hardware and model size**.

❓Question:
1. Efficient memory usage on multi-GPU with Gemma-3-4B

2. Guidelines or best practices for avoiding static graph explosion or excessive JAX compilation memory without reducing the batch size/max sequence length

Thanks so much for the amazing work on Gemma !