NSR封面文章 | 片上全光超节点,面向超低时延深度神经网络推理

问AI · 这项技术将如何重塑AI计算架构?

北京大学王兴军教授、舒浩文研究员团队提出了一种基于片上全光超节点光电融合分布式计算系统。该系统将高速硅光收发芯片、低损耗光交换(OCS)芯片与多个计算节点相结合,使深度神经网络不同层之间的中间数据能够通过片上高速全光超节点直接传递与交换,从而减少传统数据搬运、分组调度和缓存等待带来的时延开销。

相关论文将收录于《国家科学评论》National Science Review, NSR)“光电融合”专题,并作为当期封面文章发表。北京大学为主要完成单位,陶子涵周龑虞潍榛常华墐为文章共同第一作者,舒浩文研究员和王兴军教授为共同通讯作者。

随着人工智能(AI)模型规模快速增长,计算系统的性能提升正从注重单芯片性能转向多芯片协同扩展。然而,目前Scale-up网络在片间数据传输在带宽、能效和时延方面逐渐逼近瓶颈

针对这一挑战,北京大学王兴军教授、舒浩文研究员团队提出了一种基于片上全光超节点的光电融合分布式计算系统。在硬件层面,研究团队研发了400 Gbps硅光收发芯片作为高速信号电光转换接口,并研制了非阻塞16×16光交换芯片用于扩展片间连接与交换,可实现在物理层直接完成数据路由,系统总聚合交换带宽可达6.4 Tbps

图片

基于片上全光超节点的光电融合分布式计算系统示意图

尤为关键的是,该光交换芯片实现了低于5 dB的总损耗(含耦合损耗),无需外部光增益补偿即可实现高速无误码传输。且多条不同交换路径均可保持零误码性能,表明该片上全光超节点能够为分布式推理提供稳定数据流,并避免重传机制引入额外时延;同时,器件在超过100 nm的谱宽范围内保持平坦响应,为波分复用、带宽扩展和后续系统扩容提供了重要基础。

进一步,研究团队将一个五层卷积神经网络图像去噪模型按层依次部署到每个计算单元中,并将光交换芯片配置形成流水并行计算模式。每一层计算完成后,中间特征图经全光网络直接传输至下一处理节点继续处理,使网络能够对连续输入数据流并行工作,从而缓解传统架构中计算-存储往返带来的内存墙问题。所选GPU基准执行相同任务相比,该系统在计算资源占用仅为其约九分之一的情况下,实现了百余倍的推理速度提升

图片

基于五层卷积神经网络的图像去噪任务验证性能对比

该研究展示了片上全光超节点在大规模分布式智能计算中的潜力。未来,随着共封装光学、高速硅光收发器和算力芯片I/O速率的进一步提升,片上全光超节点有望成为高带宽、低时延、高能效智能计算系统的重要基础。


论文信息


On-chip large-scale all-optical interconnect for ultra-low-latency deep neural network inference 

National Science Review, nwag282, https://doi.org/10.1093/nsr/nwag282