EgoX2_exp2 Dual Expert Steps 1–3 구현

TL;DR

목표: Wan 2.2 A14B의 transformer_2 (low-noise refiner) 를 HCPT 학습/추론 경로에 통합
완료: Args.boundary_ratio 필드, Wan_Components.transformer_2 필드, sft_trainer pipeline + compute_loss expert switching 구현
backward compatible: transformer_2 subfolder 없으면 자동 fallback — 기존 single-expert 체크포인트 그대로 동작
미완: Steps 4–6 (trainer.py LoRA/DDP/save-hook, infer.py, config yaml) 별도 작업 필요

Files Changed

~35

Lines Added

0.9

boundary_ratio

1 배경 / 목적

Wan 2.2 I2V A14B 체크포인트는 transformer/ (high-noise expert, denoising 처음 90%)와 transformer_2/ (low-noise refiner, 마지막 10%) 두 개의 WanTransformer3DModel을 갖고 있다. model_index.json에 boundary_ratio: 0.9가 정의되어 있다.

기존 문제: sft_trainer.py의 load_components()가 transformer subfolder만 로드하고 transformer_2는 완전히 무시했다. 학습/추론 모두 high-noise expert 하나로만 동작 중 — low-noise refinement 단계 누락으로 디테일 품질 저하 가능성.

현재 설치된 diffusers 0.34.0에는 dual expert 지원이 없어 수동 구현이 필요하다. EgoX2_exp2는 HCPT (Height-Concat Point Map + ego Prior) 구조를 사용하므로 exp1과 다른 코드 경로를 갖는다.

2 작업 내용

Step 1: `core/finetune/schemas/args.py`

Args 모델에 boundary_ratio: float = 0.9 필드 추가 (LoRA 섹션 바로 아래). argparse에 --boundary_ratio 인자 추가. 변경 라인: +4.

Step 2: `core/finetune/schemas/components.py`

Wan_Components에 transformer_2: Any = None 필드 추가. 기존 vars(self) 순회 로직이 device 관리를 자동으로 처리하므로 추가 변경 불필요. 변경 라인: +1.

Step 3: `core/finetune/models/wan_i2v/sft_trainer.py` (5곳)

(3A) Pipeline.__init__: transformer_2, boundary_ratio 파라미터 추가 register_modules(transformer_2=transformer_2) self.boundary_ratio = boundary_ratio (3B) Pipeline.__call__ denoising loop: boundary_step = int(len(timesteps) * self.boundary_ratio) if transformer_2 is None: boundary_step = len(timesteps) active_transformer = self.transformer if i < boundary_step else self.transformer_2 (3C) load_components(): if os.path.isdir(os.path.join(model_path, "transformer_2")): components.transformer_2 = WanTransformer3DModel.from_pretrained( model_path, subfolder="transformer_2") (3D) initialize_pipeline(): pipeline 생성 시 transformer_2=..., boundary_ratio=... 전달 (3E) compute_loss(): boundary_idx = int(num_train_ts * self.args.boundary_ratio) if transformer_2 is not None and (timesteps_idx >= boundary_idx).any(): active_transformer = unwrap_model(..., transformer_2) else: active_transformer = unwrap_model(..., transformer) predicted = active_transformer(...)

설계 결정: training에서 batch_size=1이므로 매 step에서 한 expert만 forward/backward. 통계적으로 90% 확률로 transformer, 10% 확률로 transformer_2가 선택된다. inference denoising loop에서는 index 기반으로 결정론적으로 전환.

3 결과

파일	변경 라인	검증 방법	결과
args.py	+4	python -m ast	통과
components.py	+1	python -m ast	통과
sft_trainer.py	~+30 (net)	python -m ast	통과

Backward compatibility: transformer_2 subfolder가 없으면 components.transformer_2 = None, boundary_step = len(timesteps) → 기존 single-expert 동작과 100% 동일. grep 검증으로 필드 반영 확인.

동적 검증 미수행: 정적 syntax 검사만 통과. Steps 4–6 완료 후 smoke test 필요.

4 Takeaway

의미

Wan 2.2 A14B의 dual expert 구조를 HCPT 학습 파이프라인에 통합하기 위한 핵심 switching 로직이 완성됐다. transformer_2가 없는 기존 체크포인트에서는 자동 fallback되므로 기존 smoke test / checkpoint에 영향이 없다. training loop의 expert 선택(timestep_idx vs boundary_idx)이 inference denoising loop의 expert 선택(step index vs boundary_step)과 자연스럽게 대응된다는 점이 설계의 핵심.

5 Next Steps

Steps 4–6 구현 필요

trainer.py: transformer_2에 LoRA config 적용, optimizer에 포함, find_unused_parameters=True (매 step에 한 expert만 active), save/load hooks (별도 pytorch_lora_weights_transformer_2.safetensors), eval/train 모드 전환.

infer.py: transformer_2 로드 + LoRA fuse, --boundary_ratio argparse.

configs/egoexo4d_hcpt.yaml: boundary_ratio: 0.9 추가.

End-to-end smoke test

Steps 4–6 완료 후 5-clip mini manifest로 dual expert smoke test 실행. 양 expert 모두 forward/backward 확인, loss 단조 감소, checkpoint 저장 후 transformer_2 LoRA 별도 파일 존재 확인.