ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Gemma precision

Browse source
This commit is contained in:
Daniel Han-Chen 2024-03-06 04:14:11 +11:00
commit be81c07d43
3 changed files with 69 additions and 44 deletions
Download patch file Download diff file Expand all files Collapse all files

73
unsloth/kernels/rms_layernorm.py
View file

@ -18,7 +18,6 @@ import torch
from .utils import calculate_settings
@triton.heuristics({"ADD_ONE": lambda args: args["ADD_ONE"],})
@triton.jit
def _rms_layernorm_forward(
Y, Y_row_stride,
@ -26,8 +25,7 @@ def _rms_layernorm_forward(
W, W_row_stride,
r, r_row_stride,
n_cols, eps,
ADD_ONE: tl.constexpr,
BLOCK_SIZE : tl.constexpr,
BLOCK_SIZE : tl.constexpr
):
"""
Fast RMS Layernorm kernel
@ -50,20 +48,12 @@ def _rms_layernorm_forward(
tl.store(r, inv_var)
normed = X_row * inv_var
normed = normed.to(W_row.dtype) # Exact copy from HF
# For Gemma - cannot do += 1 since float16 - maybe use FMADD
if not ADD_ONE:
output = normed * W_row
else:
# Error analysis shows we need to do +1 in float32 then downcast to float16
output = normed * (W_row.to(tl.float32) + 1.0).to(W_row.dtype)
pass
output = normed * W_row
tl.store(Y + col_offsets, output, mask = mask)
pass
@triton.heuristics({"ADD_ONE": lambda args: args["ADD_ONE"],})
@triton.heuristics({"GEMMA": lambda args: args["GEMMA"],})
@triton.jit
def _rms_layernorm_backward(
dY, dY_row_stride,
@ -72,7 +62,7 @@ def _rms_layernorm_backward(
r, r_row_stride,
dW, dW_row_stride,
n_cols, eps,
ADD_ONE: tl.constexpr,
GEMMA : tl.constexpr,
BLOCK_SIZE : tl.constexpr,
):
"""
@ -96,12 +86,8 @@ def _rms_layernorm_backward(
inv_var = tl.load(r).to(tl.float32)
normed = X_row * inv_var
# For Gemma - cannot do += 1 since float16 - maybe use FMADD
if not ADD_ONE:
dY_W = dY_row * W_row
else:
dY_W = dY_row * (W_row + 1.0)
pass
if GEMMA: dY_W = dY_row * (W_row + 1.0)
else: dY_W = dY_row * W_row
rowsum_dY_normed = tl.sum(dY_W * normed, axis = 0)
output = inv_var/n_cols * (n_cols*dY_W - normed*rowsum_dY_normed)
@ -109,9 +95,42 @@ def _rms_layernorm_backward(
pass
@triton.jit
def _gemma_rms_layernorm_forward(
Y, Y_row_stride,
X, X_row_stride,
W, W_row_stride,
r, r_row_stride,
n_cols, eps,
BLOCK_SIZE : tl.constexpr,
):
# Copies https://github.com/google-deepmind/gemma/blob/main/gemma/layers.py#L31
# and https://github.com/keras-team/keras-nlp/blob/v0.8.2/keras_nlp/models/gemma/rms_normalization.py#L33
# exactly. Essentially all in float32!
row_idx = tl.program_id(0)
col_offsets = tl.arange(0, BLOCK_SIZE)
mask = col_offsets < n_cols
Y += row_idx * Y_row_stride
X += row_idx * X_row_stride
r += row_idx * r_row_stride
X_row = tl.load(X + col_offsets, mask = mask, other = 0).to(tl.float32)
W_row = tl.load(W + col_offsets, mask = mask, other = 0).to(tl.float32)
row_var = tl.sum(X_row * X_row, axis = 0) / n_cols
inv_var = tl.math.rsqrt(row_var + eps)
tl.store(r, inv_var)
normed = X_row * inv_var
output = normed * (W_row + 1.0)
tl.store(Y + col_offsets, output, mask = mask)
pass
class Fast_RMS_Layernorm(torch.autograd.Function):
@staticmethod
def forward(ctx, X, W, eps, add_one = False):
def forward(ctx, X, W, eps, gemma = False):
shape = X.shape
dim = shape[-1]
X = X.view(-1, dim)
@ -121,20 +140,20 @@ class Fast_RMS_Layernorm(torch.autograd.Function):
Y = torch.empty((n_rows, n_cols), dtype = X.dtype, device = "cuda")
r = torch.empty(n_rows, dtype = torch.float32, device = "cuda")
_rms_layernorm_forward[(n_rows,)](
fx = _gemma_rms_layernorm_forward if gemma else _rms_layernorm_forward
fx[(n_rows,)](
Y, Y.stride(0),
X, X.stride(0),
W, W.stride(0),
r, r.stride(0),
n_cols, eps,
ADD_ONE = add_one,
BLOCK_SIZE = BLOCK_SIZE,
num_warps = num_warps,
)
ctx.eps = eps
ctx.BLOCK_SIZE = BLOCK_SIZE
ctx.num_warps = num_warps
ctx.ADD_ONE = add_one
ctx.GEMMA = gemma
ctx.save_for_backward(X, W, r)
return Y.view(*shape)
pass
@ -155,7 +174,7 @@ class Fast_RMS_Layernorm(torch.autograd.Function):
r, r .stride(0),
dW, dW.stride(0),
n_cols, ctx.eps,
ADD_ONE = ctx.ADD_ONE,
GEMMA = ctx.GEMMA,
BLOCK_SIZE = ctx.BLOCK_SIZE,
num_warps = ctx.num_warps,
)
@ -165,9 +184,9 @@ class Fast_RMS_Layernorm(torch.autograd.Function):
pass
def fast_rms_layernorm(layernorm, X, add_one = False):
def fast_rms_layernorm(layernorm, X, gemma = False):
W = layernorm.weight
eps = layernorm.variance_epsilon
out = Fast_RMS_Layernorm.apply(X, W, eps, add_one)
out = Fast_RMS_Layernorm.apply(X, W, eps, gemma)
return out
pass

32
unsloth/models/gemma.py
View file

@ -58,16 +58,15 @@ def fast_geglu_inference(self, X):
pass
def fast_rms_layernorm_inference_add_one(self, X, out_weight = None):
old_dtype = X.dtype
def fast_rms_layernorm_inference_gemma(self, X, out_weight):
XX = X.to(torch.float32)
variance = XX.square().mean(-1, keepdim = True)
variance += self.variance_epsilon
XX *= variance.rsqrt_()
X = XX.to(old_dtype) # Must preserve due to residual
out_weight = torch.add(self.weight, 1.0, out = out_weight)
X *= out_weight
return X
out_weight[:] = self.weight
out_weight += 1.0
XX *= out_weight
return XX.to(X.dtype)
pass
@ -86,10 +85,11 @@ def GemmaDecoderLayer_fast_forward(
):
if past_key_value is not None:
do_prefill = not hasattr(self.self_attn, "paged_attention")
out_weight = torch.empty(self.input_layernorm.weight.shape, dtype = torch.float32, dtype = "cuda")
# Self Attention
residual = hidden_states
hidden_states = fast_rms_layernorm_inference_add_one(self.input_layernorm, hidden_states)
hidden_states = fast_rms_layernorm_inference_gemma(self.input_layernorm, hidden_states, out_weight)
hidden_states, present_key_value = LlamaAttention_fast_forward_inference(
self.self_attn,
hidden_states,
@ -101,12 +101,12 @@ def GemmaDecoderLayer_fast_forward(
# Fully Connected
residual = hidden_states
hidden_states = fast_rms_layernorm_inference_add_one(self.post_attention_layernorm, hidden_states)
hidden_states = fast_rms_layernorm_inference_gemma(self.post_attention_layernorm, hidden_states, out_weight)
hidden_states = fast_geglu_inference(self.mlp, hidden_states)
hidden_states += residual
else:
residual = hidden_states
hidden_states = fast_rms_layernorm(self.input_layernorm, hidden_states, add_one = True)
hidden_states = fast_rms_layernorm(self.input_layernorm, hidden_states, gemma = True)
# hidden_states = self.input_layernorm(hidden_states)
hidden_states, self_attn_weights, present_key_value = self.self_attn(
hidden_states=hidden_states,
@ -122,7 +122,7 @@ def GemmaDecoderLayer_fast_forward(
# Fully Connected
residual = hidden_states
hidden_states = fast_rms_layernorm(self.post_attention_layernorm, hidden_states, add_one = True)
hidden_states = fast_rms_layernorm(self.post_attention_layernorm, hidden_states, gemma = True)
# hidden_states = self.post_attention_layernorm(hidden_states)
hidden_states = self.mlp(hidden_states)
hidden_states = residual + hidden_states
@ -151,16 +151,18 @@ def GemmaModel_fast_forward_inference(
):
# Fix out of bounds tokenization
input_ids = input_ids[:,:self.max_seq_length]
out_weight = torch.empty_like(self.layers[0].input_layernorm.weight)
out_weight = torch.empty_like(self.layers[0].input_layernorm.weight, dtype = torch.float32, device = "cuda")
hidden_states = self.embed_tokens(input_ids)
hidden_states *= math_sqrt(self.config.hidden_size)
# 3072**0.5 = 55.5000 in bfloat16, whilst 55.4256 in float32
# 2048**0.5 = 45.2500 in bfloat16, whilst 45.2548 in float32
inputs_embeds *= torch.tensor(math_sqrt(self.config.hidden_size), dtype = inputs_embeds.dtype)
next_decoder_cache = []
for idx, decoder_layer in enumerate(self.layers):
# Self Attention
residual = hidden_states
hidden_states = fast_rms_layernorm_inference_add_one(decoder_layer.input_layernorm, hidden_states, out_weight)
hidden_states = fast_rms_layernorm_inference_gemma(decoder_layer.input_layernorm, hidden_states, out_weight)
hidden_states, present_key_value = LlamaAttention_fast_forward_inference(
decoder_layer.self_attn,
hidden_states,
@ -171,13 +173,13 @@ def GemmaModel_fast_forward_inference(
# Fully Connected
residual = hidden_states
hidden_states = fast_rms_layernorm_inference_add_one(decoder_layer.post_attention_layernorm, hidden_states, out_weight)
hidden_states = fast_rms_layernorm_inference_gemma(decoder_layer.post_attention_layernorm, hidden_states, out_weight)
hidden_states = fast_geglu_inference(decoder_layer.mlp, hidden_states)
hidden_states += residual
next_decoder_cache.append(present_key_value)
pass
hidden_states = fast_rms_layernorm_inference_add_one(self.norm, hidden_states, out_weight)
hidden_states = fast_rms_layernorm_inference_gemma(self.norm, hidden_states, out_weight)
return BaseModelOutputWithPast(
last_hidden_state = hidden_states,

8
unsloth/models/llama.py
View file

@ -519,7 +519,11 @@ def LlamaModel_fast_forward(
elif inputs_requires_grad:
inputs_embeds.requires_grad_(False)
pass
inputs_embeds *= math_sqrt(self.config.hidden_size)
# Match Gemma exactly by casting to bfloat16 / float16
# inputs_embeds *= math_sqrt(self.config.hidden_size)
# Ie 3072**0.5 = 55.5000 in bfloat16, whilst 55.4256 in float32
# & 2048**0.5 = 45.2500 in bfloat16, whilst 45.2548 in float32
inputs_embeds *= torch.tensor(math_sqrt(self.config.hidden_size), dtype = inputs_embeds.dtype)
if inputs_requires_grad: inputs_embeds.requires_grad_(True)
pass
@ -621,7 +625,7 @@ def LlamaModel_fast_forward(
all_self_attns += (layer_outputs[1],)
pass
hidden_states = fast_rms_layernorm(self.norm, hidden_states, add_one = IS_GEMMA)
hidden_states = fast_rms_layernorm(self.norm, hidden_states, gemma = IS_GEMMA)
# add hidden states from the last decoder layer
if output_hidden_states: