Elastic pipeline load balancing for dynamic DNNS

Training of dynamic models is gaining traction in DNNs as it reduces computational and memory requirements of large-scale training. Gradual pruning, one of the prominent approaches for dynamic training, prunes (or sparsifies) the parameters of a model during training. However, one of the side effects of gradual pruning is that sparsification introduces an imbalanced workload across accelerators, which in turn affects the pipeline parallelism efficiency. This work introduces DynPipe which dynamically load balances the stages of the pipeline to offset the negative performance effects of pruning. On top of load balancing dynamic models, DynPipe can dynamically pack work into fewer GPUs, while sustaining performance. DynPipe works on single nodes with multi-GPUs and also on systems with multinodes. Experimental results show that DynPipe can speed up the training up to 5.64% in a single node, and 8.43% in a multi-node setting, over state-of-the art solutions used in training production large language models. DynPipe is available at https://anonymous.4open.science/r/DynPipe-1EC5