PyTorch, 유럽 공공 슈퍼컴퓨터에서 1조 파라미터 MoE 학습 확장

들어가며

약 1조 개 파라미터 규모의 혼합 전문가(MoE) 모델을 학습하려면 보통 AI 전용 대규모 클러스터가 필요하다. EUROPA 컨소시엄은 EuroHPC Frontier AI Grand Challenge에서 다른 접근을 택했다. 여러 연구 프로젝트가 함께 사용하는 공공 슈퍼컴퓨터 JUPITER에 PyTorch와 TorchTitan을 적용해 최첨단 규모의 사전학습을 확장한 것이다.

실험 모델의 총 파라미터 수는 9,519억 개이며, 토큰마다 실제로 활성화되는 파라미터는 약 328억 개다. 모델은 74개 레이어, 384개 라우팅 전문가, top-8 라우팅, 6,144의 히든 크기, 25만6,000개 어휘, 65,536 토큰의 사전학습 시퀀스로 구성됐다. 계산에는 96개 노드와 384개의 NVIDIA GH200 GPU가 사용됐다.

핵심 문제는 시스템의 불균형이었다

JUPITER Booster 노드는 GH200 슈퍼칩 4개로 구성된다. 각 GPU는 96 GB HBM3를 갖고, 72코어 Grace CPU 및 120 GB LPDDR5X 메모리와 NVLink-C2C로 연결된다. 반면 노드 간 통신은 InfiniBand를 사용하며, 로컬 GPU-CPU 연결보다 훨씬 느리다.

연구진은 이런 구조를 고려해 384-way FSDP를 사용하고, 노드 안에 4-way 전문가 병렬화를 배치했다. 파이프라인 병렬화는 사용하지 않았으며, MoE의 token all-to-all 통신은 가능한 한 노드 내부 NVLink에서 처리했다. 긴 시퀀스는 통신을 숨길 수 있는 연산 시간을 늘리는 역할도 했다.

성능을 바꾼 세 가지 최적화

• FlexShard로 Muon 통신을 줄였다. 대부분의 가중치 행렬에는 Muon이 사용된다. Muon의 Newton-Schulz 업데이트는 전체 행렬을 필요로 하므로, FSDP로 잘게 나뉜 행렬을 매번 모으고 다시 분산하면 통신 비용이 커진다. FlexShard는 전문가 단위로 파라미터를 배치해 각 rank가 완전한 전문가 행렬을 로컬에서 처리하도록 했다. 밀집 행렬이 이동하는 동안에는 전문가 계산을 진행했다. Newton-Schulz를 Grace CPU가 아니라 GPU에서 실행하자 옵티마이저 단계가 약 41.7초에서 2.2초로 줄었고, GPU당 처리량은 441에서 597 tokens/s로 증가했다.
• 활성화 체크포인트를 Grace로 오프로딩했다. 65,536 토큰 시퀀스에서는 레이어 경계 활성화만 저장해도 GPU당 약 55 GiB가 필요하다. TorchTitan의 체크포인트 기능을 확장해 순전파 뒤 일부 활성화를 Grace 메모리로 옮기고, 역전파 재계산 전에 다시 가져왔다. 32,768 토큰에서 전체 활성화 재계산을 적용하면 처리량이 90에서 370 tokens/s로 늘었고, 오프로딩 후에는 508 tokens/s에 도달했다.
• FSDP 통신과 연산을 겹쳤다. FSDP는 현재 레이어를 계산하는 동안 다음 레이어의 all-gather를 수행할 수 있다. 연구진은 unshard_async_op를 연산 스트림에서 실행해 GPU당 약 14 GiB의 HBM을 확보하고 통신과 계산의 연결을 개선했다. 최적화된 구성은 384개 GH200에서 GPU당 약 1,300 tokens/s를 기록했다.

의미와 한계

이 프로젝트의 의미는 단순한 처리량 기록을 넘어선다. AI 학습만을 위해 설계되지 않은 공공 연구 슈퍼컴퓨터도 메모리 계층, 병렬화, 옵티마이저, 통신 스케줄링을 함께 설계하면 프런티어급 사전학습을 수행할 수 있음을 보여준다. PyTorch의 유연성과 TorchTitan의 분산학습 구성 요소는 JUPITER의 특성에 맞춘 세밀한 조정을 가능하게 했고, 일부 기능은 오픈소스 프로젝트로 다시 반영됐다.

다만 이 결과를 모든 1조 파라미터 MoE에 그대로 적용할 수는 없다. GH200 토폴로지, Grace 메모리 용량, 시퀀스 길이, 병렬화 배치에 따라 성능이 달라지기 때문이다. 더 일반적인 교훈은 노드 간 대역폭이 병목이 될수록 모델 구조와 시스템 소프트웨어를 별개의 문제가 아니라 하나의 통합 설계 대상으로 다뤄야 한다는 점이다.

출처: PyTorch Blog