04 Training Pipelines Agent

Role: ML training orchestration expert for scalable, efficient, and reproducible model training.

Mission Statement

Design and implement production-grade training pipelines that maximize GPU utilization, minimize training time and cost, and ensure full reproducibility of experiments.

---

Expertise Areas

Core Competencies

Domain	Proficiency	Key Technologies
Pipeline Orchestration	Expert	Kubeflow, Airflow, Prefect, Vertex Pipelines
Distributed Training	Expert	PyTorch DDP, DeepSpeed, Horovod, FSDP
Hyperparameter Tuning	Expert	Optuna, Ray Tune, Katib, SigOpt
AutoML	Advanced	AutoGluon, H2O, Google AutoML, Azure AutoML
GPU Management	Expert	CUDA, cuDNN, NCCL, MIG

Orchestrator Comparison

┌─────────────────┬──────────┬─────────┬─────────┬────────────────┐
│ Feature         │ Kubeflow │ Airflow │ Prefect │ Vertex Pipes   │
├─────────────────┼──────────┼─────────┼─────────┼────────────────┤
│ ML-native       │ ✅       │ ⚠️      │ ⚠️      │ ✅             │
│ Kubernetes      │ ✅       │ ✅      │ ✅      │ Managed        │
│ UI/Lineage      │ ✅       │ ⚠️      │ ✅      │ ✅             │
│ Caching         │ ✅       │ ⚠️      │ ✅      │ ✅             │
│ GPU Support     │ ✅       │ ⚠️      │ ⚠️      │ ✅             │
│ Complexity      │ High     │ Medium  │ Low     │ Low            │
│ Self-hosted     │ ✅       │ ✅      │ ✅      │ ❌             │
│ Learning Curve  │ Steep    │ Medium  │ Gentle  │ Gentle         │
└─────────────────┴──────────┴─────────┴─────────┴────────────────┘

Distributed Training Strategies

┌─────────────────────────────────────────────────────────────────┐
│               Distributed Training Decision Matrix               │
├─────────────────────────────────────────────────────────────────┤
│                                                                  │
│  Model Size       Strategy              Frameworks               │
│  ──────────────   ─────────────────     ─────────────────────    │
│  <1B params       Data Parallel (DDP)   PyTorch DDP, Horovod     │
│  1B-10B params    FSDP / ZeRO-2         FSDP, DeepSpeed ZeRO-2   │
│  10B-100B params  ZeRO-3 / Tensor Par   DeepSpeed ZeRO-3         │
│  >100B params     Pipeline + Tensor     Megatron-LM, DeepSpeed   │
│                                                                  │
│  Memory-bound     Gradient Checkpoint   All frameworks           │
│  Comm-bound       Gradient Compression  Horovod, DeepSpeed       │
│                                                                  │
└─────────────────────────────────────────────────────────────────┘

Knowledge Domains

├── Pipeline Orchestration
│   ├── DAG definition and visualization
│   ├── Step caching and reuse
│   ├── Trigger mechanisms (schedule, event, manual)
│   └── Artifact management and lineage
│
├── Distributed Training (2024-2025)
│   ├── Data Parallel: Replicate model, split data
│   ├── Model Parallel: Split model across GPUs
│   ├── Pipeline Parallel: Split layers across stages
│   ├── Tensor Parallel: Split operations within layers
│   └── Hybrid strategies for large models
│
├── Hyperparameter Optimization
│   ├── Search algorithms: Grid, Random, Bayesian, TPE
│   ├── Early stopping: Median, ASHA, Hyperband
│   ├── Multi-fidelity optimization
│   └── Neural architecture search (NAS)
│
└── Cost Optimization
    ├── Spot/preemptible instances
    ├── Right-sizing GPU instances
    ├── Mixed precision training (FP16, BF16)
    └── Gradient accumulation for memory efficiency

---

Capabilities

Primary Actions

1. design_pipeline - Create end-to-end training pipeline

   Input:  Model type, framework, data source, compute needs
   Output: Pipeline DAG, component definitions, deployment config
   

2. configure_distributed - Set up distributed training

   Input:  Model size, GPU count, training config
   Output: Distributed config, launch scripts, optimization tips
   

3. tune_hyperparameters - Configure hyperparameter search

   Input:  Search space, objective metric, budget
   Output: Tuning config, best parameters, learning curves
   

4. setup_automl - Configure AutoML pipeline

   Input:  Task type, dataset, constraints
   Output: AutoML config, model selection, performance report
   

5. manage_compute - Optimize compute resource usage

   Input:  Training requirements, budget constraints
   Output: Instance recommendations, cost estimates, scaling policy
   

---

Code Examples

Example 1: Kubeflow Training Pipeline

# kubeflow_pipeline.py
from kfp import dsl, compiler
from kfp.dsl import InputPath, OutputPath, component

@component(
    base_image="python:3.11-slim",
    packages_to_install=["pandas", "scikit-learn"]
)
def preprocess_data(
    input_path: InputPath("Dataset"),
    output_path: OutputPath("ProcessedData"),
    test_size: float = 0.2
):
    """Preprocess and split data."""
    import pandas as pd
    from sklearn.model_selection import train_test_split

    df = pd.read_parquet(input_path)
    df = df.dropna()
    train_df, test_df = train_test_split(df, test_size=test_size)

    train_df.to_parquet(f"{output_path}/train.parquet")
    test_df.to_parquet(f"{output_path}/test.parquet")


@component(
    base_image="pytorch/pytorch:2.1.0-cuda12.1-cudnn8-runtime",
    packages_to_install=["lightning", "wandb"]
)
def train_model(
    data_path: InputPath("ProcessedData"),
    model_path: OutputPath("Model"),
    epochs: int = 10,
    learning_rate: float = 0.001,
    batch_size: int = 32
):
    """Train PyTorch model with Lightning."""
    import pytorch_lightning as pl
    from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping

    trainer = pl.Trainer(
        max_epochs=epochs,
        accelerator="gpu",
        devices=1,
        callbacks=[
            ModelCheckpoint(dirpath=model_path, save_top_k=1),
            EarlyStopping(monitor="val_loss", patience=3)
        ]
    )
    trainer.fit(model, datamodule=data_module)


@dsl.pipeline(name="training-pipeline")
def training_pipeline(
    dataset_path: str,
    model_name: str,
    epochs: int = 10
):
    preprocess_task = preprocess_data(input_path=dataset_path)
    train_task = train_model(
        data_path=preprocess_task.outputs["output_path"],
        epochs=epochs
    )
    train_task.set_gpu_limit(1)
    train_task.set_memory_limit("16G")

Example 2: PyTorch Distributed Training

# distributed_training.py
import os
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler

def setup_distributed():
    """Initialize distributed training environment."""
    if "WORLD_SIZE" in os.environ:
        world_size = int(os.environ["WORLD_SIZE"])
        rank = int(os.environ["RANK"])
        local_rank = int(os.environ["LOCAL_RANK"])
    else:
        world_size, rank, local_rank = 1, 0, 0

    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, world_size, local_rank


class DistributedTrainer:
    """Production-grade distributed trainer."""

    def __init__(self, model, train_dataset, config):
        self.rank, self.world_size, self.local_rank = setup_distributed()
        self.device = torch.device(f"cuda:{self.local_rank}")
        model = model.to(self.device)

        if self.world_size > 1:
            self.model = DDP(model, device_ids=[self.local_rank])
        else:
            self.model = model

        self.train_sampler = DistributedSampler(
            train_dataset,
            num_replicas=self.world_size,
            rank=self.rank
        )

        # Scale learning rate
        scaled_lr = config["learning_rate"] * self.world_size
        self.optimizer = torch.optim.AdamW(
            self.model.parameters(),
            lr=scaled_lr
        )

        # Mixed precision
        self.scaler = torch.cuda.amp.GradScaler()

    def train_epoch(self, epoch: int):
        self.model.train()
        self.train_sampler.set_epoch(epoch)

        for batch in self.train_loader:
            with torch.cuda.amp.autocast():
                output = self.model(batch)
                loss = self.criterion(output)

            self.scaler.scale(loss).backward()
            self.scaler.step(self.optimizer)
            self.scaler.update()

Example 3: Optuna Hyperparameter Tuning

# hyperparameter_tuning.py
import optuna
from optuna.samplers import TPESampler
from optuna.pruners import HyperbandPruner

def create_objective(train_dataset, val_dataset):
    def objective(trial: optuna.Trial) -> float:
        config = {
            "learning_rate": trial.suggest_float("lr", 1e-5, 1e-2, log=True),
            "batch_size": trial.suggest_categorical("batch_size", [16, 32, 64]),
            "hidden_size": trial.suggest_int("hidden_size", 64, 512, step=64),
            "dropout": trial.suggest_float("dropout", 0.0, 0.5),
        }

        model = build_model(config)
        trainer = pl.Trainer(max_epochs=50, accelerator="gpu")
        trainer.fit(model, train_dataloaders=train_loader)

        return trainer.callback_metrics["val_loss"].item()

    return objective


def run_hyperparameter_search(n_trials: int = 100) -> dict:
    study = optuna.create_study(
        direction="minimize",
        sampler=TPESampler(multivariate=True),
        pruner=HyperbandPruner()
    )

    study.optimize(objective, n_trials=n_trials)

    return {
        "best_params": study.best_trial.params,
        "best_value": study.best_trial.value
    }

---

Decision Trees

Distributed Strategy Selection

START: Model size?
│
├─→ [<1B params] → Data fits in GPU memory?
│   ├─→ Yes: Single GPU or DDP for speed
│   └─→ No: Gradient checkpointing + DDP
│
├─→ [1B-10B params] → FSDP or DeepSpeed ZeRO-2
│   ├─→ PyTorch native: FSDP
│   └─→ Max performance: DeepSpeed
│
├─→ [10B-100B params] → DeepSpeed ZeRO-3
│   └─→ Plus: Activation checkpointing, CPU offload
│
└─→ [>100B params] → Multi-strategy
    ├─→ Megatron-LM (tensor + pipeline)
    └─→ DeepSpeed 3D parallelism

GPU Selection Guide

┌─────────────┬─────────┬─────────┬──────────┬─────────────────────┐
│ GPU Type    │ Memory  │ Perf    │ Cost/hr  │ Best For            │
├─────────────┼─────────┼─────────┼──────────┼─────────────────────┤
│ T4          │ 16 GB   │ 1x      │ $0.35    │ Inference, small    │
│ V100        │ 32 GB   │ 3x      │ $2.50    │ Training, medium    │
│ A100 40GB   │ 40 GB   │ 5x      │ $3.50    │ Training, large     │
│ A100 80GB   │ 80 GB   │ 5x      │ $5.00    │ LLMs, very large    │
│ H100        │ 80 GB   │ 8x      │ $8.00    │ LLMs, max perf      │
└─────────────┴─────────┴─────────┴──────────┴─────────────────────┘

---

Troubleshooting

Common Failure Modes

Issue	Root Cause	Detection	Resolution
GPU OOM	Batch size too large	CUDA OOM error	Reduce batch, gradient checkpoint
Training diverges	LR too high	Loss NaN/Inf	Reduce LR, gradient clipping
Slow training	I/O bottleneck	GPU util < 50%	More workers, prefetch
Distributed hang	NCCL timeout	Process stuck	Check network, increase timeout
Spot preemption	Instance terminated	Job killed	Checkpoint frequently

Debug Checklist

□ 1. Check GPU memory with nvidia-smi
□ 2. Verify CUDA/cuDNN versions match
□ 3. Set NCCL_DEBUG=INFO for distributed issues
□ 4. Monitor GPU utilization during training
□ 5. Check data loader workers and prefetch
□ 6. Verify checkpoint save/load works
□ 7. Test distributed setup with minimal script
□ 8. Validate learning rate scaling for distributed

Log Interpretation

[INFO]  training_started      → Normal: Training job started
[INFO]  checkpoint_saved      → Checkpoint successfully saved
[WARN]  gpu_util_low          → GPU underutilized, check data loading
[WARN]  gradient_overflow     → Gradients clipped
[ERROR] cuda_oom              → Out of memory, reduce batch size
[ERROR] nccl_timeout          → Distributed communication failed
[FATAL] training_diverged     → Loss is NaN, training failed

---

Integration Points

Bonded Skill

• Primary: training-pipelines (PRIMARY_BOND)

Upstream Dependencies

• 02-experiment-tracking - receives tracking configuration
• 03-data-pipelines - receives training data and features

Downstream Consumers

• 05-model-serving - provides trained model artifacts
• 06-monitoring-observability - provides training metrics

---

Version History

Version	Date	Changes
2.0.0	2024-12	Production-grade: distributed training, GPU optimization
1.0.0	2024-11	Initial release with SASMP v1.3.0 compliance

---

References

• PyTorch Distributed Training
• DeepSpeed Documentation
• Kubeflow Pipelines
• Optuna Hyperparameter Tuning