Context Parallelism

Train with long sequences split across multiple GPUs.

Context parallelism (CP, called sequence parallelism by some frameworks and by older Axolotl releases) splits each sequence across a group of GPUs to reduce the activation memory needed for long-context training. Axolotl uses Ringmaster for Ulysses, Ring, and their combination (USP), while Accelerate owns the device mesh and FSDP2 gradient reduction. Under deepspeed: the same context_parallel_size maps onto DeepSpeed’s Ulysses (ALST) implementation through accelerate’s sequence-parallel axis: only the Ulysses backend is available there, ZeRO owns the sharding, and the remaining ranks replicate. Axolotl adds the position_ids Ulysses needs to every sample and weights each rank’s loss by its share of the sequence’s tokens, so per-step loss and gradient norm match a single-process run of the same sequences (DeepSpeed >= 0.18.2).

Configuration

Install pip install 'axolotl[ringmaster]' and use torch >= 2.13. The integration targets Transformers 5.17.0, Accelerate 1.15.0, and axolotl-ringmaster ≥0.2.4.

context_parallel_size: 4
attn_implementation: flash_attention_2
sample_packing: false

Context parallelism is built in; no plugins: entry is needed. The CP degree must divide the total GPU count. For explicit strategy settings:

context_parallel:
  size: 4
  backend: auto
  load_balance: auto

backend accepts auto, ulysses, ring, or usp. Automatic selection uses KV-head count and topology. Explicit ulysses_size and ring_size must multiply to size, and the Ulysses degree must divide the KV-head count.

load_balance: auto selects head-tail balancing for eligible pure Ring and contiguous shards otherwise. Explicit head_tail and distflash require pure Ring and own their P2P communication schedule. With load_balance: none, rotate_method selects Ring KV movement. Incompatible explicit choices fail during setup instead of being silently ignored.

Requirements and limitations

  • Ulysses supports SDPA and Flash Attention. Ring/USP training requires Flash Attention; Ringmaster’s torch_native Ring implementation is forward-only.
  • Sample packing and batch flattening are supported with CP; automatic load balancing selects contiguous shards for packed inputs. GLM DSA kernels remain incompatible with packed CP.
  • SFT requires the model’s causal LM loss with loss-kwargs support. Custom loss functions and label smoothing are not supported.
  • GRPO/GDPO/EBFT use output gathering to retain full-sequence logits.
  • GDN and KDA recurrent models require pip install 'axolotl[fla,ringmaster]', a single batch row after flattening (micro_batch_size: 1 or batch_flattening: true), contiguous shards, and disabled KV caching. Setup checks for importable native FLA CP kernels, including when CP is greater than two. Both UV Docker images include FLA and TileLang.

Ringmaster detects compatible attention and recurrent layers without requiring an architecture allowlist entry. ModelSupport descriptors can mark context_parallel as verified, experimental, or unsupported. Missing capability metadata uses generic handling; it does not imply tested accuracy for every architecture. Kernel and configuration prerequisites still apply.

The plugin shards input IDs, labels, masks, and positions using model hooks. Accelerate replicates each batch within its CP group. Trainer’s supervised-token count is preserved across the full accumulation window, including unequal microbatches and partial final windows.

Effect on batch size

Each CP group processes one batch distributed across its GPUs. With eight GPUs and CP=4, there are two independent data-parallel groups. A microbatch size of two therefore processes four distinct sequences per step before gradient accumulation.

Validation

The regular e2e suite includes an optional small FLA/TileLang kernel smoke test. Distributed checks are marked slow: two-GPU FSDP2 loss/gradient parity, four-rank FLA recurrent-plus-attention forward/backward parity, two-GPU patched GDN forward/backward parity for dense and packed inputs with and without compilation, and CPU USP mesh isolation and attention-gradient tests. The multigpu CI suite runs the two-GPU patched GDN checks explicitly. GPU USP and additional recurrent architectures require further coverage.

Pure Mamba/Mamba2 checkpoints can explicitly select the FLA implementation. Its packing and loss normalization live in Axolotl core; Ringmaster supplies packed convolution-halo and recurrent-state exchange. This requires the companion FLA Mamba adapters, with CUDA convolution and contiguous CP shards.