SmolVLA LIBERO — QNN HTP W8A16 (INT8 权重量化)

在 Qualcomm HTP (Hexagon 张量处理器) 上运行的 SmolVLA 三模型流水线 (视觉编码器 + LLM 主干 + 动作头),采用 INT8 权重量化 (W8A16)。 在 Snapdragon X Elite 上实现 约 4.8 秒端到端延迟 (实际耗时),成功完成 LIBERO 任务 3,有效控制频率为 0.84 Hz (n_action_steps=4) 或 1.70 Hz (n_action_steps=8)。

硬件 / 软件要求

注意: Python 必须是 x86-64 版本 (而非 ARM64),因为 LIBERO / robosuite 需要 x86 兼容性。

安装

1. Python 包

pip install \"numpy<2.0\" torch torchvision transformers safetensors onnx imageio pillow

2. LIBERO 基准测试

git clone https://github.com/Lifelong-Robot-Learning/LIBERO.git
cd LIBERO && pip install -e .

3. QNN SDK

下载并安装 QAIRT 2.43.0.260128。设置环境变量:

export QNN_SDK=C:/Users/<user>/qualcomm/qairt/2.43.0.260128

4. 克隆本仓库及 SmolVLA 权重

git clone https://huggingface.co/xpuenabler/smolvla-libero-QNN-HTP-W8A16
cd smolvla-libero-QNN-HTP-W8A16

# 下载 SmolVLA 权重 (lerobot/smolvla_base)
python -c \"
from huggingface_hub import snapshot_download
snapshot_download('lerobot/smolvla_base', local_dir='smolvla_weights')
\"

模型转换流水线

步骤 0: 导出 ONNX (如需要)

如果您没有 ONNX 模型,请从 LeRobot 导出:

python scripts/export_onnx.py --weights smolvla_weights/
# 输出: onnx_models/vision_encoder.onnx
#        onnx_models/llm_backbone.onnx
#        onnx_models/action_head.onnx

步骤 1: ONNX 预处理 — 消除布尔值

QNN 转换器无法处理 Python bool 链。应用以下修复:

# LLM 主干: 折叠布尔常量 + 消除布尔运算
python scripts/fold_and_eliminate_bool.py onnx_models/llm_backbone.onnx onnx_models/llm_backbone_v7.onnx

# 动作头: 消除布尔值 + 修复 Expand 算子
python scripts/fold_and_eliminate_bool.py onnx_models/action_head.onnx onnx_models/action_head_v5.onnx
python scripts/fix_ah_expand_bool.py onnx_models/action_head_v5.onnx onnx_models/action_head_qnn_v6.onnx

# LLM 主干 v7: 移除 embedding Gather (改为在 CPU 端运行)
python scripts/fix_llm_embed_on_cpu.py onnx_models/llm_backbone_v7.onnx
# 输出: onnx_models/llm_backbone_v7.onnx (原地更新)

步骤 2a: 构建 FP32 QNN 模型 (HTP 后端)

python scripts/build_all_models.py --precision fp32
# 输出: qnn_models/windows-aarch64-fp32/lib*.dll

步骤 2b: 构建 INT8W QNN 模型 (W8A16)

python scripts/build_int8w_models.py
# 输出: qnn_models/windows-aarch64-int8w/lib*.dll

需要在仓库根目录放置 qnn_convert_patched.py (禁用 matmul_to_fc 优化):

cp $QNN_SDK/lib/python/qnn_convert.py qnn_convert_patched.py
# 应用补丁: 禁用 matmul_to_fc (见 scripts/build_all_models.py 头部)

步骤 3: 创建展开的动作头 (内置 10 步去噪)

python scripts/create_ah_unrolled.py
# 输入:  onnx_models/action_head_qnn_v6.onnx
# 输出: onnx_models/action_head_unrolled10.onnx (393 MB)
#         qnn_context_cache/libaction_head_unrolled_htp.dll (387 MB)
# 构建时间: 约 24 分钟 (图排序阶段)

步骤 4: 生成上下文二进制文件 (预编译 HTP 图)

# FP32 模型
python scripts/build_all_models.py --context-only

# INT8W 模型
python scripts/gen_int8w_context_binaries.py
# 输出: qnn_context_cache/lib*_fp32.serialized.bin
#         qnn_context_cache/lib*_int8w.serialized.bin

上下文二进制文件可为每次模型调用节省约 0.5-1 秒,通过跳过图编译实现。

运行推理

基本用法 (HTP 上 FP32)

python inference/infer_libero_episode_qnn_htp.py \
  --precision htp \
  --task-id 3 \
  --trial-id 3 \
  --n-action-steps 4 \
  --use-unrolled-ah \
  --output output/result.mp4

HTP 上 INT8W

python inference/infer_libero_episode_qnn_htp.py \
  --precision int8w \
  --task-id 3 \
  --trial-id 3 \
  --n-action-steps 4 \
  --use-unrolled-ah \
  --output output/result_int8w.mp4

关键参数

推理速度

单模型延迟 (Snapdragon X Elite)

模型执行时间单独测量 (不包含子进程启动开销 约 390ms/推理)。

实际耗时包括子进程开销 (约 130ms × 3 次调用)、临时文件 I/O,以及模型执行时间之外的 NumPy 转置操作。

端到端推理 (任务 3, 试验 3)

• 端到端延迟: 从观测捕获到首个动作输出的实际耗时
• 控制频率: 有效动作执行速率 = n_action_steps / 端到端延迟

CPU vs HTP 输出相似度 (余弦)

LLM 相似度较低是由于 32 层 Transformer 中的 FP16 内部累积。

架构

观测 (图像 + 状态 + 语言)
        │
        ▼
[视觉编码器]  pixel_values[1,3,512,512] → image_embeddings[1,64,960]
        │
        ▼
[LLM 主干]    image_embs×2 + lang_emb[1,48,960] + state[1,32]
                  → kv_keys/values[32,1,177,5,64] + prefix_pad_masks[1,177]
        │
        ▼
[动作头]     10 步流匹配 (欧拉法, t:1.0→0.1, dt=-0.1)
  (展开)      noisy_actions[1,50,32] → denoised_actions[1,50,32]
        │
        ▼
动作 (7-DoF: 末端执行器位置 + 轴角 + 夹爪)

关键实现说明

• HTP 上 FP32 DLL (而非 FP16): --float_bitwidth 16 会导致 ±512 饱和。FP32 DLL + HTP 内部自动转换为 FP16,精度更好。
• CPU 端嵌入: LLM v7 移除了 Gather 算子 (HTP INT32 输入 bug)。嵌入查找在 Python 中使用 embed_tokens.weight 完成。
• 消除布尔值: QNN 会剥离布尔链。fold_and_eliminate_bool.py 用 int32/float 运算替换。
• 上下文二进制: 预序列化的 HTP 图,每个模型加载快约 0.5 秒。
• 张量转置: HTP 自动转置某些维度;见推理脚本中的 run_vision_encoder、run_llm_backbone_htp。

文件结构

├── inference/
│   └── infer_libero_episode_qnn_htp.py   # 主要 LIBERO 推理脚本
├── scripts/
│   ├── build_all_models.py               # FP32 QNN 构建
│   ├── build_int8w_models.py             # W8A16 QNN 构建
│   ├── create_ah_unrolled.py             # 展开 10 步动作头
│   ├── gen_int8w_context_binaries.py     # 上下文二进制生成
│   ├── fold_and_eliminate_bool.py        # ONNX 布尔预处理
│   ├── fix_ah_expand_bool.py             # 动作头布尔修复
│   ├── fix_llm_embed_on_cpu.py           # LLM v7 嵌入移除
│   └── fix_cpp_for_htp.py                # C++ 输出 HTP 补丁
├── qnn_models/
│   ├── windows-aarch64-fp32/             # FP32 DLL (使用 HTP 后端运行)
│   └── windows-aarch64-int8w/            # W8A16 DLL
├── qnn_context_cache/
│   ├── lib*_fp32.serialized.bin          # 预编译 HTP 上下文
│   └── libaction_head_unrolled_htp.*     # 展开动作头 (10 步)
├── onnx_models/                          # 预处理 ONNX 模型
├── policy_preprocessor_step_5_normalizer_processor.safetensors
└── policy_postprocessor_step_1_unnormalizer_processor.safetensors

参考资料

• SmolVLA (HuggingFaceTB)
• LIBERO 基准测试
• Qualcomm AI Engine Direct SDK
• CPU 参考实现