xbe.at

Miscellany / Data engineering and tools

Response:

8 min read, 18 sections

Contents
  1. Introduction to PyTorch Fabric
  2. Setting Up Fabric Environment
  3. Model Configuration with Fabric
  4. Distributed Data Loading
  5. Training Loop Implementation
  6. Response:
  7. Mixed Precision Training
  8. Implementing FSDP Strategy
  9. Checkpoint Management
  10. Performance Monitoring
  11. Response:
  12. Real-world Example - BERT Fine-tuning
  13. Training Loop for BERT Fine-tuning
  14. Implementing Gradient Accumulation
  15. Response:
  16. Advanced Metrics Tracking
  17. Results and Performance Analysis
  18. Additional Resources

Introduction to PyTorch Fabric

PyTorch Fabric represents a revolutionary approach to distributed training, bridging the gap between PyTorch's flexibility and Lightning's simplified distributed capabilities. It enables seamless scaling from single-device training to multi-GPU setups while maintaining pure PyTorch-style control over training loops.

# Traditional PyTorch training loop
import torch
from torch import nn
from torch.utils.data import DataLoader

# Standard training setup
model = MyModel().cuda()
optimizer = torch.optim.Adam(model.parameters())
train_loader = DataLoader(dataset, batch_size=32)

# Training loop with device management
for batch in train_loader:
    x, y = batch
    x, y = x.cuda(), y.cuda()
    loss = model(x, y)
    loss.backward()
    optimizer.step()

Setting Up Fabric Environment

Fabric initialization requires minimal code changes while providing immediate access to advanced distributed training features. The Fabric object becomes the central manager for all training-related operations, handling device placement and strategy coordination.

import lightning.fabric as fabric
from lightning.fabric import Fabric

# Initialize Fabric with desired settings
fabric = Fabric(
    accelerator="cuda",
    devices=2,
    strategy="ddp",
    precision="16-mixed"
)

# Launch the fabric environment
fabric.launch()

# Fabric manages model and optimizer setup
model, optimizer = fabric.setup(model, optimizer)
train_loader = fabric.setup_dataloaders(train_loader)

Model Configuration with Fabric

The transformation to Fabric involves restructuring how models and optimizers are initialized and managed. This slide demonstrates the core setup pattern that replaces traditional PyTorch device management while maintaining full control over the training process.

import torch.nn as nn
from lightning.fabric import Fabric

class ModelArchitecture(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(784, 512),
            nn.ReLU(),
            nn.Linear(512, 10)
        )
    
    def forward(self, x):
        return self.layers(x)

# Initialize Fabric and model
fabric = Fabric(accelerator="cuda", devices=2)
with fabric.device:
    model = ModelArchitecture()
    optimizer = torch.optim.Adam(model.parameters())

# Setup model and optimizer with Fabric
model, optimizer = fabric.setup(model, optimizer)

Distributed Data Loading

Fabric streamlines data loading in distributed environments by automatically handling sharding and worker initialization. This removes the complexity of manual distributed data management while ensuring efficient data processing across multiple GPUs.

from torch.utils.data import DataLoader, Dataset
import lightning.fabric as fabric

class CustomDataset(Dataset):
    def __init__(self, data, labels):
        self.data = data
        self.labels = labels
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        return self.data[idx], self.labels[idx]

# Create dataset and dataloader
dataset = CustomDataset(data, labels)
dataloader = DataLoader(
    dataset,
    batch_size=32,
    num_workers=4,
    shuffle=True
)

# Let Fabric handle distributed data loading
dataloader = fabric.setup_dataloaders(dataloader)

Training Loop Implementation

The training loop with Fabric maintains PyTorch's intuitive structure while abstracting away device management and distributed training complexities. This implementation shows how Fabric simplifies the training process without sacrificing control.

def training_step(batch, model, optimizer, fabric):
    # Fabric handles device placement automatically
    inputs, labels = batch
    
    # Forward pass
    outputs = model(inputs)
    loss = nn.functional.cross_entropy(outputs, labels)
    
    # Backward pass using Fabric
    fabric.backward(loss)
    
    optimizer.step()
    optimizer.zero_grad()
    
    return loss.item()

# Training loop
for epoch in range(num_epochs):
    for batch in train_loader:
        loss = training_step(batch, model, optimizer, fabric)
        
        # Fabric handles distributed logging
        fabric.log("train_loss", loss)

[Continuing with the remaining slides...]

Response:

Mixed Precision Training

Fabric simplifies the implementation of mixed precision training, automatically handling the conversion between float32 and float16 datatypes. This optimization technique significantly reduces memory usage and accelerates training while maintaining numerical stability.

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

# Training loop with automatic mixed precision
def training_step_amp(batch, model, optimizer, fabric):
    inputs, labels = batch
    
    # Forward pass - Fabric handles precision casting
    outputs = model(inputs)
    loss = nn.functional.cross_entropy(outputs, labels)
    
    # Backward pass with automatic scaling
    fabric.backward(loss)
    optimizer.step()
    optimizer.zero_grad()
    
    return loss.item()

Implementing FSDP Strategy

Fully Sharded Data Parallel (FSDP) is a powerful distributed training strategy that Fabric makes accessible through simple configuration. This implementation demonstrates how to enable FSDP for training large models efficiently.

from lightning.fabric import Fabric
from torch.distributed.fsdp import (
    FullyShardedDataParallel,
    CPUOffload
)

# Configure Fabric with FSDP strategy
fabric = Fabric(
    accelerator="cuda",
    devices=4,
    strategy="fsdp",
    precision="16-mixed"
)

# FSDP-specific configuration
fsdp_config = {
    "sharding_strategy": "FULL_SHARD",
    "cpu_offload": CPUOffload(offload_params=True),
    "mixed_precision": True
}

# Setup model with FSDP
with fabric.init():
    model = LargeTransformerModel()
    model = fabric.setup_module(model)

Checkpoint Management

Fabric provides robust checkpoint management capabilities that work seamlessly across different distributed strategies. This implementation shows how to save and load model states efficiently in a distributed environment.

def save_checkpoint(fabric, model, optimizer, epoch, path):
    checkpoint = {
        "model": model.state_dict(),
        "optimizer": optimizer.state_dict(),
        "epoch": epoch
    }
    
    # Fabric handles distributed checkpoint saving
    fabric.save(path, checkpoint)

def load_checkpoint(fabric, model, optimizer, path):
    # Load checkpoint with Fabric
    checkpoint = fabric.load(path)
    
    model.load_state_dict(checkpoint["model"])
    optimizer.load_state_dict(checkpoint["optimizer"])
    
    return checkpoint["epoch"]

# Usage in training loop
if epoch % save_frequency == 0:
    save_checkpoint(fabric, model, optimizer, epoch, f"checkpoint_{epoch}.ckpt")

Performance Monitoring

Fabric integrates seamlessly with various monitoring tools while providing built-in performance metrics tracking. This implementation demonstrates how to set up comprehensive monitoring in a distributed training environment.

from lightning.fabric.loggers import TensorBoardLogger
import time

# Initialize Fabric with logger
logger = TensorBoardLogger("logs/")
fabric = Fabric(loggers=logger)

# Performance monitoring implementation
class PerformanceMonitor:
    def __init__(self, fabric):
        self.fabric = fabric
        self.step_times = []
        self.start_time = None
    
    def start_batch(self):
        self.start_time = time.time()
    
    def end_batch(self, loss, batch_size):
        duration = time.time() - self.start_time
        self.step_times.append(duration)
        
        # Log metrics using Fabric
        self.fabric.log({
            "batch_time": duration,
            "samples_per_second": batch_size / duration,
            "loss": loss
        })

# Usage in training loop
monitor = PerformanceMonitor(fabric)
for batch in train_loader:
    monitor.start_batch()
    loss = training_step(batch, model, optimizer, fabric)
    monitor.end_batch(loss, batch.size(0))

[Continuing with the remaining slides...]

Response:

Real-world Example - BERT Fine-tuning

This implementation demonstrates how to use Fabric for fine-tuning a BERT model on a text classification task, showcasing the framework's capabilities with large language models and complex architectures.

from transformers import BertForSequenceClassification, BertTokenizer
from torch.utils.data import Dataset, DataLoader
import lightning.fabric as fabric

class TextClassificationDataset(Dataset):
    def __init__(self, texts, labels, tokenizer, max_length=512):
        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'] = self.labels[idx]
        return item

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

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

# Model and training setup
with fabric.device:
    model = BertForSequenceClassification.from_pretrained('bert-base-uncased')
    optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5)

# Setup with Fabric
model, optimizer = fabric.setup(model, optimizer)
train_loader = fabric.setup_dataloaders(train_loader)

Training Loop for BERT Fine-tuning

The training implementation shows how Fabric simplifies distributed training while maintaining full control over the training process, including gradient accumulation and learning rate scheduling.

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

def train_epoch(fabric, model, train_loader, optimizer, scheduler):
    model.train()
    total_loss = 0
    
    for batch in tqdm(train_loader):
        # Move batch to device (Fabric handles this automatically)
        input_ids = batch['input_ids']
        attention_mask = batch['attention_mask']
        labels = batch['labels']
        
        # Forward pass
        outputs = model(input_ids=input_ids,
                       attention_mask=attention_mask,
                       labels=labels)
        
        loss = outputs.loss
        
        # Backward pass through Fabric
        fabric.backward(loss)
        
        optimizer.step()
        scheduler.step()
        optimizer.zero_grad()
        
        # Log metrics
        fabric.log("train_loss", loss.item())
        total_loss += loss.item()
    
    return total_loss / len(train_loader)

# Training execution
scheduler = LinearLR(optimizer, start_factor=1.0, end_factor=0.0,
                    total_iters=num_epochs * len(train_loader))

for epoch in range(num_epochs):
    avg_loss = train_epoch(fabric, model, train_loader, optimizer, scheduler)
    print(f"Epoch {epoch + 1}, Average Loss: {avg_loss:.4f}")

Implementing Gradient Accumulation

This implementation shows how to incorporate gradient accumulation with Fabric, allowing for effective training with larger batch sizes than what would fit in GPU memory.

def train_with_gradient_accumulation(fabric, model, train_loader, optimizer,
                                   accumulation_steps=4):
    model.train()
    optimizer.zero_grad()
    accumulated_loss = 0
    
    for idx, batch in enumerate(train_loader):
        # Forward pass
        outputs = model(**batch)
        # Scale loss by accumulation steps
        loss = outputs.loss / accumulation_steps
        
        # Backward pass with Fabric
        fabric.backward(loss)
        accumulated_loss += loss.item()
        
        # Step optimizer only after accumulating gradients
        if (idx + 1) % accumulation_steps == 0:
            optimizer.step()
            optimizer.zero_grad()
            
            # Log accumulated loss
            fabric.log("train_loss", accumulated_loss * accumulation_steps)
            accumulated_loss = 0
            
    return accumulated_loss

# Usage example
accumulation_steps = 4
train_with_gradient_accumulation(fabric, model, train_loader, optimizer,
                               accumulation_steps)

[Continuing with the remaining slides...]

Response:

Advanced Metrics Tracking

This implementation demonstrates how to integrate advanced metrics tracking with Fabric, including custom metrics computation and distributed synchronization for accurate multi-GPU training statistics.

from torchmetrics import Accuracy, F1Score
import torch.distributed as dist

class MetricsTracker:
    def __init__(self, fabric, num_classes):
        self.fabric = fabric
        # Initialize metrics
        with self.fabric.device:
            self.accuracy = Accuracy(task="multiclass", num_classes=num_classes)
            self.f1_score = F1Score(task="multiclass", num_classes=num_classes)
    
    def update(self, logits, labels):
        preds = torch.argmax(logits, dim=1)
        self.accuracy.update(preds, labels)
        self.f1_score.update(preds, labels)
    
    def compute_and_log(self):
        # Synchronize metrics across processes
        accuracy = self.accuracy.compute()
        f1 = self.f1_score.compute()
        
        # Log metrics using Fabric
        self.fabric.log({
            "accuracy": accuracy,
            "f1_score": f1
        })
        
        return accuracy, f1

# Usage in training loop
metrics = MetricsTracker(fabric, num_classes=10)
for batch in train_loader:
    outputs = model(batch)
    metrics.update(outputs.logits, batch['labels'])
    
accuracy, f1 = metrics.compute_and_log()

Results and Performance Analysis

Analysis of training results across different distributed strategies, demonstrating the performance benefits of using Fabric for large-scale model training.

def analyze_training_performance(fabric, model_size="base", batch_size=32,
                               num_gpus=2):
    # Initialize performance metrics
    training_stats = {
        "throughput": [],
        "memory_usage": [],
        "time_per_epoch": []
    }
    
    def get_gpu_memory():
        return torch.cuda.max_memory_allocated() / 1e9  # Convert to GB
    
    # Training loop with performance monitoring
    start_time = time.time()
    for epoch in range(num_epochs):
        epoch_start = time.time()
        
        for batch in train_loader:
            batch_start = time.time()
            loss = training_step(batch, model, optimizer, fabric)
            
            # Calculate throughput (samples/second)
            samples_per_second = batch_size / (time.time() - batch_start)
            training_stats["throughput"].append(samples_per_second)
        
        # Record memory usage
        training_stats["memory_usage"].append(get_gpu_memory())
        training_stats["time_per_epoch"].append(time.time() - epoch_start)
    
    # Log final statistics
    fabric.log({
        "avg_throughput": np.mean(training_stats["throughput"]),
        "peak_memory_usage": max(training_stats["memory_usage"]),
        "avg_epoch_time": np.mean(training_stats["time_per_epoch"])
    })
    
    return training_stats

# Run performance analysis
stats = analyze_training_performance(fabric)

Additional Resources

More in Data engineering and tools

  1. 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...

  2. PyTorch Tutorial5 min

    PyTorch is an open-source machine learning library developed by Facebook's AI Research lab. It provides a flexible and efficient framework for building and...

  3. Streamlining PyTorch Training Loops3 min

    Instead of writing lengthy training loops, Skorch allows us to define our model architecture just as in PyTorch, but handles the training complexity...

  4. PyTorch Tensor Creation and Parallelization6 min

    In PyTorch, tensors are the fundamental data structures used for computations. Here's how to create a tensor from a Python list or array. Code Example:

  5. Streamlining PyTorch with Fabric8 min

    PyTorch Fabric represents a revolutionary approach to distributed training, combining PyTorch's flexibility with Lightning's simplified distributed features....

  6. PyTorch Research Reproducibility7 min

    Research reproducibility is a crucial aspect of scientific progress, especially in the field of machine learning. PyTorch, a popular deep learning framework...