Skip to content

GitLab

Commit cbf780d0 authored by Jared Casper's avatar Jared Casper
Browse files

Update exports.

parent e00a1ca1
Show whitespace changes
Inline Side-by-side
Showing with 3 additions and 0 deletions
megatron/core/tensor_parallel/__init__.py
View file @ cbf780d0
...@@ -42,6 +42,7 @@ __all__ = [ ...@@ -42,6 +42,7 @@ __all__ = [
"ColumnParallelLinear", "ColumnParallelLinear",
"RowParallelLinear", "RowParallelLinear",
"VocabParallelEmbedding", "VocabParallelEmbedding",
"set_tensor_model_parallel_attributes",
"set_defaults_if_not_set_tensor_model_parallel_attributes", "set_defaults_if_not_set_tensor_model_parallel_attributes",
"copy_tensor_model_parallel_attributes", "copy_tensor_model_parallel_attributes",
"param_is_not_tensor_parallel_duplicate", "param_is_not_tensor_parallel_duplicate",
...@@ -59,4 +60,6 @@ __all__ = [ ...@@ -59,4 +60,6 @@ __all__ = [
"model_parallel_cuda_manual_seed", "model_parallel_cuda_manual_seed",
# utils.py # utils.py
"split_tensor_along_last_dim", "split_tensor_along_last_dim",
"split_tensor_into_1d_equal_chunks",
"gather_split_1d_tensor",
] ]
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment