LLM 병렬 학습 총정리: 데이터·ZeRO·파이프라인·텐서 병렬과 라마3, 딥시크 사례

핵심 메시지

• 집합 통신에서 올리듀스는 리듀스-스캐터와 올게더를 이어 붙인 것과 결과가 같지만 통신량은 훨씬 적어, 거의 모든 병렬화 기법이 이 성질을 활용한다.
• 단순한 데이터 병렬은 모든 GPU가 파라미터·그래디언트·옵티마이저 상태를 통째로 들고 있어야 해서, GPU를 아무리 늘려도 올릴 수 있는 모델 크기가 늘지 않는다.
• ZeRO는 옵티마이저 상태·그래디언트·파라미터를 차례로 쪼개 같은 하드웨어에서 학습 가능한 모델 크기를 열 배 가까이 늘리지만, 3단계에서는 통신량이 1.5배로 늘어나는 대가를 치른다.
• 파이프라인 병렬은 버블과 활성값 메모리 문제를 낳고, 1F1B 스케줄링은 백워드를 앞당겨 메모리를 일찍 해제하는 방식으로 이를 완화한다.
• 텐서 병렬이 나눌 수 없는 정규화·드롭아웃 같은 연산은 시퀀스 병렬이 맡으며, 실제 대형 모델들은 대부분 두 가지 이상을 섞은 다차원 병렬화를 쓴다.

쉽게 이해하기

강의는 GPU 한 장을 잘 쓰는 문제를 다뤘던 앞 회차를 짧게 복습하며 시작한다. 공유 메모리는 20사이클 남짓이면 읽고 쓸 수 있지만 L1·L2는 200사이클, HBM이나 DRAM 같은 글로벌 메모리는 300사이클에 가까워 열 배 안팎의 차이가 난다. 플래시어텐션이나 타일링 같은 기법이 등장한 배경이며, 연산량 대비 통신량을 보는 산술 강도가 낮으면 통신에 지배당하는 구조가 된다는 점도 함께 짚는다.

여러 GPU와 여러 노드로 확장할 때 먼저 배워야 하는 것은 집합 통신이다. 브로드캐스트와 스캐터, 게더, 올게더, 리듀스, 올리듀스, 리듀스-스캐터를 정리한 뒤 강의는 한 가지 성질을 특히 강조한다. 올리듀스는 리듀스-스캐터를 한 번 돌리고 그 결과를 올게더하는 것과 결과가 같지만, 통신 비용을 계산해 보면 전자는 GPU 수에 비례해 늘어나는 반면 후자는 조각 단위로만 주고받아 그 항이 상쇄된다. 대역폭이 제한된 환경에서 뒤쪽이 훨씬 유리하며, 이후 등장하는 거의 모든 기법이 이 분해를 재활용한다.

데이터 병렬은 전역 배치를 GPU 수만큼 잘라 나눠주고 각자 계산한 그래디언트를 올리듀스로 동기화하며, 통신 비용은 배치마다 파라미터 개수의 두 배다. 문제는 메모리로, 혼합 정밀도에서 파라미터와 그래디언트가 각 2바이트라 해도 안정성을 위한 FP32 마스터 가중치와 아담의 두 모멘트가 더해져 파라미터 하나당 최소 열여섯 바이트 이상이 필요하고 이 사본을 모든 GPU가 들고 있어야 한다. 2019년 제안된 ZeRO는 상태 자체를 쪼개, 1단계는 옵티마이저 상태만, 2단계는 그래디언트까지, 3단계는 파라미터까지 각 GPU가 자기 몫만 보관하게 한다. 강의의 예시에서 120GB를 요구하던 설정이 각각 30GB, 16GB, 2GB 수준으로 줄고, GPU 여덟 장 기준 학습 가능한 모델 규모는 DDP의 6B에서 53B까지 올라간다.

대신 3단계는 순전파와 역전파에서 파라미터를 각각 올게더하고 그래디언트를 리듀스-스캐터해야 해 통신량이 세 배로 늘어난다. 데이터 병렬 자체의 한계도 있는데, GPU 수를 늘리면 전역 배치도 함께 커져야 하지만 임계 배치 크기를 넘으면 학습 효율이 더 나아지지 않고, ZeRO 3단계는 GPU를 768장에서 2천 장 가까이 늘렸을 때 장당 처리량이 급격히 떨어진다. 그래서 등장하는 파이프라인 병렬은 레이어를 GPU에 나누고 배치를 마이크로 배치로 쪼개 흘려보내지만, 마이크로 배치가 너무 작으면 통신이 지배하고 앞쪽 GPU는 역전파 시점까지 활성값을 들고 있어야 해 메모리가 터진다. 파이프드림이 제안한 1F1B는 순전파 한 번과 역전파 한 번을 번갈아 수행해 오래된 마이크로 배치부터 메모리를 곧바로 해제한다.

텐서 병렬은 어텐션과 피드포워드의 행렬 곱을 여러 GPU로 쪼개지만 레이어 정규화나 잔차 연결, 드롭아웃 같은 요소별 연산은 나눌 수 없어, 시퀀스 차원으로 쪼개는 시퀀스 병렬이 번갈아 그 몫을 맡는다. 활성값 재계산은 역설적인데, 재계산을 하지 않으면 같은 배치에서 처리량이 높지만 메모리 때문에 배치를 8 이상 키우지 못하는 반면 재계산을 쓰면 256까지 늘려 전체 처리량이 더 높아진다. 사례를 보면 라마 3는 밀집 구조라 전문가 병렬을 빼고 4차원 병렬화를 썼고, 마지막 학습 단계에서 10만 토큰대 컨텍스트를 다루려 컨텍스트 병렬을 넣으며 데이터 병렬을 줄였다. 딥시크 V3는 텐서 병렬 없이 파이프라인·전문가·데이터 병렬을 조합했고, 강의의 결론은 단일 해법은 없으며 두 가지 이상을 섞는 것이 사실상의 표준이라는 것이다.

주요 인사이트

• 병렬화 설계는 결국 메모리와 통신의 교환이다. ZeRO 3단계처럼 메모리를 가장 많이 아끼는 선택이 통신량을 늘려 대규모에서는 오히려 처리량을 깎아먹을 수 있다.
• GPU를 더 꽂는 것이 곧 더 큰 모델을 뜻하지 않는다. 단순 데이터 병렬에서는 모든 장비가 모델 전체 사본을 들고 있어야 해서, 장비 수와 모델 크기는 애초에 연결돼 있지 않다.
• 임계 배치 크기라는 개념이 병렬화 전략을 제약한다. GPU 수만큼 전역 배치를 키워야 하는 데이터 병렬은 학습 효율이 더 나아지지 않는 지점을 필연적으로 만난다.
• 메모리를 아끼려 연산을 다시 하는 활성값 재계산이 전체 처리량을 높이는 결과는, 지연 시간만 보고 최적화하면 오히려 손해라는 점을 보여준다.
• 만 장 단위 학습에서는 알고리즘만큼 인프라가 문제다. 라마 3에서 개별 GPU가 죽는 일이 150회 넘게 일어나고 전원 공급 장치 문제까지 세 차례 발생했다는 보고는, 복구와 모니터링 설계가 학습 파이프라인의 일부임을 뜻한다.

자주 묻는 질문

관련 AI 소식