Checkpointing and Resuming Training

Full parameter fine-tuning often involves training for extended periods, potentially hours, days, or even weeks, consuming significant computational resources. Interruptions during these long runs are almost inevitable, whether due to hardware failures, cluster job preemption, network issues, or simply the need to pause and restart. Without a mechanism to save and restore progress, such interruptions could lead to a complete loss of the training performed up to that point, wasting valuable time and compute budget. Checkpointing provides this essential safety net.

Effective checkpointing involves periodically saving the complete state of the training process to persistent storage. This allows you to resume training precisely from where it left off, minimizing wasted effort.

What Constitutes a Training State?

To resume training, you need to save more than just the model's parameters ($\theta$). A comprehensive checkpoint should capture:

1. Model State Dictionary: This contains all the learnable parameters (weights and biases) of your model. In frameworks like PyTorch, this is typically accessed via model.state_dict(). Saving this ensures you retain the learned knowledge.
2. Optimizer State Dictionary: Modern optimizers like Adam or AdamW maintain internal states, such as moving averages of gradients and squared gradients (moments). Simply restoring model weights and initializing a new optimizer would reset these states, potentially disrupting the learning dynamics and convergence behavior achieved so far. Saving optimizer.state_dict() is necessary to preserve this momentum.
3. Learning Rate Scheduler State: If you are using a learning rate scheduler (e.g., linear warmup followed by decay), its internal state (like the current learning rate multiplier, step count, or last epoch) must be saved using scheduler.state_dict(). Failing to do so would reset the schedule, leading to incorrect learning rates upon resumption.
4. Training Progress Indicators: Essential metadata like the current epoch number, the total number of training steps completed, or the number of samples processed. This allows the training loop to correctly determine where to restart data loading and epoch counting.
5. Random Number Generator (RNG) States: For strict reproducibility, especially if your data loading, shuffling, or augmentations involve randomness, saving the state of RNGs (e.g., torch.get_rng_state(), numpy.random.get_state(), random.getstate()) can be beneficial. Restoring these ensures the data pipeline behaves identically after resumption.
6. (Optional) Loss History and Metrics: Saving performance metrics tracked during training (e.g., training loss, validation accuracy) up to the checkpoint point can be useful for monitoring and analysis, though not strictly required for resumption.
7. (Optional) Training Configuration: Saving the configuration arguments or script parameters used for the run helps ensure consistency if resuming in a slightly different environment or after code changes.

Implementing Checkpoint Saving

Checkpointing logic is typically integrated directly into the training loop. You need to decide on a checkpointing frequency. Common strategies include:

• Saving every N steps: Provides fine-grained recovery points but can lead to higher I/O overhead if N is small.
• Saving every epoch: Simpler to manage but can mean losing up to an entire epoch's work if an interruption occurs late in the epoch.
• Saving based on time intervals: E.g., every hour.
• Saving the best performing checkpoint: Based on validation metrics, ensuring you always retain the model state that performed best so far. This is often combined with saving the latest checkpoint.

To manage storage, you might keep only the latest checkpoint, the best checkpoint, or a rolling window of the last few checkpoints.

Here's a PyTorch-like example of saving a checkpoint:

# Assume model, optimizer, scheduler, epoch, global_step are defined

checkpoint_path = f"./model_checkpoint_epoch_{epoch}_step_{global_step}.pt"
save_state = {
    'epoch': epoch,
    'global_step': global_step,
    'model_state_dict': model.state_dict(), # Or model.module.state_dict() if using DDP
    'optimizer_state_dict': optimizer.state_dict(),
    'scheduler_state_dict': scheduler.state_dict(),
    # Optionally add RNG states, config, loss history etc.
    # 'rng_state': torch.get_rng_state(),
    # 'config': training_args,
}

# Best practice: save to temporary file then rename for atomicity
temp_path = checkpoint_path + ".tmp"
torch.save(save_state, temp_path)
os.rename(temp_path, checkpoint_path) # Atomic rename

print(f"Checkpoint saved to {checkpoint_path}")

# Optionally, manage older checkpoints (e.g., keep only last 3)
# manage_checkpoints(checkpoint_dir, keep_last=3)

Using a temporary file and renaming ensures that if the saving process is interrupted, you don't end up with a corrupted partial checkpoint file.

Resuming from a Checkpoint

When starting a training run, your script should check if a valid checkpoint exists. If so, it should load the saved state before commencing training.

Here's how you might load the state:

# Assume model, optimizer, scheduler are initialized
# Determine the checkpoint path to load from (e.g., latest one)
checkpoint_path = find_latest_checkpoint("./") # Function to find the checkpoint file

if checkpoint_path:
    print(f"Resuming training from checkpoint: {checkpoint_path}")
    checkpoint = torch.load(checkpoint_path, map_location='cpu') # Load to CPU first

    model.load_state_dict(checkpoint['model_state_dict'])
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    scheduler.load_state_dict(checkpoint['scheduler_state_dict'])

    start_epoch = checkpoint['epoch']
    global_step = checkpoint['global_step']
    # Optionally restore RNG states, etc.
    # torch.set_rng_state(checkpoint['rng_state'])

    # Important: Move optimizer state to the correct device(s) if using GPU
    # This step might be needed depending on how the optimizer state was saved
    # and if you loaded the checkpoint to CPU initially.
    for state in optimizer.state.values():
        for k, v in state.items():
            if torch.is_tensor(v):
                state[k] = v.to(device) # device could be 'cuda:0'

    print(f"Resumed from epoch {start_epoch}, global step {global_step}")
else:
    print("No checkpoint found, starting training from scratch.")
    start_epoch = 0
    global_step = 0

# --- Start training loop from start_epoch, tracking global_step ---

Loading the optimizer and scheduler states is essential for maintaining the training dynamics. Loading the checkpoint to the CPU first (map_location='cpu') before loading into the model/optimizer can prevent GPU memory issues if the checkpoint is large. If using GPUs, ensure the optimizer state tensors are moved to the correct device after loading.

Approaches to Distributed Training

When using distributed training frameworks like PyTorch's Distributed Data Parallel (DDP), checkpointing requires slight modifications:

• Saving: Typically, only the main process (rank 0) should save the checkpoint to avoid redundant writes and potential race conditions. The model state dictionary might need to be accessed differently (e.g., model.module.state_dict() for DDP).
• Loading: All processes need to load the model weights to ensure consistency. However, metadata like epoch and step count might only be needed by the main process. Use synchronization primitives (like torch.distributed.barrier()) to ensure all processes have loaded the state before proceeding with training.

Best Practices

• Frequency: Choose a frequency that balances recovery capability with I/O overhead. For very large models and long training times, saving every few hundred or thousand steps might be appropriate.
• Atomicity: Always save to a temporary file and then rename to prevent corruption.
• Storage Location: Use reliable storage. Network file systems (NFS) or cloud storage buckets (S3, GCS, Azure Blob Storage) are often preferred over local disk, especially in cluster environments where nodes can be ephemeral.
• Verification: Consider adding simple checks after saving (e.g., confirming file size or attempting a quick reload) if corruption is a major concern.
• Code Versioning: Be mindful that significant changes in model architecture or library versions between saving and loading a checkpoint can cause compatibility issues. Store relevant version information within the checkpoint if possible.

Training loop incorporating checkpoint loading at the start and periodic saving during the process. Interruptions trigger a restart, which then loads the most recent checkpoint.

Mastering checkpointing and resumption is not just a convenience but a necessity for reliably executing the resource-intensive process of full parameter fine-tuning on large language models. It ensures that progress is preserved against interruptions, making long training runs feasible and manageable.