stable-pretraining: a Lightning wrapper that makes every intermediate tensor loggable
Reliable, minimal and scalable library for pretraining foundation and world models
At a glance
- What is it?
- stable-pretraining is an MIT-licensed PyTorch library that assembles Lightning, HuggingFace and TorchMetrics into a pretraining harness where components exchange dicts instead of tuples. The design makes online evaluation cheap to attach, but it also means you inherit Lightning's and timm's constraints wholesale.
- Who is it for?
- Adopt stable-pretraining if you are already running Lightning and want online probes, KNN evaluation and SLURM requeue without writing your own callback layer; the dict-shaped forward makes that attachment point explicit. Do not adopt it if you need the JAX path for production work, since the README labels that backend experimental, or if your training loop cannot be expressed as backbone plus a forward function returning a state dict.
- Can I use it commercially?
- Yes. MIT is a permissive licence: you can use, modify and sell software built on it, as long as you keep its copyright and licence notices.
- Is it still maintained?
- Yes. The repository last received commits 62 days ago.
- What is it written in?
- Mainly Python, according to GitHub's language statistics.
Answers come from the project's GitHub data, last synced on September 15, 2026, and from our analysis. They are not legal advice.
DEEP OPEN-SOURCE ANALYSIS
The problem: pretraining code where nothing is observable until the run ends
Self-supervised pretraining runs are long and their loss curves are close to uninformative. Contrastive and joint-embedding methods can show a falling training loss while the representation collapses, and the only honest signal is a downstream probe. In most research codebases that probe is bolted on after the fact: a separate script loads a checkpoint, extracts features, trains a linear head, and reports a number hours later. By then the run is finished and the GPU time is spent.
stable-pretraining targets that gap. The README describes the project as a library for pretraining foundation and world models, and lists 30+ ready recipes spanning SSL, supervised and multi-modal pretraining, naming SimCLR, DINO/DINOv2, MAE, BYOL, VICReg, Barlow Twins, LeJEPA and CLIP. The intended user is a researcher or small team that wants those recipes without maintaining a private fork of someone's training loop, and that wants evaluation to happen while training is still running. The pitch is closer to a harness than to a model zoo: the repository ships recipes, but the design work is in the plumbing that connects data, module, callbacks and trainer.
Dicts as the interface between the four components
The README states the whole framework is four components that pass dicts to each other: DataModule, Module, Callbacks and the Lightning Trainer, with a Manager wrapping the outside. The diagram in the README shows batches flowing from the DataModule into the Trainer and then into the Module, the Module emitting state into callbacks and loggers, and the Manager sitting alongside the Trainer rather than inside it.
The consequence of the dict convention is that any intermediate tensor the forward function puts into its returned state becomes addressable by name. The README gives the example of a forward function returning {"loss": ..., "embedding": ...}, and the callback examples then reference those keys directly: OnlineProbe takes input="embedding" and target="label", OnlineKNN takes the same input key plus queue_length, input_dim and k. This is the mechanism that makes the evaluation callbacks work without editing the training step. A callback subscribes to a named key rather than monkey-patching the module or re-deriving features from the backbone.
The same convention is the main constraint. If your forward does not return a dict, or returns keys the callbacks do not expect, the hooks have nothing to attach to. Anything you want to monitor has to be surfaced deliberately by the code you write, which is a small but real tax on every new method you add.
Getting a run started: pip install, HFDataset, DataModule, Manager
The install path in the README is a clone plus an editable install: git clone https://github.com/galilai-group/stable-pretraining.git, then cd stable-pretraining and pip install -e . There is also a tutorial notebook at examples/simclr_cifar10_tutorial.ipynb, opened with jupyter notebook, which the README presents as an interactive walkthrough covering data loading, Module, callbacks, training and evaluation in one place.
The README also documents an spt CLI, listed in the table of contents as Quick Start with spt CLI, though the truncated material does not show its subcommands or flags. Treat the CLI as a thing that exists rather than something you can plan around until you read the full docs.
For wiring, the tour constructs datasets with spt.data.HFDataset, for example spt.data.HFDataset("cifar10", split="train", transform=...), wraps them in ordinary torch.utils.data.DataLoader objects, and passes the resulting loaders to spt.data.DataModule as train= and val=. The Module is built from spt.backbone.from_timm("resnet18", num_classes=0), an spt.backbone.MLP projector, a forward function such as spt.forward.simclr, and a loss such as spt.losses.NTXEntLoss(temperature=0.5). Callbacks are attached by passing them to the Lightning Trainer's callbacks list, and the run is launched by constructing spt.Manager(trainer=trainer, module=module, data=dm) and calling the resulting object. The README notes that on a workstation the Manager is a thin call, while on a cluster it adds preempt and resume plus run tracking. The configuration layer is described as global configuration with an output directory keyed by cache_dir, and a queryable run registry; the truncated README does not spell out the config file format or the registry query API, so check the documentation site before assuming either.
Online evaluation is the part worth copying
The callback set is the clearest expression of the project's thesis. OnlineProbe trains a linear head against a named embedding during the run, and OnlineKNN keeps a queue of embeddings (queue_length=10_000 in the example) and evaluates with k=20 neighbours. The README also names RankMe and LiDAR among the callbacks, both of which are representation-quality measures that need only the embedding and not the labels.
This matters because it changes when you learn that a run is failing. A probe or KNN score that stays flat across epochs is visible while the job is still occupying the node, which is when you can still kill it. The cost is that the probe and the KNN queue add work to the training step, and the KNN queue in particular holds embeddings in memory; the example's 10,000-entry queue is a knob you will want to size against your embedding dimension and GPU memory rather than copy blindly.
A second detail worth noting is that the README mentions GPU-side batched augmentation via a gpu_transform= argument, with CPU transforms handling decode and resize. That split is a reasonable default for dataloader-bound pipelines, but it also means augmentation code has to be written twice or written to run on tensors, and the README does not describe how the two paths are kept consistent.
Cluster behaviour: requeue, atomic checkpoints, run registry
The Manager is the piece aimed at shared clusters. The README describes it as wrapping fit() with SLURM-grade requeue, atomic checkpoints and a queryable run registry, and the tour comment says that on a cluster it adds preempt and resume plus run tracking for free. Atomic checkpoints matter on preemptible partitions because a job killed mid-write can leave a truncated file that a naive resume will happily load; the README asserts atomicity but the truncated material does not describe the write-then-rename mechanism, so verify it in the source if your resume logic depends on it.
The run registry is the other half: a record of runs you can query rather than infer from directory names. The README lists it as a top-level section and the Manager as the component that maintains it, but the schema and the query interface are not in the supplied text. If you run many short jobs, that registry is what stops you from reconstructing experiment history from filenames, and it is worth checking whether it stores the config alongside the metrics before you build tooling on top of it.
Where it is the wrong tool
The README is explicit that the JAX / Flax-NNX backend is experimental. It mirrors the same design, with forward-dict, callbacks and the Manager, and the badge in the README marks JAX as experimental rather than supported. If your team is standardising on JAX, this project is not yet a safe default, and the README's own framing says as much.
The second boundary is architectural. The library is built by assembling PyTorch, Lightning, HuggingFace and TorchMetrics, and the README says it adopts a modular design for integrating components from external libraries including architectures, loss functions, evaluation metrics and augmentations. That is a deliberate choice to not reimplement what already exists, and it means the upgrade surface is not just stable-pretraining. A breaking change in Lightning's Trainer or in a timm backbone signature reaches you through this library. If your training loop is a custom multi-stage pipeline with unusual control flow, or if you need to modify the optimizer step itself, the Lightning Trainer underneath becomes something you fight rather than something that helps.
The third case is scale of a different kind. If you are pretraining one model once, the recipes and callbacks may be more structure than you need; the value here compounds when you run many experiments and need comparable evaluation across them.
Alternatives and the actual difference in approach
The obvious comparison is to the Lightning ecosystem's own template repositories and to Lightning Fabric, which gives you the distributed and precision plumbing without the Trainer abstraction. Fabric's approach is the opposite of this project's: you keep your own training loop and add distributed primitives to it, so there is no fixed four-component structure and no callback contract to satisfy. The trade is that evaluation hooks, checkpoint atomicity and SLURM requeue are yours to write. stable-pretraining's bet is that a fixed structure plus a dict contract is worth more than loop-level freedom for pretraining specifically.
The other comparison is to Lightning-Bolts-style collections of self-supervised methods, which typically ship a model and a loss but leave the training and evaluation loop to you. Here the recipes are the smaller part of the offer; the callbacks, the Manager and the registry are the parts that do not have an obvious equivalent in a model collection. If you already have a training loop you trust and only want implementations of SimCLR or BYOL losses, you are paying for a lot of harness you will not use.
Maintenance cost and licence
The repository is MIT licensed, which is permissive and imposes no copyleft obligation on your own code. That is a statement about the licence text, not legal advice; if you are shipping a product that bundles the library, have someone check the notices you need to carry.
The release cadence visible in the supplied material is three releases in roughly four months: v0.1.6 in March 2026, v0.1.7 in May 2026 and v0.1.8 in July 2026, all 0.1.x. A 0.1.x line with that cadence means the public API is still moving, and the README's breadth (30+ recipes, a CLI, a registry, an experimental second backend) suggests a project that is expanding faster than it is freezing. Pin the version you validate against and read the release notes before bumping. The last push recorded is 2026-07-16, so the project is active rather than dormant, but activity is not the same as API stability.
Your practical upgrade cost is the intersection of three dependency trees: PyTorch, Lightning and timm, plus HuggingFace for the dataset layer. Budget for reading release notes on all of them, not just this one.
Editorial conclusion
Adopt stable-pretraining if you are already running Lightning and want online probes, KNN evaluation and SLURM requeue without writing your own callback layer; the dict-shaped forward makes that attachment point explicit. Do not adopt it if you need the JAX path for production work, since the README labels that backend experimental, or if your training loop cannot be expressed as backbone plus a forward function returning a state dict. Before committing, verify the exact signatures of spt.Module, spt.forward.simclr and spt.Manager in the version you install, and confirm that the checkpoint and registry files land where your cluster expects.
Community notes