Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Zhixin Ou, Peng Liang, Jianchen Han, Baihui Liu, Linbo Qiao

TL;DR

ParaDySe introduces a novel, layer-wise adaptive parallel-strategy switching framework for Transformer training with dynamic sequence lengths. It unifies tensor layouts across mainstream parallel methods, builds sequence-aware hybrid time-memory cost models, and employs a heuristic switching module to select optimal strategies on-the-fly without tensor redistribution. The approach demonstrates increased maximum trainable sequence lengths (up to 624K tokens) and substantial speedups across representative LLMs, effectively mitigating CPC and OOM bottlenecks. This work enables redistribution-free, memory-efficient scaling of LLMs in environments with extremely long sequences and dynamic inputs, advancing practical training of large 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.

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

TL;DR

ParaDySe introduces a novel, layer-wise adaptive parallel-strategy switching framework for Transformer training with dynamic sequence lengths. It unifies tensor layouts across mainstream parallel methods, builds sequence-aware hybrid time-memory cost models, and employs a heuristic switching module to select optimal strategies on-the-fly without tensor redistribution. The approach demonstrates increased maximum trainable sequence lengths (up to 624K tokens) and substantial speedups across representative LLMs, effectively mitigating CPC and OOM bottlenecks. This work enables redistribution-free, memory-efficient scaling of LLMs in environments with extremely long sequences and dynamic inputs, advancing practical training of large 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.

Paper Structure

This paper contains 18 sections, 4 equations, 2 figures, 5 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. LLM Training Frameworks
  4. Parallel Methods on 1D Device Grids
  5. Problem Formulation
  6. ParaDySe Framework
  7. Overview
  8. Switchable Functional Parallelism Module
  9. Hybrid Cost Estimation Module
  10. Adaptive Parallel Strategy Switching Module
  11. Implementation
  12. Experiment
  13. Setting
  14. Overall Experiment
  15. Ablation Study
  16. ...and 3 more sections

Figures (2)

  • Figure 1: Overview of ParaDySe Framework for Dynamic sequences with varying lengths. These three core Modules (namely Hybrid Cost Estimation Module, Switchable Functional Parallelism Module, and Adaptive Parallel Strategy Switching Module) are collaborative adaptively to achieve almost optimal parallel strategy according to the input sequence. The details of these modules are presented in the following subsection.
  • Figure 2: Comparing ParaDySe with baselines across LLMs and long-sequence datasets on cumulative time (s) and sequence length (K tokens). Curve termination indicates OOM failure.