MIT 6.S191 대규모 병렬 학습 강의 정리: 데이터·텐서·파이프라인 병렬과 샤딩, DeepSpeed·FSDP·Megatron 비교

핵심 메시지

• 학습 데이터와 모델 크기를 함께 키우면 손실이 낮아진다는 스케일링 법칙이 대규모 학습의 출발점이지만, 모델 수명 전체로 보면 추론 비용이 학습 비용을 압도하기 때문에 업계는 오히려 작고 효율적인 모델로 향하고 있다.
• 10억 파라미터 모델 하나를 학습시키는 데만 약 28GB가 필요하다. 가중치 외에 그래디언트 버퍼, 옵티마이저 상태, 역전파용 활성값까지 GPU에 올려야 하기 때문이다.
• 활성값 체크포인팅은 몇 개 층마다만 활성값을 남기고 나머지는 역전파 때 다시 계산하는 방식으로, 순전파 한 번을 더 하는 대가로 메모리를 크게 줄인다.
• 데이터·텐서·파이프라인·시퀀스·전문가 병렬은 택일하는 선택지가 아니라 서로 다른 축이며, 실제 클러스터에서는 이들을 동시에 겹쳐 쓰는 4D·5D 병렬로 구성한다.
• DeepSpeed(마이크로소프트), FSDP(메타), Megatron-LM(엔비디아)은 기능은 비슷하지만 설계 철학이 다르다. 강연자는 진입 장벽이 낮은 DeepSpeed를 먼저 권하면서, 자사는 FSDP에 Megatron 포크를 결합해 쓴다고 밝혔다.

쉽게 이해하기

MIT의 딥러닝 입문 강의 6.S191에 Liquid AI의 공동창업자 겸 최고기술책임자 마티아스 레히너가 초청 강연자로 나섰다. 그는 MIT CSAIL에서 리퀴드 신경망을 공동 고안한 뒤 회사를 창업했고, 이날 강연에서는 자사가 수천 장 규모의 GPU로 학습을 돌릴 때 실제로 쓰는 기법들을 순서대로 풀어놨다. 주제는 왜 규모를 키우는가, 무엇을 어떻게 쪼개는가, 그리고 어떤 프레임워크로 구현하는가 세 갈래였다.

강연자는 먼저 규모를 키우는 이유를 데이터로 보여줬다. 메타가 공개한 라마 2의 학습 곡선을 보면 같은 구조라도 학습 토큰이 늘수록 손실이 내려가고, 같은 데이터라도 파라미터가 많은 모델일수록 손실이 낮았다. 업계 흐름도 그 방향이어서 GPT-3는 1750억 파라미터를 3천억 토큰으로, 라마 2의 최대 모델은 700억 파라미터를 2조 토큰으로 학습했고, 자사가 최근 공개한 소형 모델은 12억 파라미터를 28조 토큰으로 학습했다고 소개했다. 그런데 모델 크기 자체는 오히려 줄고 있다. 학습은 한 번이지만 추론은 수백만 사용자가 반복하기 때문에, 모델 수명 전체의 연산 비용에서 학습이 차지하는 몫은 작다는 것이 그의 설명이다.

다음은 메모리 문제였다. 파라미터를 16비트로 저장한다고 해도 그래디언트 버퍼와 옵티마이저 상태, 역전파에 필요한 활성값이 함께 올라가기 때문에 10억 파라미터 모델의 학습에만 약 28GB가 든다. 700억 파라미터로 늘리면 2TB가 필요한데 현재 단일 GPU의 메모리는 수백 GB 수준이라 애초에 올라가지 않는다. 해법으로 제시된 것이 활성값 체크포인팅과 CPU 오프로딩, 그리고 샤딩이다. 체크포인팅은 연산을 더 써서 메모리를 아끼는 맞교환인데, GPU에서는 대체로 연산이 메모리보다 싸기 때문에 업계에서 널리 쓰인다. 반대로 오프로딩은 PCI 버스 대역폭이 병목이라 단일 GPU 파인튜닝처럼 다른 선택지가 없을 때 주로 쓰인다고 선을 그었다.

강연의 핵심은 병렬화 축을 하나씩 분해한 대목이다. 데이터 병렬은 모델을 복제해 배치를 나누는 가장 단순한 방식이지만 유효 배치 크기가 커지면 일반화 성능이 떨어져 한계가 있다. 파이프라인 병렬은 층을 깊이 방향으로 쪼개는데, 앞 GPU가 일할 때 뒤 GPU가 노는 파이프라인 버블이 생겨 마이크로배칭으로 겹쳐 돌린다. 텐서 병렬은 가중치 행렬 자체를 열 또는 행으로 쪼개는 방식이고, 트랜스포머의 MLP처럼 열 분할과 행 분할을 이어 붙이면 두 층을 통신 한 번으로 처리할 수 있다. 시퀀스·컨텍스트 병렬은 긴 문맥을 토큰 구간으로 나누는데 MLP나 레이어 정규화는 시간축에 독립적이라 통신이 필요 없고 어텐션만 별도 처리가 필요하다. 전문가 혼합 모델은 전문가를 서로 다른 장치에 두는 것이 자연스러운 분할이 된다.

마지막으로 도구 이야기가 나왔다. DeepSpeed의 ZeRO는 옵티마이저, 그래디언트, 파라미터를 단계적으로 쪼개고, PyTorch 네이티브인 FSDP는 층에 들어갈 때 필요한 가중치만 모았다가 나갈 때 다시 흩는 방식이다. Megatron-LM은 메모리보다 처리량과 병렬화에 무게를 두고 텐서·파이프라인 병렬을 자체 계층으로 제공한다. 강연자는 이 셋의 기능이 대체로 겹치며 가장 큰 차이는 만든 회사라고 농담하듯 정리하면서도, 코드 한 줄만 바꾸면 되는 DeepSpeed, PyTorch답게 모듈을 정의하는 FSDP, 통신 그룹을 직접 관리하는 Megatron으로 사용성과 제어권의 맞교환이 갈린다고 설명했다.

주요 인사이트

• 메모리와 연산 중 무엇이 더 비싼지가 기법 선택을 좌우한다. 활성값 체크포인팅이 널리 쓰이는 이유도, 오프로딩이 니치에 머무는 이유도 결국 GPU에서 연산은 상대적으로 싸고 메모리와 대역폭은 비싸다는 사실에서 나온다.
• 병렬화 전략은 클러스터의 배선 구조를 그대로 반영해야 한다. 통신이 잦은 텐서 병렬은 NVLink로 묶인 노드 한 대 안에 두고, 통신이 드문 데이터 병렬은 InfiniBand 너머로 펼치는 식이다. 강연자는 캐시, HBM, 노드 내 링크, 스케일아웃 망으로 이어지는 계층을 기억하는 것이 핵심 원칙이라고 말했다.
• 전문가 혼합 구조의 진짜 난제는 로드 밸런싱이다. 라우터가 특정 전문가로 쏠리면 학습이 불안정해질 뿐 아니라, 전문가를 장치별로 흩어놓은 상황에서는 일부 GPU만 일하고 나머지가 노는 결과로 이어진다. 강연자는 이 문제가 여전히 활발한 연구 주제라고 밝혔다.
• 작은 모델은 오히려 병렬화로 속도를 올리기 어렵다. 층 수도 적고 행렬도 작아 쪼갤 여지가 없기 때문이다. 이 경우에는 토큰을 더 빨리 흘리는 대신 교사 모델의 확률 분포를 함께 배우는 증류처럼 학습 신호 자체를 풍부하게 만드는 접근이 쓰인다.
• 데이터의 병목은 양이 아니라 품질이다. 초반 수조 토큰은 논문과 백과사전처럼 질 좋은 출처로 채울 수 있지만 그 이상은 웹 크롤 데이터를 걸러내야 하고, 합성 데이터로만 채우면 학습 동역학이 무너지는 현상이 나타난다고 강연자는 경고했다.

자주 묻는 질문

관련 AI 소식