`nemo_rl.algorithms.utils`#

Module Contents#

Functions#

`get_gdpo_reward_component_keys`	Return batch keys that are named reward components (e.g. reward/correctness) in sorted order.
`calculate_kl`	Calculates a per-token estimate of the KL Divergence between two logprobs.
`calculate_baseline_and_std_per_prompt`	Function to compute a baseline for each (prompt, response) pair in the batch.
`calculate_trivial_reward_distributions`	Return an all-or-nothing exact-equality mask for each prompt group.
`surpress_user_warnings`
`compute_seq_logprob_errors`	Return mean multiplicative absolute error and validity per sequence.
`masked_mean`	Computes the mean of a microbatch, using a global statistic as the normalization factor.
`mask_filtered_logprobs_outside_tokens`	Zero top-k/top-p filtered log probabilities outside `mask`.
`masked_var`
`set_seed`	Sets the seed for python, numpy, and pytorch.
`get_tokenizer`	Get the tokenizer and set pad token to eos token if it is not already set.
`maybe_pad_last_batch`	Pads the given batch so that its size is divisible by (mbs * dp_size).
`print_performance_metrics`	Print performance metrics for an RL training step.
`log_generation_metrics`	Log generation metric timelines to every configured logger backend.
`print_efficiency_summary`	Print a summary table of efficiency metrics and return loggable dict.
`finalize_draft_ratio_metrics`	Turn summed `draft__num`/`draft__den` metric pairs back into ratios, in place.

Data#

`WALL_CLOCK_EFFICIENCY_CATEGORIES`
`THREAD_ACCUMULATED_EFFICIENCY_CATEGORIES`
`EFFICIENCY_CATEGORIES`
`STEP_WINDOW_WALL_CLOCK_CATEGORIES`
`_DRAFT_RATIO_NUM_RE`

API#

nemo_rl.algorithms.utils.get_gdpo_reward_component_keys(batch) → list[str]#: Return batch keys that are named reward components (e.g. reward/correctness) in sorted order.

nemo_rl.algorithms.utils.calculate_kl( logprobs: torch.Tensor, logprobs_reference: torch.Tensor, kl_type: str = 'k3', input_clamp_value: float | None = 20.0, output_clamp_value: float | None = 10.0, importance_sampling_weights: torch.Tensor | None = None, ) → torch.Tensor#

Calculates a per-token estimate of the KL Divergence between two logprobs.

From Schulman 2020, http://joschu.net/blog/kl-approx.html.

Parameters:

logprobs – torch.Tensor (b, s)
logprobs_reference – torch.Tensor (b, s)
kl_type – Type of KL approximation to use. Valid values: “k1”, “k2”, “k3”.
input_clamp_value – Optional clamping value for logr to prevent numerical instability. If None, no clamping is applied.
output_clamp_value – Optional clamping value for kl to prevent numerical instability. If None, no clamping is applied.
importance_sampling_weights – Optional per-token importance weights. When provided, weights are multiplied into the KL before output clamping. If input clamping applies to a token, the corresponding weight is detached so the clamp suppresses both KL and sampling-weight gradients.

Returns:

Per-token KL penalty values (b, s)

Return type:

torch.Tensor

nemo_rl.algorithms.utils.calculate_baseline_and_std_per_prompt( prompts: torch.Tensor, rewards: torch.Tensor, valid_mask: torch.Tensor, leave_one_out_baseline: bool = True, std_rewards: torch.Tensor | None = None, ) → tuple[torch.Tensor, torch.Tensor, torch.Tensor]#

Function to compute a baseline for each (prompt, response) pair in the batch.

The same baseline is calculated for each prompt. Samples set to 0 in ‘valid_mask’ are not included in the baseline calculation.

prompts: tensor (b, k) Row equality keys defining prompt groups. GRPO uses explicit rollout-group ids with shape (b, 1), not prompt tokens. May be on any device. rewards: tensor (b,) Float-valued rewards. May be on any device valid_mask: tensor (b,) Vector of 0/1, where 0 is to ignore and 1 is to keep leave_one_out_baseline: bool Compute an unbiased baseline by leaving out the sample that the baseline is for (from RLOO https://arxiv.org/abs/2402.14740) std_rewards: tensor (b,) Optional separate reward tensor used only for the std calculation. Defaults to rewards. Useful for DAPO, which needs std on the raw task metric for dynamic sampling filtering while keeping baseline on the shaped reward.

Returns: tensor (b,), tensor (b,), bool tensor (b,) of baselines, std, and a per-sample boolean that is True if the sample group distribution associated with the sample is trivial, i.e. all the rewards are a single value.

Any non-zero std computed from a trivial distribution is float32 rounding noise which can create explosive advantages during normalization that should not be used for GRPO training.

nemo_rl.algorithms.utils.calculate_trivial_reward_distributions( prompts: torch.Tensor, rewards: torch.Tensor, valid_mask: torch.Tensor, ) → torch.Tensor#

Return an all-or-nothing exact-equality mask for each prompt group.

Unlike the per-sample mask returned by calculate_baseline_and_std_per_prompt, this always compares the full valid reward group. It is therefore independent of leave-one-out baseline semantics and safe to use for prompt-level dynamic sampling.

nemo_rl.algorithms.utils.surpress_user_warnings(f)#

nemo_rl.algorithms.utils.compute_seq_logprob_errors( *, policy_logprobs: torch.Tensor, generation_logprobs: torch.Tensor, token_mask: torch.Tensor, sample_mask: torch.Tensor, ) → tuple[torch.Tensor, torch.Tensor]#

Return mean multiplicative absolute error and validity per sequence.

Inputs must already be aligned to predicted tokens (without the first input token). Padding and previously masked samples do not participate. Nonfinite errors on valid tokens fail any finite threshold.

nemo_rl.algorithms.utils.masked_mean( values: torch.Tensor, mask: torch.Tensor, dim: Optional[int] = None, global_normalization_factor: Optional[torch.Tensor | float] = None, )#: Computes the mean of a microbatch, using a global statistic as the normalization factor.

nemo_rl.algorithms.utils.mask_filtered_logprobs_outside_tokens( logprobs: torch.Tensor, mask: torch.Tensor, ) → torch.Tensor#

Zero top-k/top-p filtered log probabilities outside mask.

Prompt and padding tokens are often outside the filtered support and would otherwise carry -inf (or NaN after multiplying by a mask) to every consumer. Valid tokens keep -inf so the actor loss can exclude policy support mismatches instead of treating them as log-probability zero.

Parameters:

logprobs – Filtered log probabilities.
mask – Valid-token mask (token_mask * sample_mask).

Returns:

Log probabilities with zeros outside mask.

nemo_rl.algorithms.utils.masked_var( values: torch.Tensor, mask: torch.Tensor, mean: Optional[torch.Tensor | float] = None, unbiased: bool = True, ) → torch.Tensor#

nemo_rl.algorithms.utils.set_seed(seed: int) → None#: Sets the seed for python, numpy, and pytorch.

nemo_rl.algorithms.utils.get_tokenizer( tokenizer_config: nemo_rl.models.policy.TokenizerConfig, get_processor: bool = False, ) → transformers.PreTrainedTokenizerBase#

Get the tokenizer and set pad token to eos token if it is not already set.

This function initializes a tokenizer from the Hugging Face transformers library and configures it with appropriate chat templates and padding tokens.

Parameters:

tokenizer_config – A dictionary containing tokenizer configuration. Required keys: - name: The name or path of the pretrained tokenizer Optional keys: - chat_template: The chat template to use. Can be: - None: Uses a passthrough template that just returns message content - “default”: Uses the tokenizer’s default template - A custom jinja2 template string If not specified, the tokenizer’s default template will be used. - tokenizer_kwargs: Extra keyword arguments forwarded to tokenizer loading, e.g. {“fix_mistral_regex”: False}. When get_processor=True, these are passed through AutoProcessor.from_pretrained().
get_processor – Whether to return a processor (via AutoProcessor) instead of a tokenizer.

Returns:

The configured tokenizer instance

Return type:

PreTrainedTokenizerBase

.. rubric:: Examples

>>> from transformers import AutoTokenizer
>>> from nemo_rl.algorithms.utils import get_tokenizer
>>> # not specifying a chat template uses the tokenizer's default
>>> config = {"name": "meta-llama/Llama-3.2-1B-Instruct"}
>>> tokenizer = get_tokenizer(config)
No chat template provided, using tokenizer's default
>>> messages = [
...     {"role": "system", "content": "You are a helpful AI assistant."},
...     {"role": "user", "content": "Hello!"}
... ]
>>> formatted = tokenizer.apply_chat_template(messages, tokenize=False)
>>> assert formatted == AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B-Instruct").apply_chat_template(messages, tokenize=False)

>>> # Using a passthrough template
>>> config = {
...     "name": "meta-llama/Llama-3.2-1B-Instruct",
...     "chat_template": None
... }
>>> tokenizer = get_tokenizer(config)
Using passthrough chat template
>>> formatted = tokenizer.apply_chat_template(messages, tokenize=False)
>>> assert formatted == "".join(msg["content"] for msg in messages)

>>> # Using a custom template
>>> config = {
...     "name": "meta-llama/Llama-3.2-1B-Instruct",
...     "chat_template": "{% for message in messages %}{{ ' START: ' + message['content'] + ' END.' }}{% endfor %}"
... }
>>> tokenizer = get_tokenizer(config)
Using custom chat template
>>> formatted = tokenizer.apply_chat_template(messages, tokenize=False)
>>> assert formatted == " START: You are a helpful AI assistant. END. START: Hello! END."

>>> # Requesting a processor (for multimodal models like Qwen-VL)
>>> config = {"name": "Qwen/Qwen2.5-VL-3B-Instruct"}
>>> processor = get_tokenizer(config, get_processor=True)
No chat template provided, using tokenizer's default
>>> messages = [
...     {"role": "system", "content": "You are a helpful AI assistant."},
...     {"role": "user", "content": "Hello!"}
... ]
>>> formatted = processor.tokenizer.apply_chat_template(messages, tokenize=False)
>>> assert formatted == AutoTokenizer.from_pretrained(
...     "Qwen/Qwen2.5-VL-3B-Instruct", trust_remote_code=True
... ).apply_chat_template(messages, tokenize=False)
>>> assert processor.pad_token_id == processor.tokenizer.pad_token_id
>>>

nemo_rl.algorithms.utils.maybe_pad_last_batch( batch: BatchedDataDict[Any], dp_size: int, mbs: int, ) → BatchedDataDict[Any]#

Pads the given batch so that its size is divisible by (mbs * dp_size).

Parameters:

batch (dict) – The batch to pad.
dp_size (int) – Data parallel size.
mbs (int) – Micro batch size.

Returns:

The padded batch.

Return type:

dict

nemo_rl.algorithms.utils.print_performance_metrics( train_results: dict[str, float], metrics: dict[str, Any], timing_metrics: dict[str, float], master_config: dict, *, num_prompts_per_step: int, num_generations_per_prompt: int, is_async_rl: bool, ) → dict[str, float]#: Print performance metrics for an RL training step.

nemo_rl.algorithms.utils.log_generation_metrics( generation_logger_metrics: dict[str, dict[int, list[Any]]], step: int, timeline_interval: float, logger: nemo_rl.utils.logger.Logger, ) → None#

Log generation metric timelines to every configured logger backend.

Parameters:

generation_logger_metrics – Dictionary of generation logger metrics
step – Global step value
timeline_interval – Interval between timeline points (in seconds)
logger – Logger instance

nemo_rl.algorithms.utils.WALL_CLOCK_EFFICIENCY_CATEGORIES#: [‘init/total’, ‘idle/buffer_starvation’, ‘idle/refit_bubble’, ‘idle/validation’]

nemo_rl.algorithms.utils.THREAD_ACCUMULATED_EFFICIENCY_CATEGORIES#: [‘idle/buffer_full_backoff’, ‘idle/generation_limit_pause’, ‘idle/refit_event_wait’, ‘wasted/failed_…

nemo_rl.algorithms.utils.EFFICIENCY_CATEGORIES#: None

nemo_rl.algorithms.utils.STEP_WINDOW_WALL_CLOCK_CATEGORIES#: None

nemo_rl.algorithms.utils.print_efficiency_summary( efficiency_metrics: dict[str, float], total_wall_time_s: float, step: int, step_wall_time_s: Optional[float] = None, ) → dict[str, float]#

Print a summary table of efficiency metrics and return loggable dict.

Driver-side wall-clock categories drive the efficiency percentage. Collector-side categories are summed across concurrent threads and reported separately as thread-seconds so they are not compared directly against a single wall-clock denominator.

The efficiency percentage is a per-step ratio, because the numerator is: the driver resets its Timer every step, so its idle categories are per-step deltas. Dividing them by the cumulative total_wall_time_s would make the run look monotonically more efficient the longer it ran, whatever the idle time actually did. init/total is excluded from that numerator for the same reason in reverse – it is a run-long constant (see

Data:

RUN_WINDOW_WALL_CLOCK_CATEGORIES) and would otherwise add the whole startup cost to every single step’s waste.

Parameters:

efficiency_metrics – Dict mapping category labels to total seconds spent.
total_wall_time_s – Total wall-clock time in seconds since training began. Used for the per-category % of Wall column.
step – Current training step number.
step_wall_time_s – Wall-clock seconds in the step being reported, the denominator of the efficiency percentage. Defaults to total_wall_time_s for callers with no per-step measurement.

Returns:

Dict of metrics suitable for logging to WandB/TensorBoard, including per-category seconds and percentages, total waste, and efficiency_pct.

nemo_rl.algorithms.utils._DRAFT_RATIO_NUM_RE#: ‘compile(…)’

nemo_rl.algorithms.utils.finalize_draft_ratio_metrics(metrics: dict[str, Any]) → None#

Turn summed draft_*_num/draft_*_den metric pairs back into ratios, in place.

DSparkRuntime/Eagle3Runtime (nemo_rl/models/automodel/draft/runtime.py) emit per-microbatch accuracy/rate metrics as raw num/den pairs instead of pre-divided ratios, because the training loop’s default per-key metric reduction sums every metric not on its small mean-reduction allowlist across all microbatches and DP ranks – summing a pre-divided per-microbatch ratio would produce a meaningless value (e.g. a “draft_full_acc@0” of 240 instead of a rate in [0, 1]). Call this AFTER that per-key reduction has summed the raw num/den pairs into global totals, so the ratio it computes is the correct token-weighted global rate. Deliberately dependency-free (no nemo_automodel import) so the driver process can call it unconditionally.

nemo_rl.algorithms.utils#

Module Contents#

Functions#

Data#

API#

`nemo_rl.algorithms.utils`#