Fine-tunning Gemma 3 on financial dataset

Introduction

Recently, I decided to dive into Group Relative Policy Optimization (GRPO) to enhance the reasoning capabilities of Large Language Models. Using the Tunix and Flax/NNX libraries, I set out to train the Gemma-3–1B-IT model on the GSM8K math reasoning benchmark, using this example as base.

My hardware of choice? A Google Cloud TPU v5e, which are available at Google Colab's free tier. While TPUs are incredibly fast, training RL algorithms like GRPO are notoriously memory-intensive. In this post, I want to share my journey, the Out-of-Memory (OOM) hurdles I faced, and the exact parameter tweaks that finally led to a stable and improving model.

The First Hurdle: Out of Memory (OOM) on TPUv5e

GRPO is much more memory-efficient than standard Proximal Policy Optimization (PPO) because it eliminates the need for a separate value function model. However, it still generates multiple responses for a given prompt during training to calculate a relative advantage. This “group” generation quickly ate up my TPUv5e’s High Bandwidth Memory (HBM).

My initial runs crashed almost immediately with OOM errors. To fix this, I had to aggressively tune several generation and batching parameters:

1. Micro Batch Size (`TRAIN_MICRO_BATCH_SIZE`): I had to drop this all the way down to 1. Pushing too many prompts through the model at once during the reinforcement learning rollout phase was the biggest culprit.

2. Group Generations (`NUM_GENERATIONS`): GRPO (Group Relative Policy Optimization) relies on comparing the rewards of multiple generations against each other to calculate an advantage. If you only generate one response, there is no “group” to compare against, the standard deviation is zero, and you get division-by-zero (NaN) errors. You must have a minimum of num_generations=2so I used this value.

3. Prompt & Generation Lengths: Language models allocate memory based on sequence length (kv cache). I capped the `MAX_PROMPT_LENGTH` at 512 and `TOTAL_GENERATION_STEPS` at 768.

By carefully balancing these four parameters, the TPU memory usage stabilized, and the training loop successfully started spinning (finally!).

Defining the Reward

The beauty of GRPO in this context is the reward function. I didn’t just reward the final correct answer; it's important that the reward also encourage the correct format:

• Exact Format Match: +3.0 points for properly utilizing `<reasoning>` and `<answer>` tags.
• Partial Format: Small incremental points for including the right tags.
• Correct Answer: +3.0 points for extracting the exact number.
• Close Answers: Ratios between 0.9 and 1.1 received partial credit to encourage the model when it was on the right track.

First Results

Before training, the baseline Gemma-3–1B-IT struggled with the strict formatting and reasoning constraints, hovering around a 37.5% accuracy on my test batches with a 0% exact format match.

Get Fernanda Wanderley’s stories in your inbox

Join Medium for free to get updates from this writer.

Remember me for faster sign in

After dialing in the parameters and letting the GRPO cluster run through the batches, the improvements were clear. Although the accuracy didn't improved, the model learned to consistently wrap its thoughts in `<reasoning>` tags and output the final numerical value cleanly in `<answer>` tags, heavily boosting its reliability on grade-school math problems.

Before fine tuning the model get lost at the reasoning

After, even tough the answer is incorrect the model finishes the reasoning

Changing datasets

After the promising results with GSM8K I decided to try something more complex. Fine tuning Gemma-3–1B for financial reasoning using the dataset FinQA. FinQA (Financial Question Answering) is a specialized dataset designed to evaluate a model’s numerical reasoning over real-world financial data. It features complex questions based on earnings reports, requiring models to read both dense text and financial tables, extract the correct figures, and perform multi-step arithmetic to find the answer. While it’s a gold standard for financial AI, its massive context requirements make it incredibly resource-intensive. Since it was already challenging using the TPU v5e in the first experiment, I upgraded my colab account to use the v6e.

I immediately hit Out-Of-Memory (OOM) errors due to the complexity of data and massive context lengths meant.

I'm going to have nightmares with this error message

I tried a couple of initial fixes:

• Lowering TRAIN_MICRO_BATCH_SIZE to 1.
• Setting offload_to_cpu=True to move the reference model out of the TPU's memory.
• Aggressively tuning MAX_PROMPT_LENGTH and TOTAL_GENERATION_STEPS.
• Truncating the dataset to reduce the need of a big context window.
• Selecting only samples with a maximum of 2000 characters.

Eventually, I realized that long-context tables + GRPO + 1B model + Single TPU v6e is a mathematically impossible fit.

Conclusion

I’d be lying if I said this process wasn’t incredibly frustrating. Hitting constant Out-Of-Memory walls and being forced to abandon (for now) the rich, complex reasoning of the FinQA dataset was a tough pill to swallow. Fine-tuning LLMs with RL techniques like GRPO is incredibly powerful, but the infrastructure tax — especially on a single TPU v6e — is brutal.

That said, this isn’t the end of the road. Now that the training loop and reward functions are validated and stable, I am hopeful and excited to run these experiments again. The next step is clear: throw more compute at it. With a multi-core TPU setup (like a v5e-8), I’ll finally have the memory capacity to tackle FinQA’s massive context and see what Gemma 3 can really do with complex financial reasoning! Stay tuned for more!

“Google Cloud credits are provided for this project.” #TPUSprint