웨이퍼 스케일 엔진(Wafer-Scale Engine, WSE)은 2015년 앤드루 펠드먼(Andrew Feldman), 게리 라우터바흐(Gary Lauterbach) 등이 설립한 세레브라스 시스템즈(Cerebras Systems)가 개발한 대규모 프로세서 시리즈이다. 실리콘 웨이퍼에서 잘라낸 기존의 칩과 달리, WSE는 웨이퍼 전체를 단일 다이로 사용하여 전례 없는 수의 코어와 온칩 메모리를 제공한다. 이 설계는 다중 칩 시스템을 괴롭히는 통신 병목 현상을 줄이는 것을 목표로 하며, 특히 대규모 언어 모델 및 기타 딥러닝 워크로드 훈련에 적합하다.
1세대인 WSE-1은 2019년 8월에 발표되었다. 1.2조 개의 트랜지스터와 40만 개의 AI 최적화 코어를 포함하며, TSMC의 16나노미터 공정으로 제조되었다. 칩 크기는 46,225제곱밀리미터로, 당시 가장 큰 그래픽 처리 장치보다 50배 이상 컸다. 2세대인 WSE-2는 2021년 4월에 출시되었으며, 7나노미터 공정을 사용하여 코어 수를 85만 개로 두 배로 늘리고 온칩 메모리를 40기가바이트로 증가시켰다. 2024년 3월, 세레브라스는 5나노미터 공정으로 제작된 WSE-3를 발표했으며, 90만 개의 코어와 44기가바이트의 온칩 SRAM을 갖추고 최대 24조 개의 매개변수를 가진 모델 훈련을 목표로 한다.
아키텍처 및 설계
WSE의 핵심 혁신은 웨이퍼 스케일 통합이다. 웨이퍼에 많은 소형 다이를 제조하여 개별적으로 패키징하는 대신, 세레브라스는 웨이퍼 전체를 그대로 유지하고 제조 결함을 우회하기 위해 중복 회로를 사용한다. 이 접근 방식은 분산 시스템에서 성능 병목 현상이 자주 발생하는 인쇄 회로 기판을 통한 칩 간 통신의 필요성을 제거한다. 칩은 각 코어에 자체 메모리가 있는 메시 연결 네트워크를 사용하여 고대역폭, 저지연 데이터 이동을 가능하게 한다.
각 코어는 신경망에서 흔히 사용되는 희소 선형 대수 연산에 최적화된 단순화된 프로세서이다. 코어는 FP16, BF16, FP32를 포함한 다양한 데이터 형식을 지원하며, 밀집 및 희소 연산을 효율적으로 실행하도록 설계되었다. 온칩 메모리는 코어 전체에 분산되어 있어 기존 GPU 메모리 시스템을 훨씬 초과하는 총 대역폭을 제공한다. 이 아키텍처는 AI 훈련에서 자주 제한 요소가 되는 외부 메모리에 대한 의존도를 줄인다.
WSE는 독립형 제품이 아니라 전원 공급 시스템과 냉각 솔루션을 포함하는 세레브라스의 CS-2 및 CS-3 시스템에 통합된다. 2020년에 출시된 CS-2는 WSE-2를, 2024년에 발표된 CS-3는 WSE-3를 탑재한다. 이 시스템은 기존 데이터 센터에 연결되도록 설계되었으며 대규모 훈련을 위해 클러스터로 자주 사용된다.
성능 및 벤치마크
세레브라스는 WSE에 대한 여러 성능 벤치마크를 발표했다. 2021년, 회사는 단일 CS-2 시스템이 GPT 아키텍처를 사용하여 18억 매개변수 모델을 훈련할 수 있음을 입증했으며, 여러 시스템을 연결했을 때 거의 선형 확장을 달성했다. 2022년, 세레브라스와 샌디아 국립 연구소는 CS-2를 사용하여 2억 뉴런 스파이킹 신경망을 실시간으로 시뮬레이션했는데, 이는 수천 개의 GPU를 갖춘 슈퍼컴퓨터가 필요했던 작업이었다.
2023년, 세레브라스는 자사 시스템이 통신 오버헤드 감소 덕분에 16개의 NVIDIA A100 GPU 클러스터보다 짧은 시간에 70억 매개변수 모델을 훈련할 수 있다고 보고했다. 회사에 따르면 WSE-3는 최대 125페타플롭스의 AI 컴퓨팅 성능을 제공할 수 있어 AI 분야에서 가장 강력한 단일 칩 프로세서 중 하나이다. 그러나 독립적인 벤치마크는 제한적이며, 대부분의 성능 주장은 세레브라스 자체에서 나온 것이다.
소프트웨어 생태계
WSE를 프로그래밍하기 위해 세레브라스는 칩 아키텍처에 머신러닝 모델을 매핑하는 컴파일러를 포함하는 세레브라스 소프트웨어 플랫폼(CSoft)을 개발했다. CSoft는 TensorFlow 및 PyTorch와 같은 인기 있는 프레임워크를 지원하여 연구자들이 큰 코드 변경 없이 기존 모델을 실행할 수 있게 한다. 컴파일러는 데이터 흐름 스케줄링, 메모리 할당 및 웨이퍼 스케일 설계에 대한 최적화를 처리한다.
세레브라스는 또한 트랜스포머, 컨볼루션 네트워크 및 순환 네트워크를 포함한 사전 최적화된 모델 구현 라이브러리를 제공한다. 소프트웨어 스택에는 여러 CS 시스템에서 분산 훈련을 위한 도구가 포함되어 있으며, 가중치 스트리밍이라는 기술을 사용하여 모델 가중치를 칩 메모리를 통해 파이프라인 방식으로 전달한다. 이 접근 방식은 외부 저장소에서 가중치를 스트리밍하여 온칩 메모리보다 큰 모델을 훈련할 수 있게 한다.
응용 분야 및 사용 사례
WSE는 상업 및 연구 목적으로 다양한 조직에서 채택되었다. 2021년, 아랍에미리트의 G42 그룹은 AI 연구용 슈퍼컴퓨터를 구축하기 위해 세레브라스와 파트너십을 맺었으며, 이후 이를 사용하여 Jais 및 Falcon 모델을 훈련했다. 2023년, 세레브라스는 Qualcomm과 온디바이스 AI 모델 개발 협력을 발표했지만, 이 파트너십의 세부 사항은 부족하다.
의료 분야에서 세레브라스는 제약 회사와 협력하여 약물 발견 시뮬레이션을 가속화했다. 칩이 대규모 그래프 신경망을 처리할 수 있는 능력은 분자 특성 예측에 사용되었다. 또한 WSE는 기상 예보에도 사용되었으며, GPU 클러스터보다 고해상도 기후 모델 훈련이 더 빠르다는 실험이 있었다.
다른 AI 하드웨어와의 비교
WSE는 Groq의 언어 처리 장치 및 SambaNova의 재구성 가능한 데이터플로우 장치와 같은 다른 특수 AI 프로세서와 경쟁한다. 범용 병렬 프로세서인 GPU와 달리 WSE는 AI 워크로드에 최적화된 도메인 특정 아키텍처이다. 주요 장점은 대규모 훈련에서 상당한 오버헤드가 될 수 있는 칩 간 통신을 제거한다는 점이다.
그러나 WSE의 크기와 전력 소비(CS-2의 경우 최대 15킬로와트)는 적절한 냉각 시설을 갖춘 데이터 센터로 배포를 제한한다. 반면, AWS Trainium 및 Google의 TPU는 클라우드 서비스로 제공되어 많은 조직에서 더 접근하기 쉬울 수 있다. WSE는 또한 더 전통적인 설계이지만 성숙한 소프트웨어 생태계의 이점을 누리는 AMD의 Instinct 가속기 및 Intel의 Gaudi 프로세서와도 경쟁한다.
과제 및 한계
WSE의 주요 과제 중 하나는 수율이다. 결함 없는 웨이퍼를 제조하는 것은 어렵지만, 세레브라스는 결함이 있는 코어를 비활성화하고 통신을 우회하여 결함을 허용하는 "중복성"이라는 기술을 사용한다. 이 접근 방식은 회사가 폐기되었을 웨이퍼를 사용할 수 있게 하지만, 칩 간에 유효 코어 수가 달라진다는 것을 의미한다.
또 다른 한계는 메모리 용량이다. 온칩 메모리는 크지만 GPU에서 사용 가능한 고대역폭 메모리보다 여전히 작다. 큰 임베딩 테이블이나 어텐션 행렬이 필요한 모델의 경우 WSE는 외부 메모리에 의존해야 하며, 이는 성능 이점을 감소시킨다. 또한 소프트웨어 생태계는 CUDA보다 덜 성숙하며, 일부 연구자들은 사용자 정의 커널을 이식하는 데 어려움을 보고했다.
CS-2 시스템의 비용은 공개되지 않았지만 수백만 달러로 추정되어 많은 소규모 연구 그룹이 접근하기 어렵다. 세레브라스는 자체 데이터 센터 및 Oracle Cloud 및 Azure와의 파트너십을 통해 클라우드 액세스를 제공함으로써 이를 해결했다.
향후 방향
세레브라스는 WSE 아키텍처를 계속 개선하고 있다. 2024년에 발표된 WSE-3는 최대 24조 개의 매개변수를 가진 모델 훈련을 지원하도록 설계되었으며, 이를 위해서는 여러 CS-3 시스템 클러스터가 필요하다. 회사는 또한 성능을 더욱 높이고 비용을 줄이기 위해 칩렛과 같은 새로운 패키징 기술을 탐색하고 있다.
2023년, 세레브라스는 IPO를 신청했지만 시장 상황으로 인해 계획이 연기되었다. 회사는 OpenAI 공동 창립자 샘 알트만(공식 자격은 아님) 및 Benchmark Capital을 포함한 투자자로부터 7억 달러 이상의 자금을 조달했다. 2024년 현재 WSE는 AI 하드웨어 환경에서 틈새 시장이지만 영향력 있는 제품으로 남아 있으며, 단일 칩 컴퓨팅에서 가능한 경계를 넓히고 있다.
수용 및 영향
WSE는 AI 커뮤니티에서 엇갈린 반응을 얻었다. 지지자들은 메모리 대역폭 문제를 해결하는 혁신적인 접근 방식을 강조하는 반면, 회의론자들은 GPU 클러스터와 비교하여 실용성과 비용 효율성에 의문을 제기한다. 독립적인 연구에 따르면 희소 어텐션 및 그래프 처리와 같은 특정 워크로드의 경우 WSE가 GPU를 상당한 차이로 능가할 수 있지만, 다른 워크로드에서는 이점이 덜 명확하다.
논쟁에도 불구하고 WSE는 특히 온칩 메모리 및 상호 연결 분야에서 다른 AI 칩 설계에 영향을 미쳤다. 그 성공은 이전에 비실용적이라고 여겨졌던 웨이퍼 스케일 통합의 개념을 검증했다. 생성형 AI 모델이 계속 커짐에 따라 WSE와 같은 특수 하드웨어에 대한 수요는 증가할 가능성이 높지만, 경쟁 환경은 여전히 역동적이다.
같이 보기
참고 문헌
- Cerebras Systems. "Wafer-Scale Engine: An Introduction." 2019.
- Feldman, A. et al. "The Cerebras Wafer-Scale Engine." IEEE Micro, 2021.
- Cerebras Systems. "CS-3 System Overview." 2024.
- Sandia National Laboratories. "Real-Time Spiking Neural Network Simulation." 2022.
- G42. "Building the World's Largest AI Supercomputer." 2021.