nemo_rl.algorithms.grpo_sync#

GRPO trainer — TransferQueue-mediated path (sync).

Sibling fork of nemo_rl.algorithms.grpo. Each file has zero internal branching on whether TQ is engaged; the example script chooses one or the other based on data_plane.enabled.

Setup and helpers are re-imported from grpo; the training loop body is duplicated here so the per-step lifecycle hooks (register / seed-put / per-rank fetch / clear) can live in straight sequential code. Validation is implemented locally as :func:validate_sync — a TQ-mediated sibling of :func:nemo_rl.algorithms.grpo.validate that routes val rollouts through SyncRolloutActor.rollout_to_tq into a per-batch "val" partition.

Parity with the legacy path is verified by running the same config against both entrypoints and diffing the wandb runs.

Module Contents#

Functions#

_raise_if_message_level_advantage_penalties_enabled

Raise if message-level advantage penalties are set in the sync trainer.

_train_fields_for_step

Fields workers fetch this step; prev_logprobs is dropped when skipped.

_apply_dynamic_sampling

Process one dynamic-sampling iteration.

validate_sync

TQ-mediated counterpart to :func:nemo_rl.algorithms.grpo.validate.

_compute_seq_logprob_error_metrics

_log_data_plane_metrics

Log this step’s data-plane cost. Never raises.

_log_data_plane_metrics_impl

Log this step’s data-plane cost. No-op unless observability is enabled.

grpo_train_sync

Run GRPO training algorithm — TransferQueue-mediated.

API#

nemo_rl.algorithms.grpo_sync._raise_if_message_level_advantage_penalties_enabled(
master_config: nemo_rl.algorithms.grpo.MasterConfig,
) → None#

Raise if message-level advantage penalties are set in the sync trainer.

Message-level advantage penalties are not supported with data_plane.enabled=true. Raises NotImplementedError listing the offending keys so the user can disable them or switch to the legacy GRPO trainer.

nemo_rl.algorithms.grpo_sync._train_fields_for_step(skip_prev_logprobs: bool) → tuple[str, ...]#

Fields workers fetch this step; prev_logprobs is dropped when skipped.

nemo_rl.algorithms.grpo_sync._apply_dynamic_sampling(
*,
meta: nemo_rl.data_plane.interfaces.KVBatchMeta,
driver_carry: nemo_rl.distributed.batched_data_dict.BatchedDataDict,
pending_meta: Optional[nemo_rl.data_plane.interfaces.KVBatchMeta],
pending_carry: Optional[nemo_rl.distributed.batched_data_dict.BatchedDataDict],
pending_unfiltered_rewards: list[torch.Tensor],
train_prompts_size: int,
num_gen_batches: int,
max_gen_batches: int,
policy: nemo_rl.models.policy.tq_policy.TQPolicy,
) → tuple[Optional[nemo_rl.data_plane.interfaces.KVBatchMeta], Optional[nemo_rl.distributed.batched_data_dict.BatchedDataDict], list[torch.Tensor], bool, dict[str, Any], Optional[torch.Tensor]]#

Process one dynamic-sampling iteration.

Drops zero-std (filtered) keys, merges survivors into the running pending cache, and reports whether the cache has reached train_prompts_size. When complete, the returned pending_* IS the training batch.

Parameters:

  • meta – This iteration’s KVBatchMeta.

  • driver_carry – Per-row driver-local tensors for this iteration (rewards, masks, prompt_ids_for_adv, baseline/std, …).

  • pending_meta – Survivors accumulated from prior iterations.

  • pending_carry – driver_carry rows aligned to pending_meta.

  • pending_unfiltered_rewards – All iterations’ rewards pre-filter, for legacy reward metric parity.

  • train_prompts_size – Target batch size.

  • num_gen_batches – Iteration counter (1-based).

  • max_gen_batches – Upper bound on iterations before raising.

  • policy – TQPolicy whose discard_samples is used to drop filtered keys.

Returns:

(pending_meta, pending_carry, pending_rewards, is_complete, ds_metrics, unfiltered_for_log).

nemo_rl.algorithms.grpo_sync.validate_sync(
*,
rollout_actor: nemo_rl.experience.sync_rollout_actor.SyncRolloutActor,
policy: nemo_rl.models.policy.tq_policy.TQPolicy,
val_dataloader: Optional[torchdata.stateful_dataloader.StatefulDataLoader],
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,
partition_id: str = 'val',
) → tuple[dict[str, Any], dict[str, Any]]#

TQ-mediated counterpart to :func:nemo_rl.algorithms.grpo.validate.

Per-batch: register the val partition → rollout_to_tq → policy.read_from_dataplane for message logs → policy.finish_step. Caller owns policy_generation.prepare_for_generation / finish_generation around the call; the actor’s per-rollout finish_generation is suppressed so inference state stays warm across batches.

nemo_rl.algorithms.grpo_sync._compute_seq_logprob_error_metrics(
*,
token_mask: torch.Tensor,
sample_mask: torch.Tensor,
prev_logprobs: torch.Tensor,
generation_logprobs: torch.Tensor,
rewards: torch.Tensor,
seq_logprob_error_threshold: Optional[float],
) → tuple[torch.Tensor, dict[str, Any]]#
nemo_rl.algorithms.grpo_sync._log_data_plane_metrics(
policy: Any,
logger: nemo_rl.utils.logger.Logger,
step: int,
total_step_time: float,
) → None#

Log this step’s data-plane cost. Never raises.

On by default, so this runs every step of every recipe.

nemo_rl.algorithms.grpo_sync._log_data_plane_metrics_impl(
policy: Any,
logger: nemo_rl.utils.logger.Logger,
step: int,
total_step_time: float,
) → None#

Log this step’s data-plane cost. No-op unless observability is enabled.

The policy computes both the metrics and the scope they cover, because the baselines they are differenced against belong with the client whose counters they baseline. This end owns only where they are logged.

The prefix names the scope because the two differ by a lot: the driver issues about one op of each kind per step while the bulk traffic is the workers’ per-DP-rank get_samples. Note that even the cluster view omits the rollout actor, which builds its own client and is not on the worker group – so kv_first_write is not in these totals.

nemo_rl.algorithms.grpo_sync.grpo_train_sync(
policy: nemo_rl.models.policy.interfaces.ColocatablePolicyInterface,
policy_generation: nemo_rl.models.generation.interfaces.GenerationInterface,
wrapped_dataloader,
val_dataloader: Optional[torchdata.stateful_dataloader.StatefulDataLoader],
tokenizer,
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 — TransferQueue-mediated.

Body mirrors :func:nemo_rl.algorithms.grpo.grpo_train with TQ-mediated Policy methods substituting the in-memory dispatch. The TQ lifecycle (controller bootstrap, worker attach, partition register, fan-out, drain, close) is fully encapsulated in

Class:

nemo_rl.models.policy.tq_policy.TQPolicy — this trainer just calls policy.prepare_step, policy.get_logprobs, policy.get_reference_policy_logprobs, and policy.train.

Parity with the legacy path is verified by running the same config against both entrypoints and diffing the wandb runs.