diff --git a/.gitignore b/.gitignore index b8116bbd..61435d60 100755 --- a/.gitignore +++ b/.gitignore @@ -116,6 +116,7 @@ uv.lock /benchmarks/aiperf/ /models /data +model_zoo/ .vscode .idea *.pkl diff --git a/docs/en/attention.md b/docs/en/attention.md index 757e53ca..c4aec615 100644 --- a/docs/en/attention.md +++ b/docs/en/attention.md @@ -246,6 +246,10 @@ pipe_config.dit_config.attention_config = config **Note on Flash Attention 4**: Flash Attention 4 is optimized for **Hopper (SM90, H100)** and **Blackwell (SM100+, B100/B200)** GPUs. It provides significant performance improvements on these architectures. For older GPUs (Ampere, Ada Lovelace), use Flash Attention 2 or 3 instead. +**Note on AMD ROCm**: only `TORCH_SDPA` is natively available on ROCm hosts (`tf-kernel`, SageAttention, and +`flash_attn` are CUDA-only), so ROCm examples default to SDPA explicitly. The per-platform backend matrix lives in +[Hardware Platforms](platforms.md). + ### Sparse Attention Backends | Backend | Description | Requirements | diff --git a/docs/en/installation.md b/docs/en/installation.md index c9e53e9c..a13c2e10 100644 --- a/docs/en/installation.md +++ b/docs/en/installation.md @@ -11,6 +11,7 @@ distribution are installed separately. | Python | 3.10 through 3.13 | | PyTorch | 2.6 or newer | | CUDA toolkit | 12.8 or newer for the maintained CUDA development path | +| ROCm | 7.x with a PyTorch `+rocm` build for AMD GPUs; see the ROCm note under verification | | GPU | Depends on the selected model; check its Cookbook guide | An example may impose stricter versions or GPU architecture requirements. In particular, locally built `tf-kernel` @@ -55,6 +56,12 @@ telefuser --help Model execution expects `torch.cuda.is_available()` to print `True`. If it does not, verify the installed PyTorch build and visible NVIDIA driver before diagnosing TeleFuser. +On AMD ROCm hosts, install a PyTorch `+rocm` build instead of the CUDA toolkit path. A HIP build also prints `True` +for `torch.cuda.is_available()` (check `torch.version.hip` to distinguish it), and TeleFuser's platform layer detects +ROCm before CUDA. Examples ending in `_rocm.py` (for example `examples/wan_video/wan21_1_3b_text_to_video_rocm.py`) +are the validated entry points; see [Hardware Platforms](platforms.md) for per-platform capabilities and backend +availability. + ## Model Checkpoints TeleFuser does not bundle model weights. The [Supported Models](supported_models.md) page links to each Cookbook diff --git a/docs/en/ops.md b/docs/en/ops.md index 8cad047a..36fd1a86 100644 --- a/docs/en/ops.md +++ b/docs/en/ops.md @@ -62,6 +62,9 @@ TeleFuser follows a strict layered architecture for operations: - **Performance**: ops layer uses optimized Triton kernels in eager mode - **Separation of concerns**: kernel layer focuses on pure kernel implementation, ops layer handles dispatch logic +On AMD ROCm, dispatch reuses the `forward_cuda` Triton path when no `forward_rocm` kernel is defined. See +[Hardware Platforms](platforms.md) for the full per-platform dispatch behavior and `tf-kernel` gating. + ### torch.compile Strategy by Operator Type TeleFuser uses a **mixed strategy** for torch.compile compatibility, optimizing based on operator characteristics: diff --git a/docs/en/platforms.md b/docs/en/platforms.md new file mode 100644 index 00000000..50de62b7 --- /dev/null +++ b/docs/en/platforms.md @@ -0,0 +1,71 @@ +# Hardware Platforms + +TeleFuser abstracts execution hardware behind the platform layer in `telefuser/platforms/`. Each process resolves a +single `current_platform` object at import time, and the [Ops](ops.md) dispatch layer uses it to select operator +implementations. Pipelines and examples configure a `torch.device`-style device and do not branch on the vendor. + +## Platform selection + +`_resolve_current_platform()` in `telefuser/platforms/__init__.py` probes the environment in a fixed order and +instantiates the first match: + +1. **ROCm** — a PyTorch `+rocm` (HIP) build with at least one visible AMD GPU +2. **CUDA** — a CUDA build with at least one visible NVIDIA GPU +3. **NPU** — a `torch_npu` installation with a visible Ascend device +4. **CPU** — the fallback when no accelerator is detected + +HIP builds expose the `torch.cuda` API, so the ROCm platform reuses it (`device_type` stays `cuda`); check +`torch.version.hip` to distinguish a ROCm host from an NVIDIA one. A HIP or CUDA build without a visible GPU falls +back to CPU. `CUDA_VISIBLE_DEVICES` controls device visibility on both GPU platforms, and +`ASCEND_RT_VISIBLE_DEVICES` on NPU. + +## Platform matrix + +| Platform | `device_type` | Distributed backend | Operator dispatch | `tf-kernel` | torch.compile | +|----------|---------------|---------------------|-------------------|-------------|---------------| +| CUDA (NVIDIA) | `cuda` | NCCL | Optimized `forward_cuda` paths | Supported | Experimental ([details](torch_compile_compatibility.md)) | +| ROCm (AMD) | `cuda` | RCCL (`nccl`) | Triton `forward_cuda` paths and native fallbacks | Not available | Not validated | +| NPU (Ascend) | `npu` | HCCL | Native fallbacks | Not available | Not validated | +| CPU | `cpu` | Gloo | Native fallbacks | Not available | Native path | + +## Attention backends by platform + +Attention backend availability is resolved at import time (see [Attention](attention.md)). A backend whose +dependencies are missing falls back to `TORCH_SDPA` with a one-time warning. + +- **CUDA**: all dense backends — `TORCH_SDPA`, FlashAttention 2/3/4, SageAttention through `tf-kernel` or + `sageattention`, and cuDNN — subject to GPU architecture; sparse backends likewise. +- **ROCm**: `TORCH_SDPA` only. `flash_attn` has no official ROCm wheels for consumer RDNA GPUs, and `tf-kernel`, + SageAttention, and SpargeAttn are CUDA-only. The AOTriton-backed SDPA path is the fast attention kernel on + RDNA4. +- **NPU / CPU**: `TORCH_SDPA` through the native fallback paths; no vendor attention kernels are integrated. + +## CUDA + +CUDA is the primary validated path: Python 3.10–3.13, PyTorch 2.6 or newer, CUDA toolkit 12.8 or newer, with H100 +as the validated target for optimized kernels. Optional `tf-kernel` provides fused elementwise operations, +quantized GEMM, SageAttention, and block-sparse attention — see [tf-kernel](tf_kernel.md) for build and artifact +compatibility. Multi-GPU inference uses NCCL; see [Parallel Inference](parallel.md). + +## ROCm + +ROCm support targets AMD GPUs with ROCm 7.x and a PyTorch `+rocm` build; see [Installation](installation.md) for the +setup path. + +- Attention uses `TORCH_SDPA`; no `tf-kernel`, `flash_attn`, or `sageattention` installation is required. +- The ops layer selects `forward_rocm` where a kernel defines one and otherwise reuses the CUDA Triton path (Triton + supports ROCm), falling back to native PyTorch. +- `tf-kernel` imports are gated to `CudaPlatform`, so no CUDA-only extension is loaded on ROCm hosts. +- Multi-GPU inference uses RCCL, AMD's NCCL-compatible collectives library. PyTorch's ROCm build exposes it through + the `nccl` backend string, so the platform layer requires no special configuration. +- `torch.compile` is not validated on ROCm; ROCm examples run eager. +- Validated entry points are the `*_rocm.py` examples, for example + [Wan2.1 1.3B text-to-video](https://github.com/Tele-AI/TeleFuser/tree/main/examples/wan_video) on a Radeon RX 9070 + (gfx1201, ROCm 7.2). Multi-GPU branches reuse the `_h100.py` parallel configuration but are not yet validated. + +## NPU and CPU + +The NPU platform targets Huawei Ascend devices through `torch_npu` with the HCCL distributed backend. It is wired +into the platform and ops dispatch layers, but the maintained examples are validated on CUDA and, for select +examples, ROCm — validate on your target NPU before production use. The CPU platform is the fallback when no +accelerator is detected; it is intended for tests and for pipelines that explicitly request CPU execution. diff --git a/docs/zh/attention.md b/docs/zh/attention.md index 7722b790..18e53759 100644 --- a/docs/zh/attention.md +++ b/docs/zh/attention.md @@ -244,6 +244,9 @@ pipe_config.dit_config.attention_config = config **Flash Attention 4 说明**: Flash Attention 4 针对 **Hopper (SM90, H100)** 和 **Blackwell (SM100+, B100/B200)** GPU 架构进行了优化，在这些架构上提供显著的性能提升。对于旧版 GPU（Ampere、Ada Lovelace），请使用 Flash Attention 2 或 3。 +**AMD ROCm 说明**：ROCm 主机上仅 `TORCH_SDPA` 原生可用（`tf-kernel`、SageAttention 与 `flash_attn` 仅支持 +CUDA），因此 ROCm 示例显式默认 SDPA。各平台后端矩阵见[硬件平台](platforms.md)。 + ### 稀疏注意力后端 | 后端 | 描述 | 依赖 | diff --git a/docs/zh/installation.md b/docs/zh/installation.md index 4957096a..a4df9c7e 100644 --- a/docs/zh/installation.md +++ b/docs/zh/installation.md @@ -10,6 +10,7 @@ | Python | 3.10 至 3.13 | | PyTorch | 2.6 或更高版本 | | CUDA Toolkit | 当前 CUDA 开发路径要求 12.8 或更高版本 | +| ROCm | AMD GPU 使用 ROCm 7.x 与 PyTorch `+rocm` 构建，详见验证安装一节的说明 | | GPU | 取决于所选模型，以对应 Cookbook 为准 | 具体示例可能要求更严格的软件版本或 GPU 架构。特别是本地构建的 `tf-kernel` 产物与其记录的 PyTorch、 @@ -54,6 +55,12 @@ telefuser --help 模型执行要求 `torch.cuda.is_available()` 输出 `True`。否则应先检查 PyTorch CUDA 构建和 NVIDIA 驱动， 再排查 TeleFuser。 +AMD ROCm 主机应安装 PyTorch `+rocm` 构建，而非 CUDA Toolkit 路径。HIP 构建下 +`torch.cuda.is_available()` 同样输出 `True`（可通过 `torch.version.hip` 区分），TeleFuser 平台层会先 +检测 ROCm 再检测 CUDA。以 `_rocm.py` 结尾的示例（例如 +`examples/wan_video/wan21_1_3b_text_to_video_rocm.py`）是已验证的入口；各平台能力与后端可用性见 +[硬件平台](platforms.md)。 + ## 模型权重 TeleFuser 不随软件包分发模型权重。[支持的模型](supported_models.md)页面会链接到各模型的 Cookbook， diff --git a/docs/zh/ops.md b/docs/zh/ops.md index a2963711..1d144c0e 100644 --- a/docs/zh/ops.md +++ b/docs/zh/ops.md @@ -62,6 +62,9 @@ TeleFuser 遵循严格的分层架构： - **性能优化**：ops 层在 eager 模式下使用优化的 Triton 内核 - **关注点分离**：kernel 层专注纯内核实现，ops 层处理分发逻辑 +AMD ROCm 上，未定义 `forward_rocm` 内核时分发会复用 `forward_cuda` 的 Triton 路径。各平台的完整分发行为与 +`tf-kernel` 门控见[硬件平台](platforms.md)。 + ### 不同算子类型的 torch.compile 策略 TeleFuser 采用**混合策略**处理 torch.compile 兼容性，根据算子特性优化： diff --git a/docs/zh/platforms.md b/docs/zh/platforms.md new file mode 100644 index 00000000..28d27424 --- /dev/null +++ b/docs/zh/platforms.md @@ -0,0 +1,66 @@ +# 硬件平台 + +TeleFuser 通过 `telefuser/platforms/` 中的平台层屏蔽底层硬件差异。每个进程在导入时解析出唯一的 +`current_platform` 对象，[算子](ops.md)分发层根据它选择具体实现。Pipeline 与示例只需配置 +`torch.device` 风格的设备，无需针对厂商编写分支。 + +## 平台选择 + +`telefuser/platforms/__init__.py` 中的 `_resolve_current_platform()` 按固定顺序探测环境，并实例化第一个 +命中的平台： + +1. **ROCm** — PyTorch `+rocm`（HIP）构建且至少有一块可见的 AMD GPU +2. **CUDA** — CUDA 构建且至少有一块可见的 NVIDIA GPU +3. **NPU** — 安装了 `torch_npu` 且有可见的昇腾设备 +4. **CPU** — 未检测到加速器时的回退 + +HIP 构建同样暴露 `torch.cuda` API，因此 ROCm 平台复用它（`device_type` 仍为 `cuda`）；可通过 +`torch.version.hip` 区分 ROCm 主机与 NVIDIA 主机。HIP 或 CUDA 构建在没有可见 GPU 时同样回退到 CPU。 +两块 GPU 平台通过 `CUDA_VISIBLE_DEVICES` 控制设备可见性，NPU 使用 `ASCEND_RT_VISIBLE_DEVICES`。 + +## 平台矩阵 + +| 平台 | `device_type` | 分布式后端 | 算子分发 | `tf-kernel` | torch.compile | +|------|---------------|------------|----------|-------------|---------------| +| CUDA（NVIDIA） | `cuda` | NCCL | 优化的 `forward_cuda` 路径 | 支持 | 实验性（[详情](torch_compile_compatibility.md)） | +| ROCm（AMD） | `cuda` | RCCL（`nccl`） | Triton `forward_cuda` 路径与原生回退 | 不支持 | 未验证 | +| NPU（昇腾） | `npu` | HCCL | 原生回退 | 不支持 | 未验证 | +| CPU | `cpu` | Gloo | 原生回退 | 不支持 | 原生路径 | + +## 各平台的注意力后端 + +注意力后端可用性在导入时解析（见[注意力机制](attention.md)）。依赖缺失的后端会带一次性警告回退到 +`TORCH_SDPA`。 + +- **CUDA**：全部稠密后端 —— `TORCH_SDPA`、FlashAttention 2/3/4、经 `tf-kernel` 或 `sageattention` 提供的 + SageAttention，以及 cuDNN —— 取决于 GPU 架构；稀疏后端同理。 +- **ROCm**：仅 `TORCH_SDPA`。`flash_attn` 没有面向消费级 RDNA GPU 的官方 ROCm 轮子，`tf-kernel`、 + SageAttention 与 SpargeAttn 仅支持 CUDA。AOTriton 支持的 SDPA 路径是 RDNA4 上的快速注意力内核。 +- **NPU / CPU**：通过原生回退路径使用 `TORCH_SDPA`；未集成厂商注意力内核。 + +## CUDA + +CUDA 是主要的已验证路径：Python 3.10–3.13、PyTorch 2.6 及以上、CUDA Toolkit 12.8 及以上，优化内核以 +H100 为验证目标。可选的 `tf-kernel` 提供融合逐元素算子、量化 GEMM、SageAttention 与块稀疏注意力 —— +构建与制品兼容性见 [tf-kernel](tf_kernel.md)。多卡推理使用 NCCL，参见[并行推理](parallel.md)。 + +## ROCm + +ROCm 支持面向使用 ROCm 7.x 与 PyTorch `+rocm` 构建的 AMD GPU；安装路径见[安装指南](installation.md)。 + +- 注意力使用 `TORCH_SDPA`；无需安装 `tf-kernel`、`flash_attn` 或 `sageattention`。 +- 算子层在内核定义了 `forward_rocm` 时优先选择，否则复用 CUDA Triton 路径（Triton 支持 ROCm），并回退到 + PyTorch 原生实现。 +- `tf-kernel` 导入被限定在 `CudaPlatform`，ROCm 主机不会加载任何 CUDA-only 扩展。 +- 多卡推理使用 RCCL（AMD 与 NCCL 兼容的集合通信库）。PyTorch 的 ROCm 构建通过 `nccl` 后端字符串暴露它， + 平台层无需特殊配置。 +- `torch.compile` 在 ROCm 上未验证；ROCm 示例以 eager 模式运行。 +- 已验证入口为 `*_rocm.py` 示例，例如在 Radeon RX 9070（gfx1201，ROCm 7.2）上运行的 + [Wan2.1 1.3B 文生视频](https://github.com/Tele-AI/TeleFuser/tree/main/examples/wan_video)。多卡分支复用 + `_h100.py` 的并行配置，但尚未在 ROCm 上验证。 + +## NPU 与 CPU + +NPU 平台通过 `torch_npu` 与 HCCL 分布式后端支持华为昇腾设备。平台层与算子分发层均已接入 NPU，但现有 +示例在 CUDA 上验证、部分示例在 ROCm 上验证 —— 生产使用前请先在目标 NPU 上完成验证。CPU 平台是未检测到 +加速器时的回退，面向测试以及显式请求 CPU 执行的 Pipeline。 diff --git a/examples/wan_video/README.md b/examples/wan_video/README.md index f3a80f88..e46cb48e 100644 --- a/examples/wan_video/README.md +++ b/examples/wan_video/README.md @@ -32,6 +32,9 @@ Video generation using Wan2.1 and Wan2.2 models for Text-to-Video and Image-to-V - GPU: CUDA GPUs with enough memory for the selected 1.3B, 5B, or 14B checkpoint; H100 is the validated target for scripts ending in `_h100.py` +- GPU: AMD ROCm GPUs for scripts ending in `_rocm.py`; validated on a Radeon RX 9070 (ROCm 7.2, `torch` built with + `+rocm`). These examples use the PyTorch SDPA attention backend and need no tf-kernel, flash-attn, or SageAttention + installation - Software: the standard TeleFuser installation; optional attention, FP8, Ray, and RIFE paths require their respective dependencies - Input assets: a readable image for I2V/FL2V and optional LoRA, distillation, cache, or RIFE weights for those variants @@ -120,6 +123,44 @@ python examples/wan_video/wan21_1_3b_text_to_video_h100.py --resolution 480p --a **Features:** - Video Frame Interpolation (VFI) with RIFE model for 30fps output - CFG parallel when cfg_scale > 1 + +#### `wan21_1_3b_text_to_video_rocm.py` + +T2V on AMD ROCm GPUs. + +**Purpose:** Wan2.1 1.3B text-to-video for ROCm hosts, loading the official (non-Diffusers) checkpoint layout. + +**Usage:** +```bash +TELEAI_EXAMPLE_OUTPUT_DIR=work_dirs \ +python examples/wan_video/wan21_1_3b_text_to_video_rocm.py \ + --model_root "$TF_MODEL_ZOO_PATH/Wan2.1-T2V-1.3B" \ + --prompt "A sailboat crosses a calm lake at sunrise" +``` + +**Features:** +- PyTorch SDPA attention backend (natively available on ROCm; no flash-attn, SageAttention, or tf-kernel needed) +- Eager execution (`torch.compile` disabled by default; not validated on ROCm) +- 2-tile VAE decode geometry (`tile_size=(60, 62)`, `tile_stride=(30, 54)`): covers the 480p 16:9 latent with + ~1.4x redundant compute instead of the default 12-tile layout's ~2.9x, cutting VAE decode from ~61s to ~34s on a + Radeon RX 9070 at ~7.7GiB peak VRAM +- Text encoder CPU offloading with pageable (non-pinned) host copies: the ~10.6GB bf16 T5 encoder is only + moved to the GPU during text encoding, and page-locked copies are avoided because they exceed a 16GB + host RAM budget together with the DiT/VAE weights +- Validated single-GPU on Radeon RX 9070 (gfx1201) with ROCm 7.2 +- Multi-GPU branches follow the `_h100.py` parallel configuration and are not yet validated on ROCm +- VFI (RIFE) is disabled by default; enable it in `PPL_CONFIG` to add the interpolation model + +**ROCm performance notes (Radeon RX 9070, gfx1201, ROCm 7.2, 832x480, 81 frames):** + +- The DiT denoiser is at the operator-level hardware limit in eager mode: hipBLASLt serves the MLP GEMMs at + ~99 TFLOPS (RDNA4 bf16 peak) and the AOTriton-backed flash SDPA is the only fast attention kernel on this GPU, + so there is no faster ROCm operator to switch to. Measured alternatives are slower: `torch.compile` warm steps + (~14.6s vs ~14.0s), TunableOp autotuned GEMMs (~15.2s plus a ~400s autotune pass), and MIOpen's fused + attention has no gfx1201 kernel ("No available kernel" error) +- The first MIOpen conv run per shape pays one-time JIT compilation, cached cross-process under + `~/.cache/miopen`; subsequent runs (including in new processes) reuse it + #### `wan21_1_3b_text_to_video_hf.py` T2V with HuggingFace format loading. diff --git a/examples/wan_video/wan21_1_3b_text_to_video_rocm.py b/examples/wan_video/wan21_1_3b_text_to_video_rocm.py new file mode 100644 index 00000000..38fe5c96 --- /dev/null +++ b/examples/wan_video/wan21_1_3b_text_to_video_rocm.py @@ -0,0 +1,226 @@ +import os +import time + +import click +import torch + +from telefuser.core.config import AttentionConfig, AttnImplType, WeightOffloadType +from telefuser.core.module_manager import ModuleManager +from telefuser.pipelines.wan_video.wan21_video import ( + Wan21VideoPipeline, + Wan21VideoPipelineConfig, +) +from telefuser.utils.logging import logger +from telefuser.utils.utils import get_example_name +from telefuser.utils.video import get_target_video_size_from_ratio, save_video + +TF_MODEL_ZOO_PATH = os.environ.get("TF_MODEL_ZOO_PATH", "model_zoo") +PPL_CONFIG = dict( + name="wan21_1.3B_t2v_rocm", + model_root=TF_MODEL_ZOO_PATH + "/Wan2.1-T2V-1.3B", + negative_prompt="Camera shake, overly saturated colors, overexposed, static, blurry details, subtitles, style, artwork, painting, frame, still, overall grayish, worst quality, low quality, JPEG compression artifacts, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn face, deformed, disfigured, malformed limbs, fused fingers, static frames, cluttered background, three legs, crowded background, walking backwards", + num_inference_steps=40, + num_frames=81, + resolution="480p", + cfg_scale=5.0, + tiled=True, + target_fps=16, + # 2-tile VAE decode geometry: on a 60x108 (480p 16:9) latent this covers the + # full frame with only ~1.4x redundant compute instead of the ~2.9x of the + # default (30, 52)/(15, 26) 12-tile layout, cutting decode time by ~45% on a + # 16GB gfx1201 while staying within ~7.7GiB of VRAM. + vae_tile_size=(60, 62), + vae_tile_stride=(30, 54), + sample_solver="euler", + # ROCm has no flash-attn / SageAttention / tf-kernel support; SDPA is natively available. + attn_impl=AttnImplType.TORCH_SDPA, + model_type="Wan2.1-I2V-1.3B-720P", + sigma_shift=8.0, + enable_vfi=False, +) + + +def get_pipeline(parallelism=1, model_root=PPL_CONFIG["model_root"]): + """ + Args: + parallelism (int): Number of parallel GPUs for inference (REQUIRED) + model_root (str): Root directory of the model files (REQUIRED) + """ + # Load models from the official (non-Diffusers) checkpoint layout. + # low_cpu_mem_usage avoids duplicating the state dict in host RAM while loading. + module_manager = ModuleManager(device="cpu") + module_manager.load_models( + [f"{model_root}/Wan2.1_VAE.pth"], + torch_dtype=torch.bfloat16, + low_cpu_mem_usage=True, + ) + module_manager.load_models( + [[f"{model_root}/diffusion_pytorch_model.safetensors"]], + torch_dtype=torch.bfloat16, + low_cpu_mem_usage=True, + ) + module_manager.load_models( + [ + f"{model_root}/models_t5_umt5-xxl-enc-bf16.pth", + ], + torch_dtype=torch.bfloat16, + low_cpu_mem_usage=True, + ) + pipe = Wan21VideoPipeline(device="cuda", torch_dtype=torch.bfloat16) + pipe_config = Wan21VideoPipelineConfig() + pipe_config.dit_config.attention_config = AttentionConfig.dense_attention(PPL_CONFIG["attn_impl"]) + pipe_config.sample_solver = PPL_CONFIG["sample_solver"] + pipe_config.enable_clip_stage = False + pipe_config.enable_vfi = PPL_CONFIG["enable_vfi"] + pipe_config.enable_metrics = True + # The ~10.6GB bf16 text encoder stays in host RAM and moves to the GPU only + # during text encoding, keeping resident VRAM within a 16GB budget. + pipe_config.text_encoding_config.offload_config.offload_type = WeightOffloadType.MODEL_CPU_OFFLOAD + # Pageable (non-pinned) CPU copies: the default pinned offload would page-lock + # ~10.6GB of host RAM, which exceeds a 16GB WSL2 VM together with the DiT/VAE + # and triggers the kernel OOM killer. + pipe_config.text_encoding_config.offload_config.pin_cpu_memory = False + # torch.compile is not validated on ROCm; run eager by default. + pipe_config.dit_config.compile_config.enabled = False + if parallelism > 1: + # Configure parallel based on cfg_scale + # For cfg_scale > 1: cfg_degree=2, sp_ulysses_degree=parallelism//2 + # For cfg_scale == 1: cfg_degree=1, sp_ulysses_degree=parallelism + cfg_scale = PPL_CONFIG["cfg_scale"] + + if cfg_scale > 1: + pipe_config.dit_config.parallel_config.cfg_degree = 2 + pipe_config.dit_config.parallel_config.sp_ulysses_degree = parallelism // 2 + else: + pipe_config.dit_config.parallel_config.sp_ulysses_degree = parallelism + + pipe_config.dit_config.parallel_config.device_ids = list(range(parallelism)) + pipe_config.enable_denoising_parallel = True + pipe.init(module_manager, pipe_config) + return pipe + + +def run( + pipeline, + prompt, + negative_prompt="", + seed=42, + resolution=PPL_CONFIG["resolution"], + aspect_ratio="16:9", +): + """ + Convert text prompts to video sequences using the Wan2.1 1.3B model. + Args: + pipeline (Wan21VideoPipeline): Preloaded pipeline object + prompt (str): Positive guidance text prompt + negative_prompt (str, optional): Negative guidance prompt merged with the base negative prompt + seed (int, optional): Random seed. Default is 42 + resolution (str, optional): Resolution such as "480p", "720p" + aspect_ratio (str, optional): Aspect ratio such as "16:9" + + Returns: + List[PIL.Image]: Generated video sequence + """ + width, height = get_target_video_size_from_ratio( + aspect_ratio, + resolution=resolution, + height_division_factor=2, + width_division_factor=2, + ) + video = pipeline( + prompt=prompt, + negative_prompt=f"{negative_prompt} {PPL_CONFIG['negative_prompt']}", + num_inference_steps=PPL_CONFIG["num_inference_steps"], + num_frames=PPL_CONFIG["num_frames"], + cfg_scale=PPL_CONFIG["cfg_scale"], + seed=seed, + tiled=PPL_CONFIG["tiled"], + tile_size=PPL_CONFIG["vae_tile_size"], + tile_stride=PPL_CONFIG["vae_tile_stride"], + height=height, + width=width, + sigma_shift=PPL_CONFIG["sigma_shift"], + target_fps=PPL_CONFIG["target_fps"] if PPL_CONFIG["enable_vfi"] else None, + ) + return video + + +def run_with_file( + pipeline, + prompt, + negative_prompt, + seed, + resolution, + output_path, + aspect_ratio: str = "16:9", + **kwargs, +): + video = run( + pipeline, + prompt, + aspect_ratio=aspect_ratio, + negative_prompt=negative_prompt, + seed=seed, + resolution=resolution, + ) + logger.info(f"save target video to {output_path}") + save_video( + video, + output_path, + fps=PPL_CONFIG["target_fps"], + quality=6, + ) + + +@click.command() +@click.option("--gpu_num", default=1, help="Number of GPUs to use, default is 1") +@click.option( + "--prompt", + default="A stylish little girl gently caressing her dog while they relax in a sunny, beautiful backyard. Perfect for pet and family content, or videos aiming to showcase love, style, and the bond between kids and their pets.", + help="Positive guidance text prompt", +) +@click.option("--negative_prompt", default="", help="Negative guidance prompt") +@click.option("--seed", default=42, help="Random seed") +@click.option("--resolution", default=PPL_CONFIG["resolution"], help="Resolution (480p, 720p)") +@click.option("--aspect_ratio", default="16:9", help="Aspect ratio") +@click.option("--model_root", default=PPL_CONFIG["model_root"], help="Root directory of the model files") +def main( + gpu_num, + prompt, + negative_prompt, + seed, + resolution, + aspect_ratio, + model_root, +): + """Text to video conversion using Wan2.1 1.3B on AMD ROCm GPUs""" + pipe = get_pipeline(gpu_num, model_root) + + # Run inference + start = time.time() + video = run( + pipe, + prompt, + negative_prompt, + seed, + resolution, + aspect_ratio, + ) + elapsed_time = time.time() - start + print(pipe.get_prometheus_metrics()) + + print(f"Video generation time: {elapsed_time:.2f} seconds") + + # Save results + output_dir = os.getenv("TELEAI_EXAMPLE_OUTPUT_DIR", "./") + filename = get_example_name(__file__).replace(".py", f"_{gpu_num}gpu.mp4") + output_path = os.path.join(output_dir, filename) + + save_video(video, output_path, fps=PPL_CONFIG["target_fps"], quality=6) + print(f"Video saved to: {output_path}") + + del pipe + + +if __name__ == "__main__": + main() diff --git a/mkdocs.yml b/mkdocs.yml index 66a6772d..9f26946f 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -141,6 +141,7 @@ plugins: CUDA IPC Ulysses: CUDA IPC Ulysses Parallel Inference: 并行推理 Communication Architecture: 通信架构 + Hardware Platforms: 硬件平台 Developer Guide: 开发者指南 Configuration: 配置 Tools: 工具 @@ -198,6 +199,7 @@ nav: - Model Loading: model_loading.md - Parallel Inference: parallel.md - Communication Architecture: communication.md + - Hardware Platforms: platforms.md - Attention: attention.md - Quantization: quantization.md - Feature Cache: feature_cache.md diff --git a/telefuser/platforms/__init__.py b/telefuser/platforms/__init__.py index a23fdbe2..7d517faa 100644 --- a/telefuser/platforms/__init__.py +++ b/telefuser/platforms/__init__.py @@ -28,7 +28,10 @@ def _is_cuda_available() -> bool: def _is_rocm_available() -> bool: """Check if ROCm is available.""" - return hasattr(torch.version, "hip") and torch.version.hip is not None + if not (hasattr(torch.version, "hip") and torch.version.hip is not None): + return False + # A HIP build without a visible GPU falls back to CPU like the CUDA check below. + return torch.cuda.is_available() def _is_npu_available() -> bool: diff --git a/telefuser/platforms/cuda.py b/telefuser/platforms/cuda.py index 3dd890b4..055a62c6 100644 --- a/telefuser/platforms/cuda.py +++ b/telefuser/platforms/cuda.py @@ -76,6 +76,18 @@ def get_device_properties(device: int | str | torch.device | None = None) -> Any def set_device(device: int | str | torch.device) -> None: return torch.cuda.set_device(device) + @staticmethod + def device_count() -> int: + return torch.cuda.device_count() + + @staticmethod + def is_accelerator_available() -> bool: + return torch.cuda.is_available() and torch.cuda.device_count() > 0 + + @staticmethod + def current_device() -> int: + return torch.cuda.current_device() + @staticmethod def get_device_capability(device: int | str | torch.device | None = None) -> tuple[int, int]: return torch.cuda.get_device_capability(device) diff --git a/telefuser/platforms/rocm.py b/telefuser/platforms/rocm.py index 1a2c60fb..fc38c716 100644 --- a/telefuser/platforms/rocm.py +++ b/telefuser/platforms/rocm.py @@ -62,6 +62,18 @@ def get_device_properties(device: int | str | torch.device | None = None) -> Any def set_device(device: int | str | torch.device) -> None: return torch.cuda.set_device(device) + @staticmethod + def device_count() -> int: + return torch.cuda.device_count() + + @staticmethod + def is_accelerator_available() -> bool: + return torch.cuda.is_available() and torch.cuda.device_count() > 0 + + @staticmethod + def current_device() -> int: + return torch.cuda.current_device() + @staticmethod def get_device_capability(device: int | str | torch.device | None = None) -> tuple[int, int]: return torch.cuda.get_device_capability(device) diff --git a/tests/unit/pipelines/wan_video/test_rocm_example.py b/tests/unit/pipelines/wan_video/test_rocm_example.py new file mode 100644 index 00000000..cbb62cd0 --- /dev/null +++ b/tests/unit/pipelines/wan_video/test_rocm_example.py @@ -0,0 +1,87 @@ +from __future__ import annotations + +import importlib.util +from pathlib import Path +from unittest.mock import MagicMock, patch + +from telefuser.core.config import AttnImplType, WeightOffloadType + + +def _load_module(path: Path): + spec = importlib.util.spec_from_file_location(path.stem, path) + assert spec is not None + assert spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _example_path() -> Path: + return Path(__file__).resolve().parents[4] / "examples/wan_video/wan21_1_3b_text_to_video_rocm.py" + + +def test_rocm_example_defaults_to_sdpa_and_eager() -> None: + module = _load_module(_example_path()) + + assert module.PPL_CONFIG["attn_impl"] == AttnImplType.TORCH_SDPA + assert module.PPL_CONFIG["enable_vfi"] is False + + +def test_rocm_example_uses_overlapping_two_tile_vae_geometry() -> None: + module = _load_module(_example_path()) + + tile_size = module.PPL_CONFIG["vae_tile_size"] + tile_stride = module.PPL_CONFIG["vae_tile_stride"] + + assert tile_size == (60, 62) + assert tile_stride == (30, 54) + # Zero overlap (tile_size == tile_stride) breaks the tiled-decode border mask. + assert all(s > d for s, d in zip(tile_size, tile_stride)) + + +def test_rocm_example_get_pipeline_loads_official_layout() -> None: + module = _load_module(_example_path()) + pipeline = MagicMock() + + with ( + patch.object(module, "ModuleManager") as manager_cls, + patch.object(module, "Wan21VideoPipeline", return_value=pipeline), + ): + result = module.get_pipeline(1, "/models/Wan2.1-T2V-1.3B") + + assert result is pipeline + manager = manager_cls.return_value + loaded = [call.args[0] for call in manager.load_models.call_args_list] + assert loaded == [ + ["/models/Wan2.1-T2V-1.3B/Wan2.1_VAE.pth"], + [["/models/Wan2.1-T2V-1.3B/diffusion_pytorch_model.safetensors"]], + ["/models/Wan2.1-T2V-1.3B/models_t5_umt5-xxl-enc-bf16.pth"], + ] + for call in manager.load_models.call_args_list: + assert call.kwargs["low_cpu_mem_usage"] is True + pipeline.init.assert_called_once() + + pipe_config = pipeline.init.call_args.args[1] + assert pipe_config.dit_config.attention_config.attn_impl == AttnImplType.TORCH_SDPA + assert pipe_config.dit_config.compile_config.enabled is False + assert pipe_config.enable_clip_stage is False + assert pipe_config.enable_vfi is False + assert pipe_config.text_encoding_config.offload_config.offload_type == WeightOffloadType.MODEL_CPU_OFFLOAD + assert pipe_config.text_encoding_config.offload_config.pin_cpu_memory is False + + +def test_rocm_example_multi_gpu_enables_denoising_parallel() -> None: + module = _load_module(_example_path()) + pipeline = MagicMock() + + with ( + patch.object(module, "ModuleManager"), + patch.object(module, "Wan21VideoPipeline", return_value=pipeline), + ): + module.get_pipeline(2, "/models/Wan2.1-T2V-1.3B") + + pipe_config = pipeline.init.call_args.args[1] + assert pipe_config.enable_denoising_parallel is True + assert pipe_config.dit_config.parallel_config.device_ids == [0, 1] + assert pipe_config.dit_config.parallel_config.cfg_degree == 2 + assert pipe_config.dit_config.parallel_config.sp_ulysses_degree == 1 diff --git a/tests/unit/platforms/__init__.py b/tests/unit/platforms/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit/platforms/test_rocm_platform.py b/tests/unit/platforms/test_rocm_platform.py new file mode 100644 index 00000000..3fa852d3 --- /dev/null +++ b/tests/unit/platforms/test_rocm_platform.py @@ -0,0 +1,75 @@ +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from telefuser.platforms import ( + CpuPlatform, + CudaPlatform, + RocmPlatform, + _resolve_current_platform, +) +from telefuser.platforms.rocm import RocmPlatform as RocmPlatformClass + + +def test_resolve_prefers_rocm_over_cuda() -> None: + with ( + patch.object(torch.version, "hip", "test-hip", create=True), + patch("telefuser.platforms.torch.cuda.is_available", return_value=True), + ): + platform = _resolve_current_platform() + assert isinstance(platform, RocmPlatformClass) + + +def test_resolve_rocm_without_visible_gpu_falls_back_to_cpu() -> None: + with ( + patch.object(torch.version, "hip", "test-hip", create=True), + patch("telefuser.platforms.torch.cuda.is_available", return_value=False), + ): + platform = _resolve_current_platform() + assert isinstance(platform, CpuPlatform) + + +def test_resolve_cuda_when_no_hip() -> None: + with ( + patch.object(torch.version, "hip", None, create=True), + patch("telefuser.platforms.torch.cuda.is_available", return_value=True), + ): + platform = _resolve_current_platform() + assert isinstance(platform, CudaPlatform) + + +@pytest.mark.parametrize("platform_cls", [CudaPlatform, RocmPlatform]) +def test_accelerator_methods_delegate_to_torch_cuda(platform_cls: type) -> None: + with ( + patch("telefuser.platforms.cuda.torch.cuda.device_count", return_value=3), + patch("telefuser.platforms.cuda.torch.cuda.current_device", return_value=1), + ): + assert platform_cls.device_count() == 3 + assert platform_cls.current_device() == 1 + + +@pytest.mark.parametrize("platform_cls", [CudaPlatform, RocmPlatform]) +def test_is_accelerator_available_requires_visible_device(platform_cls: type) -> None: + for available, count, expected in ((True, 2, True), (True, 0, False), (False, 0, False)): + with ( + patch("telefuser.platforms.cuda.torch.cuda.is_available", return_value=available), + patch("telefuser.platforms.cuda.torch.cuda.device_count", return_value=count), + ): + assert platform_cls.is_accelerator_available() is expected + + +def test_rocm_platform_implements_full_interface() -> None: + assert {"device_count", "is_accelerator_available", "current_device"} <= set(RocmPlatformClass.__dict__) + + +def test_is_accelerator_available_matches_mocked_torch_cuda() -> None: + torch_cuda = MagicMock() + torch_cuda.is_available.return_value = True + torch_cuda.device_count.return_value = 1 + with patch("telefuser.platforms.rocm.torch.cuda", torch_cuda): + assert RocmPlatform.is_accelerator_available() is True + assert RocmPlatform.device_count() == 1 + assert RocmPlatform.current_device() is torch_cuda.current_device()