TT-NN和TT-Metalium:Tenstorrent的神经网络库和低级编程模型

tt-metal

TT-NN:强大的神经网络操作库

TT-NN是Tenstorrent公司开发的一个Python和C++神经网络操作库,专门为其AI加速芯片设计。该库提供了丰富的神经网络操作API,让开发者可以方便地在Tenstorrent硬件上构建和运行各种深度学习模型。

支持多种模型

TT-NN支持在Tenstorrent的Grayskull(GS)和Wormhole(WH)系列芯片上运行多种流行的深度学习模型,包括:

ResNet-50图像分类模型
BERT-Large自然语言处理模型
Falcon 7B大语言模型
Mistral 7B大语言模型
LLaMA系列大语言模型
Stable Diffusion图像生成模型

这些模型在Tenstorrent硬件上都展现出了优异的性能。例如,ResNet-50在单个GS芯片上可以达到5,600 fps的端到端吞吐量,BERT-Large可以达到370 sentences/s。在WH芯片上,性能进一步提升,ResNet-50可达4,100 fps,Falcon 7B的token生成速度可达534 tokens/s。

易用的API

TT-NN提供了简洁易用的Python API,让开发者可以轻松地使用Tenstorrent硬件加速张量计算。以下是一个简单的示例:

import ttnn
import torch

with ttnn.manage_device(device_id=0) as device:
   a = torch.ones((5, 7))
   b = torch.ones((1, 7))

   a = ttnn.from_torch(a, device=device, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT)
   b = ttnn.from_torch(b, device=device, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT)

   output = a + b
   output = ttnn.to_torch(output)

print(output)

这个例子展示了如何将PyTorch张量转换为TT-NN张量,在Tenstorrent设备上执行加法运算,然后将结果转回PyTorch张量。TT-NN的API设计贴近PyTorch,使得现有PyTorch用户可以快速上手。

TT-Metalium:低级内核编程模型

除了高级的TT-NN库,Tenstorrent还提供了TT-Metalium这一低级编程模型,让开发者可以直接为Tenstorrent硬件开发定制化内核。

TT-Metalium提供了一套完整的API,涵盖了内存管理、张量操作、数学函数等多个方面。通过TT-Metalium,开发者可以充分发挥Tenstorrent硬件的计算能力,实现高度优化的自定义算子。

以下是TT-Metalium的一些主要特性:

细粒度的内存控制
灵活的张量布局定义
丰富的数学函数库
支持多种数据类型,包括BFloat16
硬件感知的优化提示

TT-Metalium的编程指南和API参考文档提供了详细的使用说明,帮助开发者快速入门并掌握这一强大的工具。

性能表现

Tenstorrent发布的性能数据显示,其硬件在运行各种AI模型时都展现出了卓越的性能:

在单个Wormhole N300卡上,Falcon 7B模型可以达到534 tokens/s的端到端吞吐量。
在8卡TT-QuietBox系统上,LLaMA-2 70B模型可以达到333 tokens/s的吞吐量。
ResNet-50在单个Wormhole卡上可达4,100 fps,在8卡系统上可达31,250 fps。

这些数据充分展示了Tenstorrent硬件与软件栈的强大性能。值得注意的是,Tenstorrent还在不断优化其软硬件,未来有望实现更高的性能。

开源与社区

TT-NN和TT-Metalium都是开源项目,代码托管在GitHub上。Tenstorrent积极鼓励社区参与,欢迎开发者贡献代码、报告问题或提出建议。项目遵循Apache 2.0开源协议,保障了使用的自由性。

此外,Tenstorrent还建立了Discord社区,方便开发者交流讨论。公司也定期发布更新日志,及时向用户通报新特性和性能改进。

结语

TT-NN和TT-Metalium为AI开发者提供了强大而灵活的工具,充分发挥了Tenstorrent AI加速芯片的性能潜力。无论是构建高层神经网络还是开发底层优化内核,这套软件栈都能满足开发者的需求。随着Tenstorrent硬件和软件的不断进步,相信会有越来越多的AI应用在这一平台上落地,推动AI技术的进一步发展。