Cosmos-Policy-RoboCasa Latency 비교 추가

← 목록으로

TL;DR

NVIDIA Cosmos Policy (2B, RoboCasa 체크포인트)를 isolated 1×B200에서 측정 → action call 353 ms (eager), per-action 11.0 ms.
2B인데도 per-action이 5B X-WAM의 eager(11.0)와 동률 — X-WAM의 우위는 backbone 크기가 아니라 compiled(7.4) + ANS 디커플링에서 나옴을 정량 확인.
Cosmos 계열 공통 약점: action-only fast path 부재. 5-step generate_samples_from_batch가 action+future+value를 공동 denoise → action call도 전액 지불, VAE decode만 skip(+106ms→459ms).
이미 등록된 16B Cosmos3-Nano-Policy-DROID와는 다른 모델. 같은 RoboCasa suite라 X-WAM과 per-action 직접 비교 가능.
compiled도 측정 → torch.compile 이득 ~0 (350 vs 353 ms). DiT가 TransformerEngine 커널(rmsnorm_fwd)을 써서 Dynamo graph break 발생 → 융합 불가. X-WAM(~1.5×)/Cosmos3-Nano(~1.25×)와 대조적.

353 ms

action call (eager)

350 ms

action call (compiled)

11.0 ms

per-action (eager=comp)

459 ms

full (+future VAE)

~1.0×

compile speedup

1배경 / 목적

X-WAM을 다른 VLA/world 모델과 input·latency 기준으로 비교하는 comparison/ 프레임워크(SCHEMA.md + models/*.yaml + build_table.py)에 모델을 추가하는 작업. 기존 등록: X-WAM, Cosmos3-Nano-DROID(16B), Pi0.5-DROID. 이번엔 NVIDIA Cosmos Policy(arXiv:2601.16163, Robotics/cosmos-policy OSS)의 RoboCasa 2B 체크포인트를 측정해 추가한다.

중요: 이 2B Cosmos Policy는 이미 등록돼 있는 16B Cosmos3-Nano-Policy-DROID와 다른 모델이다(전자=Cosmos-Predict2-2B 비디오 디퓨전 정책, 후자=Cosmos3-Nano MoT). 학습 없이 공식 weight latency만 측정 — 같은 RoboCasa suite라 X-WAM과 per-action 직접 비교가 핵심.

2작업 내용 (어떻게)

모델 구조 파악

inference 진입점은 get_action()(cosmos_utils.py:851). 핵심: model.generate_samples_from_batch(num_steps=5)가 action + future state + value를 공동(joint) denoise하고, generate_future_state_and_value_in_parallel=True이면 추가로 future image VAE decode까지 수행. 즉 Cosmos3-Nano처럼 action-only fast path가 없고, VAE decode만 토글된다(ANS류 step 디커플링 아님).

벤치 하네스 — 타 모델과 동일 프로토콜

cosmos-policy/tmp/bench_cosmos_policy_latency.py warmup 5 / timed 30, cuda-synced, median, batch=1, eager bf16 두 경로: action_only : generate_future_state_and_value_in_parallel=False (제어 latency) full : =True (action + future-image VAE decode + value, as-run) 합성 obs: 224×224×3 RGB ×3 view + 9-dim proprio + 실제 T5 임베딩(pkl→ndarray 경로 주입) 상수: COSMOS_IMAGE_SIZE=224, ACTION_DIM=7, PROPRIO_DIM=9, chunk=32, open_loop=16, steps_action=5

실행: 로그인 GPU0 idle, CUDA_VISIBLE_DEVICES=0 timeout 5400 .venv/bin/python ... (총 70 calls, <2분 → 디버그 임계 이내). T5 캐시 lock/on-demand 우회를 위해 pkl에서 임베딩 1개를 .float().numpy()로 로드해 ndarray 경로로 주입(latency는 shape 결정이라 대표성 OK).

3결과 (수치)

isolated 1×B200, eager, cuda-synced, warmup 5 / timed 30, batch 1:

경로	median	mean	p10/p90	min/max	per-action ÷32	per-action ÷16
action_only (5-step joint, VAE decode 없음)	353 ms	351	333 / 364	304 / 381	11.0 ms	22.1 ms
full as-run (+ future-image VAE decode + value)	459 ms	455	430 / 469	427 / 482	14.3 ms	28.7 ms

future-image VAE decode = +106 ms (353→459). action call은 이미 5-step joint denoise 전액을 지불하고 VAE decode만 skip한다.

compiled vs eager — `torch.compile(net, dynamic=False)`

레포 자체 recipe(text2world_model_rectified_flow.py:355)와 동일하게 model.net을 compile (FSDP setup 경로라 단일 GPU inference엔 미적용 → 벤치에서 명시). warmup 10 / timed 30.

경로	eager	compiled	speedup
action_only	353 ms	350 ms	~1.0× (노이즈 내)
full	459 ms	434 ms	~1.06×
per-action (÷32)	11.0 ms	11.0 ms	동일

torch.compile 이득 ~0. 로그에 원인 명시: DiT가 TransformerEngine 커널(transformer_engine_torch.PyCapsule.rmsnorm_fwd)을 호출 → Dynamo가 트레이스 못 해 graph break → 융합 불가. flex_attention은 eager에서도 이미 torch.compile로 감싸짐(minimal_v4_dit.py:352). pi0.5(GEMM-bound)와는 다른 원인이지만 결과는 같음(compile 무의미). X-WAM(~1.5×)/Cosmos3-Nano(~1.25×)와 대조적.

cross-model per-action (eager, normalize 기준)

Model	Params	Suite	Call latency	per-action	fast path?
X-WAM	5B	RoboCasa	237(comp)/352(eager)	7.4(comp) / 11.0(eager)	✅ ANS
Cosmos-Policy-RoboCasa	2B	RoboCasa	353(eager)	11.0(eager)	❌ joint
Pi0.5-DROID	4B	DROID	165(eager)	11.0(eager)	✅ action-only
Cosmos3-Nano-Policy-DROID	16B	DROID	961(comp)	30.0(comp)	❌ joint

4Takeaway

2B Cosmos Policy(eager 11.0 ms/action) = 5B X-WAM(eager 11.0). 같은 RoboCasa suite·eager 컬럼에서 per-action 동률 — X-WAM의 속도 우위는 backbone 크기(5B)가 아니라 (a) compiled(7.4) 경로와 (b) ANS 디커플링(action 10-step / video 50-step)에서 나온다. Cosmos Policy는 backbone이 작아도(2B) action+future+value를 5-step 공동 denoise하느라 per-action 이득이 상쇄된다.

Cosmos 계열(2B Policy, 16B Nano) 모두 action-only fast path 부재가 공통 약점 — ANS류 decoupling의 가치를 보여주는 직접 증거(X-WAM/pi0.5는 fast path 보유).

World-model 축에서 Cosmos-Policy는 1-step future obs(rollout video 아님)라 X-WAM 9프레임/Nano 33프레임과 성격이 다름 — 테이블/footnote에 명시.

5보완점 / 다음

compiled 측정 완료 (위 §3): torch.compile 이득 ~0 (TE graph break). eager/compiled 양쪽 모두 per-action 11.0 ms로 기재 완료. 추가 가속은 TE op을 custom-op으로 래핑하거나 compile 우회 attention 백엔드가 필요.
success_rate: 이번 add는 latency 전용. cosmos-policy smoke는 TurnOffMicrowave 3-seed 98%. 24-task sweep 후 paper 평균(~91%)과 SR 비교 가능(RoboCasa라 X-WAM과만).
open-loop 정규화 합의: 헤드라인은 ÷chunk(32)=11.0(X-WAM 관례), 배포-정직 ÷16=22.1도 병기. X-WAM 실제 open-loop 주기 확인 시 더 엄밀.

산출물: comparison/models/cosmos_policy_robocasa.yaml, comparison/comparison_table.md(행+footnote³+World-model 행), 벤치 cosmos-policy/tmp/bench_cosmos_policy_latency.py.

Cosmos-Policy-RoboCasa (2B) Latency — cross-model 비교 추가