add docs

xrsrke · Aug 31, 2023 · 90d479c · 90d479c
1 parent d2dc575
commit 90d479c
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/pipegoose/nn/tensor_parallel/parallelize.py b/pipegoose/nn/tensor_parallel/parallelize.py
@@ -39,6 +39,7 @@ def deparallelize(self):
 class ParallelizeEmbedding(ParallelizeModule):
     # TODO: refactor to staticmethod
     def parallelize(self) -> nn.Module:
+        """Parallelize nn.Embedding module."""
         assert isinstance(self.module, nn.Embedding), "only parallelize nn.Embedding"
         self._resize_vocab_size()
         self._split_weight()
@@ -48,6 +49,7 @@ def deparallelize(self):
         pass
 
     def _split_weight(self):
+        """Split weight into chunks and assign to each process."""
         world_size = self.parallel_context.get_world_size(ParallelMode.TENSOR)
         rank = self.parallel_context.get_local_rank(ParallelMode.TENSOR)