# torch.compile 在 AMD K100 (gfx928) 上编译失败问题报告 ## 问题概述 在 AMD K100 GPU 上使用 `torch.compile` 加速扩散模型时，Triton kernel 编译失败，导致无法获得性能加速。 --- ## 环境信息 ### 硬件配置 | 项目 | 规格 | |------|------| | GPU 型号 | AMD K100_AI | | GPU 架构 | gfx928 | | 显存 | 64GB HBM2e | | GPU 数量 | 8 卡 | ### 软件环境 | 项目 | 版本 | 来源 | |------|------|------| | 操作系统 | Ubuntu 22.04 LTS | | | DTK | 25.04.1 (基于 ROCm 6.1) | 源舟 | | PyTorch | 2.5.1+das.opt1.dtk25041 | 源舟 | | Triton | 3.0.0+das.opt1.dtk25041 | 源舟 | | Python | 3.10.12 | | | diffusers | 0.32.0.dev0 | | ### 安装来源 ```bash # Triton 下载地址 https://download.sourcefind.cn:65024/file/4/triton/DAS1.6/triton-3.0.0+das.opt1.dtk25041-cp310-cp310-manylinux_2_28_x86_64.whl # PyTorch 开发者社区下载安装于环境中 /dev/shm/ai-models/venv/ ``` --- ## 问题描述 ### 错误信息 使用 `torch.compile` 编译 Z-Image-Turbo 模型时出现以下错误： ``` backend='inductor' raised: CompilationError: at 1:0: def triton_(in_ptr0, in_ptr1, out_ptr0, ynumel, xnumel, YBLOCK : tl.constexpr, XBLOCK : tl.constexpr): ^ IndexError('list index out of range') Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information ``` ### 复现步骤 1. 加载环境 ```bash source /root/dtk-25.04.1/env.sh source /dev/shm/ai-models/venv/bin/activate ``` 2. 运行测试代码 ```python import torch from diffusers import ZImagePipeline MODEL_PATH = "/dev/shm/ai-models/z-image/models/Tongyi-MAI/Z-Image-Turbo" pipe = ZImagePipeline.from_pretrained( MODEL_PATH, torch_dtype=torch.bfloat16, low_cpu_mem_usage=False, use_safetensors=True, ) pipe.to("cuda") # 编译 transformer pipe.transformer = torch.compile( pipe.transformer, mode="max-autotune", fullgraph=True ) # 触发编译 result = pipe( prompt="a cat", width=1024, height=1024, num_inference_steps=8, ) ``` 3. 观察错误 - 第一次生成时触发编译 - 编译过程中抛出 `IndexError` --- ## 诊断测试 ### 测试 1: Triton 基础功能 ✅ 通过 ```python import torch import triton print(f'PyTorch: {torch.__version__}') print(f'Triton: {triton.__version__}') print(f'Triton Backend: {triton.runtime.driver.active}') # 输出: HIPDriver @torch.compile def test_fn(x, y): return x + y * 2 x = torch.randn(10, 10, device='cuda') y = torch.randn(10, 10, device='cuda') result = test_fn(x, y) print(f'简单编译测试：通过') # 编译耗时 ~0.9 秒 ``` ### 测试 2: 复杂模型编译 ❌ 失败 ```python pipe.transformer = torch.compile(pipe.transformer) # CompilationError: IndexError('list index out of range') ``` ### 测试 3: 可用后端检查 ```python print(f'Available backends: {torch._dynamo.list_backends()}') # 输出: ['cudagraphs', 'inductor', 'onnxrt', 'openxla', 'tvm'] ``` ### 测试 4: 性能基线 | 配置 | 生成时间 (8 步) | 状态 | |------|----------------|------| | Eager 模式 (默认) | ~20 秒 | ✅ 正常 | | torch.compile | ❌ 编译失败 | | | suppress_errors=True | ~20 秒 (回退 eager) | ✅ 正常 | --- ## 已尝试的解决方案 ### 方案 1: 修改编译模式 ```python # 尝试不同模式 pipe.transformer = torch.compile(pipe.transformer, mode="max-autotune") # ❌ pipe.transformer = torch.compile(pipe.transformer, mode="reduce-overhead") # ❌ pipe.transformer = torch.compile(pipe.transformer, dynamic=False, fullgraph=False) # ❌ ``` ### 方案 2: 环境变量 ```bash export HSA_OVERRIDE_GFX_VERSION=9.4.2 export PYTORCH_ROCM_ARCH=gfx928 export TORCHINDUCTOR_COMPILE_THREADS=1 ``` **结果**: 无明显改善 ### 方案 3: 禁用定制 kernel ```python torch._inductor.config.triton.disable_custom_kernels = True ``` **结果**: 仍然失败 ### 方案 4: 回退机制 ```python torch._dynamo.config.suppress_errors = True ``` **结果**: 可以运行，但完全回退到 eager 模式，无性能提升 ### 方案 5: diffusers 源码修改 ✅ 修改 `/dev/shm/ai-models/venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_z_image.py:277` **修改前**: ```python @staticmethod def precompute_freqs_cis(dim, end, theta=256.0): with torch.device("cpu"): # ← 问题：Dynamo 无法追踪 ... ``` **修改后**: ```python @staticmethod def precompute_freqs_cis(dim, end, theta=256.0): # 移除 with torch.device("cpu") 上下文 freqs_cis = [] for i, (d, e) in enumerate(zip(dim, end)): freqs = 1.0 / (theta ** (torch.arange(0, d, 2, dtype=torch.float64, device="cpu") / d)) ... ``` **结果**: 解决了 `torch.device` 上下文问题，但无法解决 Triton kernel 编译错误 --- ## 额外观察 ### 警告信息 ``` UserWarning: Torch was not compiled with memory efficient attention. (Triggered internally at /home/pytorch/aten/src/ATen/native/transformers/hip/sdpa_attention.py:663) ``` 这表明 PyTorch DAS 版本的 SDPA (Scaled Dot Product Attention) 没有启用高效注意力机制，这也是性能瓶颈之一。 ### 性能对比 | 平台 | 生成时间 (1024x1024, 8 步) | |------|--------------------------| | AMD K100 64GB (当前) | 20 秒 | | NVIDIA RTX 4090 24GB | 5-8 秒 (官方参考) | | 预期加速后 | 8-12 秒 | --- ## 核心问题 1. **gfx928 架构支持不完整** - Triton kernel 编译时 `IndexError: list index out of range` - 可能是硬件参数查找表缺失 2. **inductor 后端问题** - 简单 kernel 可以编译 - 复杂模型 (如扩散 transformer) 失败 3. **SDPA 优化缺失** - PyTorch 没有编译 memory efficient attention - 注意力计算效率低 --- ## 请求支持 ### 问题 1: gfx928 Triton 支持 - gfx928 架构的 Triton kernel 是否已完整支持？ - `IndexError` 是否已知问题？ - 是否需要额外的配置文件？ ### 问题 2: inductor 配置 - 是否有推荐的 `torch.compile` 配置用于 gfx928？ - 是否需要设置特定的环境变量？ ### 问题 3: SDPA 优化 - PyTorch DAS 版本是否启用了 SDPA？ - 如何启用 memory efficient attention？ ### 问题 4: 替代方案 - 是否有其他加速途径 (ONNX Runtime, OpenXLA 等)？ - 是否有针对 AMD GPU 的优化建议？