光学计算有望彻底改变AI性能的游戏规则

shuszhao

从麻省理工学院(MIT)分离出来的初创公司Lightmatter专注于开发针对AI加速的光学计算处理器，在第32届Hot Chips大会上展示了一款测试芯片。该处理器利用硅光子学和MEMS技术，通过毫瓦级激光光源供电，可以光速（在硅片中）执行矩阵矢量乘法。其计算速度比现今晶体管结构的芯片（包括最新的GPU）要快几个数量级，而且功耗极低。
Lightmatter意欲通过展示该测试芯片来证明其处理器设计方法是可靠的。该公司是最早推出专门用于AI推理的光学计算（硅光子学）芯片的公司之一。
Lightmatter将于2021年秋季推出其首个商用产品，它是基于此次测试芯片的后继产品，是一款带光学计算芯片的PCIe卡，专为数据中心AI推理工作负载而设计。

广告

                               
登录/注册后可看大图

Lightmatter PCIe卡上的硅光子芯片，光源通过光纤进入（图片来源：Lightmatter）
硅光子技术的进步（通过硅芯片传播光）使复杂的片上结构成为可能，而这种结构可以被控制，以一种与传统晶体管电子结构完全不同的方式来执行MAC操作。由于晶体管结构的芯片已到达Dennard缩放定律的极限，其单位面积内的功耗不断增加，而冷却技术的发展却无法满足更大芯片的需求。因此，具有节能优势的其它技术还有很大的发展空间。
Lightmatter 的首席执行官 Nick Harris在Hot Chips大会前接受EE Times采访时表示：“
“我们通过一种完全不同的物理方法回避了整个能量缩放问题 – 使用光。这意味着我们可以利用一套不同的规则进行扩展，因此光学计算可以更快，而且能耗更低。”
那么，光学计算究竟有多快，能耗有多低？
Harris说：“我们可以将现有AI数据中心的能耗减少20倍，占板面积减少5倍。而这只是我们正在开发的第一代产品，接下来还有很长的规划路线图。”
Harris强调说，这次的测试芯片只是用于演示这项技术，它在基准测试中表现还不算好。不过他坚信，在实际应用中，Lightmatter的量产芯片将击败AI加速器市场领导者Nvidia的Ampere A100。Harris说，在BERT和Resnet-50推理等工作负载上，Lightmatter的芯片能效是A100的20倍，吞吐量至少是其5倍。


芯片设计

实际上，Lightmatter的芯片包含两个垂直堆叠的裸片。上面是一个12nm ASIC，用于存储并协调对位于下方的90nm光学计算机芯片的控制。这两颗裸片均由GlobalFoundries采用标准CMOS工艺制造。
这种光子处理器具有一个64 x 64光子矩阵矢量乘积计算器；数据可以在不到200皮秒的时间内在整个芯片内传输，比耗费多个时钟周期的晶体管计算要快几个数量级。其计算引擎由50 mW激光器驱动。
根据Harris提供的信息，这种低功耗光学计算芯片的优点之一是可以与控制/存储ASIC进行3D堆叠。晶体管结构的计算芯片会散发很多热量。Harris指出，堆叠的芯片缩短了ASIC上操作数存储器与光子芯片上计算元件之间的迹线长度（从数据转换器到光学计算引擎的总路由距离不到1毫米），延迟和功耗也因此得到改善。
Harris说：“它具有良好的正反馈环路，低功耗使堆叠成为可能，而堆叠又更加省电。”

                               
登录/注册后可看大图

Lightmatter的“芯片”实际上是两个以3D封装堆叠的裸片。上面是具有存储功能和控制功能的12nm ASIC，下面是充当计算引擎的90nm硅光子器件（图片来源：Lightmatter）
DAC接收数字输入信号，将其转换为模拟电压，并采用该电压来驱动激光器（该技术已在光纤发送器中广泛使用）。来自该激光器的光进入计算阵列，其计算单元为马赫曾德尔干涉仪（MZI）。进入MZI的相干光被分为两部分，每部分的相位调整幅度不同。将具有不同相位的信号组合会导致相长或相消干涉，从而有效调制通过MZI的光亮度（该调制可被视为乘法运算）。在波导（承载光的“缆线”）相遇的地方，信号被有效叠加。这就是光学MAC的基本原理。从计算阵列输出的光到达光电二极管，其信号再通过ADC馈送，以便与其余数字电路接口。

                               
登录/注册后可看大图

Lightmatter的光学计算阵列由DAC和ADC模块组成，再与其余数字电路对接（图片来源：Lightmatter）
MZI中的关键操作是通过机械方式实现光相位的改变。Lightmatter副总工程师Carl Ramey在Hot Chips演讲中解释到，光子芯片采用了纳米光学机电系统（NOEMS）。与微机电系统（MEMS）器件类似，其波导结构通过下方蚀刻悬挂，然后通过向其上方和下方的电容器板添加电荷来偏转，从而成功地按照所需数量改变光相位。
“NOEMS器件具有一些惊人的特性，”Ramey说。“它们损耗极低，静态功耗几乎为零。简单地加载一些电子在小型电容器上，几乎不会有任何泄漏。因为电容足够小，致动所需的动态功率也非常小。这种结构也可以相对较高的速度驱动，最高可达数百兆赫。”


节能

Ramey说，Lightmatter的演示版芯片具有64 x 64个计算单元，而且很容易扩展。
他说：“类似于晶体管收缩阵列，计算量随面积线性增长，延迟也随阵列尺寸而变化。在典型的流水线晶体管设计中，从左到右共需要64个时钟周期来执行操作。我们的延迟也会随阵列尺寸而定，但速度要快三个数量级。因此，即使是一千乘一千的阵列，其延迟也远低于1纳秒。”
有趣的是，光学计算阵列消耗的功率与面积的平方根成比例。这是因为功耗主要归因于数据转换。
“每添加一个新单元到阵列中，只需付出正常功率的平方根代价，就能获得更高性能。”Ramey说。“因此，我们的芯片越大，实际上效率就会越高。这与电子系统完全不同：遵循线性缩放规律的电子系统性能越高，功耗越大。”
除了与计算相关的能量外，还有与在芯片周围移动数据相关的能量（当今大型的晶体管AI芯片在传输数据时可能会消耗50W至100W的功耗）。通过光学计算，以光学方式移动数据意味着不再需要电源，这是巨大的节能。
其结果是，计算设备工作功率低于3W，这只是其他计算方法每次推理操作所消耗能量的一小部分。


多路复用

光学计算的另一个有趣之处是其多路复用能力。与光通信中使用的技术类似，光学计算可以将多个独立的数据流编码到不同波长的光路上，并同时输入到计算引擎中。这意味着光学计算芯片可以同时执行多个AI推理。
“这是光学计算一个非常独特的特性，”Lightmatter CEO Nick Harris说。“这意味着，只需一个物理资源，即一个处理器，其功效却如同一组处理器。”
Harris表示，尽管指定光谱（1310至1600nm）理论上可以支持至少一千个通道，但激光技术的不成熟限制了其目前只能支持8个通道。


展示光学计算的可行性

Lightmatter当前的目标客户是数据中心，包括高性能计算等大规模扩展系统，但将来还会扩展到企业领域。自动驾驶也是其远期目标，但Harris承认，进入这一领域所需的可靠性工程将是“一项艰巨的任务”。
Lightmatter有一套完整的软件栈，可以与TensorFlow或Pytorch集成。Harris说，他们的目标是在两种机器学习框架中实现即插即用。
Lightmatter成立于2017年，目前在马萨诸塞州波士顿市有46位员工，并拥有30项专利。该公司已经从Google Ventures等风投机构筹集了3300万美元的资金。
对于这个初创企业而言，首要的挑战可能是让那些持怀疑态度的客户接受光学计算的概念。Harris计划如何做到这一点呢？
他说：“这是一项艰巨的挑战。自上世纪60年代以来，计算历史上从没有一种技术取代过电子晶体管，从来没有过。人们做过尝试，但没成功。我想，这将是人们第一次看到它的实现，而我们的推销方式就是展示它能够正常运行。”