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

examples/configs/grpo_math_1B.yaml

@@ -520,6 +520,10 @@ data_plane:
   storage_capacity: 1000000          # max samples retained per partition
   num_storage_units: 2               # storage shards
   claim_meta_poll_interval_s: 0.5    # blocking-claim poll cadence
+  # SingleController only: save native TQ state. Required when trainer
+  # checkpointing is enabled with a replay-checkpoint-capable sampler;
+  # supported samplers restore from metadata-only replay indexes.
+  checkpointing_enabled: false
   global_segment_size: 549755813888  # 512 GiB — used when backend == "mooncake_cpu"
   local_buffer_size:   68719476736   # 64 GiB  — used when backend == "mooncake_cpu"
   # observability:                    # NotRequired

examples/configs/grpo_math_1B_megatron_single_controller.yaml

@@ -58,6 +58,8 @@ logger:
 
 data_plane:
   enabled: true
+  # Required shadow snapshot: a save failure aborts checkpoint finalization.
+  checkpointing_enabled: true
 
 cluster:
   gpus_per_node: 2

nemo_rl/algorithms/async_utils/staleness_sampler.py

@@ -44,6 +44,7 @@
 from typing import (
     Annotated,
     Callable,
+    ClassVar,
     Literal,
     Optional,
     Protocol,
@@ -109,6 +110,11 @@ def is_on_policy(self) -> bool:
         """True when the policy admits zero staleness (sync mode)."""
         ...
 
+    @property
+    def supports_buffer_checkpoint(self) -> bool:
+        """Whether completed buffered groups can be restored safely."""
+        ...
+
     def required_buffer_capacity(self, groups_per_step: int) -> Optional[int]:
         """Buffer-capacity the policy needs, or ``None`` if unconstrained."""
         ...
@@ -196,6 +202,10 @@ def should_abort_inflight(
     def is_on_policy(self) -> bool:
         return self._eviction_window() == 0
 
+    @property
+    def supports_buffer_checkpoint(self) -> bool:
+        return False
+
     def required_buffer_capacity(self, groups_per_step: int) -> Optional[int]:
         return None
 
@@ -266,6 +276,11 @@ def __init__(
     def _eviction_window(self) -> int:
         return self.max_staleness_versions
 
+    @property
+    def supports_buffer_checkpoint(self) -> bool:
+        # Ungated restored groups are ordinary in-window candidates.
+        return True
+
     def should_abort_inflight(
         self,
         *,
@@ -443,6 +458,7 @@ async def evict(self, *, current_train_weight: int) -> int:
 
 
 class WindowedSamplerConfig(BaseModel, extra="allow"):
+    supports_buffer_checkpoint: ClassVar[bool] = True
     name: Literal["windowed"] = "windowed"
     # Max weight-version gap a selected group may have from the trainer.
     max_staleness_versions: NonNegativeInt = 1
@@ -451,18 +467,22 @@ class WindowedSamplerConfig(BaseModel, extra="allow"):
 
 
 class WeightFifoSamplerConfig(BaseModel, extra="allow"):
+    supports_buffer_checkpoint: ClassVar[bool] = False
     name: Literal["weight_fifo"] = "weight_fifo"
     # Lookahead + selectable weight window, in trainer versions.
     max_staleness_versions: NonNegativeInt = 1
 
 
 class InOrderSamplerConfig(BaseModel, extra="allow"):
+    supports_buffer_checkpoint: ClassVar[bool] = False
     name: Literal["in_order"] = "in_order"
     # How far generation may run ahead of the trainer, in dispatch batches.
     max_lookahead_versions: NonNegativeInt = 1
 
 
 class CustomSamplerConfig(BaseModel, extra="allow"):
+    # A custom implementation's capability is known only after construction.
+    supports_buffer_checkpoint: ClassVar[Optional[bool]] = None
     name: Literal["custom"] = "custom"
     # "module:ClassName" of a PromptGroupSampler defined outside this repo.
     # Extra keys are forwarded to the constructor (after ``buffer``).
@@ -505,20 +525,30 @@ def create_sampler(
     buffer: TQReplayBuffer,
     cfg: SamplerConfig,
 ) -> PromptGroupSampler:
-    """Build a sampler from its config (or import one by FQN)."""
+    """Build a sampler from its config (or import one by FQN).
+
+    Args:
+        buffer: Shared TQReplayBuffer holding the candidate slots.
+        cfg: Discriminated sampler config selecting the policy.
+    """
+    sampler: PromptGroupSampler
     if isinstance(cfg, WindowedSamplerConfig):
-        return WindowedSampler(
+        sampler = WindowedSampler(
             buffer,
             max_staleness_versions=cfg.max_staleness_versions,
             sample_freshest_first=cfg.sample_freshest_first,
         )
-    if isinstance(cfg, WeightFifoSamplerConfig):
-        return WeightFifoSampler(
-            buffer, max_staleness_versions=cfg.max_staleness_versions
+    elif isinstance(cfg, WeightFifoSamplerConfig):
+        sampler = WeightFifoSampler(
+            buffer,
+            max_staleness_versions=cfg.max_staleness_versions,
         )
-    if isinstance(cfg, InOrderSamplerConfig):
-        return InOrderSampler(buffer, max_lookahead_versions=cfg.max_lookahead_versions)
-    if isinstance(cfg, CustomSamplerConfig):
+    elif isinstance(cfg, InOrderSamplerConfig):
+        sampler = InOrderSampler(
+            buffer,
+            max_lookahead_versions=cfg.max_lookahead_versions,
+        )
+    elif isinstance(cfg, CustomSamplerConfig):
         module_name, sep, class_name = cfg.target.partition(":")
         if not sep:
             raise ValueError(
@@ -532,5 +562,17 @@ def create_sampler(
                 f"interface (needs admit/select/evict/should_abort_inflight, "
                 f"set_dispatch_index, is_on_policy, required_buffer_capacity)"
             )
-        return sampler
-    raise ValueError(f"unknown sampler config {type(cfg).__name__}")
+    else:
+        raise ValueError(f"unknown sampler config {type(cfg).__name__}")
+
+    expected_capability = cfg.supports_buffer_checkpoint
+    if (
+        expected_capability is not None
+        and sampler.supports_buffer_checkpoint != expected_capability
+    ):
+        raise RuntimeError(
+            f"{type(cfg).__name__}.supports_buffer_checkpoint="
+            f"{expected_capability} disagrees with {type(sampler).__name__}."
+            f"supports_buffer_checkpoint={sampler.supports_buffer_checkpoint}"
+        )
+    return sampler

nemo_rl/algorithms/grpo.py

@@ -315,6 +315,9 @@ class GRPOSaveState:
     total_steps: int
     total_valid_tokens: int  # Track total number of non-padding tokens during training
     val_reward: float  # May be removed when no validation metrics are available
+    # SC may advance the policy version independently from the optimizer-step
+    # counter. None preserves compatibility with checkpoints predating it.
+    trainer_version: Optional[int] = None
     # SingleController only: name of the sampler that wrote the replay buffer,
     # used to gate the SC buffer restore. None on checkpoints from the other
     # algorithms and from SC runs that predate this field.
@@ -329,6 +332,7 @@ def _initial_grpo_save_state() -> GRPOSaveState:
         total_steps=0,
         total_valid_tokens=0,
         val_reward=-99999999.0,
+        trainer_version=None,
         sampler_name=None,
     )