Add Ascend NPU Support for Trinity-RFT

trinity/cli/launcher.py

@@ -250,7 +250,7 @@ def both(config: Config) -> StageStatus:
         # This must happen after both sides are prepared (Trainer has model
         # meta cached, Explorer has rollout models created) and before the
         # first weight sync.
-        if config.synchronizer.sync_method == SyncMethod.NCCL:
+        if config.synchronizer.sync_method in (SyncMethod.NCCL, SyncMethod.HCCL):
             synchronizer = Synchronizer.get_actor(namespace=config.ray_namespace)
             ray.get(synchronizer.coordinate_weight_sync_setup.remote())
         ray.get(
@@ -303,7 +303,7 @@ def both(config: Config) -> StageStatus:
     finally:
         # Tear down the NCCL weight sync group before shutting down actors.
         # Best-effort: if actors or Synchronizer are already dead, skip.
-        if config.synchronizer.sync_method == SyncMethod.NCCL:
+        if config.synchronizer.sync_method in (SyncMethod.NCCL, SyncMethod.HCCL):
             try:
                 synchronizer = Synchronizer.get_actor(namespace=config.ray_namespace)
                 ray.get(synchronizer.coordinate_weight_sync_teardown.remote(), timeout=30)

trinity/common/config.py

@@ -25,6 +25,7 @@
     SyncStyle,
 )
 from trinity.utils.annotations import Experimental
+from trinity.utils.device import is_npu
 from trinity.utils.log import get_logger
 
 logger = get_logger(__name__)
@@ -829,6 +830,7 @@ class TrainerConfig:
     max_token_len_per_gpu: Optional[int] = None
     ulysses_sequence_parallel_size: int = 1  # sp size
     fix_actor_microbatch_loss_scale: bool = False  # EXPERIMENTAL
+    use_torch_compile: bool = True  # whether to use torch.compile in actor/ref; set to false on NPU
 
     # offloading
     param_offload: bool = False
@@ -867,7 +869,9 @@ class MonitorConfig:
 class SynchronizerConfig:
     """Configs for model weight synchronization."""
 
-    sync_method: SyncMethod = SyncMethod.NCCL
+    sync_method: SyncMethod = field(
+        default_factory=lambda: SyncMethod.HCCL if is_npu() else SyncMethod.NCCL
+    )
     sync_style: SyncStyle = SyncStyle.FIXED
     # sync weights every `sync_interval` steps
     sync_interval: int = 1

trinity/common/config_validator.py

@@ -17,6 +17,7 @@
 from trinity.common.constants import StorageType, SyncMethod, SyncStyle
 from trinity.common.dataclass_utils import build_dataclass_from_mapping
 from trinity.common.patch import kimi_vl_monkey_patch_decorator
+from trinity.utils.device import get_ray_resource_key
 from trinity.utils.log import get_logger
 from trinity.utils.lora_utils import create_dummy_lora
 
@@ -192,8 +193,9 @@ def _set_cluster_info(self, config: Config) -> None:
         # set gpu_per_node
         if not config.cluster.gpu_per_node:
             gpu_per_node = 0
+            resource_key = get_ray_resource_key()
             for node in alive_nodes:
-                node_gpus = node.get("Resources", {}).get("GPU")
+                node_gpus = node.get("Resources", {}).get(resource_key)
                 if node_gpus and node_gpus > 0:
                     gpu_per_node = int(node_gpus)
                     break
@@ -548,10 +550,10 @@ def _check_tinker(self, config: Config) -> None:
             config.trainer.trainer_type = "tinker"
             self.logger.warning("Trainer type is set to `tinker`.")
 
-        if config.synchronizer.sync_method == SyncMethod.NCCL:
+        if config.synchronizer.sync_method in (SyncMethod.NCCL, SyncMethod.HCCL):
             config.synchronizer.sync_method = SyncMethod.CHECKPOINT
             self.logger.warning(
-                "Tinker do not support NCCL, `synchronizer.sync_method` is set to `checkpoint`."
+                "Tinker does not support NCCL/HCCL, `synchronizer.sync_method` is set to `checkpoint`."
             )
 
     def _check_model_len(self, config: Config) -> None:
@@ -841,23 +843,23 @@ def validate(self, config: Config) -> None:
         """
         config.synchronizer.ray_namespace = config.ray_namespace
         config.synchronizer.explorer_world_size = config.cluster.rollout_gpu_num
-        if config.synchronizer.sync_method == SyncMethod.NCCL:
+        if config.synchronizer.sync_method in (SyncMethod.NCCL, SyncMethod.HCCL):
             if config.mode in ["train", "explore", "bench", "serve"]:
                 config.synchronizer.sync_method = SyncMethod.CHECKPOINT
                 self.logger.warning(
-                    f"`{config.mode}` mode does not support NCCL synchronization, "
+                    f"`{config.mode}` mode does not support NCCL/HCCL synchronization, "
                     "set `synchronizer.sync_method` to `checkpoint`."
                 )
             if config.model.lora_configs is not None:
                 config.synchronizer.sync_method = SyncMethod.CHECKPOINT
                 self.logger.warning(
-                    "LoRA is not supported with NCCL synchronization, "
+                    "LoRA is not supported with NCCL/HCCL synchronization, "
                     "set `synchronizer.sync_method` to `checkpoint`."
                 )
             if config.mode == "colocate":
                 config.synchronizer.sync_method = SyncMethod.MEMORY
                 self.logger.warning(
-                    "Colocate mode can't use NCCL synchronization. "
+                    "Colocate mode can't use NCCL/HCCL synchronization. "
                     "Set `synchronizer.sync_method` to `memory` instead."
                 )
 

trinity/common/constants.py

@@ -79,6 +79,7 @@ class SyncMethodEnumMeta(CaseInsensitiveEnumMeta):
 class SyncMethod(CaseInsensitiveEnum, metaclass=SyncMethodEnumMeta):
     """Sync Method."""
 
+    HCCL = "hccl"
     NCCL = "nccl"
     CHECKPOINT = "checkpoint"
     MEMORY = "memory"

trinity/common/models/allocator.py

@@ -17,6 +17,8 @@
 
 from trinity.common.config import ExplorerConfig, InferenceModelConfig
 from trinity.common.models.model import ModelWrapper
+from trinity.utils.device import get_ray_resource_key
+from trinity.utils.device import is_npu
 from trinity.utils.log import get_logger
 
 
@@ -61,7 +63,7 @@ def allocate_bundles(self) -> BundleResult:
             gpus_per_bundle = config.gpu_per_engine // config.nnodes
             for engine_id in range(config.engine_num):
                 for node_id in range(config.nnodes):
-                    bundles.append({"GPU": float(gpus_per_bundle), "CPU": 1})
+                    bundles.append({get_ray_resource_key(): float(gpus_per_bundle), "CPU": 1})
                     actor_name = self.get_actor_name(role, engine_id, node_id)
                     actor_bundle_map[actor_name] = bundle_id
                     bundle_actor_map[bundle_id] = actor_name
@@ -240,20 +242,36 @@ async def get_model_wrapper(
         engine_config = deepcopy(config)
         engine_config.ray_actor_name = actor_name
         engine_config.node_rank = i
-        handlers.append(
-            ray.remote(actor_cls)
-            .options(
-                name=actor_name,
-                num_gpus=engine_config.gpu_per_engine / engine_config.nnodes,
-                namespace=engine_config.ray_namespace,
-                scheduling_strategy=PlacementGroupSchedulingStrategy(
-                    placement_group=pg,
-                    placement_group_capture_child_tasks=True,
-                    placement_group_bundle_index=bundle_id,
-                ),
+        if is_npu():
+            handlers.append(
+                ray.remote(actor_cls)
+                .options(
+                    name=actor_name,
+                    resources={"NPU": engine_config.gpu_per_engine / engine_config.nnodes},
+                    namespace=engine_config.ray_namespace,
+                    scheduling_strategy=PlacementGroupSchedulingStrategy(
+                        placement_group=pg,
+                        placement_group_capture_child_tasks=True,
+                        placement_group_bundle_index=bundle_id,
+                    ),
+                )
+                .remote(config=engine_config)
+            )
+        else:
+            handlers.append(
+                ray.remote(actor_cls)
+                .options(
+                    name=actor_name,
+                    num_gpus=engine_config.gpu_per_engine / engine_config.nnodes,
+                    namespace=engine_config.ray_namespace,
+                    scheduling_strategy=PlacementGroupSchedulingStrategy(
+                        placement_group=pg,
+                        placement_group_capture_child_tasks=True,
+                        placement_group_bundle_index=bundle_id,
+                    ),
+                )
+                .remote(config=engine_config)
             )
-            .remote(config=engine_config)
-        )
     if len(actor_bundle_list) > 1:
         # get master address and port from the first handler and set it to all handlers for distributed communication
         master_addr, master_port = await handlers[0].get_available_address.remote(random_port=True)

trinity/common/models/vllm_model.py

@@ -19,6 +19,8 @@
 )
 from trinity.common.models.model import BaseInferenceModel
 from trinity.common.models.vllm_patch import get_vllm_version
+from trinity.common.models.vllm_worker import get_trinity_worker_cls_name
+from trinity.utils.device import get_collective_backend
 
 
 # V0 engine is deprecated since vLLM v0.10.2, related code will be removed in the future.
@@ -126,7 +128,7 @@ async def prepare(self) -> None:
                 return
 
             weight_transfer_config = WeightTransferConfig(
-                backend="nccl" if self.config.sync_method == SyncMethod.NCCL else "checkpoint"
+                backend=get_collective_backend() if self.config.sync_method in (SyncMethod.NCCL, SyncMethod.HCCL) else "checkpoint"
             )
 
             rope_params = defaultdict(dict)
@@ -142,7 +144,7 @@ async def prepare(self) -> None:
             engine_args = vllm.AsyncEngineArgs(
                 model=self.config.model_path,
                 enforce_eager=self.config.enforce_eager,
-                worker_cls="trinity.common.models.vllm_worker.TrinityGPUWorker",
+                worker_cls=get_trinity_worker_cls_name(),
                 tensor_parallel_size=self.config.tensor_parallel_size,
                 pipeline_parallel_size=self.config.pipeline_parallel_size,
                 data_parallel_size=self.config.data_parallel_size,
@@ -589,7 +591,7 @@ async def sync_model_weights(
 
         await self.async_llm.start_weight_update(is_checkpoint_format=True)
         update_info = {}
-        if method == SyncMethod.NCCL:
+        if method in (SyncMethod.NCCL, SyncMethod.HCCL):
             update_info = dict(
                 names=[meta[0] for meta in self.state_dict_meta],
                 dtype_names=[meta[1] for meta in self.state_dict_meta],
@@ -614,11 +616,13 @@ async def init_process_group(
         rank_offset: int,
         world_size: int,
         group_name: str,
-        backend: str = "nccl",
+        backend: Optional[str] = None,
         timeout: float = 1200,
     ):
         from vllm.distributed.weight_transfer.base import WeightTransferInitRequest
 
+        if backend is None:
+            backend = get_collective_backend()
         if self.config.node_rank != 0:
             self.logger.warning(
                 "init_process_group should only be called on the main node (node_rank=0). "
@@ -638,7 +642,7 @@ async def init_process_group(
             rank_offset=rank_offset,
             world_size=world_size,
         )
-        if self.config.sync_method != SyncMethod.NCCL:
+        if self.config.sync_method not in (SyncMethod.NCCL, SyncMethod.HCCL):
             init_info["namespace"] = self.ray_namespace
             init_info["sync_method"] = self.config.sync_method.value
         await self.async_llm.init_weight_transfer_engine(

trinity/common/models/vllm_worker.py

@@ -1,10 +1,17 @@
 # -*- coding: utf-8 -*-
-"""Custom vLLM worker classes for Trinity."""
+"""Custom vLLM worker classes for Trinity.
+
+Provides device-aware worker subclasses for both GPU (vLLM) and NPU
+(vllm-ascend). The appropriate class is selected at runtime via
+:func:`get_trinity_worker_cls_name` based on the detected device type.
+"""
 import logging
 from typing import Any
 
 from vllm.v1.worker.gpu_worker import Worker as VLLMGPUWorker
 
+from trinity.utils.device import is_npu
+
 
 def _suppress_layerwise_reload_warnings() -> None:
     """Silence benign vLLM layerwise reload warnings during weight sync."""
@@ -16,24 +23,78 @@ def _suppress_layerwise_reload_warnings() -> None:
         pass
 
 
+def _apply_trinity_patches(model_runner: Any) -> None:
+    """Apply Trinity-specific patches shared by GPU and NPU workers.
+
+    - Patches the vLLM fused-MoE weight loader (workaround for missing
+      ``weight_loader`` on MoE params, also handles ACLGraphWrapper on NPU).
+    - Patches prompt logprobs computation to apply temperature scaling.
+    - Suppresses benign layerwise reload warnings during weight sync.
+    """
+    from verl.utils.vllm.patch import patch_vllm_moe_model_weight_loader
+
+    from trinity.common.models.vllm_patch.worker_patch import (
+        patch_vllm_prompt_logprobs,
+    )
+
+    patch_vllm_moe_model_weight_loader(model_runner.model)
+    patch_vllm_prompt_logprobs(model_runner)
+    _suppress_layerwise_reload_warnings()
+
+
+def _register_trinity_weight_transfer_engine() -> None:
+    """Register Trinity's checkpoint weight transfer backend with vLLM."""
+    from trinity.common.models.vllm_extension import (
+        register_checkpoint_weight_transfer_engine,
+    )
+
+    register_checkpoint_weight_transfer_engine()
+
+
 class TrinityGPUWorker(VLLMGPUWorker):
     def apply_patches(self) -> None:
         """Apply necessary patches to vLLM."""
-        from verl.utils.vllm.patch import patch_vllm_moe_model_weight_loader
-
-        from trinity.common.models.vllm_patch.worker_patch import (
-            patch_vllm_prompt_logprobs,
-        )
-
-        patch_vllm_moe_model_weight_loader(self.model_runner.model)
-        patch_vllm_prompt_logprobs(self.model_runner)
-        _suppress_layerwise_reload_warnings()
+        _apply_trinity_patches(self.model_runner)
 
     def load_model(self, *args: Any, **kwargs: Any) -> None:
         """Register Trinity weight-transfer engines before vLLM loads them."""
-        from trinity.common.models.vllm_extension import (
-            register_checkpoint_weight_transfer_engine,
-        )
-
-        register_checkpoint_weight_transfer_engine()
+        _register_trinity_weight_transfer_engine()
         return super().load_model(*args, **kwargs)
+
+
+if is_npu():
+    from vllm_ascend.worker.worker import NPUWorker
+
+    class TrinityNPUWorker(NPUWorker):
+        """Trinity worker for Ascend NPU, based on vllm-ascend's NPUWorker.
+
+        vllm-ascend's own patches are applied via ``adapt_patch()`` in
+        ``NPUWorker.__init__``; this subclass only adds Trinity-specific
+        patches (MoE weight loader + prompt logprobs) and registers
+        Trinity's checkpoint weight transfer engine before model loading.
+        """
+
+        def apply_patches(self) -> None:
+            """Apply Trinity-specific patches after vllm-ascend patches.
+
+            ``adapt_patch()`` is already invoked in ``NPUWorker.__init__``,
+            so here we only apply the Trinity-specific MoE weight loader and
+            prompt logprobs patches.
+            """
+            _apply_trinity_patches(self.model_runner)
+
+        def load_model(self, *args: Any, **kwargs: Any) -> None:
+            """Register Trinity weight-transfer engines before vLLM loads them."""
+            _register_trinity_weight_transfer_engine()
+            return super().load_model(*args, **kwargs)
+
+
+def get_trinity_worker_cls_name() -> str:
+    """Return the Trinity worker class qualified name for the current device.
+
+    Returns ``"trinity.common.models.vllm_worker.TrinityNPUWorker"`` on NPU
+    and ``"trinity.common.models.vllm_worker.TrinityGPUWorker"`` otherwise.
+    """
+    if is_npu():
+        return "trinity.common.models.vllm_worker.TrinityNPUWorker"
+    return "trinity.common.models.vllm_worker.TrinityGPUWorker"

trinity/explorer/explorer.py

@@ -71,7 +71,10 @@ def __init__(self, config: Config):
         else:
             self.min_wait_num = None
         self.rollout_coordinator = None
-        self.use_nccl_sync = self.config.synchronizer.sync_method == SyncMethod.NCCL
+        self.use_nccl_sync = self.config.synchronizer.sync_method in (
+            SyncMethod.NCCL,
+            SyncMethod.HCCL,
+        )
         self.pending_eval_tasks = deque()
 
         # For checkpoint weights update
@@ -367,7 +370,7 @@ async def _watch_trainer_sync_signal(self) -> None:
         """
         while not self._async_watch_stopped:
             try:
-                if self.sync_method == SyncMethod.NCCL:
+                if self.sync_method in (SyncMethod.NCCL, SyncMethod.HCCL):
                     ready = await self.synchronizer.trainer_requires_weight_sync.remote()
                     if ready is None:
                         break  # trainer stopped
@@ -394,7 +397,10 @@ async def need_sync(self) -> bool:
             return False
         if (self.explore_step_num - self.sync_offset) % self.sync_interval == 0:
             await self.finish_current_steps()
-            if self.sync_style == SyncStyle.TRAINER_DRIVEN and self.sync_method == SyncMethod.NCCL:
+            if (
+                self.sync_style == SyncStyle.TRAINER_DRIVEN
+                and self.sync_method in (SyncMethod.NCCL, SyncMethod.HCCL)
+            ):
                 require_sync = bool(await self.synchronizer.trainer_requires_weight_sync.remote())
             else:
                 require_sync = True