Wan 2.2 I2V A14B 마이그레이션 구현 + Smoke 통과

TL;DR

Backbone 교체 완료: Wan 2.1 T2V 1.3B → Wan 2.2 I2V A14B. 36-ch native I2V input (noisy(16) + mask(4) + cond(16)), dual-expert MoE (transformer + transformer_2, boundary_ratio=0.875).
RGB-only smoke S1~S5 통과: 4 step, loss 0.85→1.38→0.82→0.92, peak 134 GB, 총 7분. Validation 50-step CFG mp4 628 KB 생성.
HCPT smoke 통과: 4 step, peak 156 GB (+22 GB), step 시간 22s (RGB-only 9s의 2.4×). HCPT loss가 RGB-only 대비 일관되게 높음 (PT row 추가 학습 부담).
사전 검증 (Phase A): T2V RGB-only smoke에서 Wan transformer spatial patch=2 제약 발견 — H_lat 짝수 필요 (432→H_lat=54 OK, 216→27 NG).

134 GB

RGB-only Peak VRAM

156 GB

HCPT Peak VRAM

105M

Trainable Params (rank=16)

9s/step

RGB-only 학습 속도

1 배경 / 목적

기존 robocasa 학습 backbone은 Wan 2.1 T2V 1.3B + frame-0 clean clamp trick으로 i2v를 흉내내는 구조였다.

기존 한계 세 가지:
1. Frame-0 leak에만 의존 → conditioning channel 없이 anchor drift 가능성.
2. 1.3B 용량 한계 → robocasa fine-grain manipulation + PT XYZ regression 디테일 부족.
3. Wan 2.2의 dual-expert MoE benefit 미수령.

Phase A (RGB-only T2V smoke)는 i2v 마이그레이션 전에 현 코드 RGB-only path가 살아있는지 확인하기 위한 사전 검증. Phase B에서 Wan 2.2 I2V A14B로 backbone을 교체하고, Phase C (cache 3-way split)는 Phase B 완료 후 적용하는 A→B→C 순서를 확정했다.

2 작업 내용

Phase A — RGB-only T2V Smoke (사전 검증)

신규 config configs/robocasa_t2v_smoke_rgbonly.yaml (use_pointmap: false)로 T2V 1.3B RGB-only 경로 검증. 코드 수정 0줄.

함정 발견: 첫 시도 train_resolution: 81x216x768에서 step 0 shape mismatch — target[H]=27, input[H]=26. Wan transformer spatial patch_size=2 → H_lat 홀수면 patch embed가 truncate. 216/8=27(홀수). 수정: 81x432x768 (H_lat=54, 짝수) 로 변경.

Phase B — Wan 2.2 I2V A14B 마이그레이션 (7개 Phase)

Phase 1: 모델 다운로드 — Wan2.2-I2V-A14B-Diffusers (118 GB) subfolder: transformer/ transformer_2/ vae/ text_encoder/ tokenizer/ scheduler/ image_encoder/ 부재 (CLIP 없음 — load_components None 분기 자동 처리) Phase 2: build_wan_i2v_full_grid_mask 추가 core/finetune/models/_base/width_concat.py frame 0 = 1 (전 spatial), 나머지 = 0 → shape [B, vae_scale_temporal, F_lat, H_lat, W_lat] Phase 3: build_robocasa_first_frame_latent에 start_frame 파라미터 추가 core/finetune/utils/robocasa_grid.py 기존 start_frame=0 hardcoded → random temporal crop의 s와 동기화 Phase 4: dataset first_frame_latent slot 추가 core/finetune/datasets/robocasa_hdf5.py monolithic safetensors에 first_frame_latent 추가 저장 Phase 5: Args.model_type Literal에 "wan-i2v-robocasa" 추가 core/finetune/schemas/args.py Phase 6: 신규 trainer 작성 (신규 모듈) core/finetune/models/wan_i2v_robocasa/sft_trainer.py (~370 lines) WanRobocasaI2VPipeline: 50-step denoising, cat([noisy, mask, cond]) 36-ch input dual-expert dispatch: i < boundary_step ? transformer : transformer_2 WanI2VRobocasaSftTrainer: load_components/encode/collate_fn/compute_loss/validation_step core/finetune/models/wan_i2v_robocasa/lora_trainer.py (SFT 상속) Phase 7: config + launcher configs/robocasa_i2v.yaml (full HCPT, rank=64, boundary_ratio=0.875) configs/robocasa_i2v_smoke_rgbonly.yaml / smoke_hcpt.yaml scripts/run_robocasa_i2v.sh

HCPT shape mismatch 수정

HCPT smoke 첫 시도에서 first_frame_latent shape (1,16,1,54,96) ≠ expected (B,16,1,108,96) 에러. 원인: dataset은 self.height=432를 per-cell 픽셀 높이로 사용 (H_lat=54), first_frame 헬퍼는 height를 total grid 높이로 해석해 per-row=216→H_lat=27로 계산. Fix: dataset이 헬퍼 호출 시 effective_height = self.height * 2 if use_pointmap 전달.

3 결과

검증 항목	RGB-only	HCPT
4-step loss 범위	0.85 / 1.38 / 0.82 / 0.92	1.29 / 1.73 / 1.38 / 1.64
grad_norm 범위	0.13 / 0.29 / 0.09 / 0.22	0.24 / 0.37 / 0.28 / 0.35
NaN 발생	없음	없음
Peak GPU memory	134 GB / 188 GB	156 GB / 188 GB
Step 시간	9 s/step	22 s/step (2.4×)
Validation 시간	~3분 (50-step CFG)	~3분
Trainable params	105M (rank=16, transformer + transformer_2 양쪽)
Total wall-clock	7분	12분
Validation mp4	628 KB, 81 frame, 432×768	567 KB

Dual-expert dispatch 정상 동작: boundary_ratio=0.875 → train 시 timestep > 0.875면 transformer_2, validation 시 step 인덱스 > 43면 transformer_2 사용 (50 step × 0.875 = 43.75).

HCPT loss가 RGB-only 대비 일관되게 높음: PT row 추가 학습 부담. 둘 다 NaN 없이 finite — 정상 동작으로 판단.

신규/수정 파일 요약

파일	종류
`core/finetune/models/_base/width_concat.py`	build_wan_i2v_full_grid_mask 추가
`core/finetune/utils/robocasa_grid.py`	start_frame 파라미터 추가
`core/finetune/datasets/robocasa_hdf5.py`	first_frame_latent slot, effective_height 패치
`core/finetune/models/wan_i2v_robocasa/sft_trainer.py`	신규 ~370 lines
`core/finetune/models/wan_i2v_robocasa/lora_trainer.py`	신규 (SFT 상속)
`configs/robocasa_i2v*.yaml`	신규 3개 (full, smoke_rgbonly, smoke_hcpt)
`scripts/run_robocasa_i2v.sh`	신규 launcher

4 Takeaway

의미

Wan 2.2 I2V A14B가 robocasa 2-view 데이터에 대해 RGB-only + HCPT 양 모드에서 36-ch 입력 forward, 50-step CFG inference, checkpoint 저장까지 e2e 동작을 확인했다. 기존 T2V 1.3B 대비 capacity 14× 증가와 native first-frame conditioning이 추가됐다.

HCPT step 시간이 RGB-only 대비 2.4× (22s vs 9s)지만 peak VRAM이 156 GB로 4-GPU B200 (188 GB × 4) 범위 내. 본격 LoRA rank=64 학습에서 rank=16 대비 trainable params 4× 증가가 예상되지만 OOM 위험은 낮다.

Validation 비용은 50-step × CFG × dual-expert로 3분/sample. 학습 1 step ~9s 기준 500 step마다 val 3min × 2 sample = 6min으로 전체 학습 시간의 ~8%.

5 Next Steps

본격 4-GPU LoRA 학습 sbm 제출

Phase C (cache 3-way split) 완료 후 sbm 제출:

conda activate <env> sbm "bash scripts/run_robocasa_i2v.sh" --gres=gpu:4 -c 56 --mem 800GB --qos=extra # configs/robocasa_i2v.yaml: rank=64, boundary_ratio=0.875, use_pointmap=true

DDP find_unused_parameters=True 오버헤드

dual-expert dispatch는 timestep마다 한쪽 expert만 사용 → unused param 발생 → DDP find_unused_parameters=True 필수. ~10-15% throughput 손실 예상.

Validation half-resolution 한계

HCPT 학습은 H_lat=108에서 일어나지만 validation pipe는 H_lat=54에서 sample. crash는 없음, visualization 품질만 약간 손해. 우선순위 낮음.