Update cross_entropy_loss.py
This commit is contained in:
Daniel Han-Chen
parent
c1dbf6710e
commit
48736a0755
1 changed files with 63 additions and 49 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,29 +36,38 @@ 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
|
||||
|
||||
|
|
@ -120,10 +129,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,24 +149,27 @@ 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)
|
||||
# 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
|
||||
|
||||
|
||||
|
|
@ -166,38 +178,38 @@ 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)
|
||||
n_chunks = div + (mod != 0)
|
||||
|
||||
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)
|
||||
losses = torch.empty(n_rows, dtype = torch.float32, device = "cuda")
|
||||
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")
|
||||
losses = torch.empty((n_chunks, n_rows), dtype = torch.float32, device = "cuda")
|
||||
logsumexp = torch.empty((n_chunks, n_rows), dtype = torch.float32, device = "cuda")
|
||||
|
||||
_large_cross_entropy_forward[(n_rows, n_splits,)](
|
||||
_large_cross_entropy_forward[(n_rows, n_chunks,)](
|
||||
logits, logits.stride(0),
|
||||
losses,
|
||||
logsumexp,
|
||||
labels,
|
||||
n_rows,
|
||||
n_cols,
|
||||
vocab_size,
|
||||
BLOCK_SIZE = MAX_FUSED_SIZE,
|
||||
num_warps = 32,
|
||||
)
|
||||
|
|
@ -214,17 +226,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
|
||||
|
||||
print(logits.stride(), dlosses.stride(), dlosses.shape, dlosses)
|
||||
_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 = vocab_size,
|
||||
BLOCK_SIZE = BLOCK_SIZE,
|
||||
num_warps = 8,
|
||||
)
|
||||
return logits, None, None,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue