OFC 2025 Google报告： OCS技术使能AI集群的高效扩展，显著提升系统可用性和能效

1. 引言

• TPU 发展历程：
  • 2013 年，Google 启动 TPU（张量处理单元）项目，首代 TPU v1 通过 PCIe 接口实现每秒千次推理，推动语音识别、图像搜索等 AI 应用落地。

• TPU v2 首次支持模型训练，256 个芯片通过 2D Torus 硬连线连接，适应矩阵运算的天然并行性。

• TPU v3 扩展至 1024 个芯片，性能是 v2 的 10 倍，但硬连线 Torus 架构在 4096 芯片规模时面临可靠性和部署难题(线缆采用不同颜色区分，但仍然难以管理)。

• TPU v4 引入光交换（OCS）技术，构建由 64 个机架组成的 Super Pod，支持 3D Torus 逻辑拓扑，提供超过 1 Exaflop（百亿亿次浮点运算）的计算能力。

• 挑战与创新： TPU v3 的硬连线 Torus 在扩展时因物理连接限制，导致单点故障影响整个集群。OCS 通过动态调整网络拓扑，实现灵活的逻辑分组，显著提升系统可用性。

2. OCS 技术细节与优势

• 架构设计：
  • 3D Torus 模块：每个机架内的 4x4x4 Torus 模块通过 OCS 连接，形成 16x16x16 的逻辑架构，支持灵活任务分割（如 8x8x4 或 16x8x16）。

• 波分复用（WDM）：将多个波长MUX到单根光纤，减少光纤和 OCS 数量，降低成本和复杂度。

• 环形器（Circulators）：实现双向信号传输，节省 50% 光纤资源，同时控制插入损耗（<1dB）和回波损耗（>50dB）。

• OCS 优势：
  • 可用性：在 1024 芯片系统中，可用性从接近 0% 提升至 75%，通过光调度切片，动态绕过故障节点实现高可靠性。

• 能效与成本：相比电交换，OCS 降低功耗 3.5 倍，成本仅增加约 10%，支持物理隔离的安全机制。

• 灵活性：最高支持 9 种并行度，适应不同规模的 AI 模型（如 LLM-0 到 LLM-3），通过 Torus 拓扑优化矩阵乘法通信效率。芯片间互连采用共享内存架构，比基于协议栈的分组网络更强大。

3. OCS 在 AI 集群中的应用

• 模型训练案例：
  • Palm 模型训练6 个Super Pod动态分配资源，50天内通过 OCS 调整拓扑，应对节点故障和维护，确保训练连续性。

• 推理优化通过共享内存架构和 Torus 拓扑，减少协议栈开销，提升推理吞吐量。

• 性能数据：

• 扩展性：在 3000 芯片规模下，系统效率接近 100%.
• 带宽优化：通过 OCS 动态分配，消除跨机架流量瓶颈，降低延迟。

4. 未来展望与行业影响

• 行业趋势：
  • 超大规模数据中心需求推动 OCS 技术主流化，解决铜缆连接的能效和扩展性瓶颈。
  • 与英伟达硅光子学技术互补，OCS 更侧重动态拓扑和高可用性，适用于高密度 AI 集群。

5. 总结 Google 通过光交换技术实现了 AI 集群的高效扩展，显著提升系统可用性和能效。OCS 的动态拓扑调整、WDM 复用和环形器设计，使其成为大规模 AI 训练的关键技术。未来，随着超大规模数据中心的发展，OCS 技术将在行业中发挥更重要作用，推动 AI 基础设施的革新。