파이토치 분산 학습 완전 정리: 데이터 병렬화, 그래디언트 누적, All-Reduce와 장애 복구

핵심 메시지

• 학습 규모를 키우는 방법은 더 좋은 장비로 바꾸는 수직 확장과 여러 대를 연결하는 수평 확장으로 나뉘며, 수평 확장은 코드 수정이 필요하지만 그 양은 매우 적다.
• 모델이 GPU 한 장에 올라간다면 데이터를 쪼개 나눠 학습하는 데이터 병렬화를, 모델 자체가 안 올라간다면 계층을 쪼개는 모델 병렬화를 쓴다.
• 그래디언트 누적은 매 데이터마다 파라미터를 갱신하는 대신 여러 스텝의 기울기를 더해 두었다가 한 번에 갱신하는 기법으로, 배치 크기를 키우지 않고도 큰 배치의 효과를 낸다.
• 노드들은 기울기를 한 노드로 모아 합산하고 다시 모두에게 돌려받는 All-Reduce 연산으로 동기화하며, 이 방식은 분할 정복 구조라 통신 시간이 노드 수에 로그로 비례한다.
• 노드가 죽는 상황에 대비해 공유 저장소에 주기적으로 체크포인트를 저장하되, 서로 덮어쓰지 않도록 글로벌 랭크 0인 노드만 저장을 담당하게 만든다.

쉽게 이해하기

영상은 학습이 느려지는 세 가지 상황을 출발점으로 삼는다. 파라미터가 너무 많아 모델이 GPU 메모리에 올라가지 않거나, 배치 크기를 조금만 키워도 메모리 부족 오류가 나거나, 데이터가 너무 많아 학습에 하염없이 시간이 걸리는 경우다. 이 가운데 하나라도 해당한다면 학습 환경을 키워야 하는데, 선택지는 장비 자체를 더 좋은 것으로 바꾸는 수직 확장과 여러 대를 묶는 수평 확장 두 가지다. 수직 확장은 코드를 그대로 쓸 수 있는 대신 한계가 뚜렷하고, 수평 확장은 코드 수정이 필요하지만 파이토치의 분산 데이터 병렬 구현 덕분에 그 수정량이 작다는 점이 강조된다.

분산 방식은 크게 둘로 나뉜다. 모델이 GPU 한 장에 들어간다면 데이터를 겹치지 않는 부분집합으로 쪼개 각 GPU가 서로 다른 데이터를 학습하고 역전파 시점에 기울기를 동기화하는 데이터 병렬화를 쓴다. 반대로 모델 자체가 한 장에 안 들어가면 모델을 계층 단위로 쪼개 각 GPU가 순전파와 역전파의 일부만 담당하는 모델 병렬화를 쓰고, 두 방식을 섞은 하이브리드도 가능하다. 영상이 다루는 것은 모델은 올라가지만 데이터가 너무 많은 경우, 즉 데이터 병렬화다.

본격적인 분산 설명에 앞서 그래디언트 누적 개념이 먼저 나온다. 일반적인 학습 루프는 데이터 하나마다 손실을 구하고 역전파한 뒤 곧바로 파라미터를 갱신하고 기울기를 0으로 되돌리는데, 갱신과 초기화를 매번 하지 않으면 파이토치는 새로 계산한 기울기를 이전 기울기에 덮어쓰지 않고 더해서 쌓아 둔다. 이렇게 몇 스텝치를 누적한 뒤 한 번에 갱신하면 배치 크기를 물리적으로 키우지 않고도 큰 배치로 학습한 것과 비슷한, 더 부드러운 학습 곡선을 얻을 수 있다.

분산 데이터 병렬 학습의 흐름은 네 단계로 정리된다. 먼저 한 노드가 초기 가중치를 만들어 브로드캐스트로 모두에게 뿌려 같은 지점에서 출발하게 하고, 각 노드는 서로 다른 데이터 부분집합으로 순전파와 역전파를 돌려 자기만의 지역 기울기를 만든다. 그 기울기들을 한 노드로 모아 합산하는 리듀스와, 그 합을 다시 모두에게 돌려주는 브로드캐스트를 거치면 모든 노드가 동일한 합계를 갖게 된다. 마지막으로 각 노드가 같은 합계로 자기 모델의 파라미터를 갱신하므로, 갱신이 끝난 뒤에도 모든 노드의 가중치는 여전히 똑같이 유지된다.

여기서 통신 방식이 성능을 좌우한다. 5메가바이트 파일을 초당 1메가바이트 속도로 친구 일곱 명에게 보낸다고 하면, 한 명씩 차례로 보내도 35초, 일곱 명에게 동시에 보내도 대역폭이 나뉘어 역시 35초가 걸린다. 반면 파일을 받은 사람이 다시 다른 사람에게 전달하는 분할 정복 방식을 쓰면 받은 노드 수가 매 단계 두 배로 늘어 15초로 줄어든다. 실제 브로드캐스트와 리듀스가 이런 구조로 동작하며, 리듀스와 브로드캐스트를 이어 붙인 All-Reduce는 두 연산을 따로 하는 것보다 빠르게 하나의 연산으로 구현되어 있다.

주요 인사이트

• 분산 학습에서 실제로 오가는 것은 데이터가 아니라 기울기다. 각 노드는 처음에 같은 가중치를 받고 같은 합계로 갱신하기 때문에, 학습이 진행되는 내내 모든 노드의 모델은 완전히 동일한 상태를 유지한다.
• 통신을 분할 정복으로 설계하면 노드 수가 늘어도 통신 시간이 로그 스케일로만 증가한다. 노드를 늘릴 때 통신 비용이 선형으로 폭발하지 않는 이유가 여기에 있다.
• 체크포인트를 누가 저장할지는 반드시 코드로 정해야 한다. 모든 노드가 같은 코드를 돌리므로 조건을 걸지 않으면 서로의 체크포인트를 덮어쓴다. 영상은 글로벌 랭크가 0인 노드만 저장하게 하되, 재시작 시 어느 노드가 랭크 0이 될지는 가정하지 말라고 조언한다.
• 로컬 랭크와 글로벌 랭크의 쓰임새가 다르다. 진행 표시줄 출력처럼 같은 컴퓨터 안에서만 충돌하는 작업은 로컬 랭크로, 실험 추적 서비스 초기화처럼 클러스터 전체에서 한 번만 해야 하는 작업은 글로벌 랭크로 제어한다.
• 파이토치는 역전파가 끝나기를 기다리지 않고 계산이 끝난 기울기부터 먼저 전송한다. 게다가 하나씩 보내지 않고 약 25메가바이트 단위의 버킷으로 묶어 보내는데, 너무 작으면 통신 오버헤드가 커지고 너무 크면 계산과 통신을 겹칠 여지가 줄어들기 때문이다.

자주 묻는 질문

관련 AI 소식