/ AI, ARCHITECTURE, SYSTEM

Tenstorrent 논문 요약

기존의 Shared Memory 방식과는 전혀 다른 설계 구조로 만들어진 Tenstorrent의 AI 가속기가 어떤 물건인지 대략적으로 알아볼 수 있는 논문 Compute Substrate for Software 2.0 을 읽고 핵심 내용을 요약해 보겠습니다.

Introduction

컴퓨팅 능력이 무어의 법칙에 따라 계속 발전하였지만 그동안 오직 단조적인 shared memory model에 의존해 왔습니다. Tenstorrent에서는 인간의 뇌와 유사하게 “불필요한 연산을 포기”하고 “오직 입력과 관련 있는 것만 계산”할 수 있는 두 가지 접근 방식인 dynamic computation 및 massive scaleout을 디자인했다고 합니다. 구체적으로는 shared memory model에서 multicore private memory model로의 패러다임 시프트를 통해, 또한 push-based data movement의 구축을 통해 필요한 요구를 충족시킬 수 있는 하드웨어 - 소프트웨어 구조 설계를 해냈다고 합니다.

논문은 Hardware, Software, Full Stack Performance Optimizations, Dynamic Execution, Results, Conclusion 섹션으로 구성되어 있습니다.

논문이 작성된 당시 출시된 Tenstorrent 칩은 Wormhole이 최신 세대였고, 시장에 널리 풀린 Tenstorrnet 디바이스가 Grayskull 세대였기 때문에 Hardware 섹션에서는 이 Grayskull 칩을 중심으로 하드웨어 디테일을 설명하고 있습니다. Software 섹션에서는 소프트웨어 스택을, Full Stack Performance Optimizations 섹션에서는 하드웨어와 소프트웨어의 연계로 이루어진 풀 스택 성능 최적화를 분석하고 있습니다. 또 Dynamic Execution 섹션에서는 Tenstorrnet 아키텍쳐가 가져다 주는 다양한 동적 실행 방식에 대해 기술합니다. Results에서는 성능 평가 결과가, Conclusion 에서는 결론이 적혀 있습니다.

Hardware

Grayskull 칩은 칩 내에 120개의 컴퓨팅 코어(=Tensix 코어)가 들어 있으며, 각 코어가 독립적으로 동작한다고 합니다. 즉 각각의 코어가 고유한 instruction queue를 갖고 있으며 자신의 페이스대로 처리를 수행합니다. 그리고 각 코어간의 데이터 전송을 위해 Network-on-Chip (NoC) 기술이 사용되었다고 합니다. NoC는 2D의 양방향 torus topology로 구성되어 있으며 노드당 bandwidth가 192GBps라고 합니다. NoC는 각 코어들을 연결할 뿐만 아니라 off-chip communication 및 memory controller block들 까지 연결해 주기 때문에 PCIe와 DRAM에 120개 코어가 각각 접근할 수 있다고 합니다.

2D Torus topology

2D Torus

위 그림과 같이 각 노드가 그리드 구조로 연결된 것에 더해, 끝에 있는 노드가 반대편 끝에 있는 노드와도 연결되도록 하여 순환이 가능한 형태로 된 것이 Torus 구조입니다.

이러한 통신 기반 아키텍쳐가 단일 디바이스에서의 AI 컴퓨팅과 converged networking이라는 비전을 실현시킨다고 합니다. 예를 들어 Wormhole에는 독점 통신 링크가 많이 탑재되어 있기 때문에 중앙 호스트 CPU나 이더넷 스위치의 도움 없이 직접 Wormhole 디바이스를 연결하기만 함으로써 many-device scale out이 쉽게 가능한 것입니다.

Tensix Core

Tenstorrent 아키텍쳐는 패킷 기반으로 동작합니다. NoC를 타고 메모리 간에 이동하는 데이터의 단위는 패킷입니다. 특히 텐서를 ‘미니 텐서’들로 작게 나눠서 작은 크기의 패킷으로 보내는 것이 이후 언급할 Software 중 Graph Compiler 의 핵심 기능입니다. 또한 Tensix 코어(특히 내부의 Packet Compute Engine)의 계산 역시 패킷에 대해 직접적으로 실행됩니다.

단일 Tensix 코어는 packer processor, packet manager, SRAM, 5개의 RISC 프로세서로 구성됩니다. RISC 프로세서들은 런타임 소프트웨어 (펌웨어) 를 실행하는데, 이 런타임 소프트웨어가 packet manager와 packet processor에게 instruction들을 디스패치합니다. Grayskull 기준 SRAM의 용량은 1MB이며, 384GBps bandwidth를 갖고 있습니다. 물론 이후의 세대에서는 이 SRAM 용량이 더욱 늘어납니다.

여기서 알 수 있는 것은 RISC-V 코어는 tensor 및 벡터 기반 연산에 직접 관여하지는 않는다는 것입니다. 여기서 RISC-V 코어는 철저히 control plane만 담당하고 있습니다. 물론 그렇다고 해서 RISC-V 코어가 성능에 기여하는 영향이 작은 것은 아닙니다. 이후 언급할 런타임 소프트웨어(펌웨어)의 실행을 RISC-V 코어가 담당하고 있기 때문이죠. 한편, Tensor 및 벡터(를 잘개 쪼갠 ‘패킷’)에 대한 계산은 packet compute engine이 담당하게 됩니다. 그렇다면 이 packet compute engine의 ISA는 무엇이 채택되었을까요?: Gemini에 의하면,

Tenstorrent 고유의 최적화된 명령어 세트를 가집니다. 이 명령어들은 딥러닝 워크로드에 특화되어 있으며, 텐서 및 벡터 연산을 매우 효율적으로 수행하도록 설계되었습니다. 외부에는 이 고유 ISA에 대한 상세한 문서가 많이 공개되어 있지는 않지만, 내부적으로는 컴파일러가 PyTorch/TensorFlow 연산을 이 저수준의 고유 명령어로 변환하여 발행합니다.

라고 합니다.

Packet Compute Engine

Packet compute engine는 SIMD 기반 매트릭스(텐서) 및 벡터 엔진이며 높은 유연성과 programmability를 갖고 있습니다. Packet compute engine은 연관된 RISC 코어를 통해 software-programmable 하다고 합니다. 그리고 이 RISC 코어들이 앞서 언급했듯이 펌웨어 및 C언어로 작성된 커널을 돌리며 matrix, vector 연산 명령어를 compute engine으로 보냅니다. 여기서 많은 수의 PyTorch 및 TensorFlow 명령어가 지원된다고 합니다 (TT-forge 등을 통한 컴파일 과정이 필요하지만).

Packet compute engine에서는 기본적으로 멀티코어 시스템 전반에 걸친 실행의 큰 흐름이 보이지 않습니다. packet compute engine이 수행하는 연산은 주로 packet manager에 의해 결정됩니다. Packet manager에서 보낸 패킷을 계산하고 다시 packet manager로 보내 데이터 전송으로 이어집니다.

Packet Manager Engine

Packet manager engine은 data transfer engine, router, tensor manipulation engine으로 구성되어 있습니다. 각 구성 요소에 대해 간단히 알아보겠습니다.

  • Data transfer engine: On-chip SRAM 및 off-chip memory, I/O 를 포함한 모든 compute engine 간의 모든 데이터 이동과 synchronization을 관리합니다.

Packet manager와 compute engine은 각각 컴파일러로부터 자신에게 고유한 instruction queue를 받고 동시에 실행합니다. Packet manager 덕분에 데이터 이동과 멀티코어 동기화에 대한 부담을 compute engine은 덜게 됩니다.

  • Router: NoC를 따라 실제로 패킷을 이동시킵니다. ordering을 보장하고, flow control 및 backpressure를 관리합니다. 또한 deadlock-free하게 작동합니다. 병렬화된 AI 워크로드에 최적화되어 있어 효율적인 멀티캐스트 및 수집 능력(gather capabilities)을 갖고 있습니다.

Backpressure란?

노드의 수신 버퍼가 꽉 차서 패킷 손실이 발생하는 것을 방지하기 위해, 받는 쪽이 보내는 쪽에게 데이터 전송을 멈추거나 속도를 늦춰달라는 신호를 보냅니다. 이것이 바로 backpressure signal입니다. 이 신호의 송수신 및 동작 처리를 tenstorrent에서는 packet manager engine 내 router가 담당한다는 것을 알 수 있습니다.

Gather Capabilities

여러 분산된 source에 있는 데이터를 하나의 특정 목적지(sink)로 효율적으로 모으는 능력을 의미합니다.

AI 워크로드의 병렬 분산 처리 중 부분 결과(partial results)들이 하나의 지점으로 모여서 최종 값을 계산하거나 결과를 취합할 필요가 있는데 router가 이러한 작업을 효율적으로 수행한다는 뜻입니다.

  • Tensor manipulation engine: Dynamic packet compression을 수행할 수 있습니다. 이 압축 과정을 통해 memory footprint를 줄여 빠른 전송과 데이터 locality를 증가시킵니다. 또한 tensor manipulation instruction이 이 엔진에 의해 수행됩니다. Optimization of Tensor Manipulation Instructions 섹션에서 자세히 설명됩니다.

Software

소프트웨어는 3가지의 메인 요소로 구성됩니다.

  1. ML 프레임워크 통합 및 플러그인 소프트웨어

  2. RISC 프로세서에서 작동하는 런타임 소프트웨어

  3. ahead-of-time 그래프 컴파일러

Framework Integration

Tenstorrnet의 컴파일러와 런타임은 PyTorch에 natively integrated 되었다고 하며, inference와 training flow 둘 다를 지원한다고 합니다. 유저는 단일 디바이스 또는 여러 디바이스의 클러스터 환경에서 이러한 AI 워크로드를 실행할 수 있지만, 추상화 덕분에 하드웨어는 유저에게 단일 디바이스로 보여지게 됩니다. Multidevice scheduling 및 parallelization은 소프트웨어 스택에서 알아서 수행됩니다. 소프트웨어 스택은 ONNX 네트워크 (ML 모델을 표현하는 표준 형식) 또한 실행할 수 있어 프레임워크에서 ONNX 포맷으로 export 하는 것도 가능합니다.

요약: Tenstorrent 컴파일러 및 런타임이 PyTorch랑 ONNX 지원한다.

Graph Compiler

그래프 컴파일러는 프론트엔드, optimizer, 백엔드의 3요소로 구성됩니다. 먼저 프론트엔드의 주된 역할은 넓은 범위의 instruction들을 하드웨어에서 지원하는 low-level의 최적화된 instruction들로 변환하는 것입니다. 이후 optimizer에 의해 compute, data locality, data movement의 균형을 맞춰 성능이 최대화됩니다. 마지막으로 백엔드는 컴파일된 그래프를 각 코어의 instruction queue로 변환해 줍니다.

딥러닝 모델에서의 그래프

노드: 딥러닝 모델의 개별 연산자(Operations)를 나타냅니다.

엣지: 연산들 사이에서 데이터(텐서)가 흐르는 방향을 나타냅니다.

한 마디로, PyTorch 등의 딥러닝 프레임워크에서 정의된 딥러닝 모델의 추상적인 연산 흐름과 데이터 종속정을 나타내는 구조가 그래프에 해당합니다.

그림5FIGURE 5. Tensor manipulation instructions executed on Grayskull and a GPU. (a) Application graph received from PyTorch. (b) Execution trace on a Grayskull Tensix core. (c) Execution trace on a GPU.

위 Figure 5.의 (a)가 바로 딥러닝 모델의 그래프에 해당합니다.

Packet manager들과 코어들을 연결하는 NoC는 소프트웨어에 visible하며, 컴파일러에 의해 데이터 이동 및 동기화가 명시적으로 이루어집니다. 데이터 이동을 스케줄링하기 위해, 컴파일러는 각 텐서를 미니 텐서로 나누어 패킷화하며, 각 미니 텐서는 패킷 헤더가 붙습니다. 각 패킷 헤더는 고유의 패킷 ID를 가지며, 모든 데이터는 이 패킷 ID를 통해 참조됩니다. 헤더는 routing information도 포함하고 있기에 packet manager의 router가 NoC를 통해 코어 간 원하는 데이터 이동을 수행할 수 있도록 합니다.

요약: Graph compiler가 텐서를 패킷화하며 instruction 변환도 수행한다.

Runtime Software

런타임 소프트웨어(펌웨어)는 모든 코어에 존재하는 RISC 프로세서 안에서 동시에 실행됩니다. 컴파일된 실행 파일은 각 코어의 packet processor와 packet manager에서 쓰이는 instruction queue를 포함하고 있습니다. 이 queue들을 런타임 소프트웨어가 관리하며, instruction들을 packet compute 및 packet manager에 디스패치합니다.

패킷들이 들어오는 버퍼는 런타임에 동적으로 할당 및 해제됩니다. 런타임 소프트웨어는 packet manager와 긴밀히 협력하여 할당된 버퍼에 패킷을 저장하도록 합니다. 런타임은 또한 storage target 결정을 통제하기도 합니다. 예를 들어 코어의 local SRAM에 맞지 않는 버퍼를 remote SRAM 또는 off-chip 메모리로 내보내기도 합니다.

RISC-V ISA를 사용하고 있으므로 conditional statement의 실행이 가능하며 instruction queue를 해석해서 control flow decision을 반영하기 위해 instruction queue 내부의 특정 instruction으로 jump를 실행하는 것 역시 가능합니다.

Full Stack Performance Optimizations

Optimization of Data Transfers: The Push-Based Model

Pull-based 데이터 전송 모델을 도입한 전통적인 멀티코어 디바이스들과 다르게 Tenstorrent는 push-based 모델을 사용합니다. 프로듀서 코어 (output buffer를 생성하는 코어)가 consumer 코어를 인지하고 능동적으로 버퍼를 consumer 코어에 복사합니다. 이러한 접근 방식이 consumer core의 idle time(리퀘스트를 보내고 기다리는 시간)을 최소화할 수 있게 합니다. 앞서 언급되었던 data transfer engine이 push-based data transfer model에 필요한 모든 flow control을 실행합니다. 이는 Router가 수행하는 물리/데이터 링크 계층 레벨의 flow control과는 별개로, 더 상위 레벨에서 데이터의 전송 전략 및 스케줄링을 관리합니다. Data transfer engine은 또한 그래프 컴파일러가 생성한 instruction queue를 받아 producer-consumer 연결성 정보를 확인합니다.

Optimization of Tensor Manipulation Instructions

DNN의 주된 연산은 크게 두 메인 카테고리로 나눌 수 있는데, 1. 수학적 연산 (math instructions) 및 2. Tensor Manipulation (TM) 연산입니다. TM 연산은 tensor 내의 데이터는 변경하지 않고, tensor의 모양을 바꾸는, 즉 데이터들의 위치를 바꾸는 연산입니다. 일반적인 TM 연산은 reshape, transpose, flatten, permute가 있습니다. Tenstorrent는 기존 GPU와 달리 두 연산을 병렬로, 파이프라이닝하여 수행이 가능하도록 합니다. 수학적 연산은 packet compute engine이 수행하고, TM 연산은 packet manager engine 내부에 있는 tensor manipulation engine이 수행하도록 하는 것이죠. Figure 5(b)를 보면 이 과정이 잘 나타나 있습니다.

Tensor manipulation engine은 programmable 합니다. 컴파일러로부터 고유한 instruction queue를 받아 옵니다. TM 연산은 두 가지 방법을 조합하여 수행됩니다. 첫 번째로, tensor manipulation engine은 패킷을 로드하고 그 자리에서 reshuffle을 수행할 작은 저장 공간을 갖고 있습니다. 두 번째로는, 복잡한 memory read/write 패턴을 실행할 수 있습니다. 이 두 가지를 조합하면, 패킷은 packet compute engine 외부에서 스트리밍되어 local SRAM에 쓰여지므로 TM 연산은 inline으로 구현될 수 있습니다.

위와 같이 컴파일러로부터 각 코어가 자신만의 instruction queue를 받아 자신만의 pace로 처리를 수행하므로 tenstorrent는 동시성(concurrency)을 매우 크게 확보할 수 있는 것이 특징이라고 합니다.

Dynamic Execution

런타임에 네트워크의 computational complexity를 높이는 다양한 접근 방식을 dynamic execution이라고 뭉뚱그려 설명하고 있습니다. 어떤 방식은 Mixture-of-Experts (MoE) 와 같이 네트워크의 구조 그 자체 내에서 표현될 수 있고, 다른 방식은 network execution을 augment하기 위해 사용될 수 있습니다. Tenstorrent 아키텍처가 제공하는 4가지 방식이 아래에서 소개됩니다.

Block Sparsity

Block Sparsity in Deep Neural Network

원소 단위가 아니라, 텐서 내의 특정 ‘블록(특정 크기를 가진 sub-region)’들이 통째로 0이거나 거의 0이라 계산에서 제외되어도 무방한 것들을 뜻합니다.

네트워크에서 연산에 사용되는 tensor들은 각각의 서로 다른 sparsity를 갖고 있습니다. 어떤 모델은 training된 parameter의 sparsity를 이용하거나, model parameter 내 block sparsity를 이용합니다. 하지만, 이러한 접근 방식은 런타임에서 야기되거나 내재적일 수 있는 activation의 sparsity는 간과하고 있습니다. Tenstorrent의 아키텍쳐는 이러한 activation의 block sparsity를 지원한다고 합니다.

Dynamic Precision

각 패킷당 수의 precision을 세세한 level로 설정할 수 있다고 합니다. 이는 컴파일러 및 런타임 둘 다에 의해 명시 가능하다고 합니다.

Runtime Compression

텐서들은 런타임에 압축 및 압축 해제되며, parameter들은 컴파일 타임에 압축되고 런타임에 하드웨어에 의해 압축 해제될 수 있다고 합니다. 비슷하게, 런타임 동안, math layer의 output activation 역시 생성 후에 inline으로 압축이 가능하다고 합니다. 이런 접근방식은 memory footprint와 전력소모를 줄이고 data transfer 속도를 높일 수 있다고 합니다.

Conditional Computation

Host CPU 대신에 Tenstorrent의 아키텍처가 conditional computation을 수행함으로써 불필요한 데이터 전송을 줄일 수 있다고 합니다.

Results

Dynamic Execution 섹션에서 설명된 최적화 기법을 통해 불필요한 연산을 줄여 activation의 sparsity가 높은 경우 BERT 레이어에서 4배의 속도 향상을 이뤄냈다고 합니다. 또한 Dynamic Precision덕분에 대부분의 모델 parameter들을 chip에 맞출 수 있었다고 합니다.

Conclusion

호스트 CPU fallback의 필요성을 최소화하는 데에 성공했기 때문에, Tenstorrent의 아키텍처는 neural net뿐만 아니라 전통적인 프로그램에도 잘 맞는 구조라고 합니다.

  • 단일 칩과 멀티 칩 구조 모두에 똑같은 프로그래밍 모델이 사용된다고 합니다.

  • 다양한 동적 및 conditional execution techniques 덕분이라고 합니다.

  • 커널 및 런타임 시스템(펌웨어)가 표준 C언어로 작성되어 fully programmable하다고 합니다.

  • Scaleout capability and flexibility: 표준 Ethernet을 통해 scale out이 가능하며 Tenstorrent 내 NoC 구조에 seamless하게 통합이 가능하다고 합니다.

참고문헌

J. Vasiljevic et al., “Compute Substrate for Software 2.0,” in IEEE Micro, vol. 41, no. 2, pp. 50-55, 1 March-April 2021, doi: 10.1109/MM.2021.3061912. keywords: {Artificial intelligence;Software;Computer architecture;Data transfer;Runtime;Random access memory;Computational modeling;Substrates;machine learning;compilers;multi-core architecture},