Running GRPO training on Gemma 2 2B-IT

[cool-RR]

Hi guys!

My name is Ram Rachum, and I'm a research intern at CHAI, UCB. I'm working on research on RL explainability, and I want to use Tunix as case study.

I'm trying to replicate your GPRO training of Gemma 2 2B-IT. I ran your grpo_demo.ipynb notebook on a v5litepod-4 machine on GCP. I had to stop it because it was filling the entire browser with output, and Jupyter isn't very good in handling that.

Questions:

1. I understand that you ran this on v5litepod-8. How long did running the entire notebook take?
2. Do you think this would run on v5litepod-4? I'm still waiting for GCP staff to allow me to use v5litepod-8. If you have any way to help with, or could connect me to a program that grants GCP compute credits, that would be helpful.

I already ran a simplified version of grpo_demo.ipynb locally on my RTX 5080, but I could not get it to improve the GSM8K score of any LLM.

[lc5211]

Hi Ram, you should be able to run the notebook on v5litepod-4. The notebook defaults to have a mesh as [(1, 4), ("fsdp", "tp")] which indicates that it uses only 4 v5e chips. We will update the descriptions in the notebook to reflect that. Sorry for the confusion.

[cool-RR]

Thank you!

[lc5211]

The notebook should be finished in around 1 hour on v5litepod-4.

[cool-RR]

Thank you!

[cool-RR]

I've attempted to reproduce your results, and ran into the following bugs:

• "RuntimeError: cannot enter context" in grpo_demo.ipynb #686
• grpo_demo.ipynb: Excessive output #687
• grpo_demo.ipynb: Training reduces accuracy from 53.0% to 49.0% #688

I hope you can help me with these.