Implementing Checkpointing and Recovery for Long-Running Training Jobs in AReaL

AReaL implements a dual-path checkpointing architecture that separates lightweight HuggingFace exports for model evaluation from full Distributed Checkpoint (DCP) snapshots for fault-tolerant recovery, using frequency-controlled triggers that run asynchronously for the Archon engine to prevent I/O blocking during training.

The AReaL repository solves the operational challenge of implementing checkpointing and recovery for long-running training jobs by decoupling model serialization from training-state persistence. This design uses distinct Saver and RecoverHandler components governed by EpochStepTimeFreqController policies, configured to trigger on step counts, epochs, or wall-clock intervals.

Architecture Overview

AReaL’s checkpointing system operates through two separate paths that serve different operational purposes:

Component	Purpose	Format	Included State
Saver	Periodic HF export for evaluation/deployment	HuggingFace (safetensors, config.json)	Model weights only
RecoverHandler	Full training-state snapshots for crash recovery	DCP (sharded)	Model + optimizer + RNG + dataloader

Both components are driven by a frequency controller (EpochStepTimeFreqController) that supports triggers on epochs, steps, or elapsed seconds. The training loop (PPOTrainer.train()) invokes the saver and recovery handler at configured intervals:

PPOTrainer.train()
└─ training loop
   ├─ _save_hf() → Saver.save() → engine.save(weight_format="hf")
   └─ _save_recover_checkpoint() → RecoverHandler.dump()

The Saver decides between synchronous versus asynchronous persistence based on config.saver.mode (auto, sync, or async). Asynchronous mode is enabled only for the Archon engine, where heavy I/O runs in a background process using pinned memory staging via torch.distributed.checkpoint.

Configuring Checkpoint Frequency

Checkpoint triggers are managed through YAML configuration under the saver and recover blocks. The EpochStepTimeFreqController interprets freq_epochs, freq_steps, or freq_secs to determine when _save_hf() and _save_recover_checkpoint() are invoked.

saver:
  mode: auto          # auto → async for Archon, otherwise sync

  freq_epochs: 1      # save HF checkpoint each epoch

  freq_steps: null
  freq_secs: null

recover:
  mode: on            # enable fault-tolerant recovery

  freq_steps: 100     # dump recover checkpoint every 100 steps

  retries: 3

The Saver Module – Export and Async Persistence

Core logic lives in areal/utils/saver.py, which wraps HuggingFace exports and Distributed Checkpoint APIs.

Synchronous versus Asynchronous Modes

The _should_use_async() method (lines 99-114) determines engine compatibility, returning True only for the Archon engine when config.saver.mode is set to async. In auto mode, the saver probes the engine type to decide whether to spawn background processes.

• Sync path: Calls dcp.save() directly, blocking the training loop until GPU-to-CPU serialization and disk writes complete.
• Async path: Uses AsyncCheckpointManager to offload staging and consolidation, allowing the training loop to continue while I/O progresses in a separate process.

Staging and Consolidation Internals

When _async_save() is triggered (lines 55-74), it instantiates an AsyncCheckpointManager from areal/experimental/engine/archon_checkpoint and hands off the DCP async API. The implementation uses a DefaultStager that stages tensors into pinned or shared memory buffers (lines 96-124) to optimize PCIe transfers.

Background consolidation runs on a single-thread ThreadPoolExecutor. Before the optimizer step, callers must invoke maybe_wait_for_staging() (lines 140-149) to ensure the GPU-to-CPU copy has finished and it is safe to mutate model parameters. When training concludes, finalize() (lines 182-209) shuts down the executor, closes the stager, and destroys the process groups to prevent resource leaks.

The RecoverHandler – Full State Snapshots

Fault-tolerant recovery is implemented in areal/utils/recover.py, which serializes the entire training context beyond model weights.

Distributed State Serialization

RecoverInfo.dump() (lines 46-58) writes a JSON bundle containing the current step, saver configuration, evaluator state, stats-logger metadata, and pickled dataloader state. In distributed runs, only rank 0 writes the files after gathering dataloader information from all ranks via collective communication.

Resuming Training from a Crash

RecoverInfo.load() (lines 84-127) reverses the serialization process, reconstructing the dataloader state across all ranks. RecoverHandler.load() then constructs a SaveLoadMeta object with weight_format="dcp" and with_optim=True, calling engine.save() to restore the optimizer and RNG state (lines 252-263). If an inference engine is present, the handler updates its weights to the restored version before training resumes.

Training Loop Integration

Within PPOTrainer.train(), the checkpointing flow follows this sequence:

1. Every N steps/epochs/seconds → Saver.save()
   • Sync → dcp.save() (blocking)
   • Async → dcp.async_save() + background consolidation
2. Every M steps/epochs/seconds → RecoverHandler.dump()
   • Serialize RecoverInfo (step, saver, evaluator, etc.)
   • Save DCP checkpoint (model + optimizer + state)
3. On crash → New process calls RecoverHandler.load()
   • Load RecoverInfo → restore counters, RNG, dataloader
   • Load DCP checkpoint → engine.load(...)

Implementation Examples

YAML Configuration

Define checkpointing policies in your experiment file:

saver:
  mode: auto
  experiment_name: "my_exp"
  trial_name: "run1"
  file_root: "/tmp/areal"
  freq_epochs: null
  freq_steps: null
  freq_secs: null

recover:
  mode: auto
  experiment_name: "my_exp"
  trial_name: "run1"
  file_root: "/tmp/areal"
  freq_epochs: null
  freq_steps: 1000
  freq_secs: null

Manual Checkpoint Trigger

Programmatically export weights or force a recovery dump:

from areal.utils.saver import Saver
from areal.api.cli_args import SaverConfig
from areal.api.io_struct import FinetuneSpec

# Create config objects (normally read from CLI/YAML)

saver_cfg = SaverConfig(
    experiment_name="my_exp",
    trial_name="run1",
    file_root="/tmp/areal",
    mode="async",          # force async

    freq_epochs=None,
    freq_steps=None,
    freq_secs=None,
)

ft_spec = FinetuneSpec(steps_per_epoch=1000, max_steps=10000)

saver = Saver(saver_cfg, ft_spec)

# engine is an instantiated ArchonEngine or similar

saver.save(engine, epoch=3, step=999, global_step=3999)

Recovery Workflow

Handle resumption logic before starting training:

from areal.utils.recover import RecoverHandler, check_if_recover
from areal.api.cli_args import RecoverConfig
from areal.api.io_struct import FinetuneSpec

recover_cfg = RecoverConfig(
    experiment_name="my_exp",
    trial_name="run1",
    file_root="/tmp/areal",
    mode="auto",
    freq_steps=1000,
    retries=3,
)

ft_spec = FinetuneSpec(steps_per_epoch=1000, max_steps=10000)

if check_if_recover(recover_cfg, run_id=0):
    recover = RecoverHandler(recover_cfg, ft_spec)
    # Restores saver, evaluator, dataloader, and engine state

    recover.load(engine, saver, evaluator, dataloader)
else:
    # Start fresh training run

    pass

Async Staging Synchronization

When using asynchronous checkpointing, ensure staging completes before optimizer updates:

# Inside the training loop, before optimizer.step()

saver.maybe_wait_for_staging()   # ensures GPU→CPU copy finished

optimizer.step()

Summary

• Dual-path design: Saver handles lightweight HF exports for evaluation, while RecoverHandler manages full DCP snapshots for crash recovery.
• Async I/O: The Archon engine uses AsyncCheckpointManager with pinned memory staging to eliminate checkpoint-related blocking, coordinated through maybe_wait_for_staging() and finalize().
• Complete state recovery: RecoverInfo serialization captures optimizer states, RNG seeds, and distributed dataloader progress, enabling seamless resumption from hardware failures.
• Flexible triggering: EpochStepTimeFreqController supports arbitrary cadences based on epochs, steps, or wall-clock time, configured via YAML or Python APIs.

Frequently Asked Questions

What is the difference between Saver and RecoverHandler in AReaL?

The Saver (areal/utils/saver.py) exports model weights to HuggingFace format (safetensors, config.json) suitable for inference and evaluation, while the RecoverHandler (areal/utils/recover.py) creates full Distributed Checkpoint (DCP) snapshots containing optimizer states, RNG seeds, and dataloader progress for fault-tolerant training recovery. The Saver operates on SaveLoadMeta with weight_format="hf", whereas RecoverHandler uses weight_format="dcp" with with_optim=True to capture the complete training context.

How does AReaL handle asynchronous checkpointing for the Archon engine?

AReaL enables asynchronous checkpointing exclusively for the Archon engine through Saver._should_use_async() (lines 99-114), which validates engine compatibility. When active, _async_save() (lines 55-74) instantiates an AsyncCheckpointManager that stages tensors to CPU memory via a DefaultStager using pinned buffers (lines 96-124). Consolidation runs in a background ThreadPoolExecutor, allowing the training loop to proceed without blocking. Callers must invoke maybe_wait_for_staging() (lines 140-149) before parameter updates to ensure GPU-to-CPU transfers complete.

What state is included in a RecoverHandler dump versus a Saver export?

A RecoverHandler dump includes the model weights, optimizer states, RNG seeds, current training step, evaluator configuration, stats-logger metadata, and pickled dataloader state—enabling exact resumption of distributed training. A Saver export contains only model weights in HuggingFace format (safetensors + config.json), optimized for evaluation and deployment but lacking optimizer or training loop state. RecoverHandler writes JSON metadata via RecoverInfo.dump() (lines 46-58) and sharded DCP data, while Saver produces consolidated HF artifacts.

How does the training loop know whether to resume from a checkpoint or start fresh?

The training loop calls check_if_recover() (lines 26-83 in areal/utils/recover.py) to scan the checkpoint directory for existing RecoverInfo metadata. If valid recovery files are detected, the loop instantiates RecoverHandler and executes load() to reconstruct the SaveLoadMeta, restore engine weights, and reinitialize the dataloader from its pickled state. If no valid checkpoints are found, the loop proceeds with fresh initialization. This logic is typically invoked immediately after process spawn in distributed environments.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}