⏯️ Fix: handle None inputs when resuming GRPO Trainer from checkpoint

[PenutChen]

What does this PR do?

This PR refactors the _prepare_inputs method to ensure it never returns None when resuming training from checkpoints in GRPOTrainer.

Previously, if self._buffered_inputs[...] was None during the resume process, and self.state.global_step % self.num_iterations != 0, the method would return None. This caused issues downstream where non-None inputs were expected.

To fix this, the logic has been updated so that if the buffered input is None, it always falls back to generating new inputs via _generate_and_score_completions(inputs), regardless of the step count. This ensures stability and correctness when resuming from checkpoints.

In addition, the modulo expression self._step % self.args.gradient_accumulation_steps has been refactored into a dedicated variable accumulation_index to improve readability and reduce redundancy.

---

Fixes

No corresponding issue was filed, but this change addresses a potential silent failure when using resume_from_checkpoint with GRPOTrainer.

---

Motivation and context

Users resuming training from checkpoints may encounter None inputs in _prepare_inputs, leading to errors in the training loop. This fix ensures robustness by avoiding returning None in any case and improves the maintainability of the code through refactoring.

---

Before submitting

• I have read the [contributor guideline](https://github.com/huggingface/trl/blob/main/CONTRIBUTING.md#create-a-pull-request).
• I have verified that existing tests pass.
• I have run make precommit to ensure code style consistency.
• I have verified the change does not introduce breaking behavior.

---

Who can review?

Anyone familiar with GRPOTrainer, buffered input logic, or checkpoint resume behavior in TRL.

[qgallouedec]

Thanks!! Can share a simple piece of code that would fail without your fix?

[PenutChen]

@qgallouedec here's the sample code:

from datasets import load_dataset
from trl import GRPOTrainer, GRPOConfig

def reward_fn(completions, **_):
    return [1.0 for _ in completions]

# Normal training
trainer = GRPOTrainer(
    model="facebook/opt-125m",
    args=GRPOConfig(
        output_dir="save/test",
        num_generations=2,
        per_device_train_batch_size=2,
        num_iterations=4,
        save_steps=1,
        max_steps=10,
        max_prompt_length=1,
        max_completion_length=1,
    ),
    reward_funcs=reward_fn,
    train_dataset=load_dataset("trl-lib/tldr", split="train"),
)
trainer.train()

# Simulating a fresh new trainer instance after interruption
trainer = GRPOTrainer(
    model="facebook/opt-125m",
    args=GRPOConfig(
        output_dir="save/test",
        num_generations=2,
        per_device_train_batch_size=2,
        num_iterations=4,
        save_steps=1,
        max_steps=10,
        max_prompt_length=1,
        max_completion_length=1,
    ),
    reward_funcs=reward_fn,
    train_dataset=load_dataset("trl-lib/tldr", split="train"),
)

# Resume from checkpoint at step which is not divisible by num_iterations
trainer.train(resume_from_checkpoint="save/test/checkpoint-6")

error traceback:

Traceback (most recent call last):
  File "your_script.py", line 41, in <module>
    trainer.train(resume_from_checkpoint="save/test/checkpoint-6")
  File ".../transformers/trainer.py", line 2245, in train
    return inner_training_loop(
  File ".../transformers/trainer.py", line 2556, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
  File ".../transformers/trainer.py", line 3718, in training_step
    loss = self.compute_loss(model, inputs, num_items_in_batch=num_items_in_batch)
  File ".../trl/extras/profiling.py", line 87, in wrapper
    return func(self, *args, **kwargs)
  File ".../trl/trainer/grpo_trainer.py", line 905, in compute_loss
    prompt_ids, prompt_mask = inputs["prompt_ids"], inputs["prompt_mask"]
                              ~~~~~~^^^^^^^^^^^^^^
TypeError: 'NoneType' object is not subscriptable

[qgallouedec]

Nice, thank you @PenutChen! In this case, the results won't be exactly the same as if we hadn't interrupted the training (we would have to save and load this buffer), but that's not a big deal.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.