feat(sc): recover replay buffer from native TQ checkpoints

[macandro96]

What does this PR do ?

Summary

Adds authoritative recovery of completed, unconsumed Single Controller rollouts from native TransferQueue checkpoints.

TQ is the sole durable store for rollout tensor payloads. The SC checkpoint stores only a metadata replay index describing which TQ rows belong to ready prompt groups. On restart, TQ is loaded first and SC reconstructs its local replay-buffer indexes without copying tensors back into TQ.

This is a stacked change and currently depends on:

• chore: bump TQ to v0.19 #3423 — TransferQueue v0.1.9 bump and compatibility fixes
• feat(sc): add checkpoint save/restore to SingleController #3429 — Single Controller trainer/checkpoint lifecycle
• Native TQ checkpoint save/lifecycle hardening in amahishi/sc-tq-native-checkpoint

Motivation

Previously, SC checkpoint recovery either serialized rollout tensors separately in replay_buffer.pt or restarted with an empty rollout buffer.

Serializing tensors in both the replay-buffer checkpoint and TQ:

• duplicates potentially large payloads;
• creates two possible sources of truth;
• requires re-putting tensors into TQ during recovery;
• cannot guarantee that the replay index and TQ snapshot describe the same rows.

This change makes the native TQ snapshot authoritative for tensors and keeps only the controller metadata required to resume consuming those rows.

Checkpoint contents

step_N/
├── policy/
│   ├── weights/
│   └── optimizer/
├── train_dataloader.pt
├── training_info.json
├── data_plane/                 # Native TQ checkpoint; owns tensors
└── replay_buffer_metadata.pt   # Metadata-only SC replay index

For every completed prompt group, the replay sidecar stores:

• group ID;
• TQ partition and sample IDs;
• start/end policy weight versions;
• optional target step;
• sequence lengths, fields, tags, and auxiliary KVBatchMeta information.

It does not contain rollout tensors or fields_data.

Save workflow

flowchart LR
    Generation["Generation workers"] -->|"Completed prompt group"| Commit["TQReplayBuffer.commit"]
    Commit -->|"Tensor rows"| TQ["Canonical TQ partition"]
    Commit -->|"Ready-group metadata"| Index["SC replay index"]

    SC["Single Controller checkpoint"] --> Barrier["Exclusive checkpoint barrier"]
    Barrier --> TQSnapshot["Save native TQ snapshot"]
    Barrier --> Sidecar["Capture metadata-only replay index"]

    TQSnapshot --> Validate["Validate exact sample-ID inventory"]
    Sidecar --> Validate
    Validate --> Bundle["Finalize step_N bundle"]

Loading

Commits and destructive clears participate in a shared/exclusive checkpoint barrier:

• normal mutations may run concurrently;
• a checkpoint waits for active mutations to finish;
• new mutations wait while the snapshot is being captured;
• generation may continue, but completed groups wait at commit;
• the TQ snapshot and metadata sidecar therefore describe the same canonical rows.

Checkpoint finalization fails if:

• TQ save fails;
• the metadata sidecar names a missing TQ row;
• TQ contains an unexpected canonical row;
• serialization or bundle finalization fails.

The previously finalized checkpoint remains available as the fallback.

Restore workflow

flowchart LR
    Bundle["Latest finalized step_N"] --> Trainer["Restore model, optimizer, dataloader"]
    Bundle --> Bootstrap["Clean TQ bootstrap client"]
    Bootstrap -->|"Load data_plane/"| TQ["Restored native TQ state"]
    TQ --> Client["Create normal SC data-plane client"]
    Bundle -->|"Load replay metadata"| Index["Rebuild local replay index"]
    Client --> Validate["Validate digest, group count, and exact inventory"]
    Index --> Validate
    Validate --> Pumps["Start rollout and train pumps"]

Loading

Restore happens before SingleControllerActor.run() and before any ordinary data-plane operation.

Recovery validates:

• checkpoint schema versions;
• trainer step, trainer weight version, and epoch;
• sampler name and partition ID;
• replay manifest digest;
• replay group count;
• configured buffer capacity;
• exact equality between sidecar sample IDs and canonical TQ sample IDs.

No rollout tensor is re-serialized or re-put into TQ.

Sampler support

Authoritative replay recovery currently supports the ungated windowed sampler.

windowed
    restored in-window groups → selectable
    restored stale groups     → evicted by the normal sampler path

weight_fifo / in_order
    native TQ snapshot        → shadow mode only
    replay recovery           → not enabled

Gated samplers need additional durable dispatch/quota state before completed subsets can be recovered safely.

Sampler configs now declare their replay-checkpoint capability, and the sampler factory verifies that the declaration matches the runtime implementation. Custom sampler capability is checked after construction.

Configuration

checkpointing:
  enabled: true

data_plane:
  enabled: true
  impl: transfer_queue
  backend: simple
  checkpointing_enabled: true

async_rl:
  sampler:
    name: windowed
    max_staleness_versions: 1

Current native restore support requires the TQ simple backend. Unsupported configurations fail during setup rather than silently producing incomplete checkpoints.

Compatibility

• Legacy tensor-bearing replay_buffer.pt checkpoints are detected and rejected with an actionable error.
• Checkpoints without any replay artifact continue with an empty replay buffer.
• Gated samplers retain shadow-mode native TQ saves but do not advertise recoverable rollout state.
• Trainer checkpoints remain immutable; TQ and replay metadata are finalized inside the same step bundle.

Testing

• Native TQ fresh-process save/load verifier
• Native TQ parent-directory rename round trip
• Focused checkpoint, replay-buffer, data-plane lifecycle, and interface unit tests
• bash tests/functional/grpo_dp_single_controller_tq_recovery.sh
  • phase 1 trains one step and saves an authoritative TQ checkpoint;
  • phase 2 starts a fresh process;
  • native TQ state and replay metadata are restored;
  • exact inventory is validated;
  • training continues to step 2.
• bash tests/functional/grpo_dp_single_controller.sh
  • existing SC convergence regression test

Scope and limitations

This PR recovers completed, unconsumed prompt groups.

It does not yet recover:

• token-level unfinished generations;
• in-flight reservations that have not committed canonical TQ rows;
• NeMo-Gym episode/environment state;
• sandbox memory or filesystem state;
• gated-sampler dispatch ledgers.

The planned #3456 integration remains upstream of this mechanism:

Staged token records
        ↓
Receipt/finalizer recovery
        ↓
Canonical completed rollout in TQ
        ↓
This PR: replay-index restoration and trainer consumption

Review guide

The primary areas to review are:

1. TQReplayBuffer metadata-only serialization and restore validation.
2. Shared/exclusive checkpoint-barrier coverage for commits and clears.
3. Native TQ restore ordering in SC setup.
4. Fail-loud schema, digest, inventory, and configuration checks.
5. Windowed sampler recovery semantics and gated-sampler exclusions.
6. Two-process functional recovery coverage.

Issues

List issues that this PR closes (syntax):

Usage

• You can potentially add a usage example below

# Add a code snippet demonstrating how to use this

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

• ...

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.