[Python] Mixed Precision & NaN Loss 문제 해결: 행렬 연산의 수치적 안정성

배경

특정 딥러닝 학습 스크립트가 의도치않게 종료되었습니다.

Epoch 4에서 Out of Memory로 인한 예상치 못한 종료가 발생하였는데, 모델의 복잡도를 줄일 수는 없는 상황이라, 두 가지 최적화 방법을 적용하였습니다:

1. Batch size를 25%로 낮춤
2. Mixed Precision 도입
   • `torch.nn.BCELoss`를 `torch.nn.functional.binary_cross_entropy_with_logits`로 교체하여 `autocast` 호환성 확보

문제상황

코드 수정 후 loss가 NaN으로 계산되는 새로운 문제가 발생했습니다.

추론

1. Batch size 감소만으로는 문제가 발생하지 않았음을 확인
   • 이 요소만 조작한 뒤 학습하면, 상식적인 loss 값이 나오므로 배제
2. Mixed Precision 적용이 NaN loss의 원인으로 추정
   
   • FP16으로 표현 불가능한 값이 연산 중 발생하는 것으로 의심 (overflow or underflow)

(참고) Mixed Precision

Mixed Precision은 기존의 FP32에 FP16를 섞어 정밀도는 조금 떨어지지만, 메모리 사용 감소 및 연산 가속하는 방법이다.

FP32 (Single Precision)

• 비트 수: 32비트 (4바이트)
• 유효 자릿수(정확도): 약 7자리 10진수
• 최대 표현: 약 $\pm 3.4×10^{38}$
• 최소 표현: 약 $\pm 1.18×10^{−38}$

FP16 (Half Precision)

• 비트 수: 16비트 (2바이트)
• 유효 자릿수(정확도): 약 3~4자리 10진수
• 최대 표현: 약 $\pm 6.55 \times 10^4$
• 최소 표현: 약 $\pm 6.1×10^{−5}$

조사

`torch.amp`  문서를 통해 모델의 호환성 문제를 확인했습니다. 대부분의 PyTorch 모델은 호환되지만, 제가 돌린 최신 논문의 코드에서 직접 작성된 부분이 문제가 되는 것 같습니다.

현재 모델은 encoder - decoder - multilayer (Batch normalization, 8 layers) 구조를 가지며, encoder에서 수치적으로 매우 민감한 `(역)행렬 연산`을 포함합니다. Encoder 구조 내 의심스러운 행렬 연산 사이사이의 결과물을 출력하도록 코드를 수정했습니다.

DTD_inv_DT (FP16): tensor([[[nan, nan, nan,  ..., nan, nan, nan],
         [nan, nan, nan,  ..., nan, nan, nan],
         [nan, nan, nan,  ..., nan, nan, nan],
         ...,
         [nan, nan, nan,  ..., nan, nan, nan],
         [nan, nan, nan,  ..., nan, nan, nan],
         [nan, nan, nan,  ..., nan, nan, nan]]], device='cuda:3',
       dtype=torch.float16, grad_fn=<UnsafeViewBackward0>)

실제로 행렬 곱 연산 후 NaN 값이 발생함을 확인했습니다.

해결

해당 연산에 대하여만 FP32를 사용하도록 설정했습니다.

with torch.cuda.amp.autocast(enabled=False):
        DTD_inv = torch.inverse(DTD + self.lambda3 * torch.eye(self.input_dim).to(device).float())
        DTD_inv_DT = torch.matmul(DTD_inv, Z_f.float())

코드를 돌려보니 정상적인 loss 값이 나왔습니다.

 

만세!

정리

Mixed precision 활용시 overflow or underflow 문제로 인해 수치적인 안정성이 중요하다는것은 알고 있었지만, 직접 겪은건 처음입니다. 아무래도 결과값이 극도로 팽창되거나 수축되는 연산인 행렬연산이 의심스러웠고 해당 부분에 대하여 FP32를 사용하도록 처리하니 문제가 해결되었습니다. 

 

행렬 연산 말고도 mixed precision 사용시 문제가 되셨던 연산이 있다면 댓글로 알려주세요~!

참조

https://pytorch.org/docs/stable/amp.html#prefer-binary-cross-entropy-with-logits-over-binary-cross-entropy

Automatic Mixed Precision package - torch.amp — PyTorch 2.3 documentation