📄 Add GRPO batch size note in docs (#2672)

* add note for OOM error * update note * Apply suggestions from code review --------- Co-authored-by: Quentin Gallouédec <[email protected]>
huggingface · Jan 30, 2025 · 9ac8d97 · 9ac8d97
1 parent 094d51b
commit 9ac8d97
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/docs/source/grpo_trainer.md b/docs/source/grpo_trainer.md
@@ -24,6 +24,8 @@ This example demonstrates how to train a model using the GRPO method. We train a
 ></iframe>
 
 Below is the script to train the model.
+Note that the input tensor for the forward pass has a size of `num_generations * per_device_train_batch_size` because GRPO generates `num_generations` completions for each prompt in the batch. Adjusting these values appropriately can help prevent OOM errors.
+Consequently, the effective train batch size is `num_generations * per_device_train_batch_size * gradient_accumulation_steps`.
 
 ```python
 # train_grpo.py