HSAP: A Hierachical Sequence-aware Parallelism for Hybrid-Context Generative Models

2026-06-29

AI Summary

This work addresses the challenges of incorrect causal attention computation and high communication overhead arising from mixed-context packed sequences in sequence parallelism. To resolve these issues, the authors propose a hierarchical sequence-aware parallel framework that integrates existing parallelization paradigms. The framework employs a sequence-aware algorithm to guarantee correctness in causal attention while leveraging just-in-time (JIT) compilation and NCCL-based communication optimizations to jointly manage memory consumption and communication costs. Experimental results demonstrate that the proposed method significantly outperforms state-of-the-art sequence parallel strategies in training efficiency, scalability, and achievable parallelism, thereby enabling effective and efficient training of large language models with mixed-context inputs.

Abstract

In this paper, we aim to combine the advantages of existing sequence parallelism paradigms and overcomes their drawbacks, the most serious of which is the incapability to correctly compute causal attention on the hybrid-context packed sequences, in a stronger sequence parallelism framework. The practical technique of packing sequences for efficiently pretraining and fine-tuning large language models causes cross-contamination problem in attention computation, which can be effectively solved when no parallelism in the sequence length dimension is taken. However, in sequence parallelism, existing approaches either ignore the scenario of hybrid-context sequences or conversely sacrifice and limit parallelism degree for supporting the scenario. To this end, we innovatively propose an efficient Sequence-Aware Parallelism algorithm to conquer the obstacles of intensive tensor transmission and partial attention computation across multiple device groups. Our algorithm utilizes JIT (Just-In-Time) compilation to optimize the communication strategy of all device groups in NCCL level. Further, we integrate existing sequence parallelism paradigms into a Hierachical Sequence-Aware Parallelism framework which benefits from our sequence-aware algorithm. We additionally elaborate on the memory and communication overhead management of the hierachical framework to optimize its performance. Through multiple experiments, we demonstrate that our proposed approach outperform other state-of-the-arts sequence parallelism approches in multiple metrics.

Problem

• sequence parallelism
• causal attention
• hybrid-context sequences
• cross-contamination
• large language models

Innovation

• Sequence Parallelism
• Hybrid-Context
• Causal Attention
• JIT Compilation
• Hierarchical Framework