nemo_rl.utils.packed_tensor#

Module Contents#

Functions#

get_target_packed_tensor_size

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.packed_broadcast_producer(iterator, group, src, post_iter_func)#

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

Returns:

None

nemo_rl.utils.packed_tensor.packed_broadcast_consumer(iterator, group, src, post_unpack_func)#

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

Returns:

None