- Turing Post Korea
- Posts
- Topic #5: GPU 최적화 기술, FSDP와 YaFSDP는 무엇인가?
Topic #5: GPU 최적화 기술, FSDP와 YaFSDP는 무엇인가?
YaFSDP는 어떻게 FSDP보다 나은 GPU 최적화 성능을 보여주는가
‘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에도 몇 가지 중요한 한계는 있어요.
Reply