gemma
This commit is contained in:
Daniel Han-Chen
parent
df8034d185
commit
8259b16a5d
2 changed files with 131 additions and 107 deletions
|
|
@ -20,13 +20,13 @@ from transformers.models.llama.modeling_llama import logger
|
|||
|
||||
|
||||
@triton.jit
|
||||
def _small_cross_entropy_forward(
|
||||
def _cross_entropy_forward(
|
||||
logits_ptr, logits_row_stride,
|
||||
loss_ptr,
|
||||
lse_ptr,
|
||||
logsumexp_ptr,
|
||||
labels_ptr,
|
||||
n_cols,
|
||||
BLOCK_SIZE: tl.constexpr,
|
||||
VOCAB_SIZE : tl.constexpr,
|
||||
BLOCK_SIZE : tl.constexpr,
|
||||
):
|
||||
"""
|
||||
Cross Entropy Loss = 1/n sum [ -yi log(Pi) ]
|
||||
|
|
@ -36,83 +36,102 @@ def _small_cross_entropy_forward(
|
|||
= y * (log[sum(exp(x))] - x)
|
||||
If y == 0: CE_i = 0
|
||||
If y == 1: CE_i = logsumexp - x
|
||||
|
||||
logsumexp is also stable
|
||||
Take y = log[sum(exp(x))]
|
||||
exp(y) = sum(exp(x))
|
||||
exp(y) = sum(exp(x - c)*exp(c)) Since e^(x-c)*e^c = e^x
|
||||
exp(y) = exp(c)*sum(exp(x - c))
|
||||
y = log(exp(c)*sum(exp(x - c)))
|
||||
y = c + log[sum(exp(x - c))]
|
||||
This means we can set c = max(x) to make sure
|
||||
exp(x - c) always is exp(x - max(x)).
|
||||
This ensures exp(x - max(x))'s maximum is 1 as exp(0) = 1.
|
||||
"""
|
||||
row_idx = tl.program_id(0)
|
||||
logits_ptr += row_idx * logits_row_stride
|
||||
loss_ptr += row_idx
|
||||
lse_ptr += row_idx
|
||||
labels_ptr += row_idx
|
||||
logits_ptr += row_idx * logits_row_stride.to(tl.int64)
|
||||
loss_ptr += row_idx
|
||||
logsumexp_ptr += row_idx
|
||||
labels_ptr += row_idx
|
||||
|
||||
col_offsets = tl.arange(0, BLOCK_SIZE)
|
||||
mask = col_offsets < n_cols
|
||||
mask = col_offsets < VOCAB_SIZE
|
||||
|
||||
# TODO: Fixup int32 locations to int64
|
||||
label_idx = tl.load(labels_ptr).to(tl.int32)
|
||||
logits = tl.load(logits_ptr + col_offsets, mask = mask, other = -float("inf")).to(tl.float32)
|
||||
max_logits = tl.max(logits, 0)
|
||||
# Maximum stops overflow
|
||||
lse = tl.log(tl.sum(tl.exp(logits - max_logits), 0)) + max_logits
|
||||
tl.store(lse_ptr, lse)
|
||||
c = tl.max(logits, 0)
|
||||
logsumexp = c + tl.log(tl.sum(tl.exp(logits - c), 0))
|
||||
|
||||
if label_idx != -100:
|
||||
logits_label = tl.load(logits_ptr + label_idx).to(tl.float32)
|
||||
loss = lse - logits_label
|
||||
x = tl.load(logits_ptr + label_idx).to(tl.float32)
|
||||
loss = logsumexp - x
|
||||
else:
|
||||
loss = 0.0
|
||||
tl.store(logsumexp_ptr, logsumexp)
|
||||
tl.store(loss_ptr, loss)
|
||||
pass
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _large_cross_entropy_forward(
|
||||
def _chunked_cross_entropy_forward(
|
||||
logits_ptr, logits_row_stride,
|
||||
loss_ptr,
|
||||
lse_ptr,
|
||||
logsumexp_ptr,
|
||||
labels_ptr,
|
||||
n_rows,
|
||||
n_cols,
|
||||
BLOCK_SIZE: tl.constexpr,
|
||||
VOCAB_SIZE : tl.constexpr,
|
||||
N_CHUNKS : tl.constexpr,
|
||||
BLOCK_SIZE : tl.constexpr,
|
||||
):
|
||||
"""
|
||||
Cross Entropy Loss = 1/n sum [ -yi log(Pi) ]
|
||||
Pi = exp(xi) / sum(exp(xi))
|
||||
CE_i = -y log(p) = -y log[ exp(x) / sum(exp(x)) ]
|
||||
= -y [ x - log[sum(exp(x))] ]
|
||||
= y * (log[sum(exp(x))] - x)
|
||||
256K vocab divided in 4 chunks
|
||||
|
||||
|-65536-| |-65536-| |-65536-| |-65536-|
|
||||
|-------| |-------| |-------| |-------|
|
||||
|-------| |-------| |-------| |-------|
|
||||
|
||||
If y == 0: CE_i = 0
|
||||
If y == 1: CE_i = logsumexp - x
|
||||
|
||||
Notice we can do logsumexp for each chunk and then
|
||||
logsumexp[chunk_sum(logsumexp)] == logsumexp
|
||||
|
||||
chunk_sum = log[chunk_sum(logsumexp)]
|
||||
= log[exp(logsumexp(a)) + ... + exp(logsumexp(z))]
|
||||
= log[exp(log[sum(exp(a))]) + ... + exp(log[sum(exp(z))])]
|
||||
= log[sum(exp(a)) + ... + sum(exp(z))]
|
||||
= logsumexp(x)
|
||||
|
||||
This means we can perform a logsumexp for each chunk, then do a
|
||||
final logsumexp reduction!
|
||||
|
||||
Ie do: logsumexp(chunked_logsumexp) - x
|
||||
"""
|
||||
row_idx = tl.program_id(0)
|
||||
col_idx = tl.program_id(1)
|
||||
logits_ptr += row_idx * logits_row_stride.to(tl.int64)
|
||||
loss_ptr += row_idx + col_idx*n_rows
|
||||
lse_ptr += row_idx + col_idx*n_rows
|
||||
labels_ptr += row_idx
|
||||
row_idx = tl.program_id(0)
|
||||
chunk_idx = tl.program_id(1)
|
||||
logits_ptr += row_idx * logits_row_stride.to(tl.int64)
|
||||
loss_ptr += row_idx
|
||||
logsumexp_ptr += row_idx * N_CHUNKS + chunk_idx
|
||||
labels_ptr += row_idx
|
||||
|
||||
col_offsets = col_idx*BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
|
||||
mask = col_offsets < n_cols
|
||||
col_offsets = chunk_idx*BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
|
||||
mask = col_offsets < VOCAB_SIZE
|
||||
|
||||
# Get labels and logits
|
||||
label_idx = tl.load(labels_ptr).to(tl.int64)
|
||||
label_idx = tl.load(labels_ptr).to(tl.int32)
|
||||
logits = tl.load(logits_ptr + col_offsets, mask = mask, other = -float("inf")).to(tl.float32)
|
||||
max_logits = tl.max(logits, 0)
|
||||
# Maximum stops overflow
|
||||
lse = tl.log(tl.sum(tl.exp(logits - max_logits), 0)) + max_logits
|
||||
tl.store(lse_ptr, lse)
|
||||
c = tl.max(logits, 0)
|
||||
logsumexp = c + tl.log(tl.sum(tl.exp(logits - c), 0))
|
||||
|
||||
loss = 0.0
|
||||
# chained boolean operators (A or B or C) are not supported; use parentheses to split the chain.
|
||||
if (label_idx != -100):
|
||||
if (label_idx >= (col_idx+0)*BLOCK_SIZE) and \
|
||||
(label_idx < min((col_idx+1)*BLOCK_SIZE, n_cols)):
|
||||
|
||||
logits_label = tl.load(logits_ptr + label_idx).to(tl.float32)
|
||||
lse = 0.0
|
||||
loss = lse - logits_label # We add the final logsumexp after a reduction
|
||||
pass
|
||||
if chunk_idx == 0:
|
||||
# logsumexp(chunked_logsumexp) - x
|
||||
# Do the -x separately
|
||||
if label_idx != -100:
|
||||
x = tl.load(logits_ptr + label_idx).to(tl.float32)
|
||||
loss = -1.0 * x
|
||||
else:
|
||||
loss = 0.0
|
||||
tl.store(loss_ptr, loss)
|
||||
pass
|
||||
|
||||
tl.store(loss_ptr, loss)
|
||||
tl.store(logsumexp_ptr, logsumexp)
|
||||
pass
|
||||
|
||||
|
||||
|
|
@ -120,10 +139,10 @@ pass
|
|||
def _cross_entropy_backward(
|
||||
logits_ptr, logits_row_stride,
|
||||
dloss_ptr, dloss_row_stride,
|
||||
lse_ptr,
|
||||
logsumexp_ptr,
|
||||
labels_ptr,
|
||||
n_cols,
|
||||
BLOCK_SIZE: tl.constexpr,
|
||||
VOCAB_SIZE : tl.constexpr,
|
||||
BLOCK_SIZE : tl.constexpr,
|
||||
):
|
||||
"""
|
||||
CE_i = -y log(P) = y * (log[sum(exp(x))] - x)
|
||||
|
|
@ -140,25 +159,28 @@ def _cross_entropy_backward(
|
|||
If y == 1 and x == label: dC/dlabel = exp[x - logsumexp] - 1
|
||||
If y == 1 and x != label: dC/dx = exp[x - logsumexp]
|
||||
"""
|
||||
row_idx = tl.program_id(0)
|
||||
col_idx = tl.program_id(1)
|
||||
row_idx = tl.program_id(0)
|
||||
block_idx = tl.program_id(1)
|
||||
|
||||
logits_ptr += row_idx * logits_row_stride.to(tl.int64)
|
||||
dloss_ptr += row_idx * dloss_row_stride
|
||||
col_offsets = col_idx*BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
|
||||
mask = col_offsets < n_cols
|
||||
col_offsets = block_idx*BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
|
||||
mask = col_offsets < VOCAB_SIZE
|
||||
label_idx = tl.load(labels_ptr + row_idx).to(tl.int32)
|
||||
|
||||
if label_idx != -100:
|
||||
dloss = tl.load(dloss_ptr)
|
||||
else:
|
||||
dloss = 0.0
|
||||
logits = tl.load(logits_ptr + col_offsets, mask = mask, other = -float("inf")).to(tl.float32)
|
||||
lse = tl.load(lse_ptr + row_idx)
|
||||
probs = tl.exp(logits - lse)
|
||||
x = tl.load(logits_ptr + col_offsets, mask = mask, other = -float("inf")).to(tl.float32)
|
||||
logsumexp = tl.load(logsumexp_ptr + row_idx)
|
||||
y = tl.exp(x - logsumexp)
|
||||
y = tl.where(
|
||||
col_offsets == label_idx,
|
||||
y - 1.0, # exp(x - logsumexp) - 1
|
||||
y, # exp(x - logsumexp)
|
||||
)
|
||||
|
||||
probs = tl.where(col_offsets == label_idx, probs - 1.0, probs)
|
||||
tl.store(logits_ptr + col_offsets, dloss * probs, mask = mask)
|
||||
pass
|
||||
# If y == 0: dC/dx = 0 ==> we already masked it to be = 0, so dloss = 0.
|
||||
dloss = tl.load(dloss_ptr) if label_idx != -100 else 0.0
|
||||
tl.store(logits_ptr + col_offsets, dloss * y, mask = mask)
|
||||
pass
|
||||
|
||||
|
||||
MAX_FUSED_SIZE = 65536 # 2**16
|
||||
|
|
@ -166,45 +188,45 @@ MAX_FUSED_SIZE = 65536 # 2**16
|
|||
class Fast_CrossEntropyLoss(torch.autograd.Function):
|
||||
@staticmethod
|
||||
def forward(ctx, logits, labels):
|
||||
n_rows, n_cols = logits.shape
|
||||
n_rows, vocab_size = logits.shape
|
||||
|
||||
div, mod = divmod(n_cols, MAX_FUSED_SIZE)
|
||||
n_splits = div + (mod != 0)
|
||||
div, mod = divmod(vocab_size, MAX_FUSED_SIZE)
|
||||
n_chunks = div + (mod != 0)
|
||||
losses = torch.empty(n_rows, dtype = torch.float32, device = "cuda")
|
||||
|
||||
if n_splits == 1:
|
||||
if n_chunks == 1:
|
||||
# For small vocabs <= 65336 like Llama, Mistral
|
||||
BLOCK_SIZE, num_warps = calculate_settings(n_cols)
|
||||
losses = torch.empty(n_rows, dtype = torch.float32, device = "cuda")
|
||||
BLOCK_SIZE, num_warps = calculate_settings(vocab_size)
|
||||
logsumexp = torch.empty(n_rows, dtype = torch.float32, device = "cuda")
|
||||
|
||||
_small_cross_entropy_forward[(n_rows,)](
|
||||
_cross_entropy_forward[(n_rows,)](
|
||||
logits, logits.stride(0),
|
||||
losses,
|
||||
logsumexp,
|
||||
labels,
|
||||
n_cols,
|
||||
VOCAB_SIZE = vocab_size,
|
||||
BLOCK_SIZE = BLOCK_SIZE,
|
||||
num_warps = num_warps,
|
||||
)
|
||||
else:
|
||||
# For large vocabs > 65336 like Gemma 256K
|
||||
losses = torch.empty((n_splits, n_rows), dtype = torch.float32, device = "cuda")
|
||||
logsumexp = torch.empty((n_splits, n_rows), dtype = torch.float32, device = "cuda")
|
||||
logsumexp = torch.empty((n_rows, n_chunks), dtype = torch.float32, device = "cuda")
|
||||
|
||||
_large_cross_entropy_forward[(n_rows, n_splits,)](
|
||||
_chunked_cross_entropy_forward[(n_rows, n_chunks,)](
|
||||
logits, logits.stride(0),
|
||||
losses,
|
||||
logsumexp,
|
||||
labels,
|
||||
n_rows,
|
||||
n_cols,
|
||||
VOCAB_SIZE = vocab_size,
|
||||
BLOCK_SIZE = MAX_FUSED_SIZE,
|
||||
num_warps = 32,
|
||||
)
|
||||
logsumexp = torch.logsumexp(logsumexp, dim = 0) # Column sum
|
||||
losses = losses.sum(dim = 0) # Column sum
|
||||
losses += logsumexp # loss = lse - logits_label
|
||||
losses.masked_fill_(labels == -100, 0) # Padding tokens
|
||||
# logsumexp(chunked_logsumexp) - x
|
||||
# Do the -x separately
|
||||
logsumexp = torch.logsumexp(logsumexp, dim = 1) # Row sum
|
||||
losses += logsumexp
|
||||
losses.masked_fill_(labels == -100, 0) # Don't forget to mask padding out!
|
||||
pass
|
||||
|
||||
ctx.save_for_backward(logits, logsumexp, labels)
|
||||
|
|
@ -214,16 +236,19 @@ class Fast_CrossEntropyLoss(torch.autograd.Function):
|
|||
@staticmethod
|
||||
def backward(ctx, dlosses):
|
||||
logits, logsumexp, labels = ctx.saved_tensors
|
||||
n_rows, n_cols = logits.shape
|
||||
grid = lambda meta: (n_rows, triton.cdiv(n_cols, meta["BLOCK_SIZE"]))
|
||||
n_rows, vocab_size = logits.shape
|
||||
|
||||
_cross_entropy_backward[grid](
|
||||
BLOCK_SIZE = 4096
|
||||
div, mod = divmod(vocab_size, BLOCK_SIZE)
|
||||
n_blocks = div + (mod != 0)
|
||||
|
||||
_cross_entropy_backward[(n_rows, n_blocks,)](
|
||||
logits, logits.stride(0),
|
||||
dlosses, dlosses.stride(0),
|
||||
logsumexp,
|
||||
labels,
|
||||
n_cols,
|
||||
BLOCK_SIZE = 4096,
|
||||
VOCAB_SIZE = n_cols,
|
||||
BLOCK_SIZE = BLOCK_SIZE,
|
||||
num_warps = 8,
|
||||
)
|
||||
return logits, None, None,
|
||||
|
|
|
|||
|
|
@ -203,15 +203,14 @@ pass
|
|||
# https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L590
|
||||
def GemmaDecoderLayer_fast_forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
causal_mask: Optional[xformers.attn_bias.BlockDiagonalCausalMask] = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_value: Optional[Tuple[torch.Tensor]] = None,
|
||||
output_attentions: Optional[bool] = False,
|
||||
use_cache: Optional[bool] = False,
|
||||
padding_mask: Optional[torch.LongTensor] = None,
|
||||
*args, **kwargs,
|
||||
hidden_states: torch.Tensor,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_value: Optional[Tuple[torch.Tensor]] = None,
|
||||
output_attentions: Optional[bool] = False,
|
||||
use_cache: Optional[bool] = False,
|
||||
cache_position: Optional[torch.LongTensor] = None,
|
||||
**kwargs,
|
||||
):
|
||||
if False:#past_key_value is not None:
|
||||
do_prefill = not hasattr(self.self_attn, "paged_attention")
|
||||
|
|
@ -239,13 +238,13 @@ def GemmaDecoderLayer_fast_forward(
|
|||
# hidden_states = self.input_layernorm(hidden_states)
|
||||
hidden_states, self_attn_weights, present_key_value = self.self_attn(
|
||||
hidden_states=hidden_states,
|
||||
causal_mask=causal_mask,
|
||||
attention_mask=attention_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_value=past_key_value,
|
||||
output_attentions=output_attentions,
|
||||
use_cache=use_cache,
|
||||
padding_mask=padding_mask,
|
||||
cache_position=cache_position,
|
||||
**kwargs,
|
||||
)
|
||||
hidden_states = residual + hidden_states
|
||||
|
||||
|
|
@ -541,12 +540,12 @@ class FastGemmaModel(FastLlamaModel):
|
|||
|
||||
@staticmethod
|
||||
def pre_patch():
|
||||
GemmaAttention .forward = LlamaAttention_fast_forward
|
||||
GemmaSdpaAttention .forward = LlamaAttention_fast_forward
|
||||
GemmaFlashAttention2.forward = LlamaAttention_fast_forward
|
||||
GemmaAttention .forward = GemmaAttention_fast_forward
|
||||
GemmaSdpaAttention .forward = GemmaAttention_fast_forward
|
||||
GemmaFlashAttention2.forward = GemmaAttention_fast_forward
|
||||
GemmaDecoderLayer .forward = GemmaDecoderLayer_fast_forward
|
||||
GemmaModel .forward = LlamaModel_fast_forward
|
||||
GemmaForCausalLM .forward = LlamaForCausalLM_fast_forward
|
||||
GemmaModel .forward = GemmaModel_fast_forward
|
||||
GemmaForCausalLM .forward = GemmaForCausalLM_fast_forward
|
||||
PeftModelForCausalLM.forward = PeftModelForCausalLM_fast_forward
|
||||
# Solves https://github.com/unslothai/unsloth/issues/168
|
||||
# Static KV Cache was introduced in 4.38.0, causing training to be much slower.
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue