EgoX2_exp1 Dual Expert 구현 계획

TL;DR

Wan 2.2 A14B는 transformer (high-noise 90%) + transformer_2 (low-noise 10%) dual expert 구조인데, exp1 코드가 transformer만 사용 중 — low-noise refinement 단계 누락
exp1은 exp2와 코드 구조가 다름 (WanTransformer3DModel_GGA, ptmap 32ch output, WanDROIDPipeline). 별도 구현 계획 수립
수정 대상: 7개 파일 (args, components, sft_trainer, trainer, droid_pipeline, infer_droid, config)
메모리 예상: 28B params(bf16) + LoRA = ~82GB, B200 178GB 내 여유

수정 파일 수

~82GB

예상 peak VRAM

0.9

boundary_ratio

32ch

ptmap output

1 배경/목적

Wan 2.2 I2V A14B 체크포인트는 두 개의 WanTransformer3DModel을 갖는다:

transformer/ — high-noise expert (denoising 처음 90%, 구조 생성)
transformer_2/ — low-noise refiner (denoising 마지막 10%, 디테일 품질)

model_index.json의 "boundary_ratio": 0.9가 전환 시점을 정의한다.

현재 문제: sft_trainer.py:load_components()(line 715)에서 transformer subfolder만 로드. transformer_2 완전 무시 → low-noise refinement 없이 high-noise expert만으로 학습/추론 중.

exp2에서 동일 이슈의 구현을 완료했으나(Steps 1–3), exp1은 WanTransformer3DModel_GGA, ptmap 32ch, WanDROIDPipeline 등 고유 구조를 가지므로 별도 계획이 필요하다.

2 작업 내용 (구현 계획)

exp1 고유 코드 구조

특성	설명
Custom transformer	`WanTransformer3DModel_GGA` (GGA attention 포함)
ptmap 모드	`use_pointmap=True` → `expand_proj_out(32)`, 32ch output (RGB 16 + PtMap 16)
Loss 함수	`_compute_loss_original` (RGB) + `_compute_loss_with_ptmap` (RGB+PtMap) 둘 다 수정 필요
Training pipeline	`WanWidthConcatImageToVideoPipeline` (GGA validation)
Inference pipeline	`WanDROIDPipeline` (DROID ptmap inference)
DDP 설정	`find_unused_parameters=False` → True로 변경 필요 (매 step 한 expert만 사용)

수정 파일 목록 (7개)

파일	핵심 변경
`core/finetune/schemas/args.py`	`boundary_ratio: float = 0.9` 필드 + argparse 추가
`core/finetune/schemas/components.py`	`Wan_Components.transformer_2: Any = None` 추가
`core/finetune/models/wan_i2v/sft_trainer.py`	load_components / compute_loss (×2) / initialize_pipeline / pipeline denoising loop — expert switching 로직
`core/finetune/trainer.py`	LoRA/optimizer/DDP/prepare/train loop/save-load hooks — transformer_2 전체 지원
`core/finetune/models/wan_i2v/droid_pipeline.py`	denoising loop에 boundary_step 기반 expert 전환
`infer_droid.py`	transformer_2 로드 + LoRA/proj_out weights 별도 복원
`configs/droid_train.yaml`	`boundary_ratio: 0.9` 추가

핵심 설계 결정

항목	결정	이유
Expert 선택 (training)	`timesteps_idx < boundary_idx`	inference(denoising loop index)와 자연스럽게 대응
DDP	`find_unused_parameters=True`	매 step에 한 expert만 active — unused params 존재
LoRA 저장	별도 `pytorch_lora_weights_transformer_2.safetensors`	diffusers 기본 save/load는 단일 transformer만 인식
proj_out 저장	별도 `proj_out_weights_transformer_2.safetensors`	두 expert의 proj_out 독립 학습
backward 호환	transformer_2 subfolder 없으면 None → 기존 코드 유지	기존 checkpoint 사용 가능

# Training expert selection (sft_trainer.py) boundary_idx = int(num_train_timesteps * self.args.boundary_ratio) # 90% 지점 active_transformer = ( transformer if timesteps_idx < boundary_idx else unwrap_model(self.accelerator, self.components.transformer_2) ) # Inference denoising loop (droid_pipeline.py) boundary_step = int(len(self.scheduler.timesteps) * self.boundary_ratio) for i, t in enumerate(self.scheduler.timesteps): active_transformer = self.transformer if i < boundary_step else self.transformer_2 noise_pred = active_transformer(...)

구현 순서

schemas (args + components) — 의존성 없음
sft_trainer.py load_components — transformer_2 로드
trainer.py LoRA/optimizer/DDP (4A~4E)
trainer.py save/load hooks (4F)
sft_trainer.py compute_loss — training expert selection (3B, 3C)
sft_trainer.py pipeline + WanWidthConcatImageToVideoPipeline (3D, 3E)
droid_pipeline.py — inference dual expert
infer_droid.py
configs/droid_train.yaml
Smoke test — 5 steps mini run, 양 expert forward 확인

3 결과 (현황)

구현 완료 (2026-04-16/17): 위 7개 파일 수정 완료. context.md 기준 "Dual expert 구현 완료, 학습 테스트 대기 중".

메모리 예산 (예상)

항목	용량
14B × 2 transformer (bf16)	56 GB
LoRA rank 128 × 2	~2 GB
Optimizer (LoRA params)	~4 GB
Peak activations + grad	~20 GB
총 예상	~82 GB (B200 178GB 내 여유)

Smoke test 미완: 구현 후 5-step mini run으로 양 expert 모두 forward 확인 필요. exp2는 동일 구조 smoke test에서 loss 3.82→2.75 (10 steps), peak VRAM 81 GB 확인됨 — exp1도 유사 예상.

4 Takeaway

exp1은 GGA attention, ptmap 32ch output, WanDROIDPipeline이라는 세 가지 고유 구조 때문에 exp2 구현을 그대로 가져올 수 없다. 특히 _compute_loss_original과 _compute_loss_with_ptmap 두 경로 모두 expert switching이 필요하고, WanWidthConcatImageToVideoPipeline과 WanDROIDPipeline 두 pipeline 모두 수정이 필요하다는 점이 복잡도의 주요 원인이다.

기존 checkpoint 호환성

transformer_2 subfolder가 없으면 자동 fallback되도록 설계해 backward compatibility를 보장했다. 기존 checkpoint-5000에서 resume 시 transformer_2 LoRA가 없으므로 처음부터 두 expert 모두 학습하는 방향이 맞다.

5 Next Steps

Smoke test 실행 — 5 steps mini run, loss 단조 감소 및 양 expert forward 확인. CUDA_VISIBLE_DEVICES=0 python ... --max_train_steps 5 (login node debug)
find_unused_parameters=True 오버헤드 측정 — DDP 통신 비용 증가 예상, 실제 throughput (it/s) 확인
전체 학습 제출: sbm "bash scripts/finetune_droid.sh" --gres=gpu:8 -c 192 --mem 1600GB --qos=core-extra
분리 logging — transformer_2 loss 별도 tensorboard 추적 (현재 merged loss만 logging)