Ml Profile

ML Training Performance Profiling

Profile ML training to identify bottlenecks in data loading, model computation, and memory usage.

Quick Start

# Simple profiling (built-in)
python src/train.py trainer.profiler=simple trainer.max_epochs=1

# Detailed profiling (PyTorch Profiler)
python src/train.py trainer.profiler=pytorch trainer.max_steps=100

# View in TensorBoard
tensorboard --logdir=./tb_logs/profiler

# Monitor GPU utilization
watch -n 1 nvidia-smi

Profiling Tools

1. PyTorch Lightning Profilers

Simple Profiler (quick overview):

from pytorch_lightning import Trainer

trainer = Trainer(
    profiler="simple",  # Built-in simple profiler
    max_epochs=1,
)
trainer.fit(model, datamodule)

# Output: timing summary printed at end

Advanced Profiler (detailed breakdown):

from pytorch_lightning.profilers import AdvancedProfiler

profiler = AdvancedProfiler(
    dirpath="./profiler_logs",
    filename="advanced_profile",
)

trainer = Trainer(profiler=profiler, max_epochs=1)
trainer.fit(model, datamodule)

PyTorch Profiler (most detailed, visualizable):

import torch
from pytorch_lightning.profilers import PyTorchProfiler

profiler = PyTorchProfiler(
    dirpath="./profiler_logs",
    filename="pytorch_profile",
    activities=[
        torch.profiler.ProfilerActivity.CPU,
        torch.profiler.ProfilerActivity.CUDA,
    ],
    schedule=torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=1),
   
    record_shapes=True,
    profile_memory=True,
    with_stack=True,
)

trainer = Trainer(profiler=profiler, max_epochs=1)
trainer.fit(model, datamodule)

# View results:
# tensorboard --logdir=./tb_logs

2. Data Loading Profiling

Diagnose data loading bottleneck:

import time

def profile_dataloader(dataloader, num_batches=100):
    """Measure data loading speed."""
    times = []

    for i, batch in enumerate(dataloader):
        if i >= num_batches:
            break

        start = time.time()
        _ = batch  # Access data
        elapsed = time.time() - start
        times.append(elapsed)

    avg_time = sum(times) / len(times)
    throughput = 1.0 / avg_time if avg_time > 0 else 0

    print(f"Data Loading Profile:")
    print(f"  Avg time/batch: {avg_time*1000:.2f}ms")
    print(f"  Throughput: {throughput:.2f} batches/sec")

    return avg_time

# Usage
train_loader = datamodule.train_dataloader()
profile_dataloader(train_loader)

Find optimal num_workers:

from torch.utils.data import DataLoader

def find_optimal_num_workers(dataset, batch_size=32, max_workers=16):
    """Test different num_workers to find optimal."""
    results = []

    for num_workers in range(0, max_workers + 1):
        loader = DataLoader(
            dataset,
            batch_size=batch_size,
            num_workers=num_workers,
            pin_memory=True,
        )

        avg_time = profile_dataloader(loader, num_batches=50)
        results.append((num_workers, avg_time))
        print(f"num_workers={num_workers}: {avg_time*1000:.2f}ms/batch")

    # Find optimal
    workers, times = zip(*results)
    optimal_idx = times.index(min(times))
    optimal_workers = workers[optimal_idx]

    print(f"\n✓ Optimal num_workers: {optimal_workers}")
    return optimal_workers

See scripts/profile_dataloader.py for complete implementation.

3. GPU Utilization Monitoring

Real-time monitoring:

# Watch GPU usage
watch -n 1 nvidia-smi

# Log GPU usage during training
nvidia-smi dmon -s um -o TD > gpu_usage.log &
PID=$!
python src/train.py
kill $PID

GPU Monitor (programmatic):

import subprocess
import threading
import time

class GPUMonitor:
    """Monitor GPU utilization during training."""

    def __init__(self, interval=1.0):
        self.interval = interval
        self.running = False
        self.utilizations = []
        self.memories = []

    def start(self):
        self.running = True
        self.thread = threading.Thread(target=self._monitor, daemon=True)
        self.thread.start()

    def _monitor(self):
        while self.running:
            try:
                result = subprocess.run(
                    ["nvidia-smi", "--query-gpu=utilization.gpu,memory.used",
                     "--format=csv,noheader,nounits"],
                    capture_output=True,
                    text=True,
                )
                util, mem = result.stdout.strip().split(", ")
                self.utilizations.append(float(util))
                self.memories.append(float(mem))
            except Exception:
                pass
            time.sleep(self.interval)

    def stop(self):
        self.running = False
        self.thread.join()

    def report(self):
        avg_util = sum(self.utilizations) / len(self.utilizations)
        avg_mem = sum(self.memories) / len(self.memories)
        max_mem = max(self.memories)

        print(f"\nGPU Utilization Report:")
        print(f"  Average GPU util: {avg_util:.1f}%")
        print(f"  Average memory: {avg_mem:.0f} MB")
        print(f"  Peak memory: {max_mem:.0f} MB")

        if avg_util < 80:
            print(f"  ⚠️  Low GPU utilization ({avg_util:.1f}%)")
            print("     → Increase num_workers or batch size")

# Use in LightningModule
class MyModel(pl.LightningModule):
    def on_train_start(self):
        self.gpu_monitor = GPUMonitor()
        self.gpu_monitor.start()

    def on_train_end(self):
        self.gpu_monitor.stop()
        self.gpu_monitor.report()

4. Model Profiling

Layer-wise profiling:

import torch
from torch.profiler import profile, ProfilerActivity

def profile_model_layers(model, input_shape=(1, 3, 224, 224)):
    """Profile computation time of each layer."""
    model.eval()
    device = next(model.parameters()).device
    dummy_input = torch.randn(input_shape).to(device)

    with profile(
        activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
        record_shapes=True,
    ) as prof:
        with torch.no_grad():
            model(dummy_input)

    # Print layer-wise stats
    print(prof.key_averages().table(
        sort_by="cuda_time_total", row_limit=20
    ))

    return prof

Component timing:

import time
from pytorch_lightning import LightningModule

class ProfiledModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.timings = {"forward": [], "loss": [], "backward": []}

    def training_step(self, batch, batch_idx):
        x, y = batch

        # Time forward pass
        start = time.time()
        y_hat = self(x)
        self.timings["forward"].append(time.time() - start)

        # Time loss computation
        start = time.time()
        loss = self.criterion(y_hat, y)
        self.timings["loss"].append(time.time() - start)

        return loss

    def on_train_epoch_end(self):
        # Report average timings
        for key, times in self.timings.items():
            if times:
                avg = sum(times) / len(times) * 1000
                print(f"  Avg {key}: {avg:.2f}ms")

        # Reset for next epoch
        self.timings = {"forward": [], "loss": [], "backward": []}

5. Memory Profiling

Track GPU memory:

import torch

def print_memory_stats(device=0):
    """Print current GPU memory usage."""
    allocated = torch.cuda.memory_allocated(device) / 1024**2
    reserved = torch.cuda.memory_reserved(device) / 1024**2
    max_allocated = torch.cuda.max_memory_allocated(device) / 1024**2

    print(f"\nGPU Memory:")
    print(f"  Allocated: {allocated:.2f} MB")
    print(f"  Reserved: {reserved:.2f} MB")
    print(f"  Max: {max_allocated:.2f} MB")

# Use during training
class MemoryProfiledModel(pl.LightningModule):
    def training_step(self, batch, batch_idx):
        if batch_idx % 100 == 0:
            print_memory_stats()

        # Regular training
        ...

Memory callback:

from pytorch_lightning.callbacks import Callback

class MemoryProfileCallback(Callback):
    def __init__(self, log_every_n_steps=100):
        self.log_every_n_steps = log_every_n_steps

    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
        if batch_idx % self.log_every_n_steps == 0:
            allocated = torch.cuda.memory_allocated() / 1024**3
            reserved = torch.cuda.memory_reserved() / 1024**3

            pl_module.log("memory/allocated_gb", allocated)
            pl_module.log("memory/reserved_gb", reserved)

# Use in trainer
trainer = Trainer(callbacks=[MemoryProfileCallback(log_every_n_steps=50)])

Complete Profiling Script

Use the automated profiling script:

python scripts/profile_training.py --max-steps 100

See scripts/profile_training.py for implementation.

Performance Analysis

Interpreting Results

Symptoms and fixes:

Optimization Recommendations

def analyze_profile(gpu_util, data_time, model_time):
    """Provide optimization recommendations based on profile."""
    recommendations = []

    if gpu_util < 70:
        recommendations.append("🔴 Low GPU utilization")
        recommendations.append("   → Increase num_workers")
        recommendations.append("   → Use pin_memory=True")
        recommendations.append("   → Cache/preprocess data")

    if data_time > model_time * 0.5:
        recommendations.append("🟡 Data loading slow")
        recommendations.append("   → Optimize transforms")
        recommendations.append("   → Use faster format (LMDB, HDF5)")

    if gpu_util > 95:
        recommendations.append("🟢 GPU well utilized")
        recommendations.append("   → Increase batch size if possible")

    return recommendations

Common Bottlenecks

Data Loading Bottleneck

Symptoms:

• GPU utilization <80%
• Slow batch iteration
• nvidia-smi shows low GPU usage

Solutions:

data:
  num_workers: 8  # Increase workers
  pin_memory: true  # Faster GPU transfer
  persistent_workers: true  # Keep workers alive
  prefetch_factor: 2  # Prefetch batches

Advanced:

• Use faster storage (SSD vs HDD)
• Preprocess data offline
• Use LMDB or HDF5 format
• Cache augmented data

Compute Bottleneck

Symptoms:

• GPU utilization >95%
• Training very slow
• Model forward/backward takes most time

Solutions:

trainer:
  precision: "16-mixed"  # Mixed precision
  accumulate_grad_batches: 2  # Simulate larger batch

data:
  batch_size: 256  # Increase if memory allows

Model optimizations:

• Use torch.compile() (PyTorch 2.0+)
• Optimize architecture (fewer parameters)
• Use efficient operations

Memory Bottleneck

Symptoms:

• OOM errors
• Memory usage near limit
• Can't increase batch size

Solutions:

data:
  batch_size: 32  # Reduce

trainer:
  precision: "16-mixed"
  accumulate_grad_batches: 4

model:
  use_checkpoint: true  # Gradient checkpointing

Profiling Workflow

1. Run simple profiler - Quick overview

python src/train.py trainer.profiler=simple trainer.max_epochs=1

1. Monitor GPU - Check utilization

watch -n 1 nvidia-smi

1. Profile data loading - If GPU util <80%

python scripts/profile_dataloader.py

1. Detailed profiling - PyTorch Profiler + TensorBoard

python src/train.py trainer.profiler=pytorch trainer.max_steps=100
tensorboard --logdir=./tb_logs

1. Optimize - Based on findings
2. Re-profile - Verify improvements

Success Criteria

• GPU utilization >80% during training
• Data loading time <10% of total time
• No memory leaks
• Bottlenecks identified
• Optimizations applied
• Training speed improved

✅ Training optimized for maximum performance!