Fixing GRPO reward_func being a model with DeepSpeed ZeRO-3

[jamesbraza]

This PR enables GRPOTrainer's reward_func model to work with DeepSpeed ZeRO-3.

Running the new test with current main:

torchrun --nproc_per_node=1 -m pytest -sv tests/test_grpo_trainer.py::GRPOTrainerTester::test_training_with_deepspeed_zero3

tests/test_grpo_trainer.py:346:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
.venv/lib/python3.12/site-packages/transformers/trainer.py:2241: in train
    return inner_training_loop(
.venv/lib/python3.12/site-packages/transformers/trainer.py:2548: in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
.venv/lib/python3.12/site-packages/transformers/trainer.py:3692: in training_step
    inputs = self._prepare_inputs(inputs)
trl/extras/profiling.py:87: in wrapper
    return func(self, *args, **kwargs)
trl/trainer/grpo_trainer.py:692: in _prepare_inputs
    inputs = self._generate_and_score_completions(inputs)
trl/trainer/grpo_trainer.py:833: in _generate_and_score_completions
    rewards_per_func[:, i] = reward_func(**reward_inputs).logits[:, 0]  # Shape (B*G,)
.venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1736: in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
.venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1747: in _call_impl
    return forward_call(*args, **kwargs)
.venv/lib/python3.12/site-packages/transformers/models/qwen2/modeling_qwen2.py:945: in forward
    transformer_outputs = self.model(
.venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1736: in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
.venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1747: in _call_impl
    return forward_call(*args, **kwargs)
.venv/lib/python3.12/site-packages/transformers/models/qwen2/modeling_qwen2.py:535: in forward
    inputs_embeds = self.embed_tokens(input_ids)
.venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1736: in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
.venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1747: in _call_impl
    return forward_call(*args, **kwargs)
.venv/lib/python3.12/site-packages/torch/nn/modules/sparse.py:190: in forward
    return F.embedding(

...

E       RuntimeError: 'weight' must be 2-D

.venv/lib/python3.12/site-packages/torch/nn/functional.py:2551: RuntimeError

[jamesbraza]

Is there a standard solution for DeepSpeed tests in CI? I think this is the first integration test for DeepSpeed added to the repo.

In the future, we can expand it to cover #2871 and #2963.

[qgallouedec]

Why do you need to unwrap here? it seems like your loosing interest of using deepspeed

[jamesbraza]

Thanks for the question, yeah without this unwrap the unit test will crash with the stack trace in the PR description.

Try running torchrun --nproc_per_node=1 -m pytest -sv tests/test_grpo_trainer.py::GRPOTrainerTester::test_training_with_deepspeed_zero3 with this PR's unit test on main and then with the unwrap.