From 87489f0872dede2f30b6ef6e912bb804b159db27 Mon Sep 17 00:00:00 2001 From: zyang6 <92960729+zyang6@users.noreply.github.com> Date: Fri, 10 Oct 2025 07:25:08 +0800 Subject: [PATCH] Add wan2.1 functionality support for Ascend NPU platform (#810) Co-authored-by: kiritorl <1021709528@qq.com> --- .../dataset/parquet_dataset_map_style.py | 2 + .../device_communicators/npu_communicator.py | 71 ++++++ .../device_communicators/pyhccl.py | 146 ++++++++++++ .../device_communicators/pyhccl_wrapper.py | 208 ++++++++++++++++++ fastvideo/distributed/parallel_state.py | 38 +++- fastvideo/models/dits/wanvideo.py | 7 +- fastvideo/models/loader/component_loader.py | 13 ++ fastvideo/models/loader/fsdp_load.py | 13 +- fastvideo/platforms/__init__.py | 29 +++ fastvideo/platforms/cuda.py | 7 + fastvideo/platforms/interface.py | 22 +- fastvideo/platforms/npu.py | 131 +++++++++++ fastvideo/training/training_pipeline.py | 12 +- fastvideo/utils.py | 22 ++ 14 files changed, 690 insertions(+), 31 deletions(-) create mode 100644 fastvideo/distributed/device_communicators/npu_communicator.py create mode 100644 fastvideo/distributed/device_communicators/pyhccl.py create mode 100644 fastvideo/distributed/device_communicators/pyhccl_wrapper.py create mode 100644 fastvideo/platforms/npu.py diff --git a/fastvideo/dataset/parquet_dataset_map_style.py b/fastvideo/dataset/parquet_dataset_map_style.py index d6458e1ea..fa4a4460e 100644 --- a/fastvideo/dataset/parquet_dataset_map_style.py +++ b/fastvideo/dataset/parquet_dataset_map_style.py @@ -12,6 +12,7 @@ import tqdm # Dataset from torch.utils.data import Dataset, Sampler from torchdata.stateful_dataloader import StatefulDataLoader +from fastvideo.platforms import current_platform from fastvideo.dataset.utils import collate_rows_from_parquet_schema from fastvideo.distributed import (get_sp_world_size, get_world_group, @@ -342,6 +343,7 @@ def build_parquet_map_style_dataloader( collate_fn=passthrough, num_workers=num_data_workers, pin_memory=True, + pin_memory_device=current_platform.device_name, persistent_workers=num_data_workers > 0, ) return dataset, loader diff --git a/fastvideo/distributed/device_communicators/npu_communicator.py b/fastvideo/distributed/device_communicators/npu_communicator.py new file mode 100644 index 000000000..6a300270d --- /dev/null +++ b/fastvideo/distributed/device_communicators/npu_communicator.py @@ -0,0 +1,71 @@ +import torch +from torch.distributed import ProcessGroup + +from fastvideo.distributed.device_communicators.base_device_communicator import ( + DeviceCommunicatorBase) + + +class NpuCommunicator(DeviceCommunicatorBase): + + def __init__(self, + cpu_group: ProcessGroup, + device: torch.device | None = None, + device_group: ProcessGroup | None = None, + unique_name: str = ""): + super().__init__(cpu_group, device, device_group, unique_name) + + from fastvideo.distributed.device_communicators.pyhccl import ( + PyHcclCommunicator) + + self.pyhccl_comm: PyHcclCommunicator | None = None + if self.world_size > 1: + self.pyhccl_comm = PyHcclCommunicator( + group=self.cpu_group, + device=self.device, + ) + + def all_reduce(self, input_, op: torch.distributed.ReduceOp | None = None): + pyhccl_comm = self.pyhccl_comm + assert pyhccl_comm is not None, "pyhccl_comm should not be None" + out = pyhccl_comm.all_reduce(input_, op=op) + if out is None: + # fall back to the default all-reduce using PyTorch. + # this usually happens during testing. + # when we run the model, allreduce only happens for the TP + # group, where we always have either custom allreduce or pyhccl. + out = input_.clone() + torch.distributed.all_reduce(out, group=self.device_group, op=op) + return out + + def send(self, tensor: torch.Tensor, dst: int | None = None) -> None: + """Sends a tensor to the destination rank in a non-blocking way""" + """NOTE: `dst` is the local rank of the destination rank.""" + if dst is None: + dst = (self.rank_in_group + 1) % self.world_size + + pyhccl_comm = self.pyhccl_comm + if pyhccl_comm is not None and not pyhccl_comm.disabled: + pyhccl_comm.send(tensor, dst) + else: + torch.distributed.send(tensor, self.ranks[dst], self.device_group) + + def recv(self, + size: torch.Size, + dtype: torch.dtype, + src: int | None = None) -> torch.Tensor: + """Receives a tensor from the source rank.""" + """NOTE: `src` is the local rank of the source rank.""" + if src is None: + src = (self.rank_in_group - 1) % self.world_size + + tensor = torch.empty(size, dtype=dtype, device=self.device) + pyhccl_comm = self.pyhccl_comm + if pyhccl_comm is not None and not pyhccl_comm.disabled: + pyhccl_comm.recv(tensor, src) + else: + torch.distributed.recv(tensor, self.ranks[src], self.device_group) + return tensor + + def destroy(self) -> None: + if self.pyhccl_comm is not None: + self.pyhccl_comm = None diff --git a/fastvideo/distributed/device_communicators/pyhccl.py b/fastvideo/distributed/device_communicators/pyhccl.py new file mode 100644 index 000000000..88c784299 --- /dev/null +++ b/fastvideo/distributed/device_communicators/pyhccl.py @@ -0,0 +1,146 @@ +import torch +import torch.distributed as dist +from torch.distributed import ProcessGroup, ReduceOp + +from fastvideo.distributed.device_communicators.pyhccl_wrapper import ( + HCCLLibrary, aclrtStream_t, buffer_type, hcclComm_t, hcclDataTypeEnum, + hcclRedOpTypeEnum, hcclUniqueId) +from fastvideo.distributed.utils import StatelessProcessGroup +from fastvideo.logger import init_logger +from fastvideo.utils import current_stream + +logger = init_logger(__name__) + + +class PyHcclCommunicator: + + def __init__( + self, + group: ProcessGroup | StatelessProcessGroup, + device: int | str | torch.device, + library_path: str | None = None, + ): + """ + Args: + group: the process group to work on. If None, it will use the + default process group. + device: the device to bind the PyHcclCommunicator to. If None, + it will be bind to f"npu:{local_rank}". + library_path: the path to the HCCL library. If None, it will + use the default library path. + It is the caller's responsibility to make sure each communicator + is bind to a unique device. + """ + + if not isinstance(group, StatelessProcessGroup): + assert dist.is_initialized() + assert dist.get_backend(group) != dist.Backend.HCCL, ( + "PyHcclCommunicator should be attached to a non-HCCL group.") + # note: this rank is the rank in the group + self.rank = dist.get_rank(group) + self.world_size = dist.get_world_size(group) + else: + self.rank = group.rank + self.world_size = group.world_size + + self.group = group + + # if world_size == 1, no need to create communicator + if self.world_size == 1: + self.available = False + self.disabled = True + return + + try: + self.hccl = HCCLLibrary(library_path) + except Exception: + logger.warning("disable hccl because of missing HCCL library") + # disable because of missing HCCL library + # e.g. in a non-NPU environment + self.available = False + self.disabled = True + return + + self.available = True + self.disabled = False + + logger.info("FastVideo is using pyhccl") + + if isinstance(device, int): + device = torch.device(f"npu:{device}") + elif isinstance(device, str): + device = torch.device(device) + # now `device` is a `torch.device` object + assert isinstance(device, torch.device) + self.device = device + + if self.rank == 0: + # get the unique id from HCCL + with torch.npu.device(device): + self.unique_id = self.hccl.hcclGetUniqueId() + else: + # construct an empty unique id + self.unique_id = hcclUniqueId() + + if not isinstance(group, StatelessProcessGroup): + tensor = torch.ByteTensor(list(self.unique_id.internal)) + ranks = dist.get_process_group_ranks(group) + # arg `src` in `broadcast` is the global rank + dist.broadcast(tensor, src=ranks[0], group=group) + byte_list = tensor.tolist() + for i, byte in enumerate(byte_list): + self.unique_id.internal[i] = byte + else: + self.unique_id = group.broadcast_obj(self.unique_id, src=0) + + # hccl communicator and stream will use this device + # `torch.npu.device` is a context manager that changes the + # current npu device to the specified one + with torch.npu.device(device): + self.comm: hcclComm_t = self.hccl.hcclCommInitRank( + self.world_size, self.unique_id, self.rank) + + stream = current_stream() + # A small all_reduce for warmup. + data = torch.zeros(1, device=device) + self.all_reduce(data) + stream.synchronize() + del data + + def all_reduce(self, + in_tensor: torch.Tensor, + op: ReduceOp = ReduceOp.SUM, + stream=None) -> torch.Tensor: + if self.disabled: + return None + # hccl communicator created on a specific device + # will only work on tensors on the same device + # otherwise it will cause "illegal memory access" + assert in_tensor.device == self.device, ( + f"this hccl communicator is created to work on {self.device}, " + f"but the input tensor is on {in_tensor.device}") + + out_tensor = torch.empty_like(in_tensor) + + if stream is None: + stream = current_stream() + self.hccl.hcclAllReduce(buffer_type(in_tensor.data_ptr()), + buffer_type(out_tensor.data_ptr()), + in_tensor.numel(), + hcclDataTypeEnum.from_torch(in_tensor.dtype), + hcclRedOpTypeEnum.from_torch(op), self.comm, + aclrtStream_t(stream.npu_stream)) + return out_tensor + + def broadcast(self, tensor: torch.Tensor, src: int, stream=None): + if self.disabled: + return + assert tensor.device == self.device, ( + f"this hccl communicator is created to work on {self.device}, " + f"but the input tensor is on {tensor.device}") + if stream is None: + stream = current_stream() + buffer = buffer_type(tensor.data_ptr()) + self.hccl.hcclBroadcast(buffer, tensor.numel(), + hcclDataTypeEnum.from_torch(tensor.dtype), src, + self.comm, aclrtStream_t(stream.npu_stream)) diff --git a/fastvideo/distributed/device_communicators/pyhccl_wrapper.py b/fastvideo/distributed/device_communicators/pyhccl_wrapper.py new file mode 100644 index 000000000..3bf047e3f --- /dev/null +++ b/fastvideo/distributed/device_communicators/pyhccl_wrapper.py @@ -0,0 +1,208 @@ +import ctypes +import platform +from dataclasses import dataclass +from typing import Any + +import torch +from torch.distributed import ReduceOp + +from fastvideo.logger import init_logger +from fastvideo.utils import find_hccl_library + +logger = init_logger(__name__) + +hcclResult_t = ctypes.c_int +hcclComm_t = ctypes.c_void_p + + +class hcclUniqueId(ctypes.Structure): + _fields_ = [("internal", ctypes.c_byte * 4108)] + + +aclrtStream_t = ctypes.c_void_p +buffer_type = ctypes.c_void_p + +hcclDataType_t = ctypes.c_int + + +class hcclDataTypeEnum: + hcclInt8 = 0 + hcclInt16 = 1 + hcclInt32 = 2 + hcclFloat16 = 3 + hcclFloat32 = 4 + hcclInt64 = 5 + hcclUint64 = 6 + hcclUint8 = 7 + hcclUint16 = 8 + hcclUint32 = 9 + hcclFloat64 = 10 + hcclBfloat16 = 11 + hcclInt128 = 12 + + @classmethod + def from_torch(cls, dtype: torch.dtype) -> int: + if dtype == torch.int8: + return cls.hcclInt8 + if dtype == torch.uint8: + return cls.hcclUint8 + if dtype == torch.int32: + return cls.hcclInt32 + if dtype == torch.int64: + return cls.hcclInt64 + if dtype == torch.float16: + return cls.hcclFloat16 + if dtype == torch.float32: + return cls.hcclFloat32 + if dtype == torch.float64: + return cls.hcclFloat64 + if dtype == torch.bfloat16: + return cls.hcclBfloat16 + raise ValueError(f"Unsupported dtype: {dtype}") + + +hcclRedOp_t = ctypes.c_int + + +class hcclRedOpTypeEnum: + hcclSum = 0 + hcclProd = 1 + hcclMax = 2 + hcclMin = 3 + + @classmethod + def from_torch(cls, op: ReduceOp) -> int: + if op == ReduceOp.SUM: + return cls.hcclSum + if op == ReduceOp.PRODUCT: + return cls.hcclProd + if op == ReduceOp.MAX: + return cls.hcclMax + if op == ReduceOp.MIN: + return cls.hcclMin + raise ValueError(f"Unsupported op: {op}") + + +@dataclass +class Function: + name: str + restype: Any + argtypes: list[Any] + + +class HCCLLibrary: + exported_functions = [ + Function("HcclGetErrorString", ctypes.c_char_p, [hcclResult_t]), + Function("HcclGetRootInfo", hcclResult_t, + [ctypes.POINTER(hcclUniqueId)]), + Function("HcclCommInitRootInfo", hcclResult_t, [ + ctypes.c_int, + ctypes.POINTER(hcclUniqueId), + ctypes.c_int, + ctypes.POINTER(hcclComm_t), + ]), + Function("HcclAllReduce", hcclResult_t, [ + buffer_type, + buffer_type, + ctypes.c_size_t, + hcclDataType_t, + hcclRedOp_t, + hcclComm_t, + aclrtStream_t, + ]), + Function("HcclBroadcast", hcclResult_t, [ + buffer_type, + ctypes.c_size_t, + hcclDataType_t, + ctypes.c_int, + hcclComm_t, + aclrtStream_t, + ]), + Function("HcclCommDestroy", hcclResult_t, [hcclComm_t]), + ] + + # class attribute to store the mapping from the path to the library + # to avoid loading the same library multiple times + path_to_library_cache: dict[str, Any] = {} + + # class attribute to store the mapping from library path + # to the correspongding directory + path_to_dict_mapping: dict[str, dict[str, Any]] = {} + + def __init__(self, so_file: str | None = None): + + so_file = so_file or find_hccl_library() + + try: + if so_file not in HCCLLibrary.path_to_dict_mapping: + lib = ctypes.CDLL(so_file) + HCCLLibrary.path_to_library_cache[so_file] = lib + self.lib = HCCLLibrary.path_to_library_cache[so_file] + except Exception as e: + logger.error( + "Failed to load HCCL library from %s. " + "It is expected if you are not running on Ascend NPUs." + "Otherwise, the hccl library might not exist, be corrupted " + "or it does not support the current platform %s. " + "If you already have the library, please set the " + "environment variable HCCL_SO_PATH" + " to point to the correct hccl library path.", so_file, + platform.platform()) + raise e + + if so_file not in HCCLLibrary.path_to_dict_mapping: + _funcs: dict[str, Any] = {} + for func in HCCLLibrary.exported_functions: + f = getattr(self.lib, func.name) + f.restype = func.restype + f.argtypes = func.argtypes + _funcs[func.name] = f + HCCLLibrary.path_to_dict_mapping[so_file] = _funcs + self._funcs = HCCLLibrary.path_to_dict_mapping[so_file] + + def hcclGetErrorString(self, result: hcclResult_t) -> str: + return self._funcs["HcclGetErrorString"](result).decode("utf-8") + + def HCCL_CHECK(self, result: hcclResult_t) -> None: + if result != 0: + error_str = self.hcclGetErrorString(result) + raise RuntimeError(f"HCCL error: {error_str}") + + def hcclGetUniqueId(self) -> hcclUniqueId: + unique_id = hcclUniqueId() + self.HCCL_CHECK(self._funcs["HcclGetRootInfo"](ctypes.byref(unique_id))) + return unique_id + + def hcclCommInitRank(self, world_size: int, unique_id: hcclUniqueId, + rank: int) -> hcclComm_t: + comm = hcclComm_t() + self.HCCL_CHECK(self._funcs["HcclCommInitRootInfo"]( + world_size, ctypes.byref(unique_id), rank, ctypes.byref(comm))) + return comm + + def hcclAllReduce(self, sendbuff: buffer_type, recvbuff: buffer_type, + count: int, datatype: int, op: int, comm: hcclComm_t, + stream: aclrtStream_t) -> None: + self.HCCL_CHECK(self._funcs["HcclAllReduce"](sendbuff, recvbuff, count, + datatype, op, comm, + stream)) + + def hcclBroadcast(self, buf: buffer_type, count: int, datatype: int, + root: int, comm: hcclComm_t, + stream: aclrtStream_t) -> None: + self.HCCL_CHECK(self._funcs["HcclBroadcast"](buf, count, datatype, root, + comm, stream)) + + def hcclCommDestroy(self, comm: hcclComm_t) -> None: + self.HCCL_CHECK(self._funcs["HcclCommDestroy"](comm)) + + +__all__ = [ + "HCCLLibrary", + "hcclDataTypeEnum", + "hcclRedOpTypeEnum", + "hcclUniqueId", + "hcclComm_t", + "aclrtStream_t", + "buffer_type", +] diff --git a/fastvideo/distributed/parallel_state.py b/fastvideo/distributed/parallel_state.py index aa15e2904..cfa0ef713 100644 --- a/fastvideo/distributed/parallel_state.py +++ b/fastvideo/distributed/parallel_state.py @@ -203,6 +203,15 @@ class GroupCoordinator: device_group=self.device_group, unique_name=self.unique_name, ) + elif current_platform.is_npu(): + from fastvideo.distributed.device_communicators.npu_communicator import ( + NpuCommunicator) + self.device_communicator = NpuCommunicator( + cpu_group=self.cpu_group, + device=self.device, + device_group=self.device_group, + unique_name=self.unique_name, + ) else: # For MPS and CPU, use the CPU communicator self.device_communicator = CpuCommunicator( @@ -776,8 +785,13 @@ def init_distributed_environment( ): # Determine the appropriate backend based on the platform from fastvideo.platforms import current_platform - if backend == "nccl" and not current_platform.is_cuda_alike(): - # Use gloo backend for non-CUDA platforms (MPS, CPU) + backend = "nccl" + if current_platform.is_cuda_alike(): + logger.info("Using nccl backend for CUDA platform") + elif current_platform.is_npu(): + backend = "hccl" + logger.info("Using hccl backend for NPU platform") + else: backend = "gloo" logger.info("Using gloo backend for %s platform", current_platform.device_name) @@ -792,7 +806,7 @@ def init_distributed_environment( "distributed environment") # For MPS, don't pass device_id as it doesn't support device indices - if current_platform.is_mps(): + if current_platform.is_mps() or current_platform.is_npu(): torch.distributed.init_process_group( backend=backend, init_method=distributed_init_method, @@ -948,9 +962,13 @@ def get_dp_rank() -> int: def get_local_torch_device() -> torch.device: """Return the torch device for the current rank.""" - return torch.device(f"cuda:{envs.LOCAL_RANK}" - ) if current_platform.is_cuda_alike() else torch.device( - "mps") + if current_platform.is_npu(): + device = torch.device(f"npu:{envs.LOCAL_RANK}") + elif current_platform.is_cuda_alike() or current_platform.is_cuda(): + device = torch.device(f"cuda:{envs.LOCAL_RANK}") + else: + device = torch.device("mps") + return device def maybe_init_distributed_environment_and_model_parallel( @@ -979,10 +997,10 @@ def maybe_init_distributed_environment_and_model_parallel( initialize_model_parallel(tensor_model_parallel_size=tp_size, sequence_model_parallel_size=sp_size) - # Only set CUDA device if we're on a CUDA platform - if current_platform.is_cuda_alike(): - device = torch.device(f"cuda:{local_rank}") - torch.cuda.set_device(device) + # set device if we're on a CUDA/NPU platform + device_type = "npu" if current_platform.is_npu() else "cuda" + device = torch.device(f"{device_type}:{local_rank}") + current_platform.get_torch_device().set_device(device) def model_parallel_is_initialized() -> bool: diff --git a/fastvideo/models/dits/wanvideo.py b/fastvideo/models/dits/wanvideo.py index ac79e62e8..cd4e59466 100644 --- a/fastvideo/models/dits/wanvideo.py +++ b/fastvideo/models/dits/wanvideo.py @@ -685,9 +685,10 @@ class WanTransformer3DModel(CachableDiT): encoder_hidden_states = torch.concat( [encoder_hidden_states_image, encoder_hidden_states], dim=1) - encoder_hidden_states = encoder_hidden_states.to( - orig_dtype) if current_platform.is_mps( - ) else encoder_hidden_states # cast to orig_dtype for MPS + if current_platform.is_mps() or current_platform.is_npu(): + encoder_hidden_states = encoder_hidden_states.to(orig_dtype) + else: + encoder_hidden_states = encoder_hidden_states # cast to orig_dtype for MPS & NPU assert encoder_hidden_states.dtype == orig_dtype diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index 5f0186c78..8eb75fcdc 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -280,6 +280,19 @@ class TextEncoderLoader(ComponentLoader): logger.info( "Disabling FSDP sharding for MPS platform as it's not compatible" ) + elif current_platform.is_npu(): + mesh = init_device_mesh( + "npu", + mesh_shape=(1, dist.get_world_size()), + mesh_dim_names=("offload", "replicate"), + ) + shard_model( + model, + cpu_offload=True, + reshard_after_forward=True, + mesh=mesh["offload"], + fsdp_shard_conditions=model._fsdp_shard_conditions, + pin_cpu_memory=fastvideo_args.pin_cpu_memory) else: mesh = init_device_mesh( "cuda", diff --git a/fastvideo/models/loader/fsdp_load.py b/fastvideo/models/loader/fsdp_load.py index 3913f7275..f8002392a 100644 --- a/fastvideo/models/loader/fsdp_load.py +++ b/fastvideo/models/loader/fsdp_load.py @@ -108,8 +108,17 @@ def maybe_load_fsdp_model( if not training_mode and not fsdp_inference: hsdp_replicate_dim = world_size hsdp_shard_dim = 1 - - device_mesh = init_device_mesh( + + if current_platform.is_npu(): + with torch.device("cpu"): + device_mesh = init_device_mesh( + "npu", + # (Replicate(), Shard(dim=0)) + mesh_shape=(hsdp_replicate_dim, hsdp_shard_dim), + mesh_dim_names=("replicate", "shard"), + ) + else: + device_mesh = init_device_mesh( "cuda", # (Replicate(), Shard(dim=0)) mesh_shape=(hsdp_replicate_dim, hsdp_shard_dim), diff --git a/fastvideo/platforms/__init__.py b/fastvideo/platforms/__init__.py index 2decdbbf6..2f7266e6d 100644 --- a/fastvideo/platforms/__init__.py +++ b/fastvideo/platforms/__init__.py @@ -64,6 +64,29 @@ def mps_platform_plugin() -> str | None: return "fastvideo.platforms.mps.MpsPlatform" if is_mps else None +def npu_platform_plugin() -> str | None: + is_npu = False + + try: + import torch + # 导入 torch_npu 以初始化 NPU 后端 + import torch_npu # noqa: F401 + if torch.npu.is_available(): + is_npu = True + logger.info("NPU is available") + except ImportError: + logger.error( + "NPU detection failed: PyTorch or PyTorch_NPU is not installed") + except AttributeError: + logger.error( + "NPU detection failed: PyTorch has no 'npu' attribute (use Ascend-adapted PyTorch)" + ) + except Exception as e: + logger.error("NPU detection failed: unknown error - %s", str(e)) + + return "fastvideo.platforms.npu.NPUPlatform" if is_npu else None + + def cpu_platform_plugin() -> str | None: """Detect if CPU platform should be used.""" # CPU is always available as a fallback @@ -93,6 +116,7 @@ builtin_platform_plugins = { 'rocm': rocm_platform_plugin, 'mps': mps_platform_plugin, 'cpu': cpu_platform_plugin, + 'npu': npu_platform_plugin, } @@ -115,6 +139,11 @@ def resolve_current_platform_cls_qualname() -> str: if platform_cls_qualname is not None: return platform_cls_qualname + # Fall back to NPU + platform_cls_qualname = npu_platform_plugin() + if platform_cls_qualname is not None: + return platform_cls_qualname + # Fall back to CPU as last resort platform_cls_qualname = cpu_platform_plugin() if platform_cls_qualname is not None: diff --git a/fastvideo/platforms/cuda.py b/fastvideo/platforms/cuda.py index 8bb58bb7a..731a20631 100644 --- a/fastvideo/platforms/cuda.py +++ b/fastvideo/platforms/cuda.py @@ -107,6 +107,13 @@ class CudaPlatformBase(Platform): torch.cuda.reset_peak_memory_stats(device) return float(torch.cuda.max_memory_allocated(device)) + @classmethod + def get_torch_device(cls): + """ + Return torch.cuda + """ + return torch.cuda + @classmethod def get_attn_backend_cls(cls, selected_backend: AttentionBackendEnum | None, head_size: int, dtype: torch.dtype) -> str: diff --git a/fastvideo/platforms/interface.py b/fastvideo/platforms/interface.py index a8675325b..35221c563 100644 --- a/fastvideo/platforms/interface.py +++ b/fastvideo/platforms/interface.py @@ -1,6 +1,3 @@ -# SPDX-License-Identifier: Apache-2.0 -# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/platforms/interface.py - import enum import random from typing import NamedTuple @@ -32,6 +29,7 @@ class PlatformEnum(enum.Enum): MPS = enum.auto() OOT = enum.auto() UNSPECIFIED = enum.auto() + NPU = enum.auto() class CpuArchEnum(enum.Enum): @@ -62,16 +60,8 @@ class Platform: device_name: str device_type: str - # available dispatch keys: - # check https://github.com/pytorch/pytorch/blob/313dac6c1ca0fa0cde32477509cce32089f8532a/torchgen/model.py#L134 # noqa - # use "CPU" as a fallback for platforms not registered in PyTorch dispatch_key: str = "CPU" - # The torch.compile backend for compiling simple and - # standalone functions. The default value is "inductor" to keep - # the same behavior as PyTorch. - # NOTE: for the forward part of the model, vLLM has another separate - # compilation strategy. simple_compile_backend: str = "inductor" supported_quantization: list[str] = [] @@ -98,6 +88,9 @@ class Platform: def is_mps(self) -> bool: return self._enum == PlatformEnum.MPS + def is_npu(self) -> bool: + return self._enum == PlatformEnum.NPU + @classmethod def get_attn_backend_cls(cls, selected_backend: AttentionBackendEnum | None, head_size: int, dtype: torch.dtype) -> str: @@ -157,6 +150,13 @@ class Platform: """ raise NotImplementedError + @classmethod + def get_torch_device(cls): + """ + Check if the current platform supports torch device. + """ + raise NotImplementedError + @classmethod def inference_mode(cls): """A device-specific wrapper of `torch.inference_mode`. diff --git a/fastvideo/platforms/npu.py b/fastvideo/platforms/npu.py new file mode 100644 index 000000000..4980b975e --- /dev/null +++ b/fastvideo/platforms/npu.py @@ -0,0 +1,131 @@ +import gc +from datetime import timedelta + +import torch +from torch.distributed import ProcessGroup +from torch.distributed.distributed_c10d import PrefixStore + +import fastvideo.envs as envs +from fastvideo.logger import init_logger +from fastvideo.platforms.interface import (AttentionBackendEnum, Platform, + PlatformEnum) + +logger = init_logger(__name__) + + +class NPUPlatform(Platform): + + _enum = PlatformEnum.NPU + device_name: str = "npu" + device_type: str = "npu" + simple_compile_backend: str = "eager" # Disable torch.compile() + ray_device_key: str = "NPU" + device_control_env_var: str = "ASCEND_RT_VISIBLE_DEVICES" + dispatch_key: str = "PrivateUse1" + + def is_sleep_mode_available(self) -> bool: + return True + + @classmethod + def get_device_capability(cls, device_id: int = 0): + return None + + @classmethod + def get_device_name(cls, device_id: int = 0) -> str: + return torch.npu.get_device_name(device_id) + + @classmethod + def is_async_output_supported(cls, enforce_eager: bool | None) -> bool: + return True + + @classmethod + def inference_mode(cls): + return torch.inference_mode() + + @classmethod + def set_device(cls, device: torch.device): + torch.npu.set_device(device) + + @classmethod + def empty_cache(cls): + torch.npu.empty_cache() + + @classmethod + def synchronize(cls): + torch.npu.synchronize() + + @classmethod + def mem_get_info(cls) -> tuple[int, int]: + return torch.npu.mem_get_info() + + @classmethod + def clear_npu_memory(cls): + gc.collect() + torch.npu.empty_cache() + torch.npu.reset_peak_memory_stats() + + @classmethod + def get_attn_backend_cls(cls, selected_backend: AttentionBackendEnum | None, + head_size: int, dtype: torch.dtype) -> str: + logger.info("Trying FASTVIDEO_ATTENTION_BACKEND=%s", + envs.FASTVIDEO_ATTENTION_BACKEND) + if envs.FASTVIDEO_ATTENTION_BACKEND != "TORCH_SDPA": + logger.info("Ascend NPU only supports the Torch SDPA backend.") + else: + logger.info("Using Torch SDPA backend.") + return "fastvideo.attention.backends.sdpa.SDPABackend" + + @classmethod + def get_current_memory_usage(cls, + device: torch.types.Device | None = None + ) -> float: + torch.npu.reset_peak_memory_stats(device) + return torch.npu.max_memory_allocated(device) + + @classmethod + def get_device_communicator_cls(cls) -> str: + return "fastvideo.distributed.device_communicators.npu_communicator.NpuCommunicator" + + @classmethod + def is_pin_memory_available(cls): + return True + + @classmethod + def get_torch_device(cls): + """ + Return torch.npu + """ + return torch.npu + + @classmethod + def stateless_init_device_torch_dist_pg( + cls, + backend: str, + prefix_store: PrefixStore, + group_rank: int, + group_size: int, + timeout: timedelta, + ) -> ProcessGroup: + from torch.distributed import is_hccl_available + from torch_npu._C._distributed_c10d import ProcessGroupHCCL + + assert is_hccl_available() + options = ProcessGroup.Options(backend=backend) + pg: ProcessGroup = ProcessGroup( + prefix_store, + group_rank, + group_size, + options, + ) + + backend_options = ProcessGroupHCCL.Options() + backend_options._timeout = timeout + + backend_class = ProcessGroupHCCL(prefix_store, group_rank, group_size, + backend_options) + device = torch.device("npu") + backend_class._set_sequence_number_for_group() + backend_type = ProcessGroup.BackendType.CUSTOM + + pg._register_backend(device, backend_type, backend_class) + return pg diff --git a/fastvideo/training/training_pipeline.py b/fastvideo/training/training_pipeline.py index 731febe1b..da1362f58 100644 --- a/fastvideo/training/training_pipeline.py +++ b/fastvideo/training/training_pipeline.py @@ -35,6 +35,7 @@ from fastvideo.forward_context import set_forward_context from fastvideo.logger import init_logger from fastvideo.pipelines import (ComposedPipelineBase, ForwardBatch, LoRAPipeline, TrainingBatch) +from fastvideo.platforms import current_platform from fastvideo.training.activation_checkpoint import ( apply_activation_checkpointing) from fastvideo.training.training_utils import ( @@ -260,7 +261,6 @@ class TrainingPipeline(LoRAPipeline, ABC): encoder_hidden_states = batch['text_embedding'] encoder_attention_mask = batch['text_attention_mask'] infos = batch['info_list'] - training_batch.latents = latents.to(get_local_torch_device(), dtype=torch.bfloat16) training_batch.encoder_hidden_states = encoder_hidden_states.to( @@ -561,8 +561,8 @@ class TrainingPipeline(LoRAPipeline, ABC): # Set random seeds for deterministic training self.noise_random_generator = torch.Generator(device="cpu").manual_seed( self.seed) - self.noise_gen_cuda = torch.Generator(device="cuda").manual_seed( - self.seed) + self.noise_gen_cuda = torch.Generator( + device=current_platform.device_name).manual_seed(self.seed) self.validation_random_generator = torch.Generator( device="cpu").manual_seed(self.seed) logger.info("Initialized random seeds with seed: %s", self.seed) @@ -653,7 +653,8 @@ class TrainingPipeline(LoRAPipeline, ABC): training_batch, self.training_args, step) self._log_validation(self.transformer, self.training_args, step) - gpu_memory_usage = torch.cuda.memory_allocated() / 1024**2 + gpu_memory_usage = current_platform.get_torch_device( + ).memory_allocated() / 1024**2 trainable_params = round( count_trainable(self.transformer) / 1e9, 3) logger.info( @@ -702,7 +703,8 @@ class TrainingPipeline(LoRAPipeline, ABC): logger.info(" Master weight dtype: %s", self.transformer.parameters().__next__().dtype) - gpu_memory_usage = torch.cuda.memory_allocated() / 1024**2 + gpu_memory_usage = current_platform.get_torch_device().memory_allocated( + ) / 1024**2 logger.info("GPU memory usage before train_one_step: %s MB", gpu_memory_usage) logger.info("VSA validation sparsity: %s", diff --git a/fastvideo/utils.py b/fastvideo/utils.py index d27451ff0..9a8109009 100644 --- a/fastvideo/utils.py +++ b/fastvideo/utils.py @@ -80,6 +80,28 @@ def find_nccl_library() -> str: return str(so_file) +def find_hccl_library() -> str: + """ + We either use the library file specified by the `HCCL_SO_PATH` + environment variable, or we find the library file brought by PyTorch. + After importing `torch`, `libhccl.so` can be + found by `ctypes` automatically. + """ + so_file = envs.HCCL_SO_PATH + + # manually load the nccl library + if so_file: + logger.info("Found hccl from environment variable HCCL_SO_PATH=%s", + so_file) + else: + if torch.version.cann is not None: # codespell:ignore cann + so_file = "libhccl.so" + else: + raise ValueError("HCCL only supports Ascend NPU backends.") + logger.info("Found hccl from library %s", so_file) + return so_file + + prev_set_stream = torch.cuda.set_stream _current_stream = None