This commit is contained in:
Daniel Han 2024-07-07 09:49:55 -07:00
commit 5813e85bb9
5 changed files with 21 additions and 9 deletions

View file

@ -619,7 +619,11 @@ def patch_linear_scaling(
f"from {model_filepath} import logger, "\
f"{model_name.title()}Attention, {model_name.title()}Config"
function = inspect.getsource(attention_module.__init__)
try:
function = inspect.getsource(attention_module.__init__)
except:
# Most likely already patched!
return None, None
where = function.find("def")
function = function.split("\n")
function = "\n".join(x[where:] for x in function)

View file

@ -281,8 +281,10 @@ class FastGemmaModel(FastLlamaModel):
scaled_rope_module = GemmaFixedLinearScalingRotaryEmbedding,
attention_module = GemmaAttention,
)
exec(function, globals())
GemmaAttention.__init__ = eval(init_name)
if init_name is not None:
exec(function, globals())
GemmaAttention.__init__ = eval(init_name)
pass
GemmaAttention .forward = LlamaAttention_fast_forward
GemmaSdpaAttention .forward = LlamaAttention_fast_forward
GemmaFlashAttention2.forward = LlamaAttention_fast_forward

View file

@ -436,8 +436,10 @@ class FastGemma2Model(FastLlamaModel):
scaled_rope_module = GemmaFixedLinearScalingRotaryEmbedding,
attention_module = Gemma2Attention,
)
exec(function, globals())
Gemma2Attention.__init__ = eval(init_name)
if init_name is not None:
exec(function, globals())
Gemma2Attention.__init__ = eval(init_name)
pass
Gemma2Attention .forward = Gemma2Attention_fast_forward
Gemma2SdpaAttention .forward = Gemma2Attention_fast_forward
Gemma2FlashAttention2.forward = Gemma2Attention_fast_forward

View file

@ -277,8 +277,10 @@ class FastMistralModel(FastLlamaModel):
scaled_rope_module = LlamaLinearScalingRotaryEmbedding,
attention_module = MistralAttention,
)
exec(function, globals())
MistralAttention.__init__ = eval(init_name)
if init_name is not None:
exec(function, globals())
MistralAttention.__init__ = eval(init_name)
pass
MistralAttention .forward = MistralAttention_fast_forward
MistralSdpaAttention .forward = MistralAttention_fast_forward
MistralFlashAttention2.forward = MistralAttention_fast_forward

View file

@ -45,8 +45,10 @@ class FastQwen2Model(FastLlamaModel):
scaled_rope_module = LlamaLinearScalingRotaryEmbedding,
attention_module = Qwen2Attention,
)
exec(function, globals())
Qwen2Attention.__init__ = eval(init_name)
if init_name is not None:
exec(function, globals())
Qwen2Attention.__init__ = eval(init_name)
pass
Qwen2Attention .forward = LlamaAttention_fast_forward
Qwen2SdpaAttention .forward = LlamaAttention_fast_forward
Qwen2FlashAttention2.forward = LlamaAttention_fast_forward