[rocm] rlhf

[Delaunay]

* 1 x E1003 23:51:14.134000 139750057968064 torch/distributed/elastic/multiprocessing/api.py:833] failed (exitcode: 1) local_rank: 7 (pid: 776026) of binary: /home/mila/rocm/results/venv/torch/bin/python
    	| [rank5]: Traceback (most recent call last):
    	| [rank5]:   File "/home/mila/milabench/benchmarks/rlhf/main.py", line 133, in <module>
    	| [rank5]: 	main()
    	| [rank5]:   File "/home/mila/milabench/benchmarks/rlhf/main.py", line 124, in main
    	| [rank5]: 	trainer.train()
    	| [rank5]:   File "/home/mila/rocm/results/venv/torch/lib/python3.10/site-packages/trl/trainer/ppov2_trainer.py", line 309, in train
    	| [rank5]: 	query_responses, logitss = batch_generation(
    	| [rank5]:   File "/home/mila/rocm/results/venv/torch/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    	| [rank5]: 	return func(*args, **kwargs)
    	| [rank5]:   File "/home/mila/rocm/results/venv/torch/lib/python3.10/site-packages/trl/trainer/utils.py", line 1176, in batch_generation
    	| [rank5]: 	query_response, logits = generate(
    	| [rank5]:   File "/home/mila/rocm/results/venv/torch/lib/python3.10/site-packages/trl/trainer/utils.py", line 1151, in generate
    	| [rank5]: 	output = lm_backbone.generate(
    	| [rank5]:   File "/home/mila/rocm/results/venv/torch/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    	| [rank5]: 	return func(*args, **kwargs)
    	| [rank5]:   File "/home/mila/rocm/results/venv/torch/lib/python3.10/site-packages/transformers/generation/utils.py", line 2024, in generate
    	| [rank5]: 	result = self._sample(
    	| [rank5]:   File "/home/mila/rocm/results/venv/torch/lib/python3.10/site-packages/transformers/generation/utils.py", line 3020, in _sample
    	| [rank5]: 	next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
    	| [rank5]: RuntimeError: probability tensor contains either `inf`, `nan` or element < 0