Workflow:Huggingface Transformers 3D Parallel Distributed Training
| Knowledge Sources | |
|---|---|
| Domains | LLMs, Distributed_Training, Training |
| Last Updated | 2026-02-13 20:00 GMT |
Overview
End-to-end process for training causal language models using a combination of Tensor Parallelism, Data Parallelism (FSDP), and Context Parallelism across multiple GPUs.
Description
This workflow demonstrates advanced distributed training using 3D parallelism, which combines three complementary strategies: Tensor Parallelism (TP) splits individual model layers across devices, Data Parallelism (DP via FSDP) replicates the model and distributes batches, and Context Parallelism (CP) splits long sequences across devices. The approach enables training of large models that exceed single-GPU memory capacity while maximizing throughput. The workflow uses PyTorch's native distributed primitives (DeviceMesh, FSDP, DTensor) with Transformers' built-in TP support to orchestrate the parallelism strategies.
Usage
Execute this workflow when training large language models (7B+ parameters) that require multiple GPUs either due to memory constraints or for throughput optimization. This is appropriate when you need to combine model sharding (TP), batch distribution (DP), and sequence splitting (CP) to fully utilize a multi-GPU cluster. Requires NCCL-enabled GPUs and torchrun launcher.
Execution Steps
Step 1: Distributed Environment Initialization
Initialize the NCCL distributed process group and configure parallelism dimensions. Read TP_SIZE, DP_SIZE, and CP_SIZE from environment variables. Validate that world_size equals the product of all parallelism dimensions. Set each rank's CUDA device based on local rank.
Key considerations:
- World size must equal TP_SIZE x DP_SIZE x CP_SIZE
- Launched via torchrun --nproc_per_node=N with environment variables for each parallelism dimension
- NCCL backend required for GPU-to-GPU communication
Step 2: Device Mesh Construction
Create a 3D DeviceMesh that organizes GPU ranks into a logical grid with named dimensions ("dp", "tp", "cp"). Extract sub-meshes for each parallelism type to use in subsequent operations. Flatten DP and CP dimensions into a combined mesh for cross-dimension gradient synchronization.
Key considerations:
- DeviceMesh defines the communication topology for all parallelism operations
- Sub-meshes enable independent communication within each parallelism dimension
- The combined dp_cp mesh is used for gradient all-reduce operations
Step 3: Model Loading with Tensor Parallelism
Load the pretrained model using AutoModelForCausalLM with tensor parallelism configuration. Pass the TP sub-mesh and tp_plan="auto" to automatically shard model weights across TP ranks. Load with bfloat16 precision for memory efficiency.
Key considerations:
- device_mesh=tp_mesh enables automatic weight sharding across TP ranks
- tp_plan="auto" lets Transformers determine optimal sharding strategy per layer
- Model weights are distributed as DTensors (Distributed Tensors) across TP ranks
Step 4: FSDP Data Parallelism Wrapping
If data parallelism is enabled (DP > 1), wrap the tensor-parallel model with Fully Sharded Data Parallel (FSDP). Use NO_SHARD strategy for compatibility with the existing tensor parallelism sharding. This distributes batch processing across DP ranks.
Key considerations:
- FSDP wrapping happens after TP sharding to avoid conflicts
- ShardingStrategy.NO_SHARD is used when combining with TP (model already sharded)
- Each DP rank processes a different subset of the training batch
Step 5: Dataset Preparation and Packing
Load the training dataset, tokenize all examples, and pack them into fixed-length sequences to eliminate padding waste. Flatten all tokenized sequences into a continuous stream, then re-chunk into uniform blocks of the target sequence length. Create input-label pairs by shifting the packed sequences by one token.
Key considerations:
- Sequence packing maximizes GPU utilization by eliminating padding tokens
- Labels are created by shifting input_ids by one position (standard causal LM objective)
- Use distributed sampler to ensure each DP rank receives non-overlapping data
Step 6: Training Loop with Context Parallelism
Execute the training loop with forward pass wrapped in Context Parallelism context manager. CP shards the input sequence across CP ranks, with each rank processing a portion of the sequence. Use Flash Attention backend within CP for efficient cross-rank attention computation. After the forward pass, compute loss and run backward pass.
Key considerations:
- Context Parallelism shards input_ids, labels, and position_ids along the sequence dimension
- Flash Attention is required for efficient CP attention computation
- Loss is computed locally on each rank's sequence shard
Step 7: Gradient Synchronization and Update
Synchronize gradients across all DP and CP ranks using all-reduce on the combined dp_cp mesh. Handle both DTensor and regular tensor gradients appropriately. Apply gradient clipping to prevent exploding gradients. Execute the optimizer step to update model parameters.
Key considerations:
- Gradient all-reduce averages gradients across DP and CP dimensions
- DTensor gradients require special handling (extract local, all-reduce, reconstruct)
- Gradient clipping to max_norm=1.0 stabilizes training
Step 8: Distributed Checkpointing
Save the trained model and optimizer state using PyTorch's Distributed Checkpoint (DCP) protocol. DCP handles saving sharded model state across all ranks without gathering to a single rank, enabling efficient checkpointing of large distributed models.
Key considerations:
- DCP saves each rank's shard independently, avoiding memory bottleneck
- Checkpoint directory is named by parallelism configuration for easy identification
- Both model and optimizer state are saved for training resumption