clean up embeddings.py (#7)

- clean up embeddings.py (6d5377580042f614ac525febd19c641c4d456d9b)


Co-authored-by: Bo Wang <[email protected]>

embedding.py

@@ -10,59 +10,6 @@ import torch.nn as nn
 from torch import Tensor
 
 
-class GPT2Embeddings(nn.Module):
-    def __init__(
-        self,
-        embed_dim,
-        vocab_size,
-        max_position_embeddings,
-        padding_idx=None,
-        word_embed_proj_dim=None,
-        device=None,
-        dtype=None,
-    ):
-        """
-        If max_position_embeddings <= 0, there's no position embeddings
-        If word_embe_proj_dim is not None (e.g., OPT-350m), we embed to that dimension
-            the project up to embed_dim
-        """
-        factory_kwargs = {"device": device, "dtype": dtype}
-        super().__init__()
-        if word_embed_proj_dim is None:
-            self.word_embeddings = nn.Embedding(
-                vocab_size, embed_dim, padding_idx=padding_idx, **factory_kwargs
-            )
-            self.project_in = None
-        else:
-            self.word_embeddings = nn.Embedding(
-                vocab_size, word_embed_proj_dim, padding_idx=padding_idx, **factory_kwargs
-            )
-            self.project_in = nn.Linear(
-                word_embed_proj_dim, embed_dim, bias=False, **factory_kwargs
-            )
-        self.max_position_embeddings = max_position_embeddings
-        if self.max_position_embeddings > 0:
-            self.position_embeddings = nn.Embedding(
-                max_position_embeddings, embed_dim, **factory_kwargs
-            )
-
-    def forward(self, input_ids, position_ids=None):
-        """
-        input_ids: (batch, seqlen)
-        position_ids: (batch, seqlen)
-        """
-        batch_size, seqlen = input_ids.shape
-        embeddings = self.word_embeddings(input_ids)
-        if self.project_in is not None:
-            embeddings = self.project_in(embeddings)
-        if self.max_position_embeddings > 0:
-            if position_ids is None:
-                position_ids = torch.arange(seqlen, dtype=torch.long, device=input_ids.device)
-            position_embeddings = self.position_embeddings(position_ids)
-            embeddings = embeddings + position_embeddings
-        return embeddings
-
-
 class BertEmbeddings(nn.Module):
     def __init__(
         self,
@@ -111,52 +58,3 @@ class BertEmbeddings(nn.Module):
             token_type_embeddings = self.token_type_embeddings(token_type_ids)
             embeddings = embeddings + token_type_embeddings
         return embeddings
-
-
-class VocabParallelEmbedding(nn.Embedding):
-    def __init__(self, num_embeddings, *args, process_group=None, padding_idx=None, **kwargs):
-        self.process_group = process_group
-        if process_group is not None:
-            world_size = torch.distributed.get_world_size(process_group)
-            if num_embeddings % world_size != 0:
-                raise ValueError(
-                    f"num_embeddings ({num_embeddings}) must be divisible by "
-                    f"world_size ({world_size})"
-                )
-            if world_size > 1 and padding_idx is not None:
-                raise RuntimeError("ParallelEmbedding does not support padding_idx")
-        else:
-            world_size = 1
-        super().__init__(num_embeddings // world_size, *args, padding_idx=padding_idx, **kwargs)
-
-    def forward(self, input: Tensor) -> Tensor:
-        if self.process_group is None:
-            return super().forward(input)
-        else:
-            rank = torch.distributed.get_rank(self.process_group)
-            vocab_size = self.num_embeddings
-            vocab_start_index, vocab_end_index = rank * vocab_size, (rank + 1) * vocab_size
-            # Create a mask of valid vocab ids (1 means it needs to be masked).
-            input_ids_mask = (input < vocab_start_index) | (input >= vocab_end_index)
-            input = input - vocab_start_index
-            input[input_ids_mask] = 0
-            embeddings = super().forward(input)
-            embeddings[input_ids_mask] = 0.0
-            return embeddings
-
-
-class ColumnParallelEmbedding(nn.Embedding):
-    def __init__(self, num_embeddings, embedding_dim, *args, process_group=None, **kwargs):
-        self.process_group = process_group
-        if process_group is not None:
-            world_size = torch.distributed.get_world_size(process_group)
-            if embedding_dim % world_size != 0:
-                raise ValueError(
-                    f"embedding_dim ({embedding_dim}) must be divisible by "
-                    f"world_size ({world_size})"
-                )
-        else:
-            world_size = 1
-        super().__init__(num_embeddings, embedding_dim // world_size, *args, **kwargs)
-
-