nemo_rl.data_plane.driver_mixin#

Driver-side TransferQueue helpers shared by TQPolicy and TQValue.

Module Contents#

Classes#

TQDriverMixin

Pad-target minting and column read/write against the data plane.

API#

class nemo_rl.data_plane.driver_mixin.TQDriverMixin#

Pad-target minting and column read/write against the data plane.

Hosts must provide cfg, dp_client, and the use_dynamic_batches / use_sequence_packing attribute pairs that Policy and Value both set.

_packing_args(
mb_tokens_key: str,
) → tuple[Optional[dict[str, Any]], Optional[dict[str, Any]]]#

Resolve (sequence_packing_args, dynamic_batching_args) for a given stage.

The stage is identified by mb_tokens_key ("logprob_mb_tokens" or "train_mb_tokens").

_stamp_pad_seqlen(
meta: nemo_rl.data_plane.interfaces.KVBatchMeta,
) → None#

Mint GLOBAL_FORWARD_PAD_SEQLEN onto meta.extra_info (idempotent).

Cross-DP forward pad target. Preshard shards inherit it via dict(meta.extra_info) propagation.

_isolated_meta(
meta: nemo_rl.data_plane.interfaces.KVBatchMeta,
*,
fields: list[str],
task_name: str,
include_multimodal: bool = True,
) → nemo_rl.data_plane.interfaces.KVBatchMeta#

Narrow meta for one model’s dispatch and mint it a fresh pad target.

The mint is idempotent, so sharing or inheriting the target would let whichever model dispatches first decide the forward pad for the rest – and with ppo_epochs > 1 the caller’s meta is already stamped when the critic dispatches again.

fields is the dispatch’s static column list; the multimodal columns the rollout actually wrote are unioned in here rather than at each call site. Every forward-running dispatch of the policy needs them – omitting them runs a VLM forward image-blind while the sibling dispatches saw the images – and nothing in a static list forces the next caller to remember. The union is a no-op for text-only runs, since present_multimodal_fields intersects with meta.fields.

include_multimodal=False opts a dispatch out. A value model is the case that exists: its inputs are a fixed list and it never consumes pixels, so unioning them in would fetch and broadcast a column the forward discards.

read_from_dataplane(
meta: nemo_rl.data_plane.interfaces.KVBatchMeta,
*,
select_fields: list[str],
pad_value_dict: Optional[dict[str, Any]] = None,
) → nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any]#

Fetch + materialize columns from the data plane (TQ).

read_columns pads to meta.extra_info[GLOBAL_FORWARD_PAD_SEQLEN] — the same value workers pad to in their forward pass. Driver and workers thus return columns at one identical seq dim, with no driver-side knowledge of sequence_length_round.

write_to_dataplane(
meta: nemo_rl.data_plane.interfaces.KVBatchMeta,
fields: dict[str, Any],
) → None#

Write driver-computed columns to the data plane (TQ).