fix: fix sft-llama3.1-70b-8n8g-tp4pp2-long-megatron

[yuki-97]

Summary

Restore torch.distributed.init_process_group("nccl") without device_id=. device_id= was added in #2355 and empirically slows down sft-llama3.1-70b-8n8g-tp4pp2-long-megatron training. Removing it restores throughput.

torch.cuda.set_device(local_rank) is already called in MegatronPolicyWorker.__init__ before setup_distributed(), so device binding is unaffected.

Yellow: before #2355. Blue: after #2355, before this PR. Red: this PR.

Tests

Validated these tests passed:

• sft-llama3.1-70b-8n8g-tp4pp2-long-megatron
• grpo-llama3.2-1b-instruct-1n8g-megatron_generation (added in 2355)
• grpo-llama3.2-1b-instruct-2n8g-megatron_generation-noncolocated (added in 2355)
• grpo-nanov3-30BA3B-2n8g-megatron_generation (added in 2355)

This test has unrelated error:

• grpo-nanov3-30BA3B-2n8g-megatron_generation-async-gym (added in 2355)

Concrete differences from PyTorch/NCCL docs

Passing device_id= to init_process_group has two documented effects and one known runtime regression:

1. Eager NCCL communicator init: ncclCommInit* runs inside init_process_group() instead of lazily on the first collective. Only benefit is earlier NCCL error reporting.
2. Sub-groups use ncclCommSplit instead of ncclCommInitRankConfig. However, PyTorch does not set splitShare=1 on the split config, so:
   • The intended memory/resource sharing is turned off — split sub-comms still allocate independent resources.
   • Repeated new_group() calls can OOM (see pytorch#129865).
3. Known regression in torch 2.7+: device_id= causes NCCL to randomly hang during communications, tracked in pytorch#153960. torch 2.6 was fine. Workaround is to omit device_id=.

Under the current PyTorch implementation there is no measurable benefit to passing device_id= for training workloads, and it triggers the hang regression.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[yuki-97]

/ok to test 855a7f8