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技术的进一步发展。
