Z1과 Z1T로 Transformer 추론 전력 낮추기

이 요약은 AI가 원문을 분석해 생성했습니다. 정확한 내용은 원문 기준으로 확인하세요.

TL;DR

Transformer 추론의 전력 부담을 줄이기 위해 Extropic은 확률적 CMOS 칩 Z1과 FPGA를 결합한 sparse 모델 Z1T를 설계했습니다. Z1T는 Z1의 고정된 sparse connectivity에 맞춰 Gated Convolutional Attention과 4-bit weight를 사용하고, Z1에서 처리하기 어려운 연산은 FPGA로 분리합니다. Z1T가 GPT-2와 비슷한 손실에 도달하려면 약 한 자릿수 더 많은 학습 FLOPs가 필요하지만, Z1 연산 자체는 GPU보다 약 3자릿수 높은 에너지 효율을 보여 전체 모델에서는 약 2자릿수 개선을 기대할 수 있습니다. 다만 현재 구성에서는 FPGA가 에너지의 95% 이상을 차지해 병목으로 남으며, 이를 제거하면 GPU 대비 최대 1000배 에너지 효율을 목표로 할 수 있다는 추정이 제시됩니다.

섹션별 상세

01

Transformer 모델의 학습과 추론이 디지털 가속기와 데이터센터 전력 수요를 크게 키우면서, 기존 GPU와 다른 계산 기판을 함께 쓰는 접근이 필요해졌습니다. Extropic은 확률 알고리즘과 확률적 하드웨어를 공동 설계한다는 장기 방향 아래, Z1의 sparse in-memory 연산을 Transformer 계열 모델에 연결했습니다. Z1은 pbit 네트워크에서 확률 상태를 갱신하고 여러 통계의 평균을 고전적 벡터로 출력하므로, 생성 모델에 필요한 확률적 계산을 낮은 전력으로 처리하는 역할을 맡습니다. 이 시도는 Transformer를 그대로 옮기는 작업보다 하드웨어 구조에 맞춰 모델의 연결성과 연산 방식을 함께 바꾸는 첫 단계에 가깝습니다.

02

Z1T는 Z1의 pbit 사이에 존재하는 고정된 sparse connectivity를 전제로 설계됐으며, 각 출력 노드에 네 개의 입력 edge를 연결하고 weight를 4-bit로 제한하는 구성을 사용합니다. Z1에서 직접 처리하기 어려운 연산은 FPGA로 보내고, 여러 Z1 칩 사이에서는 model parallelism을 적용하며 Z1과 FPGA 네트워크 사이에서는 pipeline parallelism을 적용합니다. 모델은 Boltzmann-machine 계열의 아이디어에서 출발해 칩의 샘플을 그대로 출력하지 않고 통계 평균을 벡터로 바꾸며, Gated Convolutional Attention의 sparse·local tanh-linear 연산을 Z1 구조에 맞춥니다. 결과적으로 일반 Transformer의 dense matrix multiplication 중심 구조에서 sparse in-memory computing에 맞춘 구조로 계산 그래프가 이동합니다.

03

Z1T의 scaling 실험은 OpenWebText와 GPT-2 BPE tokenizer를 사용해 모델 크기와 학습 FLOPs를 바꿔가며 validation loss를 측정했습니다. GCA 기반 Z1T가 GPT-2와 비슷한 손실에 도달하려면 약 9.5×10¹⁹ FLOPs가 필요하다는 외삽 결과가 나왔고, 이는 같은 손실을 내는 GPT-2 모델보다 약 한 자릿수 많은 계산량입니다. 반면 Z1에서 수행되는 연산은 GPU보다 약 3자릿수 높은 에너지 효율을 보이며, 전체 Z1T 구성에서는 약 2자릿수의 에너지 효율 개선이 추정됩니다. 따라서 계산량만 비교하면 sparse 모델이 불리하지만, 연산이 실행되는 기판의 비용까지 포함하면 낮은 전력 추론의 여지가 생깁니다.

Z1T GCA 모델의 training FLOPs와 validation loss 관계를 나타낸 로그 스케일 그래프입니다. — Chart모델별 실험점은 body parameter count에 따라 색이 달라지며, GCA compute-optimal frontier가 training FLOPs 증가에 따라 validation loss를 낮춥니다. 그래프의 외삽 endpoint는 GPT-2 small과 비슷한 손실에 도달하려면 약 9.5×10¹⁹ FLOPs가 필요함을 가리킵니다. GPT-2 small과 GPT-2 XL 기준선이 각각 별표로 표시되어 Z1T의 sparse 구조가 같은 손실까지 더 많은 학습 계산을 요구하는 양상을 비교하게 합니다.

04

별도의 GPT-2 스타일 decoder Transformer 실험에서는 connectivity c를 4, 16, 32, 64, 128과 dense로 바꾸고 3×10¹⁴부터 10¹⁸ FLOPs까지 비교했습니다. c는 각 행렬 출력 노드가 참조하는 입력 수와 attention이 볼 수 있는 최근 token 수를 함께 결정하며, 고정된 연결 차수이므로 모델 폭이 커질수록 percentage sparsity는 5%에서 99.8%까지 높아집니다. 같은 FLOP 예산에서 더 높은 compute는 모든 connectivity에서 더 낮은 loss로 이어졌고, 고정된 FLOP 조건에서는 dense 연산이 같은 parameter count의 sparse 연산보다 효율적인 경향을 보였습니다. 그러나 GPU와 Z1에서 sparse multiply-accumulate의 에너지 가격이 다르기 때문에, 모델 선택은 loss뿐 아니라 토큰당 하드웨어 에너지까지 함께 기준으로 삼아야 합니다.

Connectivity c를 4, 16, 32, 64, 128과 dense로 바꾼 여섯 조건에서 validation loss의 학습 궤적을 비교한 격자 그래프입니다. — Chart각 패널의 곡선은 parameter count에 따라 색이 달라지고, 모든 connectivity에서 training FLOPs가 커질수록 더 낮은 validation loss에 도달합니다. c가 작을수록 고정된 입력 연결 수가 적어 sparse 구조가 강해지며, dense 패널은 더 많은 연결을 사용하는 기준선 역할을 합니다. 이 비교는 sparse connectivity를 scaling law에 추가 변수로 넣어 compute와 연결 차수가 손실에 미치는 영향을 분리해서 보는 실험과 연결됩니다.

Connectivity별로 parameter count와 최종 validation loss의 관계를 training FLOPs 예산별로 그린 iso-FLOP 곡선입니다. — Chart각 패널의 U자형 곡선은 고정된 training FLOPs 예산에서 최적 parameter 규모가 존재하며, compute 예산이 커질수록 그 최적점이 더 큰 모델과 더 낮은 loss 쪽으로 이동함을 나타냅니다. dense 조건과 sparse 조건을 비교하면 동일 parameter count에서 dense FLOPs가 더 효율적인 경향이 보이지만, 글은 Z1에서 sparse 연산의 에너지 단가가 GPU보다 낮다는 점을 함께 고려해야 한다고 연결합니다. 따라서 이 그림은 loss 기준의 모델 효율과 실제 하드웨어 에너지 효율이 서로 다른 최적점을 가질 수 있음을 뒷받침합니다.

05

Z1은 programmable Ising model을 실행하는 확률적 sub-threshold CMOS 칩으로, 한 칩에 269,568개 pbit와 2,135,904개의 coupling edge를 갖고 각 노드는 degree 16의 고정 연결을 가집니다. pbit는 이웃 상태를 조건으로 Gibbs sampling을 수행하며, coupling matrix의 0이 아닌 원소가 물리적 parent graph에 의해 제한되므로 dense GPU와 달리 sparse 구조가 칩에 직접 새겨져 있습니다. 내부 update clock은 50 MHz이며, 이 구조는 메모리 계층을 거쳐 dense matrix multiplication을 실행하는 GPU와 다른 에너지·데이터 이동 특성을 만듭니다. 다만 현재 Z1T 파이프라인에서는 FPGA가 전체 에너지의 95% 이상을 소비하므로, Z1 자체의 효율만으로 시스템 전체 효율이 결정되지는 않습니다.

06

현재 추정에서 Z1 연산은 약 0.46 nJ, FPGA 연산은 31.75 nJ, 전체 decoding pass는 32.21 nJ로 제시됩니다. 이 수치는 Z1이 담당하는 sparse 연산보다 FPGA에서 수행하는 비호환 연산과 dense 처리가 훨씬 큰 에너지 비중을 차지함을 보여줍니다. 글은 FPGA 병목을 제거하고 Transformer 계열 연산을 더 많이 Z1 같은 sub-threshold CMOS substrate로 이동하면 GPU 대비 최대 1000배 에너지 효율에 접근할 가능성을 남은 설계 목표로 둡니다. 현재 결과는 완성된 GPU 대체품의 성능 수치라기보다 sparse neural network와 확률적 하드웨어를 함께 설계하기 위한 초기 측정입니다.

Z1 TSU와 XPU·FPGA 사이에 decoding 연산을 분리하고 단계별 에너지 추정치를 누적하는 구조를 표현한 다이어그램입니다. — Diagram왼쪽 Z1 TSU는 pbit starburst와 sparse 연결로 표시되고, 오른쪽 XPU·FPGA는 dense tile 격자로 표시되어 두 기판이 pipeline으로 연결되는 구성을 보여줍니다. 하단에는 Z1 0.46 nJ, FPGA 31.75 nJ, total 32.21 nJ가 표시되어 현재 구성에서 FPGA가 에너지 대부분을 소비한다는 글의 병목 분석을 시각적으로 뒷받침합니다. 상단 block tracker는 decoding pass가 여러 단계로 나뉘어 처리되는 흐름을 나타내며, Z1의 sparse 연산과 FPGA의 비호환 연산을 분리하는 disaggregated inference와 직접 연결됩니다.

용어 해설

확률적 하드웨어(Probabilistic Hardware): — 확률 회로의 상태 자체를 계산 과정에 활용하는 하드웨어입니다. Z1은 pbit 간 결합을 조절해 Ising model의 확률 분포를 샘플링하고, 여러 샘플의 통계적 평균을 고전적 벡터로 변환합니다. 낮은 전력에서 확률적 연산을 수행한다는 점이 Transformer 추론의 에너지 절감과 연결됩니다.
확률 비트(pbit): — 확률적으로 0과 1 또는 두 상태 사이를 오가는 이진 CMOS 회로 단위입니다. Z1의 각 pbit는 이웃 pbit의 상태와 programmable coupling에 따라 다음 상태의 확률을 정하며, 50 MHz 내부 업데이트 클록에서 Gibbs sampling을 수행합니다. 여러 pbit의 묶음은 저정밀 정수와 비슷한 표현을 구성할 수 있습니다.
Gibbs 샘플링(Gibbs Sampling): — 주변 변수의 현재 상태를 조건으로 삼아 한 변수를 차례로 갱신하는 확률적 샘플링 방법입니다. Z1은 각 pbit가 이웃 상태에 조건부인 새 값을 메모리 내부에서 계산하도록 하며, 이를 통해 programmable Ising model의 상태 분포를 생성합니다. 글에서는 이 구조가 확률적 신경 연산을 하드웨어에 직접 배치하는 기반으로 쓰입니다.
게이트형 합성곱 어텐션(Gated Convolutional Attention): — Z1의 sparse하고 국소적인 연결 구조에 맞춰 설계한 attention 변형입니다. tanh와 선형 연산을 사용해 고정된 물리적 connectivity 안에서 Transformer와 유사한 표현을 만들며, Z1T의 핵심 구조로 적용됩니다. 글의 scaling 실험에서는 GCA 기반 Z1T 모델의 손실과 학습 FLOPs 관계를 GPT-2 기준선과 비교합니다.
분리형 추론(Disaggregated Inference): — 하나의 모델 추론 그래프를 여러 종류의 프로세서와 단계로 나눠 실행하는 방식입니다. Z1T는 Z1에서 sparse 연산을 처리하고 FPGA에서 Z1에 맞지 않는 연산을 수행하며, 두 장치 사이에 pipeline parallelism과 model parallelism을 결합합니다. 이 구조는 기존 GPU를 모두 대체하기보다 Z1과 dense matmul 가속기를 함께 사용해 토큰당 에너지 비용을 낮추려는 접근입니다.