[Hotchips 리뷰] Wafer-Scale Engine 2 Cerabras CS-2: 120 Trillion 모델을 Data Parallelism으로 처리하기
Model Parallelism이 필요없이 Data Parallelism만으로 120 Trillion 모델을 학습할 수 있는 시대가 왔다.
12인치 실리콘 디스크인 Wafer는 개별적인 칩으로 분할하여 수백에서 추천개의 컴퓨티 칩을 만들 수 있다. 반면 Cerebras는 1개의 Wafer에 수십만개의 코어를 집적하여 세계 최초의 Wafer-Scale AI 칩인 Wafter-Scale Engine(WSE)-1을 ‘19년 8월에 발표하였다.
Cerebras의 WSE-1 칩을 장착한 CS-1이라는 가속기 시스템 제품은 다수의 국립연구소, 제약회사, 생명공학 연구 센터에 구축될 예정이다. NVIDIA GPU가 그 영역을 넓히고 있는 Supercomputing 시장에서 Cerebras는 Sambanova와 더불어 가속기 스타트업으로 Supercomputing에 진출하는 몇 안되는 업체이다.
1. Cerebras WSE-2 칩
Cerebras는 이와같은 성공을 바탕으로 ‘21년 8월 Hotchips에서 WSE-2를 발표하였다. WSE-1에 비해 AI cores를 비롯하여 Fabric Bandwidth까지 2 배 이상 증가하였다.
Cerebras WSE-2와 A100을 비교하면 Chip 사이즈가 56배 증가한데 비해 On-board SRAM은 40GB으로 NVIDIA A100 GPU의 40MB에 비해 1000배가 많다. 또한 on-chip 메모리의 BW는 12,862x, 내부 core간 fabric BW는 45,833x으로 A100보다 훨씬 더 큰 BW를 제공하고 있다.
(*NOTE : A100는 Off-chip 메모리 BW와, A100 칩간 NVLINK BW를 사용하여 직접 비교가 바르지 못할 수 있음은 주의해야 한다.)
2. Cerebras CS-2
Cerebras CS-2는 WSE-2를 1대 탑재한 Super Computer로 기존 Rack의 1/3 크기이며 약 20kW가 필요하다. NVIDIA의 DGX A100 (A100 x8)이 6.5 kW의 파워를 소모하므로 CS-2가 약 3배의 파워를 소모하고 있음을 알 수 있다. 또한 NVIDIA의 DGX A100이 20만$인 반면, CS-2는 200만~300만$라고 예상한다면 CS-2가 DGX A100에 비해 최대 15배이상 비싸다.
동일 가격으로 DGX A100를 15대를 살 수 있으므로 Single CS-2가 A100 120 대의 성능이 나와야 가성비가 있다라고 판정할 수 있다. A100 GPU의 FP16 & BF16 성능이 624 TFLOPS (Sparsity enabled)이므로 GPU 120 대는 총 75 PFLOPS를 제공할 수 있다.
Cerebras가 아직 구체적인 벤치마크 결과를 제시하지 않고 있지만, Single CS-2가 75 PFLOPS 이상의 성능을 제공하는 것으로 불가능할 것으로 보여 가성비에 대한 의구심은 여전히 남아 있다.
3. Hotchips ‘21에 발표된 Cerebras의 새로운 기술
Hotchips ‘21년에서 Cerebras는 4가지 새로운 기술이 대거 발표되었다.
- Cerebras MemoryX: Single CS-2로 120 Trillion 모델을 학습하기 위한Pooled Memory 시스템
- Cerebras SwarmX: CS-2를 192개까지 연결하여 Data Parallelism 성능을 Scaling하기 위한 네트워크 Switch
- Weight Streaming: MemoryX에 있는 Weight를 CS-2으로 streaming으로 보내기 위한 SW 실행 구조
- Streaming Sparse Weights: Sparse computing을 활용하여 연산 응답성과 weight/gradient 이동 시 latency를 줄일 수 있다.
2.1 Memory X
MemoryX는 CS-2의 Off-Chip Memory로 거대한 Pooled Memory 시스템이다. 4TB ~ 2.4PB 용량을 제공하며 DRAM & Flash Hybrid 스토리지로 120 Trillion 파라미터까지 학습을 지원한다.
Note: 인간 두뇌의 스냅스가 100조 개(100 Trillion)으로 파악되어 GPT-4와 같은 AI 모델 또한 100조개의 파라미터를 목표로 개발되고 있다는 루머가 있다. MemoryX의 120 Trillion 파라미터 역시 바로 이 스냅스의 개수인 100 Trillion과 관계가 있을 것으로 생각된다.
Cerebras는 Compute(CS-2)와 Memory(MemoryX)를 decoupling하여 GPU가 Compute와 Memory가 coupling되어 발생하는 scaling 문제와 높은 학습 비용문제를 해결한다.
즉 GPU의 구조가 Compute와 Memory가 결합된 형태라서 Memory 용량 부족문제를 GPU 개수를 늘려 해결할 수 밖에 없기 때문에 GPU 투입 대수에 비해 학습 속도 증가는 매우 낮고 학습 비용은 크게 증가한다.
MemoryX는 Memory를 따로 분리함으로써 여러가지 이점을 기대할 수 있다.
- Model size에 따라 Memory 용량만을 별도로 증가시킬 수 있어 학습 비용을 절약할 수 있다. (하지만 Cerebas의 장비는 매우 고가이지만 이와 같은 아이디어를 채용한다면 얼마든지 더 저렴한 시스템을 구현할 수 있을 것으로 생각한다.)
- 더이상 Model Parallelism을 사용하지 않고, Data Parallelism만으로 학습할 수 있으므로 Model Parallelism 사용시 발생하는 성능 저하 문제가 발생하지 않아 Compute(e.g., CS-2)를 투입에 비례하여 선형적으로 학습속도를 증가시킬 수 있다.
MemoryX는 DeepSpeed ZeRO-Offload의 아이디어를 대부분 차용한 것으로 보인다. ZeRO-Offload가 gradient, optimizer states 등을 CPU 메모리로 offload하여 optimization 연산을 CPU로 실행시킨다면 MemoryX는 ZeRO-Offload가 CPU와 CPU 메모리 역할을 대체한다.
MemoryX는 자체적으로 Optimizer 연산을 하기 위한 Compute 기능을 가지고 있다. (Compute는 x86 CPU일 것으로 예상한다.)
그렇다면 MemoryX의 적절한 메모리용량은 얼마가 될 것인가? 120 Trillion 모델 학습에 필요한 총 메모리양부터 파악해보면 다음과 같다.
- 120 x 10¹² (파라미터) x 4(Bytes) x 5(학습에 필요한 메모리 scaling factor)=2.4 x 10¹⁵ B = (2.4 x 10¹⁵)/(1 x 10¹²) = 2400 TB = 2.4 PB
MemoryX는 DRAM과 Flash의 Hybrid 스토리지이며 4TB~2.4PB까지 제공한다고 하니 2.4PB의 비밀은 위와 같이 120 Trillion과 관계가 있다다. MemoryX가 DRAM & Flash를 Hybrid로 사용하는 이유는 DRAM이 워낙 고가의 디바이스이기 때문에 Flash를 사용하지 않으면 학습 비용을 낮출 수 없기 때문이다.
MemoryX는 Model의 parameter(weight)와 optimizer states들을 저장하는 반면, 학습 데이터셋을 저장하지 않는다. CS-2는 학습 데이터셋을 별도의 학습 데이터베이스로부터 공급받는다.
2.2 Swarm X
Cerebras는 Single CS-2+MemoryX를 이용하여 120T 모델을 Data Parallelism으로 학습이 가능하다라고 발표하였다. 학습 속도를 증가시키려면 다수 CS-2를 MemoryX에 연결하면 된다. 하지만 CS-2와 MemoryX만으로 학습 속도가 선형적으로 증가시킬 수 있을까?
수십개의 CS-2를 MemoryX에 직접 연결하게 되면 MemoryX의 Weight를 CS-2로 broadcast해야 하는데 이는 엄청한 네트워크 트래픽을 발생한다. 또한 CS-2로부터 gradient를 MemoryX가 직접 reduce를 한다면 MemoryX의 compute가 이를 감당하기 어려울 것이다.
따라서 Cerebras는 Swarm X라는 새로운 시스템을 도입하여 1-to-N으로 연결된 MemoryX + CS-2 네트워크의 문제점을 해결한다.
- Weight broadcast: SwarmX는 MemoryX → CS-2로 weight를 전송시 동일한 weight를 복제하여 각 CS-2에 전송하는 역할을 수행한다. MemoryX가 직접 네트워크 트래픽 발생시키지 않고 SwarmX로 offload하여 분산처리하기 때문에 네트워크 성능이 저하되는 것을 방지한다.
- Gradient reduce: CS-2 → MemoryX로 전송되는 gradient를 SwarmX가 중간에서 reduce를 실행하여 MemoryX가 직접 reduce할 때 발생할 수 있는 학습속도 저하를 방지할 수 있다.
Cerebras는 SwarmX를 사용하면 192개의 CS-2까지 성능 저하없이 학습속도를 선형적으로 증가시킬 수 있다라고 주장한다. (SwarmX는 32개의 CS-2를 연결가능하므로 총 6개의 SwarmX를 이용)
2.3 Weight Streaming
CS-2는 on-chip memory의 40GB을 주로 activation용 메모리로 사용하며 weight용 메모리로 사용하지 않는다. (하지만 layer-by-layer로 수행되는 뉴럴네트워크의 특성상 weight용 temp. buffer는 존재할 것으로 생각한다.)
그 이유는 Forward 동작 시 생성된 activation은 Backward 동작 시 즉시 접근하는 것이 필요하므로 latency에 민감하다. 따라서 activation은 CS-2의 on-chip memory에 activation을 저장하는 것이 학습속도를 향상시키기 위해 필요하다.
MemoryX에서 CS-2가 필요한 시점에 weight를 제공하는 것이 중요한데 Cerebras는 이를 weight streaming 방식을 통해 해결하였다.
weight streaming이란 off-chip memory가 on-chip memory처럼 작동하도록 weight를 CS-2로 pipelining하는 방식이다. 이때 사용되는 pipelining 방식은 2가지이다.
- Coarse-grained pipelininig: N-1번째 레이어의 forward pass가 종료되기 전에 N번째 레이어의 weight를 streaming을 시작함
- Fine-grained pipelininig: backward pass의 weight updates를 같은 레이어의 후속 forward pass와 중첩함
Cerebras는 weight streaming 기술을 이용하면 off-chip memory 접근으로 인해 추가되는 latency를 숨길 수 있기 때문에 Wafer가 on-chip memory에 있는 weight를 접근하는 경우와 동일한 성능을 달성 가능하다라고 주장한다.
2.4 Streaming Sparse Weights
Cerebras는 Data Parallelism에 sparse computing을 Large Model을 학습 시간을 크게 줄일 수 있다라고 한다. Sparse computing를 이용하면 MemoryX에서 각 CS-2로 sparse weight만을 스트리밍하고 CS-2는 sparse gradient를 MemoryX로 스트리밍할 수 있다.
Sparse computing를 CS-2, MemoryX에 적용할 경우, 장점은 다음과 같다.
- Sparsity를 이용하여 ‘0’값의 계산을 건너뛸 수 있으므로 optimization 연산의 계산 복잡성을 낮춰 더 빠른 응답을 제공하는데 도움이 된다.
- Sparse Weight와 Sparse Gradient만을 네트워크로 전송하므로 interconnect latency을 줄일 수 있다.
3. 시사점
GPU를 이용하여 Large Model을 학습하기 위해선 Model Parallelism 사용하는 것이 필수적이었다. 하지만 Model Parallelism을 사용하였을 때 성능 저하를 피할 수 없고, GPU의 총 메모리 용량를 증설하기 위해 GPU 대수를 늘려야 하므로 학습에 매우 큰 비용이 지출되어야 한다.
Cerebras는 이와 같은 Model Parallelism의 문제를 해결하기 위해 Compute와 Memory를 decoupling 시키면 Data Parallelism만으로 Large Model 학습이 가능함을 제안하였다. 이로서 기존보다 획기적으로 큰 모델(e.g., 120 Trillion)을 Single CS-2 시스템으로 학습이 가능하며 단순히 CS-2의 대수를 늘려 학습 속도를 선형적으로 scaling할 수 있는 장점을 제공한다.
여기서 Cerebras 시스템에서 추가로 확인해야 될 문제에 대해서 몇가지 짚어보기로 하자.
1)MemoryX은 대용량 메모리를 어떻게 제공할 것인가?
Cerebras는 MemoryX 시스템을 DRAM & Flash hybrid 스토리지로 설명하였다. DRAM은 매우 비싼 자원이기 때문에 DRAM과 Flash을 같이 사용할 수 밖에 없다면 DRAM과 Flash의 비중을 어떻게 결정할 수 있을까? 또한 DRAM과 Flash의 속도 차이를 어떻게 hiding할 수 있을까?
2)MemoryX ←→ CS-2간 네트워크 BW은 과연 충분한 것인가?
120 T 모델을 학습 시 필요한 메모리 용량은2.4 PB이다. 그 중 weight는 480 TB이다. Weight에 FP16를 사용하고 sparsity를 90%를 사용한다고 하면 weight streaming 시 weight를 24TB까지 줄일 수 있다. CS-2의 최대 네트워크 BW는 1.2 Tb(12x 100 Gb optical Ethernet)으로 이론상으로 데이터 전송에만 160초가 걸린다. weight streaming을 사용하여 weight updates와 forward pass를 중첩시키더라도 물리적으로 데이터 전송에 걸리는 속도를 줄이지 않은 이상 120 T 모델의 학습 시간을 단축하는 것으로 쉽지 않을 것으로 보인다.
3)MemoryX의 optimization연산을 위한 컴퓨팅은 충분할 것인가?
120T와 같은 거대 모델을 연산하는데 sparsity를 사용한다라고 하더라도 막대한 컴퓨팅 리소스가 들어갈 것으로 예상할 수 있다. 120T 모델의 optimizer state와 gradient를 이용하여 optimization 연산을 실행할 때 컴퓨팅 파워은 과연 부족하지 않을까?
4)SwarmX가 gradient reduce를 하기 위해선 대용량 메모리가 필요하다. 얼마나 필요할 것인가? gradient reduce가 끝날 때까지 gradient를 유지해야 하는데 2)과 같은 조건이다면 gradient 또한 24TB를 필요로 한다. SwarmX 또한 대용량 메모리가 필요할 것인가?
Cerebas는 CS-2/MemoryX/SwarmX를 통해 Large Model 학습에 새로운 방향을 제시한 것으로 판단된다. 비록 CS-2의 구체적인 성능이 발표되지 않아 시장에서 Cerebras가 성공하리라는 보장은 여전히 미지수이다.
하지만 Compute와 Memory를 decoupling했을 때 Scaling 시 장점을 고려한다면 미래의 Massive AI 학습 시스템은 Compute와 Memory가 분리된 disaggregated 시스템으로 발전할 것으로 전망된다.
4. 레퍼런스
[1] Cerebras unveils Wafer Scale Engine Two (WSE2): 2.6 Trillion Transistors, 100% Yield
[2] Cerebras WSE-2 and CS-2 More than Double
[3] Solving AI’s Memory Bottleneck
[4] Scaling Up and Out: Training Massive Models on Cerebras Systems using Weight Streaming
[5] Inside World’s largest Chip: Introducing Cerebras CS-2
