How to Save and Restore an ML Training Environment (Your Checkpoint Saves the Model, Not the Machine)

Run torch.save and you feel safe. The weights are on disk, the loss curve is preserved, and if the box dies you can pick up where you left off. That belief is what most training setups are built on, and it is only half true. I build the infrastructure underneath GPU workloads, not the models that run on top, and after a few weeks reading r/LocalLLaMA, r/MachineLearning, and a stack of framework and provider docs, the same quiet failure keeps surfacing. People save a checkpoint, lose the box, spin up a fresh GPU somewhere cheaper, and then spend the next hour finding out that the checkpoint alone will not bring their run back. To save and restore an ML training environment, you have to save a lot more than the model.

TL;DR: A checkpoint saves your model: weights, optimizer state, the random number generator, the data loader position. It does not save the machine that trained it. The venv, the exact CUDA build, the system libraries, your custom code at a specific commit, and the data layout all live on the box, and a fresh GPU has none of them. Restoring a training run means restoring both the model checkpoint and the whole environment around it. The reliable way to carry the environment is to snapshot the entire box, not just the volume it wrote to.

What does a checkpoint actually save?

A checkpoint saves the state of your model, not the state of your machine. At minimum that is the weights and the optimizer state, and a careful one also stores the learning-rate scheduler, the random number generator states, and the data loader position so training resumes deterministically. That is the model's memory. It says nothing about the computer.

The word "checkpoint" quietly does two jobs, and framework docs are clear on the first. A model checkpoint captures the architecture, the weights, the optimizer state, and the variables needed to resume training, which is why you can restart from the last save instead of from epoch zero (Milvus, "What is model checkpointing?"). The PyTorch reproducibility guidance goes deeper: to resume identically you have to checkpoint every stateful object, meaning the weights, the optimizer state, the RNG states, and the data-generator state (Hotegni, "Reproducibility and Training State Management in PyTorch"). Every item on that list lives inside your process. None of it is the venv, the driver, or the disk.

So the checkpoint is doing the model's job correctly. The problem is what people assume it covers.

Your checkpoint saves the model, not the machine

Here is the part the tutorials skip. A training run does not float in the air. It runs on top of a whole stack of things that took an hour to assemble and that a .pt file has no idea exist:

• the Python environment: the venv or conda env, and the exact resolved versions of torch, transformers, CUDA libraries, and a hundred transitive deps
• the system layer: the NVIDIA driver, the CUDA toolkit, apt packages that a wheel compiled against, LD_LIBRARY_PATH and other env vars
• your own code, checked out at the specific commit that produced this checkpoint, including whatever local edits never got pushed
• the data: the dataset on disk, the preprocessed cache, the tokenizer files, the layout your data loader expects at fixed paths
• the config: the exact hyperparameters, the sweep index, the shell history that documents how you actually launched the job

A checkpoint stores none of that stack. This is the distinction the reproducibility literature keeps drawing: reproducing a result is not just running the same code, it is holding the software environment and hardware constant across runs, because a different library or CUDA version changes the outcome (Neptune, "How to Solve Reproducibility in ML"). The checkpoint is the output. The environment is the machine that computed it, and the machine is what disappears when the box goes away.

That gap is why "I saved my work" and "I can get back to work" are two different claims. The checkpoint makes the first true. Only saving the environment makes the second true.

Why does a restored checkpoint still fail on a fresh box?

Because the fresh box has your model but not the environment it runs in. torch.load expects the same torch build, the same CUDA runtime, and your custom classes at the same import paths. A clean GPU has a different driver, a bare OS, and none of your code, so the load throws errors before a step runs.

This is not a rare edge case. It is the default outcome, and the dependency layer bites hardest. torch is pinned to a specific CUDA build, and half the ecosystem's wheels either compile against the host driver or expect a CUDA version the new box did not ship with. Teams lean on conda precisely because a plain pip install misses the non-Python system dependencies that make or break a GPU project, and even then the CUDA toolkit inside conda regularly needs manual LD_LIBRARY_PATH and CONDA_OVERRIDE_CUDA surgery to line up (BSWEN, "Why Deep Learning Projects Still Use Conda"; Oxford Protein Informatics Group, "CUDA Toolkit in Conda"). Stack the rest on top: research code at the wrong commit, a dataset at a path that no longer exists, a preprocessing cache to rebuild from raw files. Any one turns "restore and resume" into "debug for an hour, then rebuild."

People who do this for a living feel the tax directly. One team running spot GPUs described the loop plainly:

"each attempt to initiate a new instance often leads to a chain of continuous preemptions... time required for 'context switching', typically over 10 minutes to set up each new instance, causes significant training interruptions. Such a continuous loop can trap the training process in an endless cycle of starts and stops."

That is Lunit's engineering team, writing about spot GPUs. The ten minutes is not the model loading. It is the environment being rebuilt, over and over, because the checkpoint never carried it.

What a complete save actually includes

If a checkpoint is the model and the rest is the machine, a real "save and restore" has to account for both. It helps to lay out every layer against what each common tool actually captures.

Each conventional tool covers one band and leaves the neighbors exposed. The checkpoint nails the model and stops. A network volume holds your data and, set up carefully, your env, but it stays pinned to one datacenter and keeps billing while the GPU is off (RunPod, network volumes docs). The only column that covers every row is the one that captures the box itself.

How to save and restore the whole training environment

There is a ladder here, and most people climb it one painful rung at a time. Each rung is a real improvement, and each one still leaves something on the machine you did not save.

1. Pin and reproduce. Freeze exact versions in a requirements.txt or a conda environment.yml, script your setup, and commit it. This is the baseline and you should do it regardless. But a lockfile is a recipe, not a meal: the next box still has to re-resolve, re-download multi-gigabyte wheels, recompile CUDA extensions, and can still land a different CUDA build than the one you tested (Neptune reproducibility guide).

2. Bake an image. Put the whole environment into a Docker image so the next box pulls a fixed layer instead of rebuilding. Now your driver, CUDA, and packages travel together. The catch is that an image is a starting state, not a live one: it holds no latest checkpoint, no uncommitted code, and no data cache you built after boot, so you are still mounting or re-downloading those on top.

3. Mount a volume for the heavy files. Keep the dataset and model files on a persistent volume so a new pod attaches instead of re-fetching 20GB from Hugging Face. This is why people pay for storage they never turn off. It solves data. It does not solve the venv, the driver, or provider lock, because the volume cannot leave its datacenter.

4. Snapshot the whole box. Capture the entire filesystem at real paths, the environment, the code, the data, and the latest checkpoint together, then restore it as one unit on the next GPU. This is the only rung where a resume is a restore, not a rebuild, because nothing is left on the old machine that you forgot to carry. Historically no general tool did this across providers, which is why practitioners hand-rolled the worse version with rclone-to-S3 crons and 200-line boot scripts.

The higher you climb, the less the fresh box has to reconstruct. The full-box snapshot is the top of the ladder because it stops treating the environment as something to reassemble and starts treating it as something to move.

The environment is a versioned artifact, not a README

Step back and the real problem is a mismatch in how we version things. Your code lives in git. Your infrastructure lives in Terraform. Your container is a tagged image. Every expensive, hard-to-reproduce part of your stack is a versioned, portable artifact, except the single most expensive one: the GPU box itself, which we still rebuild from a README and a prayer every time it dies.

Treating the training environment as a first-class artifact means giving the box the same verbs you already use on code: snapshot it, restore it, carry it. That is the standard we are building in the open at Aquanode with ogre, a single-binary CLI that snapshots an entire GPU box, the full filesystem and not just a data volume, so the environment and the checkpoint restore together on whatever provider is cheapest and in stock. We have run the round-trip and it works; cross-provider reliability at scale is something we are still proving in public rather than asserting, so treat this as the direction we are building toward, not a benchmark to take on faith. The point is the shape of the fix: a saved training environment should behave like a saved commit, not a setup script you rerun and hope.

What you should actually do

Keep saving your checkpoints. They are doing their job. Just stop assuming they save your training environment, because the model and the machine are two different things and only one fits in a .pt file. Pin your dependencies and commit the lockfile. Bake an image so the driver and CUDA stop being a fresh gamble on every box. Put your data on storage you do not re-download. And when losing an hour to a rebuild actually costs you, whether the box got reclaimed, the balance ran dry, or you found a cheaper GPU, reach for a save that captures the whole machine, not only the weights. The test is simple: could you delete this box right now and be back to work in one command on a different one? If the answer is no, you saved the model, not the machine.

About the author

I am on the team at Aquanode, where we build ogre, an open-source tool for snapshotting and restoring an entire GPU box across providers. I work on the infrastructure layer under ML workloads, and I spend a lot of time reading how practitioners actually lose and rebuild their training environments so we can make that loop disappear.

Sources

• Milvus, "What is model checkpointing?". A checkpoint captures architecture, weights, optimizer state, and resume variables.
• Hotegni, "Reproducibility and Training State Management in PyTorch". Resuming identically requires checkpointing weights, optimizer, RNG, and data-generator state.
• Neptune, "How to Solve Reproducibility in ML". Reproducibility is holding the software environment and hardware constant, not just the code.
• BSWEN, "Why Deep Learning Projects Still Use Conda Despite Its Reputation Problems". Conda is used for the non-Python system dependencies that make or break GPU projects.
• Oxford Protein Informatics Group, "Tips and Tricks to correct a CUDA Toolkit installation in Conda". CUDA-in-conda mismatches and the manual env-var fixes they need.
• Lunit, "Optimizing GPU Costs by Leveraging Spot Instances". The context-switching loop: over 10 minutes to set up each new instance.
• RunPod, network volumes documentation. Volumes persist data but stay datacenter-locked and bill while the GPU is off.