Pass all tokens through both normal and prompt embedding tables. Tokens are masked so that "normal" embedding only see "normal" tokens. Same logic for "prompt" embedding. After those two embedding, combine results based on whether the token was "normal" or "promp
(self, tokens, prompt_embedding_table, tasks, task_vocab_size)
| 140 | self.vocab_size = vocab_size |
| 141 | |
| 142 | def forward(self, tokens, prompt_embedding_table, tasks, task_vocab_size): |
| 143 | """ |
| 144 | Pass all tokens through both normal and prompt embedding tables. |
| 145 | Tokens are masked so that "normal" embedding only see "normal" tokens. Same logic for "prompt" embedding. |
| 146 | After those two embedding, combine results based on whether the token was "normal" or "prompt-tuned". |
| 147 | |
| 148 | Parameters: |
| 149 | tokens : Tensor |
| 150 | the ids to embed, size [batch_size, seq_len] |
| 151 | |
| 152 | prompt_embedding_table : Tensor |
| 153 | the additional embedding table for prompt-tuned tokens, size [num_tasks * num_tokens_per_task, hidden_size] |
| 154 | |
| 155 | tasks: Tensor |
| 156 | the task required by each token, size [batch_size, seq_len] |
| 157 | |
| 158 | task_vocab_size: Tensor |
| 159 | the number of tokens used for each task, should be equal to prompt_embedding_table's num_tokens_per_task, size [1] |
| 160 | |
| 161 | Returns: |
| 162 | Tokens' embedding |
| 163 | """ |
| 164 | # do not use ">=" because internally the layer works with floating points |
| 165 | prompt_tokens_mask = tokens > (self.vocab_size - 1) |
| 166 | |
| 167 | # clip tokens in the [0, vocab_size) range |
| 168 | normal_tokens = where(prompt_tokens_mask, self.vocab_size - 1, tokens) |
| 169 | normal_embeddings = embedding(normal_tokens, self.weight.value, |
| 170 | self.tp_size, self.tp_group, |
| 171 | self.sharding_dim, self.tp_rank) |
| 172 | |
| 173 | # put virtual tokens in the [0, max_prompt_vocab_size) range |
| 174 | prompt_tokens = where(prompt_tokens_mask, tokens - self.vocab_size, 0) |
| 175 | # add offsets to match the concatenated embedding tables |
| 176 | tasks = tasks * task_vocab_size |
| 177 | |
| 178 | # tasks: [batch_size, seq_len] |
| 179 | # prompt_tokens: [batch_size, seq_len] |
| 180 | # if speculative decoding is enabled the shape of prompt_tokens is [batch_size, seq_len + max_draft_len], |
| 181 | # so we need to expand tasks to [batch_size, seq_len + max_draft_len] |
| 182 | tasks = expand(tasks, shape(prompt_tokens)) |
| 183 | prompt_tokens = prompt_tokens + tasks |
| 184 | prompt_embeddings = embedding(prompt_tokens, prompt_embedding_table) |
| 185 | |
| 186 | # prompt_tokens_mask: [batch_size, seq_len] -> [batch_size, seq_len, 1] |
| 187 | # combine the correct sources of embedding: normal/prompt |
| 188 | return where(unsqueeze(prompt_tokens_mask, -1), prompt_embeddings, |
| 189 | normal_embeddings) |
| 190 | |
| 191 | |
| 192 | class LabelEmbedding(Module): |