From 598b279a523c6a5859e6dcf3e1119549fc8c7f2e Mon Sep 17 00:00:00 2001 From: huajinguo12 Date: Tue, 14 Jul 2026 16:11:24 +0800 Subject: [PATCH] Add Ascend NPU Support for Trinity-RFT --- trinity/common/config.py | 6 +- trinity/common/config_validator.py | 4 +- trinity/common/constants.py | 23 +++- trinity/common/models/allocator.py | 7 +- trinity/common/models/vllm_model.py | 7 +- trinity/trainer/verl_legacy/fsdp_workers.py | 10 +- .../trainer/verl_legacy/megatron_workers.py | 8 +- trinity/trainer/verl_legacy/verl_config.py | 6 +- trinity/trainer/verl_legacy/verl_trainer.py | 2 + trinity/utils/device.py | 117 ++++++++++++++++++ trinity/utils/distributed.py | 6 +- 11 files changed, 175 insertions(+), 21 deletions(-) create mode 100644 trinity/utils/device.py diff --git a/trinity/common/config.py b/trinity/common/config.py index 6d27cb3af87..710d9d242d2 100644 --- a/trinity/common/config.py +++ b/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 diff --git a/trinity/common/config_validator.py b/trinity/common/config_validator.py index 24663f9e094..3e4134022ef 100644 --- a/trinity/common/config_validator.py +++ b/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 diff --git a/trinity/common/constants.py b/trinity/common/constants.py index 1c85e0c5f16..5c79a047b99 100644 --- a/trinity/common/constants.py +++ b/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" diff --git a/trinity/common/models/allocator.py b/trinity/common/models/allocator.py index 50d03eb881d..b48a265fa5f 100644 --- a/trinity/common/models/allocator.py +++ b/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, diff --git a/trinity/common/models/vllm_model.py b/trinity/common/models/vllm_model.py index b30e83e8287..7fbc20d6e2e 100644 --- a/trinity/common/models/vllm_model.py +++ b/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). " diff --git a/trinity/trainer/verl_legacy/fsdp_workers.py b/trinity/trainer/verl_legacy/fsdp_workers.py index b6935a12ba7..7508577b0f5 100644 --- a/trinity/trainer/verl_legacy/fsdp_workers.py +++ b/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), ) diff --git a/trinity/trainer/verl_legacy/megatron_workers.py b/trinity/trainer/verl_legacy/megatron_workers.py index 634720945d3..65c60b1884a 100644 --- a/trinity/trainer/verl_legacy/megatron_workers.py +++ b/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), ) diff --git a/trinity/trainer/verl_legacy/verl_config.py b/trinity/trainer/verl_legacy/verl_config.py index 2bf6c09ac01..45b110af9ee 100644 --- a/trinity/trainer/verl_legacy/verl_config.py +++ b/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( diff --git a/trinity/trainer/verl_legacy/verl_trainer.py b/trinity/trainer/verl_legacy/verl_trainer.py index cc0f1ca3b61..762bb95c9bb 100644 --- a/trinity/trainer/verl_legacy/verl_trainer.py +++ b/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) diff --git a/trinity/utils/device.py b/trinity/utils/device.py new file mode 100644 index 00000000000..ed42a5b4f05 --- /dev/null +++ b/trinity/utils/device.py @@ -0,0 +1,117 @@ +# -*- coding: utf-8 -*- +"""Device detection and abstraction layer. + +Unifies the differences among NPU / GPU / CPU devices for trinity modules. +Low-level device operations (e.g. get_torch_device / get_device_id) prefer +reusing verl.utils.device; this module only handles +"detection + initialization + metadata". + +Detection priority: + 1. Environment variable TRINITY_DEVICE forces override (for debugging/special cases) + 2. torch_npu installed and torch.npu.is_available() -> NPU + 3. torch.cuda.is_available() -> CUDA + 4. CPU fallback +""" +import functools +import importlib +import os +from enum import Enum + + +class DeviceType(str, Enum): + """Device type enum. Inherits str so it can be passed directly to APIs that + require "npu"/"cuda" strings.""" + + NPU = "npu" + CUDA = "cuda" + CPU = "cpu" + + +# ---------- Core detection API ---------- + +@functools.lru_cache(maxsize=1) +def get_device_type() -> DeviceType: + """Detect the currently available device type, with process-level caching. + + Returns: + DeviceType.NPU / DeviceType.CUDA / DeviceType.CPU + """ + env_override = os.environ.get("TRINITY_DEVICE", "").lower() + if env_override in ("npu", "cuda", "cpu"): + return DeviceType(env_override) + + if _is_torch_npu_available(): + return DeviceType.NPU + + import torch + if torch.cuda.is_available(): + return DeviceType.CUDA + + return DeviceType.CPU + + +def is_npu() -> bool: + """Whether the current process is running in an NPU environment.""" + return get_device_type() is DeviceType.NPU + + +def is_cuda() -> bool: + """Whether the current process is running in a CUDA environment.""" + return get_device_type() is DeviceType.CUDA + + +def is_cpu() -> bool: + """Whether the current process is running in a CPU environment.""" + return get_device_type() is DeviceType.CPU + + +# ---------- Ray / distributed related ---------- + +def get_ray_resource_key() -> str: + """Accelerator key name in the Ray cluster Resources dict. + + NPU nodes report as "NPU", GPU nodes report as "GPU". + Used by config_validator to auto-detect gpu_per_node. + """ + return "NPU" if is_npu() else "GPU" + + +def get_collective_backend() -> str: + """Collective communication backend name. NPU uses hccl, GPU uses nccl.""" + return "hccl" if is_npu() else "nccl" + + +def get_device_capability() -> int: + """Get major device capability version (device-agnostic). + + Used to decide whether to enable meta tensor initialization for FSDP2. + - NPU: returns 10 (supports meta tensor init, equivalent to sm90+) + - CUDA: returns the actual major compute capability from torch.cuda + - CPU: returns 0 (meta tensor not beneficial) + """ + if is_npu(): + return 10 + if is_cuda(): + import torch + major, _ = torch.cuda.get_device_capability(0) + return major + return 0 + + + +# ---------- Private helpers ---------- + +def _is_torch_npu_available() -> bool: + """Detect whether torch_npu is available, without hard-depending on the package. + + Note: in some environments torch_npu is installed but the CANN environment + variables are not sourced, causing is_available() to raise an exception; + we need to fall back to returning False here. + """ + try: + torch_npu = importlib.import_module("torch_npu") + return bool(torch_npu.npu.is_available()) + except ImportError: + return False + except Exception: + return False diff --git a/trinity/utils/distributed.py b/trinity/utils/distributed.py index 4b46218cb22..c6c7bab98c0 100644 --- a/trinity/utils/distributed.py +++ b/trinity/utils/distributed.py @@ -61,11 +61,11 @@ def init_process_group( """ This function is used to initialize the process group. It requires torch >= 2.6.0 """ - assert backend == "nccl", "Only nccl backend is supported for now." + # assert backend == "nccl", "Only nccl backend is supported for now." - from torch.distributed.distributed_c10d import is_nccl_available + # from torch.distributed.distributed_c10d import is_nccl_available - assert is_nccl_available() + # assert is_nccl_available() init_method = ( f"tcp://[{host}]:{port}" if is_ipv6_address(ip_str=host) else f"tcp://{host}:{port}"