feat(npu): colocate TMS env and CANN path for Ascend

Add TMS torch mode, training region, colocate-only PYTORCH_NPU_ALLOC_CONF override, CANN PYTHONPATH for Ray workers, and safe NPU empty_cache. Complements IPC weight sync for bridge colocate on 8x910B.

Signed-off-by: kaiyuan <kyxiezju@163.com>

Author: CalvinXKY

vime/backends/megatron_utils/__init__.py

@@ -1,6 +1,12 @@
 import logging
 
 import torch
+
+try:
+    import torch_npu  # noqa: F401
+except ImportError:
+    pass
+
 from vime.utils.common import is_npu
 
 if is_npu():

vime/backends/megatron_utils/actor.py

@@ -84,10 +84,18 @@ def init(
                 logger.info(f"Set torch_memory_saver.memory_margin_bytes to {x}")
                 torch_memory_saver.memory_margin_bytes = x
 
+        tms_region_ctx = None
+        if args.offload_train and is_npu():
+            tms_region_ctx = torch_memory_saver.region(tag="training", enable_cpu_backup=True)
+            tms_region_ctx.__enter__()
+
         self.model, self.optimizer, self.opt_param_scheduler, loaded_rollout_id = initialize_model_and_optimizer(
             args, role
         )
 
+        if tms_region_ctx is not None:
+            tms_region_ctx.__exit__(None, None, None)
+
         vpp_size = mpu.get_virtual_pipeline_model_parallel_world_size() or 1
         if vpp_size > 1:
             from megatron.core.utils import get_model_config
@@ -621,7 +629,7 @@ def update_weights(self) -> None:
             if dist.get_rank() == 0:
                 ray.get(self.rollout_manager.clear_updatable_num_new_engines.remote())
 
-        with torch_memory_saver.disable() if self.args.offload_train else nullcontext():
+        with torch_memory_saver.disable() if (self.args.offload_train and not is_npu()) else nullcontext():
             print_memory("before update_weights")
             self.weight_updater.update_weights()
             print_memory("after update_weights")

vime/backends/vllm_utils/vllm_engine.py

@@ -24,7 +24,7 @@
 
 from vime.backends.vllm_utils.arguments import SKIPPED_DESTS, get_vllm_cli_action_table
 from vime.ray.ray_actor import RayActor
-from vime.utils.common import is_npu
+from vime.utils.common import get_cann_python_site_packages, is_npu, prepend_pythonpath
 from vime.utils.http_utils import get_host_info
 
 logger = logging.getLogger(__name__)
@@ -361,6 +361,12 @@ def build_vllm_subprocess_env(server_args: dict[str, Any]) -> dict[str, str]:
     env.setdefault("NCCL_CUMEM_ENABLE", "0")
     if is_npu():
         env["ASCEND_RT_VISIBLE_DEVICES"] = server_args["visible_devices"]
+        env["VLLM_USE_AOT_COMPILE"] = "0"
+        cann_python_path = get_cann_python_site_packages()
+        if cann_python_path is not None:
+            prepend_pythonpath(env, cann_python_path)
+        if getattr(args, "colocate", False):
+            env.pop("PYTORCH_NPU_ALLOC_CONF", None)
     else:
         env["CUDA_VISIBLE_DEVICES"] = server_args["visible_devices"]
     env.setdefault("VLLM_SERVER_DEV_MODE", "1")

vime/ray/actor_group.py

@@ -5,7 +5,7 @@
 from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
 
 from vime.ray.utils import NOSET_VISIBLE_DEVICES_ENV_VARS_LIST
-from vime.utils.common import is_npu
+from vime.utils.common import get_cann_python_site_packages, is_npu, prepend_pythonpath
 
 
 class RayTrainGroup:
@@ -82,6 +82,16 @@ def _allocate_gpus_for_actor(self, pg, num_gpus_per_actor):
             env_vars["TMS_INIT_ENABLE"] = "1"
             env_vars["TMS_INIT_ENABLE_CPU_BACKUP"] = "1"
 
+            if is_npu():
+                env_vars["TMS_HOOK_MODE"] = "torch"
+                env_vars["TMS_REGION_TAG"] = "training"
+                env_vars["TMS_ENABLE_CPU_BACKUP"] = "1"
+                if self.args.colocate:
+                    env_vars.pop("PYTORCH_NPU_ALLOC_CONF", None)
+                cann_python_path = get_cann_python_site_packages()
+                if cann_python_path is not None:
+                    prepend_pythonpath(env_vars, cann_python_path)
+
         # We cannot do routing replay for critic.
         if self.args.use_routing_replay and self.role == "actor":
             env_vars["ENABLE_ROUTING_REPLAY"] = "1"

vime/ray/placement_group.py

@@ -18,10 +18,26 @@
 @ray.remote
 class InfoActor:
     def get_ip_and_gpu_id(self):
-        if is_npu():
-            return ray.util.get_node_ip_address(), ray.get_runtime_context().get_accelerator_ids()["NPU"][0]
-        else:
-            return ray.util.get_node_ip_address(), ray.get_gpu_ids()[0]
+        try:
+            import torch_npu  # noqa: F401
+
+            has_npu = True
+        except ImportError:
+            has_npu = False
+
+        if has_npu or is_npu():
+            npu_ids = ray.get_runtime_context().get_accelerator_ids().get("NPU", [])
+            if npu_ids:
+                return ray.util.get_node_ip_address(), npu_ids[0]
+
+        gpu_ids = ray.get_gpu_ids()
+        if gpu_ids:
+            return ray.util.get_node_ip_address(), gpu_ids[0]
+
+        raise RuntimeError(
+            "No GPU/NPU IDs found. "
+            f"Accelerator IDs: {ray.get_runtime_context().get_accelerator_ids()}, GPU IDs: {gpu_ids}"
+        )
 
 
 def sort_key(x):

vime/ray/rollout.py

@@ -21,7 +21,7 @@
 GPU_MEMORY_TYPE_CUDA_GRAPH = "cuda_graph"
 from vime.rollout.base_types import call_rollout_fn
 from vime.utils import logging_utils
-from vime.utils.common import is_npu
+from vime.utils.common import get_cann_python_site_packages, is_npu, prepend_pythonpath
 from vime.utils.dp_schedule import build_dp_schedule
 from vime.utils.health_monitor import RolloutHealthMonitor
 from vime.utils.http_utils import _wrap_ipv6, find_available_port, get_host_info, init_http_client
@@ -126,6 +126,13 @@ def start_engines(self, port_cursors: dict[int, int] | None = None) -> tuple[lis
             )
 
             env_vars = {name: "1" for name in NOSET_VISIBLE_DEVICES_ENV_VARS_LIST}
+            if is_npu():
+                cann_python_path = get_cann_python_site_packages()
+                if cann_python_path is not None:
+                    prepend_pythonpath(env_vars, cann_python_path)
+                if self.args.colocate:
+                    env_vars["PYTORCH_NPU_ALLOC_CONF"] = ""
+                env_vars["VLLM_USE_AOT_COMPILE"] = "0"
             rollout_engine = RolloutRayActor.options(
                 num_cpus=num_cpus,
                 scheduling_strategy=scheduling_strategy,

vime/utils/common.py

@@ -1,6 +1,37 @@
+import os
+
 import torch
 
 
+def get_cann_python_site_packages() -> str | None:
+    """Return CANN Python site-packages if ``acl`` is importable from there."""
+    candidates: list[str] = []
+    for env_key in ("ASCEND_TOOLKIT_HOME", "ASCEND_HOME_PATH"):
+        base = os.environ.get(env_key)
+        if not base:
+            continue
+        candidates.extend(
+            [
+                os.path.join(base, "python", "site-packages"),
+                os.path.normpath(os.path.join(base, "..", "python", "site-packages")),
+            ]
+        )
+    candidates.append("/usr/local/Ascend/ascend-toolkit/latest/python/site-packages")
+
+    for path in candidates:
+        if os.path.isdir(os.path.join(path, "acl")):
+            return path
+    return None
+
+
+def prepend_pythonpath(env: dict[str, str], *paths: str) -> None:
+    existing = env.get("PYTHONPATH", os.environ.get("PYTHONPATH", ""))
+    existing_parts = {part for part in existing.split(os.pathsep) if part}
+    prefix_parts = [path for path in paths if path and path not in existing_parts]
+    if prefix_parts:
+        env["PYTHONPATH"] = os.pathsep.join([*prefix_parts, existing] if existing else prefix_parts)
+
+
 def is_npu() -> bool:
     if not hasattr(torch, "npu"):
         return False

vime/utils/memory_utils.py

@@ -17,10 +17,16 @@ def clear_memory(clear_host_memory: bool = False):
     gc.collect()
     torch.cuda.empty_cache()
     if is_npu():
-        torch.npu.empty_cache()
+        try:
+            torch.npu.empty_cache()
+        except RuntimeError:
+            pass
     if clear_host_memory:
         if is_npu():
-            torch.npu.empty_cache()
+            try:
+                torch.npu.empty_cache()
+            except RuntimeError:
+                pass
         else:
             torch._C._host_emptyCache()