Checkpointing

The primary purpose of checkpointing is to provide fault tolerance for long-running, computationally expensive jobs. By periodically saving the application state—including model parameters, optimizer state, and random number generator seeds—to persistent storage, the system can recover from hardware failures, node outages, or preemptions (e.g., with cloud spot instances) without losing all progress. Recovery involves restarting the computation from the last valid checkpoint, significantly reducing wasted compute hours and cost.

Checkpointing requires the serialization of the in-memory state of an application into a format that can be written to disk. For machine learning, this typically includes:

• Model Weights: The parameters of the neural network.
• Optimizer State: Momentum buffers, variance accumulators (e.g., for Adam), and other optimizer-specific variables.
• Training Loop Metadata: The current epoch, iteration number, learning rate schedule step, and random number generator state.
• Simulation State: In physics-based training, the full state of the simulated world (object positions, velocities). Formats like PyTorch's .pt or .pth and TensorFlow's SavedModel are standard for this serialization.

The checkpoint interval is a critical trade-off between overhead and potential data loss. Frequent checkpoints minimize rework but increase I/O overhead and storage costs. Key strategies include:

• Time-based: Save every N minutes/hours.
• Iteration-based: Save every N training steps or epochs.
• Validation-based: Save when a validation metric improves (model checkpointing).
• Coordinated vs. Uncoordinated: In distributed training, coordinated checkpointing halts all processes to capture a globally consistent state, while asynchronous methods are faster but more complex to restore.

In distributed training across hundreds of GPUs, checkpointing must handle data parallelism (model replicas) and model parallelism (split models). Techniques include:

• Sharded Checkpointing: Each parallel process writes its portion of the model state to a separate file, reducing I/O burden on a single node. Frameworks like Microsoft's DeepSpeed use this.
• Global Consensus: Ensuring a consistent snapshot across all nodes, often requiring barrier synchronization.
• Leveraging Parallel File Systems: Checkpoints are written to high-bandwidth, shared storage like Lustre or GPFS that all nodes can access for fast write/read.

In parallelized robotic simulation, checkpointing is vital for reinforcement learning (RL) training that may run for days or weeks. It allows:

• Resuming Policy Training: After a system failure in the simulation cluster.
• Branching Experiments: Loading a checkpoint to start new training runs with different hyperparameters from a known-good state.
• Rollback for Safety: If a policy's behavior degrades during training, reverting to a prior, more stable checkpoint.
• Hardware-in-the-Loop (HITL): Saving the state of a simulation that is interacting with physical hardware for later analysis and replay.

Managing checkpoint storage is a key operational concern. Strategies include:

• Hierarchical Storage: Keeping recent checkpoints on fast NVMe storage and archiving older ones to slower, cheaper object storage (e.g., Amazon S3).
• Checkpoint Pruning: Automatically deleting all but the best-performing (by validation metric) and most recent checkpoints to conserve space.
• Incremental Checkpointing: Only saving the difference from the previous checkpoint to reduce I/O volume.
• Metadata Tagging: Associating checkpoints with experiment IDs, git commits, and hyperparameters for full reproducibility.

State serialization is the process of converting a complex, in-memory object or system state into a format (a byte stream) that can be stored (e.g., on disk) or transmitted and later reconstructed. It is the foundational technology enabling checkpointing.

• Key Objects Serialized: Includes the model's weights/parameters, optimizer state (momentum, variances), random number generator seeds, and the complete environment state in reinforcement learning.
• Formats: Common formats include PyTorch's .pt or .pth, TensorFlow's .ckpt, and protocol buffers (.pb).
• Challenges: Must handle custom C++ extensions, open file handles, and GPU memory tensors, often requiring specialized frameworks like PyTorch's torch.save().

Distributed training is the practice of parallelizing machine learning model training across multiple processors (GPUs/TPUs) or nodes. Checkpointing here must correctly synchronize and save the sharded state of the model and data across the entire system.

• Synchronous vs. Asynchronous: Synchronous training (e.g., using Horovod, PyTorch DDP) requires a global barrier during checkpointing to ensure a consistent snapshot. Asynchronous training poses greater consistency challenges.
• Model Parallelism: When a single model is split across devices, the checkpoint must reassemble the sharded layers correctly upon restore.
• Framework Support: PyTorch Fully Sharded Data Parallel (FSDP) and TensorFlow's tf.distribute have built-in, optimized checkpointing for their specific distribution strategies.