TL;DR
- 완료: trainer.py 6개 위치 수정 (LoRA 적용, optimizer, DDP
find_unused_parameters=True, save/load hook, eval/train 전환)
- 완료: infer.py — transformer_2 로드 + LoRA fuse +
--boundary_ratio argparse
- 완료:
configs/egoexo4d_hcpt.yaml — boundary_ratio: 0.9 추가
- 핵심: transformer_2 LoRA는 별도
pytorch_lora_weights_transformer_2.safetensors로 저장 (diffusers 기본 save/load는 단일 transformer만 인식)
1 배경 / 목적
Steps 1–3 (schemas + sft_trainer.py)에서 expert switching logic이 완성됐다. 하지만 학습이 실제로 돌아가려면 trainer.py에서 두 expert 모두에 LoRA를 얹고, optimizer에 포함하고, DDP 설정을 바꾸고, 체크포인트 저장/복원 hook을 수정해야 한다. 추론 경로(infer.py)도 transformer_2를 로드하고 LoRA를 fuse해야 완성된다.
핵심 제약: diffusers의 기본 save_pretrained / load_lora_weights는 단일 transformer만 인식한다. transformer_2의 LoRA는 별도 safetensors 파일로 수동 저장/복원이 필요하다.
2 작업 내용
Step 4: core/finetune/trainer.py (6곳)
(4A) _init_distributed():
find_unused_parameters = (self.components.transformer_2 is not None)
— 매 step에 한 expert만 forward → DDP unused parameter 오류 방지
(4A) prepare_trainable_parameters():
if training_type == "lora" and transformer_2 is not None:
transformer_2.add_adapter(transformer_lora_config)
if gradient_checkpointing: transformer_2.enable_gradient_checkpointing()
ignore_list에 "transformer_2" 추가 (accelerator.prepare()가 device 관리)
(4B) prepare_optimizer():
cast_training_params에 transformer_2 포함
trainable_params += [p for p in transformer_2.parameters() if p.requires_grad]
(4C) prepare_for_training():
transformer_2 존재 시 accelerator.prepare()에 별도 분기로 포함
(4D) train() loop:
transformer_2.train() # 에폭 시작 시
models_to_accumulate에 transformer_2 추가
grad norm clipping: clip_params += list(transformer_2.parameters())
(4E) __prepare_saving_loading_hooks():
save hook: unwrapped is unwrap_model(..., transformer_2)로 두 expert 구분
→ transformer_2 LoRA: safetensors.torch.save_file(...)
"pytorch_lora_weights_transformer_2.safetensors"
load hook: 동일 구분 → set_peft_model_state_dict(transformer_2, ...)
(4F) _maybe_run_validation():
transformer_2.eval() / transformer_2.train() 전환
Step 5: infer.py
transformer_2 subfolder 존재 시:
transformer_2 = WanTransformer3DModel.from_pretrained(..., subfolder="transformer_2")
pipeline 생성:
pipe = WanI2VPipeline(..., transformer_2=transformer_2, boundary_ratio=args.boundary_ratio)
LoRA fuse:
"pytorch_lora_weights_transformer_2.safetensors" 존재 시
→ load_lora_weights(...) + fuse_lora()
argparse: --boundary_ratio (default=0.9)
Step 6: configs/egoexo4d_hcpt.yaml
# HCPT flags 섹션 아래
boundary_ratio: 0.9 # 추가 (+1줄)
3 결과
| 파일 | 변경 규모 | 검증 | 결과 |
|---|
| trainer.py | ~70줄 추가/수정 | python -m py_compile | OK |
| infer.py | ~20줄 추가/수정 | python -m py_compile | OK |
| egoexo4d_hcpt.yaml | +1줄 | yaml load | OK |
동적 검증 미완: Steps 1–3과 합산한 end-to-end smoke test (5-clip mini manifest, dual expert 양쪽 forward/backward 확인)가 아직 미실행.
잠재적 호환성 이슈: load_lora_weights(..., transformer=transformer_2) API는 diffusers 버전에 따라 다를 수 있음. 실행 시 확인 필요.
4 Takeaway
의미
Steps 1–6이 모두 완성되어 EgoX2_exp2 (HCPT)의 dual expert 지원 코드가 완비됐다. 두 WanTransformer3DModel이 각각 독립적인 LoRA adapter를 가지며, DDP에서 unused parameter 처리, 별도 safetensors 저장/로드, validation 모드 전환까지 일관되게 처리된다. smoke test 통과 후 바로 본 학습(8 GPU, 3660 clip) 제출 가능한 상태.
LoRA 저장 분리 (별도 safetensors 파일) 전략은 diffusers의 단일-transformer 가정을 우회하는 실용적 해법이다. 향후 diffusers가 dual expert를 공식 지원하면 이 부분만 교체하면 된다.
5 Next Steps
End-to-end smoke test
5-clip mini manifest (configs/hcpt_manifest_smoke.json)로 dual expert smoke 실행. 확인 포인트: (1) transformer, transformer_2 모두 forward/backward, (2) loss 단조 감소, (3) checkpoint-N에 pytorch_lora_weights_transformer_2.safetensors 생성, (4) resume 시 transformer_2 LoRA 복원 후 loss 연속성.
transformer_2 None 체크 일관성 검증
모든 코드 경로에서 if self.components.transformer_2 is not None 분기가 빠진 곳이 없는지 실행 시 검증. load_lora_weights API 호환성도 실제 diffusers 버전에서 확인.
본 학습 제출 (smoke 통과 후)
sbm "bash scripts/finetune_egoexo4d_hcpt.sh" --gres=gpu:8 -c 192 --mem 1600GB --qos=core-extra -J hcpt_train. 첫 ~25분은 eager-precompute (8 GPU DDP로 병렬화). loss 1000 step 이후 정상 하강 + rgb_loss / pt_loss 비율 확인.