gpu training checkpointing deep learning resume reclaimable instances cloud gpu
Training large machine learning models often requires significant computational resources, and cloud-based GPU instances offer a flexible solution. However, to make these resources more accessible and cost-effective, cloud providers sometimes offer 'spot' or 'reclaimable' GPU instances. These instances come at a lower price point but with the understanding that the provider can reclaim them with short notice. For long-running training jobs, this presents a significant challenge: how do you ensure your progress isn't lost when the instance unexpectedly shuts down? This tutorial guides you through building a fault-tolerant training workflow using checkpointing and resuming, making your machine learning jobs resilient to such interruptions.
Step 1: Understanding Reclaimable GPU Instances and Their Implications
Reclaimable GPU instances, such as DigitalOcean Spot GPU Droplets, are a common offering in cloud environments. They allow users to access powerful hardware like NVIDIA HGX B300 or AMD Instinct MI350X/MI355X GPUs at a reduced cost by utilizing the provider's idle capacity. The trade-off is that these instances can be terminated by the cloud provider when the capacity is needed elsewhere, often with a short notice period (e.g., two hours via email). When an instance is reclaimed, its local storage (boot disk, scratch disk) is typically lost.
For machine learning training, this means that any progress saved only on the instance's local disk, or any training state held only in memory, will be permanently lost upon termination. This can lead to wasted computation time, increased costs, and significant delays in your project. The core implication is that your training job must be designed to withstand arbitrary interruptions and resume seamlessly from a known good state.
To achieve this, the primary strategy involves:
Periodic Checkpointing: Regularly saving the entire state of your training job.
External Storage: Storing these checkpoints on a durable, highly available object storage service (e.g., DigitalOcean Spaces, Amazon S3, Google Cloud Storage).
Resume Logic: Implementing code that can detect and load the latest checkpoint upon startup, allowing training to continue from where it left off.
Graceful Shutdown Handling: Ensuring that your application can perform a final checkpoint and upload it when it receives a termination signal.
Step 2: Designing Your Checkpointing Strategy
A robust checkpointing strategy is the foundation of fault-tolerant training. It involves deciding what data to save, when to save it, and where to store it. The goal is to capture enough information to fully restore your training process to its exact state at the time of the checkpoint.
What to Checkpoint
To resume training accurately, you typically need to save:
Model Weights: The learned parameters of your neural network.
Optimizer State: The internal state of your optimizer (e.g., AdamW's momentum and variance estimates), which is crucial for correct gradient updates.
Learning Rate Scheduler State: If you're using a learning rate scheduler, its state ensures the learning rate continues to evolve correctly.
Current Epoch/Step Number: To know exactly where to restart the training loop.
Random States: The states of Python's, NumPy's, and your deep learning framework's (PyTorch, TensorFlow) random number generators. This ensures reproducibility if you need to debug or verify runs.
When to Checkpoint
The frequency of checkpointing involves a trade-off between recovery time and overhead. More frequent checkpoints mean less re-computation upon a reclaim but incur higher I/O and storage costs. A good starting point is:
Every N steps or epochs: For example, every 500-1000 training steps, or after every epoch.
Before significant operations: Such as before validation runs or saving intermediate results.
Upon receiving a termination signal: To ensure the very latest progress is saved.
Choosing External Storage
Since the instance's local disk is ephemeral, checkpoints must be stored externally. Object storage services are ideal for this due to their durability, scalability, and accessibility from any new instance. Services like DigitalOcean Spaces provide an S3-compatible API, making integration straightforward.
Here's a conceptual Python example using a deep learning framework (e.g., PyTorch) for saving a checkpoint:
import torch
import os
def save_checkpoint(model, optimizer, scheduler, epoch, step, path):
checkpoint_data = {
‘epoch’: epoch,
‘step’: step,
‘model_state_dict’: model.state_dict(),
‘optimizer_state_dict’: optimizer.state_dict(),
‘scheduler_state_dict’: scheduler.state_dict(),
‘rng_state’: torch.get_rng_state(),
‘cuda_rng_state’: torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None,
# Add other states like numpy, python random states if needed
}
torch.save(checkpoint_data, path)
print(f”Checkpoint saved to {path} at epoch {epoch}, step {step}“)
def load_checkpoint(model, optimizer, scheduler, path):
if not os.path.exists(path):
print(f”No checkpoint found at {path}. Starting training from scratch.”)
return 0, 0 # Return initial epoch and step
def signal_handler(sig, frame):
print(“Received termination signal (SIGTERM). Performing final checkpoint…”)
# These variables (model, optimizer, etc.) would be accessible from the main scope
# For demonstration, assume current_epoch, current_step are globally accessible
global model, optimizer, scheduler, current_epoch, current_step
The KillSignal=SIGTERM and TimeoutStopSec directives are crucial. systemd will send SIGTERM to your script and wait for TimeoutStopSec before forcibly terminating it with SIGKILL if it doesn’t exit. This gives your signal handler time to save the final checkpoint.
Step 5: Orchestrating the Workflow for Resilience
The true test of fault tolerance comes when an instance is reclaimed and a new one needs to take over. This step involves automating the setup of a new instance to seamlessly resume training from the latest checkpoint.
Using Cloud-Init for Automated Setup
Cloud providers offer mechanisms like cloud-init (or similar user data scripts) to execute commands automatically when a new instance boots up. This is the ideal place to orchestrate your resume workflow:
Fetch Code: Clone your training repository from a version control system (e.g., Git).
Configure Object Storage: Set up credentials for accessing your external object storage.
Download Latest Checkpoint: Implement a script that identifies and downloads the latest checkpoint from your object storage to the instance’s local scratch disk.
Start Training Service: Enable and start your systemd service (as defined in Step 4) which will then load the downloaded checkpoint and resume training.
The download_latest_checkpoint.py script would use the boto3 library (or equivalent for your object storage) to list objects in your checkpoint bucket, find the one with the highest timestamp or step number, and download it to a known location (e.g., /tmp/checkpoints/latest_checkpoint.pt).
The Resumption Flow
When a new reclaimable GPU instance is provisioned with this cloud-init script, the following sequence occurs:
The instance boots.
cloud-init executes, installing dependencies, cloning your repository, and downloading the latest checkpoint.
Your systemd service is enabled and started.
Your training script starts, detects the downloaded checkpoint, loads the model and optimizer states, and continues training from the exact step where the previous instance left off.
Step 6: Testing Your Fault-Tolerant Setup
The final and most crucial step is to thoroughly test your fault-tolerant workflow. This involves simulating a reclaim event and verifying that your new instance successfully picks up the training without losing progress.
Simulating a Reclaim
To simulate a reclaim, you can simply destroy your running GPU instance through your cloud provider’s control panel or API. This will trigger the termination process, including sending the SIGTERM signal to your running training job.
Verification Steps
Monitor Checkpoint Uploads: Ensure that your training script is regularly uploading checkpoints to your object storage. After manually destroying an instance, verify that a final checkpoint was uploaded shortly before its termination.
Provision a New Instance: Create a new reclaimable GPU instance using the same cloud-init script you’ve configured.
Monitor New Instance Startup: Observe the logs of the new instance (e.g., using journalctl -u my-training-job.service) to confirm that it successfully downloaded the latest checkpoint and started training from the correct epoch and step.
Track Training Progress: Compare the training metrics (loss, accuracy, etc.) from before the reclaim with those after resumption. They should show a continuous trend, indicating no loss of training progress.
Measure Recovery Time: Note the time it takes from destroying the old instance to the new instance actively resuming training. This helps you understand the operational overhead of a reclaim event.
By following these steps, you can build and verify a robust, fault-tolerant machine learning training workflow that leverages the cost benefits of reclaimable GPU instances without sacrificing the integrity or continuity of your valuable training jobs.
To explore more ways to build resilient and efficient online systems, consider visiting Yammbo.
