grpo.py

# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import gc
import json
import os
import time
import warnings
from concurrent.futures import Future, ThreadPoolExecutor
from contextlib import nullcontext
from dataclasses import dataclass, fields
from typing import Any, Callable, Optional, TypeVar, cast

import numpy as np
import ray
import torch
from pydantic import BaseModel, Field, model_validator
from torchdata.stateful_dataloader import StatefulDataLoader
from transformers import AutoProcessor
from transformers.tokenization_utils_base import PreTrainedTokenizerBase

from nemo_rl.algorithms import opd as opd_module
from nemo_rl.algorithms.advantage_estimator import (
    AdvEstimatorConfig,
    GDPOAdvantageEstimator,
    GRPOAdvantageEstimator,
    OPDAdvantageEstimator,
    ReinforcePlusPlusAdvantageEstimator,
)
from nemo_rl.algorithms.logits_sampling_utils import (
    TrainingSamplingParams,
    need_top_k_or_top_p_filtering,
)
from nemo_rl.algorithms.loss import (
    ClippedPGLossConfig,
    ClippedPGLossDataDict,
    ClippedPGLossFn,
)
from nemo_rl.algorithms.loss.interfaces import LossFunction
from nemo_rl.algorithms.metric_utils import (
    GRAD_NORM_KEY,
    LOSS_KEY,
    MEAN_GEN_TOKENS_PER_SAMPLE_KEY,
    REWARD_KEY,
    SETUP_TIMING_PREFIX,
    SetupTimingMetrics,
    print_setup_timing_summary,
)
from nemo_rl.algorithms.opd import OnPolicyDistillationConfig
from nemo_rl.algorithms.reward_functions import (
    RewardShapingConfig,
    apply_reward_shaping,
)
from nemo_rl.algorithms.utils import (
    WALL_CLOCK_EFFICIENCY_CATEGORIES,
    calculate_baseline_and_std_per_prompt,
    calculate_trivial_reward_distributions,
    compute_seq_logprob_errors,
    get_gdpo_reward_component_keys,
    log_generation_metrics,
    print_efficiency_summary,
    print_performance_metrics,
    set_seed,
)
from nemo_rl.data import DataConfig
from nemo_rl.data.collate_fn import rl_collate_fn
from nemo_rl.data.dataloader import CyclingDataLoader, MultipleDataloaderWrapper
from nemo_rl.data.datasets import AllTaskProcessedDataset
from nemo_rl.data.interfaces import DatumSpec, LLMMessageLogType, VLMMessageLogType
from nemo_rl.data.llm_message_utils import (
    batched_message_log_to_flat_message,
    get_keys_from_message_log,
)
from nemo_rl.data.multimodal_utils import PACKED_MULTIMODAL_FIELDS, PackedTensor
from nemo_rl.data.utils import extract_necessary_env_names, load_dataloader_state
from nemo_rl.data_plane.interfaces import DataPlaneConfig
from nemo_rl.distributed.batched_data_dict import BatchedDataDict
from nemo_rl.distributed.virtual_cluster import (
    TOPO_RANK_UNKNOWN,
    ClusterConfig,
    RayVirtualCluster,
    get_ray_cluster_topology,
    prepare_segment_topology,
)
from nemo_rl.environments.interfaces import EnvironmentInterface
from nemo_rl.environments.nemo_gym import (
    NemoGymShardSet,
    build_nemo_gym_actors,
    should_use_nemo_gym,
    validate_dataset_agent_coverage,
)
from nemo_rl.environments.utils import shutdown_environments
from nemo_rl.experience.interfaces import (
    FRONTIER_ORDINAL_KEY,
    NEMO_GYM_TASK_INDEX_KEY,
    NEXT_NEMO_GYM_TASK_INDEX_KEY,
    PENDING_PROMPTS_KEY,
    RESUME_BASE_ORDINAL_KEY,
    RETAINED_TASK_INDICES_KEY,
    TRAINED_TASK_INDICES_KEY,
)
from nemo_rl.experience.metric_utils import is_histogram_metric
from nemo_rl.experience.rollouts import (
    EffortLevelsConfig,
    attach_initial_nemo_gym_image_payloads,
    backfill_missing_routed_experts,
    get_effort_config,
    get_nemo_gym_thinking_tags,
    run_async_multi_turn_rollout,
    run_multi_turn_rollout,
    run_nemo_gym_rollout_sync,
    should_mask_flagged_samples,
)
from nemo_rl.models.generation.dynamo import DynamoConfig, DynamoGeneration
from nemo_rl.models.generation.interfaces import (
    GenerationConfig,
    GenerationInterface,
    GenerationSamplingParams,
    should_use_async_rollouts,
)
from nemo_rl.models.generation.megatron import MegatronGeneration
from nemo_rl.models.generation.sglang.config import SGLangConfig
from nemo_rl.models.generation.sglang.sglang_generation import SGLangGeneration
from nemo_rl.models.generation.trtllm import TrtllmConfig, TrtllmGeneration
from nemo_rl.models.generation.trtllm.config import (
    replica_width,
    resolve_disagg_layout,
)
from nemo_rl.models.generation.vllm import VllmConfig, VllmGeneration
from nemo_rl.models.generation.vllm.config import (
    REFITTABLE_FP8_KV_CACHE_DTYPES,
    VLLM_SPARSE_REFIT_TRANSPORTS,
    normalize_nvfp4_pertoken_policy_config,
    normalize_vllm_refit_config,
)
from nemo_rl.models.megatron.router_replay import (
    configure_vllm_for_router_replay,
    router_replay_enabled,
)
from nemo_rl.models.megatron.zero_train_gen_mismatch import validate_zero_train_gen_kl
from nemo_rl.models.policy import PolicyConfig
from nemo_rl.models.policy.draft_config import coerce_draft_config
from nemo_rl.models.policy.interfaces import ColocatablePolicyInterface
from nemo_rl.models.policy.lm_policy import Policy
from nemo_rl.telemetry.config import TelemetryConfig
from nemo_rl.telemetry.instrumentation import (
    RL_ITERATION_ATTR,
    current_trace_carrier,
    efficiency_span,
    evaluate_span,
    iteration_scope,
    managed_span,
    umbrella_span,
    umbrella_trace_fn,
)
from nemo_rl.telemetry.setup import get_telemetry_handle
from nemo_rl.telemetry.span_groups import RLSpanGroup
from nemo_rl.telemetry.vocabulary import INIT_TOTAL_CATEGORY, WORKER_GROUP_ENV
from nemo_rl.utils.checkpoint import (
    CheckpointingConfig,
    CheckpointManager,
)
from nemo_rl.utils.length_penalty import LengthPenaltyConfig
from nemo_rl.utils.logger import (
    Logger,
    LoggerConfig,
    print_message_log_samples,
    should_log_nemo_gym_full_result_tables,
)
from nemo_rl.utils.memory_tracker import MemoryTracker
from nemo_rl.utils.multimodal_payload_metrics import (
    collect_multimodal_payload_metrics,
    drain_multimodal_payload_metrics,
    merge_multimodal_payload_metrics,
    print_multimodal_payload_metrics,
)
from nemo_rl.utils.nsys import maybe_gpu_profile_step
from nemo_rl.utils.timer import TimeoutChecker, Timer
from nemo_rl.utils.venvs import make_actor_runtime_env
from nemo_rl.weight_sync.checkpoint_engine_config import (
    checkpoint_engine_refit_config,
)
from nemo_rl.weight_sync.factory import (
    create_weight_synchronizer,
    validate_offload_policy_before_refit,
)
from nemo_rl.weight_sync.nccl_reshard_utils import check_nccl_reshard_refit_support

# ===============================================================================
# Configuration
# ===============================================================================
TokenizerType = TypeVar("TokenizerType", bound=PreTrainedTokenizerBase)


def _maybe_restore_async_replay_buffer_checkpoint(
    replay_buffer: Any,
    checkpoint_path: str,
    *,
    load_replay_buffer: bool | None,
    num_prompts_per_step: int,
    current_training_step: int,
    max_age_steps: int,
) -> dict[str, Any] | None:
    """Restore async replay state unless the config explicitly opts out.

    With ``checkpointing.load_replay_buffer=false`` the buffer starts empty
    and, on a frontier-aligned checkpoint, the whole buffered window is
    regenerated fresh from the rewound dataloader — the empty-retained-set
    case of the same resume path. This trades resume compute for an unbiased
    step composition: retained groups are the ones whose longest rollout
    happened to finish before the save, so reusing them skews the next steps
    toward short-rollout prompts.

    Returns:
        The restore metadata from ``load_from_path``, or ``None`` when the
        restore was skipped or no checkpoint file exists.
    """
    if load_replay_buffer is False:
        print(
            "📦 Skipping replay buffer restore (checkpointing.load_replay_buffer=false)"
        )
        return None

    replay_buffer_path = os.path.join(checkpoint_path, "replay_buffer.pt")
    if not os.path.exists(replay_buffer_path):
        print(
            f"⚠️ No replay buffer checkpoint found at {replay_buffer_path}. "
            "Starting with an empty replay buffer."
        )
        return None

    print(f"📦 Restoring replay buffer from checkpoint: {replay_buffer_path}")
    restore_metadata = ray.get(
        replay_buffer.load_from_path.remote(
            replay_buffer_path,
            num_prompts_per_step=num_prompts_per_step,
            current_training_step=current_training_step,
            max_age_steps=max_age_steps,
        )
    )
    print("✅ Replay buffer restored from checkpoint")
    return restore_metadata


def _save_async_replay_buffer_checkpoint(
    replay_buffer: Any,
    checkpoint_path: str,
) -> int:
    """Checkpoint replay state inside its actor."""
    print("📦 Saving replay buffer state...")
    num_buffered_trajectories = ray.get(
        replay_buffer.save_to_path.remote(
            os.path.join(checkpoint_path, "replay_buffer.pt")
        )
    )
    print(f"✅ Saved replay buffer with {num_buffered_trajectories} trajectories")
    return num_buffered_trajectories


class RewardScalingConfig(BaseModel, extra="allow"):
    """Configure linear reward scaling with clamping.

    When `enabled` is True, each reward is clamped to the source interval
    [source_min, source_max] and linearly mapped to the target interval
    [target_min, target_max]. Refer to the scale_rewards function for the implementation.
    """

    enabled: bool = False
    source_min: float = 0.0
    source_max: float = 1.0
    target_min: float = 0.0
    target_max: float = 1.0


class AsyncGRPOConfig(BaseModel, extra="allow"):
    enabled: bool = False
    # Maximum trajectory age in training steps for samples drawn from the
    # async replay buffer. Trajectories older than this are excluded during
    # sampling; buffer sizing also scales with this value.
    max_trajectory_age_steps: int = 1
    # Generation-worker failures tolerated before the AsyncTrajectoryCollector
    # aborts the run. A successful batch worker resets the count.
    # 0 makes the very first worker exception fatal.
    max_generation_failures: int = 0
    # Does the weight synchronization as soon as the training is done
    # without waiting for the pending generations to finish.
    in_flight_weight_updates: bool = False
    # Recomputes the KV cache after weight updates.
    recompute_kv_cache_after_weight_updates: bool = False


class RewardPenaltyTokenIdsConfig(BaseModel, extra="allow"):
    """Optional token IDs for reward penalties."""

    unwanted: list[int] | None = None
    think_open: int | None = None
    think_close: int | None = None


class RewardPenaltyConfig(BaseModel, extra="allow"):
    """Reward-zeroing penalties applied to NeMo-Gym rollout results."""

    penalize_duplicated_reasoning: bool = False
    penalize_empty_final_answer: bool = False
    penalize_unwanted_tokens: bool = False
    penalize_malformed_think_tag: bool = False
    # Optional token IDs. token_ids.unwanted is required when
    # penalize_unwanted_tokens is true;
    # think-tag IDs are inferred from configured tag strings when possible.
    token_ids: Optional[RewardPenaltyTokenIdsConfig] = None

    @model_validator(mode="after")
    def _require_unwanted_token_ids_when_penalized(self) -> "RewardPenaltyConfig":
        if self.penalize_unwanted_tokens and (
            self.token_ids is None or not self.token_ids.unwanted
        ):
            raise ValueError(
                "reward_penalties.token_ids.unwanted must be set when "
                "reward_penalties.penalize_unwanted_tokens is true"
            )
        return self


_REWARD_PENALTY_FLAGS = (
    "penalize_duplicated_reasoning",
    "penalize_empty_final_answer",
    "penalize_unwanted_tokens",
    "penalize_malformed_think_tag",
)


class GRPOConfig(BaseModel, extra="allow"):
    num_prompts_per_step: int = 32
    num_generations_per_prompt: int = 16
    max_num_epochs: int = 1
    max_num_steps: int = 1000000
    max_rollout_turns: int = 1
    normalize_rewards: bool = True
    # Clipping bounds for normalized advantages to prevent extreme values
    # When set, advantages are clipped to [advantage_clip_low, advantage_clip_high] after normalization
    # Default: null (no clipping)
    advantage_clip_low: float | None = None
    advantage_clip_high: float | None = None
    use_leave_one_out_baseline: bool = True
    val_period: int = 10
    # First training step eligible for periodic validation; -1 disables the delay.
    val_start_at: int = -1
    val_batch_size: int | None = 256  # None for NeMo-Gym compatibility
    val_at_start: bool = False
    # Whether to run validation on the last training step. Setting this to True ensures the
    # final checkpoint has validation metrics, which is required for get_best_checkpoint_path().
    val_at_end: bool = False
    # Counts PROMPTS, not rollouts: with val_num_generations_per_prompt = k,
    # total validation rollouts = max_val_samples * k.
    max_val_samples: int | None = 256  # None for NeMo-Gym compatibility
    # Number of independent validation rollouts generated for each prompt;
    # k > 1 additionally reports pass@k over each prompt's k rollouts as the
    # pass_k metric.
    val_num_generations_per_prompt: int = 1
    # Early stop: end training once this validation metric (e.g. accuracy,
    # always reported, or pass_k with grouped validation) reaches
    # stop_at_validation_threshold; null disables early stopping.
    stop_at_validation_metric: str | None = None
    # Threshold for the early stop; required when stop_at_validation_metric
    # is set.
    stop_at_validation_threshold: float | None = None
    skip_reference_policy_logprobs_calculation: bool = False
    seed: int = 42
    # Legacy async config block; SC reads its async knobs from `async_rl` instead.
    async_grpo: AsyncGRPOConfig | None = Field(default_factory=AsyncGRPOConfig)
    overlong_filtering: bool = False
    # whether to enable dynamic sampling, i.e.
    # whether to discard prompts whose rewards have zero standard deviation
    use_dynamic_sampling: bool = False
    # When using dynamic sampling, the maximum number of batches to generate
    # before throwing an error
    dynamic_sampling_max_gen_batches: int = 10
    # When using dynamic sampling, generation prompt batch size will equal
    # num_prompts_per_step * batch_multiplier
    batch_multiplier: float = 1.0
    reward_shaping: RewardShapingConfig = Field(default_factory=RewardShapingConfig)
    reward_scaling: RewardScalingConfig = Field(default_factory=RewardScalingConfig)
    # By default advantages are calculated on CPU. Setting this flag to true leverages GPU for their computation.
    calculate_advantages_on_gpu: bool = False
    # Sequence-level logprob error masking for training stability. If set, mask sequences with mult_prob_error exceeding this threshold (same scale as token_mult_prob_error metric, e.g., 1.5)
    # Note that this is slightly different than Masked Importance Sampling (MIS) because this uses the absolute value of the difference between the training and generation logprobs, whereas MIS just uses the difference between the training and generation logprobs.
    seq_logprob_error_threshold: float | None = None
    # Advantage value to assign to invalid tool call tokens. When set (e.g. -5.0), overwrites the
    # computed advantage for those tokens to penalize them; absent/None disables the penalty.
    invalid_tool_call_advantage: float | None = None
    # Advantage value to assign to tokens with malformed <think>/</think> tags. When set (e.g. -5.0),
    # overwrites the computed advantage for those tokens; absent/None disables the penalty.
    malformed_thinking_advantage: float | None = None
    # Advantage estimator configuration (grpo or reinforce_plus_plus)
    adv_estimator: AdvEstimatorConfig = Field(default_factory=AdvEstimatorConfig)
    # Share and compact immutable image/video/audio payload segments across
    # logical GRPO rows. Prompt identity is never used as proof of equality.
    deduplicate_multimodal_data: bool = False
    # Emit exact-boundary and logical-vs-physical payload metrics.
    debug_payload_metrics: bool = False
    # Optional per-prompt-group length penalties/bonuses applied to NeMo-Gym
    # training rollout rewards (binary single-reward envs only); see
    # docs/guides/length-penalty.md. None/absent disables the feature.
    length_penalty: Optional[LengthPenaltyConfig] = None


@dataclass
class GRPOSaveState:
    consumed_samples: int
    current_step: int
    current_epoch: int
    total_steps: int
    total_valid_tokens: int  # Track total number of non-padding tokens during training
    val_reward: float  # May be removed when no validation metrics are available
    # SC may advance the policy version independently from the optimizer-step
    # counter. None preserves compatibility with checkpoints predating it.
    trainer_version: Optional[int] = None
    # SingleController only: name of the sampler that wrote the replay buffer,
    # used to gate the SC buffer restore. None on checkpoints from the other
    # algorithms and from SC runs that predate this field.
    sampler_name: Optional[str] = None
    # SingleController only: exact last admitted dispatch batch. None preserves
    # compatibility with checkpoints that only recorded the trainer version.
    sampler_dispatch_index: Optional[int] = None


def _initial_grpo_save_state() -> GRPOSaveState:
    return GRPOSaveState(
        consumed_samples=0,
        current_step=0,
        current_epoch=0,
        total_steps=0,
        total_valid_tokens=0,
        val_reward=-99999999.0,
        trainer_version=None,
        sampler_name=None,
        sampler_dispatch_index=None,
    )


def _get_grpo_save_state(
    loaded_state: Optional[dict[str, Any]],
) -> GRPOSaveState:
    if loaded_state is None:
        return _initial_grpo_save_state()

    # Start from current defaults so partial/legacy checkpoints remain loadable.
    known_fields = {field.name for field in fields(GRPOSaveState)}
    state_values = vars(_initial_grpo_save_state()).copy()
    state_values.update(
        {key: value for key, value in loaded_state.items() if key in known_fields}
    )
    return GRPOSaveState(**state_values)


class MasterConfig(BaseModel, extra="allow"):
    policy: PolicyConfig
    loss_fn: ClippedPGLossConfig
    env: dict[str, Any]
    data: DataConfig
    grpo: GRPOConfig
    logger: LoggerConfig
    cluster: ClusterConfig
    checkpointing: CheckpointingConfig
    reward_penalties: RewardPenaltyConfig = Field(default_factory=RewardPenaltyConfig)
    data_plane: Optional[DataPlaneConfig] = None
    on_policy_distillation: Optional[OnPolicyDistillationConfig] = None
    telemetry: Optional[TelemetryConfig] = None


# ===============================================================================
# Setup & Initialization
# ===============================================================================


def _validate_seq_logprob_error_in_loss(master_config: MasterConfig) -> None:
    """Validate the single-forward threshold path before allocating workers."""
    if not master_config.loss_fn.seq_logprob_error_in_loss:
        return
    if master_config.grpo.seq_logprob_error_threshold is None:
        raise ValueError(
            "loss_fn.seq_logprob_error_in_loss requires seq_logprob_error_threshold"
        )
    loss = master_config.loss_fn
    if not loss.force_on_policy_ratio or not loss.token_level_loss:
        raise ValueError(
            "loss_fn.seq_logprob_error_in_loss requires force_on_policy_ratio=true "
            "and token_level_loss=true"
        )
    if master_config.grpo.adv_estimator.name != "grpo" or loss.use_kl_in_reward:
        raise ValueError(
            "loss_fn.seq_logprob_error_in_loss requires the grpo advantage estimator "
            "without use_kl_in_reward"
        )
    policy = master_config.policy
    if "megatron_cfg" not in policy or not policy["megatron_cfg"]["enabled"]:
        raise ValueError(
            "loss_fn.seq_logprob_error_in_loss requires the Megatron backend"
        )
    draft = coerce_draft_config(policy.get("draft"))
    if (
        policy["megatron_cfg"].get("mtp_num_layers")
        or (draft is not None and draft.enabled)
        or loss.positive_example_nll_weight != 0
        or opd_module.is_opd_enabled(master_config)
    ):
        raise ValueError(
            "loss_fn.seq_logprob_error_in_loss does not support MTP, draft, "
            "positive-example NLL, or distillation losses"
        )


def _validate_multimodal_dedup_capability(master_config: MasterConfig) -> None:
    """Reject configurations whose media transfer path is not qualified."""
    if not master_config.grpo.deduplicate_multimodal_data:
        return

    generation_config = master_config.policy["generation"]
    if generation_config.get("backend") != "vllm":
        raise NotImplementedError(
            "grpo.deduplicate_multimodal_data=true is currently qualified "
            "only with policy.generation.backend=vllm."
        )
    # The data plane accepts deduplicated payloads, so the wire format is not
    # the constraint -- but note what dedup buys there. ``to_wire`` emits one
    # row per *logical* row, so a shared segment is concatenated once per
    # generation: the saving is in driver RAM (the deepcopy memo), not in wire
    # or TQ-storage bytes, which stay O(G x images). The one gap is NeMo-Gym:
    # ``grpo_train_sync`` does not call
    # ``attach_initial_nemo_gym_image_payloads``, which supplies the initial
    # image tensors a Gym dataset omits from ``extra_env_info``. That helper is
    # itself gated on ``should_use_nemo_gym``, so non-Gym recipes never needed
    # it and are unaffected.
    if (master_config.data_plane or {}).get("enabled", False) and (
        should_use_nemo_gym(master_config)
    ):
        raise NotImplementedError(
            "grpo.deduplicate_multimodal_data=true with data_plane.enabled=true "
            "is not supported for NeMo-Gym runs: the TransferQueue trainer does "
            "not attach the initial Gym image payloads. Non-Gym recipes are "
            "supported."
        )


def _needs_hf_refit_handshake(
    generation_backend: str,
    nccl_reshard_refit_enabled: bool,
    colocated_inference: bool,
) -> bool:
    """Whether setup must run the HF-schema prepare_refit_info handshake."""
    if generation_backend == "megatron":
        return False
    return not (nccl_reshard_refit_enabled and not colocated_inference)


def _uses_managed_noncolocated_refit(
    generation_backend: str,
    nccl_reshard_refit_enabled: bool,
    offload_policy_before_refit: bool,
) -> bool:
    """Whether non-colocated setup must attach a weight synchronizer."""
    return (
        nccl_reshard_refit_enabled
        or generation_backend == "dynamo"
        or offload_policy_before_refit
    )


def _shutdown_completed_nemo_gym_startup(
    future: Future[tuple[NemoGymShardSet, float]] | None,
) -> None:
    if future is None:
        return
    try:
        shard_set, _ = future.result()
    except BaseException:
        return
    shard_set.shutdown()


def setup(
    master_config: MasterConfig,
    tokenizer: TokenizerType,
    dataset: AllTaskProcessedDataset | dict[str, AllTaskProcessedDataset],
    val_dataset: Optional[AllTaskProcessedDataset],
    processor: Optional[AutoProcessor] = None,
    policy_factory: Optional[Callable[..., ColocatablePolicyInterface]] = None,
) -> tuple[
    ColocatablePolicyInterface,
    Optional[GenerationInterface],
    Optional[EnvironmentInterface],
    tuple[RayVirtualCluster, RayVirtualCluster],
    StatefulDataLoader | MultipleDataloaderWrapper,
    Optional[StatefulDataLoader],
    ClippedPGLossFn,
    Logger,
    CheckpointManager,
    GRPOSaveState,
    MasterConfig,
    dict[str, Any],
    dict[str, str],
]:
    """Main entry point for running GRPO algorithm.

    Returns:
        A 13-tuple, in order:
            policy, policy_generation, nemo_gym (the NeMo-Gym env actor, or None
            when not enabled), cluster, dataloader, val_dataloader, loss_fn,
            logger, checkpointer, grpo_save_state, master_config,
            teacher_worker_groups, alias_to_group_alias.
    """
    # Start timing the entire setup process
    setup_start_time = time.perf_counter()

    # Extract individual configs for easier access
    policy_config = master_config.policy
    generation_config = policy_config["generation"]
    loss_config: ClippedPGLossConfig = master_config.loss_fn
    env_configs = master_config.env
    data_config = master_config.data
    grpo_config = master_config.grpo
    logger_config = master_config.logger
    cluster_config = master_config.cluster
    checkpointing_config = master_config.checkpointing

    checkpointing_pretrained = checkpointing_config.get("pretrained_checkpoint")
    if checkpointing_pretrained is not None:
        policy_config["pretrained_checkpoint"] = checkpointing_pretrained

    assert generation_config is not None, (
        "A generation config in the PolicyConfig is required for GRPO"
    )
    if generation_config["backend"] == "vllm":
        normalize_nvfp4_pertoken_policy_config(policy_config, entry_point="grpo")
        normalize_vllm_refit_config(cast(VllmConfig, generation_config))
    elif generation_config["backend"] == "dynamo":
        # Validate the complete managed-Dynamo boundary before allocating Ray
        # placement groups or starting any external services.
        if grpo_config.async_grpo.in_flight_weight_updates:
            raise ValueError(
                "grpo.async_grpo.in_flight_weight_updates must be false when "
                "policy.generation.backend='dynamo'; managed Dynamo drains "
                "rollouts before layerwise weight refit"
            )
        generation_config.setdefault("vllm_kwargs", {})["hf_overrides"] = (
            policy_config.get("hf_config_overrides") or {}
        )
        generation_config = DynamoConfig.model_validate(generation_config).model_dump()
        policy_config["generation"] = generation_config

    # generation.disaggregation is written in backend-agnostic terms, but only
    # the TRT-LLM backend implements it today: every other backend would
    # silently ignore the block and serve aggregated, so refuse to start.
    if (
        resolve_disagg_layout(generation_config).enabled
        and generation_config["backend"] != "trtllm"
    ):
        raise ValueError(
            "policy.generation.disaggregation.enabled is only supported for "
            f"backend='trtllm', got '{generation_config['backend']}'. Set it to "
            "false, or switch the backend."
        )
    # The disagg frontends are only reachable over HTTP, so without NeMo-Gym
    # the whole fleet would come up and fail at the first rollout instead.
    if resolve_disagg_layout(generation_config).enabled and not should_use_nemo_gym(
        master_config
    ):
        raise ValueError(
            "policy.generation.disaggregation.enabled requires NeMo-Gym "
            "rollouts (env.should_use_nemo_gym=true). Set it to false, or run "
            "through the NeMo-Gym entrypoint."
        )

    _validate_multimodal_dedup_capability(master_config)
    _validate_seq_logprob_error_in_loss(master_config)

    # Validation-only sampling is honored only on the NeMo-Gym vLLM rollout
    # path; everywhere else validation must sample exactly like training.
    val_sampling_overridden = (
        generation_config["val_temperature"] != generation_config["temperature"]
        or generation_config["val_top_p"] != generation_config["top_p"]
        or generation_config["val_top_k"] != generation_config["top_k"]
    )
    if val_sampling_overridden:
        assert generation_config["backend"] == "vllm" and should_use_nemo_gym(
            master_config
        ), (
            "generation.val_temperature/val_top_p/val_top_k differing from the "
            "train sampling params is only supported for vLLM NeMo-Gym rollouts."
        )
        # The NeMo-Gym path only stamps temperature/top_p onto requests and
        # rejects any top_k at rollout time, so a val_top_k override can never
        # be honored — fail here instead of at the first validation step.
        assert not generation_config["val_top_k"], (
            "generation.val_top_k is not supported: the NeMo-Gym rollout path "
            "only honors val_temperature/val_top_p. Leave val_top_k null."
        )
    assert grpo_config.val_num_generations_per_prompt >= 1, (
        "grpo.val_num_generations_per_prompt must be >= 1"
    )
    # pass_k is only reported when k > 1; catch the mismatch here instead of
    # at the first validation step.
    assert not (
        grpo_config.stop_at_validation_metric == "pass_k"
        and grpo_config.val_num_generations_per_prompt <= 1
    ), (
        "grpo.stop_at_validation_metric='pass_k' requires "
        "grpo.val_num_generations_per_prompt > 1"
    )

    # opd_full is implemented only on the SingleController runtime, which has its
    # own setup(). This entry point never reads on_policy_distillation.full, so
    # accepting it would silently train the sampled-token top-k objective while
    # the user believes they configured the full-vocabulary reverse KL.
    if opd_module.get_opd_full_config(master_config) is not None:
        raise ValueError(
            "on_policy_distillation.full is only implemented on the "
            "SingleController runtime (examples/run_grpo_single_controller.py). "
            "This entry point never reads it, so leaving it enabled here would "
            "silently train the sampled-token top-k objective instead of the "
            "full-vocabulary reverse KL."
        )

    # Set seed for all random number generators
    set_seed(grpo_config.seed)

    # ==========================
    #         Logger
    # ==========================
    logger = Logger(logger_config)
    logger.log_hyperparams(master_config.model_dump())

    # ==========================
    #      Checkpointing
    # ==========================
    checkpointer = CheckpointManager(checkpointing_config)
    last_checkpoint_path = checkpointer.get_latest_checkpoint_path()
    loaded_state = checkpointer.load_training_info(last_checkpoint_path)
    grpo_save_state = _get_grpo_save_state(loaded_state)

    # ==========================
    #           Data
    # ==========================
    # num_prompts_per_step and dataloader_batch_size will be different when using multiple dataloaders
    num_prompts_per_step = grpo_config.num_prompts_per_step
    if data_config["use_multiple_dataloader"]:
        dataloader_batch_size = data_config["num_prompts_per_dataloader"]
    else:
        dataloader_batch_size = num_prompts_per_step

    # Validate batch_multiplier
    batch_multiplier = grpo_config.batch_multiplier
    if grpo_config.use_dynamic_sampling:
        num_prompts_per_step = int(num_prompts_per_step * batch_multiplier)
        dataloader_batch_size = int(dataloader_batch_size * batch_multiplier)
    else:
        assert batch_multiplier == 1, (
            "batch_multiplier>1 can only be used if use_dynamic_sampling=True"
        )

    # Validate the early-stop pairing
    if grpo_config.stop_at_validation_metric is not None:
        assert grpo_config.stop_at_validation_threshold is not None, (
            "grpo.stop_at_validation_threshold must be set when "
            "grpo.stop_at_validation_metric is set"
        )

    # Validate number of prompts per step
    if data_config["use_multiple_dataloader"]:
        assert num_prompts_per_step % dataloader_batch_size == 0, (
            "Expected num_prompts_per_step to be a multiple of num_prompts_per_dataloader, "
            f"but got {num_prompts_per_step} and {dataloader_batch_size}. "
            "Please check the configuration of num_prompts_per_step and num_prompts_per_dataloader. "
            "If use_dynamic_sampling is enabled and batch_multiplier is used, please also check the configuration of batch_multiplier."
        )

    # Load train dataset
    def init_train_dataloader(dataset, suffix: str = ""):
        dataloader = StatefulDataLoader(
            dataset,
            batch_size=dataloader_batch_size,
            shuffle=data_config["shuffle"],
            collate_fn=rl_collate_fn,
            drop_last=True,
            num_workers=data_config["num_workers"],
        )
        if last_checkpoint_path is not None:
            load_dataloader_state(dataloader, last_checkpoint_path, data_config, suffix)
        return dataloader

    if data_config["use_multiple_dataloader"]:
        # Initialize dataloaders
        dataloaders = {}
        for task_name, task_dataset in dataset.items():
            dataloaders[task_name] = init_train_dataloader(
                task_dataset, f"_{task_name}"
            )
            print(
                f"  ✓ Training dataloader {task_name} loaded with {len(task_dataset)} samples",
                flush=True,
            )

        train_sample_count = sum(
            len(task_dataloader) for task_dataloader in dataloaders.values()
        )

        # Wrap dataloader
        dataloader = MultipleDataloaderWrapper(
            expected_num_prompts=num_prompts_per_step,
            data_config=data_config,
            dataloaders=dataloaders,
        )
    else:
        dataloader = init_train_dataloader(dataset)
        train_sample_count = len(dataloader)
        print(
            f"  ✓ Training dataloader loaded with {train_sample_count} samples",
            flush=True,
        )

    # Load validation dataset if provided
    val_dataloader: Optional[StatefulDataLoader] = None
    # If validation is enabled, load the validation dataloader
    if grpo_config.val_period > 0 or grpo_config.val_at_start or grpo_config.val_at_end:
        assert val_dataset is not None, (
            "Validation dataset is required if validation is enabled"
        )
        val_dataloader = StatefulDataLoader(
            val_dataset,
            batch_size=grpo_config.val_batch_size,
            shuffle=False,
            collate_fn=rl_collate_fn,
            num_workers=data_config["num_workers"],
        )
        print(
            f"  ✓ Validation dataloader loaded with {len(val_dataset)} samples",
            flush=True,
        )

    # ==========================
    #        Loss Function
    # ==========================
    # Fused linear logprobs compute next-token logprobs directly from hidden states
    # (chunked over the sequence) and never materialize the full
    # [batch, seq_len, vocab_size] logit tensor, which significantly reduces peak
    # memory. It is only available on the Megatron backend.
    # Both megatron_cfg and use_fused_linear_logprobs are NotRequired, and many
    # configs (e.g. nemo_gym, modelopt, non-megatron) omit them -- use .get() with
    # a {} fallback to avoid a KeyError.
    megatron_cfg = policy_config.get("megatron_cfg", {})
    use_fused_linear_logprobs = bool(
        megatron_cfg.get("enabled") and megatron_cfg.get("use_fused_linear_logprobs")
    )
    if use_fused_linear_logprobs:
        # Sequence packing is not yet validated with the fused path: the fused
        # forward rolls labels over the whole (packed) sequence and would mix
        # tokens across packed-sequence boundaries.
        assert not policy_config["sequence_packing"]["enabled"], (
            "Linear CE fusion loss is not supported with sequence packing for GRPO. "
            "The fused path has not been validated with cu_seqlens-based logprob "
            "aggregation. Set policy.megatron_cfg.use_fused_linear_logprobs=false "
            "or policy.sequence_packing.enabled=false."
        )
        # The fused forward gathers the logprob of the realized token from the raw
        # (unfiltered) logits, so top-k/top-p training-time filtering cannot be
        # applied. This also keeps prev/reference logprobs (computed via the fused
        # get_logprobs path) consistent with the actor logprobs.
        assert not need_top_k_or_top_p_filtering(
            TrainingSamplingParams(
                top_k=generation_config["top_k"],
                top_p=generation_config["top_p"],
            )
        ), (
            "Linear CE fusion loss is not supported with top-k/top-p training-time "
            "filtering for GRPO. The fused path computes logprobs from unfiltered "
            "logits. Set policy.megatron_cfg.use_fused_linear_logprobs=false, or "
            "disable filtering (policy.generation.top_k=null, "
            "policy.generation.top_p=1.0)."
        )

    loss_fn = ClippedPGLossFn(
        loss_config,
        use_fused_linear_logprobs=use_fused_linear_logprobs,
        seq_logprob_error_threshold=grpo_config.seq_logprob_error_threshold,
    )

    # Validate force_on_policy_ratio
    if loss_config.force_on_policy_ratio:
        assert (
            grpo_config.num_prompts_per_step * grpo_config.num_generations_per_prompt
            == policy_config["train_global_batch_size"]
        ), (
            "force_on_policy_ratio requires train_global_batch_size == num_prompts_per_step * num_generations_per_prompt"
        )
        os.environ["NRL_IGNORE_TP_ACCURACY_CHECK"] = "1"
        print("  ✓ force_on_policy_ratio enabled")

    # Validate skip_reference_policy_logprobs_calculation
    if grpo_config.skip_reference_policy_logprobs_calculation:
        assert loss_config.reference_policy_kl_penalty == 0, (
            "grpo.skip_reference_policy_logprobs_calculation=True requires "
            "loss_fn.reference_policy_kl_penalty == 0"
        )
        print(
            "Reference policy logprob calculation will be skipped since `grpo.skip_reference_policy_logprobs_calculation` is set to True and `loss_fn.reference_policy_kl_penalty` is 0."
        )

    _validate_use_kl_in_reward_compat(master_config)

    # Fail on an invalid zero-KL config before Ray allocates GPUs.
    validate_zero_train_gen_kl(policy_config, check_environment=False)

    # ==========================
    #          Cluster
    # ==========================
    print("\n▶ Setting up compute cluster...", flush=True)
    colocated_inference = generation_config["colocated"]["enabled"]

    env_name_list = extract_necessary_env_names(data_config)
    rm_env_enabled = "reward_model" in env_name_list

    # NeMo Gym is initialized inside setup() (rather than by the caller) so its
    # spinup can overlap with vLLM model loading via deferred model load.
    enable_nemo_gym = should_use_nemo_gym(master_config)
    _raise_if_reward_penalties_enabled_without_nemo_gym(
        master_config, enable_nemo_gym=enable_nemo_gym
    )
    _raise_if_length_penalty_enabled_without_nemo_gym(
        master_config, enable_nemo_gym=enable_nemo_gym
    )
    nemo_gym_actor = None

    def _spinup_nemo_gym(base_urls, model_name, policy_api_key=None):
        """Spin up the NeMo Gym stack against the given generation server URLs.

        Returns a shard set, which is the one actor of an unsharded job as much
        as it is the K of a sharded one. ``policy_api_key`` is the bearer token
        the generation server requires on its OpenAI routes (the vLLM backend's
        ``http_server_api_key``), or None when the server checks none.
        """
        t0 = time.perf_counter()
        shard_set = build_nemo_gym_actors(
            env_configs,
            base_urls=base_urls,
            model_name=model_name,
            tokenizer=tokenizer,
            enable_router_replay=router_replay_enabled(policy_config),
            use_fastokens=bool(policy_config["tokenizer"].get("use_fastokens")),
            policy_api_key=policy_api_key,
        )
        train_splits = (
            {f"train[{name}]": split for name, split in dataset.items()}
            if isinstance(dataset, dict)
            else {"train": dataset}
        )
        try:
            validate_dataset_agent_coverage(
                shard_set, {**train_splits, "validation": val_dataset}
            )
        except BaseException:
            shard_set.shutdown()
            raise
        return shard_set, time.perf_counter() - t0

    total_nodes = cluster_config.num_nodes
    segment_size = cluster_config.segment_size
    # Topology of nodes left over after policy/inference placement; non-colocated
    # OPD teachers are placed within it so their collectives stay on NVLink.
    teacher_segment_topology: Optional[dict[str, tuple[str, int]]] = None
    if rm_env_enabled:
        rm_resource = env_configs["reward_model"]["resources"]
        rm_nodes = rm_resource["num_nodes"]
        rm_gpus_per_node = rm_resource["gpus_per_node"]
    else:
        rm_nodes = 0
        rm_gpus_per_node = 0

    if total_nodes == 1:
        policy_nodes = total_nodes
    else:
        policy_nodes = total_nodes - rm_nodes
        assert policy_nodes > 0, (
            "policy_nodes must be > 0, but got "
            f"policy_nodes:{policy_nodes} + rm_nodes:{rm_nodes} = total_nodes:{total_nodes}"
        )

    # Reserve nodes for non-colocated OPD teachers so training doesn't claim them.
    opd_teacher_nodes = 0
    enable_opd_teachers = opd_module.is_non_colocated_teachers_enabled(master_config)
    if enable_opd_teachers:
        assert should_use_async_rollouts(generation_config), (
            "Non-colocated OPD teachers require async GRPO (vLLM backend with async_engine enabled)."
        )
        from nemo_rl.models.policy.teacher_worker_group import (
            create_teacher_configs_from_opd_config,
        )

        opd_cfg = opd_module._opd_cfg(master_config)
        teacher_configs = create_teacher_configs_from_opd_config(opd_cfg)
        for tcfg in teacher_configs:
            assert tcfg.gpus_per_node <= cluster_config.gpus_per_node, (
                f"OPD teacher '{tcfg.alias}' requests gpus_per_node={tcfg.gpus_per_node} > "
                f"cluster.gpus_per_node={cluster_config.gpus_per_node}; "
                "each teacher placement group must fit on one node."
            )
            opd_teacher_nodes += tcfg.num_nodes
        policy_nodes -= opd_teacher_nodes
        assert policy_nodes > 0, (
            "policy_nodes must be > 0 after reserving OPD teacher nodes, but got "
            f"policy_nodes:{policy_nodes} + rm_nodes:{rm_nodes} + opd_teacher_nodes:{opd_teacher_nodes} = total_nodes:{total_nodes}"
        )
        print(
            f"policy_nodes:{policy_nodes} + rm_nodes:{rm_nodes} + opd_teacher_nodes:{opd_teacher_nodes} = total_nodes:{total_nodes}",
            flush=True,
        )

    if colocated_inference:
        if total_nodes == 1:
            policy_gpus_per_node = cluster_config.gpus_per_node - rm_gpus_per_node
            assert policy_gpus_per_node > 0, (
                "policy.generation.colocated.resources.gpus_per_node must be > 0 "
                "when cluster.num_nodes = 1, "
                f"but got {policy_gpus_per_node}."
            )
        else:
            policy_gpus_per_node = cluster_config.gpus_per_node

        node_resource_constraints, policy_remaining_ids, policy_topology = (
            prepare_segment_topology(segment_size, policy_nodes)
        )
        if segment_size is not None:
            teacher_segment_topology = {
                nid: policy_topology[nid] for nid in policy_remaining_ids
            }
        cluster = RayVirtualCluster(
            name="grpo_policy_cluster",
            bundle_ct_per_node_list=[policy_gpus_per_node] * policy_nodes,
            use_gpus=True,
            num_gpus_per_node=policy_gpus_per_node,
            max_colocated_worker_groups=1
            if generation_config["backend"] == "megatron"
            else 2,
            port_range_low=cluster_config.master_port_range_low,
            port_range_high=cluster_config.master_port_range_high,
            segment_size=segment_size,
            node_resource_constraints=node_resource_constraints,
        )
        train_cluster = cluster
        inference_cluster = cluster
        # Colocated generation reuses the policy's cluster; need to decide topology here.
        if (
            node_resource_constraints is not None
            and generation_config["backend"] == "megatron"
        ):
            MegatronGeneration.init_cluster_placement_groups(cluster, policy_config)
        print(
            f"  ✓ Ray cluster for policy initialized with {policy_nodes} nodes",
            flush=True,
        )

    else:
        # train resources will be updated through overall and inference resources below
        train_gpus_per_node = cluster_config.gpus_per_node
        train_nodes = policy_nodes

        inference_resources = generation_config["colocated"]["resources"]
        inference_gpus_per_node = inference_resources["gpus_per_node"]
        inference_nodes = inference_resources["num_nodes"]

        # validate and configure resources
        if policy_nodes == 1:
            # When policy_nodes == 1, train and inference are on the same node
            assert (
                inference_gpus_per_node is not None and inference_gpus_per_node > 0
            ), (
                "policy.generation.colocated.resources.gpus_per_node must be explicitly set to a value > 0 "
                "when policy_nodes = 1 and inference is non-colocated, "
                f"but got {inference_gpus_per_node}."
            )
            assert inference_nodes is None or inference_nodes == 1, (
                "policy.generation.colocated.resources.num_nodes must be 1 or set to null "
                "when policy_nodes = 1 and inference is non-colocated, "
                f"but got {inference_nodes}."
            )

            inference_nodes = 1
            # If total_nodes == 1, reward model is also on the same node; otherwise it's on a different node
            reward_gpus_to_subtract = (
                rm_gpus_per_node if total_nodes == 1 and rm_env_enabled else 0
            )
            train_gpus_per_node -= inference_gpus_per_node + reward_gpus_to_subtract
            assert train_gpus_per_node > 0, (
                "No enough GPUs for training, "
                f"train_gpus_per_node:{train_gpus_per_node} = cluster_config.gpus_per_node:{cluster_config.gpus_per_node} - inference_gpus_per_node:{inference_gpus_per_node}"
                + (
                    f" - rm_gpus_per_node:{rm_gpus_per_node}"
                    if total_nodes == 1 and rm_env_enabled
                    else ""
                )
            )
        else:
            # train, inference, and reward model are all on different nodes
            assert inference_nodes > 0, (
                "policy.generation.colocated.resources.num_nodes must be > 0 "
                "when cluster.num_nodes > 1 and inference is non-colocated, "
                f"but got {inference_nodes}."
            )
            assert (
                inference_gpus_per_node is not None
                and inference_gpus_per_node == cluster_config.gpus_per_node
            ), (
                "policy.generation.colocated.resources.gpus_per_node must be explicitly set and equal to cluster.gpus_per_node "
                "when cluster.num_nodes > 1 and inference is non-colocated, "
                f"but got inference_gpus_per_node={inference_gpus_per_node}, cluster.gpus_per_node={cluster_config.gpus_per_node}."
            )
            train_nodes -= inference_nodes

        assert train_nodes > 0 and inference_nodes > 0, (
            f"Non-colocated mode requires train_nodes > 0 and inference_nodes > 0, "
            f"got train_nodes={train_nodes}, inference_nodes={inference_nodes}"
        )

        # Build topology-aware domain constraints for placement groups.
        # Each selected node's bundles are pinned to a specific NVLink domain so
        # that EP groups stay within high-bandwidth switch fabrics.
        #
        # NOTE: segment_size is also passed to RayVirtualCluster and used later
        # by _sort_bundle_indices_by_topology to trim incomplete domain segments
        # when ordering ranks. When constraints successfully pin nodes to
        # complete segments, that post-placement trimming is a no-op. It serves
        # as defense-in-depth for the fallback path where constraints are absent.
        node_resource_constraints = None
        inference_node_resource_constraints = None
        inference_segment_size = None
        if segment_size is not None:
            topology = get_ray_cluster_topology()
            num_alive_nodes = len(topology)
            required_nodes = train_nodes + inference_nodes
            assert num_alive_nodes >= required_nodes, (
                f"Not enough alive Ray nodes for all roles: "
                f"need {required_nodes} (train={train_nodes} + inference={inference_nodes}), "
                f"but only {num_alive_nodes} alive nodes found"
            )
            node_resource_constraints, remaining_node_ids, topology = (
                prepare_segment_topology(
                    segment_size, train_nodes, topology=topology, role="training"
                )
            )
            # Teachers default to the nodes left after training; narrowed further
            # below if a non-colocated inference cluster is also pinned.
            teacher_segment_topology = {
                nid: topology[nid] for nid in remaining_node_ids
            }
            # Warn if any selected training node lacks topo_rank — domain pinning
            # still works but intra-domain rank ordering will be arbitrary.
            if node_resource_constraints is not None:
                training_node_ids = set(topology) - set(remaining_node_ids)
                nodes_missing_topo_rank = [
                    nid
                    for nid in training_node_ids
                    if topology[nid][1] == TOPO_RANK_UNKNOWN
                ]
                if nodes_missing_topo_rank:
                    print(
                        f"  ⚠ {len(nodes_missing_topo_rank)} selected training nodes have NVLink domain "
                        f"info but no topo_rank; intra-domain rank ordering may be suboptimal",
                        flush=True,
                    )

                # Inference topology: each inference instance spans
                # nodes_per_instance nodes; keep those within one domain
                # so cross-node all-reduce uses NVLink, not InfiniBand.
                #
                # For vLLM: total GPUs per instance = TP * PP (separate dimensions).
                # For SGLang: gpus_per_server already includes all parallelism
                #   dimensions (TP, DP-attention, PP are internal subdivisions),
                #   so we use it directly without multiplying by pp_size.
                # For Megatron: the NVLink-domain span of the parallelism the
                #   generation workers actually run with.
                if generation_config["backend"] == "megatron":
                    gpus_per_instance = MegatronGeneration.nvlink_domain_span(
                        policy_config
                    )
                elif generation_config["backend"] == "vllm":
                    vllm_cfg = generation_config.get("vllm_cfg", {})
                    gpus_per_instance = vllm_cfg["tensor_parallel_size"] * vllm_cfg.get(
                        "pipeline_parallel_size", 1
                    )
                elif generation_config["backend"] == "trtllm":
                    # Under PD disaggregation the unit to keep inside one NVLink
                    # domain is the whole replica, not one engine.
                    gpus_per_instance = replica_width(
                        cast(TrtllmConfig, generation_config)
                    )
                elif generation_config["backend"] == "dynamo":
                    gpus_per_instance = DynamoConfig.model_validate(
                        generation_config
                    ).engine_world_size
                else:
                    sglang_cfg = generation_config.get("sglang_cfg", {})
                    gpus_per_instance = sglang_cfg.get("gpus_per_server", 1)
                nodes_per_instance = (
                    gpus_per_instance + inference_gpus_per_node - 1
                ) // inference_gpus_per_node
                if nodes_per_instance > 1 and inference_nodes % nodes_per_instance == 0:
                    remaining_topology = {
                        nid: topology[nid] for nid in remaining_node_ids
                    }
                    (
                        inference_node_resource_constraints,
                        inference_remaining_ids,
                        _,
                    ) = prepare_segment_topology(
                        nodes_per_instance,
                        inference_nodes,
                        topology=remaining_topology,
                        role="inference",
                    )
                    inference_segment_size = nodes_per_instance
                    teacher_segment_topology = {
                        nid: topology[nid] for nid in inference_remaining_ids
                    }
                elif nodes_per_instance > 1:
                    print(
                        f"  ⚠ inference_nodes={inference_nodes} is not divisible by "
                        f"nodes_per_instance={nodes_per_instance} (gpus_per_instance={gpus_per_instance}); "
                        f"skipping inference topology constraints",
                        flush=True,
                    )

        # initialize train cluster
        train_cluster = RayVirtualCluster(
            name="grpo_train_cluster",
            bundle_ct_per_node_list=[train_gpus_per_node] * train_nodes,
            use_gpus=True,
            num_gpus_per_node=train_gpus_per_node,
            max_colocated_worker_groups=1,
            port_range_low=cluster_config.master_port_range_low,
            port_range_high=cluster_config.master_port_range_high,
            segment_size=segment_size,
            node_resource_constraints=node_resource_constraints,
        )
        # When domain constraints are set, eagerly create placement groups
        # so training claims the constrained nodes before inference can grab them.
        if node_resource_constraints is not None:
            train_cluster.get_placement_groups()
        print(
            f"  ✓ Ray train cluster initialized with {train_nodes} nodes with {train_gpus_per_node} GPUs per node",
            flush=True,
        )

        # Create inference cluster with topology constraints so TP groups
        # stay within NVLink domains. Eagerly initialize PGs when constraints
        # are set so inference claims domain-aligned nodes first.
        inference_cluster = RayVirtualCluster(
            name="grpo_inference_cluster",
            bundle_ct_per_node_list=[inference_gpus_per_node] * inference_nodes,
            use_gpus=True,
            num_gpus_per_node=inference_gpus_per_node,
            max_colocated_worker_groups=1,
            port_range_low=cluster_config.master_port_range_low,
            port_range_high=cluster_config.master_port_range_high,
            segment_size=inference_segment_size,
            node_resource_constraints=inference_node_resource_constraints,
        )
        if inference_node_resource_constraints is not None:
            if generation_config["backend"] == "megatron":
                # Megatron inference reuses the training parallelism config.
                MegatronGeneration.init_cluster_placement_groups(
                    inference_cluster, policy_config
                )
            elif generation_config["backend"] == "dynamo":
                # Managed Dynamo creates one single-node engine per placement
                # group and does not need a backend-specific PG strategy.
                inference_cluster.get_placement_groups()
            else:
                {
                    "vllm": VllmGeneration,
                    "trtllm": TrtllmGeneration,
                }[generation_config["backend"]].init_cluster_placement_groups(
                    inference_cluster,
                    generation_config,
                )
        print(
            f"  ✓ Ray inference cluster initialized with {inference_nodes} nodes with {inference_gpus_per_node} GPUs per node",
            flush=True,
        )

    # Reserve topology-aware teacher placement groups before NeMo Gym starts
    # opportunistically placing its GPU-backed services. Worker creation and
    # model loading remain deferred until the policy is ready to avoid racing
    # Megatron-Bridge checkpoint conversion.
    teacher_clusters: dict[str, RayVirtualCluster] = {}
    teacher_reservation_time = 0.0
    if enable_opd_teachers:
        t0 = time.perf_counter()
        teacher_clusters = opd_module.reserve_teacher_clusters(
            master_config,
            segment_size=segment_size,
            teacher_segment_topology=teacher_segment_topology,
        )
        teacher_reservation_time = time.perf_counter() - t0

    # ==========================
    #   Training and Inference
    # ==========================
    print("\n▶ Setting up model and training...", flush=True)

    # vllm model loading prefers clean environment, initialize policy_generation before policy in colocated mode
    backend = generation_config["backend"]
    offload_policy_before_refit = policy_config["offload_policy_before_refit"]
    validate_offload_policy_before_refit(
        enabled=offload_policy_before_refit,
        megatron_enabled=bool(
            (policy_config.get("megatron_cfg") or {}).get("enabled", False)
        ),
        generation_backend=backend,
        colocated=colocated_inference,
        refit_transport=generation_config.get("refit_transport"),
    )
    generation_config["model_name"] = policy_config["model_name"]  # Needed for vLLM
    generation_config["_debug_payload_metrics"] = grpo_config.debug_payload_metrics
    remote_transport = None
    remote_synchronizer_cls = None
    remote_baseline_init_refs: list[Any] = []
    checkpoint_engine_config = None

    # Worker initialization timing stats — populated as each phase completes.
    setup_timing_metrics = SetupTimingMetrics()
    if teacher_reservation_time:
        setup_timing_metrics.teacher_reservation_time_s = teacher_reservation_time

    weights_path, optimizer_path = checkpointer.get_resume_paths(last_checkpoint_path)

    if policy_config.get("megatron_cfg", {}).get("enabled", False):
        ## NOTE: this is equal to the total number of scheduler steps
        total_train_iters = min(
            grpo_config.max_num_steps,
            grpo_config.max_num_epochs * train_sample_count,
        )
        policy_config["megatron_cfg"]["train_iters"] = total_train_iters

    # Megatron generation expresses recompute-after-refit engine-side via
    # `kv_cache_management_mode="recompute"`; the loop-level flag must agree.
    if generation_config["backend"] == "megatron":
        async_grpo_config = grpo_config.async_grpo
        recompute_kv_cache = bool(
            async_grpo_config is not None
            and async_grpo_config.recompute_kv_cache_after_weight_updates
        )
        kv_cache_mode = generation_config["mcore_generation_config"][
            "kv_cache_management_mode"
        ]
        if recompute_kv_cache != (kv_cache_mode == "recompute"):
            raise ValueError(
                "grpo.async_grpo.recompute_kv_cache_after_weight_updates="
                f"{recompute_kv_cache} conflicts with policy.generation."
                f"mcore_generation_config.kv_cache_management_mode={kv_cache_mode!r}: "
                "with policy.generation.backend='megatron' the two must agree. "
                "Either set the flag to true with kv_cache_management_mode="
                "'recompute', or leave the flag false with 'persist'/'offload'."
            )

    # Define initialization functions that will be used in all paths
    init_reference_model = loss_config.reference_policy_kl_penalty > 0

    # Auto-enable skip_reference_policy_logprobs_calculation when the reference model is not loaded.
    if (
        not init_reference_model
        and not grpo_config.skip_reference_policy_logprobs_calculation
    ):
        grpo_config.skip_reference_policy_logprobs_calculation = True
        print(
            "Auto-enabling `grpo.skip_reference_policy_logprobs_calculation=True` "
            "because `loss_fn.reference_policy_kl_penalty == 0` "
            "(reference model is not loaded)."
        )

    # Caller-supplied factory lets the sync trainer swap in a TQ-mediated
    # Policy subclass without this shared setup needing to know the data
    # plane exists. Default is the plain Policy class — legacy behavior.
    _make_policy = policy_factory if policy_factory is not None else Policy

    def init_policy(reserved_http_server_ports: Optional[dict[int, int]] = None):
        """Initialize policy training workers."""
        t0 = time.perf_counter()
        extra_policy_kwargs = {}
        if reserved_http_server_ports is not None:
            # Colocated Megatron generation serves HTTP from the training workers.
            extra_policy_kwargs["reserved_http_server_ports"] = (
                reserved_http_server_ports
            )
        p = _make_policy(
            cluster=train_cluster,
            config=policy_config,
            tokenizer=tokenizer,
            processor=processor,
            weights_path=weights_path,
            optimizer_path=optimizer_path,
            init_optimizer=True,
            init_reference_model=init_reference_model,
            **extra_policy_kwargs,
        )
        # Keep custom policy_factory call signatures backward compatible.
        p.debug_payload_metrics = grpo_config.debug_payload_metrics
        if remote_transport is not None:
            assert remote_synchronizer_cls is not None
            remote_baseline_init_refs.extend(
                remote_synchronizer_cls.start_baseline(p, remote_transport)
            )
        return p, time.perf_counter() - t0

    def init_vllm():
        """Initialize vLLM generation workers."""
        t0 = time.perf_counter()
        pg = VllmGeneration(cluster=inference_cluster, config=generation_config)
        pg.finish_generation()
        return pg, time.perf_counter() - t0

    def init_sglang():
        """Initialize SGLang generation workers."""
        t0 = time.perf_counter()
        pg = SGLangGeneration(
            cluster=inference_cluster,
            sglang_cfg=generation_config,
        )
        pg.finish_generation()
        return pg, time.perf_counter() - t0

    def init_megatron_generation(
        policy=None, reserved_http_server_ports: Optional[dict[int, int]] = None
    ):
        """Initialize Megatron generation."""
        t0 = time.perf_counter()
        mg = MegatronGeneration(
            config=policy_config,
            tokenizer=tokenizer,
            cluster=None if colocated_inference else inference_cluster,
            policy=policy if colocated_inference else None,
            processor=processor,
            skip_weight_load=not colocated_inference,
            reserved_http_server_ports=reserved_http_server_ports,
        )
        return mg, time.perf_counter() - t0

    def init_megatron_weight_synchronizer(
        policy: ColocatablePolicyInterface,
        policy_generation: MegatronGeneration,
    ) -> None:
        """Initialize Megatron weight synchronizer.

        For non-colocated inference, also performs the initial weight sync.
        """
        t0 = time.perf_counter()
        weight_synchronizer = create_weight_synchronizer(
            policy=policy,
            generation=policy_generation,
            generation_backend="megatron",
            colocated=colocated_inference,
            train_cluster=train_cluster,
            inference_cluster=None if colocated_inference else inference_cluster,
        )
        policy_generation.weight_synchronizer = weight_synchronizer
        weight_synchronizer.init_communicator()
        setup_timing_metrics.collective_init_time_s = time.perf_counter() - t0
        if not colocated_inference:
            # The skip-load inference engine gets its final weight buffers here.
            # Its first prepare_for_generation also starts the HTTP server only after the refit,
            # so CUDA graphs can capture those persistent buffers.
            t0 = time.perf_counter()
            weight_synchronizer.sync_weights()
            setup_timing_metrics.weight_sync_time_s = time.perf_counter() - t0

    def initialize_generation_with_policy(
        init_generation_fn,
        colocated_inference: bool,
        setup_timing_metrics: SetupTimingMetrics,
    ):
        """Initialize a generation engine along with policy, sequentially or in parallel.

        Args:
            init_generation_fn: Function that initializes the generation engine (init_vllm, ...).
            colocated_inference: Whether inference is colocated with training.
            setup_timing_metrics: SetupTimingMetrics to store timings on.

        Returns:
            Tuple of (policy_generation, policy).
        """
        # Determine if parallel initialization is possible (non-colocated mode)
        use_parallel_init = not colocated_inference

        if use_parallel_init:
            # Parallel initialization: Generation engine and Policy can initialize simultaneously
            print(
                "  ⚡ Using parallel worker initialization (non-colocated mode)",
                flush=True,
            )

            # Execute both initializations in parallel
            parallel_start_time = time.perf_counter()
            with ThreadPoolExecutor(max_workers=2) as executor:
                generation_future = executor.submit(init_generation_fn)
                policy_future = executor.submit(init_policy)
                policy_generation, generation_time = generation_future.result()
                policy, policy_time = policy_future.result()
            parallel_wall_time = time.perf_counter() - parallel_start_time

            # Store timing metrics
            setup_timing_metrics.generation_init_time_s = generation_time
            setup_timing_metrics.policy_init_time_s = policy_time
            setup_timing_metrics.parallel_wall_time_s = parallel_wall_time
            setup_timing_metrics.parallel_init_enabled = 1.0

        else:
            # Sequential initialization: colocated mode (GPU memory requires generation engine first)
            print(
                "  ⚙️  Using sequential worker initialization (colocated mode)",
                flush=True,
            )

            # Initialize generation engine first (clean GPU memory), then policy
            policy_generation, generation_time = init_generation_fn()
            setup_timing_metrics.generation_init_time_s = generation_time

            policy, policy_time = init_policy()
            setup_timing_metrics.policy_init_time_s = policy_time
            setup_timing_metrics.parallel_init_enabled = 0.0

        return policy_generation, policy

    # Handle generation-specific setup
    if backend == "megatron":
        if enable_nemo_gym:
            print(
                "  ⚡ Reserving the Megatron server address for overlapped NeMo Gym init",
                flush=True,
            )
            reserve_t0 = time.perf_counter()
            reserved_urls, reserved_http_server_ports, port_holders = (
                MegatronGeneration.reserve_http_server_addresses(
                    train_cluster if colocated_inference else inference_cluster,
                    policy_config,
                )
            )
            reserve_time = time.perf_counter() - reserve_t0
            setup_timing_metrics.generation_init_reserve_time_s = reserve_time
            print(
                f"  ✓ Reserved {len(reserved_urls)} Megatron server URL(s): "
                f"{reserved_urls}",
                flush=True,
            )

            def init_nemo_gym():
                """Spin up NeMo Gym servers against the reserved URLs."""
                return _spinup_nemo_gym(reserved_urls, generation_config["model_name"])

            # Exactly one task adopts the reserved port: the policy when colocated
            # (generation wraps it), else the dedicated generation policy.
            policy_ports, generation_ports = (
                (reserved_http_server_ports, None)
                if colocated_inference
                else (None, reserved_http_server_ports)
            )

            def init_megatron_generation_task(policy_future):
                """Colocated generation waits; non-colocated inits in parallel."""
                if colocated_inference:
                    p, _ = policy_future.result()
                    return init_megatron_generation(p)
                return init_megatron_generation(
                    reserved_http_server_ports=generation_ports
                )

            print("  ⚡ Init tasks: policy, megatron_generation, nemo_gym", flush=True)
            init_tasks_t0 = time.perf_counter()
            nemo_gym_future: Future[tuple[NemoGymShardSet, float]] | None = None
            try:
                with ThreadPoolExecutor(max_workers=3) as executor:
                    policy_future = executor.submit(
                        init_policy, reserved_http_server_ports=policy_ports
                    )
                    generation_future = executor.submit(
                        init_megatron_generation_task, policy_future
                    )
                    nemo_gym_future = executor.submit(init_nemo_gym)
                    policy, policy_time = policy_future.result()
                    policy_generation, megatron_gen_time = generation_future.result()
                    if not colocated_inference:
                        setup_timing_metrics.parallel_wall_time_s = (
                            time.perf_counter() - init_tasks_t0
                        )
                        setup_timing_metrics.parallel_init_enabled = 1.0
                        # NeMo Gym probes the pre-published endpoint before its future completes.
                        # A skip-load Megatron endpoint starts only during this initial refit,
                        # so it must happen while Gym is waiting rather than after it resolves.
                        init_megatron_weight_synchronizer(policy, policy_generation)
                    nemo_gym_actor, nemo_gym_time = nemo_gym_future.result()
            except BaseException:
                _shutdown_completed_nemo_gym_startup(nemo_gym_future)
                raise
            finally:
                for port_holder in port_holders:
                    ray.kill(port_holder)

            if colocated_inference:
                setup_timing_metrics.parallel_init_enabled = 0.0
            setup_timing_metrics.policy_init_time_s = policy_time
            setup_timing_metrics.generation_init_time_s = (
                reserve_time + megatron_gen_time
            )
            setup_timing_metrics.generation_init_load_time_s = megatron_gen_time
            setup_timing_metrics.nemo_gym_init_time_s = nemo_gym_time

        else:
            if not colocated_inference:
                policy_generation, policy = initialize_generation_with_policy(
                    init_megatron_generation,
                    colocated_inference,
                    setup_timing_metrics,
                )
            else:
                # Colocated generation wraps the training policy.
                policy, policy_time = init_policy()
                setup_timing_metrics.policy_init_time_s = policy_time

                policy_generation, megatron_gen_time = init_megatron_generation(policy)
                setup_timing_metrics.generation_init_time_s = megatron_gen_time
                setup_timing_metrics.parallel_init_enabled = 0.0

        print(
            f"  ✓ Using {backend} backend for generation with {policy_config['model_name']}",
            flush=True,
        )

    elif backend == "vllm":
        # vLLM generation: setup config, then initialize with policy
        generation_config = cast(VllmConfig, generation_config)
        refit_transport = generation_config.get("refit_transport")
        if refit_transport in VLLM_SPARSE_REFIT_TRANSPORTS:
            # Keep optional remote transport dependencies off the default path.
            from nemo_rl.weight_sync.vllm_remote_sparse_weight_synchronizer import (
                VllmRemoteSparseWeightSynchronizer,
                validate_vllm_remote_sparse_refit,
            )

            remote_transport = validate_vllm_remote_sparse_refit(
                generation_config,
                colocated=colocated_inference,
                megatron_enabled=policy_config["megatron_cfg"]["enabled"],
            )
            assert remote_transport is not None
            remote_synchronizer_cls = VllmRemoteSparseWeightSynchronizer
        elif refit_transport is not None and refit_transport != "nccl_reshard":
            # nccl_reshard is handled below via nccl_reshard_refit_enabled,
            # not via checkpoint-engine.
            checkpoint_engine_config = checkpoint_engine_refit_config(generation_config)
            assert checkpoint_engine_config is not None

        if generation_config["vllm_cfg"]["precision"] == "fp8":
            assert loss_config.use_importance_sampling_correction, (
                "Importance sampling must be enabled for vLLM FP8 generation for good convergence!"
            )
        kv_cache_dtype = generation_config["vllm_cfg"]["kv_cache_dtype"]
        if kv_cache_dtype.startswith("fp8"):
            # FP8 KV cache requires FP8 model precision
            assert generation_config["vllm_cfg"]["precision"] == "fp8", (
                f"kv_cache_dtype='{kv_cache_dtype}' requires precision='fp8'. "
                "FP8 KV cache can only be used together with FP8 model weights."
            )
        if kv_cache_dtype in REFITTABLE_FP8_KV_CACHE_DTYPES:
            assert not (policy_config.get("automodel_cfg") or {}).get("enabled"), (
                "Automodel backend is not supported with kv cache fp8 enabled."
            )
            assert not should_use_async_rollouts(generation_config), (
                "Async rollouts is not supported with kv cache fp8 enabled."
            )
            assert policy_config["megatron_cfg"]["pipeline_model_parallel_size"] == 1, (
                "Currently when using FP8 KV cache in generation, then in megatron we only support pipeline_model_parallel_size=1. We will add more support in future."
            )

        configure_vllm_for_router_replay(policy_config)
        vllm_kwargs = generation_config.setdefault("vllm_kwargs", {})

        ## make vllm hf overrides match the training policy
        vllm_kwargs["hf_overrides"] = policy_config.get("hf_config_overrides", {})

        if enable_nemo_gym:
            # ---- NeMo Gym: reserve vLLM ports up-front so we can hand the
            # server URLs to NeMo Gym and spin it up while vLLM loads weights.
            print(
                "  ⚡ Deferred model load: reserving vLLM ports for overlapped NeMo Gym init",
                flush=True,
            )
            vllm_reserve_t0 = time.perf_counter()
            deferred_vllm = VllmGeneration(
                cluster=inference_cluster,
                config=generation_config,
                defer_model_load=True,
            )
            vllm_reserve_time = time.perf_counter() - vllm_reserve_t0
            print(
                f"  ✓ Reserved {len(deferred_vllm.dp_openai_server_base_urls)} vLLM server URLs: "
                f"{deferred_vllm.dp_openai_server_base_urls}",
                flush=True,
            )

            def init_vllm_deferred():
                """Complete the deferred vLLM model load started above."""
                t0 = time.perf_counter()
                deferred_vllm.load_and_start()
                deferred_vllm.finish_generation()
                return deferred_vllm, time.perf_counter() - t0

            def init_nemo_gym():
                """Spin up NeMo Gym servers with the pre-assigned vLLM URLs."""
                return _spinup_nemo_gym(
                    deferred_vllm.dp_openai_server_base_urls,
                    generation_config["model_name"],
                    policy_api_key=deferred_vllm.http_server_api_key,
                )

            # Colocated: vLLM + policy share GPUs -> sequential; otherwise parallel.
            init_tasks = {}
            if colocated_inference:

                def init_vllm_then_policy():
                    pg, vllm_t = init_vllm_deferred()
                    p, policy_t = init_policy()
                    return pg, vllm_t, p, policy_t

                init_tasks["vllm_policy"] = init_vllm_then_policy
            else:
                init_tasks["vllm"] = init_vllm_deferred
                init_tasks["policy"] = init_policy
            init_tasks["nemo_gym"] = init_nemo_gym

            print(
                f"  ⚡ Init tasks: {', '.join(init_tasks.keys())}",
                flush=True,
            )
            submitted: dict[str, Future[Any]] = {}
            try:
                with ThreadPoolExecutor(max_workers=len(init_tasks)) as executor:
                    submitted = {
                        key: executor.submit(task) for key, task in init_tasks.items()
                    }
                    results = {
                        key: future.result() for key, future in submitted.items()
                    }
            except BaseException:
                _shutdown_completed_nemo_gym_startup(submitted.get("nemo_gym"))
                raise

            if colocated_inference:
                policy_generation, vllm_load_time, policy, policy_time = results[
                    "vllm_policy"
                ]
            else:
                policy_generation, vllm_load_time = results["vllm"]
                policy, policy_time = results["policy"]
            nemo_gym_actor, nemo_gym_time = results["nemo_gym"]
            setup_timing_metrics.generation_init_time_s = (
                vllm_reserve_time + vllm_load_time
            )
            setup_timing_metrics.generation_init_reserve_time_s = vllm_reserve_time
            setup_timing_metrics.generation_init_load_time_s = vllm_load_time
            setup_timing_metrics.policy_init_time_s = policy_time
            setup_timing_metrics.nemo_gym_init_time_s = nemo_gym_time
        else:
            policy_generation, policy = initialize_generation_with_policy(
                init_generation_fn=init_vllm,
                colocated_inference=colocated_inference,
                setup_timing_metrics=setup_timing_metrics,
            )

        print(
            f"  ✓ Using vLLM backend for generation with {policy_config['model_name']}",
            flush=True,
        )

    elif backend == "sglang":
        generation_config = cast(SGLangConfig, generation_config)

        # Set model_path if not already set
        if "model_path" not in generation_config["sglang_cfg"]:
            generation_config["sglang_cfg"]["model_path"] = policy_config["model_name"]

        policy_generation, policy = initialize_generation_with_policy(
            init_generation_fn=init_sglang,
            colocated_inference=colocated_inference,
            setup_timing_metrics=setup_timing_metrics,
        )

        print(
            f"  ✓ Using SGLang backend for generation with {policy_config['model_name']}",
            flush=True,
        )

    elif backend == "trtllm":
        generation_config = cast(TrtllmConfig, generation_config)

        def init_trtllm():
            """Initialize TRT-LLM generation workers."""
            t0 = time.perf_counter()
            pg = TrtllmGeneration(cluster=inference_cluster, config=generation_config)
            pg.finish_generation()
            return pg, time.perf_counter() - t0

        policy_generation, policy = initialize_generation_with_policy(
            init_generation_fn=init_trtllm,
            colocated_inference=colocated_inference,
            setup_timing_metrics=setup_timing_metrics,
        )

        print(
            f"  ✓ Using TRT-LLM backend for generation with {policy_config['model_name']}",
            flush=True,
        )

        if enable_nemo_gym:
            nemo_gym_actor, nemo_gym_time = _spinup_nemo_gym(
                policy_generation.dp_openai_server_base_urls,
                generation_config["model_name"],
            )
            setup_timing_metrics.nemo_gym_init_time_s = nemo_gym_time

    elif backend == "dynamo":
        # Managed Dynamo owns a fixed worker fleet on the inference virtual cluster.

        def init_dynamo():
            t0 = time.perf_counter()
            generation = DynamoGeneration(
                cluster=inference_cluster,
                config=generation_config,
                tokenizer=tokenizer,
                tokenizer_config=policy_config["tokenizer"],
            )
            return generation, time.perf_counter() - t0

        policy_generation, policy = initialize_generation_with_policy(
            init_generation_fn=init_dynamo,
            colocated_inference=False,
            setup_timing_metrics=setup_timing_metrics,
        )

        if enable_nemo_gym:
            nemo_gym_actor, nemo_gym_time = _spinup_nemo_gym(
                policy_generation.dp_openai_server_base_urls,
                generation_config["model_name"],
            )
            setup_timing_metrics.nemo_gym_init_time_s = nemo_gym_time

        print(
            f"  ✓ Using Dynamo backend (frontend: {policy_generation.frontend_url})",
            flush=True,
        )

    # Record when worker initialization completes (for calculating other setup time)
    worker_init_complete_time = time.perf_counter() - setup_start_time

    # print the node IP and GPU ID of the policy workers for debugging
    policy.print_node_ip_and_gpu_id()

    nccl_reshard_refit_enabled = (
        generation_config.get("refit_transport") == "nccl_reshard"
    )
    if nccl_reshard_refit_enabled:
        check_nccl_reshard_refit_support(master_config)

    refit_transport = generation_config.get("refit_transport")
    if refit_transport is not None and not (
        backend == "vllm"
        or (backend == "megatron" and refit_transport in ("mcore", "nccl_reshard"))
    ):
        raise NotImplementedError(
            f"refit_transport={refit_transport!r} is not supported for "
            f"policy.generation.backend={backend!r}. "
            "Set policy.generation.refit_transport=null. Megatron generation "
            "supports refit_transport='mcore' or 'nccl_reshard'."
        )

    if backend == "megatron":
        if policy_generation.weight_synchronizer is None:
            init_megatron_weight_synchronizer(policy, policy_generation)
        if enable_nemo_gym:
            MegatronGeneration.verify_served_addresses(
                policy_generation.dp_openai_server_base_urls, reserved_urls
            )
    # if it is not colocated inference, initialize collective communication for update weights
    elif (
        not colocated_inference
        and backend != "sglang"
        and remote_transport is None
        and checkpoint_engine_config is None
    ):
        t0 = time.perf_counter()
        # init collective
        if _uses_managed_noncolocated_refit(
            generation_backend=backend,
            nccl_reshard_refit_enabled=nccl_reshard_refit_enabled,
            offload_policy_before_refit=offload_policy_before_refit,
        ):
            policy_generation.weight_synchronizer = create_weight_synchronizer(
                policy=policy,
                generation=policy_generation,
                generation_backend=backend,
                colocated=False,
                train_cluster=train_cluster,
                inference_cluster=inference_cluster,
            )
            policy_generation.weight_synchronizer.init_communicator()
        else:
            ip, port = train_cluster.get_master_address_and_port()
            print(
                f"Using ip: {ip}, port: {port} for collective communication",
                flush=True,
            )
            train_world_size = train_cluster.world_size()
            inference_world_size = inference_nodes * inference_gpus_per_node
            world_size = train_world_size + inference_world_size
            futures_train = policy.init_collective(
                ip, port, world_size, train_world_size=train_world_size
            )
            futures_inference = policy_generation.init_collective(
                ip, port, world_size, train_world_size=train_world_size
            )  # type: ignore
            ray.get(futures_train + futures_inference)
        setup_timing_metrics.collective_init_time_s = time.perf_counter() - t0

    if remote_transport is not None:
        t0 = time.perf_counter()
        assert isinstance(policy_generation, VllmGeneration)
        assert remote_synchronizer_cls is not None
        refit_config = generation_config["refit_cfg"]
        assert refit_config is not None
        policy_generation.weight_synchronizer = remote_synchronizer_cls(
            policy,
            policy_generation,
            transport=remote_transport,
            api_key_env_var=generation_config["vllm_cfg"].get(
                "http_refit_api_key_env_var"
            ),
            request_timeout_s=refit_config.sparse.request_timeout_s,
            baseline_init_refs=remote_baseline_init_refs,
        )
        policy_generation.weight_synchronizer.init_communicator()
        setup_timing_metrics.extras[f"vllm_{remote_transport}_sparse_init_time_s"] = (
            time.perf_counter() - t0
        )
    elif checkpoint_engine_config is not None:
        t0 = time.perf_counter()
        assert isinstance(policy_generation, VllmGeneration)
        policy_generation.weight_synchronizer = create_weight_synchronizer(
            policy=policy,
            generation=policy_generation,
            generation_backend=backend,
            colocated=colocated_inference,
            train_cluster=train_cluster,
            inference_cluster=inference_cluster,
        )
        policy_generation.weight_synchronizer.init_communicator()
        setup_timing_metrics.vllm_checkpoint_engine_init_time_s = (
            time.perf_counter() - t0
        )
        print(
            f"Using checkpoint-engine refit backend: {checkpoint_engine_config['backend']}",
            flush=True,
        )
    elif backend == "sglang":
        t0 = time.perf_counter()
        policy_generation.weight_synchronizer = create_weight_synchronizer(
            policy=policy,
            generation=policy_generation,
            generation_backend=backend,
            colocated=colocated_inference,
            refit_buffer_size_gb=policy_config.get("refit_buffer_size_gb"),
        )
        # Only exchanges refit metadata. SGLang's own weight-update group is
        # established lazily on the first refit.
        policy_generation.weight_synchronizer.init_communicator()
        setup_timing_metrics.extras["sglang_weight_sync_init_time_s"] = (
            time.perf_counter() - t0
        )
    else:
        if getattr(
            policy_generation, "weight_synchronizer", None
        ) is None and _needs_hf_refit_handshake(
            backend, nccl_reshard_refit_enabled, colocated_inference
        ):
            state_dict_info = policy.prepare_refit_info()
            if policy_generation is not None:
                policy_generation.prepare_refit_info(state_dict_info)

    # Spin up non-colocated OPD teacher worker groups AFTER policy / vLLM are
    # ready. Parallelizing with policy init races on Megatron-Bridge's HF->mcore
    # cache (shared key when student == teacher) — both workers write to the
    # same iter_0000000/ path and the second reader gets a truncated file.
    teacher_worker_groups: dict[str, Any] = {}
    alias_to_group_alias: dict[str, str] = {}
    if enable_opd_teachers:
        t0 = time.perf_counter()
        teacher_worker_groups, alias_to_group_alias = (
            opd_module.create_teacher_worker_groups(
                master_config,
                policy_config,
                tokenizer,
                teacher_clusters=teacher_clusters,
            )
        )
        teacher_model_init_time = time.perf_counter() - t0
        setup_timing_metrics.teacher_model_init_time_s = teacher_model_init_time
        # Preserve the existing metric's end-to-end meaning while exposing the
        # newly separated reservation and model-initialization phases.
        setup_timing_metrics.teacher_init_time_s = (
            teacher_reservation_time + teacher_model_init_time
        )

    # Calculate total setup time
    total_setup_time = time.perf_counter() - setup_start_time
    setup_timing_metrics.total_setup_time_s = total_setup_time
    setup_timing_metrics.other_setup_time_s = (
        total_setup_time - worker_init_complete_time
    )

    # Log worker initialization timing metrics to logger
    print_setup_timing_summary(setup_timing_metrics)
    logger.log_metrics(
        setup_timing_metrics.to_metrics_dict(), step=0, prefix=SETUP_TIMING_PREFIX
    )

    print("\n" + "=" * 60)
    print(" " * 18 + "SETUP COMPLETE")
    print(f"  Total setup time: {total_setup_time:.1f}s")
    print("=" * 60 + "\n", flush=True)

    return (
        policy,
        policy_generation,
        nemo_gym_actor,
        (train_cluster, inference_cluster),
        dataloader,
        val_dataloader,
        loss_fn,
        logger,
        checkpointer,
        grpo_save_state,
        master_config,
        teacher_worker_groups,
        alias_to_group_alias,
    )


# ===============================================================================
# Core Algorithm Functions
# ===============================================================================


def dynamic_sampling(
    repeated_batch: BatchedDataDict[DatumSpec],
    std: torch.Tensor,
    baseline: torch.Tensor,
    dynamic_sampling_num_gen_batches: int,
    master_config: MasterConfig,
    timer: Timer,
    batch_cache: BatchedDataDict[DatumSpec] = None,
    is_trivial_prompt_distribution: torch.Tensor | None = None,
) -> BatchedDataDict[DatumSpec]:
    """Implements the dynamic sampling algorithm to select prompts with non-zero standard deviation.

    This function filters the current batch to retain only those prompts that have a non-zero standard deviation.
    If the current batch has fewer number of prompts with non-zero standard deviation than the required batch size, defined as num_prompts_per_step * num_generations_per_prompt,
    we store it in the batch_cache to be used in later iterations.
    If the current batch has more number of prompts with non-zero standard deviation than the required batch size, defined as num_prompts_per_step * num_generations_per_prompt,
    the batch is sliced to ensure batch size is num_prompts_per_step * num_generations_per_prompt.
    is_batch_complete is set to False to indicate that the current batch is not enough to meet the required batch size. This is used as a signal in the GRPO training loop
    to continue sampling or proceed to training.
    This approach is based on the dynamic sampling algorithm from the DAPO paper:
    https://arxiv.org/pdf/2503.14476.

    Args:
        repeated_batch (BatchedDataDict[DatumSpec]): The current batch of data containing prompts, responses, rewards, baselines, and std.
        std (torch.Tensor): Tensor representing the standard deviation for each prompt group.
        baseline (torch.Tensor): Baseline values for each prompt group.
        dynamic_sampling_num_gen_batches (int): Number of generation batches processed at the current step.
        master_config (MasterConfig): Configuration containing GRPO and policy settings.
        batch_cache (BatchedDataDict[DatumSpec], optional): Cache storing previously selected prompts with non-zero std.
        is_trivial_prompt_distribution (torch.Tensor, optional): Exact-equality
            mask for each sample's full prompt reward group. When provided,
            trivial groups are filtered all-or-nothing even if floating-point
            roundoff produces a positive std.

    Returns:
        tuple: A tuple containing:
            - repeated_batch (BatchedDataDict[DatumSpec]): Updated batch with selected prompts.
            - is_batch_complete (bool): Indicates if the batch has enough samples with non-zero std for training.
            - batch_cache (BatchedDataDict[DatumSpec]): Updated cache for future iterations.
    """
    # is_batch_complete is used to indicate if the current batch was able to generate enough prompts with non-zero std.
    is_batch_complete = True

    # Required batch size for training
    train_prompts_size = (
        master_config.grpo.num_prompts_per_step
        * master_config.grpo.num_generations_per_prompt
    )
    # Store the baseline, std and total_reward for the current unfiltered batch.
    repeated_batch["baseline"] = baseline
    repeated_batch["std"] = std
    total_rewards = repeated_batch["total_reward"]
    dynamic_sampling_metrics = {}

    # Dynamic sampling algorithm (used in DAPO algorithm)
    # This block implements dynamic sampling by selecting prompt groups with non-zero std.
    # If sampled prompts (with non-zero std) are fewer than num_prompts_per_step * num_generations_per_prompt, continue sampling until dynamic_sampling_max_gen_batches is reached.
    if master_config.grpo.use_dynamic_sampling:
        with timer.time("dynamic_sampling"):
            # Exact reward equality, rather than floating-point std noise, decides
            # whether a prompt has useful reward variation.
            if is_trivial_prompt_distribution is None:
                raise ValueError(
                    "dynamic_sampling: is_trivial_prompt_distribution is None -- "
                    "the caller must compute it before this call when "
                    "use_dynamic_sampling is set."
                )
            non_trivial_reward_mask = ~is_trivial_prompt_distribution

            keep_prompt_indices = torch.arange(
                len(non_trivial_reward_mask), device=std.device
            )[non_trivial_reward_mask].tolist()

            # Only select the inputs that have non-zero std
            # total_reward is already a part of repeated_batch so we don't need to add it again
            filtered_repeated_batch = repeated_batch.select_indices(keep_prompt_indices)
            filtered_repeated_batch["std"] = std[keep_prompt_indices]
            filtered_repeated_batch["baseline"] = baseline[keep_prompt_indices]

            # Store filtered and total rewards to track them separately
            filtered_rewards = filtered_repeated_batch["total_reward"]
            filtered_repeated_batch["total_reward"] = total_rewards
            filtered_repeated_batch["filtered_reward"] = filtered_rewards

            # Store the total_reward for the current filtered batch.
            # If none of the prompts in current batch have non-zero std, filtered_repeated_batch.size will be 0.
            # In this case, the current batch will be ignored and the next batch will be processed and we generate responses for it.
            if filtered_repeated_batch.size > 0:
                # Concatenate the previous partially filled batch with the current batch. This serves as a cache to store and collect the prompts with non-zero std.
                # This is used in the next iteration when the current batch is not enough to fill the buffer.
                batch_cache = (
                    filtered_repeated_batch
                    if batch_cache is None
                    else BatchedDataDict.from_batches(
                        [batch_cache, filtered_repeated_batch],
                        allow_missing_packed_tensors=(
                            master_config.grpo.deduplicate_multimodal_data
                        ),
                    )
                )
                filtered_repeated_batch = batch_cache

            filtered_prompts_size = filtered_repeated_batch.size
            print(
                f"Detected {filtered_prompts_size} prompts with non-zero std; "
                f"{train_prompts_size} are required and used for training."
            )

            # If the generation samples size is smaller than a fixed threshold (train_prompts_size), keep generating by processing the next batch
            if filtered_prompts_size < train_prompts_size:
                dynamic_sampling_max_gen_batches = (
                    master_config.grpo.dynamic_sampling_max_gen_batches
                )
                assert dynamic_sampling_max_gen_batches > 0, (
                    "When using grpo.use_dynamic_sampling, grpo.dynamic_sampling_max_gen_batches must be > 0"
                )
                if dynamic_sampling_num_gen_batches <= dynamic_sampling_max_gen_batches:
                    print(
                        f"Generation sample buffer size: {filtered_prompts_size} is smaller than train_prompts_size: {train_prompts_size}. Processed {dynamic_sampling_num_gen_batches} batches so far out of {dynamic_sampling_max_gen_batches}."
                    )
                    is_batch_complete = False
                else:
                    raise ValueError(
                        f"Dynamic sampling has reached the maximum allowed number of batches ({dynamic_sampling_max_gen_batches}). Consider evaluating the complexity of your data or adjusting the num_prompts_per_step or num_generations_per_prompt parameters to enhance the diversity of the samples."
                    )
            else:
                num_discarded_valid_samples = filtered_prompts_size - train_prompts_size
                dynamic_sampling_metrics[
                    "dynamic_sampling_num_discarded_valid_samples"
                ] = num_discarded_valid_samples

                #  Slice the batch, rewards, baselines and std to ensure batch size is train_prompts_size
                filtered_repeated_batch = filtered_repeated_batch.slice(
                    0, train_prompts_size
                )

    batch_to_return = (
        filtered_repeated_batch
        if master_config.grpo.use_dynamic_sampling
        else repeated_batch
    )
    return batch_to_return, is_batch_complete, batch_cache, dynamic_sampling_metrics


def scale_rewards(
    repeated_batch: BatchedDataDict[DatumSpec], reward_scaling_cfg: RewardScalingConfig
) -> BatchedDataDict[DatumSpec]:
    """Linearly scales rewards from a source range to a target range.

    If `reward_scaling.enabled` is True, each reward in `repeated_batch["total_reward"]`
    is clamped to the configured source interval [source_min, source_max] and then
    rescaled to the target interval [target_min, target_max].

    Default configuration:
        source_min = 0.0
        source_max = 1.0
        target_min = 0.0
        target_max = 1.0
    """
    if reward_scaling_cfg.enabled:
        rewards = repeated_batch["total_reward"]
        source_min = float(reward_scaling_cfg.source_min)
        source_max = float(reward_scaling_cfg.source_max)
        target_min = float(reward_scaling_cfg.target_min)
        target_max = float(reward_scaling_cfg.target_max)

        # Detect out-of-range values
        out_of_range_mask = (rewards < source_min) | (rewards > source_max)
        if torch.any(out_of_range_mask):
            print(
                f"[reward_scaling] WARNING: {int(out_of_range_mask.sum())} rewards "
                f"are outside the configured source range [{source_min}, {source_max}]. "
                f"Values will be clipped before scaling."
            )

        # Clamp and scale
        def _scale(reward_tensor: torch.Tensor) -> torch.Tensor:
            r = torch.clamp(reward_tensor, min=source_min, max=source_max)
            return target_min + (r - source_min) / (source_max - source_min) * (
                target_max - target_min
            )

        scaled_rewards = _scale(rewards)
        repeated_batch["total_reward"] = scaled_rewards
        for key in get_gdpo_reward_component_keys(repeated_batch):
            repeated_batch[key] = _scale(repeated_batch[key])

    return repeated_batch


def extract_initial_prompt_messages(
    message_logs: list,
    original_prompt_lengths: torch.Tensor,
) -> list:
    """Extract the original prompt messages from message logs using token length.

    This function correctly identifies original prompt messages even when the prompt
    contains assistant messages (e.g., multi-turn conversation history).

    Args:
        message_logs: List of message logs, where each log is a list of messages.
        original_prompt_lengths: Tensor of original prompt token lengths per sample.

    Returns:
        List of message logs containing only the original prompt messages.
    """
    initial_prompt_message_logs = []
    for i, message_log in enumerate(message_logs):
        initial_prompt_log = []
        cumulative_length = 0
        target_length = original_prompt_lengths[i].item()

        for message in message_log:
            if cumulative_length >= target_length:
                break
            initial_prompt_log.append(message)
            cumulative_length += len(message["token_ids"])

        initial_prompt_message_logs.append(initial_prompt_log)

    return initial_prompt_message_logs


def add_grpo_token_loss_masks_and_generation_logprobs(
    message_logs: list[LLMMessageLogType | VLMMessageLogType],
) -> None:
    """Add GRPO loss masks and ensure generation logprobs exist in message logs.

    Assistant messages can be part of the original multi-turn prompt history. Only
    generated assistant messages have generation_logprobs, so use that field as the
    trainable-token marker. This function mutates each message in-place by adding a
    token_loss_mask and, when missing, a zero-valued generation_logprobs tensor.
    Batches with at least one non-empty media segment also receive a per-token
    media-validity field. Exact masks minted by the multimodal processor are
    preserved; legacy messages fall back to PackedTensor ownership. Every other
    message, including generated assistant text, is invalid. Batches without
    media segments carry no media-validity field.
    Router-replay routes get the same treatment via
    :func:`backfill_missing_routed_experts`, so every per-token field is defined
    for every tokenized message before the batch is flattened.

    Args:
        message_logs: Batch of tokenized message logs. Each message must contain a
            ``role`` and ``token_ids`` field. Messages that already contain
            ``generation_logprobs`` are treated as rollout-generated messages.
    """
    backfill_missing_routed_experts(message_logs)

    def _owns_media(message: dict[str, Any]) -> bool:
        return any(
            key in PACKED_MULTIMODAL_FIELDS
            and isinstance(value, PackedTensor)
            and any(value.logical_segment_counts_by_row())
            for key, value in message.items()
        )

    # Only add masks if some message has non-empty media. Overlength messages keep
    # empty media, which is not stored; storing their mask anyway would make
    # reading them together with text-only messages fail on the missing mask.
    has_media = any(
        _owns_media(message) for message_log in message_logs for message in message_log
    )
    for message_log in message_logs:
        for message in message_log:
            role = cast(str, message["role"])
            token_ids = cast(torch.Tensor, message["token_ids"])

            if role == "assistant" and "generation_logprobs" in message:
                message["token_loss_mask"] = torch.ones_like(token_ids)
            else:
                message["token_loss_mask"] = torch.zeros_like(token_ids)

            if "generation_logprobs" not in message:
                message["generation_logprobs"] = torch.zeros_like(
                    token_ids, dtype=torch.float32
                )

            if not has_media:
                message.pop("media_token_validity_mask", None)
                continue
            existing_media_mask = message.get("media_token_validity_mask")
            if isinstance(existing_media_mask, torch.Tensor):
                message["media_token_validity_mask"] = existing_media_mask.bool()
            else:
                message["media_token_validity_mask"] = torch.full_like(
                    token_ids,
                    fill_value=_owns_media(message),
                    dtype=torch.bool,
                )


def _resolve_message_level_advantage_penalties(
    master_config: MasterConfig,
) -> tuple[float | None, float | None]:
    """Return configured message-level penalties and validate feature support."""
    invalid_tool_call_advantage = master_config.grpo.invalid_tool_call_advantage
    malformed_thinking_advantage = master_config.grpo.malformed_thinking_advantage
    if invalid_tool_call_advantage is None and malformed_thinking_advantage is None:
        return invalid_tool_call_advantage, malformed_thinking_advantage

    # The is_invalid_tool_call / has_malformed_thinking flags these penalties rely on
    # are only populated by the NeMo-Gym environment. Without that path the penalties
    # would silently no-op, so fail loudly instead.
    if not should_use_nemo_gym(master_config):
        raise ValueError(
            "grpo.invalid_tool_call_advantage / grpo.malformed_thinking_advantage require "
            "the NeMo-Gym path (env.should_use_nemo_gym=true); they are not supported with "
            "the native generation path."
        )
    return invalid_tool_call_advantage, malformed_thinking_advantage


def _raise_if_reward_penalties_enabled_without_nemo_gym(
    master_config: MasterConfig,
    *,
    enable_nemo_gym: bool,
) -> None:
    """Validate reward-zeroing penalties are only used with NeMo-Gym."""
    if enable_nemo_gym:
        return

    if not any(
        getattr(master_config.reward_penalties, flag) for flag in _REWARD_PENALTY_FLAGS
    ):
        return

    raise ValueError(
        "reward_penalties require the NeMo-Gym path "
        "(env.should_use_nemo_gym=true); they are not supported with the native "
        "generation path."
    )


def _raise_if_length_penalty_enabled_without_nemo_gym(
    master_config: MasterConfig,
    *,
    enable_nemo_gym: bool,
) -> None:
    """Validate grpo.length_penalty is only used with NeMo-Gym.

    The hook lives in the NeMo-Gym rollout postprocessor, so a native-path run
    with the block set would train on unadjusted rewards without any error.
    """
    if enable_nemo_gym or master_config.grpo.length_penalty is None:
        return

    raise ValueError(
        "grpo.length_penalty requires the NeMo-Gym path "
        "(env.should_use_nemo_gym=true); it is not supported with the native "
        "generation path."
    )


def _apply_message_level_advantage_penalties(
    train_data: BatchedDataDict[ClippedPGLossDataDict],
    message_logs: list[LLMMessageLogType | VLMMessageLogType],
    invalid_tool_call_advantage: float | None,
    malformed_thinking_advantage: float | None,
    log_config: bool = False,
) -> Optional[dict[str, float]]:
    """Overwrite advantages for flagged assistant-message token spans.

    For each assistant message flagged by the NeMo-Gym detector as an invalid
    tool call or malformed thinking, overwrite that message's advantage span in
    ``train_data["advantages"]`` with the configured negative value. No-op when
    neither ``grpo.invalid_tool_call_advantage`` nor
    ``grpo.malformed_thinking_advantage`` is set.

    Args:
        train_data: Training batch; ``advantages`` is modified in place.
        message_logs: Batch of message logs with per-message flags.
        invalid_tool_call_advantage: Advantage value assigned to invalid tool calls.
        malformed_thinking_advantage: Advantage value assigned to malformed thinking.
        log_config: If True, print the configured penalty values once.

    Returns:
        Dictionary of penalty metrics if penalties are applied, otherwise None.
    """
    penalty_metrics = {}
    invalid_neg_adv = invalid_tool_call_advantage
    malformed_neg_adv = malformed_thinking_advantage
    if invalid_neg_adv is None and malformed_neg_adv is None:
        return penalty_metrics

    if log_config:
        print(
            f"Invalid tool call advantage: {invalid_neg_adv}",
            flush=True,
        )
        print(
            f"Malformed thinking advantage: {malformed_neg_adv}",
            flush=True,
        )

    advantages = train_data["advantages"]
    materialized_advantages = False
    num_invalid_tool_calls = 0
    num_malformed_thinking = 0
    num_assistant_messages = 0

    for i, message_log in enumerate(message_logs):
        token_offset = 0
        for j, message in enumerate(message_log):
            token_ids = cast(torch.Tensor, message["token_ids"])
            msg_len = len(token_ids)
            is_assistant = (
                message["role"] == "assistant" and "generation_logprobs" in message
            )
            if is_assistant:
                num_assistant_messages += 1
            is_invalid = (
                is_assistant
                and invalid_neg_adv is not None
                and message.get("is_invalid_tool_call", False)
            )
            is_malformed_thinking = (
                is_assistant
                and malformed_neg_adv is not None
                and message.get("has_malformed_thinking", False)
            )
            if (is_invalid or is_malformed_thinking) and not materialized_advantages:
                # GRPO/GDPO may expand per-sample advantages into zero-stride views;
                # clone before span writes so penalties only affect targeted tokens.
                advantages = advantages.clone()
                train_data["advantages"] = advantages
                materialized_advantages = True

            if is_invalid:
                num_invalid_tool_calls += 1
                print(
                    f"Setting negative advantage ({invalid_neg_adv}) for invalid tool call in assistant message {i} {j}",
                    flush=True,
                )
                advantages[i, token_offset : token_offset + msg_len] = invalid_neg_adv
            elif is_malformed_thinking:
                num_malformed_thinking += 1
                print(
                    f"Setting negative advantage ({malformed_neg_adv}) for malformed thinking in assistant message {i} {j}",
                    flush=True,
                )
                advantages[i, token_offset : token_offset + msg_len] = malformed_neg_adv
            token_offset += msg_len

    invalid_tool_call_rate = num_invalid_tool_calls / max(num_assistant_messages, 1)
    malformed_thinking_rate = num_malformed_thinking / max(num_assistant_messages, 1)
    print(
        f"Invalid tool call rate: {invalid_tool_call_rate:.4f} ({num_invalid_tool_calls}/{num_assistant_messages})",
        flush=True,
    )
    print(
        f"Malformed thinking rate: {malformed_thinking_rate:.4f} ({num_malformed_thinking}/{num_assistant_messages})",
        flush=True,
    )
    penalty_metrics["invalid_tool_call_rate"] = invalid_tool_call_rate
    penalty_metrics["malformed_thinking_rate"] = malformed_thinking_rate
    penalty_metrics["num_invalid_tool_calls"] = num_invalid_tool_calls
    penalty_metrics["num_malformed_thinking"] = num_malformed_thinking
    penalty_metrics["num_assistant_messages"] = num_assistant_messages
    return penalty_metrics


def _apply_configured_message_level_advantage_penalties(
    train_data: BatchedDataDict[ClippedPGLossDataDict],
    message_logs: list[LLMMessageLogType | VLMMessageLogType],
    master_config: MasterConfig,
    log_config: bool = False,
) -> Optional[dict[str, float]]:
    """Resolve config and apply message-level advantage penalties."""
    (
        invalid_tool_call_advantage,
        malformed_thinking_advantage,
    ) = _resolve_message_level_advantage_penalties(master_config)
    return _apply_message_level_advantage_penalties(
        train_data=train_data,
        message_logs=message_logs,
        invalid_tool_call_advantage=invalid_tool_call_advantage,
        malformed_thinking_advantage=malformed_thinking_advantage,
        log_config=log_config,
    )


def _preserve_router_replay_routed_experts(
    target: BatchedDataDict,
    flat_messages: BatchedDataDict,
    policy_config: PolicyConfig,
) -> None:
    """Carry rollout-recorded routes into policy worker inputs when R3 is enabled."""
    if router_replay_enabled(policy_config) and "routed_experts" in flat_messages:
        target["routed_experts"] = flat_messages["routed_experts"]


def _policy_dtype(policy_config: PolicyConfig) -> torch.dtype:
    """Resolve the configured policy precision to its matching torch dtype."""
    return getattr(torch, policy_config["precision"])


def _build_async_grpo_train_data(
    flat_messages: BatchedDataDict,
    input_lengths: torch.Tensor,
    repeated_batch: BatchedDataDict,
    policy_config: PolicyConfig,
) -> BatchedDataDict[ClippedPGLossDataDict]:
    """Build the async no-TQ policy train batch from flattened rollout messages."""
    train_data = BatchedDataDict[ClippedPGLossDataDict](
        {
            "input_ids": flat_messages["token_ids"],
            "input_lengths": input_lengths,
            "generation_logprobs": flat_messages["generation_logprobs"],
            "token_mask": flat_messages["token_loss_mask"],
            "sample_mask": repeated_batch["loss_multiplier"],
        }
    )
    _preserve_router_replay_routed_experts(train_data, flat_messages, policy_config)
    # update multimodal data unconditionally
    extra_multimodal_data = flat_messages.get_multimodal_dict(
        as_tensors=False, pixel_dtype=_policy_dtype(policy_config)
    )
    train_data.update(extra_multimodal_data)
    return train_data


# Kept in the batch so row selection and dynamic-sampling caches preserve it.
PROMPT_GROUP_IDS_KEY = "prompt_group_ids"


def _prompt_group_ids(
    num_rows: int, num_generations_per_prompt: int, *, group_offset: int = 0
) -> torch.Tensor:
    """Name each repeated input prompt with a ``[B, 1]`` integer key.

    Assign before rollout, which can rewrite the message history. A cumulative
    offset keeps groups from successive dynamic-sampling batches disjoint.
    """
    if num_generations_per_prompt <= 0 or num_rows % num_generations_per_prompt:
        raise ValueError(
            f"repeated batch of {num_rows} rows is not a whole number of prompt "
            f"groups of {num_generations_per_prompt}"
        )
    num_prompts = num_rows // num_generations_per_prompt
    return (
        (group_offset + torch.arange(num_prompts, dtype=torch.long))
        .repeat_interleave(num_generations_per_prompt)
        .reshape(-1, 1)
    )


def _trajectory_group_ids(trajectories: list[dict[str, Any]]) -> torch.Tensor:
    """Name rows by their per-prompt replay entry in concatenation order."""
    if not trajectories:
        return torch.empty((0, 1), dtype=torch.long)
    return torch.cat(
        [
            torch.full((trajectory["batch"].size, 1), position, dtype=torch.long)
            for position, trajectory in enumerate(trajectories)
        ]
    )


def _apply_mask_sample_filter(repeated_batch: BatchedDataDict[DatumSpec]) -> int:
    """Zero loss_multiplier where mask_sample is True and return the count."""
    if "mask_sample" not in repeated_batch:
        return 0

    loss_multiplier = repeated_batch["loss_multiplier"].clone()
    mask_sample = repeated_batch["mask_sample"]

    if isinstance(mask_sample, list):
        mask_sample = torch.tensor(mask_sample, dtype=torch.bool)
    mask_sample_bool = mask_sample.bool()

    num_masked = int(mask_sample_bool.sum().item())
    loss_multiplier[mask_sample_bool] = 0
    repeated_batch["loss_multiplier"] = loss_multiplier
    return num_masked


def _should_log_nemo_gym_responses(master_config: MasterConfig) -> bool:
    """Whether NeMo Gym is responsible for full response logging.

    When **True**, skip the expensive per-step ``train_data_step*.jsonl`` dump.
    When **False** (the default if unset), write the local JSONL file.

    W&B full-result Tables are controlled independently by
    ``logger.wandb.log_nemo_gym_full_result_tables``.
    """
    env_config = master_config.env
    should_log_nemo_gym_responses = bool(
        env_config.get("should_log_nemo_gym_responses")
    )

    return should_log_nemo_gym_responses


def _write_latest_checkpoint_status(
    checkpointer: CheckpointManager, last_checkpoint_step: int
) -> None:
    """Write a lightweight, top-level ``latest_checkpoint_status.json`` for monitoring.

    Records the wall-clock time and step of the most recent successful checkpoint
    save so an out-of-band watchdog can poll checkpoint progress on long runs.

    Intentionally distinct from ``CheckpointManager``'s per-step
    ``step_{N}/training_info.json`` (the resume state): different schema, written
    at the checkpoint-dir root. There is no in-repo consumer yet. The read is
    deliberately unguarded so a corrupt file surfaces loudly (signalling
    corruption) instead of being silently masked.
    """
    status_path = os.path.join(
        str(checkpointer.checkpoint_dir), "latest_checkpoint_status.json"
    )
    status: dict[str, Any] = {}
    if os.path.exists(status_path):
        with open(status_path) as f:
            status = json.load(f)
    status["last_successful_ckpt_save_completion"] = time.time()
    status["last_checkpoint_step"] = last_checkpoint_step
    with open(status_path, "w") as f:
        json.dump(status, f)


def _get_effort_config(master_config: MasterConfig) -> Optional[EffortLevelsConfig]:
    """Return the effort-levels reward-shaping config from env.nemo_gym, if set."""
    return get_effort_config(master_config.env)


def _pad_teacher_logprobs(teacher_logprobs: torch.Tensor, train_S: int) -> torch.Tensor:
    """Right-zero-pad teacher logprobs ``[B, teacher_S]`` to ``train_S``.

    ``from_batches`` pads teacher logprobs to ``max(S_i)``; ``train_data`` may be
    longer due to ``make_sequence_length_divisible_by``. Zero-pad is safe because
    the mask zeros padding in advantage computation. ``teacher_S > train_S`` is
    unexpected (teacher pads to a finer grid than the student) and raises.
    """
    teacher_S = teacher_logprobs.shape[1]
    if teacher_S > train_S:
        raise ValueError(
            f"Teacher logprobs seq length ({teacher_S}) > train_data seq length ({train_S}). "
            "Teacher logprobs are padded to max(S_i) by from_batches, "
            "and train_data is padded to roundup(max(S_i), make_sequence_length_divisible_by)."
        )
    if teacher_S < train_S:
        teacher_logprobs = torch.nn.functional.pad(
            teacher_logprobs, (0, train_S - teacher_S), value=0.0
        )
    return teacher_logprobs


def _create_advantage_estimator(master_config: MasterConfig):
    """Create and return an advantage estimator based on configuration.

    Args:
        master_config: The master configuration dictionary.

    Returns:
        An advantage estimator instance (GRPO, GDPO, or ReinforcePlusPlus).

    Raises:
        ValueError: If the advantage estimator name is not recognized.
    """
    grpo_config = master_config.grpo
    loss_config = master_config.loss_fn

    adv_estimator_config = grpo_config.adv_estimator

    adv_estimator_name = adv_estimator_config.name
    if adv_estimator_name == "gdpo":
        adv_estimator = GDPOAdvantageEstimator(adv_estimator_config, loss_config)
        print("  ✓ Using GDPO advantage estimator (multi-reward)")
    elif adv_estimator_name == "grpo":
        adv_estimator = GRPOAdvantageEstimator(adv_estimator_config, loss_config)
        print("  ✓ Using GRPO advantage estimator")
    elif adv_estimator_name == "opd":
        opd_module.assert_prev_logprobs_available(master_config)
        opd_module.assert_trust_region_supported(master_config)
        adv_estimator = OPDAdvantageEstimator(adv_estimator_config, loss_config)
        print("  ✓ Using OPD advantage estimator")
        # Warn if loss_fn is not configured per MOPD paper recommendations.
        if not loss_config.disable_ppo_ratio:
            warnings.warn(
                "OPD recommends loss_fn.disable_ppo_ratio: true (REINFORCE-style, MOPD Eq. 7)"
            )
        if not loss_config.use_importance_sampling_correction:
            warnings.warn(
                "OPD recommends loss_fn.use_importance_sampling_correction: true (MOPD Eq. 8 w_t)"
            )
        if loss_config.truncated_importance_sampling_type != "icepop":
            warnings.warn(
                "OPD recommends loss_fn.truncated_importance_sampling_type: 'icepop' "
                "(hard gate, MOPD Eq. 8)"
            )
    elif adv_estimator_name == "reinforce_plus_plus":
        adv_estimator = ReinforcePlusPlusAdvantageEstimator(
            adv_estimator_config, loss_config
        )
        print("  ✓ Using Reinforce++ advantage estimator")
    else:
        raise ValueError(f"Invalid adv_estimator name: {adv_estimator_name}")

    return adv_estimator


def _clip_grpo_advantages(
    advantages: torch.Tensor,
    grpo_config: GRPOConfig,
) -> torch.Tensor:
    """Clamp normalized advantages when clip bounds are configured."""
    clip_low = grpo_config.advantage_clip_low
    clip_high = grpo_config.advantage_clip_high
    if clip_low is not None:
        advantages = advantages.clamp(min=clip_low)
    if clip_high is not None:
        advantages = advantages.clamp(max=clip_high)
    return advantages


def refit_policy_generation(
    policy: ColocatablePolicyInterface,
    policy_generation: GenerationInterface,
    colocated_inference: bool,
    _refit_buffer_size_gb: Optional[float] = None,
    timer: Optional[Timer] = None,
    kv_scales: Optional[dict[str, float]] = None,
) -> dict[str, float]:
    """Refit the policy generation interface with the latest policy weights.

    Args:
        policy: The policy to provide weights to the inference engine.
        policy_generation: The inference engine to refit.
        _refit_buffer_size_gb: Fixed refit buffer size in GiB. If it is None,
            the buffer size is computed from remaining memory.
        timer: Optional Timer used to time the prepare/transfer/update phase
        kv_scales: Optional dictionary of KV cache scales for FP8 quantization.

    Returns:
        Scalar metrics reported by the selected weight synchronizer.
    """
    synchronizer = getattr(policy_generation, "weight_synchronizer", None)
    if isinstance(policy_generation, SGLangGeneration) and synchronizer is None:
        # Fail loudly rather than falling through to the vLLM branches, which
        # would call methods the SGLang path does not implement.
        raise RuntimeError(
            "SGLang refits require policy_generation.weight_synchronizer to be "
            "set. Attach one with create_weight_synchronizer(...) during setup."
        )

    # Materialize deferred Megatron parameter all-gathers before any transport
    # reads policy weights, including synchronizers that return early below.
    sync_context = (
        timer.time("prepare_for_generation/sync_policy_params")
        if timer is not None
        else nullcontext()
    )
    with sync_context:
        policy.sync_params_before_refit()

    if synchronizer is not None:
        return synchronizer.sync_weights(timer=timer, kv_scales=kv_scales) or {}

    if colocated_inference:
        policy.offload_before_refit()
        policy_generation.prepare_for_generation(tags=["weights"])

    # Create a context manager that does nothing when timer is None
    timer_context = (
        timer.time("prepare_for_generation/transfer_and_update_weights")
        if timer is not None
        else nullcontext()
    )
    with timer_context:
        # update weights
        update_success = False
        if colocated_inference:
            # get model param keys, which is grouped by size
            if _refit_buffer_size_gb is not None:
                buffer_size_bytes = int(_refit_buffer_size_gb * (1024**3))
            else:
                # Empirically sets ratio as 30% to maximize efficiency.
                # The remaining 70% is a necessary buffer reserved for the parameter all-gathering across the expert-parallelism dimension.
                memory_ratio = os.getenv("NRL_REFIT_BUFFER_MEMORY_RATIO", "0.3")
                buffer_size_bytes = int(
                    policy.get_free_memory_bytes() * float(memory_ratio)
                )

            # ZMQ IPC path: shared by vLLM and TRT-LLM colocated. Trainer
            # streams CUDA IPC handles in chunks; receiver reconstructs
            # tensors in-place and feeds them into the inference engine's
            # loader.
            futures_train = policy.stream_weights_via_ipc_zmq(
                buffer_size_bytes=buffer_size_bytes,
                kv_scales=kv_scales,
            )
            futures_inference = policy_generation.update_weights_via_ipc_zmq()
            # wait for all futures to complete
            ray.get(futures_train)
            results = ray.get(futures_inference)
            update_success = all(result for result in results if result is not None)
        else:
            # update weights through nccl (vLLM)
            futures_train = policy.broadcast_weights_for_collective(
                kv_scales=kv_scales,
            )
            futures_inference = policy_generation.update_weights_from_collective()
            # wait for all futures to complete
            ray.get(futures_train)
            results = ray.get(futures_inference)
            update_success = all(result for result in results if result is not None)

        # check if update is successful
        if not update_success:
            error_tag = "cuda-ipc" if colocated_inference else "nccl"
            error_message = (
                "❌ Error: Updating weights for the generation policy failed during refit.\n"
                f"This often indicates an issue with {error_tag} or "
                "a problem within the generation backend (e.g., vLLM worker).\n"
            )
            raise RuntimeError(error_message)

    if colocated_inference:
        policy.offload_after_refit()
        policy_generation.prepare_for_generation(tags=["kv_cache"])

    return {}


def _initial_policy_generation_stale(
    policy_generation: GenerationInterface, completed_steps: int
) -> bool:
    """Skip a fresh run's redundant sync when the synchronizer is already current."""
    synchronizer = getattr(policy_generation, "weight_synchronizer", None)
    return completed_steps > 0 or synchronizer is None or synchronizer.is_stale


def _log_mixed_rewards_and_advantages_information(
    logger: Logger,
    total_steps: int,
    metrics: dict[str, Any],
    baseline: torch.Tensor,
    advantages: torch.Tensor,
    env_baseline: Optional[torch.Tensor] = None,
) -> None:
    # The histograms that are logged are logged with a prefix "train/" to the name, since that is what the remaining metrics will be logged with.
    logger.log_histogram(
        baseline.numpy(), total_steps + 1, "train/baseline_reward/histogram"
    )
    # The pct_* diagnostics assume binary rewards. When grpo.length_penalty
    # rewrote total_reward, read them off the env-reward baseline instead.
    pct_baseline = env_baseline if env_baseline is not None else baseline
    metrics["baseline_reward/pct_0"] = 100 * (pct_baseline == 0).float().mean().item()
    metrics["baseline_reward/pct_1"] = 100 * (pct_baseline == 1).float().mean().item()
    metrics["baseline_reward/pct_mixed"] = (
        100 - metrics["baseline_reward/pct_0"] - metrics["baseline_reward/pct_1"]
    )

    logger.log_histogram(
        advantages.numpy(), total_steps + 1, "train/advantages/histogram"
    )
    metrics["advantages/sum"] = advantages.float().sum().item()
    metrics["advantages/mean"] = advantages.float().mean().item()


def _placeholder_seq_logprob_error_metrics() -> dict[str, float]:
    """Zero-valued seq-level metrics used when the prev_logprobs forward is skipped."""
    return {
        "max_seq_mult_prob_error": 0.0,
        "mean_seq_mult_prob_error": 0.0,
        "min_seq_mult_prob_error": 0.0,
        "max_seq_mult_prob_error_after_mask": 0.0,
        "mean_seq_mult_prob_error_after_mask": 0.0,
        "min_seq_mult_prob_error_after_mask": 0.0,
        "num_masked_seqs_by_logprob_error": 0,
        "masked_correct_pct": 0.0,
    }


def _validate_use_kl_in_reward_compat(master_config: MasterConfig) -> None:
    """Reject ``use_kl_in_reward`` when the KL term would read zero placeholder logprobs.

    ``force_on_policy_ratio`` (without ``seq_logprob_error_threshold``) skips
    the prev_logprobs forward and passes a zero placeholder to the advantage
    estimator; ``use_kl_in_reward`` then applies
    ``kl_coef * calculate_kl(zeros, ref)`` which corrupts the advantage.
    ``kl_coef=0`` (``reference_policy_kl_penalty=0``) zeros the term regardless,
    so that case is allowed.
    """
    loss_config = master_config.loss_fn
    if loss_config.use_kl_in_reward and loss_config.reference_policy_kl_penalty != 0:
        assert not opd_module._skip_prev_logprobs(master_config), (
            "loss_fn.use_kl_in_reward with nonzero loss_fn.reference_policy_kl_penalty "
            "requires real prev_logprobs, but force_on_policy_ratio (without "
            "grpo.seq_logprob_error_threshold) zeros them — KL would be computed "
            "against a zero placeholder."
        )


def _resolve_logprob_skip_flags(
    master_config: MasterConfig,
) -> tuple[bool, bool | None]:
    """Return (skip_prev_logprobs, skip_reference_logprobs); warn on incompatible combos.

    Skip prev_logprobs when force_on_policy_ratio=True unless the sequence
    threshold is evaluated before training rather than inside the loss.
    Skip reference_policy_logprobs when
    ``grpo.skip_reference_policy_logprobs_calculation`` is set.
    """
    # todo @jiaqi: is there a better way to skip prev_logprobs computation while still computing the seq-level error metrics?
    if (
        master_config.loss_fn.force_on_policy_ratio
        and master_config.grpo.seq_logprob_error_threshold is not None
        and not master_config.loss_fn.seq_logprob_error_in_loss
    ):
        warnings.warn(
            "force_on_policy_ratio=True but seq_logprob_error_threshold is set. "
            "Computing prev_logprobs anyway for seq-level error masking."
        )
    return (
        opd_module._skip_prev_logprobs(master_config),
        master_config.grpo.skip_reference_policy_logprobs_calculation,
    )


def compute_and_apply_seq_logprob_error_masking(
    train_data: BatchedDataDict,
    rewards: torch.Tensor,
    seq_logprob_error_threshold: Optional[float],
    *,
    filtering_on: bool = False,
) -> dict:
    """Compute sequence-level logprob error metrics and optionally mask high-error sequences.

    This function computes the multiplicative probability error per sequence
    (same calculation as token_mult_prob_error but aggregated per-sequence) and
    optionally masks sequences that exceed the configured threshold.

    Args:
        train_data: Training data dict containing token_mask, sample_mask,
                   prev_logprobs, and generation_logprobs. If masking is applied,
                   sample_mask will be updated in-place.
        rewards: Reward tensor for computing statistics on masked sequences.
        seq_logprob_error_threshold: If set, mask sequences with mult_prob_error
                                    exceeding this threshold. If None, only compute metrics.
        filtering_on: Whether policy top-k/top-p can produce unsupported -inf tokens.

    Returns:
        Dict with keys: max_seq_mult_prob_error, mean_seq_mult_prob_error,
        min_seq_mult_prob_error, max/mean/min_seq_mult_prob_error_after_mask,
        num_masked_seqs, masked_correct_pct
    """
    # Compute sequence-level logprob error metrics (always)
    token_mask = train_data["token_mask"][:, 1:]
    sample_mask = train_data["sample_mask"]
    prev_logprobs = train_data["prev_logprobs"][:, 1:]
    generation_logprobs = train_data["generation_logprobs"][:, 1:]
    if filtering_on:
        # Narrow only the diagnostic mask; the shared mask still owns KL terms.
        token_mask = token_mask * (
            ~torch.isneginf(prev_logprobs) & ~torch.isneginf(generation_logprobs)
        ).to(token_mask.dtype)
    seq_mult_prob_error, valid_seq_mask = compute_seq_logprob_errors(
        policy_logprobs=prev_logprobs,
        generation_logprobs=generation_logprobs,
        token_mask=token_mask,
        sample_mask=sample_mask,
    )
    if valid_seq_mask.any():
        valid_errors = seq_mult_prob_error[valid_seq_mask]
        max_seq_mult_prob_error = valid_errors.max().item()
        mean_seq_mult_prob_error = valid_errors.mean().item()
        min_seq_mult_prob_error = valid_errors.min().item()
    else:
        max_seq_mult_prob_error = 0.0
        mean_seq_mult_prob_error = 0.0
        min_seq_mult_prob_error = 0.0

    # Apply sequence-level masking if configured
    num_masked_seqs = 0
    masked_correct_pct = 0.0

    # After-mask metrics (same as before if no threshold)
    max_seq_mult_prob_error_after_mask = max_seq_mult_prob_error
    mean_seq_mult_prob_error_after_mask = mean_seq_mult_prob_error
    min_seq_mult_prob_error_after_mask = min_seq_mult_prob_error

    if seq_logprob_error_threshold is not None:
        print(
            f"▶ Applying sequence-level logprob error masking (threshold={seq_logprob_error_threshold})...",
            flush=True,
        )

        original_sample_mask = sample_mask.clone()

        # Create mask for sequences below threshold
        seq_error_mask = (
            seq_mult_prob_error <= seq_logprob_error_threshold
        ).float() * original_sample_mask

        diff_mask = original_sample_mask - seq_error_mask
        num_masked_seqs = int(diff_mask.sum().item())

        if num_masked_seqs > 0:
            diff_mask_bool = diff_mask.bool()
            masked_correct_count = int(
                (rewards.view(-1)[diff_mask_bool] == 1).sum().item()
            )
            masked_correct_pct = masked_correct_count / num_masked_seqs

        # Compute after-mask metrics (only for sequences that passed the threshold)
        kept_mask = seq_error_mask.bool() & valid_seq_mask
        if kept_mask.sum() > 0:
            kept_errors = seq_mult_prob_error[kept_mask]
            max_seq_mult_prob_error_after_mask = kept_errors.max().item()
            mean_seq_mult_prob_error_after_mask = kept_errors.mean().item()
            min_seq_mult_prob_error_after_mask = kept_errors.min().item()
        else:
            # All sequences were masked
            max_seq_mult_prob_error_after_mask = 0.0
            mean_seq_mult_prob_error_after_mask = 0.0
            min_seq_mult_prob_error_after_mask = 0.0

        # Update sample_mask in train_data
        train_data["sample_mask"] = seq_error_mask

        print(
            f"  Masked {num_masked_seqs} sequences with mult_prob_error > {seq_logprob_error_threshold}",
            flush=True,
        )
        if num_masked_seqs > 0:
            print(
                f"  • {masked_correct_count}/{num_masked_seqs} masked sequences were correct (reward=1)"
                f" → {masked_correct_pct:.2%}",
                flush=True,
            )

    return {
        "max_seq_mult_prob_error": max_seq_mult_prob_error,
        "mean_seq_mult_prob_error": mean_seq_mult_prob_error,
        "min_seq_mult_prob_error": min_seq_mult_prob_error,
        "max_seq_mult_prob_error_after_mask": max_seq_mult_prob_error_after_mask,
        "mean_seq_mult_prob_error_after_mask": mean_seq_mult_prob_error_after_mask,
        "min_seq_mult_prob_error_after_mask": min_seq_mult_prob_error_after_mask,
        "num_masked_seqs": num_masked_seqs,
        "masked_correct_pct": masked_correct_pct,
    }


# ===============================================================================
# Training & Validation
# ===============================================================================


def _validation_stop_value(val_metrics: dict[str, Any], stop_metric: str) -> float:
    """Value of the early-stop metric chosen by grpo.stop_at_validation_metric."""
    assert stop_metric in val_metrics, (
        f"grpo.stop_at_validation_metric={stop_metric!r} is not a reported "
        f"validation metric; available: {sorted(val_metrics)}"
    )
    return val_metrics[stop_metric]


def _validation_early_stop_message(
    val_metrics: dict[str, Any],
    stop_threshold: float | None,
    stop_metric: str | None,
    *,
    initial: bool = False,
) -> Optional[str]:
    """Stop message when the early-stop threshold is reached, else None."""
    if stop_metric is None:
        return None
    # setup() guards this pairing at startup; keep the invariant visible here.
    assert stop_threshold is not None, (
        "grpo.stop_at_validation_threshold must be set when "
        "grpo.stop_at_validation_metric is set"
    )
    value = _validation_stop_value(val_metrics, stop_metric)
    if value < stop_threshold:
        return None
    prefix = "Initial validation" if initial else "Validation"
    return (
        f"{prefix} {stop_metric} reached the early-stop threshold "
        f"({value:.4f} >= {stop_threshold}); stopping training"
    )


def _grpo_train_impl(
    policy: ColocatablePolicyInterface,
    policy_generation: Optional[GenerationInterface],
    wrapped_dataloader: StatefulDataLoader | MultipleDataloaderWrapper,
    val_dataloader: Optional[StatefulDataLoader],
    tokenizer: TokenizerType,
    loss_fn: LossFunction,
    task_to_env: dict[str, EnvironmentInterface],
    val_task_to_env: Optional[dict[str, EnvironmentInterface]],
    logger: Logger,
    checkpointer: CheckpointManager,
    grpo_save_state: GRPOSaveState,
    master_config: MasterConfig,
    processor: Optional[AutoProcessor] = None,
) -> None:
    """Run GRPO training algorithm."""
    timer = Timer(context={"worker": "driver"})
    _telemetry = get_telemetry_handle()
    _tracer = _telemetry.tracer if _telemetry is not None else None
    timeout = TimeoutChecker(
        timeout=master_config.checkpointing["checkpoint_must_save_by"],
        fit_last_save_time=True,
    )
    timeout.start_iterations()
    memory_tracker = MemoryTracker()

    kv_scales_cache = None  # Cache reused for computed kv scales

    assert policy_generation is not None

    # Check if we need to sync KV cache scales
    # When fallback to policy as the policy_generation, we use getattr to check.
    sync_kv_scales = getattr(policy_generation, "requires_kv_scale_sync", False)

    # common config/state times
    current_step = grpo_save_state.current_step  # current step within an epoch
    total_steps = grpo_save_state.total_steps  # total steps across all epochs
    POLICY_GENERATION_STALE = _initial_policy_generation_stale(
        policy_generation, total_steps
    )
    max_num_steps = master_config.grpo.max_num_steps  # max number of steps to train for
    current_epoch = grpo_save_state.current_epoch  # current epoch
    max_num_epochs = (
        master_config.grpo.max_num_epochs
    )  # max number of epochs to train for
    consumed_samples = (
        grpo_save_state.consumed_samples
    )  # total samples consumed across all epochs
    total_valid_tokens = (
        grpo_save_state.total_valid_tokens
    )  # total valid tokens processed across all epochs
    val_at_start = master_config.grpo.val_at_start
    val_at_end = master_config.grpo.val_at_end
    val_period = master_config.grpo.val_period
    val_start_at = master_config.grpo.val_start_at
    colocated_inference = master_config.policy["generation"]["colocated"]["enabled"]
    refit_buffer_size_gb = master_config.policy.get("refit_buffer_size_gb")
    stop_at_validation_threshold = master_config.grpo.stop_at_validation_threshold
    stop_at_validation_metric = master_config.grpo.stop_at_validation_metric

    # Initialize advantage estimator
    adv_estimator = _create_advantage_estimator(master_config)

    # Run validation at the start if configured
    # TODO: Add validation with kv scales if needed
    if val_at_start and current_step == 0:
        print("\n🔍 Running initial validation...", flush=True)
        memory_tracker.snapshot_start_of_stage("Initial validation", dir())

        if POLICY_GENERATION_STALE:
            refit_policy_generation(
                policy,
                policy_generation,
                colocated_inference,
                _refit_buffer_size_gb=refit_buffer_size_gb,
            )
            POLICY_GENERATION_STALE = False
        else:
            policy_generation.prepare_for_generation()
        val_metrics, validation_timings = validate(
            policy_generation,
            val_dataloader,
            tokenizer,
            val_task_to_env,
            step=0,
            master_config=master_config,
            logger=logger,
            processor=processor,
        )
        policy_generation.finish_generation()
        logger.log_metrics(val_metrics, current_step, prefix="validation")
        logger.log_metrics(validation_timings, current_step, prefix="timing/validation")
        if master_config.grpo.debug_payload_metrics:
            validation_payload_metrics = drain_multimodal_payload_metrics()
            if validation_payload_metrics:
                logger.log_metrics(
                    validation_payload_metrics,
                    current_step,
                    prefix="validation",
                )
        stop_message = _validation_early_stop_message(
            val_metrics,
            stop_at_validation_threshold,
            stop_at_validation_metric,
            initial=True,
        )
        if stop_message is not None:
            print(stop_message, flush=True)
            # Flush pending checkpoint finalization, like the other early returns.
            checkpointer.shutdown()
            return

    if master_config.data["use_multiple_dataloader"]:
        warnings.warn(
            "When using multiple dataloaders, MultipleDataloaderWrapper operates as an infinite iterator. "
            "As a result, grpo.max_num_epochs will be ignored, and only grpo.max_num_steps will be used. "
            "See https://github.com/NVIDIA-NeMo/RL/blob/main/docs/guides/grpo.md#multiple-dataloaders for more details."
        )

    ft_save_period = master_config.checkpointing.get("ft_save_period")

    while current_epoch < max_num_epochs and total_steps < max_num_steps:
        memory_tracker.snapshot_start_of_stage("Preparing batch", dir())
        print(f"\n{'=' * 25} Epoch {current_epoch + 1}/{max_num_epochs} {'=' * 25}")
        # batch cache is used for DAPO. We store prompts with non-zero standard deviation in this cache.
        batch_cache: BatchedDataDict[DatumSpec] = None
        # This is the number of batches we processed so far at each step to generate responses whose std is non-zero. Maximum threshold is set by dynamic_sampling_max_gen_batches. Used in the case of dynamic sampling.
        dynamic_sampling_num_gen_batches = 0
        next_prompt_group_id = 0

        # Run grpo/dapo training loop (single-turn)
        for batch in wrapped_dataloader:
            refit_metrics: dict[str, float] = {}
            # A central place to store logging data that won't be deleted until the loop ends
            metrics_logging_data = dict()
            metrics = dict()

            if master_config.data["use_multiple_dataloader"]:
                print(
                    f"\n{'=' * 25} Step {current_step + 1}/{max_num_steps} {'=' * 25}",
                    flush=True,
                )
            else:
                print(
                    f"\n{'=' * 25} Step {current_step + 1}/{min(len(wrapped_dataloader), max_num_steps)} {'=' * 25}",
                    flush=True,
                )

            maybe_gpu_profile_step(policy, total_steps + 1)
            if policy != policy_generation:
                maybe_gpu_profile_step(policy_generation, total_steps + 1)
            val_metrics, validation_timings = None, None

            with (
                iteration_scope(total_steps + 1),
                timer.time("total_step_time"),
                umbrella_span(
                    RLSpanGroup.U_STEP,
                    "rl.grpo.step",
                    tracer=_tracer,
                    **{"rl.epoch": current_epoch + 1},
                ),
            ):
                # Prepare batch
                print("▶ Preparing batch...", flush=True)
                with (
                    timer.time("data_processing"),
                    managed_span(
                        RLSpanGroup.DATA_PROCESSING,
                        "rl.grpo.data_processing",
                        tracer=_tracer,
                    ),
                ):
                    if (
                        master_config.grpo.deduplicate_multimodal_data
                        and should_use_nemo_gym(master_config)
                    ):
                        attach_initial_nemo_gym_image_payloads(
                            batch,
                            processor,
                            env_config=master_config.env,
                        )
                    # Repeat batch items
                    repeated_batch: BatchedDataDict[DatumSpec] = (
                        batch.repeat_interleave(
                            master_config.grpo.num_generations_per_prompt,
                            share_immutable_media=(
                                master_config.grpo.deduplicate_multimodal_data
                            ),
                        )
                    )
                    print_multimodal_payload_metrics(
                        collect_multimodal_payload_metrics(
                            repeated_batch,
                            "prompt_repeat",
                            enabled=master_config.grpo.debug_payload_metrics,
                        )
                    )
                    prompt_group_ids = _prompt_group_ids(
                        repeated_batch.size,
                        master_config.grpo.num_generations_per_prompt,
                        group_offset=next_prompt_group_id,
                    )
                    next_prompt_group_id += (
                        repeated_batch.size
                        // master_config.grpo.num_generations_per_prompt
                    )

                # Generate responses - this updates the LLMMessageLogType in repeated_batch
                memory_tracker.snapshot_start_of_stage("Generation", dir())
                print(
                    f"▶ Generating responses for batch of size {repeated_batch.size}...",
                    flush=True,
                )
                with timer.time("prepare_for_generation/total"):
                    if POLICY_GENERATION_STALE:
                        # Compute KV scales if needed for FP8 quantization
                        if sync_kv_scales and kv_scales_cache is None:
                            print("▶ Computing KV cache scales...", flush=True)
                            policy.prepare_for_lp_inference()
                            # Align with training data processing to ensure parallel training compatibility
                            calib_flat, calib_input_lengths = (
                                batched_message_log_to_flat_message(
                                    repeated_batch["message_log"],
                                    pad_value_dict={
                                        "token_ids": tokenizer.pad_token_id
                                    },
                                    make_sequence_length_divisible_by=master_config.policy[
                                        "make_sequence_length_divisible_by"
                                    ],
                                )
                            )
                            # Create calibration data from flattened messages
                            calibration_data = BatchedDataDict[ClippedPGLossDataDict](
                                {
                                    "input_ids": calib_flat["token_ids"],
                                    "input_lengths": calib_input_lengths,
                                }
                            )
                            calibration_data.update(
                                calib_flat.get_multimodal_dict(
                                    as_tensors=False,
                                    pixel_dtype=_policy_dtype(master_config.policy),
                                )
                            )
                            calibration_data.to("cpu")
                            kv_scales_cache = policy.calibrate_qkv_fp8_scales(
                                calibration_data, include_q=True
                            )["layers"]

                        refit_metrics = refit_policy_generation(
                            policy,
                            policy_generation,
                            colocated_inference,
                            _refit_buffer_size_gb=refit_buffer_size_gb,
                            timer=timer,
                            kv_scales=kv_scales_cache if sync_kv_scales else None,
                        )
                        POLICY_GENERATION_STALE = False
                    else:
                        if colocated_inference:
                            policy.offload_after_refit()  # unload optimizer to make space for generation
                        policy_generation.prepare_for_generation()

                dynamic_sampling_num_gen_batches += 1
                if dynamic_sampling_num_gen_batches == 1 and hasattr(
                    policy_generation, "snapshot_step_metrics"
                ):
                    policy_generation.snapshot_step_metrics()
                with (
                    timer.time("generation"),
                    umbrella_span(
                        RLSpanGroup.U_ROLLOUT,
                        "rl.grpo.generation",
                        tracer=_tracer,
                        **{
                            "rl.num_generations_per_prompt": master_config.grpo.num_generations_per_prompt,
                        },
                    ),
                ):
                    # Clear logger metrics for each generation step
                    if policy_generation is not None:
                        policy_generation.clear_logger_metrics()
                    # Use NeMo-Gym rollouts if enabled. We cascade NeMo-Gym first since NeMo-Gym requires async rollouts.
                    if should_use_nemo_gym(master_config):
                        # configure_generation_config auto-fills stop_token_ids from the EOS
                        # token, but run_async_nemo_gym_rollout asserts these are unset because
                        # NeMo-Gym manages its own stop criteria. Clear them here so the
                        # assertion reflects user intent (null in YAML) rather than the auto-fill.
                        generation_config: GenerationConfig = {
                            **master_config.policy["generation"],
                            "stop_token_ids": None,
                            "stop_strings": None,
                        }
                        nemo_gym_rollout_result = run_nemo_gym_rollout_sync(
                            policy_generation=policy_generation,
                            input_batch=repeated_batch,
                            tokenizer=tokenizer,
                            task_to_env=task_to_env,
                            max_seq_len=master_config.policy[
                                "max_total_sequence_length"
                            ],
                            generation_config=generation_config,
                            num_generations_per_prompt=(
                                master_config.grpo.num_generations_per_prompt
                            ),
                            log_full_result_tables=should_log_nemo_gym_full_result_tables(
                                wandb_enabled=master_config.logger.wandb_enabled,
                                wandb_config=master_config.logger.wandb,
                            ),
                            max_rollout_turns=None,
                            greedy=False,
                            effort_config=_get_effort_config(master_config),
                            reward_penalty_config=master_config.reward_penalties,
                            length_penalty_config=master_config.grpo.length_penalty,
                            thinking_tags=get_nemo_gym_thinking_tags(master_config.env),
                            mask_env_flagged_samples=should_mask_flagged_samples(
                                master_config.env
                            ),
                            deduplicate_multimodal_data=(
                                master_config.grpo.deduplicate_multimodal_data
                            ),
                            debug_payload_metrics=(
                                master_config.grpo.debug_payload_metrics
                            ),
                        )
                        repeated_batch = nemo_gym_rollout_result.final_batch
                        rollout_metrics = nemo_gym_rollout_result.rollout_metrics
                        del nemo_gym_rollout_result

                    # Use async rollouts when enabled by config/backend defaults.
                    elif should_use_async_rollouts(master_config.policy["generation"]):
                        (
                            repeated_batch,
                            rollout_metrics,
                        ) = run_async_multi_turn_rollout(
                            policy_generation=policy_generation,
                            input_batch=repeated_batch,
                            tokenizer=tokenizer,
                            task_to_env=task_to_env,
                            max_seq_len=master_config.policy[
                                "max_total_sequence_length"
                            ],
                            max_rollout_turns=master_config.grpo.max_rollout_turns,
                            greedy=False,
                            deduplicate_multimodal_data=(
                                master_config.grpo.deduplicate_multimodal_data
                            ),
                        )
                    else:
                        repeated_batch, rollout_metrics = run_multi_turn_rollout(
                            policy_generation=policy_generation,
                            input_batch=repeated_batch,
                            tokenizer=tokenizer,
                            task_to_env=task_to_env,
                            max_seq_len=master_config.policy[
                                "max_total_sequence_length"
                            ],
                            max_rollout_turns=master_config.grpo.max_rollout_turns,
                            greedy=False,
                            deduplicate_multimodal_data=(
                                master_config.grpo.deduplicate_multimodal_data
                            ),
                        )
                    if prompt_group_ids.shape[0] != repeated_batch.size:
                        raise ValueError(
                            f"rollout returned {repeated_batch.size} rows for "
                            f"{prompt_group_ids.shape[0]} input prompt-group ids"
                        )
                    repeated_batch[PROMPT_GROUP_IDS_KEY] = prompt_group_ids
                    del prompt_group_ids
                    policy_generation.finish_generation()
                    # Collect generation logger metrics for performance reporting after each generation step
                    # inflight batch sizes and num pending samples are collected from each worker
                    if policy_generation is not None:
                        generation_logger_metrics = (
                            policy_generation.get_logger_metrics()
                        )

                    metrics_logging_data[MEAN_GEN_TOKENS_PER_SAMPLE_KEY] = (
                        rollout_metrics[MEAN_GEN_TOKENS_PER_SAMPLE_KEY]
                    )
                    logger.log_metrics(rollout_metrics, total_steps + 1, prefix="train")

                repeated_batch = scale_rewards(
                    repeated_batch, master_config.grpo.reward_scaling
                )
                # Process rewards with custom reward function
                if master_config.grpo.reward_shaping.enabled:
                    repeated_batch = apply_reward_shaping(
                        repeated_batch, master_config.grpo.reward_shaping
                    )

                # Calculate rewards & advantages
                memory_tracker.snapshot_start_of_stage("Processing rewards", dir())
                print("▶ Processing rewards...,", flush=True)
                with (
                    timer.time("reward_calculation"),
                    managed_span(
                        RLSpanGroup.REWARD, "rl.grpo.reward_calculation", tracer=_tracer
                    ),
                ):
                    # Extract rewards from final_batch
                    rewards = repeated_batch["total_reward"]

                    print("▶ Computing advantages...", flush=True)
                    # For DAPO with reward shaping, compute std on the raw
                    # pre-shaping reward so dynamic sampling filters prompt
                    # groups on the raw task metric (e.g. acc) instead of on
                    # length-dependent shaped reward variance. Baseline
                    # (which drives advantages) stays on the shaped reward.
                    std_rewards = (
                        repeated_batch["unshaped_total_reward"]
                        if master_config.grpo.use_dynamic_sampling
                        and "unshaped_total_reward" in repeated_batch
                        else None
                    )
                    is_trivial_prompt_distribution = (
                        calculate_trivial_reward_distributions(
                            repeated_batch[PROMPT_GROUP_IDS_KEY],
                            std_rewards if std_rewards is not None else rewards,
                            torch.ones_like(rewards),
                        )
                        if master_config.grpo.use_dynamic_sampling
                        else None
                    )
                    if master_config.grpo.calculate_advantages_on_gpu:
                        print("Computing advantages on GPU!")
                        # Just fix the device id for now
                        device_id = 0
                        (
                            baseline,
                            std,
                            _,
                        ) = calculate_baseline_and_std_per_prompt(
                            repeated_batch[PROMPT_GROUP_IDS_KEY].cuda(device_id),
                            rewards.cuda(device_id),
                            torch.ones_like(rewards).cuda(device_id),
                            leave_one_out_baseline=master_config.grpo.use_leave_one_out_baseline,
                            std_rewards=(
                                std_rewards.cuda(device_id)
                                if std_rewards is not None
                                else None
                            ),
                        )
                        baseline = baseline.cpu()
                        std = std.cpu()
                    else:
                        (
                            baseline,
                            std,
                            _,
                        ) = calculate_baseline_and_std_per_prompt(
                            repeated_batch[PROMPT_GROUP_IDS_KEY],
                            rewards,
                            torch.ones_like(rewards),
                            leave_one_out_baseline=master_config.grpo.use_leave_one_out_baseline,
                            std_rewards=std_rewards,
                        )
                    if "env_reward" in repeated_batch:
                        # grpo.length_penalty rewrote total_reward; keep a
                        # baseline of the env reward for the pct_* diagnostics.
                        # Stored on the batch so dynamic sampling filters it.
                        repeated_batch["env_baseline"], _, _ = (
                            calculate_baseline_and_std_per_prompt(
                                repeated_batch[PROMPT_GROUP_IDS_KEY],
                                repeated_batch["env_reward"],
                                torch.ones_like(rewards),
                                leave_one_out_baseline=master_config.grpo.use_leave_one_out_baseline,
                            )
                        )

                    # Apply dynamic sampling to filter prompts with non-zero std (DAPO algorithm)
                    repeated_batch, is_batch_complete, batch_cache, ds_metrics = (
                        dynamic_sampling(
                            repeated_batch,
                            std,
                            baseline,
                            dynamic_sampling_num_gen_batches,
                            master_config,
                            timer,
                            batch_cache,
                            is_trivial_prompt_distribution=(
                                is_trivial_prompt_distribution
                            ),
                        )
                    )
                    if ds_metrics:
                        ds_metrics["dynamic_sampling_num_gen_batches"] = (
                            dynamic_sampling_num_gen_batches
                        )
                    # Get the updated rewards and baselines. For DAPO, these rewards and baselines only correspond to the prompts with non-zero std.
                    rewards = (
                        repeated_batch["total_reward"]
                        if not master_config.grpo.use_dynamic_sampling
                        else repeated_batch["filtered_reward"]
                    )
                    baseline = repeated_batch["baseline"]
                    std = repeated_batch["std"]

                    # If the current batch is not enough to fill the buffer during dynamic sampling, we update the cache and process the next batch.
                    if not is_batch_complete:
                        continue

                    gen_step_metrics = {}
                    if hasattr(policy_generation, "get_step_metrics"):
                        gen_step_metrics = policy_generation.get_step_metrics()

                    # Save baseline for logging (before deletion)
                    baseline_for_log = baseline.clone()
                    env_baseline_for_log = (
                        repeated_batch["env_baseline"].clone()
                        if "env_baseline" in repeated_batch
                        else None
                    )

                    # Backfill before the training flatten reuses these messages.
                    backfill_missing_routed_experts(repeated_batch["message_log"])

                    prompt_ids_for_adv = repeated_batch[PROMPT_GROUP_IDS_KEY]
                    del baseline
                    del std

                with (
                    timer.time("data_processing"),
                    managed_span(
                        RLSpanGroup.DATA_PROCESSING,
                        "rl.grpo.data_processing",
                        tracer=_tracer,
                    ),
                ):
                    use_overlong_filtering = master_config.grpo.overlong_filtering
                    if use_overlong_filtering:
                        loss_multiplier = repeated_batch["loss_multiplier"].clone()
                        truncated = repeated_batch["truncated"]

                        if isinstance(truncated, list):
                            truncated = torch.tensor(truncated, dtype=torch.bool)

                        loss_multiplier[truncated] = 0
                        repeated_batch["loss_multiplier"] = loss_multiplier

                    num_mask_sample_filtered = _apply_mask_sample_filter(repeated_batch)
                    metrics["num_mask_sample_filtered"] = num_mask_sample_filtered

                    add_grpo_token_loss_masks_and_generation_logprobs(
                        repeated_batch["message_log"]
                    )

                    # Convert updated LLMMessageLogType to FlatMessagesType for training
                    flat_messages, input_lengths = batched_message_log_to_flat_message(
                        repeated_batch["message_log"],
                        pad_value_dict={"token_ids": tokenizer.pad_token_id},
                        make_sequence_length_divisible_by=master_config.policy[
                            "make_sequence_length_divisible_by"
                        ],
                    )

                    # Create training data from flattened messages
                    # Note: advantages will be computed and added after logprobs are available
                    train_data = BatchedDataDict[ClippedPGLossDataDict](
                        {
                            "input_ids": flat_messages["token_ids"],
                            "input_lengths": input_lengths,
                            "generation_logprobs": flat_messages["generation_logprobs"],
                            "token_mask": flat_messages["token_loss_mask"],
                            "sample_mask": repeated_batch["loss_multiplier"],
                        }
                    )
                    # this will be mini-batched inside the policy, so maintain the packed multimodal structure
                    # This is also used to populate part of the downstream logprob calculation data
                    extra_multimodal_data = flat_messages.get_multimodal_dict(
                        as_tensors=False,
                        pixel_dtype=_policy_dtype(master_config.policy),
                    )
                    train_data.update(extra_multimodal_data)
                    print_multimodal_payload_metrics(
                        collect_multimodal_payload_metrics(
                            train_data,
                            "rollout_to_policy",
                            enabled=master_config.grpo.debug_payload_metrics,
                        )
                    )
                    # Router replay (R3) on the legacy data_plane.enabled=false
                    # driver path: routed_experts already rides flat_messages
                    # (attached to message_log during rollout, then batched into
                    # a [B, S, L, K] tensor by batched_message_log_to_flat_message),
                    # but the train_data whitelist above drops it. Copy it back so
                    # the Megatron worker's train-stage router-replay guard finds
                    # it. Mirrors the TQ producer (sync_rollout_actor.py).
                    _preserve_router_replay_routed_experts(
                        train_data, flat_messages, master_config.policy
                    )
                    train_data.to("cpu")

                    metrics_logging_data["content"] = flat_messages["content"]

                memory_tracker.snapshot_start_of_stage("Computing logprobs", dir())
                skip_prev_logprobs, skip_reference_logprobs = (
                    _resolve_logprob_skip_flags(master_config)
                )
                seq_logprob_error_threshold = (
                    master_config.grpo.seq_logprob_error_threshold
                )

                if not (skip_prev_logprobs and skip_reference_logprobs):
                    print("▶ Preparing for logprob inference...", flush=True)
                    with timer.time("logprob_inference_prep"):
                        policy.prepare_for_lp_inference()

                print("▶ Computing logprobs...", flush=True)
                with (
                    timer.time("policy_and_reference_logprobs"),
                    managed_span(
                        RLSpanGroup.LOGPROB,
                        "rl.grpo.policy_and_reference_logprobs",
                        tracer=_tracer,
                    ),
                ):
                    # Custom create this logprob_data so we avoid Ray comm overheads sending unused data to workers.
                    logprob_data = BatchedDataDict[ClippedPGLossDataDict](
                        {
                            "input_ids": train_data["input_ids"],
                            "input_lengths": train_data["input_lengths"],
                            "token_mask": flat_messages["token_loss_mask"],
                            "sample_mask": repeated_batch["loss_multiplier"],
                            **extra_multimodal_data,
                        }
                    )
                    # Router replay (R3): the prev-logprobs forward replays the
                    # recorded routed_experts, so logprob_data must carry the
                    # field too (it is a separate whitelist from train_data). The
                    # reference-policy logprobs call reuses logprob_data but
                    # intentionally ignores routed_experts (require_router_replay
                    # =False short-circuits before the field is read), so a
                    # present-but-unused field here is safe.
                    _preserve_router_replay_routed_experts(
                        logprob_data, flat_messages, master_config.policy
                    )

                    if not skip_prev_logprobs:
                        prev_lp_result = policy.get_logprobs(logprob_data, timer=timer)
                        train_data["prev_logprobs"] = prev_lp_result["logprobs"]
                    else:
                        print(
                            "▶ Skipping prev_logprobs (force_on_policy_ratio=True)...",
                            flush=True,
                        )
                        train_data["prev_logprobs"] = torch.zeros_like(
                            train_data["generation_logprobs"]
                        )

                    if not skip_reference_logprobs:
                        train_data["reference_policy_logprobs"] = (
                            policy.get_reference_policy_logprobs(
                                logprob_data,
                                timer=timer,
                            )["reference_logprobs"]
                        )
                    else:
                        print(
                            "▶ Skipping reference_logprobs (skip_reference_policy_logprobs_calculation=True)...",
                            flush=True,
                        )
                        train_data["reference_policy_logprobs"] = torch.zeros_like(
                            train_data["prev_logprobs"]
                        )

                    del logprob_data
                    del extra_multimodal_data

                # Separate-pass seq-level metrics/masking require real prev_logprobs
                if skip_prev_logprobs:
                    # In-loss filtering reports counts through all_mb_metrics.
                    # Use {} so placeholder zeros cannot overwrite those counts
                    # when seq_logprob_error_metrics is merged after training.
                    # Otherwise, placeholder prev_logprobs cannot provide
                    # sequence-error metrics.
                    seq_logprob_error_metrics = (
                        {}
                        if master_config.loss_fn.seq_logprob_error_in_loss
                        else _placeholder_seq_logprob_error_metrics()
                    )
                else:
                    seq_error_result = compute_and_apply_seq_logprob_error_masking(
                        train_data=train_data,
                        rewards=rewards,
                        seq_logprob_error_threshold=seq_logprob_error_threshold,
                        filtering_on=need_top_k_or_top_p_filtering(
                            TrainingSamplingParams(
                                top_k=master_config.policy["generation"]["top_k"],
                                top_p=master_config.policy["generation"]["top_p"],
                            )
                        ),
                    )
                    seq_logprob_error_metrics = seq_error_result
                    if "num_masked_seqs" in seq_logprob_error_metrics:
                        seq_logprob_error_metrics[
                            "num_masked_seqs_by_logprob_error"
                        ] = seq_logprob_error_metrics.pop("num_masked_seqs")

                # Compute advantages with adv_estimator using correct mask and logprobs
                with (
                    timer.time("advantage_calculation"),
                    managed_span(
                        RLSpanGroup.ADVANTAGE,
                        "rl.grpo.advantage_calculation",
                        tracer=_tracer,
                    ),
                ):
                    print("▶ Computing advantages...", flush=True)
                    # Get token-level mask: token_mask * sample_mask
                    token_mask = train_data["token_mask"]
                    sample_mask = train_data["sample_mask"]
                    mask = token_mask * sample_mask.unsqueeze(-1)

                    train_data["advantages"] = adv_estimator.compute_advantage(
                        prompt_ids=prompt_ids_for_adv,
                        rewards=rewards,
                        mask=mask,
                        repeated_batch=repeated_batch,
                        logprobs_policy=train_data["prev_logprobs"],
                        logprobs_reference=train_data.get("reference_policy_logprobs"),
                    )
                    del prompt_ids_for_adv

                    # Log rewards and advantages information
                    _log_mixed_rewards_and_advantages_information(
                        logger=logger,
                        total_steps=total_steps,
                        metrics=metrics,
                        baseline=baseline_for_log,
                        advantages=train_data["advantages"],
                        env_baseline=env_baseline_for_log,
                    )
                    del baseline_for_log, env_baseline_for_log

                    penalty_metrics = (
                        _apply_configured_message_level_advantage_penalties(
                            train_data, repeated_batch["message_log"], master_config
                        )
                    )

                    # Clip advantages to prevent extreme values from small std normalization
                    train_data["advantages"] = _clip_grpo_advantages(
                        train_data["advantages"], master_config.grpo
                    )

                memory_tracker.snapshot_start_of_stage("Policy train", dir())
                print("▶ Preparing for training...", flush=True)
                with timer.time("training_prep"):
                    policy.prepare_for_training()  # set model train and reload optim to GPU
                    POLICY_GENERATION_STALE = True

                print("▶ Training policy...", flush=True)
                with (
                    timer.time("policy_training"),
                    managed_span(
                        RLSpanGroup.POLICY_UPDATE,
                        "rl.grpo.policy_training",
                        tracer=_tracer,
                    ),
                ):
                    train_results = policy.train(
                        train_data,
                        loss_fn,
                        timer=timer,
                    )

                # Recompute KV scales after policy training if needed
                if sync_kv_scales:
                    with timer.time("recompute_kv_scales"):
                        print(
                            "▶ Recomputing KV cache scales after policy update...",
                            flush=True,
                        )
                        kv_scales_cache = policy.calibrate_qkv_fp8_scales(
                            train_data, include_q=True
                        )["layers"]
                        # Set generation as stale to force refit with new scales
                        POLICY_GENERATION_STALE = True

                is_last_step = total_steps + 1 >= max_num_steps
                if not master_config.data["use_multiple_dataloader"]:
                    is_last_step = is_last_step or (
                        (current_epoch + 1 == max_num_epochs)
                        and (current_step + 1 == len(wrapped_dataloader))
                    )

                early_stop_message: Optional[str] = None
                should_run_validation = (
                    val_period > 0
                    and (total_steps + 1) >= val_start_at
                    and (total_steps + 1) % val_period == 0
                ) or (val_at_end and is_last_step)

                # Keep training and validation traffic in separate metric intervals.
                payload_metrics: dict[str, int | float] = {}
                if master_config.grpo.debug_payload_metrics:
                    payload_metrics = drain_multimodal_payload_metrics()

                # Run validation if it's a validation step or last step with val_at_end
                if should_run_validation:
                    memory_tracker.snapshot_start_of_stage("Validation", dir())
                    if POLICY_GENERATION_STALE:
                        refit_metrics = refit_policy_generation(
                            policy,
                            policy_generation,
                            colocated_inference,
                            _refit_buffer_size_gb=refit_buffer_size_gb,
                            kv_scales=kv_scales_cache if sync_kv_scales else None,
                        )
                        POLICY_GENERATION_STALE = False
                    else:
                        if colocated_inference:
                            policy.offload_after_refit()  # unload optimizer to make space for generation
                        policy_generation.prepare_for_generation()
                    val_metrics, validation_timings = validate(
                        policy_generation,
                        val_dataloader,
                        tokenizer,
                        val_task_to_env,
                        step=total_steps + 1,
                        master_config=master_config,
                        logger=logger,
                        processor=processor,
                    )
                    policy_generation.finish_generation()
                    logger.log_metrics(
                        validation_timings, total_steps + 1, prefix="timing/validation"
                    )
                    logger.log_metrics(
                        val_metrics, total_steps + 1, prefix="validation"
                    )
                    if master_config.grpo.debug_payload_metrics:
                        validation_payload_metrics = drain_multimodal_payload_metrics()
                        if validation_payload_metrics:
                            logger.log_metrics(
                                validation_payload_metrics,
                                total_steps + 1,
                                prefix="validation",
                            )
                    early_stop_message = _validation_early_stop_message(
                        val_metrics,
                        stop_at_validation_threshold,
                        stop_at_validation_metric,
                    )
                    if early_stop_message is not None:
                        # Exit at the end of this step, after checkpointing.
                        print(early_stop_message, flush=True)

                # Get flat advantages and token mask for masked metrics computation
                flat_advantages = train_data["advantages"]
                flat_token_mask = flat_messages["token_loss_mask"]

                # Filter advantages using token mask (only valid response tokens)
                response_advantages = torch.masked_select(
                    flat_advantages, flat_token_mask.bool()
                )

                memory_tracker.snapshot_start_of_stage("Metrics", dir())
                metrics = {
                    **metrics,
                    LOSS_KEY: train_results["loss"].numpy(),
                    GRAD_NORM_KEY: train_results["grad_norm"].numpy(),
                    REWARD_KEY: rewards.numpy(),
                    "mean_prompt_length": repeated_batch["length"].numpy(),
                    "total_num_tokens": input_lengths.numpy(),
                    # Add masked advantages tracking metrics (only for valid response tokens)
                    "advantages/mean": torch.mean(response_advantages).detach().item()
                    if response_advantages.numel() > 0
                    else 0.0,
                    "advantages/max": torch.max(response_advantages).detach().item()
                    if response_advantages.numel() > 0
                    else 0.0,
                    "advantages/min": torch.min(response_advantages).detach().item()
                    if response_advantages.numel() > 0
                    else 0.0,
                    **ds_metrics,
                }
                if "moe_metrics" in train_results:
                    metrics.update(
                        {f"moe/{k}": v for k, v in train_results["moe_metrics"].items()}
                    )
                if "mtp_metrics" in train_results:
                    metrics.update(
                        {f"mtp/{k}": v for k, v in train_results["mtp_metrics"].items()}
                    )
                if "draft_grad_norm" in train_results:
                    metrics["draft_grad_norm"] = train_results[
                        "draft_grad_norm"
                    ].numpy()
                if master_config.grpo.use_dynamic_sampling:
                    metrics["filtered_reward"] = rewards.numpy()
                    metrics["reward"] = repeated_batch["total_reward"].numpy()

                metrics.update(train_results["all_mb_metrics"])
                metrics.update(gen_step_metrics)
                metrics.update(penalty_metrics)
                for k, v in metrics.items():
                    if k in {"probs_ratio_min", "probs_ratio_clamped_min"}:
                        valid_values = [x for x in v if not np.isinf(x)]
                        metrics[k] = (
                            np.min(valid_values).item() if valid_values else -1.0
                        )
                    elif k in {"probs_ratio_max", "probs_ratio_clamped_max"}:
                        valid_values = [x for x in v if not np.isinf(x)]
                        metrics[k] = (
                            np.max(valid_values).item() if valid_values else -1.0
                        )
                    elif k in {
                        "lr",
                        "wd",
                        "reward",
                        "filtered_reward",
                        "global_valid_seqs",
                        "global_valid_toks",
                        "mean_prompt_length",
                    }:
                        metrics[k] = np.mean(v).item()
                    elif isinstance(v, (np.ndarray, list)):
                        metrics[k] = np.sum(v).item()
                    else:
                        print(f"Skipping aggregation for {k} ({type(v)})")

                metrics.update(rollout_metrics)
                metrics["generation_logger_metrics"] = generation_logger_metrics
                total_valid_tokens += metrics["global_valid_toks"]

                # Always log sequence-level error metrics (useful for deciding threshold)
                metrics.update(seq_logprob_error_metrics)

                ## Checkpointing
                consumed_samples += master_config.grpo.num_prompts_per_step
                timeout.mark_iteration()

                # +1 because step is 0-indexed
                should_save_by_step = (
                    is_last_step
                    # Early stop saves the final state like a last step.
                    or early_stop_message is not None
                    or (total_steps + 1) % master_config.checkpointing["save_period"]
                    == 0
                    or (
                        ft_save_period is not None
                        and (total_steps + 1) % ft_save_period == 0
                    )
                )
                # Check if timeout-based checkpointing is enabled in config.
                should_save_by_timeout = timeout.check_save()

                memory_tracker.snapshot_start_of_stage("Checkpointing", dir())
                if master_config.checkpointing["enabled"] and (
                    should_save_by_step or should_save_by_timeout
                ):
                    policy.prepare_for_training()

                    # +1 because step is 0-indexed
                    grpo_save_state.current_step = current_step + 1
                    grpo_save_state.total_steps = total_steps + 1
                    grpo_save_state.current_epoch = current_epoch
                    grpo_save_state.total_valid_tokens = total_valid_tokens
                    if val_metrics is not None:
                        grpo_save_state.val_reward = val_metrics["accuracy"]
                    elif hasattr(grpo_save_state, "val_reward"):
                        delattr(grpo_save_state, "val_reward")
                    grpo_save_state.consumed_samples = consumed_samples

                    full_metric_name = master_config.checkpointing["metric_name"]
                    if full_metric_name is not None:
                        prefix, metric_name = full_metric_name.split(":", 1)
                        metrics_source = metrics if prefix == "train" else val_metrics
                        if not metrics_source:
                            warnings.warn(
                                f"You asked to save checkpoints based on {metric_name} but no {prefix} metrics were collected. "
                                "This checkpoint will not be saved as top-k.",
                                stacklevel=2,
                            )
                            if hasattr(grpo_save_state, full_metric_name):
                                delattr(grpo_save_state, full_metric_name)
                        elif metric_name not in metrics_source:
                            raise ValueError(
                                f"Metric {metric_name} not found in {prefix} metrics"
                            )
                        else:
                            setattr(
                                grpo_save_state,
                                full_metric_name,
                                metrics_source[metric_name],
                            )

                    with (
                        timer.time("checkpointing"),
                        managed_span(
                            RLSpanGroup.CHECKPOINT,
                            "rl.grpo.checkpointing",
                            tracer=_tracer,
                        ),
                    ):
                        # Finalize the previous (possibly async) checkpoint before
                        # starting a new one. No-op with sync save / nothing pending.
                        checkpointer.finalize_pending()

                        print(
                            f"Saving checkpoint for step {total_steps + 1}...",
                            flush=True,
                        )
                        checkpoint_path = checkpointer.init_tmp_checkpoint(
                            total_steps + 1, vars(grpo_save_state), master_config
                        )
                        policy.save_checkpoint(
                            weights_path=os.path.join(
                                checkpoint_path, "policy", "weights"
                            ),
                            optimizer_path=os.path.join(
                                checkpoint_path, "policy", "optimizer"
                            )
                            if checkpointer.save_optimizer
                            else None,
                            tokenizer_path=os.path.join(
                                checkpoint_path, "policy", "tokenizer"
                            ),
                            is_final_checkpoint=(
                                is_last_step or early_stop_message is not None
                            ),
                        )
                        if master_config.data["use_multiple_dataloader"]:
                            for (
                                task_name,
                                task_dataloader,
                            ) in wrapped_dataloader.dataloaders.items():
                                torch.save(
                                    task_dataloader.state_dict(),
                                    os.path.join(
                                        checkpoint_path,
                                        f"train_dataloader_{task_name}.pt",
                                    ),
                                )
                        else:
                            torch.save(
                                wrapped_dataloader.state_dict(),
                                os.path.join(checkpoint_path, "train_dataloader.pt"),
                            )
                        # Finalize in the background. The directory rename is
                        # deferred until any async write completes (via wait_fn);
                        # with sync save it renames immediately. Finalization is
                        # flushed at the next save (finalize_pending) or on exit
                        # (shutdown).
                        checkpointer.begin_finalization(
                            checkpoint_path,
                            wait_fn=policy.finalize_async_save,
                        )

                        # Record last-successful-checkpoint time/step for external
                        # monitoring (parity with async_grpo_train; see
                        # _write_latest_checkpoint_status).
                        _write_latest_checkpoint_status(
                            checkpointer, last_checkpoint_step=total_steps + 1
                        )

            # Logging
            # Log training data
            memory_tracker.snapshot_start_of_stage("Logging", dir())
            if not _should_log_nemo_gym_responses(master_config):
                log_data = {}
                if "agent_ref" in repeated_batch:
                    log_data["agent_ref"] = repeated_batch["agent_ref"]
                log_data["content"] = flat_messages["content"]
                log_data["rewards"] = rewards.tolist()
                if master_config.grpo.use_dynamic_sampling:
                    log_data["filtered_rewards"] = rewards.tolist()
                    log_data["rewards"] = repeated_batch["total_reward"].tolist()
                log_data["input_lengths"] = input_lengths.tolist()
                log_data["token_ids"] = train_data["input_ids"].tolist()
                log_data["token_loss_mask"] = train_data["token_mask"].tolist()
                log_data["sample_loss_mask"] = train_data["sample_mask"].tolist()
                log_data["advantages"] = train_data["advantages"].tolist()
                log_data["generation_logprobs"] = train_data[
                    "generation_logprobs"
                ].tolist()
                log_data["prev_logprobs"] = train_data["prev_logprobs"].tolist()

                logger.log_batched_dict_as_jsonl(
                    log_data, f"train_data_step{total_steps + 1}.jsonl"
                )
                del log_data
            del flat_messages

            timing_metrics: dict[str, float] = timer.get_timing_metrics(
                reduction_op="sum"
            )  # type: ignore
            # track example with high token mult prob error above 1.05
            if metrics["token_mult_prob_error"] > 1.05:
                logger.log_plot_token_mult_prob_error(
                    {
                        "prompt_lengths": repeated_batch["length"],
                        "full_lengths": input_lengths,
                        "generation_logprobs": train_data["generation_logprobs"],
                        "prev_logprobs": train_data["prev_logprobs"],
                        "token_mask": train_data["token_mask"],
                        "sample_mask": train_data["sample_mask"],
                    },
                    total_steps + 1,
                    name="train/token_mult_prob_error_plot_sample",
                    filtering_on=need_top_k_or_top_p_filtering(
                        TrainingSamplingParams(
                            top_k=master_config.policy["generation"]["top_k"],
                            top_p=master_config.policy["generation"]["top_p"],
                        )
                    ),
                )
            del train_data
            if (
                master_config.policy["generation"]
                .get("vllm_cfg", {})
                .get("enable_vllm_metrics_logger", False)
            ):
                log_generation_metrics(
                    generation_logger_metrics,
                    total_steps + 1,
                    master_config.policy["generation"]["vllm_cfg"][
                        "vllm_metrics_logger_interval"
                    ],
                    logger,
                )

            print("\n📊 Training Results:")

            print(f"  • Loss: {metrics['loss']:.4f}")
            if "draft_loss" in metrics:
                print(f"  • Draft Loss: {metrics['draft_loss']:.4f}")
            generation_kl_error = metrics.get("gen_kl_error")
            print(
                "  • Generation KL Error: "
                + (
                    f"{generation_kl_error:.4f}"
                    if generation_kl_error is not None
                    else "not reported"
                )
            )
            if master_config.grpo.use_dynamic_sampling:
                print(f"  • Avg Filtered Reward: {np.mean(rewards.numpy()):.4f}")
                print(
                    f"  • Avg Total Reward: {np.mean(repeated_batch['total_reward'].numpy()):.4f}"
                )
            else:
                print(f"  • Avg Reward: {np.mean(rewards.numpy()):.4f}")
            print(
                f"  • Mean Generation Length: {metrics_logging_data['mean_gen_tokens_per_sample']:.4f}",
                flush=True,
            )

            print("\n⏱️  Timing:", flush=True)
            # Display total time first, separately
            total_time = timing_metrics.get("total_step_time", 0)

            number_of_samples_per_step = (
                master_config.grpo.num_prompts_per_step
                * master_config.grpo.num_generations_per_prompt
            )
            total_num_gpus = (
                master_config.cluster.num_nodes * master_config.cluster.gpus_per_node
            )

            print(f"  • Total step time: {total_time:.2f}s", flush=True)

            # Display all other timing metrics
            for k, v in sorted(
                timing_metrics.items(), key=lambda item: item[1], reverse=True
            ):
                if k != "total_step_time":
                    percent = (v / total_time * 100) if total_time > 0 else 0
                    print(f"  • {k}: {v:.2f}s ({percent:.1f}%)", flush=True)

            timing_metrics["valid_tokens_per_sec_per_gpu"] = (
                metrics["global_valid_toks"] / total_time / total_num_gpus
            )
            performance_metrics = print_performance_metrics(
                train_results,
                metrics,
                timing_metrics,
                master_config,
                num_prompts_per_step=master_config.grpo.num_prompts_per_step,
                num_generations_per_prompt=master_config.grpo.num_generations_per_prompt,
                is_async_rl=master_config.grpo.async_grpo.enabled,
            )

            if payload_metrics:
                logger.log_metrics(payload_metrics, total_steps + 1, prefix="")

            if refit_metrics:
                logger.log_metrics(refit_metrics, total_steps + 1, prefix="refit")
            logger.log_metrics(metrics, total_steps + 1, prefix="train")
            logger.log_metrics(
                performance_metrics, total_steps + 1, prefix="performance"
            )
            # step_finished=True here since this is the final log of our current step.
            logger.log_metrics(
                timing_metrics,
                total_steps + 1,
                prefix="timing/train",
                step_finished=True,
            )

            # Reset the batch and set dynamic_sampling_num_gen_batches to 0
            batch_cache = None
            dynamic_sampling_num_gen_batches = 0

            # Clear mem
            memory_tracker.snapshot_start_of_stage("After CPU memory clear", dir())

            # processing rewards
            del repeated_batch
            del rewards
            # train_data already deleted after logging above
            # logging
            del metrics
            if "val_metrics" in dir():
                del val_metrics

            timer.reset()
            current_step += 1
            total_steps += 1
            if early_stop_message is not None:
                checkpointer.shutdown()
                memory_tracker.snapshot_start_of_stage("", dir())
                return
            if should_save_by_timeout:
                checkpointer.shutdown()
                memory_tracker.snapshot_start_of_stage("", dir())
                print("Timeout has been reached, stopping training early", flush=True)
                return
            if total_steps >= max_num_steps:
                checkpointer.shutdown()
                memory_tracker.snapshot_start_of_stage("", dir())
                print(
                    "Max number of steps has been reached, stopping training early",
                    flush=True,
                )
                return

        current_epoch += 1
        current_step = 0  # Reset step counter for new epoch

    # Flush the last checkpoint's background finalization on an epoch-bounded
    # exit. Reaching max_num_epochs falls through the while loop and bypasses
    # the inline shutdown() calls at the max_num_steps / timeout early returns,
    # so without this the daemon finalization thread would be killed before the
    # final tmp_step_N is renamed.
    checkpointer.shutdown()


@umbrella_trace_fn(RLSpanGroup.U_JOB, "rl.grpo.job")
def grpo_train(
    policy: ColocatablePolicyInterface,
    policy_generation: Optional[GenerationInterface],
    wrapped_dataloader: StatefulDataLoader | MultipleDataloaderWrapper,
    val_dataloader: Optional[StatefulDataLoader],
    tokenizer: TokenizerType,
    loss_fn: LossFunction,
    task_to_env: dict[str, EnvironmentInterface],
    val_task_to_env: Optional[dict[str, EnvironmentInterface]],
    logger: Logger,
    checkpointer: CheckpointManager,
    grpo_save_state: GRPOSaveState,
    master_config: MasterConfig,
    processor: Optional[AutoProcessor] = None,
) -> None:
    """Run GRPO training and always tear down its environments."""
    try:
        _grpo_train_impl(
            policy=policy,
            policy_generation=policy_generation,
            wrapped_dataloader=wrapped_dataloader,
            val_dataloader=val_dataloader,
            tokenizer=tokenizer,
            loss_fn=loss_fn,
            task_to_env=task_to_env,
            val_task_to_env=val_task_to_env,
            logger=logger,
            checkpointer=checkpointer,
            grpo_save_state=grpo_save_state,
            master_config=master_config,
            processor=processor,
        )
    finally:
        shutdown_environments(task_to_env, val_task_to_env)


def validate(
    policy_generation: GenerationInterface,
    val_dataloader: Optional[StatefulDataLoader],
    tokenizer,
    val_task_to_env: Optional[dict[str, EnvironmentInterface]],
    step: int,
    master_config: MasterConfig,
    logger: Optional[Logger] = None,
    processor: Optional[AutoProcessor] = None,
) -> tuple[dict[str, Any], dict[str, Any]]:
    """Run validation on the validation dataset."""
    if val_dataloader is None:
        assert val_dataloader is not None or master_config.grpo.val_period == 0, (
            "val_dataloader is None, so grpo.val_period must be 0"
        )
        print("  ⚠️ No validation dataloader provided, skipping validation", flush=True)
        return {}, {}

    timer = Timer(context={"worker": "validator"})
    with (
        timer.time("total_validation_time"),
        evaluate_span("grpo", **{RL_ITERATION_ATTR: step}),
    ):
        print(f"▶ Starting validation at step {step}...", flush=True)
        # >= 1 is validated in setup().
        val_num_generations_per_prompt = (
            master_config.grpo.val_num_generations_per_prompt
        )

        total_rewards = []
        total_lengths = []
        all_message_logs = []  # Collect all message logs

        max_batches = (
            master_config.grpo.max_val_samples // master_config.grpo.val_batch_size
        )
        for batch_idx, val_batch in enumerate(val_dataloader):
            if batch_idx >= max_batches:
                break

            if val_num_generations_per_prompt > 1:
                val_batch = val_batch.repeat_interleave(val_num_generations_per_prompt)

            additional_metrics_to_report = dict()
            # Generate responses (updates the LLMMessageLogType in batch_with_msg_logs)
            # Use async rollouts when enabled by config/backend defaults.
            # We cascade NeMo-Gym first since NeMo-Gym also uses async rollouts.
            if should_use_nemo_gym(master_config):
                if master_config.grpo.deduplicate_multimodal_data:
                    attach_initial_nemo_gym_image_payloads(
                        val_batch,
                        processor,
                        env_config=master_config.env,
                    )
                generation_config = master_config.policy["generation"]
                # Validation-only sampling (e.g. near-greedy validation);
                # defaults to the train profile via the exemplar YAML
                # interpolations. Training rollouts keep policy.generation.
                val_sampling_params = GenerationSamplingParams(
                    temperature=generation_config["val_temperature"],
                    top_p=generation_config["val_top_p"],
                    top_k=generation_config["val_top_k"],
                )
                nemo_gym_rollout_result = run_nemo_gym_rollout_sync(
                    policy_generation=policy_generation,
                    input_batch=val_batch,
                    tokenizer=tokenizer,
                    task_to_env=val_task_to_env,
                    max_seq_len=master_config.policy["max_total_sequence_length"],
                    generation_config=generation_config,
                    num_generations_per_prompt=val_num_generations_per_prompt,
                    sampling_params=val_sampling_params,
                    log_full_result_tables=should_log_nemo_gym_full_result_tables(
                        wandb_enabled=master_config.logger.wandb_enabled,
                        wandb_config=master_config.logger.wandb,
                    ),
                    max_rollout_turns=None,
                    greedy=False,
                    effort_config=_get_effort_config(master_config),
                    reward_penalty_config=master_config.reward_penalties,
                    # No length_penalty_config here: validation metrics
                    # (accuracy/pass_k) must reflect the raw env reward, and the
                    # adjustment code groups by the TRAINING stride
                    # (num_generations_per_prompt), which does not match
                    # val_num_generations_per_prompt.
                    thinking_tags=get_nemo_gym_thinking_tags(master_config.env),
                    mask_env_flagged_samples=should_mask_flagged_samples(
                        master_config.env
                    ),
                    deduplicate_multimodal_data=(
                        master_config.grpo.deduplicate_multimodal_data
                    ),
                    debug_payload_metrics=master_config.grpo.debug_payload_metrics,
                )
                val_batch = nemo_gym_rollout_result.final_batch
                gen_metrics = nemo_gym_rollout_result.rollout_metrics
                additional_metrics_to_report = gen_metrics
            elif should_use_async_rollouts(master_config.policy["generation"]):
                val_batch, gen_metrics = run_async_multi_turn_rollout(
                    policy_generation,
                    val_batch,
                    tokenizer,
                    val_task_to_env,
                    max_seq_len=master_config.policy["max_total_sequence_length"],
                    max_rollout_turns=master_config.grpo.max_rollout_turns,
                    greedy=False,
                    deduplicate_multimodal_data=(
                        master_config.grpo.deduplicate_multimodal_data
                    ),
                )
            else:
                val_batch, gen_metrics = run_multi_turn_rollout(
                    policy_generation,
                    val_batch,
                    tokenizer,
                    val_task_to_env,
                    max_seq_len=master_config.policy["max_total_sequence_length"],
                    max_rollout_turns=master_config.grpo.max_rollout_turns,
                    greedy=False,
                    deduplicate_multimodal_data=(
                        master_config.grpo.deduplicate_multimodal_data
                    ),
                )

            total_rewards.extend(val_batch["total_reward"].tolist())
            total_lengths.append(gen_metrics["mean_gen_tokens_per_sample"])

            # Collect message logs for later display
            to_env = [
                get_keys_from_message_log(
                    val_batch["message_log"][i], ["role", "content"]
                )
                for i in range(len(val_batch["message_log"]))
            ]

            all_message_logs.extend(to_env)

        # Calculate validation metrics. accuracy is the mean reward over all
        # rollouts; grouped validation (val_num_generations_per_prompt > 1)
        # additionally reports pass@k over each prompt's k rollouts as pass_k.
        num_samples = len(total_rewards)
        pass_k = None
        if num_samples > 0:
            rewards_t = torch.tensor(total_rewards, dtype=torch.float32)
            accuracy = rewards_t.mean().item()
            if val_num_generations_per_prompt > 1:
                assert num_samples % val_num_generations_per_prompt == 0, (
                    "Validation rewards must be divisible by "
                    "grpo.val_num_generations_per_prompt"
                )
                pass_k = (
                    (rewards_t.view(-1, val_num_generations_per_prompt) > 0)
                    .any(dim=1)
                    .float()
                    .mean()
                    .item()
                )
        else:
            accuracy = 0.0

        avg_length = (
            sum(total_lengths) / len(total_lengths) if len(total_lengths) > 0 else 0.0
        )

        val_metrics = {
            "accuracy": accuracy,
            "avg_length": avg_length,
            **additional_metrics_to_report,
        }
        if pass_k is not None:
            val_metrics["pass_k"] = pass_k

        # Print sample conversations only once at the end of validation
        try:
            print_message_log_samples(
                all_message_logs,
                total_rewards,
                num_samples=min(
                    master_config.logger.num_val_samples_to_print,
                    len(all_message_logs),
                ),
                step=step,
            )
        except Exception as e:
            print(f"\n  ⚠️ Error displaying message samples: {str(e)}")
            print("  ⚠️ Continuing validation without displaying samples...", flush=True)

    # Get timing metrics
    timing_metrics = timer.get_timing_metrics(reduction_op="sum")
    validation_time = timing_metrics.get("total_validation_time", 0)

    # Print summary of validation results
    print("\n📊 Validation Results:")
    print(f"    • Accuracy: {accuracy:.4f}")
    print(f"    • Average response length: {avg_length:.1f} tokens")
    print(f"    • Samples processed: {len(total_rewards)}", flush=True)

    # Print timing information
    print("\n  ⏱️  Validation Timing:")
    validation_time = timing_metrics.get("total_validation_time", 0)
    print(f"    • Total validation time: {validation_time:.2f}s", flush=True)

    # Log validation data to JSONL file
    if logger is not None:
        val_log_data = {
            "content": all_message_logs,
            "rewards": total_rewards,
        }
        logger.log_batched_dict_as_jsonl(val_log_data, f"val_data_step{step}.jsonl")

    # Make sure to reset the timer after validation
    timer.reset()

    # Explicit GPU memory cleanup after validation
    gc.collect()
    torch.cuda.empty_cache()

    return val_metrics, timing_metrics


def aggregate_rollout_metrics(
    per_group_metrics: dict[str, list],
) -> dict[str, Any]:
    """Aggregate rollout metrics from multiple trajectory groups.

    Different metric types are aggregated according to their semantics:
    - Histogram observations: flattened into one step-level distribution
    - Metrics ending with "/min" or starting with "min_" (excluding "_rate" suffix): take the minimum
    - Metrics ending with "/max" or starting with "max_" (excluding "_rate" suffix): take the maximum
    - "total_turns": summed
    - Non-numeric values: passed through as-is
    - All other numeric metrics: averaged

    Args:
        per_group_metrics: A dict mapping metric names to lists of per-group values.

    Returns:
        A dict mapping metric names to their aggregated scalar values.
    """
    aggregated = {}
    for k, v in per_group_metrics.items():
        if is_histogram_metric(k):
            aggregated[k] = [observation for group in v for observation in group]
        elif not isinstance(v[0], (int, float)):
            aggregated[k] = v
        elif k.endswith("/min") or (k.startswith("min_") and not k.endswith("_rate")):
            aggregated[k] = min(v)
        elif k.endswith("/max") or (k.startswith("max_") and not k.endswith("_rate")):
            aggregated[k] = max(v)
        elif k == "total_turns":
            aggregated[k] = sum(v)
        elif k == "trajectory_duration_s":
            sorted_v = sorted(v)
            p95_idx = min(int(len(sorted_v) * 0.95), len(sorted_v) - 1)
            aggregated[k] = sum(v) / len(v)
            aggregated["trajectory_duration_s/max"] = max(v)
            aggregated["trajectory_duration_s/p95"] = (
                sorted_v[p95_idx] if sorted_v else 0
            )
        else:
            aggregated[k] = sum(v) / len(v)
    return aggregated


def _startup_pipeline_ready(
    replay_buffer: Any,
    collector_status: dict[str, Any],
    *,
    current_step_ready: bool,
    step: int,
    num_prompts_per_step: int,
    max_trajectory_age_steps: int,
    max_num_steps: int,
) -> bool:
    """Return whether async training can overlap safely with lookahead generation.

    A restored buffer may already contain ``step``, letting startup consume it
    before the collector generates ``step + 1``. After the refit that follows
    ``step``, the collector's target window advances to ``step + 2``, so an
    unclaimed ``step + 1`` would then be permanently missing. Requiring the
    next target to be complete *or* actively claimed keeps that gap closed
    while preserving the training/generation overlap used by steady-state async
    training, and without letting an unrelated reservation open the barrier.
    """
    if not current_step_ready:
        return False

    next_step = step + 1
    need_lookahead = max_trajectory_age_steps > 0 and next_step < max_num_steps
    if not need_lookahead:
        return True

    next_step_ready = ray.get(
        replay_buffer.has_complete_batch.remote(
            next_step, num_prompts_per_step, max_trajectory_age_steps
        )
    )
    return next_step_ready or next_step in collector_status.get(
        "generating_targets", ()
    )


def _raise_if_collector_stopped(
    collector_status: dict[str, Any],
    *,
    awaited_target: str,
    awaited_work: str,
    action: str,
) -> None:
    """Raise if the collector stopped terminally while training waits on it.

    Both ``data_exhausted`` and ``errored`` are terminal — neither is reset once
    set — so a stopped collector with no in-flight workers can never supply the
    awaited target. Waiting longer would hang silently instead of failing.

    Args:
        collector_status: Snapshot from ``AsyncTrajectoryCollector.get_status``.
        awaited_target: Rendered target being awaited, e.g. ``"target=5"`` or
            ``"training_step=5"``.
        awaited_work: What is being awaited, e.g. ``"lookahead claim"``.
        action: Verb for the message — ``"start"`` or ``"continue"``.
    """
    if not (
        (collector_status["data_exhausted"] or collector_status.get("errored", False))
        and not collector_status["running"]
        and collector_status["inflight_workers"] == 0
    ):
        return

    stop_reason = (
        "dataloader exhausted"
        if collector_status["data_exhausted"]
        else "collector errored"
    )
    recovery_advice = (
        "Check the training dataset and dataloader configuration."
        if collector_status["data_exhausted"]
        else "Inspect the preceding trajectory collector error."
    )
    raise RuntimeError(
        f"Trajectory collector stopped ({stop_reason}) while waiting for "
        f"{awaited_work} at {awaited_target}. "
        f"Training cannot {action} without the required trajectories. "
        f"Collector status: {collector_status}. "
        f"{recovery_advice}"
    )


@umbrella_trace_fn(RLSpanGroup.U_JOB, "rl.grpo.job")
def async_grpo_train(
    policy: ColocatablePolicyInterface,
    policy_generation: Optional[GenerationInterface],
    dataloader: StatefulDataLoader,
    val_dataloader: Optional[StatefulDataLoader],
    tokenizer: TokenizerType,
    loss_fn: LossFunction,
    task_to_env: dict[str, EnvironmentInterface],
    val_task_to_env: Optional[dict[str, EnvironmentInterface]],
    logger: Logger,
    checkpointer: CheckpointManager,
    grpo_save_state: GRPOSaveState,
    master_config: MasterConfig,
    max_trajectory_age_steps: int = 1,
    teacher_worker_groups: Optional[dict[str, Any]] = None,
    alias_to_group_alias: Optional[dict[str, str]] = None,
    processor: Optional[AutoProcessor] = None,
) -> None:
    """Run asynchronous GRPO training with replay buffer.

    Args:
        policy: Training policy
        policy_generation: Generation interface
        dataloader: Training data loader
        val_dataloader: Validation data loader
        tokenizer: Tokenizer
        loss_fn: Loss function
        task_to_env: Training environments
        val_task_to_env: Validation environments
        logger: Logger
        checkpointer: Checkpoint manager
        grpo_save_state: Training state
        master_config: Master configuration
        max_trajectory_age_steps: Maximum age (in training steps) for trajectories to be used in training
        processor: Optional multimodal processor used to attach compact policy
            media to NeMo Gym prompt rows.
    """
    # Ensure we are running with a compatible async generation backend.
    # Async GRPO supports vLLM, Megatron, TRT-LLM, and Dynamo;
    # SGLang async rollouts do not support the async GRPO replay path.
    generation_config = master_config.policy["generation"]
    backend = generation_config.get("backend", "") if generation_config else ""
    assert backend in ("vllm", "megatron", "trtllm", "dynamo"), (
        "Async GRPO supports the vLLM, Megatron, TRT-LLM, and Dynamo generation backends; "
        f"got policy.generation.backend={backend!r}."
    )
    assert should_use_async_rollouts(generation_config), (
        "Async GRPO requires Dynamo, Megatron, or an async vLLM or TRT-LLM "
        "generation engine. Set policy.generation.backend=dynamo, "
        "policy.generation.vllm_cfg.async_engine=true (vLLM), or "
        "policy.generation.trtllm_cfg.async_engine=true (TRT-LLM). "
        "Megatron Inference always uses its async engine."
    )
    assert master_config.loss_fn.use_importance_sampling_correction, (
        "Importance sampling correction must be enabled for async GRPO for good convergence due to off-policy samples!"
    )
    max_generation_failures = master_config.grpo.async_grpo.max_generation_failures

    if router_replay_enabled(master_config.policy) and (
        master_config.data_plane or {}
    ).get("enabled", False):
        raise NotImplementedError(
            "policy.router_replay.enabled=true with async GRPO on this "
            "entrypoint is supported only when data_plane.enabled=false. For "
            "async + TransferQueue, use the SingleController entrypoint: "
            "examples/run_grpo_single_controller.py with e.g. "
            "examples/configs/recipes/llm/"
            "grpo-qwen3-30ba3b-10n8g-megatron-cp2-r3-async-single-controller.yaml"
        )

    if master_config.grpo.async_grpo.max_trajectory_age_steps > 1:
        if not master_config.grpo.async_grpo.in_flight_weight_updates:
            print(
                "⚠️ WARNING: In-flight weight updates must be enabled for async GRPO with max_trajectory_age_steps > 1. "
                "Without in-flight weight updates, having more max_trajectory_age_steps will not give any performance benefit."
            )

    # Import async utilities only when needed
    from nemo_rl.algorithms.async_utils import AsyncTrajectoryCollector, ReplayBuffer

    timer = Timer(context={"worker": "driver"})
    _telemetry = get_telemetry_handle()
    _tracer = _telemetry.tracer if _telemetry is not None else None
    training_wall_start = time.perf_counter()
    timeout = TimeoutChecker(
        timeout=master_config.checkpointing["checkpoint_must_save_by"],
        fit_last_save_time=True,
    )
    timeout.start_iterations()
    assert policy_generation is not None

    # Training state
    step = grpo_save_state.current_step
    max_num_epochs = master_config.grpo.max_num_epochs
    if max_num_epochs is not None and max_num_epochs > 0:
        master_config.grpo.max_num_steps = min(
            master_config.grpo.max_num_steps,
            max_num_epochs * len(dataloader),
        )
    max_num_steps = master_config.grpo.max_num_steps
    if step >= max_num_steps:
        print(
            "Async GRPO training is already complete: "
            f"current step {step} reached the effective limit of "
            f"{max_num_steps} steps.",
            flush=True,
        )
        checkpointer.shutdown()
        return

    POLICY_GENERATION_STALE = _initial_policy_generation_stale(policy_generation, step)
    weight_version = step  # Tracks refitted weight versions
    consumed_samples = grpo_save_state.consumed_samples
    total_valid_tokens = grpo_save_state.total_valid_tokens
    val_period = master_config.grpo.val_period
    val_start_at = master_config.grpo.val_start_at
    val_at_start = master_config.grpo.val_at_start
    val_at_end = master_config.grpo.val_at_end
    colocated_inference = master_config.policy["generation"]["colocated"]["enabled"]
    stop_at_validation_threshold = master_config.grpo.stop_at_validation_threshold
    stop_at_validation_metric = master_config.grpo.stop_at_validation_metric

    assert (not colocated_inference) or (
        isinstance(policy_generation, MegatronGeneration)
    ), "Colocated async GRPO is only supported for the Megatron generation backend."

    # Initialize advantage estimator
    adv_estimator = _create_advantage_estimator(master_config)

    # Calculate minimum buffer size from training requirements
    # In per-prompt buffer mode, one buffer entry is 1 prompt * num_generations_per_prompt
    num_prompts_per_step = master_config.grpo.num_prompts_per_step
    samples_per_prompt_group = master_config.grpo.num_generations_per_prompt
    train_gbs = master_config.policy["train_global_batch_size"]

    # Ensure the buffer has at least one step worth of prompt-groups before training
    min_trajectories_needed = num_prompts_per_step

    print("📊 Buffer requirements calculation:", flush=True)
    print(f"   - num_prompts_per_step: {num_prompts_per_step}")
    print(f"   - num_generations_per_prompt: {samples_per_prompt_group}")
    print(f"   - samples_per_prompt_group: {samples_per_prompt_group}")
    print(f"   - train_global_batch_size: {train_gbs}")
    print(f"   - min_trajectories_needed: {min_trajectories_needed} (async mode)")

    _replay_runtime_env = make_actor_runtime_env(
        "nemo_rl.algorithms.async_utils.ReplayBuffer"
    )

    # Calculate optimal buffer size based on generation limits to prevent length bias
    # Each weight version generates exactly num_prompts_per_step trajectories
    # With max_age_steps, we keep trajectories from multiple weight versions
    num_prompts_per_step = master_config.grpo.num_prompts_per_step
    late_arrival_slack = 2
    optimal_buffer_size = (
        num_prompts_per_step * max_trajectory_age_steps * late_arrival_slack
    )

    replay_buffer = ReplayBuffer.options(runtime_env=_replay_runtime_env).remote(
        max_size=optimal_buffer_size,
        drop_incomplete_targets_on_restore=False,
    )

    last_checkpoint_path = checkpointer.get_latest_checkpoint_path()
    replay_buffer_restore_metadata: dict[str, Any] | None = None
    rollouts_state = None
    if last_checkpoint_path is not None:
        replay_buffer_restore_metadata = _maybe_restore_async_replay_buffer_checkpoint(
            replay_buffer,
            last_checkpoint_path,
            load_replay_buffer=master_config.checkpointing.get("load_replay_buffer"),
            num_prompts_per_step=num_prompts_per_step,
            current_training_step=step,
            max_age_steps=max_trajectory_age_steps,
        )

        rollouts_path = os.path.join(last_checkpoint_path, "rollouts.pt")
        if os.path.exists(rollouts_path):
            # weights_only=False: this is a trusted same-job checkpoint artifact.
            rollouts_state = torch.load(rollouts_path, weights_only=False)

    next_nemo_gym_task_index = max(
        int((rollouts_state or {}).get(NEXT_NEMO_GYM_TASK_INDEX_KEY, 0)),
        int(
            (replay_buffer_restore_metadata or {}).get(NEXT_NEMO_GYM_TASK_INDEX_KEY, 0)
        ),
    )

    # Frontier-aligned resume: the checkpoint saved the dataloader state at
    # the trained frontier rather than the live cursor, so the collector
    # re-yields the covered window and regenerates every prompt that is
    # neither trained nor retained in the restored buffer. Legacy checkpoints
    # (no frontier metadata) keep today's behavior.
    frontier_ordinal = (rollouts_state or {}).get(FRONTIER_ORDINAL_KEY)
    resume_base_ordinal = (rollouts_state or {}).get(RESUME_BASE_ORDINAL_KEY)
    frontier_restore = frontier_ordinal is not None and resume_base_ordinal is not None
    if frontier_restore:
        retained_task_indices = list(
            (replay_buffer_restore_metadata or {}).get(RETAINED_TASK_INDICES_KEY, [])
        )
        # Ordinals trained at/above the cut, covered like retained groups so
        # the re-yielded window regenerates only what was lost.
        trained_task_indices = [
            int(ordinal)
            for ordinal in (rollouts_state or {}).get(TRAINED_TASK_INDICES_KEY, [])
        ]
        covered_task_indices = sorted(
            set(retained_task_indices) | set(trained_task_indices)
        )
        collector_start_kwargs: dict[str, Any] = {
            "next_nemo_gym_task_index": int(resume_base_ordinal),
            "resume_frontier_ordinal": int(frontier_ordinal),
            "resume_covered_task_indices": covered_task_indices,
            # The rewound dataloader re-yields any carried-over remainder.
            "pending_batch": None,
            "ordinals_frontier_aligned": True,
        }
        print(
            "📦 Frontier-aligned resume: dataloader rewound to ordinal "
            f"{resume_base_ordinal}, trained frontier {frontier_ordinal}, "
            f"{len(retained_task_indices)} retained prompt groups, "
            f"{len(trained_task_indices)} trained above the cut"
        )
    else:
        collector_start_kwargs = {
            "next_nemo_gym_task_index": next_nemo_gym_task_index,
            "pending_batch": (rollouts_state or {}).get(PENDING_PROMPTS_KEY),
            # Ordinal == stream position only holds for runs that have used
            # frontier-aligned checkpoints from the start; a legacy resume
            # keeps live-cursor checkpoints.
            "ordinals_frontier_aligned": last_checkpoint_path is None,
        }
        if last_checkpoint_path is not None:
            print(
                "⚠️ Legacy checkpoint resume: frontier-aligned checkpointing "
                "is disabled for this run and every checkpoint descended from "
                "it. Checkpoints will save the live dataloader cursor, so a "
                "resume may skip prompts that were in flight at the save."
            )

    # High-water mark of trained group ordinals, exclusive — the checkpoint
    # frontier. consumed_samples cannot serve here: tolerated generation
    # failures leave stream holes it never sees, so it lags the true stream
    # position.
    trained_frontier_ordinal = (
        int(frontier_ordinal) if frontier_ordinal is not None else consumed_samples
    )
    # Trained ordinals at/above the last checkpoint cut, persisted so a
    # resume covers them instead of re-training them. Pruned at each save;
    # the cut never decreases, so pruning is safe.
    recent_trained_task_indices: set[int] = (
        set(trained_task_indices) if frontier_restore else set()
    )

    _tc_runtime_env = make_actor_runtime_env(
        "nemo_rl.algorithms.async_utils.AsyncTrajectoryCollector",
        extra_env_vars={
            # Names this actor's spans the way RayWorkerGroup names its groups'.
            WORKER_GROUP_ENV: "trajectory_collector",
        },
    )

    # Captured inside rl.grpo.job so the collector's spans join this trace.
    # Empty if the job group is off, which degrades to roots.
    _tc_trace_carrier = current_trace_carrier()

    # Initialize trajectory collector with synchronized collection
    trajectory_collector = AsyncTrajectoryCollector.options(
        runtime_env=_tc_runtime_env
    ).remote(
        policy_generation=policy_generation,
        tokenizer=tokenizer,
        task_to_env=task_to_env,
        master_config=master_config,
        replay_buffer=replay_buffer,
        start_step=step,
        teacher_worker_groups=teacher_worker_groups,
        alias_to_group_alias=alias_to_group_alias,
        on_policy_distillation_cfg=opd_module._opd_cfg(master_config),
        processor=processor,
        trace_carrier=_tc_trace_carrier,
        **collector_start_kwargs,
    )

    def _flush_collector_telemetry() -> None:
        """Export the collector's buffered spans before the actor is reaped.

        ``ray.kill`` runs no atexit handler, so whatever the span processor has
        not sent yet goes with the actor -- including the last rollout batches
        of the run. Every path that reaps the collector needs this, not only the
        normal one, and collection is already running by the time this is
        defined. The timeout covers the callee's quiesce budget *plus* its
        export budget; too short and it gives up mid-export, dropping the very
        spans it exists to save. ``shutdown_telemetry`` enforces the latter
        itself -- the SDK ignores the timeout it is handed -- so this stays
        ahead of it only as long as that remains true.
        """
        try:
            ray.get(
                trajectory_collector.flush_telemetry.remote(quiesce_timeout_s=3.0),
                timeout=15,
            )
        except Exception as e:
            print(f"Error flushing trajectory collector telemetry: {e}")

    print(
        f"🚀 Starting async GRPO training with buffer_size={optimal_buffer_size}, "
        f"max_age={max_trajectory_age_steps} steps, "
        f"max_generation_failures={max_generation_failures}"
    )

    print("⏳ Preparing policy generation for training...", flush=True)
    if POLICY_GENERATION_STALE:
        print("🔄 Refitting policy generation with actual model weights...", flush=True)
        try:
            refit_policy_generation(
                policy,
                policy_generation,
                colocated_inference,
            )
            print("✅ Policy generation refit completed successfully", flush=True)
            POLICY_GENERATION_STALE = False
        except Exception as e:
            print(f"❌ Policy generation refit failed: {e}")
            import traceback

            traceback.print_exc()
            _flush_collector_telemetry()
            raise
    else:
        print("🔄 Preparing policy generation for inference...")
        try:
            policy_generation.prepare_for_generation()
            print("✅ Policy generation preparation completed successfully")
        except Exception as e:
            print(f"❌ Policy generation preparation failed: {e}")
            import traceback

            traceback.print_exc()
            _flush_collector_telemetry()
            raise

    # Generation must hold the policy's real weights before any backend starts
    # collecting. In particular, vLLM and Dynamo start with dummy weights when
    # the first refit supplies model parameters.
    ray.get(trajectory_collector.set_weight_version.remote(weight_version))
    ray.get(trajectory_collector.start_collection.remote(CyclingDataLoader(dataloader)))
    print("📦 Started continuous background trajectory collection")

    print("✅ Policy generation setup complete, proceeding to validation...")

    # Run validation at start if configured
    if val_at_start and step == 0:
        print("\n🔍 Running initial validation...")
        # Pause trajectory collection during initial validation
        ray.get(trajectory_collector.pause.remote())

        initial_val_metrics: Optional[dict[str, Any]] = None
        try:
            val_metrics, validation_timings = validate(
                policy_generation,
                val_dataloader,
                tokenizer,
                val_task_to_env,
                step=0,
                master_config=master_config,
                logger=logger,
                processor=processor,
            )
            initial_val_metrics = val_metrics
            # A colocated engine keeps serving between phases (preserves its
            # KV/prefix cache); the backend makes that call, not the loop.
            policy_generation.finish_generation(release_gpu=False)
            logger.log_metrics(val_metrics, step, prefix="validation")
            logger.log_metrics(validation_timings, step, prefix="timing/validation")
            if master_config.grpo.debug_payload_metrics:
                validation_payload_metrics = drain_multimodal_payload_metrics()
                if validation_payload_metrics:
                    logger.log_metrics(
                        validation_payload_metrics,
                        step,
                        prefix="validation",
                    )
            print("✅ Initial validation completed successfully")
        except Exception as e:
            print(f"❌ Initial validation failed: {e}")
            import traceback

            traceback.print_exc()
            # Continue anyway since validation is optional
        finally:
            # Resume trajectory collection after initial validation
            trajectory_collector.resume.remote()

        stop_message = (
            _validation_early_stop_message(
                initial_val_metrics,
                stop_at_validation_threshold,
                stop_at_validation_metric,
                initial=True,
            )
            if initial_val_metrics is not None
            else None
        )
        if stop_message is not None:
            print(stop_message, flush=True)
            # Flush pending checkpoint finalization and stop rollout
            # generation; the remaining actors are reaped when the driver
            # exits right after this return.
            checkpointer.shutdown()
            _flush_collector_telemetry()
            try:
                ray.kill(trajectory_collector)
            except Exception as e:
                print(f"Error stopping trajectory collector: {e}")
            try:
                ray.kill(replay_buffer)
            except Exception as e:
                print(f"Error stopping replay buffer: {e}")
            return

    print("✅ All setup complete, starting buffer wait...")
    # Clear logger metrics at start of training
    if policy_generation is not None:
        policy_generation.clear_logger_metrics()

    # Wait for initial buffer fill for the current training step.
    print(
        f"⏳ Waiting for replay buffer to have sufficient trajectories for step {step}..."
    )
    # Spanned as well as timed so the wait shows up in the startup waterfall
    # rather than only as a scalar. Unbucketed by construction (see
    # UNBUCKETED_SPAN_CATEGORIES): the generation fleet is busy for this whole
    # window, and its spans join this trace.
    with efficiency_span(INIT_TOTAL_CATEGORY, tracer=_tracer):
        timer.start(INIT_TOTAL_CATEGORY)
        wait_iterations = 0
        while True:
            buffer_size_current = ray.get(replay_buffer.size.remote())
            ray.get(trajectory_collector.check_health.remote())
            current_step_ready = ray.get(
                replay_buffer.has_complete_batch.remote(
                    step, num_prompts_per_step, max_trajectory_age_steps
                )
            )

            print(
                f"  Wait iteration {wait_iterations}: buffer_size={buffer_size_current}, "
                f"step {step} ready={current_step_ready}",
                flush=True,
            )

            collector_status = ray.get(trajectory_collector.get_status.remote())
            pipeline_ready = _startup_pipeline_ready(
                replay_buffer,
                collector_status,
                current_step_ready=current_step_ready,
                step=step,
                num_prompts_per_step=num_prompts_per_step,
                max_trajectory_age_steps=max_trajectory_age_steps,
                max_num_steps=master_config.grpo.max_num_steps,
            )
            if current_step_ready and not pipeline_ready:
                print(
                    f"  Pipeline barrier: step {step} ready but "
                    f"step {step + 1} is not yet claimed — waiting for lookahead "
                    f"to prevent resume deadlock"
                )

            if pipeline_ready:
                break

            trajectories_needed = ray.get(
                replay_buffer.get_trajectories_needed.remote(
                    step, num_prompts_per_step, max_trajectory_age_steps
                )
            )
            if (
                buffer_size_current >= min_trajectories_needed
                and trajectories_needed > 0
            ):
                print(
                    f"  ⏳ Gap-filling in progress: need {trajectories_needed} more "
                    f"trajectories for step {step}"
                )

            awaited_target = step + 1 if current_step_ready else step
            _raise_if_collector_stopped(
                collector_status,
                awaited_target=f"target={awaited_target}",
                awaited_work="lookahead claim" if current_step_ready else "buffer fill",
                action="start",
            )

            wait_iterations += 1
            time.sleep(1.0)

        # Retained because the per-step timer.reset() below discards it; the
        # efficiency snapshot re-supplies it every step.
        init_total_s = timer.stop(INIT_TOTAL_CATEGORY)
    print(f"✅ Buffer ready for step {step}! Starting training loop...")

    ft_save_period = master_config.checkpointing.get("ft_save_period")

    # Main training loop
    try:
        while step < max_num_steps:
            ray.get(trajectory_collector.check_health.remote())
            refit_metrics: dict[str, float] = {}
            early_stop_message: Optional[str] = None
            print(f"\n{'=' * 25} Step {step + 1}/{max_num_steps} {'=' * 25}")
            maybe_gpu_profile_step(policy, step + 1)
            if policy != policy_generation:
                maybe_gpu_profile_step(policy_generation, step + 1)

            with (
                iteration_scope(step + 1),
                timer.time("total_step_time"),
                umbrella_span(
                    RLSpanGroup.U_STEP,
                    "rl.grpo.step",
                    tracer=_tracer,
                ),
            ):
                num_mask_sample_filtered = 0

                # Sample trajectories from replay buffer
                print("📦 Sampling from replay buffer...")
                with timer.time("exposed_generation"):
                    buffer_size_current = ray.get(replay_buffer.size.remote())
                    print(
                        f"📊 Step coordination: training_step={step}, max_age={max_trajectory_age_steps}, buffer_size={buffer_size_current}",
                        flush=True,
                    )

                    # Sample the required number of per-prompt groups.
                    num_prompt_groups_needed = master_config.grpo.num_prompts_per_step
                    sample_result = ray.get(
                        replay_buffer.sample.remote(
                            num_prompt_groups=num_prompt_groups_needed,
                            current_weight_version=weight_version,
                            max_age_steps=max_trajectory_age_steps,
                        )
                    )
                    if sample_result is not None:
                        print_multimodal_payload_metrics(
                            collect_multimodal_payload_metrics(
                                sample_result,
                                "replay_sample",
                                enabled=master_config.grpo.debug_payload_metrics,
                            )
                        )

                    if (
                        sample_result is None
                        or len(sample_result["trajectories"])
                        != num_prompt_groups_needed
                    ):
                        print(
                            "⏳ Buffer empty or not enough groups to form a full step, waiting..."
                        )

                        # Get buffer debug info to help diagnose the issue
                        buffer_debug = ray.get(replay_buffer.get_debug_info.remote())
                        buffer_size = buffer_debug["total_trajectories"]

                        if buffer_size > 0:
                            print(
                                f"🔍 Debug: Buffer has {buffer_size} trajectories but sampling requires exactly {num_prompt_groups_needed}."
                            )
                            print(f"   Current weight version: {weight_version}")
                            print(f"   Max trajectory age: {max_trajectory_age_steps}")
                            print(
                                f"   Trajectory versions in buffer: {buffer_debug['trajectory_versions']}"
                            )
                            diag = buffer_debug.get("starvation_diagnostics")
                            if diag:
                                print(
                                    "   📊 Buffer starvation diagnostics (long-tail root cause):"
                                )
                                print(
                                    f"      trajectory_duration_s: mean={diag['trajectory_duration_s']['mean']:.1f}s, "
                                    f"median={diag['trajectory_duration_s']['median']:.1f}s, "
                                    f"max={diag['trajectory_duration_s']['max']:.1f}s, "
                                    f"p95={diag['trajectory_duration_s']['p95']:.1f}s"
                                )
                                print(
                                    f"      max_gen_tokens_per_turn: mean={diag['max_gen_tokens_per_turn_in_buffer']['mean']:.0f}, "
                                    f"median={diag['max_gen_tokens_per_turn_in_buffer']['median']:.0f}, "
                                    f"max={diag['max_gen_tokens_per_turn_in_buffer']['max']:.0f}, "
                                    f"p95={diag['max_gen_tokens_per_turn_in_buffer']['p95']:.0f} "
                                    "(high = long single generations per turn)"
                                )
                                print(
                                    f"      turns_per_sample: mean={diag['turns_per_sample_in_buffer']['mean']:.1f}, "
                                    f"median={diag['turns_per_sample_in_buffer']['median']:.1f}, "
                                    f"max={diag['turns_per_sample_in_buffer']['max']:.0f}, "
                                    f"p95={diag['turns_per_sample_in_buffer']['p95']:.1f} "
                                    "(high = many turns per trajectory)"
                                )

                        collector_status = ray.get(
                            trajectory_collector.get_status.remote()
                        )
                        awaited_target = step
                        print(
                            f"   Awaiting target {awaited_target}; claimed by collector: "
                            f"{awaited_target in collector_status.get('generating_targets', ())}"
                        )
                        _raise_if_collector_stopped(
                            collector_status,
                            awaited_target=f"training_step={step}",
                            awaited_work="a full batch",
                            action="continue",
                        )

                        with (
                            timer.time("idle/buffer_starvation"),
                            efficiency_span("idle/buffer_starvation", tracer=_tracer),
                        ):
                            time.sleep(0.5)
                        continue

                    # Extract trajectories and metadata from sample result
                    trajectories = sample_result["trajectories"]
                    avg_trajectory_age = sample_result["avg_trajectory_age"]

                    # Advance the trained frontier from the sampled groups'
                    # own stream ordinals.
                    sampled_ordinals = [
                        trajectory.get(NEMO_GYM_TASK_INDEX_KEY)
                        for trajectory in trajectories
                        if isinstance(trajectory, dict)
                    ]
                    if sampled_ordinals and all(
                        ordinal is not None for ordinal in sampled_ordinals
                    ):
                        trained_frontier_ordinal = max(
                            trained_frontier_ordinal,
                            max(int(ordinal) for ordinal in sampled_ordinals) + 1,
                        )
                        recent_trained_task_indices.update(
                            int(ordinal) for ordinal in sampled_ordinals
                        )

                    print(
                        f"✅ Sampled {len(trajectories)} trajectory groups from buffer (avg age: {avg_trajectory_age:.2f} steps)"
                    )

                    # Concatenate per-prompt groups into a single training batch
                    per_prompt_batches = [t["batch"] for t in trajectories]
                    repeated_batch = BatchedDataDict.from_batches(
                        per_prompt_batches,
                        allow_missing_packed_tensors=(
                            master_config.grpo.deduplicate_multimodal_data
                        ),
                    )

                    # Teacher logprobs are stored in batch dict by collection-time
                    # computation and padded by from_batches. Extract here.
                    trajectory_teacher_logprobs = None
                    if opd_module.is_opd_enabled(master_config):
                        if "teacher_reference_logprobs" in repeated_batch:
                            trajectory_teacher_logprobs = repeated_batch[
                                "teacher_reference_logprobs"
                            ]

                    # Aggregate rollout metrics across groups with proper aggregation per metric type
                    per_group_metrics = {}
                    for t in trajectories:
                        for k, v in t["rollout_metrics"].items():
                            per_group_metrics.setdefault(k, []).append(v)
                    rollout_metrics = aggregate_rollout_metrics(per_group_metrics)

                # Enforce fixed training batch: num_prompts_per_step * num_generations_per_prompt
                expected_batch_size = (
                    master_config.grpo.num_prompts_per_step
                    * master_config.grpo.num_generations_per_prompt
                )
                if repeated_batch.size != expected_batch_size:
                    print(
                        f"❌ Unexpected training batch size: got {repeated_batch.size}, expected {expected_batch_size}. Skipping step and waiting for correct buffer content."
                    )
                    time.sleep(0.5)
                    continue

                # Optional sanity: ensure DP divisibility to avoid sharding issues
                dp_size = policy.sharding_annotations.get_axis_size("data_parallel")
                if expected_batch_size % dp_size != 0:
                    raise AssertionError(
                        f"Configuration error: (num_prompts_per_step * num_generations_per_prompt) = {expected_batch_size} must be divisible by data_parallel size {dp_size}."
                    )

                print(f"Got trajectory batch (size: {repeated_batch.size})")

                # Baseline spec-decode counters; the delta read at metrics time gives
                # MTP acceptance over this step's generation window (async generation
                # runs continuously in the background collector).
                if hasattr(policy_generation, "snapshot_step_metrics"):
                    policy_generation.snapshot_step_metrics()

                print("▶ Processing rewards...")
                with (
                    timer.time("reward_calculation"),
                    managed_span(
                        RLSpanGroup.REWARD, "rl.grpo.reward_calculation", tracer=_tracer
                    ),
                ):
                    # Backfill before the training flatten reuses these messages.
                    backfill_missing_routed_experts(repeated_batch["message_log"])

                    prompt_ids_for_adv = _trajectory_group_ids(trajectories)
                    if prompt_ids_for_adv.shape[0] != repeated_batch.size:
                        raise ValueError(
                            f"training batch has {repeated_batch.size} rows but replay "
                            f"groups account for {prompt_ids_for_adv.shape[0]}"
                        )

                    rewards = repeated_batch["total_reward"]

                    print(
                        f"  📊 Rewards stats: min={rewards.min():.4f}, max={rewards.max():.4f}, mean={rewards.mean():.4f}, std={rewards.std():.4f}"
                    )

                # Prepare training data (same as sync version)
                with (
                    timer.time("data_processing"),
                    managed_span(
                        RLSpanGroup.DATA_PROCESSING,
                        "rl.grpo.data_processing",
                        tracer=_tracer,
                    ),
                ):
                    # Apply overlong filtering - mask out truncated sequences from loss computation
                    with timer.time("overlong_filter"):
                        use_overlong_filtering = master_config.grpo.overlong_filtering
                        if use_overlong_filtering:
                            loss_multiplier = repeated_batch["loss_multiplier"].clone()
                            truncated = repeated_batch["truncated"]

                            if isinstance(truncated, list):
                                truncated = torch.tensor(truncated, dtype=torch.bool)

                            loss_multiplier[truncated] = 0
                            repeated_batch["loss_multiplier"] = loss_multiplier

                    with timer.time("mask_sample_filter"):
                        num_mask_sample_filtered = _apply_mask_sample_filter(
                            repeated_batch
                        )

                    # Add loss mask to each message
                    # Only unmask assistant messages that were actually generated (have generation_logprobs),
                    # not assistant messages that were part of the prompt history
                    add_grpo_token_loss_masks_and_generation_logprobs(
                        repeated_batch["message_log"]
                    )

                    # Convert to flat format for training
                    flat_messages, input_lengths = batched_message_log_to_flat_message(
                        repeated_batch["message_log"],
                        pad_value_dict={"token_ids": tokenizer.pad_token_id},
                        make_sequence_length_divisible_by=master_config.policy[
                            "make_sequence_length_divisible_by"
                        ],
                    )

                    # Create training data. Advantages are added after logprobs.
                    train_data = _build_async_grpo_train_data(
                        flat_messages,
                        input_lengths,
                        repeated_batch,
                        master_config.policy,
                    )
                    print_multimodal_payload_metrics(
                        collect_multimodal_payload_metrics(
                            train_data,
                            "rollout_to_policy_async",
                            enabled=master_config.grpo.debug_payload_metrics,
                        )
                    )
                    train_data.to("cpu")

                generation_logger_metrics = None
                if policy_generation.blocks_training():
                    print("⏸️ Pausing colocated engine + collector for training...")
                    with timer.time("exposed_generation"):
                        ray.get(trajectory_collector.prepare_for_refit.remote())
                    generation_logger_metrics = policy_generation.get_logger_metrics()
                    policy_generation.finish_generation(release_gpu=True)

                # Training phase (same as sync version)
                skip_prev_logprobs, skip_reference_logprobs = (
                    _resolve_logprob_skip_flags(master_config)
                )
                seq_logprob_error_threshold = (
                    master_config.grpo.seq_logprob_error_threshold
                )

                if not (skip_prev_logprobs and skip_reference_logprobs):
                    print("▶ Preparing for logprob inference...", flush=True)
                    with timer.time("logprob_inference_prep"):
                        policy.prepare_for_lp_inference()

                print("▶ Computing logprobs...", flush=True)
                with (
                    timer.time("policy_and_reference_logprobs"),
                    managed_span(
                        RLSpanGroup.LOGPROB,
                        "rl.grpo.policy_and_reference_logprobs",
                        tracer=_tracer,
                    ),
                ):
                    if not skip_prev_logprobs:
                        prev_lp_result = policy.get_logprobs(train_data, timer=timer)
                        train_data["prev_logprobs"] = prev_lp_result["logprobs"]
                    else:
                        train_data["prev_logprobs"] = torch.zeros_like(
                            train_data["generation_logprobs"]
                        )

                    if not skip_reference_logprobs:
                        train_data["reference_policy_logprobs"] = (
                            policy.get_reference_policy_logprobs(
                                train_data,
                                timer=timer,
                            )["reference_logprobs"]
                        )
                    else:
                        print(
                            "▶ Skipping reference_logprobs (skip_reference_policy_logprobs_calculation=True)...",
                            flush=True,
                        )
                        train_data["reference_policy_logprobs"] = torch.zeros_like(
                            train_data["prev_logprobs"]
                        )

                # Separate-pass seq-level metrics/masking require real prev_logprobs
                if skip_prev_logprobs:
                    # In-loss filtering reports counts through all_mb_metrics.
                    # Use {} so placeholder zeros cannot overwrite those counts
                    # when seq_logprob_error_metrics is merged after training.
                    # Otherwise, placeholder prev_logprobs cannot provide
                    # sequence-error metrics.
                    seq_logprob_error_metrics = (
                        {}
                        if master_config.loss_fn.seq_logprob_error_in_loss
                        else _placeholder_seq_logprob_error_metrics()
                    )
                else:
                    seq_error_result = compute_and_apply_seq_logprob_error_masking(
                        train_data=train_data,
                        rewards=rewards,
                        seq_logprob_error_threshold=seq_logprob_error_threshold,
                        filtering_on=need_top_k_or_top_p_filtering(
                            TrainingSamplingParams(
                                top_k=master_config.policy["generation"]["top_k"],
                                top_p=master_config.policy["generation"]["top_p"],
                            )
                        ),
                    )
                    seq_logprob_error_metrics = seq_error_result
                    if "num_masked_seqs" in seq_logprob_error_metrics:
                        seq_logprob_error_metrics[
                            "num_masked_seqs_by_logprob_error"
                        ] = seq_logprob_error_metrics.pop("num_masked_seqs")

                # Pad teacher logprobs to match train_data sequence length.
                if trajectory_teacher_logprobs is not None:
                    trajectory_teacher_logprobs = _pad_teacher_logprobs(
                        trajectory_teacher_logprobs, train_data["input_ids"].shape[1]
                    )

                # Compute advantages with adv_estimator using correct mask and logprobs
                with (
                    timer.time("advantage_calculation"),
                    managed_span(
                        RLSpanGroup.ADVANTAGE,
                        "rl.grpo.advantage_calculation",
                        tracer=_tracer,
                    ),
                ):
                    print("▶ Computing advantages...", flush=True)
                    # Get token-level mask: token_mask * sample_mask
                    token_mask = train_data["token_mask"]
                    sample_mask = train_data["sample_mask"]
                    mask = token_mask * sample_mask.unsqueeze(-1)

                    train_data["advantages"] = adv_estimator.compute_advantage(
                        prompt_ids=prompt_ids_for_adv,
                        rewards=rewards,
                        mask=mask,
                        repeated_batch=repeated_batch,
                        logprobs_policy=train_data["prev_logprobs"],
                        logprobs_reference=train_data.get("reference_policy_logprobs"),
                        # OPD kwargs (ignored by non-OPD estimators via **kwargs)
                        teacher_logprobs=trajectory_teacher_logprobs.to(
                            train_data["prev_logprobs"].device
                        )
                        if trajectory_teacher_logprobs is not None
                        else None,
                        prev_logprobs=train_data["prev_logprobs"],
                        generation_logprobs=train_data["generation_logprobs"],
                        sample_mask=train_data["sample_mask"],
                    )
                    if (
                        hasattr(adv_estimator, "last_metrics")
                        and adv_estimator.last_metrics
                    ):
                        rollout_metrics.update(adv_estimator.last_metrics)
                    del prompt_ids_for_adv

                    # Log advantages stats
                    # Note: For GRPOAdvantageEstimator with normalize_rewards=True, these are
                    # already normalized advantages (equivalent to "Normalized advantages stats"
                    # in older versions). For ReinforcePlusPlusAdvantageEstimator, advantages
                    # are globally normalized across valid tokens.
                    advantages = train_data["advantages"]
                    print(
                        f"  📊 Advantages stats: min={advantages.min():.4f}, max={advantages.max():.4f}, mean={advantages.mean():.4f}, std={advantages.std():.4f}"
                    )

                    penalty_metrics = (
                        _apply_configured_message_level_advantage_penalties(
                            train_data,
                            repeated_batch["message_log"],
                            master_config,
                            log_config=True,
                        )
                    )

                    # Clip advantages to prevent extreme values from small std normalization
                    train_data["advantages"] = _clip_grpo_advantages(
                        train_data["advantages"], master_config.grpo
                    )

                print("▶ Preparing for training...")
                with timer.time("training_prep"):
                    policy.prepare_for_training()
                    POLICY_GENERATION_STALE = True

                print("▶ Training policy...")
                with (
                    timer.time("policy_training"),
                    managed_span(
                        RLSpanGroup.POLICY_UPDATE,
                        "rl.grpo.policy_training",
                        tracer=_tracer,
                    ),
                ):
                    train_results = policy.train(
                        train_data,
                        loss_fn,
                        timer=timer,
                    )

                is_last_step = step + 1 == max_num_steps
                should_save_by_step = (
                    is_last_step
                    or (step + 1) % master_config.checkpointing["save_period"] == 0
                    or (ft_save_period is not None and (step + 1) % ft_save_period == 0)
                )
                # Checked pre-validation so the wake-deferral below can see it.
                # A crossing during refit/validation is caught by the lookahead in check_save.
                should_save_by_timeout = timeout.check_save()
                will_save_checkpoint = master_config.checkpointing["enabled"] and (
                    should_save_by_step or should_save_by_timeout
                )
                # An early stop (known only after validation) also saves.
                saving_this_step = will_save_checkpoint
                # Save-bound colocated steps leave the engine asleep through save with no transfer.
                defer_wake_for_save = (
                    policy_generation.blocks_training()
                    and will_save_checkpoint
                    and policy_generation.wake_carries_weight_updates()
                )

                print("🔄 Synchronizing policy weights to trajectory collector…")
                if defer_wake_for_save:
                    # Wake-deferral (checkpoint scheduling, which the backend
                    # cannot see): the engine is about to be saved, so leave it
                    # asleep; just drop training-only buffers and version-stamp
                    # the weights. The post-save block wakes it and resumes
                    # collection.
                    print("⏸️ Keeping colocated engine asleep for checkpointing...")
                    # Seed the category with 0.0 (no refit wake happens on
                    # save-bound steps) so efficiency summaries, which skip
                    # missing keys, stay comparable across modes.
                    with timer.time("idle/refit_bubble"):
                        pass
                    with timer.time("offload_before_refit"):
                        policy.offload_before_refit()
                    POLICY_GENERATION_STALE = False
                    weight_version += 1
                    ray.get(
                        trajectory_collector.set_weight_version.remote(weight_version)
                    )
                else:
                    # A context manager rather than start/stop, so the timer is
                    # stopped and the span closed even if the refit raises.
                    with (
                        timer.time("idle/refit_bubble"),
                        efficiency_span("idle/refit_bubble", tracer=_tracer),
                    ):
                        # Measure pending-generation wait as exposed_generation time
                        print(
                            "🔄 Coordinating with trajectory collector before refit..."
                        )
                        with timer.time("exposed_generation"):
                            ray.get(trajectory_collector.prepare_for_refit.remote())

                        # Collect generation logger metrics for performance reporting
                        # inflight batch sizes and num pending samples are collected from each worker
                        # (colocated collects them before the engine sleeps for training).
                        if generation_logger_metrics is None:
                            generation_logger_metrics = (
                                policy_generation.get_logger_metrics()
                            )

                        # Only the actual refit/weight transfer should be counted as weight_sync
                        print("🔄 Performing policy generation refit...")
                        with timer.time("weight_sync"):
                            refit_metrics = refit_policy_generation(
                                policy,
                                policy_generation,
                                colocated_inference,
                            )
                            POLICY_GENERATION_STALE = False

                            # Update weight version before resuming trajectory collection so that all trajectories are updated with the new correct weight version
                            weight_version += 1
                            ray.get(
                                trajectory_collector.set_weight_version.remote(
                                    weight_version
                                )
                            )
                            ray.get(trajectory_collector.resume_after_refit.remote())

                # Clear logger metrics after each refit (weight sync), starting a new logging cycle
                if policy_generation is not None:
                    policy_generation.clear_logger_metrics()

                # Validation
                val_metrics, validation_timings = None, None
                should_run_validation = (
                    val_period > 0
                    and (step + 1) >= val_start_at
                    and (step + 1) % val_period == 0
                ) or (val_at_end and is_last_step)

                payload_metrics: dict[str, int | float] = {}
                if should_run_validation:
                    # Stop new dispatch before separating the training and
                    # validation payload-metric intervals.
                    ray.get(trajectory_collector.pause.remote())
                    if master_config.grpo.debug_payload_metrics:
                        payload_metrics = merge_multimodal_payload_metrics(
                            [
                                drain_multimodal_payload_metrics(),
                                ray.get(
                                    trajectory_collector.drain_payload_metrics.remote()
                                ),
                            ]
                        )

                # Run validation if it's a validation step or last step with val_at_end
                if should_run_validation:
                    # Timer only, no efficiency_span: validate() accounts this
                    # window as overhead (see the bucket_scope in validate), so
                    # an idle-bucketed span over the same interval would both
                    # contradict that label and, on the sync rollout path where
                    # the generate spans below carry it, be double-counted by a
                    # rollup that sums durations by rl.bucket. The metric has no
                    # such hierarchy and stays correct.
                    with timer.time("idle/validation"):
                        # No-op on an already-running engine;
                        # wakes the colocated engine when it stayed asleep for a save-bound step.
                        policy_generation.prepare_for_generation()
                        val_metrics, validation_timings = validate(
                            policy_generation,
                            val_dataloader,
                            tokenizer,
                            val_task_to_env,
                            step=step + 1,
                            master_config=master_config,
                            logger=logger,
                            processor=processor,
                        )
                        # An early stop triggers a save; must note before engine wake/resume.
                        early_stop_message = _validation_early_stop_message(
                            val_metrics,
                            stop_at_validation_threshold,
                            stop_at_validation_metric,
                        )
                        saving_this_step = will_save_checkpoint or (
                            master_config.checkpointing["enabled"]
                            and early_stop_message is not None
                        )
                        # Save-bound steps need the GPUs for checkpointing,
                        # so the engine must stand down; otherwise a colocated
                        # engine keeps serving (backend's call).
                        policy_generation.finish_generation(
                            release_gpu=saving_this_step
                        )
                        logger.log_metrics(
                            validation_timings, step + 1, prefix="timing/validation"
                        )
                        logger.log_metrics(val_metrics, step + 1, prefix="validation")
                        if master_config.grpo.debug_payload_metrics:
                            validation_payload_metrics = (
                                drain_multimodal_payload_metrics()
                            )
                            if validation_payload_metrics:
                                logger.log_metrics(
                                    validation_payload_metrics,
                                    step + 1,
                                    prefix="validation",
                                )
                        if early_stop_message is not None:
                            # Exit at the end of this step, after checkpointing.
                            print(early_stop_message, flush=True)

                        # Explicit GPU memory cleanup after validation in async mode
                        gc.collect()
                        torch.cuda.empty_cache()

                        if early_stop_message is None:
                            # Resume trajectory collection after validation
                            trajectory_collector.resume.remote()
                # Get flat advantages and token mask for masked metrics computation
                flat_advantages = train_data["advantages"]
                flat_token_mask = flat_messages["token_loss_mask"]
                # Save content for logging before deleting flat_messages
                flat_messages_content = flat_messages.get("content", [])
                del flat_messages

                # Filter advantages using token mask (only valid response tokens)
                response_advantages = torch.masked_select(
                    flat_advantages, flat_token_mask.bool()
                )

                metrics = {
                    LOSS_KEY: train_results["loss"].numpy(),
                    REWARD_KEY: rewards.numpy(),
                    "num_mask_sample_filtered": num_mask_sample_filtered,
                    GRAD_NORM_KEY: train_results["grad_norm"].numpy(),
                    "mean_prompt_length": repeated_batch["length"].numpy(),
                    "total_num_tokens": input_lengths.numpy(),
                    # Add masked advantages tracking metrics (only for valid response tokens)
                    "advantages/mean": torch.mean(response_advantages).detach().item()
                    if response_advantages.numel() > 0
                    else 0.0,
                    "advantages/max": torch.max(response_advantages).detach().item()
                    if response_advantages.numel() > 0
                    else 0.0,
                    "advantages/min": torch.min(response_advantages).detach().item()
                    if response_advantages.numel() > 0
                    else 0.0,
                }
                if "moe_metrics" in train_results:
                    metrics.update(
                        {f"moe/{k}": v for k, v in train_results["moe_metrics"].items()}
                    )
                if "mtp_metrics" in train_results:
                    metrics.update(
                        {f"mtp/{k}": v for k, v in train_results["mtp_metrics"].items()}
                    )
                if "draft_grad_norm" in train_results:
                    metrics["draft_grad_norm"] = train_results[
                        "draft_grad_norm"
                    ].numpy()
                metrics.update(train_results["all_mb_metrics"])
                metrics.update(penalty_metrics)
                for k, v in metrics.items():
                    if k in {"probs_ratio_min", "probs_ratio_clamped_min"}:
                        valid_values = [x for x in v if not np.isinf(x)]
                        metrics[k] = (
                            np.min(valid_values).item() if valid_values else -1.0
                        )
                    elif k in {"probs_ratio_max", "probs_ratio_clamped_max"}:
                        valid_values = [x for x in v if not np.isinf(x)]
                        metrics[k] = (
                            np.max(valid_values).item() if valid_values else -1.0
                        )
                    elif k in {
                        "lr",
                        "wd",
                        "reward",
                        "global_valid_seqs",
                        "global_valid_toks",
                        "mean_prompt_length",
                    }:
                        metrics[k] = np.mean(v).item()
                    else:
                        metrics[k] = np.sum(v).item()
                metrics.update(rollout_metrics)
                if generation_logger_metrics is not None:
                    metrics["generation_logger_metrics"] = generation_logger_metrics
                total_valid_tokens += metrics["global_valid_toks"]

                # Always log sequence-level error metrics (useful for deciding threshold)
                metrics.update(seq_logprob_error_metrics)

                # Speculative-decoding (MTP) acceptance metrics for this step.
                if hasattr(policy_generation, "get_step_metrics"):
                    metrics.update(policy_generation.get_step_metrics())

                # Checkpointing (same as sync version)
                consumed_samples += master_config.grpo.num_prompts_per_step
                timeout.mark_iteration()

                if saving_this_step:
                    grpo_save_state.current_step = step + 1
                    grpo_save_state.total_steps = step + 1
                    grpo_save_state.total_valid_tokens = total_valid_tokens
                    if val_metrics is not None:
                        grpo_save_state.val_reward = val_metrics["accuracy"]
                    elif hasattr(grpo_save_state, "val_reward"):
                        delattr(grpo_save_state, "val_reward")
                    grpo_save_state.consumed_samples = consumed_samples

                    full_metric_name = master_config.checkpointing["metric_name"]
                    if full_metric_name is not None:
                        prefix, metric_name = full_metric_name.split(":", 1)
                        metrics_source = metrics if prefix == "train" else val_metrics
                        if not metrics_source:
                            warnings.warn(
                                f"You asked to save checkpoints based on {metric_name} but no {prefix} metrics were collected. "
                                "This checkpoint will not be saved as top-k.",
                                stacklevel=2,
                            )
                            if hasattr(grpo_save_state, full_metric_name):
                                delattr(grpo_save_state, full_metric_name)
                        elif metric_name not in metrics_source:
                            raise ValueError(
                                f"Metric {metric_name} not found in {prefix} metrics"
                            )
                        else:
                            setattr(
                                grpo_save_state,
                                full_metric_name,
                                metrics_source[metric_name],
                            )

                    with (
                        timer.time("checkpointing"),
                        managed_span(
                            RLSpanGroup.CHECKPOINT,
                            "rl.grpo.checkpointing",
                            tracer=_tracer,
                        ),
                    ):
                        # Finalize the previous (possibly async) checkpoint before
                        # starting a new one. No-op with sync save / nothing pending.
                        checkpointer.finalize_pending()

                        print(f"Saving checkpoint for step {step + 1}...")
                        checkpoint_path = checkpointer.init_tmp_checkpoint(
                            step + 1, vars(grpo_save_state), master_config
                        )
                        policy.save_checkpoint(
                            weights_path=os.path.join(
                                checkpoint_path, "policy", "weights"
                            ),
                            optimizer_path=os.path.join(
                                checkpoint_path, "policy", "optimizer"
                            )
                            if checkpointer.save_optimizer
                            else None,
                            tokenizer_path=os.path.join(
                                checkpoint_path, "policy", "tokenizer"
                            ),
                            is_final_checkpoint=(
                                is_last_step or early_stop_message is not None
                            ),
                        )
                        # Save the dataloader state at the checkpoint cut
                        # rather than the live cursor; a resume re-yields the
                        # covered window and regenerates what the restored
                        # buffer does not account for. One actor call returns
                        # the snapshot and rollout state as a consistent pair
                        # (separate reads would race the collection loop).
                        collector_checkpoint = ray.get(
                            trajectory_collector.get_checkpoint_state.remote(
                                trained_frontier_ordinal
                            )
                        )
                        dataloader_snapshot = collector_checkpoint["dataloader"]
                        rollouts_state = collector_checkpoint["rollouts"]
                        torch.save(
                            dataloader_snapshot["dataloader_state"],
                            os.path.join(checkpoint_path, "train_dataloader.pt"),
                        )
                        _save_async_replay_buffer_checkpoint(
                            replay_buffer,
                            checkpoint_path,
                        )
                        if dataloader_snapshot["frontier_aligned"]:
                            # Persist the (possibly lowered) cut as the
                            # resume filter threshold, not the trained
                            # frontier.
                            cut_ordinal = int(dataloader_snapshot["frontier_ordinal"])
                            rollouts_state[FRONTIER_ORDINAL_KEY] = cut_ordinal
                            rollouts_state[RESUME_BASE_ORDINAL_KEY] = (
                                dataloader_snapshot["base_ordinal"]
                            )
                            # Ordinals trained at/above the cut: the resume
                            # must not regenerate these. Prune below the cut
                            # (it never decreases).
                            recent_trained_task_indices = {
                                ordinal
                                for ordinal in recent_trained_task_indices
                                if ordinal >= cut_ordinal
                            }
                            rollouts_state[TRAINED_TASK_INDICES_KEY] = sorted(
                                recent_trained_task_indices
                            )
                        torch.save(
                            rollouts_state,
                            os.path.join(checkpoint_path, "rollouts.pt"),
                        )

                        # Defer the directory rename until any async write
                        # completes; flushed at the next save or on training exit.
                        checkpointer.begin_finalization(
                            checkpoint_path,
                            wait_fn=policy.finalize_async_save,
                        )

                        # Record last-successful-checkpoint time/step for external
                        # monitoring (see _write_latest_checkpoint_status).
                        _write_latest_checkpoint_status(
                            checkpointer, last_checkpoint_step=step + 1
                        )

                    # On save-bound steps, engine stayed asleep after training;
                    # wake it unless the loop exits right below (last step, timeout, early stop),
                    # where a wake would only feed the teardown.
                    # The intervening logging runs with the collector paused either way.
                    if defer_wake_for_save and not (
                        is_last_step
                        or should_save_by_timeout
                        or early_stop_message is not None
                    ):
                        # The save onloaded model+optimizer;
                        # generation windows must start from the offloaded state.
                        policy.offload_after_refit()
                        policy_generation.prepare_for_generation()
                        ray.get(trajectory_collector.resume_after_refit.remote())

            # Logging
            # Log training data (match sync GRPO logging payload for parity).
            # NeMo Gym responses can be very large and expensive to log; when
            # env.should_log_nemo_gym_responses is true, skip this jsonl (see
            # _should_log_nemo_gym_responses).
            if not _should_log_nemo_gym_responses(master_config):
                log_data = {}
                if "agent_ref" in repeated_batch:
                    log_data["agent_ref"] = repeated_batch["agent_ref"]
                log_data["content"] = flat_messages_content
                log_data["rewards"] = rewards.tolist()
                if master_config.grpo.use_dynamic_sampling:
                    # In dynamic sampling, `rewards` corresponds to filtered rewards
                    log_data["filtered_rewards"] = rewards.tolist()
                    log_data["rewards"] = repeated_batch["total_reward"].tolist()
                log_data["input_lengths"] = input_lengths.tolist()
                log_data["token_ids"] = train_data["input_ids"].tolist()
                log_data["token_loss_mask"] = train_data["token_mask"].tolist()
                log_data["sample_loss_mask"] = train_data["sample_mask"].tolist()
                log_data["advantages"] = train_data["advantages"].tolist()
                log_data["generation_logprobs"] = train_data[
                    "generation_logprobs"
                ].tolist()
                log_data["prev_logprobs"] = train_data["prev_logprobs"].tolist()
                logger.log_batched_dict_as_jsonl(
                    log_data, f"train_data_step{step + 1}.jsonl"
                )
                del log_data
            del train_data
            del flat_messages_content

            timing_metrics: dict[str, float] = timer.get_timing_metrics(
                reduction_op="sum"
            )

            # Add buffer stats
            buffer_size_current = ray.get(replay_buffer.size.remote())
            metrics["buffer_size"] = buffer_size_current
            metrics["avg_trajectory_age"] = avg_trajectory_age

            if (
                master_config.policy["generation"]
                .get("vllm_cfg", {})
                .get("enable_vllm_metrics_logger", False)
            ):
                log_generation_metrics(
                    generation_logger_metrics,
                    step + 1,
                    master_config.policy["generation"]["vllm_cfg"][
                        "vllm_metrics_logger_interval"
                    ],
                    logger,
                )

            print("\n📊 Training Results:")
            print(f"  • Loss: {metrics['loss']:.4f}")
            if "draft_loss" in metrics:
                print(f"  • Draft Loss: {metrics['draft_loss']:.4f}")
            generation_kl_error = metrics.get("gen_kl_error")
            print(
                "  • Generation KL Error: "
                + (
                    f"{generation_kl_error:.4f}"
                    if generation_kl_error is not None
                    else "not reported"
                )
            )
            print(f"  • Avg Reward: {np.mean(rewards.numpy()):.4f}")
            print(f"  • Buffer Size: {buffer_size_current}")
            print(f"  • Avg Trajectory Age: {avg_trajectory_age:.2f} steps")

            print("\n⏱️  Timing:")
            total_time = timing_metrics.get("total_step_time", 0)
            print(f"  • Total step time: {total_time:.2f}s")
            for k, v in sorted(
                timing_metrics.items(), key=lambda item: item[1], reverse=True
            ):
                if k != "total_step_time":
                    percent = (v / total_time * 100) if total_time > 0 else 0
                    print(f"  • {k}: {v:.2f}s ({percent:.1f}%)")

            total_num_gpus = (
                master_config.cluster.num_nodes * master_config.cluster.gpus_per_node
            )
            timing_metrics["valid_tokens_per_sec_per_gpu"] = (
                metrics["global_valid_toks"] / total_time / total_num_gpus
            )
            performance_metrics = print_performance_metrics(
                train_results,
                metrics,
                timing_metrics,
                master_config,
                num_prompts_per_step=master_config.grpo.num_prompts_per_step,
                num_generations_per_prompt=master_config.grpo.num_generations_per_prompt,
                is_async_rl=master_config.grpo.async_grpo.enabled,
            )

            collector_efficiency = ray.get(
                trajectory_collector.get_efficiency_metrics.remote()
            )
            driver_efficiency = {
                cat: timer.reduce(cat, "sum")
                for cat in WALL_CLOCK_EFFICIENCY_CATEGORIES
                if cat in timer._timers
            }
            # init/total is measured once, before the loop, and the timer.reset()
            # at the end of every step drops it -- so re-supply the captured
            # value, or the series reports the real startup cost at step 1 and
            # zero for the rest of the run.
            driver_efficiency[INIT_TOTAL_CATEGORY] = init_total_s
            merged_efficiency = {**driver_efficiency}
            for cat, dur in collector_efficiency.items():
                merged_efficiency[cat] = merged_efficiency.get(cat, 0.0) + dur

            total_wall_time = time.perf_counter() - training_wall_start
            efficiency_loggable = print_efficiency_summary(
                merged_efficiency,
                total_wall_time,
                step + 1,
                # The driver's idle categories are per-step (timer.reset()
                # below), so the efficiency ratio needs a per-step denominator;
                # against the run's elapsed time it would climb toward 100%
                # whatever the idle time did.
                step_wall_time_s=total_time,
            )

            if master_config.grpo.debug_payload_metrics and not should_run_validation:
                payload_metrics = merge_multimodal_payload_metrics(
                    [
                        drain_multimodal_payload_metrics(),
                        ray.get(trajectory_collector.drain_payload_metrics.remote()),
                    ]
                )
            if payload_metrics:
                logger.log_metrics(payload_metrics, step + 1, prefix="")

            if refit_metrics:
                logger.log_metrics(refit_metrics, step + 1, prefix="refit")
            logger.log_metrics(performance_metrics, step + 1, prefix="performance")
            logger.log_metrics(metrics, step + 1, prefix="train")
            logger.log_metrics(efficiency_loggable, step + 1, prefix="")
            # step_finished=True here since this is the final log of our current step.
            logger.log_metrics(
                timing_metrics,
                step + 1,
                prefix="timing/train",
                step_finished=True,
            )

            timer.reset()
            step += 1
            if early_stop_message is not None:
                checkpointer.shutdown()
                return
            if should_save_by_timeout:
                checkpointer.shutdown()
                print("Timeout has been reached, stopping training early", flush=True)
                return
            if step >= max_num_steps:
                checkpointer.shutdown()
                print(
                    "Effective max number of steps has been reached, stopping training",
                    flush=True,
                )
                return

    except Exception as e:
        print(f"❌ Error in async loop: {e}")
        import traceback

        traceback.print_exc()
        raise

    finally:
        # Finalize any pending async checkpoint before tearing down workers.
        try:
            checkpointer.shutdown()
        except Exception as e:
            print(f"Error finalizing pending checkpoint: {e}")

        print("🛑 Stopping trajectory collection...")
        _flush_collector_telemetry()
        try:
            ray.kill(trajectory_collector)
        except Exception as e:
            print(f"Error stopping trajectory collector: {e}")

        try:
            ray.kill(replay_buffer)
        except Exception as e:
            print(f"Error stopping replay buffer: {e}")

        # Environments can have in-flight HTTP requests to generation workers.
        shutdown_environments(task_to_env, val_task_to_env)

        print("🛑 Shutting down generation workers...")
        try:
            policy_generation.shutdown()
        except Exception as e:
            print(f"Error shutting down generation workers: {e}")

        if policy is not policy_generation:
            print("🛑 Shutting down policy workers...")
            try:
                policy.shutdown()
            except Exception as e:
                print(f"Error shutting down policy workers: {e}")

        print("Async GRPO training complete!")