RDT-1B 昇腾 NPU 适配

robotics-diffusion-transformer/rdt-1b(RDT-1B: A Diffusion Foundation Model for Bimanual Manipulation)在华为昇腾 Ascend 910 NPU 上的推理适配与 CPU 精度/性能对比。

模型简介

RDT-1B 是清华大学 TSAIL 组提出的 1B 参数模仿学习扩散 Transformer,基于 100 万+ 条多机器人 episode 预训练。给定语言指令与最多 3 视角 RGB 图像,RDT 预测未来 64 步机器人动作;统一动作空间兼容单臂/双臂、关节/末端执行器、位置/速度等多种机器人本体。

• 输入:语言 token + 图像 token + 本体感觉 state(128 维)+ 动作掩码 + 控制频率
• 输出:64 步 × 128 维动作序列
• 预训练数据:RT-1、RH20T、DROID、BridgeData V2、RoboSet 及 Open X-Embodiment 子集等 46 个数据集

环境要求

部署步骤

# 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('robotics-diffusion-transformer/rdt-1b', local_dir='/data/rdt-1b/models')"

# 3. 官方模型代码: 本仓库 rdt_code/ 已内置(rdt_pytorch, 来自 thu-ml/RoboticsDiffusionTransformer)

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

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

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

关键适配点

1. 官方代码本地化(rdt_code/):RDT-1B 无 transformers remote code,官方 rdt_pytorch 从 GitHub 仓库获取,sys.path 引入后即可加载。模型权重 pytorch_model.bin 的 state_dict 前缀(model.*、lang_adaptor.*、img_adaptor.*、state_adaptor.*)与官方 RDTRunner 完全一致,load_state_dict 零缺失零多余。
2. 精度选择 fp32:官方默认 dtype=torch.bfloat16,昇腾 910 原生支持 bf16;本脚本以 fp32 加载以获得 CPU/NPU 逐位可比的数值结果(用户如需 bf16 可改 dtype=torch.bfloat16)。
3. SDPA 注意力:RDT Block 自注意力(timm Attention)与交叉注意力均使用 F.scaled_dot_product_attention,昇腾上由 torch_npu 原生支持,无需修改。
4. 确定性对比:初始噪声 noisy_action 用 CPU 固定种子 torch.Generator 生成后再搬运到目标设备,DPM-Solver 采样流程完全确定性,保证两端输入逐位一致,输出可直接做数值对比。
5. 合成观测:语言/图像/本体感觉均为 token 级条件输入,固定 seed 合成即可完成端到端策略推理(无需外部编码器权重)。

推理示例

import torch, torch_npu  # NPU 环境
from models.rdt_runner import RDTRunner  # rdt_code/ 已加入 sys.path

runner = RDTRunner(action_dim=128, pred_horizon=64, config=cfg, ...)  # 见 inference.py
runner.load_state_dict(torch.load("models/pytorch_model.bin", map_location="cpu"))
runner = runner.to("npu:0").eval()

action = runner.predict_action(
    lang_tokens=...,   # [1, L, 4096]
    lang_attn_mask=...,  # [1, L] bool
    img_tokens=...,    # [1, 4374, 1152]
    state_tokens=...,  # [1, 1, 128]
    action_mask=...,   # [1, 1, 128] 0-1 float
    ctrl_freqs=...,    # [1] 控制频率
)  # -> [1, 64, 128] 动作序列

CPU vs NPU 精度对比

输入:固定种子合成观测(64 语言 token + 4374 图像 token + 128 维 state),初始噪声 seed=42,DPM-Solver 5 步采样;计时为 3 次迭代均值(预热 1 次后)。实测数据见 results_*.json。

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

产出文件

许可证

MIT(遵循上游 RDT,见 models/README.md)。