xbe.at

Miscellany / Data engineering and tools

Streamlining PyTorch with Fabric

8 min read, 15 sections

Contents
  1. Introduction to PyTorch Fabric
  2. Basic Configuration and Model Setup
  3. DataLoader Configuration with Fabric
  4. Training Loop Implementation
  5. Mixed Precision Training
  6. Implementing FSDP with Fabric
  7. Real-world Example - BERT Fine-tuning
  8. Data Processing for BERT Example
  9. Training Loop for BERT Fine-tuning
  10. DeepSpeed Integration
  11. Advanced Gradient Accumulation with Fabric
  12. Custom Metrics and Logging
  13. Production-Ready Checkpointing
  14. Results Analysis and Visualization
  15. Additional Resources

Introduction to PyTorch Fabric

PyTorch Fabric represents a revolutionary approach to distributed training, combining PyTorch's flexibility with Lightning's simplified distributed features. It enables seamless scaling from single-device training to multi-device distributed environments while maintaining clean, maintainable code structure.

import lightning as L
import torch
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

# Initialize Fabric with specific accelerator and strategy
fabric = L.Fabric(accelerator="cuda", devices=2, strategy="ddp")
fabric.launch()

# Basic training configuration
model = YourModel()
optimizer = torch.optim.Adam(model.parameters())

Basic Configuration and Model Setup

Fabric requires minimal changes to existing PyTorch code, primarily focusing on initialization and device management. The framework automatically handles device placement and distributed training setup, eliminating boilerplate device management code.

# Traditional PyTorch setup vs Fabric setup
class ConvNet(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = torch.nn.Conv2d(1, 32, 3)
        self.conv2 = torch.nn.Conv2d(32, 64, 3)
        self.fc = torch.nn.Linear(64 * 12 * 12, 10)
    
    def forward(self, x):
        x = torch.relu(self.conv1(x))
        x = torch.relu(self.conv2(x))
        return self.fc(x.view(x.size(0), -1))

# Fabric handles device placement automatically
model, optimizer = fabric.setup(model, optimizer)

DataLoader Configuration with Fabric

Fabric seamlessly integrates with PyTorch's DataLoader, automatically handling distributed sampling and batch distribution across multiple devices. This eliminates the need for manual DistributedSampler setup and device placement logic.

# Define dataset and dataloader
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

train_dataset = datasets.MNIST('data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=32)

# Setup dataloader with Fabric
train_loader = fabric.setup_dataloaders(train_loader)

Training Loop Implementation

The training loop in Fabric maintains PyTorch's flexibility while abstracting away distributed training complexities. The key difference is the replacement of manual loss.backward() calls with fabric.backward(loss).

def train_epoch(fabric, model, train_loader, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        optimizer.zero_grad()
        output = model(data)
        loss = torch.nn.functional.cross_entropy(output, target)
        
        # Use fabric.backward instead of loss.backward()
        fabric.backward(loss)
        optimizer.step()
        
        if batch_idx % 100 == 0:
            fabric.print(f'Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)}'
                        f' ({100. * batch_idx / len(train_loader):.0f}%)]\tLoss: {loss.item():.6f}')

Mixed Precision Training

Fabric simplifies the implementation of mixed precision training, which can significantly reduce memory usage and increase training speed on modern GPUs. This feature is enabled through simple configuration parameters.

# Initialize Fabric with mixed precision
fabric = L.Fabric(
    accelerator="cuda",
    devices=2,
    strategy="ddp",
    precision="16-mixed"  # Enable mixed precision training
)

# Training loop remains the same, Fabric handles precision automatically
def train_with_mixed_precision(fabric, model, train_loader, optimizer):
    model.train()
    for data, target in train_loader:
        optimizer.zero_grad()
        output = model(data)
        loss = torch.nn.functional.cross_entropy(output, target)
        fabric.backward(loss)
        optimizer.step()

Implementing FSDP with Fabric

Fully Sharded Data Parallel (FSDP) is a powerful distributed training strategy that shards model parameters across multiple devices. Fabric makes implementing FSDP straightforward through simple configuration, enabling training of larger models with limited memory.

# Initialize Fabric with FSDP strategy
fabric = L.Fabric(
    accelerator="cuda",
    devices=4,
    strategy="fsdp",
    precision="16-mixed",
    fsdp_config={
        "sharding_strategy": "FULL_SHARD",
        "min_num_params": 1e8,
        "cpu_offload": False
    }
)

# Model and optimizer setup remains unchanged
model = LargeTransformerModel()
optimizer = torch.optim.AdamW(model.parameters())
model, optimizer = fabric.setup(model, optimizer)

Real-world Example - BERT Fine-tuning

This example demonstrates how to use Fabric for fine-tuning a BERT model on a text classification task. The implementation showcases Fabric's ability to handle complex models while maintaining clean code structure.

from transformers import BertModel, BertTokenizer
import torch.nn as nn

class BERTClassifier(nn.Module):
    def __init__(self, num_classes=2):
        super().__init__()
        self.bert = BertModel.from_pretrained('bert-base-uncased')
        self.classifier = nn.Linear(768, num_classes)
        
    def forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        pooled_output = outputs.pooler_output
        return self.classifier(pooled_output)

# Fabric setup for BERT fine-tuning
fabric = L.Fabric(accelerator="cuda", devices=2, strategy="ddp")
model = BERTClassifier()
optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5)
model, optimizer = fabric.setup(model, optimizer)

Data Processing for BERT Example

The data processing pipeline demonstrates how Fabric integrates with the Hugging Face transformers library while maintaining efficient distributed training capabilities. This setup handles tokenization and batching for the BERT model.

from torch.utils.data import Dataset, DataLoader
from transformers import BertTokenizer

class TextClassificationDataset(Dataset):
    def __init__(self, texts, labels, tokenizer, max_length=128):
        self.encodings = tokenizer(texts, truncation=True, padding=True, 
                                 max_length=max_length, return_tensors='pt')
        self.labels = labels

    def __getitem__(self, idx):
        item = {key: val[idx] for key, val in self.encodings.items()}
        item['labels'] = torch.tensor(self.labels[idx])
        return item

    def __len__(self):
        return len(self.labels)

# Setup tokenizer and dataset
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
train_dataset = TextClassificationDataset(train_texts, train_labels, tokenizer)
train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)
train_loader = fabric.setup_dataloaders(train_loader)

Training Loop for BERT Fine-tuning

This implementation shows how Fabric handles the training loop for BERT fine-tuning, including gradient accumulation and learning rate scheduling, while maintaining distributed training efficiency.

from torch.optim.lr_scheduler import LinearLR
from tqdm import tqdm

def train_bert(fabric, model, train_loader, optimizer, num_epochs=3):
    scheduler = LinearLR(optimizer, start_factor=1.0, end_factor=0.0, 
                        total_iters=len(train_loader) * num_epochs)
    scheduler = fabric.setup_scheduler(scheduler)
    
    for epoch in range(num_epochs):
        model.train()
        total_loss = 0
        
        for batch in tqdm(train_loader):
            input_ids = batch['input_ids']
            attention_mask = batch['attention_mask']
            labels = batch['labels']
            
            outputs = model(input_ids, attention_mask)
            loss = torch.nn.functional.cross_entropy(outputs, labels)
            
            fabric.backward(loss)
            optimizer.step()
            scheduler.step()
            optimizer.zero_grad()
            
            total_loss += loss.item()
        
        avg_loss = total_loss / len(train_loader)
        fabric.print(f"Epoch {epoch+1}, Average Loss: {avg_loss:.4f}")

DeepSpeed Integration

Fabric provides seamless integration with DeepSpeed, enabling advanced optimization techniques like ZeRO optimization and pipeline parallelism. This implementation shows how to configure DeepSpeed with minimal code changes.

# DeepSpeed configuration
deepspeed_config = {
    "train_batch_size": 32,
    "fp16": {
        "enabled": True
    },
    "zero_optimization": {
        "stage": 2,
        "offload_optimizer": {
            "device": "cpu"
        }
    }
}

# Initialize Fabric with DeepSpeed
fabric = L.Fabric(
    accelerator="cuda",
    devices=4,
    strategy="deepspeed",
    precision="16-mixed",
    deepspeed_config=deepspeed_config
)

# Setup remains the same
model, optimizer = fabric.setup(model, optimizer)

Advanced Gradient Accumulation with Fabric

Fabric simplifies the implementation of gradient accumulation for training larger batch sizes than what fits in GPU memory. This technique maintains training stability while enabling efficient memory usage across distributed environments.

def train_with_gradient_accumulation(fabric, model, train_loader, optimizer, 
                                   accumulation_steps=4):
    model.train()
    optimizer.zero_grad()
    accumulated_loss = 0
    
    for batch_idx, (data, target) in enumerate(train_loader):
        # Forward pass
        output = model(data)
        loss = torch.nn.functional.cross_entropy(output, target)
        # Scale loss by accumulation steps
        scaled_loss = loss / accumulation_steps
        
        # Backward pass with scaled loss
        fabric.backward(scaled_loss)
        accumulated_loss += loss.item()
        
        # Update weights after accumulation steps
        if (batch_idx + 1) % accumulation_steps == 0:
            optimizer.step()
            optimizer.zero_grad()
            fabric.print(f'Accumulated Loss: {accumulated_loss/accumulation_steps:.4f}')
            accumulated_loss = 0

Custom Metrics and Logging

Fabric integrates seamlessly with various logging systems while providing distributed-aware metric computation. This implementation demonstrates how to implement custom metrics and logging in a distributed training setup.

from torchmetrics import Accuracy, F1Score
import time

class MetricTracker:
    def __init__(self, fabric):
        self.fabric = fabric
        self.accuracy = Accuracy(task="multiclass", num_classes=10).to(fabric.device)
        self.f1 = F1Score(task="multiclass", num_classes=10).to(fabric.device)
        self.start_time = time.time()
        
    def update(self, predictions, targets):
        self.accuracy.update(predictions, targets)
        self.f1.update(predictions, targets)
    
    def compute_and_log(self, epoch):
        metrics = {
            'accuracy': self.accuracy.compute().item(),
            'f1_score': self.f1.compute().item(),
            'time_elapsed': time.time() - self.start_time
        }
        
        # Fabric ensures proper metric aggregation across devices
        self.fabric.print(f"Epoch {epoch} Metrics:")
        for name, value in metrics.items():
            self.fabric.print(f"{name}: {value:.4f}")
        
        self.accuracy.reset()
        self.f1.reset()

Production-Ready Checkpointing

Efficient model checkpointing is crucial for production environments. Fabric provides distributed-aware checkpointing that handles model states, optimizer states, and training progress across multiple devices.

class CheckpointManager:
    def __init__(self, fabric, model, optimizer, save_dir='checkpoints'):
        self.fabric = fabric
        self.model = model
        self.optimizer = optimizer
        self.save_dir = save_dir
        os.makedirs(save_dir, exist_ok=True)
        
    def save_checkpoint(self, epoch, metrics):
        checkpoint = {
            'epoch': epoch,
            'model_state_dict': self.model.state_dict(),
            'optimizer_state_dict': self.optimizer.state_dict(),
            'metrics': metrics
        }
        
        # Fabric handles distributed saving
        checkpoint_path = os.path.join(self.save_dir, f'checkpoint_epoch_{epoch}.pt')
        self.fabric.save(checkpoint_path, checkpoint)
        
    def load_checkpoint(self, checkpoint_path):
        # Fabric handles distributed loading
        checkpoint = self.fabric.load(checkpoint_path)
        self.model.load_state_dict(checkpoint['model_state_dict'])
        self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        return checkpoint['epoch'], checkpoint['metrics']

Results Analysis and Visualization

This implementation shows how to collect and visualize training metrics across distributed processes, ensuring accurate performance monitoring in multi-GPU environments.

import matplotlib.pyplot as plt
import numpy as np

class TrainingAnalyzer:
    def __init__(self, fabric):
        self.fabric = fabric
        self.training_history = {
            'loss': [],
            'accuracy': [],
            'learning_rate': []
        }
    
    def update_metrics(self, loss, accuracy, lr):
        # Gather metrics from all processes
        all_losses = self.fabric.all_gather(torch.tensor(loss))
        all_accuracies = self.fabric.all_gather(torch.tensor(accuracy))
        
        # Update history on rank 0
        if self.fabric.global_rank == 0:
            self.training_history['loss'].append(all_losses.mean().item())
            self.training_history['accuracy'].append(all_accuracies.mean().item())
            self.training_history['learning_rate'].append(lr)
    
    def plot_metrics(self):
        if self.fabric.global_rank == 0:
            fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 10))
            
            epochs = range(1, len(self.training_history['loss']) + 1)
            
            ax1.plot(epochs, self.training_history['loss'], 'b-', label='Training Loss')
            ax1.set_title('Training Loss over Epochs')
            ax1.set_xlabel('Epochs')
            ax1.set_ylabel('Loss')
            ax1.legend()
            
            ax2.plot(epochs, self.training_history['accuracy'], 'r-', 
                    label='Training Accuracy')
            ax2.set_title('Training Accuracy over Epochs')
            ax2.set_xlabel('Epochs')
            ax2.set_ylabel('Accuracy')
            ax2.legend()
            
            plt.tight_layout()
            plt.savefig('training_metrics.png')

Additional Resources

More in Data engineering and tools

  1. Top 5 Deployment Strategies7 min

    Big Bang Deployment is a strategy where the entire application is updated at once. This approach involves replacing the old version with the new version in a...

  2. Secure Web API Access Design9 min

    Authentication is a critical security measure that verifies the identity of users or systems making API requests. It ensures that only authorized entities...

  3. Top Authentication Mechanisms for Secure Digital Interactions6 min

    A secure user authentication system using password hashing with salt to prevent rainbow table attacks and ensure secure credential storage in a database.

  4. Scaling Data-Intensive Applications18 min

    Partitioning data involves breaking large datasets into smaller, more manageable chunks called partitions or shards. This fundamental concept enables...

  5. Top System Integration Patterns for Seamless Systems12 min

    Peer-to-Peer integration enables direct communication between systems without intermediaries. This pattern is ideal for small networks where systems can...

  6. Response:8 min

    PyTorch Fabric represents a revolutionary approach to distributed training, bridging the gap between PyTorch's flexibility and Lightning's simplified...