I think using accumulated_local_batch directly might lead to OOM when calling _get_per_token_logps, maybe splitting _generate_and_score_completions into two parts would be better.

Good point. I haven't encountered such OOM yet. In fact, here, we do the forward pass with no grad, so maybe it prevents this problem? I'd have to test it with a large gradient accumulation step to confirm.

I tested it and it will be OOM when max_completion_length is large.

Using the same setting, you don't get OOM with main?

I train 1.5B model on two gpus with GAS=8, BS=16 and max_completion_length=2048. Everything is fine with main.

I think using accumulated_local_batch directly might lead to OOM when calling _get_per_token_logps, maybe splitting _generate_and_score_completions into two parts would be better.

In the first part, use vLLM to generate all the completion_ids and calculate the advantages of all these completions. In the second part, compute old_per_token_logps and ref_per_token_logps with original per_device_batchsize (perhaps old_per_token_logps can be obtained directly in the vLLM generation phase).

One benefit of this is that instead of (per_device_batchsize * num_processes) % num_generations == 0, we need (per_device_batchsize * num_processes * gradient_accumulation_steps) % num_generations == 0. This solves (#3017) and (#3288) without introducing additional parameters in config, just adjust the per_device_batchsize and gradient_accumulation_steps to adapt to various scenarios.

A good model to test would be any of the DeepSeek distilled ones: they easily produce 32k tokens per prompt, so you'll know quickly if improvements to the generation/scoring still OOM or not

That's a very good point @I-l-l-I, I didn't even realised that with this PR, we're releasing this requirement (per_device_batchsize * num_processes) % num_generations == 0

This should avoid OOM: eedaab5

Now the memory usage is the same as with main:

from datasets import load_dataset
from trl import GRPOTrainer, GRPOConfig
import random

dataset = load_dataset("trl-lib/tldr", split="train[:200]")


def reward_random(completions, **kwargs):
    return [random.random() for _ in range(len(completions))]

trainer = GRPOTrainer(
    model="Qwen/Qwen2.5-1.5B",
    reward_funcs=reward_random,
    train_dataset=dataset,
    args=GRPOConfig(
        per_device_train_batch_size=16,
        gradient_accumulation_steps=8,
        max_completion_length=2048,
        use_vllm=True,
        gradient_checkpointing=True,
        bf16=True,
    ),
)
trainer.train()

2 GPUs

CUDA_VISIBLE_DEVICES=0,1 accelerate launch --config_file examples/accelerate_configs/deepspeed_zero3.yaml --num_processes 2 3283.py 

Side note: with this script, for example, training time is reduced from 30 min to 10 min. 🤯

Probably not for this PR, but would it be worth having a per_device_mega_batch_size and completely decoupling the gradient accumulation from the size of the big generation batch? Some batch shuffling would be required to ensure that the minibatches vary from one iteration to the next, but I don't think it would be a huge change.

@edbeeching Do you mean sth like this in simple terms?

maga_batch_size = 524288  # ~0.5M, in number of tokens if T = 1024
B = 64 # micro batch size
T = 1024 
assert maga_batch_size % (B * T * ddp_world_size) == 0
grad_accum_steps = maga_batch_size // (B * T * ddp_world_size)

where we later do;

for micro_step in range(grad_accum_steps):
        x, y = train_loader.next_batch()
        x, y = x.to(device), y.to(device)
        logits, loss = model(x, y)
        loss = loss / grad_accum_steps
        loss_accum += loss.detach()
        loss.backward()

I think what Ed is suggesting is to align the GRPOTrainer partially with our other RL trainers where one defines a large effective batch size of N_prompts x N_generations and then performing K steps of mini-batch optimisation (i.e. the num_minibatches arg in the PPOTrainer). This is also what other RL frameworks like verl do.

In that context, gradient accumulation is not involved in defining the large effective batch size, but I realise it is not easy to force the transformers.Trainer logic to do mini-batch optimisation and should be done in a separate PR (if at all)

Generally for training (independent of this PR) per_device_eval_batch_size is 2*per_device_train_batch_size

As we are in no_grad mode you can probably set the per_device_train_batch_size here to be 2*self.args.per_device_train_batch_size as there will be no memory used for the storage of activations etc.

It would be worth validating in a scenario with significant memory pressure, so do not worry about it if you want to get the PR merged.