(
self,
input_ids,
position_ids,
attention_mask,
labels=None,
tokentype_ids=None,
layer_past=None,
get_key_value=False,
forward_method_parallel_output=None,
prompt_length=None,
context_length=None,
)
| 46 | self.language_model.set_input_tensor(input_tensor) |
| 47 | |
| 48 | def forward( |
| 49 | self, |
| 50 | input_ids, |
| 51 | position_ids, |
| 52 | attention_mask, |
| 53 | labels=None, |
| 54 | tokentype_ids=None, |
| 55 | layer_past=None, |
| 56 | get_key_value=False, |
| 57 | forward_method_parallel_output=None, |
| 58 | prompt_length=None, |
| 59 | context_length=None, |
| 60 | ): |
| 61 | |
| 62 | # Language model. |
| 63 | lm_output = self.language_model(input_ids, |
| 64 | position_ids, |
| 65 | attention_mask, |
| 66 | tokentype_ids=tokentype_ids, |
| 67 | layer_past=layer_past, |
| 68 | get_key_value=get_key_value, |
| 69 | prompt_length=prompt_length, |
| 70 | context_length=context_length) |
| 71 | |
| 72 | if get_key_value: |
| 73 | lm_output, presents = lm_output |
| 74 | |
| 75 | lm_output = torch.add(lm_output, 0) |
| 76 | # Output. |
| 77 | parallel_output = self.parallel_output |
| 78 | if forward_method_parallel_output is not None: |
| 79 | parallel_output = forward_method_parallel_output |
| 80 | output = parallel_lm_logits( |
| 81 | lm_output, |
| 82 | self.language_model.embedding.word_embeddings.weight, |
| 83 | parallel_output) |
| 84 | |
| 85 | if get_key_value: |
| 86 | output = [output, presents] |
| 87 | |
| 88 | if labels is None: |
| 89 | return output |
| 90 | else: |
| 91 | if self.fp16_lm_cross_entropy: |
| 92 | assert output.dtype == torch.half |
| 93 | loss = mpu.vocab_parallel_cross_entropy(output, labels) |
| 94 | else: |
| 95 | loss = mpu.vocab_parallel_cross_entropy(output.float(), labels) |
| 96 | |
| 97 | return loss |
| 98 | |
| 99 | def state_dict_for_save_checkpoint(self, destination=None, prefix='', |
| 100 | keep_vars=False): |
nothing calls this directly
no test coverage detected