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_batches

Yield unpacked batches while retaining ownership of their receive buffers.

packed_broadcast_consumer

Consume a packed tensor as callbacks or a lazy weight iterator.

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,
use_legacy_packing: bool = False,
)#

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.

  • use_legacy_packing – pack an oversized tensor with the ones before it, as vLLM’s own consumer expects, instead of sending it in its own broadcast.

Returns:

None

nemo_rl.utils.packed_tensor._packed_broadcast_consumer_batches(
iterator: collections.abc.Iterator[Any],
group: Any,
src: int,
*,
num_buffers: int | None = None,
) → collections.abc.Iterator[list[tuple[str, torch.Tensor]]]#

Yield unpacked batches while retaining ownership of their receive buffers.

nemo_rl.utils.packed_tensor.packed_broadcast_consumer(
iterator: collections.abc.Iterator[Any],
group: Any,
src: int,
post_unpack_func: collections.abc.Callable[[list[tuple[str, torch.Tensor]]], None] | None,
return_iterator: bool = False,
*,
num_buffers: int | None = None,
) → collections.abc.Iterator[tuple[str, torch.Tensor]] | None#

Consume a packed tensor as callbacks or a lazy weight iterator.

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 unpacked batch. May be None when return_iterator is True.

  • return_iterator – return a lazy, flattened weight iterator for consumers such as vLLM’s reload_weights API.

  • 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:

A lazy iterator when return_iterator is True, otherwise None.