RYDE-WORK/MiniCPM
1
0
Fork 0
mirror of https://github.com/RYDE-WORK/MiniCPM.git synced 2026-02-05 22:53:34 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Fix: flash_attn and cpu inference

Browse Source
This commit is contained in:
huangyuxiang03 2024-02-03 22:13:15 +08:00
parent c8ad2ec0d3
commit 5977a63783
2 changed files with 9 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
model/configuration_minicpm.py
View File

@ -174,6 +174,11 @@ class MiniCPMConfig(PretrainedConfig):
tie_word_embeddings=tie_word_embeddings, tie_word_embeddings=tie_word_embeddings,
**kwargs, **kwargs,
) )
try:
import flash_attn
self._attn_implementation = "flash_attention_2"
except:
pass
def _rope_scaling_validation(self): def _rope_scaling_validation(self):
""" """

8
model/modeling_minicpm.py
View File

@ -51,10 +51,11 @@ from transformers.utils.import_utils import is_torch_fx_available
from .configuration_minicpm import MiniCPMConfig from .configuration_minicpm import MiniCPMConfig
import re import re
try:
if is_flash_attn_2_available():
from flash_attn import flash_attn_func, flash_attn_varlen_func from flash_attn import flash_attn_func, flash_attn_varlen_func
from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa
except:
pass
# This makes `_prepare_4d_causal_attention_mask` a leaf function in the FX graph. # This makes `_prepare_4d_causal_attention_mask` a leaf function in the FX graph.
@ -125,7 +126,7 @@ ALL_LAYERNORM_LAYERS.append(MiniCPMRMSNorm)
class MiniCPMRotaryEmbedding(nn.Module): class MiniCPMRotaryEmbedding(nn.Module):
def __init__(self, dim, max_position_embeddings=2048, base=10000, device="cuda"): def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
super().__init__() super().__init__()
self.dim = dim self.dim = dim
@ -763,7 +764,6 @@ class MiniCPMDecoderLayer(nn.Module):
def __init__(self, config: MiniCPMConfig, layer_idx: int): def __init__(self, config: MiniCPMConfig, layer_idx: int):
super().__init__() super().__init__()
self.hidden_size = config.hidden_size self.hidden_size = config.hidden_size
self.self_attn = MINICPM_ATTENTION_CLASSES[config._attn_implementation](config=config, layer_idx=layer_idx) self.self_attn = MINICPM_ATTENTION_CLASSES[config._attn_implementation](config=config, layer_idx=layer_idx)
self.mlp = MiniCPMMLP(config) self.mlp = MiniCPMMLP(config)