Checkpointing PyTorch Models to Google Drive in Colab: Surviving Ephemeral Runtimes

The problem: runtimes disappear

Google Colab gives you a free GPU with a few clicks, but the virtual machine behind the notebook is ephemeral. After an idle timeout or when the maximum session length is reached, Colab reclaims the VM and its /content filesystem. Anything you have not copied elsewhere — model weights, optimizer state, training logs — vanishes instantly. A multi‑hour training run can be lost because you stepped away for coffee.

Mounting Google Drive as durable storage

The standard way to persist data across runtime resets is the google.colab.drive module. Running the following in a notebook cell mounts your personal Google Drive under /content/drive after a short OAuth consent screen:

from google.colab import drive
drive.mount('/content/drive')

Where to run: any Colab notebook cell. Permissions: you must approve the OAuth prompt once per session. Result: a directory /content/drive/MyDrive that survives runtime deletion. Risk: the mount behaves like a network filesystem; latency and throughput are lower than the local SSD.

Checkpoint‑and‑resume pattern (PyTorch)

Treat /content as scratch space and write checkpoints to a Drive folder. The snippet below shows a minimal training loop that saves the model, optimizer, and current epoch each epoch, then resumes from the newest checkpoint on startup.

import os, glob, torch
from google.colab import drive

drive.mount('/content/drive')
CKPT_DIR = '/content/drive/MyDrive/checkpoints'
os.makedirs(CKPT_DIR, exist_ok=True)

def save_checkpoint(model, optimizer, epoch, path):
    torch.save({
        'epoch': epoch,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
    }, path)

def load_latest_checkpoint(model, optimizer):
    checkpoints = glob.glob(os.path.join(CKPT_DIR, 'epoch_*.pt'))
    if not checkpoints:
        return 0  # start from epoch 0
    latest = max(checkpoints, key=os.path.getmtime)
    ckpt = torch.load(latest, map_location='cpu')
    model.load_state_dict(ckpt['model_state_dict'])
    optimizer.load_state_dict(ckpt['optimizer_state_dict'])
    print(f'Resumed from {latest} at epoch {ckpt["epoch"]}')
    return ckpt['epoch']

# Example model / optimizer
model = torch.nn.Linear(10, 2).cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

start_epoch = load_latest_checkpoint(model, optimizer)

for epoch in range(start_epoch, 100):
    # ... training step ...
    if epoch % 5 == 0:          # checkpoint every 5 epochs
        ckpt_path = os.path.join(CKPT_DIR, f'epoch_{epoch:03d}.pt')
        save_checkpoint(model, optimizer, epoch, ckpt_path)

Placeholders: replace the model/optimizer definitions with your own. Expected check: after a forced disconnect (Runtime → Disconnect and delete runtime), re‑run the notebook; the load_latest_checkpoint call should print the resumed epoch and the checkpoint file should still exist under MyDrive/checkpoints. Risk: frequent large checkpoints add wall‑clock time because each torch.save writes over the network mount.

I/O trade‑offs and mitigations

• Local‑first workflow: copy large datasets from Drive to /content at session start, train entirely on the local SSD, and push only checkpoints and final artifacts back to Drive.
• Checkpoint interval: saving every epoch is safe but slow; increase the interval (e.g., every 5–10 epochs) or save locally then shutil.copy to Drive in a background thread.
• Unattended runs: drive.mount requires interactive consent, so fully automated scheduled notebooks need a service‑account or Colab Pro+ background execution setup — outside the scope of this pattern.

Verify it yourself

1. Open a fresh notebook and run the mount cell above.
2. Create a test file: !echo 'hello' > /content/drive/MyDrive/test_persist.txt.
3. From the Runtime menu choose Disconnect and delete runtime.
4. Re‑open the notebook, remount Drive, and run !cat /content/drive/MyDrive/test_persist.txt. The file should still contain "hello".
5. Optional: time a representative checkpoint with %timeit -n 1 -r 3 torch.save(model.state_dict(), '/content/drive/MyDrive/checkpoints/bench.pt') and compare to /content/bench.pt to quantify the overhead.

Closing: add resume logic today

Ephemeral runtimes are a fact of Colab life, but a few lines of checkpoint code turn a catastrophic loss into a minute‑long restart. Mount Drive, pick a checkpoint directory, and wrap your training loop with the save/load helpers shown above. The next time Colab reclaims your VM, you’ll pick up exactly where you left off.