ROCm平台下Radeon 7800XT运行YuE音乐生成模型的实践指南

背景介绍

YuE是一款基于深度学习技术的音乐生成模型，能够根据文本提示生成音乐片段。该模型依赖于PyTorch框架和Flash Attention优化技术。在AMD Radeon 7800XT显卡上运行YuE时，由于ROCm平台对RDNA架构的支持特性，需要进行特定的环境配置才能正常工作。

环境准备

Python环境配置

推荐使用Python 3.11版本，因为Python 3.12在某些情况下可能存在兼容性问题。可以通过以下命令安装：

sudo add-apt-repository ppa:deadsnakes/ppa
sudo apt update
sudo apt install python3.11 python3.11-dev python3.11-venv

创建并激活虚拟环境：

python3.11 -m venv yue
source yue/bin/activate

ROCm和PyTorch安装

安装适配ROCm 6.3.4的PyTorch版本：

pip3 install torch==2.4.0 torchaudio torchvision==0.19.0 pytorch_triton -f https://repo.radeon.com/rocm/manylinux/rocm-rel-6.3.4

验证安装是否成功：

import torch
print(torch.__version__)
print(torch.cuda.is_available())
print(torch.cuda.get_device_name(0))

Flash Attention编译与安装

由于Radeon 7800XT使用RDNA3架构(gfx1100)，需要启用Triton后端支持：

git clone https://github.com/Dao-AILab/flash-attention.git
cd flash-attention
FLASH_ATTENTION_TRITON_AMD_ENABLE="TRUE" python setup.py install

关键环境变量说明：

• FLASH_ATTENTION_TRITON_AMD_ENABLE="TRUE"：强制使用Triton后端而非CUDA后端
• HSA_OVERRIDE_GFX_VERSION=11.0.0：让系统将7800XT识别为7900XT(gfx1100)

YuE模型部署

依赖安装

sudo apt install git-lfs
git clone https://github.com/multimodal-art-projection/YuE.git
cd YuE
git lfs install
git lfs pull
pip install -r requirements.txt

模型下载

YuE采用两阶段模型架构：

1. 第一阶段模型(m-a-p/YuE-s1-7B-anneal-en-cot)
2. 第二阶段模型(m-a-p/YuE-s2-1B-general)

这些模型会自动从HuggingFace下载，确保网络连接正常。

运行优化技巧

提示词优化

原始示例中的提示词较长，可以简化为一个主歌和一个副歌部分，显著减少推理时间：

[Verse]
风格：电子流行
节奏：中速
情绪：欢快

[Chorus]
和声：丰富层次
乐器：合成器主导

代码修改建议

YuE的infer.py脚本需要两处修改以适应短提示：

1. 确保最小输出时长：

output_duration = max(output_duration, 6)  # 添加在L355附近

2. 修改时长检查逻辑：

if output_duration*50 < prompt.shape[-1]:  # 替换L382的!=检查

性能调优

在Radeon 7800XT上运行时，可能会遇到内存不足警告：

[drm:amdgpu_cs_ioctl [amdgpu]] *ERROR* Not enough memory for command submission!

建议解决方案：

1. 减少stage2_batch_size参数值（默认为4，可尝试2或1）
2. 缩短max_new_tokens（默认为3000，可适当减少）
3. 确保系统交换空间充足

典型运行时间

在优化后的配置下：

• 10秒音乐片段：约15分钟
• 完整长度生成：可能需要数小时

总结

通过合理的环境配置和参数调整，可以在Radeon 7800XT上成功运行YuE音乐生成模型。关键点在于正确编译Flash Attention的ROCm版本，适当调整模型参数以避免内存问题，以及优化输入提示词结构。这些实践不仅适用于YuE模型，也可为其他基于Transformer的大模型在ROCm平台上的部署提供参考。

对于更复杂的音乐生成需求，建议分阶段生成后再进行后期处理，而非一次性生成过长片段，这样可以提高成功率和效率。随着ROCm对RDNA架构支持的不断完善，未来在消费级AMD显卡上运行此类模型的体验将会进一步改善。