tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

v2 trial

Browse Source
This commit is contained in:
Wing Lian
2024-12-21 13:43:48 -05:00
parent c73acd7de0
commit 119d586cf4
2 changed files with 252 additions and 261 deletions
Download Patch File Download Diff File Expand all files Collapse all files

49
src/axolotl/core/trainers/kd.py
View File

@@ -7,10 +7,7 @@ from typing import Optional
import torch import torch
from axolotl.core.trainers.base import AxolotlTrainer from axolotl.core.trainers.base import AxolotlTrainer
from axolotl.integrations.kd.kernels.kd import ( from axolotl.integrations.kd.kernels.kd import kd_loss_triton
forward_kl_topk,
prepare_topk_student_teacher,
)
def kd_loss_function( def kd_loss_function(
@@ -100,43 +97,29 @@ class AxolotlKDTrainer(AxolotlTrainer):
outputs = model(**inputs) outputs = model(**inputs)
student_logits = outputs["logits"] student_logits = outputs["logits"]
# Slice or gather student logits to match teacher seq len
# e.g.:
teacher_seq_len = target_token_ids.shape[1]
student_logits_for_kd = student_logits[:, :teacher_seq_len, :] # [B, seq_len, vocab_size]
# Gather & flatten to [N, K] # GATHER top-K from student
stud_lp_f, teach_lp_f, mask_f = prepare_topk_student_teacher( student_logits_topk = torch.gather(
student_logits, student_logits_for_kd, dim=-1, index=target_token_ids # same shape [B, seq_len, K]
target_token_ids,
target_logprobs,
target_mask,
) )
loss_kd = forward_kl_topk(teach_lp_f, stud_lp_f, mask_f, reduction="none") # Now call the Triton-based KD loss
loss_kd = kd_loss_triton(
# Normalize by number of items or mean over valid tokens student_logits_topk,
if num_items_in_batch is not None: target_logprobs, # teacher logprobs [B, seq_len, K]
# If you know how many items should be considered in the batch target_mask, # mask [B, seq_len, K]
loss_kd = loss_kd / num_items_in_batch num_items_in_batch=num_items_in_batch,
else: )
# Otherwise, just average over all valid tokens
# count number of unmasked tokens in teacher_mask
kd_loss_per_token = target_mask.sum(dim=1).unsqueeze(-1)
# Normalize by number of unmasked tokens in teacher_mask
loss_kd = loss_kd / kd_loss_per_token.float()
# optionally combine with CE loss
if self.args.kd_ce_alpha > 0: if self.args.kd_ce_alpha > 0:
loss = self.args.kd_ce_alpha * outputs["loss"] + loss_kd loss = self.args.kd_ce_alpha * outputs["loss"] + loss_kd
else: else:
loss = loss_kd loss = loss_kd
# Save past state if it exists
# TODO: this needs to be fixed and made cleaner later.
if self.args.past_index >= 0:
self._past = outputs[ # pylint: disable=attribute-defined-outside-init
self.args.past_index
]
if self.args.average_tokens_across_devices and self.model_accepts_loss_kwargs:
loss *= self.accelerator.num_processes
torch.cuda.empty_cache()
return (loss, outputs) if return_outputs else loss return (loss, outputs) if return_outputs else loss

464
src/axolotl/integrations/kd/kernels/kd.py
View File

@@ -2,267 +2,275 @@ import torch
import triton import triton
import triton.language as tl import triton.language as tl
configs = [
triton.Config({"BLOCK_SIZE": 32}, num_warps=1, num_stages=1), # --------------------------------------------------------
triton.Config({"BLOCK_SIZE": 64}, num_warps=1, num_stages=1), # Triton Kernel for forward pass
triton.Config({"BLOCK_SIZE": 128}, num_warps=2, num_stages=2), # --------------------------------------------------------
# Add more if needed # We'll assume:
] # - B * seq_len threads in 1D dimension
# - Each thread handles K tokens (the top-K from teacher).
# - For large K, you might want a more 2D approach to keep good occupancy.
#
# Pseudocode steps inside kernel:
# 1) compute index for [batch, seq_position]
# 2) read top-K token IDs from teacher_token_ids
# 3) gather student_logits_topk
# 4) compute logsumexp for those K logits
# 5) compute student_logprobs_topk
# 6) read teacher_logprobs
# 7) compute teacher_probs = exp(teacher_logprobs)
# 8) compute partial KL = sum(teacher_probs * (teacher_logprobs - student_logprobs_topk))
# 9) store partial KL in a buffer
#
# Later, we'll do a reduction on partial KL across all threads.
#
# NOTE: This is a reference skeleton. You must adapt indexing carefully.
#
@triton.autotune(configs=configs, key=["N", "K"])
@triton.jit @triton.jit
def fwd_kl_topk_kernel( def kd_forward_kernel(
teacher_lp_ptr, # float32 [N, K] student_logits_ptr, # float32[B, seq_len, K] after gather
student_lp_ptr, # float32 [N, K] teacher_logprobs_ptr, # float32[B, seq_len, K]
mask_ptr, # bool [N, K] mask_ptr, # bool[B, seq_len, K] or int8
loss_out_ptr, # float32 [N] partial_kd_ptr, # float32[B, seq_len] (accumulator)
stride_tn, B, # total batch size
stride_tk, seq_len, # total sequence length from teacher
stride_sn, K, # top-K from teacher
stride_sk, BLOCK_SIZE: tl.constexpr # how many tokens per block in dimension0
stride_mn,
stride_mk,
stride_loss_n,
N: tl.constexpr,
K: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
): ):
""" # program_id is the global index for each block
Each kernel instance: row_id = tl.program_id(0). We'll tile the K dimension in chunks of BLOCK_SIZE. pid = tl.program_id(0)
Summation => store into loss_out[row_id].
"""
row_id = tl.program_id(0)
if row_id >= N:
return
# Base pointers for teacher, student, mask rows # Each block handles a range of seq positions in [0..B*seq_len)
t_row_ptr = teacher_lp_ptr + row_id * stride_tn block_start = pid * BLOCK_SIZE
s_row_ptr = student_lp_ptr + row_id * stride_sn block_end = tl.min((pid+1)*BLOCK_SIZE, B * seq_len)
m_row_ptr = mask_ptr + row_id * stride_mn length = block_end - block_start
# We'll accumulate KL in local variable # Offsets for indexing
kl_sum = 0.0 # We want to interpret a linear index in [0..B*seq_len) as (batch_idx, seq_idx)
# E.g.:
# batch_idx = block_start // seq_len
# seq_idx = block_start % seq_len
# but we must do this for each element in the block. We'll do that inside a loop.
# tile the K dimension # We'll store a running partial KL sum in registers
num_tiles = (K + BLOCK_SIZE - 1) // BLOCK_SIZE # We do a for-loop for each position in the block, then do a thread-level reduction
for tile_id in range(num_tiles): kd_reg = 0.0
k_offset = tile_id * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
mask = k_offset < K
# load teacher logprobs # We'll iterate over each item in [block_start, block_end).
t_lp = tl.load(t_row_ptr + k_offset * stride_tk, mask=mask, other=-float("inf")) # A more advanced approach can use vectorization / warp-based parallelism inside the block.
# load student logprobs for offset in range(length):
s_lp = tl.load(s_row_ptr + k_offset * stride_sk, mask=mask, other=-float("inf")) # Convert offset -> actual index in [0..B*seq_len)
linear_idx = block_start + offset
# batch index and sequence index
b_idx = linear_idx // seq_len
s_idx = linear_idx % seq_len
# load mask => bool (0 or 1) # For K top tokens, read the relevant student logits and teacher logprobs
valid = tl.load(m_row_ptr + k_offset * stride_mk, mask=mask, other=0) # We'll load them in a small loop:
valid_f32 = valid.to(tl.float32) logsumexp_val = float('-inf')
# We'll store them in a local array for a second pass
student_logits_k = [0.0 for _ in range(K)]
teacher_logprobs_k = [0.0 for _ in range(K)]
valid_k = [0 for _ in range(K)]
# teacher probs # gather the top-K logits & teacher logprobs
t_p = tl.exp(t_lp) for k in range(K):
# load student logit
student_val = tl.load(
student_logits_ptr
+ b_idx*seq_len*K
+ s_idx*K
+ k,
mask=(b_idx < B) and (s_idx < seq_len)
)
teacher_val = tl.load(
teacher_logprobs_ptr
+ b_idx*seq_len*K
+ s_idx*K
+ k,
mask=(b_idx < B) and (s_idx < seq_len)
)
# get mask
mask_val = tl.load(
mask_ptr
+ b_idx*seq_len*K
+ s_idx*K
+ k,
mask=(b_idx < B) and (s_idx < seq_len)
)
# local_kl = p^T * (lp^T - lp^S) student_logits_k[k] = student_val
local_kl = t_p * (t_lp - s_lp) teacher_logprobs_k[k] = teacher_val
# multiply by valid_f32 to ignore padded or invalid positions valid_k[k] = mask_val
local_kl *= valid_f32
# sum over the chunk # track max for logsumexp (naive approach)
kl_sum += tl.sum(local_kl, where=mask) if student_val > logsumexp_val:
logsumexp_val = student_val
# store rowwise result # now compute logsumexp for the K student logits
tl.store(loss_out_ptr + row_id * stride_loss_n, kl_sum) # logsumexp = max_val + log(sum( exp(student_val - max_val) ))
exp_sum = 0.0
for k in range(K):
if valid_k[k] != 0: # if valid
exp_sum += float(torch.exp(student_logits_k[k] - logsumexp_val))
# safe check
if exp_sum == 0.0:
exp_sum = 1e-8
logsumexp_val = logsumexp_val + float(torch.log(torch.tensor(exp_sum)))
# compute partial kl
# sum_{k in valid} p^T_k (log p^T_k - log p^S_k)
# teacher_probs_k = exp(teacher_logprobs_k)
for k in range(K):
if valid_k[k] != 0: # only valid tokens
teacher_prob = float(torch.exp(teacher_logprobs_k[k]))
student_logprob = student_logits_k[k] - logsumexp_val
kd_val = teacher_prob * (teacher_logprobs_k[k] - student_logprob)
kd_reg += kd_val
# Write out partial kd for this block. We store a single partial sum in partial_kd_ptr
# We'll store it at partial_kd_ptr[pid]
# In real code, you might do an atomic add into partial_kd_ptr or a parallel reduction pass
# for now, let's just store it at index=pid
tl.store(partial_kd_ptr + pid, kd_reg)
@triton.autotune(configs=configs, key=["N", "K"]) class _KLDivergenceTritonFn(torch.autograd.Function):
@triton.jit
def bwd_kl_topk_kernel(
teacher_lp_ptr, # float32 [N, K]
mask_ptr, # bool [N, K]
grad_stud_ptr, # float32 [N, K], output
stride_tn,
stride_tk,
stride_mn,
stride_mk,
stride_gn,
stride_gk,
N: tl.constexpr,
K: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
):
"""
For forward KL, d/d(student_lp) = - exp(teacher_lp), if mask=1, else 0.
Each kernel instance processes one row [K].
"""
row_id = tl.program_id(0)
if row_id >= N:
return
t_row_ptr = teacher_lp_ptr + row_id * stride_tn
m_row_ptr = mask_ptr + row_id * stride_mn
g_row_ptr = grad_stud_ptr + row_id * stride_gn
num_tiles = (K + BLOCK_SIZE - 1) // BLOCK_SIZE
for tile_id in range(num_tiles):
k_offset = tile_id * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
mask = k_offset < K
t_lp = tl.load(t_row_ptr + k_offset * stride_tk, mask=mask, other=-float("inf"))
valid = tl.load(m_row_ptr + k_offset * stride_mk, mask=mask, other=0).to(
tl.int1
)
grad_val = -tl.exp(t_lp) # derivative
grad_val = tl.where(valid, grad_val, 0.0)
tl.store(g_row_ptr + k_offset * stride_gk, grad_val, mask=mask)
class FwdKLTopKFunction(torch.autograd.Function):
@staticmethod @staticmethod
def forward( def forward(ctx, student_logits, teacher_logprobs, mask, num_items_in_batch):
ctx,
teacher_lp_topk: torch.Tensor,
student_lp_topk: torch.Tensor,
mask_topk: torch.Tensor,
reduction: str = "batchmean",
) -> torch.Tensor:
""" """
teacher_lp_topk: [N, K] Inputs shape assumptions (after gather!):
student_lp_topk: [N, K] - student_logits: [B, seq_len, K]
mask_topk: [N, K] bool - teacher_logprobs: [B, seq_len, K]
returns either scalar (if batchmean) or [N] if 'none' - mask: [B, seq_len, K] (bool or 0/1) for valid tokens
""" """
assert teacher_lp_topk.shape == student_lp_topk.shape B, seq_len, K = student_logits.shape
assert teacher_lp_topk.shape == mask_topk.shape
N, K = teacher_lp_topk.shape # Prepare output buffer for partial sums
dev = teacher_lp_topk.device # We'll have BLOCK_SIZE define how many (batch*seq_len) items each block processes
dtype = teacher_lp_topk.dtype # For simplicity, let's aim for one block per 1024 positions
BLOCK_SIZE = 1024
# compute how many blocks we need
total_positions = B * seq_len
grid = ( (total_positions + BLOCK_SIZE - 1) // BLOCK_SIZE , )
# Contiguous partial_kd = torch.empty(
t_lp_c = teacher_lp_topk.contiguous() grid[0], dtype=student_logits.dtype, device=student_logits.device
s_lp_c = student_lp_topk.contiguous()
m_c = mask_topk.contiguous()
# [N] rowwise sums
loss_out = torch.empty(N, dtype=torch.float32, device=dev)
grid = (N,)
fwd_kl_topk_kernel[grid](
t_lp_c,
s_lp_c,
m_c,
loss_out,
# strides
t_lp_c.stride(0),
t_lp_c.stride(1),
s_lp_c.stride(0),
s_lp_c.stride(1),
m_c.stride(0),
m_c.stride(1),
loss_out.stride(0),
N=N,
K=K
# BLOCK_SIZE, warps, stages => autotune
) )
if reduction == "batchmean": # Launch kernel
loss_val = loss_out.mean() kd_forward_kernel[grid](
elif reduction == "none": student_logits,
loss_val = loss_out teacher_logprobs,
mask,
partial_kd,
B, seq_len, K,
BLOCK_SIZE=BLOCK_SIZE
)
# Sum partials on CPU or GPU
kd_sum = partial_kd.sum()
# normalize
if num_items_in_batch is not None:
kd_loss = kd_sum / num_items_in_batch
else: else:
raise ValueError("reduction must be 'batchmean' or 'none'") # Just average over all valid tokens; in practice you'd need the count of valid tokens
# For a quick approximation, let's do kd_sum / total_positions (or do a separate reduction on mask)
# This is a simplification. For correctness, you should count valid tokens in the kernel.
kd_loss = kd_sum / (total_positions * K)
# Save for backward # Save context for backward
ctx.save_for_backward(t_lp_c, m_c) # Typically, you'd need to save the raw student_logits, teacher_logprobs, etc. for grad
ctx.reduction = reduction # But be mindful of memory usage. We’ll demonstrate the minimal approach here:
ctx.shape = (N, K) ctx.save_for_backward(student_logits, teacher_logprobs, mask, torch.tensor(num_items_in_batch or 0))
ctx.B = B
ctx.seq_len = seq_len
ctx.K = K
ctx.total_positions = total_positions
ctx.BLOCK_SIZE = BLOCK_SIZE
return loss_val return kd_loss
@staticmethod @staticmethod
def backward(ctx, grad_output): def backward(ctx, grad_output):
# grad_output is either scalar ([1]) if batchmean, or shape=[N] if 'none' """
t_lp_c, m_c = ctx.saved_tensors grad_output is dLoss/dOut (a scalar).
(N, K) = ctx.shape We want dLoss/dStudentLogits.
Recall that:
# We'll create a grad for the student's top-K logprobs Loss = sum_{valid k} p^T_k ( log p^T_k - (student_logits_k - logsumexp(student_logits_all_k)) )
grad_stud = torch.empty_like(t_lp_c) # [N, K] = sum_{valid k} p^T_k log p^T_k - sum_{valid k} p^T_k student_logits_k + sum_{valid k} p^T_k logsumexp(...)
grid = (N,) Let’s break down the derivative wrt student_logits_k. More precisely, from:
bwd_kl_topk_kernel[grid]( d/d student_logits_k [ - p^T_k student_logprobs_k ]
t_lp_c, you get:
m_c, - p^T_k * ( d/d student_logits_k [ student_logits_k - logsumexp(...) ] )
grad_stud, = - p^T_k * (1 - p^S_k)
t_lp_c.stride(0), = p^T_k * p^S_k - p^T_k
t_lp_c.stride(1), = p^S_k * p^T_k - p^T_k
m_c.stride(0), = p^T_k( p^S_k - 1 )
m_c.stride(1),
grad_stud.stride(0),
grad_stud.stride(1),
N=N,
K=K,
)
# Multiply by grad_output In practice, we also must handle the mask.
# If batchmean => scalar A real implementation typically re-runs the gather & logsumexp calculations or caches them in forward().
# If none => shape=[N] For brevity, we do a naive approach in PyTorch (not Triton) for the backward.
if grad_output.numel() == 1: For maximum speed, you'd do a second Triton kernel.
grad_stud *= grad_output
else:
# shape=[N], broadcast over K
grad_stud *= grad_output.unsqueeze(1)
return grad_stud, None, None, None We'll do a minimal approach here: recompute everything on the host side or a pure PyTorch pass.
"""
student_logits, teacher_logprobs, mask, num_items_in_batch_t = ctx.saved_tensors
num_items_in_batch = int(num_items_in_batch_t.item())
B, seq_len, K = ctx.B, ctx.seq_len, ctx.K
# We can either replicate the entire forward logic in PyTorch for gradient
# or do a second Triton pass. Here, let's do it in PyTorch for clarity.
# 1) compute logsumexp of student_logits_k for each [b, s]
# 2) compute p^S_k
# 3) compute p^T_k from teacher_logprobs
# 4) dLoss/dStudentLogits = grad_output * p^T_k ( p^S_k - 1 ), masked
# 5) sum or gather the final gradient
with torch.enable_grad():
# treat student_logits as if it requires grad
stl = student_logits.clone().detach().requires_grad_(True)
# compute logsumexp along K
logsumexp_val = torch.logsumexp(stl, dim=-1, keepdim=True) # [B, seq_len, 1]
student_logprobs_topk = stl - logsumexp_val
teacher_probs = teacher_logprobs.exp()
# p^S_k
p_s = student_logprobs_topk.exp()
# forward kl = sum p^T_k ( teacher_logprobs_k - student_logprobs_topk )
# derivative wrt stl = p^T_k( p^S_k - 1 )
grad_stl = teacher_probs * (p_s - 1.0)
# respect the mask
grad_stl = grad_stl * mask # zero out invalid
# sum or average
if num_items_in_batch != 0:
grad_stl = grad_stl / num_items_in_batch
else:
grad_stl = grad_stl / (B * seq_len * K) # fallback
# multiply by upstream grad_output
grad_stl = grad_stl * grad_output
return grad_stl, None, None, None
def forward_kl_topk( def kd_loss_triton(
teacher_lp_topk: torch.Tensor, student_logits, # [B, teacher_seq_len, vocab_size], but typically we gather for top-K
student_lp_topk: torch.Tensor, teacher_logprobs,
mask_topk: torch.Tensor, mask,
reduction: str = "batchmean", num_items_in_batch=None
) -> torch.Tensor: ):
""" """
Calls the autograd function that launches Triton kernels for forward + backward. Wrapper that calls our Triton-based forward+backward for KD.
For production, you likely want to do the gather (teacher top-K) outside
or inside a separate kernel. This function expects that you've *already*
called gather on student_logits -> shape [B, seq_len, K].
""" """
return FwdKLTopKFunction.apply( return _KLDivergenceTritonFn.apply(student_logits, teacher_logprobs, mask, num_items_in_batch)
teacher_lp_topk, student_lp_topk, mask_topk, reduction
)
def prepare_topk_student_teacher(
student_logits: torch.Tensor, # [B, teacher_seq_len, vocab_size]
target_token_ids: torch.Tensor, # [B, teacher_seq_len, K]
target_logprobs: torch.Tensor, # [B, teacher_seq_len, K], teacher logprobs
target_mask: torch.Tensor, # [B, teacher_seq_len, K], bool or 0/1
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""
Gathers student logits for the teacher's top-K tokens and flattens the first 2 dims => N = B * teacher_seq_len.
Returns:
(student_lp_topk, teacher_lp_topk, valid_mask) each shape = [N, K].
"""
B, S, K = target_token_ids.shape
# Gather the student logits => [B, S, K]
# 1) slice or use the entire student_logits if it matches teacher_seq_len
student_logits_for_kd = student_logits[:, :S, :] # ensure alignment if needed
# 2) gather top-k => [B, S, K]
student_logits_topk = torch.gather(
student_logits_for_kd, dim=-1, index=target_token_ids
)
# 3) convert student logits to logprobs => [B, S, K]
student_logprobs_topk = student_logits_topk - torch.logsumexp(
student_logits_topk, dim=-1, keepdim=True
)
# Flatten batch dimension
N = B * S
student_logprobs_topk_f = student_logprobs_topk.view(N, K) # [N, K]
teacher_logprobs_topk_f = target_logprobs.view(N, K) # [N, K]
mask_f = target_mask.view(N, K).bool() # [N, K]
return student_logprobs_topk_f, teacher_logprobs_topk_f, mask_f