Fix Gemma activation function (#214)
* Update save.py * Update save.py * Update save.py * save * trainer * spaces * original * Gemma * Update pyproject.toml * Update mapper.py * Update fast_lora.py * FastGemmaModel * model_type * Update llama.py * Update llama.py * Update gemma.py * Update gemma.py * Update gemma.py * Update llama.py * Update llama.py * Update fast_lora.py * Update llama.py * Update llama.py * Update cross_entropy_loss.py * Update llama.py * Update llama.py * gemma * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update fast_lora.py * Update fast_lora.py * Fast CE Loss * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * CE * Update llama.py * Update llama.py * Update cross_entropy_loss.py * Update geglu.py * Update cross_entropy_loss.py * revert * Update llama.py * Update llama.py * norm * Update gemma.py * Update gemma.py * position_ids * Update gemma.py * Update gemma.py * pos * Update llama.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update cross_entropy_loss.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update llama.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update llama.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * revert * revert * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update llama.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update cross_entropy_loss.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * rope * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * llama * Update llama.py * gemma * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update save.py * RoPE * Update llama.py * Update llama.py * Update llama.py * Update gemma.py * correct_dtype * Update gemma.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Chat Templates * Update README.md * Update README.md * Update llama.py * DoRA * Update _utils.py * Update chat_templates.py * Update llama.py * Hotfix - fix DoRA, Gemma prompt template (#202) (#203) * Update save.py * saving * Update save.py * Update save.py * Update save.py * Update save.py * Update save.py * Update save.py * Update save.py * Update save.py * Update save.py * Update save.py * Update save.py * Update save.py * Update save.py * Update __init__.py * Update save.py * Update save.py * Update save.py * save * trainer * spaces * original * Gemma * Update pyproject.toml * Update mapper.py * Update fast_lora.py * FastGemmaModel * model_type * Update llama.py * Update llama.py * Update gemma.py * Update gemma.py * Update gemma.py * Update llama.py * Update llama.py * Update fast_lora.py * Update llama.py * Update llama.py * Update cross_entropy_loss.py * Update llama.py * Update llama.py * gemma * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update fast_lora.py * Update fast_lora.py * Fast CE Loss * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * CE * Update llama.py * Update llama.py * Update cross_entropy_loss.py * Update geglu.py * Update cross_entropy_loss.py * revert * Update llama.py * Update llama.py * norm * Update gemma.py * Update gemma.py * position_ids * Update gemma.py * Update gemma.py * pos * Update llama.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update cross_entropy_loss.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update llama.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update llama.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * revert * revert * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update llama.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update cross_entropy_loss.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * rope * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * llama * Update llama.py * gemma * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update save.py * RoPE * Update llama.py * Update llama.py * Update llama.py * Update gemma.py * correct_dtype * Update gemma.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Chat Templates * Update README.md * Update README.md * Update llama.py * DoRA * Update _utils.py * Update chat_templates.py * Update pyproject.toml * Small fixes * Update pyproject.toml * Approx gelu * Update geglu.py * Approx gelu * Update llama.py * Update __init__.py * Update __init__.py * Update _utils.py * Update geglu.py
This commit is contained in:
Daniel Han
GitHub
parent
851ab78fd2
commit
05a2d7d75f
8 changed files with 180 additions and 31 deletions
|
|
@ -33,8 +33,8 @@ exclude = ["images*"]
|
|||
|
||||
[project.optional-dependencies]
|
||||
huggingface = [
|
||||
"transformers>=4.38.0",
|
||||
"datasets",
|
||||
"transformers>=4.38.2",
|
||||
"datasets>=2.16.0",
|
||||
"sentencepiece",
|
||||
"accelerate>=0.26.1",
|
||||
"trl>=0.7.9",
|
||||
|
|
@ -64,6 +64,16 @@ cu121onlytorch211 = [
|
|||
"xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.23-cp310-cp310-manylinux2014_x86_64.whl ; python_version=='3.10'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.23-cp311-cp311-manylinux2014_x86_64.whl ; python_version=='3.11'",
|
||||
]
|
||||
cu118onlytorch212 = [
|
||||
"xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.23.post1%2Bcu118-cp39-cp39-manylinux2014_x86_64.whl ; python_version=='3.9'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.23.post1%2Bcu118-cp310-cp310-manylinux2014_x86_64.whl ; python_version=='3.10'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.23.post1%2Bcu118-cp311-cp311-manylinux2014_x86_64.whl ; python_version=='3.11'",
|
||||
]
|
||||
cu121onlytorch212 = [
|
||||
"xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.23.post1-cp39-cp39-manylinux2014_x86_64.whl ; python_version=='3.9'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.23.post1-cp310-cp310-manylinux2014_x86_64.whl ; python_version=='3.10'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.23.post1-cp311-cp311-manylinux2014_x86_64.whl ; python_version=='3.11'",
|
||||
]
|
||||
cu118onlytorch220 = [
|
||||
"xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.24%2Bcu118-cp39-cp39-manylinux2014_x86_64.whl ; python_version=='3.9'",
|
||||
"xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.24%2Bcu118-cp310-cp310-manylinux2014_x86_64.whl ; python_version=='3.10'",
|
||||
|
|
|
|||
|
|
@ -16,11 +16,17 @@ from .cross_entropy_loss import fast_cross_entropy_loss
|
|||
from .rms_layernorm import fast_rms_layernorm
|
||||
from .rope_embedding import fast_rope_embedding, inplace_rope_embedding
|
||||
from .swiglu import swiglu_fg_kernel, swiglu_DWf_DW_dfg_kernel
|
||||
from .geglu import geglu_forward_kernel, geglu_backward_kernel
|
||||
from .geglu import (
|
||||
geglu_exact_forward_kernel,
|
||||
geglu_exact_backward_kernel,
|
||||
geglu_approx_forward_kernel,
|
||||
geglu_approx_backward_kernel,
|
||||
)
|
||||
from .fast_lora import (
|
||||
get_lora_parameters,
|
||||
apply_lora_mlp_swiglu,
|
||||
apply_lora_mlp_geglu,
|
||||
apply_lora_mlp_geglu_exact,
|
||||
apply_lora_mlp_geglu_approx,
|
||||
apply_lora_qkv,
|
||||
apply_lora_o,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -183,8 +183,8 @@ def apply_lora_mlp_swiglu(self, X):
|
|||
pass
|
||||
|
||||
|
||||
from .geglu import geglu_forward_kernel, geglu_backward_kernel
|
||||
def apply_lora_mlp_geglu(self, X):
|
||||
from .geglu import geglu_exact_forward_kernel, geglu_exact_backward_kernel
|
||||
def apply_lora_mlp_geglu_exact(self, X):
|
||||
gateW, gateW_quant, gateA, gateB, gateS = get_lora_parameters(self.gate_proj)
|
||||
upW, upW_quant, upA, upB, upS = get_lora_parameters(self. up_proj)
|
||||
downW, downW_quant, downA, downB, downS = get_lora_parameters(self.down_proj)
|
||||
|
|
@ -192,7 +192,21 @@ def apply_lora_mlp_geglu(self, X):
|
|||
gateW, gateW_quant, gateA, gateB, gateS,
|
||||
upW, upW_quant, upA, upB, upS,
|
||||
downW, downW_quant, downA, downB, downS,
|
||||
geglu_forward_kernel, geglu_backward_kernel,)
|
||||
geglu_exact_forward_kernel, geglu_exact_backward_kernel,)
|
||||
return out
|
||||
pass
|
||||
|
||||
|
||||
from .geglu import geglu_approx_forward_kernel, geglu_approx_backward_kernel
|
||||
def apply_lora_mlp_geglu_approx(self, X):
|
||||
gateW, gateW_quant, gateA, gateB, gateS = get_lora_parameters(self.gate_proj)
|
||||
upW, upW_quant, upA, upB, upS = get_lora_parameters(self. up_proj)
|
||||
downW, downW_quant, downA, downB, downS = get_lora_parameters(self.down_proj)
|
||||
out = LoRA_MLP.apply(X,
|
||||
gateW, gateW_quant, gateA, gateB, gateS,
|
||||
upW, upW_quant, upA, upB, upS,
|
||||
downW, downW_quant, downA, downB, downS,
|
||||
geglu_approx_forward_kernel, geglu_approx_backward_kernel,)
|
||||
return out
|
||||
pass
|
||||
|
||||
|
|
|
|||
|
|
@ -19,7 +19,7 @@ from .utils import calculate_settings
|
|||
|
||||
|
||||
@triton.jit
|
||||
def _forward_kernel(e, g, h, n_elements, BLOCK_SIZE : tl.constexpr,):
|
||||
def _exact_forward_kernel(e, g, h, n_elements, BLOCK_SIZE : tl.constexpr,):
|
||||
block_idx = tl.program_id(0)
|
||||
offsets = block_idx*BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
|
||||
mask = offsets < n_elements
|
||||
|
|
@ -38,18 +38,18 @@ def _forward_kernel(e, g, h, n_elements, BLOCK_SIZE : tl.constexpr,):
|
|||
pass
|
||||
|
||||
|
||||
def geglu_forward_kernel(gate, up):
|
||||
def geglu_exact_forward_kernel(gate, up):
|
||||
batch, seq_len, hd = gate.shape
|
||||
n_elements = gate.numel()
|
||||
out = torch.empty((batch, seq_len, hd), dtype = gate.dtype, device = "cuda")
|
||||
grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),)
|
||||
_forward_kernel[grid](gate, up, out, n_elements, BLOCK_SIZE = 1024,)
|
||||
_exact_forward_kernel[grid](gate, up, out, n_elements, BLOCK_SIZE = 1024,)
|
||||
return out
|
||||
pass
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _backward_kernel(DW, e, g, n_elements, BLOCK_SIZE : tl.constexpr,):
|
||||
def _exact_backward_kernel(DW, e, g, n_elements, BLOCK_SIZE : tl.constexpr,):
|
||||
"""
|
||||
f = 1/2 * e * (1 + erf(1/sqrt(2) * e))
|
||||
h = f * up
|
||||
|
|
@ -95,10 +95,109 @@ def _backward_kernel(DW, e, g, n_elements, BLOCK_SIZE : tl.constexpr,):
|
|||
pass
|
||||
|
||||
|
||||
def geglu_backward_kernel(DW, e, g):
|
||||
def geglu_exact_backward_kernel(DW, e, g):
|
||||
batch_seq_len, hd = e.shape
|
||||
n_elements = e.numel()
|
||||
grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),)
|
||||
_backward_kernel[grid](DW, e, g, n_elements, BLOCK_SIZE = 1024,)
|
||||
_exact_backward_kernel[grid](DW, e, g, n_elements, BLOCK_SIZE = 1024,)
|
||||
return DW, e, g
|
||||
pass
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _approx_forward_kernel(e, g, h, n_elements, BLOCK_SIZE : tl.constexpr,):
|
||||
block_idx = tl.program_id(0)
|
||||
offsets = block_idx*BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
|
||||
mask = offsets < n_elements
|
||||
|
||||
# f = 1/2 * e * (1 + tanh( sqrt(2/pi) * (x + 0.044715 * x^3 ) ))
|
||||
# f = 1/2 * e * (1 + tanh( sqrt(2/pi) * x * (1 + 0.044715 * x^2 ) ))
|
||||
# h = f * up
|
||||
s = 0.7978845608028654 # math.sqrt(2 / math.pi)
|
||||
|
||||
e_row = tl.load(e + offsets, mask = mask, other = 0).to(tl.float32)
|
||||
g_row = tl.load(g + offsets, mask = mask, other = 0)#.to(tl.float32)
|
||||
|
||||
f_row = 0.5 * e_row * (
|
||||
tl.math.tanh(s * e_row * (1.0 + 0.044715 * e_row * e_row)) \
|
||||
+ 1.0
|
||||
)
|
||||
f_row = f_row.to(g_row.dtype) # Exact copy from HF
|
||||
h_row = f_row * g_row
|
||||
|
||||
# Store h
|
||||
tl.store(h + offsets, h_row, mask = mask)
|
||||
pass
|
||||
|
||||
|
||||
def geglu_approx_forward_kernel(gate, up):
|
||||
batch, seq_len, hd = gate.shape
|
||||
n_elements = gate.numel()
|
||||
out = torch.empty((batch, seq_len, hd), dtype = gate.dtype, device = "cuda")
|
||||
grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),)
|
||||
_approx_forward_kernel[grid](gate, up, out, n_elements, BLOCK_SIZE = 1024,)
|
||||
return out
|
||||
pass
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _approx_backward_kernel(DW, e, g, n_elements, BLOCK_SIZE : tl.constexpr,):
|
||||
"""
|
||||
f = 1/2 * e * (1 + tanh( sqrt(2/pi) * x * (1 + 0.044715 * x^2 ) ))
|
||||
h = f * up
|
||||
|
||||
df/de (with help from https://arxiv.org/pdf/2305.12073.pdf :))
|
||||
df/de = 1/2 * [1 + tanh( sqrt(2/pi) * x * (1 + 0.044715 * x^2 ) )] +
|
||||
1/2 * sech^2 [ sqrt(2/pi) * x * (1 + 0.044715 * x^2 ) ] * \
|
||||
( sqrt(2/pi) * x * (1 + 0.044715 * x^2 * 3 ) )
|
||||
|
||||
Notice sech^2(x) = 1 - tanh^2(x)
|
||||
So reuse tanh( sqrt(2/pi) * x * (1 + 0.044715 * x^2 ) )
|
||||
|
||||
See https://www.desmos.com/calculator/nqprfoni6x
|
||||
"""
|
||||
block_idx = tl.program_id(0)
|
||||
offsets = block_idx*BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
|
||||
mask = offsets < n_elements
|
||||
|
||||
DW_row = tl.load(DW + offsets, mask = mask, other = 0)#.to(tl.float32)
|
||||
e_row = tl.load(e + offsets, mask = mask, other = 0).to(tl.float32)
|
||||
g_row = tl.load(g + offsets, mask = mask, other = 0)#.to(tl.float32)
|
||||
|
||||
# See https://www.desmos.com/calculator/nqprfoni6x
|
||||
s = 0.7978845608028654 # math.sqrt(2 / math.pi)
|
||||
a = s * e_row # a = sqrt(2 / pi) * x
|
||||
b = a * 0.044715 * e_row * e_row # b = a * 0.044715 * x^2
|
||||
T = 1.0 + tl.math.tanh(a + b)
|
||||
T2 = 0.5 * T
|
||||
# Q = 0.5 * -T * (T - 2.0) * (a + 3.0 * b)
|
||||
Q2 = -T2 * (T - 2.0) * (a + 3.0 * b)
|
||||
df_de = T2 + Q2 # 1/2 * (T + Q)
|
||||
|
||||
# f = 1/2 * e * (1 + tanh( sqrt(2/pi) * (x + 0.044715 * x^3 ) ))
|
||||
f_row = T2 * e_row
|
||||
f_row = f_row.to(DW_row.dtype)
|
||||
# h = f * g
|
||||
h_row = f_row * g_row
|
||||
# df = DW * f
|
||||
df_row = DW_row * f_row
|
||||
# dg = DW * g
|
||||
dg_row = DW_row * g_row
|
||||
|
||||
de_row = dg_row.to(tl.float32) * df_de
|
||||
de_row = de_row.to(DW_row.dtype)
|
||||
|
||||
# Store derivatives in buffers
|
||||
tl.store(DW + offsets, h_row, mask = mask) # h = f * g
|
||||
tl.store(e + offsets, df_row, mask = mask) # df = DW * f
|
||||
tl.store(g + offsets, de_row, mask = mask) # de
|
||||
pass
|
||||
|
||||
|
||||
def geglu_approx_backward_kernel(DW, e, g):
|
||||
batch_seq_len, hd = e.shape
|
||||
n_elements = e.numel()
|
||||
grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),)
|
||||
_approx_backward_kernel[grid](DW, e, g, n_elements, BLOCK_SIZE = 1024,)
|
||||
return DW, e, g
|
||||
pass
|
||||
|
|
|
|||
|
|
@ -25,7 +25,7 @@ from platform import system as platform_system
|
|||
platform_system = platform_system()
|
||||
import math
|
||||
|
||||
__version__ = "2024.2"
|
||||
__version__ = "2024.3"
|
||||
|
||||
# Get Flash Attention v2 if Ampere (RTX 30xx, A100)
|
||||
major_version, minor_version = torch.cuda.get_device_capability()
|
||||
|
|
|
|||
|
|
@ -12,11 +12,16 @@
|
|||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from transformers.utils.notebook import (
|
||||
IntervalStrategy,
|
||||
NotebookTrainingTracker,
|
||||
NotebookProgressCallback,
|
||||
)
|
||||
try:
|
||||
from transformers.utils.notebook import (
|
||||
IntervalStrategy,
|
||||
NotebookTrainingTracker,
|
||||
NotebookProgressCallback,
|
||||
)
|
||||
HAS_NOTEBOOK = True
|
||||
except:
|
||||
HAS_NOTEBOOK = False
|
||||
pass
|
||||
|
||||
DPOTrainer_metrics = [
|
||||
"rewards/chosen",
|
||||
|
|
@ -101,13 +106,15 @@ pass
|
|||
|
||||
|
||||
def PatchDPOTrainer():
|
||||
from transformers.trainer import is_in_notebook
|
||||
if is_in_notebook():
|
||||
# Patch DPO notebook printing
|
||||
NotebookTrainingTracker.write_line = NotebookTrainingTracker_write_line
|
||||
from transformers.trainer import DEFAULT_PROGRESS_CALLBACK
|
||||
DEFAULT_PROGRESS_CALLBACK.on_train_begin = NotebookProgressCallback_on_train_begin
|
||||
DEFAULT_PROGRESS_CALLBACK.on_log = NotebookProgressCallback_on_log
|
||||
if HAS_NOTEBOOK:
|
||||
from transformers.trainer import is_in_notebook
|
||||
if is_in_notebook():
|
||||
# Patch DPO notebook printing
|
||||
NotebookTrainingTracker.write_line = NotebookTrainingTracker_write_line
|
||||
from transformers.trainer import DEFAULT_PROGRESS_CALLBACK
|
||||
DEFAULT_PROGRESS_CALLBACK.on_train_begin = NotebookProgressCallback_on_train_begin
|
||||
DEFAULT_PROGRESS_CALLBACK.on_log = NotebookProgressCallback_on_log
|
||||
pass
|
||||
pass
|
||||
pass
|
||||
|
||||
|
|
|
|||
|
|
@ -48,7 +48,7 @@ def fast_geglu_inference(self, X):
|
|||
|
||||
gate = fast_linear_forward(self.gate_proj, X, out = temp[0])
|
||||
up = fast_linear_forward(self. up_proj, X, out = temp[1])
|
||||
gate = torch.nn.functional.gelu(gate)
|
||||
gate = torch.nn.functional.gelu(gate, approximate = "tanh")
|
||||
gate *= up
|
||||
|
||||
# X = self.down_proj(gate)
|
||||
|
|
@ -70,7 +70,7 @@ def GemmaDecoderLayer_fast_forward(
|
|||
padding_mask: Optional[torch.LongTensor] = None,
|
||||
*args, **kwargs,
|
||||
):
|
||||
if False:#past_key_value is not None:
|
||||
if past_key_value is not None:
|
||||
do_prefill = not hasattr(self.self_attn, "paged_attention")
|
||||
|
||||
# Self Attention
|
||||
|
|
@ -267,6 +267,9 @@ class FastGemmaModel(FastLlamaModel):
|
|||
# Patch RMS Layernorm
|
||||
for name, module in model.named_modules():
|
||||
if isinstance(module, GemmaRMSNorm):
|
||||
# Must be in float32
|
||||
# https://github.com/keras-team/keras-nlp/blob/v0.8.2/keras_nlp/models/gemma/rms_normalization.py#L36
|
||||
module = module.to(torch.float32)
|
||||
module.weight += 1.0 # return output * (1 + self.weight)
|
||||
if not hasattr(module, "variance_epsilon"):
|
||||
module.variance_epsilon = module.eps # Gemma doesn't use variance_epsilon
|
||||
|
|
|
|||
|
|
@ -511,7 +511,12 @@ def LlamaModel_fast_forward(
|
|||
# Mormalized from Gemma
|
||||
if self.config.model_type == "gemma":
|
||||
inputs_requires_grad = inputs_embeds.requires_grad
|
||||
if inputs_requires_grad: inputs_embeds.requires_grad_(False)
|
||||
if not inputs_embeds.is_leaf:
|
||||
inputs_embeds = inputs_embeds.detach()
|
||||
inputs_requires_grad = True
|
||||
elif inputs_requires_grad:
|
||||
inputs_embeds.requires_grad_(False)
|
||||
pass
|
||||
inputs_embeds *= math_sqrt(self.config.hidden_size)
|
||||
if inputs_requires_grad: inputs_embeds.requires_grad_(True)
|
||||
pass
|
||||
|
|
@ -522,7 +527,12 @@ def LlamaModel_fast_forward(
|
|||
# Careful for inference the attention_mask is size (1, kv_seq_len)
|
||||
# Whilst the input_embeds is size (1, 1, 4096)
|
||||
inputs_requires_grad = inputs_embeds.requires_grad
|
||||
if inputs_requires_grad: inputs_embeds.requires_grad_(False)
|
||||
if not inputs_embeds.is_leaf:
|
||||
inputs_embeds = inputs_embeds.detach()
|
||||
inputs_requires_grad = True
|
||||
elif inputs_requires_grad:
|
||||
inputs_embeds.requires_grad_(False)
|
||||
pass
|
||||
inputs_embeds *= attention_mask.unsqueeze(0).transpose(0, 1).transpose(1, 2)
|
||||
if inputs_requires_grad: inputs_embeds.requires_grad_(True)
|
||||
pass
|
||||
|
|
@ -1335,7 +1345,7 @@ class FastLlamaModel:
|
|||
|
||||
if model_type == "llama": apply_lora_mlp = apply_lora_mlp_swiglu
|
||||
elif model_type == "mistral": apply_lora_mlp = apply_lora_mlp_swiglu
|
||||
elif model_type == "gemma": apply_lora_mlp = apply_lora_mlp_geglu
|
||||
elif model_type == "gemma": apply_lora_mlp = apply_lora_mlp_geglu_approx
|
||||
else:
|
||||
raise NotImplementedError(f"Unsloth: {model_type} is not yet implemented!")
|
||||
pass
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue