bf16 파인튜닝, 95% 가중치 동결로 옵티마이저 불안정성 가려

• bf16 가중치와 torch.optim.AdamW를 fp32 마스터 복사본 없이 사용하면 큰 크기의 가중치 업데이트가 동결되어 사실상 스파스 파인튜너로 작동함
• Mistral-7B에서 Quiet-STaR 레시피를 재실행했을 때, bf16에서 제안된 업데이트의 97~99.7%가 0으로 반올림됨. 반면 동일한 학습률의 올바른 fp32 실행은 발산하며, 손실이 안정적인 bf16 실행보다 훨씬 높게 피크를 기록함
• bf16 반올림 오류는 모델을 사전학습 가중치에 고정시켜 불안정성을 가림. 더 낮은 학습률(3e-7)의 올바른 fp32 옵티마이저는 더 높은 학습률(1e-6)의 bf16 성능과 일치함
• 이 현상은 PaFi와 NanoAdam과 같은 스파스 업데이트 방식이 잘 작동하는 이유를 설명함. bf16 버그가 우연히 크기에 기반한 스파시티를 구현하기 때문임

메커니즘 및 영향

• bf16은 7비트 맨티사를 가지며, 이는 fp32보다 65,536배 넓은 간격을 생성함. 간격의 절반보다 작은 업데이트는 원래 값으로 반올림됨
• 학습률 η에서 약 512·η보다 작은 크기의 가중치만 업데이트 가능. 학습률 1e-6일 경우 약 4.88e-4를 초과하는 가중치는 동결됨
• AdamW의 가중치 감쇠와 2차 모멘트(v)도 상대적 변화가 반올림되어 버려지므로 bf16에서 동결됨

진단 및 수정

• 진단: 전체 학습률로 단일 스텝을 실행. 가중치의 10% 미만만 변경되면 옵티마이저가 스파스 파인튜너로 작동 중임
• 권장 수정: fp32 마스터 가중치 사용. 대안: torchao의 bf16_stochastic_round=True와 같은 확률적 반올림을 사용하여 업데이트를 시간 평균화
• 재현 조언: DeepSpeed나 FSDP와 같은 올바른 fp32 스택으로 bf16 레시피를 이동할 때 학습률을 3~10으로 나누기