Trainer

class RayPPOTrainer:
    def __init__(
        self,
        cfg: SkyRLTrainConfig,
        tracker: Tracking,
        tokenizer: AutoTokenizer,
        train_dataset: Optional[PromptDataset],
        inference_engine_client: InferenceEngineInterface,
        generator: GeneratorInterface,
        colocate_pg: Optional[ResolvedPlacementGroup] = None,
        eval_dataset: Optional[PromptDataset] = None,
        callbacks: Optional[List[TrainingCallback]] = None,
    ):
        self.cfg = cfg
        self.colocate_all = cfg.trainer.placement.colocate_all
        self.tracker = tracker
        self.tokenizer = tokenizer
        self.train_dataset = train_dataset
        self.eval_dataset = eval_dataset
        self.inference_engine_client = inference_engine_client
        self.generator = generator
        self.train_dataloader = None
        self.total_training_steps = None
        self._build_train_dataloader_and_compute_training_steps()

        self.eval_dataloader = (
            build_dataloader(self.cfg, eval_dataset, is_train=False) if eval_dataset is not None else None
        )
        self.colocate_pg = colocate_pg

        self.resume_mode = ResumeMode(cfg.trainer.resume_mode)

        self.all_metrics = {}
        self.all_timings = {}
        self.global_step = 0

        self._vllm_metrics_scraper: Optional[VLLMMetricsScraper] = (
            VLLMMetricsScraper() if cfg.generator.inference_engine.enable_ray_prometheus_stats else None
        )

        self._ray_gpu_monitor = RayGpuMonitor() if cfg.trainer.enable_ray_gpu_monitor else None

        # trajectory logger is installed after construction if needed
        self.trajectory_logger: TrajectoryLogger = None

        # initialized in `build_models`
        self.policy_model: PPORayActorGroup = None
        self.critic_model: Optional[PPORayActorGroup] = None
        self.ref_model: Optional[PPORayActorGroup] = None
        # used for checkpoint cleanup
        self._node_ids: Optional[List[str]] = None

        self.dynamic_sampling_state: Optional[DynamicSamplingState] = None

        self.reward_kl_controller: Optional[Union[FixedKLController, AdaptiveKLController]] = None
        self.dispatch: WorkerDispatch = None

        self._callback_handler = CallbackHandler(callbacks)
        self._training_control = TrainingControl()
        self._current_epoch: int = 0

        configure_ray_worker_logging()

        self._num_training_gpus = (
            cfg.trainer.placement.policy_num_gpus_per_node * cfg.trainer.placement.policy_num_nodes
        )

    def add_callback(self, callback: TrainingCallback) -> None:
        """Register a callback. Events fired after this call reach the new callback."""
        self._callback_handler.add(callback)

    def _build_callback_input(self, **fields) -> CallbackInput:
        """Snapshot loop counters + per-event fields into a CallbackInput."""
        steps_per_epoch = len(self.train_dataloader) if self.train_dataloader is not None else 0
        total_steps = self.total_training_steps or 0
        return CallbackInput(
            global_step=self.global_step,
            epoch=self._current_epoch,
            total_steps=total_steps,
            steps_per_epoch=steps_per_epoch,
            **fields,
        )

    def _fire(self, event_name: str, **fields) -> None:
        """Build a CallbackInput and dispatch the given event to all callbacks."""
        cb_input = self._build_callback_input(**fields)
        getattr(self._callback_handler, event_name)(self, cb_input, self._training_control)

    @property
    def has_critic(self) -> bool:
        """Check if critic model is configured."""
        return bool(self.cfg.trainer.critic.model.path)

    @property
    def _torch_profiler_enabled(self) -> bool:
        """Whether to dispatch policy profiler RPCs."""
        return self.cfg.trainer.policy.torch_profiler_config.enable

    def _profiler_start(self) -> None:
        """Start policy profiling when enabled."""
        if self._torch_profiler_enabled:
            self.dispatch.start_profile("policy")

    def _profiler_step(self) -> None:
        """Advance policy profiling by one global step."""
        if self._torch_profiler_enabled:
            self.dispatch.profile_step("policy")

    def _profiler_stop(self) -> None:
        """Stop policy profiling when enabled."""
        if self._torch_profiler_enabled:
            self.dispatch.stop_profile("policy")

    def _build_train_dataloader_and_compute_training_steps(self):
        """
        Hook for constructing the training dataloader. Subclasses can override
        this to customize dataloader behavior. For instance, fully async training
        needs a batch size of 1, among other features.
        Defaults to `trainer_utils.build_dataloader` with `is_train=True`.
        When train_dataset is None (e.g. Tinker backend provides data externally),
        the dataloader is not built.
        """
        if self.train_dataset is not None:
            self.train_dataloader = build_dataloader(self.cfg, self.train_dataset, is_train=True)
            self.total_training_steps = len(self.train_dataloader) * self.cfg.trainer.epochs
            if self.cfg.trainer.max_training_steps is not None:
                self.total_training_steps = min(self.total_training_steps, self.cfg.trainer.max_training_steps)

    @torch.no_grad()
    async def eval(self, vllm_metrics_scraper: Optional[VLLMMetricsScraper] = None) -> Dict[str, float]:
        """
        Run generation and scoring on the evaluation dataset.

        The eval metrics are recorded after having finished training `self.global_step` steps.
        Metrics recorded in global_step 0 corresponds to evaluations before training.

        Args:
            vllm_metrics_scraper: when provided, the eval loop calls
                ``resume()``/``pause()`` around each generation so the scraper
                attributes only generation time to the open ``vllm/eval`` window.

        Returns:
            A dictionary of evaluation metrics.
        """
        if self.cfg.generator.step_wise_trajectories:
            eval_metrics = await evaluate_step_wise(
                eval_dataloader=self.eval_dataloader,
                generator=self.generator,
                cfg=self.cfg,
                global_step=self.global_step,
                tokenizer=self.tokenizer,
                trajectory_logger=self.trajectory_logger,
                tracker=self.tracker,
                vllm_metrics_scraper=vllm_metrics_scraper,
            )
        else:
            eval_metrics = await evaluate(
                eval_dataloader=self.eval_dataloader,
                generator=self.generator,
                cfg=self.cfg,
                global_step=self.global_step,
                tokenizer=self.tokenizer,
                trajectory_logger=self.trajectory_logger,
                tracker=self.tracker,
                vllm_metrics_scraper=vllm_metrics_scraper,
            )
        return eval_metrics

    async def train(self):
        """
        Main training loop for PPO
        """
        if self._ray_gpu_monitor is not None:
            self._ray_gpu_monitor.start()

        # Initialize weight sync state between policy model and inference engines.
        with Timer("init_weight_sync_state"):
            self.init_weight_sync_state()

        # Load checkpoint state if resumption is enabled.
        if self.resume_mode != ResumeMode.NONE:
            with Timer("load_checkpoints"):
                self.global_step, _ = self.load_checkpoints()

        # Prepare weights for sampling
        with Timer("sync_weights"):
            await self.dispatch.save_weights_for_sampler()

        # Compute start_epoch up-front so callback metadata is ready before
        # any event fires (including the baseline eval below).
        start_epoch = self.global_step // len(self.train_dataloader)
        self._current_epoch = start_epoch
        self._training_control.reset()

        self._fire("on_train_start")

        # Eval before training. Wrapped in eval callbacks + on_log so that e.g.
        # a best-checkpoint callback sees the baseline reading.
        if self.cfg.trainer.eval_interval > 0 and self.cfg.trainer.eval_before_train:
            self._fire("on_eval_start")
            with Timer("eval", self.all_timings):
                eval_metrics = await self.eval()
            self._fire("on_eval_end", metrics=eval_metrics)
            self._fire("on_log", logs=eval_metrics)
            self.tracker.log(eval_metrics, step=self.global_step, commit=True)

        # initialize kl controller
        if self.cfg.trainer.algorithm.use_kl_in_reward:
            self.reward_kl_controller = get_kl_controller(self.cfg.trainer.algorithm)

        # main training loop
        pbar = tqdm(total=self.total_training_steps, initial=self.global_step, desc="Training Batches Processed")
        self.global_step += 1  # start training at global_step 1
        stop_training = False

        # booleans tracking whether we save ckpts
        # as well as hf model at step end
        will_save_ckpts = False
        hf_model_save = False
        self._profiler_start()
        try:
            for epoch in range(start_epoch, self.cfg.trainer.epochs):
                self._current_epoch = epoch
                self._fire("on_epoch_start")
                # ``step_started`` tracks the on_step_start/on_step_end pairing taking
                # dynamic-sampling into account (which span multiple inner iterations
                # before completing a logical step).
                step_started = False
                for _, rand_prompts in enumerate(self.train_dataloader):
                    if not step_started:
                        self._fire("on_step_start")
                        step_started = True
                        # Open the train-rollout metrics window once per logical
                        # step; paused so only the generation spans count toward the
                        # throughput denominator (dynamic sampling may generate more
                        # than once before the step completes).
                        if self._vllm_metrics_scraper is not None:
                            await self._vllm_metrics_scraper.start("vllm/train")
                            self._vllm_metrics_scraper.pause()
                    with Timer("step", self.all_timings):
                        # for colocate_all=true, inference engine is always on GPU when starting the training step

                        # 0. truncate data to have even shards
                        rand_prompts = self._remove_tail_data(rand_prompts)
                        generator_input, uids = prepare_generator_input(
                            rand_prompts,
                            self.cfg.generator.n_samples_per_prompt,
                            get_sampling_params_for_backend(
                                self.cfg.generator.inference_engine.backend, self.cfg.generator.sampling_params
                            ),
                            self.cfg.environment.env_class,
                            "train",
                            self.global_step,
                        )

                        # 1.1. generation phase
                        if self._vllm_metrics_scraper is not None:
                            self._vllm_metrics_scraper.resume()
                        with Timer("generate", self.all_timings):
                            generator_output: GeneratorOutput = await self.generate(generator_input)
                        if self._vllm_metrics_scraper is not None:
                            self._vllm_metrics_scraper.pause()

                        if self.cfg.generator.step_wise_trajectories:
                            # NOTE: We use instance_ids from `trajectory_ids` here instead of re-using `uids`
                            # this is because in step-wise training, len(uids) != len(generator_output["response_ids"])
                            uids = [trajectory_id.instance_id for trajectory_id in generator_output["trajectory_ids"]]

                        # dynamic sampling
                        if self.cfg.trainer.algorithm.dynamic_sampling.type is not None:
                            generator_output, uids, keep_sampling = self.handle_dynamic_sampling(generator_output, uids)
                            if keep_sampling:  # continue sampling
                                # update progress bar for current batch (but not global step)
                                pbar.update(1)
                                continue

                        if self.colocate_all:
                            # if we are not continuing sampling, we sleep the inference engine
                            await self.inference_engine_client.sleep()

                        # The train rollout for this step is done generating; close
                        # its metrics window. ``vllm/eval/*`` is collected separately
                        # around eval below.
                        vllm_metrics: Dict[str, float] = {}
                        if self._vllm_metrics_scraper is not None:
                            vllm_metrics = await self._vllm_metrics_scraper.stop()

                        # 1.2 postprocess rewards (and merge step-wise turns if enabled)
                        with Timer("postprocess_generator_output", self.all_timings):
                            generator_output, uids = self.postprocess_generator_output(generator_output, uids)

                        # 2.1 print example just for debugging
                        print_interval = self.cfg.trainer.print_example_interval
                        if print_interval > 0 and self.global_step % print_interval == 0:
                            vis = self.tokenizer.decode(generator_output["response_ids"][0])
                            pretty_print_example(
                                logger,
                                prompt=generator_input["prompts"][0],
                                response=vis,
                                reward=generator_output["rewards"][0],
                            )

                        # 2.2 Optionally upload up to `num_logger_train_samples` samples to tracker
                        if self.trajectory_logger is not None:
                            with Timer("log_train_results"):
                                self.trajectory_logger.log(
                                    tracker=self.tracker,
                                    num_samples=self.cfg.trainer.num_logger_train_samples,
                                    prompts=generator_input["prompts"],
                                    generator_output=generator_output,
                                    tokenizer=self.tokenizer,
                                    global_step=self.global_step,
                                    wandb_key="trajectories/train",
                                    include_idx=False,
                                )

                        # 3. Convert GeneratorOutput to TrainingInputBatch
                        with Timer("convert_to_training_input", self.all_timings):
                            training_input: TrainingInputBatch = self.convert_to_training_input(generator_output, uids)

                        # 4. Inference and calculate values, log probs, rewards, kl divergence
                        with Timer("fwd_logprobs_values_reward", self.all_timings):
                            training_input = self.fwd_logprobs_values_reward(training_input)

                        # 5. apply kl divergence penalty to rewards
                        if self.cfg.trainer.algorithm.use_kl_in_reward:
                            with Timer("apply_reward_kl_penalty", self.all_timings):
                                training_input = self.apply_reward_kl_penalty(training_input)

                        # 6. calculate advantages and returns
                        with Timer("compute_advantages_and_returns", self.all_timings):
                            training_input = self.compute_advantages_and_returns(training_input)
                            # remove some unwanted keys
                            for key in ["rewards"]:
                                training_input.pop(key)
                            training_input.metadata.pop("uids")
                            training_input.metadata.pop("is_last_step", None)

                        if self.cfg.trainer.dump_data_batch:
                            # dump data to file
                            with Timer("dump_data_batch"):
                                self.dump_data(
                                    training_input, file_name=f"global_step_{self.global_step}_training_input"
                                )

                        # 7. train policy/critic model
                        # Policy model is backloaded to GPU during training
                        with Timer("train_critic_and_policy", self.all_timings):
                            status = self.train_critic_and_policy(training_input)

                            # One profiler step per RL global step.
                            self._profiler_step()

                        self._fire("on_step_end", batch=training_input, metrics=status)
                        step_started = False

                        # Capture callback-driven triggers, then reset.
                        force_save = self._training_control.should_save
                        force_eval = self._training_control.should_evaluate
                        self._training_control.should_save = False
                        self._training_control.should_evaluate = False

                        # 8. conditionally save checkpoints and hf model
                        is_epoch_end = self.global_step % len(self.train_dataloader) == 0
                        hf_model_save = self.cfg.trainer.hf_save_interval > 0 and (
                            is_epoch_end or self.global_step % self.cfg.trainer.hf_save_interval == 0
                        )
                        ckpt_interval_save = self.cfg.trainer.ckpt_interval > 0 and (
                            is_epoch_end or self.global_step % self.cfg.trainer.ckpt_interval == 0
                        )
                        will_save_ckpts = force_save or ckpt_interval_save
                        if will_save_ckpts:
                            with Timer("save_checkpoints", self.all_timings):
                                ckpt_path = self.save_checkpoints()
                            self._fire("on_save", ckpt_path=ckpt_path)
                        if hf_model_save:
                            with Timer("save_hf_model", self.all_timings):
                                self.save_models()

                        # 9. conditionally sync policy and ref at the end of the epoch
                        if (
                            self.cfg.trainer.update_ref_every_epoch
                            and self.ref_model is not None
                            and is_epoch_end
                            and epoch != self.cfg.trainer.epochs - 1  # skip updating ref at the end of the last epoch
                        ):
                            with Timer("update_ref_with_policy", self.all_timings):
                                self.update_ref_with_policy()

                        # 10. Prepare weights for sampling
                        with Timer("sync_weights", self.all_timings):
                            await self.dispatch.save_weights_for_sampler()

                    # 11. set logs
                    logger.info(status)
                    # Throughput metrics
                    train_time = self.all_timings.get("train_critic_and_policy", 0.0)
                    if train_time > 0 and training_input.get("attention_mask") is not None:
                        total_tokens = int(training_input["attention_mask"].sum().item())
                        self.all_metrics["trainer/tokens_per_second_per_gpu"] = total_tokens / (
                            train_time * self._num_training_gpus
                        )
                    # log epoch info
                    self.all_metrics.update({"trainer/epoch": epoch, "trainer/global_step": self.global_step})
                    interval_eval = self.cfg.trainer.eval_interval > 0 and (
                        self.global_step % self.cfg.trainer.eval_interval == 0
                        or self.global_step == self.total_training_steps
                    )
                    if force_eval or interval_eval:
                        # Open the eval-rollout window; the scraper itself measures
                        # the generation spans via resume()/pause() inside eval().
                        if self._vllm_metrics_scraper is not None:
                            await self._vllm_metrics_scraper.start("vllm/eval")
                            self._vllm_metrics_scraper.pause()
                        self._fire("on_eval_start")
                        with Timer("eval", self.all_timings):
                            eval_metrics = await self.eval(vllm_metrics_scraper=self._vllm_metrics_scraper)
                            self.all_metrics.update(eval_metrics)
                        self._fire("on_eval_end", metrics=eval_metrics)
                        if self._vllm_metrics_scraper is not None:
                            vllm_metrics.update(await self._vllm_metrics_scraper.stop())

                    log_payload = {
                        **self.all_metrics,
                        **{f"timing/{k}": v for k, v in self.all_timings.items()},
                        # vllm/train/* = train rollout, vllm/eval/* = eval rollout,
                        # each over its own generation time (owned by the scraper).
                        **vllm_metrics,
                    }

                    if self._ray_gpu_monitor is not None:
                        log_payload.update(self._ray_gpu_monitor.flush())

                    self._fire("on_log", logs=log_payload)

                    self.tracker.log(log_payload, step=self.global_step, commit=True)
                    self.all_metrics = {}
                    self.all_timings = {}

                    # update progress bar after logging
                    pbar.update(1)

                    self.global_step += 1

                    if (
                        self.cfg.trainer.max_training_steps is not None
                        and self.global_step > self.cfg.trainer.max_training_steps
                    ):
                        logger.info(
                            f"Reached max_training_steps={self.cfg.trainer.max_training_steps}, stopping early."
                        )
                        stop_training = True
                        break

                    del training_input, generator_output

                self._fire("on_epoch_end")

                if stop_training:
                    break
        finally:
            self._profiler_stop()

        pbar.close()
        if self.colocate_all:
            await self.inference_engine_client.sleep()

        # Decrement global step by 1 to stop at the last global step
        # We use the global step value in callbacks when training finishes,
        # as well as for a final checkpoint save
        self.global_step -= 1

        # Safety net: always save final checkpoint at end of training.
        # Skip if we already saved at the last step
        if self.cfg.trainer.ckpt_interval > 0 and not will_save_ckpts:
            with Timer("save_checkpoints", self.all_timings):
                ckpt_path = self.save_checkpoints()
                logger.info("Saved final checkpoint.")
            self._fire("on_save", ckpt_path=ckpt_path)
        if self.cfg.trainer.hf_save_interval > 0 and not hf_model_save:
            with Timer("save_hf_model", self.all_timings):
                self.save_models()
                logger.info("Saved final model.")
        if self._vllm_metrics_scraper is not None:
            await self._vllm_metrics_scraper.aclose()

        if self._ray_gpu_monitor is not None:
            self._ray_gpu_monitor.stop()

        self._fire("on_train_end")
        self.tracker.finish()
        logger.info("Training done!")

    def flush_pending_metrics(self):
        """Best-effort flush of metrics accumulated for the in-flight step.

        Idempotent: the accumulators are cleared after the flush attempt, so a
        second call is a no-op. Never raises.
        """
        if not self.all_metrics and not self.all_timings:
            return
        log_payload = {
            **self.all_metrics,
            **{f"timing/{k}": v for k, v in self.all_timings.items()},
        }
        try:
            self.tracker.log(log_payload, step=self.global_step, commit=True)
        except Exception as e:
            logger.warning(f"Failed to flush pending metrics at step {self.global_step}: {e}")
        self.all_metrics = {}
        self.all_timings = {}

    def _remove_tail_data(self, entries: List[Any]) -> List[Any]:
        """Remove tail data to have even shards in terms of *effective* samples.

        Each prompt produces `n_samples_per_prompt` samples. For data-parallel
        training we care that the total number of samples is nicely splittable
        across the (combined) data-parallel size of all enabled models.
        """
        lcm_dp_size = self.dispatch.get_lcm_dp_size()

        n_samples_per_prompt = self.cfg.generator.n_samples_per_prompt

        # We want the largest m <= len(entries) such that:
        #   (m * n_samples_per_prompt) % lcm_dp_size == 0
        #
        # Let g = gcd(lcm_dp_size, n_samples_per_prompt). Then this is equivalent
        # to requiring m to be a multiple of (lcm_dp_size / g).
        stride = lcm_dp_size // math.gcd(lcm_dp_size, n_samples_per_prompt)
        if stride <= 1:
            # Every prompt count is valid, keep all entries.
            return entries

        kept_prompts = (len(entries) // stride) * stride
        return entries[:kept_prompts]

    def build_models(self, PolicyWorker, CriticWorker, RefWorker):
        """
        Initialize the actors for training, and handle colocation logic
        """
        cfg = self.cfg
        pg = None

        use_ref_model = cfg.trainer.algorithm.use_kl_loss or cfg.trainer.algorithm.use_kl_in_reward

        if cfg.trainer.placement.colocate_all:
            num_policy_gpus = cfg.trainer.placement.policy_num_gpus_per_node * cfg.trainer.placement.policy_num_nodes
            num_critic_gpus = cfg.trainer.placement.critic_num_gpus_per_node * cfg.trainer.placement.critic_num_nodes
            num_ref_gpus = cfg.trainer.placement.ref_num_gpus_per_node * cfg.trainer.placement.ref_num_nodes
            ie_cfg = cfg.generator.inference_engine
            num_rollout_gpus = (
                ie_cfg.num_engines
                * ie_cfg.tensor_parallel_size
                * ie_cfg.pipeline_parallel_size
                * ie_cfg.data_parallel_size
            )
            assert (
                num_policy_gpus == num_rollout_gpus
            ), "num_policy_gpus and num_rollout_gpus must be the same when colocating all models"
            pg = self.colocate_pg

            policy_model = PPORayActorGroup(
                cfg.trainer,
                cfg.trainer.placement.policy_num_nodes,
                cfg.trainer.placement.policy_num_gpus_per_node,
                PolicyWorker,
                pg=pg,
                num_gpus_per_actor=0.2 if pg else 1,
                colocate_all=True,
                sequence_parallel_size=cfg.trainer.policy.sequence_parallel_size,
                record_memory=cfg.trainer.policy.record_memory,
            )
            if use_ref_model:
                assert (
                    num_policy_gpus == num_ref_gpus
                ), "num_policy_gpus and num_ref_gpus must be the same when colocating policy and ref model"
                ref_model = PPORayActorGroup(
                    cfg.trainer,
                    cfg.trainer.placement.ref_num_nodes,
                    cfg.trainer.placement.ref_num_gpus_per_node,
                    RefWorker,
                    pg=pg,
                    num_gpus_per_actor=0.2 if pg else 1,
                    colocate_all=True,
                    sequence_parallel_size=cfg.trainer.ref.sequence_parallel_size,
                )
            else:
                ref_model = None

            if cfg.trainer.critic.model.path:
                assert (
                    num_policy_gpus == num_critic_gpus
                ), "num_policy_gpus and num_critic_gpus must be the same when colocating policy and critic model"
                critic_model = PPORayActorGroup(
                    cfg.trainer,
                    cfg.trainer.placement.critic_num_nodes,
                    cfg.trainer.placement.critic_num_gpus_per_node,
                    CriticWorker,
                    pg=pg,
                    num_gpus_per_actor=0.2,
                    colocate_all=True,
                    sequence_parallel_size=cfg.trainer.critic.sequence_parallel_size,
                )
            else:
                critic_model = None

        else:
            if cfg.trainer.placement.colocate_policy_ref and use_ref_model:
                assert (
                    cfg.trainer.placement.policy_num_nodes == cfg.trainer.placement.ref_num_nodes
                    and cfg.trainer.placement.policy_num_gpus_per_node == cfg.trainer.placement.ref_num_gpus_per_node
                ), "num_nodes and num_gpus_per_node must be the same when colocate policy and ref model."

                bundles = [
                    {
                        "GPU": cfg.trainer.placement.policy_num_gpus_per_node,
                        "CPU": cfg.trainer.placement.policy_num_gpus_per_node,
                    }
                    for _ in range(cfg.trainer.placement.policy_num_nodes)
                ]
                raw_pg = placement_group(bundles, strategy="PACK")
                get_ray_pg_ready_with_timeout(raw_pg, timeout=SKYRL_RAY_PG_TIMEOUT_IN_S)
                pg = ResolvedPlacementGroup(raw_pg)

            policy_model = PPORayActorGroup(
                cfg.trainer,
                cfg.trainer.placement.policy_num_nodes,
                cfg.trainer.placement.policy_num_gpus_per_node,
                PolicyWorker,
                pg=pg,
                num_gpus_per_actor=0.75 if pg else 1,
                colocate_all=False,
                sequence_parallel_size=cfg.trainer.policy.sequence_parallel_size,
            )
            if use_ref_model:
                ref_model = PPORayActorGroup(
                    cfg.trainer,
                    cfg.trainer.placement.ref_num_nodes,
                    cfg.trainer.placement.ref_num_gpus_per_node,
                    RefWorker,
                    pg=pg,
                    num_gpus_per_actor=0.25 if pg else 1,
                    colocate_all=False,
                    sequence_parallel_size=cfg.trainer.ref.sequence_parallel_size,
                )
                if pg is not None:
                    # The shared policy/ref placement group `pg` is set only when colocate_policy_ref is enabled
                    logger.info(
                        "Colocating policy and ref on the same GPUs across "
                        f"{cfg.trainer.placement.policy_num_nodes} node(s)."
                    )
            else:
                ref_model = None

            if cfg.trainer.critic.model.path:
                critic_model = PPORayActorGroup(
                    cfg.trainer,
                    cfg.trainer.placement.critic_num_nodes,
                    cfg.trainer.placement.critic_num_gpus_per_node,
                    CriticWorker,
                    num_gpus_per_actor=1,
                    colocate_all=False,
                    sequence_parallel_size=cfg.trainer.critic.sequence_parallel_size,
                )
            else:
                critic_model = None

        policy_steps_per_train_batch = (
            cfg.trainer.train_batch_size // cfg.trainer.policy_mini_batch_size * cfg.trainer.update_epochs_per_batch
        )
        critic_steps_per_train_batch = 0
        if cfg.trainer.critic.model.path:
            critic_steps_per_train_batch = (
                cfg.trainer.train_batch_size // cfg.trainer.critic_mini_batch_size * cfg.trainer.update_epochs_per_batch
            )
        policy_num_training_steps = (
            self.total_training_steps * policy_steps_per_train_batch if self.total_training_steps is not None else None
        )
        critic_num_training_steps = (
            self.total_training_steps * critic_steps_per_train_batch if self.total_training_steps is not None else None
        )
        if not cfg.trainer.placement.colocate_all:
            refs = []
            if ref_model is not None:
                refs.extend(ref_model.async_init_model(cfg.trainer.ref.model.path))
            refs.extend(
                policy_model.async_init_model(
                    cfg.trainer.policy.model.path,
                    num_training_steps=policy_num_training_steps,
                )
            )
            if cfg.trainer.critic.model.path:
                refs.extend(
                    critic_model.async_init_model(
                        cfg.trainer.critic.model.path,
                        num_training_steps=critic_num_training_steps,
                    )
                )
            ray.get(refs)
            ray.get(policy_model.async_run_ray_method("pass_through", "_set_pad_token_id", self.tokenizer.pad_token_id))
        else:
            if ref_model is not None:
                ray.get(ref_model.async_init_model(cfg.trainer.ref.model.path))
                ref_model.offload_to_cpu()
            ray.get(
                policy_model.async_init_model(
                    cfg.trainer.policy.model.path,
                    num_training_steps=policy_num_training_steps,
                )
            )
            ray.get(policy_model.async_run_ray_method("pass_through", "_set_pad_token_id", self.tokenizer.pad_token_id))
            policy_model.offload_to_cpu()
            if cfg.trainer.critic.model.path:
                ray.get(
                    critic_model.async_init_model(
                        cfg.trainer.critic.model.path,
                        num_training_steps=critic_num_training_steps,
                    )
                )
                critic_model.offload_to_cpu()

        self.policy_model: PPORayActorGroup = policy_model
        self.critic_model: Optional[PPORayActorGroup] = critic_model
        self.ref_model: Optional[PPORayActorGroup] = ref_model

        # Create unified dispatch that manages all actor groups
        self.dispatch = WorkerDispatch(
            cfg=self.cfg,
            policy_actor_group=policy_model,
            critic_actor_group=critic_model,
            ref_actor_group=ref_model,
            inference_engine_client=self.inference_engine_client,
        )

        # Mark all models as offloaded if colocate_all (they were offloaded above)
        if self.colocate_all:
            self.dispatch.mark_all_offloaded()

        logger.info("init policy/ref/critic models done")

    def init_weight_sync_state(self):
        """
        Setup the connection between policy model and inference engine for weight syncing.
        """
        self.dispatch.init_weight_sync_state(self.inference_engine_client)
        logger.info("Initialized weight sync state for policy model and inference engines.")

    def convert_to_training_input(self, generator_output: GeneratorOutput, uids: List[str]) -> TrainingInputBatch:
        """Converts lists to a padded batch of tensors for training

        Args:
            generator_output (GeneratorOutput): Generated rollouts and associated data.
            uids (List[str]): List of prompt-unique identifiers for each generator ouput in the same
                order as `generator_output`. Used to identify which prompt each generated rollout belongs to.
        Returns:
            training_input (TrainingInputBatch): Padded batch of tensors for training. It preserves the
                order of `generator_output` and hence `uids`.
        """
        # 1. Extract generator output fields.
        prompt_ids: List[List[int]] = generator_output["prompt_token_ids"]
        response_ids: List[List[int]] = generator_output["response_ids"]
        rewards: List[List[float]] = generator_output["rewards"]
        loss_masks: List[List[int]] = generator_output["loss_masks"]

        logprobs: Optional[List[List[float]]] = generator_output.get("rollout_logprobs", None)
        rollout_expert_indices: Optional[List[List[List[List[int]]]]] = generator_output.get(
            "rollout_expert_indices", None
        )

        pixel_values = generator_output.get("pixel_values", None)
        image_grid_thw = generator_output.get("image_grid_thw", None)
        if pixel_values is not None:
            assert (
                pixel_values is not None and image_grid_thw is not None
            ), "Both pixel_values and image_grid_thw must exist for multi-modal inputs"
            assert len(pixel_values) == len(
                image_grid_thw
            ), "Number of pixel values should match number of image grid thw"
            pixel_values = TensorList(pixel_values)
            image_grid_thw = TensorList(image_grid_thw)

        # 2. Convert to tensors.
        (
            sequences_tensor,
            attention_masks_tensor,
            response_masks_tensor,
            rewards_tensor,
            loss_masks_tensor,
            rollout_logprobs_tensor,
            rollout_expert_indices_tensor,
        ) = convert_prompts_responses_to_batch_tensors(
            self.tokenizer,
            prompt_ids,
            response_ids,
            rewards,
            loss_masks,
            logprobs,
            rollout_expert_indices,
            max_seq_len=self.cfg.trainer.algorithm.max_seq_len,
        )

        # sanity check for off_policy_correction
        off_policy_correction = self.cfg.trainer.algorithm.off_policy_correction
        tis_ratio_type = off_policy_correction.tis_ratio_type
        sequence_mask_metric = off_policy_correction.sequence_mask_metric
        if tis_ratio_type is not None or sequence_mask_metric is not None:
            assert (
                rollout_logprobs_tensor is not None
            ), "expected non-null rollout logprobs tensor when off_policy_correction is enabled"
            assert rollout_logprobs_tensor.shape == loss_masks_tensor.shape, "Logprobs should look like responses"

        # 3. Create training input batch.
        training_input = TrainingInputBatch(
            {
                "sequences": sequences_tensor,  # Full trajectories (padded and concatenated prompts and responses)
                "attention_mask": attention_masks_tensor,
                "response_mask": response_masks_tensor,
                "rewards": rewards_tensor,
                "loss_mask": loss_masks_tensor,
                "rollout_logprobs": rollout_logprobs_tensor,
                "rollout_expert_indices": rollout_expert_indices_tensor,
                "pixel_values": pixel_values,
                "image_grid_thw": image_grid_thw,
            },
        )
        training_input.metadata = {"uids": uids}
        if generator_output.get("is_last_step", None) is not None:
            training_input.metadata["is_last_step"] = generator_output["is_last_step"]

        # 4. Compute mini-batch boundaries for train_critic_and_policy(). It excludes the ones
        # we will add in pad_training_input_batch().
        train_batch_size = self.cfg.trainer.train_batch_size
        n_samples_per_prompt = self.cfg.generator.n_samples_per_prompt
        is_stepwise = self.cfg.generator.step_wise_trajectories
        training_input.metadata["policy_mini_batch_boundaries"] = compute_prompt_mini_batch_boundaries(
            uids, self.cfg.trainer.policy_mini_batch_size, train_batch_size, is_stepwise, n_samples_per_prompt
        )
        # Per-prompt boundaries (used by the `prompt_mean` loss reduction). Policy-only,
        # since advantage normalization only applies to the policy.
        training_input.metadata["policy_prompt_boundaries"] = compute_prompt_boundaries(uids)
        if self.cfg.trainer.critic.model.path is not None:
            training_input.metadata["critic_mini_batch_boundaries"] = compute_prompt_mini_batch_boundaries(
                uids, self.cfg.trainer.critic_mini_batch_size, train_batch_size, is_stepwise, n_samples_per_prompt
            )

        # 5. Record metadata and metrics.
        training_input.metadata["response_length"] = response_masks_tensor.shape[1]
        batch_num_seq, batch_padded_seq_len = sequences_tensor.shape
        logger.info(f"batch_num_seq: {batch_num_seq}, batch_padded_seq_len: {batch_padded_seq_len}")
        self.all_metrics.update(
            {
                "generate/batch_num_seq": batch_num_seq,
                "generate/batch_padded_seq_len": batch_padded_seq_len,
            }
        )
        training_input.metadata["avg_response_length"] = sum(
            len(sample_response_ids) for sample_response_ids in response_ids
        ) / len(response_ids)

        # 6. Pad the batch, only needed for step-wise training's `fwd_logprobs_values_reward()`.
        logger.info(f"Number of sequences before padding: {len(training_input['sequences'])}")
        dp_size = self.dispatch.get_lcm_dp_size()
        pad_size = math.ceil(training_input.batch_size / dp_size) * dp_size - training_input.batch_size
        training_input = pad_training_input_batch(training_input, pad_size)
        logger.info(f"Number of sequences after padding: {len(training_input['sequences'])}")

        return training_input

    @torch.no_grad()
    async def generate(
        self,
        input_batch: GeneratorInput,
    ) -> GeneratorOutput:
        """
        Generate rollouts.

        If colocate_all is enabled:
        - before calling this method, the policy model should be on CPU and inference engine should
            be awake (i.e. on GPU).
        - after calling this method, the same model placement still holds.
        """
        # NOTE: we assume that .generate returns samples in the same order as passed in
        generator_output: GeneratorOutput = await self.generator.generate(input_batch)

        # add rollout metrics to self.all_metrics
        if generator_output["rollout_metrics"] is not None:
            self.all_metrics.update(generator_output["rollout_metrics"])
        generator_output.pop("rollout_metrics", None)

        validate_generator_output(
            len(input_batch["prompts"]),
            generator_output,
            step_wise=self.cfg.generator.step_wise_trajectories,
        )

        return generator_output

    @torch.no_grad()
    def postprocess_generator_output(
        self,
        generator_output: GeneratorOutput,
        uids: List[str],
        metrics_generator_output: Optional[GeneratorOutput] = None,
        metrics_uids: Optional[List[str]] = None,
    ) -> Tuple[GeneratorOutput, List[str]]:
        """
        Converts to per token rewards and computes pass@N.

        For step-wise training with ``merge_stepwise_output=true``, also collapses
        consecutive turns sharing a common prefix into a single sequence; ``uids``
        is shortened to match.

        In the future algorithm specific reward or loss mask post processing should be done here.

        Reward metrics are computed over ``metrics_generator_output`` / ``metrics_uids`` when provided
        (a superset of the trained output -- e.g. sample_full_batch passes the dropped groups so metrics
        stay comparable), otherwise over ``generator_output`` / ``uids``. The per-token / loss-mask
        conversion always applies to ``generator_output`` only.

        Returns:
            (generator_output, uids) — uids may be shorter than the input when merging.
        """
        metrics_output = metrics_generator_output if metrics_generator_output is not None else generator_output
        metrics_output_uids = metrics_uids if metrics_uids is not None else uids
        generator_output_for_metrics = metrics_output
        uids_for_metrics = metrics_output_uids
        if self.cfg.generator.step_wise_trajectories:
            generator_output_for_metrics = defaultdict(list)
            for key in metrics_output:
                if isinstance(metrics_output[key], list):
                    generator_output_for_metrics[key] = [
                        metrics_output[key][i]
                        for i in range(len(metrics_output[key]))
                        if metrics_output["is_last_step"][i]
                    ]
            uids_for_metrics = [
                uid for uid, is_last_step in zip(metrics_output_uids, metrics_output["is_last_step"]) if is_last_step
            ]

        # only use `generator_output_for_metrics` for metrics calculation
        # For step-wise training, we only calculate metrics for the last step of each trajectory
        overall_metrics = get_metrics_from_generator_output(
            generator_output_for_metrics,
            uids_for_metrics,
        )

        # Prefix-aware merging of step-wise turns.
        if self.cfg.generator.merge_stepwise_output:
            assert self.cfg.generator.step_wise_trajectories, "merge_stepwise_output requires step-wise training"
            num_seq_before_merge = len(generator_output["response_ids"])
            generator_output = merge_stepwise_output(generator_output)
            num_seq_after_merge = len(generator_output["response_ids"])
            logger.info(f"Merged step wise: {num_seq_before_merge} sequences -> {num_seq_after_merge} sequences")
            self.all_metrics.update(
                {
                    "generate/num_seq_before_merge": num_seq_before_merge,
                    "generate/num_seq_after_merge": num_seq_after_merge,
                }
            )
            uids = [tid.instance_id for tid in generator_output["trajectory_ids"]]

        # these use the full generator output
        rewards: Union[List[float], List[List[float]]] = generator_output["rewards"]
        responses: List[List[int]] = generator_output["response_ids"]
        per_token_rewards: List[List[float]] = []

        # Check if rewards are already token-level (List[List[float]]) or response-level (List[float])
        if rewards and isinstance(rewards[0], list):
            # Token-level rewards: rewards is List[List[float]]
            per_token_rewards = rewards
        else:
            if self.cfg.trainer.algorithm.zero_variance_filter:
                kept_indices_set = set(
                    zero_variance_filter(
                        rewards,
                        uids,
                        loss_masks=generator_output["loss_masks"],
                        tol=self.cfg.trainer.algorithm.zero_variance_filter_tol,
                    )
                )
                num_groups = len(set(uids))
                num_kept_groups = len({uids[i] for i in kept_indices_set})
                self.all_metrics["reward/num_zero_variance_filtered"] = num_groups - num_kept_groups
                generator_output["loss_masks"] = [
                    [0] * len(mask) if i not in kept_indices_set else mask
                    for i, mask in enumerate(generator_output["loss_masks"])
                ]
            # Response-level rewards: rewards is List[float], convert to per-token rewards
            for reward, response in zip(rewards, responses):
                per_token_reward = [0.0] * len(response)
                per_token_reward[-1] = float(reward)
                per_token_rewards.append(per_token_reward)

        n_samples_per_prompt = self.cfg.generator.n_samples_per_prompt

        reward_metrics = {
            f"reward/avg_pass_at_{n_samples_per_prompt}": overall_metrics["pass_at_n"],
            "reward/avg_raw_reward": overall_metrics["avg_score"],
            "reward/mean_positive_reward": overall_metrics["mean_positive_reward"],
        }
        self.all_metrics.update(reward_metrics)
        logger.info(
            f"reward/avg_pass_at_{n_samples_per_prompt}: {overall_metrics['pass_at_n']}, reward/avg_raw_reward: {overall_metrics['avg_score']}, reward/mean_positive_reward: {overall_metrics['mean_positive_reward']}"
        )
        # re-assign reward but now it's per token rewards
        generator_output["rewards"] = per_token_rewards
        return generator_output, uids

    @torch.no_grad()
    def compute_advantages_and_returns(self, data: TrainingInputBatch) -> TrainingInputBatch:
        """Calculate advantages and returns for the data batch.

        Expects:
            - `["sequences"]`: Integer[torch.Tensor, "batch_size seqlen"]
            - `["response_mask"]`: Integer[torch.Tensor, "batch_size seqlen"]
            - `["loss_mask"]`: Integer[torch.Tensor, "batch_size seqlen"]
            - `["values"]`: Float[torch.Tensor, "batch_size seqlen"]
            - `["rewards"]`: Float[torch.Tensor, "batch_size seqlen"]
            - `.metadata["uids"]`: List[str]
            - `.metadata["is_last_step"]`: List[bool] for step-wise training

        Adds:
            - `["advantages"]`: Float[torch.Tensor, "batch_size seqlen"]
            - `["returns"]`: Float[torch.Tensor, "batch_size seqlen"]
        """
        token_level_rewards = data["rewards"]

        if self.cfg.generator.step_wise_trajectories:
            is_last_step = torch.tensor(data.metadata["is_last_step"], dtype=torch.bool)
            index = np.array(data.metadata["uids"])
            values = data["values"]
            # Step-wise only supports outcome-based estimators (GRPO, RLOO, MAXRL); ensured by `validate_cfg`.
            # We use the last step of each trajectory to compute advantages and broadcast them to
            # all steps of that trajectory, so we ignore per-step rewards in step-wise training.
            # We pass an all-ones mask here so the estimator returns the scalar advantage at every
            # position. The real per-step `response_mask` is re-applied on broadcast below.
            # Shapes:
            #   traj_ids, (batch_size,):         trajectory id per step (cumsum of shifted is_last_step)
            #   last_step_advantages/returns,
            #       (num_traj, seqlen):          scalar advantage/return per trajectory at every position
            #   last_step_advantages/returns[traj_ids],
            #       (batch_size, seqlen):        broadcast to every step of the owning trajectory
            #   response_mask_float,
            #       (batch_size, seqlen):        per-step response mask
            last_step_response_mask = data["response_mask"][is_last_step]
            last_step_advantages, last_step_returns = ppo_utils.compute_advantages_and_returns(
                token_level_rewards=token_level_rewards[is_last_step],
                response_mask=torch.ones_like(last_step_response_mask, dtype=torch.float),
                index=index[is_last_step.cpu().numpy()],
                adv_estimator=self.cfg.trainer.algorithm.advantage_estimator,
                values=values[is_last_step] if values is not None else None,
                config=self.cfg.trainer.algorithm,
                gamma=self.cfg.trainer.algorithm.gamma,
                lambd=self.cfg.trainer.algorithm.lambd,
                grpo_norm_by_std=self.cfg.trainer.algorithm.grpo_norm_by_std,
            )
            traj_ids = (
                torch.cat([torch.tensor([False], device=is_last_step.device), is_last_step[:-1]]).int().cumsum(dim=0)
            )
            num_traj = traj_ids[-1].item() + 1
            assert num_traj == len(
                last_step_advantages
            ), f"num_traj {num_traj} doesn't match the number of trajectories as given by `is_last_step` {len(last_step_advantages)}. The `is_last_step` tensor is likely malformed"
            response_mask_float = data["response_mask"].to(last_step_advantages.dtype)
            advantages = last_step_advantages[traj_ids] * response_mask_float
            returns = last_step_returns[traj_ids] * response_mask_float
        else:
            advantages, returns = ppo_utils.compute_advantages_and_returns(
                token_level_rewards=token_level_rewards,
                response_mask=data["response_mask"],
                index=data.metadata["uids"],
                adv_estimator=self.cfg.trainer.algorithm.advantage_estimator,
                config=self.cfg.trainer.algorithm,
                values=data["values"],
                gamma=self.cfg.trainer.algorithm.gamma,
                lambd=self.cfg.trainer.algorithm.lambd,
                grpo_norm_by_std=self.cfg.trainer.algorithm.grpo_norm_by_std,
            )
        data["returns"] = returns
        data["advantages"] = advantages

        # remove padding while calculating metrics
        pad_size = data.metadata.get("pad_size", 0)
        num_samples = len(token_level_rewards)

        return_sums = token_level_rewards.sum(dim=-1)[: num_samples - pad_size]
        if self.cfg.generator.step_wise_trajectories:
            avg_rewards: float = return_sums[is_last_step[: num_samples - pad_size]].mean().item()
        else:
            avg_rewards: float = return_sums.mean().item()

        avg_response_length = data.metadata["avg_response_length"]
        data = data.to("cpu")

        valid_advantages = torch.masked_select(
            data["advantages"][: num_samples - pad_size, ...], data["response_mask"][: num_samples - pad_size].bool()
        )
        avg_advantages: float = valid_advantages.mean().item()
        avg_advantages_abs: float = valid_advantages.abs().mean().item()

        if "metrics" not in data.metadata:
            data.metadata["metrics"] = {}
        data.metadata["metrics"].update(
            {
                "avg_final_rewards": avg_rewards,
                "avg_response_length": avg_response_length,
                "avg_advantages": avg_advantages,
                "avg_advantages_abs": avg_advantages_abs,
            }
        )

        logger.info(f"avg_final_rewards: {avg_rewards}, avg_response_length: {avg_response_length}")
        self.all_metrics.update(
            {
                "loss/avg_final_rewards": avg_rewards,
                "loss/avg_raw_advantages": avg_advantages,
                "loss/avg_raw_advantages_abs": avg_advantages_abs,
            }
        )
        return data

    def dump_data(self, data: TrainingInputBatch, file_name: str):
        """
        Dump data to pickle file
        """
        data_save_dir = Path(self.cfg.trainer.export_path) / "dumped_data"
        data_save_dir.mkdir(parents=True, exist_ok=True)
        data.save(data_save_dir / f"{file_name}.pkl")

    def _execute_forward_pass(
        self,
        model: str,
        data_fwd_pass: TrainingInputBatch,
        key: str,
        mini_batch_boundaries: Optional[List[Tuple[int, int]]],
    ) -> torch.Tensor:
        """Executes forward pass that produces to produce the "old" logprobs/values.

        With ``trainer.recompute_old_logprobs_per_minibatch`` set (and mini-batch boundaries
        available), the forward is run per mini-batch — matching the mini-batch + DP partition
        that the training step (``_execute_training_step`` / ``stage_data``) will use — so the
        microbatch packing, and therefore the resulting logprobs/values, are identical to what
        ``forward_backward`` recomputes. This makes the PPO ratio (and critic value clipping)
        exact at the first inner step. Otherwise a single full-batch forward is run.

        Per-sample outputs are concatenated in mini-batch order, which matches the global sample
        order (mini-batches are contiguous and in order), so the result aligns with the full-batch
        forward's ordering. Tensorizing the combined ``loss_fn_outputs`` once pads uniformly.
        """
        if self.cfg.trainer.recompute_old_logprobs_per_minibatch and mini_batch_boundaries:
            # Pre-stage all per-DP mini-batch chunks once (same as the training step), so chunk
            # serialization is amortized off the dispatch critical path across mini-batches. The
            # staged chunks use the same partition as `_execute_training_step`'s `stage_data`, so
            # the packing — and resulting logprobs/values — match what forward_backward recomputes.
            all_chunk_refs = self.dispatch.stage_data(model, data_fwd_pass, mini_batch_boundaries)
            combined_outputs: List[Dict[str, Any]] = []
            for chunk_refs in all_chunk_refs:
                mb_output = self.dispatch.forward_from_staged(model, chunk_refs)
                combined_outputs.extend(mb_output.loss_fn_outputs)
            return loss_fn_outputs_to_tensor(combined_outputs, key=key)

        output = self.dispatch.forward(model, data_fwd_pass)
        return loss_fn_outputs_to_tensor(output.loss_fn_outputs, key=key)

    def _skip_policy_forward(self, training_input: TrainingInputBatch) -> bool:
        """Whether the policy forward pass producing the "old" logprobs can be skipped.

        Safe only when the loss optimizes against rollout logprobs and nothing else reads the
        old logprobs: rollout logprobs are present (these losses fall back to old logprobs
        without them), the KL reward penalty is off, and off-policy correction is disabled.
        """
        algorithm = self.cfg.trainer.algorithm
        return (
            algorithm.policy_loss_type in LOSSES_WITHOUT_OLD_LOGPROBS
            and training_input.get("rollout_logprobs", None) is not None
            and not algorithm.use_kl_in_reward
            and not off_policy_correction_enabled(algorithm.off_policy_correction)
        )

    @torch.no_grad()
    def fwd_logprobs_values_reward(
        self,
        training_input: TrainingInputBatch,
    ):
        """
        Calculate values from the critic, log probs from the policy and ref model.

        Dispatch handles offload/backload automatically for all colocation configurations.

        Expects:
            - `["sequences"]`: Integer[torch.Tensor, "batch_size seqlen"]
            - `["attention_mask"]`: Integer[torch.Tensor, "batch_size seqlen"]
            - `.metadata["response_length"]`: Int

        Adds:
            - `["base_action_log_probs"]`: Float[torch.Tensor, "batch_size seqlen"]
            - `["action_log_probs"]`: Float[torch.Tensor, "batch_size seqlen"]
            - `["values"]`: Float[torch.Tensor, "batch_size seqlen"]
        """
        fwd_keys = ["sequences", "attention_mask"]
        if training_input.get("rollout_expert_indices") is not None:
            fwd_keys.append("rollout_expert_indices")
        if training_input.get("pixel_values") is not None:
            fwd_keys.append("pixel_values")
        if training_input.get("image_grid_thw") is not None:
            fwd_keys.append("image_grid_thw")
        data_fwd_pass = training_input.select(keys=fwd_keys, metadata_keys=["response_length"])

        values = None
        base_log_probs = None
        action_log_probs = None

        # Critic forward (dispatch handles offload/backload automatically)
        if self.has_critic:
            values = self._execute_forward_pass(
                "critic",
                data_fwd_pass,
                key="values",
                mini_batch_boundaries=training_input.metadata.get("critic_mini_batch_boundaries"),
            )

        # Ref forward. The ref model is not trained, so there is no forward_backward to match
        # its packing against -> always a single full-batch forward (boundaries=None).
        if self.ref_model is not None:
            base_log_probs = self._execute_forward_pass(
                "ref", data_fwd_pass, key="logprobs", mini_batch_boundaries=None
            )
            self.dispatch.empty_cache("ref")

        # Policy forward. Skipped for losses that optimize against rollout logprobs (see
        # `_skip_policy_forward`), where the resulting logprobs are never read.
        if self._skip_policy_forward(training_input):
            action_log_probs = None
        else:
            action_log_probs = self._execute_forward_pass(
                "policy",
                data_fwd_pass,
                key="logprobs",
                mini_batch_boundaries=training_input.metadata.get("policy_mini_batch_boundaries"),
            )

        # Empty cache after all forward passes
        self.dispatch.empty_cache()

        sequences_all: torch.Tensor = training_input["sequences"]
        # NOTE (sumanthrh): The slicing is needed to make sure that the batch dimension doesn't change for the tensordict.
        base_log_probs = base_log_probs[: len(sequences_all)] if base_log_probs is not None else None
        action_log_probs = action_log_probs[: len(sequences_all)] if action_log_probs is not None else None
        values = values[: len(sequences_all)] if values is not None else None

        training_input["base_action_log_probs"] = base_log_probs
        training_input["action_log_probs"] = action_log_probs
        training_input["values"] = values

        if training_input.get("rollout_logprobs", None) is not None and action_log_probs is not None:
            # Abs diff between rollout and forward-pass logprobs, over response tokens. When the
            # forward pass is skipped, the worker's `minibatch_rollout_logprobs_abs_diff_*` is used.
            logprobs_diff = (
                training_input["rollout_logprobs"][training_input["loss_mask"] > 0]
                - action_log_probs[training_input["loss_mask"] > 0]
            ).abs()

            # Guard: a batch with no trainable response tokens (loss_mask all zero, e.g. every
            # response dropped by overlong filtering) leaves logprobs_diff empty, and .max()/.min()
            # on a 0-element tensor raises. Skip the diagnostic metrics in that case.
            if logprobs_diff.numel() > 0:
                logprobs_diff_max = logprobs_diff.max().item()
                logprobs_diff_min = logprobs_diff.min().item()
                logprobs_diff_mean = logprobs_diff.mean().item()
                logprobs_diff_std = logprobs_diff.std().item()
                self.all_metrics.update(
                    {
                        "policy/rollout_train_logprobs_abs_diff_max": logprobs_diff_max,
                        "policy/rollout_train_logprobs_abs_diff_min": logprobs_diff_min,
                        "policy/rollout_train_logprobs_abs_diff_mean": logprobs_diff_mean,
                        "policy/rollout_train_logprobs_abs_diff_std": logprobs_diff_std,
                    }
                )
        return training_input

    def apply_reward_kl_penalty(
        self,
        data: TrainingInputBatch,
    ) -> TrainingInputBatch:
        """Applies a penalty for KL divergence between the policy log probs and the base model log probs to the rewards."""
        loss_masks_all: torch.Tensor = data["loss_mask"]
        rewards: torch.Tensor = data["rewards"]
        base_action_log_probs: torch.Tensor = data["base_action_log_probs"]
        action_log_probs: torch.Tensor = data["action_log_probs"]

        # single batched computation
        with torch.no_grad():
            kl: Float[torch.Tensor, "batch_size seqlen"] = compute_approx_kl(  # type: ignore
                action_log_probs,
                base_action_log_probs,
                loss_mask=loss_masks_all,
                kl_estimator_type=self.cfg.trainer.algorithm.kl_estimator_type,
            )
        kl_max: Float[torch.Tensor, "batch_size"] = torch.max(kl.abs(), dim=-1)[0]  # noqa: F821
        kl_mean: Float[torch.Tensor, "batch_size"] = masked_mean(kl, loss_masks_all, dim=-1)  # noqa: F821

        # NOTE (erictang000): only supporting custom rewards currently
        kl_loss_coef = (
            self.reward_kl_controller.value
            if self.reward_kl_controller is not None
            else self.cfg.trainer.algorithm.kl_loss_coef
        )
        rewards = rewards - kl * max(0, kl_loss_coef)
        data["rewards"] = rewards

        avg_kl: float = kl_mean.mean().item()
        avg_kl_max: float = kl_max.mean().item()

        # update the kl controller
        if self.reward_kl_controller is not None:
            self.reward_kl_controller.update(current=avg_kl, n_steps=kl.shape[0])  # n_steps is just the batch size
        if "metrics" not in data.metadata:
            data.metadata["metrics"] = {}

        data.metadata["metrics"].update(
            {
                "avg_kl": avg_kl,
                "avg_kl_max": avg_kl_max,
                "kl_loss_coef": kl_loss_coef,
            }
        )

        self.all_metrics.update(
            {
                "loss/avg_kl": avg_kl,
                "loss/avg_kl_max": avg_kl_max,
                "loss/kl_loss_coef": kl_loss_coef,
            }
        )

        return data

    @torch.no_grad()
    def _normalize_advantages(
        self,
        data: TrainingInputBatch,
        mini_batch_boundaries: List[Tuple[int, int]],
        prompt_boundaries: Optional[List[Tuple[int, int]]] = None,
    ) -> TrainingInputBatch:
        advantages = data["advantages"]
        response_mask = data["response_mask"]

        # Step 1: Z-score normalization (if enabled)
        if self.cfg.trainer.algorithm.advantage_batch_normalize:
            num_actions = response_mask.sum()
            mean = advantages.mean()
            std = ((advantages - mean).pow(2) * response_mask).sum()
            rstd = (std / num_actions).clamp(min=1e-8).rsqrt()
            data["advantages"] = (advantages - mean) * rstd

        # Step 2: Loss reduction normalization per mini-batch
        normalized_advantages = torch.zeros_like(advantages)
        for start_idx, end_idx in mini_batch_boundaries:
            mini_batch = data[start_idx:end_idx]
            # For prompt_mean, select the prompt boundaries falling within this mini-batch
            # and rebase them to mini-batch-relative indices.
            mb_prompt_boundaries = None
            if prompt_boundaries is not None:
                mb_prompt_boundaries = [
                    (p_start - start_idx, p_end - start_idx)
                    for p_start, p_end in prompt_boundaries
                    if start_idx <= p_start < end_idx
                ]
            normalized_advantages[start_idx:end_idx] = apply_loss_reduction_to_advantages_minibatch(
                advantages=mini_batch["advantages"],
                loss_mask=mini_batch["loss_mask"],
                loss_reduction=self.cfg.trainer.algorithm.loss_reduction,
                micro_batch_size=self.cfg.trainer.micro_train_batch_size_per_gpu,
                max_seq_len=self.cfg.trainer.algorithm.max_seq_len,
                prompt_boundaries=mb_prompt_boundaries,
            )

        data["advantages"] = normalized_advantages
        return data

    def _execute_training_step(self, model: str, data: TrainingInputBatch) -> Dict[str, float]:
        """
        Execute training step using forward_backward + optim_step.

        The trainer loops over epochs and mini-batches. Workers handle micro-batching
        internally for gradient accumulation (memory efficiency).

        All per-DP mini-batch chunks are pre-staged in the Ray object store before
        the training loop so serialization stays off the GPU critical path.

        Args:
            model: Model name ("policy" or "critic")
            data: Training data batch

        Returns:
            Dict of reduced metrics from training
        """
        boundaries = data.metadata[f"{model}_mini_batch_boundaries"]

        if model == "policy":
            # Normalize advantages for policy training; critic training does not need this
            prompt_boundaries = data.metadata.get("policy_prompt_boundaries")
            data = self._normalize_advantages(data, boundaries, prompt_boundaries)

        all_metrics: Dict[str, List[float]] = defaultdict(list)

        # Pre-stage all per-DP mini-batch chunks in the object store so that
        # serialization is fully off the critical path during training.
        all_chunk_refs = self.dispatch.stage_data(model, data, boundaries)

        # Training loop over epochs and mini-batches
        for _epoch in range(self.cfg.trainer.update_epochs_per_batch):
            for chunk_refs in all_chunk_refs:
                status = self.dispatch.forward_backward_from_staged(model, chunk_refs)
                for k, v in status.metrics.items():
                    all_metrics[k].append(v)

                # Optimizer step after each mini batch
                grad_norm = self.dispatch.optim_step(model)
                if grad_norm is not None:
                    all_metrics["grad_norm"].append(grad_norm)

        # Reduce metrics across all mini-batches and epochs
        reduced_metrics = reduce_metrics(all_metrics, sum_loss_metrics=False)
        finalize_minibatch_rollout_logprob_diff_std(reduced_metrics)
        return reduced_metrics

    def train_critic_and_policy(self, data: TrainingInputBatch):
        """
        Run the training step for the policy and critic models.

        Uses forward_backward + optim_step for both FSDP and Megatron strategies.
        """
        data.metadata["global_step"] = self.global_step
        critic_status = None

        # Unified training interface for both FSDP and Megatron
        if self.has_critic:
            with Timer("critic_train", self.all_timings):
                critic_status = self._execute_training_step("critic", data)
        with Timer("policy_train", self.all_timings):
            policy_status = self._execute_training_step("policy", data)

        # Update metrics
        if critic_status is not None:
            for k, v in critic_status.items():
                self.all_metrics.update({f"critic/{k}": v})

        for k, v in policy_status.items():
            self.all_metrics.update({f"policy/{k}": v})

        self.dispatch.empty_cache()

        return policy_status

    def handle_dynamic_sampling(
        self, generator_output: GeneratorOutput, uids: List[str]
    ) -> Tuple[GeneratorOutput, List[str], bool]:
        """
        Handle dynamic sampling for the current batch.

        Accumulates the generator output and UIDs across batches if we are sampling repeatedly
        and applies the dynamic sampling strategy (i.e. filter, replace) to the current batch.
        If we hit the limit of max sample batches, we raise an error.

        Args:
            generator_output: Current batch generator output
            uids: Current batch UIDs

        Returns:
            processed_output: Filtered generator output
            processed_uids: Filtered UIDs
            keep_sampling: Whether to keep sampling
        """
        # Prepare sampling configuration
        max_sample_batches = self.cfg.trainer.algorithm.dynamic_sampling.max_sample_batches
        dynamic_sampling_config = {
            "type": self.cfg.trainer.algorithm.dynamic_sampling.type,
            "max_sample_batches": max_sample_batches,
            "min_replace_ratio": self.cfg.trainer.algorithm.dynamic_sampling.min_replace_ratio,
            "train_batch_size": self.cfg.trainer.train_batch_size,
            "n_samples_per_prompt": self.cfg.generator.n_samples_per_prompt,
        }

        if self.dynamic_sampling_state is None:
            self.dynamic_sampling_state: DynamicSamplingState = {
                "sample_batch_count": 1,
            }
        else:
            self.dynamic_sampling_state["sample_batch_count"] += 1

        # Handle dynamic sampling using utilities
        processed_output, processed_uids, keep_sampling, updated_state = trainer_utils.handle_dynamic_sampling(
            generator_output, uids, dynamic_sampling_config, self.dynamic_sampling_state
        )

        # Check max resample limit, and if we hit it, raise an error
        if (
            keep_sampling
            and max_sample_batches > 0
            and self.dynamic_sampling_state["sample_batch_count"] >= max_sample_batches
        ):
            raise RuntimeError(
                f"Exiting training loop due to hitting dynamic sampling limit for "
                f"{self.cfg.trainer.algorithm.dynamic_sampling.type} strategy with "
                f"{self.cfg.trainer.algorithm.dynamic_sampling.max_sample_batches} max sample batches. "
                f"Please check your data difficulty distribution."
            )
        # Update state
        self.dynamic_sampling_state = updated_state

        if not keep_sampling:
            # Reset state when sampling is complete
            self.dynamic_sampling_state = None

        return processed_output, processed_uids, keep_sampling

    def _get_dp_group_models(self, rank: int, model_type: str = ""):
        model = getattr(self, model_type)
        return model._actor_handlers[rank]

    def _get_mesh_rank(self, rank: int, model_type: str = "") -> MeshRank:
        model: PPORayActorGroup = getattr(self, model_type)
        actor_info: ActorInfo = model.actor_infos[rank]
        return actor_info.rank

    def save_checkpoints(self) -> str:
        """
        Save the model, optimizer, and training states to disk. Returns the
        checkpoint folder path.

        Dispatch handles offload/backload automatically for all colocation configurations.
        """
        # Create global step folder structure
        global_step_folder = os.path.join(self.cfg.trainer.ckpt_path, f"global_step_{self.global_step}")
        policy_save_dir = os.path.join(global_step_folder, "policy")
        critic_save_dir = os.path.join(global_step_folder, "critic")

        io.makedirs(global_step_folder, exist_ok=True)

        # Save policy checkpoint (dispatch handles offload/backload)
        self.dispatch.save_checkpoint("policy", policy_save_dir, self.tokenizer)

        # Save critic checkpoint (if it exists)
        if self.has_critic:
            self.dispatch.save_checkpoint("critic", critic_save_dir, self.tokenizer)

        # Save dataloader state
        dataloader_save_path = os.path.join(global_step_folder, "data.pt")
        try:
            dataloader_state_dict = self.train_dataloader.state_dict()
            with io.open_file(dataloader_save_path, "wb") as f:
                torch.save(dataloader_state_dict, f)
            logger.info(f"Saved dataloader state to {dataloader_save_path}")
        except Exception as e:
            logger.warning(f"Failed to save dataloader state: {e}")

        # Save additional trainer state
        trainer_state = {
            "global_step": self.global_step,
            "config": asdict(self.cfg),
        }
        trainer_state_path = os.path.join(global_step_folder, "trainer_state.pt")
        with io.open_file(trainer_state_path, "wb") as f:
            torch.save(trainer_state, f)
        logger.info(f"Saved trainer state to {trainer_state_path}")

        # Atomic tracking - write this last after all saves succeed
        latest_checkpoint_file = os.path.join(self.cfg.trainer.ckpt_path, "latest_ckpt_global_step.txt")
        with io.open_file(latest_checkpoint_file, "w") as f:
            f.write(str(self.global_step))

        logger.info(f"Successfully saved checkpoint for global_step_{self.global_step} to: {global_step_folder}")

        # Clean up old checkpoints after successful save
        with Timer("cleanup_old_checkpoints", self.all_timings):
            self._cleanup_old_checkpoints()

        return global_step_folder

    def _cleanup_old_checkpoints(self):
        if not self._node_ids:
            self._node_ids = self.dispatch.get_node_ids()
        run_on_each_node(
            self._node_ids,
            cleanup_old_checkpoints,
            self.cfg.trainer.ckpt_path,
            self.cfg.trainer.max_ckpts_to_keep,
        )
        # run on driver as well
        # NOTE (sumanthrh): the function will get called twice on the node with driver process, but it's ok because it's idempotent
        cleanup_old_checkpoints(self.cfg.trainer.ckpt_path, self.cfg.trainer.max_ckpts_to_keep)

    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.

        If colocate_all is True, assumes that the policy model is currently on GPU.

        Returns:
            global_step: The global step to resume from.
            checkpoint_path: The path to the checkpoint.
        """
        checkpoint_path = None
        # Check if resumption is enabled
        if self.resume_mode == ResumeMode.NONE:
            logger.info("Checkpoint resumption disabled, starting training from scratch")
            return 0, None
        # first, let's get resume_path
        elif self.resume_mode == ResumeMode.LATEST:
            latest_checkpoint_file = os.path.join(self.cfg.trainer.ckpt_path, "latest_ckpt_global_step.txt")
            if not io.exists(latest_checkpoint_file):
                logger.info("No checkpoint found, starting training from scratch")
                return 0, None
            with io.open_file(latest_checkpoint_file, "r") as f:
                ckpt_iteration = int(f.read().strip())
            checkpoint_path = os.path.join(self.cfg.trainer.ckpt_path, f"{GLOBAL_STEP_PREFIX}{ckpt_iteration}")
            # Run validation: Make sure ckpt folder is consistent with latest_ckpt_global_step.txt
            validate_consistency_for_latest_checkpoint(
                self.cfg.trainer.ckpt_path,
                ckpt_iteration,
                checkpoint_path,
                latest_checkpoint_file,
                self.cfg.trainer.ckpt_interval,
            )
        else:
            # Get and validate resume path
            checkpoint_path = Path(self.cfg.trainer.resume_path)
            if not checkpoint_path:
                raise ValueError("`trainer.resume_path` must be specified when resume_mode is 'from_path'")

            # Validate that it's a global_step directory
            if GLOBAL_STEP_PREFIX not in checkpoint_path.name:
                raise ValueError(
                    f"`trainer.resume_path` must point to a directory whose name starting with {GLOBAL_STEP_PREFIX}, got: {checkpoint_path}"
                )

        # Validate that the path exists
        if not io.exists(str(checkpoint_path)):
            raise FileNotFoundError(f"Checkpoint path not found: {checkpoint_path}")

        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:
            raise ValueError(f"Checkpoint path {checkpoint_path} is not a valid checkpoint path")
        logger.info(f"Resuming from global_step: {global_step}")

        # Define paths for different checkpoint components
        policy_ckpt_dir = os.path.join(checkpoint_path, "policy")
        critic_ckpt_dir = os.path.join(checkpoint_path, "critic")
        trainer_state_path = os.path.join(checkpoint_path, "trainer_state.pt")
        dataloader_state_path = os.path.join(checkpoint_path, "data.pt")

        # Validate that required checkpoint files exist
        if not io.exists(trainer_state_path):
            raise FileNotFoundError(f"Trainer state file not found: {trainer_state_path}")

        # 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.")

        # 2. Load dataloader state if available
        if 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)
                self.train_dataloader.load_state_dict(dataloader_state)
                logger.info("Successfully loaded dataloader state")
            except Exception as e:
                logger.warning(f"Failed to load dataloader state: {e}. Dataloader will start from beginning.")
        else:
            logger.warning(
                f"No dataloader state found at {dataloader_state_path}. Dataloader will start from beginning."
            )

        # 3. Load policy checkpoint (dispatch handles offload/backload)
        logger.info(f"Loading policy checkpoint from {policy_ckpt_dir}")
        self.dispatch.load_checkpoint(
            "policy",
            policy_ckpt_dir,
            load_optimizer_states=True,
            load_lr_scheduler_states=True,
        )
        logger.info("Successfully loaded policy checkpoint")

        # 4. Load critic checkpoint if it exists and we have a critic model
        if self.has_critic:
            logger.info(f"Loading critic checkpoint from {critic_ckpt_dir}")
            self.dispatch.load_checkpoint(
                "critic",
                critic_ckpt_dir,
                load_optimizer_states=True,
                load_lr_scheduler_states=True,
            )
            logger.info("Successfully loaded critic checkpoint")

        logger.info(f"Successfully loaded complete checkpoint state from global_step_{global_step}")
        return global_step, str(checkpoint_path)

    def save_models(self):
        """
        Save the model parameters in HF format at `cfg.trainer.export_path`.

        Dispatch handles offload/backload automatically for all colocation configurations.
        """
        policy_export_dir = os.path.join(self.cfg.trainer.export_path, f"global_step_{self.global_step}", "policy")
        self.dispatch.save_hf_model("policy", policy_export_dir, self.tokenizer)

        if self.has_critic:
            critic_export_dir = os.path.join(self.cfg.trainer.export_path, f"global_step_{self.global_step}", "critic")
            self.dispatch.save_hf_model("critic", critic_export_dir, self.tokenizer)

        logger.info("Successfully saved model weights.")

    def update_ref_with_policy(self):
        """
        Update the reference model with the policy model weights (required by some algorithms).

        Dispatch handles offload/backload automatically for all colocation configurations.
        After this method, save_weights_for_sampler() should be called to sync weights.
        """
        # TODO(tgriggs): Make policy-to-ref sync faster.
        policy_export_dir = os.path.join(self.cfg.trainer.export_path, f"global_step_{self.global_step}", "policy")

        # Save policy model (dispatch handles GPU state)
        self.dispatch.save_hf_model("policy", policy_export_dir, self.tokenizer)

        # Re-initialize ref model from saved policy (dispatch handles offloading policy first)
        self.dispatch.init_model("ref", policy_export_dir)

        # Clean up temporary saved model files
        try:
            shutil.rmtree(policy_export_dir)
            logger.info(f"Cleaned up temporary policy export directory: {policy_export_dir}")
        except Exception as e:
            logger.warning(f"Failed to clean up temporary policy export directory {policy_export_dir}: {e}")

        logger.info("Successfully updated ref model with policy model, training continues.")

    def add_callback(self, callback: TrainingCallback) -> None:
        """Register a callback. Events fired after this call reach the new callback."""
        self._callback_handler.add(callback)

    @torch.no_grad()
    async def eval(self, vllm_metrics_scraper: Optional[VLLMMetricsScraper] = None) -> Dict[str, float]:
        """
        Run generation and scoring on the evaluation dataset.

        The eval metrics are recorded after having finished training `self.global_step` steps.
        Metrics recorded in global_step 0 corresponds to evaluations before training.

        Args:
            vllm_metrics_scraper: when provided, the eval loop calls
                ``resume()``/``pause()`` around each generation so the scraper
                attributes only generation time to the open ``vllm/eval`` window.

        Returns:
            A dictionary of evaluation metrics.
        """
        if self.cfg.generator.step_wise_trajectories:
            eval_metrics = await evaluate_step_wise(
                eval_dataloader=self.eval_dataloader,
                generator=self.generator,
                cfg=self.cfg,
                global_step=self.global_step,
                tokenizer=self.tokenizer,
                trajectory_logger=self.trajectory_logger,
                tracker=self.tracker,
                vllm_metrics_scraper=vllm_metrics_scraper,
            )
        else:
            eval_metrics = await evaluate(
                eval_dataloader=self.eval_dataloader,
                generator=self.generator,
                cfg=self.cfg,
                global_step=self.global_step,
                tokenizer=self.tokenizer,
                trajectory_logger=self.trajectory_logger,
                tracker=self.tracker,
                vllm_metrics_scraper=vllm_metrics_scraper,
            )
        return eval_metrics

    async def train(self):
        """
        Main training loop for PPO
        """
        if self._ray_gpu_monitor is not None:
            self._ray_gpu_monitor.start()

        # Initialize weight sync state between policy model and inference engines.
        with Timer("init_weight_sync_state"):
            self.init_weight_sync_state()

        # Load checkpoint state if resumption is enabled.
        if self.resume_mode != ResumeMode.NONE:
            with Timer("load_checkpoints"):
                self.global_step, _ = self.load_checkpoints()

        # Prepare weights for sampling
        with Timer("sync_weights"):
            await self.dispatch.save_weights_for_sampler()

        # Compute start_epoch up-front so callback metadata is ready before
        # any event fires (including the baseline eval below).
        start_epoch = self.global_step // len(self.train_dataloader)
        self._current_epoch = start_epoch
        self._training_control.reset()

        self._fire("on_train_start")

        # Eval before training. Wrapped in eval callbacks + on_log so that e.g.
        # a best-checkpoint callback sees the baseline reading.
        if self.cfg.trainer.eval_interval > 0 and self.cfg.trainer.eval_before_train:
            self._fire("on_eval_start")
            with Timer("eval", self.all_timings):
                eval_metrics = await self.eval()
            self._fire("on_eval_end", metrics=eval_metrics)
            self._fire("on_log", logs=eval_metrics)
            self.tracker.log(eval_metrics, step=self.global_step, commit=True)

        # initialize kl controller
        if self.cfg.trainer.algorithm.use_kl_in_reward:
            self.reward_kl_controller = get_kl_controller(self.cfg.trainer.algorithm)

        # main training loop
        pbar = tqdm(total=self.total_training_steps, initial=self.global_step, desc="Training Batches Processed")
        self.global_step += 1  # start training at global_step 1
        stop_training = False

        # booleans tracking whether we save ckpts
        # as well as hf model at step end
        will_save_ckpts = False
        hf_model_save = False
        self._profiler_start()
        try:
            for epoch in range(start_epoch, self.cfg.trainer.epochs):
                self._current_epoch = epoch
                self._fire("on_epoch_start")
                # ``step_started`` tracks the on_step_start/on_step_end pairing taking
                # dynamic-sampling into account (which span multiple inner iterations
                # before completing a logical step).
                step_started = False
                for _, rand_prompts in enumerate(self.train_dataloader):
                    if not step_started:
                        self._fire("on_step_start")
                        step_started = True
                        # Open the train-rollout metrics window once per logical
                        # step; paused so only the generation spans count toward the
                        # throughput denominator (dynamic sampling may generate more
                        # than once before the step completes).
                        if self._vllm_metrics_scraper is not None:
                            await self._vllm_metrics_scraper.start("vllm/train")
                            self._vllm_metrics_scraper.pause()
                    with Timer("step", self.all_timings):
                        # for colocate_all=true, inference engine is always on GPU when starting the training step

                        # 0. truncate data to have even shards
                        rand_prompts = self._remove_tail_data(rand_prompts)
                        generator_input, uids = prepare_generator_input(
                            rand_prompts,
                            self.cfg.generator.n_samples_per_prompt,
                            get_sampling_params_for_backend(
                                self.cfg.generator.inference_engine.backend, self.cfg.generator.sampling_params
                            ),
                            self.cfg.environment.env_class,
                            "train",
                            self.global_step,
                        )

                        # 1.1. generation phase
                        if self._vllm_metrics_scraper is not None:
                            self._vllm_metrics_scraper.resume()
                        with Timer("generate", self.all_timings):
                            generator_output: GeneratorOutput = await self.generate(generator_input)
                        if self._vllm_metrics_scraper is not None:
                            self._vllm_metrics_scraper.pause()

                        if self.cfg.generator.step_wise_trajectories:
                            # NOTE: We use instance_ids from `trajectory_ids` here instead of re-using `uids`
                            # this is because in step-wise training, len(uids) != len(generator_output["response_ids"])
                            uids = [trajectory_id.instance_id for trajectory_id in generator_output["trajectory_ids"]]

                        # dynamic sampling
                        if self.cfg.trainer.algorithm.dynamic_sampling.type is not None:
                            generator_output, uids, keep_sampling = self.handle_dynamic_sampling(generator_output, uids)
                            if keep_sampling:  # continue sampling
                                # update progress bar for current batch (but not global step)
                                pbar.update(1)
                                continue

                        if self.colocate_all:
                            # if we are not continuing sampling, we sleep the inference engine
                            await self.inference_engine_client.sleep()

                        # The train rollout for this step is done generating; close
                        # its metrics window. ``vllm/eval/*`` is collected separately
                        # around eval below.
                        vllm_metrics: Dict[str, float] = {}
                        if self._vllm_metrics_scraper is not None:
                            vllm_metrics = await self._vllm_metrics_scraper.stop()

                        # 1.2 postprocess rewards (and merge step-wise turns if enabled)
                        with Timer("postprocess_generator_output", self.all_timings):
                            generator_output, uids = self.postprocess_generator_output(generator_output, uids)

                        # 2.1 print example just for debugging
                        print_interval = self.cfg.trainer.print_example_interval
                        if print_interval > 0 and self.global_step % print_interval == 0:
                            vis = self.tokenizer.decode(generator_output["response_ids"][0])
                            pretty_print_example(
                                logger,
                                prompt=generator_input["prompts"][0],
                                response=vis,
                                reward=generator_output["rewards"][0],
                            )

                        # 2.2 Optionally upload up to `num_logger_train_samples` samples to tracker
                        if self.trajectory_logger is not None:
                            with Timer("log_train_results"):
                                self.trajectory_logger.log(
                                    tracker=self.tracker,
                                    num_samples=self.cfg.trainer.num_logger_train_samples,
                                    prompts=generator_input["prompts"],
                                    generator_output=generator_output,
                                    tokenizer=self.tokenizer,
                                    global_step=self.global_step,
                                    wandb_key="trajectories/train",
                                    include_idx=False,
                                )

                        # 3. Convert GeneratorOutput to TrainingInputBatch
                        with Timer("convert_to_training_input", self.all_timings):
                            training_input: TrainingInputBatch = self.convert_to_training_input(generator_output, uids)

                        # 4. Inference and calculate values, log probs, rewards, kl divergence
                        with Timer("fwd_logprobs_values_reward", self.all_timings):
                            training_input = self.fwd_logprobs_values_reward(training_input)

                        # 5. apply kl divergence penalty to rewards
                        if self.cfg.trainer.algorithm.use_kl_in_reward:
                            with Timer("apply_reward_kl_penalty", self.all_timings):
                                training_input = self.apply_reward_kl_penalty(training_input)

                        # 6. calculate advantages and returns
                        with Timer("compute_advantages_and_returns", self.all_timings):
                            training_input = self.compute_advantages_and_returns(training_input)
                            # remove some unwanted keys
                            for key in ["rewards"]:
                                training_input.pop(key)
                            training_input.metadata.pop("uids")
                            training_input.metadata.pop("is_last_step", None)

                        if self.cfg.trainer.dump_data_batch:
                            # dump data to file
                            with Timer("dump_data_batch"):
                                self.dump_data(
                                    training_input, file_name=f"global_step_{self.global_step}_training_input"
                                )

                        # 7. train policy/critic model
                        # Policy model is backloaded to GPU during training
                        with Timer("train_critic_and_policy", self.all_timings):
                            status = self.train_critic_and_policy(training_input)

                            # One profiler step per RL global step.
                            self._profiler_step()

                        self._fire("on_step_end", batch=training_input, metrics=status)
                        step_started = False

                        # Capture callback-driven triggers, then reset.
                        force_save = self._training_control.should_save
                        force_eval = self._training_control.should_evaluate
                        self._training_control.should_save = False
                        self._training_control.should_evaluate = False

                        # 8. conditionally save checkpoints and hf model
                        is_epoch_end = self.global_step % len(self.train_dataloader) == 0
                        hf_model_save = self.cfg.trainer.hf_save_interval > 0 and (
                            is_epoch_end or self.global_step % self.cfg.trainer.hf_save_interval == 0
                        )
                        ckpt_interval_save = self.cfg.trainer.ckpt_interval > 0 and (
                            is_epoch_end or self.global_step % self.cfg.trainer.ckpt_interval == 0
                        )
                        will_save_ckpts = force_save or ckpt_interval_save
                        if will_save_ckpts:
                            with Timer("save_checkpoints", self.all_timings):
                                ckpt_path = self.save_checkpoints()
                            self._fire("on_save", ckpt_path=ckpt_path)
                        if hf_model_save:
                            with Timer("save_hf_model", self.all_timings):
                                self.save_models()

                        # 9. conditionally sync policy and ref at the end of the epoch
                        if (
                            self.cfg.trainer.update_ref_every_epoch
                            and self.ref_model is not None
                            and is_epoch_end
                            and epoch != self.cfg.trainer.epochs - 1  # skip updating ref at the end of the last epoch
                        ):
                            with Timer("update_ref_with_policy", self.all_timings):
                                self.update_ref_with_policy()

                        # 10. Prepare weights for sampling
                        with Timer("sync_weights", self.all_timings):
                            await self.dispatch.save_weights_for_sampler()

                    # 11. set logs
                    logger.info(status)
                    # Throughput metrics
                    train_time = self.all_timings.get("train_critic_and_policy", 0.0)
                    if train_time > 0 and training_input.get("attention_mask") is not None:
                        total_tokens = int(training_input["attention_mask"].sum().item())
                        self.all_metrics["trainer/tokens_per_second_per_gpu"] = total_tokens / (
                            train_time * self._num_training_gpus
                        )
                    # log epoch info
                    self.all_metrics.update({"trainer/epoch": epoch, "trainer/global_step": self.global_step})
                    interval_eval = self.cfg.trainer.eval_interval > 0 and (
                        self.global_step % self.cfg.trainer.eval_interval == 0
                        or self.global_step == self.total_training_steps
                    )
                    if force_eval or interval_eval:
                        # Open the eval-rollout window; the scraper itself measures
                        # the generation spans via resume()/pause() inside eval().
                        if self._vllm_metrics_scraper is not None:
                            await self._vllm_metrics_scraper.start("vllm/eval")
                            self._vllm_metrics_scraper.pause()
                        self._fire("on_eval_start")
                        with Timer("eval", self.all_timings):
                            eval_metrics = await self.eval(vllm_metrics_scraper=self._vllm_metrics_scraper)
                            self.all_metrics.update(eval_metrics)
                        self._fire("on_eval_end", metrics=eval_metrics)
                        if self._vllm_metrics_scraper is not None:
                            vllm_metrics.update(await self._vllm_metrics_scraper.stop())

                    log_payload = {
                        **self.all_metrics,
                        **{f"timing/{k}": v for k, v in self.all_timings.items()},
                        # vllm/train/* = train rollout, vllm/eval/* = eval rollout,
                        # each over its own generation time (owned by the scraper).
                        **vllm_metrics,
                    }

                    if self._ray_gpu_monitor is not None:
                        log_payload.update(self._ray_gpu_monitor.flush())

                    self._fire("on_log", logs=log_payload)

                    self.tracker.log(log_payload, step=self.global_step, commit=True)
                    self.all_metrics = {}
                    self.all_timings = {}

                    # update progress bar after logging
                    pbar.update(1)

                    self.global_step += 1

                    if (
                        self.cfg.trainer.max_training_steps is not None
                        and self.global_step > self.cfg.trainer.max_training_steps
                    ):
                        logger.info(
                            f"Reached max_training_steps={self.cfg.trainer.max_training_steps}, stopping early."
                        )
                        stop_training = True
                        break

                    del training_input, generator_output

                self._fire("on_epoch_end")

                if stop_training:
                    break
        finally:
            self._profiler_stop()

        pbar.close()
        if self.colocate_all:
            await self.inference_engine_client.sleep()

        # Decrement global step by 1 to stop at the last global step
        # We use the global step value in callbacks when training finishes,
        # as well as for a final checkpoint save
        self.global_step -= 1

        # Safety net: always save final checkpoint at end of training.
        # Skip if we already saved at the last step
        if self.cfg.trainer.ckpt_interval > 0 and not will_save_ckpts:
            with Timer("save_checkpoints", self.all_timings):
                ckpt_path = self.save_checkpoints()
                logger.info("Saved final checkpoint.")
            self._fire("on_save", ckpt_path=ckpt_path)
        if self.cfg.trainer.hf_save_interval > 0 and not hf_model_save:
            with Timer("save_hf_model", self.all_timings):
                self.save_models()
                logger.info("Saved final model.")
        if self._vllm_metrics_scraper is not None:
            await self._vllm_metrics_scraper.aclose()

        if self._ray_gpu_monitor is not None:
            self._ray_gpu_monitor.stop()

        self._fire("on_train_end")
        self.tracker.finish()
        logger.info("Training done!")