diff --git a/trinity/cli/launcher.py b/trinity/cli/launcher.py index 7b0977803ba..685789eff63 100644 --- a/trinity/cli/launcher.py +++ b/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) 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..abc38b1332f 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 @@ -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." ) diff --git a/trinity/common/constants.py b/trinity/common/constants.py index 1c85e0c5f16..0e97a49cc2a 100644 --- a/trinity/common/constants.py +++ b/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" diff --git a/trinity/common/models/allocator.py b/trinity/common/models/allocator.py index 50d03eb881d..d4f7c485fba 100644 --- a/trinity/common/models/allocator.py +++ b/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) diff --git a/trinity/common/models/vllm_model.py b/trinity/common/models/vllm_model.py index b30e83e8287..8028e2278fd 100644 --- a/trinity/common/models/vllm_model.py +++ b/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( diff --git a/trinity/common/models/vllm_worker.py b/trinity/common/models/vllm_worker.py index 14a5aefe506..aeb37c91335 100644 --- a/trinity/common/models/vllm_worker.py +++ b/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" diff --git a/trinity/explorer/explorer.py b/trinity/explorer/explorer.py index 0188b14ee93..e572df1b1c5 100644 --- a/trinity/explorer/explorer.py +++ b/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 diff --git a/trinity/trainer/trainer.py b/trinity/trainer/trainer.py index e85dd83f15d..202fa1295c9 100644 --- a/trinity/trainer/trainer.py +++ b/trinity/trainer/trainer.py @@ -211,7 +211,10 @@ async def need_sync(self) -> bool: ) else: # explorer driven # for memory & checkpoint; TODO: apply to nccl sync - if self.last_sync_step == self.train_step_num and self.sync_method != SyncMethod.NCCL: + if ( + self.last_sync_step == self.train_step_num + and self.sync_method not in (SyncMethod.NCCL, SyncMethod.HCCL) + ): await self.synchronizer.notify_no_new_model_state_dict.remote() return False return await self.synchronizer.explorer_requires_sync.remote() @@ -227,20 +230,20 @@ async def sync_weight(self) -> Dict: if self.last_sync_time is not None: metrics["time/trainer_sync_interval"] = time.time() - self.last_sync_time with Timer(metrics, "time/sync_weight"): - if self.sync_method == SyncMethod.NCCL: + if self.sync_method in (SyncMethod.NCCL, SyncMethod.HCCL): result = await self.synchronizer.ready_to_nccl_sync.remote( "trainer", self.train_step_num ) if result is None: self.logger.warning( - "NCCL weight sync skipped: Explorer has stopped or is unreachable." + "Online weight sync skipped: Explorer has stopped or is unreachable." ) else: try: self.engine.sync_weight_nccl() except Exception: self.logger.warning( - "NCCL weight sync failed (Explorer may have exited);" + "Online weight sync failed (Explorer may have exited);" f" continuing with stale weights:\n{traceback.format_exc()}" ) elif self.train_step_num > 0: diff --git a/trinity/trainer/verl/fsdp_engine.py b/trinity/trainer/verl/fsdp_engine.py index c9740014d75..57aa3d5ce76 100644 --- a/trinity/trainer/verl/fsdp_engine.py +++ b/trinity/trainer/verl/fsdp_engine.py @@ -30,6 +30,7 @@ from verl.utils.fsdp_utils import get_fsdp_full_state_dict, get_fsdp_state_ctx from trinity.trainer.verl.checkpoint import CheckpointCoordinator +from verl.utils.device import get_device_name def _save_checkpoint_metadata(engine, local_path: str, logger): @@ -187,4 +188,4 @@ def fsdp_sync_weight_nccl(engine, model_update_group): torch.distributed.broadcast(full_param, 0, group=model_update_group) if torch.distributed.get_rank() == 0: - torch.cuda.synchronize() + getattr(torch, get_device_name()).synchronize() diff --git a/trinity/trainer/verl/megatron_engine.py b/trinity/trainer/verl/megatron_engine.py index 5c67d09a8f4..61ea672a7f6 100644 --- a/trinity/trainer/verl/megatron_engine.py +++ b/trinity/trainer/verl/megatron_engine.py @@ -20,6 +20,7 @@ from verl.utils.memory_utils import aggressive_empty_cache from trinity.trainer.verl.checkpoint import CheckpointCoordinator +from verl.utils.device import get_device_name def megatron_save_state_dict( @@ -115,7 +116,7 @@ def megatron_upload_state_dict(engine, synchronizer, global_step: int, logger): ray.get(synchronizer.set_model_state_dict.remote(state_dict, global_step)) torch.distributed.barrier() - torch.cuda.empty_cache() + getattr(torch, get_device_name()).empty_cache() logger.info(f"[Megatron] state_dict uploaded to Synchronizer: step={global_step}") @@ -134,4 +135,4 @@ def megatron_sync_weight_nccl(engine, model_update_group): torch.distributed.broadcast(param, src=0, group=model_update_group) if torch.distributed.get_rank() == 0: - torch.cuda.synchronize() + getattr(torch, get_device_name()).synchronize() diff --git a/trinity/trainer/verl/monkey_patch.py b/trinity/trainer/verl/monkey_patch.py index 40db776ae1c..ee36d00a0be 100644 --- a/trinity/trainer/verl/monkey_patch.py +++ b/trinity/trainer/verl/monkey_patch.py @@ -20,6 +20,7 @@ from verl.workers.utils.padding import build_attention_mask_from_nested from trinity.trainer.verl_legacy.monkey_patch import apply_monkey_patch +from trinity.utils.device import get_device_capability AutoModelForVision2Seq = get_auto_model_for_vision2seq() @@ -108,7 +109,7 @@ def _build_module(self): torch_dtype = PrecisionType.to_dtype(torch_dtype) - major_capability, _ = torch.cuda.get_device_capability(0) + major_capability = get_device_capability() use_meta = ( (self.rank != 0 if self.device_mesh is None else self.device_mesh.get_coordinate()[-1] != 0) if self.engine_config.strategy == "fsdp2" and major_capability >= 9 diff --git a/trinity/trainer/verl/trainer.py b/trinity/trainer/verl/trainer.py index 582786eebbc..84d51742ba3 100644 --- a/trinity/trainer/verl/trainer.py +++ b/trinity/trainer/verl/trainer.py @@ -57,6 +57,7 @@ from trinity.algorithm import ADVANTAGE_FN, ALGORITHM_TYPE, KL_FN from trinity.algorithm.utils import prefix_metrics from trinity.common.config import Config +from verl.utils.device import get_device_name from trinity.common.constants import SaveStrategy from trinity.common.experience import Experience from trinity.trainer.trainer import TrainEngineWrapper @@ -297,8 +298,8 @@ def __init__(self, global_config: Config): # Training steps config self.total_training_steps = global_config.trainer.total_steps or sys.maxsize - # we only support cuda for now - self.device_name = "cuda" + # Device-aware: 'npu' on Ascend NPU, 'cuda' on GPU. + self.device_name = get_device_name() checkpoint_monitor = CheckpointMonitor.get_actor( namespace=global_config.synchronizer.ray_namespace, save_strategy=global_config.trainer.save_strategy, diff --git a/trinity/trainer/verl/workers.py b/trinity/trainer/verl/workers.py index 6baa2967f0e..4e26b2aa209 100644 --- a/trinity/trainer/verl/workers.py +++ b/trinity/trainer/verl/workers.py @@ -32,6 +32,7 @@ from trinity.manager.synchronizer import Synchronizer from trinity.trainer.verl.checkpoint import CheckpointCoordinator from trinity.trainer.verl.losses import build_trinity_loss +from verl.utils.device import get_device_name from trinity.utils.distributed import WeightTransferEngine from trinity.utils.log import get_logger from trinity.utils.stream_saver import save_safetensors_streaming @@ -178,7 +179,7 @@ def _save_lora(self, local_path: str): try: if fsdp_version(model) > 0: - model = model.to(torch.cuda.current_device()) + model = model.to(getattr(torch, get_device_name()).current_device()) lora_params = layered_summon_lora_params(model) if rank == 0: save_file( @@ -350,7 +351,7 @@ def sync_weight_nccl(self): 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: 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..4411c6e63d2 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"]) ) @@ -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..938f043d325 --- /dev/null +++ b/trinity/utils/device.py @@ -0,0 +1,86 @@ +# -*- coding: utf-8 -*- +"""Device detection and abstraction layer. + +Unifies the differences among NPU / GPU / CPU devices for trinity modules. + +""" +import functools +import os +import torch +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 hasattr(torch, "npu") and torch.npu.is_available(): + return DeviceType.NPU + elif torch.cuda.is_available(): + return DeviceType.CUDA + else: + 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". + """ + 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(): + major, _ = torch.cuda.get_device_capability(0) + return major + return 0 \ No newline at end of file diff --git a/trinity/utils/distributed.py b/trinity/utils/distributed.py index 4b46218cb22..68573c39884 100644 --- a/trinity/utils/distributed.py +++ b/trinity/utils/distributed.py @@ -61,11 +61,6 @@ 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." - - from torch.distributed.distributed_c10d import 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}" @@ -134,37 +129,72 @@ def create( class VLLMWeightTransferEngine(WeightTransferEngine): - """A helper class to manage NCCL weight synchronization using vLLM's API.""" + """A helper class to manage weight synchronization using vLLM's API. + + Device-aware: uses HCCL engine on NPU and NCCL engine on GPU. Both engines + expose the same ``trainer_init`` / ``trainer_send_weights`` API surface, so + the only difference is the engine class and args dataclass used. + """ def __init__(self, master_address: str, master_port: int, world_size: int, group_name: str): - """Initialize the NCCL process group for weight sync with vLLM's API.""" - from vllm.distributed.weight_transfer.nccl_engine import ( - NCCLWeightTransferEngine, - ) + """Initialize the process group for weight sync with vLLM's API.""" + from trinity.utils.device import is_npu + + del group_name # vLLM's weight transfer engines do not require a group name + self._is_npu = is_npu() + if self._is_npu: + from vllm_ascend.distributed.weight_transfer.hccl_engine import ( + HCCLWeightTransferEngine, + ) - del group_name # vLLM's NCCL engine does not require a group name - self._model_update_group = NCCLWeightTransferEngine.trainer_init( - dict( - master_address=master_address, - master_port=master_port, - world_size=world_size, + self._engine_cls = HCCLWeightTransferEngine + self._model_update_group = HCCLWeightTransferEngine.trainer_init( + dict( + master_address=master_address, + master_port=master_port, + world_size=world_size, + ) + ) + else: + from vllm.distributed.weight_transfer.nccl_engine import ( + NCCLWeightTransferEngine, + ) + + self._engine_cls = NCCLWeightTransferEngine + self._model_update_group = NCCLWeightTransferEngine.trainer_init( + dict( + master_address=master_address, + master_port=master_port, + world_size=world_size, + ) ) - ) def sync_weight(self, iterator): - """Perform the NCCL weight sync using vLLM's API.""" - from vllm.distributed.weight_transfer.nccl_engine import ( - NCCLTrainerSendWeightsArgs, - NCCLWeightTransferEngine, - ) + """Perform the weight sync using vLLM's API (HCCL on NPU, NCCL on GPU).""" + if self._is_npu: + from vllm_ascend.distributed.weight_transfer.hccl_engine import ( + HCCLTrainerSendWeightsArgs, + ) - NCCLWeightTransferEngine.trainer_send_weights( - iterator=iterator, - trainer_args=NCCLTrainerSendWeightsArgs( - group=self._model_update_group, - packed=True, - ), - ) + self._engine_cls.trainer_send_weights( + iterator=iterator, + trainer_args=HCCLTrainerSendWeightsArgs( + group=self._model_update_group, + packed=True, + ), + ) + else: + from vllm.distributed.weight_transfer.nccl_engine import ( + NCCLTrainerSendWeightsArgs, + ) + + self._engine_cls.trainer_send_weights( + iterator=iterator, + trainer_args=NCCLTrainerSendWeightsArgs( + group=self._model_update_group, + packed=True, + ), + ) def teardown(self): self._model_update_group.destroy()