composer_replication/distillation/tests/test_distillation_losses.py

"""Distillation-loss unit tests — SimPO + TAID + Entropy-Aware OPD."""
from __future__ import annotations

import math

import pytest
import torch
import torch.nn.functional as F

from composer_replication.distillation import (
    entropy_aware_opd_loss,
    simpo_loss,
    taid_loss,
)
from composer_replication.distillation.simpo import avg_sequence_logprob
from composer_replication.distillation.taid import TAIDScheduler
from composer_replication.distillation.entropy_aware_opd import teacher_entropy


# ---------------------------------------------------------------------
# SimPO
# ---------------------------------------------------------------------

def test_simpo_loss_returns_scalar():
    chosen = torch.tensor([0.5, 0.4, 0.3])
    rejected = torch.tensor([0.1, 0.0, -0.2])
    loss = simpo_loss(chosen, rejected, beta=2.0, gamma=1.0)
    assert loss.dim() == 0
    assert torch.isfinite(loss)


def test_simpo_loss_lower_for_better_separation():
    """Larger margin between chosen and rejected → lower loss."""
    # Same setup, two batches with different separations
    small_sep_loss = simpo_loss(
        torch.tensor([0.1]), torch.tensor([0.05]),
    )
    large_sep_loss = simpo_loss(
        torch.tensor([1.0]), torch.tensor([-1.0]),
    )
    assert large_sep_loss < small_sep_loss, (
        f"large separation should give smaller loss; "
        f"got small_sep={small_sep_loss}, large_sep={large_sep_loss}"
    )


def test_simpo_loss_differentiable():
    chosen = torch.tensor([0.5], requires_grad=True)
    rejected = torch.tensor([0.0], requires_grad=True)
    loss = simpo_loss(chosen, rejected)
    loss.backward()
    assert chosen.grad is not None
    assert rejected.grad is not None
    assert torch.isfinite(chosen.grad).all()
    assert torch.isfinite(rejected.grad).all()


def test_simpo_loss_shape_mismatch_raises():
    with pytest.raises(ValueError, match="same shape"):
        simpo_loss(torch.zeros(3), torch.zeros(5))


def test_avg_sequence_logprob():
    """Helper averages over response tokens, ignoring prompt + padding."""
    # B=2, T=4
    logprobs = torch.tensor([
        [-10.0, -10.0, -1.0, -2.0],   # response is last 2 tokens, avg=-1.5
        [-1.0, -3.0, -1.0, -10.0],    # response is first 3 tokens, avg=-5/3
    ])
    mask = torch.tensor([
        [0, 0, 1, 1],
        [1, 1, 1, 0],
    ])
    avg = avg_sequence_logprob(logprobs, mask)
    expected = torch.tensor([-1.5, -5.0 / 3.0])
    torch.testing.assert_close(avg, expected, atol=1e-5, rtol=1e-5)


# ---------------------------------------------------------------------
# TAID
# ---------------------------------------------------------------------

def test_taid_loss_returns_scalar_and_differentiable():
    """Basic shape + grad check at t=0.5."""
    B, T, V = 2, 4, 8
    student_logits = torch.randn(B, T, V, requires_grad=True)
    teacher_logits = torch.randn(B, T, V)
    mask = torch.ones(B, T)
    loss = taid_loss(student_logits, teacher_logits, mask, t=0.5)
    assert loss.dim() == 0
    assert torch.isfinite(loss)
    loss.backward()
    assert student_logits.grad is not None
    assert torch.isfinite(student_logits.grad).all()


def test_taid_loss_t_zero_target_matches_detached_student():
    """At t=0, p_t = softmax(student.detach()), so the forward-KL target is
    the detached student. The loss is then the entropy of that detached
    distribution against itself — finite, but more importantly the gradient
    flowing into student_logits comes only through the log_softmax term, not
    through the target (because of the .detach()).
    """
    B, T, V = 1, 2, 4
    s1 = torch.randn(B, T, V, requires_grad=True)
    teacher_a = torch.zeros(B, T, V); teacher_a[..., 0] = 10.0
    teacher_b = torch.zeros(B, T, V); teacher_b[..., 3] = 10.0
    mask = torch.ones(B, T)
    # At t=0 the teacher is completely ignored — same student detach anchor.
    loss_a = taid_loss(s1, teacher_a, mask, t=0.0)
    loss_b = taid_loss(s1, teacher_b, mask, t=0.0)
    assert abs(float(loss_a) - float(loss_b)) < 1e-6


def test_taid_loss_t_one_is_pure_forward_kl():
    """At t=1, target = softmax(teacher_logits), so taid_loss reduces to
    upstream forward_kl on the masked tokens.
    """
    B, T, V = 2, 3, 5
    student = torch.randn(B, T, V, requires_grad=True)
    teacher = torch.randn(B, T, V)
    mask = torch.ones(B, T)

    loss_taid = taid_loss(student, teacher, mask, t=1.0)

    # Reference forward-KL: -mean_token sum_v p_teacher(v) * log_q(v)
    p_teacher = F.softmax(teacher, dim=-1, dtype=torch.float32)
    log_q = F.log_softmax(student, dim=-1, dtype=torch.float32)
    per_token = -(p_teacher * log_q).sum(dim=-1)
    ref = per_token.mean()

    torch.testing.assert_close(loss_taid, ref, atol=1e-5, rtol=1e-5)


def test_taid_loss_mask_is_token_mean():
    """Mask zeros out tokens; loss = sum(per_token * mask) / sum(mask)."""
    B, T, V = 1, 4, 6
    s = torch.randn(B, T, V)
    t_logits = torch.randn(B, T, V)
    full_mask = torch.ones(B, T)
    half_mask = torch.tensor([[1.0, 1.0, 0.0, 0.0]])

    loss_full = taid_loss(s, t_logits, full_mask, t=0.7)
    loss_half = taid_loss(s, t_logits, half_mask, t=0.7)

    # Manually: token-mean over only the first 2 positions
    blended = (1 - 0.7) * s.detach() + 0.7 * t_logits
    p_t = F.softmax(blended, dim=-1, dtype=torch.float32)
    log_q = F.log_softmax(s, dim=-1, dtype=torch.float32)
    per_token = -(p_t * log_q).sum(dim=-1)
    expected_half = per_token[:, :2].mean()
    torch.testing.assert_close(loss_half, expected_half, atol=1e-5, rtol=1e-5)
    # Sanity: full vs half differ when teacher has structure.
    assert not torch.allclose(loss_full, loss_half)


def test_taid_loss_shape_mismatch_raises():
    s = torch.randn(2, 3, 5)
    t_logits = torch.randn(2, 3, 6)
    with pytest.raises(ValueError, match="shape mismatch"):
        taid_loss(s, t_logits, t=0.5)


def test_taid_loss_invalid_mask_raises():
    s = torch.randn(2, 3, 5)
    t_logits = torch.randn(2, 3, 5)
    bogus_mask = torch.ones(2, 4)  # wrong T
    with pytest.raises(ValueError, match="mask shape"):
        taid_loss(s, t_logits, bogus_mask, t=0.5)


# ---------------------------------------------------------------------
# TAIDScheduler
# ---------------------------------------------------------------------

def test_taid_scheduler_initial_state():
    sched = TAIDScheduler(num_train_steps=1000, t_start=0.4)
    assert sched.t == pytest.approx(0.4)


def test_taid_scheduler_first_update_seeds():
    """First update_t() with finite loss only sets prev_loss, returns None,
    leaves t at t_start.
    """
    sched = TAIDScheduler(num_train_steps=100, t_start=0.4)
    delta = sched.update_t(torch.tensor(2.0), global_step=0)
    assert delta is None
    assert sched.t == pytest.approx(0.4)


def test_taid_scheduler_monotonic_non_decreasing():
    """Even with noisy/oscillating loss, t is non-decreasing."""
    sched = TAIDScheduler(num_train_steps=1000, t_start=0.4)
    losses = [3.0, 2.5, 2.7, 2.3, 2.4, 2.0, 1.8, 1.85, 1.7, 1.5]
    prev_t = sched.t
    for step, loss in enumerate(losses):
        sched.update_t(torch.tensor(loss), global_step=step)
        assert sched.t >= prev_t - 1e-6, (
            f"t decreased at step {step}: {prev_t} -> {sched.t}"
        )
        prev_t = sched.t


def test_taid_scheduler_t_end_clamp():
    """t never exceeds t_end."""
    sched = TAIDScheduler(num_train_steps=10, t_start=0.4, t_end=0.9)
    # Push global_step past num_train_steps so the linear floor would exceed t_end.
    for step in range(0, 100):
        sched.update_t(torch.tensor(2.0 - 0.01 * step), global_step=step)
    assert sched.t <= 0.9 + 1e-6


def test_taid_scheduler_disable_adaptive_is_linear():
    """With disable_adaptive=True, t = t_start + progress * (t_end - t_start)."""
    sched = TAIDScheduler(
        num_train_steps=100, t_start=0.0, t_end=1.0, disable_adaptive=True
    )
    # Seed prev_loss
    sched.update_t(torch.tensor(2.0), global_step=0)
    sched.update_t(torch.tensor(1.5), global_step=50)
    assert sched.t == pytest.approx(0.5, abs=1e-6)
    sched.update_t(torch.tensor(1.0), global_step=100)
    assert sched.t == pytest.approx(1.0, abs=1e-6)


# ---------------------------------------------------------------------
# Entropy-Aware OPD
# ---------------------------------------------------------------------

def test_teacher_entropy_one_hot_is_zero():
    """Argmax-1 distribution has entropy 0."""
    logits = torch.zeros(1, 1, 4)
    logits[..., 0] = 100.0  # essentially one-hot
    H = teacher_entropy(logits)
    assert float(H[0, 0]) < 1e-3


def test_teacher_entropy_uniform_is_log_v():
    """Uniform distribution over V symbols has entropy = log(V)."""
    logits = torch.zeros(1, 1, 5)
    H = teacher_entropy(logits)
    assert float(H[0, 0]) == pytest.approx(math.log(5), rel=1e-5)


def test_entropy_aware_opd_returns_scalar_and_differentiable():
    B, T, V = 2, 3, 8
    student_logits = torch.randn(B, T, V, requires_grad=True)
    teacher_logits = torch.randn(B, T, V)
    loss = entropy_aware_opd_loss(student_logits, teacher_logits)
    assert loss.dim() == 0
    assert torch.isfinite(loss)
    loss.backward()
    assert student_logits.grad is not None
    assert torch.isfinite(student_logits.grad).all()


def test_entropy_aware_opd_with_label_mask():
    """Label mask should zero out per-token loss on labels==0 positions."""
    B, T, V = 1, 4, 6
    student_logits = torch.randn(B, T, V, requires_grad=True)
    teacher_logits = torch.randn(B, T, V)
    full_loss = entropy_aware_opd_loss(student_logits, teacher_logits)
    half_mask = torch.tensor([[1, 1, 0, 0]])
    half_loss = entropy_aware_opd_loss(
        student_logits, teacher_logits, labels=half_mask,
    )
    # half_loss should be ~half of the unmasked sum (modulo the entropy gating
    # being position-dependent — but it should at least be < full_loss)
    assert float(half_loss) < float(full_loss)


def test_entropy_aware_opd_zero_when_distributions_match():
    """When student and teacher are identical, both KLs are 0 → loss is 0."""
    logits = torch.randn(1, 2, 4)
    loss = entropy_aware_opd_loss(logits, logits)
    assert float(loss) < 1e-5


def test_entropy_aware_opd_reduction_modes():
    student_logits = torch.randn(2, 3, 4, requires_grad=True)
    teacher_logits = torch.randn(2, 3, 4)
    none_loss = entropy_aware_opd_loss(student_logits, teacher_logits, reduction="none")
    mean_loss = entropy_aware_opd_loss(student_logits, teacher_logits, reduction="mean")
    sum_loss = entropy_aware_opd_loss(student_logits, teacher_logits, reduction="sum")
    batchmean_loss = entropy_aware_opd_loss(student_logits, teacher_logits, reduction="batchmean")
    assert none_loss.shape == (2, 3)
    assert mean_loss.dim() == 0
    assert sum_loss.dim() == 0
    assert batchmean_loss.dim() == 0
    # batchmean = sum / batch_size
    assert abs(float(batchmean_loss) - float(sum_loss) / 2) < 1e-4