Silences NCCL warning

[finbarrtimbers]

Previously, we would get this warning:

2025-10-03T20:49:10.479Z (PolicyTrainerRayProcess pid=1720) [rank0]:[W1003 20:49:10.123654124
  ProcessGroupNCCL.cpp:4715] [PG ID 0 PG GUID 0 Rank 0]  using GPU 0 as device used by this process is currently
  unknown. This can potentially cause a hang if this rank to GPU mapping is incorrect. You can pecify device_id in
  init_process_group() to force use of a particular device.

By specifying the device_id in init_process_group, we don't get it any more.

Example run: Beaker

[cursor]

You have run out of free Bugbot PR reviews for this billing cycle. This will reset on November 1.

To receive reviews on all of your PRs, visit the Cursor dashboard to activate Pro and start your 14-day free trial.

[saurabh111233212]

nice ty

[finbarrtimbers]

/gemini what do you think?

[gemini-code-assist]

Thanks for asking! I've reviewed the changes in this pull request.

The addition of device_id=torch.device("cuda", self.local_rank) to torch.distributed.init_process_group in open_instruct/grpo_fast.py is a good and necessary change. As you've noted in the PR description, explicitly setting the device_id helps to silence the NCCL warning regarding unknown device usage.

This is a best practice in distributed training with PyTorch and NCCL, as it ensures that each process is correctly assigned to its intended GPU, preventing potential issues like device mismatches or hangs that the warning indicates. It also improves the clarity and robustness of the distributed setup.

Overall, this is a solid improvement that addresses a common warning in distributed PyTorch applications.