fix(grpo): ReplayBuffer checkpointing and fix reservation leaks (#2651)

Signed-off-by: Anish Mahishi <amahishi@cw-dfw-cs-001-vscode-02.cm.cluster>
Signed-off-by: Anish Mahishi <amahishi@cw-dfw-cs-001-vscode-01.cm.cluster>
Signed-off-by: Anish Mahishi <20884035+macandro96@users.noreply.github.com>
Co-authored-by: Anish Mahishi <amahishi@cw-dfw-cs-001-vscode-02.cm.cluster>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: Anish Mahishi <amahishi@cw-dfw-cs-001-vscode-01.cm.cluster>
Co-authored-by: Terry Kong <terryk@nvidia.com>

Author: 5 people

docs/guides/async-grpo.md

@@ -151,6 +151,31 @@ sequenceDiagram
     end
 ```
 
+## Checkpointing
+
+Async GRPO checkpoints the replay buffer alongside the rest of training state so that in-progress trajectory generation is not lost across restarts.
+
+### What is saved
+
+On each checkpoint, a `replay_buffer.pt` file is written next to the other checkpoint artifacts. It contains all trajectories currently in the buffer together with their weight and target versions, and the `last_target_weight_already_generated` watermark.
+
+### Restore behaviour
+
+On resume, the buffer is restored before the trajectory collector starts, then cleaned up as follows:
+
+1. **Past targets dropped** — trajectories whose target step is earlier than the resume step are removed.
+2. **Stale trajectories evicted** — if `max_trajectory_age_steps` is set, trajectories too old for their target step are removed.
+3. **Incomplete targets kept** — target steps that still lack a full batch are kept in the buffer. The collector will *gap-fill* only the missing trajectories for those targets before moving on.
+4. **Buffer truncated** — if the restored count exceeds `max_size`, the buffer is truncated, prioritising entries closest to the resume step.
+
+### Gap-filling after restore
+
+After a restore, `last_target_weight_already_generated` is reset to `current_training_step - 1` so the collector re-evaluates every target from the resume step onward. For each target it queries `get_trajectories_needed` and spawns only the workers required to complete the batch — previously buffered trajectories are reused and the collector does not regenerate them.
+
+### Disabling replay-buffer restore
+
+If no `replay_buffer.pt` file is found in the latest checkpoint directory, training starts with an empty buffer and waits for the collector to fill it before the first training step.
+
 ## Usage Tips
 
 1. **Buffer Sizing**: The replay buffer size is automatically calculated as:

nemo_rl/algorithms/async_utils/interfaces.py

@@ -15,7 +15,7 @@
 from typing import Any, Optional, Protocol
 
 
-class ReplayBufferProtocol(Protocol):
+class ReplayBufferProtocol(Protocol):  # pragma: no cover
     """Interface for the replay buffer used in async RL training."""
 
     def add(
@@ -58,3 +58,35 @@ def size(self) -> int:
     def clear(self) -> None:
         """Clear the buffer."""
         ...
+
+    def state_dict(self) -> dict[str, Any]:
+        """Return serializable state for checkpointing."""
+        ...
+
+    def load_state_dict(
+        self,
+        state: dict[str, Any],
+        num_prompts_per_step: int | None = None,
+        current_training_step: int | None = None,
+        max_age_steps: int | None = None,
+    ) -> None:
+        """Restore state produced by ``state_dict``."""
+        ...
+
+    def get_trajectories_needed(
+        self,
+        target_step: int,
+        num_prompts_per_step: int,
+        max_age_steps: int | None = None,
+    ) -> int:
+        """Return additional trajectories needed for ``target_step``."""
+        ...
+
+    def has_complete_batch(
+        self,
+        target_step: int,
+        num_prompts_per_step: int,
+        max_age_steps: int | None = None,
+    ) -> bool:
+        """Return whether ``target_step`` has enough trajectories to train."""
+        ...