Functions#

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)#

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