My training seems to be failing with the following message. This only occurs for me on bigger models (32B, 70B), and I don't encounter this in smaller models such as 1.5B. This seems to happen exactly after a batch of rewards being calculated. Does anyone have a guess at why this may be happening?
2025-07-21 05:50:50
0%| | 0/31 [00:00<?, ?it/s][WARNING|logging.py:328] 2025-07-21 05:50:50,487 >> `use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`.
2025-07-21 06:01:54
Traceback (most recent call last):
2025-07-21 06:01:54
File "/home/user_name/open-r1/src/open_r1/grpo.py", line 234, in <module>
2025-07-21 06:01:54
main(script_args, training_args, model_args)
2025-07-21 06:01:54
File "/home/user_name/open-r1/src/open_r1/grpo.py", line 183, in main
2025-07-21 06:01:54
train_result = trainer.train(resume_from_checkpoint=checkpoint)
2025-07-21 06:01:54
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2025-07-21 06:01:54
File "/home/user_name/miniconda3/envs/openr1/lib/python3.11/site-packages/transformers/trainer.py", line 2231, in train
2025-07-21 06:01:54
return inner_training_loop(
2025-07-21 06:01:54
^^^^^^^^^^^^^^^^^^^^
2025-07-21 06:01:54
File "/home/user_name/miniconda3/envs/openr1/lib/python3.11/site-packages/transformers/trainer.py", line 2555, in _inner_training_loop
2025-07-21 06:01:54
tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
2025-07-21 06:01:54
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2025-07-21 06:01:54
File "/home/user_name/miniconda3/envs/openr1/lib/python3.11/site-packages/transformers/trainer.py", line 3791, in training_step
2025-07-21 06:01:54
self.accelerator.backward(loss, **kwargs)
2025-07-21 06:01:54
File "/home/user_name/miniconda3/envs/openr1/lib/python3.11/site-packages/accelerate/accelerator.py", line 2321, in backward
2025-07-21 06:01:54
self.deepspeed_engine_wrapped.backward(loss, **kwargs)
2025-07-21 06:01:54
File "/home/user_name/miniconda3/envs/openr1/lib/python3.11/site-packages/accelerate/utils/deepspeed.py", line 266, in backward
2025-07-21 06:01:54
self.engine.backward(loss, **kwargs)
2025-07-21 06:01:54
File "/home/user_name/miniconda3/envs/openr1/lib/python3.11/site-packages/deepspeed/utils/nvtx.py", line 20, in wrapped_fn
2025-07-21 06:01:54
ret_val = func(*args, **kwargs)
2025-07-21 06:01:54
^^^^^^^^^^^^^^^^^^^^^
2025-07-21 06:01:54
File "/home/user_name/miniconda3/envs/openr1/lib/python3.11/site-packages/deepspeed/runtime/engine.py", line 2216, in backward
2025-07-21 06:01:54
self._do_optimizer_backward(loss, retain_graph)
2025-07-21 06:01:54
File "/home/user_name/miniconda3/envs/openr1/lib/python3.11/site-packages/deepspeed/runtime/engine.py", line 2162, in _do_optimizer_backward
2025-07-21 06:01:54
self.optimizer.backward(loss, retain_graph=retain_graph)
2025-07-21 06:01:54
File "/home/user_name/miniconda3/envs/openr1/lib/python3.11/site-packages/deepspeed/utils/nvtx.py", line 20, in wrapped_fn
2025-07-21 06:01:54
ret_val = func(*args, **kwargs)
2025-07-21 06:01:54
^^^^^^^^^^^^^^^^^^^^^
2025-07-21 06:01:54
File "/home/user_name/miniconda3/envs/openr1/lib/python3.11/site-packages/deepspeed/runtime/zero/stage3.py", line 2280, in backward
2025-07-21 06:01:54
self.loss_scaler.backward(loss.float(), retain_graph=retain_graph)
2025-07-21 06:01:54
File "/home/user_name/miniconda3/envs/openr1/lib/python3.11/site-packages/deepspeed/runtime/fp16/loss_scaler.py", line 63, in backward
2025-07-21 06:01:54
scaled_loss.backward(retain_graph=retain_graph)
2025-07-21 06:01:54
File "/home/user_name/miniconda3/envs/openr1/lib/python3.11/site-packages/torch/_tensor.py", line 626, in backward
2025-07-21 06:01:54
torch.autograd.backward(
2025-07-21 06:01:54
File "/home/user_name/miniconda3/envs/openr1/lib/python3.11/site-packages/torch/autograd/__init__.py", line 347, in backward
2025-07-21 06:01:54
_engine_run_backward(
2025-07-21 06:01:54
File "/home/user_name/miniconda3/envs/openr1/lib/python3.11/site-packages/torch/autograd/graph.py", line 823, in _engine_run_backward
2025-07-21 06:01:54
return Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
2025-07-21 06:01:54
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2025-07-21 06:01:54
RuntimeError: The size of tensor a (0) must match the size of tensor b (5120) at non-singleton dimension 1
2025-07-21 06:01:54
[rank0]: Traceback (most recent call last):
2025-07-21 06:01:54
[rank0]: File "/home/user_name/open-r1/src/open_r1/grpo.py", line 234, in <module>
2025-07-21 06:01:54
[rank0]: main(script_args, training_args, model_args)
2025-07-21 06:01:54
[rank0]: File "/home/user_name/open-r1/src/open_r1/grpo.py", line 183, in main
2025-07-21 06:01:54
[rank0]: train_result = trainer.train(resume_from_checkpoint=checkpoint)
2025-07-21 06:01:54
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2025-07-21 06:01:54
[rank0]: File "/home/user_name/miniconda3/envs/openr1/lib/python3.11/site-packages/transformers/trainer.py", line 2231, in train
2025-07-21 06:01:54
[rank0]: return inner_training_loop(
2025-07-21 06:01:54
[rank0]: ^^^^^^^^^^^^^^^^^^^^
2025-07-21 06:01:54
[rank0]: File "/home/user_name/miniconda3/envs/openr1/lib/python3.11/site-packages/transformers/trainer.py", line 2555, in _inner_training_loop
2025-07-21 06:01:54
[rank0]: tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
2025-07-21 06:01:54
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2025-07-21 06:01:54
[rank0]: File "/home/user_name/miniconda3/envs/openr1/lib/python3.11/site-packages/transformers/trainer.py", line 3791, in training_step
2025-07-21 06:01:54
[rank0]: self.accelerator.backward(loss, **kwargs)
2025-07-21 06:01:54
[rank0]: File "/home/user_name/miniconda3/envs/openr1/lib/python3.11/site-packages/accelerate/accelerator.py", line 2321, in backward
2025-07-21 06:01:54
[rank0]: self.deepspeed_engine_wrapped.backward(loss, **kwargs)
2025-07-21 06:01:54
[rank0]: File "/home/user_name/miniconda3/envs/openr1/lib/python3.11/site-packages/accelerate/utils/deepspeed.py", line 266, in backward
2025-07-21 06:01:54
[rank0]: self.engine.backward(loss, **kwargs)
2025-07-21 06:01:54
[rank0]: File "/home/user_name/miniconda3/envs/openr1/lib/python3.11/site-packages/deepspeed/utils/nvtx.py", line 20, in wrapped_fn
2025-07-21 06:01:54
[rank0]: ret_val = func(*args, **kwargs)
2025-07-21 06:01:54
[rank0]: ^^^^^^^^^^^^^^^^^^^^^
2025-07-21 06:01:54
[rank0]: File "/home/user_name/miniconda3/envs/openr1/lib/python3.11/site-packages/deepspeed/runtime/engine.py", line 2216, in backward
2025-07-21 06:01:54
[rank0]: self._do_optimizer_backward(loss, retain_graph)
2025-07-21 06:01:54
[rank0]: File "/home/user_name/miniconda3/envs/openr1/lib/python3.11/site-packages/deepspeed/runtime/engine.py", line 2162, in _do_optimizer_backward
2025-07-21 06:01:54
[rank0]: self.optimizer.backward(loss, retain_graph=retain_graph)
2025-07-21 06:01:54
[rank0]: File "/home/user_name/miniconda3/envs/openr1/lib/python3.11/site-packages/deepspeed/utils/nvtx.py", line 20, in wrapped_fn
2025-07-21 06:01:54
[rank0]: ret_val = func(*args, **kwargs)
2025-07-21 06:01:54
[rank0]: ^^^^^^^^^^^^^^^^^^^^^
2025-07-21 06:01:54
[rank0]: File "/home/user_name/miniconda3/envs/openr1/lib/python3.11/site-packages/deepspeed/runtime/zero/stage3.py", line 2280, in backward
2025-07-21 06:01:54
[rank0]: self.loss_scaler.backward(loss.float(), retain_graph=retain_graph)
2025-07-21 06:01:54
[rank0]: File "/home/user_name/miniconda3/envs/openr1/lib/python3.11/site-packages/deepspeed/runtime/fp16/loss_scaler.py", line 63, in backward
2025-07-21 06:01:54
[rank0]: scaled_loss.backward(retain_graph=retain_graph)
2025-07-21 06:01:54
[rank0]: File "/home/user_name/miniconda3/envs/openr1/lib/python3.11/site-packages/torch/_tensor.py", line 626, in backward
2025-07-21 06:01:54
[rank0]: torch.autograd.backward(
2025-07-21 06:01:54
[rank0]: File "/home/user_name/miniconda3/envs/openr1/lib/python3.11/site-packages/torch/autograd/__init__.py", line 347, in backward
2025-07-21 06:01:54
[rank0]: _engine_run_backward(
2025-07-21 06:01:54
[rank0]: File "/home/user_name/miniconda3/envs/openr1/lib/python3.11/site-packages/torch/autograd/graph.py", line 823, in _engine_run_backward
2025-07-21 06:01:54
[rank0]: return Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
2025-07-21 06:01:54
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2025-07-21 06:01:54
[rank0]: RuntimeError: The size of tensor a (0) must match the size of tensor b (5120) at non-singleton dimension 1
My training seems to be failing with the following message. This only occurs for me on bigger models (32B, 70B), and I don't encounter this in smaller models such as 1.5B. This seems to happen exactly after a batch of rewards being calculated. Does anyone have a guess at why this may be happening?