再開ファイルには重みのほかに3つが入る

한국어 원문으로 표시합니다.

한 줄 요약

학습은 끊긴다 — 세션이 끝나고, 노드가 재시작하고, 비용 때문에 멈춘다. 멈춘 자리에서 한 번에 끝까지 간 것과 똑같이 이으려면 가중치만으로는 부족하다. MiniMind 의 lm_checkpoint 는 가중치와 따로 재개 파일을 두고, 거기에 옵티마이저 상태·걸음 수·(분산 학습이면) 프로세스 수를 함께 넣는다. 이 모듈에서는 60걸음을 한 번에 한 결과와 30걸음에서 멈췄다 이은 결과가 한 비트도 다르지 않게 만들고, 하나씩 빠뜨렸을 때 얼마나 벌어지는지 잰다.

왜 이게 필요했나

가중치만 저장했다가 이어 학습하면 결과가 조용히 달라진다. 오류는 나지 않는다. 손실 곡선도 그럴듯하게 이어진다. 그런데 두 가지가 어긋나 있다.

옵티마이저 상태. AdamW 는 파라미터마다 기울기의 이동 평균(exp_avg)과 제곱의 이동 평균(exp_avg_sq)을 들고 다니며, 걸음 크기를 그것으로 정한다. 새 옵티마이저로 이으면 이 기억이 0 에서 다시 시작해, 이은 직후 몇 걸음이 전혀 다른 크기로 움직인다. 이 상태는 파라미터의 두 배 크기라 재개 파일이 가중치 파일보다 세 배쯤 크다.
자료의 위치. 에폭 중간에 멈췄다면 그 에폭에서 이미 본 배치를 건너뛰어야 한다. 처음부터 다시 뽑으면 앞부분을 두 번 보고 뒷부분은 못 본다. MiniMind 는 SkipBatchSampler 로 이미 본 배치 수만큼 건너뛰고, 배치 순서는 seed + epoch 로 다시 만들어 같은 순서를 재현한다.

학습률 일정은 따로 저장하지 않아도 된다 — 걸음 번호로 계산하는 함수(get_lr)라 걸음 수만 있으면 된다. 분산 학습에서 GPU 수가 바뀌면 한 걸음이 보는 자료의 양이 달라지므로, MiniMind 는 저장한 world_size 와 지금의 수로 걸음 수를 환산한다.

어떻게 동작하나

MiniMind 의 저장은 두 파일을 쓴다.

state_dict = {k: v.half().cpu() for k, v in raw_model.state_dict().items()}
torch.save(state_dict, ckp_path + ".tmp");  os.replace(ckp_path + ".tmp", ckp_path)
resume = {"model": state_dict, "optimizer": optimizer.state_dict(),
          "epoch": epoch, "step": step, "world_size": ..., "wandb_id": ...}
torch.save(resume, resume_path + ".tmp");  os.replace(resume_path + ".tmp", resume_path)

임시 파일 → 이름 바꾸기. 저장하는 도중에 프로세스가 죽으면 반쯤 쓴 파일이 남는다. 그 파일이 원래 이름을 덮어쓰고 있었다면 멀쩡하던 이전 체크포인트까지 잃는다. .tmp 에 다 쓴 뒤 os.replace 로 한 번에 바꾸면 같은 파일 시스템 안에서는 원자적으로 바뀌어, 언제 죽어도 온전한 파일 하나가 남는다.

반정밀도. 가중치는 .half() 로 저장해 파일을 절반으로 줄인다. 불러서 fp32 로 펴면 로짓이 조금 달라진다(이 코스의 모델에서 최대 0.003 안팎). 추론에는 문제가 없지만, 재개 파일의 가중치까지 반정밀도라면 이어 한 학습이 한 번에 한 것과 비트 단위로 같을 수는 없다. 이 실습은 비교를 위해 재개 파일에 fp32 가중치를 넣는다.

안전하게 읽기. torch.load 는 파이썬 객체를 되살리는 pickle 을 쓴다. 남이 준 체크포인트를 그냥 읽으면 그 안의 코드가 실행될 수 있다. weights_only=True 로 읽으면 텐서·수·문자열·dict 만 되살린다 — 이 실습의 재개 파일(난수 생성기 상태 포함)도 그대로 읽힌다.

현장에서 만나는 모습

선점형(spot) 인스턴스로 학습하면 몇 시간마다 노드를 빼앗긴다. 재개가 정확하지 않으면 매번 조금씩 다른 학습이 되어, 최종 모델을 다시 만들 수 없고 문제가 생겨도 원인을 좁힐 수 없다. 그래서 '이어 한 것 = 한 번에 한 것' 을 짧은 학습으로 먼저 확인해 둔다. 저장이 원자적이지 않으면 반대 방향의 사고가 난다 — 저장 도중 노드가 죽어 가장 최근 체크포인트가 깨지고, 그 전 것은 이미 덮어써져 없다.

MiniMind 원본과 이 코스가 다른 점

MiniMind 는 걸음마다 새로 뽑는 대신 에폭 단위의 순서를 쓴다. 에폭마다 setup_seed(seed + epoch) 로 torch.randperm 순서를 만들고, 재개하면 같은 순서를 다시 만든 뒤 SkipBatchSampler 로 이미 본 배치 수만큼 건너뛴다. 이 코스의 루프는 걸음마다 임의 위치를 뽑으므로, 같은 일을 배치 난수 생성기의 상태(get_state·set_state)를 저장·복원하는 것으로 한다. 방법은 달라도 지키는 것은 같다 — '다음에 볼 자료' 가 멈춘 자리에서 이어진다. 또 MiniMind 는 가중치를 반정밀도로 저장하므로 이어 한 학습이 비트 단위로 같지는 않다. 이 코스는 그 차이까지 없애 확인하려고 재개 파일에 fp32 가중치를 넣는다.

다음 실습에서 할 것

60걸음을 한 번에 학습한 결과를 기준으로 두고, 30걸음에서 멈춰 재개 파일을 원자적으로 저장한 뒤 이어 60걸음까지 가서 가중치가 같은지 본다. 옵티마이저 상태를 버리고 이은 것, 배치 순서를 되감은 것이 얼마나 벌어지는지 재고, 반정밀도 저장의 크기와 로짓 차이, 재개 파일에서 옵티마이저 상태가 차지하는 몫을 센다.