Skip to main content

N-D Parallelism (Beta)

Axolotl enables training models at scale by composing different parallelism techniques. This is essential when:

or combinations of the above!

Core Concepts

Parallelism strategies can be combined. The key is understanding how each one divides the workload. PyTorch’s DeviceMesh is the modern way to manage these combinations, creating a logical grid of your GPUs and assigning different parallel strategies to different dimensions of the grid.

Data Parallelism

Data Parallelism focuses on splitting the global data batch across GPUs.

  • Distributed Data Parallel (DDP): The classic approach. The full model is replicated on every GPU. Each GPU processes a different slice of the data batch. Gradients are then averaged across all GPUs after the backward pass to keep the models synchronized. This can substantially improve data throughput compared to single-device training, but requires that each GPU is able to hold the entire model, its gradients, and optimizer states.

  • Fully Sharded Data Parallel (FSDP): A highly memory-efficient form of data parallelism (inspired by DeepSpeed’s ZeRO). Instead of replicating the model, FSDP shards the model’s parameters, gradients, and optimizer states across the GPUs in the data-parallel group. During computation, each GPU receives the specific parameters it needs via an all_gather operation just before they are used, and they can be discarded immediately after (reshard-after-forward).

    • FSDP maps to ZeRO stages:
      • ZeRO-2 (reshard_after_forward=False): Shards gradients and optimizer states. Model weights are replicated on each GPU.
      • ZeRO-3 (reshard_after_forward=True): Shards gradients, optimizer states, AND model parameters. This provides the most memory savings at the cost of more communication (re-gathering parameters for both forward and backward passes).

[Experimental] Tensor Parallelism (TP)

Also known as “horizontal model parallelism,” as described in the Megatron-LM paper. Instead of splitting the batch, TP splits the model’s layers themselves across GPUs.

  • How it works: For a linear layer Y = XA, the weight matrix A is split column-wise (A = [A_1, A_2]). The computation becomes Y_1 = XA_1 and Y_2 = XA_2, which can happen in parallel on different GPUs. The final output Y is simply the concatenation of Y_1 and Y_2. Check this comment for more detailed info.
  • Requirement: TP involves frequent, small communications within a forward/backward pass. It requires a very fast interconnect between GPUs (e.g., NVLink) and is typically not recommended across different nodes.
  • Limits: full-parameter training, or LoRA on a native NVFP4 base (the torchao_tp_lora path). Adapters on a bf16 base are rejected at model load because PEFT sizes the LoRA factors from the already TP-sharded base weight, so each rank would train a different shard-local adapter. fsdp_config.cpu_ram_efficient_loading is turned off under TP (each rank loads its own shard directly; a rank-0-only load would leave the other ranks with zeros). Every TP rank takes part in the final save_pretrained, and the logged grad_norm is the global norm.

Context Parallelism (CP)

Context Parallelism (CP; the term we use for what some frameworks call sequence parallelism, see the guide) addresses the memory bottleneck from long sequences. The input sequence itself is split along the sequence length dimension and distributed across the GPUs of a cp group; every GPU in the group sees the same batch.

  • How it works: If you have a sequence of 8192 tokens and a context_parallel_size of 4, each GPU will only handle a chunk of 2048 tokens.
  • The Challenge: Attention is not local; every token needs to “attend to” every other token. Splitting the sequence breaks this.
  • The Solution (Ringmaster): CP is built in via axolotl-ringmaster and needs no plugins: entry. With context_parallel.backend: auto (the default) it prefers Ulysses, which all-to-alls the attention heads across the group so each GPU runs ordinary attention on a subset of heads over the full sequence, and covers any remainder of the CP degree with Ring attention (KV blocks passed around the group), a combination known as USP. Ulysses runs on SDPA or Flash Attention; Ring/USP training needs Flash Attention. Set backend: ulysses, ring, or usp (with ulysses_size / ring_size) to override.
  • Previous releases implemented CP with ring-flash-attention under sequence_parallel_degree; that key is deprecated and maps to context_parallel_size, and ring_attn_func / heads_k_stride are ignored.
  • Mamba/SSM and linear-attention hybrids (Nemotron-H, Falcon-H1, Granite MoE Hybrid, GDN/KDA): Ringmaster shards the recurrent layers too, passing state across ranks. See Context Parallelism for requirements.

Expert Parallelism (EP)

Expert Parallelism shards the experts of a Mixture-of-Experts (MoE) model across GPUs. Each rank holds num_experts / ep_size experts; tokens are routed to the rank that owns the chosen expert via a token dispatch/combine collective.

  • How it works: the EP plugin slices each Experts module’s gate_up_proj / down_proj along the experts dimension and replaces the dispatch/combine path with the chosen backend. Each rank sees a different batch (so EP is data-parallel for data) but holds different expert shards (model-parallel for expert weights).
  • When to use: MoE training where the experts dominate the parameter count and don’t fit under FSDP alone, or where you want to compose with FSDP on orthogonal mesh axes.
  • Requires FSDP2 (fsdp_version: 2 + fsdp_config) with fsdp_config.state_dict_type: FULL_STATE_DICT; EP under DDP or with a sharded state dict is rejected at config validation (a sharded checkpoint would keep only EP group 0’s experts). Composes with FSDP, CP and HSDP: dp_replicate_size × expert_parallel_size × dp_shard_size × context_parallel_size = world_size. EP shards experts on the ep axis; the dense parameters FSDP-shard across every other axis; the expert slices FSDP-shard across the dp_shard and cp ranks of their EP group (both see different tokens, so expert gradients reduce over both), with dp_replicate as an outer replication dim under HSDP (with no dp_shard/cp axis each replica keeps its expert slices whole and the replicas reduce expert gradients with each other). When expert_parallel_size equals the world size (pure EP) the trainable experts stay whole on each rank. Full-parameter and LoRA (lora_target_parameters on the experts) training both work; the saved checkpoint or adapter carries every expert. EP × TP is not supported.
  • Two dispatch backends, set with expert_parallel_backend (default auto): deep_ep uses the DeepEP library’s fused kernels (Ampere or Hopper, all-pairs NVLink/RDMA); torch dispatches over a plain all_to_all_single on the EP process group and runs on any NCCL/gloo fabric, including PCIe-only nodes with no NVLink. auto picks deep_ep when it’s importable, else torch. See Expert Parallelism Integration for backend details, install steps, and the torch backend’s checkpointing note.
  • Checkpoints and resume: mid-training checkpoints hold every expert of every EP group, with its optimizer state, for full-parameter training and for LoRA on the experts. Incomplete legacy checkpoints containing only EP group 0’s experts are refused on resume. See Training checkpoints and resume for the file contents and layout, how to recover from an old checkpoint, and limits.

Hybrid Sharding Data Parallel (HSDP)

HSDP is a 2D strategy that intelligently combines FSDP and DDP, typically for multi-node training.

  • Intra-Node (within a machine): Use FSDP. This is efficient because GPUs on the same node have fast interconnects (NVLink), making the all_gather operations for sharded parameters fast.
  • Inter-Node (across machines): Use DDP. The gradient synchronization between nodes is less frequent than FSDP’s parameter gathering, making it a better fit for the slower node-to-node network (e.g., Ethernet/Infiniband).
  • Example: With 2 nodes of 8 GPUs each (16 total), you could have dp_shard_size=8 (FSDP within each node) and dp_replicate_size=2 (DDP across the two nodes).

Usage

# FSDP config. See https://docs.axolotl.ai/docs/multi-gpu.html#sec-fsdp
fsdp_version: 2
fsdp_config:
  # ...

# The number of GPUs to shard the model parameters across (FSDP dimension).
dp_shard_size: 4

# The number of times to replicate the sharded model (DDP dimension).
dp_replicate_size: 2

# Number of GPUs for Tensor Parallelism.
tensor_parallel_size: 1  # (default is 1, no TP)

# Number of GPUs for Context Parallelism.
context_parallel_size: 1 # (default is 1, no CP)

# Number of GPUs to shard MoE experts across (Expert Parallel).
expert_parallel_size: 1  # (default is 1, no EP)

Note: We recommend FSDP. DeepSpeed is only compatible with tensor_parallel_size.

Examples

Tip

See our example configs here.

  1. HSDP on 2 nodes with 4 GPUs each (8 GPUs total):
    • You want FSDP within each node and DDP across nodes.
    • Set dp_shard_size: 4 and dp_replicate_size: 2.
  2. FSDP + TP on a single 8-GPU node:
    • You want to split the model across 4 GPUs using FSDP, and further split each layer across 2 GPUs with TP.
    • Set dp_shard_size: 4 and tensor_parallel_size: 2.
  3. FSDP + CP on a single 8-GPU node for long context:
    • You want to shard the model across all 8 GPUs and also split the sequence length across all 8 GPUs.
    • Set dp_shard_size: 8 and context_parallel_size: 8. Note: this means the data parallel group and context parallel group are the same. A more common setup might be to shard across a smaller group.
  4. FSDP + EP on a 4-GPU MoE training run:
    • You want EP to shard the experts and FSDP to shard non-expert params on orthogonal mesh axes.
    • Set expert_parallel_size: 2 and dp_shard_size: 2 (ep × dp_shard == world_size).
  5. EP + CP on a 4-GPU long-context MoE run:
    • You want the experts sharded and each sequence split across 2 GPUs.
    • Set expert_parallel_size: 2 and context_parallel_size: 2. Add dp_shard_size: 2 on 8 GPUs to also shard the dense weights (ep × cp × dp_shard == world_size).
    • Axis order is dp_replicate, dp_shard, cp, ep, tp (outer to inner): ep groups are consecutive ranks and stay node-local whenever ep × tp fits on a node, while dp_shard groups stride across nodes. The startup log prints the mesh and which axes cross nodes.
  6. HSDP + EP on 2 nodes of 4 GPUs:
    • Set dp_replicate_size: 2, expert_parallel_size: 2, dp_shard_size: 2.

Minimal MoE configs for the two EP compositions (4 and 8 GPUs; everything else is an ordinary FSDP2 config):

# EP x CP on 4 GPUs: experts sharded 2 ways, each sequence split across 2 GPUs
base_model: Qwen/Qwen3-30B-A3B
expert_parallel_size: 2
context_parallel_size: 2
fsdp_version: 2
fsdp_config:
  auto_wrap_policy: TRANSFORMER_BASED_WRAP
  transformer_layer_cls_to_wrap: Qwen3MoeDecoderLayer
  state_dict_type: FULL_STATE_DICT
  cpu_ram_efficient_loading: true
attn_implementation: flash_attention_2
sample_packing: false
# HSDP x EP on 2 nodes x 4 GPUs: replicate across nodes, EP x FSDP inside each node
base_model: Qwen/Qwen3-30B-A3B
dp_replicate_size: 2
expert_parallel_size: 2
dp_shard_size: 2
fsdp_version: 2
fsdp_config:
  auto_wrap_policy: TRANSFORMER_BASED_WRAP
  transformer_layer_cls_to_wrap: Qwen3MoeDecoderLayer
  state_dict_type: FULL_STATE_DICT
  cpu_ram_efficient_loading: true

With EP active, dp_shard_size is only the extra FSDP split applied to the expert modules inside an EP group; the dense (non-expert) weights always shard over the flattened dp_shard x cp x ep axes. On a node of 8 GPUs, expert_parallel_size: 8, dp_shard_size: 1 therefore means dense weights sharded 8 ways within the node and each rank keeping its experts whole, and adding dp_replicate_size: 2 replicates that layout on a second node (HSDP) so only gradients cross the inter-node fabric. The startup log prints this effective layout.

expert_parallel_backend defaults to auto (DeepEP when installed, else the torch all-to-all backend); LoRA on the experts works with lora_target_parameters in either layout. A fresh LoRA adapter is drawn per module from seed, so the init (and a one-step gradient check) is identical whether or not the experts are sharded.

Support Matrix

This matrix describes how different parallelism methods can be combined in Axolotl.

Combination dp_replicate_size dp_shard_size tp_size cp_size ep_size Status & Notes
FSDP (ZeRO-3) 1 >1 1 1 1 ✅ Fully supported. Shards model across all GPUs.
HSDP >1 >1 1 1 1 ✅ Fully supported. FSDP intra-node, DDP inter-node.
FSDP + TP 1 >1 >1 1 1 ✅ 2D Parallelism. Shards the model across a dp_shard group, and TP-splits layers within the tp group. Verified on dense and MoE (grouped_mm) models against FSDP-only; full-parameter, or LoRA on native NVFP4 bases.
HSDP + TP >1 >1 >1 1 1 ⚠️ 3D Parallelism. Same code path as FSDP + TP; not yet verified for parity.
FSDP + CP 1 >1 1 >1 1 ✅ 2D Parallelism. Combines FSDP with context parallelism; gradients verified identical to plain FSDP at the same data-parallel degree (Ulysses, 8 GPUs).
HSDP + CP >1 >1 1 >1 1 ✅ 3D Parallelism. Verified with FSDP + CP; needs dp_shard_size > 1.
FSDP + TP + CP 1 >1 >1 >1 1 ⚠️ 3D Parallelism. Same code path as FSDP + TP; not yet verified for parity.
EP 1 1 1 1 =world ✅ Supported (MoE only). Shards experts across all ranks; dense weights FSDP-sharded across the world. Every EP layout requires FSDP2 with state_dict_type: FULL_STATE_DICT.
FSDP + EP 1 >1 1 1 >1 ✅ 2D Parallelism (MoE). EP shards experts on the ep axis; FSDP shards non-experts on dp_shard.
EP + CP 1 1 1 >1 >1 ✅ 2D Parallelism (MoE, long context). Experts on ep, sequences on cp; expert slices FSDP-shard across cp.
FSDP + EP + CP 1 >1 1 >1 >1 ✅ 3D Parallelism (MoE). Expert slices shard across dp_shard × cp.
HSDP + EP >1 >1 1 1 >1 ✅ 3D Parallelism (MoE, multi-node). dp_replicate replicates the whole EP × FSDP layout; needs dp_shard_size > 1 (or cp_size > 1, next row).
HSDP + EP, no shard axis >1 1 1 1 >1 ✅ 2D Parallelism (MoE, multi-node). dp_replicate replicates the pure-EP layout; the dense weights FSDP-shard over ep inside each replica and the expert slices stay whole per rank, reducing across replicas.
HSDP + EP + CP >1 1 1 >1 >1 ✅ 3D Parallelism (MoE, multi-node, long context). dp_replicate outermost; the dense weights FSDP-shard over ep × cp and the expert slices over cp, so no dp_shard axis is needed (accelerate’s own DDP + CP rejection is lifted when ep supplies the shard group).
EP + TP 1 * >1 * >1 ❌ Not supported; rejected at config validation.
DDP + TP/CP >1 1 >1 >1 1 ❌ Not Supported. The ParallelismConfig explicitly prevents this, as composing pure DDP with TP or CP is currently not supported. You should use FSDP + TP/CP instead (dp_shard_size > 1).
Just TP / CP 1 1 >1 >1 1 ✅ Supported. Useful for inference or when the model fits on one GPU but context is too long.
Note

Axolotl normalizes the loss per rank (average_tokens_across_devices is off), so without CP each rank’s tokens are averaged separately and the ranks are then averaged. Under CP the count is the global one for the CP group, so the same batch is token-weighted slightly differently than an FSDP-only run with the same data-parallel degree. Both are valid mean losses; they only coincide when every rank holds the same number of supervised tokens.

  • tp_size refers to tensor_parallel_size
  • cp_size refers to context_parallel_size
  • sp is accelerate’s DeepSpeed Ulysses axis; axolotl sets it only when context_parallel_size runs under deepspeed:. Older configs’ sequence_parallel_degree is a deprecated alias of context_parallel_size.
  • ep_size refers to expert_parallel_size