src/losses.py

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,
    )