aigroup-workflow 2.2.1 → 2.2.2

package/skills/ai-ml/ml-pipeline/references/training-pipelines.md

@@ -1,782 +1,782 @@
-# 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
+# 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