aigroup-workflow 2.1.2 → 2.2.0

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

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