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

fix: assign linear head instead of loading state dict

Browse Source
This commit is contained in:
NanoCode012
2025-02-05 18:24:31 +07:00
parent 2d5f692fc0
commit 9e1c4de13c
1 changed files with 5 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
src/axolotl/integrations/lolcats/linear_llama/modeling_linear_llama.py
View File

@@ -113,15 +113,17 @@ class LinearLlamaForCausalLM(LlamaForCausalLM):
# initialize the model with prior weights # initialize the model with prior weights
new_model = cls(config=config) new_model = cls(config=config)
del new_model.model # remove the default model # remove the default model and lm_head
del new_model.model
del new_model.lm_head
new_model.model = convert_attention( new_model.model = convert_attention(
model.model, model.model,
attention_config=config.attention_config, attention_config=config.attention_config,
train_attention=train_attention, train_attention=train_attention,
remove_base_attn=remove_base_attn, remove_base_attn=remove_base_attn,
) )
new_model.lm_head = model.lm_head
new_model.lm_head.load_state_dict(model.lm_head.state_dict())
return new_model return new_model