Tutorials/10-distributed-training-scale.md

Tutorial 10: Distributed Training at Scale

Overview

When models exceed the memory capacity of a single GPU (e.g., Llama-3-70B, Mixtral), or when training time needs to be reduced from months to days, we must distribute the workload across multiple GPUs and nodes. This tutorial covers NTF's distributed training capabilities via Accelerate, along with the three pillars of distributed training: Data Parallelism, Tensor Parallelism, and Pipeline Parallelism, along with ZeRO optimization.

Note: NTF currently supports multi-GPU training on a single node via Accelerate. Multi-node distributed training is planned for future releases. For large-scale training, consider using external orchestration tools like Kubernetes or cloud provider solutions.

Prerequisites

• Completion of Tutorial 03 (Full Fine-Tuning)
• Multi-GPU environment (2+ GPUs, ideally 8+ for TP/PP)
• NCCL backend installed (torch.distributed support)

---

1. NTF Distributed Training with Accelerate

NTF automatically handles distributed training via Accelerate. Simply configure your training parameters and let NTF manage the complexity.

Basic Distributed Configuration

from ntf.config import NTFConfig, TrainingConfig, ModelConfig
from ntf.finetuning import FullFinetuneTrainer

# NTF automatically handles distributed training via Accelerate
config = NTFConfig(
    model=ModelConfig(name="meta-llama/Llama-2-7b-hf"),
    training=TrainingConfig(
        per_device_train_batch_size=4,
        gradient_accumulation_steps=4,
        # Accelerate auto-detects distributed setup
        fp16=False,
        bf16=True,
        gradient_checkpointing=True
    )
)

# Trainer automatically uses all available GPUs
trainer = FullFinetuneTrainer(config=config, train_dataset=train_dataset)
trainer.train()  # Distributed training handled internally

Launch Script Example (Single Node Multi-GPU)

# Using accelerate launch
accelerate launch --num_processes=4 train_script.py

# Or with explicit config
accelerate config  # Interactive setup
accelerate launch train_script.py

Checking Your Distributed Setup

import torch
from accelerate import Accelerator

accelerator = Accelerator()

print(f"Number of processes: {accelerator.num_processes}")
print(f"Local process index: {accelerator.local_process_index}")
print(f"Device: {accelerator.device}")
print(f"Distributed type: {accelerator.distributed_type}")

---

2. Concepts: The 3D Parallelism Strategy

To train massive models, we slice the computation in three dimensions:

1. Data Parallelism (DP): Replicate the model on all GPUs. Split the batch across GPUs. Each GPU computes gradients on its slice, then averages them.
   • Limit: Model must fit on one GPU.
2. Tensor Parallelism (TP): Split individual matrix multiplications within a layer across GPUs.
   • Example: A linear layer $Y = XA$ is split so GPU1 computes $X A_1$ and GPU2 computes $X A_2$. They sync to get $Y$.
   • Requirement: High-speed interconnect (NVLink) within a node.
3. Pipeline Parallelism (PP): Split the model by layers. GPU1 holds layers 1-10, GPU2 holds 11-20. Micro-batches flow through the pipeline like an assembly line.
   • Requirement: Handles inter-node communication better than TP.

Hybrid Approach: Modern training uses all three: DP x TP x PP.

---

2. Data Parallelism (DDP) Deep Dive

PyTorch's DistributedDataParallel (DDP) is the baseline. Unlike DataParallel (which is Python-threaded and slow), DDP spawns separate processes per GPU.

Step 2.1: Initializing the Process Group

You must initialize the distributed context before building your model.

# train_ddp.py
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP
from transformers import AutoModelForCausalLM, AutoTokenizer
import os

def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'
    
    # Initialize the process group
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)

def cleanup():
    dist.destroy_process_group()

def train_rank(rank, world_size):
    print(f"Running DDP on rank {rank}.")
    setup(rank, world_size)

    # 1. Load Model (Sharded or Replicated depending on strategy)
    # For standard DDP, the full model loads on every GPU
    model_name = "microsoft/phi-2" 
    model = AutoModelForCausalLM.from_pretrained(model_name).to(rank)
    
    # 2. Wrap with DDP
    # bucket_cap_mb tunes gradient bucketing for communication overlap
    ddp_model = DDP(model, device_ids=[rank], bucket_cap_mb=25)

    # 3. Use DistributedSampler to split data uniquely per rank
    # dataset = MyDataset(...)
    # sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True)
    # dataloader = DataLoader(dataset, sampler=sampler, batch_size=32)

    # Dummy training loop
    optimizer = torch.optim.AdamW(ddp_model.parameters(), lr=1e-4)
    
    for step in range(100):
        # input_ids = next(dataloader) # shape: [batch/world_size, seq_len]
        # inputs = input_ids.to(rank)
        
        # Fake input for demo
        inputs = torch.randint(0, 1000, (4, 128)).to(rank)
        labels = inputs.clone()
        
        outputs = ddp_model(inputs, labels=labels)
        loss = outputs.loss
        
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
        
        if rank == 0 and step % 10 == 0:
            print(f"Step {step}, Loss: {loss.item()}")

    cleanup()

if __name__ == "__main__":
    world_size = torch.cuda.device_count()
    print(f"Starting DDP with {world_size} GPUs")
    mp.spawn(train_rank, args=(world_size,), nprocs=world_size, join=True)

Key DDP Concepts:

• Gradient Bucketing: DDP groups gradients into buckets to overlap communication (all-reduce) with backprop computation.
• Static Graph: DDP expects the computation graph to be static. If you have dynamic control flow (e.g., variable depth), set find_unused_parameters=True (slower).

---

3. ZeRO (Zero Redundancy Optimizer)

Standard DDP replicates Model, Gradients, and Optimizer States on every GPU.

• For a 7B model (fp16):
  • Parameters: 14GB
  • Gradients: 14GB
  • Optimizer States (Adam): 28GB (fp32 master weights + momentum + variance)
  • Total per GPU: ~56GB. This limits us to small models even on 8x80GB GPUs.

ZeRO (implemented in DeepSpeed and FSDP) shards these states across GPUs.

ZeRO Stages:

1. ZeRO-1: Shard Optimizer States. Gradients and Params replicated.
   • Memory Save: Significant (Adam states are huge).
2. ZeRO-2: Shard Optimizer States + Gradients. Params replicated.
   • Memory Save: Massive. Communication overhead increases slightly.
3. ZeRO-3: Shard Optimizer States + Gradients + Parameters.
   • Memory Save: Maximum. Allows training models 10x larger than VRAM.
   • Cost: High communication. Requires gathering parameters before every forward pass.

Using PyTorch FSDP (Fully Sharded Data Parallel)

FSDP is PyTorch's native implementation of ZeRO-3.

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import MixedPrecision, ShardingStrategy
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
from transformers.models.llama.modeling_llama import LlamaDecoderLayer

def get_fsdp_policy():
    return transformer_auto_wrap_policy(
        transformer_layer_cls={LlamaDecoderLayer}
    )

def train_fsdp(rank, world_size):
    setup(rank, world_size)
    
    model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
    
    # Configure FSDP
    fsdp_model = FSDP(
        model,
        device_id=rank,
        sharding_strategy=ShardingStrategy.FULL_SHARD, # ZeRO-3
        auto_wrap_policy=get_fsdp_policy(),
        mixed_precision=MixedPrecision(
            param_dtype=torch.float16,
            reduce_dtype=torch.float16,
            buffer_dtype=torch.float16,
        ),
        use_orig_params=True, # Important for optimizer compatibility
    )
    
    # Optimizer must be created AFTER wrapping
    optimizer = torch.optim.AdamW(fsdp_model.parameters(), lr=1e-4)
    
    # Training loop similar to DDP, but FSDP handles parameter gathering internally
    # ...
    cleanup()

---

4. Tensor Parallelism (TP)

When a single layer's weights don't fit on one GPU, we split the matrix. In Transformers, the heavy parts are:

1. QKV Projection ($W_{qkv}$)
2. Output Projection ($W_{out}$)
3. MLP Up/Down Projections

How TP Works (Column/Row Splitting)

Consider $Y = XA$.

• Split $A$ by columns: $A = [A_1, A_2]$.
• GPU1 computes $Y_1 = X A_1$.
• GPU2 computes $Y_2 = X A_2$.
• Result $Y = [Y_1, Y_2]$ is split naturally.
• For the next layer $Z = YB$, we split $B$ by rows to match.

Implementation Note: Implementing TP manually is error-prone. We use DeepSpeed or Megatron-LM patterns.

DeepSpeed Configuration for TP

Create ds_config_tp.json:

{
  "train_batch_size": 64,
  "gradient_accumulation_steps": 1,
  "fp16": { "enabled": true },
  "zero_optimization": {
    "stage": 0, 
    "offload_optimizer": { "device": "none" }
  },
  "tensor_parallel": {
    "tp_size": 4, 
    "enabled": true
  },
  "pipeline_parallel": {
    "pp_size": 1,
    "enabled": false
  }
}

Note: TP usually disables ZeRO sharding of parameters because the parameters are already split.

---

5. Pipeline Parallelism (PP)

PP splits layers across devices.

• GPU0: Layers 1-8
• GPU1: Layers 9-16
• GPU2: Layers 17-24

The Bubble Problem

In naive PP, GPU1 waits for GPU0 to finish the forward pass before starting. This creates idle time ("bubble"). Solution: Micro-batching. Split a global batch into 4 micro-batches.

1. GPU0 processes MicroBatch 1 (Forward).
2. GPU0 sends activation to GPU1, then starts MicroBatch 2 (Forward).
3. GPU1 processes MicroBatch 1 (Forward) while GPU0 does MB2. This overlaps computation and reduces the bubble size.

GPipe vs 1F1B

• GPipe: Simple forward-all, backward-all. Large bubble.
• 1F1B (One-Forward-One-Backward): Interleaves forward and backward passes to minimize bubble. Used in modern schedulers.

---

6. Putting It All Together: Hybrid Parallelism

For a 70B model on a 64-GPU cluster:

• TP=4: Split each layer across 4 GPUs (within a node).
• PP=4: Split 80 layers into 4 chunks of 20 layers.
• DP=4: Replicate this TP+PP setup 4 times to process different data batches.
• Total GPUs: $4 \times 4 \times 4 = 64$.

Launching with DeepSpeed

DeepSpeed unifies these configurations.

deepspeed --num_gpus=8 train.py \
  --deepspeed ds_config_hybrid.json

ds_config_hybrid.json:

{
  "fp16": { "enabled": true },
  "bf16": { "enabled": false },
  "zero_optimization": {
    "stage": 1, 
    "contiguous_gradients": true,
    "overlap_comm": true
  },
  "tensor_parallel": { "tp_size": 2, "enabled": true },
  "pipeline_parallel": { "pp_size": 2, "enabled": true, "schedule": "1F1B" },
  "train_micro_batch_size_per_gpu": 2,
  "gradient_accumulation_steps": 4
}

---

7. Practical Exercise: Scaling a 7B Model

Scenario: You have 4x A100 (40GB). You want to fine-tune Llama-2-7B.

• Full model (fp16): 14GB.
• Optimizer states (fp32): 28GB.
• Total needed: 42GB. OOM on 40GB cards.

Solution A: ZeRO-2 (Recommended for 7B) Shard optimizer and gradients.

• Memory per GPU $\approx$ Params (14) + Shard(Opt+Grad) $\approx$ 14 + 10 = 24GB. Fits easily.

Solution B: ZeRO-3 + Offload If you had 2 GPUs, ZeRO-3 shards params too.

• Memory per GPU $\approx$ Shard(Params) + Small Buffer.
• If still tight, offload optimizer to CPU (offload_optimizer: {device: "cpu"}).

Code Snippet for ZeRO-2 Config:

{
  "zero_optimization": {
    "stage": 2,
    "allgather_partitions": true,
    "allgather_bucket_size": 2e8,
    "overlap_comm": true,
    "reduce_scatter": true,
    "reduce_bucket_size": 2e8,
    "contiguous_gradients": true
  },
  "optimizer": {
    "type": "AdamW",
    "params": { "lr": 1e-4, "betas": [0.9, 0.999], "eps": 1e-8 }
  }
}

---

8. Debugging Distributed Training

1. Hangs/Timeouts: Usually caused by mismatched tensor shapes causing one rank to wait forever for a broadcast.
   • Fix: Ensure DistributedSampler drops last batch consistently or pad data. Check find_unused_parameters.
2. NCCL Errors: unhandled system error.
   • Cause: Network issues, P2P access disabled.
   • Fix: export NCCL_P2P_DISABLE=1, check nvidia-smi topo -m.
3. Slow Scaling:
   • Check: Are you using overlap_comm? Is the network bandwidth sufficient for ZeRO-3?
   • Metric: Monitor "MFU" (Model Flops Utilization). >40% is good for H100s.

---

9. Summary Checklist

Strategy	When to use	Memory Efficiency	Comm Overhead	Complexity
DDP	Model fits on 1 GPU	Low	Low	Low
ZeRO-1	Optimizer states too big	Medium	Low-Med	Low
ZeRO-2	Gradients too big	High	Med	Low
ZeRO-3	Model params too big	Very High	High	Med
TP	Single layer too big	High (Intra-node)	Very High	High
PP	Total layers too many	High (Inter-node)	Med (Bubble)	High

Next Steps

In Tutorial 11, we will explore Inference Optimization, covering quantization (INT8/FP4), KV Caching, and serving engines like vLLM and TGI to deploy these massive models efficiently.