(self, num_embeddings, embedding_dim, params_dtype=torch.float, init_method=unscaled_init_method(0.02), skip_init=False, device=torch.device('cpu'))
| 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. |
nothing calls this directly
no test coverage detected