ML Training

Overview

Use this skill for model training across PyTorch, TensorFlow/Keras, JAX/Flax, Hugging Face Transformers/Diffusers/Accelerate/PEFT, scikit-learn, XGBoost, LightGBM, CatBoost, Spark MLlib, and Ray. Optimize for correctness first: validated data, leakage-safe splits, reproducible configuration, meaningful metrics, and a simple baseline before complex distributed or accelerator-heavy runs.

Training Readiness Checklist

1. Define task, target, metric, baseline, and acceptance threshold.
2. Validate data schema, label quality, missingness, duplicates, class balance, and train/serving feature parity.
3. Choose split strategy: random stratified for iid classification, grouped for correlated entities, time-based for temporal data, nested CV for model-selection claims.
4. Pin environment: framework versions, CUDA/cuDNN, drivers, dataset snapshot, preprocessing code, model config, and hardware.
5. Set seeds where meaningful and record nondeterministic operations. Do not promise bitwise reproducibility across GPUs or distributed kernels unless deterministic modes are verified.
6. Start with a small overfit test on a tiny batch, then a full baseline, then tuning or scale-out.

Training Loop Essentials

For deep learning, every training loop should make the forward pass, loss computation, backward pass, optimizer step, scheduler step, gradient zeroing, metric logging, validation, checkpointing, and early stopping explicit.