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
| 53 | |
| 54 | |
| 55 | class LLaMaEmbedding(nn.Module): |
| 56 | """Language model embeddings. |
| 57 | |
| 58 | Arguments: |
| 59 | hidden_size: hidden size |
| 60 | vocab_size: vocabulary size |
| 61 | max_sequence_length: maximum size of sequence. This |
| 62 | is used for positional embedding |
| 63 | embedding_dropout_prob: dropout probability for embeddings |
| 64 | init_method: weight initialization method |
| 65 | num_tokentypes: size of the token-type embeddings. 0 value |
| 66 | will ignore this embedding |
| 67 | """ |
| 68 | |
| 69 | def __init__(self, |
| 70 | cfg, |
| 71 | ): |
| 72 | super().__init__() |
| 73 | self.hidden_size = cfg.hidden_size |
| 74 | self.params_dtype = cfg.params_dtype |
| 75 | self.fp32_residual_connection = cfg.fp32_residual_connection |
| 76 | self.embedding_weights_in_fp32 = cfg.embedding_weights_in_fp32 |
| 77 | self.word_embeddings = torch.nn.Embedding( |
| 78 | cfg.padded_vocab_size, self.hidden_size, |
| 79 | ) |
| 80 | self.embedding_dropout = torch.nn.Dropout(cfg.hidden_dropout) |
| 81 | |
| 82 | def forward(self, input_ids): |
| 83 | # Embeddings. |
| 84 | if self.embedding_weights_in_fp32: |
| 85 | self.word_embeddings = self.word_embeddings.to(torch.float32) |
| 86 | embeddings = self.word_embeddings(input_ids) |
| 87 | if self.embedding_weights_in_fp32: |
| 88 | embeddings = embeddings.to(self.params_dtype) |
| 89 | self.word_embeddings = self.word_embeddings.to(self.params_dtype) |
| 90 | |
| 91 | # Data format change to avoid explicit transposes : [b s h] --> [s b h]. |
| 92 | embeddings = embeddings.transpose(0, 1).contiguous() |
| 93 | |
| 94 | # If the input flag for fp32 residual connection is set, convert for float. |
| 95 | if self.fp32_residual_connection: |
| 96 | embeddings = embeddings.float() |
| 97 | |
| 98 | # Dropout. |
| 99 | embeddings = self.embedding_dropout(embeddings) |
| 100 | |
| 101 | return embeddings |
| 102 | |
| 103 | |
| 104 |