Embedding parallelized in the vocabulary dimension. This is mainly adapted from torch.nn.Embedding and all the default values are kept. Arguments: num_embeddings: vocabulary size. embedding_dim: size of hidden state. init_method: method to initialize weights.
| 68 | |
| 69 | |
| 70 | class VocabParallelEmbedding(torch.nn.Module): |
| 71 | """Embedding parallelized in the vocabulary dimension. |
| 72 | |
| 73 | This is mainly adapted from torch.nn.Embedding and all the default |
| 74 | values are kept. |
| 75 | Arguments: |
| 76 | num_embeddings: vocabulary size. |
| 77 | embedding_dim: size of hidden state. |
| 78 | init_method: method to initialize weights. |
| 79 | """ |
| 80 | def __init__(self, num_embeddings, embedding_dim, params_dtype=torch.float, init_method=unscaled_init_method(0.02), skip_init=False, device=torch.device('cpu')): |
| 81 | super(VocabParallelEmbedding, self).__init__() |
| 82 | # Keep the input dimensions. |
| 83 | self.num_embeddings = num_embeddings |
| 84 | self.embedding_dim = embedding_dim |
| 85 | # Set the detauls for compatibility. |
| 86 | self.padding_idx = None |
| 87 | self.max_norm = None |
| 88 | self.norm_type = 2. |
| 89 | self.scale_grad_by_freq = False |
| 90 | self.sparse = False |
| 91 | self._weight = None |
| 92 | # Divide the weight matrix along the vocaburaly dimension. |
| 93 | self.vocab_start_index, self.vocab_end_index = \ |
| 94 | VocabUtility.vocab_range_from_global_vocab_size( |
| 95 | self.num_embeddings, get_model_parallel_rank(), |
| 96 | get_model_parallel_world_size()) |
| 97 | self.num_embeddings_per_partition = self.vocab_end_index - \ |
| 98 | self.vocab_start_index |
| 99 | |
| 100 | # Allocate weights. |
| 101 | self.weight = Parameter(torch.empty(self.num_embeddings_per_partition, |
| 102 | self.embedding_dim, dtype=params_dtype, |
| 103 | device=device)) |
| 104 | self.weight.model_parallel = True |
| 105 | # And initialize. |
| 106 | if not skip_init: |
| 107 | _initialize_affine_weight( |
| 108 | self.weight, self.num_embeddings, self.embedding_dim, |
| 109 | self.num_embeddings_per_partition, 0, init_method, self=self) |
| 110 | |
| 111 | def forward(self, input_): |
| 112 | # Build the mask. |
| 113 | input_mask = (input_ < self.vocab_start_index) | \ |
| 114 | (input_ >= self.vocab_end_index) |
| 115 | # Mask the input. |
| 116 | masked_input = input_.clone() - self.vocab_start_index |
| 117 | masked_input[input_mask] = 0 |
| 118 | # Get the embeddings. |
| 119 | output_parallel = F.embedding(masked_input, self.weight, |
| 120 | self.padding_idx, self.max_norm, |
| 121 | self.norm_type, self.scale_grad_by_freq, |
| 122 | self.sparse) |
| 123 | # Mask the output embedding. |
| 124 | output_parallel[input_mask, :] = 0.0 |
| 125 | # Reduce across all the model parallel GPUs. |
| 126 | output = reduce_from_model_parallel_region(output_parallel) |
| 127 | return output |