Resume 시 LR 스케줄 재웜업 버그 수정

TL;DR

Part 1의 사각지대: checkpoint 파일명 기준 step 카운터 resume(Part 1)은 정상 동작했지만, self.lr_scheduler는 매 프로세스 재시작마다 새 LambdaLR로 생성되며 내부 스케줄 진행도가 0부터 다시 시작 — furniture_all_e128_pack 런이 step 750 / 1000 / 1500에서 200-step 재웜업을 3회 반복.
수정: lr_lambda를 checkpoint step(step0)만큼 offset해 즉시 목표 LR(constant phase)로 복귀시키고, optimizer moment가 실제로 초기화된 상태를 완충하기 위해 resume_warmup_steps(기본 50) 동안만 짧게 재웜업.
미검증: 이 패치는 아직 uncommitted 상태이며, 적용 후 첫 requeue가 발생하지 않아 실측 재현 확인 전.

3회

관측된 재웜업

200→50

재웜업 step 단축

750/1000/1500

재현 지점

1 배경 / 목적

Part 1(260907-resume_reliability_fixes)에서 resume_from이 transformer weight만 복원하고 optimizer/step 상태를 저장하지 않아 requeue마다 step counter가 0부터 시작하던 버그를, checkpoint 디렉토리명(checkpoint_step_N)에서 N을 역산해 self.step을 초기화하는 패치로 고쳤다. 그 카드의 Next Steps는 이미 이렇게 예고하고 있었다:

Part 1 예고(Next Steps 원문): "optimizer/scheduler 상태 자체는 여전히 checkpoint에 저장되지 않음 — resume 시 Adam의 momentum/variance가 초기화되므로, 매 preemption마다 사실상 warm-restart에 가까운 손실 스파이크 위험이 남아 있다"

실제로 이번에 그 위험이 LR 스케줄 쪽에서 구체적으로 확인됐다. self.step(체크포인트 번호)은 Part 1 이후 정확히 이어지지만, self.lr_scheduler = torch.optim.lr_scheduler.LambdaLR(...)는 매 프로세스 재시작마다 새로 생성되므로 내부 스케줄 진행도는 항상 0부터 다시 계산됐다. 그 결과 warmup_constant_lambda가 이미 warmup(200 step)을 끝낸 시점(step 750/1000/1500)에서도 다시 0-step 기준 낮은 LR부터 끌어올리는 재웜업을 반복했다.

2 작업 내용

wan_va/train.py의 Trainer.__init__ resume 분기를 확장해, checkpoint step(step0)을 lr_lambda에 오프셋으로 반영하고, optimizer moment 리셋을 완충하기 위한 짧은 재웜업만 남겼다.

if hasattr(config, 'resume_from') and config.resume_from: m = re.search(r"checkpoint_step_(\d+)$", str(config.resume_from).rstrip("/")) if m: self.step = int(m.group(1)) step0 = self.step resume_warmup = max(1, int(getattr(config, 'resume_warmup_steps', 50))) self.lr_scheduler = torch.optim.lr_scheduler.LambdaLR( self.optimizer, lr_lambda=lambda s: warmup_constant_lambda(s + step0, warmup_steps=config.warmup_steps) * min(1.0, (s + 1) / resume_warmup))

s + step0로 warmup_constant_lambda를 평가하므로 원래 스케줄상 이미 warmup을 끝낸 시점이면 즉시 constant phase LR을 돌려준다. 다만 optimizer state(Adam m/v)는 실제로 fresh 상태이므로, 곱셈항 min(1.0, (s+1)/resume_warmup)으로 resume_warmup_steps(기본 50, config로 override 가능) 동안만 0→1로 짧게 완충한다. 이 파라미터는 va_furniture_all_e128_train_cfg.py와 va_furniture_all_e128_pack_cfg.py 양쪽에 기본값 50으로 추가했다.

3 결과 (수치)

재현 지점 (step)	수정 전 관측	수정 후 기대 동작
750	200-step 재웜업 반복 (코드 주석 근거: "the e128 run re-warmed 200 steps at 750, 1000 and 1500")	즉시 constant LR 근접, 50-step 짧은 완충만
1000
1500

미검증: 이 패치는 아직 커밋 전(uncommitted) 상태이고, 적용된 이후의 requeue 사이클이 아직 발생하지 않아 재웜업이 실제로 50-step으로 줄었는지는 로그로 확인되지 않았다. "수정 전 3회 재현"의 근거도 이번 세션에서 직접 step-단위 LR curve를 그려 검증한 것이 아니라, 코드에 남겨진 주석(과거 세션이 furniture_all_e128_pack의 반복 preemption 로그를 보고 남긴 관측)에 의존한다 — 다음 requeue에서 wandb의 lr 곡선을 직접 확인해야 한다.

4 Takeaway

의미

Part 1이 "체크포인트 파일명의 단조 증가"라는 시스템적 정합성 문제를 해결했다면, Part 2는 "학습 신호 자체의 일관성"(LR 스케줄)을 다룬다. 두 수정 모두 근본 원인(optimizer/scheduler state가 checkpoint에 저장되지 않음)을 없애지 못하고 증상을 완화하는 패치라는 점은 동일하다. extra QOS로 preemption이 상시 발생하는 이 프로젝트 환경에서는, resume 경로에 숨어 있는 상태 불일치를 하나씩 찾아 고치는 작업이 반복적으로 필요하다는 것을 재확인했다.

5 Next Steps

미해결 한계

optimizer moment(Adam m/v) 자체는 여전히 checkpoint에 저장되지 않는다 — resume_warmup_steps=50이 실제로 충분한 완충인지 grad_norm curve로 검증되지 않았다.

다음 실험

이 패치가 적용된 상태로 다음 preemption 사이클을 관찰해 wandb의 lr 곡선에서 재웜업이 사라졌는지 (또는 50-step으로 줄었는지) 직접 확인. 장기적으로는 Part 1 카드에서도 동일하게 제안된 대로, _load_training_state를 활성화해 optimizer state까지 checkpoint에 포함하는 근본 수정을 우선순위로 검토한다.

Resume 시 LR 스케줄 재웜업 버그 수정 — Part 2