(self, num_embeddings, embedding_dim, init_method=init.xavier_normal_)
| 150 | """ |
| 151 | |
| 152 | def __init__(self, num_embeddings, embedding_dim, init_method=init.xavier_normal_): |
| 153 | super(VocabParallelEmbedding, self).__init__() |
| 154 | # Keep the input dimensions. |
| 155 | self.num_embeddings = num_embeddings |
| 156 | self.embedding_dim = embedding_dim |
| 157 | # Set the detauls for compatibility. |
| 158 | self.padding_idx = None |
| 159 | self.max_norm = None |
| 160 | self.norm_type = 2.0 |
| 161 | self.scale_grad_by_freq = False |
| 162 | self.sparse = False |
| 163 | self._weight = None |
| 164 | self.tensor_model_parallel_size = get_tensor_model_parallel_world_size() |
| 165 | # Divide the weight matrix along the vocaburaly dimension. |
| 166 | ( |
| 167 | self.vocab_start_index, |
| 168 | self.vocab_end_index, |
| 169 | ) = VocabUtility.vocab_range_from_global_vocab_size( |
| 170 | self.num_embeddings, |
| 171 | get_tensor_model_parallel_rank(), |
| 172 | self.tensor_model_parallel_size, |
| 173 | ) |
| 174 | self.num_embeddings_per_partition = ( |
| 175 | self.vocab_end_index - self.vocab_start_index |
| 176 | ) |
| 177 | |
| 178 | # Allocate weights and initialize. |
| 179 | args = get_args() |
| 180 | if args.use_cpu_initialization: |
| 181 | self.weight = Parameter( |
| 182 | torch.empty( |
| 183 | self.num_embeddings_per_partition, |
| 184 | self.embedding_dim, |
| 185 | dtype=args.params_dtype, |
| 186 | # dtype=torch.float32, |
| 187 | ) |
| 188 | ) |
| 189 | _initialize_affine_weight_cpu( |
| 190 | self.weight, |
| 191 | self.num_embeddings, |
| 192 | self.embedding_dim, |
| 193 | self.num_embeddings_per_partition, |
| 194 | 0, |
| 195 | init_method, |
| 196 | ) |
| 197 | else: |
| 198 | self.weight = Parameter( |
| 199 | torch.empty( |
| 200 | self.num_embeddings_per_partition, |
| 201 | self.embedding_dim, |
| 202 | device=torch.cuda.current_device(), |
| 203 | dtype=args.params_dtype, |
| 204 | # dtype=torch.float32, |
| 205 | ) |
| 206 | ) |
| 207 | _initialize_affine_weight_gpu( |
| 208 | self.weight, init_method, partition_dim=0, stride=1 |
| 209 | ) |
nothing calls this directly
no test coverage detected