Checkpoints

Learn how to use Baseten’s checkpointing feature to manage model checkpoints and avoid disk errors during training.

With checkpointing enabled, you can manage your model checkpoints seamlessly and avoid common training issues.

Benefits of checkpointing

• Avoid catastrophic out of disk errors: We mount additional storage at the checkpointing directory to help avoid out of disk errors during your training run.
• Maximize GPU utilization: A separate process uploads the checkpoint directory to cloud storage in the background, so your GPUs keep training instead of waiting on uploads.
• Seamless checkpoint management: Uploaded checkpoints are ready to browse, deploy, or load into later jobs.

Enable checkpointing

To enable checkpointing, add a CheckpointingConfig to the Runtime and set enabled to True:

from truss_train import definitions

training_runtime = definitions.Runtime(
    # ... other configuration options
    checkpointing_config=definitions.CheckpointingConfig(enabled=True)
)

Use the checkpoint directory

Baseten exports the $BT_CHECKPOINT_DIR environment variable in your job. Write your checkpoints to this directory so Baseten can back them up.

Writing a checkpoint to $BT_CHECKPOINT_DIR stages it for upload. A background process syncs the checkpoint directory to cloud storage on a recurring cycle, and a checkpoint becomes durable when its sync completes. If your training script deletes a checkpoint before it syncs, that checkpoint never reaches cloud storage: it doesn’t appear in your checkpoint list, and errors aren’t raised.

Fast checkpoint rotation is the common way to hit this. Rolling-checkpoint setups that keep only the last N checkpoints, and frameworks that write to a temporary path and then rename, can both remove a checkpoint within seconds of writing it. To keep every checkpoint:

• Keep enough checkpoint history that each checkpoint stays on disk through at least one full sync cycle.
• Confirm a checkpoint appears in truss train checkpoints list before your script deletes it locally.

Once a checkpoint has synced, deleting it locally is safe. Synced checkpoints persist in cloud storage even after you remove them from $BT_CHECKPOINT_DIR, so you can still reclaim disk space during long runs.

Browse checkpoints

Use the CLI to list and interactively explore checkpoint files for a job:

truss train checkpoints list --job-id abc123

In interactive mode, you can fuzzy-search checkpoints, navigate their directory tree, and inspect file contents, including tensor summaries for .safetensors files. See checkpoints list for all options.

Resume training from a checkpoint

To resume training from a saved checkpoint or initialize a new job from a previous run, configure a LoadCheckpointConfig on the Runtime. Baseten downloads the referenced checkpoints into $BT_LOAD_CHECKPOINT_DIR before your start_commands run:

from truss_train import definitions

load_checkpoint_config = definitions.LoadCheckpointConfig(
    enabled=True,
    checkpoints=[
        definitions.BasetenCheckpoint.from_latest_checkpoint(
            project_name="my-training-project",
        ),
    ],
)

For more information, see loading checkpoints.

Serve checkpoints

Serve your model checkpoints using Baseten’s serving infrastructure. Reference training checkpoints in your weights configuration using the bt:// URI scheme. See serving checkpoints for deployment details.

When you delete a job or project, all undeployed checkpoints are permanently deleted with no archival or recovery option. Deployed checkpoints aren’t affected. See Management for details.