我对TPU芯片的兴趣不仅在于其高效的张量处理能力,更因为它正在成为国内应对AI大模型挑战的重要策略。随着外部供应链的不确定性,国内企业开始加大对TPU芯片技术的研发和应用。

自从AI大模型来了,英伟达喝汤喝到撑,GPU自然也就成了香饽饽。但在地缘政治局势愈发紧张的现如今,国内高端AI芯片不断被围追堵截。

就比如,7月22日就出现了戏剧性的一幕,英伟达出现一正一反的消息:一方面,美政府正考虑新的贸易限制,阻止英伟达向中国市场推出“特供版”HGX-H20 AI GPU,如果限制正式实施,英伟达可能会损失约120亿美元的收入;另一方面,英伟达正在为中国市场打造全新的特供版GPU,以刚推出的“Blackwell”为基础打造B20。

可以说,美国方面的态度非常鲜明,就是要全面围堵中国获取高端AI芯片的渠道,以此占领AI领域的高地。

在这种情况下,国内又该如何应对?最近一段时间内,国内开始瞄准TPU(张量处理单元),另辟蹊径。

国内也有厂商做TPU了

众所周知,AI大模型主要分为两个阶段,一是训练,二是推理。

推理芯片常见,而训练芯片不常见,这是因为训练不仅消耗巨大的算力资源,同时需要处理大量的并行任务,所以GPU才会成为当前的主流。

TPU全称Tensor Processing Unit,是一种专为处理张量运算而设计的ASIC芯片,由谷歌自研在2016年推出首款产品。在深度学习的世界里,张量(多维数组)是无处不在的。TPU就是为了高效处理这些张量运算而诞生的。

TPU内置大量矩阵运算单元,使得其能够并行处理大量的矩阵运算,大大提高计算效率。

不过相比GPU或者说GPGPU,TPU太专用了,但是应付AI训练还是绰绰有余。

简单粗暴对比起来就是:TPU与同期的CPU和GPU相比,可以提供15~30倍的性能提升,以及30~80倍的效率(性能/瓦特)提升。

早在2018年,就有一家AGM Micro国内公司提供TPU推理技术授权,不过,后来这家公司基本不怎么发布关于TPU相关的消息了。

而最近,一家名为中昊芯英的国产公司,就展出了其首枚高性能TPU(张量处理器)AI训练芯片。

据了解,中昊芯英TPU“刹那”于去年成功量产,已在全国多地千卡集群规模的智算中心交付落地。该芯片以1024片芯片高速片间互联的能力构建了大规模智算集群“泰则”,系统集群性能远超传统GPU数十倍,可支撑超千亿参数AIGC大模型训练与推理。

公开资料显示,中昊芯英创始人杨龚轶凡曾在谷歌作为芯片研发核心人员,深度参与了谷歌TPU 2/3/4的设计与研发,在他看来,TPU是为AI大模型而生的优势架构。

碳纳米管和TPU,牵手了

昨日,也传出另一个与TPU相关的消息。

消息显示,北京大学电子学院碳基电子学研究中心的彭练矛-张志勇团队,在下一代芯片技术领域取得突破,成功研发出世界首个基于碳纳米管的张量处理器芯片(TPU)。

官方表示,高能效计算芯片的发展有两个重大瓶颈:一是传统冯诺依曼架构已经无法满足高速、高带宽的数据搬运和处理需求;二是构建芯片的硅基互补金属氧化物半导体晶体管,进入了尺寸缩减、功耗剧增的困境,亟需发展超薄、高载流子迁移率的半导体作为沟道材料。

而碳纳米管具有优异的电学特性和超薄结构,碳纳米管晶体管已经展现出超越商用硅基晶体管的性能和功耗潜力。不过,为了最大化发挥芯片算力和能效,必须将新材料与器件结合,北大的这一个成果就主要围绕这方面进行了研究。

作为世界首个碳纳米管基的张量处理器(TPU)芯片,可实现高能效的卷积神经网络运算。省流版总结如下:

工艺:该芯片采用2bit MAC(乘累加单元),3微米工艺技术节点,集成3000个碳基晶体管,可实现图像轮廓识别、提取等功能,图像轮廓提取正确率达100%;

架构:该芯片采用脉动阵列架构设计,可实现高效地数据复用,大大节约张量运算所需的数据存储、搬运等操作,精准匹配了神经网络的运算特点;

识别率:其上搭建了5层卷积神经网络,实现手写数字识别的应用,理论正确率90%,实际正确率可达86%;

功耗:仅为295µW,器件总数也为新型卷积加速硬件中的最低值;

实际应用效果:该芯片可使用180 nm碳基技术进行流片加工,仿真结果表示,碳基神经网络加速芯片可在1 V电压下工作,可运行的最高主频为850 MHz,能效可以达到1TOPS/w。

TPU芯片:国内AI大模型的新选择

“群殴”英伟达

相比其它AI芯片来说,TPU的关注度的确高很多。今年6月,就有消息称,生成式AI技术大厂OpenAI为了自研AI芯片,新招募的研究人员几乎为谷歌TPU团队的前员工。可以说,在大规模训练和推理上,TPU是相对成熟的方案。

作为TPU的发明者,谷歌之所以推出TPU,其目标便是为企业提供Nvidia GPU的替代品。前阵子的Google I/O 2024上,谷歌推出第六代TPU,性能有显著提升。

与TPU v5e相比,Trillium TPU峰值计算性能提高了4.7倍。为了实现更高的性能,谷歌投入了大量精力扩展执行计算的矩阵乘法单元或MXU的大小,并提高了其整体时钟速度。此外,Trillium GPU的高带宽内存容量和带宽是原来的两倍,而芯片间互连带宽也增加了一倍。

为了让客户更放心地替代英伟达,谷歌也在谋划用TPU替换掉GPU——可能在今年底停止外部AI算力芯片的采购,转而完全依赖自研的TPU。谷歌的算力总量,结合自研TPU和先前的芯片采购,预计可达全球算力总量的25%。

总之,现在的英伟达四面楚歌,谁都想取代他。随着特供版接连被围堵,国内市场似乎对英伟达不买账了,转向采购国产芯片。但在AI芯片市场上,从来没有什么稳赚不赔,前两天就有一家日本AI芯片厂商宣布解散。可以从此看出,TPU的立足之本,便是更好的能效比和软件生态。在此方向上,国产在路上。

因此,TPU芯片不仅是技术进步的体现,也是国内AI产业自主创新的重要一环。随着国内厂商不断推出高性能的TPU芯片,我相信未来在AI芯片领域,国内将有更多精彩的发展与突破。