nemo_rl.utils.packed_tensor#

Module Contents#

Functions#

get_target_packed_tensor_size

get_num_buffers

packed_broadcast_producer

Broadcast a list of tensors in a packed manner.

packed_broadcast_consumer

Consume a packed tensor and unpack it into a list of tensors.

API#

nemo_rl.utils.packed_tensor.get_target_packed_tensor_size()#
nemo_rl.utils.packed_tensor.get_num_buffers()#
nemo_rl.utils.packed_tensor.packed_broadcast_producer(
iterator,
group,
src,
post_iter_func,
*,
buffer_size_bytes: int | None = None,
num_buffers: int | None = None,
)#

Broadcast a list of tensors in a packed manner.

Parameters:

  • iterator – iterator of model parameters. Returns a tuple of (name, tensor)

  • group – process group (vllm PyNcclCommunicator)

  • src – source rank (0 in current implementation)

  • post_iter_func – function to apply to each tensor before packing, should return a tensor

  • buffer_size_bytes – packed-buffer target. Uses the NeMo-RL default when unset.

  • num_buffers – number of alternating CUDA buffers. Uses the default when unset.

Returns:

None

nemo_rl.utils.packed_tensor.packed_broadcast_consumer(
iterator,
group,
src,
post_unpack_func,
*,
num_buffers: int | None = None,
)#

Consume a packed tensor and unpack it into a list of tensors.

Parameters:

  • iterator – iterator of model parameters. Returns a tuple of (name, tensor)

  • group – process group (vllm PyNcclCommunicator)

  • src – source rank (0 in current implementation)

  • post_unpack_func – function to apply to each tensor after unpacking

  • num_buffers – number of alternating CUDA buffers/streams. Uses the NRL_REFIT_NUM_BUFFERS default when unset. Chunk boundaries only depend on the packed-buffer target size, so the producer and consumer may use different buffer counts.

Returns:

None