galilai-group/stable-pretraining

What it solves

Training large neural networks for "foundation models" (models like CLIP, DINO, or SimCLR that learn general visual features) is notoriously fiddly: researchers must juggle data loading, augmentation, logging, evaluation, and cluster job management, and small mistakes can silently corrupt a long run. This project is a framework that wraps PyTorch Lightning to make that whole workflow simpler, more flexible, and more stable — letting researchers focus on the model and loss while the framework handles the plumbing.

How it works

The framework is built around four components that pass Python dictionaries to each other: a DataModule (data), a Module (model + forward pass), Callbacks (monitoring/evaluation hooks), and a Trainer/Manager (orchestration). The key idea is that your forward function just returns a dictionary of tensors (like {"loss": ..., "embedding": ...}), and everything else — loggers, evaluation probes, checkpoints — reads from that dictionary without you having to modify the training loop. It ships with pre-built "forward" recipes for many self-supervised learning methods (SimCLR, DINO, MAE, BYOL, CLIP, etc.), plus callbacks like OnlineProbe (trains a small linear classifier on frozen features to track accuracy live) and OnlineKNN (a zero-training nearest-neighbor evaluator). It also moves data augmentation onto the GPU (via kornia) to remove a CPU bottleneck, and provides a Manager that adds SLURM cluster features like job requeue/resume, atomic checkpoints, and a queryable run registry. An experimental JAX/Flax-NNX backend mirrors the same design.

Who it’s for

Researchers and engineers doing foundation-model or self-supervised learning research — especially those training large vision models on GPU clusters and wanting live evaluation, robust checkpointing, and less boilerplate. It’s also useful for anyone who wants to quickly prototype or benchmark SSL methods (SimCLR, DINOv2, MAE, etc.) with minimal code.

Highlights

• Dict-based design: forward functions return a state dict, so any intermediate tensor is automatically loggable and callbacks attach without touching the training loop.
• 30+ built-in recipes spanning SSL, supervised, and multi-modal pretraining (SimCLR, DINO/DINOv2, MAE, BYOL, VICReg, Barlow Twins, LeJEPA, CLIP, …).
• Live evaluation callbacks like OnlineProbe and OnlineKNN that monitor representation quality during training.
• GPU-side batched augmentation (via kornia) that vectorizes augmentation across the batch, with measured throughput improvements across model sizes and precisions.
• SLURM-grade orchestration: Manager handles preempt/requeue, atomic checkpoints, and a queryable run registry.
• Experimental JAX/Flax-NNX backend mirroring the torch design, with parity regression tests for numerical claims.