rdt-1b 昇腾 NPU 适配

模型简介

RDT-1B（robotics-diffusion-transformer/rdt-1b）是清华 THU-ML 开源的 10 亿参数机器人扩散 Transformer：以扩散建模机器人动作轨迹，语言指令（T5-xxl 4096 维 token）与多相机视觉观测（SigLIP 1152 维 token）经 adaptor 进入 DiT 主干（28 层 × 2048 隐层 × 32 头），输出 64 步 × 128 维动作块（action chunk），DDPM 5 步去噪采样。

• 原始权重地址（HuggingFace）：https://huggingface.co/robotics-diffusion-transformer/rdt-1b
• 参数量：1.229B（fp16 权重约 2.4GB）
• 官方实现：thu-ml/robotics_diffusion_transformer（toy_rdt.py 布局）

验证环境

推理引擎：torch_npu（npu:0，fp16）。

主干实现说明

外网代理对该 codebase 拉取失败，故按 checkpoint 键名逐张量自建主干（inference.py 内 RDT/RDTCore/Block 等，与官方 toy_rdt.py 布局一一对应）：model.*（DiT 核心：x_pos_embed(1,67,2048)、t_embedder.mlp、freq_embedder.mlp、28×blocks（自注意力 fused-QKV+QK-Norm / 交叉注意力 / FFN，LayerNorm 无 bias）、final_layer）+ 顶层 lang/img/state_adaptor；DDPM 采用 squaredcos_cap_v2、prediction_type=sample、5 步推理，与 config.json 的 noise_scheduler 一致。

推理操作流程

source /usr/local/Ascend/cann-8.5.1/set_env.sh
export RDT1B_MODEL_DIR=/tmp/bigw/rdt-1b     # 已下载权重目录
python inference.py --device npu

inference.py 自动完成：按 config 建主干 → strict 加载审计 → 合成观测（1024×4096 语言 token + 4374×1152 图像 token + 2×128 本体状态）+ 语言指令 → 5 步 DDPM 去噪 → 动作块统计与功能检查 → 性能统计，结果写入 /tmp/rdt1b_npu_results.json。

加载正确性门禁

• 门禁 1（strict load）：torch.load + load_state_dict(strict=True) 实测 0 missing / 0 unexpected / 0 mismatched（0/0/0），618 个张量逐一对齐（含全部 28 层 block、adaptor、pos embed）。
• 门禁 2（config 落地审计）：config.json 全部维度字段（action_dim=128、pred_horizon=64、img_cond_len=4374、max_lang_cond_len=1024、rdt.depth/hidden_size/num_heads=28/2048/32、lang/img/state_token_dim、mlp2x/mlp3x_gelu adaptor 类型、DDPM 调度参数）均直接驱动主干构建，无静默默认值。
• 功能检查（真跑证据）：① 输出动作块非恒定（std=0.2009，范围 [-1.55, 1.73]，64 步连续平滑）；② 语言条件通路有效（lang token ×5 放大后输出 mean|Δx0|=0.0377，交叉注意力通路真实生效）。如实披露：随机合成 conditioning 在训练分布外，跨指令差异（0.0033）与噪声敏感性（0.0014）偏小，x0 预测在此条件下由条件主导。

昇腾 NPU 真实输出（测试用例）

用例：合成观测 + 语言指令 'pick_up_the_red_block'（25Hz 控制频率，5 步 DDPM）。

action chunk shape: (1, 64, 128) step0 (t=0) first 6 dims: [0.2122, 0.2952, 0.2106, -1.4561, -0.0635, 1.6924] step1 (t=1) first 6 dims: [0.2151, 0.2905, 0.2155, -1.5156, -0.0536, 1.7100] step32 (t=32) first 6 dims: [0.1788, 0.3865, 0.1683, -1.3564, 0.0059, 1.6152] chunk stats: mean=0.0214 std=0.2009 min=-1.5469 max=1.7305

指令 'open_the_drawer'（不同观测种子）输出同样非恒定（std=0.1919）且数值不同。

NPU 性能（fp16, npu:0）

截图

• assets/agent_workflow.png：模型上卡日志（config 建主干 → strict 加载 → npu:0）
• assets/npu_device_call.png：npu-smi 实时占用 + torch_npu 设备断言
• assets/model_result.png：真机 NPU 动作预测输出

#NPU