transformers: Add eager attention option to make Gemma-2 work properly (#6188)

This commit is contained in:
GralchemOz 2024-07-01 23:08:08 +08:00 committed by GitHub
parent 19a56dd538
commit 8a39f579d8
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 7 additions and 0 deletions

View file

@ -146,6 +146,9 @@ def huggingface_loader(model_name):
if shared.args.force_safetensors:
params['force_safetensors'] = True
if shared.args.use_eager_attention:
params['attn_implementation'] = 'eager'
config = AutoConfig.from_pretrained(path_to_model, trust_remote_code=shared.args.trust_remote_code)
if 'chatglm' in model_name.lower():