Training Trillion-Parameter LLMs With 3-D DRAM-Based Accelerator and Co-Packaged Optical Interconnects

Large language models (LLMs) have become ubiquitous for text-based generative applications. The model size has been increasing rapidly to reach a trillion parameters to improve the accuracy of such models. However, training such large models requires huge memory, and the current DRAM technologies (high-bandwidth memory with die stacking) are insufficient. Various device technologies have been proposed for monolithically integrated 3D stackable DRAM to improve memory density by enabling bit-cost scaling in the vertical direction. However, given the reticle limit, a single accelerator (even equipped with 3D DRAM) is not sufficient to store and train a complete trillion-parameter model. As a result, a cluster of such accelerators is required to be interconnected with high bandwidth. Here we explore the use of advanced interconnect technologies such as co-packaged optics in a 3D integration scheme to aggregate such a cluster with optimized dataflow. Finally, we perform a system-level benchmarking on different model size and obtain an energy efficiency of 6.7 TFLOPS/W and a throughput of 592 output tokens per second for training one trillion-parameter model.