aicloud
/
chatglm-6b
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Close CPU fusion on Mac

This commit is contained in:
duzx16 2023-03-23 22:43:06 +08:00
parent d2bbc82a2c
commit 4a9b711e61
1 changed files with 8 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
modeling_chatglm.py
View File

@ -5,6 +5,7 @@ import copy
import os
import warnings
import re
import sys
import torch
import torch.utils.checkpoint
@ -32,6 +33,8 @@ from transformers.generation.utils import LogitsProcessorList, StoppingCriteriaL
from .configuration_chatglm import ChatGLMConfig
# flags required to enable jit fusion kernels
if sys.platform != 'darwin':
torch._C._jit_set_profiling_mode(False)
torch._C._jit_set_profiling_executor(False)
torch._C._jit_override_can_fuse_on_cpu(True)
@ -266,7 +269,7 @@ def attention_fn(
if not (attention_mask == 0).all():
# if auto-regressive, skip
attention_scores.masked_fill_(attention_mask, -10000.0)
dtype = attention_scores.type()
dtype = attention_scores.dtype
attention_scores = attention_scores.float()
attention_scores = attention_scores * query_key_layer_scaling_coeff