def collective_fn(input_tensor, output_tensor, context):
pygloo.allreduce(
context, gloo_util.get_tensor_ptr(input_tensor),
gloo_util.get_tensor_ptr(output_tensor),
gloo_util.get_tensor_n_elements(input_tensor),
gloo_util.get_gloo_tensor_dtype(input_tensor),
gloo_util.get_gloo_reduce_op(allreduce_options.reduceOp))
def collective_fn(input_tensor, output_tensor, context):
pygloo.allgather(
context,
gloo_util.get_tensor_ptr(input_tensor),
gloo_util.get_tensor_ptr(output_tensor),
gloo_util.get_tensor_n_elements(input_tensor),
gloo_util.get_gloo_tensor_dtype(input_tensor),
)
def collective_fn(input_tensor, output_tensor, context):
size = gloo_util.get_tensor_n_elements(input_tensor)
world_size = self._gloo_context.size
pygloo.reduce_scatter(
context, gloo_util.get_tensor_ptr(input_tensor),
gloo_util.get_tensor_ptr(output_tensor), size,
[size // world_size for _ in range(world_size)],
gloo_util.get_gloo_tensor_dtype(output_tensor),
gloo_util.get_gloo_reduce_op(reducescatter_options.reduceOp))
def collective_fn(input_tensor, output_tensor, context):
pygloo.broadcast(
context,
gloo_util.get_tensor_ptr(input_tensor),
gloo_util.get_tensor_ptr(output_tensor),
gloo_util.get_tensor_n_elements(input_tensor),
gloo_util.get_gloo_tensor_dtype(input_tensor),
root_rank,
)
def p2p_fn(tensor, context, peer):
pygloo.recv(
context,
gloo_util.get_tensor_ptr(tensor),
gloo_util.get_tensor_n_elements(tensor),
gloo_util.get_gloo_tensor_dtype(tensor),
peer,
)