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

> Learn how AReaL implements checkpointing and recovery for long-running training jobs with a dual-path architecture. Prevent I/O blocking and ensure fault tolerance.

- Repository: [inclusionAI/areal](https://github.com/inclusionai/areal)
- Tags: how-to-guide
- Published: 2026-03-04

---

**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`](https://github.com/inclusionai/areal/blob/main/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.

```yaml
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`](https://github.com/inclusionai/areal/blob/main/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`](https://github.com/inclusionai/areal/blob/main/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:

```yaml
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:

```python
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:

```python
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:

```python

# 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`](https://github.com/inclusionai/areal/blob/main/areal/utils/saver.py)) exports model weights to HuggingFace format (`safetensors`, [`config.json`](https://github.com/inclusionai/areal/blob/main/config.json)) suitable for inference and evaluation, while the **RecoverHandler** ([`areal/utils/recover.py`](https://github.com/inclusionai/areal/blob/main/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`](https://github.com/inclusionai/areal/blob/main/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`](https://github.com/inclusionai/areal/blob/main/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.