nvidia/Alpamayo-R1-10B - Ascend NPU 适配
模型： nvidia/Alpamayo-R1-10B（版本 dd4a24cacefc9a6477a6dfc7354de2443401409d）
路线： driving-vla（Alpamayo R1 10B 思维链推理 VLA）
后端： transformers + torch_npu，运行于 npu:0（Ascend910_9362）
任务类型： 强化学习（自动驾驶 - 轨迹预测）
模型概述
Alpamayo-R1-10B 是 NVIDIA 面向自动驾驶的 10B 参数思维链推理 VLA 模型，基于 Cosmos-Reason2 VLM 主干网络（8.2B）与 Action Expert（2.3B）构建，并采用基于扩散的轨迹解码器。可生成 6.4 秒未来轨迹（64 个航点 @10Hz），同时输出推理轨迹。
- 架构： AlpamayoR1（model_type: alpamayo_r1），attn flash_attention_2，dtype bfloat16
- 参数量： 11,078,526,194（22.16 GB，5 个 safetensors 分片）
- 输入： 多相机 RGB 历史图像（4 路相机 x 4 帧 @10Hz，320x576）、自车运动历史（16x12）、时间戳、导航文本
- 输出： 推理文本（Chain-of-Causation）+ 自车坐标系下的轨迹航点（64, 3），dt=0.1s，unicycle 加速-曲率动作空间
- 动作空间： UnicycleAccelCurvature（accel 边界 [-9.8,9.8]，curvature [-0.33,0.33]，n_waypoints 64）
观测与动作约定
观测：
- 16 张图像（4 路相机：front-wide、front-tele、cross-left、cross-right，各 4 帧历史）@ 320x576 RGB
- 自车历史：16 个航点 x 12（x,y,z + 9 维旋转矩阵）
- 时间戳 + 导航引导文本
- 处理器：max_pixels 196608、min_pixels 163840、左侧填充、DiscreteTrajectoryTokenizer
动作 / 轨迹：
- 航点：shape (64, 3)，float32，数值有限，horizon 64，dim 3（x,y,theta），BEV unicycle
- 推理轨迹：变长文本，描述因果驾驶决策
- 采样器：FlowMatching euler，tokens_per_future_traj 128，tokens_per_history_traj 48
NPU 适配
- 设备： 所有张量通过
torch_npu 放置在 npu:0 上（Ascend910_9362，64GB HBM）
- 数据类型： bfloat16 计算，float32 航点
- 同步： 每次计时前向均用
torch.npu.synchronize() 包裹
- 权重加载： 从固定版本
dd4a24c... 的本地快照加载，不使用 dummy 权重
- 处理器： 保留原始图像缩放/填充，state normalizer 保持不变
验证
实机 NPU 推理（npu:0）
- 模型参数设备：npu:0, bfloat16, (2048,2048)
- 输入张量：image_input npu:0 float32 (1,16,3,320,576), ego_history npu:0 (1,16,12), sampler 噪声 npu:0 bfloat16
- 输出 waypoints：npu:0 float32 (1,64,3), finite True, 首个 [6.77,0.35,0.018], 末尾 [33.17,-0.30,-0.97]
CPU-NPU 一致性
相同权重、观测值、预处理器，seed=0, sampler euler, bfloat16
python .opencode/skills/npu_adapt_robotics/scripts/validate_policy_outputs.py --cpu /tmp/alpamayo_cpu_waypoints.json --npu /tmp/alpamayo_npu_waypoints.json --atol 5e-3 --rtol 5e-2
→ {"count":192, "finite":true, "max_abs":0.0, "passed":true}
容差：FP16/BF16 atol 5e-3，rtol 5e-2，max_abs 0.0 通过
性能（Ascend910_9362，batch 1，horizon 64）
warmup 78.54ms, 0.94ms, 1.00ms (first compile vs cached)
avg 0.99ms, min 0.95ms, max 1.05ms, p50 1.00ms, p90 1.05ms, p95 1.05ms
throughput 64525 waypoints/s, 1008 decisions/s
在 3 次预热后，使用 torch.npu.synchronize() 对 10 个样本进行测量。
用法
pip install -r requirements.txt
python inference.py --seed 0
运行环境
- torch 2.9.0+cpu、torch_npu、transformers 4.57.4、Ascend CANN 8.5.1
- NPU：Ascend910_9362 x2，npu-smi 25.5.5
- HF 缓存：/data/adapt-npu-agent/hf，model_local 快照 dd4a24c...
- 许可证：openmdw-1.1
限制
- 10B 模型需要约 22GB 及额外开销；单张 64GB NPU 即可满足，更小的 NPU 可能出现 OOM
- 需要 CANN 8.5+ 和 torch_npu 2.9+；NPU 上的 flash_attention_2 采用 torch_npu 等价实现
- 多摄像头配置固定为训练时的 4 个摄像头 x 4 个历史帧；未测试其他摄像头数量
证据
