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]
|
[project.optional-dependencies]
|
||||||
huggingface = [
|
huggingface = [
|
||||||
"transformers>=4.38.0",
|
"transformers>=4.38.2",
|
||||||
"datasets",
|
"datasets>=2.16.0",
|
||||||
"sentencepiece",
|
"sentencepiece",
|
||||||
"accelerate>=0.26.1",
|
"accelerate>=0.26.1",
|
||||||
"trl>=0.7.9",
|
"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-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'",
|
"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 = [
|
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-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'",
|
"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 .rms_layernorm import fast_rms_layernorm
|
||||||
from .rope_embedding import fast_rope_embedding, inplace_rope_embedding
|
from .rope_embedding import fast_rope_embedding, inplace_rope_embedding
|
||||||
from .swiglu import swiglu_fg_kernel, swiglu_DWf_DW_dfg_kernel
|
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 (
|
from .fast_lora import (
|
||||||
get_lora_parameters,
|
get_lora_parameters,
|
||||||
apply_lora_mlp_swiglu,
|
apply_lora_mlp_swiglu,
|
||||||
apply_lora_mlp_geglu,
|
apply_lora_mlp_geglu_exact,
|
||||||
|
apply_lora_mlp_geglu_approx,
|
||||||
apply_lora_qkv,
|
apply_lora_qkv,
|
||||||
apply_lora_o,
|
apply_lora_o,
|
||||||
)
|
)
|
||||||
|
|
|
||||||
|
|
@ -183,8 +183,8 @@ def apply_lora_mlp_swiglu(self, X):
|
||||||
pass
|
pass
|
||||||
|
|
||||||
|
|
||||||
from .geglu import geglu_forward_kernel, geglu_backward_kernel
|
from .geglu import geglu_exact_forward_kernel, geglu_exact_backward_kernel
|
||||||
def apply_lora_mlp_geglu(self, X):
|
def apply_lora_mlp_geglu_exact(self, X):
|
||||||
gateW, gateW_quant, gateA, gateB, gateS = get_lora_parameters(self.gate_proj)
|
gateW, gateW_quant, gateA, gateB, gateS = get_lora_parameters(self.gate_proj)
|
||||||
upW, upW_quant, upA, upB, upS = get_lora_parameters(self. up_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)
|
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,
|
gateW, gateW_quant, gateA, gateB, gateS,
|
||||||
upW, upW_quant, upA, upB, upS,
|
upW, upW_quant, upA, upB, upS,
|
||||||
downW, downW_quant, downA, downB, downS,
|
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
|
return out
|
||||||
pass
|
pass
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -19,7 +19,7 @@ from .utils import calculate_settings
|
||||||
|
|
||||||
|
|
||||||
@triton.jit
|
@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)
|
block_idx = tl.program_id(0)
|
||||||
offsets = block_idx*BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
|
offsets = block_idx*BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
|
||||||
mask = offsets < n_elements
|
mask = offsets < n_elements
|
||||||
|
|
@ -38,18 +38,18 @@ def _forward_kernel(e, g, h, n_elements, BLOCK_SIZE : tl.constexpr,):
|
||||||
pass
|
pass
|
||||||
|
|
||||||
|
|
||||||
def geglu_forward_kernel(gate, up):
|
def geglu_exact_forward_kernel(gate, up):
|
||||||
batch, seq_len, hd = gate.shape
|
batch, seq_len, hd = gate.shape
|
||||||
n_elements = gate.numel()
|
n_elements = gate.numel()
|
||||||
out = torch.empty((batch, seq_len, hd), dtype = gate.dtype, device = "cuda")
|
out = torch.empty((batch, seq_len, hd), dtype = gate.dtype, device = "cuda")
|
||||||
grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),)
|
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
|
return out
|
||||||
pass
|
pass
|
||||||
|
|
||||||
|
|
||||||
@triton.jit
|
@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))
|
f = 1/2 * e * (1 + erf(1/sqrt(2) * e))
|
||||||
h = f * up
|
h = f * up
|
||||||
|
|
@ -95,10 +95,109 @@ def _backward_kernel(DW, e, g, n_elements, BLOCK_SIZE : tl.constexpr,):
|
||||||
pass
|
pass
|
||||||
|
|
||||||
|
|
||||||
def geglu_backward_kernel(DW, e, g):
|
def geglu_exact_backward_kernel(DW, e, g):
|
||||||
batch_seq_len, hd = e.shape
|
batch_seq_len, hd = e.shape
|
||||||
n_elements = e.numel()
|
n_elements = e.numel()
|
||||||
grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),)
|
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
|
return DW, e, g
|
||||||
pass
|
pass
|
||||||
|
|
|
||||||
|
|
@ -25,7 +25,7 @@ from platform import system as platform_system
|
||||||
platform_system = platform_system()
|
platform_system = platform_system()
|
||||||
import math
|
import math
|
||||||
|
|
||||||
__version__ = "2024.2"
|
__version__ = "2024.3"
|
||||||
|
|
||||||
# Get Flash Attention v2 if Ampere (RTX 30xx, A100)
|
# Get Flash Attention v2 if Ampere (RTX 30xx, A100)
|
||||||
major_version, minor_version = torch.cuda.get_device_capability()
|
major_version, minor_version = torch.cuda.get_device_capability()
|
||||||
|
|
|
||||||
|
|
@ -12,11 +12,16 @@
|
||||||
# See the License for the specific language governing permissions and
|
# See the License for the specific language governing permissions and
|
||||||
# limitations under the License.
|
# limitations under the License.
|
||||||
|
|
||||||
from transformers.utils.notebook import (
|
try:
|
||||||
IntervalStrategy,
|
from transformers.utils.notebook import (
|
||||||
NotebookTrainingTracker,
|
IntervalStrategy,
|
||||||
NotebookProgressCallback,
|
NotebookTrainingTracker,
|
||||||
)
|
NotebookProgressCallback,
|
||||||
|
)
|
||||||
|
HAS_NOTEBOOK = True
|
||||||
|
except:
|
||||||
|
HAS_NOTEBOOK = False
|
||||||
|
pass
|
||||||
|
|
||||||
DPOTrainer_metrics = [
|
DPOTrainer_metrics = [
|
||||||
"rewards/chosen",
|
"rewards/chosen",
|
||||||
|
|
@ -101,13 +106,15 @@ pass
|
||||||
|
|
||||||
|
|
||||||
def PatchDPOTrainer():
|
def PatchDPOTrainer():
|
||||||
from transformers.trainer import is_in_notebook
|
if HAS_NOTEBOOK:
|
||||||
if is_in_notebook():
|
from transformers.trainer import is_in_notebook
|
||||||
# Patch DPO notebook printing
|
if is_in_notebook():
|
||||||
NotebookTrainingTracker.write_line = NotebookTrainingTracker_write_line
|
# Patch DPO notebook printing
|
||||||
from transformers.trainer import DEFAULT_PROGRESS_CALLBACK
|
NotebookTrainingTracker.write_line = NotebookTrainingTracker_write_line
|
||||||
DEFAULT_PROGRESS_CALLBACK.on_train_begin = NotebookProgressCallback_on_train_begin
|
from transformers.trainer import DEFAULT_PROGRESS_CALLBACK
|
||||||
DEFAULT_PROGRESS_CALLBACK.on_log = NotebookProgressCallback_on_log
|
DEFAULT_PROGRESS_CALLBACK.on_train_begin = NotebookProgressCallback_on_train_begin
|
||||||
|
DEFAULT_PROGRESS_CALLBACK.on_log = NotebookProgressCallback_on_log
|
||||||
|
pass
|
||||||
pass
|
pass
|
||||||
pass
|
pass
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -48,7 +48,7 @@ def fast_geglu_inference(self, X):
|
||||||
|
|
||||||
gate = fast_linear_forward(self.gate_proj, X, out = temp[0])
|
gate = fast_linear_forward(self.gate_proj, X, out = temp[0])
|
||||||
up = fast_linear_forward(self. up_proj, X, out = temp[1])
|
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
|
gate *= up
|
||||||
|
|
||||||
# X = self.down_proj(gate)
|
# X = self.down_proj(gate)
|
||||||
|
|
@ -70,7 +70,7 @@ def GemmaDecoderLayer_fast_forward(
|
||||||
padding_mask: Optional[torch.LongTensor] = None,
|
padding_mask: Optional[torch.LongTensor] = None,
|
||||||
*args, **kwargs,
|
*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")
|
do_prefill = not hasattr(self.self_attn, "paged_attention")
|
||||||
|
|
||||||
# Self Attention
|
# Self Attention
|
||||||
|
|
@ -267,6 +267,9 @@ class FastGemmaModel(FastLlamaModel):
|
||||||
# Patch RMS Layernorm
|
# Patch RMS Layernorm
|
||||||
for name, module in model.named_modules():
|
for name, module in model.named_modules():
|
||||||
if isinstance(module, GemmaRMSNorm):
|
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)
|
module.weight += 1.0 # return output * (1 + self.weight)
|
||||||
if not hasattr(module, "variance_epsilon"):
|
if not hasattr(module, "variance_epsilon"):
|
||||||
module.variance_epsilon = module.eps # Gemma doesn't use variance_epsilon
|
module.variance_epsilon = module.eps # Gemma doesn't use variance_epsilon
|
||||||
|
|
|
||||||
|
|
@ -511,7 +511,12 @@ def LlamaModel_fast_forward(
|
||||||
# Mormalized from Gemma
|
# Mormalized from Gemma
|
||||||
if self.config.model_type == "gemma":
|
if self.config.model_type == "gemma":
|
||||||
inputs_requires_grad = inputs_embeds.requires_grad
|
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)
|
inputs_embeds *= math_sqrt(self.config.hidden_size)
|
||||||
if inputs_requires_grad: inputs_embeds.requires_grad_(True)
|
if inputs_requires_grad: inputs_embeds.requires_grad_(True)
|
||||||
pass
|
pass
|
||||||
|
|
@ -522,7 +527,12 @@ def LlamaModel_fast_forward(
|
||||||
# Careful for inference the attention_mask is size (1, kv_seq_len)
|
# Careful for inference the attention_mask is size (1, kv_seq_len)
|
||||||
# Whilst the input_embeds is size (1, 1, 4096)
|
# Whilst the input_embeds is size (1, 1, 4096)
|
||||||
inputs_requires_grad = inputs_embeds.requires_grad
|
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)
|
inputs_embeds *= attention_mask.unsqueeze(0).transpose(0, 1).transpose(1, 2)
|
||||||
if inputs_requires_grad: inputs_embeds.requires_grad_(True)
|
if inputs_requires_grad: inputs_embeds.requires_grad_(True)
|
||||||
pass
|
pass
|
||||||
|
|
@ -1335,7 +1345,7 @@ class FastLlamaModel:
|
||||||
|
|
||||||
if model_type == "llama": apply_lora_mlp = apply_lora_mlp_swiglu
|
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 == "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:
|
else:
|
||||||
raise NotImplementedError(f"Unsloth: {model_type} is not yet implemented!")
|
raise NotImplementedError(f"Unsloth: {model_type} is not yet implemented!")
|
||||||
pass
|
pass
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue