`nemo_rl.algorithms.grpo`#

Module Contents#

Classes#

`RewardScalingConfig`	Configure linear reward scaling with clamping.
`AsyncGRPOConfig`
`AdvEstimatorConfig`	Configuration for advantage estimator (GRPO, GDPO, or Reinforce++).
`RewardPenaltyTokenIdsConfig`	Optional token IDs for reward penalties.
`RewardPenaltyConfig`	Reward-zeroing penalties applied to NeMo-Gym rollout results.
`GRPOConfig`
`GRPOSaveState`
`GRPOLoggerConfig`
`MasterConfig`

Functions#

`_default_grpo_save_state`
`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.
`_should_use_async_rollouts`	Determine if async rollouts should be used based on the configuration.
`_preserve_router_replay_routed_experts`	Carry rollout-recorded routes into policy worker inputs when R3 is enabled.
`_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_use_nemo_gym`	Determine if NeMo-Gym should be used for rollouts and validation based on the configuration.
`_should_log_nemo_gym_responses`	Whether NeMo Gym is responsible for full response logging (wandb/metrics paths).
`_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.
`_log_mixed_rewards_and_advantages_information`
`compute_and_apply_seq_logprob_error_masking`	Compute sequence-level logprob error metrics and optionally mask high-error sequences.
`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(…)’

class nemo_rl.algorithms.grpo.RewardScalingConfig#

Bases: typing.TypedDict

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.

Defaults: source_min=0.0, source_max=1.0, target_min=0.0, target_max=1.0

Initialization

Initialize self. See help(type(self)) for accurate signature.

enabled: bool#: None

source_min: NotRequired[float]#: None

source_max: NotRequired[float]#: None

target_min: NotRequired[float]#: None

target_max: NotRequired[float]#: None

class nemo_rl.algorithms.grpo.AsyncGRPOConfig#

Bases: typing.TypedDict

enabled: bool#: None

max_trajectory_age_steps: int#: None

in_flight_weight_updates: NotRequired[bool]#: None

recompute_kv_cache_after_weight_updates: NotRequired[bool]#: None

class nemo_rl.algorithms.grpo.AdvEstimatorConfig#

Bases: typing.TypedDict

Configuration for advantage estimator (GRPO, GDPO, or Reinforce++).

Initialization

Initialize self. See help(type(self)) for accurate signature.

name: str#: None

normalize_rewards: NotRequired[bool]#: None

use_leave_one_out_baseline: NotRequired[bool]#: None

minus_baseline: NotRequired[bool]#: None

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: typing.TypedDict

num_prompts_per_step: int#: None

num_generations_per_prompt: int#: None

max_num_epochs: int#: None

max_num_steps: int#: None

max_rollout_turns: int#: None

normalize_rewards: bool#: None

advantage_clip_low: NotRequired[float | None]#: None

advantage_clip_high: NotRequired[float | None]#: None

use_leave_one_out_baseline: bool#: None

val_period: int#: None

val_batch_size: int | None#: None

val_at_start: bool#: None

val_at_end: bool#: None

max_val_samples: int | None#: None

skip_reference_policy_logprobs_calculation: NotRequired[bool]#: None

seed: int#: None

async_grpo: NotRequired[nemo_rl.algorithms.grpo.AsyncGRPOConfig]#: None

overlong_filtering: NotRequired[bool]#: None

use_dynamic_sampling: bool#: None

dynamic_sampling_max_gen_batches: NotRequired[int]#: None

batch_multiplier: NotRequired[float]#: None

reward_shaping: nemo_rl.algorithms.reward_functions.RewardShapingConfig#: None

reward_scaling: nemo_rl.algorithms.grpo.RewardScalingConfig#: None

calculate_advantages_on_gpu: NotRequired[bool]#: None

seq_logprob_error_threshold: float | None#: None

invalid_tool_call_advantage: NotRequired[float | None]#: None

malformed_thinking_advantage: NotRequired[float | None]#: None

adv_estimator: NotRequired[nemo_rl.algorithms.grpo.AdvEstimatorConfig]#: None

class nemo_rl.algorithms.grpo.GRPOSaveState#

Bases: typing.TypedDict

consumed_samples: int#: None

current_step: int#: None

current_epoch: int#: None

total_steps: int#: None

total_valid_tokens: int#: None

val_reward: NotRequired[float]#: None

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

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.

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._should_use_async_rollouts( master_config: nemo_rl.algorithms.grpo.MasterConfig, ) → bool#

Determine if async rollouts should be used based on the configuration.

SGLang only uses async rollouts when explicitly configured with policy.generation.use_async_rollouts. vLLM and Megatron use async rollouts when their respective async_engine config is enabled.

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._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_use_nemo_gym( master_config: nemo_rl.algorithms.grpo.MasterConfig, ) → bool#: Determine if NeMo-Gym should be used for rollouts and validation based on the configuration.

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 (wandb/metrics paths).

When True, we skip the expensive per-step train_data_step*.jsonl dump and keep full_result-style keys in rollout metrics (large payloads).

When False (default if unset), we strip full_result from rollout metrics and write the train_data_step*.jsonl file (can be very large for Gym).

Set via env.should_log_nemo_gym_responses in the master config.

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: dict[str, Any], ) → 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, ) → None#

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.

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.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.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, ) → 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, ) → 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:

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, ) → 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

nemo_rl.algorithms.grpo#

Module Contents#

Classes#

Functions#

Data#

API#

`nemo_rl.algorithms.grpo`#