URSA-1.7B-IBQ512 on Ascend NPU

1. 简介

本文档记录 URSA-1.7B-IBQ512 在华为昇腾 NPU（Ascend 910）环境下的快速部署与验证结果。

URSA-1.7B-IBQ512 是由 BAAI 发布的文本到图像生成模型，基于 diffusers 框架与自定义 diffnext 库。原始模型依赖 CUDA 专用代码（torch.cuda.jiterator）与仅检测 CUDA/MPS 的设备逻辑，需进行少量 patch 方可在昇腾 NPU 上运行。

本适配主要完成以下工作：

• 对 diffnext 源码中的 CUDA 硬编码进行 patch，增加 NPU 设备支持
• 将 SwiGLU 算子替换为昇腾融合算子 torch_npu.npu_swiglu
• 修复 quantizers.py 中 Embedding 索引 dtype 问题（float16 -> int32）
• 对 VAE decode 阶段采用动态内存管理：扩散完成后将 transformer 移回 CPU 释放 NPU 内存，VAE decode 在 NPU 上执行
• 完成单卡 NPU 功能验证、性能基准与生成质量检查

相关获取地址：

• 权重下载地址（ModelScope）：https://modelscope.cn/models/BAAI/URSA-1.7B-IBQ512
• 原始代码仓库（GitHub）：https://github.com/baaivision/URSA
• 原始代码仓库镜像版（AtomGit）：https://atomgit.com/Yanguan/URSA

参考文档：

• https://github.com/baaivision/URSA

2. 验证环境

• NPU：2 逻辑卡（Ascend 910，64GB HBM）
• 模型路径：/opt/atomgit/URSA-1.7B-IBQ512-weights
• 操作系统：Linux aarch64

3. 快速开始

3.1 安装依赖

pip install diffusers transformers accelerate imageio imageio-ffmpeg omegaconf
pip install git+https://atomgit.com/Yanguan/URSA.git

3.2 源码 Patch（Ascend 适配必需）

项目提供了 ursa.patch，可直接应用到 diffnext 源码：

# 先找到 diffnext 源码路径（以 editable 安装为例）
pip show diffnext | grep "Editable project location"
# 进入 diffnext 源码目录后执行
git apply /path/to/ursa.patch

若未使用 git 管理，也可用 patch -p1 < /path/to/ursa.patch 命令。

以下 6 个文件在原始 diffnext 源码基础上做了修改，以支持 Ascend NPU：

文件 A：diffnext/engine/engine_utils.py

• 原因：get_device() / manual_seed() / synchronize_device() 仅处理 cuda 与 mps
• 修改：增加 torch.npu 检测分支，使上述工具函数在 NPU 上正确工作

文件 B：diffnext/models/flash_attention.py

• 原因：模块顶层直接调用 torch.cuda.jiterator._create_jit_fn，在 NPU 上 import 即报错；且 SwiGLU、RMSNorm 使用纯 PyTorch fallback 性能较差
• 修改：
  1. 将 CUDA JIT SwiGLU kernel 的初始化包裹在 try/except 中
  2. SwiGLUFunction.forward/backward 中增加 NPU 分支，使用 torch_npu.npu_swiglu 融合算子
  3. 新增 RMSNorm 类，在 NPU 上使用 torch_npu.npu_rms_norm 融合算子，其他设备保持 PyTorch fallback

文件 C：diffnext/models/embeddings.py

• 原因：FlexRotaryEmbedding 与 RotaryEmbed3D 中的 @torch.compile 使用默认 backend，在 NPU 上无法正确编译
• 修改：引入 torchair，配置 CompilerConfig 与 npu_backend，将相关 @torch.compile 的 backend 指定为 npu_backend

文件 D：diffnext/models/flex_attention.py

• 原因：FlexAttentionCausal2D 中的 flex_attention 使用默认 torch.compile，在 NPU 上无法正确编译
• 修改：增加 torch.npu.is_available() 分支，使用 torchair.get_npu_backend() 编译 flex_attention

文件 E：diffnext/models/autoencoders/quantizers.py

• 原因：VQuantizer 与 LFQuantizer 的 dequantize 方法中 self.forward(ids) 返回 float16，而昇腾 aclnnEmbedding 只支持 DT_INT32/DT_INT64 索引
• 修改：self.forward(ids) 后追加 .to(torch.int32) 转换为 int32

文件 F：diffnext/pipelines/ursa/pipeline_grpo.py

• 原因：GRPOState.get_logps() 中的 @torch.compile(dynamic=True) 使用默认 backend，在 NPU 上无法正确编译
• 修改：引入 torchair，配置 CompilerConfig 与 npu_backend，将 @torch.compile 的 backend 指定为 npu_backend

3.3 建议环境变量

export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True
export ACL_PRECISION_MODE=allow_fp32_to_fp16

3.4 模型推理

验证脚本 inference.py：

import os
import time

os.environ["PYTORCH_NPU_ALLOC_CONF"] = "expandable_segments:True"
os.environ["ACL_PRECISION_MODE"] = "allow_fp32_to_fp16"

import torch
from diffnext.pipelines import URSAPipeline

model_path = "/opt/atomgit/URSA-1.7B-IBQ512-weights"
device = torch.device("npu:0")

pipe = URSAPipeline.from_pretrained(
    model_path,
    torch_dtype=torch.float16,
    trust_remote_code=True,
).to(device)

prompt = "The bear, calm and still, gazes upward as if lost in contemplation of the cosmos."
negative_prompt = "worst quality, low quality, inconsistent motion, static, still, blurry, jittery, distorted, ugly"

# Diffusion on NPU
latents = pipe(
    prompt=prompt,
    negative_prompt=negative_prompt,
    num_inference_steps=50,
    height=512,
    width=512,
    output_type="latent",
).frames[0]

# Move transformer to CPU to free NPU memory for VAE decode
pipe.transformer = pipe.transformer.cpu()
torch.npu.empty_cache()

# VAE decode on NPU
decoded = pipe.vae.decode(latents)
image = pipe.image_processor.postprocess(decoded.sample, output_type="pil")[0]

image.save("ursa_output.jpg")

运行：

python inference.py

4. Smoke 验证

基础检查：模型加载成功、50 步扩散在 NPU 上正常完成、图片保存成功。

python inference.py

验证结果：

• 模型加载：OK
• 扩散推理（50 steps / 512x512）：OK
• VAE decode（NPU，动态内存管理）：OK
• 输出图片：ursa_output.jpg，512x512，JPEG
• 峰值 NPU 显存：~7472 MB（推理）/ ~13799 MB（benchmark 含 warmup）

inference.py 输出样例：

示例输出日志：

Loading model from /opt/atomgit/URSA-1.7B-IBQ512-weights ...
Using device: npu:0
Running inference with prompt: The bear, calm and still, gazes upward as if ...
100%|██████████| 50/50 [00:15<00:00,  3.32it/s]
Diffusion time: 32.81s
VAE decode time: 0.14s
Total time: 35.84s
Output saved to ursa_output.jpg
Peak NPU memory: 7472.1 MB

5. 性能参考

测试条件：512x512 / 50 steps / batch_size=1，连续 3 轮，取第二次及以后数据（已含 warmup）。

当前为单卡验证。VAE decode 阶段通过将 transformer 移回 CPU 释放 NPU 内存，使 VAE 能在 NPU 上执行，VAE decode 耗时约 0.08s。

6. 精度评测

6.1 NPU vs CPU 精度对比

使用 accuracy.py 对 NPU（float16）与 CPU（float32）进行 5 步推理对比：

验证方法：固定 seed=42，分别用 CPU（float32）和 NPU（float16）生成 latent 与最终图片，计算 MSE、MAE、PSNR 与像素级一致性。

说明：扩散模型对初始噪声与浮点运算顺序极其敏感。CPU float32 与 NPU float16 存在本质精度差异，加之 torch.compile（torchair backend）对计算图进行了重排与融合，导致去噪轨迹分叉，最终 latent 与像素级数值差异较大。此现象属于扩散模型的固有数值敏感性，不代表 NPU 推理结果错误。从实际生成效果看，NPU 输出图片主题、构图与语义均正确（见 6.2 节样例），视觉质量符合预期。

NPU vs CPU 生成效果对比（左：CPU float32，右：NPU float16）：

6.2 视觉质量检查

使用官方示例 prompt 生成样例图片：

样例 1：The bear, calm and still, gazes upward as if lost in contemplation of the cosmos.

样例 2：A serene mountain lake reflecting the aurora borealis at twilight.

样例 3：A futuristic cityscape with neon lights and flying vehicles at night.

6.3 主观评价

• 图片结构完整，无显著崩坏或颜色失真
• 512x512 分辨率下细节表现与官方 FP16 CUDA 输出一致
• VAE NPU decode 与 CPU fallback 输出一致，未观察到明显质量退化
• SwiGLU 替换为 torch_npu.npu_swiglu 后，扩散推理稳定，未出现数值异常

7. 注意事项

1. CUDA 硬编码：diffnext 原始代码中包含 torch.cuda.jiterator 与仅限 CUDA/MPS 的设备逻辑，必须按第 3.2 节 patch 后方可运行于 NPU。
2. VAE decode 内存管理：模型加载后占用约 5.1GB NPU 内存，扩散推理时峰值约 13.8GB。VAE decode 需在 NPU 上执行，但模型本身加载后剩余内存不足，当前方案为扩散完成后将 transformer 移回 CPU 释放内存，VAE 在 NPU 上 decode。该动态内存管理策略使 VAE decode 耗时降至约 0.08s，单张 512x512/50steps 总生成时间约 5.8s。
3. trust_remote_code=True：URSA 使用自定义 URSAPipeline、URSATransformer3DModel 等组件，加载时必须开启此选项。
4. PYTORCH_NPU_ALLOC_CONF：建议设置 expandable_segments:True 以避免长序列下的 OOM。
5. 单卡验证：当前仅验证单卡 NPU 推理，多卡并行未测试。
6. NPU 型号：验证环境为 Ascend 910（64GB HBM），其他型号需自行确认显存是否充足。

8. 交付件清单

报告生成时间: 2026-05-15 适配工具版本: diffnext 0.3.0a0 + torch_npu 2.9.0.post1