فهرست منبع

修改设备名称

lxylxy123321 1 هفته پیش
والد
کامیت
57275945e4
3فایلهای تغییر یافته به همراه8 افزوده شده و 11 حذف شده
  1. 0 2
      backend/app/core/remote_executor.py
  2. 2 2
      backend/app/engines/remote_train.py
  3. 6 7
      backend/app/engines/text_engine.py

+ 0 - 2
backend/app/core/remote_executor.py

@@ -158,7 +158,6 @@ def run_training_remote(
         f"docker exec "
         f"-e MACA_MPS_MODE=1 "
         f"-e METAX_VISIBLE_DEVICES=2,3 "
-        f"-e CUDA_VISIBLE_DEVICES=2,3 "
         f"-w {settings.compute_node_workdir} "
         f"{settings.compute_node_docker_container} "
         f"bash -c '"
@@ -222,7 +221,6 @@ def run_inference_remote(
         f"docker exec "
         f"-e MACA_MPS_MODE=1 "
         f"-e METAX_VISIBLE_DEVICES=2,3 "
-        f"-e CUDA_VISIBLE_DEVICES=2,3 "
         f"{settings.compute_node_docker_container} "
         f"{settings.compute_node_python} -c \""
         "import asyncio, json; "

+ 2 - 2
backend/app/engines/remote_train.py

@@ -18,9 +18,9 @@ os.environ["FLASH_ATTENTION_ENABLED"] = "0"
 os.environ["PT2_COMPILE"] = "0"
 os.environ["TORCHINDUCTOR_MAX_WORKERS"] = "1"
 # 限制训练只用 GPU 2 和 3(GPU 0/1 被 VLLM 占用)
-# 沐曦 GPU 优先用 METAX_VISIBLE_DEVICES,同时设 CUDA_VISIBLE_DEVICES 兜底
+# 沐曦 MPS 模式下 CUDA_VISIBLE_DEVICES 可能干扰设备映射,
+# 只设 METAX_VISIBLE_DEVICES,device_map 里用物理 GPU 号手动指定。
 os.environ["METAX_VISIBLE_DEVICES"] = "2,3"
-os.environ["CUDA_VISIBLE_DEVICES"] = "2,3"
 # 启用 MPS 多进程服务,允许与 VLLM 共享 GPU
 os.environ["MACA_MPS_MODE"] = "1"
 

+ 6 - 7
backend/app/engines/text_engine.py

@@ -9,9 +9,9 @@ os.environ["TORCH_FLASH_ATTN"] = "0"
 os.environ["PT2_COMPILE"] = "0"
 os.environ["TORCHINDUCTOR_MAX_WORKERS"] = "1"
 # 限制训练只用 GPU 2 和 3(GPU 0/1 被 VLLM 占用)
-# 沐曦 GPU 优先用 METAX_VISIBLE_DEVICES,同时设 CUDA_VISIBLE_DEVICES 兜底
+# 沐曦 MPS 模式下 CUDA_VISIBLE_DEVICES 可能干扰设备映射,
+# 只设 METAX_VISIBLE_DEVICES,device_map 里用物理 GPU 号手动指定。
 os.environ["METAX_VISIBLE_DEVICES"] = "2,3"
-os.environ["CUDA_VISIBLE_DEVICES"] = "2,3"
 # 启用 MPS 多进程服务,允许与 VLLM 共享 GPU
 os.environ["MACA_MPS_MODE"] = "1"
 
@@ -78,12 +78,11 @@ class TextEngine(BaseEngine):
         else:
             raise RuntimeError("No GPU detected! Training requires GPU.")
 
-        # 沐曦 MPS 模式下 device_map="auto" 会导致 tensor 分散到不同设备
-        # 引发跨 GPU 计算错误。强制用 "cuda:0"(对应 METAX_VISIBLE_DEVICES 的第一个)。
+        # 沐曦 MPS 模式下 CUDA_VISIBLE_DEVICES 可能不被遵守
+        # 直接用物理 GPU 号(METAX_VISIBLE_DEVICES 的第一个)。
         visible_devices = os.environ.get("METAX_VISIBLE_DEVICES", "0")
-        device_map = {
-            "": int(visible_devices.split(",")[0]) if visible_devices else 0
-        }
+        first_gpu = int(visible_devices.split(",")[0])  # 物理 GPU 2
+        device_map = {"": first_gpu}
 
         load_kwargs: dict[str, Any] = {
             "torch_dtype": torch.float16,