Language model embeddings. Arguments: hidden_size: hidden size vocab_size: vocabulary size max_sequence_length: maximum size of sequence. This is used for positional embedding embedding_dropout_prob: dropout probability for embeddings
| 305 | |
| 306 | |
| 307 | class QueryEmbedding(MegatronModule): |
| 308 | """Language model embeddings. |
| 309 | |
| 310 | Arguments: |
| 311 | hidden_size: hidden size |
| 312 | vocab_size: vocabulary size |
| 313 | max_sequence_length: maximum size of sequence. This |
| 314 | is used for positional embedding |
| 315 | embedding_dropout_prob: dropout probability for embeddings |
| 316 | init_method: weight initialization method |
| 317 | num_tokentypes: size of the token-type embeddings. 0 value |
| 318 | will ignore this embedding |
| 319 | """ |
| 320 | |
| 321 | def __init__(self, |
| 322 | hidden_size, |
| 323 | vocab_size, |
| 324 | max_sequence_length, |
| 325 | embedding_dropout_prob, |
| 326 | init_method, |
| 327 | num_tokentypes=0): |
| 328 | super(QueryEmbedding, self).__init__() |
| 329 | |
| 330 | self.hidden_size = hidden_size |
| 331 | self.init_method = init_method |
| 332 | self.num_tokentypes = num_tokentypes |
| 333 | self.max_sequence_length = max_sequence_length |
| 334 | |
| 335 | # Top query position embedding (serial). |
| 336 | self.top_query_embeddings = mpu.VocabParallelEmbedding( |
| 337 | max_sequence_length, self.hidden_size, init_method=self.init_method) |
| 338 | self.top_query_embeddings = self.top_query_embeddings.half() |
| 339 | self._top_query_embeddings_key = 'top_query_embeddings' |
| 340 | |
| 341 | # Initialize the top query position embeddings. |
| 342 | self.init_method(self.top_query_embeddings.weight) |
| 343 | |
| 344 | # Token type embedding. |
| 345 | # Add this as an optional field that can be added through |
| 346 | # method call so we can load a pretrain model without |
| 347 | # token types and add them as needed. |
| 348 | self._tokentype_embeddings_key = 'tokentype_embeddings' |
| 349 | if self.num_tokentypes > 0: |
| 350 | self.tokentype_embeddings = torch.nn.Embedding(self.num_tokentypes, |
| 351 | self.hidden_size) |
| 352 | # Initialize the token-type embeddings. |
| 353 | self.init_method(self.tokentype_embeddings.weight) |
| 354 | else: |
| 355 | self.tokentype_embeddings = None |
| 356 | |
| 357 | # Embeddings dropout |
| 358 | self.embedding_dropout = torch.nn.Dropout(embedding_dropout_prob) |
| 359 | |
| 360 | def add_tokentype_embeddings(self, num_tokentypes): |
| 361 | """Add token-type embedding. This function is provided so we can add |
| 362 | token-type embeddings in case the pretrained model does not have it. |
| 363 | This allows us to load the model normally and then add this embedding. |
| 364 | """ |