Accelerating Frontier MoE Training with 3D Integrated Optics

TL;DR

The paper tackles the scalability bottlenecks of frontier MoE training by replacing copper-based scale-up interconnects with 3D integrated optics (Passage). It introduces Passage as a 3D photonics platform that stacks electrical and photonic layers to deliver high bandwidth density, large radix, and energy efficiency, enabling multi-rack scale-up within datacenters. Through analytical modeling of MoE workloads and comparisons across LPO, 2.5D CPO, and Passage interposers, it demonstrates up to an 8x increase in scale-up bandwidth and a 2.7x reduction in time-to-train, with substantial area and power advantages over legacy optical solutions. The work shows that maintaining expert communications within high-bandwidth domains dramatically improves MoE scaling, suggesting a practical path to training models well beyond 1 trillion parameters with more efficient data-center infrastructure.

Abstract

The unabated growth in AI workload demands is driving the need for concerted advances in compute, memory, and interconnect performance. As traditional semiconductor scaling slows, high-speed interconnects have emerged as the new scaling engine, enabling the creation of larger logical GPUs by linking many GPUs into a single, low-latency, high-bandwidth compute domain. While initial scale-up fabrics leveraged copper interconnects for their power and cost advantages, the maximum reach of passive electrical interconnects (approximately 1 meter) effectively limits the scale-up domain to within a single rack. The advent of 3D-stacked optics and logic offers a transformative, power-efficient scale-up solution for connecting hundreds of GPU packages (thousands of GPUs) across multiple data center racks. This work explores the design tradeoffs of scale-up technologies and demonstrates how frontier LLMs necessitate novel photonic solutions to achieve aggressive power and performance targets. We model the benefits of 3D CPO (Passage) enabled GPUs and switches within the scale-up domain when training Frontier Mixture of Experts (MoE) models exceeding one trillion parameters. Our results show that the substantial increases in bandwidth and radix enabled by 3D CPO allow for an 8X increase in scale-up capability. This affords new opportunities for multi-dimensional parallelism within the scale-up domain and results in a 2.7X reduction in time-to-train, unlocking unprecedented model scaling.

Figures (11)

• Figure 1: Transformer architectures: (a) Dense model with self-attention and FFN. (b) Sparse MoE with top-k=2 routing selecting experts E2 and E4 based on highest scores shown for token x2.
• Figure 2: Single-layer Switch electrical scale-up topology. A single layer of switches (top) is connected to every GPU (bottom) in the pod (only three GPU-to-switch connections shown for brevity). This provides full bandwidth connectivity between any two GPUs in the pod over multiple rails.
• Figure 3: A GPU package in a 4 × 1 reticle configuration. Four logic reticles surrounded by HBM stacks on the north and south side in black, intra-package I/O in the middle and inter-package I/O on the east and west side. SerDes Shoreline is highlighted in orange.
• Figure 4: Difference between 2D and 2.5D integration of optical engines (OEs). The left most approach shows larger 2.5D OEs with 2D host integration over an organic substrate and the resulting beachfront expansion. The rightmost approach shows smaller 3D OEs that are 2.5D-integrated in close proximity to the host on an interposer or bridge.
• Figure 5: Exploded view of a Passage Interposer Solution: (1) EIC, (2) PIC, (3) Fiber Attach Unit, (4) fibers and (5) substrate. Blue vertical lines distributed throughout package area represent I/O enabled without SerDes shoreline constraints.