rl_utils.py

# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.

import gc

import copy
from functools import partial
# Keep this to make the env registered.
import itertools
import math
import logging
import json
import os
from collections import Counter, defaultdict
from contextlib import contextmanager, nullcontext
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, Iterator, List, Optional 

import numpy as np
import torch
import torch.distributed as dist
import yaml
from torch.utils.data import DataLoader, TensorDataset
from torch.utils.tensorboard import SummaryWriter

from megatron.core import mpu
from megatron.core.packed_seq_params import PackedSeqParams
from megatron.core.full_cuda_graph import FullCudaGraphWrapper
from megatron.core.models.common.language_module.language_module import LanguageModule
from megatron.core.num_microbatches_calculator import reconfigure_num_microbatches_calculator
from megatron.core.optimizer import MegatronOptimizer
from megatron.core.pipeline_parallel import get_forward_backward_func
from megatron.core.pipeline_parallel.utils import is_pp_last_stage, get_pp_last_rank
from megatron.core.rerun_state_machine import RerunDataIterator
from megatron.core.tokenizers import MegatronTokenizer
from megatron.core.tokenizers.text.libraries.huggingface_tokenizer import HuggingFaceTokenizer
from megatron.core.transformer.cuda_graphs import _CudagraphGlobalRecord
from megatron.core.transformer.enums import CudaGraphScope
from megatron.core.transformer.utils import (
    toggle_cuda_graphs,
    transition_moe_cudagraphs,
)
from megatron.core.inference.utils import set_decode_expert_padding
from megatron.core.resharding.refit import swap_model_weights
from megatron.core.inference.unified_memory import (
    advise_managed_module_parameters_preferred_location,
    prefetch_managed_module_parameters,
)
from megatron.core.inference.utils import device_memory_summary
from megatron.core.utils import get_asyncio_loop, log_single_rank
from megatron.rl.sequence_packing_utils import (
    get_microbatch_dataloader,
    pack_inference_logprobs,
    compute_packed_inference_logprobs_stats,
    pack_all_trajectories,
    load_packed_data_by_index,
    get_sequence_packing_tensorboard_metrics,
    get_sequence_packing_log_info,
    get_default_packed_seq_params,
    update_microbatch_calculator,
)
from megatron.rl.agent.api import (
    EvaluationRequest,
    EvaluationResponse,
    GroupedRolloutRequest,
    RewardEvaluationResult,
    Rollout,
    TokenRollout,
)
from megatron.rl.agent.weighted_multi_task import WeightedMultiTask
from megatron.rl.inference.megatron import MegatronLocal
from megatron.rl.logging import LOG_DIR as lang_rl_log_dir
from megatron.rl.logging import log as lang_rl_log
from megatron.rl.server.inference.inference_interface_server import InferenceInterfaceServer
from megatron.training.global_vars import (
    get_args,
    get_tensorboard_writer,
    get_tokenizer,
    get_wandb_writer,
)
from megatron.training.utils import (
    get_ltor_masks_and_position_ids,
    get_nvtx_range,
    print_rank_0,
    unwrap_model,
)
from megatron.core.utils import get_pg_rank, get_pg_size, get_attr_wrapped_model
from megatron.core.process_groups_config import ProcessGroupCollection
from wandb import wandb_run
from megatron.core.transformer.custom_layers.batch_invariant_kernels import (
    is_batch_invariant_mode_enabled,
)

from megatron.core.inference.contexts.dynamic_context import HAVE_TORCH_MEMORY_SAVER
if HAVE_TORCH_MEMORY_SAVER:
    from torch_memory_saver import torch_memory_saver

logger = logging.getLogger(__name__)

# Global variable to store packing context for forward_step
_GLOBAL_PACKING_CONTEXT = None


# Track whether the inference model is currently paused (offloaded to CPU).
# Model starts on GPU after creation and is used immediately, so starts as False.
_INFERENCE_MODEL_IS_PAUSED = False


def _torch_saver_swap_inference_model(*, to_cpu: bool) -> None:
    """Swap RL inference model weights between CPU and GPU using torch_memory_saver.

    Uses torch_memory_saver.pause()/resume() to transfer inference model weights
    that were allocated within a torch_memory_saver.region() context.

    Args:
        to_cpu: If True, move weights to CPU (pause). If False, restore weights to GPU (resume).
    """
    global _INFERENCE_MODEL_IS_PAUSED

    if not HAVE_TORCH_MEMORY_SAVER:
        raise RuntimeError(
            "torch_memory_saver is required for inference model offloading when not using UVM. "
            "Please install it: pip install torch_memory_saver "
            "(see https://github.com/fzyzcjy/torch_memory_saver)"
        )

    tag = "rl_inference_model"
    if to_cpu:
        if not _INFERENCE_MODEL_IS_PAUSED:
            print_rank_0(f"torch_memory_saver: pausing {tag}, before: {device_memory_summary()}")
            torch_memory_saver.pause(tag)
            _INFERENCE_MODEL_IS_PAUSED = True
            print_rank_0(f"torch_memory_saver: paused  {tag}, after:  {device_memory_summary()}")
    else:
        if _INFERENCE_MODEL_IS_PAUSED:
            print_rank_0(f"torch_memory_saver: resuming {tag}, before: {device_memory_summary()}")
            torch_memory_saver.resume(tag)
            _INFERENCE_MODEL_IS_PAUSED = False
            print_rank_0(f"torch_memory_saver: resumed  {tag}, after:  {device_memory_summary()}")


def _maybe_prefetch_separate_inference_model_weights(model_core, *, to_cpu: bool) -> None:
    """Prefetch RL *separate inference model* weights to CPU/GPU.

    Supports two modes:
    1. UVM-based offloading (when --rl-inference-model-unified-memory-level=1)
    2. torch_memory_saver-based offloading (when offloading is enabled but UVM is not)

    Gated by user args; this assumes the separate inference model was allocated
    with UVM or torch_memory_saver when enabled.
    """
    args = get_args()
    if not args.rl_offload_inference_model_weights_when_idle:
        return

    # Check for torch_memory_saver path (when offloading is enabled but UVM is not)
    if args.rl_inference_model_unified_memory_level != 1:
        _torch_saver_swap_inference_model(to_cpu=to_cpu)
        return

    # UVM-based path (when UVM level is 1)
    device = -1 if to_cpu else int(torch.cuda.current_device())
    # Note: include_buffers=False because buffers created with explicit device= in register_buffer()
    # are not allocated via the UVM mempool and will fail UVM operations. Only parameters are UVM-allocated.
    advise_managed_module_parameters_preferred_location(model_core, device=device, include_buffers=False)
    nbytes = prefetch_managed_module_parameters(model_core, device=device, include_buffers=False)
    # Ensure pages are resident before we enter CUDA-graph capture / inference, or before training continues.
    torch.cuda.synchronize()

    if to_cpu:
        print_rank_0(f"[Rank 0] offloaded {nbytes / 1024**2:.2f} MB of separate RL inference model weights to CPU (other ranks may vary)")
    else:
        print_rank_0(f"[Rank 0] prefetched {nbytes / 1024**2:.2f} MB of separate RL inference model weights to GPU (other ranks may vary)")


def verify_model_weights_swap(
    train_model: LanguageModule,
    inference_model: LanguageModule,
    seq_len: int = 8,
    batch_size: int = 2,
    atol: float = 1e-4,
    rtol: float = 1e-4,
) -> None:
    """Verify that the inference model produces the same forward pass outputs
    as the training model after the weights have been swapped.

    This function should be called after swap_model_weights to ensure the weight
    transfer was successful. It runs a forward pass on both models and asserts
    the outputs match.  This is meant for debugging purposes only.

    Args:
        train_model: The training model (source of weights).
        inference_model: The inference model (target of weights).
        seq_len: Sequence length for test input.
        batch_size: Batch size for test input.
        atol: Absolute tolerance for comparing outputs.
        rtol: Relative tolerance for comparing outputs.

    Raises:
        AssertionError: If forward pass outputs do not match within tolerance.
    """
    args = get_args()

    # Unwrap models to get the core module
    train_lm = train_model[0] if isinstance(train_model, (list, tuple)) else train_model
    inf_lm = inference_model[0] if isinstance(inference_model, (list, tuple)) else inference_model

    train_core = unwrap_model(train_lm)
    inf_core = unwrap_model(inf_lm)

    actual_vocab_size = getattr(args, 'padded_vocab_size', 128256)
    actual_seq_len = min(seq_len, getattr(args, 'seq_length', seq_len))
    device = torch.device(f"cuda:{torch.cuda.current_device()}")

    # Generate deterministic test input - same across ALL ranks
    torch.manual_seed(1234)
    test_tokens = torch.randint(
        low=0, high=actual_vocab_size, size=(batch_size, actual_seq_len),
        device=device, dtype=torch.long
    )
    test_position_ids = (
        torch.arange(actual_seq_len, device=device, dtype=torch.long)
        .unsqueeze(0)
        .expand(batch_size, -1)
    )
    test_attention_mask = torch.ones(
        (batch_size, 1, actual_seq_len, actual_seq_len), device=device, dtype=torch.bool
    )

    # Save and restore training state
    train_was_training = train_core.training
    inf_was_training = inf_core.training

    train_core.eval()
    inf_core.eval()

    try:
        with torch.no_grad():
            train_output = train_lm(
                test_tokens, test_position_ids, test_attention_mask,
                runtime_gather_output=True
            )

            inf_output = inf_lm(
                test_tokens, test_position_ids, test_attention_mask,
                runtime_gather_output=True
            )

        # Only check on ranks that have output (last PP stage)
        if train_output is not None and inf_output is not None:
            assert train_output.shape == inf_output.shape, (
                f"Output shape mismatch: train={train_output.shape}, infer={inf_output.shape}"
            )
            
            max_diff = (train_output - inf_output).abs().max().item()
            assert torch.allclose(train_output, inf_output, atol=atol, rtol=rtol), (
                f"Forward pass outputs do not match: max_diff={max_diff:.6e}, atol={atol}, rtol={rtol}"
            )

    finally:
        # Restore training state
        if train_was_training:
            train_core.train()
        if inf_was_training:
            inf_core.train()

Rollouts = list[TokenRollout | Rollout]
GroupedRollouts = list[Rollouts]


@dataclass(slots=True)
class RolloutStats:
    rewards: list[list[float]] # inner list is for a group
    env_ids: list[str] # same length as len(rewards)
    turn_lens: list[list[int]] # token lengths of turns, grouped.
    traj_lens: list[list[int]] # all turns comprise one trajectory.
    num_turns: None | list[list[int]] # num_turns per traj
    advantages: None | list[list[float]]
    min_piold_to_inf_prob: None | float
    max_piold_to_inf_prob: None | float
    mean_piold_to_inf_prob: None | float
    min_inf_train_prob_abs_diff: None | float
    max_inf_train_prob_abs_diff: None | float
    mean_inf_train_prob_abs_diff: None | float
    min_inf_prob: None | float
    max_inf_prob: None | float
    mean_inf_prob: None | float
    policy_staleness: list[list[int]]
    kv_cache_staleness: list[list[int]]
    completed_at_steps: list[list[int]]
    num_evictions: list[list[int]]


# Runtime state container for RL-specific data that shouldn't be checkpointed
class RLRuntimeState:
    """Container for runtime state that is not checkpointed, tracking state between rollout collections"""

    def __init__(self):
        self.packing_context = None
        self.last_collection_iteration = 0
        self.sequences_this_iteration_on_rank = 0
        self.latest_batch_num_sequences = 0

    def reset_iteration_counters(self, iteration):
        """Reset per-iteration counters."""
        self.sequences_this_iteration_on_rank = 0
        self.last_collection_iteration = iteration

    def increment_sequences(self, count):
        """Increment the sequence counter."""
        self.sequences_this_iteration_on_rank += count
        self.latest_batch_num_sequences = count


# Global runtime state instance
_rl_runtime_state = RLRuntimeState()


def get_rl_runtime_state():
    """Get the global RL runtime state."""
    return _rl_runtime_state


def update_inference_logprobs_group_stats(
    old_logprobs: torch.Tensor,
    inference_logprobs: torch.Tensor,
    mask: torch.Tensor,
    group_stats: Any,
) -> None:
    """Update group statistics with inference/train logprobs comparison metrics.

    This is the common statistics computation used by both packed and unpacked cases.

    Args:
        old_logprobs: Old logprobs tensor (train side)
        inference_logprobs: Inference logprobs tensor (aligned to match old_logprobs shape)
        mask: Boolean mask indicating valid positions for statistics
        group_stats: Statistics object to update with computed metrics
    """
    n_elems = mask.sum()
    if n_elems > 0:
        ratios = (old_logprobs - inference_logprobs).exp()[mask]
        abs_diffs = (old_logprobs.exp() - inference_logprobs.exp()).abs()[mask]

        group_stats.min_piold_to_inf_prob = ratios.min().item()
        group_stats.max_piold_to_inf_prob = ratios.max().item()
        group_stats.mean_piold_to_inf_prob = (ratios.sum() / n_elems).item()
        group_stats.min_inf_train_prob_abs_diff = abs_diffs.min().item()
        group_stats.max_inf_train_prob_abs_diff = abs_diffs.max().item()
        group_stats.mean_inf_train_prob_abs_diff = (abs_diffs.sum() / n_elems).item()

        inf_probs = inference_logprobs.exp()[mask]
        group_stats.min_inf_prob = inf_probs.min().item()
        group_stats.max_inf_prob = inf_probs.max().item()
        group_stats.mean_inf_prob = inf_probs.mean().item()


def align_unpacked_inference_logprobs(
    inference_logprobs: List[torch.Tensor],
    old_logprobs_for_data: torch.Tensor,
    generation_masks: torch.Tensor,
    group_stats: Any,
) -> torch.Tensor:
    """Align inference logprobs with old_logprobs for unpacked sequences and compute statistics.

    Args:
        inference_logprobs: List of inference logprobs tensors for each sequence
        old_logprobs_for_data: Template tensor with correct shape for alignment
        generation_masks: Tensor indicating which tokens were generated
        group_stats: Statistics object to update with computed metrics

    Returns:
        Aligned inference logprobs tensor
    """
    # Get first occurrence of a generation token
    # In get_logprobs() we chop off the first token -> the generation mask is shifted by one
    gen_masks_for_alignment = generation_masks
    first_gen_tok = gen_masks_for_alignment.int().argmax(dim=1) - 1

    # Align inference logprobs with old_logprobs
    # Note: We use old_logprobs_for_data as template since it has correct shape
    padded_inference_logprobs = old_logprobs_for_data.clone()

    # We need to align old_logprobs and inference logprobs as the latter are only for generations
    for i, inf_logprobs in enumerate(inference_logprobs):
        first_gen_idx = first_gen_tok[i]
        # We subtract -1 here because we append eod token on the train side, and we do not
        # get it from the inference. For the eod token, we reuse old_logprobs value.
        end_idx = min(first_gen_idx + len(inf_logprobs), padded_inference_logprobs.shape[1])
        actual_len = end_idx - first_gen_idx
        if actual_len > 0:
            padded_inference_logprobs[i, first_gen_idx:end_idx] = inf_logprobs[:actual_len]

    # Create truncated mask for statistics
    if old_logprobs_for_data.shape[1] + 1 < gen_masks_for_alignment.shape[1]:
        gen_masks_for_alignment = gen_masks_for_alignment[:, : old_logprobs_for_data.shape[1] + 1]

    truncated_mask = gen_masks_for_alignment[:, 1:].bool()

    # Final safety check
    if truncated_mask.shape != old_logprobs_for_data.shape:
        if truncated_mask.shape[1] > old_logprobs_for_data.shape[1]:
            truncated_mask = truncated_mask[:, : old_logprobs_for_data.shape[1]]
        elif truncated_mask.shape[1] < old_logprobs_for_data.shape[1]:
            pad_size = old_logprobs_for_data.shape[1] - truncated_mask.shape[1]
            truncated_mask = torch.nn.functional.pad(truncated_mask, (0, pad_size), value=False)

    # Sanity check: Two probability values cannot be more than 1.0 apart
    abs_diffs = (old_logprobs_for_data.exp() - padded_inference_logprobs.exp()).abs()[truncated_mask]
    assert all(abs_diffs <= 1.0)

    # Update group statistics using common helper
    update_inference_logprobs_group_stats(
        old_logprobs=old_logprobs_for_data,
        inference_logprobs=padded_inference_logprobs,
        mask=truncated_mask,
        group_stats=group_stats,
    )

    return padded_inference_logprobs


def get_agent(args, parallel_generation_tasks: int | None = None):
    """Get an agent based on environment configuration.

    If args.langrl_env_config is provided, uses weighted environment selection.
    Otherwise falls back to legacy single environment selection.
    """
    with open(args.langrl_env_config, 'r') as f:
        config = yaml.safe_load(f)

    return WeightedMultiTask.from_config(
        config,
        parallel_generation_tasks=parallel_generation_tasks,
    )


_INFERENCE_INTERFACE = None


def get_inference_interface(args, loop, model):
    global _INFERENCE_INTERFACE
    if _INFERENCE_INTERFACE is None:
        _INFERENCE_INTERFACE = loop.run_until_complete(
            MegatronLocal.launch(
                model[0],
                host='0.0.0.0',
                port=8294,
                verbose=args.inference_text_gen_server_logging)
        )
    return _INFERENCE_INTERFACE


_ROLLOUT_GENERATOR = None


def get_rollout_generator(args, inference_interface, n_prompts, samples_per_group):
    global _ROLLOUT_GENERATOR
    if not args.rl_partial_rollouts or _ROLLOUT_GENERATOR is None:
        agent = get_agent(args, parallel_generation_tasks=args.rl_parallel_generation_tasks)
        # Collect Rollouts
        request = GroupedRolloutRequest(
            num_groups=-1 if args.rl_partial_rollouts else n_prompts,
            rollouts_per_group=samples_per_group,
            inference_interface=inference_interface,
            generation_args={
                'temperature': args.rl_default_temperature,
                'max_tokens': args.inference_max_seq_length,
                'top_p': args.rl_default_top_p,
                'top_k': args.rl_default_top_k,
            },
            filter_groups_with_same_reward=args.grpo_filter_groups_with_same_reward,
        )
        _ROLLOUT_GENERATOR = agent.get_grouped_rollouts(request)
    return _ROLLOUT_GENERATOR


def get_environment_rollouts(
    model: LanguageModule, inference_model: LanguageModule, optimizer: MegatronOptimizer, n_prompts: int, samples_per_group: int
):
    """Sample environment rollouts from an LLM.

    Args:
        model: Model to sample from.
        inference_model: Inference model to use for inference.
        n_prompts: Number of prompts to sample for across *all* data parallel workers.
        samples_per_group: Amount of trajectories per prompt.

    Returns:
        GroupedRollouts object which is a nested list with each element being a list of rollouts of a group.
    """
    args = get_args()
    nvtx_range = get_nvtx_range()

    if args.rl_offload_optimizer_during_inference:
        with nvtx_range("rl/offload-optimizer-before-inference", time=True):
            if not args.rl_training_cuda_graphs:
                with nvtx_range("rl/offload/grad-buffers", time=True):
                    model[0].offload_grad_buffers()
            else:
                logger.warning(
                    "Gradient buffers will not be offloaded when training cudagraphs are enabled!")
            with nvtx_range("rl/offload/optimizer-state", time=True):
                optimizer.offload_to_cpu()

    # If we have separate training and inference models we to refit weights from the training model to the inference model.
    has_separate_inference_model = inference_model is not None
    if has_separate_inference_model:
        # If the separate inference model weights were prefetched to CPU while idle, bring them
        # back to GPU before refit/copy and before any CUDA-graph'd inference.
        with nvtx_range("rl/prefetch-weights-to-gpu", time=True):
            inf_core = unwrap_model(inference_model[0])
            _maybe_prefetch_separate_inference_model_weights(inf_core, to_cpu=False)
        swap_model_weights(model, inference_model, args.refit_method)
        if args.rl_verify_model_weights_swap:
            verify_model_weights_swap(
                train_model=model,
                inference_model=inference_model,
                atol=.1,
                rtol=5e-4,
            )
    else:
        inference_model = model

    inference_pg_collection = get_attr_wrapped_model(inference_model[0], "pg_collection")
    pg_size = get_pg_size(inference_pg_collection.ep)
    assert (n_prompts % pg_size == 0), f"{n_prompts=} must be divisible by {pg_size=}"

    with nvtx_range("rl/rollout-collection", time=True):
        loop = get_asyncio_loop()
        with megatron_rl_inference_mode(
            inference_model,
            optimizer,
            args.cuda_graph_impl,
            False, # offload optimizer during rollout collection is handled above
            training_model=model if has_separate_inference_model else None,
            increment_staleness_on_suspend=True,
        ) as inference_interface:

            with nvtx_range("rl/inference-setup", time=True):
                # Asyncronously run inference and rollout collection
                rollout_generator = get_rollout_generator(
                    args, inference_interface, n_prompts, samples_per_group
                )

            # NOTE(jbarker): we need to double check this when using PP>1
            rank = torch.distributed.get_rank()
            with nvtx_range("rl/collect-rollouts", time=True):
                if rank == 0:
                    log_single_rank(
                        logger,
                        logging.INFO,
                        f"Collecting rollouts, Iteration {args.curr_iteration}...",
                    )
                    rollouts = [
                        loop.run_until_complete(anext(rollout_generator)) for _ in range(n_prompts)
                    ]
                    # In deterministic mode, sort rollouts by problem_id for consistent ordering
                    # regardless of completion order due to system timing jitter.
                    if torch.are_deterministic_algorithms_enabled():
                        rollouts.sort(key=lambda group: group[0].problem_id if group and group[0].problem_id else "")
                    if not args.rl_partial_rollouts:
                        while True:
                            try:
                                loop.run_until_complete(anext(rollout_generator))
                                assert False, "Unexpected group left in generator."
                            except StopAsyncIteration:
                                break
                else:
                    # Just set up space to collect the rollouts
                    rollouts = [[None for _ in range(samples_per_group)] for _ in range(n_prompts)]

        with nvtx_range("rl/sync-rollouts", time=True):
            # Wait for Rollouts to be collected
            # TODO(jbarker): double check why this isn't causing rank 0 memory allocations
            torch.distributed.broadcast_object_list(rollouts, src=0)
        logger.debug(f"Got rollouts on rank {rank}")

    if args.rl_offload_optimizer_during_inference:
        with nvtx_range("rl/restore-optimizer-after-inference", time=True):
            with nvtx_range("rl/restore/grad-buffers", time=True):
                model[0].restore_grad_buffers()
            with nvtx_range("rl/restore/optimizer-state", time=True):
                optimizer.restore_from_cpu()

    if lang_rl_log_dir and rank == get_pg_rank(inference_pg_collection.tp):
        with open(
            lang_rl_log_dir
            + f'/rollouts_rank{rank}_iteration{args.curr_iteration}_'
            + f'{Path(args.langrl_env_config).stem}.json',
            'w',
        ) as f:
            json.dump([[r.model_dump() for r in group] for group in rollouts], f)

    return rollouts


def selective_log_softmax(logits, index):
    """Taken from: https://github.com/huggingface/trl/blob/26d86757a7c7e24e397ea44f57ecce6031dfac01/trl/trainer/utils.py#L1659.

    A memory-efficient implementation of the common `log_softmax -> gather` operation.

    This function is equivalent to the following naive implementation:
    ```python
    logps = torch.gather(logits.log_softmax(-1), dim=-1, index=index.unsqueeze(-1)).squeeze(-1)
    ```

    Args:
        logits (`torch.Tensor`):
            Logits tensor of shape `(..., num_classes)`.
        index (`torch.Tensor`):
            Index tensor of shape `(...)`, specifying the positions to gather from the log-softmax output.

    Returns:
        `torch.Tensor`:
            Gathered log probabilities with the same shape as `index`.
    """
    use_bik_logsoftmax = is_batch_invariant_mode_enabled()
    if logits.dtype in [torch.float32, torch.float64] and not use_bik_logsoftmax:
        selected_logits = torch.gather(logits, dim=-1, index=index.unsqueeze(-1)).squeeze(-1)
        # loop to reduce peak mem consumption
        logsumexp_values = torch.stack([torch.logsumexp(lg, dim=-1) for lg in logits])
        per_token_logps = (
            selected_logits - logsumexp_values
        )  # log_softmax(x_i) = x_i - logsumexp(x)
    else:
        # logsumexp approach is unstable with bfloat16, fall back to slightly less efficent approach
        per_token_logps = []
        for row_logits, row_labels in zip(logits, index):  # loop to reduce peak mem consumption
            row_logps = torch.nn.functional.log_softmax(row_logits, dim=-1)
            row_per_token_logps = row_logps.gather(dim=-1, index=row_labels.unsqueeze(-1)).squeeze(
                -1
            )
            per_token_logps.append(row_per_token_logps)
        per_token_logps = torch.stack(per_token_logps)
    return per_token_logps


def get_logprobs(model, tokens, position_ids, no_grad=False, sequence_packing=False, packed_seq_params=None):
    """Get sequence logprobs from their token ids.

    Args:
        model: model to predict with.
        tokens: inputs for which we want to get logprobs.
        position_ids: position ids that come with tokens.
        attention_mask: attention mask that comes with tokens.
        no_grad: whether to run in no_grad mode.
        packed_seq_params: Optional PackedSeqParams for sequence packing with TE.
            When provided with qkv_format='thd', the input tokens are sliced to
            remove padding before the forward pass, and outputs are padded back.
        packed_seq_len: Optional length of the packed sequence (excluding padding).
            Required when packed_seq_params is provided to avoid CPU-GPU synchronization.

    Returns:
        Logprobs of input sequences.

    """

    args = get_args()
    # Ensure packed_seq_params is always provided for CUDA graph signature consistency.
    # When sequence_packing is enabled, construct from packing config (max_sequences_per_bin).
    # When sequence_packing is disabled, construct a single-sequence default so the CUDA
    # graph signature matches the training forward_step in train_rl.py.
    # This is necessary because reference logprobs steps will reuse the training forward graph.
    if packed_seq_params is None:
        if sequence_packing:
            packed_seq_params = get_default_packed_seq_params(
                seq_length=tokens.shape[1],
                max_sequences_per_bin=args.rl_sequence_packing_max_sequences_per_bin,
                device=tokens.device,
            )
        else:
            cu_seqlens = torch.tensor([0, tokens.shape[1]], dtype=torch.int32, device=tokens.device)
            packed_seq_params = PackedSeqParams(
                qkv_format='thd',
                cu_seqlens_q=cu_seqlens,
                cu_seqlens_kv=cu_seqlens,
                max_seqlen_q=tokens.shape[1],
                max_seqlen_kv=tokens.shape[1],
                total_tokens=tokens.shape[1],
            )

    nvtx_range = get_nvtx_range()

    with nvtx_range("rl/get-logprobs", time=True):
        with nvtx_range("rl/forward-pass", time=True):
            # TODO(vitalyk): use fp16/bf16 as a function argument. Do not use args.

            attention_mask_for_forward = None

            # This is a hack to fix megatron's behaviour when flash-decode affects the training code flow.
            flash_decode = model.config.flash_decode
            model.config.flash_decode = False
            fp32_output = not (args.fp16 or args.bf16)
            with torch.no_grad() if no_grad else nullcontext():
                logits_or_hidden_states = model(
                    tokens,
                    position_ids,
                    attention_mask_for_forward,
                    packed_seq_params=packed_seq_params,
                    runtime_gather_output=True,
                    fp32_output=fp32_output,
                )
            model.config.flash_decode = flash_decode

        pg_collection = get_attr_wrapped_model(model, "pg_collection")
        pp_group = pg_collection.pp

        if not is_pp_last_stage(pp_group):
            return logits_or_hidden_states
        else:
            logits = logits_or_hidden_states
            with nvtx_range("rl/log-softmax", time=True):
                # We do not need logprobs for the n+1 token.
                logprobs = selective_log_softmax(logits[:, :-1, :], tokens[:, 1:])
            return logprobs


def calculate_grpo_advantages(rewards: list[list[float]], num_turns: list[list[int]]) -> np.ndarray:
    """Calculate GRPO advantages from rewards/num_turns.

    For multiturn rollouts, the logic is a bit more involved.
    # For training, we'll be turning each turn into a trajectory with the same reward
    # within a trajectory, e.g. if [[a,b],[c,d,e]] trajectory has reward 1.0, we will
    # get [a,b] with 1.0 and [c,d,e] with 1.0 when doing updates.
    """

    rewards = np.array(rewards)

    num_turns = np.array(num_turns)
    # Each outer dimension of num_turns is a group. Sum of those gives total num_turns per group.
    # Let's use this to calculate advantage.
    # mean/std should be repeated based on group lens
    group_turns = num_turns.sum(axis=-1)
    reward_means = rewards.mean(axis=1, keepdims=True).repeat(group_turns)
    reward_stds = rewards.std(axis=1, keepdims=True).repeat(group_turns)

    # rewards are originally [g, group_size]
    # Making an assumption that all groups are of the same size!
    # @vitalyk: this will go away when we start sending env-based sample reqs.
    rewards = rewards.flatten().repeat(num_turns.flatten())

    return ((rewards - reward_means) / (1e-4 + reward_stds)).tolist()


def compute_group_stats(
    rollouts: GroupedRollouts, tokenizer: MegatronTokenizer, seq_len: int,
) -> RolloutStats:
    """Add group-based rollout stats for logging.

    Args:
        rollouts: Rollouts to generate the stats for. Each inner list is a group (as in GRPO group), i.e. all rollouts are for the same prompt.
        tokenizer: Tokenizer to tokenize the rollouts in case they are raw strings.
        seq_len: Maximum sequence length.

    Returns:
       RolloutStats object containing all the stats.
    """
    # TODO (rkirby) Maybe do some of this after the tensor building
    group_reward_means = []
    group_reward_stds = []
    turn_lens = []
    traj_lens = []
    rewards = []
    env_ids = []
    group_reward_ids = []
    num_turns = [] # num_turns per traj
    all_policy_staleness = []
    all_kv_cache_staleness = []
    all_completed_at_steps = []
    all_num_evictions = []
    for group in rollouts:
        group_rewards = []
        group_traj_lengths = []
        group_turn_lengths = []
        group_num_turns = []
        group_policy_staleness = []
        group_kv_staleness = []
        group_completed_at_steps = []
        group_num_evictions = []
        for rollout in group:
            if isinstance(rollout, TokenRollout):
                for turn_traj in rollout.trajectory:
                    detokenized_traj = tokenizer.detokenize(turn_traj)
                    lang_rl_log(
                        f"Rollout: [{rollout.env_id}] [{rollout.reward} : {len(rollout.trajectory)} tokens] {detokenized_traj}"
                    )
                    # TODO(vitalyk): how does multiturn change EOD/EOT?
                    assert (len(turn_traj) == seq_len) or (
                        turn_traj[-1] == tokenizer.eod
                    ), f"Rollout is not the correct length: {len(turn_traj)} {turn_traj[-1]}\n{detokenized_traj}"
            else:
                lang_rl_log(
                    f"Rollout: [{rollout.env_id}] [{rollout.reward} : {len(rollout.trajectory)} chars] {rollout.trajectory}"
                )
            group_num_turns.append(len(rollout.trajectory))
            group_rewards.append(rollout.reward)
            roll_turn_lens = [len(t) for t in rollout.trajectory]
            group_turn_lengths.extend(roll_turn_lens)
            group_traj_lengths.append(sum(roll_turn_lens))
            group_policy_staleness.extend(s for turn in rollout.policy_staleness for s in turn)
            group_kv_staleness.extend(s for turn in rollout.kv_cache_staleness for s in turn)
            group_completed_at_steps.extend(rollout.completed_at_step)
            group_num_evictions.append(sum(rollout.num_evictions))
        all_policy_staleness.append(group_policy_staleness)
        all_kv_cache_staleness.append(group_kv_staleness)
        all_completed_at_steps.append(group_completed_at_steps)
        all_num_evictions.append(group_num_evictions)
        traj_lens.append(group_traj_lengths)
        turn_lens.append(group_turn_lengths)
        env_ids.append(group[0].env_id) # All rollouts in a group share the env_id by design.
        rewards.append(group_rewards)
        # https://arxiv.org/abs/2504.21233 reports that lens variance hurts.
        # Let's track this.
        num_turns.append(group_num_turns)

    stats = RolloutStats(
        traj_lens=traj_lens,
        turn_lens=turn_lens,
        rewards=rewards,
        # --------
        # Everything above is per-group, i.e. it is a list of lists,
        # with the inner list being the group data.
        env_ids=env_ids,
        num_turns=num_turns,
        advantages=calculate_grpo_advantages(rewards, num_turns),
        min_piold_to_inf_prob=None,
        max_piold_to_inf_prob=None,
        mean_piold_to_inf_prob=None,
        min_inf_train_prob_abs_diff=None,
        max_inf_train_prob_abs_diff=None,
        mean_inf_train_prob_abs_diff=None,
        min_inf_prob=None,
        max_inf_prob=None,
        mean_inf_prob=None,
        policy_staleness=all_policy_staleness,
        kv_cache_staleness=all_kv_cache_staleness,
        completed_at_steps=all_completed_at_steps,
        num_evictions=all_num_evictions,
    )
    return stats


def compute_true_staleness(
    per_token_staleness: list[list[int]],
    completed_at_steps: list[list[int]],
    turn_lens: list[list[int]],
    current_iteration: int,
) -> list[int]:
    """Compute true per-token staleness by adding the completion gap.

    Args:
        per_token_staleness: Grouped flat list of per-token raw staleness values.
        completed_at_steps: Grouped list of per-turn completion steps.
        turn_lens: Grouped list of per-turn token counts.
        current_iteration: Current training iteration.

    Returns:
        Flat list of true staleness values (one per token across all groups).
    """
    result = []
    for group_staleness, group_completed, group_turn_lens in zip(
        per_token_staleness, completed_at_steps, turn_lens
    ):
        token_idx = 0
        for completed_at, num_tokens in zip(group_completed, group_turn_lens):
            gap = current_iteration - completed_at
            for _ in range(num_tokens):
                result.append(group_staleness[token_idx] + gap)
                token_idx += 1
    return result


def prep_wandb_metrics(
        wandb_writer: wandb_run.Run,
        traj_lens: List[List[int]],
        turn_lens: List[List[int]],
        rewards: List[List[float]],
        num_turns: List[List[int]],
        advantages: List[float],
        policy_staleness: List[List[int]],
        kv_cache_staleness: List[List[int]],
        num_evictions: List[List[int]],
        completed_at_steps: List[List[int]],
        current_iteration: int,
        example_group: list[TokenRollout | Rollout] | None = None,
        tokenizer: MegatronTokenizer | None = None,
    ):

    """Make a wandb-parseable dictionary of metrics for logging.

    Args:
        wandb_writer: Wandb run to log to.
        traj_lens: Grouped list of trajectory lengths.
        turn_lens: Grouped list of turn lengths.
        rewards: Grouped list of rewards.
        num_turns: Grouped list of number of turns in the trajectories.
        advantages: Flattened list of advantages.
        policy_staleness: Grouped list of per-token policy staleness.
        kv_cache_staleness: Grouped list of per-token KV cache staleness.
        num_evictions: Grouped list of per-rollout number of evictions.
        completed_at_steps: Grouped list of per-turn completed at steps.
        current_iteration: Current training iteration.
        example_group: A list of rollouts of one group to log examples of trajectories.
        tokenizer: Tokenizer to untokenize trajectories for logging.
    """

    group_table = wandb_writer.Table(
        columns=['group_means', 'group_stds'],
        data=[[np.mean(g), np.std(g)] for g in rewards],
    )

    true_policy_staleness = compute_true_staleness(
        policy_staleness, completed_at_steps, turn_lens, current_iteration)
    true_kv_staleness = compute_true_staleness(
        kv_cache_staleness, completed_at_steps, turn_lens, current_iteration)

    metrics = {
            'group_means_hist': wandb_writer.plot.histogram(
                group_table, 'group_means', 'Group Means'
            ),
            'group_stds_hist': wandb_writer.plot.histogram(
                group_table, 'group_stds', 'Group STDs'
            ),
            'rewards_hist': wandb_writer.plot.histogram(
                wandb_writer.Table(
                    columns=['reward'], data=[[r] for g in rewards for r in g]
                ),
                'reward', 'All Rewards'
            ),
            'advantages_hist': wandb_writer.plot.histogram(
                wandb_writer.Table(
                    columns=['advantages'], data=[[x] for x in advantages]
                ),
                'advantages', 'Advantages'
            ),
            'rollout_table': wandb_writer.Table(
                columns=['reward', 'traj_length', 'num_evictions'],
                data=list(zip(
                    [r for g in rewards for r in g],
                    [l for g in traj_lens for l in g],
                    [e for g in num_evictions for e in g],
                )),
            ),
            'mean_turn_length': np.mean([np.mean(g) for g in turn_lens]),
            'mean_turn_length_std': np.mean([np.std(g) for g in turn_lens]),
            'max_turn_length': max([max(g) for g in turn_lens]),
            'min_turn_length': min([min(g) for g in turn_lens]),
            'mean_traj_length': np.mean([np.mean(g) for g in traj_lens]),
            'mean_traj_length_std': np.mean([np.std(g) for g in traj_lens]),
            'max_traj_length': max([max(g) for g in traj_lens]),
            'min_traj_length': min([min(g) for g in traj_lens]),
            'mean_num_turns': np.mean([np.mean(g) for g in num_turns]),
            'max_num_turns': max([max(g) for g in num_turns]),
            'min_num_turns': min([min(g) for g in num_turns]),
            'mean_reward': np.mean([np.mean(g) for g in rewards]),
            'mean_advantage': np.mean(advantages),
            'nonzero_groups_ratio': np.count_nonzero(advantages)
            / len(advantages),
            'mean_policy_staleness': np.mean(true_policy_staleness),
            'max_policy_staleness': max(true_policy_staleness),
            'min_policy_staleness': min(true_policy_staleness),
            'mean_kv_cache_staleness': np.mean(true_kv_staleness),
            'max_kv_cache_staleness': max(true_kv_staleness),
            'min_kv_cache_staleness': min(true_kv_staleness),
            'total_eviction_count': sum([sum(g) for g in num_evictions]),
            'max_num_evictions': max([max(g) for g in num_evictions]),
            'mean_completion_gap': np.mean([current_iteration - s for g in completed_at_steps for s in g]),
    }
    if example_group:
        if tokenizer is None:
            raise ValueError("If you provide an example group to log, you need to provide a tokenizer too.")
        metrics['rollouts'] = wandb_writer.Table(
            columns=['Trajectories', 'Tokens', 'Rewards'],
            rows=[
                [
                    tokenizer.detokenize(turn) if isinstance(r, TokenRollout) else turn,
                    r.trajectory,
                    r.reward,
                ]
                for r in example_group for turn in r.trajectory
            ],
        )
    return metrics


def maybe_log_training_metrics(
    group_stats: RolloutStats,
    current_iteration: int,
    tokenizer: MegatronTokenizer,
    example_groups: dict[str, list[TokenRollout | Rollout]],
):
    """Log training metrics if writers are available.

    Args:
        group_stats: RolloutStats object to pass to writers.
        current_iteration: Current training iteration.
        tokenizer: Tokenizer to untokenize trajectories for logging.
        example_groups: A dict with values as list of rollouts of one group to log examples of trajectories. Keys are env names.
    """

    wandb_writer = get_wandb_writer()
    tb_writer = get_tensorboard_writer()
    if tb_writer:
        tb_writer.add_scalar('mean_reward', np.mean([np.mean(g) for g in group_stats.rewards]), current_iteration)
    if not wandb_writer:
        return

    # We log these metrics for the aggregated data, no split per env.
    metrics = {
        'min_piold_to_inf_prob': group_stats.min_piold_to_inf_prob,
        'max_piold_to_inf_prob': group_stats.max_piold_to_inf_prob,
        'mean_piold_to_inf_prob': group_stats.mean_piold_to_inf_prob,
        'min_inf_train_prob_abs_diff': group_stats.min_inf_train_prob_abs_diff,
        'max_inf_train_prob_abs_diff': group_stats.max_inf_train_prob_abs_diff,
        'mean_inf_train_prob_abs_diff': group_stats.mean_inf_train_prob_abs_diff,
        'min_inf_prob': group_stats.min_inf_prob,
        'max_inf_prob': group_stats.max_inf_prob,
        'mean_inf_prob': group_stats.mean_inf_prob,
    }

    traj_lens = group_stats.traj_lens
    turn_lens = group_stats.turn_lens
    rewards = group_stats.rewards
    num_turns = group_stats.num_turns
    advantages = group_stats.advantages
    policy_staleness = group_stats.policy_staleness
    kv_cache_staleness = group_stats.kv_cache_staleness
    num_evictions = group_stats.num_evictions
    completed_at_steps = group_stats.completed_at_steps

    metrics = metrics | prep_wandb_metrics(wandb_writer=wandb_writer,
        traj_lens=traj_lens, turn_lens=turn_lens, rewards=rewards, num_turns=num_turns, advantages=advantages,
        policy_staleness=policy_staleness, kv_cache_staleness=kv_cache_staleness, num_evictions=num_evictions,
        completed_at_steps=completed_at_steps, current_iteration=current_iteration)
    env_stats = lambda cont, idx: [cont[i] for i in idx]
    group_turn_counts = [sum(nt) for nt in num_turns]

    for env_id in set(group_stats.env_ids):
        env_idx = [i for i, eidx in enumerate(group_stats.env_ids) if eidx == env_id]

        # Advantages are flattened, we need to be more careful with those.
        env_advantages = []
        for i in env_idx:
            st = sum(group_turn_counts[:i])
            end = st + group_turn_counts[i]
            env_advantages.extend(advantages[st:end])

        env_metrics = prep_wandb_metrics(wandb_writer=wandb_writer, traj_lens=env_stats(traj_lens, env_idx),
            turn_lens=env_stats(turn_lens, env_idx),
            rewards=env_stats(rewards, env_idx),
            num_turns=env_stats(num_turns, env_idx),
            advantages=env_advantages,
            policy_staleness=env_stats(policy_staleness, env_idx),
            kv_cache_staleness=env_stats(kv_cache_staleness, env_idx),
            num_evictions=env_stats(num_evictions, env_idx),
            completed_at_steps=env_stats(completed_at_steps, env_idx),
            current_iteration=current_iteration,
            example_group=example_groups[env_id],
            tokenizer=tokenizer,
        )
        for k, v in env_metrics.items():
            metrics[f"{env_id}_{k}"] = v

    wandb_writer.log(metrics, step=current_iteration)


def prepare_trajectories(
    rollouts: Rollouts, tokenizer: MegatronTokenizer, seq_length: int, sequence_packing: bool, skip_bos_token: bool
):
    """Pad trajectories and extract the generation masks.
    Args:
        rollouts: Rollouts to extract trajectories from.
        tokenizer: Tokenizer to get the padding token and potentially tokenize.
        seq_length:  Maximum sequence length to pad to.

    Returns:
        Trajectories and their generation masks.

    Raises:
        ValueError:
    """
    # Track counts for each environment ID
    env_id_counts = Counter()

    DEFAULT_PAD_TOKENS = ['<|finetune_right_pad_id|>', '<SPECIAL_999>']

    if tokenizer.library == "huggingface":
        tokenizer : HuggingFaceTokenizer
        if not tokenizer.pad:
            for pad_token in DEFAULT_PAD_TOKENS:
                if pad_token in tokenizer._tokenizer.tokenizer.get_vocab():
                    log_single_rank(
                        logger, logging.INFO, f"Updating tokenizer pad token to {pad_token}"
                    )
                    tokenizer._tokenizer.pad_token = pad_token
                    break
            else:
                raise ValueError("No pad token found in tokenizer vocabulary")
    elif tokenizer.library == "tiktoken":
        assert "<SPECIAL_233>" in tokenizer.vocab, "Pad token is NOT in the tokenizer"
        tokenizer._pad_id = tokenizer.vocab["<SPECIAL_233>"]

    log_single_rank(logger, logging.INFO, f"Tokenizer vocab size: {tokenizer.vocab_size}")
    log_single_rank(
        logger,
        logging.INFO,
        f"Tokenizer PAD: '{tokenizer.detokenize([tokenizer.pad])} ({tokenizer.pad})'",
    )
    log_single_rank(
        logger,
        logging.INFO,
        f"Tokenizer EOD: '{tokenizer.detokenize([tokenizer.eod])} ({tokenizer.eod})'",
    )

    trajs = []
    generation_masks = []
    inference_logprobs = []
    for rollout in rollouts:
        # traj, gen mask and logprobs are lists now.
        # each list entry is a turn, single-turn environments just have a single-element list.
        # We assume that all lengths of the structs above have the same lengths (number of turns).

        all_turns_trajectories = (
            copy.deepcopy(rollout.trajectory)
            if isinstance(rollout, TokenRollout)
            else tokenizer.tokenize(rollout.trajectory)
        )
        for turn_idx, trajectory in enumerate(all_turns_trajectories):
            inf_logprobs = rollout.logprobs[turn_idx]
            generation_mask = rollout.generation_mask[turn_idx] if isinstance(rollout, TokenRollout) else None
            length = len(trajectory)
            assert length <= seq_length, "Rollout too long, how did this happen?"
            if len(trajectory) < seq_length:
                assert (
                    trajectory[-1] == tokenizer.eod
                ), "Trajectories under a seq_length limit should have eod token at the end."

            if length < seq_length:
                trajectory.extend([tokenizer.pad] * (seq_length - length))
                if generation_mask:
                    generation_mask.extend([False] * (seq_length - length))
            trajs.append(trajectory)
            generation_masks.append(generation_mask)

            if inf_logprobs is not None:
                inf_logprobs_tensor = torch.Tensor(inf_logprobs)
                # Don't pad individual logprobs here - padding happens later if needed
                inference_logprobs.append(inf_logprobs_tensor)
            else:
                inference_logprobs.append(None)

        env_id_counts[rollout.env_id] += 1

    if torch.distributed.is_initialized():
        logger.info(f"[{dist.get_rank()}] Rollout counts:")
        for env_id, count in env_id_counts.items():
            logger.info(f"[{dist.get_rank()}] \t{env_id}: {count}")

    generation_masks = torch.tensor(generation_masks, dtype=torch.bool, device='cpu')
    trajs = torch.tensor(trajs, device='cpu')

    # Only process if we have inference_logprobs
    if inference_logprobs and any(lp is not None for lp in inference_logprobs):
        # We need to pad all logprobs to the same size for sequence packing.
        # For non-packing mode, keep as list of tensors (unpadded)
        # This preserves the original behavior where each sequence can have different lengths
        if sequence_packing:
            inference_logprobs = _pad_nonnull_with_zeros(inference_logprobs, seq_length)
    else:
        inference_logprobs = None

    # Some sanity checks regarding the tokenization
    if not skip_bos_token:
        assert (
            tokenizer.bos is None or (trajs[:, 0] == tokenizer.bos).all()
        ), "First token should be bos"
    else:
        assert (
            tokenizer.bos is None or (trajs[:, 0] != tokenizer.bos).all()
        ), "First token should not be bos"  
    assert (
        tokenizer.bos is None or (trajs[:, 1] != tokenizer.bos).all()
    ), "Second token should not be bos"
    assert (
        (trajs * generation_masks.int() == tokenizer.eod).sum(axis=1) <= 1
    ).all(), "Only one eod per trajectory in generated tokens."
    # TODO(rkirby):
    # We should avoid the tokenizer pad token being the same as the eod token for proper loss masking,
    # But now the deepseek tokenizer has the pad token set to eod, we need to handle this.
    # assert (tokenizer.pad != tokenizer.eod), "Pad and eod should be different"
    return trajs, generation_masks, inference_logprobs


def logprobs_forward_step(data_iterator, model, is_correction, packing_context=None):
    # Avoid self.training checks which will trigger cudagraph capture; this path reuses
    # the forward pass from training after it has been captured on the 1st iteration.
    model.eval()

    if packing_context is not None:
        # When using sequence packing, the data iterator returns a tuple with a single element, the bin index.
        bin_tensor = next(data_iterator)[0]
        #TODO(jalbericiola): change for named tuple
        (b_trajs, _, _, _, b_posids, _, _, _, _, _, b_packed_seq_params) = (
            load_packed_data_by_index(bin_tensor.item(), packing_context, is_correction)
        )
    else:
        b_trajs, b_posids = next(data_iterator)
        b_packed_seq_params = None

    logprobs = (
        get_logprobs(
            model,
            b_trajs.cuda(),
            b_posids.cuda(),
            no_grad=True,
            sequence_packing=packing_context is not None,
            packed_seq_params=b_packed_seq_params,
        ),
        None,
    )
    model.train()
    return logprobs


def compute_logprobs_batch(
    model,
    data_loader,
    forward_backward_func,
    packing_context,
    trajs_batch_size, # n_bins for seq packing, and batch_size for non seq packing
    seq_length,
    logprobs_batch_size,
    decoder_seq_length,
    dtype,
    pp_group,
    is_correction,
    collect_non_loss_data=False,
):
    """Compute logprobs for all batches in the data loader."""
    logprobs_list = []
    data_iterator = iter(data_loader)
    for i in range(len(data_loader)):
        output_tensor = forward_backward_func(
            forward_step_func=partial(logprobs_forward_step, is_correction=is_correction, packing_context=packing_context),
            data_iterator=data_iterator,
            model=model,
            num_microbatches=1,
            seq_length=seq_length,
            micro_batch_size=logprobs_batch_size,
            decoder_seq_length=decoder_seq_length,
            forward_only=True,
            adjust_tensor_shapes_fn=None,
            collect_non_loss_data=collect_non_loss_data,
        )
        if is_pp_last_stage(pp_group):
            logprobs_list.append(output_tensor[0].detach())

    if is_pp_last_stage(pp_group):
        logprobs = torch.concat(logprobs_list, dim=0)
        assert logprobs.dtype == dtype
    else:
        logprobs = torch.empty(
            trajs_batch_size,
            seq_length-1,
            dtype=dtype,
            device=torch.cuda.current_device(),
        )

    # Only PP>1 needs a broadcast from the last stage; for PP=1 the output is already local.
    if get_pg_size(pp_group) > 1:
        dist.broadcast(logprobs, src=get_pp_last_rank(pp_group), group=pp_group)
    return logprobs.cpu()


def prepare_data_for_update(
    model: list[LanguageModule],
    ref_state_dict: Dict[str, Any],
    rollouts: GroupedRollouts,
    tokenizer: MegatronTokenizer,
    sequence_packing: bool,
    is_correction: bool,
) -> tuple[RerunDataIterator, RolloutStats, dict]:
    """Extract data for the update from raw rollouts.

    Args:
        model: Current policy as the zero-eth element.
        ref_state_dict: Reference policy state dict.
        rollouts: Rollouts to extract the data from.
        tokenizer: Tokenizer to pad/tokenize data.
        sequence_packing: Use sequence packing if True.
        is_correction: Prepare data for IS correction if True.

    Returns:
        Tuple of (cycled iterator over dataset batches, group stats, example groups per env).
    """
    args = get_args()
    nvtx_range = get_nvtx_range()
    runtime_state = get_rl_runtime_state()

    if args.cuda_graph_impl != "none" and not args.rl_training_cuda_graphs:
        lang_module = (
            model[0].module.module if hasattr(model[0].module, "module") else model[0].module
        )
        toggle_cuda_graphs(lang_module, "none")

    model = model[0]
    dtype = torch.bfloat16 if args.bf16 else (torch.float16 if args.fp16 else torch.float32)

    with nvtx_range("rl/prepare-data-for-update", time=True):
        with nvtx_range("rl/compute-group-stats", time=True):
            group_stats = compute_group_stats(rollouts, tokenizer, args.seq_length)
            # TODO(vitalyk): why do we need global_advantages here? go inside packing
            advantages = global_advantages = torch.tensor(group_stats.advantages, dtype=dtype).cuda()

        # Now split the rollouts across the data parallel ranks for training
        # This needs to be done at this point because we are about to calculate logprobs
        # Note :- For EP, do not use the expert data parallel group here. Always 
        # use the regular data parallel group. 

        # Get example group per environment to log their rollouts.
        example_groups = {}
        for g in rollouts:
            if g[0].env_id not in example_groups:
                example_groups[g[0].env_id] = g

        # Let's expand rollouts getting rid of the groups.
        # We need this to correctly split the rollouts across dp groups.
        # And we do not actually need them grouped in anything below anyways.
        rollouts = [r for g in rollouts for r in g]
        num_turns = [nt for g in group_stats.num_turns for nt in g]
        total_turns_sampled = len(rollouts)

        # We might sample more than we consume in one step.
        samples_ratio_per_step = args.global_batch_size / (args.grpo_prompts_per_step * args.grpo_group_size)
        assert samples_ratio_per_step <= 1, "You cannot use more data than you sampled."

        if (data_parallel_world_size := mpu.get_data_parallel_world_size()) > 0:
            data_split_size = len(rollouts) // data_parallel_world_size
            data_split_range = (
                mpu.get_data_parallel_rank() * data_split_size,
                (mpu.get_data_parallel_rank() + 1) * data_split_size,
            )
            rollouts = rollouts[data_split_range[0] : data_split_range[1]]
            local_num_turns = sum(num_turns[data_split_range[0] : data_split_range[1]])
            steps_before = sum(num_turns[:data_split_range[0]])
            advantages = advantages[steps_before:steps_before+local_num_turns]
            # First we calculate them on a global level and then we split and recalculate on a local level.
            # Sequence packing and reporting needs it global but non-packing wants it local.

        with nvtx_range("rl/prepare-trajectories", time=True):
            trajs, generation_masks, inference_logprobs = prepare_trajectories(
                rollouts, tokenizer, args.seq_length, sequence_packing, args.rl_skip_bos_token
            )

        packing_context = None
        # Build trajectories based on sequence packing or standard processing
        if args.rl_use_sequence_packing:
            with nvtx_range("rl/sequence-packing", time=True):
                runtime_state.packing_context = packing_context = pack_all_trajectories(
                    trajs, 
                    generation_masks, 
                    inference_logprobs, 
                    global_advantages, 
                    args.seq_length, 
                    args.rl_sequence_packing_max_sequences_per_bin,
                    args.rl_sequence_packing_algo
                    )
    
                compute_trajs = packing_context.packed_trajs
                compute_position_ids = packing_context.packed_position_ids
                # Use batch_size=1 for packed computation to enable proper attention masking
                # via PackedSeqParams (TE needs cu_seqlens per bin)
                dataset = TensorDataset(torch.arange(len(compute_trajs)))
                data_loader = DataLoader(dataset, batch_size=1)
                logprobs_batch_size = 1
        else:
            # Always compute standard masks for the original data (we'll need them later)
            with nvtx_range("rl/get-ltor-masks", time=True):
                _, original_loss_mask, original_position_ids = get_ltor_masks_and_position_ids(
                    trajs,
                    tokenizer.eod,
                    tokenizer.pad,
                    args.reset_position_ids,
                    args.reset_attention_mask,
                    eod_mask_loss=False,
                    pad_mask_loss=True,
                )
                original_loss_mask[~generation_masks] = 0.0
                compute_trajs = trajs
                compute_position_ids = original_position_ids
                data_loader = DataLoader(
                    TensorDataset(compute_trajs, compute_position_ids),
                    batch_size=args.micro_batch_size,
                )
                logprobs_batch_size = args.micro_batch_size


        with torch.no_grad(), nvtx_range("rl/compute-logprobs", time=True):
            # Before we can update the model, we need to get the logprobs for the \pi_{old} model.

            forward_backward_func = get_forward_backward_func()
            if args.cuda_graph_impl == "local" and CudaGraphScope.full_iteration in args.cuda_graph_scope:
                forward_backward_func = FullCudaGraphWrapper(
                    forward_backward_func, cuda_graph_warmup_steps=args.cuda_graph_warmup_steps
                )

            dtype = (
                torch.bfloat16 if args.bf16 else (torch.float16 if args.fp16 else torch.float32)
            )

            pg_collection = get_attr_wrapped_model(model, "pg_collection")
            pp_group = pg_collection.pp

            with torch.no_grad(), nvtx_range("rl/compute-old-logprobs", time=True):
                old_logprobs = compute_logprobs_batch(
                    model=model,
                    data_loader=data_loader,
                    forward_backward_func=forward_backward_func,
                    packing_context=packing_context,
                    trajs_batch_size=len(compute_trajs),
                    seq_length=args.seq_length,
                    logprobs_batch_size=logprobs_batch_size,
                    decoder_seq_length=args.decoder_seq_length,
                    dtype=dtype,
                    pp_group=pp_group,
                    is_correction=args.rl_inference_logprobs_is_correction,
                )

            with torch.no_grad(), nvtx_range("rl/compute-ref-logprobs", time=True):
                # We need to load the ref model state dict and compute the logprobs for the ref model
                cur_st_dict = {
                    k: (v.cpu() if v is not None else v) for k, v in model.state_dict().items()
                }
                model.load_state_dict(ref_state_dict)
                ref_logprobs = compute_logprobs_batch(
                    model=model,
                    data_loader=data_loader,
                    forward_backward_func=forward_backward_func,
                    packing_context=packing_context,
                    trajs_batch_size=len(compute_trajs),
                    seq_length=args.seq_length,
                    logprobs_batch_size=logprobs_batch_size,
                    decoder_seq_length=args.decoder_seq_length,
                    dtype=dtype,
                    pp_group=pp_group,
                    is_correction=args.rl_inference_logprobs_is_correction,
                )

                # logprobs are [b, seq, h] now.
                model.load_state_dict(cur_st_dict)

                with nvtx_range("rl/synchronize-cuda-and-collect-garbage", time=True):
                    torch.cuda.synchronize()
                    gc.collect()
                    torch.cuda.empty_cache()


        if sequence_packing:
            with nvtx_range("rl/pack-logprobs", time=True):
                # Store logprobs on gpu in packing context
                # Since PackingContext is a dataclass, we add these as new attributes
                packing_context.old_logprobs = old_logprobs.cuda()
                packing_context.ref_logprobs = ref_logprobs.cuda()

                if inference_logprobs is not None:
                    # Pack the inference logprobs using the helper function
                    # We do this for logging purposes even if is_correction is disabled
                    packed_inference_logprobs = pack_inference_logprobs(
                        inference_logprobs=packing_context.original_inference_logprobs,
                        packing_info=packing_context.packing_info,
                        generation_masks=packing_context.original_generation_masks,
                        bin_size=args.seq_length,
                    )

                    # Compute statistics for logging using packed data
                    compute_packed_inference_logprobs_stats(
                        old_logprobs=old_logprobs,
                        packed_inference_logprobs=packed_inference_logprobs,
                        packed_loss_mask=packing_context.packed_loss_mask,
                        group_stats=group_stats,
                    )

                    # Store packed inference logprobs in packing context
                    packing_context.packed_inference_logprobs = packed_inference_logprobs.cuda()
                    # Only mark as having inference logprobs for IS correction if enabled
                    packing_context.has_inference_logprobs = args.rl_inference_logprobs_is_correction
            with nvtx_range("rl/create-dataloader", time=True):
                # @vitalyk: This function also reconfigures the data loader to count the
                # global_batch_size in the bins frame of reference.
                # I think it will be a better design if we split the data loader creating and logic
                # that reconfigures the microbatch calculator.

                update_microbatch_calculator(
                    samples_ratio_per_step=samples_ratio_per_step,
                    num_bins_this_rank = len(packing_context.packed_trajs),
                    bin_seq_indices = packing_context.packing_info.bin_seq_indices,
                    global_batch_size=args.global_batch_size, 
                    rampup_batch_size=args.rampup_batch_size, 
                    micro_batch_size=args.micro_batch_size, 
                    decrease_batch_size_if_needed=args.decrease_batch_size_if_needed,
               )
                loader = get_microbatch_dataloader(len(packing_context.packed_trajs), args.micro_batch_size)
        else:
            with nvtx_range("rl/align-inference-logprobs", time=True):
                if inference_logprobs is not None:
                    inference_logprobs = align_unpacked_inference_logprobs(
                        inference_logprobs=inference_logprobs,
                        old_logprobs_for_data=old_logprobs,
                        generation_masks=generation_masks,
                        group_stats=group_stats,
                    )
                    # We run the above to fill in the inference/train side mismatch stats.
                    # We do the above for logging purposes.
                    # Nullify logprobs if not used in IS correction,
                    if not args.rl_inference_logprobs_is_correction:
                        inference_logprobs = None
            with nvtx_range("rl/create-dataloader", time=True):
                # Because of multiturn, our batch sizes for non-sequence packed trajectories are not fixed anymore.
                # As in sequence packing above, we need to reconfigure it too.
                runtime_state.packing_context = None

                reconfigure_num_microbatches_calculator(
                    rank=torch.distributed.get_rank() if torch.distributed.is_initialized() else 0,
                    global_batch_size=math.ceil(samples_ratio_per_step*total_turns_sampled), 
                    rampup_batch_size=args.rampup_batch_size, 
                    micro_batch_size=args.micro_batch_size, 
                    decrease_batch_size_if_needed=args.decrease_batch_size_if_needed,
                    data_parallel_size=mpu.get_data_parallel_world_size(),
                )

                dataset_tensors = [
                    compute_trajs,
                    advantages,
                    old_logprobs,
                    original_loss_mask,
                    original_position_ids,
                    ref_logprobs,
                ]
                if is_correction and inference_logprobs is not None:
                    dataset_tensors.append(inference_logprobs)
                else:
                    dataset_tensors.append(torch.zeros_like(old_logprobs))
                data = TensorDataset(*dataset_tensors)
                loader = DataLoader(data, batch_size=args.micro_batch_size)

        with nvtx_range("rl/log-wandb-tb", time=True):
            maybe_log_training_metrics(
                group_stats=group_stats,
                current_iteration=args.curr_iteration,
                tokenizer=tokenizer,
                example_group=example_group,
                wandb_writer=wandb_writer,
                tb_writer=tb_writer,
            )

    return RerunDataIterator(itertools.cycle(loader)), group_stats, example_groups


def get_grpo_data_iterator(
    model: LanguageModule,
    inference_model: LanguageModule | None,
    optimizer: MegatronOptimizer,
    iteration: int,
    ref_state_dict: Dict[str, torch.Tensor],
    grpo_iterations: int,
    grpo_prompts_per_step: int,
    grpo_group_size: int,
    global_batch_size: int,
    sequence_packing: bool,
    is_correction: bool,
    buffered_rollouts: RerunDataIterator | None = None,
) -> RerunDataIterator:
    """
    Get the data iterator for GRPO training.

    Depending on the sampling parameters either performs data collections or returns
    the buffered_rollouts as is.

    Args:
        model: The language model
        optimizer: The Megatron optimizer
        iteration: Current training iteration
        ref_state_dict: Reference model state dict for GRPO
        grpo_iterations: How many steps we reuse the sampled data for.
        grpo_prompts_per_step: How many prompts we sample per data collection.
        grpo_group_size: How many samples we do per prompt.
        global_batch_size: Global batch size.
        sequence_packing: Use sequence packing if True.
        is_correction: Use IS correction if True.
        buffered_rollouts: Previously collected rollouts (if any)

    Returns:
        RerunDataIterator for the current training step
    """
    runtime_state = get_rl_runtime_state()
    tokenizer = get_tokenizer()

    # We collect new rollouts when we've gone over the collected data 'grpo_iterations' times.
    global_batches_per_collection = (grpo_prompts_per_step * grpo_group_size) // global_batch_size
    if (
        buffered_rollouts is None or
        iteration == runtime_state.last_collection_iteration +
        (grpo_iterations * global_batches_per_collection)
    ):

        rollouts = get_environment_rollouts(
            model, inference_model, optimizer, grpo_prompts_per_step, grpo_group_size
        )
        buffered_rollouts, group_stats, example_groups = prepare_data_for_update(
            model=model,
            ref_state_dict=ref_state_dict,
            rollouts=rollouts,
            tokenizer=tokenizer,
            sequence_packing=sequence_packing,
            is_correction=is_correction,
        )
        runtime_state.group_stats = group_stats
        runtime_state.example_groups = example_groups
        runtime_state.reset_iteration_counters(iteration)

    maybe_log_training_metrics(
        group_stats=runtime_state.group_stats,
        current_iteration=iteration,
        tokenizer=tokenizer,
        example_groups=runtime_state.example_groups,
    )

    return buffered_rollouts


def evaluate_and_print_results_rl(
    data_iterator: Iterator[TensorDataset],
    model: list[LanguageModule],
    optimizer: MegatronOptimizer,
    iteration: int,
    write_to_tensorboard: bool = True,
    training_model: Optional[list[LanguageModule]] = None,
):
    """Helper function to evaluate and dump results on screen.

    Args:
        data_iterator: Iterator over batches of evaluation dataset.
        model: Model to evaluate with (may be separate inference model).
        iteration: Current training iteration.
        write_to_tensorboard: Dumpt stuff to tensorboard or not.
        training_model: Training model (if separate from inference model). Used to offload
            grad buffers and restore to train mode. If None, uses model parameter.
    """
    args = get_args()

    # TODO(vitalyk): I do not track eval loss as in training. We probably should.
    # megatron-lm uses forward_step_func to do the above.

    # Use context manager to temporarily disable sequence parallelism for evaluation

    with torch.no_grad():
        with megatron_rl_inference_mode(
            model,
            optimizer,
            args.cuda_graph_impl,
            args.rl_offload_optimizer_during_inference,
            training_model,
        ) as inference_interface:

            loop = get_asyncio_loop()

            rank = torch.distributed.get_rank()
            if rank == 0:
                logger.info("Collecting evaluation results...")
                agent = get_agent(args)
                request = EvaluationRequest(
                    inference_interface=inference_interface,
                    num_prompts=args.rl_prompts_per_eval,
                    validation=True,
                    rank_info=None,
                    generation_args={
                        'temperature': args.rl_default_temperature,
                        'max_tokens': args.seq_length,
                        'top_p': args.rl_default_top_p,
                        'top_k': args.rl_default_top_k,
                    },
                )
                with get_nvtx_range()("rl/run-evaluation", time=True):
                    evaluation_responses = loop.run_until_complete(agent.run_evaluation(request))
                    if not isinstance(evaluation_responses, list):
                        evaluation_responses = [evaluation_responses]
            else:
                evaluation_responses = None

        dp_eval_results: list[None | list[EvaluationResponse]] = [
            None for _ in range(args.world_size)
        ]
        dist.gather_object(
            evaluation_responses,
            dp_eval_results if dist.get_rank() == (args.world_size - 1) else None,
            dst=args.world_size - 1,
        )

        if dist.get_rank() == args.world_size - 1:
            dp_eval_results = [x for x in dp_eval_results if x is not None]
            # TODO(rkirby): maybe factor this out into a function?
            eval_metrics = defaultdict(list)
            for responses in dp_eval_results:
                for response in responses:
                    if response is None:
                        continue
                    for k, v in response.metrics().items():
                        eval_metrics[f"{response.env_id}_eval_mean_{k}"].extend(v)
                    for result in response.results:
                        if isinstance(result, RewardEvaluationResult):
                            try:
                                lang_rl_log(
                                    f"Evaluation: [{response.env_id}] [{result.reward}] {result.prompt} {result.response}"
                                )
                            except Exception as e:
                                lang_rl_log(f"Error: {e}")
                                lang_rl_log(f"Result: {result}")
            logger.info(
                "Collected metrics:"
                + "".join([f"\n\t{k} count: {len(v)}" for k, v in eval_metrics.items()])
            )
            eval_metrics = {k: np.mean(v) for k, v in eval_metrics.items()}
            if write_to_tensorboard:
                tb_writer = get_tensorboard_writer()
                if tb_writer:
                    for k, v in eval_metrics.items():
                        tb_writer.add_scalar(k, v, iteration)
            wandb_writer = get_wandb_writer()
            if wandb_writer:
                wandb_writer.log(eval_metrics, step=iteration)
            logger.info(
                "Evaluation results:"
                + "".join([f"\n\t{k}: {v:0.4f}" for k, v in eval_metrics.items()])
            )
            if lang_rl_log_dir:
                with open(
                    lang_rl_log_dir
                    + f'/eval_rank{rank}_iteration{args.curr_iteration}_'
                    + f'{Path(args.langrl_env_config).stem}.json',
                    'w',
                ) as f:
                    json.dump([[r.model_dump() for r in group] for group in dp_eval_results], f)


def calculate_grpo_loss(
    current_logprobs: torch.Tensor,
    old_logprobs: torch.Tensor,
    ref_logprobs: torch.Tensor,
    advantages: torch.Tensor,
    clamp_eps_lower: float,
    clamp_eps_upper: float,
    kl_beta: float,
    entropy_weight: float,
    inference_logprobs: torch.Tensor | None = None,
    is_truncation_coef: float | None = None,
    seq_starts: list | None = None,
    seq_lengths: list | None = None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Get GRPO loss, the kl term of the loss and the pi/pi_{old} ratios.

    Args:
        current_logprobs: pi logprobs, [batch, seq] for unpacked or [1, bin_size] for packed.
        old_logprobs: pi_{old} logprobs, [batch, seq] for unpacked or [1, bin_size] for packed.
        ref_logprobs: pi_{ref} logprobs, [batch, seq] for unpacked or [1, bin_size] for packed.
        advantages: advantages tensor, [batch,] for unpacked or [num_sequences_in_bin,] for packed.
        clamp_eps_lower: eps to clamp ratios from below.
        clamp_eps_upper: eps to clamp ratios from above, if vanilla GRPO, this should be equal to clamp_eps_lower.
        kl_beta: weight for the KL penalty term measuring the distance between pi and pi_{ref}.
        entropy_weight: weight for the entropy term.
        inference_logprobs: pi_{old} logprobs calculated by the inference engine.
            If not None, importance sampling correction will be applied.
        is_truncation_coef: importance sampling truncation coefficient. Will be applied if it is not None and inference_logprobs are present.
        seq_starts: (optional) For packed sequences: start positions of each sequence in the bin.
        seq_lengths: (optional) For packed sequences: original lengths of each sequence.

    Returns:
        total per-token GRPO loss [batch, seq] or [1, bin_size],
        kl_term of the loss [batch, seq] or [1, bin_size],
        pi/pi_{old} ratios [batch, seq] or [1, bin_size],
        entropy_term of the loss [batch, seq] or [1, bin_size],
        truncated_from_above [batch, seq] or [1, bin_size] (whether we clamped the ratios or not),
        truncated_from_below [batch, seq] or [1, bin_size] (whether we clamped the ratios or not).
    """
    # Ensure shapes match before computation
    if current_logprobs.shape != old_logprobs.shape:
        log_single_rank(
            logger,
            logging.WARNING,
            f"WARNING: Shape mismatch - current_logprobs: {current_logprobs.shape}, old_logprobs: {old_logprobs.shape}",
        )

    ratios = (current_logprobs - old_logprobs).exp()
    clamped_ratios = ratios.clamp(1 - clamp_eps_lower, 1 + clamp_eps_upper)
    truncated_from_above = torch.gt(ratios, 1 + clamp_eps_upper)
    truncated_from_below = torch.lt(ratios, 1 - clamp_eps_lower)

    # Handle advantages based on whether this is packed or unpacked
    if seq_starts is not None and seq_lengths is not None:
        # Packed sequences: map each sequence's advantage to its tokens
        bin_size = current_logprobs.shape[1]
        packed_advantages = torch.zeros(
            (1, bin_size), device=current_logprobs.device, dtype=current_logprobs.dtype
        )

        for seq_idx, (start, seq_len) in enumerate(zip(seq_starts, seq_lengths)):
            # Logprobs are 1 token shorter than sequences
            end = min(start + seq_len - 1, bin_size)
            if end > start:
                packed_advantages[0, start:end] = advantages[seq_idx].item()

        advantages = packed_advantages
    else:
        # Unpacked sequences: broadcast single advantage per sequence
        # Reshape to [batch, 1] to match logprobs shape [batch, seq]
        advantages = advantages.view(-1, 1)

    ref_diff = ref_logprobs - current_logprobs
    kl_term = ref_diff.exp() - ref_diff - 1
    entropy_term = -current_logprobs.exp() * current_logprobs

    is_weights = torch.tensor(1.0, dtype=old_logprobs.dtype).to(old_logprobs.device)
    if inference_logprobs is not None:
        is_weights = (old_logprobs - inference_logprobs).exp()
        if is_truncation_coef is not None:
            is_weights = torch.min(
                is_weights,
                torch.tensor(is_truncation_coef, dtype=old_logprobs.dtype).to(old_logprobs.device),
            )

    loss = (
        -is_weights * torch.min(ratios * advantages, clamped_ratios * advantages)
        + kl_beta * kl_term
        - entropy_weight * entropy_term
    )

    return loss, kl_term, ratios, entropy_term, truncated_from_above, truncated_from_below


@contextmanager
def megatron_rl_inference_mode(
    model: list[LanguageModule],
    optimizer: MegatronOptimizer,
    cuda_graph_impl: str,
    offload_optimizer_during_inference: bool,
    training_model: Optional[list[LanguageModule]] = None,
    increment_staleness_on_suspend: bool = False,
):
    """Manage the model inference context when collecting rollouts.

    Args:
        model: model to prepare for inference (may be separate inference model).
        optimizer: optimizer used to train the model.
        cuda_graph_impl: which cuda graph implementation to use.
        offload_optimizer_during_inference: move optimizer to cpu during inference or not.
        training_model: training model (if separate from inference model). Used to offload
            grad buffers and restore to train mode. If None, uses model parameter.

    Yields:
        None: this context manager does not return a value.

    """
    args = get_args()
    loop = get_asyncio_loop()
    nvtx_range = get_nvtx_range()

    logger.debug(f"[{dist.get_rank()}] Entering inference mode")

    # Change cudagraph scope for inference (empty list = full-layer capture)
    model[0].config.cuda_graph_scope = []
    model[0].config.cuda_graph_impl = "local"

    # If we get a lower precision wrapper, we go one object deeper.
    lang_module = model[0].module.module if hasattr(model[0].module, "module") else model[0].module

    # Switch MoE layers to full CUDA graph capture for inference
    if args.rl_training_cuda_graphs and args.num_experts is not None:
        transition_moe_cudagraphs(lang_module, 'full')

    lang_module.eval()
    # If this is a separate RL inference model with offloading enabled, ensure weights are on GPU
    # before any CUDA-graph capture/replay or inference. This is a no-op if already on GPU.
    model_core = unwrap_model(model[0])
    with nvtx_range("rl/prefetch-weights-to-gpu", time=True):
        _maybe_prefetch_separate_inference_model_weights(model_core, to_cpu=False)

    rotary_module = getattr(lang_module, "rotary_pos_emb", None)
    # Vanilla RotaryEmbedding module has lru_cache decorator which breaks RL training
    # as it tries to reuse frequences tensors cached in inference mode.
    has_lru_cache = rotary_module is not None and hasattr(rotary_module.forward, "cache_parameters")
    if has_lru_cache:
        rotary_module.forward.cache_clear()

    with torch.no_grad():

        if offload_optimizer_during_inference:
            with nvtx_range("rl/offload-optimizer-before-inference", time=True):
                if not args.rl_training_cuda_graphs:
                    with nvtx_range("rl/offload/grad-buffers", time=True):
                        model_for_grad_offload = training_model if training_model is not None else model
                        model_for_grad_offload[0].offload_grad_buffers()
                else:
                    logger.warning(
                        "Gradient buffers will not be offloaded when training cudagraphs are enabled!")
                with nvtx_range("rl/offload/optimizer-state", time=True):
                    optimizer.offload_to_cpu()

        if cuda_graph_impl != "none" and not args.rl_training_cuda_graphs:
            toggle_cuda_graphs(lang_module, cuda_graph_impl)

        inference_interface = get_inference_interface(args, loop, model)
        loop.run_until_complete(inference_interface.resume())

        logger.debug(f"[{dist.get_rank()}] Entered inference mode")
        yield inference_interface

        with nvtx_range("rl/suspend-engine", time=True):
            loop.run_until_complete(inference_interface.suspend())
            if increment_staleness_on_suspend:
                inference_interface.increment_staleness()

        if cuda_graph_impl != "none" and not args.rl_training_cuda_graphs:
            toggle_cuda_graphs(lang_module, 'none')

        # Reset drop_and_pad leaked from inference decode
        set_decode_expert_padding(unwrap_model(model[0]), set_to=False)

        # Restore partial capture cudagraph scope for training if this is MoE
        if args.num_experts is not None:
            model[0].config.cuda_graph_scope = [
                CudaGraphScope.mamba,
                CudaGraphScope.attn,
                CudaGraphScope.moe_router,
                CudaGraphScope.moe_preprocess,
            ]

        # Switch MoE layers to partial CUDA graph capture for training
        if args.rl_training_cuda_graphs and args.num_experts is not None:
            transition_moe_cudagraphs(lang_module, 'partial')

        # If this is a separate RL inference model, prefetch weights back to CPU so they
        # don't consume GPU memory during training.
        with nvtx_range("prefetch-inference-model-weights-to-cpu"):
            _maybe_prefetch_separate_inference_model_weights(model_core, to_cpu=True)

        if offload_optimizer_during_inference:
            with nvtx_range("rl/onload-optimizer-after-inference", time=True):
                with nvtx_range("rl/onload/grad-buffers", time=True):
                    model_for_grad_offload = training_model if training_model is not None else model
                    model_for_grad_offload[0].restore_grad_buffers()
                with nvtx_range("rl/onload/optimizer-state", time=True):
                    optimizer.restore_from_cpu()

        # Set training model back to train mode (not inference model if they're separate)
        training_lang_module = unwrap_model(training_model[0]) if training_model is not None else lang_module
        training_lang_module.train()

        if has_lru_cache:
            rotary_module.forward.cache_clear()

        logger.debug(f"[{dist.get_rank()}] Exiting inference mode")


def rl_inference_interface_shutdown():
    global _INFERENCE_INTERFACE
    global _ROLLOUT_GENERATOR

    if _ROLLOUT_GENERATOR is not None:
        loop = get_asyncio_loop()
        loop.run_until_complete(_ROLLOUT_GENERATOR.aclose())
        _ROLLOUT_GENERATOR = None

    if _INFERENCE_INTERFACE is not None:
        loop = get_asyncio_loop()
        loop.run_until_complete(_INFERENCE_INTERFACE.kill())
        _INFERENCE_INTERFACE = None
    else:
        logger.warning("No inference interface to shutdown. This should not happen.")

    # TODO(rkirby): This is a hack to hard exit. There is a bug that is preventing us from using sys.exit(0).
    # It seem the Flask server has non-daemon threads that are preventing the program from exiting.
    # We need to find a way to gracefully complete all in progress requests and shutdown the Flask server.
    import os
    os._exit(0)


def get_iteration_sequence_count(args):
    """Get the total number of sequences processed in this iteration across all ranks."""
    runtime_state = get_rl_runtime_state()
    sequences_tensor = torch.tensor(
        runtime_state.sequences_this_iteration_on_rank, device='cuda', dtype=torch.long
    )
    if torch.distributed.is_initialized():
        torch.distributed.all_reduce(sequences_tensor, group=mpu.get_data_parallel_group())
    return int(sequences_tensor.item())
    
def _pad_nonnull_with_zeros(data: list[Optional[torch.Tensor]], max_len: int) -> torch.Tensor:
    """Pad each element of a list of tensors to the length required.
    Args:
        data: List of tensors to pad.
        max_len: Maximum length to pad to. Must be higher or equal than the max len of the data tensors.
    Returns:
        A padded tensor which is a stacked list of padded input tensors.

    """
    if all([el is None for el in data]):
        raise ValueError("At least one element of the data list should be not None.")
    padded_data = []
    for chunk in data:
        if chunk is not None:
            padding_size = max_len - len(chunk)
            if padding_size > 0:
                # Pad with zeros (these positions will be masked anyway)
                padded = torch.nn.functional.pad(chunk, (0, padding_size), value=0.0)
                padded_data.append(padded)
            elif padding_size == 0:
                padded_data.append(chunk)
            else:
                raise ValueError("One of the input tensors has larger length than padding max len.")
        else:
            # Create zero tensor for None logprobs
            padded_data.append(torch.zeros(max_len))
    return torch.stack(padded_data)