GRPO unbalanced-memory

[mdy666]

Reproduction

the memory in each rank(0-6) is not same, and i find when the train steps increase, the memory will increase much
step 0 use the origin code

Then i write a efficient grpo loss kernel by triton。
step 0

step 5

step 20

System Info

trl = 0.14.0
torch = 2.5.1+cuda12.4
vllm = 0.7.1

Checklist

• I have checked that my issue isn't already filed (see open issues)
• I have included my system information
• Any code provided is minimal, complete, and reproducible (more on MREs)
• Any code provided is properly formatted in code blocks, (no screenshot, more on code blocks)
• Any traceback provided is complete

[mdy666]

I do lots of exps, I finally find it's caused by accuracy_reward
if set

def accuracy_reward(completions, solution, **kwargs):
    return [0.5 for _ in range(len(completions))]

with the steps increase, the memory in my code (use triton write the grpo loss) does not increase. But I don't know why this would happen?

[nomadlx]

Is it convenient for you to provide your implementation? Currently, on an 8*80G graphics card, I can only run a 7B-sized model, and the maximum length for completion can only be set to 1024. This length is far from sufficient for an inference model. I don't quite understand why it leads to such high VRAM consumption.

[mdy666]

Is it convenient for you to provide your implementation? Currently, on an 8*80G graphics card, I can only run a 7B-sized model, and the maximum length for completion can only be set to 1024. This length is far from sufficient for an inference model. I don't quite understand why it leads to such high VRAM consumption.

Sure! Here is my code. please click a star, thank you.
https://github.com/mdy666/mdy_triton/tree/main/others/grpo

[leonardtang]

I do lots of exps, I finally find it's caused by accuracy_reward if set

def accuracy_reward(completions, solution, **kwargs):
return [0.5 for _ in range(len(completions))]
with the steps increase, the memory in my code (use triton write the grpo loss) does not increase. But I don't know why this would happen?

You're saying that accuracy_reward is eating up GPU memory?

[mdy666]

I do lots of exps, I finally find it's caused by accuracy_reward if set
def accuracy_reward(completions, solution, **kwargs):
return [0.5 for _ in range(len(completions))]
with the steps increase, the memory in my code (use triton write the grpo loss) does not increase. But I don't know why this would happen?

You're saying that accuracy_reward is eating up GPU memory?

this sound incredible，but my exps tell me it caused by the math_verfiy package. The GPU memory will increase from 60G to 120G. And it increase 20 G each time, it equals 2 * bs * num_generations * seq_len * vocab_size, I think the logits dose not allow the memory