# Checkpointing

This document describes AReaL's checkpointing system, which handles model saving for
evaluation and fault-tolerant recovery during distributed RL training.

## Overview

AReaL provides two complementary checkpointing mechanisms:

| Mechanism          | Purpose                                    | Format                       | Includes Optimizer/DataLoader State |
| ------------------ | ------------------------------------------ | ---------------------------- | ----------------------------------- |
| **Saver**          | Export models for evaluation or publishing | HuggingFace                  | No                                  |
| **RecoverHandler** | Resume training after failures             | DCP (Distributed Checkpoint) | Yes                                 |

Both mechanisms are invoked automatically during training and can be configured via
`config.saver` and `config.recover` respectively.

## Checkpoint Formats

### HuggingFace Format

Used by `Saver` for model export:

- Standard HuggingFace model format (safetensors + config.json)
- Compatible with `transformers.AutoModel.from_pretrained()`
- Can be uploaded to HuggingFace Hub
- Does not include optimizer state

### DCP Format (Distributed Checkpoint)

Used by `RecoverHandler` for fault tolerance:

- Backend's native distributed checkpoint format (`torch.distributed.checkpoint` or
  Megatron distributed checkpoint)
- Sharded across all ranks for efficient parallel I/O
- Includes model weights, optimizer state, RNG state, etc
- Backend-specific: checkpoints are only compatible with the same parallelism
  configuration
- In single-controller mode, writes an immutable generation and atomically publishes it
  through a `LATEST` pointer
- Retains the legacy overwrite layout in SPMD mode

## Architecture

```
PPOTrainer.train()
│
├── Training loop
│   ├── Rollout, compute values, PPO update...
│   │
│   ├── _save_hf()                          # HuggingFace export
│   │   └── Saver.save()
│   │       └── engine.save(weight_format="hf")
│   │
│   └── _save_recover_checkpoint()          # Fault tolerance
│       └── RecoverHandler.dump()
│           └── engine.save(weight_format="dcp", with_optim=True)
│
└── On restart
    └── RecoverHandler.load()
        ├── Restore dataloader, saver, evaluator states
        └── engine.load(weight_format="dcp", with_optim=True)
```

## Saver: HuggingFace Model Export

The [`Saver`](https://github.com/areal-project/AReaL/blob/main/areal/utils/saver.py)
periodically exports model weights in HuggingFace format for evaluation or deployment.

### Save Mode

The `mode` parameter controls how checkpoints are written:

| Mode    | Behavior                                                                                                                                                                                                                              |
| ------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
| `auto`  | Use async for Archon engine, sync for others (default). Zero-config optimal for all engines.                                                                                                                                          |
| `sync`  | Always synchronous `dcp.save()`.                                                                                                                                                                                                      |
| `async` | Always process-based async with pinned-memory staging. Archon engine only; other engines fall back to sync with a warning. Extra CPU pinned memory proportional to per-rank model shard size (e.g., ~17.5 GB/rank for 70B on 8 GPUs). |

With the default `auto` mode, Archon engine users get async checkpoint saving
automatically - the training loop blocks only while the checkpoint is staged to pinned
CPU memory, and the actual disk I/O happens in a background process.

### Configuration

Configure via `config.saver`:

| Parameter     | Type        | Default  | Description                          |
| ------------- | ----------- | -------- | ------------------------------------ |
| `mode`        | str         | `"auto"` | Save mode (see above).               |
| `freq_epochs` | int \| None | None     | Save every N epochs. None disables.  |
| `freq_steps`  | int \| None | None     | Save every N steps. None disables.   |
| `freq_secs`   | int \| None | None     | Save every N seconds. None disables. |

Example configuration:

```yaml
saver:
  freq_epochs: 1      # Save at end of each epoch
  freq_steps: null    # Disabled
  freq_secs: null     # Disabled
  # mode defaults to "auto" - Archon users get async automatically
```

Saving is triggered when any of epoch/step/time condition is met.

### Output Location

Checkpoints are saved to:

```
{fileroot}/checkpoints/{user}/{experiment_name}/{trial_name}/default/
└── epoch{E}epochstep{S}globalstep{G}/
    ├── config.json
    ├── model.safetensors (or model-00001-of-00002.safetensors, etc.)
    ├── tokenizer.json
    └── ...
```

### Usage

Load saved checkpoints with standard HuggingFace APIs:

```python
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained(
    "/path/to/checkpoint/epoch0epochstep99globalstep99"
)
tokenizer = AutoTokenizer.from_pretrained(
    "/path/to/checkpoint/epoch0epochstep99globalstep99"
)
```

## RecoverHandler: Fault Tolerance

The
[`RecoverHandler`](https://github.com/areal-project/AReaL/blob/main/areal/utils/recover.py)
enables resuming training after failures by saving complete training state.

### Configuration

Configure via `config.recover`:

| Parameter     | Type        | Default    | Description                                      |
| ------------- | ----------- | ---------- | ------------------------------------------------ |
| `mode`        | str         | "disabled" | Recovery mode: "on"/"auto" or "off"/"disabled"   |
| `freq_epochs` | int \| None | None       | Checkpoint every N epochs                        |
| `freq_steps`  | int \| None | None       | Checkpoint every N steps                         |
| `freq_secs`   | int \| None | None       | Checkpoint every N seconds                       |
| `retries`     | int         | 3          | Number of recovery retries when recovery enabled |

#### Recovery Modes

| Mode                | Behavior                                        |
| ------------------- | ----------------------------------------------- |
| `on` or `auto`      | Automatically resume if valid checkpoint exists |
| `off` or `disabled` | No checkpointing or recovery                    |

When recovery is enabled (`on`/`auto`), the system will:

1. Periodically save recovery checkpoints (model weights, optimizer state, dataloader
   position)
1. Automatically resume from the last valid checkpoint on restart
1. Retry up to `retries` times on failure

Example configuration:

```yaml
recover:
  mode: on            # or "auto" for backward compatibility
  freq_steps: 100     # Checkpoint every 100 steps
  retries: 3
```

### What Gets Saved

RecoverHandler saves complete training state:

| Component         | Contents                                           |
| ----------------- | -------------------------------------------------- |
| Model weights     | DCP format, sharded across ranks                   |
| Optimizer state   | Momentum, variance (Adam), learning rate scheduler |
| RNG state         | Python, NumPy, PyTorch, CUDA random states         |
| Dataloader state  | Current position in dataset                        |
| Training progress | Epoch, step, global_step counters                  |
| Auxiliary states  | Saver, Evaluator, StatsLogger states               |

### Output Location

In single-controller mode, recovery checkpoints are saved as immutable generations:

```
{fileroot}/checkpoints/{user}/{experiment_name}/{trial_name}/
├── LATEST                      # Atomic JSON pointer to one complete generation
└── checkpoint_generations/
    └── generation_step00000123/
        ├── manifest/           # Dataloader and training progress metadata
        └── payloads/
            ├── default/        # Model + optimizer (DCP format)
            └── critic/         # Present when a critic is checkpointed
```

The pointer is updated only after every payload is durable. When a generation contains
multiple train engines, `RecoverHandler` drains all but one asynchronous save before
scheduling the final asynchronous engine. Publication is appended only to that engine's
MCore finalize callback, so one background save can still overlap training without
exposing a partially written actor/critic generation. A crash may leave an unpublished
generation on disk, but recovery continues to use the generation selected by `LATEST`.
The previous published generation is removed after the pointer switches successfully.

SPMD trainers keep the pre-pointer `recover_info/` and `<engine>/recover_checkpoint/`
layout for compatibility. Recovery can read both layouts.

### Recovery Process

When training resumes:

1. `RecoverHandler.load()` restores all saved state (if any)
1. Training continues from `last_step_info.next().global_step`
1. Inference engine weights are synchronized to match recovered state

## Best Practices

### Frequency Guidelines

| Scenario           | Recommended Setting                       |
| ------------------ | ----------------------------------------- |
| Long training runs | `freq_epochs: 1` or `freq_steps: 1000`    |
| Unpredictable time | `freq_secs: 7200`                         |
| Unstable clusters  | `freq_steps: 100` with `recover.mode: on` |
| Limited disk space | Lower frequency, rely on final checkpoint |
| Debugging          | `freq_steps: 1` for quick iteration       |

### Disk Space Considerations

- **Saver**: Each save creates a new directory. High frequency consumes significant
  space.
- **RecoverHandler**: Normally keeps one published generation plus any in-progress or
  crash-abandoned generation. SPMD mode continues to overwrite the legacy paths.

### Recovery Tips

1. **Verify checkpoint validity**: Read `LATEST`, then inspect the selected generation's
   `manifest/step_info.json`
1. **Same config required**: DCP checkpoints require identical parallelism
   configuration, experiment name, and trial name
1. **One writer per trial**: Do not run concurrent checkpoint writers against the same
   experiment and trial directory
1. **Clean restart**: Delete both `LATEST` and `checkpoint_generations/`. For a legacy
   SPMD checkpoint, delete `recover_info/` and each engine's `recover_checkpoint/`