Skip to content

[FEATURE] Add AMD ROCm platform support with validated Wan2.1 example - #47

Merged
lzx1413 merged 1 commit into
mainfrom
rocm
Sep 10, 2026
Merged

[FEATURE] Add AMD ROCm platform support with validated Wan2.1 example#47
lzx1413 merged 1 commit into
mainfrom
rocm

Conversation

@lzx1413

@lzx1413 lzx1413 commented Sep 10, 2026

Copy link
Copy Markdown
Collaborator

Description

Adds first-class AMD ROCm support to TeleFuser: platform-layer detection fixes and interface completion, a validated
Wan2.1 1.3B text-to-video example for ROCm hosts (wan21_1_3b_text_to_video_rocm.py), unit tests, and a new
hardware-platforms documentation page (en/zh) that consolidates per-platform capabilities previously scattered across
installation/ops/attention docs.

Motivation

  • A HIP (+rocm) PyTorch build without a visible GPU was misdetected as a ROCm platform; a HIP build also sets
    torch.cuda.is_available() to True, so the CUDA path could never be distinguished correctly.
  • BasePlatform declares device_count / is_accelerator_available / current_device and existing callers
    (worker/ray_worker.py, service/core/pipeline_service.py) already use them, but CudaPlatform and
    RocmPlatform raised NotImplementedError — any CUDA/ROCm process reaching those call sites crashed.
  • ROCm knowledge (backend availability, dispatch behavior, install path) was duplicated across three doc pages; the
    new example needed a single authoritative reference. RCCL is now correctly documented as the collectives library
    behind the nccl backend string on ROCm.

Type of Change

  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • Performance improvement
  • Code refactoring
  • Documentation update
  • Other (please describe):

Changes Made

  • telefuser/platforms/__init__.py: _is_rocm_available() now additionally requires torch.cuda.is_available(),
    so a HIP build without a visible GPU falls back to CPU symmetrically with the CUDA check.
  • telefuser/platforms/cuda.py, telefuser/platforms/rocm.py: implement device_count(),
    is_accelerator_available(), and current_device() per the existing BasePlatform contract (mirrors the
    NPUPlatform implementations; no behavior change on healthy CUDA hosts).
  • examples/wan_video/wan21_1_3b_text_to_video_rocm.py: Wan2.1 1.3B T2V on ROCm using the official checkpoint
    layout — TORCH_SDPA attention (no flash-attn / SageAttention / tf-kernel), eager execution (compile not
    validated on ROCm), T5 encoder CPU offload with non-pinned host copies for a 16GB host budget, 2-tile VAE decode
    geometry, CLIP stage disabled for T2V. Multi-GPU branches reuse the _h100.py parallel configuration (not yet
    validated on ROCm).
  • Tests: tests/unit/platforms/test_rocm_platform.py (platform resolution order, visible-GPU fallback, delegation,
    interface completeness) and tests/unit/pipelines/wan_video/test_rocm_example.py (example defaults, VAE tile
    geometry, official-layout loading, config assertions).
  • Docs: new docs/en/platforms.md + docs/zh/platforms.md (platform selection order, capability matrix, per-platform
    attention backends, RCCL note); trimmed cross-references in installation/ops/attention docs (en/zh); registered
    under Runtime & Optimization in mkdocs.yml with zh nav translation; documented the example in
    examples/wan_video/README.md with measured performance notes; ignore local model_zoo/ checkpoints.

Testing

  • Unit tests pass (pytest tests/) — focused selection: 13 new tests pass (full suite not run in this
    environment; GPU-marked tests skip via existing markers on CPU-only hosts)
  • Manual testing performed — end-to-end T2V validated on Radeon RX 9070 (gfx1201, ROCm 7.2, 832x480x81 frames);
    measured numbers recorded in examples/wan_video/README.md
  • Benchmarks added/updated (if applicable) — performance observations are documented in the example README; no
    benchmark harness changes

Test commands:

pytest tests/unit/platforms/ tests/unit/pipelines/wan_video/test_rocm_example.py
pre-commit run --files <changed files>   # ruff + ruff-format pass
python scripts/docs/prepare_cookbook.py build   # docs build + link/anchor validation

Checklist

  • Code follows the project's coding standards (ruff)
  • Pre-commit hooks pass (pre-commit run --all-files — verified on all files changed by this PR)
  • All tests pass (pytest tests/ — new and affected unit tests; full suite requires GPU hardware)
  • New tests added for new functionality
  • Documentation updated (README, docstrings)
  • Commit messages are clear and descriptive
  • PR title follows the convention: [TYPE] Brief description

Related Issues

Fixes #

Additional Notes

  • The ROCm example intentionally sets enable_clip_stage = False (T2V loads no CLIP weights; the _h100.py
    baseline leaves it True) — flagged here as a deliberate deviation from the baseline example.
  • No new framework interfaces, environment variables, loaders, registries, CLI options, or service schemas were
    introduced; the new platform methods complete an existing BasePlatform contract with existing callers.
  • Preview caveat for reviewers: prepare_cookbook.py serve serves the .build/docs copy; re-run
    prepare after editing docs/ sources.

GPU Architecture Support

No kernels added or modified (tf-kernel untouched; ROCm uses Triton forward_cuda paths and PyTorch-native
fallbacks via the existing ops dispatch).

  • SM80 (Ampere, Ada Lovelace)
  • SM90 (Hopper H100)
  • SM100+ (Blackwell)
  • Other: AMD RDNA4 (gfx1201, Radeon RX 9070) validated for the new example; other ROCm targets untested

Performance Impact

Measured on Radeon RX 9070 (gfx1201, ROCm 7.2, 832x480, 81 frames), recorded in the example README:

  • VAE tiled-decode re-geometry (2 tiles instead of the default 12) cuts decode from ~61s to ~34s at ~7.7GiB peak
    VRAM
  • Text encoder CPU offload with pageable (non-pinned) copies keeps resident VRAM within a 16GB host/GPU budget
  • DiT denoiser runs at the operator-level eager limit (~99 TFLOPS hipBLASLt on RDNA4 bf16); torch.compile,
    TunableOp, and MIOpen fused attention measured slower or unsupported on this GPU
  • No CUDA-path behavior change: platform methods only replace NotImplementedError stubs with correct
    implementations

- Guard ROCm platform detection with torch.cuda.is_available() so a HIP
  build without a visible GPU falls back to CPU like the CUDA check
- Implement device_count/is_accelerator_available/current_device on
  CudaPlatform and RocmPlatform, completing the BasePlatform contract
  already consumed by ray_worker and pipeline_service
- Add wan21_1_3b_text_to_video_rocm.py: TORCH_SDPA attention, eager
  execution, T5 CPU offload with non-pinned memory, 2-tile VAE decode
  geometry, and CLIP stage disabled for T2V; validated single-GPU on a
  Radeon RX 9070 (gfx1201, ROCm 7.2)
- Add a hardware platforms doc page (en/zh) covering platform selection
  order, capability matrix, and per-platform attention backends with
  RCCL noted as the collectives library behind the nccl backend on
  ROCm; cross-link from installation/ops/attention docs and register
  the page in the mkdocs nav
- Document the ROCm example in examples/wan_video/README.md and ignore
  the local model_zoo/ checkpoint directory

Verification: 13 new unit tests (platform dispatch and example config)
pass; pre-commit passes on all touched files; docs site build passes
with rendered page, asset, and anchor validation; preview server
rendered the new en/zh pages with correct titles.
@lzx1413
lzx1413 merged commit 5e3e609 into main Sep 10, 2026
7 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant