From acffa4f871aa51a7a52571bd0c1e6badc62c4bda Mon Sep 17 00:00:00 2001 From: root Date: Tue, 28 Jul 2026 11:36:15 +0800 Subject: [PATCH 1/2] feat(npu): add Ascend 950 platform abstraction and Qwen-Image support 1. Add NPU 950 platform abstraction layer with device registry and auto_detect_device 2. Support mindiesd_attention and mindiesd_compile with capability detection and compile_backend 3. Adapt Qwen-Image model to use platform compile kwargs and NPU attention backend --- diffsynth_engine/configs/pipeline.py | 1 + diffsynth_engine/models/basic/attention.py | 19 ++ diffsynth_engine/pipelines/qwen_image.py | 4 +- diffsynth_engine/platforms/__init__.py | 148 +++++++++++++ diffsynth_engine/platforms/ascend.py | 228 +++++++++++++++++++++ diffsynth_engine/platforms/base.py | 68 ++++++ diffsynth_engine/utils/flag.py | 5 + diffsynth_engine/utils/platform.py | 49 +++-- 8 files changed, 505 insertions(+), 17 deletions(-) create mode 100644 diffsynth_engine/platforms/__init__.py create mode 100644 diffsynth_engine/platforms/ascend.py create mode 100644 diffsynth_engine/platforms/base.py diff --git a/diffsynth_engine/configs/pipeline.py b/diffsynth_engine/configs/pipeline.py index 37da60a1..5ea5f7f4 100644 --- a/diffsynth_engine/configs/pipeline.py +++ b/diffsynth_engine/configs/pipeline.py @@ -35,6 +35,7 @@ class AttnImpl(Enum): SAGE = "sage" # Sage Attention SPARGE = "sparge" # Sparge Attention VSA = "vsa" # Video Sparse Attention + MINDIE = "mindie" # Mindie Attention @dataclass diff --git a/diffsynth_engine/models/basic/attention.py b/diffsynth_engine/models/basic/attention.py index 3d3d49cf..d472b5fd 100644 --- a/diffsynth_engine/models/basic/attention.py +++ b/diffsynth_engine/models/basic/attention.py @@ -15,6 +15,7 @@ SPARGE_ATTN_AVAILABLE, VIDEO_SPARSE_ATTN_AVAILABLE, AITER_AVAILABLE, + MINDIE_AVAILABLE, ) from diffsynth_engine.utils.platform import DTYPE_FP8 @@ -107,6 +108,16 @@ def sparge_attn( distributed_video_sparse_attn, ) +if MINDIE_AVAILABLE: + from mindiesd.layers.flash_attn.attention_forward import attention_forward + + def mindie_attn(q, k, v, attn_mask=None, scale=None): + return attention_forward( + query=q, key=k, value=v, + attn_mask=attn_mask, scale=scale, + fused=True, head_first=False, + ) + def eager_attn(q, k, v, attn_mask=None, scale=None): q = q.transpose(1, 2) @@ -152,6 +163,7 @@ def attention( "sage", "sparge", "vsa", + "mindie", ] flash_attn3_compatible = q.shape[-1] <= FA3_MAX_HEADDIM if attn_impl is None or attn_impl == "auto": @@ -192,6 +204,8 @@ def attention( ) if XFORMERS_AVAILABLE: return xformers_attn(q, k, v, attn_mask=attn_mask, scale=scale) + if MINDIE_AVAILABLE: + return mindie_attn(q, k, v, attn_mask=attn_mask, scale=scale) if SDPA_AVAILABLE: return sdpa_attn(q, k, v, attn_mask=attn_mask, scale=scale) if FLASH_ATTN_2_AVAILABLE: @@ -263,6 +277,8 @@ def attention( cdfthreshd=kwargs.get("cdfthreshd", 0.98), pvthreshd=kwargs.get("pvthreshd", 50), ) + if attn_impl == "mindie": + return mindie_attn(q, k, v, attn_mask=attn_mask, scale=scale) if attn_impl == "vsa": return video_sparse_attn( q, @@ -354,7 +370,10 @@ def long_context_attention( "sage", "sparge", "vsa", + "mindie", ] + if attn_impl == "mindie": + raise RuntimeError("mindie long-context attention is not supported yet") assert attn_mask is None, "long context attention does not support attention mask" flash_attn3_compatible = q.shape[-1] <= FA3_MAX_HEADDIM if attn_impl is None or attn_impl == "auto": diff --git a/diffsynth_engine/pipelines/qwen_image.py b/diffsynth_engine/pipelines/qwen_image.py index 342cb70b..79741547 100644 --- a/diffsynth_engine/pipelines/qwen_image.py +++ b/diffsynth_engine/pipelines/qwen_image.py @@ -376,7 +376,9 @@ def update_weights(self, state_dicts: QwenImageStateDicts) -> None: self.update_component(self.vae, state_dicts.vae, self.config.device, self.config.vae_dtype) def compile(self): - self.dit.compile_repeated_blocks() + from diffsynth_engine.platforms import resolve_platform + platform_cls = resolve_platform(self.config.device) + self.dit.compile_repeated_blocks(**platform_cls.compile_kwargs()) def load_loras(self, lora_list: List[Tuple[str, float]], fused: bool = True, save_original_weight: bool = False): assert self.config.tp_degree is None or self.config.tp_degree == 1, ( diff --git a/diffsynth_engine/platforms/__init__.py b/diffsynth_engine/platforms/__init__.py new file mode 100644 index 00000000..c90a953b --- /dev/null +++ b/diffsynth_engine/platforms/__init__.py @@ -0,0 +1,148 @@ +from __future__ import annotations + +import platform as host_platform +from functools import lru_cache +from typing import Type + +import torch + +from .ascend import ( + AscendPlatform, + probe_ascend_capabilities, + probe_ascend_feature, + reset_ascend_capability_cache, +) +from .base import PlatformBackend, PlatformCapabilities + + +class CPUPlatform(PlatformBackend): + name = "cpu" + device_type = "cpu" + + +class CUDAPlatform(PlatformBackend): + name = "cuda" + device_type = "cuda" + + @classmethod + def is_available(cls) -> bool: + return torch.cuda.is_available() + + @classmethod + def set_device(cls, index: int | str | torch.device) -> None: + torch.cuda.set_device(index) + + @classmethod + def synchronize(cls) -> None: + torch.cuda.synchronize() + + @classmethod + def empty_cache(cls) -> None: + torch.cuda.empty_cache() + + @classmethod + def distributed_backend(cls) -> str: + return "nccl" + + +class ROCmPlatform(CUDAPlatform): + name = "rocm" + + +class MPSPlatform(PlatformBackend): + name = "mps" + device_type = "mps" + + @classmethod + def is_available(cls) -> bool: + return torch.backends.mps.is_available() + + @classmethod + def synchronize(cls) -> None: + torch.mps.synchronize() + + @classmethod + def empty_cache(cls) -> None: + torch.mps.empty_cache() + + +_PLATFORM_REGISTRY: dict[str, Type[PlatformBackend]] = { + "cpu": CPUPlatform, + "cuda": ROCmPlatform if torch.version.hip else CUDAPlatform, + "mps": MPSPlatform, + "npu": AscendPlatform, +} + + +def register_platform(device_type: str, platform_cls: Type[PlatformBackend], *, overwrite: bool = False) -> None: + if device_type in _PLATFORM_REGISTRY and not overwrite: + raise ValueError(f"Platform for device type {device_type!r} is already registered") + _PLATFORM_REGISTRY[device_type] = platform_cls + + +@lru_cache(maxsize=None) +def auto_detect_device() -> str: + """auto detect device type in order of cuda(gpu/rocm), npu, mps, cpu""" + for device_type in ("cuda", "npu", "mps", "cpu"): + try: + if resolve_platform(device_type).is_available(): + return device_type + except Exception: + continue + return "cpu" + + +def get_device_type(device: str | torch.device | None = None) -> str: + if device is None or (isinstance(device, str) and device.lower() in ("auto", "")): + return auto_detect_device() + if isinstance(device, torch.device): + return device.type + return str(device).split(":", 1)[0].lower() + + +def resolve_platform(device: str | torch.device) -> Type[PlatformBackend]: + device_type = get_device_type(device) + try: + return _PLATFORM_REGISTRY[device_type] + except KeyError as exc: + available = ", ".join(sorted(_PLATFORM_REGISTRY)) + raise ValueError(f"Unsupported device type {device_type!r}. Registered device types: {available}") from exc + + +def get_preferred_fp8_dtype(device: str | torch.device = "cuda") -> torch.dtype: + platform_cls = resolve_platform(device) + if platform_cls is ROCmPlatform and platform_cls.is_available(): + properties = torch.cuda.get_device_properties(0) + if "gfx94" in properties.gcnArchName: + return torch.float8_e4m3fnuz + return torch.float8_e4m3fn + + +def pin_memory( + tensor: torch.Tensor, + device: str | torch.device | None = None, +) -> torch.Tensor: + if host_platform.system() != "Linux": + return tensor + platform_cls = resolve_platform(get_device_type(device)) + return platform_cls.pin_memory(tensor) + + +__all__ = [ + "AscendPlatform", + "CPUPlatform", + "CUDAPlatform", + "MPSPlatform", + "PlatformBackend", + "PlatformCapabilities", + "ROCmPlatform", + "auto_detect_device", + "get_device_type", + "get_preferred_fp8_dtype", + "pin_memory", + "probe_ascend_capabilities", + "probe_ascend_feature", + "register_platform", + "reset_ascend_capability_cache", + "resolve_platform", +] diff --git a/diffsynth_engine/platforms/ascend.py b/diffsynth_engine/platforms/ascend.py new file mode 100644 index 00000000..bf3964a6 --- /dev/null +++ b/diffsynth_engine/platforms/ascend.py @@ -0,0 +1,228 @@ +from __future__ import annotations + +import importlib +from functools import lru_cache +from typing import Any + +import torch + +from .base import PlatformBackend, PlatformCapabilities + + +def _import_torch_npu(): + try: + return importlib.import_module("torch_npu") + except (ImportError, OSError) as exc: + raise RuntimeError( + "Ascend device requested, but torch_npu is not installed. " + "Install the torch_npu wheel matching the PyTorch and CANN versions." + ) from exc + + +def _import_mindie_sd(): + try: + return importlib.import_module("mindiesd") + except (ImportError, OSError) as exc: + raise RuntimeError( + "This Ascend feature requires MindIE-SD. Install a MindIE-SD 3.x wheel " + "matching the current torch_npu and CANN versions." + ) from exc + + +def _has_callable(obj: Any, name: str) -> bool: + return callable(getattr(obj, name, None)) + + +def _probe_npu_runtime(npu: Any) -> bool: + try: + probe = torch.zeros(1, device="npu:0") + probe.add_(1) + npu.synchronize() + return True + except Exception: + return False + + +@lru_cache(maxsize=1) +def _probe_ascend_device() -> bool: + try: + torch_npu = _import_torch_npu() + npu = getattr(torch_npu, "npu", getattr(torch, "npu", None)) + device_available = bool(npu is not None and _has_callable(npu, "is_available") and npu.is_available()) + except Exception: + return False + + if device_available: + # is_available() can stay true while ACL initialization is failing. + device_available = _probe_npu_runtime(npu) + return device_available + + +@lru_cache(maxsize=1) +def _probe_mindie_installation() -> bool: + if not _probe_ascend_device(): + return False + + try: + _import_mindie_sd() + except Exception: + return False + return True + + +def _feature_api_available(feature: str) -> bool: + if feature == "mindie_attention": + module = importlib.import_module("mindiesd.layers.flash_attn.attention_forward") + return _has_callable(module, "attention_forward") + if feature == "mindie_compile": + module = importlib.import_module("mindiesd.compilation") + return callable(getattr(module, "MindieSDBackend", None)) + + raise ValueError(f"Unknown Ascend capability: {feature}") + + +def _tensor_probe_succeeded(output: torch.Tensor) -> bool: + torch_npu = _import_torch_npu() + torch_npu.npu.synchronize() + return bool(torch.isfinite(output).all().cpu().item()) + + +def _probe_mindie_attention_operation() -> bool: + module = importlib.import_module("mindiesd.layers.flash_attn.attention_forward") + query = torch.randn(1, 128, 8, 128, device="npu:0", dtype=torch.bfloat16) + with torch.no_grad(): + output = module.attention_forward( + query=query, + key=query, + value=query, + attn_mask=None, + scale=None, + fused=True, + head_first=False, + ) + return output.shape == query.shape and _tensor_probe_succeeded(output) + + +def _probe_mindie_compile_operation() -> bool: + compilation_module = importlib.import_module("mindiesd.compilation") + + def probe_fn(value): + return torch.nn.functional.gelu(value + 1) + + compiled_fn = torch.compile(probe_fn, backend=compilation_module.MindieSDBackend(), fullgraph=False) + value = torch.randn(8, 32, device="npu:0", dtype=torch.bfloat16) + with torch.no_grad(): + output = compiled_fn(value) + return output.shape == value.shape and _tensor_probe_succeeded(output) + + + + +_OPERATION_PROBES = { + "mindie_attention": _probe_mindie_attention_operation, + "mindie_compile": _probe_mindie_compile_operation, +} + + +@lru_cache(maxsize=None) +def probe_ascend_feature(feature: str) -> bool: + if feature == "device": + return _probe_ascend_device() + if feature == "mindie": + return _probe_mindie_installation() + if feature not in _OPERATION_PROBES: + raise ValueError(f"Unknown Ascend capability: {feature}") + if not _probe_mindie_installation(): + return False + + try: + if not _feature_api_available(feature): + return False + return bool(_OPERATION_PROBES[feature]()) + except Exception: + return False + + +def probe_ascend_capabilities() -> PlatformCapabilities: + device = probe_ascend_feature("device") + if not device: + return PlatformCapabilities() + mindie = probe_ascend_feature("mindie") + if not mindie: + return PlatformCapabilities(device=True) + + return PlatformCapabilities( + device=True, + mindie=True, + mindie_attention=probe_ascend_feature("mindie_attention"), + mindie_compile=probe_ascend_feature("mindie_compile"), + ) + + +def reset_ascend_capability_cache() -> None: + _probe_ascend_device.cache_clear() + _probe_mindie_installation.cache_clear() + probe_ascend_feature.cache_clear() + + +class AscendPlatform(PlatformBackend): + name = "ascend" + device_type = "npu" + + @classmethod + def is_available(cls) -> bool: + return probe_ascend_feature("device") + + @classmethod + def normalize_device(cls, device: str | torch.device) -> torch.device: + _import_torch_npu() + return torch.device(device) + + @classmethod + def set_device(cls, index: int | str | torch.device) -> None: + torch_npu = _import_torch_npu() + torch_npu.npu.set_device(index) + + @classmethod + def synchronize(cls) -> None: + torch_npu = _import_torch_npu() + torch_npu.npu.synchronize() + + @classmethod + def empty_cache(cls) -> None: + torch_npu = _import_torch_npu() + torch_npu.npu.empty_cache() + + @classmethod + def pin_memory(cls, tensor: torch.Tensor) -> torch.Tensor: + _import_torch_npu() + try: + return tensor.pin_memory(device="npu") + except (RuntimeError, TypeError): + # Pageable CPU memory is slower but remains correct on runtimes that + # do not expose an NPU-specific pinned allocator. + return tensor + + @classmethod + def distributed_backend(cls) -> str: + return "hccl" + + @classmethod + def compile_backend(cls): + if not cls.supports("mindie_compile"): + raise RuntimeError( + "MindIE-SD compilation was requested, but MindieSDBackend is unavailable " + "in the installed MindIE-SD package." + ) + from mindiesd.compilation import CompilationConfig, MindieSDBackend + + CompilationConfig.fusion_patterns.enable_fast_gelu = False + return MindieSDBackend() + + @classmethod + def capabilities(cls) -> PlatformCapabilities: + return probe_ascend_capabilities() + + @classmethod + def supports(cls, capability: str) -> bool: + return probe_ascend_feature(capability) diff --git a/diffsynth_engine/platforms/base.py b/diffsynth_engine/platforms/base.py new file mode 100644 index 00000000..955c55d9 --- /dev/null +++ b/diffsynth_engine/platforms/base.py @@ -0,0 +1,68 @@ +from __future__ import annotations + +from abc import ABC +from dataclasses import dataclass +from typing import Any + +import torch + + +@dataclass(frozen=True) +class PlatformCapabilities: + device: bool = False + mindie: bool = False + mindie_attention: bool = False + mindie_compile: bool = False + + +class PlatformBackend(ABC): + name = "unknown" + device_type = "cpu" + + @classmethod + def is_available(cls) -> bool: + return True + + @classmethod + def normalize_device(cls, device: str | torch.device) -> torch.device: + return torch.device(device) + + @classmethod + def set_device(cls, index: int | str | torch.device) -> None: + return None + + @classmethod + def synchronize(cls) -> None: + return None + + @classmethod + def empty_cache(cls) -> None: + return None + + @classmethod + def pin_memory(cls, tensor: torch.Tensor) -> torch.Tensor: + return tensor.pin_memory() + + @classmethod + def distributed_backend(cls) -> str: + return "gloo" + + @classmethod + def compile_backend(cls) -> Any | None: + return None + + @classmethod + def compile_kwargs(cls) -> dict[str, Any]: + backend = cls.compile_backend() + return {} if backend is None else {"backend": backend} + + @classmethod + def supports(cls, capability: str) -> bool: + capabilities = cls.capabilities() + if not hasattr(capabilities, capability): + raise ValueError(f"Unknown platform capability: {capability}") + return bool(getattr(capabilities, capability)) + + @classmethod + def capabilities(cls) -> PlatformCapabilities: + return PlatformCapabilities(device=cls.is_available()) diff --git a/diffsynth_engine/utils/flag.py b/diffsynth_engine/utils/flag.py index 51aabc01..d0ee0e24 100644 --- a/diffsynth_engine/utils/flag.py +++ b/diffsynth_engine/utils/flag.py @@ -2,6 +2,7 @@ import torch from diffsynth_engine.utils import logging +from diffsynth_engine.utils.platform import is_mindie_sd_available logger = logging.get_logger(__name__) @@ -41,6 +42,10 @@ def check_module_available(module_path: str, module_name: str = None) -> bool: SPARGE_ATTN_AVAILABLE = check_module_available("spas_sage_attn", "Sparge attention") VIDEO_SPARSE_ATTN_AVAILABLE = check_module_available("vsa", "Video sparse attention") +# NPU +MINDIE_AVAILABLE = is_mindie_sd_available() + + NUNCHAKU_AVAILABLE = check_module_available("nunchaku", "Nunchaku") NUNCHAKU_IMPORT_ERROR = None if not NUNCHAKU_AVAILABLE: diff --git a/diffsynth_engine/utils/platform.py b/diffsynth_engine/utils/platform.py index 49a69680..060d7e7d 100644 --- a/diffsynth_engine/utils/platform.py +++ b/diffsynth_engine/utils/platform.py @@ -1,25 +1,42 @@ # cross-platform definitions and utilities -import torch import gc -import platform +import torch + +from diffsynth_engine.platforms import ( + AscendPlatform, + get_device_type as _get_device_type, + get_preferred_fp8_dtype, + pin_memory as _pin_memory, + resolve_platform, +) -# data type -# AMD only supports float8_e4m3fnuz -# https://onnx.ai/onnx/technical/float8.html -if torch.version.hip and "gfx94" in torch.cuda.get_device_properties(0).gcnArchName: - DTYPE_FP8 = torch.float8_e4m3fnuz -else: - DTYPE_FP8 = torch.float8_e4m3fn +DTYPE_FP8 = get_preferred_fp8_dtype("cuda") -def empty_cache(): +def empty_cache(device: str | torch.device | None = None): gc.collect() - if torch.cuda.is_available(): - torch.cuda.empty_cache() - if torch.mps.is_available(): - torch.mps.empty_cache() + resolve_platform(get_device_type(device)).empty_cache() + + +def pin_memory( + tensor: torch.Tensor, + device: str | torch.device | None = None, +): + return _pin_memory(tensor, device) + + +def is_npu_available() -> bool: + return AscendPlatform.supports("device") + + +def is_mindie_sd_available() -> bool: + return AscendPlatform.supports("mindie") + + +def get_device_type(device: str | torch.device | None = None) -> str: + return _get_device_type(device) -def pin_memory(tensor: torch.Tensor): - return tensor.pin_memory() if platform.system() == "Linux" else tensor +def get_torch_distributed_backend(device: str | torch.device) -> str: + return resolve_platform(device).distributed_backend() From ef3599800a8127aa4b73d84f9274b20c9414bc3b Mon Sep 17 00:00:00 2001 From: hammer Date: Fri, 31 Jul 2026 09:50:00 +0800 Subject: [PATCH 2/2] feat(npu): enable Ulysses SP with CFG parallel on Ascend Adapt ParallelWrapper for multi-card Ascend (HCCL/FSDP), route long-context attention through SeqAllToAll + MindIE, and align config.device at parallel load entry while workers bind rank-local devices. Co-authored-by: Cursor --- diffsynth_engine/models/basic/attention.py | 45 ++++++++- diffsynth_engine/utils/parallel.py | 110 ++++++++++++++++++--- 2 files changed, 139 insertions(+), 16 deletions(-) diff --git a/diffsynth_engine/models/basic/attention.py b/diffsynth_engine/models/basic/attention.py index d472b5fd..c49ae411 100644 --- a/diffsynth_engine/models/basic/attention.py +++ b/diffsynth_engine/models/basic/attention.py @@ -340,6 +340,44 @@ def forward( return self.to_out(out) +def _npu_ulysses_mindie_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + attn_mask: Optional[torch.Tensor] = None, + scale: Optional[float] = None, +): + """Ulysses SP on NPU: SeqAllToAll4D + MindIE local attn. ring_degree>1 not supported.""" + from yunchang.comm.all_to_all import SeqAllToAll4D + + from diffsynth_engine.utils.process_group import get_sp_ring_world_size, get_sp_ulysses_group + + if q.device.type != "npu": + raise RuntimeError("mindie long-context attention is only supported on NPU") + if not MINDIE_AVAILABLE: + raise RuntimeError( + "NPU Ulysses sequence parallel requires MindIE attention, but MindIE-SD is not available" + ) + if get_sp_ring_world_size() > 1: + raise RuntimeError( + "NPU long-context attention currently supports Ulysses only " + f"(sp_ring_degree must be 1, got {get_sp_ring_world_size()})" + ) + assert attn_mask is None, "long context attention does not support attention mask" + + # scatter heads (dim=2), gather sequence (dim=1) — same as video_sparse / v1 USP + scatter_idx, gather_idx = 2, 1 + group = get_sp_ulysses_group() + q = SeqAllToAll4D.apply(group, q, scatter_idx, gather_idx) + k = SeqAllToAll4D.apply(group, k, scatter_idx, gather_idx) + v = SeqAllToAll4D.apply(group, v, scatter_idx, gather_idx) + + # Must not call attention() here — it is patched to long_context_attention under SP. + out = mindie_attn(q, k, v, attn_mask=attn_mask, scale=scale) + out = SeqAllToAll4D.apply(group, out, gather_idx, scatter_idx) + return out + + def long_context_attention( q: torch.Tensor, k: torch.Tensor, @@ -372,11 +410,12 @@ def long_context_attention( "vsa", "mindie", ] - if attn_impl == "mindie": - raise RuntimeError("mindie long-context attention is not supported yet") assert attn_mask is None, "long context attention does not support attention mask" flash_attn3_compatible = q.shape[-1] <= FA3_MAX_HEADDIM if attn_impl is None or attn_impl == "auto": + # NPU has no FA/yunchang TORCH_EFFICIENT kernel; pick MindIE Ulysses when available. + if q.device.type == "npu" and MINDIE_AVAILABLE: + return _npu_ulysses_mindie_attention(q, k, v, attn_mask=attn_mask, scale=scale) if FLASH_ATTN_3_AVAILABLE: if flash_attn3_compatible: return LongContextAttention(attn_type=AttnType.FA3)(q, k, v, softmax_scale=scale) @@ -397,6 +436,8 @@ def long_context_attention( return LongContextAttention(attn_type=AttnType.FA)(q, k, v, softmax_scale=scale) raise ValueError("No available long context attention implementation") else: + if attn_impl == "mindie": + return _npu_ulysses_mindie_attention(q, k, v, attn_mask=attn_mask, scale=scale) if attn_impl == "fa3" or attn_impl == "fa3_fp8": if not flash_attn3_compatible: raise RuntimeError( diff --git a/diffsynth_engine/utils/parallel.py b/diffsynth_engine/utils/parallel.py index fc8e2c57..45a2582e 100644 --- a/diffsynth_engine/utils/parallel.py +++ b/diffsynth_engine/utils/parallel.py @@ -11,7 +11,6 @@ from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy from torch.distributed.device_mesh import DeviceMesh from torch.distributed.tensor.parallel.style import ParallelStyle -from torch.distributed.tensor.parallel._utils import _validate_tp_mesh_dim from contextlib import contextmanager from datetime import timedelta from functools import partial @@ -19,7 +18,13 @@ from queue import Empty import diffsynth_engine.models.basic.attention as attention_ops -from diffsynth_engine.utils.platform import empty_cache +from diffsynth_engine.utils.platform import ( + empty_cache, + get_device_type, + get_torch_distributed_backend, + is_npu_available, +) +from diffsynth_engine.platforms import resolve_platform from diffsynth_engine.utils import logging from diffsynth_engine.utils.process_group import ( PROCESS_GROUP, @@ -74,35 +79,44 @@ def make_parallel_groups(blocks: List[List[int]], degree: int): groups.extend(list(zip(*chunk))) return groups, chunks + # NOTE: every rank must call dist.new_group() for every subgroup, in the + # same order — even if it is not a member. Skipping non-member ranks + # breaks HCCL/NCCL communicators and surfaces later as all_to_all errors + # (common when use_cfg_parallel=True creates multiple CFG/SP subgroups). blocks = [list(range(world_size))] cfg_groups, cfg_blocks = make_parallel_groups(blocks, cfg_degree) for cfg_ranks in cfg_groups: + group = dist.new_group(cfg_ranks) if rank in cfg_ranks: - PROCESS_GROUP.CFG_GROUP = dist.new_group(cfg_ranks) + PROCESS_GROUP.CFG_GROUP = group PROCESS_GROUP.CFG_RANKS = cfg_ranks sp_groups, sp_blocks = make_parallel_groups(cfg_blocks, sp_degree) for sp_ranks in sp_groups: + group = dist.new_group(sp_ranks) if rank in sp_ranks: - PROCESS_GROUP.SP_GROUP = dist.new_group(sp_ranks) + PROCESS_GROUP.SP_GROUP = group PROCESS_GROUP.SP_RANKS = sp_ranks sp_ulysses_groups, sp_ulysses_blocks = make_parallel_groups(cfg_blocks, sp_ulysses_degree) for sp_ulysses_ranks in sp_ulysses_groups: + group = dist.new_group(sp_ulysses_ranks) if rank in sp_ulysses_ranks: - PROCESS_GROUP.SP_ULYSSUES_GROUP = dist.new_group(sp_ulysses_ranks) + PROCESS_GROUP.SP_ULYSSUES_GROUP = group PROCESS_GROUP.SP_ULYSSUES_RANKS = sp_ulysses_ranks sp_ring_groups, _ = make_parallel_groups(sp_ulysses_blocks, sp_ring_degree) for sp_ring_ranks in sp_ring_groups: + group = dist.new_group(sp_ring_ranks) if rank in sp_ring_ranks: - PROCESS_GROUP.SP_RING_GROUP = dist.new_group(sp_ring_ranks) + PROCESS_GROUP.SP_RING_GROUP = group PROCESS_GROUP.SP_RING_RANKS = sp_ring_ranks tp_groups, _ = make_parallel_groups(sp_blocks, tp_degree) for tp_ranks in tp_groups: + group = dist.new_group(tp_ranks) if rank in tp_ranks: - PROCESS_GROUP.TP_GROUP = dist.new_group(tp_ranks) + PROCESS_GROUP.TP_GROUP = group PROCESS_GROUP.TP_RANKS = tp_ranks set_seq_parallel_pg(sp_ulysses_degree, sp_ring_degree, rank, world_size) @@ -149,6 +163,16 @@ def parallelize_module( device_mesh: DeviceMesh, parallelize_plan: Optional[Union[ParallelStyle, Dict[str, ParallelStyle]]] = None, ): + # TP relies on a private torch API removed in PyTorch >= 2.9; adapt separately later. + try: + from torch.distributed.tensor.parallel._utils import _validate_tp_mesh_dim + except ModuleNotFoundError as e: + raise RuntimeError( + "Tensor parallel (tp_degree > 1) is not adapted for this PyTorch build yet " + "(missing torch.distributed.tensor.parallel._utils on torch>=2.9). " + "Use tp_degree=1 with Ulysses/FSDP for now." + ) from e + _validate_tp_mesh_dim(device_mesh) if parallelize_plan is None: return module @@ -169,6 +193,55 @@ def parallelize_module( PARALLEL_FWD_TIMEOUT_SEC = int(os.environ.get("PARALLEL_FWD_TIMEOUT_SEC", 3600)) +def _resolve_parallel_device_type(device_type: Optional[str]) -> str: + """Keep historical CUDA default; only auto-pick NPU when CUDA is absent.""" + if device_type is not None: + return device_type + if is_npu_available() and not torch.cuda.is_available(): + return "npu" + return "cuda" + + +def _ensure_parallel_device_ready(device_type: str) -> None: + """Availability gate: cuda/npu must both be selected and actually usable. + + After this returns, the worker may branch on ``device_type`` only. + """ + if device_type not in ("cuda", "npu"): + raise RuntimeError( + f"parallelism is only supported on CUDA or NPU devices, got {device_type!r}" + ) + if not resolve_platform(device_type).is_available(): + raise RuntimeError( + f"parallelism requested {device_type} but {device_type} is not available" + ) + + +def _align_config_device_type(config_device: str, target_type: str) -> str: + """Main-process entry: make config.device backend match ParallelWrapper workers. + + Rewrites only the historical CUDA placeholder when workers run on NPU. + Explicit mismatched devices raise instead of being silently rewritten. + """ + current_type = get_device_type(config_device) + if current_type == target_type: + return config_device + device_str = str(config_device) + if target_type == "npu" and (device_str == "cuda" or device_str.startswith("cuda:")): + return target_type + raise RuntimeError( + f"config.device={config_device!r} does not match parallel device_type={target_type!r}" + ) + + +def _bind_config_device_rank(kwargs: dict, device: torch.device) -> None: + """Worker-only: bind config.device to this rank's local device (e.g. npu:0).""" + config = kwargs.get("config") + if config is None or not hasattr(config, "device"): + return + config.device = str(device) + + def _worker_loop( rank: int, world_size: int, @@ -185,8 +258,7 @@ def _worker_loop( """ https://pytorch.org/docs/stable/multiprocessing.html#sharing-cuda-tensors """ - if device_type != "cuda" or not torch.cuda.is_available(): - raise RuntimeError("parallelism is only supported on CUDA devices") + _ensure_parallel_device_ready(device_type) try: os.environ["RANK"] = str(rank) @@ -194,11 +266,16 @@ def _worker_loop( os.environ["MASTER_ADDR"] = "localhost" os.environ["MASTER_PORT"] = str(master_port) device = torch.device(type=device_type, index=rank) - torch.cuda.set_device(rank) + # NPU-only distributed env; keep CUDA path free of LOCAL_RANK side effects. + if device_type == "npu": + os.environ["LOCAL_RANK"] = str(rank) + # Bind device before init_process_group (required by HCCL on Ascend). + resolve_platform(device_type).set_device(rank) + backend = get_torch_distributed_backend(device_type) timeout = timedelta(seconds=NCCL_TIMEOUT_SEC) dist.init_process_group( - backend="nccl", + backend=backend, init_method="env://", timeout=timeout, world_size=world_size, @@ -247,6 +324,7 @@ def wrap_for_parallel(module): empty_cache() elif name == "load_module": init_fn, kwargs = data[1:] + _bind_config_device_rank(kwargs, device) module = wrap_for_parallel(init_fn(**kwargs)) elif module is None: res = RuntimeError("module is not initialized") @@ -271,7 +349,7 @@ def wrap_for_parallel(module): queue_out.put(err) # any exception caught in the worker will be raised to the main process finally: del module, data, args, kwargs - torch.cuda.synchronize() + resolve_platform(device_type).synchronize() empty_cache() dist.destroy_process_group() @@ -285,10 +363,11 @@ def __init__( tp_degree: int, use_fsdp: bool = False, master_port: int = 29500, - device_type: str = "cuda", + device_type: Optional[str] = None, ): super().__init__() self._module_name = None + self.device_type = _resolve_parallel_device_type(device_type) self.world_size = cfg_degree * sp_ulysses_degree * sp_ring_degree * tp_degree spawn_ctx = mp.get_context("spawn") @@ -306,13 +385,16 @@ def __init__( tp_degree, use_fsdp, master_port, - device_type, + self.device_type, ), nprocs=self.world_size, join=False, ) def load_module(self, init_fn, **kwargs): + config = kwargs.get("config") + if config is not None and hasattr(config, "device"): + config.device = _align_config_device_type(config.device, self.device_type) data = ["load_module", init_fn, kwargs] for q in self.queue_in: q.put(data)