Workflow:Deepspeedai DeepSpeed AutoTP Training
| Knowledge Sources | |
|---|---|
| Domains | Distributed_Training, Tensor_Parallelism, LLMs |
| Last Updated | 2026-02-09 00:00 GMT |
Overview
End-to-end process for training HuggingFace transformer models with automatic tensor parallelism (AutoTP), partitioning model layers across GPUs without manual sharding code.
Description
This workflow covers training large language models using DeepSpeed's Automatic Tensor Parallelism (AutoTP), which automatically detects and partitions model layers across multiple GPUs for tensor-parallel training. Unlike pipeline parallelism which splits layers sequentially, tensor parallelism splits individual layers (attention heads, MLP blocks) across GPUs, allowing each GPU to participate in every layer's computation. AutoTP eliminates the need for manual model modification by automatically identifying compatible layers (linear projections, attention, MLP) and applying the correct sharding strategy. It integrates seamlessly with HuggingFace Transformers and can be combined with ZeRO optimization for additional memory savings.
Usage
Execute this workflow when training HuggingFace transformer models (LLaMA, Phi, GPT-NeoX, Mistral, etc.) where you want to use tensor parallelism to enable larger per-device batch sizes or to fit models that are too large for data parallelism alone. AutoTP is especially useful when you want tensor parallelism benefits without writing custom model sharding code. Use this when you have 2+ GPUs per node connected with high-bandwidth interconnect (NVLink/NVSwitch).
Execution Steps
Step 1: Model Selection and Loading
Select a HuggingFace transformer model architecture that is supported by AutoTP. Load the model using standard HuggingFace APIs. AutoTP supports common architectures including LLaMA, Phi, GPT-NeoX, Mistral, Mixtral, and others where attention and MLP layers follow standard patterns.
Key considerations:
- Model must use standard transformer patterns (self-attention + MLP)
- AutoTP auto-detects QKV projections, output projections, and MLP layers
- Custom models can be supported through custom pattern specifications
- Load model in the desired training precision (float16, bfloat16)
Step 2: DeepSpeed Configuration with AutoTP
Create a DeepSpeed configuration that includes the tensor_parallel section with autotp_size set to the desired number of tensor parallel ranks. Optionally combine with ZeRO optimization (typically Stage 1 with AutoTP). The configuration must account for the relationship between tensor parallel size and data parallel size.
Key considerations:
- autotp_size specifies how many GPUs participate in each tensor parallel group
- Data parallel size is automatically computed as total_gpus / autotp_size
- Global batch size calculation changes: per_device_batch * (num_gpus / tp_size) * grad_accum
- Enable gather_16bit_weights_on_model_save for proper checkpoint saving
- ZeRO Stage 1 is typically used alongside AutoTP
Step 3: Engine Initialization with AutoTP
Call deepspeed.initialize() with the model and AutoTP configuration. During initialization, DeepSpeed analyzes the model architecture, identifies tensor-parallel-compatible layers, and applies the appropriate sharding (column-wise for QKV/MLP-up, row-wise for output/MLP-down). Communication groups for allreduce and allgather are established automatically.
Key considerations:
- AutoTP modifies the model in-place during initialization
- Weight sharding happens automatically based on detected layer patterns
- Communication groups are created for tensor-parallel allreduce operations
- The engine internally calls set_autotp_mode(training=True) to enable training-specific behavior
Step 4: Training with Tensor Parallelism
Execute the training loop using the DeepSpeed engine. During forward and backward passes, tensor-parallel communication (allreduce for row-parallel outputs, allgather for column-parallel inputs) happens automatically. Each GPU computes its shard of every layer and synchronizes results as needed.
Key considerations:
- Forward pass includes automatic allreduce after row-parallel layers
- Backward pass handles gradient synchronization across TP groups
- Data parallelism handles gradient averaging across DP groups
- Per-GPU memory usage decreases with TP size for model parameters and activations
Step 5: Model Saving and Deployment
Save the trained model using DeepSpeed's checkpoint mechanisms. With gather_16bit_weights_on_model_save enabled, the engine gathers tensor-parallel shards back into full weights for saving. The resulting checkpoint is a standard format that can be loaded without tensor parallelism for deployment.
Key considerations:
- Enable gather_16bit_weights_on_model_save in ZeRO config for unified checkpoints
- Saved weights are in standard HuggingFace format for easy deployment
- Training can be resumed with different TP sizes using Universal Checkpointing
- HuggingFace Trainer integration handles save/load automatically