[ROCm][CI] Fix rlhf_nccl.py on ROCm (#46851)
Signed-off-by: charlifu <charlifu@amd.com>
Charlie Fu committed
Jun 26, 2026 at 15:41 UTC
6e2fb02fe5bcd3990c6ecd2663a5468c27d130f9
1 file changed
+16
-1
examples/rl/rlhf_nccl.py
+16
-1
@@ -29,6 +29,7 @@ causes unexpected behavior.
29
import os
30
31
import ray
32
+import torch
33
from ray.util.placement_group import placement_group
34
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
35
from transformers import AutoModelForCausalLM
@@ -39,12 +40,24 @@ from vllm.distributed.weight_transfer.nccl_engine import (
40
NCCLTrainerSendWeightsArgs,
41
NCCLWeightTransferEngine,
42
)
43
+from vllm.platforms import current_platform
44
from vllm.utils.network_utils import get_ip, get_open_port
45
46
MODEL_NAME = "facebook/opt-125m"
47
# MODEL_NAME = "inference-optimization/Qwen3-0.6B-W4A16-G128"
48
49
50
+def get_assigned_gpu():
51
+ """This is a temporary workaround for a runtime bug in RCCL on ROCm."""
52
+ if not current_platform.is_rocm():
53
+ return 0
54
+ assigned_gpu = int(ray.get_gpu_ids()[0])
55
+ os.environ.pop("CUDA_VISIBLE_DEVICES", None)
56
+ os.environ.pop("HIP_VISIBLE_DEVICES", None)
57
+ torch.accelerator.set_device_idx(assigned_gpu)
58
+ return assigned_gpu
59
+
60
+
61
class MyLLM(LLM):
62
"""Configure the vLLM worker for Ray placement group execution."""
63
@@ -58,9 +71,11 @@ class TrainModel:
71
"""Ray actor that wraps the training model on a dedicated GPU."""
72
73
def __init__(self, model_name: str):
74
+ assigned_gpu = get_assigned_gpu()
75
+
76
self.model = AutoModelForCausalLM.from_pretrained(
77
model_name,
63
- ).to("cuda:0")
78
+ ).to(f"cuda:{assigned_gpu}")
79
80
self.port = get_open_port()
81
self.master_address = get_ip()