megatron.text_generation.communication

Description

Communications utilities.

Functions

broadcast_float_list(size[, float_list, rank])	Broadcast a list of float values.
broadcast_from_last_pipeline_stage(size, dtype)	Broadcast a tensor from last pipeline stage to all ranks.
broadcast_from_last_to_first_pipeline_stage(...)	Broadcast tensor values from last stage into the first stage.
broadcast_int_list(size[, int_list, rank])	Broadcast a list of interger values.
broadcast_list(size, dtype[, list_values, rank])	Broadcast a list of values with a given type.
broadcast_tensor(size, dtype[, tensor, rank])	Given size and type of a tensor on all ranks and the tensor value only on a specific rank, broadcast from that rank to all other ranks.
copy_from_last_to_first_pipeline_stage(size, ...)	Copy tensor values from last stage into the first stage.
recv_from_prev_pipeline_rank_([recv_buffer])	Receive from previous pipeline stage and update the input buffer inplace.
send_to_next_pipeline_rank([tensor])	Send output to the next pipeline stage.