‘AI의 미래’를 이야기할 때, 회의적인 관점을 가진 사람들은 때로 ‘파운데이션 모델’을 트레이닝하는데 드는 막대한 자원, 그리고 ‘최적화 기술’의 현재 한계를 지적한다는 걸 아실 겁니다. 실제로 파운데이션 모델을 트레이닝하는데 드는 비용을 지속적으로 감당할 만한 회사나 조직이 많지는 않죠.
‘최적화 기술’의 한계와 관련해서는, FSDP (Fully Sharded Data Parellel)과 이것의 향상된 버전이라고 할 수 있는 YaFSDP (Yet another FSDP) 같은 기법이 유망한 해결책으로 이야기되고 있는데요. (이 용어들에 대해서는 잠시 후에 살펴볼게요)
FSDP 기법을 사용하면 모델을 여러 개의 GPU로 분할해서 메모리의 오버헤드를 줄이고 트레이닝 속도를 빠르게 할 수 있습니다. YaFSDP는 그 이름처럼 FSDP를 기반으로 훨씬 더 나은 효율성, 확장성을 제공한다고 해요. FSDP와 YaFSDP 메커니즘을 한 번 자세히 살펴보고, 이 솔루션이 과연 AI의 미래에 큰 도움이 될지 어떨지 생각해 보시죠!
오늘 에피소드에서는 아래의 내용을 다뤄보겠습니다:
(오늘 포스트는 평소보다도 조금 더 기술적인 배경지식이 있다면 좋은 내용이긴 하지만, 가벼운 마음으로 ‘개념적으로 그렇구나’ 하고 읽어보셔도 좋겠습니다!)
FSDP (Fully Sharded Data Parallel)란 게 뭔가요?
먼저 ‘데이터 병렬화 (Data Parallelism)’의 개념을 짚고 넘어갈까요? 간단히 이야기하면,딥러닝 모델이 발전하면서 모델의 규모가 커지고 - 모델의 계층 (Layer)가 더 깊어지고, 계산할 뉴런이 많아지죠 - 따라서 학습해야 하는 파라미터도 함께 엄청나게 증가하게 됩니다.특히 트레이닝 및 추론 단계에서 계산량이 증가해서, 점점 고성능의 메모리/CPU/GPU가 필요하게 되죠. 그래서, 딥러닝 모델의 트레이닝과 추론을 효율적으로 하도록 하는 테크닉들이 다양하게 고민되었는데, 그 중 하나가 학습할 데이터를 분산시키는 ‘데이터 병렬화 (Data Parallelism)’입니다. (이 외에도 ‘모델 병렬화’, ‘파이프라인 병렬화’ 등의 방법도 있는데, 여기서는 다루지 않겠습니다)
PyTorch에서 지원하는 병렬처리 패키지인 DDP (Distributed Data Parallel)에서 볼 수 있듯이, 기존의 데이터 병렬화 방식은 여러 개의 GPU에 전체 모델을 복제하고 Gradient를 동기화하는 방법으로 모델의 일관성을 보장합니다. 그런데 이런 방법은 개별 GPU의 메모리 용량이 제약 조건이 되기 때문에 점점 더 큰 모델을 트레이닝하는데는 어려운 점이 있습니다.
2023년 9월, Meta AI 연구원들이 개별 GPU의 메모리 용량을 초과하는 거대 신경망 모델을 트레이닝하기 어려운 문제를 해결하기 위한, PyTorch FSDP (Fully Sharded Data Parallel)에 대한 연구를 발표했습니다. 거대 신경망 모델이 여러가지 영역에서 뛰어난 성능을 보여주는 만큼, 그 기술적 장벽 때문에 소수의 회사와 조직만이 접근하게 되는 상황을 방지하려면 이런 기술이 아주 중요합니다.
데이터 샤딩 (Data Sharding)은 뭘까요? 이건 원래는 데이터베이스 영역에서 나온 테크닉인데, 전체 데이터셋을 여러 부분으로 나눠서 각각의 샤드 (Shard)를 서로 다른 DB 서버나 클러스터에서 독점적으로 관리하는 방법을 이야기합니다. 물론 이렇게 하면 확장성이라든가 성능 측면의 장점이 있기 때문에 생긴 방법이죠.
딥러닝에서의 데이터 샤딩은 데이터를 GPU에 담는 분산 시스템을 이야기하고, 그냥 데이터셋 뿐 아니라 모델 파라미터, 그레디언트, 최적화 상태 등 다양한 데이터가 분산 대상이 됩니다.
Meta AI 연구팀은 Tensor Implementation, Dispatcher System, CUDA Memory Caching Allocator 같은 파이토치의 핵심적 구성 요소와 밀접하게 통합되는 FSDP를 설계했습니다. 이렇게 함으로써 기존의 파이토치 사용자 경험에 잘 연계되면서도 트레이닝 효율성을 높일 수 있게 한 것이죠.
Image Credit: The original FSDP paper
FSDP는 모델 파라미터를 샤딩 (Sharding)해서 다양한 하드웨어 구성 조건에서도 자원을 최적으로 활용하게 합니다. 모델을 더미 디바이스 (Dummy Device)에서 생성한 다음 실제 GPU에서 개별적으로 샤딩하는 ‘초기화 지연 (Deferred Initialization)’ 같은 기술을 포함해서 다양한 최적화 기술을 포함하고 있습니다. 종합적인 평가를 보자면, FSDP가 파이토치의 DDP 대비 훨씬 더 큰 모델을 잘 다루면서도 성능은 비슷하게 유지하고, 거의 선형적인 확장성을 보여주는 장점이 있습니다. 물론, FSDP에도 몇 가지 중요한 한계는 있어요.
FSDP는 어떤 한계점이 있을까요?
FSDP의 대표적인 한계점들은 아래와 같습니다:
샤딩되지 않았던 값이나 특정한 텐서 구조에 의존하는 옵티마이저를 사용하는 경우에는 로컬 트레이닝의 결과가 항상 정확히 일치하지 않을 수도 있습니다.
공유 파라미터를 처리하는게 까다로울 수 있고, 잘 처리되지 않는 경우에 오류가 발생하거나 필요 이상의 메모리를 사용하게 될 수 있습니다.
최적화를 하더라도, 특히 더 큰 모델이나 더 많은 GPU를 사용하는 경우에는 여전히 많은 통신 오버헤드가 발생할 수 있으므로 신중한 조정이 필요합니다.
특히 대형 모델의 경우에 FSDP를 설정하는 것이 복잡할 수 있으며, 경우에 따라서는 다른 방법이 나을 수도 있습니다.
파이프라인이나 텐서 병렬 처리와 같은 다른 병렬 처리 방법과 FSDP를 혼합하는 것은 쉽지 않고, 지나친 오버헤드를 피하기 위해 설정을 신중하게 해야 합니다.
FSDP는 PyTorch와 긴밀하게 통합되어 있기 때문에 단순한 데이터 병렬 처리 방식에 비해 디버깅 및 문제 해결이 더 어려울 수도 있습니다.
또, YaFSDP의 개발자들은 ‘FSDP가 주는 여러가지 장점에도 불구하고 여전히 해결해야 할 문제가 있다’면서 아래와 같은 점들을 지적하기도 했습니다:
FSDP는 레이어에 메모리를 동적으로 할당하는데, 때로 실제 필요한 것보다 훨씬 더 많은 메모리를 필요로 합니다.
역전파 (Backpropagation; Backward Pass)를 하는 동안 ‘Give-way 효과'라는 현상이 나타날 수 있습니다. 밑의 그림을 보죠:
그림의 첫 번째 줄은 ‘Computation Stream’이고, 다른 줄들은 ‘Communication Stream’인데, 파란색으로 표시된 reduce_scatter 연산 전에 많은 준비 작업들이 발생하고 있음을 알 수 있습니다. 이 작은 계산들이 주된 Computation Stream과 동시에 병렬적으로 실행되기 때문에 통신 속도가 크게 느려지게 됩니다.
YaFSDP의 등장
이런 FSDP의 문제를 해결하겠다고 등장한 것이 YaFSDP라고 보시면 되겠습니다. YaFSDP (Yet another Fully Sharded Data Parallel)는 얀덱스 (Yandex)의 연구자들이 2024년 5월 오픈소스로 발표했는데요. (얀덱스는 러시아에서 60%의 시장 점유율을 차지하는 러시아 최대 검색엔진 기업입니다)
YaFSDP는 FSDP보다 더 효율적으로 메모리를 관리하고, 중복 계산을 줄이고, 거대 언어모델을 트레이닝하는 중의 통신 및 동기화를 최적화해 줍니다.
YaFSDP의 작동 방식을 알아봅시다.
레이어 샤딩 (Layer Sharding): 개별 파라미터를 샤당하는 대신, YaFSDP는 전체 레이어를 샤딩하는 방식으로 통신을 효율적으로 하고 중복성을 줄입니다. 각 GPU는 모델의 다른 샤드를 처리해서 메모리 사용량을 최소화합니다.
Image Credit: Yandex at Medium
버퍼 사전 할당 (Buffer Pre-allocation): YaFSD는 필요한 모든 데이터에 대한 버퍼를 미리 할당해서 Torch Allocator가 효율적으로 메모리 관리를 할 수 있도록 합니다. 버퍼를 사전 할당할 때 홀수층과 짝수층을 번갈아 가면서 Weights (가중치)와 그레디언트 (Gradient)에 두 개의 버퍼를 사용합니다.
Image Credit: Yandex at Medium
메모리 소비 최적화
YaFSDP를 사용하면 메모리 소비를 크게 줄일 수 있습니다:
Efficient Buffer Use (버퍼의 효율적 사용): 버퍼에 중간 값들을 저장하고 일정량의 메모리를 소비하도록 합니다.
Activation Checkpointing (활성화 체크포인트): 순전파 (Forward Pass; Forward Propagation) 동안 필수적인 활성화값만 저장하고 역전파 시 이를 다시 계산하여 메모리 사용량을 크게 줄입니다. 예를 들어서, 배치 사이즈가 8,192 토큰인 라마-2-70B 모델을 트레이닝하는 경우라면 Activation의 저장 공간을 110GB 이상에서 5GB로 줄일 수 있습니다. 이 때, 계산의 오버헤드가 생기게 되는데, YaFSDP는 메모리 사용량을 최적화하고 일부 레이어에 대한 Activation Checkpointing을 하지 않는 방법으로 오버헤드를 줄입니다.
가중치, 그레디언트, 옵티마이저 상태를 샤딩 (Sharding): 이런 구성 요소의 메모리 소비량은 프로세스 수가 증가함에 따라 0에 가까워지는 경향이 나타나, 중복을 최소화합니다.
통신 (Communication) 최적화
YaFSDP는 아래와 같은 방법으로 통신 효율성을 높여줍니다:
통신 (Communication)과 연산 (Computation)을 오버래핑 (Overlapping): YaFSDP는 CUDA 스트림을 사용해서 동시 연산 (Concurrent Computation)과 통신 (Communication)을 효과적으로 관리합니다. 연산용 스트림과 통신용 스트림 두 개의 스트림을 사용하고, 이벤트에 따라 동기화되어서 올바른 작동 순서를 보장합니다.
통신 오버헤드 (Communication Overhead) 감소: 필요한 경우에만 데이터를 전송하게 하고, 중복된 작업을 최소화하는 기술을 사용해서 YaFSDP는 전반적인 통신의 효율성을 높입니다.
YaFSDP는 어떤 장점이 있는 걸까요?
YaFSDP가 보여주는 성능 향상 수치는 상당합니다. 700억 개의 파라미터가 있는 모델의 경우에 약 150개의 GPU 리소스를 절약할 수 있는데, 이건 매월 50~150만 달러 정도의 비용 에 해당합니다. 트레이닝 시간은 FSDP 같은 방식에 비해서 최대 26%까지 단축할 수 있다고 합니다.
라마-2와 라마-3에서 YaFSDP가 보여준 최종적인 속도 향상 수치를 봐도, FSDP와 비교했을 때 트레이닝 효율성이 크게 개선되었다는 걸 알 수 있습니다.
YaFSDP는 HuggingFace 워크플로우와 함께 사용할 수 있는데, FSDP 대비 최대 25% 더 빠르다고 합니다.
보너스 자료
GitHub: YaFSDP
읽어주셔서 감사합니다. 친구와 동료 분들에게도 뉴스레터 추천해 주세요!
