Wan 2.2 I2V A14B 마이그레이션 계획

← 목록으로

TL;DR

현재 robocasa 학습 backbone(Wan2.1-T2V-1.3B)의 근본 한계(frame-0 clamp trick, 1.3B 용량, HCPT 불가)를 해소하기 위해 Wan2.2-I2V-A14B로 마이그레이션 계획 수립
모델 capacity 1.3B → 28B (14×), native 36ch I2V conditioning, dual-expert MoE(boundary_ratio=0.875) 모두 활성화
CLIP 의존 제거 — image_encoder=None 분기 이미 trainer에 존재, dataset 변경 최소
새 trainer wan_i2v_robocasa + dataset first_frame_latent slot + config robocasa_i2v.yaml 신설, 7단계 구현 후 S1~S5 smoke test 완료 시 본 학습 제출

14×

capacity 증가

28B

A14B dual expert

36ch

native I2V input

구현 Phase

0.875

boundary_ratio

1 배경/목적 (왜)

현재 robocasa 학습은 Wan2.1-T2V-1.3B-Diffusers backbone에서 동작한다. WanT2VSftTrainer.compute_loss는 16ch input의 vanilla T2V 모델에 temporal frame-0 clean clamp trick(apply_first_frame_clean_noise)을 끼워서 I2V-스러운 효과만 흉내내는 구조 — 모델 자체는 unconditional 생성을 학습하면서 frame-0 anchor만 후처리로 박는다.

T2V backbone 4가지 근본 한계

Frame-0 leak에만 의존: conditioning channel이 없어서 첫 프레임 정보가 latent attention에 자연스럽게 주입 안 됨. denoising 후반부로 갈수록 anchor가 흐려지는 reanimation drift 관찰.
1.3B 용량 한계: A14B(14B×2 expert) 대비 표현력 매우 제한적. fine-grain robot manipulation + PT XYZ regression 디테일 부족.
HCPT conditioning 미지원: 2×2 grid의 exo column을 clean conditioning으로 두는 native 구조(wan_i2v/sft_trainer.py에 이미 구현됨)가 T2V backbone에서는 불가능.
Wan 2.2 dual-expert benefit 미수령: high-noise/low-noise expert MoE 구조를 1.3B single expert로는 활용 불가.

선택 variant: Wan2.2-I2V-A14B — Wan 2.1 VAE 호환(기존 latent cache dtype/shape 동일), CLIP 없음(image_encoder=None 분기 이미 trainer에 존재), A14B 라인 최대 capacity.

2 작업 내용 (어떻게) — 7-Phase 설계

Variant 비교 (최종 선택 근거)

Variant	Params	VAE	CLIP	비고
T2V-A14B	14B×2	Wan 2.1 호환	None	T2V 유지 시 I2V 목표와 안 맞음
I2V-A14B	14B×2	Wan 2.1 호환	None	✓ 선택 — CLIP 없음, dual expert, 최대 capacity
I2V-14B (2.1)	14B×1	Wan 2.1	CLIP-G 필요	dual expert 없음, CLIP 필요
TI2V-5B	5B×1	새 16× VAE	None	latent cache 전체 invalidate

7-Phase 구현 계획

Phase 1

모델 다운로드

HuggingFace Wan-AI/Wan2.2-I2V-A14B-Diffusers — transformer/, transformer_2/, vae/, text_encoder/ 한 트리. ~56GB(28B×bf16).

Phase 2

신규 trainer 작성

core/finetune/models/wan_i2v_robocasa/sft_trainer.py (~350 lines). compute_loss에서 native 36ch I2V conditioning + dual expert dispatch + frame-0 loss skip.

Phase 3

Dataset 확장

RobocasaHDF5Dataset에 first_frame_latent [16,1,H,W] slot 추가. safetensors에 key 병기. cache_root 분리(robocasa_i2v_a14b).

Phase 4

Config 신설

configs/robocasa_i2v.yaml. model_name: wan-i2v-robocasa, LoRA rank=64, boundary_ratio=0.875, cache_root 분리.

Phase 5

Trainer dispatch 라우팅

Args.model_type literal에 wan-i2v-robocasa 추가. finetune.py의 get_model_cls에서 신규 trainer로 분기.

Phase 6

Validation 적응

WanRobocasaI2VPipeline 신설 — 36ch input, dual expert loop, decode 구조. trainer.py의 robocasa_hdf5 분기에서 pipe만 교체.

Phase 7

Launcher 스크립트

scripts/run_robocasa_i2v.sh — run_robocasa_t2v.sh에서 CONFIG default만 i2v로 변경. 나머지 동일.

핵심 compute_loss 구조 (Phase 2)

latent_model_input = cat([noisy_latents, condition], dim=1) # [B, 36, F_lat, H, W] # condition = cat([mask(4ch), cond_latent(16ch)], dim=1) # cond_latent = [first_frame_latent | zeros × (F_lat-1)] # Dual expert dispatch (timestep-boundary 기준) boundary_idx = int(num_train_ts * boundary_ratio) # 0.875 active = transformer_2 if timesteps_idx >= boundary_idx else transformer # Loss: frame-0 제외 (clean clamp으로 degenerate) loss = F.mse_loss(pred[:, :, 1:], target[:, :, 1:])

Cache 분리 구조 (Phase 3)

# 기존 T2V cache (보존) ./datasets/robocasa_t2v/cache/video_latent/robocasa_hdf5/81x216x768_hcpt/ └── encoded_exo_ego_gt_video [16, 21, 54, 96] # 신규 I2V cache (분리) ./datasets/robocasa_i2v_a14b/cache/video_latent/robocasa_hdf5/81x216x768_hcpt/ ├── encoded_exo_ego_gt_video [16, 21, 54, 96] ← VAE 동일, 재인코딩 └── first_frame_latent [16, 1, 54, 96] ← 신규 key

3 결과 (수치) — 계획 기준 기대치

구현 전 계획 단계이므로 실측 수치 없음. 아래는 설계 분석 기반 예상치.

항목	현재 (T2V 1.3B)	목표 (I2V A14B)
모델 파라미터	1.3B (single expert)	28B (14B×2 dual expert)
Conditioning	frame-0 clamp trick (16ch)	native I2V (36ch = 16 noise + 4 mask + 16 cond)
CLIP 의존	없음	없음 (A14B CLIP-free)
VRAM (smoke test 기준)	~40GB (1.3B LoRA)	~81GB (A14B LoRA rank=64, HCPT Option F 기준)
Input channels	16ch	36ch
LoRA rank	128	64 (A14B base param 많아 동등 효과)
Cache invalidation	—	신규 cache_root — 기존 T2V cache 보존

위험 요소: LoRA dual-expert 메모리 (A14B×2 + adapter×2 + optimizer state). 8 GPU×B200(180GB)로 rank=64 가능할 것으로 추정. fail 시 rank=32 또는 transformer_2 frozen 대응.

DDP overhead: dual-expert dispatch는 timestep마다 한쪽만 사용 → unused param 발생 → find_unused_parameters=True 강제 (trainer.py:101 자동 활성화). ~10-15% 속도 손실 불가피.

4 Takeaway

설계 핵심 인사이트

Wan2.2-I2V-A14B 선택의 핵심은 VAE 호환성이다 — Wan 2.1과 동일한 VAE weight를 써서 latent shape/scale이 그대로 유지되므로, 기존 T2V cache의 encoded_exo_ego_gt_video를 이론상 재사용 가능하다(안전을 위해 별도 cache_root로 분리). TI2V-5B는 새 16× VAE를 써서 latent 전체 invalidate 가 불가피했으므로 기각.

CLIP-free 선택(A14B)은 dataset 구조 단순화에도 기여한다 — image_embedding cache slot 불필요, collate_fn에서 image_embedding=None으로 처리. 이미 sft_trainer.py:751-769에 None 분기 존재.

14× capacity: 1.3B→28B effective. robocasa fine-grain manipulation + PT XYZ regression 정확도 향상 기대.
Native I2V conditioning: frame-0 anchor가 transformer attention에 자연 흐름 → reanimation drift 감소 기대.
기존 코드 재사용 극대화: dual expert hook(trainer.py:101, 298-321), CLIP-free 분기, validation grid composition 모두 그대로. 신규 작성 최소화.
T2V rollback 가능: cache_root 분리 + 기존 trainer 보존으로 언제든 T2V로 복귀 가능.

5 Next Steps

Smoke Test 체크리스트 (S1~S5)

단계	확인 항목	성공 기준
S1 모델 로드	load_components 동작, dual expert 인식	`transformer_2 is not None`, `image_encoder is None`, `in_channels==36`
S2 Forward shape	batch_size=1, compute_loss 1 step	`latent_model_input.shape == [1, 36, 21, 54, 96]`
S3 Cache 빌드	task 1개 latent + first_frame cache 생성	safetensors에 두 key 모두 존재
S4 Validation pipeline	1-step checkpoint로 50-step inference	output mp4 생성 (`F,H_total,W_total,3` 형태)
S5 학습 안정성	4 GPU LoRA 100 steps	loss 발산 없음, grad_norm ~0.1-1.0 범위

이후 단계

Phase 1-7 구현 → S1~S5 smoke test → 본 학습(sbm "bash scripts/run_robocasa_i2v.sh" --gres=gpu:8 -c 112 --mem 1600GB)
Phase 1 직후 latent cache RGB/PT 분리 plan(plan-latent_cache_rgb_pt_split.md) 적용 — I2V 마이그레이션 이후 cache 구조가 확정되면 split 적용.
transformer_2 weight diff 검토: high-noise vs low-noise expert weight 분포 유사하면 transformer_2 frozen + transformer만 LoRA로 메모리 절감 가능.
HCPT mode 보존 검증: PT row 포함 4-cell 구조가 동일하게 동작하는지, Wan 2.2 VAE에서 4DNeX 통계가 그대로인지 확인.

미결 한계: build_robocasa_first_frame_latent이 현재 start_frame=0 고정 → Phase 3에서 start_frame 파라미터 추가 필요. 또한 T2V cache를 I2V cache로 symlink 최적화는 VAE 동일성 smoke 확인 후 결정.