56 lines
1.6 KiB
Python
56 lines
1.6 KiB
Python
|
from .layers import (
|
||
|
|
ColumnParallelLinear,
|
||
|
|
RowParallelLinear,
|
||
|
|
VocabParallelEmbedding,
|
||
|
|
set_tensor_model_parallel_attributes,
|
||
|
|
set_defaults_if_not_set_tensor_model_parallel_attributes,
|
||
|
|
copy_tensor_model_parallel_attributes,
|
||
|
|
param_is_not_tensor_parallel_duplicate,
|
||
|
|
)
|
||
|
|
|
||
|
|
from .mappings import (
|
||
|
|
copy_to_tensor_model_parallel_region,
|
||
|
|
gather_from_tensor_model_parallel_region,
|
||
|
|
gather_from_sequence_parallel_region,
|
||
|
|
scatter_to_tensor_model_parallel_region,
|
||
|
|
scatter_to_sequence_parallel_region,
|
||
|
|
)
|
||
|
|
|
||
|
|
from .random import (
|
||
|
|
checkpoint,
|
||
|
|
get_cuda_rng_tracker,
|
||
|
|
model_parallel_cuda_manual_seed,
|
||
|
|
)
|
||
|
|
|
||
|
|
from .utils import (
|
||
|
|
split_tensor_along_last_dim,
|
||
|
|
split_tensor_into_1d_equal_chunks,
|
||
|
|
gather_split_1d_tensor,
|
||
|
|
)
|
||
|
|
|
||
|
|
__all__ = [
|
||
|
|
#layers.py
|
||
|
|
"ColumnParallelLinear",
|
||
|
|
"RowParallelLinear",
|
||
|
|
"VocabParallelEmbedding",
|
||
|
|
"set_tensor_model_parallel_attributes",
|
||
|
|
"set_defaults_if_not_set_tensor_model_parallel_attributes",
|
||
|
|
"copy_tensor_model_parallel_attributes",
|
||
|
|
"param_is_not_tensor_parallel_duplicate",
|
||
|
|
# mappings.py
|
||
|
|
"copy_to_tensor_model_parallel_region",
|
||
|
|
"gather_from_tensor_model_parallel_region",
|
||
|
|
"gather_from_sequence_parallel_region",
|
||
|
|
# "reduce_from_tensor_model_parallel_region",
|
||
|
|
"scatter_to_tensor_model_parallel_region",
|
||
|
|
"scatter_to_sequence_parallel_region",
|
||
|
|
# random.py
|
||
|
|
"checkpoint",
|
||
|
|
"get_cuda_rng_tracker",
|
||
|
|
"model_parallel_cuda_manual_seed",
|
||
|
|
# utils.py
|
||
|
|
"split_tensor_along_last_dim",
|
||
|
|
"split_tensor_into_1d_equal_chunks",
|
||
|
|
"gather_split_1d_tensor",
|
||
|
|
]
|