Automatic Tensor Parallelism (Training)

This tutorial covers Automatic Tensor Parallelism for combining tensor parallelism with ZeRO optimization during training. For inference-only tensor parallelism, see Automatic Tensor Parallelism (Inference).

Contents

Introduction

The AutoTP Training API enables hybrid parallelism by combining:

  • Tensor Parallelism (TP): Split model weights across GPUs within a node
  • Data Parallelism (DP): Replicate model across GPU groups
  • ZeRO Optimization: Memory-efficient optimizer states (Stage 0, 1, or 2)

Tensor parallelism (TP) splits the computations and parameters of large layers across multiple GPUs so each rank holds only a shard of the weight matrix. This is an efficient way to train large-scale transformer models by reducing per-GPU memory pressure while keeping the layer math distributed across the TP group.

Quick Start

Basic Usage

AutoTP training can be enabled entirely through the DeepSpeed config. When tensor_parallel is set in the config, deepspeed.initialize(...) applies AutoTP sharding during engine initialization, so the training loop itself does not change.

import torch
import deepspeed

# 1. Create your model
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.1-8B")

# 2. Define the DeepSpeed config with tensor_parallel settings
ds_config = {
    "train_micro_batch_size_per_gpu": 1,
    "zero_optimization": {"stage": 2},
    "bf16": {"enabled": True},
    "tensor_parallel": {"autotp_size": 4},
}

# 3. Initialize DeepSpeed with AutoTP + ZeRO
engine, optimizer, _, _ = deepspeed.initialize(
    model=model,
    optimizer=optimizer,
    config=ds_config,
    mpu=mpu  # Model parallel unit (optional if you provide tp_group elsewhere)
)

# 4. Train as usual
for batch in dataloader:
    outputs = engine(input_ids=batch["input_ids"], labels=batch["labels"])
    engine.backward(outputs.loss)
    engine.step()

Compatibility note: For backward compatibility, you can still call set_autotp_mode(training=True) and deepspeed.tp_model_init(...), but they are not required when the DeepSpeed config provides the necessary tensor_parallel settings.

Preset-based Sharding

If your model matches a built-in preset, set tensor_parallel.preset_model in the DeepSpeed config:

{
    "train_batch_size": 8,
    "train_micro_batch_size_per_gpu": 1,
    "bf16": { "enabled": true },
    "zero_optimization": { "stage": 2 },
    "tensor_parallel": {
        "autotp_size": 4,
        "preset_model": "llama"
    }
}

For the list of available presets, see supported models.

HuggingFace tp_plan Support

Many HuggingFace models (e.g. Llama, Qwen, Gemma2) ship with a built-in base_model_tp_plan in their model config that describes how each layer should be partitioned for tensor parallelism. DeepSpeed can automatically detect and use this plan, so you do not need to configure preset_model or partition_config for these models.

When tensor_parallel is set in the DeepSpeed config, the initialization follows this priority:

  1. Custom partition_config (highest): User-defined regex patterns.
  2. HuggingFace tp_plan: Automatically extracted from model._tp_plan or model.config.base_model_tp_plan.
  3. AutoTP heuristics (lowest): Built-in parser based on module structure.

For models that define a tp_plan, you only need a minimal config:

{
    "train_micro_batch_size_per_gpu": 1,
    "zero_optimization": { "stage": 2 },
    "bf16": { "enabled": true },
    "tensor_parallel": { "autotp_size": 4 }
}

DeepSpeed will read the model’s tp_plan at initialization and convert it to internal partition rules. The supported types are colwise, rowwise, and colwise_gather_output(colwise_rep). The gathered column styles shard the linear weight along its output dimension and AllGather the local output shards so every tensor-parallel rank receives the complete output. For untied output layers, the output dimension does not need to be divisible by autotp_size; DeepSpeed uses uneven local shards and gathers back to the original logical output size.

Gathered column parallelism currently supports untied output layers. If an output layer such as lm_head shares the same runtime Parameter object with an embedding, DeepSpeed leaves both modules replicated and applies tensor parallelism to the remaining matched layers. This preserves the tie without silently cloning the weight, but does not reduce the embedding or output-layer memory footprint. A coupled vocabulary-parallel embedding is required to shard the tied weight and is not yet implemented. This fallback uses actual Parameter identity rather than model configuration metadata such as tie_word_embeddings.

Additional HuggingFace types such as local_colwise and local_rowwise are not yet handled and fall back to AutoTP preset-based partitioning.

For an lm_head or embed_out, an explicit row partition rule also supports eager training with a replicated input. The head slices the input to match its weight shard, reduces the complete output, and reconstructs the full input gradient across tensor-parallel ranks. Bias stays replicated; uneven hidden dimensions and both flattened and sequence-shaped inputs are supported. The default training output-head layout remains column parallel. A head tied to its input embedding (tie_word_embeddings) is also supported: DeepSpeed shards every tied nn.Embedding along the same hidden dimension, so the pair keeps reading from and accumulating gradients into a single per-rank weight shard, and all-gathers each token’s embedding slices into the full hidden activation. A conflicting explicit spec on the tied embedding is superseded with a warning. The embedding options max_norm and sparse, custom embedding forwards other than Gemma3’s scaled embedding, weights tied to modules other than nn.Embedding, and reshaped/non-input-dimension shards are rejected rather than silently breaking a parameter tie. Deferred DeepCompile collectives are not supported for this output-head path. Note that a row-parallel head all-reduces the complete vocabulary logits; for tied models, vocab_parallel_lm_head avoids that and keeps logits sharded.

If you need to override the model’s built-in tp_plan, provide a partition_config in the DeepSpeed config – it takes precedence.

Vocabulary-parallel LM Loss

Causal language models normally gather the complete lm_head output before computing cross entropy. To keep an output vocabulary sharded, enable vocab_parallel_lm_head:

{
    "train_micro_batch_size_per_gpu": 1,
    "zero_optimization": { "stage": 2 },
    "tensor_parallel": {
        "autotp_size": 4,
        "vocab_parallel_lm_head": true
    }
}

DeepSpeed then keeps lm_head (or embed_out) local to each TP rank and installs a pure-PyTorch distributed causal-LM loss through the model’s loss_function hook. The loss computes a numerically stable distributed log-sum-exp and target lookup without gathering vocabulary logits. Uneven vocabulary shards are supported.

For optional fused CE, install liger-kernel>=0.8.1 and set tensor_parallel.vocab_parallel_ce_backend to "liger" alongside vocab_parallel_lm_head. The default remains "torch", with no Liger dependency. The Liger backend accepts nonempty accelerator logits in FP32, FP16, or BF16 with equal vocabulary shards behind an explicit tensor-parallel group, on any accelerator whose Triton support the DeepSpeed accelerator reports. Uneven shards, a missing tensor-parallel group, and unsupported layouts use the PyTorch backend on every TP rank. This fuses CE, not the linear projection: local logits are still materialized. Liger supports one first-order backward per forward; retained-graph second backward and higher-order gradients raise an error rather than reusing its overwritten gradient buffer. Select "torch" for those workloads.

Only an nn.Linear whose final name segment is lm_head or embed_out is supported. With vocab_parallel_lm_head: true, if no such head exists, initialization fails instead of quietly falling back to an ordinary gathered head. When partition_config or a HuggingFace tp_plan also describes that head, vocab_parallel_lm_head takes precedence and a warning names the superseded partitioning.

Models that tie the output head to the input embedding (tie_word_embeddings) are also supported: DeepSpeed detects the shared nn.Embedding and vocabulary- shards it jointly with lm_head, including all registered aliases of the embedding and distinct embedding modules sharing the same weight. These modules keep reading from and accumulating gradients into a single physical, per-rank-sharded weight Parameter instead of falling back to a gathered, replicated head. This also covers the HuggingFace tp_plan embedding_rowwise style that newer transformers releases inject for tied-embedding models. That style automatically enables the same vocabulary-parallel embedding/head path when vocab_parallel_lm_head is omitted or null, so no separate enable flag is required. This returns rank-local outputs.logits; set vocab_parallel_lm_head: false to opt out and keep the replicated tied embedding/head with full-vocabulary logits. Tied models whose plan does not contain embedding_rowwise retain the previous replicated behavior unless the flag is enabled explicitly. A conflicting explicit partition_config spec on the tied embedding is likewise superseded, with a warning.

Vocabulary-parallel embeddings preserve padding_idx gradient suppression and Gemma3’s scaled-embedding forward behavior. Other custom embedding forwards and the max_norm, scale_grad_by_freq, and sparse options are unsupported. Implicit automatic sharding checks compatibility before mutating weights; an unsupported embedding, missing or unwritable loss hook, empty vocabulary shard, or ambiguous output head retains the replicated embedding/head and logs a warning. Explicit true requests raise instead of falling back. Tied embeddings retain the full vocabulary shape in universal-checkpoint metadata, including uneven shards.

With compile.deepcompile: true and "autotp" in compile.passes, automatic sharding from embedding_rowwise is disabled: the tied embedding and output head remain replicated with gathered logits, while the other supported layers are still tensor-parallel. An explicit vocab_parallel_lm_head: true request is not downgraded; vocabulary-parallel embeddings are not supported by the AutoTP compile pass.

This option requires a model with a writable loss_function hook. The head’s vocabulary must also produce a nonempty shard on every TP rank; smaller vocabularies or a grain size that leaves empty shards trigger the compatibility behavior above instead of leaving ranks with empty logits. An explicit true flag or a supported tied HuggingFace embedding_rowwise plan entry triggers this path. A plain colwise lm_head specification with local output keeps the previous behavior of returning rank-local logits without installing the distributed loss.

The lower-level vocab_parallel_cross_entropy API also accepts an explicit sequence-parallel group. With reduction="none", callers may return local token losses or gather them along sequence dimension 0. sum and mean reduce over the supplied SP group. TP and SP may be combined with explicit orthogonal process groups, but AutoTP does not currently construct a combined TP x SP mesh automatically.

Gathered sequence losses have two backward conventions. The general vocab_parallel_cross_entropy(..., gather_sequence_loss=True) API sums the gradient contributions from every SP rank before returning each rank’s local slice. The compatibility wrapper vocab_sequence_parallel_cross_entropy preserves its legacy behavior and returns only the corresponding local slice of the gathered loss gradient. New callers that consume or reduce the gathered loss on every SP rank should use the general API; existing callers can retain the wrapper without changing their gradient scale. Both gathered forms require an explicit sp_group.

Under DeepSpeed’s Ulysses sequence-parallel engine, the installed loss must keep its default sequence-parallel settings (no sp_group): the engine aggregates each shard’s mean itself, weighted by the shard’s valid-token count, so an additional SP reduction would double-count tokens. Pass an sp_group only when this loss is the sole aggregation over a manually constructed TP x SP mesh.

Custom Patterns

If you are training a custom model, define regex-based patterns and partition rules in tensor_parallel.partition_config:

{
    "tensor_parallel": {
        "autotp_size": 4,
        "partition_config": {
            "use_default_specs": false,
            "layer_specs": [
                {
                    "patterns": [".*\\.o_proj\\.weight$", ".*\\.down_proj\\.weight$"],
                    "partition_type": "row"
                },
                {
                    "patterns": [".*\\.[qkv]_proj\\.weight$"],
                    "partition_type": "column"
                },
                {
                    "patterns": [".*\\.gate_up_proj\\.weight$"],
                    "partition_type": "column",
                    "shape": [2, -1],
                    "partition_dim": 0
                }
            ]
        }
    }
}

Custom Layer Specifications

For models not covered by presets, define custom layer specs:

{
    "tensor_parallel": {
        "autotp_size": 4,
        "partition_config": {
            "use_default_specs": false,
            "layer_specs": [
                {
                    "patterns": [".*\\.o_proj\\.weight$", ".*\\.down_proj\\.weight$"],
                    "partition_type": "row"
                },
                {
                    "patterns": [".*\\.[qkv]_proj\\.weight$"],
                    "partition_type": "column"
                },
                {
                    "patterns": [".*\\.gate_up_proj\\.weight$"],
                    "partition_type": "column",
                    "shape": [2, -1],
                    "partition_dim": 0
                }
            ]
        }
    }
}

Fused Layers with Unequal Sub-parameters (GQA)

For Grouped Query Attention with different Q/K/V sizes:

{
    "tensor_parallel": {
        "partition_config": {
            "layer_specs": [
                {
                    "patterns": [".*\\.qkv_proj\\.weight$"],
                    "partition_type": "column",
                    "shape": [[q_size, kv_size, kv_size], -1],
                    "partition_dim": 0
                }
            ]
        }
    }
}

Limitations

  1. Ranks beyond the key/value head count stay idle: Attention heads are distributed whole, and the distribution may be uneven – 6 key/value heads over 4 ranks becomes 2/2/1/1, and a fused QKV weight is cut on the same head boundaries rather than inside a head. With more ranks than key/value heads, for example an 8-head model at autotp_size=16, the surplus ranks receive no attention weights at all. The result is still correct, because those ranks contribute zeros to the row-parallel all-reduce, but they do no attention work; AutoTP logs a warning instead of replicating heads to fill them. Hidden and vocabulary dimensions do not need to be divisible by the tensor parallel size: uneven shards are carried through save, conversion and restore via per-TP-rank shapes and widths.

  2. Vocabulary-parallel LM loss requires a writable loss_function hook: vocab_parallel_lm_head requires a model with a writable loss_function hook, and a vocabulary at least as large as the tensor-parallel size. Both tied and untied lm_head/embed_out heads are supported; a tied embedding is jointly vocabulary-sharded with its output head.

  3. Combined TP and SP orchestration: The vocab-parallel loss supports explicit orthogonal TP and SP groups, but AutoTP does not currently construct a combined TP x SP process mesh automatically.

  4. Cross-topology universal restore: Loading a universal checkpoint back into a topology with a different tensor-parallel degree goes through DeepSpeed’s Megatron-style model-state loader, which is not AutoTP-aware; prefer same-topology restore when changing world size.

See Also

Updated: