혼합 정밀도 학습: autocast, 손실 스케일링, 수치적 한계

정의

혼합 정밀도 학습은 한 번의 학습에서 여러 부동소수점 형식을 함께 사용한다. 비용이 큰 행렬 연산은 FP16이나 BF16을 쓰고 민감한 연산은 FP32를 유지할 수 있다. 자동 혼합 정밀도는 연산과 장치 규칙에 따라 형식을 고르며, 모든 텐서를 반정밀도로 변환하는 것과는 다르다. 효과는 모델, 하드웨어, 작업 부하에 달려 있다. [1][2]

작동 원리

Autocast는 연산 정밀도를 선택한다. 손실 스케일링은 매우 작은 FP16 기울기가 0이 되는 별도의 문제를 다룬다. 손실에 스케일을 곱하면 역전파 기울기가 커지고, 옵티마이저 업데이트 전에 스케일을 복원하면 원래 의도한 크기를 유지한다. 동적 스케일러는 유한하지 않은 기울기를 감지해 해당 업데이트를 건너뛰고 스케일을 조정한다. 하지만 모든 수치적으로 불안정한 모델을 정상화할 수는 없다. [1][2]

실무 절차

고정된 데이터와 평가 규칙으로 FP32 기준을 만든다. 일반적인 PyTorch eager 학습에서는 순전파와 손실 계산을 torch.autocast로 감싸고, 역전파는 그 문맥 밖에서 실행한다. FP16에서는 필요에 따라 torch.amp.GradScaler를 사용한다. Autocast를 켜기 위해 모델 전체를 수동 변환하지 않는다. [1]

기울기를 클리핑한다면 먼저 스케일을 복원한다. 문서화된 순서대로 스케일러의 step과 update를 호출하고 다음 업데이트를 위해 기울기를 지운다. 모델 및 옵티마이저 체크포인트와 함께 스케일러 상태도 저장한다. [2] 누적 시에는 유효 배치 전체에 같은 스케일을 유지한다. 컴파일 실행과 사용자 정의 연산은 버전에 맞는 정밀도 검증이 필요하다.

설명용 예시

합성곱이 실행 시간 대부분을 차지하지만 최종 손실은 더 넓은 수치 범위가 필요한 분류기를 생각해 보자. Autocast는 대상 합성곱을 낮은 정밀도로 실행하고, 지원되는 민감한 손실은 FP32로 계산할 수 있다. 한 스텝에서 무한대 기울기가 나오면 스케일러가 업데이트를 건너뛸 수 있다. 이는 진단 사건이며 학습 수렴의 증거가 아니다. 이 예시는 동작을 설명할 뿐 측정된 속도 향상을 제시하지 않는다.

한계와 검증

FP16과 BF16은 수치 범위가 다르다. BF16을 위해 학습한 모델은 최대 유한 값이 65504인 FP16에서 오버플로가 생길 수 있다. 손실 스케일을 높여도 순전파의 오버플로를 고치지 못한다. BF16은 흔히 FP16식 기울기 스케일링을 필요로 하지 않지만 여전히 반올림을 유발하므로 검증해야 한다. [1]

검증 지표와 손실 추이를 FP32 기준과 비교하고, 유한하지 않은 값과 건너뛴 업데이트를 점검하며, 워밍업 후 최대 메모리와 처리량을 측정한다. 작은 모델이나 데이터 입력이 병목인 모델은 이득이 작을 수 있다. [2] 정밀도, 하드웨어, 소프트웨어 버전과 관찰된 차이를 기록한다. 이 글은 로컬 실험을 보고하지 않는다.

출처

[1] PyTorch: Automatic Mixed Precision package [2] PyTorch: Automatic Mixed Precision recipe