Generations after few training steps

[iraj465]

I see after 6-7 steps, the generations fails in the sense the probability scores becomes very low (0 actually) with bf16 precision enabled and gives this error:

Example:

0:   with torch.enable_grad(), device_autocast_ctx, torch.cpu.amp.autocast(**ctx.cpu_autocast_kwargs):  # type: ignore[attr-defined]
  0%|          | 4/962 [07:33<27:12:27, 102.24s/it][rank1]: Traceback (most recent call last):
0: [rank1]:   File "/home/sapmajum/codebases/open-r1-multimodal/src/open_r1/grpo.py", line 183, in <module>
0: [rank1]:     main(script_args, training_args, model_args)
0: [rank1]:   File "/home/sapmajum/codebases/open-r1-multimodal/src/open_r1/grpo.py", line 172, in main
0: [rank1]:     trainer.train()
0: [rank1]:   File "/home/sapmajum/.local/lib/python3.10/site-packages/transformers/trainer.py", line 2171, in train
0: [rank1]:     return inner_training_loop(
0: [rank1]:   File "/home/sapmajum/.local/lib/python3.10/site-packages/transformers/trainer.py", line 2531, in _inner_training_loop
0: [rank1]:     tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
0: [rank1]:   File "/home/sapmajum/.local/lib/python3.10/site-packages/transformers/trainer.py", line 3678, in training_step
0: [rank1]:     loss = self.compute_loss(model, inputs, num_items_in_batch=num_items_in_batch)
0: [rank1]:   File "/home/sapmajum/codebases/open-r1-multimodal/src/open_r1/trainer/grpo_trainer.py", line 369, in compute_loss
0: [rank1]:     prompt_completion_ids = unwrapped_model.generate(**prompt_inputs, generation_config=self.generation_config)
0: [rank1]:   File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
0: [rank1]:     return func(*args, **kwargs)
0: [rank1]:   File "/home/sapmajum/.local/lib/python3.10/site-packages/transformers/generation/utils.py", line 2255, in generate
0: [rank1]:     result = self._sample(
0: [rank1]:   File "/home/sapmajum/.local/lib/python3.10/site-packages/transformers/generation/utils.py", line 3300, in _sample
0: [rank1]:     next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
0: [rank1]: RuntimeError: probability tensor contains either `inf`, `nan` or element < 0

Libraries:

transformers  == 4.48.0
trl == 0.15.0.dev0
bitsandbytes  == 0.43.2.dev0
torch ==  2.4.0
accelerate  == 1.3.0
deepspeed   ==  0.14.4

Please suggest how to resolve this. Thanks!