x

2026-03-28 14:07:42 +08:00
parent b563d85519
commit f6f11afda4
2 changed files with 20 additions and 1 deletions
@@ -161,6 +161,18 @@ def launch_model(model_id, config, model_path, gpu_count):
    env = os.environ.copy()
    env.update(config.get("env", {}))
    
+    # Add AMD GPU tensor parallel environment variables
+    if tp_size > 1:
+        env["NCCL_DEBUG"] = "INFO"
+        env["NCCL_SOCKET_IFNAME"] = "^lo,docker0"
+        env["NCCL_P2P_DISABLE"] = "0"
+        env["NCCL_SHM_DISABLE"] = "0"
+        env["NCCL_IB_DISABLE"] = "1"
+        env["NCCL_P2P_LEVEL"] = "SYS"
+        env["HIP_VISIBLE_DEVICES"] = os.getenv("HIP_VISIBLE_DEVICES", "0,1")
+        log("Added NCCL environment variables for AMD GPU tensor parallel")
+        log(f"HIP_VISIBLE_DEVICES: {env['HIP_VISIBLE_DEVICES']}")
+    
    # Launch vLLM
    log("Starting vLLM server...")
    try: