`nemo_rl.algorithms.grpo`#

Module Contents#

Classes#

`RewardScalingConfig`	Configure linear reward scaling with clamping.
`AsyncGRPOConfig`
`RewardPenaltyTokenIdsConfig`	Optional token IDs for reward penalties.
`RewardPenaltyConfig`	Reward-zeroing penalties applied to NeMo-Gym rollout results.
`GRPOConfig`
`GRPOSaveState`
`GRPOLoggerConfig`
`MasterConfig`

Functions#

`_maybe_restore_async_replay_buffer_checkpoint`	Restore async replay state unless the config explicitly opts out.
`_save_async_replay_buffer_checkpoint`	Checkpoint replay state inside its actor.
`_initial_grpo_save_state`
`_get_grpo_save_state`
`_validate_multimodal_dedup_capability`	Reject configurations whose media transfer path is not qualified.
`_needs_hf_refit_handshake`	Whether setup must run the HF-schema prepare_refit_info handshake.
`shutdown_environments`	Shut down each unique environment actor before generation stops.
`setup`	Main entry point for running GRPO algorithm.
`dynamic_sampling`	Implements the dynamic sampling algorithm to select prompts with non-zero standard deviation.
`scale_rewards`	Linearly scales rewards from a source range to a target range.
`extract_initial_prompt_messages`	Extract the original prompt messages from message logs using token length.
`add_grpo_token_loss_masks_and_generation_logprobs`	Add GRPO loss masks and ensure generation logprobs exist in message logs.
`_resolve_message_level_advantage_penalties`	Return configured message-level penalties and validate feature support.
`_raise_if_reward_penalties_enabled_without_nemo_gym`	Validate reward-zeroing penalties are only used with NeMo-Gym.
`_apply_message_level_advantage_penalties`	Overwrite advantages for flagged assistant-message token spans.
`_apply_configured_message_level_advantage_penalties`	Resolve config and apply message-level advantage penalties.
`_preserve_router_replay_routed_experts`	Carry rollout-recorded routes into policy worker inputs when R3 is enabled.
`_policy_dtype`	Resolve the configured policy precision to its matching torch dtype.
`_build_async_grpo_train_data`	Build the async no-TQ policy train batch from flattened rollout messages.
`_apply_mask_sample_filter`	Zero loss_multiplier where mask_sample is True and return the count.
`_should_log_nemo_gym_responses`	Whether NeMo Gym is responsible for full response logging.
`_write_latest_checkpoint_status`	Write a lightweight, top-level `latest_checkpoint_status.json` for monitoring.
`_get_effort_config`	Return the effort-levels reward-shaping config from env.nemo_gym, if set.
`_pad_teacher_logprobs`	Right-zero-pad teacher logprobs `[B, teacher_S]` to `train_S`.
`_create_advantage_estimator`	Create and return an advantage estimator based on configuration.
`_clip_grpo_advantages`	Clamp normalized advantages when clip bounds are configured.
`refit_policy_generation`	Refit the policy generation interface with the latest policy weights.
`_initial_policy_generation_stale`	Skip a fresh run’s redundant sync when the synchronizer is already current.
`_log_mixed_rewards_and_advantages_information`
`_placeholder_seq_logprob_error_metrics`	Zero-valued seq-level metrics used when the prev_logprobs forward is skipped.
`_validate_use_kl_in_reward_compat`	Reject `use_kl_in_reward` when the KL term would read zero placeholder logprobs.
`_resolve_logprob_skip_flags`	Return (skip_prev_logprobs, skip_reference_logprobs); warn on incompatible combos.
`compute_and_apply_seq_logprob_error_masking`	Compute sequence-level logprob error metrics and optionally mask high-error sequences.
`_validation_stop_value`	Value of the early-stop metric chosen by grpo.stop_at_validation_metric.
`_validation_early_stop_message`	Stop message when the early-stop threshold is reached, else None.
`grpo_train`	Run GRPO training algorithm.
`validate`	Run validation on the validation dataset.
`aggregate_rollout_metrics`	Aggregate rollout metrics from multiple trajectory groups.
`async_grpo_train`	Run asynchronous GRPO training with replay buffer.

Data#

`TokenizerType`
`_REWARD_PENALTY_FLAGS`

API#

nemo_rl.algorithms.grpo.TokenizerType#: ‘TypeVar(…)’

nemo_rl.algorithms.grpo._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.

nemo_rl.algorithms.grpo._save_async_replay_buffer_checkpoint( replay_buffer: Any, checkpoint_path: str, ) → int#: Checkpoint replay state inside its actor.

class nemo_rl.algorithms.grpo.RewardScalingConfig#

Bases: pydantic.BaseModel

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 nemo_rl.algorithms.grpo.AsyncGRPOConfig#

Bases: pydantic.BaseModel

enabled: bool#: False

max_trajectory_age_steps: int#: 1

max_generation_failures: int#: 0

in_flight_weight_updates: bool#: False

recompute_kv_cache_after_weight_updates: bool#: False

class nemo_rl.algorithms.grpo.RewardPenaltyTokenIdsConfig#

Bases: pydantic.BaseModel

Optional token IDs for reward penalties.

unwanted: list[int] | None#: None

think_open: int | None#: None

think_close: int | None#: None

class nemo_rl.algorithms.grpo.RewardPenaltyConfig#

Bases: pydantic.BaseModel

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

token_ids: Optional[nemo_rl.algorithms.grpo.RewardPenaltyTokenIdsConfig]#: None

_require_unwanted_token_ids_when_penalized() → nemo_rl.algorithms.grpo.RewardPenaltyConfig#

nemo_rl.algorithms.grpo._REWARD_PENALTY_FLAGS#: (‘penalize_duplicated_reasoning’, ‘penalize_empty_final_answer’, ‘penalize_unwanted_tokens’, ‘penali…

class nemo_rl.algorithms.grpo.GRPOConfig#

Bases: pydantic.BaseModel

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

advantage_clip_low: float | None#: None

advantage_clip_high: float | None#: None

use_leave_one_out_baseline: bool#: True

val_period: int#: 10

val_start_at: int#: None

val_batch_size: int | None#: 256

val_at_start: bool#: False

val_at_end: bool#: False

max_val_samples: int | None#: 256

val_num_generations_per_prompt: int#: 1

stop_at_validation_metric: str | None#: None

stop_at_validation_threshold: float | None#: None

skip_reference_policy_logprobs_calculation: bool#: False

seed: int#: 42

async_grpo: nemo_rl.algorithms.grpo.AsyncGRPOConfig | None#: ‘Field(…)’

overlong_filtering: bool#: False

use_dynamic_sampling: bool#: False

dynamic_sampling_max_gen_batches: int#: 10

batch_multiplier: float#: 1.0

reward_shaping: nemo_rl.algorithms.reward_functions.RewardShapingConfig#: ‘Field(…)’

reward_scaling: nemo_rl.algorithms.grpo.RewardScalingConfig#: ‘Field(…)’

calculate_advantages_on_gpu: bool#: False

seq_logprob_error_threshold: float | None#: None

invalid_tool_call_advantage: float | None#: None

malformed_thinking_advantage: float | None#: None

adv_estimator: nemo_rl.algorithms.advantage_estimator.AdvEstimatorConfig#: ‘Field(…)’

deduplicate_multimodal_data: bool#: False

debug_payload_metrics: bool#: False

class nemo_rl.algorithms.grpo.GRPOSaveState#

consumed_samples: int#: None

current_step: int#: None

current_epoch: int#: None

total_steps: int#: None

total_valid_tokens: int#: None

val_reward: float#: None

trainer_version: Optional[int]#: None

sampler_name: Optional[str]#: None

sampler_dispatch_index: Optional[int]#: None

nemo_rl.algorithms.grpo._initial_grpo_save_state() → nemo_rl.algorithms.grpo.GRPOSaveState#

nemo_rl.algorithms.grpo._get_grpo_save_state( loaded_state: Optional[dict[str, Any]], ) → nemo_rl.algorithms.grpo.GRPOSaveState#

class nemo_rl.algorithms.grpo.GRPOLoggerConfig#

Bases: nemo_rl.utils.logger.LoggerConfig

num_val_samples_to_print: int#: None

class nemo_rl.algorithms.grpo.MasterConfig#

Bases: pydantic.BaseModel

policy: nemo_rl.models.policy.PolicyConfig#: None

loss_fn: nemo_rl.algorithms.loss.ClippedPGLossConfig#: None

env: dict[str, Any]#: None

data: nemo_rl.data.DataConfig#: None

grpo: nemo_rl.algorithms.grpo.GRPOConfig#: None

logger: nemo_rl.algorithms.grpo.GRPOLoggerConfig#: None

cluster: nemo_rl.distributed.virtual_cluster.ClusterConfig#: None

checkpointing: nemo_rl.utils.checkpoint.CheckpointingConfig#: None

reward_penalties: nemo_rl.algorithms.grpo.RewardPenaltyConfig#: ‘Field(…)’

data_plane: Optional[nemo_rl.data_plane.interfaces.DataPlaneConfig]#: None

on_policy_distillation: Optional[nemo_rl.algorithms.opd.OnPolicyDistillationConfig]#: None

telemetry: Optional[nemo_rl.telemetry.config.TelemetryConfig]#: None

nemo_rl.algorithms.grpo._validate_multimodal_dedup_capability( master_config: nemo_rl.algorithms.grpo.MasterConfig, ) → None#: Reject configurations whose media transfer path is not qualified.

nemo_rl.algorithms.grpo._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.

nemo_rl.algorithms.grpo.shutdown_environments( task_to_env: dict[str, nemo_rl.environments.interfaces.EnvironmentInterface] | None, val_task_to_env: dict[str, nemo_rl.environments.interfaces.EnvironmentInterface] | None, ) → None#: Shut down each unique environment actor before generation stops.

nemo_rl.algorithms.grpo.setup( master_config: nemo_rl.algorithms.grpo.MasterConfig, tokenizer: nemo_rl.algorithms.grpo.TokenizerType, dataset: nemo_rl.data.datasets.AllTaskProcessedDataset | dict[str, nemo_rl.data.datasets.AllTaskProcessedDataset], val_dataset: Optional[nemo_rl.data.datasets.AllTaskProcessedDataset], processor: Optional[transformers.AutoProcessor] = None, policy_factory: Optional[Callable[..., nemo_rl.models.policy.interfaces.ColocatablePolicyInterface]] = None, ) → tuple[nemo_rl.models.policy.interfaces.ColocatablePolicyInterface, Optional[nemo_rl.models.generation.interfaces.GenerationInterface], Optional[nemo_rl.environments.interfaces.EnvironmentInterface], tuple[nemo_rl.distributed.virtual_cluster.RayVirtualCluster, nemo_rl.distributed.virtual_cluster.RayVirtualCluster], torchdata.stateful_dataloader.StatefulDataLoader | nemo_rl.data.dataloader.MultipleDataloaderWrapper, Optional[torchdata.stateful_dataloader.StatefulDataLoader], nemo_rl.algorithms.loss.ClippedPGLossFn, nemo_rl.utils.logger.Logger, nemo_rl.utils.checkpoint.CheckpointManager, nemo_rl.algorithms.grpo.GRPOSaveState, nemo_rl.algorithms.grpo.MasterConfig, dict[str, Any], dict[str, str]]#

Main entry point for running GRPO algorithm.

Returns:: 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.
Return type:: A 13-tuple, in order

nemo_rl.algorithms.grpo.dynamic_sampling( repeated_batch: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.data.interfaces.DatumSpec], std: torch.Tensor, baseline: torch.Tensor, dynamic_sampling_num_gen_batches: int, master_config: nemo_rl.algorithms.grpo.MasterConfig, timer: nemo_rl.utils.timer.Timer, batch_cache: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.data.interfaces.DatumSpec] = None, ) → nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.data.interfaces.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.

Parameters:

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.

Returns:

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.

Return type:

tuple

nemo_rl.algorithms.grpo.scale_rewards( repeated_batch: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.data.interfaces.DatumSpec], reward_scaling_cfg: nemo_rl.algorithms.grpo.RewardScalingConfig, ) → nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.data.interfaces.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

nemo_rl.algorithms.grpo.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).

Parameters:

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.

nemo_rl.algorithms.grpo.add_grpo_token_loss_masks_and_generation_logprobs( message_logs: list[nemo_rl.data.interfaces.LLMMessageLogType | nemo_rl.data.interfaces.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. 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.
Parameters:: 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.

nemo_rl.algorithms.grpo._resolve_message_level_advantage_penalties( master_config: nemo_rl.algorithms.grpo.MasterConfig, ) → tuple[float | None, float | None]#: Return configured message-level penalties and validate feature support.

nemo_rl.algorithms.grpo._raise_if_reward_penalties_enabled_without_nemo_gym( master_config: nemo_rl.algorithms.grpo.MasterConfig, *, enable_nemo_gym: bool, ) → None#: Validate reward-zeroing penalties are only used with NeMo-Gym.

nemo_rl.algorithms.grpo._apply_message_level_advantage_penalties( train_data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.algorithms.loss.ClippedPGLossDataDict], message_logs: list[nemo_rl.data.interfaces.LLMMessageLogType | nemo_rl.data.interfaces.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.

Parameters:

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.

nemo_rl.algorithms.grpo._apply_configured_message_level_advantage_penalties( train_data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.algorithms.loss.ClippedPGLossDataDict], message_logs: list[nemo_rl.data.interfaces.LLMMessageLogType | nemo_rl.data.interfaces.VLMMessageLogType], master_config: nemo_rl.algorithms.grpo.MasterConfig, log_config: bool = False, ) → Optional[dict[str, float]]#: Resolve config and apply message-level advantage penalties.

nemo_rl.algorithms.grpo._preserve_router_replay_routed_experts( target: nemo_rl.distributed.batched_data_dict.BatchedDataDict, flat_messages: nemo_rl.distributed.batched_data_dict.BatchedDataDict, policy_config: nemo_rl.models.policy.PolicyConfig, ) → None#: Carry rollout-recorded routes into policy worker inputs when R3 is enabled.

nemo_rl.algorithms.grpo._policy_dtype( policy_config: nemo_rl.models.policy.PolicyConfig, ) → torch.dtype#: Resolve the configured policy precision to its matching torch dtype.

nemo_rl.algorithms.grpo._build_async_grpo_train_data( flat_messages: nemo_rl.distributed.batched_data_dict.BatchedDataDict, input_lengths: torch.Tensor, repeated_batch: nemo_rl.distributed.batched_data_dict.BatchedDataDict, policy_config: nemo_rl.models.policy.PolicyConfig, ) → nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.algorithms.loss.ClippedPGLossDataDict]#: Build the async no-TQ policy train batch from flattened rollout messages.

nemo_rl.algorithms.grpo._apply_mask_sample_filter( repeated_batch: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.data.interfaces.DatumSpec], ) → int#: Zero loss_multiplier where mask_sample is True and return the count.

nemo_rl.algorithms.grpo._should_log_nemo_gym_responses( master_config: nemo_rl.algorithms.grpo.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.

nemo_rl.algorithms.grpo._write_latest_checkpoint_status( checkpointer: nemo_rl.utils.checkpoint.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.

nemo_rl.algorithms.grpo._get_effort_config( master_config: nemo_rl.algorithms.grpo.MasterConfig, ) → Optional[nemo_rl.experience.rollouts.EffortLevelsConfig]#: Return the effort-levels reward-shaping config from env.nemo_gym, if set.

nemo_rl.algorithms.grpo._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.

nemo_rl.algorithms.grpo._create_advantage_estimator( master_config: nemo_rl.algorithms.grpo.MasterConfig, )#

Create and return an advantage estimator based on configuration.

Parameters:: 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.

nemo_rl.algorithms.grpo._clip_grpo_advantages( advantages: torch.Tensor, grpo_config: nemo_rl.algorithms.grpo.GRPOConfig, ) → torch.Tensor#: Clamp normalized advantages when clip bounds are configured.

nemo_rl.algorithms.grpo.refit_policy_generation( policy: nemo_rl.models.policy.interfaces.ColocatablePolicyInterface, policy_generation: nemo_rl.models.generation.interfaces.GenerationInterface, colocated_inference: bool, _refit_buffer_size_gb: Optional[float] = None, timer: Optional[nemo_rl.utils.timer.Timer] = None, kv_scales: Optional[dict[str, float]] = None, ) → dict[str, float]#

Refit the policy generation interface with the latest policy weights.

Parameters:

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.

nemo_rl.algorithms.grpo._initial_policy_generation_stale( policy_generation: nemo_rl.models.generation.interfaces.GenerationInterface, completed_steps: int, ) → bool#: Skip a fresh run’s redundant sync when the synchronizer is already current.

nemo_rl.algorithms.grpo._log_mixed_rewards_and_advantages_information( logger: nemo_rl.utils.logger.Logger, total_steps: int, metrics: dict[str, Any], baseline: torch.Tensor, advantages: torch.Tensor, ) → None#

nemo_rl.algorithms.grpo._placeholder_seq_logprob_error_metrics() → dict[str, float]#: Zero-valued seq-level metrics used when the prev_logprobs forward is skipped.

nemo_rl.algorithms.grpo._validate_use_kl_in_reward_compat( master_config: nemo_rl.algorithms.grpo.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.

nemo_rl.algorithms.grpo._resolve_logprob_skip_flags( master_config: nemo_rl.algorithms.grpo.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 seq_logprob_error_threshold is set (which requires prev_logprobs). Skip reference_policy_logprobs when grpo.skip_reference_policy_logprobs_calculation is set.

nemo_rl.algorithms.grpo.compute_and_apply_seq_logprob_error_masking( train_data: nemo_rl.distributed.batched_data_dict.BatchedDataDict, rewards: torch.Tensor, seq_logprob_error_threshold: Optional[float], ) → 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.

Parameters:

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.

Returns:

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

Return type:

Dict with keys

nemo_rl.algorithms.grpo._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.

nemo_rl.algorithms.grpo._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.

nemo_rl.algorithms.grpo.grpo_train( policy: nemo_rl.models.policy.interfaces.ColocatablePolicyInterface, policy_generation: Optional[nemo_rl.models.generation.interfaces.GenerationInterface], wrapped_dataloader: torchdata.stateful_dataloader.StatefulDataLoader | nemo_rl.data.dataloader.MultipleDataloaderWrapper, val_dataloader: Optional[torchdata.stateful_dataloader.StatefulDataLoader], tokenizer: nemo_rl.algorithms.grpo.TokenizerType, loss_fn: nemo_rl.algorithms.loss.interfaces.LossFunction, task_to_env: dict[str, nemo_rl.environments.interfaces.EnvironmentInterface], val_task_to_env: Optional[dict[str, nemo_rl.environments.interfaces.EnvironmentInterface]], logger: nemo_rl.utils.logger.Logger, checkpointer: nemo_rl.utils.checkpoint.CheckpointManager, grpo_save_state: nemo_rl.algorithms.grpo.GRPOSaveState, master_config: nemo_rl.algorithms.grpo.MasterConfig, processor: Optional[transformers.AutoProcessor] = None, ) → None#: Run GRPO training algorithm.

nemo_rl.algorithms.grpo.validate( policy_generation: nemo_rl.models.generation.interfaces.GenerationInterface, val_dataloader: Optional[torchdata.stateful_dataloader.StatefulDataLoader], tokenizer, val_task_to_env: Optional[dict[str, nemo_rl.environments.interfaces.EnvironmentInterface]], step: int, master_config: nemo_rl.algorithms.grpo.MasterConfig, logger: Optional[nemo_rl.utils.logger.Logger] = None, processor: Optional[transformers.AutoProcessor] = None, ) → tuple[dict[str, Any], dict[str, Any]]#: Run validation on the validation dataset.

nemo_rl.algorithms.grpo.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

Parameters:: 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.

nemo_rl.algorithms.grpo.async_grpo_train( policy: nemo_rl.models.policy.interfaces.ColocatablePolicyInterface, policy_generation: Optional[nemo_rl.models.generation.interfaces.GenerationInterface], dataloader: torchdata.stateful_dataloader.StatefulDataLoader, val_dataloader: Optional[torchdata.stateful_dataloader.StatefulDataLoader], tokenizer: nemo_rl.algorithms.grpo.TokenizerType, loss_fn: nemo_rl.algorithms.loss.interfaces.LossFunction, task_to_env: dict[str, nemo_rl.environments.interfaces.EnvironmentInterface], val_task_to_env: Optional[dict[str, nemo_rl.environments.interfaces.EnvironmentInterface]], logger: nemo_rl.utils.logger.Logger, checkpointer: nemo_rl.utils.checkpoint.CheckpointManager, grpo_save_state: nemo_rl.algorithms.grpo.GRPOSaveState, master_config: nemo_rl.algorithms.grpo.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[transformers.AutoProcessor] = None, ) → None#

Run asynchronous GRPO training with replay buffer.

Parameters:

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.

nemo_rl.algorithms.grpo#

Module Contents#

Classes#

Functions#

Data#

API#

`nemo_rl.algorithms.grpo`#