refactor

xrsrke · Aug 31, 2023 · cd24b87 · cd24b87
1 parent b617c8e
commit cd24b87
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 16 deletions.
diff --git a/pipegoose/nn/tensor_parallel/_utils.py b/pipegoose/nn/tensor_parallel/_utils.py
@@ -22,12 +22,6 @@ def _is_row_parallel(module: nn.Module) -> bool:
     pass
 
 
-def get_vocab_range_idx(partition_size: int, rank: int) -> Tuple[int, int]:
-    start_idx = rank * partition_size
-    end_idx = start_idx + partition_size
-    return start_idx, end_idx
-
-
 class VocabUtility:
     @staticmethod
     def get_vocab_range_idx_from_partition_size(partition_size: int, rank: int) -> Tuple[int, int]:
@@ -39,8 +33,3 @@ def get_vocab_range_idx_from_partition_size(partition_size: int, rank: int) -> T
     def get_vocab_range_from_global_vocab_size(world_size, rank, vocab_size):
         partition_size = vocab_size // world_size
         return VocabUtility.get_vocab_range_idx_from_partition_size(partition_size, rank)
-
-
-def _update_model_arguments(module: nn.Module, **kwargs):
-    for key, value in kwargs.items():
-        setattr(module, key, value)
diff --git a/pipegoose/nn/tensor_parallel/parallelize.py b/pipegoose/nn/tensor_parallel/parallelize.py
@@ -5,14 +5,15 @@
 
 from pipegoose.distributed.parallel_context import ParallelContext
 from pipegoose.distributed.parallel_mode import ParallelMode
-from pipegoose.nn.tensor_parallel._utils import (
-    VocabUtility,
-    _update_model_arguments,
-    is_splitable,
-)
+from pipegoose.nn.tensor_parallel._utils import VocabUtility, is_splitable
 from pipegoose.nn.tensor_parallel.embedding import ParallelEmbedding
 
 
+def _update_model_arguments(module: nn.Module, **kwargs):
+    for key, value in kwargs.items():
+        setattr(module, key, value)
+
+
 class ParallelizeModule(ABC):
     def __init__(self, module: nn.Module, parallel_context: ParallelContext):
         self.module = module