# Copyright (c) ModelScope Contributors. All rights reserved.
# Part of the implementation is borrowed from huggingface/trl.

# fmt: off
# apply patch before importing trl, which may internally reference GuidedDecodingParams
try:
    import vllm
    try:
        from vllm.sampling_params import GuidedDecodingParams
    except ImportError:
        import vllm.sampling_params

        # removed in https://github.com/vllm-project/vllm/pull/22772
        vllm.sampling_params.GuidedDecodingParams = vllm.sampling_params.StructuredOutputsParams
except ImportError:
    pass

# https://github.com/modelscope/ms-swift/pull/8280
try:
    import trl.import_utils as _trl_import_utils
    _orig = _trl_import_utils.is_vllm_ascend_available
    if not isinstance(_orig(), bool):
        _trl_import_utils.is_vllm_ascend_available = lambda: bool(_orig()[0])
except Exception:
    pass
# fmt: on

import concurrent.futures
import inspect
import os
import time
import torch
import torch.distributed as dist
import torch.nn as nn
import transformers
from accelerate.utils import gather, gather_object, is_peft_model, set_seed
from collections import defaultdict, deque
from contextlib import contextmanager, nullcontext
from copy import deepcopy
from packaging import version
from transformers import PreTrainedModel
from transformers.trainer import Trainer as HfTrainer
from trl import GRPOTrainer as HFGRPOTrainer
from trl.models import prepare_deepspeed
from trl.trainer import grpo_trainer
from trl.trainer.grpo_trainer import RepeatSampler, nanmax, nanmin
from trl.trainer.utils import selective_log_softmax
from typing import Any, Callable, Dict, List, Optional, Tuple, Union

from swift.dataset import RowPreprocessor
from swift.rewards import orms, rm_plugins
from swift.rl_core.advantage import (compute_advantages, compute_advantages_dynamic, compute_reward_metrics,
                                     compute_teacher_kl_per_token, expand_advantage_to_per_token)
from swift.rl_core.data import GRPOBatch, GRPOSample
from swift.rl_core.grpo_algorithm import score_completions
from swift.rlhf_trainers.gkd_helpers import (assemble_teacher_completion_logprobs, build_opsd_samples,
                                             build_teacher_requests, encode_teacher_view,
                                             fetch_teacher_parsed_by_routing, remap_teacher_logps_to_student_frame,
                                             should_compute_local_teacher_logps)
from swift.sequence_parallel import GatherLoss, sequence_parallel
from swift.template import Template, TemplateInputs
from swift.trainers import SwiftMixin, disable_gradient_checkpointing
from swift.utils import (JsonlWriter, get_cu_seqlens_from_position_ids, get_logger, is_swanlab_available,
                         is_wandb_available, nanstd, remove_response, seed_worker, to_device,
                         unwrap_model_for_generation)
from .arguments import GRPOConfig
from .rollout_mixin import DataType, RolloutTrainerMixin, SyncRefModelCallback
from .utils import (_ForwardRedirection, collate_to_grpo_micro_batch, compute_chord_loss, encode_sample,
                    get_even_process_data, identity_data_collator, load_pil_img, make_chord_sft_dataset,
                    pad_logps_back_to_batch, patch_save_last_checkpoint, profiling_context, profiling_decorator,
                    replace_assistant_response_with_ids, swanlab_get_run)

try:
    from trl.trainer.utils import entropy_from_logits
except ImportError:
    from .utils import entropy_from_logits

del HFGRPOTrainer.__init__
del HFGRPOTrainer.log
grpo_trainer.seed_worker = seed_worker  # fix transformers 4.51.3

logger = get_logger()
if is_wandb_available():
    import wandb
if is_swanlab_available():
    import swanlab


class GRPOTrainer(RolloutTrainerMixin, SwiftMixin, HFGRPOTrainer):
    executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
    sample_cls = GRPOSample

    def __init__(self,
                 model: Optional[Union[PreTrainedModel, nn.Module]] = None,
                 ref_model: Optional[Union[PreTrainedModel, nn.Module]] = None,
                 reward_model: Optional[List[Union[PreTrainedModel, nn.Module]]] = None,
                 reward_funcs: Optional[List[Union[str, Callable]]] = None,
                 *_args,
                 **kwargs):
        patch_save_last_checkpoint()
        args: GRPOConfig = kwargs['args']
        self.args = args
        self.ref_adapter_name = getattr(args, 'ref_adapter_name', None)
        self.model_adapter_name = None
        self.is_multimodal = model.model_meta.is_multimodal
        self.model_kwarg_keys = (
            inspect.signature(model.forward).parameters.keys() if not hasattr(model, 'get_base_model') else
            inspect.signature(model.get_base_model().forward).parameters.keys())

        self.vllm_client = kwargs.pop('vllm_client', None)
        self.chord_sft_dataset = kwargs.pop('chord_sft_dataset', None)
        reward_templates = kwargs.pop('reward_template', None)
        # Teacher kwargs (same set as GKD; setting one turns GRPO into OPD-RL: teacher KL as advantage).
        self._pop_teacher_kwargs(kwargs)
        self._prepare_algorithm_params()
        super().__init__(model, ref_model, *_args, **kwargs)
        self._prepare_chord_dataset()
        self.prepare_rollout()
        self._prepare_rewards(reward_funcs, reward_model, reward_templates)

        # A configured teacher enables OPD-RL (teacher KL as advantage); with no reward
        # functions it is the sole signal. Same gating as GKD: teacher set -> teacher branch.
        self._setup_teacher()
        self.teacher_kl_coef = args.teacher_kl_coef
        if not self.reward_funcs and not self.use_gym_env and not self._has_teacher:
            raise ValueError('You must specify reward_funcs or reward_model')

        if self.args.eval_strategy != 'no' and not self.args.eval_use_evalscope:
            total_eval_batch_size = self.args.per_device_eval_batch_size * \
                self.accelerator.num_processes // self.num_generations_eval
            assert len(self.eval_dataset) >= total_eval_batch_size, (
                f'eval_dataset size {len(self.eval_dataset)} is smaller than '
                f'total_eval_batch_size {total_eval_batch_size}. '
                f'Please increase the size of eval_dataset or set a larger value for split_dataset_ratio.')

        self._prepare_liger_loss()
        self._prepare_metrics()

        # Ensure each process receives a unique seed to prevent duplicate completions when generating with
        # transformers if num_generations exceeds per_device_train_batch_size. We could skip it if we use vLLM, but
        # it's safer to set it in all cases.
        set_seed(args.seed, device_specific=True)

        # Gradient accumulation requires scaled loss. Normally, loss scaling in the parent class depends on whether the
        # model accepts loss-related kwargs. Since we compute our own loss, this check is irrelevant. We set
        # self.model_accepts_loss_kwargs to False to enable scaling.
        self.model_accepts_loss_kwargs = False

        if args.sync_ref_model:
            self.add_callback(SyncRefModelCallback(self))

        if self.args.dynamic_sample or self.template.truncation_strategy == 'raise':
            self._prepare_resample_data_iterator()
        # flag indicating whether the evaluation has started
        self.eval_flag = False

        if self.template.sequence_parallel_size > 1:
            self.args.gradient_accumulation_steps = self.args.gradient_accumulation_steps * sequence_parallel.world_size

        # for multi-turn server, maybe the num of rollout outputs is not equal to the num of rollout inputs
        self.dynamic_num_samples = False
        # Record the number of samples that need to be padded for even distribution across processes
        self.rollout_pad_count = 0

        # Tracks the number of iterations (forward + backward passes), including those within a gradient accumulation cycle. # noqa
        self._step = 0
        # Buffer the batch to reuse generated outputs across multiple updates. For more details, see
        # `_get_train_sampler` and `_prepare_inputs`.
        self._buffered_inputs = None
        self._current_train_step_time = 0.0

    def _get_data_collator(self, args, template):
        return identity_data_collator

    def _get_train_sampler(self, train_dataset=None):
        if self.template.sequence_parallel_size > 1:
            return RepeatSampler(
                data_source=train_dataset or self.train_dataset,
                mini_repeat_count=self.num_generations,
                batch_size=self.args.generation_batch_size // self.num_generations,
                repeat_count=self.num_iterations * self.args.steps_per_generation * sequence_parallel.world_size,
                shuffle=self.shuffle_dataset,
                seed=self.args.seed,
            )
        else:
            return super()._get_train_sampler(train_dataset)

    @profiling_decorator
    def _prepare_inputs(self, inputs: DataType) -> Dict[str, Union[torch.Tensor, Any]]:
        # Prepares inputs for model training/evaluation by managing completion generation and batch handling.
        # During training:
        #   - Receives the local generation batch (Per-GPU batch size × steps per generation)
        #     from the modified training dataloader instead of the standard local batch
        #   - Generates completions once for the entire generation batch and splits it into batches of size
        #     `per_device_train_batch_size`
        #   - Buffers these completions and returns the appropriate slice for the current accumulation step
        #   - Optimizes by regenerating completions only periodically (every steps_per_generation * num_iterations)
        # During evaluation:
        #   - The input is treated as a standard local batch (no accumulation, no multiple iterations)
        #   - Completions are generated for each batch without buffering or reuse
        # Returns a single local batch in both cases.

        mode = 'train' if self.model.training else 'eval'
        if mode == 'train':
            num_rollout_samples = self.args.steps_per_generation * self.template.sequence_parallel_size
            generate_every = num_rollout_samples * self.num_iterations
            if self._step % generate_every == 0 or self._buffered_inputs is None:
                micro_batches = self._generate_and_score_completions(inputs)
                self._buffered_inputs = micro_batches
            micro_batch = self._buffered_inputs[self._step % num_rollout_samples]
        else:
            micro_batch = self._generate_and_score_completions(inputs)
        return micro_batch

    def _rollout_samples(self, inputs: DataType) -> List[GRPOSample]:
        """Convert raw rows to samples and generate completions."""
        if self.template.truncation_strategy == 'raise':
            inputs = self.resample_encode_failed_inputs(inputs)
        samples: List[GRPOSample] = self.to_samples(inputs)
        samples = self._generate_completions(samples)
        return samples

    @profiling_decorator
    def _score_completions(self, samples: List[GRPOSample]) -> List[GRPOSample]:
        """Score completions with reward functions and run DAPO dynamic sampling.

        Rewards are stashed on ``self._rewards_per_func`` for ``_postprocess_batch``
        (advantages depend on the encoded batches produced after this stage).
        Returns the (possibly resampled) samples.
        """
        total_rewards_per_func = self._compute_rewards_per_func(samples)
        mode = 'train' if self.model.training else 'eval'
        if self.dynamic_sample and mode == 'train':
            # dynamic sampling for std=0 groups
            samples, total_rewards_per_func = self._dynamic_sampling(samples, total_rewards_per_func)  # noqa
        self._rewards_per_func = total_rewards_per_func
        return samples

    def _postprocess_batch(self, samples: List[GRPOSample], batch_encoded_inputs: List[Dict[str, Any]]) -> None:
        """Compute advantages and write them into each ``grpo_batch``."""
        # OPD-RL teacher API: fetch single-token teacher logps before advantage computation.
        if self._has_teacher and self.use_teacher_api:
            self._assemble_teacher_api_logps(samples, batch_encoded_inputs)
        total_advantages = self._compute_advantages(samples, self._rewards_per_func, batch_encoded_inputs)

        local_advantages = get_even_process_data(self, total_advantages)
        assert len(local_advantages) == len(samples)
        for i, advantage in enumerate(local_advantages):
            samples[i].advantages = advantage
        # log metrics in samples
        self._logs['advantages'].extend(total_advantages.tolist())

        # Add advantages to each batch in batch_encoded_inputs
        gas_chunks = self.split_by_mini_batches(samples)
        assert len(gas_chunks) == len(batch_encoded_inputs), \
            f'Mismatch: {len(gas_chunks)} chunks vs {len(batch_encoded_inputs)} batches'

        for batch, batch_encoded in zip(gas_chunks, batch_encoded_inputs):
            # Under sequence parallel, split_by_mini_batches gathers samples across the SP group via
            # all_gather_object, so per-sample advantages may carry tensors from different ranks/devices;
            # move them onto the current device before stacking.
            device = self.accelerator.device
            grpo_batch: GRPOBatch = batch_encoded['grpo_batch']
            base_advantages = torch.stack([data.advantages.to(device) for data in batch])
            # Expand the per-sequence base advantage to per-token [B, T] here (not by broadcast
            # in the loss), so the OPD-RL signed teacher log-ratio is added per token
            # (adv_t = base + coef * (teacher_logp - student_logp)).
            grpo_batch.advantages = expand_advantage_to_per_token(
                base_advantages,
                grpo_batch.completion_mask,
                teacher_per_token_logps=grpo_batch.teacher_per_token_logps,
                policy_per_token_logps=grpo_batch.old_per_token_logps
                if grpo_batch.teacher_per_token_logps is not None else None,
                teacher_kl_coef=self.teacher_kl_coef if grpo_batch.teacher_per_token_logps is not None else 0.0,
            )
            if grpo_batch.teacher_per_token_logps is not None:
                mode = 'train' if self.model.training else 'eval'
                # Monitoring uses the non-negative k3 estimator (a "distance from the teacher" gauge
                # that should decrease over training); the advantage above uses the signed k1 log-ratio.
                k3 = compute_teacher_kl_per_token(grpo_batch.teacher_per_token_logps, grpo_batch.old_per_token_logps,
                                                  grpo_batch.completion_mask)
                self._metrics[mode]['teacher_kl'].append(
                    (k3.sum() / grpo_batch.completion_mask.sum().clamp(min=1.0)).item())

    def _log_rollout(self, samples: List[GRPOSample]) -> None:
        """Log prompts/completions and extra gathered metrics (solution, num_turns)."""
        with profiling_context(self, 'log_metrics'):
            # --- logs (prompts + completions) ---
            messages = [s.messages[:-1] for s in samples]
            completions = [deepcopy(s.messages[-1]['content']) for s in samples]
            for i, completion in enumerate(completions):
                if isinstance(completion, str):
                    continue
                if isinstance(completion, list):
                    token_ids = completion
                elif isinstance(completion, dict):
                    token_ids = completion['token_ids']
                completions[i] = self.processing_class.decode(token_ids)
            valid_messages = self._gather_and_flatten(messages, flatten_level=0)
            valid_completions = self._gather_and_flatten(completions, flatten_level=0)
            self._logs['prompt'].extend(self._apply_chat_template_to_messages_list(valid_messages))
            self._logs['completion'].extend(valid_completions)

            # Example: if you want to log extra data in the wandb / swanlab table,
            #          add them to metrics_to_gather
            # NOTE: every key you register must appear in ALL rollout outputs
            #       to avoid potential communication / synchronization issues
            metrics_for_logs_to_gather = {}

            if all('solution' in s.extra for s in samples):
                metrics_for_logs_to_gather['solution'] = [s.extra['solution'] for s in samples]

            if all(s.rollout_infos and 'num_turns' in s.rollout_infos for s in samples):
                metrics_for_logs_to_gather['num_turns'] = [s.rollout_infos['num_turns'] for s in samples]

            if metrics_for_logs_to_gather:
                for key, value in metrics_for_logs_to_gather.items():
                    if key not in self._logs:
                        self._logs[key] = deque(maxlen=self.args.generation_batch_size)
                    self._logs[key].extend(self._gather_and_flatten(value, flatten_level=0))

    @profiling_decorator
    def _compute_rewards_per_func(self, samples: List[GRPOSample]) -> torch.Tensor:
        """Score completions using all reward functions.

        Args:
            samples: List of on-policy samples, each carrying messages with conversation history.

        Returns:
            rewards_per_func: Tensor of shape (num_examples, num_reward_funcs) with all reward values
        """
        extra_reward_kwargs = None
        if self.enable_server_multi_turn:
            extra_reward_kwargs = {'trajectory_inputs': self._get_trajectory_inputs(samples)}

        with self._disable_sp_context():
            local_rewards_per_func = score_completions(
                samples,
                reward_funcs=self.reward_funcs,
                reward_model_plugins=self.reward_model_plugins,
                use_gym_env=self.use_gym_env,
                device=self.accelerator.device,
                trainer_state=self.state,
                extra_reward_kwargs=extra_reward_kwargs,
            )

        # OPD-RL pure distillation: no reward_funcs -> a [N, 0] tensor. accelerate.gather
        # can't reshape a 0-column tensor (view(-1, 0) is ambiguous), so gather only the
        # global sample count and rebuild an empty [N_global, 0] reward matrix. Downstream
        # advantages then come solely from the teacher KL injection.
        if local_rewards_per_func.shape[1] == 0:
            global_count = sum(gather_object([local_rewards_per_func.shape[0]]))
            return torch.zeros((global_count, 0), dtype=torch.float32, device=self.accelerator.device)

        # gather rewards
        if not self.dynamic_num_samples:
            total_rewards_per_func = gather(local_rewards_per_func)
        else:
            # gather_object to avoid shape mismatch
            local_rewards_list = [row.tolist() for row in local_rewards_per_func]
            total_rewards_per_func = gather_object(local_rewards_list)
            total_rewards_per_func = torch.tensor(
                total_rewards_per_func, dtype=torch.float32, device=self.accelerator.device)

        return total_rewards_per_func

    def _compute_advantages(self, samples: List[GRPOSample], rewards_per_func: torch.Tensor,
                            batch_encoded_inputs: List[Dict[str, Any]]) -> torch.Tensor:
        """
        Compute advantages for RL training.

        Supports two modes:
        1. **Default grouped mode** (no prompt_ids / request_ids provided)
        - Assumes rewards are grouped by prompt and each prompt has exactly
            `self.num_generations` completions.
        - Computes advantages relative to group mean.
        2. **Request-aware mode** (multi-turn conversations, variable number of samples)
        - Groups rewards by unique `request_id` and computes statistics per prompt.
        - Handles dynamic sample sizes where multiple request_ids may share the same prompt.

        Args:
            samples (List[GRPOSample]):
                On-policy samples used to derive prompt_id / request_id grouping.
            rewards_per_func (torch.Tensor):
                Reward values for each reward function, shape `(N, num_reward_funcs)`.

        Returns:
            **advantages** (torch.Tensor):
                Computed advantages, shape `(N,)`.
        """

        def log_rewards_metrics(rewards: torch.Tensor, rewards_per_func_for_metrics: torch.Tensor):
            """Log reward statistics for monitoring. Only log once per unique request_id."""
            # rewards: [prompt_batch_size, num_generations]
            # rewards_per_func_for_metrics: [prompt_batch_size*num_generations, self.num_reward_funcs]
            mode = 'train' if self.model.training else 'eval'
            num_generations = self.num_generations if mode == 'train' else self.num_generations_eval
            group_rewards = rewards.view(-1, num_generations)
            rewards_mean = group_rewards.mean(-1).mean().item()
            if self.scale_rewards in ['group', 'none', 'gdpo']:
                # Handle edge case when num_generations_eval=1
                if num_generations > 1:
                    rewards_std = group_rewards.std(-1).mean().item()
                else:
                    rewards_std = 0.0
            elif self.scale_rewards == 'batch':
                rewards_std = rewards.std().item() if rewards.numel() > 1 else 0.0
            if num_generations > 1:
                is_std_zero = torch.isclose(group_rewards.std(dim=1), torch.zeros_like(group_rewards.std(dim=1)))
            else:
                is_std_zero = torch.ones(group_rewards.size(0), dtype=torch.bool, device=group_rewards.device)

            self._metrics[mode]['reward'].append(rewards_mean)
            self._metrics[mode]['reward_std'].append(rewards_std)
            self._metrics[mode]['frac_reward_zero_std'].append(is_std_zero.float().mean().item())

            # Log per-reward-function statistics using deduplicated rewards_per_func
            for i, name in enumerate(self.reward_func_names):
                col = rewards_per_func_for_metrics[:, i]
                self._metrics[mode][f'rewards/{name}/mean'].append(torch.nanmean(col).item())
                self._metrics[mode][f'rewards/{name}/std'].append(nanstd(col).item())

        def log_rewards_all(rewards_per_func: torch.Tensor):
            """Log all rewards for debugging."""
            for i, name in enumerate(self.reward_func_names):
                self._logs['rewards'][name].extend(rewards_per_func[:, i].tolist())

        # Compute per-sample KL values if kl_in_reward is enabled.
        kl_values = None
        if self.kl_in_reward and self.beta != 0.0:
            kl_list = []
            for batch_encoded in batch_encoded_inputs:
                grpo_batch = batch_encoded['grpo_batch']
                per_token_kl = grpo_batch.old_per_token_logps - grpo_batch.ref_per_token_logps
                kl = (per_token_kl * grpo_batch.completion_mask).sum(-1)
                kl_list.append(kl)
            kl_values = torch.cat(kl_list, dim=0)
            kl_values = gather(kl_values)

        # Keep weighted rewards for the request-aware (multi-turn) path below.
        rewards = (rewards_per_func * self.reward_weights.unsqueeze(0)).nansum(dim=1)
        if self.kl_in_reward and self.beta != 0.0:
            rewards = rewards - self.beta * kl_values

        # --------------------------------------------------
        # Case 1: Default grouped mode
        # --------------------------------------------------
        mode = 'train' if self.model.training else 'eval'
        num_generations = self.num_generations if mode == 'train' else self.num_generations_eval
        if not self.dynamic_num_samples:
            advantages, weighted_rewards = compute_advantages(
                rewards_per_func=rewards_per_func,
                reward_weights=self.reward_weights,
                num_generations=num_generations,
                advantage_estimator=self.advantage_estimator,
                scale_rewards=self.scale_rewards,
                kl_in_reward=self.kl_in_reward,
                beta=self.beta,
                kl_values=kl_values,
            )

            reward_metrics = compute_reward_metrics(
                rewards=weighted_rewards,
                rewards_per_func=rewards_per_func,
                reward_func_names=self.reward_func_names,
                num_generations=num_generations,
                scale_rewards=self.scale_rewards,
            )
            self._metrics[mode]['reward'].append(reward_metrics.reward_mean)
            self._metrics[mode]['reward_std'].append(reward_metrics.reward_std)
            self._metrics[mode]['frac_reward_zero_std'].append(reward_metrics.frac_reward_zero_std)
            if kl_values is not None:
                self._metrics[mode]['kl'].append(kl_values.nanmean().item())
            for name in self.reward_func_names:
                self._metrics[mode][f'rewards/{name}/mean'].append(reward_metrics.per_func_mean[name])
                self._metrics[mode][f'rewards/{name}/std'].append(reward_metrics.per_func_std[name])
            log_rewards_all(rewards_per_func)
            return advantages

        # --------------------------------------------------
        # Case 2: Request-aware mode
        # --------------------------------------------------
        else:
            prompt_ids = gather_object([s.prompt_id for s in samples])
            request_ids = gather_object([s.request_id for s in samples])
            assert rewards.shape[0] == len(prompt_ids) == len(request_ids)

            # Validate rewards consistency within the same request_id (fail-fast).
            for rid in set(request_ids):
                idxs = [i for i, r in enumerate(request_ids) if r == rid]
                if not torch.allclose(rewards[idxs], rewards[idxs[0]].expand(len(idxs)), atol=1e-6):
                    raise ValueError(f'Inconsistent rewards detected for request_id={rid}.')

            advantages, _ = compute_advantages_dynamic(
                rewards_per_func=rewards_per_func,
                reward_weights=self.reward_weights,
                prompt_ids=prompt_ids,
                request_ids=request_ids,
                advantage_estimator=self.advantage_estimator,
                scale_rewards=self.scale_rewards,
                kl_in_reward=self.kl_in_reward,
                beta=self.beta,
                kl_values=kl_values,
            )

            # Metrics are logged on request-deduplicated rewards.
            unique_indices = self._get_last_indices(request_ids)
            log_rewards_metrics(
                rewards=rewards[unique_indices], rewards_per_func_for_metrics=rewards_per_func[unique_indices])
            log_rewards_all(rewards_per_func)

            return advantages

    def _collate_teacher_opsd_batch(self, batch: List[GRPOSample], template) -> Tuple[Dict[str, Any], GRPOBatch]:
        """Encode + collate the OPSD teacher view (teacher_prompt + same response).

        Produces a teacher micro-batch whose completion frame may differ in length from
        the student's; ``_compute_teacher_logps`` remaps the teacher logps back onto the
        student frame. ``build_teacher_view`` is already done by ``build_opsd_samples``.
        """
        for s in batch:
            s.encoded = encode_teacher_view(s, template)
        teacher_model_inputs, teacher_grpo_batch = collate_to_grpo_micro_batch(
            batch, template, device=self.model.device, use_logits_to_keep=True)
        teacher_model_inputs.pop('labels', None)
        # Restore the student encoding so downstream code (advantages/loss) keeps the student frame.
        for s in batch:
            s.encoded = encode_sample(s, template)
            s.encoded.pop('_extra_kwargs', None)
        return teacher_model_inputs, teacher_grpo_batch

    def _compute_teacher_logps(self,
                               model_inputs,
                               grpo_batch,
                               origin_data=None,
                               teacher_model_inputs=None,
                               teacher_grpo_batch=None) -> torch.Tensor:
        """OPD-RL: per-token teacher logp on the sampled tokens via a local teacher forward.

        Reuses ``_get_per_token_logps_and_entropies`` so the teacher logp frame matches
        the policy's old/ref logps (token-in-token-out, single token per position).
        LoRA same-model teacher runs the student under ``disable_adapter``. For OPSD,
        the teacher forwards its own ``teacher_model_inputs`` (teacher_prompt + same
        response) and the result is remapped onto the student's completion frame.
        """
        is_opsd = teacher_model_inputs is not None
        # OPSD remaps via completion masks (per-token), which the chunked/multimodal logps path
        # does not preserve; OPSD self-distillation uses single-token text batches so this holds.
        fwd_inputs = teacher_model_inputs if is_opsd else model_inputs
        fwd_batch = teacher_grpo_batch if is_opsd else grpo_batch
        fwd_origin = None if is_opsd else origin_data
        with torch.no_grad():
            if self._teacher_use_disable_adapter:
                with self.accelerator.unwrap_model(self.model).disable_adapter(), disable_gradient_checkpointing(
                        self.model, self.args.gradient_checkpointing_kwargs):
                    teacher_logps = self._get_per_token_logps_and_entropies(
                        self.model, fwd_inputs, fwd_batch, origin_data=fwd_origin)[0]
            else:
                model = self.teacher_model if self.teacher_model else self.model
                with self.load_teacher_model_context(), disable_gradient_checkpointing(
                        model, self.args.gradient_checkpointing_kwargs):
                    teacher_logps = self._get_per_token_logps_and_entropies(
                        model, fwd_inputs, fwd_batch, origin_data=fwd_origin)[0]
        if is_opsd:
            teacher_logps = remap_teacher_logps_to_student_frame(teacher_logps, teacher_grpo_batch.completion_mask,
                                                                 grpo_batch.completion_mask)
        return teacher_logps

    def _assemble_teacher_api_logps(self, samples: List[GRPOSample], batch_encoded_inputs: List[Dict[str,
                                                                                                     Any]]) -> None:
        """OPD-RL teacher API: fetch the *sampled* token's logp at each response position
        (``prompt_logprobs=0`` -> ``parse_prompt_logprobs(topk=0)``, the actually-present
        token, not the top-1) as a completion-frame ``TeacherOutput``, then read its single
        column as ``teacher_per_token_logps``. Future top-k RL reuses the same ``TeacherOutput``.

        Each sample routes to exactly one teacher by tag (single teacher = all samples). Routing
        runs once over the flattened batch, so a single teacher is one DP gather → infer → slice.

        Samples/requests come from each chunk's ``_chunk_samples`` (the SP-gathered batch that its
        ``grpo_batch`` was collated from), not the local ``samples`` argument: under SP>1 the local
        samples are shorter/reordered, which would misalign the teacher logps.
        """
        flat_samples: List[GRPOSample] = []
        chunk_samples_list: List[List[GRPOSample]] = []
        for batch_encoded in batch_encoded_inputs:
            chunk = batch_encoded.pop('_chunk_samples', [])
            chunk_samples_list.append(chunk)
            flat_samples.extend(chunk)

        # OPSD: populate teacher_messages so build_teacher_requests scores the teacher prompt.
        build_opsd_samples(flat_samples)
        requests = build_teacher_requests(flat_samples, self.template)
        parsed = fetch_teacher_parsed_by_routing(
            flat_samples,
            requests,
            self.teacher_configs,
            self.teacher_clients,
            gather_fn=self._gather_teacher_requests,
            infer_fn=lambda handle, client: self._infer_teacher_requests(handle, topk=0, teacher_client=client),
            scatter_fn=self._scatter_teacher_parsed,
            is_main_process=self.accelerator.is_main_process,
            tag_key=self.args.teacher_tag_key)

        offset = 0
        for batch_encoded, chunk in zip(batch_encoded_inputs, chunk_samples_list):
            grpo_batch: GRPOBatch = batch_encoded['grpo_batch']
            device = grpo_batch.completion_mask.device
            n = len(chunk)
            teacher_out = assemble_teacher_completion_logprobs(
                parsed[offset:offset + n],
                grpo_batch.completion_mask,
                device,
                response_token_ids=[s.response_token_ids for s in chunk])
            grpo_batch.teacher_per_token_logps = teacher_out.topk_logprobs[..., 0]
            offset += n

    @profiling_decorator
    def _dynamic_sampling(self, samples: List[GRPOSample],
                          rewards_per_func: torch.Tensor) -> Tuple[List[GRPOSample], torch.Tensor]:
        """
        Perform dynamic sampling to replace samples with zero-reward-variance groups.

        This method implements DAPO (https://arxiv.org/abs/2503.14476) by replacing
        samples from groups with zero reward variance (std=0) through resampling.

        Args:
            samples: local on-policy samples
            rewards_per_func: reward per function for global samples

        Returns:
            tuple: (samples, rewards_per_func) with zero-variance groups replaced by resampled data
        """
        # DAPO https://arxiv.org/abs/2503.14476
        # Replaces samples with zero-reward-variance groups (std=0)
        resample_count = 0
        valid_samples = []
        valid_rewards_per_func = []
        origin_data = (samples, rewards_per_func)

        while resample_count < self.max_resample_times:
            rewards_std = self.compute_std(samples, rewards_per_func)
            valid_mask = (rewards_std > 0)
            all_samples = gather_object(samples)
            valid_samples.extend([s for s, mask in zip(all_samples, valid_mask) if mask])
            valid_rewards_per_func.append(rewards_per_func[valid_mask])
            if len(valid_samples) >= self.args.generation_batch_size:
                break

            inputs = next(self.dynamic_resample_iterator)
            if self.template.truncation_strategy == 'raise':
                inputs = self.resample_encode_failed_inputs(inputs)
            samples = self.to_samples(inputs)
            samples = self._generate_completions(samples)
            rewards_per_func = self._compute_rewards_per_func(samples)
            resample_count += 1

        if len(valid_samples) >= self.args.generation_batch_size:
            process_slice = slice(
                self.accelerator.process_index * len(samples),
                (self.accelerator.process_index + 1) * len(samples),
            )
            samples = valid_samples[:self.args.generation_batch_size][process_slice]
            rewards_per_func = torch.cat(valid_rewards_per_func)[:self.args.generation_batch_size]
        else:
            logger.warning(f'There are still std=0 groups present after {self.max_resample_times} retries.')
            samples, rewards_per_func = origin_data

        return samples, rewards_per_func

    def compute_std(self, samples: List[GRPOSample], rewards_per_func: torch.Tensor) -> torch.Tensor:
        """Compute the standard deviation of the rewards per function."""
        device = self.accelerator.device
        rewards = (rewards_per_func * self.reward_weights.unsqueeze(0)).nansum(dim=1)

        mode = 'train' if self.model.training else 'eval'
        num_generations = self.num_generations if mode == 'train' else self.num_generations_eval
        if not self.dynamic_num_samples:
            grouped_rewards = rewards.view(-1, num_generations)
            # Handle edge case when num_generations_eval=1
            if num_generations > 1:
                group_rewards_std = grouped_rewards.std(dim=1).repeat_interleave(num_generations)
            else:
                group_rewards_std = torch.zeros_like(rewards)
            return group_rewards_std
        else:
            prompt_ids = gather_object([s.prompt_id for s in samples])
            request_ids = gather_object([s.request_id for s in samples])
            device = self.accelerator.device
            unique_indices = self._get_last_indices(request_ids)
            unique_request_ids = [request_ids[i] for i in unique_indices.cpu()]
            unique_prompt_ids = [prompt_ids[i] for i in unique_indices.cpu()]

            unique_rewards = rewards[unique_indices]
            prompt_to_indices = {}
            for idx, pid in enumerate(unique_prompt_ids):
                prompt_to_indices.setdefault(pid, []).append(idx)

            prompt_stds = torch.zeros(len(unique_rewards), device=device)
            for pid, idxs in prompt_to_indices.items():
                idx_tensor = torch.tensor(idxs, device=device)
                r_group = unique_rewards[idx_tensor]
                # Edge case: when group size is 1
                prompt_stds[idx_tensor] = r_group.std() if len(idxs) > 1 else 0.0
            rid_to_idx = {rid: idx for idx, rid in enumerate(unique_request_ids)}
            indices_in_unique = torch.tensor([rid_to_idx[r] for r in request_ids], device=device)
            rewards_std = prompt_stds[indices_in_unique]

            return rewards_std

    @profiling_decorator
    def _prepare_batch_inputs(self, samples: List[GRPOSample]) -> List[Dict[str, Any]]:
        """
        Prepare the final batch inputs with ref/old_policy logps and other fields for RL training.

        Args:
            samples (List[GRPOSample]): List of local on-policy samples.

        Returns:
            List[Dict[str, Any]]: A list of prepared batch dicts (model inputs + ``grpo_batch``),
            organized as [steps_per_generation][batch_size]
        """
        template = self.template
        gas_chunks = self.split_by_mini_batches(samples)
        ga_batch_encoded_inputs: List[Dict[str, Any]] = []
        for batch in gas_chunks:
            teacher_model_inputs = teacher_grpo_batch = None
            with self._template_context(template):
                for s in batch:
                    encoded_inputs = encode_sample(s, template)
                    encoded_inputs.pop('_extra_kwargs', None)  # pop add_eos
                    s.encoded = encoded_inputs
                model_inputs, grpo_batch = collate_to_grpo_micro_batch(
                    batch, template, device=self.model.device, use_logits_to_keep=True)
                # OPSD: the local teacher forwards its own (teacher_prompt + same response)
                # encoding, so collate a separate teacher micro-batch (different length).
                has_opsd_batch = build_opsd_samples(batch)
                is_opsd = (has_opsd_batch and (self._has_teacher_explicit() or self._is_dynamic_self_distillation))
                if is_opsd:
                    teacher_model_inputs, teacher_grpo_batch = self._collate_teacher_opsd_batch(batch, template)

            model_inputs.pop('labels', None)
            batch_encoded_inputs = {'model_inputs': model_inputs, 'grpo_batch': grpo_batch}
            if self.dynamic_num_samples and self.is_multimodal:
                batch_encoded_inputs['_origin_data'] = batch
            if self._has_teacher and self.use_teacher_api:
                # OPD-RL API teacher: keep the SP-gathered chunk samples so teacher routing/fetch
                # aligns with each grpo_batch (local `samples` are shorter/reordered under SP>1).
                batch_encoded_inputs['_chunk_samples'] = batch
            origin_data = batch if (self.dynamic_num_samples and self.is_multimodal) else None

            with torch.no_grad(), disable_gradient_checkpointing(self.model, self.args.gradient_checkpointing_kwargs):
                grpo_batch.old_per_token_logps = (
                    self._get_per_token_logps_and_entropies(
                        self.model, model_inputs, grpo_batch, origin_data=origin_data)[0])
                if self.beta == 0.0:
                    ref_per_token_logps = None
                elif self.ref_model is not None:
                    with disable_gradient_checkpointing(self.ref_model, self.args.gradient_checkpointing_kwargs):
                        ref_per_token_logps = \
                            self._get_per_token_logps_and_entropies(
                                self.ref_model, model_inputs, grpo_batch, origin_data=origin_data)[0]
                else:
                    with self.null_ref_context():
                        ref_per_token_logps = \
                            self._get_per_token_logps_and_entropies(
                                self.model, model_inputs, grpo_batch, origin_data=origin_data)[0]
                grpo_batch.ref_per_token_logps = ref_per_token_logps
                # OPD-RL: local teacher logp on the sampled tokens (API path filled later).
                if should_compute_local_teacher_logps(
                        has_teacher_explicit=self._has_teacher_explicit(),
                        is_dynamic_self_distillation=self._is_dynamic_self_distillation,
                        use_teacher_api=self.use_teacher_api,
                        has_opsd_batch=is_opsd,
                ):
                    grpo_batch.teacher_per_token_logps = self._compute_teacher_logps(
                        model_inputs,
                        grpo_batch,
                        origin_data=origin_data,
                        teacher_model_inputs=teacher_model_inputs,
                        teacher_grpo_batch=teacher_grpo_batch)
            ga_batch_encoded_inputs.append(batch_encoded_inputs)

        # --- log completion lengths ---
        mode = 'train' if self.model.training else 'eval'
        device = self.accelerator.device
        local_lengths = [inp['grpo_batch'].completion_mask.sum(1).tolist() for inp in ga_batch_encoded_inputs]
        total_lengths = self._gather_and_flatten(local_lengths, dtype=torch.float32, device=device, flatten_level=1)

        # Store num_items_in_batch for DAPO loss (total completion tokens across all processes)
        num_items_in_batch = total_lengths.sum()
        for batch_encoded in ga_batch_encoded_inputs:
            batch_encoded['grpo_batch'].num_items_in_batch = num_items_in_batch

        self._metrics[mode]['completions/mean_length'].append(total_lengths.mean().item())
        self._metrics[mode]['completions/min_length'].append(total_lengths.min().item())
        self._metrics[mode]['completions/max_length'].append(total_lengths.max().item())

        # --- log completion clipped ratio ---
        local_trunc_masks = [inp['grpo_batch'].truncated_mask.tolist() for inp in ga_batch_encoded_inputs]
        total_trunc_masks = self._gather_and_flatten(
            local_trunc_masks, dtype=torch.bool, device=device, flatten_level=1)

        if not self.dynamic_num_samples:
            clipped_ratio = total_trunc_masks.sum().item() / total_lengths.shape[0]
            self._metrics[mode]['completions/clipped_ratio'].append(clipped_ratio)

            if all(s.rollout_infos and 'num_turns' in s.rollout_infos for s in samples):
                num_turns = torch.tensor(gather_object([s.rollout_infos['num_turns'] for s in samples]), device=device)
                self._metrics[mode]['num_turns'].append(num_turns.float().mean().item())
        else:
            request_ids = gather_object([s.request_id for s in samples])
            last_indices = self._get_last_indices(request_ids)

            final_trunc_masks = total_trunc_masks[last_indices]
            clipped_ratio = final_trunc_masks.sum().item() / final_trunc_masks.shape[0]
            self._metrics[mode]['completions/clipped_ratio'].append(clipped_ratio)

            if all(s.rollout_infos and 'num_turns' in s.rollout_infos for s in samples):
                num_turns_all = torch.tensor(
                    gather_object([s.rollout_infos['num_turns'] for s in samples]), device=device)
                final_num_turns = num_turns_all[last_indices]
                self._metrics[mode]['num_turns'].append(final_num_turns.float().mean().item())

        return ga_batch_encoded_inputs

    def _apply_chat_template_to_messages_list(self, messages_list: DataType):
        prompts_text = []
        for messages in messages_list:
            remove_response(messages)
            template_inputs = TemplateInputs.from_dict({'messages': messages})
            res = self.template.encode(template_inputs)
            prompts_text.append(self.template.safe_decode(res['input_ids']))
        return prompts_text

    @profiling_decorator
    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
        # Compute the per-token log probabilities for the model, return_outputs=True in mini-batch training
        if isinstance(inputs, list):
            assert len(inputs) == 1
            inputs = inputs[0]
        model_inputs = inputs['model_inputs']
        grpo_batch = inputs['grpo_batch']
        origin_data = inputs.get('_origin_data')
        if self.use_liger_loss:
            unwrapped_model = self.accelerator.unwrap_model(model)
            return self._forward_redirection(
                model, unwrapped_model,
                lambda *_, **__: self.compute_liger_loss(unwrapped_model, model_inputs, grpo_batch), **model_inputs)
        else:
            return self._compute_loss(model, model_inputs, grpo_batch, origin_data)

    def _compute_loss(self, model, model_inputs, grpo_batch, origin_data=None):
        mode = 'train' if self.model.training else 'eval'

        # Check batch size and decide processing strategy
        batch_size = grpo_batch.seq_lengths.shape[0] if self.template.padding_free else model_inputs['input_ids'].shape[
            0]
        expected_bs = self.args.per_device_train_batch_size if mode == 'train' else self.args.per_device_eval_batch_size

        should_chunk = self.dynamic_num_samples and any(gather_object([batch_size > expected_bs]))
        if not should_chunk:
            return self._compute_loss_single(model, model_inputs, grpo_batch)
        else:
            # maybe dynamic rollout num for multi-turn training
            return self._compute_loss_chunked(model, model_inputs, grpo_batch, origin_data)

    def _compute_loss_single(self, model, model_inputs, grpo_batch):
        """Original loss computation logic for single batch processing."""
        loss, metrics_data = self._compute_loss_and_metrics(model, model_inputs, grpo_batch)
        self._update_metrics(metrics_data)
        return loss

    def _compute_fipo_influence(self, log_ratio: torch.Tensor, coef_1: torch.Tensor, advantages: torch.Tensor,
                                completion_mask: torch.Tensor) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
        """Compute FIPO token-level influence weight from Future-KL divergence."""
        future_kl_delta = log_ratio.masked_fill(~completion_mask, 0.0)

        # Dual-Clip participation mask: high-ratio tokens do not contribute to Future-KL.
        if self.args.delta is not None:
            delta = torch.as_tensor(self.args.delta, dtype=log_ratio.dtype, device=log_ratio.device)
            high_ratio_mask = coef_1 > delta
            future_kl_delta = torch.where(high_ratio_mask, torch.zeros_like(future_kl_delta), future_kl_delta)

        seq_len = future_kl_delta.shape[1]
        future_kl = torch.zeros_like(future_kl_delta)
        positions = torch.arange(seq_len, device=log_ratio.device).unsqueeze(1)
        gamma = torch.as_tensor(self.fipo_gamma, dtype=log_ratio.dtype, device=log_ratio.device)
        chunk_size = 128
        for chunk_start in range(0, seq_len, chunk_size):
            chunk_end = min(seq_len, chunk_start + chunk_size)
            chunk_positions = torch.arange(chunk_start, chunk_end, device=log_ratio.device).unsqueeze(0)
            distance = chunk_positions - positions
            future_mask = distance >= 0
            decay_block = torch.pow(gamma, distance.clamp(min=0)) * future_mask.to(log_ratio.dtype)
            future_kl += torch.matmul(future_kl_delta[:, chunk_start:chunk_end], decay_block.t())
        future_kl = future_kl.masked_fill(~completion_mask, 0.0)

        influence_weight = torch.exp(future_kl)

        if self.fipo_clip_range:
            high = 1 + self.fipo_clip_range
            low = 1.0 if self.fipo_clip_high_only else 1 - self.fipo_clip_range
            influence_weight = torch.clamp(influence_weight, min=low, max=high)
        influence_weight = influence_weight.detach()

        # avoid amplifying negative-advantage tokens with very high IS ratios.
        safety_mask = torch.ones_like(completion_mask, dtype=torch.bool)
        if self.fipo_safety_threshold is not None:
            negative_advantage = advantages < 0
            high_is_ratio = coef_1 > self.fipo_safety_threshold
            safety_mask = ~(negative_advantage & high_is_ratio)
            influence_weight = torch.where(safety_mask, influence_weight,
                                           torch.clamp(influence_weight, min=0.8, max=1.0))

        metrics = {
            'future_kl': future_kl,
            'influence_weight': influence_weight,
            'safety_mask': safety_mask,
        }
        return influence_weight, metrics

    def _compute_loss_and_metrics(self, model, model_inputs: Dict[str, Any], grpo_batch: GRPOBatch):
        """Core loss computation without metrics recording."""
        mode = 'train' if self.model.training else 'eval'
        completion_mask = grpo_batch.completion_mask
        truncated_mask = grpo_batch.truncated_mask
        per_token_logps, entropies = self._get_per_token_logps_and_entropies(
            model, model_inputs, grpo_batch, compute_entropy=self.compute_entropy)

        entropy_mask = None
        entropy_metrics = {}

        if self.compute_entropy:
            # fill the padded token with NaN
            entropies = entropies.masked_fill(completion_mask == 0, float('nan'))
            if self.args.log_entropy:
                per_completion_entropies_mean = torch.nanmean(entropies, dim=1)
                global_per_completion_entropies_mean = gather(per_completion_entropies_mean)
                entropy_metrics = {
                    'entropy_logs': global_per_completion_entropies_mean.tolist(),
                    'entropy_mean': global_per_completion_entropies_mean.nanmean().item(),
                    'entropy_max': nanmax(global_per_completion_entropies_mean).item(),
                    'entropy_min': nanmin(global_per_completion_entropies_mean).item()
                }

            # compute the entropy threshold across all tokens in the batch
            if self.args.top_entropy_quantile < 1.0:
                entropy_threshold = torch.nanquantile(entropies.flatten().float(), 1 - self.top_entropy_quantile)
                entropy_metrics['entropy_threshold'] = entropy_threshold.item()
                entropy_mask = entropies >= entropy_threshold

        # apply the completion_mask to exclude loss and metrics for overlong completions
        if self.overlong_filter and any(truncated_mask):
            if all(truncated_mask):
                logger.info('All completions are overlong and truncated, '
                            'resulting in NaN some values for some metrics (e.g., KL)')
            truncated_mask = truncated_mask.unsqueeze(-1).expand_as(completion_mask)
            completion_mask = completion_mask & (~truncated_mask)

        per_token_kl = None
        if self.beta != 0.0 and not self.kl_in_reward:
            ref_per_token_logps = grpo_batch.ref_per_token_logps
            safe_ratio = torch.clamp(ref_per_token_logps - per_token_logps, min=-20, max=20)
            per_token_kl = torch.clamp(torch.exp(safe_ratio) - safe_ratio - 1, min=-10, max=10)

        advantages = grpo_batch.advantages
        old_per_token_logps = (
            per_token_logps.detach() if grpo_batch.old_per_token_logps is None else grpo_batch.old_per_token_logps)

        # Compute rollout diagnostic metrics and apply IS correction if enabled
        rollout_correction_metrics = {}
        should_compute_rollout_metrics = (
            self.rollout_importance_sampling_mode is not None or self.log_rollout_offpolicy_metrics)

        local_has_rollout = grpo_batch.rollout_per_token_logps is not None
        should_compute_rollout_metrics = should_compute_rollout_metrics and all(gather_object([local_has_rollout]))
        rollout_is_weights = None
        if (not self.disable_rollout_importance_sampling and should_compute_rollout_metrics):
            rollout_per_token_logps = grpo_batch.rollout_per_token_logps

            # Compute diagnostic metrics (KL, PPL, etc.) for monitoring off-policy gap
            rollout_correction_metrics = self._compute_rollout_offpolicy_metrics(old_per_token_logps,
                                                                                 rollout_per_token_logps,
                                                                                 completion_mask)

            rollout_log_ratio, rollout_is_weights = self._get_rollout_is_correction(old_per_token_logps,
                                                                                    rollout_per_token_logps,
                                                                                    completion_mask)
            if rollout_log_ratio is not None:
                is_metrics = self._compute_is_correction_metrics(rollout_log_ratio, rollout_is_weights, completion_mask)
                rollout_correction_metrics.update(is_metrics)

            pass
        # rollout_is_weights is a local variable, initialized to None above

        log_ratio = per_token_logps - old_per_token_logps
        if self.importance_sampling_level == 'token':
            log_importance_weights = log_ratio
        elif self.importance_sampling_level in ['sequence', 'sequence_token']:
            seq_level_log_weights = ((log_ratio * completion_mask).sum(-1)
                                     / completion_mask.sum(-1).clamp(min=1.0)).unsqueeze(-1)
            if self.importance_sampling_level == 'sequence':
                log_importance_weights = seq_level_log_weights
            else:
                # GSPO-token: sg[si(θ)] * πθ(yi,t)/sg[πθ(yi,t)]
                seq_level_log_weight = seq_level_log_weights.detach()
                log_importance_weights = per_token_logps - per_token_logps.detach() + seq_level_log_weight
        else:
            raise ValueError(
                f"Unknown importance sampling level: {self.importance_sampling_level}. Possible values are 'token' "
                "and 'sequence'.")

        coef_1 = torch.exp(log_importance_weights)

        # advantages is per-token [B, T] (expanded at batch construction so the OPD-RL signed
        # teacher log-ratio is added per token). Edge loss types that need a per-sequence
        # advantage (real / fipo / off_policy_sequence_mask) are not supported with a teacher.
        if self._has_teacher and (self.loss_type in ['real', 'fipo']
                                  or self.off_policy_sequence_mask_delta is not None):
            raise ValueError(f'OPD-RL (teacher) does not support loss_type={self.loss_type!r} / '
                             'off_policy_sequence_mask.')

        fipo_metrics = None
        if self.loss_type == 'cispo':
            clamped_ratios = torch.clamp(coef_1, max=self.epsilon_high).detach()
            per_token_loss = -clamped_ratios * advantages * per_token_logps
        elif self.loss_type == 'sapo':
            gate_pos = torch.sigmoid(self.tau_pos * (coef_1 - 1)) * (4.0 / self.tau_pos)
            gate_neg = torch.sigmoid(self.tau_neg * (coef_1 - 1)) * (4.0 / self.tau_neg)
            is_positive = advantages > 0
            soft_gate = torch.where(is_positive, gate_pos, gate_neg)

            per_token_loss = -soft_gate * advantages
        elif self.loss_type == 'real':
            per_token_loss = torch.zeros_like(per_token_logps)
        elif self.loss_type in ['grpo', 'bnpo', 'dr_grpo', 'dapo', 'fipo']:
            if self.loss_type == 'fipo':
                fipo_weight, fipo_metrics = self._compute_fipo_influence(log_ratio, coef_1, advantages, completion_mask)

            coef_2 = torch.clamp(coef_1, 1 - self.epsilon_low, 1 + self.epsilon_high)
            if self.args.delta is not None:
                coef_1 = torch.clamp(coef_1, max=self.args.delta)

            per_token_loss1 = coef_1 * advantages
            per_token_loss2 = coef_2 * advantages
            per_token_loss = -torch.min(per_token_loss1, per_token_loss2)
            if self.loss_type == 'fipo':
                per_token_loss = per_token_loss * fipo_weight
        if entropy_mask is not None:
            per_token_loss = per_token_loss * entropy_mask
        if per_token_kl is not None:
            per_token_loss = per_token_loss + self.beta * per_token_kl

        if rollout_is_weights is not None and self.rollout_importance_sampling_mode is not None:
            per_token_loss = per_token_loss * rollout_is_weights

        if self.off_policy_sequence_mask_delta is not None:
            old_policy_per_token_logps = (
                grpo_batch.rollout_per_token_logps
                if grpo_batch.rollout_per_token_logps is not None else old_per_token_logps)
            # advantages is per-token [B, T] (constant across tokens without a teacher); the mask
            # needs a per-sequence scalar.
            seq_advantages = (advantages * completion_mask).sum(-1) / completion_mask.sum(-1).clamp(min=1.0)
            off_policy_seq_mask = self._compute_off_policy_sequence_mask(per_token_logps, old_policy_per_token_logps,
                                                                         completion_mask, seq_advantages)
            # Expand sequence mask to token level and apply to completion_mask
            off_policy_seq_mask_expanded = off_policy_seq_mask.unsqueeze(-1).expand_as(completion_mask)
            completion_mask = completion_mask & off_policy_seq_mask_expanded

        if self.loss_type in ['grpo', 'sapo']:
            # completion_mask is now always [batch_size, seq_len] after pad_back
            loss = ((per_token_loss * completion_mask).sum(-1) / completion_mask.sum(-1).clamp(min=1.0)).mean()
        elif self.loss_type == 'bnpo':
            loss = (per_token_loss * completion_mask).sum() / completion_mask.sum().clamp(min=1.0)
        elif self.loss_type == 'dr_grpo':
            batch_size = completion_mask.shape[0]
            loss = (per_token_loss * completion_mask).sum() / (batch_size * self.max_completion_length)
        elif self.loss_type == 'real':
            global_scores = (log_ratio * completion_mask).sum(-1) / completion_mask.sum(-1).clamp(min=1.0)

            group_scores = global_scores.view(-1, self.num_generations)
            # advantages is per-token [B, T] (constant across tokens without a teacher); reduce
            # to a per-sequence scalar for the group pos/neg split.
            seq_advantages = (advantages * completion_mask).sum(-1) / completion_mask.sum(-1).clamp(min=1.0)
            group_rewards = seq_advantages.view(-1, self.num_generations)

            pos_mask = (group_rewards > 0)
            neg_mask = (group_rewards <= 0)
            valid_mask = (pos_mask.sum(dim=1) != 0) & (neg_mask.sum(dim=1) != 0)

            if not valid_mask.any():
                loss = torch.tensor(0., device=global_scores.device) * global_scores.mean()
            else:
                batch_scores = group_scores[valid_mask]
                batch_pos_mask = pos_mask[valid_mask]
                batch_neg_mask = neg_mask[valid_mask]

                scaled_scores = batch_scores / self.real_tau
                zeros = torch.zeros(batch_scores.size(0), 1, device=batch_scores.device, dtype=batch_scores.dtype)

                # Negative Loss: log(1 + sum(e^{S_neg}))
                neg_input = scaled_scores.masked_fill(~batch_neg_mask, float('-inf'))
                neg_loss = torch.logsumexp(torch.cat([neg_input, zeros], dim=1), dim=1)

                # Positive Loss: log(1 + sum(e^{-S_pos}))
                pos_input = (-scaled_scores).masked_fill(~batch_pos_mask, float('-inf'))
                pos_loss = torch.logsumexp(torch.cat([pos_input, zeros], dim=1), dim=1)

                loss = (neg_loss + pos_loss).sum() / group_rewards.size(0)

            if self.beta != 0.0:
                kl_loss = (per_token_kl * completion_mask).sum() / completion_mask.sum().clamp(min=1.0)
                loss = loss + kl_loss * self.beta
        elif self.loss_type in ['cispo', 'dapo', 'fipo']:
            # CISPO, DAPO, and FIPO: Normalize by total completion tokens across all processes
            normalizer = grpo_batch.num_items_in_batch / self.accelerator.num_processes
            loss = (per_token_loss * completion_mask).sum() / normalizer
        else:
            raise ValueError(f'Unknown loss type: {self.loss_type}')

        completion_token_count = completion_mask.sum().clamp(min=1.0)

        def masked_batch_mean(x):
            # compute for token-level average
            if x.shape[1] == 1:  # when importance_sampling_level == "sequence"
                return x.mean()
            else:
                return (x * completion_mask).sum() / completion_token_count

        # Prepare metrics data
        metrics_data = {
            'mode': mode,
            'entropy': entropy_metrics,
            'completion_mask': completion_mask,
            'completion_token_count': completion_token_count,
        }

        if fipo_metrics is not None:
            fipo_future_kl = masked_batch_mean(fipo_metrics['future_kl'])
            fipo_influence_weight = masked_batch_mean(fipo_metrics['influence_weight'])
            fipo_safety_keep = masked_batch_mean(fipo_metrics['safety_mask'].float())
            metrics_data['fipo'] = {
                'future_kl_mean': self.accelerator.gather_for_metrics(fipo_future_kl).nanmean().item(),
                'influence_weight_mean': self.accelerator.gather_for_metrics(fipo_influence_weight).nanmean().item(),
                'safety_keep_ratio': self.accelerator.gather_for_metrics(fipo_safety_keep).nanmean().item(),
            }

        if per_token_kl is not None:
            mean_kl = masked_batch_mean(per_token_kl)
            metrics_data['kl'] = self.accelerator.gather_for_metrics(mean_kl).nanmean().item()

        # Add rollout correction metrics
        if rollout_correction_metrics:
            metrics_data['rollout_correction'] = rollout_correction_metrics

        # Compute the clipped probability ratios
        if self.loss_type == 'cispo':
            # CISPO: Only track upper bound clipping
            is_cispo_clipped = (coef_1 > self.epsilon_high) & (advantages > 0)
            cispo_clip_ratio = masked_batch_mean(is_cispo_clipped.float())
            gathered_cispo_clip_ratio = self.accelerator.gather_for_metrics(cispo_clip_ratio)
            metrics_data['clipping'] = {'cispo_clip_ratio': gathered_cispo_clip_ratio.nanmean().item()}
        elif self.loss_type in ['sapo', 'real']:
            pass
        else:
            is_low_clipped = (coef_1 < 1 - self.epsilon_low) & (advantages < 0)
            is_high_clipped = (coef_1 > 1 + self.epsilon_high) & (advantages > 0)
            is_region_clipped = is_low_clipped | is_high_clipped

            low_clip = masked_batch_mean(is_low_clipped.float())
            high_clip = masked_batch_mean(is_high_clipped.float())
            clip_ratio = masked_batch_mean(is_region_clipped.float())

            gathered_low_clip = self.accelerator.gather_for_metrics(low_clip)
            gathered_high_clip = self.accelerator.gather_for_metrics(high_clip)
            gathered_clip_ratio = self.accelerator.gather_for_metrics(clip_ratio)

            metrics_data['clipping'] = {
                'low_clip_mean': gathered_low_clip.nanmean().item(),
                'low_clip_min': nanmin(gathered_low_clip).item(),
                'high_clip_mean': gathered_high_clip.nanmean().item(),
                'high_clip_max': nanmax(gathered_high_clip).item(),
                'region_clip_mean': gathered_clip_ratio.nanmean().item()
            }
        if mode == 'train' and self.chord_sft_iterator is not None:
            loss = compute_chord_loss(self, grpo_loss=loss)

        return loss, metrics_data

    def _update_metrics(self, metrics_data):
        """Update metrics from metrics_data."""
        mode = metrics_data['mode']

        # Update entropy metrics
        if metrics_data['entropy']:
            entropy_metrics = metrics_data['entropy']
            if 'entropy_logs' in entropy_metrics:
                self._logs['entropy'].extend(entropy_metrics['entropy_logs'])
                self._metrics[mode]['entropy/mean'].append(entropy_metrics['entropy_mean'])
                self._metrics[mode]['entropy/max'].append(entropy_metrics['entropy_max'])
                self._metrics[mode]['entropy/min'].append(entropy_metrics['entropy_min'])
            if 'entropy_threshold' in entropy_metrics:
                self._metrics[mode]['entropy/threshold'].append(entropy_metrics['entropy_threshold'])

        # Update KL metrics
        if 'kl' in metrics_data:
            self._metrics[mode]['kl'].append(metrics_data['kl'])

        # Update vLLM correction metrics
        if 'rollout_correction' in metrics_data:
            rollout_metrics = metrics_data['rollout_correction']
            for key, value in rollout_metrics.items():
                self._metrics[mode][f'rollout_correction/{key}'].append(value)

        # Update FIPO metrics
        if 'fipo' in metrics_data:
            for key, value in metrics_data['fipo'].items():
                self._metrics[mode][f'fipo/{key}'].append(value)

        # Update clipping metrics
        if 'clipping' in metrics_data:
            clipping = metrics_data['clipping']
            if 'cispo_clip_ratio' in clipping:
                # CISPO
                self._metrics[mode]['cispo_clip_ratio'].append(clipping['cispo_clip_ratio'])
            else:
                self._metrics[mode]['clip_ratio/low_mean'].append(clipping['low_clip_mean'])
                self._metrics[mode]['clip_ratio/low_min'].append(clipping['low_clip_min'])
                self._metrics[mode]['clip_ratio/high_mean'].append(clipping['high_clip_mean'])
                self._metrics[mode]['clip_ratio/high_max'].append(clipping['high_clip_max'])
                self._metrics[mode]['clip_ratio/region_mean'].append(clipping['region_clip_mean'])

    def _compute_loss_chunked(self, model, model_inputs: Dict[str, Any], grpo_batch: GRPOBatch, origin_data=None):
        """
        Compute loss in **fixed-size chunks** to reduce peak GPU memory.

        The function guarantees that **all ranks step through the same number of
        chunks**, so that collective communication remain synchronized
        even when local ``batch_size`` differs.
        """
        mode = 'train' if self.model.training else 'eval'
        chunk_size = self.args.per_device_train_batch_size if mode == 'train' else self.args.per_device_eval_batch_size
        batch_size = grpo_batch.seq_lengths.shape[0] if self.template.padding_free else model_inputs['input_ids'].shape[
            0]

        # Decide how many chunks every rank must run
        batch_sizes = gather_object([batch_size])
        chunks_per_device = [(bs + chunk_size - 1) // chunk_size for bs in batch_sizes]
        max_chunks = max(chunks_per_device)

        # Re-compute chunk size so that max_chunks * new_chunk_size covers entire batch
        new_chunk_size = (batch_size + max_chunks - 1) // max_chunks

        losses, weights = [], []
        all_metrics_data = []
        chunk_model_inputs, chunk_grpo_batch = {}, None
        for chunk_idx in range(max_chunks):
            start_idx = chunk_idx * new_chunk_size
            end_idx = min(start_idx + new_chunk_size, batch_size)

            is_dummy = False
            if start_idx < batch_size:
                chunk_model_inputs, chunk_grpo_batch = self.get_chunked_inputs(model_inputs, grpo_batch, start_idx,
                                                                               end_idx, origin_data)
                chunk_weight = end_idx - start_idx
            else:
                is_dummy = True
                chunk_weight = 0

            # Compute loss and metrics for this chunk
            chunk_loss, chunk_metrics_data = self._compute_loss_and_metrics(model, chunk_model_inputs, chunk_grpo_batch)

            if not is_dummy:
                losses.append(chunk_loss * chunk_weight)
                weights.append(chunk_weight)
                all_metrics_data.append((chunk_metrics_data, chunk_weight))
            else:
                # # Add dummy loss to computation graph to trigger ZeRO-3 backward hooks
                losses.append(chunk_loss * 0.0)

        # Compute weighted average loss
        total_weight = sum(weights)
        if total_weight > 0:
            final_loss = torch.stack(losses).sum() / total_weight
        else:
            final_loss = torch.stack(losses).sum()

        # Aggregate metrics across all chunks
        self._aggregate_and_update_metrics(all_metrics_data, mode)

        return final_loss

    def _aggregate_and_update_metrics(self, all_metrics_data, mode):
        """Aggregate metrics from multiple chunks and update global metrics."""
        if not all_metrics_data:
            return

        # Separate metrics by type for aggregation
        entropy_logs, entropy_stats, kl_values = [], [], []
        clip_values = {'low': [], 'high': [], 'region': [], 'low_min': [], 'high_max': []}
        cispo_clip_values = []
        entropy_thresholds = []
        fipo_values = {}

        for chunk_metrics, chunk_weight in all_metrics_data:
            chunk_tokens = chunk_metrics['completion_token_count']

            # Collect entropy metrics
            if chunk_metrics['entropy']:
                entropy_metrics = chunk_metrics['entropy']
                if 'entropy_logs' in entropy_metrics:
                    entropy_logs.extend(entropy_metrics['entropy_logs'])
                    entropy_stats.append({
                        'mean': entropy_metrics['entropy_mean'],
                        'max': entropy_metrics['entropy_max'],
                        'min': entropy_metrics['entropy_min']
                    })
                if 'entropy_threshold' in entropy_metrics:
                    entropy_thresholds.append(entropy_metrics['entropy_threshold'])

            # Collect KL metrics
            if 'kl' in chunk_metrics:
                kl_values.append(chunk_metrics['kl'])

            # Collect FIPO metrics (weighted by tokens)
            if 'fipo' in chunk_metrics:
                weight = chunk_tokens.item() if hasattr(chunk_tokens, 'item') else chunk_tokens
                for key, value in chunk_metrics['fipo'].items():
                    fipo_values.setdefault(key, []).append((value, weight))

            # Collect clipping metrics (weighted by tokens)
            if 'clipping' in chunk_metrics:
                clipping = chunk_metrics['clipping']
                weight = chunk_tokens.item() if hasattr(chunk_tokens, 'item') else chunk_tokens
                if 'cispo_clip_ratio' in clipping:
                    cispo_clip_values.append((clipping['cispo_clip_ratio'], weight))
                else:
                    clip_values['low'].append((clipping['low_clip_mean'], weight))
                    clip_values['high'].append((clipping['high_clip_mean'], weight))
                    clip_values['region'].append((clipping['region_clip_mean'], weight))
                    clip_values['low_min'].append(clipping['low_clip_min'])
                    clip_values['high_max'].append(clipping['high_clip_max'])

        # Build aggregated metrics
        aggregated_metrics = {'mode': mode, 'entropy': {}}

        # Aggregate entropy
        if entropy_logs:
            # Directly update entropy logs
            aggregated_metrics['entropy'] = {
                'entropy_logs': entropy_logs,
                'entropy_mean': sum(s['mean'] for s in entropy_stats) / len(entropy_stats),
                'entropy_max': max(s['max'] for s in entropy_stats),
                'entropy_min': min(s['min'] for s in entropy_stats)
            }
        if entropy_thresholds:
            aggregated_metrics['entropy']['entropy_threshold'] = sum(entropy_thresholds) / len(entropy_thresholds)

        # Aggregate KL
        if kl_values:
            aggregated_metrics['kl'] = sum(kl_values) / len(kl_values)

        # Aggregate clipping (token-weighted averages)
        def weighted_avg(values):
            return sum(v * w for v, w in values) / sum(w for _, w in values)

        if cispo_clip_values:
            # CISPO specific metric
            aggregated_metrics['clipping'] = {'cispo_clip_ratio': weighted_avg(cispo_clip_values)}
        elif clip_values['low']:
            # Two-sided clipping metrics
            aggregated_metrics['clipping'] = {
                'low_clip_mean': weighted_avg(clip_values['low']),
                'low_clip_min': min(clip_values['low_min']),
                'high_clip_mean': weighted_avg(clip_values['high']),
                'high_clip_max': max(clip_values['high_max']),
                'region_clip_mean': weighted_avg(clip_values['region'])
            }

        if fipo_values:
            aggregated_metrics['fipo'] = {key: weighted_avg(values) for key, values in fipo_values.items()}

        # Update metrics
        self._update_metrics(aggregated_metrics)

    def _unpad_logps_and_entropies(self,
                                   logps: torch.Tensor,
                                   entropies: Optional[torch.Tensor],
                                   logits_to_keep: int,
                                   batch_size: int,
                                   seq_lengths: torch.Tensor,
                                   compute_entropy: bool = False) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        """
        Restore logps and entropies from rmpad format [1, total_nnz] to batch format [batch_size, max_seq_len].

        Args:
            logps: Per-token log probabilities in rmpad format [1, total_nnz]
            entropies: Per-token entropies in rmpad format [1, total_nnz] or None
            logits_to_keep: Number of tokens to keep per sequence
            batch_size: Number of sequences in the batch
            seq_lengths: Actual sequence lengths [batch_size]
            compute_entropy: Whether entropy was computed

        Returns:
            logps: Restored log probabilities [batch_size, logits_to_keep]
            entropies: Restored entropies [batch_size, logits_to_keep] or None
        """
        logps, _ = pad_logps_back_to_batch(
            logps_rmpad=logps, logits_to_keep=logits_to_keep, batch_size=batch_size, seq_lengths=seq_lengths)

        if compute_entropy and entropies is not None:
            entropies, _ = pad_logps_back_to_batch(
                logps_rmpad=entropies, logits_to_keep=logits_to_keep, batch_size=batch_size, seq_lengths=seq_lengths)

        return logps, entropies

    def _get_logps_via_sp(self,
                          model: torch.nn.Module,
                          model_inputs: Dict[str, Any],
                          grpo_batch: GRPOBatch,
                          logits_to_keep: int,
                          input_ids: torch.Tensor,
                          compute_entropy: bool = False) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        """Get per token logps via sequence parallel, returns rmpad format [1, total_nnz] for padding_free mode"""
        sequence_parallel.prepare_inputs(model_inputs)
        with self._template_context(self.template, model_inputs):
            output = model(**model_inputs)
            logits = output.logits
        # split input_ids to labels
        position_ids = sequence_parallel.real_position_ids
        _, _, labels, _, _, _, _ = sequence_parallel.pad_and_split_inputs(
            None, None, input_ids.clone(), None, None, None, real_position_ids=position_ids)

        labels = torch.where(labels == -100, self.processing_class.pad_token_id, labels)
        logits = logits / self.temperature
        per_token_logps = selective_log_softmax(logits, labels)
        entropies = None
        per_token_logps, _ = GatherLoss.apply(per_token_logps, labels, 1, position_ids)
        if compute_entropy:
            entropies = entropy_from_logits(logits)
            entropies, _ = GatherLoss.apply(entropies, labels, 1, position_ids)

        if self.template.padding_free:
            # In padding_free mode, we need to extract completion tokens from gathered data.
            # The behavior differs based on rp_world_size:
            # - rp_world_size > 1: Each sequence is padded to world_size * 2 multiple (per-sequence padding)
            # - rp_world_size == 1: Entire data is padded to world_size multiple (end padding only)
            seq_lengths = grpo_batch.seq_lengths
            batch_size = seq_lengths.shape[0]
            rp_world_size = sequence_parallel.rp_world_size

            if rp_world_size > 1:
                # With ring parallel: GatherLoss pads each sequence to world_size * 2 multiple
                # Data layout after gather: [seq1_data, seq1_padding, seq2_data, seq2_padding, ...]
                # - Original data is at [offset:offset+orig_len]
                # - Padding is at [offset+orig_len:offset+padded_len]

                # Get original sequence boundaries (before padding)
                cu_seqlens_orig = get_cu_seqlens_from_position_ids(position_ids)

                # Get padded sequence boundaries (for offset calculation)
                padded_position_ids = sequence_parallel.pad(position_ids, padding_value=-1, position_ids=position_ids)
                cu_seqlens_padded = get_cu_seqlens_from_position_ids(padded_position_ids)

                result_logps = []
                result_entropies = [] if compute_entropy else None
                gathered_logps = per_token_logps.squeeze(0)
                gathered_entropies = entropies.squeeze(0) if compute_entropy else None

                offset = 0
                for i in range(batch_size):
                    # Original sequence length (before SP padding)
                    orig_len = (cu_seqlens_orig[i + 1] - cu_seqlens_orig[i]).item()
                    # Padded sequence length (multiple of world_size * 2)
                    padded_len = (cu_seqlens_padded[i + 1] - cu_seqlens_padded[i]).item()
                    # Actual completion tokens for this sequence
                    actual_len = seq_lengths[i].item()

                    # Extract the last `actual_len` tokens from this sequence's ORIGINAL data region
                    # Due to label shifting (roll -1), per_token_logps[i] predicts token i+1
                    # So completion tokens [prompt_len, total_len) have logps at [prompt_len-1, total_len-1)
                    seq_start = offset + orig_len - actual_len - 1
                    seq_end = offset + orig_len - 1
                    result_logps.append(gathered_logps[seq_start:seq_end])
                    if compute_entropy:
                        result_entropies.append(gathered_entropies[seq_start:seq_end])

                    # Use padded_len for offset because gathered data includes padding
                    offset += padded_len

                per_token_logps = torch.cat(result_logps).unsqueeze(0)
                if compute_entropy:
                    entropies = torch.cat(result_entropies).unsqueeze(0)
            else:
                # Without ring parallel (rp_world_size == 1): Simple gather with end padding only
                # Use input_ids length directly as the authoritative original length
                original_total_len = input_ids.shape[-1]
                # Due to label shifting (roll -1), per_token_logps[i] predicts token i+1.
                start_idx = original_total_len - logits_to_keep - 1
                end_idx = original_total_len - 1
                per_token_logps = per_token_logps[:, start_idx:end_idx]
                if compute_entropy:
                    entropies = entropies[:, start_idx:end_idx]
        else:
            per_token_logps = per_token_logps[:, -logits_to_keep - 1:-1]
            if compute_entropy:
                entropies = entropies[:, -logits_to_keep - 1:-1]

        return per_token_logps, entropies

    def _get_logps_via_local_forward(self,
                                     model: torch.nn.Module,
                                     model_inputs: Dict[str, Any],
                                     logits_to_keep: int,
                                     input_ids: torch.Tensor,
                                     compute_entropy: bool = False) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        """Get per token logps via local forward pass, returns rmpad format [1, total_nnz] for padding_free mode"""
        if 'logits_to_keep' in self.model_kwarg_keys:
            model_inputs['logits_to_keep'] = logits_to_keep + 1

        # Forward pass
        logits = model(**model_inputs).logits

        logits = logits[:, -(logits_to_keep + 1):-1, :]
        logits.div_(self.temperature)

        input_ids_for_logps = input_ids[:, -logits_to_keep:]

        is_padding_free = self.template.padding_free
        if is_padding_free:
            # In padding_free mode, compute logps on flattened tensors
            logits_rmpad = logits.squeeze(0)  # [total_nnz, vocab_size]
            input_ids_rmpad = input_ids_for_logps.squeeze(0)  # [total_nnz]

            # Compute logps on rmpad tensors
            logps = selective_log_softmax(logits_rmpad, input_ids_rmpad)  # [total_nnz]
            logps = logps.unsqueeze(0)  # [1, total_nnz]

            # Compute entropy if needed
            if compute_entropy:
                entropies = entropy_from_logits(logits_rmpad)  # [total_nnz]
                entropies = entropies.unsqueeze(0)  # [1, total_nnz]
            else:
                entropies = None
        else:
            logps = selective_log_softmax(logits, input_ids_for_logps)
            if compute_entropy:
                entropies = entropy_from_logits(logits)
            else:
                entropies = None

        return logps, entropies

    @profiling_decorator
    def _get_per_token_logps_and_entropies(self,
                                           model,
                                           model_inputs: Dict[str, Any],
                                           grpo_batch: GRPOBatch,
                                           compute_entropy=False,
                                           origin_data=None) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        """
        Compute per-token log probabilities and entropies with memory-efficient batching.

        When rollout count is larger than expected, we process in smaller batches
        to control memory usage.
        """
        batch_size = grpo_batch.seq_lengths.shape[0] if self.template.padding_free else model_inputs['input_ids'].shape[
            0]
        mode = 'train' if self.model.training else 'eval'
        expected_bs = self.args.per_device_train_batch_size if mode == 'train' else self.args.per_device_eval_batch_size  # noqa
        should_chunk = self.dynamic_num_samples and any(gather_object([batch_size > expected_bs]))
        if not should_chunk:
            return self._get_per_token_logps_and_entropies_single(
                model, model_inputs, grpo_batch, compute_entropy=compute_entropy)
        else:
            return self._get_per_token_logps_and_entropies_chunked(
                model, model_inputs, grpo_batch, compute_entropy=compute_entropy, origin_data=origin_data)

    def _get_per_token_logps_and_entropies_single(self,
                                                  model,
                                                  model_inputs: Dict[str, Any],
                                                  grpo_batch: GRPOBatch,
                                                  compute_entropy=False) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        logits_to_keep = grpo_batch.logits_to_keep
        input_ids = model_inputs['input_ids']
        is_padding_free = self.template.padding_free
        use_sp = self.template.sequence_parallel_size > 1

        # Store metadata for padding_free restoration
        if is_padding_free:
            original_seq_lengths = grpo_batch.seq_lengths
            batch_size = original_seq_lengths.shape[0]

        if use_sp:
            # Sequence parallel path.
            # In padding_free mode: returns [1, logits_to_keep] format (rmpad, needs unpad)
            # In non-padding_free mode: returns [batch_size, logits_to_keep] format
            logps, entropies = self._get_logps_via_sp(
                model, model_inputs, grpo_batch, logits_to_keep, input_ids, compute_entropy=compute_entropy)
        else:
            # Local forward pass (sole non-SP path; the TRL super() path is equivalent and
            # has been removed to keep the shift logic in one place).
            # Returns [1, total_nnz] in padding_free mode, or [batch_size, logits_to_keep] otherwise
            logps, entropies = self._get_logps_via_local_forward(
                model, model_inputs, logits_to_keep, input_ids, compute_entropy=compute_entropy)

        # Unpad for padding_free mode (both SP and non-SP paths need this)
        if is_padding_free:
            logps, entropies = self._unpad_logps_and_entropies(logps, entropies, logits_to_keep, batch_size,
                                                               original_seq_lengths, compute_entropy)

        return logps, entropies

    def _get_per_token_logps_and_entropies_chunked(self,
                                                   model,
                                                   model_inputs: Dict[str, Any],
                                                   grpo_batch: GRPOBatch,
                                                   compute_entropy=False,
                                                   origin_data=None) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        """
        Compute per-token log-probabilities (and optionally entropies) in **fixed-size
        chunks** to bound peak GPU memory.

        This routine **guarantees that every rank executes the same number of
        chunks**, even when the local batch sizes differ.

        Parameters
        ----------
        model : torch.nn.Module
            The model used to compute log-probs and entropies.
        model_inputs : Dict[str, Any]
            Clean model-forward kwargs (input_ids, attention_mask, etc.).
        grpo_batch : GRPOBatch
            Batch-level RL signals (completion_mask, advantages, etc.).
        compute_entropy : bool, optional
            Whether to compute per-token entropies as well (default: False).
        origin_data : list, optional
            Original samples for multimodal re-encoding during chunking.

        Returns
        -------
        final_logps : torch.Tensor
            Concatenated per-token log-probabilities for the **entire batch**.
        final_entropies : torch.Tensor or None
            Concatenated per-token entropies, or ``None`` if ``compute_entropy`` is
            ``False``.
        """
        batch_size = grpo_batch.seq_lengths.shape[0] if self.template.padding_free else model_inputs['input_ids'].shape[
            0]
        mode = 'train' if self.model.training else 'eval'
        chunk_size = self.args.per_device_train_batch_size if mode == 'train' else self.args.per_device_eval_batch_size

        batch_sizes = gather_object([batch_size])  # list[int]
        chunks_per_device = [(bs + chunk_size - 1) // chunk_size for bs in batch_sizes]
        max_chunks = max(chunks_per_device)

        new_chunk_size = (batch_size + max_chunks - 1) // max_chunks

        all_logps, all_entropies = [], [] if compute_entropy else None

        # Process in chunks
        chunk_model_inputs, chunk_grpo_batch = {}, None
        for chunk_idx in range(max_chunks):
            start_idx = chunk_idx * new_chunk_size
            end_idx = min(start_idx + new_chunk_size, batch_size)

            if start_idx < end_idx:
                chunk_model_inputs, chunk_grpo_batch = self.get_chunked_inputs(model_inputs, grpo_batch, start_idx,
                                                                               end_idx, origin_data)

            chunk_logps, chunk_entropies = self._get_per_token_logps_and_entropies_single(
                model, chunk_model_inputs, chunk_grpo_batch, compute_entropy)

            if start_idx < end_idx:
                all_logps.append(chunk_logps)
                if compute_entropy and chunk_entropies is not None:
                    all_entropies.append(chunk_entropies)

        # Concatenate results
        final_logps = torch.cat(all_logps, dim=0)
        final_entropies = torch.cat(all_entropies, dim=0) if all_entropies else None

        return final_logps, final_entropies

    @profiling_decorator
    def _get_last_hidden_state(self, unwrapped_model, model_inputs, logits_to_keep):
        # unwrap the model to access the model.model
        if is_peft_model(unwrapped_model):
            unwrapped_model = unwrapped_model.base_model.model
        if not self.is_multimodal:
            last_hidden_state = unwrapped_model.model(
                input_ids=model_inputs['input_ids'], attention_mask=model_inputs['attention_mask']).last_hidden_state
        else:
            forward_inputs = dict(model_inputs)
            if 'logits_to_keep' in self.model_kwarg_keys:
                forward_inputs['logits_to_keep'] = logits_to_keep + 1

            last_hidden_state = unwrapped_model.model(**forward_inputs).last_hidden_state

        last_hidden_state = last_hidden_state[:, :-1, :]  # (B, L-1, H)
        if logits_to_keep is not None:
            last_hidden_state = last_hidden_state[:, -logits_to_keep:, :]  # (B, logits_to_keep, H)
        return last_hidden_state

    def _get_rollout_is_correction(self, old_per_token_logps, rollout_per_token_logps, completion_mask):
        """Compute rollout importance sampling log-ratio and IS weights.

        Returns:
            (rollout_log_ratio, rollout_is_weights) if rollout IS correction is applicable,
            (None, None) otherwise.
        """
        if self.rollout_importance_sampling_mode is None or self.disable_rollout_importance_sampling:
            return None, None

        rollout_log_ratio = old_per_token_logps - rollout_per_token_logps
        rollout_is_weights = self._apply_rollout_importance_sampling(rollout_log_ratio, completion_mask)
        return rollout_log_ratio, rollout_is_weights

    def compute_liger_loss(self, unwrapped_model, model_inputs: Dict[str, Any], grpo_batch: GRPOBatch):
        assert not self.template.padding_free
        assert self.advantage_estimator == 'grpo'
        input_ids = model_inputs['input_ids']
        logits_to_keep = grpo_batch.logits_to_keep
        completion_ids = input_ids[:, -logits_to_keep:]
        completion_mask = grpo_batch.completion_mask

        last_hidden_state = self._get_last_hidden_state(unwrapped_model, model_inputs, logits_to_keep)

        old_per_token_logps = grpo_batch.old_per_token_logps
        local_has = grpo_batch.rollout_per_token_logps is not None
        vllm_is_ratio = None
        if all(gather_object([local_has])):
            _, vllm_is_ratio = self._get_rollout_is_correction(old_per_token_logps, grpo_batch.rollout_per_token_logps,
                                                               completion_mask)

        # LigerFusedLinearGRPOLoss expects [B]; expand_advantage_to_per_token always yields [B, T].
        advantages = grpo_batch.advantages
        advantages = (advantages * completion_mask).sum(-1) / completion_mask.sum(-1).clamp(min=1.0)

        loss, metrics = self.liger_grpo_loss(
            _input=last_hidden_state,
            lin_weight=unwrapped_model.lm_head.weight,
            selected_token_ids=completion_ids,
            attention_mask=completion_mask,
            advantages=advantages,
            bias=unwrapped_model.lm_head.bias,
            old_per_token_logps=old_per_token_logps,
            ref_per_token_logps=grpo_batch.ref_per_token_logps,
            vllm_is_ratio=vllm_is_ratio,
        )

        mean_kl = metrics[0] if self.beta != 0.0 else None
        clip_ratio = metrics[-1]

        mode = 'eval' if self.control.should_evaluate else 'train'
        if self.beta != 0.0:
            self._metrics[mode]['kl'].append(self.accelerator.gather_for_metrics(mean_kl).mean().item())
        self._metrics[mode]['clip_ratio'].append(self.accelerator.gather_for_metrics(clip_ratio).mean().item())
        return loss

    def evaluation_loop(self, dataloader, *args, **kwargs):
        # Wait for the training rollout to complete
        if self.args.async_generate:
            while not self.is_async_generate_train_rollout_done():
                time.sleep(0.1)
        if self._queue.empty() and self.args.async_generate:
            self._prefetch(dataloader)
        output = super().evaluation_loop(dataloader, *args, **kwargs)
        self.eval_flag = True
        return output

    def training_step(self, model: nn.Module, inputs: DataType, num_items_in_batch=None) -> torch.Tensor:
        if self.args.async_generate:
            # Wait for the eval rollout to complete
            while not self.is_async_generate_eval_rollout_done():
                time.sleep(0.1)
        return super().training_step(model, inputs, num_items_in_batch)

    def old_policy(self):
        if self.template.sequence_parallel_size == 1:
            return (self.num_iterations > 1
                    or self.args.gradient_accumulation_steps % self.args.steps_per_generation != 0)
        else:
            return (self.num_iterations > 1 or self.args.gradient_accumulation_steps %
                    (self.args.steps_per_generation * sequence_parallel.world_size) != 0)

    @contextmanager
    def offload_context(self):
        if self.args.offload_model:
            self.offload_model(self.accelerator.unwrap_model(self.model))
            if self.ref_model:
                self.offload_model(self.ref_model)
        if getattr(self, 'optimizer', None) and self.args.offload_optimizer:
            self.offload_optimizer()

        try:
            yield
        finally:
            # reload (load back) model when exiting context
            if self.args.offload_model:
                self.load_model(self.accelerator.unwrap_model(self.model))
                if self.ref_model:
                    self.load_model(self.ref_model)
            if getattr(self, 'optimizer', None) and self.args.offload_optimizer:
                self.load_optimizer()

    def log(self, logs: Dict[str, float], start_time: Optional[float] = None) -> None:
        mode = 'train' if self.model.training else 'eval'
        metrics = {key: sum(val) / len(val) for key, val in self._metrics[mode].items()}  # average the metrics

        # This method can be called both in training and evaluation. When called in evaluation, the keys in `logs`
        # start with "eval_". We need to add the prefix "eval_" to the keys in `metrics` to match the format.
        if mode == 'eval':
            metrics = {f'eval_{key}': val for key, val in metrics.items()}

        logs.update(metrics)
        if version.parse(transformers.__version__) >= version.parse('4.47.0.dev0'):
            super().log(logs, start_time)
        else:  # transformers<=4.46
            super().log(logs)
        self._metrics[mode].clear()

        # - entropy only includes samples that went through training (computed in _compute_loss)
        # - Other fields (e.g., prompt/completion/reward) are collected from rollout (in _prepare_inputs)
        # Therefore, if entropy exists, to ensure length consistency across fields,
        # we align all data based on the number of samples in entropy.
        seen_nums = len(self._logs['entropy']) \
            if 'entropy' in self._logs else len(self._logs['prompt'])
        if self.accelerator.is_main_process and self.log_completions:
            table = {
                'step': [str(self.state.global_step)] * seen_nums,
                'prompt': list(self._logs['prompt'])[:seen_nums],
                'completion': list(self._logs['completion'])[:seen_nums],
                **{
                    k: list(v)[:seen_nums]
                    for k, v in self._logs['rewards'].items()
                },
                'advantages': list(self._logs['advantages'])[:seen_nums],
            }
            for key, value in self._logs.items():
                if key not in table and key not in ['image', 'rewards']:
                    table[key] = list(value)[:seen_nums]

            if self.args.log_entropy:
                table.update({'entropy': list(self._logs['entropy'])[:seen_nums]})

            report_to_wandb = self.args.report_to and 'wandb' in self.args.report_to and wandb.run is not None
            report_to_swanlab = self.args.report_to and 'swanlab' in self.args.report_to and swanlab_get_run(
            ) is not None

            self.jsonl_writer.append(table)

            if report_to_wandb:
                import pandas as pd

                # Create a copy to avoid modifying the original table used by other loggers.
                wandb_table = table.copy()
                if self._logs.get('image'):
                    wandb_table['image'] = [
                        wandb.Image(load_pil_img(img)) if img is not None else None for img in self._logs['image']
                    ]
                df = pd.DataFrame(wandb_table)
                if self.wandb_log_unique_prompts:
                    df = df.drop_duplicates(subset=['prompt'])
                wandb.log({'completions': wandb.Table(dataframe=df)})

            if report_to_swanlab:
                headers = list(table.keys())
                rows = []
                for i in range(len(table['step'])):
                    row = []
                    for header in headers:
                        row.append(table[header][i])
                    rows.append(row)
                swanlab.log({'completions': swanlab.echarts.Table().add(headers, rows)})

    def is_async_generate_eval_rollout_done(self):
        return not self.eval_flag or not self.eval_queue.empty()

    def is_async_generate_train_rollout_done(self):
        return not self.train_queue.empty()

    def _gather_and_flatten(self, local_list, dtype=None, device=None, flatten_level: int = 1):
        """
        Gather data from all ranks with `gather_object` and flatten as required.

        Args
        ----
        local_list : List[Any]
            The per-rank data to be gathered. Can be any picklable structure.
        dtype : torch.dtype, optional
            If provided, the flattened result is converted to a tensor with this dtype.
        device : torch.device, optional
            Target device for the resulting tensor. Ignored if dtype is None.
        flatten_level : int
            0  keep the outer list of per-rank results: List[rank_data]
            1  flatten across ranks: List[element]
            2  flatten one more level (assumes rank_data is iterable): List[sub_element]

        Returns
        -------
        Any
            - List[Any] when dtype is None
            - torch.Tensor when dtype is given
        """
        gathered = gather_object(local_list)  # List[rank][...] returned by gather_object

        if flatten_level == 0:
            flat = gathered
        elif flatten_level == 1:  # flatten over ranks
            flat = [elem for rank_data in gathered for elem in rank_data]
        elif flatten_level == 2:  # flatten one additional level
            flat = [item for rank_data in gathered for sublist in rank_data for item in sublist]
        else:
            raise ValueError(f'Invalid flatten_level: {flatten_level}')

        if dtype is not None:
            try:
                return torch.tensor(flat, dtype=dtype, device=device)
            except (TypeError, ValueError) as e:
                raise RuntimeError(f'Cannot convert gathered+flattened data to tensor: {e}') from e
        return flat

    def _group_inputs_by_request_id(self, inputs: DataType) -> Dict[str, List[Dict]]:
        """
        Group inputs by request_id for multi-turn reward computation.

        Args:
            inputs: List of input dictionaries, each containing a 'request_id' field

        Returns:
            Dict[str, List[Dict]]: A dictionary where keys are request_ids and values are
                                  lists of input dictionaries with the same request_id
        """
        inputs_by_request_id = {}

        for input_data in inputs:
            request_id = input_data.get('request_id')
            if request_id is None:
                # Skip inputs without request_id
                continue

            if request_id not in inputs_by_request_id:
                inputs_by_request_id[request_id] = []

            inputs_by_request_id[request_id].append(input_data)

        return inputs_by_request_id

    def _get_trajectory_inputs(self, samples: List[GRPOSample]) -> Dict[str, List[Dict]]:
        current_request_ids = {sample.request_id for sample in samples}
        local_inputs = [sample.to_reward_row() for sample in samples]
        total_inputs = gather_object(local_inputs)
        filtered_total_inputs = []
        for rank_inputs in total_inputs:
            candidate_inputs = [rank_inputs] if isinstance(rank_inputs, dict) else rank_inputs
            for input_data in candidate_inputs:
                if input_data.get('request_id') in current_request_ids:
                    filtered_total_inputs.append(input_data)
        return self._group_inputs_by_request_id(filtered_total_inputs)

    def _get_last_indices(self, request_ids: List[str]) -> torch.Tensor:
        seen = {}
        for i, rid in enumerate(request_ids):
            seen[rid] = i
        return torch.tensor(list(seen.values()), dtype=torch.long, device=self.accelerator.device)

    def get_chunked_inputs(self, model_inputs, grpo_batch, start_idx, end_idx, origin_data=None):
        """Slice ``model_inputs`` and ``grpo_batch`` by batch dimension."""
        chunk_model_inputs = {}
        batch_size = grpo_batch.seq_lengths.shape[0] if self.template.padding_free else model_inputs['input_ids'].shape[
            0]
        for key, val in model_inputs.items():
            if isinstance(val, torch.Tensor):
                if val.ndim == 0:
                    chunk_model_inputs[key] = val
                elif self.is_multimodal and val.shape[0] != batch_size:
                    continue
                else:
                    chunk_model_inputs[key] = val[start_idx:end_idx]
            elif isinstance(val, list) and len(val) == batch_size:
                chunk_model_inputs[key] = val[start_idx:end_idx]
            else:
                chunk_model_inputs[key] = val

        chunk_grpo_batch = GRPOBatch(
            completion_mask=grpo_batch.completion_mask[start_idx:end_idx],
            truncated_mask=grpo_batch.truncated_mask[start_idx:end_idx],
            seq_lengths=(grpo_batch.seq_lengths[start_idx:end_idx]
                         if grpo_batch.seq_lengths.numel() > 0 else grpo_batch.seq_lengths),
            old_per_token_logps=(grpo_batch.old_per_token_logps[start_idx:end_idx]
                                 if grpo_batch.old_per_token_logps is not None else None),
            ref_per_token_logps=(grpo_batch.ref_per_token_logps[start_idx:end_idx]
                                 if grpo_batch.ref_per_token_logps is not None else None),
            rollout_per_token_logps=(grpo_batch.rollout_per_token_logps[start_idx:end_idx]
                                     if grpo_batch.rollout_per_token_logps is not None else None),
            advantages=grpo_batch.advantages[start_idx:end_idx] if grpo_batch.advantages is not None else None,
            num_items_in_batch=grpo_batch.num_items_in_batch,
            logits_to_keep=grpo_batch.logits_to_keep,
        )

        if self.is_multimodal and origin_data is not None:
            chunk_origin_data = origin_data[start_idx:end_idx]
            template = self.template
            current_length = model_inputs['input_ids'].shape[1]
            with self._template_context(template):
                encoded_data = [template.encode(data.to_template_dict()) for data in chunk_origin_data]
                for ed in encoded_data:
                    ed.pop('_extra_kwargs', None)
                chunk_model_inputs.update(
                    to_device(template.data_collator(encoded_data, padding_to=current_length), self.model.device))
                chunk_model_inputs.pop('labels', None)

        return chunk_model_inputs, chunk_grpo_batch

    def _prepare_liger_loss(self):
        self.use_liger_loss = self.args.use_liger_kernel
        if self.use_liger_loss:
            from liger_kernel.chunked_loss import LigerFusedLinearGRPOLoss
            self.liger_grpo_loss = LigerFusedLinearGRPOLoss(
                beta=self.beta,
                compiled=False,
                epsilon_low=self.epsilon_low,
                epsilon_high=self.epsilon_high,
                temperature=self.temperature,
                use_ref_model=self.beta != 0.0,
                loss_type=self.loss_type,
                max_completion_length=self.max_completion_length,
                importance_sampling_level=self.importance_sampling_level,
                sapo_temperature_pos=self.tau_pos,
                sapo_temperature_neg=self.tau_neg,
            )
            self._forward_redirection = _ForwardRedirection()

    def _prepare_metrics(self):
        args = self.args
        self._metrics = {'train': defaultdict(list), 'eval': defaultdict(list)}
        self.log_completions = args.log_completions
        self.wandb_log_unique_prompts = args.wandb_log_unique_prompts
        self.num_completions_to_print = args.num_completions_to_print
        self.jsonl_writer = JsonlWriter(os.path.join(self.args.output_dir, 'completions.jsonl'))
        self._logs = {
            'prompt': deque(maxlen=args.generation_batch_size),
            'completion': deque(maxlen=args.generation_batch_size),
            'rewards': defaultdict(lambda: deque(maxlen=args.generation_batch_size)),
            'advantages': deque(maxlen=args.generation_batch_size),
        }
        self.compute_entropy = self.args.log_entropy or self.top_entropy_quantile < 1.0
        if self.args.log_entropy:
            self._logs.update({'entropy': deque(maxlen=args.generation_batch_size)})

    def _collect_config_info(self) -> Dict[str, str]:
        config = {
            'dynamic_sample': str(self.dynamic_sample),
            'importance_sampling_level': str(self.importance_sampling_level),
            'advantage_estimator': str(self.advantage_estimator),
            'chord_sft_enabled': str(self.chord_sft_dataset is not None),
            'offpolicy_sequence_mask': 'enable' if self.off_policy_sequence_mask_delta is not None else 'disable',
            'rollout_importance_sampling': 'enable' if self.rollout_importance_sampling_mode is not None else 'disable',
            'loss_type': str(self.loss_type),
        }
        return config

    def _prepare_algorithm_params(self):
        args = self.args
        self.shuffle_dataset = args.dataset_shuffle

        self.loss_type = args.loss_type  # loss normalization
        self.scale_rewards = args.scale_rewards

        # GRPO, https://arxiv.org/abs/2402.03300
        self.num_iterations = args.num_iterations  # = 𝜇 in the GRPO paper, Multi-step

        # DAPO, https://arxiv.org/abs/2503.14476
        self.epsilon_low = args.epsilon
        self.epsilon_high = args.epsilon_high if args.epsilon_high is not None else args.epsilon
        self.dynamic_sample = args.dynamic_sample
        self.max_resample_times = args.max_resample_times
        self.overlong_filter = args.overlong_filter

        # Entropy Mask, https://arxiv.org/abs/2506.01939
        self.top_entropy_quantile = args.top_entropy_quantile

        # GSPO, https://arxiv.org/abs/2507.18071
        self.importance_sampling_level = args.importance_sampling_level

        # SAPO, https://arxiv.org/abs/2511.20347
        self.tau_pos = args.tau_pos
        self.tau_neg = args.tau_neg

        # REAL, https://arxiv.org/abs/2602.05630
        self.real_tau = args.real_tau

        # FIPO, https://arxiv.org/abs/2603.19835
        self.fipo_gamma = 2**(-1 / args.fipo_decay_rate)
        self.fipo_clip_range = args.fipo_clip_range
        self.fipo_clip_high_only = args.fipo_clip_high_only
        self.fipo_safety_threshold = args.fipo_safety_threshold

        # RLOO,
        self.advantage_estimator = args.advantage_estimator
        self.kl_in_reward = args.kl_in_reward
        if self.scale_rewards == 'gdpo' and self.kl_in_reward:
            logger.warning('GDPO mode does not support kl_in_reward=True. Setting kl_in_reward=False.')
            self.kl_in_reward = False

        # Rollout Importance Sampling Correction
        self.rollout_importance_sampling_mode = args.rollout_importance_sampling_mode
        self.rollout_importance_sampling_threshold = args.rollout_importance_sampling_threshold
        self.log_rollout_offpolicy_metrics = args.log_rollout_offpolicy_metrics

        # Off-Policy Sequence Masking
        self.off_policy_sequence_mask_delta = args.off_policy_sequence_mask_delta

    def _prepare_chord_dataset(self):
        # CHORD, https://arxiv.org/abs/2508.11408
        self.chord_sft_iterator = None
        if self.chord_sft_dataset:
            self.chord_sft_iterator = make_chord_sft_dataset(self, self.chord_sft_dataset)

    def _prepare_rewards(self, reward_funcs, reward_model=None, reward_templates=None):
        args = self.args
        device = self.accelerator.device

        if not isinstance(reward_funcs, list):
            reward_funcs = [reward_funcs]

        if reward_funcs:
            for i, reward_func in enumerate(reward_funcs):
                if reward_func in orms:
                    reward_func_class = orms[reward_func]
                    reward_funcs[i] = reward_func_class(args=args)
                elif not callable(reward_func):
                    raise ValueError(f'reward_function {reward_func} is not implemented in swift.rewards')

        self.reward_funcs = reward_funcs
        self.reward_func_names = []
        for reward_func in reward_funcs:
            if inspect.isfunction(reward_func):
                reward_func_name = reward_func.__name__
            else:
                reward_func_name = reward_func.__class__.__name__
            self.reward_func_names.append(reward_func_name)

        self.reward_model_plugins = [None] * len(self.reward_funcs)

        if reward_model is not None:
            reward_plugins = args.reward_model_plugin
            if reward_plugins is None:
                reward_plugins = ['default'] * len(reward_model)
            assert len(reward_plugins) == len(reward_model), (
                f"The number of 'reward_model_plugin' ({len(reward_plugins)}) does not match "
                f"the number of 'reward_model' ({len(reward_model)}). "
                "Please provide a corresponding 'reward_model_plugin' for each 'reward_model'.")
            for rm, rm_plugin, rm_template in zip(reward_model, reward_plugins, reward_templates):
                # Set encoding mode train(see details in Template.encode).
                # Set max_length to None to disable truncation, as the input length has already been truncated earlier.
                rm_template.set_mode('train')
                rm_template.max_length = None
                if rm_plugin not in rm_plugins:
                    raise ValueError(f'rm_plugin {rm_plugin} is not implemented in swift.rewards')
                self.reward_model_plugins.append(rm_plugins[rm_plugin](model=rm, template=rm_template))
                self.reward_funcs.append(rm)
                self.reward_func_names.append(rm.config._name_or_path.split('/')[-1])

        # use_gym_env: gym total_reward is appended as an extra reward column so it can
        # blend with reward_funcs via reward_weights. When reward_funcs is empty, it becomes
        # the single reward source.
        if self.use_gym_env:
            self.reward_func_names.append('gym_reward')

        # Reward weights
        if args.reward_weights is not None:
            if len(args.reward_weights) != len(self.reward_func_names):
                raise ValueError(f'Number of reward weights ({len(args.reward_weights)}) must match number of reward '
                                 f'functions ({len(self.reward_func_names)})')
            self.reward_weights = torch.tensor(args.reward_weights, dtype=torch.float32).to(device)
        else:
            self.reward_weights = torch.ones(len(self.reward_func_names), dtype=torch.float32).to(device)

        # after init trainer
        for i, reward_func in enumerate(self.reward_funcs):
            if isinstance(reward_func, PreTrainedModel):
                if self.is_deepspeed_enabled:
                    self.reward_funcs[i] = prepare_deepspeed(reward_func, self.accelerator)
                elif self.is_fsdp_enabled:
                    from .utils import prepare_fsdp
                    self.reward_funcs[i] = prepare_fsdp(reward_func, self.accelerator)
                else:
                    self.reward_funcs[i] = self.accelerator.prepare_model(
                        reward_func, evaluation_mode=True, device_placement=True)

    def _prepare_resample_data_iterator(self):

        def cyclic_iter(iterable):
            while True:
                for x in iterable:
                    yield x

        @contextmanager
        def seed_context():
            # Use a different seed to ensure the resample dataset does not overlap with train_dataset
            seed = self.args.seed
            self.args.seed = seed + 1
            yield
            self.args.seed = seed

        with seed_context():
            if self.args.dynamic_sample:
                self.dynamic_resample_iterator = cyclic_iter(self.get_train_dataloader())

            if self.template.truncation_strategy == 'raise':

                @contextmanager
                def single_sample_context():
                    # Patch generation-related parameters to ensure that only one sample is processed per iteration
                    # when resampling truncated data.
                    origin_ng = self.num_generations
                    origin_gbs = self.args.generation_batch_size
                    origin_spg = self.args.steps_per_generation
                    try:
                        self.num_generations = 1
                        self.args.generation_batch_size = 1
                        self.args.steps_per_generation = 1
                        yield
                    finally:
                        self.num_generations = origin_ng
                        self.args.generation_batch_size = origin_gbs
                        self.args.steps_per_generation = origin_spg

                with single_sample_context():
                    self.truncated_resample_iterator = cyclic_iter(self.get_train_dataloader())

    def _compute_sequence_level_ratios(self, is_ratio: torch.Tensor, completion_mask: torch.Tensor) -> torch.Tensor:
        """
        Helper function to compute sequence-level importance sampling ratios.

        Args:
            is_ratio: Token-level IS ratios, shape [B, T]
            completion_mask: Boolean mask for completion tokens, shape [B, T]

        Returns:
            Sequence-level ratios as geometric mean of token-level ratios
        """
        log_ratio = torch.log(is_ratio.clamp(min=1e-10))
        seq_log_ratios = (log_ratio * completion_mask).sum(-1) / completion_mask.sum(-1).clamp(min=1.0)
        seq_ratios = torch.exp(seq_log_ratios)

        return seq_ratios

    def _apply_rollout_importance_sampling(self, rollout_log_ratio: torch.Tensor,
                                           completion_mask: torch.Tensor) -> torch.Tensor:
        """
        Apply vLLM importance sampling correction using one of four modes.

        Args:
            rollout_log_ratio: log(π_θ / π_rollout) per token, shape [B, T]
            completion_mask: Boolean mask for completion tokens, shape [B, T]

        Returns:
            IS weights to multiply with loss, same shape as rollout_log_ratio
        """
        mode = self.rollout_importance_sampling_mode
        threshold = self.rollout_importance_sampling_threshold

        # Clamp log_ratio to prevent numerical overflow from padding values (-1e10)
        # A log_ratio of 20 corresponds to exp(20) ≈ 485 million, which is already extreme
        SAFETY_BOUND = 20.0
        rollout_log_ratio_safe = torch.clamp(rollout_log_ratio, min=-SAFETY_BOUND, max=SAFETY_BOUND)

        # Compute importance sampling ratios: exp(log_ratio)
        is_ratio = torch.exp(rollout_log_ratio_safe)

        if mode == 'token_truncate':
            # Token-level truncated IS: clip ratios from above at threshold
            is_weights = torch.clamp(is_ratio, max=threshold)

        elif mode == 'token_mask':
            # Token-level masked IS: mask out tokens with ratio > threshold
            is_weights = torch.where(is_ratio <= threshold, is_ratio, torch.zeros_like(is_ratio))

        elif mode == 'sequence_truncate':
            # Sequence-level truncated IS: compute sequence-level ratio and clip
            seq_ratios = self._compute_sequence_level_ratios(is_ratio, completion_mask)
            clipped_seq_ratios = torch.clamp(seq_ratios, max=threshold)

            is_weights = clipped_seq_ratios.unsqueeze(-1).expand_as(is_ratio)

        elif mode == 'sequence_mask':
            # Sequence-level masked IS: mask entire sequences with ratio > threshold
            seq_ratios = self._compute_sequence_level_ratios(is_ratio, completion_mask)
            seq_mask = (seq_ratios <= threshold).float()

            # Apply mask to original token-level ratios
            is_weights = is_ratio * seq_mask.unsqueeze(-1)
        else:
            return is_ratio

        return is_weights

    def _compute_off_policy_sequence_mask(
        self,
        per_token_logps: torch.Tensor,
        old_policy_per_token_logps: torch.Tensor,
        completion_mask: torch.Tensor,
        advantages: torch.Tensor,
    ) -> torch.Tensor:
        """
        Compute off-policy sequence mask to filter out sequences that deviate too much
        from the old/rollout policy AND have negative advantage.

        This implements the Off-Policy Sequence Masking technique from DeepSeek-V3.2
        (https://arxiv.org/abs/2512.02556). The mask filters sequences where:
        1. mean(old_policy_logps - policy_logps) > off_policy_sequence_mask_delta
        2. AND advantage < 0

        Args:
            per_token_logps: Log probs from current policy, shape [B, T]
            old_policy_per_token_logps: Log probs from old/rollout policy, shape [B, T].
                Uses rollout_per_token_logps if available, otherwise old_per_token_logps.
            completion_mask: Boolean mask for completion tokens, shape [B, T]
            advantages: Advantage values per sample, shape [B]

        Returns:
            Sequence mask, shape [B], True = keep sequence, False = mask out
        """
        # Compute per-token log ratio: log(π_old / π_current)
        # Following DeepSeek-V3.2: positive delta means old policy assigns higher prob
        log_ratio = old_policy_per_token_logps - per_token_logps

        # Compute sequence-level mean of log ratio
        seq_mean_log_ratio = (log_ratio * completion_mask).sum(-1) / completion_mask.sum(-1).clamp(min=1.0)

        # Mask condition: delta > threshold AND advantage < 0
        # Keep sequences that do NOT meet this condition
        exceeds_threshold = seq_mean_log_ratio > self.off_policy_sequence_mask_delta
        negative_advantage = advantages < 0
        should_mask = exceeds_threshold & negative_advantage

        # Return mask: True = keep, False = mask out
        return ~should_mask

    def _compute_rollout_offpolicy_metrics(
        self,
        per_token_logps: torch.Tensor,
        rollout_per_token_logps: torch.Tensor,
        completion_mask: torch.Tensor,
    ) -> Dict[str, float]:
        """
        Compute off-policy diagnostic metrics (always computed for monitoring).

        These metrics help diagnose the off-policy gap between rollout and training policies,
        which can arise from policy mismatch (e.g., vLLM BF16 vs FSDP FP32), model staleness,
        or general distribution shifts.

        Key metrics:
        - kl: Direct KL divergence estimator KL(π_rollout || π_training)
        - k3_kl: K3 KL estimator for stability (more stable for small KL)
        - training_ppl: Perplexity of training policy
        - rollout_ppl: Perplexity of rollout policy
        - log_ppl_diff: Difference in log perplexities
        - ppl_ratio: Ratio of training PPL to rollout PPL
        - chi2_token: Token-level χ² divergence E[ρ²] - 1
        - chi2_seq: Sequence-level χ² divergence E[(∏ρ_t)²] - 1

        Args:
            per_token_logps: Log probs from training policy model, shape [B, T]
            rollout_per_token_logps: Log probs from rollout policy, shape [B, T]
            completion_mask: Boolean mask for completion tokens, shape [B, T]

        Returns:
            Dictionary with off-policy diagnostic metrics
        """
        SAFETY_BOUND = 20.0
        metrics = {}

        # Helper function for masked mean
        def masked_mean(x, mask, axis=None):
            if axis is None:
                return (x * mask).sum() / mask.sum().clamp(min=1.0)
            else:
                return (x * mask).sum(axis) / mask.sum(axis).clamp(min=1.0)

        # 1. Training policy perplexity (always computed)
        # Formula: exp(-1/|T| * Σ log π_training(y_t|y_<t))
        mean_log_prob_training = masked_mean(per_token_logps, completion_mask, axis=-1)  # (batch_size,)
        training_ppl = torch.exp(-mean_log_prob_training).mean()  # Batch mean of per-sequence PPL
        metrics['training_ppl'] = self.accelerator.gather_for_metrics(training_ppl).nanmean().item()

        # Also log log-ppl for easier analysis (avoids exponential scale)
        metrics['training_log_ppl'] = self.accelerator.gather_for_metrics(
            (-mean_log_prob_training).mean()).nanmean().item()

        # 2. Compute rollout off-policy metrics
        # All KL metrics estimate KL(π_rollout || π_training), which measures how much
        # the training policy deviates from the rollout policy. This is directly related
        # to the importance sampling ratio ρ = π_training / π_rollout.

        # log_ratio = log(π_training / π_rollout), used for IS weights and KL estimators
        log_ratio = per_token_logps - rollout_per_token_logps
        log_ratio *= completion_mask

        # 2a. kl: Direct estimator for KL(π_rollout || π_training)
        # Formula: KL(P||Q) = E_P[log(P/Q)] where P=π_rollout, Q=π_training
        # = E_rollout[log(π_rollout) - log(π_training)] = E[-log_ratio]
        kl = masked_mean(-log_ratio, completion_mask)
        metrics['kl'] = self.accelerator.gather_for_metrics(kl).nanmean().item()

        # 2b. k3_kl: K3 estimator for KL(π_rollout || π_training)
        # More stable for small KL values
        log_ratio_safe = torch.clamp(log_ratio, min=-20, max=20)
        k3_kl_matrix = torch.clamp(torch.exp(log_ratio_safe) - log_ratio_safe - 1, min=-10, max=10)
        k3_kl = masked_mean(k3_kl_matrix, completion_mask)
        metrics['k3_kl'] = self.accelerator.gather_for_metrics(k3_kl).nanmean().item()

        # 2c. Rollout policy perplexity
        mean_log_prob_rollout = masked_mean(rollout_per_token_logps, completion_mask, axis=-1)  # (batch_size,)
        rollout_ppl = torch.exp(-mean_log_prob_rollout).mean()  # Batch mean of per-sequence PPL
        metrics['rollout_ppl'] = self.accelerator.gather_for_metrics(rollout_ppl).nanmean().item()
        metrics['rollout_log_ppl'] = self.accelerator.gather_for_metrics(
            (-mean_log_prob_rollout).mean()).nanmean().item()

        # 2d. Log PPL difference (sequence-level perplexity difference)
        # log_ppl_diff = mean_log_prob_rollout - mean_log_prob_training
        # Since ppl = exp(-log_prob), we have:
        #   log(ppl_ratio) = log(training_ppl/rollout_ppl) = log_ppl_diff
        # Positive value means training assigns lower probability (higher PPL) than rollout
        log_ppl_diff = mean_log_prob_rollout - mean_log_prob_training
        metrics['log_ppl_diff'] = self.accelerator.gather_for_metrics(log_ppl_diff.mean()).nanmean().item()
        metrics['log_ppl_abs_diff'] = self.accelerator.gather_for_metrics(log_ppl_diff.abs().mean()).nanmean().item()
        metrics['log_ppl_diff_max'] = self.accelerator.gather_for_metrics(log_ppl_diff.max()).max().item()
        metrics['log_ppl_diff_min'] = self.accelerator.gather_for_metrics(log_ppl_diff.min()).min().item()

        # 2e. PPL ratio (how much higher is training PPL vs rollout PPL)
        # IMPORTANT: Compute per-sequence ratio first, then average
        # For numerical stability, compute in log space using log_ppl_diff
        # Note: log_ppl_diff = log(ppl_ratio), so ppl_ratio = exp(log_ppl_diff)
        ppl_ratio = torch.exp(log_ppl_diff).mean()
        metrics['ppl_ratio'] = self.accelerator.gather_for_metrics(ppl_ratio).nanmean().item()

        # 2f. Chi-squared divergence: χ²(π_training || π_rollout) = E_μ[ρ²] - 1
        # where ρ = π_training / π_rollout and μ = π_rollout (rollout distribution)
        # This measures the variance of importance sampling weights
        # Token-level: E_token[ρ²] - 1 (averaged over all tokens)
        log_ratio_safe = torch.clamp(log_ratio, min=-SAFETY_BOUND, max=SAFETY_BOUND)
        rho_token = torch.exp(log_ratio_safe)  # ρ = π_training / π_rollout (token-level)
        rho_squared_token = rho_token.square()
        chi2_token = masked_mean(rho_squared_token, completion_mask) - 1.0
        metrics['chi2_token'] = self.accelerator.gather_for_metrics(chi2_token).nanmean().item()

        # Sequence-level (geometric mean): E_seq[ρ_geo²] - 1
        # where ρ_geo = exp(mean(log ρ_t)) is the geometric mean of token-level ratios
        # This is more interpretable than the product-based chi2_seq, as it's normalized by sequence length
        # and comparable to other per-token metrics like chi2_token
        log_ratio_mean = masked_mean(log_ratio, completion_mask, axis=-1)  # mean(log ρ_t) per sequence
        log_ratio_mean_safe = torch.clamp(log_ratio_mean, min=-SAFETY_BOUND, max=SAFETY_BOUND)
        rho_geo = torch.exp(log_ratio_mean_safe)  # geometric mean of ρ_t
        chi2_seq = (rho_geo.square().mean() - 1.0)
        metrics['chi2_seq'] = self.accelerator.gather_for_metrics(chi2_seq).nanmean().item()

        return metrics

    def _compute_is_correction_metrics(
        self,
        rollout_log_ratio: torch.Tensor,
        is_weights: torch.Tensor,
        completion_mask: torch.Tensor,
    ) -> Dict[str, float]:
        """
        Compute importance sampling correction metrics (ess, clipped_frac, is_weight_mean).
        Only called when rollout_importance_sampling_mode is enabled.

        Args:
            rollout_log_ratio: Log ratio log(π_policy / π_rollout), shape [B, T]
            is_weights: Importance sampling weights after correction, shape [B, T]
            completion_mask: Boolean mask for completion tokens, shape [B, T]

        Returns:
            Dictionary with IS-specific metrics:
                - is_weight_mean: Mean of IS weights
                - ess: Effective Sample Size = 1 / E[(w_i / E[w_i])²]
                - clipped_frac: Fraction of clipped/masked samples
        """
        metrics = {}
        SAFETY_BOUND = 20.0
        threshold = self.rollout_importance_sampling_threshold
        threshold_lower = 1.0 / threshold  # Default lower threshold (reciprocal of upper)

        # Helper function for masked mean
        def masked_mean(x, mask):
            return (x * mask).sum() / mask.sum().clamp(min=1.0)

        # Compute IS ratio with safety bounds
        log_ratio_safe = torch.clamp(rollout_log_ratio, min=-SAFETY_BOUND, max=SAFETY_BOUND)
        is_ratio = torch.exp(log_ratio_safe)

        # 1. IS weight statistics
        mean_is_weight = masked_mean(is_weights, completion_mask)
        metrics['is_weight_mean'] = self.accelerator.gather_for_metrics(mean_is_weight).nanmean().item()

        # 2. Compute Effective Sample Size (ESS) for IS weights
        # ESS = 1 / E[(w_i / E[w_i])²] (using clamped weights for stability)
        # This measures how many "effective" independent samples we have after IS weighting
        weights_for_ess = is_weights.clamp(min=threshold_lower, max=threshold)
        mean_for_ess = masked_mean(weights_for_ess, completion_mask)
        is_weights_normalized = weights_for_ess / (mean_for_ess + 1e-8)  # Avoid division by zero
        ess = 1.0 / masked_mean(is_weights_normalized.square(), completion_mask).clamp(min=1e-10)
        metrics['ess'] = self.accelerator.gather_for_metrics(ess).nanmean().item()

        # 3. Fraction of clipped/masked samples
        if self.rollout_importance_sampling_mode in ['token_truncate', 'token_mask']:
            # Token-level
            if self.rollout_importance_sampling_mode == 'token_truncate':
                clipped_frac = masked_mean((is_ratio > threshold).float(), completion_mask)
            else:  # token_mask
                clipped_frac = masked_mean((is_weights == 0).float(), completion_mask)
            metrics['clipped_frac'] = self.accelerator.gather_for_metrics(clipped_frac).nanmean().item()
        else:
            # Sequence-level (both truncate and mask)
            seq_ratios = self._compute_sequence_level_ratios(is_ratio, completion_mask)
            clipped_frac = (seq_ratios > threshold).float().mean()
            metrics['clipped_frac'] = self.accelerator.gather_for_metrics(clipped_frac).nanmean().item()

        return metrics

    def _get_eval_sampler(self, eval_dataset):
        return RepeatSampler(
            data_source=eval_dataset,
            mini_repeat_count=self.num_generations_eval,
            seed=self.args.seed,
        )
