Το PyTorch προσθέτει εγγενή υποστήριξη AsyncTP για GPU ROCm

Τι έγινε

The PyTorch development team merged pull request #177961, which introduces a native AsyncTP implementation for the fused_all_gather_matmul operation on ROCm GPUs (gfx942 / MI300X and gfx950 / MI355X). The change replaces the previous CUDA‑only async path with a ROCm‑specific kernel built from the Composable Kernel (CK) library. The kernel overlaps the all‑gather of activation shards with the subsequent matrix multiplication, using a persistent GEMM that waits on per‑chunk signals. The implementation includes a hand‑written HIP file, new build flags, and runtime checks that enable the path only on supported architectures. Tests added to the PyTorch CI suite verify correctness on multiple GPU counts and ROCm versions, confirming that the native path runs without deadlock and matches reference outputs within tight error bounds.

The PR adds a HIP source file (AsyncMM.hip) that implements c10d::cuda::detail::async_input_mm for ROCm, mirroring the existing CUDA implementation that uses a CUTLASS kernel. The new kernel is built from the Composable Kernel (CK) tile PersistentAsyncInputScheduler, supporting bf16 inputs with fp32 accumulation and a block tile of 256×256×32. The implementation includes runtime checks for supported GPU architectures (gfx942 and gfx950) and enforces tensor size limits (INT32_MAX elements) because CK indexes with 32‑bit integers. It also adds a selector that requires the activation shard’s row dimension to be a multiple of 256, matching the kernel’s tile granularity. Extensive CI tests were added: native operation tests on MI300X and MI355X across 2, 4, and 8 GPU ranks, graph‑capture replay tests, and fallback correctness checks against a float64 reference. The tests confirm that the native path runs without deadlock, respects hardware queue limits, and produces numerically accurate results. Build integration ensures the HIP file replaces the CUDA version when PYTORCH_ROCM_ARCH contains the supported architectures and the CK headers are present. If the conditions are not met, the operation falls back to the existing CUDA‑only path, raising an error similar to the CUDA behavior on unsupported platforms.

Στοιχεία πηγής: github.com ↗

Γιατί έχει σημασία

AsyncTP is a key performance technique for large‑scale model training, allowing communication and computation to proceed concurrently. By providing a native ROCm implementation, PyTorch now lets developers on AMD hardware achieve similar overlap benefits that were previously limited to NVIDIA GPUs. Early benchmarks in the PR show 4‑33 % speed‑ups on multi‑GPU runs, and the kernel’s design avoids deadlocks that could arise with CUDA‑style ordering on AMD GPUs. This expands PyTorch’s high‑performance training capabilities to a broader hardware ecosystem, potentially lowering cost and diversifying compute options for researchers and enterprises that rely on AMD GPUs. However, the feature is currently gated behind the TORCH_SYMM_MEM_ENABLE_NATIVE_ASYNC_TP flag and is only available in the development (trunk) branch; its inclusion in an official PyTorch release, supported Python versions, and documentation timeline remain unknown.

AsyncTP reduces the time spent waiting for all‑gather communication by overlapping it with computation, a critical optimization for training very large models that are split across many GPUs. By bringing this capability to AMD GPUs, PyTorch removes a performance gap that could have forced users to stick with NVIDIA hardware for certain workloads. The performance gains reported (up to 33 % faster GEMM execution on 8‑GPU runs) suggest that training pipelines on MI300X and MI355X can achieve higher throughput or lower cost per training step, especially in high‑rank configurations where communication overhead dominates. The implementation also addresses a known deadlock issue on AMD hardware caused by CUDA‑style stream ordering, demonstrating that the ROCm runtime can handle more complex scheduling patterns when kernels are designed with ROCm’s execution model in mind. Because the feature is gated behind an environment flag and resides in the trunk branch, its immediate impact is limited to developers who can build PyTorch from source. The timeline for inclusion in a stable release, as well as any potential changes to the API or default behavior, remain uncertain.

Τι να παρακολουθήσετε στη συνέχεια

Future PyTorch releases that promote the native AsyncTP path to default, broader hardware support beyond gfx942 and gfx950, and performance benchmarks from the community comparing AMD and NVIDIA training pipelines. Watch for any updates to the TORCH_SYMM_MEM_ENABLE_NATIVE_ASYNC_TP flag handling, documentation changes, and downstream libraries (e.g., DeepSpeed, FairScale) that may adopt the new kernel. Also monitor AMD’s ROCm roadmap for any driver or library updates that could affect the stability or performance of the CK‑based GEMM.

Announcements from the PyTorch release team about the promotion of the native AsyncTP path to default in a forthcoming stable version. Community benchmarks that compare end‑to‑end training times on AMD versus NVIDIA GPUs using the new kernel. Updates to AMD’s ROCm driver or CK library that could further improve kernel efficiency or broaden architecture support. Adoption of the AsyncTP kernel by higher‑level distributed training frameworks (e.g., DeepSpeed, FairScale) that may expose the performance benefit to a wider user base.

Σχετικοί οδηγοί και κουίζ