galilai-group/stable-pretraining

Reliable, minimal and scalable library for pretraining foundation and world models

stable‑pretraining – A flexible framework for self‑supervised learning

What it is

  • An open‑source Python library built on top of PyTorch Lightning (with an experimental JAX/Flax‑NNX backend) that lets researchers prototype, train, and evaluate self‑supervised learning (SSL) methods with minimal boilerplate.
  • The core idea is to only implement a forward function that returns a dictionary containing a loss (for training) and any other tensors you want to log. The library takes care of the usual Lightning plumbing – data loading, checkpointing, resume, logging, and a local experiment registry.

Main components

Component Role
Data Handles dictionary‑style datasets. Supports HuggingFace datasets, conversion from standard PyTorch datasets, and custom dict‑returning datasets. Provides a gpu_transform hook that moves heavy multi‑view augmentations onto the GPU via Kornia for higher throughput.
Module Thin wrapper around a backbone, projector, and a user‑supplied forward callable (e.g., forward.simclr). No need to write training_step/validation_step; the wrapper extracts loss from the returned dict.
Callbacks Ready‑made utilities for online evaluation (linear probe, K‑NN, RankMe, LiDAR, CLIP zero‑shot, image retrieval, latent visualisation), teacher‑student EMA, weight‑decay schedules, checkpoint handling (WandB, HuggingFace, sklearn models), and system logging (SLURM info, environment dump, etc.).
Trainer / Manager Thin façade over Lightning’s Trainer. Manager bundles the trainer, module, and data module, runs the training loop, writes logs, and registers the run in a local cache_dir registry.
JAX backend (experimental) Mirrors the PyTorch API under stable_pretraining.jax. Uses Flax‑NNX, Optax, and JAX‑native primitives while keeping the same configuration files, callbacks, and logging format.

Built‑in SSL methods

  • SimCLR, BYOL, VICReg, Barlow Twins, SimSiam, SwAV‑Sinkhorn, NT‑Xent, InfoNCE, and the newer LeJEPA and Jet invertible encoder. New methods can be added by providing a custom forward function.

Backbones & optimizers

  • MLP, ResNet‑{9,18,34,50,101,152}, ConvMixer, ViT‑{tiny,small,base,large}.
  • Optimizers include Adam, LARS, etc.; schedulers such as CosineAnnealingLR are supported.

Quick start

# Install the latest release
pip install stable-pretraining

# Or install the development version
git clone https://github.com/galilai-group/stable-pretraining.git
cd stable-pretraining
pip install -e .
python -m stable_pretraining.quickstart   # runs one CPU epoch of SimCLR on synthetic data

The command launches a tiny training run and opens a local web viewer (spt web ./spt-quickstart/runs) that displays loss curves and the accuracy of the online linear probe.

Typical workflow

  1. Define data – use spt.data.HFDataset or wrap a PyTorch dataset; optionally add a gpu_transform for fast multi‑view augmentations.
  2. Pick a method – choose a pre‑implemented forward (e.g., spt.forward.simclr) or write your own.
  3. Create a Module – pass backbone, projector, loss, optimizer config, and the forward.
  4. Add callbacks – e.g., OnlineProbe to monitor downstream performance.
  5. Run – spt.Manager(trainer=trainer, module=module, data=dm)().

Why it matters

  • Stability: All training loops go through Lightning, so you get robust checkpointing, multi‑GPU (DDP) support, and automatic resume.
  • Flexibility: Changing the SSL objective is as simple as swapping the forward callable; the rest of the pipeline (data handling, logging, evaluation) stays unchanged.
  • Speed: GPU‑side augmentations can dramatically increase throughput on modern hardware, as demonstrated in the repo’s benchmark tables.
  • Cross‑framework parity: The experimental JAX path lets you compare Torch vs. JAX implementations side‑by‑side with identical configs and logging.

Documentation & community

Who should use it

  • Researchers building new SSL objectives who want a reliable training scaffold.
  • Engineers needing a reproducible, checkpoint‑friendly pipeline for large‑scale pre‑training.
  • Anyone interested in comparing Torch and JAX implementations of the same method.

All details above are taken directly from the repository’s README; no additional features have been inferred.

Related

  • Project
    Lightning-AI/pytorch-lightning

    A deep learning framework that abstracts away the engineering boilerplate of PyTorch, enabling researchers to pretrain and finetune AI models at scale across various hardware.

  • Project
    sktime/pytorch-forecasting

    A PyTorch-based library for time series forecasting that provides a high-level API and state-of-the-art deep learning architectures like Temporal Fusion Transformers and N-BEATS.

  • Project
    open-mmlab/mmengine

    A foundational PyTorch library that serves as a training engine for deep learning models, providing standardized components for training loops, configuration, and monitoring.

  • Project
    kengz/SLM-Lab

    A modular PyTorch framework for Deep Reinforcement Learning that allows users to configure and run RL experiments via JSON specs across various Gymnasium environments.

  • Project
    facebookresearch/multimodal

    A PyTorch library for training state-of-the-art multimodal models at scale, providing modular building blocks and pretrained weights for models like CLIP and BLIP-2.