deepfold.distributed.model_parallel

Functions

all_to_all(tensor, dim0, dim1)
broadcast(tensor, src_rank)
broadcast_on_model_parallel_region(tensor, ...)
col_to_row(tensor)
copy_to_model_parallel_reigon(tensor)
disable()
enable()
gather(tensor, dim[, bwd])
gather_from_model_parallel_region(tensor, dim)
group()
group_rank()
initialize(dap_size)	Initialize Dynamic Axial Parallelism (DAP).
is_enabled()
is_initialized()
rank()
reduce(tensor)
reduce_from_model_parallel_region(tensor)
row_to_col(tensor)
scatter(tensor, dim)
scatter_to_model_parallel_region(tensor, dim)
size()
transpose_on_model_parallel_region(tensor, ...)

Classes

Disable(*args, **kwargs)
Enable(*args, **kwargs)