Fix llama2 weight loader (#1317)

This commit is contained in:
Lianmin Zheng
2024-09-03 05:32:14 -07:00
committed by GitHub
parent c500f96bb1
commit 12cb115d38
2 changed files with 8 additions and 29 deletions

View File

@@ -334,13 +334,13 @@ class LlamaForCausalLM(nn.Module):
for name, loaded_weight in weights:
if "rotary_emb.inv_freq" in name or "projector" in name:
return
continue
if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:
# Models trained using ColossalAI may include these tensors in
# the checkpoint. Skip them.
return
continue
if name.startswith("model.vision_tower") and name not in params_dict:
return
continue
for param_name, weight_name, shard_id in stacked_params_mapping:
if weight_name not in name:
@@ -356,7 +356,7 @@ class LlamaForCausalLM(nn.Module):
else:
# Skip loading extra bias for GPTQ models.
if name.endswith(".bias") and name not in params_dict:
return
continue
param = params_dict[name]
weight_loader = getattr(param, "weight_loader", default_weight_loader)
weight_loader(param, loaded_weight)