Latent-IDM C N=4 구현 완료 & Smoke PASS

TL;DR

N=12 dense routing은 step당 I/O 808 GB → ~100 sec/step → 60K = 70일+로 비현실적. N=4 sparse(blocks 5/14/23/32)로 pivot — mmap=True + 4-layer slice로 264 MB/sample (13.6×↓).
Jobs 5961·6054·6108·6158·5845·6262·6321 7차 실패 각각 root-cause 분석 후 fix 완료 (preempt/ENOSPC/persistent_workers/wrong-class 순서로 버그 누적).
1-GPU end-to-end smoke PASS: 242.6M params, peak 5.46 GB, loss 1.62→1.14, 3 step finite. 8 GPU core-own 60K 학습 제출 대기.

N=4

Cosmos layers

264 MB

I/O per sample

242.6M

params

5.46 GB

peak (1 GPU B=2)

7회

실패 job

PASS

smoke test

1배경/목적 (왜)

전 카드 요약: Latent-IDM C는 π0.5 KI/π0 MoE 영감의 Layer-Routed ActionDiT — ActionDiT block i가 Cosmos hidden layer f(i)만 cross-attn KV로 받는 1:1 layer↔depth 정렬 구조. 설계 단계에서는 Phase 1을 N=12 dense(12 Cosmos layer × 12 ActionDiT block 1:1)로 잡았었다.

N=12 I/O 불가 판정: per sample 792 MB (12 layer × 66 MB) × batch 8 × 8 GPU × num_workers = /dev/shm ENOSPC + step당 I/O 808 GB / NVMe 8 GB/s = 100 sec/step → 60K = 70일+. mmap으로도 데이터양 자체는 줄지 않음.

목표를 바꿨다: N=12 dense는 Phase 2 이후 참고 baseline으로만 두고, N=4 sparse (Cosmos block 5/14/23/32 = early/mid-early/mid-late/late 균등 선택)를 Phase 1 본 학습으로 확정. mmap=True + 4-layer slice로 264 MB/sample (13.6× I/O 절감). 이미 빌드된 N=12 cache에서 인덱스 1,4,7,10만 읽어 재사용 — cache 재빌드 불필요.

이 진행 카드는 N=4 pivot 결정 이후 코드 구현 → 7차 실패 디버깅 → 1-GPU smoke PASS까지의 과정을 기록한다.

2작업 내용 (어떻게)

N=4 Layer 선택 및 mmap 설계

N=12 cache (/home/nas_main/taewoongkang/cache/cosmos_hidden_mg30/multi12_sigma30/, 6.6 TB, 1871 .pt files)는 이미 빌드 완료. 파일 구조: bf16 dict {"hidden_multi": (12, 24, 27, 48, 5120), "layers": [2,5,...,35], ...}.

N=4 selected layer indices → cache indices: cache idx 1 → Cosmos block 5 (early, depth 13.9%) cache idx 4 → Cosmos block 14 (mid-early, depth 38.9%) cache idx 7 → Cosmos block 23 (mid-late, depth 63.9%) cache idx 10 → Cosmos block 32 (late, depth 88.9%) I/O per sample: N=12 no-mmap: 12 layer × 5 frames × 27 × 48 × 5120 × 2B = 3.6 GB/sample N=4 mmap: 4 layer × 5 frames × 27 × 48 × 5120 × 2B = 264 MB/sample (13.6× ↓) Step time 추정: N=12: ~100 sec/step → 60K = 70일+ (blocked) N=4 : ~25-30 sec/step → 60K = 18-22일 (feasible)

주요 코드 변경

파일	변경 내용
`dataset_latent_idm.py`	`RobocasaLatentIDMDatasetHiddenMultiLayer` (line 568+) — `selected_layer_indices` 파라미터 + `mmap=True` 추가. (이전 버그: line 415의 HiddenRaw class에 잘못 추가했었음)
`cosmos_mirror_dit.py`	`LayerRoutedActionDiT` 신규 — block i가 `encoder_hidden_states[:, i]`만 cross-attn. 기존 ActionDiT 변경 없음.
`flow_matching_action_head_idm.py`	`cosmos_hidden_multilayer_cache` feature_source 분기 — `kv_projectors = nn.ModuleList × N`, encode_visual "cosmos_latent_routed" mode 반환. `dataloader_persistent_workers=(num_workers>0)` 자동 처리.
`base_latent_C_layer_routed_n4.yaml`	N=4 본 학습용 YAML. `multilayer_n_cosmos_layers: 4`, `multilayer_selected_indices: "1,4,7,10"`.
`run_idm_train_robocasa_latent_C_layer_routed_n4.sh`	8 GPU 본 학습 entry. `--dataloader_num_workers 0` (shm 압박 해소). SEL=1,4,7,10.
`tmp/debug_latent_C_n4.py`	NAS cache 직접 읽는 1-GPU end-to-end smoke — 마지막 실행 PASS 확인.

Model 파라미터 구성

LayerRoutedActionDiT (N=4)

84.2M

kv_projectors ×4 (5120→1024)

20.5M

action enc/dec/timestep_emb 등

~138M

Total trainable

242.6M

3결과 (디버깅 회고 + Smoke PASS)

7차 실패 job 전체 분석

Job	설정	실패 원인	Fix
5961	8 GPU core-extra, N=12	8 GPU core-extra slot 6일 대기 → CANCELLED	own quota로 전환
6054·6108·6158	4 GPU sub, N=12	sub preempt cycle 1h < stage time 2h → PREEMPTED ×3	own quota / extra 전환 (sub는 long stage job에 부적합)
5845	4 GPU sub smoke, N=12	792 MB/sample × prefetch queue × 8 ranks → `/dev/shm` ENOSPC	`--dataloader_num_workers 0` 전체 적용
6262	8 GPU core-own, N=12	`idm_training.py:283` `dataloader_persistent_workers=True` 하드코딩 + `num_workers=0` 충돌 → ValueError	`persistent_workers=(num_workers > 0)` 자동 처리
6321	8 GPU core-own, N=4	`selected_layer_indices`를 HiddenRaw class(line 415)에 추가했으나 실제 사용 class는 HiddenMultiLayer(line 568) → TypeError	올바른 class의 `__init__` + `__getitem__` 수정. Python introspection으로 확인.

1-GPU end-to-end smoke PASS (tmp/debug_latent_C_n4.py): dataset construct 23s, 1 sample 0.7s, model 39s, 3 step backward all finite. Peak 5.46 GB @ B=2. Loss 1.62→1.14. 로그인 노드 1 GPU smoke가 2시간 stage 후 dataset instantiation에서 죽는 패턴(6262, 6321)을 1분 안에 catch하는 가장 효율적인 디버깅 방법임을 검증.

비교 baseline (학습 완료 기준)

모델	학습	EEF mm (5ep)	EEF mm (24ep)	비고
Pixel v2_8gpu	60K, B=1024	134.8	—	MG-30 baseline
Latent-A v1 (L27 sliced)	60K, B=1024	131.8	114.2
Latent-A v3 (3-layer concat)	60K	—	110.9	throughput 26%↓
Latent-A multistream	55K (preempt)	—	104.99	현재 최고
Latent-C N=4 (예정)	60K (미시작)	—	목표 ≤125mm	H-LR 채택 기준

학습 제출 명령어 (minhopark 셸)

conda activate gr00t-idm sbmr 5 "bash scripts/run_idm_train_robocasa_latent_C_layer_routed_n4.sh" \ --gres=gpu:8 -c 112 --mem 1600GB \ --qos=core-own --exclude=worker-5 \ -J latentC_n4_60k # 기대치: # Stage: NAS → /tmp 6.6 TB cp (~2시간) # Step time: ~25-30 sec/step # 60K wall: ~18-22일 # own quota → preempt 없음, sbmr 5 retries = OOM/NaN/하드웨어 방어용 # output: outputs/idm_robocasa_latent_C_n4_60k_8gpu/

4Takeaway

mmap=True가 결정적: N=4 slice + mmap으로 264 MB/sample (N=12 no-mmap 3.6 GB 대비 13.6×↓). 동일 6.6 TB N=12 cache를 그대로 재사용하면서 Phase 2에서 sparse N 변경 시에도 cache 재빌드 불필요 — 6.6 TB는 영구 자산.

sub partition은 long stage job에 부적합: preempt cycle (~1h) < stage time (~2h). latent-C처럼 stage가 긴 job은 반드시 own/extra quota 사용. 3개 job을 낭비한 교훈.

N=12 dense는 실질적으로 불가

6262·6321 시도 때 비로소 명확해진 것: N=12 I/O는 mmap으로도 해결이 안 된다. 데이터양 자체가 12×. 설계 문서에 "8 GPU, eff batch 1024, ~11 sec/step"으로 추정했지만 실제 I/O 한계를 과소평가. N=4가 sweet spot — π0.5 KI의 block-level 1:1 alignment 가설을 검증하기에 충분하고, 추가 Phase 2에서 layer 수를 늘리려면 cache만 subset 선택하면 됨.

7번의 실패 중 마지막 2개(6262, 6321)는 단일 파일이 두 개의 서로 다른 dataset class를 가진 구조에서 잘못된 class에 파라미터를 추가하는 전형적인 copy-paste 버그. 로그인 노드 1-GPU smoke가 sbatch 전에 반드시 선행되어야 한다는 교훈을 재확인.

5Next Steps

8 GPU core-own 60K 학습 제출 (minhopark 셸) — 예상 wall ~18-22일. outputs/idm_robocasa_latent_C_n4_60k_8gpu/ 확인.
5-ep simulator eval (60K 도달 후): latent-A 131.8 mm vs latent-A multistream 104.99 mm 사이 어느 위치인지 확인. Pass 기준 ≤125 mm → H-LR 채택.
Phase 2 sparse routing 결정: Phase 1 결과에 따라 N=4+depth=12 (1:3 routing) / N=6+depth=12 (1:2) / N=4+depth=4 (1:1 sparse) 중 선택. 모두 기존 N=12 cache 재사용.
heldout 24-ep eval용 multi-layer cache: 학습 720 ep는 이미 빌드됨. heldout ep도 cache_cosmos_hidden_mg30_multi.py로 추가 빌드 필요 (~7분).
latent-A multistream 60K 완료 확인 (Job 6216 진행 중) — 104.99 mm (55K) 대비 수렴 확인.