nemo_rl.models.value.tq_value#

TQ-mediated Value: meta-driven 1-hop counterpart to Value, mirroring TQPolicy.

Module Contents#

Classes#

TQValue

TQ-mediated counterpart to Value, taking an extra dp_cfg.

Functions#

_aggregate_train_results

Assemble per-rank value-train results into Value.train’s return shape.

Data#

_REPLICATED_AXES

API#

nemo_rl.models.value.tq_value._REPLICATED_AXES#

[‘context_parallel’, ‘tensor_parallel’, ‘pipeline_parallel’]

nemo_rl.models.value.tq_value._aggregate_train_results(
results: list[dict[str, Any]],
) → dict[str, Any]#

Assemble per-rank value-train results into Value.train’s return shape.

class nemo_rl.models.value.tq_value.TQValue(
*args: Any,
dp_cfg: nemo_rl.data_plane.DataPlaneConfig,
**kwargs: Any,
)#

Bases: nemo_rl.data_plane.driver_mixin.TQDriverMixin, nemo_rl.models.value.lm_value.Value

TQ-mediated counterpart to Value, taking an extra dp_cfg.

Attaches to the TQ controller rather than bootstrapping it: the TQPolicy built alongside this critic already did that. Partition lifecycle stays with the caller.

TODO(#2625): the value workers have no split begin/microbatch/finish train API yet, so one train_from_meta call is one optimizer step and the SingleController requires a PPO step to be a single streaming chunk.

Initialization

Initialize the Value model.

Parameters:

  • cluster – Ray virtual cluster for distributed training

  • config – Configuration for the value model

  • tokenizer – Tokenizer for the model

  • name_prefix – Prefix for worker names

  • workers_per_node – Number of workers per node

  • init_optimizer – Whether to initialize the optimizer

  • weights_path – Path to load model weights from

  • optimizer_path – Path to load optimizer state from

shutdown() → bool#

Close the TQ client before shutting down the worker group.

get_values_from_meta(
meta: nemo_rl.data_plane.KVBatchMeta,
micro_batch_size: Optional[int] = None,
timer: Optional[nemo_rl.utils.timer.Timer] = None,
) → None#

1-hop counterpart to get_values.

Returns nothing: the per-token prediction lands in TQ under values via the worker-side leader write-back, so the GAE stage reads it from there rather than through Ray.

Parameters:

  • meta – Full-step batch metadata consumed by all DP ranks.

  • micro_batch_size – Inference micro batch size; None uses the config default.

  • timer – Optional timer for nested get_values measurements.

train_from_meta(
meta: nemo_rl.data_plane.KVBatchMeta,
loss_fn: nemo_rl.algorithms.loss.interfaces.LossFunction,
eval_mode: bool = False,
gbs: Optional[int] = None,
mbs: Optional[int] = None,
timer: Optional[nemo_rl.utils.timer.Timer] = None,
) → dict[str, Any]#

1-hop counterpart to train. One call is one optimizer step.

Parameters:

  • meta – Full-step batch metadata consumed by all DP ranks.

  • loss_fn – Value loss; MseValueLossFn in the PPO path.

  • eval_mode – Run forward only, without an optimizer step.

  • gbs – Global batch size; defaults to the config’s train_global_batch_size.

  • mbs – Micro batch size; defaults to the config’s train_micro_batch_size.

  • timer – Optional timer for nested value_training measurements.

Returns:

Aggregated training-step output dict.