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

ensure flash-attn fixes happen in both adapter/lora modes, and use torch_dtype

Browse Source
This commit is contained in:
Aman Karmani
2023-08-02 20:15:03 +00:00
parent 77085ea24e
commit 248bf90f89
1 changed files with 8 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
src/axolotl/utils/models.py
View File

@@ -331,6 +331,14 @@ def load_model(
model, use_gradient_checkpointing=cfg.gradient_checkpointing model, use_gradient_checkpointing=cfg.gradient_checkpointing
) )
if cfg.flash_attention:
for name, module in model.named_modules():
if "norm" in name:
module.to(torch_dtype)
if "lm_head" in name or "embed_tokens" in name:
if hasattr(module, "weight"):
module.to(torch_dtype)
model, lora_config = load_adapter(model, cfg, adapter) model, lora_config = load_adapter(model, cfg, adapter)
if cfg.ddp and not load_in_8bit: if cfg.ddp and not load_in_8bit:
@@ -407,14 +415,6 @@ def load_llama_adapter(model, cfg):
else: else:
model = get_peft_model(model, peft_config) model = get_peft_model(model, peft_config)
if cfg.flash_attention:
for name, module in model.named_modules():
if "norm" in name:
module.to(torch.float16)
if "lm_head" in name or "embed_tokens" in name:
if hasattr(module, "weight"):
module.to(torch.float16)
model.print_trainable_parameters() model.print_trainable_parameters()
return model, peft_config return model, peft_config