X-VLA-Libero 昇腾 NPU 适配

2toINF/X-VLA-Libero(X-VLA 0.9B Foundation Edition)在华为昇腾 Ascend 910 NPU 上的推理适配与 CPU 精度/性能对比。

模型简介

X-VLA 是一个 0.9B 参数的跨本体视觉-语言-动作(Vision-Language-Action)通用机器人策略模型,架构为:

• 输入:3 视角图像(224×224)+ 语言指令 + 本体感觉(20 维)+ domain_id
• 输出:30 步 × 20 维末端执行器动作序列(10 步流匹配去噪)
• X-VLA 三个权重(Pt / WidowX / Libero)同构,本仓库复用 Pt 的适配模式(见 weasonlee/X-VLA-Pt)

环境要求

部署步骤

# 1. 安装依赖(华为镜像源)
pip install -r requirements.txt -i https://repo.huaweicloud.com/repository/pypi/simple/

# 2. 下载权重到 models/(本仓库 .gitignore 已排除 models/)
HF_ENDPOINT=https://hf-mirror.com HF_HUB_DOWNLOAD_TIMEOUT=120 python -c \
  "from huggingface_hub import snapshot_download; snapshot_download('2toINF/X-VLA-Libero', local_dir='/data/X-VLA-Libero/models')"

# 3. source CANN 环境
source /usr/local/Ascend/ascend-toolkit/set_env.sh

# 4. NPU 推理
python inference.py --device npu

# 5. CPU 推理(对比基准)
python inference.py --device cpu

关键适配点

1. _supports_sdpa 属性安全回退(models/modeling_florence2.py):transformers ≥ 4.57 在 __init__ 早期(子模块尚未创建)即访问 _supports_sdpa / _supports_flash_attn_2,原实现直接读 self.language_model 抛 AttributeError。改为 getattr(self, "language_model", None) 判空回退。
2. 禁用 tie_weights 钩子(models/modeling_xvla.py):Florence2 为 encoder-only,删除 decoder/lm_head 后须禁用自动权重绑定,避免加载时访问已删除模块。
3. 重新 tying 文本嵌入(inference.py):XVLA 禁用了 tie_weights 以避免访问已删除的 decoder,而新版 transformers 的 assign 加载会解开 __init__ 中 shared → encoder.embed_tokens 的参数共享,导致 embed_tokens 随机初始化(checkpoint 中存为 model.shared.weight)。加载后手动 copy_ 恢复,保证权重完整且 CPU/NPU 结果可比。
4. SDPA 注意力:action head 使用 F.scaled_dot_product_attention,昇腾上由 torch_npu 原生支持,无需修改。
5. 确定性对比:初始噪声 x1 用 CPU 固定种子 torch.Generator 生成后再搬运到目标设备,保证两端输入逐位一致。

推理示例

import torch, torch_npu  # NPU 环境
from transformers import AutoModel, AutoProcessor

model = AutoModel.from_pretrained("models/", trust_remote_code=True).to("npu:0").eval()
# 注意: 需按上文适配点 3 重新 tying encoder.embed_tokens
processor = AutoProcessor.from_pretrained("models/", trust_remote_code=True)

action = model.generate_actions(
    input_ids=..., image_input=..., image_mask=...,   # processor(images, text) 输出
    domain_id=torch.zeros(1, dtype=torch.long, device="npu:0"),
    proprio=torch.zeros(1, 20, device="npu:0"),
    steps=10,
)  # -> [1, 30, 20] ee6d 动作序列

CPU vs NPU 精度对比

输入:固定种子合成观测(3 视角 256×256 图像 + 指令 "pick up the red block and place it into the bin" + 20 维本体感觉),初始噪声 seed=42,10 步流匹配去噪;计时为 3 次迭代均值(预热 1 次后)。实测数据见 results_*.json。

差异来源说明:NPU 与 CPU 的 fp32 算子实现(LayerNorm / Softmax / SDPA)归约顺序不同,误差在 1e-4 量级属正常范围,cos_sim ≈ 1.0 表明数值对齐良好。

产出文件

许可证

Apache 2.0(遵循上游 2toINF/X-VLA)。