파이토치 학습 속도 높이는 법: 혼합 정밀도와 FSDP 다중 GPU 분산 학습을 최소 코드 수정으로 적용하기

핵심 메시지

• 처음부터 학습하는 대신 사전학습된 비전 트랜스포머를 미세조정하자 같은 데이터셋에서 정확도가 62%에서 95%로 올랐고 시간도 절반 이하로 줄었다.
• 패브릭은 파이토치 코드 구조를 바꾸지 않고 고급 기능을 쓰게 해 주는 얇은 래퍼로, 필요한 수정은 여섯 군데 정도에 그친다.
• 혼합 정밀도 학습은 인자 하나만 추가하면 되고, 예측 성능을 유지한 채 학습 시간을 18분에서 6분으로 줄였다.
• bfloat16은 float16보다 정밀도는 조금 낮지만 표현 범위가 넓어, 대규모 언어모델 학습에서 손실에 NaN이 뜨는 문제에 더 안정적이었다.
• 완전 분할 데이터 병렬(FSDP)은 데이터 병렬과 텐서 병렬을 함께 쓰는 방식으로, GPU 4장에 적용하자 학습 시간이 2분까지 내려갔다.

쉽게 이해하기

출발점은 밑바닥부터 학습한 비전 트랜스포머다. 약 60분을 들여 테스트 정확도 62%를 얻었고, 발표자는 모델과 데이터셋의 세부 사항은 중요하지 않다고 전제하며 이 수치를 기준선으로 삼는다. 첫 번째 교훈은 기법이 아니라 태도에 가깝다. 요즘은 밑바닥부터 학습하는 것이 대개 합리적이지 않다는 것인데, 실제로 사전학습된 비전 트랜스포머를 세 에폭 미세조정하자 약 20분 만에 정확도 95%에 도달했다.

그다음은 코드를 얼마나 고쳐야 하는가의 문제다. 발표자가 쓰는 도구는 패브릭으로, 파이토치 라이트닝을 만든 쪽에서 내놓은 더 가벼운 형태의 래퍼다. 무료 오픈소스이며, 라이트닝 라이브러리를 설치한 뒤 패브릭을 가져다 쓰면 된다. 바뀌는 부분은 패브릭을 임포트하고, 객체를 만들고, 모델·옵티마이저와 데이터로더를 설정에 통과시키고, 손실의 역전파를 패브릭 쪽 함수로 호출하는 정도다.

이 상태에서 성능을 재면 순수 파이토치와 사실상 같다. 패브릭 자체는 래퍼일 뿐이므로 당연한 결과다. 차이는 고급 기능을 켤 때 나타난다. 혼합 정밀도 학습 인자를 넣자 예측 성능은 유지한 채 학습 시간이 18분에서 6분으로 줄었다.

혼합 정밀도가 통하는 이유는 32비트 가중치를 16비트로 바꿔 기울기를 계산한 뒤 다시 32비트로 되돌려 파라미터를 갱신하기 때문이다. 16비트 연산이 훨씬 빠르고 메모리도 아끼면서 정확도 손실은 피한다. 발표자는 GPU가 지원한다면 bfloat16 쪽을 권하는데, 정밀도는 조금 낮아도 float32와 같은 동적 범위를 가져 더 안정적이었기 때문이다.

마지막 단계는 다중 GPU다. 데이터 병렬은 미니배치를 나눠 각 GPU에 모델 사본을 두므로 메모리를 아끼지 못하는 반면, 텐서 병렬은 행렬 곱을 행이나 열 단위로 쪼개 여러 장치에 나눠 계산한 뒤 결과를 이어 붙인다. 덕분에 한 장에 올라가지 않는 모델도 학습할 수 있다. FSDP는 둘을 함께 쓰는 방식이며, 전략 이름과 장치 수만 지정하면 켜진다.

주요 인사이트

• 가장 큰 성능 개선은 정교한 병렬화가 아니라 '밑바닥 학습을 그만두는 것'에서 나왔다. 60분에 62%였던 결과가 사전학습 모델 미세조정으로 20분에 95%가 됐다.
• 혼합 정밀도는 투입 대비 효과가 가장 좋은 선택지다. 사실상 인자 하나를 추가하는 수준의 변경으로 학습 시간이 3분의 1로 줄었다.
• float16과 bfloat16의 차이는 정밀도 대 범위의 맞바꿈이다. 순수 float16 학습에서 손실에 NaN이 생기는 일이 있는데, bfloat16은 그 점에서 더 견고했다는 것이 발표자의 경험이다.
• FSDP를 고른 이유가 성능만은 아니다. 딥스피드나 콜로설 AI 같은 대안도 있지만 FSDP는 이미 파이토치에 들어 있어 별도 라이브러리가 필요 없다는 점이 실용적인 장점으로 꼽혔다.
• 분산 학습을 직접 짜면 코드 변경이 상당히 많아진다는 점이 이 도구들의 존재 이유다. 성능 개선폭보다도 '기존 코드를 재구성하지 않아도 된다'는 점이 핵심 주장이다.

자주 묻는 질문

관련 AI 소식