(
model,
tokens,
position_ids,
attention_mask,
tokentype_ids,
layer_past=None,
get_key_value=None,
forward_method_parallel_output=None,
prompt_length=None,
context_length=None,
)
| 799 | |
| 800 | |
| 801 | def forward_step( |
| 802 | model, |
| 803 | tokens, |
| 804 | position_ids, |
| 805 | attention_mask, |
| 806 | tokentype_ids, |
| 807 | layer_past=None, |
| 808 | get_key_value=None, |
| 809 | forward_method_parallel_output=None, |
| 810 | prompt_length=None, |
| 811 | context_length=None, |
| 812 | ): |
| 813 | # Hidden size changes when not using recompute, need to tell p2p_communicate |
| 814 | # functions the correct size |
| 815 | args = get_args() |
| 816 | orig_seq_length = args.seq_length |
| 817 | args.seq_length = tokens.shape[1] |
| 818 | |
| 819 | # Forward pass through the model. |
| 820 | output_tensor = model( |
| 821 | tokens, |
| 822 | position_ids, |
| 823 | attention_mask, |
| 824 | tokentype_ids=tokentype_ids, |
| 825 | layer_past=layer_past, |
| 826 | get_key_value=get_key_value, |
| 827 | prompt_length=prompt_length, |
| 828 | context_length=context_length, |
| 829 | ) |
| 830 | |
| 831 | if get_key_value: |
| 832 | output_tensor, layer_past = output_tensor |
| 833 | |
| 834 | args.seq_length = orig_seq_length |
| 835 | if get_key_value: |
| 836 | return output_tensor, layer_past |
| 837 | |
| 838 | return output_tensor |
| 839 | |
| 840 | |
| 841 | def get_token_stream( |
no test coverage detected