在 RTX 4090 上部署并加速 Cosmos 3 机器人策略

2026/08/21

Embodied AI Quantization Edge AI

用量化装进显存
用采样器调优少算一点
用 kernel 与编译再省一点
用 rollout 拓扑和录制扩展到整个评测
哪些没成
接下来

代码：cosmos-lite v0.3.0
模型：Nano GenW8A8 · Edge GenW8A8

NVIDIA 的 Cosmos 3¹ 带来了在 DROID² 上训练的机器人策略。不过想自己把它跑起来并不容易：16B 的 Nano 策略光是加载就要约 33 GB 显存，官方工作流也默认你有 A100 或 H100 这样的数据中心 GPU。单张 RTX 4090 之前完全跑不动。

Cosmos Lite 就是为此做的。它构建在 Cosmos Framework 之上，是一个面向 Cosmos 3 策略的高效推理 runtime，覆盖两类用法：一是本地部署，以 batch-one 方式提供策略推理，可直接对接真机；二是并行 rollout，由一个共享的推理服务同时驱动多个仿真环境，并把 rollout 轨迹导出用于后续训练。一台 8 卡主机上，最快的 profile 跑完全部 1,200 个 RoboLab episode 不到五小时。

Edge BF16（官方默认配置）对比 Edge GenW8A8（优化 runtime），在 RoboLab 中闭环 rollout。

评测覆盖 RoboLab³ 全部 120 个任务，每任务 10 个 episode；所有 profile 用同一组采样参数（guidance 3、两步 UniPC 去噪、shift 5）。延迟统计的是单张 RTX 4090 上每个动作请求的耗时。

模型	产物	显存 (GB)	请求 p50 (ms)	SR (%)
Edge 4B	BF16（NVIDIA 报告⁴）	—	—	22.90
Edge 4B	BF16	9.20	582.0	20.92
Edge 4B	GenW8A8	8.79	331.4	19.25
Nano 16B	BF16（NVIDIA 报告⁴）	—	—	36.80
Nano 16B	W8A16	21.42	2,403.0	31.50
Nano 16B	GenW8A8	15.51	958.5	31.67

Nano 这两行能出现，靠的是量化：BF16 装不进 24 GB 显卡。GenW8A8 在保住任务成功率的前提下，把 Nano 相对 W8A16 提速 2.5×，把 Edge 相对 BF16 提速 1.76×。有一点要说清楚：表里的成功率全部低于 NVIDIA 报告值（Nano 31.7% 对 36.8%，Edge 19.3% 对 22.9%）。所谓“保住成功率“，指的是统计上不可区分，不是零损失。真要上自己的任务，请先验证成功率。

用量化装进显存

第一个问题是把 16B 模型装进 24 GB 显存，怎么做取决于模型结构。Cosmos 3 策略采用 Mixture-of-Transformers（MoT）结构：每个解码层都有两条分支，一条是处理语言、状态和视觉条件的理解分支，另一条是把 diffusion latents 解成动作的生成分支。两条分支的 token 数、矩阵形状、精度敏感度都不一样，量化时必须区别对待。量化对象只限 MoT decoder 中的 Linear 层，vision、embedding、action adapter 与 VAE 全部保持 BF16。

产物	MoT linears	校准	峰值显存 (GB)
Nano 16B BF16	BF16（上游基线）	—	~33（放不下）
Nano 16B W8A16	全部 W8A16（Marlin⁵）	无	21.42
Nano 16B GenW8A8	生成分支 FP8 W8A8，其余 W4A16	128 个 DROID episode	15.51
Edge 4B BF16	BF16（上游基线）	—	9.20
Edge 4B GenW8A8	生成分支 FP8 W8A8，其余 W4A16	128 个 DROID episode	8.79

在测过的所有方案里，混合式的 GenW8A8 tradeoff 最好。全模型 W8A8 我同样构建并评测过：Nano 上它只比 GenW8A8 快约 10 ms，显存却多占 3.5 GB；Edge 上干脆更慢。道理不复杂：大 GEMM 集中在生成分支，那正是 FP8 tensor core 的用武之地；理解分支用 128 个训练 episode 做过 AWQ 式校准⁶之后，W4A16 的精度就够用了。

这一阶段还有条经验值得单独记下：开环动作误差预测不了闭环成功率。按 replay 误差挑出的最优产物，在配对 rollout 里反而比一个更简单的方案低了 12 个百分点；校准能改善误差分布，但压不掉偶发的动作尖峰。因此发布里的每一个精度决策，最终标准都是 rollout 成功率。

用采样器调优少算一点

显存解决之后，下一个杠杆是每个请求的计算量。上游 RoboLab 服务器的默认采样配置是 guidance 3 加四步 UniPC⁷，v0.1.0 作为第一个可部署版本沿用了它。把去噪步数从四改成二，改的只是一个 YAML 覆盖项，每个请求的去噪器前向就少了一半。

Guidance / steps	请求 p50 (ms)	BananaOnPlate SR
g3 / s4	4,110	43/50
g3 / s2	2,403	45/50
g1 / s4	2,431	32/50
g1 / s2	1,565	40/50

两步把 Nano 的请求延迟降低 41%，而在上面的配对 rollout 里，两个模型的成功率持平甚至更好。这件事没那么反直觉：UniPC 是数值求解器，步数多意味着积分更细，不意味着策略更好。学出的向量场与推理 schedule 只要有一点错配，多算的求解步就可能把动作推离好的轨迹。Edge 表现一致，从四步到两步，成功率从 21/50 涨到 34/50。

表格同时回答了 guidance 为什么留在 3：降到 1 能省掉 CFG⁸ 的双份前向，延迟上看似白赚，闭环成功率却损失过大，过不了 gate。当然，这些都是单任务上的配对 gate；全量评测自带统计不确定性，换一个环境或机器人，最优设置未必相同。

用 kernel 与编译再省一点

再下一个杠杆，是把剩下的计算本身做便宜。项目初期的 profiling 推翻了两个普遍预期：attention 只占 GPU kernel 时间的约 4%；而通常被认为能提速的 weight-only 量化，在 Nano 的真实形状上基本持平，在 Edge 上甚至是负收益。下面是通过测试留下的改动，按落地先后：

改动	实测效果
生成分支上的 CUTLASS⁹ FP8 W8A8 GEMM	Nano −20%、Edge −12% request p50（vs. W8A16）
整个 language block 级别的 `torch.compile`	Nano −14.5%、Edge −12.2% request p50
Q/K/V 与 gate/up 共享 FP8 激活量化	−2%，动作逐位一致
按形状选 attention：长序列生成 attention 用 SageAttention¹⁰，其余用 FlashAttention2；条件 K/V cache	Nano 约 −12%；cache 仅 Nano 启用
Edge SM89 形状专属 Triton FP8 GEMM	仅 Edge：−8.4%；Nano 上被否（见下）
稀疏输入变换（只 resize 观测帧，其余零填充）	sample build 约 37 ms → 2 ms

单独说说 Triton 那一行。CUTLASS 在 Edge 的主导 SM89 形状上 occupancy 偏低，换成按形状调优的 Triton kernel 后端到端快了约 8%。同一套 kernel 在 Nano 上也更快，但配对 rollout 从 49/50 掉到了 47/50。两个 episode 的差距多半是噪声（还需进一步验证），只是对照本身已经站在 49/50，没有质量余量去换这约 9% 的延迟，Nano 最终留在了 CUTLASS。

用 rollout 拓扑和录制扩展到整个评测

前面三个杠杆都作用在单个请求内部，而一次 RoboLab-120 全量评测是 1,200 个闭环 rollout，仿真时间不会因为模型变快就消失：量化、采样器、kernel 全部做完之后，策略推理在 lane 总时间里还是只占一小部分¹¹。所以最后的杠杆不在模型里，而在一台多卡主机上，仿真和策略服务各用哪几张卡。

实际布局是 2 张卡跑策略服务、6 张卡跑仿真 lane，各 profile 的具体分配见基准报告。每条 lane 是一个 Isaac Sim 进程，内部向量化 10 个环境，一个任务的 10 个 episode 因此是同时跑的，仅这一项就把单卡吞吐从 3.51 提到 18.24 environment steps/s。此外，同一个 scene 的任务会在同一条 lane 里连续执行，因为跨 scene 重建要约 100 秒，同 scene 切换只要约 10 秒，差着一个数量级。

布局固定之后，前面推理侧的收益就能一路叠加到全量评测上——采样器、量化、kernel 全用上，一个完整 profile 从 9 小时 21 分降到 4 小时 51 分（1.93×）：

配置	Wall time	GPU-hours	vs. BF16 g3/s4
Edge BF16, g3/s4	9 h 21 m	74.8	—
Edge BF16, g3/s2	5 h 35 m	44.7	1.67×
Edge GenW8A8, g3/s2	4 h 51 m	38.8	1.93×

评测跑完还要变成训练数据：成功的 episode 会导出成 LeRobot v3 数据集（数据生成指南）。这里有个现实的坑：直接录 10 路并行 rollout 会把显存撑爆，因为上游录制栈把每一步的相机 tensor 都留在 GPU 上¹²，显存随 episode 长度线性增长，一个 600 步的 episode 自己就要约 19 GB。发布版 recorder 的做法是把数据及时移到有界的 CPU 侧缓冲，配合分块 HDF5 落盘和流式 H.264 编码，10 路并行 rollout 就能稳定落在 24 GB 显存之内，代价是相对不录制多 11.0% 的 wall time。

哪些没成

这轮探索大部分由 coding agent 驱动，失败尝试来得便宜，死胡同也就攒了一长串，值得记下来。按失败发生在哪一层分组：

倒在算子或端到端层面：

全模型 torch.compile：自定义 GEMM 处图断裂，没有收益，首个请求要 20 秒。
静态形状编译与整请求 CUDA Graphs：新 prompt 长度触发重编译；capture 的开销换不来稳定收益。
全局统一 attention 后端：SageAttention 在短条件 attention 上反而更慢；FlashInfer 始终没跑赢 FlashAttention2。
拼接 projection 权重：个别算子更快，但 Nano 峰值显存飙到 19 GB，端到端表现不稳定；后来被共享激活量化取代。
Edge 的条件 K/V cache：数值上正确，但收益无法复现；cache 因此只随 Nano 发布。
torchao INT8 weight-only：在模型真实形状上比 BF16 慢 3–4 倍。
CFG 两个分支合并 batch：attention 调用次数更少、单次更大，端到端反而更慢。
原生 batch=2 推理：四个 profile 上都比两次串行请求慢 9–21%，动作还有偏差（最大元素误差约 0.1）。
更短的动作 chunk（32→16/8）：名义计算量更小，实测延迟更差，控制语义也变了。

kernel 赢了，rollout 输了：