spikes/005-integrated-trainer-skeleton/tests/test_opsd_loss.py

"""test_opsd_loss.py — unit test for the lifted OPSD loss.

Verifies:
  1. Loss is differentiable.
  2. Loss is 0 when student == teacher (sanity).
  3. Loss is positive when student != teacher.
  4. Forward KL (beta=0), reverse KL (beta=1), and JSD (beta=0.5) all run
     and produce finite values.
  5. Label masking zeros out ignored positions.
  6. top_k restriction reduces compute and gives a valid result.

Run:  pytest spikes/005-integrated-trainer-skeleton/tests/test_opsd_loss.py -v
"""

from __future__ import annotations

import sys
from pathlib import Path

import pytest
import torch

# Make sibling modules importable without packaging the skeleton
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from opsd_loss import generalized_jsd_loss  # noqa: E402


# ----------------------------------------------------------------------------
# Test fixtures
# ----------------------------------------------------------------------------

@pytest.fixture
def small_logits():
    """B=2, T=4, V=8 — small enough to debug if anything fails."""
    torch.manual_seed(0)
    return torch.randn(2, 4, 8, requires_grad=True), torch.randn(2, 4, 8)


# ----------------------------------------------------------------------------
# Tests
# ----------------------------------------------------------------------------

def test_loss_is_finite_and_positive(small_logits):
    student, teacher = small_logits
    loss = generalized_jsd_loss(student, teacher, beta=0.5)
    assert torch.isfinite(loss).all(), "JSD loss is NaN or Inf"
    assert loss.item() > 0, "JSD loss should be positive when distributions differ"


def test_loss_is_zero_when_student_equals_teacher():
    """If student_logits == teacher_logits, JSD == 0 (within numeric tolerance)."""
    torch.manual_seed(1)
    logits = torch.randn(2, 4, 8, requires_grad=True)
    loss = generalized_jsd_loss(logits, logits.detach().clone(), beta=0.5)
    # Some tiny float noise from log_softmax round-trips → tolerance, not exact
    assert loss.abs().item() < 1e-5, f"Expected ~0 loss, got {loss.item()}"


def test_loss_is_differentiable(small_logits):
    student, teacher = small_logits
    loss = generalized_jsd_loss(student, teacher, beta=0.5)
    loss.backward()
    assert student.grad is not None
    assert torch.isfinite(student.grad).all(), "Gradient has NaN/Inf"
    # Teacher should NOT receive gradient (it had requires_grad=False from fixture)
    assert teacher.grad is None or teacher.requires_grad is False


@pytest.mark.parametrize("beta", [0.0, 0.5, 1.0])
def test_all_betas_run(small_logits, beta):
    student, teacher = small_logits
    loss = generalized_jsd_loss(student, teacher, beta=beta)
    assert torch.isfinite(loss).all(), f"Loss not finite at beta={beta}"
    assert loss.item() > 0, f"Loss not positive at beta={beta}"


def test_label_mask_excludes_ignored_positions():
    """Positions with label == -100 should not contribute to the loss."""
    torch.manual_seed(2)
    student = torch.randn(2, 4, 8, requires_grad=True)
    teacher = torch.randn(2, 4, 8)

    # Mask: include only position 0 in batch element 0; nothing else.
    labels = torch.full((2, 4), -100, dtype=torch.long)
    labels[0, 0] = 1  # one valid token

    loss_with_mask = generalized_jsd_loss(student, teacher, labels=labels, reduction="sum")

    # Compare to unmasked
    loss_unmasked = generalized_jsd_loss(student, teacher, labels=None, reduction="sum")

    # Masked loss must be strictly smaller (ignored positions zero out)
    assert loss_with_mask < loss_unmasked, (
        "Masked loss should be smaller than unmasked when most positions are masked"
    )
    assert loss_with_mask.item() > 0, "At least one valid token should give positive loss"


def test_top_k_restriction(small_logits):
    """top_k restricts the KL to the teacher's top-k tokens."""
    student, teacher = small_logits
    loss_full = generalized_jsd_loss(student, teacher, beta=0.5)
    loss_topk = generalized_jsd_loss(student, teacher, beta=0.5, top_k=4)
    assert torch.isfinite(loss_topk).all()
    # top-k loss should typically be smaller (fewer terms in the sum) but not strictly so
    # because the renormalization can flip relative magnitudes. Just check finite + positive.
    assert loss_topk.item() > 0


def test_token_clip(small_logits):
    """Per-token clip caps individual token contributions."""
    student, teacher = small_logits
    loss_unclipped = generalized_jsd_loss(student, teacher, beta=0.5)
    loss_clipped = generalized_jsd_loss(student, teacher, beta=0.5, token_clip=0.001)
    assert loss_clipped <= loss_unclipped, "Clipping should reduce or equal loss"