---
Scanned 5/31/2026
Install via CLI
Skills Directoryopenskills install tools-only/X-Skills# Training Pipelines
---
## Overview
Training pipelines orchestrate the end-to-end model training process including data loading, distributed training, hyperparameter optimization, and artifact management. Production pipelines require reproducibility, scalability, and proper resource management.
## When to Use This Reference
- Setting up distributed training with PyTorch/TensorFlow
- Implementing hyperparameter tuning (Optuna, Ray Tune)
- Managing GPU/TPU resources for training
- Building reproducible training environments
- Creating checkpointing and fault-tolerant training
## When NOT to Use
- Quick model prototyping (use notebooks)
- Small models that fit in memory on single GPU
- One-off experiments without production requirements
---
## PyTorch Training Pipeline
### Complete Training Script
```python
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR
import logging
from pathlib import Path
from dataclasses import dataclass
from typing import Optional
import json
logger = logging.getLogger(__name__)
@dataclass
class TrainingConfig:
"""Training hyperparameters and settings."""
model_name: str
batch_size: int = 32
learning_rate: float = 1e-4
weight_decay: float = 0.01
epochs: int = 10
warmup_steps: int = 100
max_grad_norm: float = 1.0
seed: int = 42
checkpoint_dir: str = "./checkpoints"
log_every_n_steps: int = 100
eval_every_n_steps: int = 500
save_every_n_steps: int = 1000
mixed_precision: bool = True
gradient_accumulation_steps: int = 1
def to_dict(self) -> dict:
return {k: v for k, v in self.__dict__.items()}
@classmethod
def from_dict(cls, d: dict) -> "TrainingConfig":
return cls(**d)
class Trainer:
"""Production-grade PyTorch trainer."""
def __init__(
self,
model: nn.Module,
config: TrainingConfig,
train_dataloader: DataLoader,
eval_dataloader: Optional[DataLoader] = None,
experiment_tracker=None,
):
self.model = model
self.config = config
self.train_dataloader = train_dataloader
self.eval_dataloader = eval_dataloader
self.tracker = experiment_tracker
self._setup_device()
self._setup_training()
self._setup_checkpointing()
def _setup_device(self) -> None:
"""Configure device and move model."""
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.model = self.model.to(self.device)
if self.config.mixed_precision and self.device.type == "cuda":
self.scaler = torch.amp.GradScaler("cuda")
else:
self.scaler = None
logger.info(f"Training on device: {self.device}")
def _setup_training(self) -> None:
"""Initialize optimizer and scheduler."""
self.optimizer = AdamW(
self.model.parameters(),
lr=self.config.learning_rate,
weight_decay=self.config.weight_decay,
)
total_steps = len(self.train_dataloader) * self.config.epochs
self.scheduler = CosineAnnealingLR(
self.optimizer,
T_max=total_steps,
eta_min=self.config.learning_rate * 0.01,
)
self.global_step = 0
self.best_eval_loss = float("inf")
def _setup_checkpointing(self) -> None:
"""Create checkpoint directory."""
self.checkpoint_dir = Path(self.config.checkpoint_dir)
self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
def _set_seed(self) -> None:
"""Set random seeds for reproducibility."""
import random
import numpy as np
torch.manual_seed(self.config.seed)
torch.cuda.manual_seed_all(self.config.seed)
np.random.seed(self.config.seed)
random.seed(self.config.seed)
torch.backends.cudnn.deterministic = True
def train(self) -> dict:
"""Run training loop."""
self._set_seed()
self.model.train()
metrics_history = []
for epoch in range(self.config.epochs):
epoch_loss = 0.0
num_batches = 0
for batch_idx, batch in enumerate(self.train_dataloader):
loss = self._training_step(batch)
epoch_loss += loss
num_batches += 1
if self.global_step % self.config.log_every_n_steps == 0:
self._log_metrics({
"train/loss": loss,
"train/lr": self.scheduler.get_last_lr()[0],
"train/epoch": epoch,
})
if (
self.eval_dataloader
and self.global_step % self.config.eval_every_n_steps == 0
):
eval_metrics = self.evaluate()
self._log_metrics(eval_metrics)
if eval_metrics["eval/loss"] < self.best_eval_loss:
self.best_eval_loss = eval_metrics["eval/loss"]
self.save_checkpoint("best")
if self.global_step % self.config.save_every_n_steps == 0:
self.save_checkpoint(f"step_{self.global_step}")
avg_epoch_loss = epoch_loss / num_batches
logger.info(f"Epoch {epoch}: avg_loss={avg_epoch_loss:.4f}")
metrics_history.append({"epoch": epoch, "loss": avg_epoch_loss})
self.save_checkpoint("final")
return {
"best_eval_loss": self.best_eval_loss,
"final_train_loss": avg_epoch_loss,
"total_steps": self.global_step,
"metrics_history": metrics_history,
}
def _training_step(self, batch: dict) -> float:
"""Execute single training step."""
batch = {k: v.to(self.device) for k, v in batch.items()}
if self.scaler:
with torch.amp.autocast("cuda"):
outputs = self.model(**batch)
loss = outputs.loss / self.config.gradient_accumulation_steps
self.scaler.scale(loss).backward()
else:
outputs = self.model(**batch)
loss = outputs.loss / self.config.gradient_accumulation_steps
loss.backward()
if (self.global_step + 1) % self.config.gradient_accumulation_steps == 0:
if self.scaler:
self.scaler.unscale_(self.optimizer)
torch.nn.utils.clip_grad_norm_(
self.model.parameters(),
self.config.max_grad_norm,
)
if self.scaler:
self.scaler.step(self.optimizer)
self.scaler.update()
else:
self.optimizer.step()
self.scheduler.step()
self.optimizer.zero_grad()
self.global_step += 1
return loss.item() * self.config.gradient_accumulation_steps
@torch.no_grad()
def evaluate(self) -> dict:
"""Run evaluation loop."""
self.model.eval()
total_loss = 0.0
num_batches = 0
for batch in self.eval_dataloader:
batch = {k: v.to(self.device) for k, v in batch.items()}
if self.scaler:
with torch.amp.autocast("cuda"):
outputs = self.model(**batch)
else:
outputs = self.model(**batch)
total_loss += outputs.loss.item()
num_batches += 1
self.model.train()
return {
"eval/loss": total_loss / num_batches,
"eval/step": self.global_step,
}
def save_checkpoint(self, name: str) -> Path:
"""Save model checkpoint."""
checkpoint_path = self.checkpoint_dir / name
torch.save({
"model_state_dict": self.model.state_dict(),
"optimizer_state_dict": self.optimizer.state_dict(),
"scheduler_state_dict": self.scheduler.state_dict(),
"global_step": self.global_step,
"best_eval_loss": self.best_eval_loss,
"config": self.config.to_dict(),
}, checkpoint_path / "checkpoint.pt")
# Save config separately for easy loading
with open(checkpoint_path / "config.json", "w") as f:
json.dump(self.config.to_dict(), f, indent=2)
logger.info(f"Saved checkpoint: {checkpoint_path}")
return checkpoint_path
def load_checkpoint(self, checkpoint_path: Path) -> None:
"""Load model checkpoint."""
checkpoint = torch.load(checkpoint_path / "checkpoint.pt", map_location=self.device)
self.model.load_state_dict(checkpoint["model_state_dict"])
self.optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
self.scheduler.load_state_dict(checkpoint["scheduler_state_dict"])
self.global_step = checkpoint["global_step"]
self.best_eval_loss = checkpoint["best_eval_loss"]
logger.info(f"Loaded checkpoint from step {self.global_step}")
def _log_metrics(self, metrics: dict) -> None:
"""Log metrics to tracker and console."""
if self.tracker:
self.tracker.log_metrics(metrics, step=self.global_step)
logger.info(f"Step {self.global_step}: {metrics}")
```
---
## Distributed Training
### PyTorch Distributed Data Parallel
```python
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler
import os
def setup_distributed() -> tuple[int, int, int]:
"""Initialize distributed training environment."""
if "RANK" in os.environ:
rank = int(os.environ["RANK"])
local_rank = int(os.environ["LOCAL_RANK"])
world_size = int(os.environ["WORLD_SIZE"])
else:
rank = 0
local_rank = 0
world_size = 1
if world_size > 1:
dist.init_process_group(
backend="nccl",
init_method="env://",
world_size=world_size,
rank=rank,
)
torch.cuda.set_device(local_rank)
return rank, local_rank, world_size
def cleanup_distributed() -> None:
"""Cleanup distributed training."""
if dist.is_initialized():
dist.destroy_process_group()
class DistributedTrainer(Trainer):
"""Trainer with DDP support."""
def __init__(self, *args, **kwargs):
self.rank, self.local_rank, self.world_size = setup_distributed()
super().__init__(*args, **kwargs)
def _setup_device(self) -> None:
"""Configure device for distributed training."""
if self.world_size > 1:
self.device = torch.device(f"cuda:{self.local_rank}")
self.model = self.model.to(self.device)
self.model = DDP(
self.model,
device_ids=[self.local_rank],
output_device=self.local_rank,
find_unused_parameters=False,
)
else:
super()._setup_device()
if self.config.mixed_precision and self.device.type == "cuda":
self.scaler = torch.amp.GradScaler("cuda")
else:
self.scaler = None
def save_checkpoint(self, name: str) -> Path:
"""Only save on rank 0."""
if self.rank == 0:
return super().save_checkpoint(name)
return None
def _log_metrics(self, metrics: dict) -> None:
"""Only log on rank 0."""
if self.rank == 0:
super()._log_metrics(metrics)
def create_distributed_dataloader(
dataset: Dataset,
batch_size: int,
world_size: int,
rank: int,
shuffle: bool = True,
) -> DataLoader:
"""Create DataLoader with distributed sampler."""
sampler = DistributedSampler(
dataset,
num_replicas=world_size,
rank=rank,
shuffle=shuffle,
)
return DataLoader(
dataset,
batch_size=batch_size,
sampler=sampler,
num_workers=4,
pin_memory=True,
drop_last=True,
)
```
### Launch Script
```bash
#!/bin/bash
# launch_distributed.sh
NUM_GPUS=4
MASTER_PORT=29500
torchrun \
--nproc_per_node=$NUM_GPUS \
--master_port=$MASTER_PORT \
train.py \
--config config/training_config.yaml
```
---
## Hyperparameter Tuning
### Optuna Integration
```python
import optuna
from optuna.trial import Trial
from optuna.integration import PyTorchLightningPruningCallback
import mlflow
def create_objective(
train_dataset: Dataset,
eval_dataset: Dataset,
model_class: type,
) -> callable:
"""Create Optuna objective function."""
def objective(trial: Trial) -> float:
# Sample hyperparameters
config = TrainingConfig(
model_name="tuned_model",
learning_rate=trial.suggest_float("lr", 1e-5, 1e-3, log=True),
batch_size=trial.suggest_categorical("batch_size", [16, 32, 64]),
weight_decay=trial.suggest_float("weight_decay", 1e-5, 1e-2, log=True),
epochs=trial.suggest_int("epochs", 3, 10),
warmup_steps=trial.suggest_int("warmup_steps", 0, 500),
)
# Create data loaders
train_loader = DataLoader(train_dataset, batch_size=config.batch_size, shuffle=True)
eval_loader = DataLoader(eval_dataset, batch_size=config.batch_size)
# Create model
model = model_class(
hidden_size=trial.suggest_categorical("hidden_size", [128, 256, 512]),
num_layers=trial.suggest_int("num_layers", 2, 6),
dropout=trial.suggest_float("dropout", 0.1, 0.5),
)
# Train
trainer = Trainer(
model=model,
config=config,
train_dataloader=train_loader,
eval_dataloader=eval_loader,
)
# Report intermediate values for pruning
for epoch in range(config.epochs):
trainer.train_epoch()
eval_loss = trainer.evaluate()["eval/loss"]
trial.report(eval_loss, epoch)
if trial.should_prune():
raise optuna.TrialPruned()
return trainer.best_eval_loss
return objective
def run_hyperparameter_search(
train_dataset: Dataset,
eval_dataset: Dataset,
model_class: type,
n_trials: int = 100,
study_name: str = "hpo_study",
) -> optuna.Study:
"""Run hyperparameter optimization with Optuna."""
# Create study with pruning
pruner = optuna.pruners.MedianPruner(
n_startup_trials=5,
n_warmup_steps=3,
interval_steps=1,
)
study = optuna.create_study(
study_name=study_name,
direction="minimize",
pruner=pruner,
storage=f"sqlite:///{study_name}.db",
load_if_exists=True,
)
objective = create_objective(train_dataset, eval_dataset, model_class)
study.optimize(
objective,
n_trials=n_trials,
timeout=3600 * 12, # 12 hours
n_jobs=1, # Sequential for GPU
show_progress_bar=True,
)
# Log best results
logger.info(f"Best trial: {study.best_trial.params}")
logger.info(f"Best value: {study.best_value}")
return study
```
### Ray Tune Integration
```python
from ray import tune
from ray.tune.schedulers import ASHAScheduler
from ray.tune.search.optuna import OptunaSearch
from ray.air import RunConfig, CheckpointConfig
def train_fn(config: dict) -> None:
"""Training function for Ray Tune."""
from ray.train import report, get_checkpoint
training_config = TrainingConfig(
model_name="ray_tune_model",
learning_rate=config["lr"],
batch_size=config["batch_size"],
weight_decay=config["weight_decay"],
epochs=config["epochs"],
)
# Build model and dataloaders
model = build_model(config["hidden_size"], config["num_layers"])
train_loader, eval_loader = build_dataloaders(config["batch_size"])
trainer = Trainer(
model=model,
config=training_config,
train_dataloader=train_loader,
eval_dataloader=eval_loader,
)
# Resume from checkpoint if available
checkpoint = get_checkpoint()
if checkpoint:
with checkpoint.as_directory() as checkpoint_dir:
trainer.load_checkpoint(Path(checkpoint_dir))
for epoch in range(training_config.epochs):
trainer.train_epoch()
metrics = trainer.evaluate()
# Report metrics to Ray Tune
report(
{"loss": metrics["eval/loss"], "epoch": epoch},
checkpoint=Checkpoint.from_directory(trainer.checkpoint_dir),
)
def run_ray_tune(num_samples: int = 50) -> tune.ResultGrid:
"""Run hyperparameter search with Ray Tune."""
search_space = {
"lr": tune.loguniform(1e-5, 1e-3),
"batch_size": tune.choice([16, 32, 64]),
"weight_decay": tune.loguniform(1e-5, 1e-2),
"hidden_size": tune.choice([128, 256, 512]),
"num_layers": tune.randint(2, 7),
"epochs": 10,
}
scheduler = ASHAScheduler(
metric="loss",
mode="min",
max_t=10,
grace_period=2,
reduction_factor=3,
)
tuner = tune.Tuner(
tune.with_resources(train_fn, {"gpu": 1}),
param_space=search_space,
tune_config=tune.TuneConfig(
num_samples=num_samples,
scheduler=scheduler,
search_alg=OptunaSearch(),
),
run_config=RunConfig(
name="hpo_experiment",
checkpoint_config=CheckpointConfig(
num_to_keep=3,
checkpoint_frequency=1,
),
),
)
results = tuner.fit()
best_result = results.get_best_result("loss", "min")
logger.info(f"Best config: {best_result.config}")
logger.info(f"Best loss: {best_result.metrics['loss']}")
return results
```
---
## Resource Management
### GPU Memory Optimization
```python
import torch
from contextlib import contextmanager
@contextmanager
def gpu_memory_manager():
"""Context manager for GPU memory cleanup."""
try:
yield
finally:
torch.cuda.empty_cache()
torch.cuda.synchronize()
def get_gpu_memory_usage() -> dict:
"""Get current GPU memory statistics."""
if not torch.cuda.is_available():
return {"available": False}
return {
"allocated": torch.cuda.memory_allocated() / 1e9,
"reserved": torch.cuda.memory_reserved() / 1e9,
"max_allocated": torch.cuda.max_memory_allocated() / 1e9,
}
class GradientCheckpointing:
"""Enable gradient checkpointing for memory efficiency."""
@staticmethod
def enable(model: nn.Module, checkpoint_layers: list[str] = None) -> None:
"""Enable gradient checkpointing on specified layers."""
if hasattr(model, "gradient_checkpointing_enable"):
model.gradient_checkpointing_enable()
return
# Manual checkpointing for custom models
from torch.utils.checkpoint import checkpoint
def create_custom_forward(module):
def custom_forward(*inputs):
return checkpoint(module._original_forward, *inputs, use_reentrant=False)
return custom_forward
for name, module in model.named_modules():
if checkpoint_layers and name not in checkpoint_layers:
continue
if hasattr(module, "forward"):
module._original_forward = module.forward
module.forward = create_custom_forward(module)
```
### Batch Size Finder
```python
def find_optimal_batch_size(
model: nn.Module,
sample_batch: dict,
device: torch.device,
min_batch_size: int = 1,
max_batch_size: int = 256,
) -> int:
"""Find maximum batch size that fits in GPU memory."""
model = model.to(device)
optimal_batch_size = min_batch_size
for batch_size in [2**i for i in range(int(np.log2(max_batch_size)) + 1)]:
if batch_size < min_batch_size:
continue
try:
# Create batch of target size
batch = {
k: v.repeat(batch_size // v.size(0) + 1, *[1] * (v.dim() - 1))[:batch_size]
for k, v in sample_batch.items()
}
batch = {k: v.to(device) for k, v in batch.items()}
# Forward pass
with torch.amp.autocast("cuda"):
outputs = model(**batch)
loss = outputs.loss
# Backward pass
loss.backward()
model.zero_grad()
torch.cuda.empty_cache()
optimal_batch_size = batch_size
except RuntimeError as e:
if "out of memory" in str(e):
torch.cuda.empty_cache()
break
raise
logger.info(f"Optimal batch size: {optimal_batch_size}")
return optimal_batch_size
```
---
## Best Practices
### Training Configuration Management
```yaml
# config/training_config.yaml
model:
name: transformer
hidden_size: 512
num_layers: 6
dropout: 0.1
training:
batch_size: 32
learning_rate: 1e-4
weight_decay: 0.01
epochs: 10
mixed_precision: true
gradient_accumulation_steps: 4
distributed:
enabled: true
backend: nccl
checkpointing:
save_every_n_steps: 1000
keep_n_checkpoints: 3
logging:
log_every_n_steps: 100
eval_every_n_steps: 500
```
### Reproducibility Checklist
```python
def ensure_reproducibility(seed: int) -> None:
"""Set all random seeds for reproducibility."""
import random
import numpy as np
import os
# Python
random.seed(seed)
# NumPy
np.random.seed(seed)
# PyTorch
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
# CUDA
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
# Environment
os.environ["PYTHONHASHSEED"] = str(seed)
logger.info(f"Set all random seeds to {seed}")
```
---
## Related References
- `feature-engineering.md` - Feature preparation for training
- `experiment-tracking.md` - Logging training metrics
- `pipeline-orchestration.md` - Orchestrating training pipelines
- `model-validation.md` - Validating trained models
## Cross-Reference Skills
- **DevOps Engineer** - CI/CD for training pipelines
- **Kubernetes Specialist** - K8s-based training infrastructure
No comments yet. Be the first to comment!
