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

reorg a bit

Browse Source
This commit is contained in:
Aman Karmani
2023-09-05 02:21:24 +00:00
parent 72a6fe1c1f
commit fc8766e502
1 changed files with 1 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
src/axolotl/monkeypatch/llama_attn_hijack_flash.py
View File

@@ -64,14 +64,13 @@ def replace_llama_attn_with_flash_attn(packed: Optional[bool] = False):
try: try:
from flash_attn.ops.rms_norm import RMSNorm from flash_attn.ops.rms_norm import RMSNorm
LOG.info("patching with flash_attn.ops.rms_norm")
class LlamaRMSNorm(RMSNorm): class LlamaRMSNorm(RMSNorm):
"""Patched LLamaRMSNorm""" """Patched LLamaRMSNorm"""
def __init__(self, hidden_size, eps=1e-6): def __init__(self, hidden_size, eps=1e-6):
super().__init__(hidden_size, eps=eps) super().__init__(hidden_size, eps=eps)
LOG.info("patching with flash_attn.ops.rms_norm")
transformers.models.llama.modeling_llama.LlamaRMSNorm = LlamaRMSNorm transformers.models.llama.modeling_llama.LlamaRMSNorm = LlamaRMSNorm
except ImportError: except ImportError:
LOG.info( LOG.info(