--- title: "Automatic Tensor Parallelism (Training)" tags: training tensor-parallelism --- 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)](/tutorials/automatic-tensor-parallelism/). ## Contents - [Introduction](#introduction) - [Quick Start](#quick-start) - [HuggingFace tp_plan Support](#huggingface-tp_plan-support) - [Custom Layer Specifications](#custom-layer-specifications) - [Limitations](#limitations) ## 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. ```python 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: ```json { "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](/code-docs/training#autotp-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: ```json { "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. Currently `colwise` and `rowwise` partition types are supported. Additional types defined by HuggingFace (such as `colwise_rep`, `local_colwise`, `local_rowwise`, etc.) are not yet handled and will raise an error if encountered. If you need to override the model's built-in `tp_plan`, provide a `partition_config` in the DeepSpeed config -- it takes precedence. ## Custom Patterns If you are training a custom model, define regex-based patterns and partition rules in `tensor_parallel.partition_config`: ```json { "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: ```json { "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: ```json { "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. **ZeRO Stage 3 not supported**: AutoTP currently only works with ZeRO stages 0, 1, and 2. 2. **TP size must divide model dimensions**: The tensor parallel size must evenly divide the attention head count and hidden dimensions. ## See Also - [Automatic Tensor Parallelism (Inference)](/tutorials/automatic-tensor-parallelism/) - [ZeRO Optimization](/tutorials/zero/) - [DeepSpeed Configuration](/docs/config-json/)