from __future__ import annotations import torch import torch.nn.functional as F PROB_EPS = 1.0e-12 def _normalize_support_logprobs( topk_logprobs: torch.Tensor, other_logprob: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: topk_probs = topk_logprobs.float().exp() other_prob = other_logprob.float().exp().unsqueeze(-1) support_probs = torch.cat([topk_probs, other_prob], dim=-1).clamp_min(PROB_EPS) support_probs = support_probs / support_probs.sum(dim=-1, keepdim=True).clamp_min(PROB_EPS) return support_probs, support_probs.log() def _masked_mean(values: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: mask = mask.float() return (values * mask).sum() / mask.sum().clamp(min=1.0) def compute_sft_ce(logits: torch.Tensor, labels: torch.Tensor, loss_mask: torch.Tensor) -> torch.Tensor: batch_size = logits.size(0) shift_labels = labels[:, 1:].contiguous() shift_loss_mask = ((loss_mask[:, 1:] > 0) & shift_labels.ne(-100)).contiguous().float() total_loss = torch.tensor(0.0, device=logits.device, dtype=torch.bfloat16) total_weight = torch.tensor(0.0, device=logits.device, dtype=torch.bfloat16) for i in range(batch_size): b_logits = logits[i, :-1, :] b_labels = shift_labels[i] b_mask = shift_loss_mask[i] ce = F.cross_entropy( b_logits, b_labels, ignore_index=-100, reduction="none", ) total_loss += (ce * b_mask).sum() total_weight += b_mask.sum() return total_loss / total_weight.clamp(min=1.0) def _compute_masked_ce_with_logits(logits, labels, loss_mask): loss_ce = compute_sft_ce(logits, labels, loss_mask) shift_logits = logits[:, :-1, :] shift_labels = labels[:, 1:].contiguous() shift_loss_mask = ((loss_mask[:, 1:] > 0) & shift_labels.ne(-100)).contiguous().float() return loss_ce, shift_logits, shift_loss_mask def compute_distillation_loss( student_logits: torch.Tensor, labels: torch.Tensor, teacher_logprobs: torch.Tensor, teacher_ids: torch.Tensor, teacher_other_logprob: torch.Tensor, loss_mask: torch.Tensor, alpha: float, temperature: float, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: vocab_size = student_logits.size(-1) loss_ce, shift_logits, shift_loss_mask = _compute_masked_ce_with_logits(student_logits, labels, loss_mask) shift_teacher_logprobs = teacher_logprobs[:, :-1, :].contiguous() shift_teacher_ids = teacher_ids[:, :-1, :].contiguous() shift_teacher_other_logprob = teacher_other_logprob[:, :-1].contiguous() shift_student = shift_logits topk_ids_clamped = shift_teacher_ids.clamp(0, vocab_size - 1) student_log_z = torch.logsumexp(shift_student / temperature, dim=-1, keepdim=True).float() student_topk_logprobs = shift_student.gather(-1, topk_ids_clamped).float() / temperature - student_log_z student_topk_probs = student_topk_logprobs.float().exp() student_other_prob = (1.0 - student_topk_probs.sum(dim=-1)).clamp_min(PROB_EPS) student_other_logprob = student_other_prob.log() teacher_support_probs, teacher_support_logprobs = _normalize_support_logprobs( shift_teacher_logprobs, shift_teacher_other_logprob, ) _, student_support_logprobs = _normalize_support_logprobs( student_topk_logprobs, student_other_logprob, ) positive_teacher = teacher_support_probs > 0 kl_terms = torch.where( positive_teacher, teacher_support_probs * (teacher_support_logprobs - student_support_logprobs), torch.zeros_like(teacher_support_probs), ) kl_per_token = kl_terms.sum(-1) loss_kd = _masked_mean(kl_per_token, shift_loss_mask) * (temperature**2) loss_total = alpha * loss_ce + (1.0 - alpha) * loss_kd return loss_total, loss_ce.detach(), loss_kd.detach() def compute_online_kd_loss( student_logits: torch.Tensor, teacher_logits: torch.Tensor, labels: torch.Tensor, loss_mask: torch.Tensor, alpha: float, temperature: float, token_chunk_size: int = 2048, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: loss_ce = compute_sft_ce(student_logits, labels, loss_mask) shift_labels = labels[:, 1:].contiguous() shift_loss_mask = ( (loss_mask[:, 1:] > 0) & shift_labels.ne(-100) ).contiguous().float() s_shifted = student_logits[:, :-1, :] t_shifted = teacher_logits[:, :-1, :] seq_len = s_shifted.size(1) total_kl = torch.tensor(0.0, device=student_logits.device, dtype=torch.float32) total_weight = torch.tensor(0.0, device=student_logits.device, dtype=torch.float32) for tok_start in range(0, seq_len, token_chunk_size): tok_end = min(tok_start + token_chunk_size, seq_len) s_chunk = s_shifted[:, tok_start:tok_end, :].float() t_chunk = t_shifted[:, tok_start:tok_end, :].float() mask_chunk = shift_loss_mask[:, tok_start:tok_end] chunk_weight = mask_chunk.sum() t_probs = F.softmax(t_chunk / temperature, dim=-1) s_log_probs = F.log_softmax(s_chunk / temperature, dim=-1) kl_tokens = F.kl_div( s_log_probs, t_probs, log_target=False, reduction="none" ).sum(dim=-1) total_kl += (kl_tokens * mask_chunk).sum() total_weight += chunk_weight del s_chunk, t_chunk, t_probs, s_log_probs, kl_tokens, mask_chunk loss_kd = (total_kl / total_weight.clamp(min=1.0)) * (temperature ** 2) loss_kd = loss_kd.to(dtype=student_logits.dtype) loss_total = alpha * loss_ce + (1.0 - alpha) * loss_kd return loss_total, loss_ce.detach(), loss_kd.detach() def compute_loss_for_phase( phase: str, logits: torch.Tensor, labels: torch.Tensor, loss_mask: torch.Tensor, batch: dict, alpha: float, temperature: float, teacher_logits: torch.Tensor | None = None, online_kd_token_chunk_size: int = 2048, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: if phase == "sft": loss_ce = compute_sft_ce(logits, labels, loss_mask) return loss_ce, loss_ce.detach(), torch.tensor(0.0, device=logits.device) if phase == "online_kd": return compute_online_kd_loss( logits, teacher_logits, labels, loss_mask, alpha, temperature, token_chunk_size=online_kd_token_chunk_size, ) return compute_distillation_loss( logits, labels, batch["teacher_logprobs"], batch["teacher_ids"], batch["teacher_other_logprob"], loss_mask, alpha, temperature, )