Torch Pipeline Parallelism

PyTorch Pipeline Parallelism

Overview

Pipeline parallelism distributes model layers across multiple GPUs/ranks, enabling training of models too large for a single device. This skill provides procedural guidance for implementing pipeline parallel training with proper model partitioning, communication patterns, and gradient handling.

Key Concepts

Pipeline Parallelism Architecture

• Stage: A subset of model layers assigned to one rank
• Microbatch: Input data split into smaller chunks processed sequentially
• AFAB (All-Forward-All-Backward): Schedule where all forward passes complete before any backward passes begin
• Activation caching: Storing intermediate outputs for backward pass computation

Critical Implementation Components

1. Model partitioning across ranks
2. Inter-rank tensor communication (send/recv)
3. Activation caching for backward pass
4. Loss computation and gradient scaling
5. Handling of model-specific components (embeddings, normalization, output heads)

Implementation Approach

Phase 1: Model Analysis and Planning

Before writing any code:

1. Examine the target model architecture

   • Identify the layer structure (e.g., model.model.layers for LLaMA)
   • Locate embedding layers, normalization layers, and output heads
   • Understand position embedding mechanisms (rotary, absolute, etc.)
   • Note any model-specific forward pass requirements

2. Plan the partitioning strategy

   • Determine how layers will be distributed across ranks
   • Decide which rank handles embeddings (typically rank 0)
   • Decide which rank handles the output head and loss (typically last rank)
   • Account for any shared parameters

3. Design the communication protocol

   • Shape communication before tensor data
   • Dtype preservation across ranks
   • Device placement consistency

Phase 2: Implement Core Functions

Implement these functions in order, testing each before proceeding:

1. partition_model(model, rank, world_size)

   • Extract layer subsets for the given rank
   • Handle embedding layers for rank 0
   • Handle output head for final rank
   • Return a clean partition object or module list

2. forward_stage(partition, hidden_states, ...)

   • Process hidden states through partition layers
   • Handle position embeddings correctly
   • Return output hidden states

3. train_step_pipeline_afab(model, inputs, targets, ...)

   • Implement the full AFAB schedule
   • Manage microbatch iteration
   • Handle inter-rank communication
   • Compute and scale loss appropriately

Phase 3: Communication Implementation

Shape Communication Pattern:

# Sending shape (more efficient as single tensor)
shape_tensor = torch.tensor(list(output.shape), dtype=torch.long, device=device)
dist.send(shape_tensor, dst=next_rank)

# Receiving shape
shape_tensor = torch.zeros(num_dims, dtype=torch.long, device=device)
dist.recv(shape_tensor, src=prev_rank)
recv_shape = tuple(shape_tensor.tolist())

Tensor Communication:

# Ensure contiguous memory and correct dtype
output = output.contiguous()
dist.send(output, dst=next_rank)

# Receive with matching dtype
recv_tensor = torch.zeros(recv_shape, dtype=dtype, device=device)
dist.recv(recv_tensor, src=prev_rank)

Phase 4: Gradient Flow Management

Critical Pattern for Pipeline Parallelism:

# During forward pass, cache inputs for backward
input_detached = input_tensor.detach().requires_grad_(True)
input_cache.append(input_detached)
output = forward_stage(partition, input_detached, ...)

# During backward pass, use cached inputs
for i in reversed(range(num_microbatches)):
    cached_input = input_cache[i]
    output = output_cache[i]
    output.backward(grad_output, retain_graph=False)
    grad_to_send = cached_input.grad

Verification Strategy

Level 1: Syntax and Import Validation

import py_compile
py_compile.compile('pipeline_parallel.py', doraise=True)

# Also verify imports work
import pipeline_parallel

Level 2: Unit Testing Components

Test partitioning independently:

def test_partition_model():
    model = load_small_test_model()
    for rank in range(world_size):
        partition = partition_model(model, rank, world_size)
        assert partition is not None
        assert len(partition.layers) > 0

Test forward stage in isolation:

def test_forward_stage():
    partition = get_test_partition()
    dummy_input = torch.randn(batch, seq_len, hidden_dim)
    output = forward_stage(partition, dummy_input, ...)
    assert output.shape == expected_shape
    assert not torch.isnan(output).any()

Level 3: Single-Rank Integration Test

Before distributed testing, verify with world_size=1:

def test_single_rank():
    # Should handle the degenerate case without communication
    os.environ['WORLD_SIZE'] = '1'
    os.environ['RANK'] = '0'
    dist.init_process_group(backend='gloo', ...)

    loss = train_step_pipeline_afab(model, inputs, targets, ...)
    assert loss is not None
    assert not torch.isnan(loss)

Level 4: Multi-Rank Distributed Test

# Launch with torchrun or mp.spawn
# Verify gradients are computed on all ranks
# Check that loss decreases over iterations

Common Pitfalls

1. Truncated File Writes

Problem: Large code blocks may be truncated during file operations. Prevention:

• Write functions incrementally, verifying each write completes
• After writing, read back the file and verify syntax
• Use smaller, complete edits rather than large file rewrites

2. Position Embedding Mishandling

Problem: Transformer models (especially LLaMA) use rotary position embeddings that require specific handling. Prevention:

• Read the model's forward pass implementation before coding
• Verify how rotary_emb or position embeddings are computed
• Pass position IDs explicitly through stages if needed
• Test with the actual model's embedding mechanism

3. Incorrect Loss Scaling

Problem: Loss scaled incorrectly when using microbatches. Prevention:

• Scale loss by 1/num_microbatches once, not per iteration
• Apply scaling before backward pass: (loss / num_microbatches).backward()
• Verify gradient magnitudes match non-pipeline baseline

4. Communication Deadlocks

Problem: Send/recv operations block indefinitely. Prevention:

• Ensure matching send/recv pairs across ranks
• Handle world_size=1 case without communication calls
• Use non-blocking operations with proper synchronization if needed
• Test communication patterns in isolation first

5. Broken Gradient Graph

Problem: Using detach() incorrectly breaks gradient flow. Prevention:

• Understand that detach().requires_grad_(True) creates a new leaf tensor
• Cache these leaf tensors for backward pass
• Verify gradients flow correctly with simple test cases

6. Architecture Assumptions

Problem: Code assumes specific model structure that may vary. Prevention:

• Validate model structure before accessing attributes
• Use hasattr() checks for optional components
• Document which model architectures are supported

7. Dtype and Device Mismatches

Problem: Tensors on wrong device or with wrong dtype after communication. Prevention:

• Explicitly specify dtype when creating receive buffers
• Verify device placement after each operation
• Use .to(device=device, dtype=dtype) consistently

8. Missing Return Values

Problem: Functions don't return expected values for testing/integration. Prevention:

• Define function contracts before implementation
• Return loss values for training functions
• Return gradients or other diagnostics as needed

References

For detailed implementation guidance, see references/pipeline_parallelism_guide.md.