[feat][checkpoint] Allow selective state restore on resume

docs/content/docs/checkpointing-logging/checkpointing.mdx

@@ -127,6 +127,37 @@ Checkpointing behavior is controlled by several parameters in the YAML configura
   - **Purpose**: Specific checkpoint directory to resume from (only used when `resume_mode: "from_path"`)
   - **Format**: Must point to a `global_step_N` directory
 
+`resume_load_global_step`
+  - **Default**: `true`
+  - **Purpose**: Restore the checkpoint's global step. Set to `false` to start a new training phase at step zero.
+
+`resume_load_dataloader_state`
+  - **Default**: `true`
+  - **Purpose**: Restore dataloader position. Set to `false` when starting a new training phase or changing datasets.
+
+`resume_load_optimizer_states`
+  - **Default**: `true`
+  - **Purpose**: Restore policy and critic optimizer states. Set to `false` to use optimizers initialized from the current configuration.
+
+`resume_load_lr_scheduler_states`
+  - **Default**: `true`
+  - **Purpose**: Restore policy and critic learning rate scheduler states. Set to `false` to use schedulers initialized from the current configuration.
+
+Disable selected state-loading options when using checkpoint weights to start a new training phase or dataset:
+
+```yaml
+trainer:
+  resume_mode: from_path
+  resume_path: /path/to/ckpts/global_step_100
+  ckpt_path: /path/to/new-phase-ckpts
+  resume_load_global_step: false
+  resume_load_dataloader_state: false
+  resume_load_optimizer_states: false
+  resume_load_lr_scheduler_states: false
+```
+
+Use a different `ckpt_path` when resetting the global step so new checkpoints cannot overwrite the source run.
+
 ## HuggingFace Model Export
 
 In addition to checkpointing, users can optionally save the policy model in HuggingFace safetensors format at regular intervals.
@@ -166,4 +197,4 @@ filesystem path visible to all Megatron ranks. Node-local paths and cloud paths
 across per-rank local work directories, and Megatron-Bridge builds `model.safetensors.index.json`
 from the shard files visible to rank 0 — so the index comes out incomplete if rank 0 cannot see
 every shard written by the other ranks.
-</Callout>
+</Callout>

skyrl/train/config/config.py

@@ -1360,6 +1360,14 @@ class TrainerConfig(BaseConfig):
     See https://docs.skyrl.ai/docs/checkpointing-logging/checkpointing"""
     resume_path: Optional[str] = None
     """Checkpoint directory to resume from. Only used when ``resume_mode="from_path"``."""
+    resume_load_global_step: bool = True
+    """Restore the global step when resuming. Disable to start a new training phase at step zero."""
+    resume_load_dataloader_state: bool = True
+    """Restore dataloader position when resuming. Disable when starting a new phase or using different data."""
+    resume_load_optimizer_states: bool = True
+    """Restore optimizer state when resuming. Disable to initialize a fresh optimizer from the current config."""
+    resume_load_lr_scheduler_states: bool = True
+    """Restore LR scheduler state when resuming. Disable to initialize a fresh scheduler from the current config."""
     log_path: str = "/tmp/skyrl-logs"
     """Path for infrastructure log files.
     vLLM engine startup, model loading, and worker initialization logs are written to

skyrl/train/fully_async_trainer.py

@@ -429,13 +429,40 @@ def _build_train_dataloader_and_compute_training_steps(self):
         logger.info(f"Number of steps per epoch: {self.num_steps_per_epoch}")
         logger.info(f"Total training steps: {self.total_training_steps}")
 
+    def _restore_async_training_state(
+        self,
+        consumed_data_uids: Optional[Set[str]],
+        filtered_data_uids: Optional[Set[str]],
+        epoch: Optional[int],
+    ) -> Tuple[int, int]:
+        self._staleness_manager.load_state_from_checkpoint(self.global_step + 1)
+        if consumed_data_uids is None:
+            return (
+                self.global_step // self.num_steps_per_epoch,
+                self.global_step % self.num_steps_per_epoch,
+            )
+
+        filtered_data_uids = filtered_data_uids or set()
+        self.async_train_dataloader.load_state_from_checkpoint(consumed_data_uids, filtered_data_uids)
+        # Only trained UIDs map to completed steps (filtered UIDs are extra consumption),
+        # so validate the trained count is a whole number of steps.
+        num_trained_loaded = len(consumed_data_uids) - len(filtered_data_uids)
+        assert num_trained_loaded % self.mini_batch_size == 0, (
+            "Loaded trained (consumed minus filtered) data UIDs must be a multiple of "
+            f"mini_batch_size={self.mini_batch_size}. Got: {num_trained_loaded}"
+        )
+        # Fall back to deriving the epoch for checkpoints that predate epoch tracking.
+        start_epoch = epoch if epoch is not None else self.global_step // self.num_steps_per_epoch
+        return start_epoch, num_trained_loaded // self.mini_batch_size
+
     async def train(self):
         """
         Main fully async training loop for PPO
         """
         self.global_step = 0
         self.epoch = 0
         resumed_start_epoch = None
+        resumed_steps_into_epoch = None
 
         # Load checkpoint state if resumption is enabled. Also load the data UIDs that are already trained on.
         if self.resume_mode != ResumeMode.NONE:
@@ -449,24 +476,10 @@ async def train(self):
                 ) = self.load_checkpoints()
                 logger.info(f"Resumed training from global_step {self.global_step}")
                 if self.global_step > 0:
-                    # Set async dataloader manager and staleness manager to the loaded state.
-                    self.async_train_dataloader.load_state_from_checkpoint(
-                        loaded_consumed_data_uids_set, loaded_filtered_data_uids_set
-                    )
-                    self._staleness_manager.load_state_from_checkpoint(
-                        self.global_step + 1
-                    )  # +1 due to we haven't incremented yet
-                    # Only trained UIDs map to completed steps (filtered UIDs are extra consumption),
-                    # so validate the trained count is a whole number of steps.
-                    num_trained_loaded = len(loaded_consumed_data_uids_set) - len(loaded_filtered_data_uids_set)
-                    assert num_trained_loaded % self.mini_batch_size == 0, (
-                        "Loaded trained (consumed minus filtered) data UIDs must be a multiple of "
-                        f"mini_batch_size={self.mini_batch_size}. Got: {num_trained_loaded}"
-                    )
-                    # Use the persisted epoch; fall back to deriving it for pre-sample_full_batch
-                    # checkpoints (where global_step stays aligned to epoch boundaries).
-                    resumed_start_epoch = (
-                        loaded_epoch if loaded_epoch is not None else self.global_step // self.num_steps_per_epoch
+                    resumed_start_epoch, resumed_steps_into_epoch = self._restore_async_training_state(
+                        loaded_consumed_data_uids_set,
+                        loaded_filtered_data_uids_set,
+                        loaded_epoch,
                     )
 
         # Initialize weight sync state
@@ -521,10 +534,14 @@ async def _watch_generators_done(tasks=generator_tasks, event=all_generators_don
 
                     generators_done_watcher = asyncio.create_task(_watch_generators_done())
 
-                # Steps trained in THIS epoch (not global_step % num_steps_per_epoch: sample_full_batch can
-                # end an epoch early, drifting global_step out of epoch alignment). On resume the dataloader
-                # already reflects this epoch's trained steps. The range below is just an upper bound.
+                # Track actual trained data separately from logical epoch progress. They differ when resume
+                # keeps global_step but intentionally skips the dataloader cursor.
                 trained_steps_this_epoch = self.async_train_dataloader.num_trained() // self.mini_batch_size
+                steps_into_epoch = (
+                    resumed_steps_into_epoch
+                    if epoch == start_epoch and resumed_steps_into_epoch is not None
+                    else trained_steps_this_epoch
+                )
                 for _step_idx in range(self.global_step, (1 + epoch) * self.num_steps_per_epoch + 1):
                     with Timer("step", self.all_timings):
                         self._loop_gauges.set(
@@ -599,6 +616,7 @@ async def _watch_generators_done(tasks=generator_tasks, event=all_generators_don
 
                     # A training step completed: count it for this epoch's bookkeeping.
                     trained_steps_this_epoch += 1
+                    steps_into_epoch += 1
 
                     # One profiler step per async global step.
                     self._profiler_step()
@@ -623,7 +641,7 @@ async def _watch_generators_done(tasks=generator_tasks, event=all_generators_don
                     self.all_metrics = {}
 
                     # 7. Checkpointing. At interval and at the last step of each epoch.
-                    is_epoch_end = trained_steps_this_epoch == self.num_steps_per_epoch
+                    is_epoch_end = steps_into_epoch == self.num_steps_per_epoch
                     if self.cfg.trainer.ckpt_interval > 0:
                         if is_epoch_end or self.global_step % self.cfg.trainer.ckpt_interval == 0:
                             with self._phase_gauge.timed_phase("save_checkpoints", self.all_timings):
@@ -1184,8 +1202,8 @@ def load_checkpoints(self) -> Tuple[int, str, Optional[Set[str]], Optional[Set[s
         checkpoint predates epoch tracking, in which case the caller derives it from global_step).
         """
         global_step, checkpoint_path = super().load_checkpoints()
-        if global_step == 0:
-            return 0, checkpoint_path, None, None, None
+        if global_step == 0 or not self.cfg.trainer.resume_load_dataloader_state:
+            return global_step, checkpoint_path, None, None, None
         fully_async_state_path = os.path.join(checkpoint_path, "fully_async_state.pt")
         assert io.exists(fully_async_state_path), f"Fully-async state file not found at {fully_async_state_path}"
         with io.open_file(fully_async_state_path, "rb") as f:

skyrl/train/trainer.py

@@ -1721,8 +1721,8 @@ def _cleanup_old_checkpoints(self):
 
     def load_checkpoints(self) -> Tuple[int, str]:
         """
-        Load complete checkpoint state and return the global_step to resume from.
-        Returns 0 if no checkpoint is loaded.
+        Load checkpoint state and return the configured starting global step.
+        Returns 0 if no checkpoint is loaded or global-step restore is disabled.
 
         If colocate_all is True, assumes that the policy model is currently on GPU.
 
@@ -1771,10 +1771,14 @@ def load_checkpoints(self) -> Tuple[int, str]:
         logger.info(f"Loading checkpoint from: {checkpoint_path}")
 
         # Extract global step from checkpoint path
-        global_step = extract_step_from_path(Path(checkpoint_path))
-        if global_step == -1:
+        checkpoint_global_step = extract_step_from_path(Path(checkpoint_path))
+        if checkpoint_global_step == -1:
             raise ValueError(f"Checkpoint path {checkpoint_path} is not a valid checkpoint path")
-        logger.info(f"Resuming from global_step: {global_step}")
+        global_step = checkpoint_global_step if self.cfg.trainer.resume_load_global_step else 0
+        if self.cfg.trainer.resume_load_global_step:
+            logger.info(f"Resuming from global_step: {global_step}")
+        else:
+            logger.info(f"Loading checkpoint weights from global_step_{checkpoint_global_step}; global step reset to 0")
 
         # Define paths for different checkpoint components
         policy_ckpt_dir = os.path.join(checkpoint_path, "policy")
@@ -1789,13 +1793,19 @@ def load_checkpoints(self) -> Tuple[int, str]:
         # 1. Load and validate trainer state
         with io.open_file(trainer_state_path, "rb") as f:
             trainer_state = torch.load(f, map_location="cpu", weights_only=False)
-        saved_global_step = trainer_state.get("global_step", global_step)
-        logger.info("Successfully loaded trainer state")
-        if saved_global_step != global_step:
-            logger.warning(f"Global step mismatch: path={global_step}, saved={saved_global_step}. Using path value.")
+        saved_global_step = trainer_state.get("global_step", checkpoint_global_step)
+        logger.info("Successfully loaded trainer metadata")
+        if saved_global_step != checkpoint_global_step:
+            logger.warning(
+                f"Global step mismatch: path={checkpoint_global_step}, saved={saved_global_step}. Using path value."
+            )
 
-        # 2. Load dataloader state if available
-        if io.exists(dataloader_state_path):
+        # 2. Load dataloader state if requested and available
+        if self.train_dataloader is None:
+            logger.info("No train dataloader initialized; skipping dataloader state restore")
+        elif not self.cfg.trainer.resume_load_dataloader_state:
+            logger.info("Skipping dataloader state restore; dataloader will start from beginning")
+        elif io.exists(dataloader_state_path):
             try:
                 with io.open_file(dataloader_state_path, "rb") as f:
                     dataloader_state = torch.load(f, map_location="cpu", weights_only=False)
@@ -1813,8 +1823,8 @@ def load_checkpoints(self) -> Tuple[int, str]:
         self.dispatch.load_checkpoint(
             "policy",
             policy_ckpt_dir,
-            load_optimizer_states=True,
-            load_lr_scheduler_states=True,
+            load_optimizer_states=self.cfg.trainer.resume_load_optimizer_states,
+            load_lr_scheduler_states=self.cfg.trainer.resume_load_lr_scheduler_states,
         )
         logger.info("Successfully loaded policy checkpoint")
 
@@ -1824,12 +1834,12 @@ def load_checkpoints(self) -> Tuple[int, str]:
             self.dispatch.load_checkpoint(
                 "critic",
                 critic_ckpt_dir,
-                load_optimizer_states=True,
-                load_lr_scheduler_states=True,
+                load_optimizer_states=self.cfg.trainer.resume_load_optimizer_states,
+                load_lr_scheduler_states=self.cfg.trainer.resume_load_lr_scheduler_states,
             )
             logger.info("Successfully loaded critic checkpoint")
 
-        logger.info(f"Successfully loaded complete checkpoint state from global_step_{global_step}")
+        logger.info(f"Successfully loaded checkpoint state from global_step_{checkpoint_global_step}")
         return global_step, str(checkpoint_path)
 
     def save_models(self):