Add wan2.1 functionality support for Ascend NPU platform (#810)

Co-authored-by: kiritorl <1021709528@qq.com>

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

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

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))

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",
+]

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:

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
 

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",

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),

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:

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:

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`.

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

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",

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