止めて再開しても一度に行ったものと同じに — 再開ファイルに何を入れるか

목표

60걸음 학습을 한 번에 한 결과와, 30걸음에서 재개 파일(가중치·옵티마이저 상태·걸음 수·배치 난수 상태)을 원자적으로 저장했다 이어 한 결과가 한 비트도 다르지 않게 만든다. 옵티마이저 상태나 자료 위치를 빠뜨리면 얼마나 벌어지는지, 반정밀도 저장이 무엇을 바꾸는지 잰다.

왜 중요한가

가중치만 저장했다가 이어 학습하면 오류 없이 결과가 달라진다. AdamW 는 파라미터마다 모멘트 두 개를 기억하고 있어, 그 기억을 잃으면 이은 직후 몇 걸음이 전혀 다른 크기로 움직인다. 에폭 중간에서 배치 순서를 되감으면 앞부분을 두 번 보고 뒷부분은 못 본다. 이런 차이는 손실 곡선에 거의 드러나지 않아 모르고 지나가기 쉽다. MiniMind 의 lm_checkpoint 는 그래서 가중치와 재개 파일을 따로 두고, 둘 다 임시 파일에 쓴 뒤 이름을 바꾼다. 이 실습은 그 설계가 정확히 무엇을 지키는지 한 번에 한 결과와 비교해 확인한다.

단계

1. /root/mm/ckpt/run.py 를 쓰고 60걸음(시드 7·배치 8·길이 128·lr 3e-3)을 한 번에 학습해 /root/mm/ckpt/straight.pth 와 /root/mm/ckpt/straight_log.csv(step,train_loss)를 남기세요.
2. 같은 학습을 30걸음에서 멈추며 재개 파일 /root/mm/ckpt/resume.pth 를 .tmp 에 쓴 뒤 os.replace 로 바꿔치기해 저장하세요(model·optimizer·step·gen_state).
3. 재개 파일에서 31걸음부터 60걸음까지 이어 /root/mm/ckpt/resumed.pth 와 /root/mm/ckpt/resumed_log.csv 를 남기세요. 1단계와 같아야 합니다.
4. 옵티마이저 상태만 버리고 이어 /root/mm/ckpt/noopt.pth 를 남기세요.
5. 배치 난수 상태만 버리고(처음 시드로 다시) 이어 /root/mm/ckpt/noskip.pth 를 남기세요.
6. straight.pth 를 반정밀도로 /root/mm/ckpt/straight_fp16.pth 에 저장하고, 두 파일 크기와 로짓 차이를 /root/mm/ckpt/fp16.json 에 적으세요.
7. 재개 파일에서 가중치와 옵티마이저 상태의 원소 수를 세어 /root/mm/ckpt/optim.json 에 적으세요.
8. /root/mm/ckpt/report.md 에 ## 무엇을 저장하나 ## 빠뜨리면 ## 반정밀도 세 절을 쓰고, 7단계의 배율과 6단계의 fp16 파일 크기(바이트)를 넣으세요.

참고

• 60걸음은 노드에서 5초 안팎입니다. 한 스크립트에 --stop-at·--resume·--no-optim·--no-skip 을 두면 3–5단계가 한 줄씩입니다.
• 배치를 torch.randint(..., generator=g) 로 뽑으면 g.get_state()·g.set_state() 로 '다음에 뽑을 배치' 를 저장·복원할 수 있습니다. MiniMind 의 SkipBatchSampler 가 하는 일을 이 한 줄이 합니다.
• 재개 파일은 torch.load(경로, weights_only=True) 로 읽으세요. 남이 준 체크포인트를 그냥 읽으면 그 안의 파이썬 코드가 실행될 수 있습니다.
• 흔한 실수: 모델을 만들기 전에 시드를 고정하지 않아 재개한 쪽의 초기화가 달라지는 것(곧 덮어쓰니 괜찮아 보이지만 생성기 상태가 어긋납니다), 학습률 일정을 이은 걸음 기준으로 다시 0 부터 세는 것.
• 원문: trainer_utils.py — lm_checkpoint·SkipBatchSampler · PyTorch — Saving and Loading · torch.load weights_only · os.replace

한 번에 60걸음 — 기준

/root/mm/ckpt/run.py 를 쓰고 시드 7·배치 8·길이 128·lr 3e-3(MiniMind 코사인, 전체 60걸음 기준)으로 60걸음을 한 번에 학습해 /root/mm/ckpt/straight.pth(state_dict)와 /root/mm/ckpt/straight_log.csv(step,train_loss, 1–60)를 남기세요.

앞 모듈의 학습 루프와 같습니다. 이번에는 손실을 소수점 여섯째 자리까지 적어 두세요 — 3단계에서 이은 기록과 한 자리씩 견줍니다.

30걸음에서 원자적으로 저장

같은 학습을 30걸음에서 멈추고 {"model": state_dict, "optimizer": optimizer.state_dict(), "step": 30, "gen_state": 배치 생성기.get_state()} 를 /root/mm/ckpt/resume.pth.tmp 에 쓴 뒤 os.replace 로 /root/mm/ckpt/resume.pth 로 바꾸세요. 학습률은 여전히 전체 60걸음 기준으로 계산합니다.

임시 파일에 다 쓴 다음 이름을 바꾸면, 저장 중에 죽어도 온전한 파일 하나가 남습니다. MiniMind 의 lm_checkpoint 가 ckp_tmp → os.replace 로 하는 일입니다. .tmp 가 남아 있으면 이름 바꾸기를 빼먹은 것입니다.

이어서 60걸음 — 기준과 같아야

새 프로세스에서 resume.pth 를 읽어 가중치·옵티마이저 상태·배치 생성기 상태를 되살리고 31걸음부터 60걸음까지 학습해 /root/mm/ckpt/resumed.pth 와 /root/mm/ckpt/resumed_log.csv(31–60)를 남기세요.

모델과 옵티마이저를 1단계와 똑같이 만든 뒤 load_state_dict 로 덮고, g.set_state(…) 로 생성기를 되돌립니다. 채점기는 두 체크포인트의 모든 텐서를 견줘 최대 차이가 1e-6 이하인지 봅니다 — 같은 노드에서 돌리면 0 입니다.

옵티마이저 상태를 버리면

재개 파일에서 가중치와 배치 생성기 상태만 되살리고 옵티마이저는 새로 만든 채 31–60걸음을 이어 /root/mm/ckpt/noopt.pth 를 남기세요.

AdamW 의 걸음 크기는 기울기를 그 이동 평균의 제곱근으로 나눈 값에 비례합니다. 기억이 0 에서 다시 시작하면 이은 직후 몇 걸음이 전혀 다른 크기로 움직입니다. 채점기는 이 결과가 기준과 확실히 다른지 봅니다.

배치 순서를 되감으면

재개 파일에서 가중치와 옵티마이저 상태만 되살리고 배치 생성기는 처음 시드(7)로 다시 만든 채 31–60걸음을 이어 /root/mm/ckpt/noskip.pth 를 남기세요.

생성기를 처음 상태로 두면 31걸음에 1걸음의 배치를 다시 봅니다. 에폭 중간에서 재개할 때 MiniMind 가 SkipBatchSampler 로 이미 본 배치를 건너뛰는 이유입니다.

반정밀도로 저장하면

straight.pth 의 모든 텐서를 .half() 로 바꿔 /root/mm/ckpt/straight_fp16.pth 에 저장하고, 두 파일의 크기(바이트)와 /opt/mm/ref/val.npy 앞 128토큰에 대한 로짓의 최대 절대 차이를 /root/mm/ckpt/fp16.json 에 fp32_bytes·fp16_bytes·max_logit_diff 로 적으세요.

mmkit.load_model 은 fp16 가중치도 fp32 로 펴서 넣습니다. 파일은 절반 가까이 줄지만(텐서 말고도 이름 같은 것이 들어 있어 정확히 절반은 아닙니다) 로짓은 소수점 셋째 자리쯤에서 달라집니다.

재개 파일이 큰 이유

resume.pth 에서 가중치의 원소 수(lm_head.weight 는 임베딩과 공유라 빼고)와 옵티마이저 상태 텐서(차원이 1 이상인 것)의 원소 수, 그 배율을 /root/mm/ckpt/optim.json 에 model_numel·optimizer_numel·ratio 로 적으세요.

optimizer.state_dict()['state'] 는 파라미터마다 step(0차원)·exp_avg·exp_avg_sq 를 담습니다. 배율이 2 면 재개 파일은 가중치 파일의 약 세 배입니다 — 큰 모델의 체크포인트가 무거운 이유입니다.

무엇을 왜 저장하는지 남기기

/root/mm/ckpt/report.md 에 ## 무엇을 저장하나 ## 빠뜨리면 ## 반정밀도 세 절을 쓰고, 7단계의 ratio 와 6단계의 fp16_bytes 를 숫자로 넣으세요.

4·5단계에서 벌어진 크기를 함께 적으면 좋습니다. 반정밀도 절에는 '추론용 가중치' 와 '재개용 가중치' 를 어떻게 나눌지 한 줄 적어 보세요.