Cerebras가 3조 파라미터 모델을 59개 웨이퍼에 배치하는 방법

조 단위 파라미터 모델을 배포할 때 가장 먼저 부딪히는 문제는 연산량보다 메모리다. 방대한 가중치와 요청별 상태를 어디에 저장할지가 핵심이기 때문이다. Cerebras는 기술 글에서 약 3조 파라미터 오픈 모델 Kimi K3를 사례로 들어, 모델을 여러 WSE-3T 웨이퍼 스케일 프로세서에 분산하는 방식을 설명했다. 원문에 언급된 성능과 벤치마크 비교는 Cerebras의 설명이며, 독립적으로 검증된 결과는 아니다.

모델을 계층 단위로 분할

WSE-3T 한 장의 온칩 메모리는 44GB다. 반면 조 단위 모델의 가중치는 테라바이트 규모이므로 단일 칩에 담을 수 없다. Cerebras는 4비트 가중치 양자화를 가정하고 모델 계층별 저장 공간을 계산한다. Kimi K3의 3조 파라미터를 93개 계층으로 나누면 계층 하나당 평균 약 16GB가 필요하다는 계산이 나온다.

모든 계층을 같은 방식으로 배치할 수 있는 것은 아니다. Kimi K3는 선형 어텐션 계층 69개와 밀집 어텐션 계층 24개로 구성된다. 선형 어텐션은 요청마다 비교적 작고 고정적인 상태를 유지하므로, 사례 구성에서는 웨이퍼 한 장에 두 계층을 배치할 수 있다. 밀집 어텐션은 컨텍스트 길이에 따라 커지는 KV 캐시를 저장해야 하므로, 한 장에 계층 하나를 놓고 남은 메모리를 캐시에 할당한다.

이 방식으로 계산하면 선형 계층에 35장, 밀집 계층에 24장이 필요해 총 59개 웨이퍼가 된다.

웨이트 대신 활성값을 이동

모델을 매핑한 뒤 각 웨이퍼에 배치된 가중치는 그 위치에 남는다. 요청은 조립 라인처럼 여러 처리 단계를 순서대로 통과하며, 한 단계에서 계산한 활성값이 다음 단계로 전달된다. Cerebras는 웨이퍼 간 통신에 RDMA over Ethernet을 사용하며, 한 홉의 지연을 약 2마이크로초로 제시한다.

이 설계의 중요한 점은 완전한 계층을 웨이퍼 한 장에 담을 수 있는 경우가 많다는 것이다. 가중치를 칩 사이에서 반복적으로 이동할 필요가 없으므로, 주요 통신 대상은 앞 단계가 생성한 활성값이 된다. 즉, 대규모 모델을 여러 칩에 배치하면서도 가중치 이동에 따른 부담을 줄이는 구조다.

여러 요청으로 파이프라인 채우기

파이프라인 방식이 한 번에 하나의 요청만 처리한다는 뜻은 아니다. 한 요청이 후반 단계에서 계산되는 동안 다른 요청은 중간 단계에, 또 다른 요청은 전반 단계에 있을 수 있다. Cerebras는 실제 운영 환경에서 단순한 계층 단위 예시보다 더 세밀한 단계로 나눈다고 설명한다. 각 단계는 단일 요청 대신 작은 마이크로배치를 처리할 수도 있다.

작은 마이크로배치는 개별 요청의 지연시간에 유리하고, 큰 마이크로배치는 전체 처리량을 높이는 데 유리할 수 있다. 동시에 밀집 계층 웨이퍼에 남겨둔 KV 캐시 공간은 파이프라인에 동시에 들어갈 수 있는 요청 수를 결정한다. 따라서 배포 설계에서는 가중치 용량뿐 아니라 캐시, 동시성, 통신량, 지연시간을 함께 조정해야 한다.

의미와 한계

이 사례는 웨이퍼 수를 늘려 단일 가속기의 메모리와 연산 한계를 넘어선 뒤, 파이프라인으로 여러 요청의 작업을 겹치는 방법을 보여준다. 다만 59개라는 수치는 모든 3조 파라미터 모델에 적용되는 규칙이 아니다. 컨텍스트 길이, 양자화 방식, 계층 구조, 캐시 형식, 요청 분포에 따라 필요한 구성이 달라진다.

출처: Hacker News