Latent-IDM C — Layer-Routed ActionDiT 설계 (π0.5 KI 영감)

TL;DR

가설 H-LR: ActionDiT block i가 Cosmos hidden layer f(i)만 KV로 받으면 (block-level 1:1 routing), 모든 block이 같은 KV를 보는 현재 방식보다 좋다
동기: π0.5 KI에서 VLM layer i ↔ Action expert block i가 joint self-attn에서 자동 대응되는 구조에서 영감. Multistream은 token concat으로 fuse; Layer-routed는 각 block에서 독립적으로 해당 layer KV만 봄
Phase 1: dense N=12 ([2,5,8,...,35]), ActionDiT 12 block 1:1 routing, cache 6.7 TB, ~7.6일, 목표 ≤125 mm EEF
코드 인프라 완료: RobocasaLatentIDMDatasetHiddenMultiLayer, cosmos_hidden_multilayer_cache feature_source 분기, dataset collate 추가됨

Dense Layers

6.7 TB

Cache Size

7.6일

학습 ETA

≤125 mm

Phase 1 목표

1 배경 / 목적

현재 Latent-IDM 트랙들의 공통 한계: 모든 ActionDiT block이 동일한 Cosmos hidden을 KV로 본다.

v3 multi-layer concat 한계: 3 layer hidden을 dim-level로 concat (5120×3 → 1024 Linear). 모든 block이 동일 fused KV — layer별 정보 차별성 없음. 24-ep EEF 110.9 mm, multistream 104.99 mm로 개선했지만 block-level alignment는 아직 미탐색.

π0.5 KI 구조에서의 영감: VLM (depth=18) + Action expert (depth=18)가 동일 depth MoE transformer로 구성. 매 transformer block에서 modality별 Q/K/V 분리 + joint self-attn으로 VLM layer i의 hidden ↔ Action expert block i가 자동으로 대응된다. Cosmos가 frozen + cached인 우리 setting에서 이 correspondence를 수동으로 구현한다.

가설 H-LR: ActionDiT block i(얕은 layer: 시각 디테일 처리) ↔ Cosmos hidden layer f(i)(얕은 layer: 저수준 특징) 정렬이 자연스럽게 이루어지면 EEF가 개선된다. 반가설 H0: Cosmos hidden 27(late layer)이 이미 충분히 condensed → routing 무효.

2 작업 내용

Phase 1 — Dense Layer Routing (N=12)

Layer 선택: Cosmos DiT 36 blocks 중 균등 12개.

LAYERS = [2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35] ActionDiT block i (0-11) ↔ LAYERS[i] (1:1) layer 27 (현재 Latent-B baseline) → LAYERS[9]=29로 근접 매핑

캐시 빌드

스크립트: scripts/cache_cosmos_hidden_mg30_multi.py 저장 schema: torch.save({"hidden": List[bf16 (T,H,W,D)] of len=12, "layers": LAYERS}) - 단일 파일 per (ep, w_start): 1회 read로 12 layer 전부 - 파일 크기: 12 × 318 MB = 3.8 GB / window (bf16) - 총 디스크: 3.8 GB × 1871 windows = 6.7 TB (NAS 영구화) 빌드 시간: ~7 min (8 GPU, 1871 windows)

Dataset 클래스 (완료)

gr00t/data/dataset_latent_idm.py에 RobocasaLatentIDMDatasetHiddenMultiLayer + collate_latent_idm_hidden_multi 추가됨.

__getitem__: torch.load → 12 layer list, LAT_LEN=5 slice 반환: hidden_multi: bf16 (N_layers=12, LAT_LEN=5, H=27, W=48, D=5120) per-sample: 12 × 66 MB = 792 MB/sample (bf16) workers=2 (latent-B의 4 → 줄임, shm 압박)

모델 — LayerRoutedActionDiT

class LayerRoutedActionDiT(nn.Module): # ActionDiT와 동일 block 구조 # forward만 다름: block i는 encoder_hidden_states[:, i]만 cross-attn KV로 받음 def forward(self, hidden_states, encoder_hidden_states, ...): for i, block in enumerate(self.blocks): x = block(x, encoder_hidden_states[:, i], t_emb) ...

Action head 분기 (완료)

flow_matching_action_head_idm.py에 feature_source = "cosmos_hidden_multilayer_cache" 분기 추가. kv_projectors = nn.ModuleList([nn.Linear(5120, 1024) for _ in range(12)]) — per-layer projector.

Phase 2 — Sparse Routing (Phase 1 결과 후 결정)

후보	Layer 수	ActionDiT depth	Cache	트리거 조건
(a) 4-layer-4-block	4 [5,14,23,32]	4	2.2 TB	Phase 1 ≤120 mm (큰 win)
(b) 4-layer-12-block	4 [5,14,23,32]	12 (3 block per layer)	2.2 TB	Phase 1 ∈[130,140] mm
(c) 6-layer-12-block	6 [5,11,17,23,29,35]	12 (2 block per layer)	3.3 TB	Phase 1 ∈[120,130] mm

3 비용 및 목표 수치

Run	Cache	학습 cost	EEF 목표
Latent-A multistream (현재 best)	84 GB (v3 재사용)	60K × ~11 s/step	104.99 mm ✨
Phase 1 dense N=12	6.7 TB	60K × ~11 s/step ≈ 7.6일	≤125 mm
Phase 2 (a) 4-layer-4-block	2.2 TB	60K × ~5 s/step ≈ 3.5일	≤ Phase 1
Phase 2 (b) 4-layer-12-block	2.2 TB	60K × ~9 s/step ≈ 6.3일	≤ Phase 1
Phase 2 (c) 6-layer-12-block	3.3 TB	60K × ~10 s/step ≈ 6.9일	≤ Phase 1

Phase 1 통과 기준: avg EEF ≤ 125 mm (5-ep simulator replay). 125-135 mm → 24-ep 확장 평가. >140 mm 또는 loss 발산 → H0 채택, Phase 2 작게 시도.

메모리 사전 검증 필수: per-sample 792 MB × batch 8 = 6.3 GB/rank hidden_multi. 8 ranks × 6.3 = 50 GB 합 (B200 140 GB → OK). activation: 12 block × cross-attn KV (B=8, T=6480, D=1024) ~25 GB/rank → gradient_checkpointing 필수.

4 Takeaway

H-LR 채택 시 의미

단순 KV 확장(v3 concat, multistream)이 아닌 layer ↔ depth 정렬이 latent-IDM의 sweet spot임이 검증됨. π0.5 KI의 MoE per-block 가설이 frozen Cosmos backbone + cached KV setting에서도 transfer됨을 보임. 향후 모든 latent-IDM 실험의 default 구조가 될 수 있다.

H-LR 기각 시 의미

Cosmos block 27은 이미 충분히 condensed한 representation. extra layer 정보는 marginal noise. → latent track은 batch + sliced KV로 충분. 다음 투자는 OOD/domain transfer 시나리오 검증으로 전환.

Multistream과의 차별점: multistream은 token 수를 3×(192)로 늘려 모든 block이 보게 함. Layer-routed는 token 수를 64 그대로 유지하되 각 block이 보는 layer를 다르게 함. 두 접근이 상보적 — 결과에 따라 hybrid 설계 가능.

5 Next Steps

즉시 — 캐시 빌드 제출

scripts/run_cache_hidden_multi_mg30.sh sbm 제출 (8 GPU, core-extra). mg30 + heldout 각각 ~7 min. NAS 영구화 경로: /home/nas_main/taewoongkang/cache/cosmos_hidden_mg30/multi12_sigma30/ (6.7 TB).

Smoke test — single GPU 10 step

tmp/test_latent_C_routed_forward.py: dummy hidden_multi (2, 12, 5, 27, 48, 5120), B=2, 10 step backward. peak <30 GB, loss finite 확인.

8 GPU 100-step smoke + 60K 본 학습

step time ~11 s/step (v3 대비 ~30% slow — kv_projectors 12개 + 12 block routing overhead). 60K 본 학습 sbmr 5 제출. Multistream 60K 완료 후 병렬 제출 가능.

잠재 위험 대비

per-layer projector 12개가 over-parameterize로 small-data overfit 위험 → train_loss vs eval gap 모니터링. early layer hidden norm 불균형 → trunc_normal init scale=1/√(in_dim) + 1k step warmup. 필요 시 per-layer RMSNorm 추가.