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

Revert "run PR e2e docker CI tests in Modal" (#1220) [skip ci]

Browse Source
This commit is contained in:
Wing Lian
2024-01-26 16:50:44 -05:00
committed by GitHub
parent 36d053f6f0
commit 8da1633124
6 changed files with 42 additions and 141 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -645,10 +645,7 @@ def load_model(
if not cfg.fsdp:
# FSDP doesn't like mixed Float and BFloat16
for name, module in model.named_modules():
if (
any(m in name for m in ["norm", "gate"])
or "LayerNorm" in module.__class__.__name__
):
if any(m in name for m in ["norm", "gate"]):
module.to(torch.float32)
if model_config.model_type == "btlm":
# don't upcast lm_head for btlm
@@ -687,7 +684,7 @@ def load_model(
if needs_fa2_dtype or cfg.flash_attention:
LOG.info("converting modules to %s for flash attention", cfg.torch_dtype)
for name, module in model.named_modules():
if "norm" in name or "LayerNorm" in module.__class__.__name__:
if "norm" in name:
module.to(cfg.torch_dtype)
if any(m in name for m in embedding_modules):
if hasattr(module, "weight"):