刷屏的清华AI光芯片

coffee198375

最近，清华大学传出了好消息。首创AI光芯片架构，研制全新AI“光芯片”——太极（Taichi），可以实现160 TOPS/W通用智能计算，能效是H100的1000倍。


训练下一代万亿级参数大模型的高效芯片诞生了。目前，相关研究论文以“Large-scale photonic chiplet Taichi empowers 160-TOPS/W artificial general intelligence”为题，已发表在权威科学期刊 Science 上。


01 成果是什么？


当前，越来越多迹象表明，LLM不会是通往AGI的最终路径。


计算机早已经成为世界能耗巨头，随着越来越多耗电量大的人工智能投入使用，计算机的能源需求也飞速上涨。


以英伟达H100为例，其峰值功耗为700 瓦，按照 61% 的年利用率计算，相当于一个美国家庭的平均功耗（假设每个家庭 2.51 人）。有专家预测，在大量部署H100后，总功耗将于一座美国大城市不相上下，甚至超过一些欧洲小国。


若是能够发明一种，节省大量能耗的芯片，LLM的性能或在未来实现更大的提升。


而太极，可能会让通用人工智能（AGI）成为现实。


根据清华大学官网介绍，清华团队设计了基于集成衍射干涉异构设计和通用分布式计算架构的大规模光芯片——太极，该架构具有上千万个神经元的能力，实现160万亿次/秒·瓦（TOPS/W）的通用智能计算。


此外，在太极光芯片在实验中实现了芯片上1000个类别级别的分类（在1623类别的Omniglot数据集上准确率为91.89%）和高保真的人工智能生成的内容，效率提高了两个数量级。


研究人员表示，“太极”为大规模的光子计算和高级任务铺平了道路，进一步发掘了光子学在现代AGI中的灵活性和潜力。


Science对这个研究有高度评论：“来自清华的团队探索了分布式衍射干涉混合光计算架构，有效地将光神经网络（ONN）的规模提高到百万神经元级别。通过实验实现了一个芯片上1396万个神经元的ONN，用于复杂的、千类级的分类和人工智能生成的内容任务。这项工作是向现实世界的光计算迈出的有希望的一步，支持人工智能中的各种应用。”




02 什么东西，有什么用？


电子芯片的瓶颈




当电子通过晶体管和其他传统集成电路元件时，会遇到阻力并产生热量。随着设计者不断将各种元件添加到芯片上，芯片产生的热量自然会升高。电子这一特性甚至成为了微型芯片性能提升的障碍，同时也是计算机能耗如此之高的主要原因。


以电子为载体的技术发展已趋近物理极限，芯片尺寸降到极致时出现的“功耗墙”难题，访存瓶颈下大量信息存储不过来、计算不过来，以及电子芯片性能提升的同时性价比降低。


在电路上，用光子替代电子的设想由来已久。


20 世纪六七十年代，研究者就已经开始开发光子芯片了。那时候，部分专家预计光子芯片会像传统集成芯片一样迅速微型化。


电路能耗降低，还是归功于光的性质。光子芯片不存在电阻问题。因为由镭射产生的光子能快速通过波导、调制器、反射器等原件阵列。因此，光子芯片产生热量更少，能耗也更低。


光为载体的计算芯片




光计算，顾名思义是将计算载体从电变为光，利用光在芯片中的传播进行计算。


人工智能时代是由算力支撑起来的，从能耗的角度来看，硅光技术能够有效提升GPU的整体性能、大幅降低其功耗，有效解决目前的算力瓶颈。也就是说，下一代算力很可能会是光子计算甚至量子计算。


问题在于，当前的集成光子计算，特别是光学神经网络（ONN），通常包含数百到数千个参数，其中数十个是可调参数，仅支持基本任务，如简单的模式识别和元音识别。能够进行一些简单任务和浅层模型，无法支撑亟需高算力与高能效的复杂大模型智能计算。




光计算对于集成度的要求会更高，但其技术难点其实并不只在集成，在计算单元循环使用和非线性处理方面挑战更大。


据了解，一个巨大的AI模型中，每一层网络都需要进行矩阵运算，一个大矩阵还可能需要拆成几个小矩阵进行计算。也就是说，在大模型中，矩阵乘加运算是一个反复、循环的计算过程，每次循环，矩阵上的元素权重都会被更新。


与电子相比，光的矩阵乘加运算非常快，但一涉及到权重的更新，光的速度就会变慢。


要实现大规模、高能效的光子计算，简单地扩大现有的光子神经网络芯片是不现实的，因为随着神经网络层数的增加，不可避免的模拟噪声会呈指数级扩大。放大现有架构的规模并不能成比例地提高性能。


新架构：分布式衍射-干涉混合光子计算架构




清华团队进行了架构方面的研究，“从0到1”重新设计适合光计算的新架构。


根据论文介绍，清华团队为采用分布式计算的太极，构建了一个深度较浅但宽度较广的网络结构。


与为深度计算堆叠一系列层的传统方法不同，Taichi 将计算资源分布到多个独立的集群中，为子任务分别组织集群，并最终合成这些子任务，从而完成复杂的高级任务。


具体地说，光学衍射层的完全连通特性，可以提供比传统DNN中的卷积层更大的变形能力。


这意味着光学网络有可能用比电子系统更少的层数实现相同的变换。


图中（B）中展示了「太极」芯片，包括用于大规模输入和输出数据的双衍射单元，以及用于可重构特征嵌入和硬件多路复用的MZI阵列的可调矩阵乘法。


论文第一作者、电子系博士生徐智昊介绍：“在“太极”架构中，自顶向下的编码拆分-解码重构机制，将复杂智能任务化繁为简，拆分为多通道高并行的子任务，构建的分布式‘大感受野’浅层光网络对子任务分而治之，突破物理模拟器件多层深度级联的固有计算误差。”


AI光芯片：干涉-衍射融合计算芯片




在这项工作中，团队设计了一种具有灵活分布式计算架构的大规模衍射-干涉混合型光子AI芯片——「太极」。


据论文报道：“太极”光芯片具备879 T MACS/mm2的面积效率与160 TOPS/W的能量效率，实现了高达两个数量级的能效提升。首次赋能光计算实现自然场景千类对象识别、跨模态内容生成等人工智能复杂任务。


03 国内光芯片企业情况如何？


在面向“后摩尔时代”的潜在颠覆性技术里，光子芯片已进入人们的视野。其所具有的高速度、低能耗、工艺技术相对成熟等优势，能够有效突破传统集成电路物理极限上的瓶颈，满足新一轮科技革命中人工智能、物联网、云计算等产业对信息获取、传输、计算、存储、显示的技术需求。


目前，全球光子芯片产业刚刚起步，作为独立于电子集成技术的新集成技术，其技术壁垒还没有形成。我国光子产业发展水平与世界处于并跑阶段，在光子基础理论研究和技术发展方面具有一定的优势。


目前中国本土的高功率激光芯片、部分高速率激光芯片（10G、25G等）等已处于国产化加速突破阶段，而光探测芯片、25G以上高速率激光芯片刚刚起步。


值得注意的是，光芯片方面华为也有所布局。


今年3月，华为公布了一项“光芯片及其制备方法、通信设备”发明专利。申请公布号为：CN117616316A，该专利申请日期为2021年9月18日。摘要显示，本申请的实施例提供一种光芯片及其制备方法、通信设备，涉及光通信技术领域，解决现有的光芯片中光波导在制备过程中尖端易断裂的问题。