FSDP 분산 학습 입문, 파라미터 샤딩과 all-gather·reduce-scatter 그리고 Ray Train

핵심 메시지

• 레이어를 GPU마다 나눠 얹는 소박한 방식은 한 번에 한 GPU만 일하고 나머지는 대기하기 때문에 좋은 분할이 아니다.
• FSDP는 레이어를 논리적 단위로 묶는 수직 분할과, 단위 안의 파라미터·그래디언트·옵티마이저 상태를 GPU 수만큼 쪼개는 수평 샤딩을 함께 쓴다.
• 메모리를 아끼는 대가로 순전파 전 all-gather, 역전파 전 all-gather, 그래디언트 reduce-scatter라는 세 번의 통신 비용을 치른다.
• FSDP2 적용은 참여 GPU를 정의하는 디바이스 메시 생성, 블록별 샤딩, 루트 모델 샤딩의 세 단계로 이뤄지고 CPU 오프로딩과 혼합 정밀도를 선택할 수 있다.
• 학습 로직은 PyTorch가, 여러 머신에 걸친 스케줄링과 내결함성·모니터링은 Ray가 맡는 역할 분담이 이 웨비나의 핵심 메시지다.

쉽게 이해하기

이 웨비나는 모델이 GPU 한 장에 들어가지 않는다는 전제에서 출발한다. 가장 먼저 떠오르는 방법은 레이어를 네 개의 GPU에 나눠 얹는 것이지만, 그렇게 하면 첫 GPU가 자기 몫의 레이어를 다 통과하기 전까지 나머지 세 장은 그냥 기다린다. 역전파에서도 같은 문제가 되풀이되므로, 한 번에 한 장만 일하는 이 방식은 원하는 그림이 아니다.

FSDP는 데이터 병렬의 한 갈래여서 각 GPU가 서로 다른 마이크로배치를 받는다. 여기에 두 단계의 분할이 겹친다. 먼저 레이어들을 세 개씩 묶는 식으로 논리적 단위를 만드는 수직 분할이 있고, 그다음 각 단위 안의 파라미터와 그래디언트, 옵티마이저 상태를 GPU 수만큼 쪼개는 수평 샤딩이 있다. 한 단위가 통틀어 32GB를 차지한다면 네 장에 8GB씩 나뉘어 놓이는 식이다.

다만 이 상태로는 순전파를 할 수 없다. 연산을 하려면 그 단위의 파라미터가 전부 필요하기 때문에, 단위마다 all-gather로 흩어진 조각을 모아 잠시 전체 파라미터를 갖춘 뒤 자기 마이크로배치로 순전파를 하고, 끝나는 즉시 남에게서 빌려온 파라미터를 버린다. 계속 들고 있으면 애초에 쪼갠 의미가 없기 때문이며, 이것이 메모리를 아끼기 위해 치르는 대가다. 역전파는 같은 일을 거꾸로 하는데, 다시 all-gather로 파라미터를 모으되 마지막 단위는 순전파 때 갖고 있던 것을 유지할 수 있어 생략되기도 하고, 각 GPU가 서로 다른 배치로 계산한 그래디언트는 reduce-scatter로 평균을 낸 뒤 각자가 책임진 몫만 남긴다. 정리하면 순전파 전 all-gather, 역전파 전 all-gather, 그리고 reduce-scatter까지 세 번의 통신이 추가로 붙는 셈이다.

실습은 최신 버전인 FSDP2를 쓰는데, 해야 할 일은 세 가지다. 어떤 GPU가 학습에 참여하는지 정의하는 디바이스 메시를 만들고, 트랜스포머 블록마다 샤딩을 적용한 다음, 임베딩과 나머지 레이어를 포함한 루트 모델에도 한 번 더 적용한다. 여기에 선택 사항으로 자원이 빠듯할 때 파라미터와 그래디언트를 CPU로 내리는 오프로딩과, 파라미터 dtype과 리덕션 dtype을 각각 지정하는 혼합 정밀도가 붙는다. 예제 모델은 12개 블록짜리 비전 트랜스포머이고, 흑백 이미지를 다루기 위해 입력 채널만 1로 바꿔 쓴다.

Ray가 들어오는 자리는 명확하게 구분된다. 모델과 샤딩, 체크포인팅 같은 학습 로직은 계속 PyTorch의 몫이고, 그것을 수많은 머신에 어떻게 배치하고 장애에 어떻게 대응하며 어떻게 관찰할지가 Ray의 몫이다. 바닐라 PyTorch였다면 분산 환경 설정과 분산 샘플러, 여러 머신에서의 동시 실행과 내결함성을 손수 챙겨야 하지만, Ray Train에서는 모델과 데이터 로더를 준비 함수로 감싸고 트레이너로 묶어 실행하면 스케일링 설정에 맞춰 학습이 워커들에 분배되고 지정한 위치에 체크포인트가 저장된다. 재개할 때는 저장된 체크포인트가 있으면 이어받고 없으면 처음부터 시작하도록 분기하며, 분산 샘플러에 에폭을 알려주지 않으면 매 에폭 같은 순서의 데이터가 반복된다는 점도 코드로 확인시켜 준다. 제목에 함께 적힌 DeepSpeed는 별도의 노트북으로 제공하되 한 번에 다루기에는 분량이 많다며 다음 기회로 미뤘다.

주요 인사이트

• FSDP의 핵심은 '쪼개 두었다가 필요할 때만 잠깐 모은다'는 한 문장으로 요약된다. 파라미터를 계속 들고 있으면 샤딩의 의미가 사라지므로, 순전파가 끝나는 즉시 빌려온 조각을 버리는 동작이 설계의 전부라 해도 좋다.
• 메모리 절약은 공짜가 아니라 통신과의 맞바꿈이다. 단위마다 all-gather가 두 번, reduce-scatter가 한 번 붙는다는 사실을 알고 있어야 어디서 성능이 새는지 짐작할 수 있다.
• 레이어를 GPU에 나눠 얹는 직관적인 방법이 왜 나쁜지를 먼저 보여주고 시작하는 구성이 좋다. GPU를 여러 장 쓴다는 것과 여러 장이 동시에 일한다는 것은 전혀 다른 이야기다.
• 학습 로직은 PyTorch, 분산 실행은 Ray라는 경계가 분명하다. 프레임워크를 갈아타는 것이 아니라 기존 학습 함수를 거의 그대로 두고 감싸는 방식이라는 점이 진입 장벽을 낮춘다.
• 분산 샘플러에 에폭을 전달하지 않으면 매 에폭 동일한 데이터 순서가 반복된다는 지적처럼, 분산 학습의 실수는 오류가 아니라 조용한 성능 저하로 나타나는 경우가 많다.

자주 묻는 질문

관련 AI 소식