PaperlayerView exact revision on arXiv

Source-linked AI summary

BASP: Communication-Efficient Batch-Aware Sequence Parallelism for LLM Training

Bigyan Ghimire, Jon C. Calhoun

arXiv:2609.03151v1cs.DCcs.LG

TL;DR

Long-context training is challenged by memory demands, and existing sequence-parallel communication does not adapt to micro-batch structure. BASP assigns sequences to disjoint GPU subgroups to reduce all-to-all group size. Across Llama and Qwen variants, it outperforms Ulysses-SP, with reported speedups from 1.17× to 1.32×.

  • Problem

    Long-context training faces memory pressure, while existing sequence parallelism uses batch-agnostic partitioning and its communication inefficiency has been understudied.

    Source: §I · p. 1

  • Method

    When B > 1 and N = KB, BASP assigns each sequence to a disjoint group of K = N/B GPUs, reducing communication peers per GPU from N−1 to K−1.

    Source: §I · p. 2

  • Results

    1.17×–1.32× speedups: BASP consistently outperforms Ulysses-SP across Llama and Qwen variants.

    Source: A. End-to-end execution time · p. 7

  • Takeaways & Limitations

    BASP reduces communication phases while preserving per-GPU memory footprint and sequence partition size.

    Source: §VIII · p. 10

    BASP requires N ≥ B and N = KB for integer K; supporting non-divisible configurations and more flexible group formation remains future work.

    Source: §VI · p. 9

Abstract

from arXiv · show

Long-context reasoning for large language models (LLMs) is becoming increasingly important, but training over long sequences remains challenging due to massive memory and communication requirements. Sequence parallelism has emerged as an essential technique for addressing bottlenecks in long sequence LLM training. However, we observe that existing sequence parallelism methods are batch-agnostic and apply uniform sequence partitioning across all batch sizes, resulting in inefficient communication. In this paper, we introduce Batch- Aware Sequence Parallelism (BASP), a sequence parallelism approach that leverages batch structure to reduce communication overhead. BASP exploits batch structure by partitioning GPUs into disjoint sequence-parallel groups according to the micro- batch size. This design reduces the all-to-all communication group size, thereby localizing communication and improving training efficiency. Experimental results on an NVIDIA A100 cluster show that BASP improves end-to-end training time by up to 1.17 - 1.31x in Llama and Qwen models compared to standard sequence parallel baselines, while preserving identical model accuracy and memory usage.

I. INTRODUCTION

Long-context training faces memory pressure and communication inefficiency, while existing sequence parallelism partitions sequences without adapting communication to micro-batch structure. BASP uses disjoint GPU groups to reduce collective size; the paper reports faster Llama and Qwen training than Ulysses-SP.

I. INTRODUCTION

  • Long-context workloads require reasoning over thousands to millions of tokens, while attention memory and communication challenge distributed training.Sequence parallelism distributes sequence length and activation memory, but communication inefficiency in sequence parallelism has been understudied.Source: §I · p. 1
  • Standard Ulysses performs global all-to-all regardless of micro-batch size, and profiling finds its time grows with batch size and consumes an increasing share of iteration time.The profiling used an 8-GPU NVIDIA A100 cluster.Source: §I · p. 1, Figure 1 · p. 2
  • BASP assigns each of B sequences to a disjoint group of K = N/B GPUs when B > 1 and N = KB, reducing peers per GPU from N−1 to K−1.The design preserves per-GPU memory footprint and sequence shard size.Source: §I · p. 2
  • On typical clusters, choosing K equal to GPUs per node confines all-to-all traffic to the faster intra-node network.The paper contrasts intra-node NVLink bandwidth with inter-node InfiniBand bandwidth.Source: §I · p. 2
  • 1.21× speedup on Llama 3.1-8B and 1.32× on Qwen 1.5-1.8B versus Ulysses-SP; all-to-all time falls by up to 85× at micro-batch size 8.These are the reported evaluation results on Llama and Qwen variants.Source: §I · p. 2

II. BACKGROUND

The background reviews transformer attention and the memory pressure that grows with sequence length. It describes FlashAttention's reduced intermediate-activation memory, while noting that long-context training still faces memory and communication bottlenecks.

A. Transformer Architecture

  • Self-attention lets each token attend to all others, and transformer layers stack self-attention with feed-forward networks.The section introduces scaled dot-product attention for understanding transformer parallelization strategies.Source: A. Transformer Architecture · p. 2
  • Scaled dot-product attention takes token inputs X and X′, projects queries, keys, and values, and uses the query-key scores to weight values.In self-attention, X = X′.Source: A. Transformer Architecture · p. 2
  • Multi-head attention computes attention over h low-dimensional projections of Q, K, and V, concatenates the head outputs, and applies an output projection.The section identifies W_O as the final output projection matrix.Source: A. Transformer Architecture · p. 2
  • The attention score matrix QK^T requires O(N^2) memory, constraining training, and activation memory grows rapidly with sequence length in distributed long-context training.The intermediate attention activations may exceed a single device's memory capacity.Source: A. Transformer Architecture · pp. 2–3
  • FlashAttention avoids materializing the full attention matrix with tiling and online softmax, reducing intermediate-activation memory complexity from O(N^2) to O(N).The passage says long-context training still encounters memory and communication bottlenecks at scale.Source: A. Transformer Architecture · p. 3

B. Modes of parallelism

The section surveys parallelism strategies that address different training bottlenecks. It contrasts model-state or layer partitioning with sequence-oriented approaches and their communication or utilization costs.

B. Modes of parallelism

  • Data parallelism partitions the input batch but replicates model parameters, gradients, and optimizer states on every GPU, limiting feasibility for large models.Each device computes on its local mini-batch and synchronizes gradients.Source: B. Modes of parallelism · p. 3
  • ZeRO and FSDP shard model states across devices, reducing per-GPU memory use and allowing models larger than a single device's capacity.ZeRO progressively shards optimizer states, gradients, and weights across its stages.Source: B. Modes of parallelism · p. 3
  • Tensor parallelism partitions computation within layers and exchanges intermediate activations through collectives, adding communication overhead.For attention, Megatron-LM partitions sequence dimensions and uses all-gather and reduce-scatter for QKV projections.Source: B. Modes of parallelism · p. 3
  • Pipeline parallelism assigns layer stages to GPUs and pipelines microbatches, but bubbles and scheduling complexity can underutilize GPUs.The passage describes pipeline parallelism as effective for alleviating memory pressure in very deep transformers.Source: B. Modes of parallelism · p. 3
  • Megatron-SP parallelizes non-tensor-parallelized layers across sequence dimensions, while context parallelism distributes the full sequence for attention computation.Megatron-SP communication volume rises linearly with sequence length; context parallelism maintains constant communication volume as sequence length and device counts increase.Source: B. Modes of parallelism · p. 3

C. DeepSpeed Ulysses

DeepSpeed-Ulysses partitions sequences across GPUs and uses all-to-all operations to redistribute projections for attention and return its outputs. The section motivates reducing communication participants because all-to-all takes a substantial and growing share of execution time.

C. DeepSpeed Ulysses

  • Ulysses partitions each sequence across N GPUs, computes local Q, K, and V, then uses all-to-all operations to distribute projections across heads and return attention outputs.The design keeps attention mathematically equivalent to full self-attention while reducing per-device memory through sequence partitioning.Source: C. DeepSpeed Ulysses · p. 3

A. Ulysses-SP Communication pattern

  • In Ulysses-SP, all N GPUs participate in the attention all-to-all, with each GPU holding token chunks from all micro-batch sequences.This is an N-way all-to-all communication pattern.Source: A. Ulysses-SP Communication pattern · p. 3
  • For two sequences and four GPUs, Ulysses uses one 4-GPU all-to-all, whereas BASP uses two parallel 2-GPU all-to-alls for attention.After attention, the output is redistributed to the original GPUs using another all-to-all.Source: Figure 2 · p. 4
  • All-to-all takes 34% of iteration execution time in the profiled sequence-parallel run, and its time increases with micro-batch size.The profiling passage notes that network-constrained cluster communication cannot easily be hidden by overlapping communication and computation.Source: A. Ulysses-SP Communication pattern · p. 3, Figure 3 · p. 4
  • BASP's motivation is to reduce communication participants using micro-batch structure, limiting synchronization to smaller groups while preserving sequence-parallel correctness.The stated rationale invokes latency and bandwidth components in classical collective communication models.Source: A. Ulysses-SP Communication pattern · p. 4

IV. BATCH AWARE SEQUENCE PARALLELISM

BASP uses microbatch structure to replace global sequence-parallel communication with smaller, disjoint GPU groups. Its group formation assumes an integer N/B and can keep communication within a node when group size matches the GPUs per node.

IV. BATCH AWARE SEQUENCE PARALLELISM

  • When B > 1 and N = KB, BASP assigns each sequence to a disjoint K = N/B GPU group, reducing per-GPU communication peers from N −1 to K −1.Each group performs an independent K-way all-to-all on its assigned sequence.Source: §IV · p. 4

A. Batch-aware group formation

  • At initialization, BASP partitions the SP degree into independent batch-aware groups, each responsible for a distinct subset of the batch.The group size is K=N/B; each of the B sequences is assigned to exactly K GPUs, with no GPU in multiple groups.Source: A. Batch-aware group formation · p. 5
  • BASP currently assumes integer K=N/B, falls back to standard data parallelism when B>N, and leaves non-divisible configurations for future work.The group-formation logic validates that N is divisible by B.Source: A. Batch-aware group formation · p. 5
  • When K equals the GPUs per node, contiguous rank assignment maps each group to one node and keeps its all-to-all traffic on NVLink.This confines communication to the intra-node interconnect when possible.Source: A. Batch-aware group formation · p. 5, §IV · p. 5

B. Batch Aware Sequence Partioning

BASP jointly partitions batch and sequence dimensions: each sequence is sharded only across its assigned GPU group, rather than across every GPU. This localizes processing to batch elements and avoids cross-sequence aggregation within the group.

B. Batch Aware Sequence Partioning

  • Unlike Ulysses-SP, which shards every sequence across all N GPUs, BASP partitions GPUs into SP groups based on microbatch size and assigns each sequence to its group.Within its group, a sequence is distributed across only those GPUs.Source: B. Batch Aware Sequence Partioning · p. 5
  • BASP preserves per-GPU computational load while eliminating cross-sequence aggregation inside the SP group.Communication is localized to the GPUs assigned to each batch element.Source: B. Batch Aware Sequence Partioning · p. 5
  • For N=4 and B=2, the BASP partition assigns each complete sequence to fewer GPUs than Ulysses-SP, reducing all-to-all participants during self-attention.The figure compares sequence partitioning with each sequence of length S.Source: Figure 5 · p. 6

C. Subgroup all-to-all

BASP restricts attention all-to-all operations to the GPUs in each subgroup while preserving per-GPU workload and memory relative to Ulysses-SP.

C. Subgroup all-to-all

  • BASP restricts each attention all-to-all to the K GPUs in a subgroup, instead of spanning all N GPUs globally.Independent K-way collectives run across groups for their assigned batch subsets.Source: C. Subgroup all-to-all · p. 6

1) Workload and Memory Equivalence.:

  • Despite processing more attention heads, each BASP GPU retains the same computational workload and memory as in Ulysses-SP.BASP divides the GPUs into groups while adjusting sequence-shard size; the paper states that per-GPU workload and memory are preserved.Source: 1) Workload and Memory Equivalence · p. 6

C. Subgroup all-to-all

  • For N=4 and B=2, BASP performs two parallel 2-GPU all-to-alls, compared with one 4-GPU all-to-all in Ulysses-SP.Each BASP pair holds a complete sequence.Source: Figure 6 · p. 6

D. Combining with ZeRO

BASP combines subgroup communication with sequence sharding, retaining the memory efficiency of Ulysses-SP at Sp=N while reducing all-to-all participants. It also integrates ZeRO to partition model states across sequence and data-parallel groups.

D. Combining with ZeRO

  • BASP integrates ZeRO, which partitions optimizer states, gradients, and parameters across data-parallel processes, across both sequence- and data-parallel groups.The described ZeRO stages partition optimizer states (Zero-1), gradients (Zero-2), and parameters (Zero-3).Source: D. Combining with ZeRO · p. 6

E. Comparison with Sequence Parallel degree

  • Setting Ulysses-SP to Sp=N/B reduces all-to-all participants but assigns multiple microbatch sequences per GPU, sacrificing sequence-sharding memory efficiency and limiting maximum trainable sequence length.With Sp=N, Ulysses-SP instead shards the sequence across all N GPUs and uses a global all-to-all.Source: E. Comparison with Sequence Parallel… · p. 7
  • BASP limits all-to-all participants to K=N/B while retaining Ulysses-SP's BS/N tokens per GPU and memory-efficient sequence sharding.This preserves maximum sequence-length capacity while reducing communication overhead.Source: E. Comparison with Sequence Parallel… · p. 7

V. EXPERIMENTAL EVALUATION

The evaluation compares BASP with DeepSpeed-Ulysses on Llama and Qwen models, measuring end-to-end time and communication across model, microbatch, and sequence settings. BASP reduces end-to-end step time and its benefits scale with microbatch size, while communication cuts can exceed overall speedups.

V. EXPERIMENTAL EVALUATION

  • The evaluation compares BASP with DeepSpeed-Ulysses on Llama and Qwen variants using two nodes of four NVIDIA A100 40GB GPUs.The nodes are interconnected by 400Gbps InfiniBand; reported metrics average 30 iterations.Source: §V · p. 7

A. End-to-end execution time

  • All-to-all time falls 2.23×–3.10× across models; for Llama 3.2-1B, a 2.45× reduction yields 1.17× end-to-end speedup.The communication benefit exceeds the end-to-end gain because all-to-all accounts for approximately 13.4% of Ulysses-SP iteration time in that example.Source: A. End-to-end execution time · pp. 7–8
  • Across all model variants shown, BASP consistently outperforms Ulysses-SP, achieving 1.17×–1.32× speedups; it reduces Qwen step time by 18.2–24.1% and Llama step time by 14.7–20.5%.Source: A. End-to-end execution time · p. 7

B. Scaling with Microbatch Size

  • At B=8, all-to-all speedup reaches 85× but end-to-end speedup is 1.25×; ZeRO-related communication becomes the bottleneck.All-to-all falls to 0.5% of iteration time at B=8.Source: B. Scaling with Microbatch Size · p. 8
  • 1.10×–1.26× speedup: on Llama 3.2-1B, BASP’s advantage grows as microbatch size increases from 2 to 8.At batch size 1, BASP and Ulysses-SP have almost identical step time because both perform an 8-way all-to-all.Source: Figure 8 · p. 8, B. Scaling with Microbatch Size · p. 8

C. End-to-End Performance for Sequence Scaling

BASP’s step-time advantage over Ulysses-SP grows with sequence length, reaching its largest reported gains at 32K tokens. A loss comparison on Llama 3.2-3B also shows overlapping curves for the two methods.

C. End-to-End Performance for Sequence Scaling

D. Loss Comparison Plot

  • After 800 iterations on the same dataset, BASP and Ulysses-SP loss curves completely overlap for Llama 3.2-3B.The authors report this as validation that the method preserves accuracy.Source: Figure 10 · p. 9, D. Loss Comparison Plot · p. 9

VI. LIMITATIONS

BASP reduces Ulysses-SP’s global all-to-all communication by forming independent sequence-parallel groups from the micro-batch structure. The conclusion reports faster Llama and Qwen training, while the paper’s design scope excludes non-divisible configurations.

VI. LIMITATIONS

  • BASP is designed for N ≥ B and N = KB with integer K; support for non-divisible configurations and more flexible group formation remains future work.Source: §VI · p. 9

VII. RELATED WORKS

  • Megatron-SP partitions dropout and layernorm across sequence dimensions, while DeepSpeed-Ulysses uses all-to-all to gather QKV projections but does not address larger-batch all-to-all bottlenecks.Source: §VII · p. 9
  • Ring Attention computes exact attention blockwise but can bottleneck on KV-chunk transfers; Striped Attention improves pipeline utilization by scheduling those transfers more effectively.Source: §VII · p. 9
  • Unified SP switches between all-to-all and ring communication but does not optimize larger-batch all-to-all, while FlexSP adapts strategies at runtime with selection overhead and possible load imbalance.Source: §VII · p. 9
  • MiCS reduces data-parallel communication through heterogeneous sharding and selective mixed precision, but does not address sequence parallelism or topology-aware all-to-all scheduling.Source: §VII · pp. 9–10

VIII. CONCLUSION

  • 1.21× on Llama 3.1-8B and 1.32× on Qwen 1.5-1.8B versus Ulysses-SP, with all-to-all time reduced by up to 67.7% in Qwen 1.5-1.8B.BASP decomposes one N-way collective into independent K-way collectives, reducing communication phases from N − 1 to K − 1 without increasing per-GPU memory footprint or sequence shard size.Source: §VIII · p. 10
Loading 2609.03151v1…