tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

don't load the actual model when pre-loading to load modeling code

Browse Source
This commit is contained in:
Wing Lian
2023-09-20 13:37:32 -04:00
parent ec0958f4f8
commit ca476d7f8e
1 changed files with 3 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
src/axolotl/monkeypatch/btlm_attn_hijack_flash.py
View File

@@ -6,6 +6,7 @@ import importlib
import logging import logging
from typing import Optional, Tuple from typing import Optional, Tuple
import accelerate
import torch import torch
from flash_attn.flash_attn_interface import flash_attn_func from flash_attn.flash_attn_interface import flash_attn_func
from transformers import AutoConfig, AutoModelForCausalLM from transformers import AutoConfig, AutoModelForCausalLM
@@ -17,7 +18,8 @@ def replace_btlm_attn_with_flash_attn(model_name="cerebras/btlm-3b-8k-base"):
# this is a wonky hack to get the remotely loaded module # this is a wonky hack to get the remotely loaded module
model_config = AutoConfig.from_pretrained(model_name, trust_remote_code=True) model_config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
# we need to load the model here in order for modeling_btlm to be available # we need to load the model here in order for modeling_btlm to be available
AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True) with accelerate.init_empty_weights():
AutoModelForCausalLM(model_config)
module_name = model_config.__class__.__module__.replace( module_name = model_config.__class__.__module__.replace(
".configuration_btlm", ".modeling_btlm" ".configuration_btlm", ".modeling_btlm"
) )