애플 실리콘 분산 추론과 학습, 맥 스튜디오 2대로 딥시크를 돌린 GPU MODE 강연 정리

핵심 메시지

• 데이터를 외부로 보낼 수 없는 환경에서는 성능이 조금 떨어져도 기기 안에서 돌아가는 모델이 필요하다는 문제의식이 출발점이다.
• M3 울트라 맥 스튜디오는 연산량은 데이터센터 GPU에 크게 못 미치지만 512GB 통합 메모리 덕분에 큰 모델을 담는 데는 압도적으로 저렴하다.
• 딥시크 원본 정밀도를 돌리는 데 맥 스튜디오 두 대면 충분했고, 같은 일을 하는 엔비디아 서버 대비 가격은 약 15분의 1, 전력은 두 자릿수 배 적었다.
• 썬더볼트 구간에 RDMA가 열리면서 통신 지연이 0.5밀리초대에서 마이크로초대로 떨어졌고, 그제서야 텐서 병렬이 현실적인 선택지가 됐다.
• 발표자는 맥 두 대에서 로라 어댑터로 딥시크를 실제로 학습시켜 초당 30토큰 처리량을 얻었다고 밝혔다.

쉽게 이해하기

GPU MODE의 99번째 강연에 오른 매트 베턴은 옥스퍼드에서 출발한 스타트업 엑소랩스에서 소비자용 기기로 대형 언어모델을 돌리는 문제를 연구해 왔다. 그는 GPT-3 시절부터 굳어진 전제, 즉 거대 모델은 데이터센터 GPU에서만 돌릴 수 있고 그러려면 내 데이터를 남의 서버에 넘겨야 한다는 전제를 다시 따져 보자고 제안했다. 이 전제 때문에 아예 시도조차 되지 않은 용도가 시장에 남아 있다는 것이다.

구체적인 예로 그는 고객 자료를 외부로 반출할 수 없는 법률 사무소를 들었다. 상자 하나 안에서 모델이 돌고 데이터가 그 상자를 벗어나지 않는다면 쓸 수 있는 곳이 많아진다는 이야기다. 최근 개인용 노트북에서 로컬 에이전트를 띄워 쓰는 흐름 역시 같은 수요를 보여 준다고 봤다.

선택한 하드웨어는 M3 울트라 맥 스튜디오였다. 1만 달러에 512GB 통합 메모리를 얹지만 연산 성능은 반정밀도 기준 26테라플롭스 남짓으로, 900테라플롭스급 데이터센터 GPU와는 비교가 되지 않는다. 대신 기가바이트당 가격이 약 20배 저렴하고 메모리 대역폭 대비 연산량 비율이 낮아, 메모리에 묶이는 토큰 생성 작업에는 오히려 잘 맞는다는 것이 그의 계산이다.

8비트 원본 정밀도 딥시크는 약 700GB가 필요한데, 맥 스튜디오 두 대(2만 달러)면 담긴다. 같은 모델을 H200 8장짜리 서버에 올리면 가격이 약 15배로 뛴다. 대신 초당 토큰 수는 10분의 1 수준으로 떨어지고, 전력 소모는 맥 쪽이 두 자릿수 배 적다. 그는 애플 사무실에서 이 기기 16대를 한꺼번에 돌리다 8킬로와트를 끌어써 사무실 차단기를 내려 버린 일화도 소개했다.

소프트웨어 쪽에서는 애플의 딥러닝 프레임워크 MLX가 파이토치와 어떻게 다른지가 강연의 절반을 차지했다. MLX는 JAX처럼 지연 실행 모델을 쓴다. 연산을 적어 두기만 하고 eval을 부르는 순간 그래프를 거슬러 올라가 필요한 커널만 한꺼번에 스케줄한다. 덕분에 프리필 단계에서는 최종 출력 대신 KV 캐시만 실체화하도록 호출을 바꾸는 식으로, 파이토치라면 모델 순전파 구조를 손봐야 할 최적화를 공짜로 얻는다.

주요 인사이트

• 여러 대로 쪼개는 방식은 통신 횟수가 갈랐다. 파이프라인 병렬은 기기 수만큼만 주고받으면 되지만, 텐서 병렬은 층마다 두 번씩 통신해야 해서 61개 층짜리 모델이면 토큰 하나에 122번을 오간다. 썬더볼트 홉당 0.5밀리초였을 때는 토큰당 61밀리초가 통신에만 들어가 계산할 시간이 남지 않았다.
• 맥OS 업데이트로 썬더볼트 위에서 RDMA가 가능해지자 홉당 지연이 5마이크로초 수준으로 떨어졌다. 같은 텐서 병렬이 갑자기 쓸 만해졌고, 기기를 더 붙이는 목적이 '더 큰 모델을 담기'에서 '초당 토큰 수 늘리기'로 바뀌었다.
• 성격이 다른 기기를 섞는 것도 이득이 된다. 연산량이 많은 프리필은 플롭스가 높은 소형 AI 워크스테이션에, 메모리에 묶이는 디코드는 대역폭이 높은 맥에 맡기면 전체 지연이 줄어든다. 다만 두 기종 사이에는 전용 연결이 없어 10기가비트 이더넷으로 KV 캐시를 흘려보내야 하는 제약이 남는다.
• MLX에는 텐서마다 기울기를 달아 두는 개념이 없어, 중간 장비가 받은 기울기에서 역전파를 시작할 방법이 없었다. 발표자는 받은 기울기와 중간 활성값의 내적을 가짜 손실로 삼으면 같은 효과가 난다는 우회로를 찾아냈다.
• 분산 환경의 버그는 대부분 조용한 교착 상태로 나타났다. MLX가 그래프의 송신·수신 노드 순서를 임의로 고르는 바람에 순전파를 기다리는 장비와 기울기를 기다리는 장비가 서로 멈춰 섰고, 인위적인 의존 관계를 그래프에 심어야 풀렸다. 성숙한 개발 도구가 얼마나 큰 자산인지 드러난 대목이다.

자주 묻는 질문

관련 AI 소식