Tenstorrent芯片架构浅谈
时间：2023-05-07 10:37:00
1. 背景
近年来，在市场上AI芯片层出不穷，初创公司和科技巨头都在积极推出AI芯片，从最常见的CPU、GPU到各类Processing Unit（TPU、NPU、APU、DPU等)，可谓百花齐放。根本原因是，由于相关算法、模型和应用场景的深度学习处于快速迭代发展阶段，尚未完全收敛，因此相应的计算硬件基础自然需要不断更新以适应上层应用。芯片公司Tenstorrent芯片架构独特，本文试图探索。
Tenstorrent它成立于2016年，总部设在加拿大多伦多AI自2021年传奇芯片设计大神以来，芯片公司在成立初期就很少受到业界的关注。Jim Keller以CTO加入身份后，公司逐渐公布了其独特的芯片架构设计理念，越来越受到业界的关注。本文结合了一些公开演讲和相关演讲paper解读其芯片架构。
2. 产品
Tenstorrent共设计了三个芯片，其中包括Jawbridge是小型测试芯片，Grayskull和Wormhole对外商用芯片，可覆盖训练和推理场景。
| Products | Jawbridge | Grayskull | Wormhole |
|---|---|---|---|
| Manufactured | 2019 | 2020 | 2021 |
| IC Process | GF 14nm | GF 12nm | GF 12nm |
| Peak Performance | 1TFLOPS at FP16, 4TOPS at INT8 | 92TFLOPS at FP16, 368TOPS at INT8 | 110TFLOPS at FP16, 430TOPS at INT8 |
| Compute cores | 6 | 120 | 80 |
| SRAM | 6MB total - 1MB/core | 120MB total - 1MB/core | 120MB total - 1.5MB/core |
| DRAM Capacity | unknown | 8GB | 12GB |
| DRAM Bandwidth | unknown | 100GB/s | 384GB/s |
| I/O Interface | 1ch. LPDDR4, PCIe Gen4x4 | 8ch. LPDDR4, PCIe Gen4x16 | 16 ports of 100G Ethernet, 6ch. GDDR6, PCIe Gen4x16 |
| Scale-out Bandwidth | unknown | 192GB/s for NoC | 400GB/s for Ethernet |
| Board Power(TDP) | 1.5W | 65W / 75W | 150W |
以最新的Wormhole以芯片为例，主要包括：
计算核心Tensix core：
5个标量RISC-V CPU：发射调度指令和数据
SRAM：作为private memory，用于保存本地计算结果
Packet Manager数据包管理器：包括Tensor数据传输和数据路由功能
Packet Compute Engine数据包计算引擎：INT8算力3TOPS，FP16算力0.75TFLOPS
DRAM：存储在片上，可被Tensix全局访问
Ethernet接口:用于芯片间互联，方便scale out
PCIE接口：连接Host，与Host传输数据
Network on Chip(NoC)：数据互联接口，在Ethernet中也含有NoC
System Manager系统管理单元：4个RISC-V CPU，负责中断、时钟、温度控制等基本功能，以及数据路由功能
3. 技术特点
Tenstorrent的芯片架构设计目标是解决模型在训练或推理过程中无法有效灵活扩展的问题（scale out）的问题，提出两个核心技术点：
3.1 摒弃传统的核间共享内存架构，采用Multicore Private Memory Model
这样，核间和芯片间的数据通信模式在软件栈上是一致的，大大降低了模型分布式培训或推理的开发难度。从用户的角度来看，硬件是一个Device。由于多核间的数据通信不依赖于共享内存，因此芯片的重点是加强硬件和软件层面的数据通信能力。在硬件层自研芯片上进行互联Network on Chip(NoC), NoC是2D双向环路结构，支持芯片所有核心和片外芯片之间的通信。在最新的Wormhole在架构芯片中，增加了Ethernet接口专门用于芯片间互联。为实现芯片与芯片间软件层面通信接口的一致性，Tenstorrent基于软件层的设计Pipe由Graph Compiler在当前模型下生成最佳数据路由策略。对比各大AI芯片公司官方公告将大量介绍计算部件，如各种DSA架构，Tenstorrent芯片架构介绍文章更多地描述了非计算硬件部件。从这里也可以看出，数据通信是一等公民（This communication oriented architecture realizes our vision of converged networking and accelerated AI compute on a single device.），片内和片间高效的数据通信能力是Tenstorrent一直强调的重点。
3.2 动态执行
运行时数据压缩：Tenstorrent芯片中有专门用于数据压缩和解压缩的硬件部件，称为Packet Manager Engine，个人认为Packet Manager Engine引入是芯片的重要特性。为了最大限度地利用计算资源，最小限度地利用内存，数据在计算单元外以压缩的方式存储和流动，从而设计更小的容量private memory，芯片内多核间和芯片间的数据通信量也相应降低，从而从整体上获得更高的性能功耗比。另外，Packet Manager还可以处理reshape/flatten等tensor形状变换操作，操作可与Compute Engine并行执行，可以可以overlap。
实施条件：特别是稀疏门控专家混合模型（Sparsely-Gated Mixture of Experts）应用效果显著。这种模型只通过门网络激活模型的部分结构，因此是一种显著增加模型容量和能力而不成比例增加计算量的方法。复杂的逻辑控制是正确的CPU它很容易处理，但对AI就芯片而言，大量资源用于计算部件，逻辑控制相的部件功能较弱，因此在编程时一般是需要尽量避免出现大量逻辑控制语句。在Tenstorrent芯片架构中，片上的逻辑控制单元和计算单元均可以高效运行，避免了CPU fallback的问题，此处也是该芯片的重要特色。
稀疏计算：除了支持常规的对权重进行稀疏化之外，还支持对激活值进行分块稀疏，进而降低计算量。
动态混合精度：可以在运行时或AOT阶段设置每个算子的计算精度。该技术在主流深度学习框架中都有实现，因此个人认为这项技术并无特色。
Tenstorrent通过软硬协同设计方式，将数据并行和模型并行的部分功能实现下沉到硬件层，有效解决了横向扩展问题，这样就可以替代当前主流深度学习框架在分布式实现方面的大量编码工作，进而降低了深度学习框架的开发和使用门槛。同时，硬件的变化并没有降低软件栈的通用性，其软件栈支持PyTorch等主流框架。另一方面，芯片具有高度模块化，多个芯片可通过标准以太网端口连接在一起，进而扩展成大型AI网络。由于芯片内已集成NoC，因此这种扩展并不需要额外的交换机，因此扩展灵活度很高。
对比GPU，Tenstorrent官网提到其芯片编程更友好、扩展性更好和更擅长处理稀疏计算和控制流计算场景。（Compared to GPUs, our processors are easier to program, scale better, and excel at handling run-time sparsity and conditional computation.）
4. Tenstorrent官网示例分析
由于无法拿到芯片进行实际测试，网上也几乎没有第三方对Tenstorrent芯片的评测数据。因此我们只能查看Tenstorrent官网示例了解具体如何使用其芯片。Tenstorrent官网示例包含Grayskull板卡、Turnkey工作站/服务器和DevCloud云资源三种，Grayskull板卡的使用说明相对详细。
Grayskull板卡：基于Ubuntu20.04，按常规流程插入该板卡，并加电，然后下载并安装驱动，不过目前驱动下载链接尚未开放，只好先跳过该步骤。官方提供了Dockerfile用于构建编译运行环境。在/tenstorrent/releases/tests下有三个文件：
hello_tt.yaml：配置文件，包括硬件架构信息、编译器配置信息、运行时配置信息等hello_tt.py：模型定义文件，利用 Torch 定义一个模型，并设置该模型的并行策略set_scalar_input.py：生成模型输入数据
hello_tt.yaml文件内容如下：
test:
activations-config:
arguments:
input-shapes: &input-shapes [[1, 1]]
file-name: releases/tests/activations/set_scalar_input.py # 生成输入数据
arch-config: releases/tests/config/grayskull_arch.yaml # Grayskull 芯片架构配置文件
comparison-config: releases/tests/config/grayskull_comparison.yaml
compiler-config:
- releases/tests/config/grayskull_compiler.yaml # 编译器配置文件
- num-fidelity-phases-default: 4
graph-config:
file-name: releases/models/pytorch/hello-tt.py # 模型定义文件
onnx: false
runtime-config:
- releases/tests/config/grayskull_runtime.yaml # 运行时配置文件
- l1-format: FP16 # 运行时精度hello_tt.py文件内容如下，主要是定义了一个2层网络，计算过程为 (((0.8 * 1.58) - 0.14) * 2.45) - 0.11 = 2.6438。
# hello world inspired from
# https://nestedsoftware.com/2019/08/15/pytorch-hello-world-37mo.156165.html
# just one step of the forward pass
# one input, one hidden, and one output
import torch
import torch.nn as nn
import torch.nn.functional as F
import sage.api as tt_api # Import Tenstorrent Sage compiler's public API
def get_model(par_strategy="RowParallel", core_count=4,):
# 1. 指定计算核分区方式，此处指定为行并行，即模型中的算子会在行方向的 tensix cores上执行
# 2. 指定使用 4 个 tensix cores
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
# 设置隐含层数据
self.hidden_layer = nn.Linear(1, 1)
self.hidden_layer.weight = torch.nn.Parameter(torch.tensor([[1.58]]))
self.hidden_layer.bias = torch.nn.Parameter(torch.tensor([-0.14]))
# 设置输出层数据
self.output_layer = nn.Linear(1, 1)
self.output_layer.weight = torch.nn.Parameter(torch.tensor([[2.45]]))
self.output_layer.bias = torch.nn.Parameter(torch.tensor([-0.11]))
def forward(self, x):
# 构建前向计算图
x = F.relu(self.hidden_layer(x))
x = F.relu(self.output_layer(x))
print('output', x)
return x
model = Net().eval()
# Set parallelization strategy
# 指定模型运行在哪些 tensix cores上，此处为在(0,0)到(3,0)这 4 个行方向的 tensix cores 上执行计算，假定一个 core 计算一层
tt_api.set_parallelization_strat(model, cores=((0, 0), (core_count - 1, 0)), strategy=par_strategy)
return model从如上代码中可以看出，用户可以灵活设定整个模型使用的计算核数以及并行方式。示例模型共4个节点，分配了芯片中列方向上的4个连续计算核心来执行（如下图所示），剩余计算核心可以并发执行其他模型的计算任务。
我们可以分析更常见的模型结构，如下图，左侧是多层级联结构，右侧是多分支结构。由于其在定义每个算子时，除了定义算法实现外，还需要提供算子的性能信息，如内存占用、计算资源占用、与临近节点的连接关系等信息。通过这些信息，编译器就可以生成一个当前最优的算子到硬件的映射图。例如，下图(a)中每个pipeline对应的计算核心是不同的；下图(b)中，由于模型中的4个分支是无数据依赖关系的，因此这4个节点可以并行执行。
tether run -d grayskull -i hello_tt.yaml
tether 是一个运行集成工具，可以用来编译和运行模型。 更多细节可以查阅
/tenstorrent/releases/UserGuide.md
5. 疑惑
受限于当前有限公开的技术资料，我们暂时无法深层次剖析Tenstorrent芯片架构。目前关于该芯片的设计架构，笔者还存在一些疑惑：
在Tenstorrent芯片构成的网络中，存在4种拓扑结构，比如芯片内多核间、芯片间、服务器间和机架间。这4种拓扑结构对应的通信带宽和网络延迟肯定各不相同。因此需要软件栈中的Graph Compiler生成一个多种约束条件下的cost model，基于该cost model生成一个在当前运行环境下最优的算子与硬件的映射关系图。笔者认为Graph Compiler的实现颇为重要，但鲜有公开资料介绍这方面信息。官方表示会开源其软件栈，但未给出具体时间点。
关于开发者生态，当前深度学习相关软件大都是将NV GPU作为第一优先级进行适配支持。Tenstorrent芯片的软件栈相对独立，虽然可以支持PyTorch、TensorFlow等框架，但如何推动这些深度学习框架在Tenstorrent芯片上适配开发，并针对性优化，路途漫漫。
6. 总结
深度学习AI芯片设计需要兼顾性能、功耗、可编程性、可扩展性和成本等多方面的因素，因此如果对比单点指标，无法科学评价芯片的整体水平。我们分析了Tenstorrent芯片产品和技术特点，并尝试从官网示例教程中寻找更多细节信息，最后给出三个关心且存疑的问题，希望与各位业界同仁一同探讨。
最后，依然期待大家支持和关注Adlik的Github仓库哦。
Adlik Github仓库
参考
[1] L. Bajić and J. Vasiljević, “Compute Substrate for Software 2.0”, Hot Chips 32, 2020.
[2] D. Ignjatović, Daniel W. Bailey, Ljubisa Bajić, “The Wormhole AI Training Processor”, ISSCC 2022 Session 21: Highlighted Chip Releases: Digital/ML
[3] L. Bajić, “Finding the Right Compute Substrate for Software 2.0”, Linley Spring Processor Conference, 2020.
[4] D. Ignjatović and D. Capalija, “Scale-out First Microarchitecture for Efficient AI Training”, Linley Spring Processor Conference, 2021.
[5] Dylan Patel, https://semianalysis.substack.com/p/tenstorrent-wormhole-analysis-a-scale?s=r
[6] OneFlow, Tenstorrent虫洞分析：挑战英伟达的新玩家？
