ParaDySe: A Parallel-Strategy Switching Framework for Dynamic Sequence Lengths in Transformer

TL;DR

ParaDySe is an adaptive framework that dynamically switches parallel strategies during Transformer training to optimize performance and memory usage for sequences of varying lengths.

Contribution

It introduces a novel on-the-fly strategy switching mechanism with cost models and heuristic algorithms for efficient Transformer training on dynamic sequence lengths.

Findings

Addresses out-of-memory issues in long sequences

Achieves communication-parallelization cancellation on short sequences

Improves training efficiency for large language models

Abstract

Dynamic sequences with varying lengths have been widely used in the training of Transformer-based large language models (LLMs). However, current training frameworks adopt a pre-defined static parallel strategy for these sequences, causing neither communication-parallelization cancellation on short sequences nor out-of-memory on long sequences. To mitigate these issues, we propose ParaDySe, a novel adaptive Parallel strategy switching framework for Dynamic Sequences. ParaDySe enables on-the-fly optimal strategy adoption according to the immediate input sequence. It first implements the modular function libraries for parallel strategies with unified tensor layout specifications, and then builds sequence-aware memory and time cost models with hybrid methods. Guided by cost models, ParaDySe selects optimal layer-wise strategies for dynamic sequences via an efficient heuristic algorithm. By integrating these techniques together, ParaDySe achieves seamless hot-switching of optimal strategies through its well-designed function libraries. We compare ParaDySe with baselines on representative LLMs under datasets with sequence lengths up to 624K. Experimental results indicate that ParaDySe addresses OOM and CPC bottlenecks in LLM training by systematically integrating long-sequence optimizations with existing frameworks.