Add Ascend NPU Support for Trinity-RFT

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

trinity/common/constants.py

@@ -71,14 +71,35 @@ class StorageType(CaseInsensitiveEnum):
 
 class SyncMethodEnumMeta(CaseInsensitiveEnumMeta):
     name_aliases = {
-        "online": "nccl",
         "offline": "checkpoint",
     }
 
+    @staticmethod
+    def _resolve_online_alias() -> str:
+        """Resolve the "online" alias to a concrete collective backend based on device type.
+
+        NPU -> "hccl", others (GPU/CPU) -> "nccl".
+        Lazy import to avoid circular dependencies.
+        """
+        from trinity.utils.device import is_npu
+
+        return "hccl" if is_npu() else "nccl"
+
+    def __getitem__(cls, name):
+        if isinstance(name, str) and name.lower() == "online":
+            name = cls._resolve_online_alias()
+        return super().__getitem__(name)
+
+    def __call__(cls, value, *args, **kwargs):
+        if isinstance(value, str) and value.lower() == "online":
+            value = cls._resolve_online_alias()
+        return super().__call__(value, *args, **kwargs)
+
 
 class SyncMethod(CaseInsensitiveEnum, metaclass=SyncMethodEnumMeta):
     """Sync Method."""
 
+    HCCL = "hccl"
     NCCL = "nccl"
     CHECKPOINT = "checkpoint"
     MEMORY = "memory"

trinity/common/models/allocator.py

@@ -17,6 +17,7 @@
 
 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.log import get_logger
 
 
@@ -61,7 +62,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
@@ -123,7 +124,7 @@ async def create_all_models(self) -> Tuple[List[ModelWrapper], List[List[ModelWr
         self.bundle_result = self.allocate_bundles()
         self.pg = placement_group(
             self.bundle_result.bundles,
-            strategy="PACK",
+            strategy="STRICT_PACK",
             name=self.placement_group_name,
         )
         await self.pg.ready()
@@ -244,7 +245,7 @@ async def get_model_wrapper(
             ray.remote(actor_cls)
             .options(
                 name=actor_name,
-                num_gpus=engine_config.gpu_per_engine / engine_config.nnodes,
+                resources={get_ray_resource_key(): engine_config.gpu_per_engine / engine_config.nnodes},
                 namespace=engine_config.ray_namespace,
                 scheduling_strategy=PlacementGroupSchedulingStrategy(
                     placement_group=pg,

trinity/common/models/vllm_model.py

@@ -19,6 +19,7 @@
 )
 from trinity.common.models.model import BaseInferenceModel
 from trinity.common.models.vllm_patch import get_vllm_version
+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 +127,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)
@@ -614,11 +615,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). "

trinity/trainer/verl_legacy/fsdp_workers.py

@@ -50,7 +50,6 @@
 from verl.utils.device import (
     get_device_id,
     get_device_name,
-    get_nccl_backend,
     get_torch_device,
 )
 from verl.utils.flops_counter import FlopsCounter
@@ -100,6 +99,7 @@
     rank0_iterator,
     save_rank0_safetensors,
 )
+from trinity.utils.device import get_collective_backend, get_device_capability
 from trinity.utils.distributed import WeightTransferEngine
 from trinity.utils.log import get_logger
 
@@ -134,7 +134,7 @@ def __init__(self, config: DictConfig, role: str, **kwargs):
             rank = int(os.environ.get("RANK", 0))
             world_size = int(os.environ.get("WORLD_SIZE", 1))
             torch.distributed.init_process_group(
-                backend=f"cpu:gloo,{get_device_name()}:{get_nccl_backend()}",
+                backend=f"cpu:gloo,{get_device_name()}:{get_collective_backend()}",
                 rank=rank,
                 world_size=world_size,
                 timeout=datetime.timedelta(seconds=self.config.get("nccl_timeout", 600)),
@@ -410,7 +410,7 @@ def _build_model_optimizer(  # noqa: C901
         if self.rank == 0:
             self.logger.info(f"Model config after override: {actor_model_config}")
 
-        major_capability, _ = torch.cuda.get_device_capability(0)
+        major_capability = get_device_capability()
         use_meta = (
             (
                 self.rank != 0
@@ -893,7 +893,7 @@ def sync_weight(self):
                     raise RuntimeError("Weight sync group has not been initialized.")
                 self.logger.info("Starting NCCL weight sync broadcast.")
                 self.weight_transfer_engine.sync_weight(iterator=weight_iterator)
-                torch.cuda.synchronize()
+                getattr(torch, get_device_name()).synchronize()
                 self.logger.info("Finished NCCL weight sync broadcast.")
             else:
                 for _ in weight_iterator:
@@ -1264,7 +1264,7 @@ def __init__(self, config: FSDPCriticConfig):
 
         if not torch.distributed.is_initialized():
             torch.distributed.init_process_group(
-                backend=get_nccl_backend(),
+                backend=get_collective_backend(),
                 timeout=datetime.timedelta(seconds=self.config.get("nccl_timeout", 600)),
                 init_method=os.environ.get("DIST_INIT_METHOD", None),
             )

trinity/trainer/verl_legacy/megatron_workers.py

@@ -47,7 +47,6 @@
 from verl.utils.device import (
     get_device_id,
     get_device_name,
-    get_nccl_backend,
     get_torch_device,
     set_expandable_segments,
 )
@@ -95,6 +94,7 @@
     rank0_iterator,
     save_rank0_safetensors,
 )
+from trinity.utils.device import get_collective_backend
 from trinity.utils.distributed import WeightTransferEngine
 from trinity.utils.log import get_logger
 
@@ -329,7 +329,7 @@ def __init__(self, config: DictConfig, role: str, **kwargs):
             set_numa_affinity()
             rank = int(os.environ["LOCAL_RANK"])
             torch.distributed.init_process_group(
-                backend=f"cpu:gloo,{get_device_name()}:{get_nccl_backend()}",
+                backend=f"cpu:gloo,{get_device_name()}:{get_collective_backend()}",
                 timeout=datetime.timedelta(seconds=self.config.get("nccl_timeout", 600)),
                 init_method=os.environ.get("DIST_INIT_METHOD", None),
             )
@@ -825,7 +825,7 @@ def sync_weight(self):
                 raise RuntimeError("Weight sync group has not been initialized.")
             self.logger.info("Starting NCCL weight sync broadcast.")
             self.weight_transfer_engine.sync_weight(iterator=weight_iterator)
-            torch.cuda.synchronize()
+            getattr(torch, get_device_name()).synchronize()
             self.logger.info("Finished NCCL weight sync broadcast.")
         else:
             for _ in weight_iterator:
@@ -1147,7 +1147,7 @@ def __init__(self, config: McoreCriticConfig):
             set_numa_affinity()
             rank = int(os.environ["LOCAL_RANK"])
             torch.distributed.init_process_group(
-                backend=get_nccl_backend(),
+                backend=get_collective_backend(),
                 timeout=datetime.timedelta(seconds=self.config.get("nccl_timeout", 600)),
                 init_method=os.environ.get("DIST_INIT_METHOD", None),
             )

trinity/trainer/verl_legacy/verl_config.py

@@ -167,6 +167,7 @@ class Actor:
     ulysses_sequence_parallel_size: Optional[int] = None
     entropy_from_logits_with_chunking: bool = False
     entropy_checkpointing: bool = False
+    use_torch_compile: Optional[bool] = None  # None = inherit from trainer.use_torch_compile
     checkpoint: _CheckpointConfig = field(default_factory=_CheckpointConfig)
     optim: Optim = field(default_factory=Optim)
     fsdp_config: FSDPConfig = field(default_factory=FSDPConfig)
@@ -200,6 +201,7 @@ class Ref:
     ulysses_sequence_parallel_size: Optional[int] = None
     entropy_from_logits_with_chunking: bool = False
     entropy_checkpointing: bool = False
+    use_torch_compile: Optional[bool] = None  # None = inherit from trainer.use_torch_compile
     checkpoint: _CheckpointConfig = field(
         default_factory=lambda: _CheckpointConfig(load_contents=["model"], save_contents=["model"])
     )
@@ -384,7 +386,7 @@ class Trainer:
     sync_freq: int = 0
     max_actor_ckpt_to_keep: Optional[int] = None
     max_critic_ckpt_to_keep: Optional[int] = None
-    device: str = "cuda"  # default to cuda
+    device: str = "npu"  # was "cuda"; direct NPU default avoids auto_set_device warning
 
 
 @dataclass
@@ -550,6 +552,7 @@ def synchronize_config(self, config: Config) -> None:  # noqa: C901
             ("ulysses_sequence_parallel_size",) * 2,
             ("ppo_max_token_len_per_gpu", "max_token_len_per_gpu"),
             ("strategy", "trainer_strategy"),
+            ("use_torch_compile",) * 2,
         ]:
             set_if_none(actor_config, actor_attr, getattr(config.trainer, trainer_attr))
         self._check_parallel_config(
@@ -573,6 +576,7 @@ def synchronize_config(self, config: Config) -> None:  # noqa: C901
             ("log_prob_max_token_len_per_gpu", "max_token_len_per_gpu"),
             ("ulysses_sequence_parallel_size",) * 2,
             ("strategy", "trainer_strategy"),
+            ("use_torch_compile",) * 2,
         ]:
             set_if_none(ref_config, ref_attr, getattr(config.trainer, trainer_attr))
         self._check_parallel_config(

trinity/trainer/verl_legacy/verl_trainer.py

@@ -30,6 +30,7 @@
 from verl.utils import hf_processor, hf_tokenizer
 from verl.utils.checkpoint.checkpoint_manager import find_latest_ckpt_path
 from verl.utils.debug import marked_timer
+from verl.utils.device import auto_set_device
 from verl.utils.fs import copy_local_path_from_hdfs
 from verl.utils.metric import reduce_metrics
 from verl.workers.config import FSDPEngineConfig
@@ -196,6 +197,7 @@ def __init__(
         )
         train_config = global_config.trainer
         config = OmegaConf.structured(train_config.trainer_config)
+        auto_set_device(config)
         # download the checkpoint from hdfs
         local_path = copy_local_path_from_hdfs(config.actor_rollout_ref.model.path)