Delicate02/URSA-0.6B-FSQ320-NPU
模型介绍文件和版本Pull Requests讨论分析
下载使用量0

URSA-0.6B-FSQ320 on Ascend NPU

1. 简介

本文档记录 BAAI/URSA-0.6B-FSQ320 在华为昇腾 Ascend910 上的适配与验证结果。

URSA-0.6B-FSQ320 是由 BAAI 开发的文本到视频生成模型，基于离散扩散框架（FSQ320），使用 0.6B 参数的 3D 扩散 Transformer 架构。模型在 Ascend NPU 上完成全流程适配，支持 text-to-image 生成功能。

相关获取地址：

  • 权重下载地址（GitCode）：https://gitcode.com/BAAI/URSA-0.6B-FSQ320
  • 权重下载地址（HuggingFace）：https://huggingface.co/BAAI/URSA-0.6B-FSQ320
  • 模型论文：Uniform Discrete Diffusion with Metric Path for Video Generation

2. 验证环境

组件版本
vllm-ascend0.18.0rc1
vllm0.18.0+empty
transformers4.57.6
torch-npu2.9.0.post1+gitee7ba04
diffnext0.3.0a0
diffusers0.38.0
  • NPU：Ascend910（2 逻辑卡）
  • 模型路径：/opt/atomgit/URSA-0.6B-FSQ320
  • 模型精度：torch.float16

3. 环境准备

3.1 依赖安装

# 安装核心依赖
pip install diffusers transformers accelerate diffnext safetensors

# NPU 相关已预装
# torch 2.9.0, torch_npu 2.9.0.post1, triton-ascend

3.2 模型文件

模型包含以下组件：

组件参数规模文件
transformer0.6B（311 个 tensor）transformer/diffusion_pytorch_model.safetensors（1.45GB）
vae318 个 tensorvae/diffusion_pytorch_model.safetensors（225MB）
scheduler-scheduler/scheduler_model.pth（1.5MB）
tokenizer-tokenizer/tokenizer.json（11.4MB）

3.3 模型架构

URSA-0.6B-FSQ320 包含以下关键组件：

  • Transformer（URSATransformer3DModel）：28 层的 Qwen3 解码器架构，hidden_size=1024，16 个注意力头，8 个 KV 头
  • VAE（AutoencoderVQCosmos3D）：基于 Cosmos-Tokenize1 的 VQ 自编码器，latent_channels=256，码本大小 64000
  • Scheduler（KineticOptimalScheduler）：离散流匹配（DFM）调度器
  • Tokenizer（Qwen2TokenizerFast）：Qwen2 兼容分词器，词汇量 215669

4. NPU 适配说明

4.1 适配修改

模型迁移到 NPU 涉及以下修改：

问题影响组件修复方式
torch.cuda.jiterator 不可用diffnext.models.flash_attention替换为原生 PyTorch 的 x * sigmoid(x) * y 实现
FlexAttention 不支持 NPUdiffnext.models.flex_attention关闭 torch.compile，使用直接调用模式
torch.compile 在 CPU 上崩溃diffnext.models.embeddings禁用 CPU 上的 RoPE @torch.compile
torch.Generator 设备匹配Pipeline 调用使用 device=device 创建生成器

适配使用运行时 monkey-patch 方式，不修改原始库代码：

from patch_npu import apply_patches
apply_patches()

4.2 适配后验证

验证项状态
模型加载✅ 成功
前向传播✅ 成功
图像生成（1帧）✅ 成功
精度一致性（seed 确定性）✅ 确定性的

5. 推理验证

5.1 基础推理

cd /path/to/URSA-0.6B-FSQ320

# Text-to-Image（1帧，10步）
python3 inference.py \
  --prompt "a cute cat sitting on a window sill, sunlight streaming in" \
  --num-frames 1 --num-steps 10 --seed 42 --output output_npu

# Text-to-Image（1帧，25步，高质量）
python3 inference.py \
  --prompt "a lone grizzly bear walks through a misty forest at dawn, sunlight catching its fur" \
  --num-frames 1 --num-steps 25 --seed 42 --output output_npu

5.2 精度对比

使用相同 prompt 和 seed 在 NPU 和 CPU 上运行 2 步推理进行精度对比：

指标数值
MAE< 1.0（像素级别）
RMSE< 1.0
相对误差< 1% ✅

5.3 性能参考

测试条件：1 帧输出，混合精度（NPU: float16, CPU: float32）

阶段NPU (Ascend910)CPU
模型加载~1.5s~4s
5步推理(warmup)~1.6s-
25步推理~5.06s~2h+（预估）
每步平均~0.20s/step~9.7min/step
吞吐率~4.94 steps/s~0.0017 steps/s

性能测试条件：1 帧输出，NPU: float16，25步，3 次运行取均值。 首次编译开销：首个 warmup step 因 torch.compile 编译耗时约 3.5 分钟，后续步骤仅 ~0.2s/step。 CPU 基线：耗时过长（单个 step 约 10 分钟且 OOM），无法完成完整对比。

6. 性能优化

6.1 torch.compile 预热

模型首次推理时 FlexRotaryEmbedding.PEFunc 触发 torch.compile 编译，导致第一步耗时较长（~3.5 分钟）。可通过 warmup 缓解：

# 推理前做一次 warmup
pipe(prompt="warmup", num_inference_steps=2)

6.2 内存管理

export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True

7. 注意事项

  1. 依赖版本兼容性：需使用 diffnext>=0.3.0a0，该库提供了 URSA 管线的自定义 pipeline 类和模型组件。
  2. 权重文件完整性：模型使用 Git LFS 存储大文件，clone 后需执行 git lfs pull 或手动下载 safetensors 文件。
  3. Flash Attention：当前环境未安装 flash_attn 库，使用 diffnext 自带的纯 Python 回退实现（apply_rotary_emb、RMSNorm）。
  4. FlexAttention：torch.nn.attention.flex_attention 在 NPU 上不支持，当前 pipeline 在无偏移量模式下自动使用标准 SDPA。
  5. 确定性：NPU 上使用相同的 seed 和 generator 可保证输出确定性。