Key Points
Why Machine Learning on HPC?
- HPC becomes essential when training takes > 4 hours or datasets exceed 10 GB
- Hyperparameter optimization and multi-GPU training turn weeks into days
- Always test code on small data first before scaling to full datasets
- Not every ML project benefits from HPC; evaluate before investing setup time
ML Workflow Overview
- GPUs accelerate deep learning 15-20x; CPUs are better for preprocessing and tree models
- Sagehen offers 10 GPUs total: A100 (80 GB), L40S (48 GB), and RTX PRO 6000 (96 GB ECC)
- Storage tiers (/rhome, /bigdata, /scratch) have different performance characteristics
- Successful ML on HPC combines local development with remote training
- CPU preprocessing + GPU training is more efficient than trying everything on GPU
Setting Up Your ML Environment
-
module load miniconda3, thenconda activatean existing environment, is the quickest way to start -
module show NAMEreveals what a module actually does (CUDA version, dependencies) - Always re-load modules inside SLURM scripts; the login-node environment does not transfer
- Sagehen’s GPU driver reports CUDA 12.7; the newest CUDA
toolkit module is
cuda/12.2.1. They are different things and do not need to match - Always test GPU access with small verification scripts before large jobs
- Roll your own conda environment when you need pinned versions or extra packages
- Conda environments can be moved to /bigdata if /rhome quota gets tight
Loading ML Frameworks
- Load miniconda3, then create environments with
conda create - Use
pytorch-cuda=12.1when installing PyTorch for Sagehen’s CUDA version - Channel order in
conda creatematters; put-c pytorch -c nvidiafirst - Mix conda and pip safely: conda first for frameworks, pip last for extras
- Export environments with
--from-historyfor portability, full export for exact reproduction - Common errors (missing CUDA, wrong versions, missing packages) are usually environment setup issues
- Watch /rhome quota; redirect conda to /bigdata if needed
- Store environment files in git repos for team consistency
Data Preparation Basics
- Use /scratch for intermediate ML data during jobs (4-7x faster than /bigdata)
- Specify dtypes when loading large CSVs to halve memory and load time
- Always split fit/transform between training and test data to avoid leakage
- Save preprocessing objects (scaler, encoder, imputer) alongside model checkpoints
- Check class balance before training; imbalance silently corrupts accuracy metrics
- pandas works well up to ~50 GB; switch to Dask above that
- Document the rationale for missing-value strategy, scaler choice, and encoding choice
Data Pipelines for HPC
- A good data pipeline overlaps storage I/O, CPU transforms, and GPU compute
- Use Dask for datasets larger than RAM that need parallel preprocessing
- Stage active training data on /scratch; keep /bigdata for source-of-truth
- /scratch is node-local and cleared at job end — anything a later job needs must be written to /bigdata before the job exits
- DataLoader settings that matter most: num_workers, pin_memory, persistent_workers
- Run preprocessing as a CPU job on the amd partition, training as a GPU job
- If GPU utilization oscillates, you have a data pipeline problem, not a model problem
Training with PyTorch
- Sagehen has 10 GPUs across three types (A100 80 GB, L40S 48 GB, RTX PRO 6000 96 GB); pick the right one for your job
- Move data and models to GPU using
.to(device, non_blocking=True)paired with pinned-memory DataLoader - Forward activations dominate GPU memory and scale with batch size
- Mixed precision gives ~2x speed and ~50% memory on A100/L40S; bfloat16 is more stable than float16
- Gradient accumulation simulates large batches when memory is tight
- Always checkpoint to /scratch and copy final results to /bigdata
- Test on CPU or a tiny dataset before submitting a real GPU job
Training with TensorFlow
- TensorFlow with Keras provides a high-level API for fast prototyping; tf.GradientTape gives full control
- Always set memory_growth=True on shared GPUs; otherwise TF grabs the whole device
- tf.data with AUTOTUNE is the equivalent of PyTorch DataLoader for parallel I/O
- Transfer learning is the default starting point: train the head first, then fine-tune
- Callbacks (ModelCheckpoint, EarlyStopping, ReduceLROnPlateau) save GPU-hours on long runs
- MirroredStrategy enables 2-GPU training but the speedup is sub-linear
- Choose framework based on ecosystem fit, not micro-benchmarks
Checkpointing Long Jobs
- Save checkpoints periodically to recover from interruptions, time limits, or hardware faults
- Include model state, optimizer state, epoch, and best validation loss in checkpoints
- Use
map_locationwhen loading checkpoints across different GPU nodes - A “resume if checkpoint exists” pattern makes scripts restartable from any point
- Save checkpoints to /scratch during the run, copy final to /bigdata at end
- DistributedDataParallel is the recommended approach for multi-GPU PyTorch
- TensorFlow MirroredStrategy handles distribution with much less boilerplate
- Save checkpoints only on rank 0 in distributed training
- Multi-GPU is worth ~1.7x; try mixed precision first for similar speedup with no complexity
SLURM Scripts for ML
- Use
--partition=gpuand--gres=gpu:Nto request GPU resources - Request 4 CPUs per GPU as a rule of thumb
- Use torchrun for launching distributed PyTorch training
- Test interactively with
srunbefore submitting batch jobs - Copy data to /scratch for faster I/O during training
Hyperparameter Tuning
- Grid search is fine for 2 to 3 dimensions, intractable beyond that
- Random search is competitive in higher dimensions and scales linearly
- Bayesian optimization wins when each trial is expensive (hours)
- Search log-uniformly for learning rate and weight decay; linear for dropout
- Early stopping inside each trial multiplies the effective number of trials
- SLURM job arrays parallelize sweeps with a single submission
- After a sweep, check that the best values are not at the edge of the range; expand if so
ML Best Practices
- Track experiments with TensorBoard, config files, and version control
- Always use separate train/val/test sets; touch test exactly once at the end
- Scale learning rate with batch size (sqrt or linear)
- Monitor GPU memory to avoid leaks during sweeps
- Save complete checkpoints including hyperparameters, git hash, and date
- Export models with TorchScript or SavedModel for deployment
- Keep a per-project run log so you can answer “which run produced figure 3?”