STON: Scaling Torus-Based AI Training Clusters via Optical Circuit Switches

Abstract

Torus networks are deployed in production AI training clusters for their path diversity and low latency, but 2D Torus scales poorly: electrical packet switches compromise latency, and high-dimensional Torus introduces excessive routing complexity. We present STON (Scalable TOrus Network), a hierarchical architecture that treats a 2D Torus as a supernode and interconnects supernodes with a reconfigurable Optical Circuit Switch (OCS) for AlltoAll-dominated large-scale training networks. STON comprises three coordinated modules: (1) fragmentaware task placement, which minimizes inter-supernode traffic by reducing job fragmentation; (2) non-disruptive logical topology mapping, governed by two principles that prevent OCS reconfiguration from disrupting running tasks or partitioning multisupernode jobs; and (3) compute-phase traffic forwarding, which ensures reachability when direct OCS circuits are unavailable. STON reduces average FCT by 42.2%-61.1% across synthetic workloads and by 52.6% on a one-day Kalos production trace (under an AlltoAll traffic model for all jobs), with 95th-percentile tail latency reduced by up to 74.5%, versus a static direct-connect baseline using the same OCS hardware.