Memory-Safe Hierarchical Planning for Hybrid-Parallel Transformer Training

Abstract

Training decoder-style Transformer models at scale requires composing data, tensor, pipeline, context, and, for sparse models, expert parallelism under tight device memory limits. Without a reliable pre-profiling admission model, algebraically legal but memory-infeasible configurations can consume profiling budget without yielding useful throughput measurements. Choosing a pipeline degree also leaves a structured layer-to-stage mapping problem unresolved. This paper presents Shockwave, a hierarchical planner that operates through the exposed launcher and partition interfaces of Megatron-style dense and sparse Transformer runtimes. Shockwave first builds a conservative memory-feasible region with an analytical model and backend-calibrated safety margins. It then profiles a balanced baseline and a few single-dimension anchors to learn which parallel degrees improve throughput most on the target runtime and uses that ranking to select a global strategy. Finally, it solves a mixed-integer program for stage-aware layer allocation under a fixed 1F1B pipeline schedule. Implemented as a thin control layer on top of Megatron-LM and MindSpeed, Shockwave matches or exceeds expert-curated plans, achieving geometric mean throughput gains of 1.08x over expert-tuned Megatron hybrid configurations on GPUs and 1.04x over expert MindSpeed recipes on NPUs. Relative to the evaluated automatic baselines, it improves throughput by 1.33x on GPUs and 1.10x on NPUs while reducing configuration search time by 4.59x and 5.60x.

Publication
ACM Transactions on Architecture and Code Optimization