The large-scale model training system, Baguafurnace 2, co-developed by Professor Zhai Jidong's team from Tsinghua University and associated organizations, has had its research paper accepted by SC26. This system tackles the issues of video memory spikes that occur due to occasional token concentrations directed towards a subset of experts during the training of large Mixture of Experts (MoE) models. These spikes can lead to task Out-of-Memory (OOM) crashes and pose challenges in balancing operational efficiency and stability when relying on traditional static memory configurations. To address these issues, Baguafurnace 2 introduces a flexible memory provisioning mechanism that activates dynamically.

During the critical time window between the MoE gating computation and the commencement of expert computation, the system employs an extensible cost model to evaluate the advantages and disadvantages of various memory strategies. It then dynamically selects the most suitable solutions, such as saving, recomputing, or offloading data, at the operator-level granularity. Additionally, the system incorporates supplementary operations into pipeline hardware gaps through dynamic planning and scheduling, ensuring minimal overall decision-making overhead.

When tested on a cluster comprising 10,000 domestically produced accelerators, Baguafurnace 2 exhibited an average throughput enhancement of 1.14 times compared to the baseline in training scenarios involving 128 to 12,288 accelerators and model sizes ranging from 188 billion to 1.04 trillion parameters. Under the configuration of 12,288 accelerators, the system achieved a total throughput of 2.012 EFLOPS, with an average Model FLOPS Utilization (MFU) of 32.9%. The core code related to this system has been made openly available.