Flattening Every Memory Peak in Long-Context MoE Training

Long-context or large-batch MoE training fails the moment any single component's peak allocation exceeds device memory, so the target is every peak at once. Salesforce identifies four peaks that common parallelism plans leave unbounded, expert dispatch (with the routing matrix), the vocabulary projection (tokens times vocabulary), gradient-checkpoint boundaries (depth times sequence length), and optimiser state (parameter count), and notes that which one fails first changes with model, context, and device count. The paper bounds all four with schedules whose GPU working set is fixed at launch, including a pipelined expert-parallel dispatch, and reports 1M-token context training on a 667B model at 10.4 times the throughput of the FSDP2 baseline.