nemo_rl.algorithms.async_utils.replay_buffer#

Module Contents#

Classes#

ReplayBufferImpl

Replay buffer storing per-prompt groups.

ReplayBuffer

ReplayBufferNew

Staleness-window replay buffer.

API#

class nemo_rl.algorithms.async_utils.replay_buffer.ReplayBufferImpl(max_size: int)#

Bases: nemo_rl.algorithms.async_utils.interfaces.ReplayBufferProtocol

Replay buffer storing per-prompt groups.

A single entry corresponds to 1 prompt repeated by grpo.num_generations_per_prompt (required to compute per-prompt advantages).

Initialization

static _rollout_metrics_turn_count_for_diagnostics(
rm: dict[str, Any],
) → Optional[float]#

One scalar turn-depth per buffered trajectory for starvation diagnostics.

Supports sync multi-turn rollouts (max_turns_per_sample / avg_turns_per_sample) and NeMo Gym (turns_per_sample/max / turns_per_sample/mean).

add(
trajectory: dict[str, Any],
weight_version: int,
target_weight_version: int,
) → str#

Add a per-prompt trajectory group with metadata.

Parameters:

  • trajectory – data dict

  • weight_version – version of the model weights used for generation

  • target_weight_version – version of the model weights this trajectory is intended for training

get_debug_info() → dict#

Get debug information about buffer state.

get_last_target_weight_already_generated() → int#
get_existing_target_weights() → set[int]#

Get set of target weight versions that already have trajectories.

_remove_indices(indices: Iterable[int]) → None#

Remove trajectories at the given indices.

sample(
num_prompt_groups: int,
current_weight_version: int,
max_age_steps: int,
) → Optional[dict[str, Any]]#

Sample per-prompt trajectory groups intended for the current training step.

Only returns trajectories with target_weight_version == current_weight_version. If insufficient trajectories are available, returns None to stall training until the remaining trajectories are generated. This ensures no trajectory loses its last chance to be used for its intended training step.

Returns:

Dictionary with ‘trajectories’ and ‘avg_trajectory_age’ keys, or None if insufficient data

size() → int#

Return current buffer size.

clear() → None#

Clear the buffer.

state_dict() → dict[str, Any]#

Return serializable state for checkpointing.

load_state_dict(
state: dict[str, Any],
num_prompts_per_step: int | None = None,
current_training_step: int | None = None,
max_age_steps: int | None = None,
) → None#

Restore replay buffer state from a checkpoint.

Parameters:

  • state – State returned by state_dict.

  • num_prompts_per_step – Number of prompt groups required for one training step. When provided, incomplete target steps can be removed or prepared for gap filling.

  • current_training_step – Step being resumed. When provided with num_prompts_per_step, past target steps are dropped and incomplete current/future target steps are kept for gap filling.

  • max_age_steps – Maximum allowed age for restored trajectories. When provided, stale trajectories are removed during restore.

Raises:

ValueError – If the checkpoint is missing required fields or has inconsistent parallel list lengths.

_prepare_for_training_step(
current_step: int,
num_prompts_per_step: int,
) → None#

Prepare restored state so training can resume at current_step.

static _is_valid_for_target(
trajectory_version: int,
target_step: int,
max_age_steps: int | None,
) → bool#
_remove_stale_trajectories(max_age_steps: int) → None#

Remove restored trajectories that are stale for their target step.

Must be called while holding self._lock.

_count_for_target(
target_step: int,
max_age_steps: int | None = None,
) → int#

Count trajectories usable for target_step.

Must be called while holding self._lock.

_truncate_to_max_size(
current_training_step: int | None = None,
) → None#

Truncate restored state to max_size after resume cleanup.

Must be called while holding self._lock.

get_trajectories_needed(
target_step: int,
num_prompts_per_step: int,
max_age_steps: int | None = None,
) → int#

Return additional trajectories needed for target_step.

has_complete_batch(
target_step: int,
num_prompts_per_step: int,
max_age_steps: int | None = None,
) → bool#

Return whether target_step has enough trajectories to train.

_remove_incomplete_target_steps(num_prompts_per_step: int) → None#

Remove target steps without a complete batch.

Must be called while holding self._lock.

class nemo_rl.algorithms.async_utils.replay_buffer.ReplayBuffer(max_size: int)#

Bases: nemo_rl.algorithms.async_utils.replay_buffer.ReplayBufferImpl

class nemo_rl.algorithms.async_utils.replay_buffer.ReplayBufferNew(
max_size: int,
max_staleness: int,
sample_freshest_first: bool = True,
)#

Bases: nemo_rl.algorithms.async_utils.replay_buffer.ReplayBufferImpl

Staleness-window replay buffer.

– WIP: DO NOT USE – This class is WIP and may be changed without notice, please DO NOT USE it.

Differences from ReplayBuffer:

  • _evict(): Stale rows (trainer_version - weight_version > max_staleness) are evicted at the start of every sample() call.

  • sample(): selects trajectories in freshest-first order (default) or FIFO order, controlled by the sample_freshest_first flag, from whatever remains in the buffer after eviction.

TODO: remove when cleaning up

  • max_age_steps won’t be used in ReplayBufferNew;

  • self.target_weight_versions won’t be used in ReplayBufferNew and will be removed when cleaning up. target_weight_versions gates generation on specific trainer steps, which causes generation pauses; ReplayBufferNew intentionally avoids this.

  • add this class to nemo_rl/algorithms/async_utils/init.py

Initialization

_evict(current_weight_version: int) → None#

Evict rows where trainer_version - weight_version > max_staleness.

Must be called with self._lock held.

sample(
num_prompt_groups: int,
current_weight_version: int,
max_age_steps: int,
) → Optional[dict[str, Any]]#

Sample num_prompt_groups trajectories, freshest-first.

Will evict stale rows before sampling, so we will get [current_weight_version - self.max_staleness, current_weight_version] valid trajectories.

Returns:

Dictionary with ‘trajectories’ and ‘avg_trajectory_age’ keys, or None.