Blame for 1.3B/code/losses/wan_kv_mmd.py

"""RBF-MMD on Wan KV matrices with a mathematically valid EMA distribution.

An RBF kernel has an (infinite-dimensional) feature map ``phi`` and

    MMD²(P, Q) = || E_P[phi(x)] - E_Q[phi(y)] ||².

We approximate ``phi`` with fixed random Fourier features (RFF), then keep an
EMA of the *kernel mean embeddings*.  This represents an EMA mixture of whole
historical distributions; averaging raw KV rows across unrelated videos does
not.  Historical embeddings are detached, so gradients only reach the latest
student KV matrix.  The natural ``1-beta`` factor is intentionally preserved.

History updates are transactional: ``forward(stage=True)`` only stages detached
current embeddings.  ``commit()`` must be called after a successful optimizer
step; ``discard()`` is used when the optimizer step is skipped.
"""

from __future__ import annotations

import math
from typing import Iterable, Optional, Tuple

import torch
import torch.distributed as dist
import torch.distributed.nn.functional as dist_nn
from torch import nn


def _sq_dist(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    x2 = x.square().sum(-1, keepdim=True)
    y2 = y.square().sum(-1).unsqueeze(0)
    return (x2 + y2 - 2.0 * x @ y.t()).clamp_min(0)


def exact_rbf_mmd(
    student: torch.Tensor,
    teacher: torch.Tensor,
    *,
    bandwidth: float,
    bandwidth_multipliers: Iterable[float] = (0.5, 1.0, 2.0),
) -> torch.Tensor:
    """Biased multi-bandwidth RBF-MMD² reference implementation."""
    if student.ndim != 2 or teacher.ndim != 2:
        raise ValueError("exact_rbf_mmd expects [samples, features] matrices")
    if student.shape[1] != teacher.shape[1] or min(student.shape[0], teacher.shape[0]) < 1:
        raise ValueError("student/teacher matrices need matching feature dims and samples")
    if not math.isfinite(float(bandwidth)) or float(bandwidth) <= 0:
        raise ValueError("bandwidth must be finite and positive")
    multipliers = tuple(float(v) for v in bandwidth_multipliers)
    if not multipliers or any((not math.isfinite(v) or v <= 0) for v in multipliers):
        raise ValueError("bandwidth multipliers must be finite and positive")

    x, y = student.float(), teacher.float()
    dxx, dyy, dxy = _sq_dist(x, x), _sq_dist(y, y), _sq_dist(x, y)
    bandwidth_sq = x.new_tensor(float(bandwidth) ** 2)
    terms = []
    for multiplier in multipliers:
        denom = 2.0 * bandwidth_sq * multiplier * multiplier
        terms.append(
            torch.exp(-dxx / denom).mean()
            + torch.exp(-dyy / denom).mean()
            - 2.0 * torch.exp(-dxy / denom).mean())
    return torch.stack(terms).mean()


class EMARBFMMD(nn.Module):
    """EMA RBF kernel-mean MMD using fixed random Fourier features.

    ``sample_mode="rows"`` preserves the original behavior, where every
    ``[D]`` row in a ``[B, N, D]`` input is one sample.  In
    ``sample_mode="video_matrix"``, the only sample axis is ``B`` and every
    position-sensitive ``[N, D]`` matrix is one sample.  Its RBF distance is
    the Frobenius distance, equivalently the Euclidean distance after
    vectorizing the matrix.
    """

    _STATE_VERSION = 2
    _SAMPLE_MODES = ("rows", "video_matrix")

    def __init__(
        self,
        beta: float = 0.99,
        num_rff_features: int = 256,
        bandwidth: Optional[float] = None,
        bandwidth_multipliers: Iterable[float] = (0.5, 1.0, 2.0),
        seed: int = 0,
        row_normalize: bool = False,
        sample_mode: str = "rows",
    ):
        super().__init__()
        if not 0.0 <= float(beta) < 1.0:
            raise ValueError("EMA beta must be in [0, 1)")
        if int(num_rff_features) <= 0:
            raise ValueError("num_rff_features must be positive")
        multipliers = tuple(float(v) for v in bandwidth_multipliers)
        if not multipliers or any((not math.isfinite(v) or v <= 0) for v in multipliers):
            raise ValueError("bandwidth multipliers must be finite and positive")
        if bandwidth is not None and (
            not math.isfinite(float(bandwidth)) or float(bandwidth) <= 0
        ):
            raise ValueError("bandwidth must be finite and positive")
        if sample_mode not in self._SAMPLE_MODES:
            raise ValueError(
                f"sample_mode must be one of {self._SAMPLE_MODES}, "
                f"got {sample_mode!r}")

        self.beta = float(beta)
        self.num_rff_features = int(num_rff_features)
        self.bandwidth_multipliers = multipliers
        self.seed = int(seed)
        self.row_normalize = bool(row_normalize)
        self.sample_mode = sample_mode

        initial_bw_sq = (
            float(bandwidth) ** 2 if bandwidth is not None else float("nan"))
        self.register_buffer(
            "bandwidth_sq", torch.tensor(initial_bw_sq, dtype=torch.float32),
            persistent=True)
        self.register_buffer("rff_weight", torch.empty(0), persistent=True)
        self.register_buffer("rff_phase", torch.empty(0), persistent=True)
        self.register_buffer("student_ema", torch.empty(0), persistent=True)
        self.register_buffer("teacher_ema", torch.empty(0), persistent=True)
        self.register_buffer(
            "num_updates", torch.zeros((), dtype=torch.long), persistent=True)

        # Deliberately not buffers: an uncommitted optimizer step must never enter
        # a checkpoint.
        self._pending_student: Optional[torch.Tensor] = None
        self._pending_teacher: Optional[torch.Tensor] = None

    @property
    def embedding_dim(self) -> int:
        return self.num_rff_features * len(self.bandwidth_multipliers)

    @staticmethod
    def _distributed() -> bool:
        return (
            dist.is_available()
            and dist.is_initialized()
            and dist.get_world_size() > 1
        )

    @classmethod
    @torch.no_grad()
    def _all_ranks_true(cls, local_value: bool, device: torch.device) -> bool:
        """Return one consensus bit before entering later collectives.

        A rank-local exception here would leave the other ranks blocked in the
        RFF broadcast/all-reduce.  Every input-validation failure therefore
        becomes an all-rank failure.
        """
        if not cls._distributed():
            return bool(local_value)
        flag = torch.tensor(
            int(bool(local_value)), device=device, dtype=torch.int32)
        dist.all_reduce(flag, op=dist.ReduceOp.MIN)
        return bool(flag.item())

    @classmethod
    def _global_mean(
        cls,
        local_sum: torch.Tensor,
        local_count: int,
        *,
        return_count: bool = False,
    ):
        """Autograd-aware, sample-weighted mean across all ranks.

        The differentiable all-reduce is intentional.  Since every rank
        backpropagates the same global loss, its backward supplies the right
        scaling before FSDP averages the replicated parameter gradients.
        """
        if not cls._distributed():
            count = local_sum.new_tensor(float(local_count))
            mean = local_sum / float(local_count)
        else:
            global_sum = dist_nn.all_reduce(
                local_sum, op=dist.ReduceOp.SUM)
            count = local_sum.new_tensor(float(local_count))
            dist.all_reduce(count, op=dist.ReduceOp.SUM)
            mean = global_sum / count
        if return_count:
            return mean, count
        return mean

    @torch.no_grad()
    def _estimate_bandwidth_sq(
        self, student: torch.Tensor, teacher: torch.Tensor, max_samples: int = 256
    ) -> torch.Tensor:
        if self.sample_mode == "rows":
            merged = torch.cat([
                student.detach().flatten(0, 1),
                teacher.detach().flatten(0, 1),
            ], dim=0).float()
        else:
            # Keep each [N, D] video matrix intact: B is the sole sample axis.
            merged = torch.cat([
                student.detach(),
                teacher.detach(),
            ], dim=0).float()
        if merged.shape[0] > max_samples:
            indices = torch.linspace(
                0, merged.shape[0] - 1, max_samples, device=merged.device
            ).round().long()
            merged = merged.index_select(0, indices)
        if self._distributed():
            # Use all ranks for the one-time median heuristic.  Rows are
            # padded only for the collective and unpadded immediately after.
            local_size = torch.tensor(
                [merged.shape[0]], device=merged.device, dtype=torch.long)
            gathered_sizes = [
                torch.zeros_like(local_size) for _ in range(dist.get_world_size())
            ]
            dist.all_gather(gathered_sizes, local_size)
            sizes = [int(size.item()) for size in gathered_sizes]
            max_size = max(sizes)
            padded = merged.new_zeros((max_size, *merged.shape[1:]))
            padded[:merged.shape[0]].copy_(merged)
            gathered_samples = [
                torch.empty_like(padded) for _ in range(dist.get_world_size())
            ]
            dist.all_gather(gathered_samples, padded)
            merged = torch.cat([
                samples[:size]
                for samples, size in zip(gathered_samples, sizes)
            ], dim=0)

        if merged.shape[0] < 2:
            estimate = merged.new_tensor(1.0)
        else:
            if not self._distributed() or dist.get_rank() == 0:
                if self.sample_mode == "rows":
                    distances_sq = torch.pdist(merged, p=2).square()
                else:
                    # ||X-Y||_F^2 without collapsing [N, D] into a token axis.
                    squared_norm = merged.square().sum(dim=(1, 2))
                    gram = torch.einsum(
                        "ind,jnd->ij", merged, merged)
                    pairwise_sq = (
                        squared_norm[:, None]
                        + squared_norm[None, :]
                        - 2.0 * gram
                    ).clamp_min(0)
                    upper = torch.triu_indices(
                        merged.shape[0],
                        merged.shape[0],
                        offset=1,
                        device=merged.device,
                    )
                    distances_sq = pairwise_sq[upper[0], upper[1]]
                positive = distances_sq[distances_sq > 0]
                estimate = (
                    positive.median()
                    if positive.numel()
                    else distances_sq.new_tensor(1.0)
                ).clamp_min(1e-6)
            else:
                estimate = merged.new_tensor(1.0)
        if self._distributed():
            # All ranks must use the same fixed, checkpointed RFF map.
            dist.broadcast(estimate, src=0)
        return estimate

    @torch.no_grad()
    def _initialize_rff(self, student: torch.Tensor, teacher: torch.Tensor) -> None:
        if self.sample_mode == "rows":
            sample_shape = (int(student.shape[-1]),)
        else:
            sample_shape = (
                int(student.shape[-2]),
                int(student.shape[-1]),
            )
        weight_shape = (*sample_shape, self.num_rff_features)
        if self.rff_weight.numel() > 0:
            if self.rff_weight.shape != weight_shape:
                raise ValueError(
                    "KV sample shape changed: RFF map expects "
                    f"{tuple(self.rff_weight.shape[:-1])}, got {sample_shape}")
            return
        if not bool(torch.isfinite(self.bandwidth_sq).item()):
            self.bandwidth_sq.copy_(self._estimate_bandwidth_sq(student, teacher))

        generator = torch.Generator(device="cpu")
        generator.manual_seed(self.seed)
        weight = torch.randn(
            weight_shape, generator=generator,
            dtype=torch.float32)
        phase = 2.0 * math.pi * torch.rand(
            self.num_rff_features, generator=generator, dtype=torch.float32)
        self.rff_weight.resize_(weight.shape).copy_(
            weight.to(device=student.device))
        self.rff_phase.resize_(phase.shape).copy_(
            phase.to(device=student.device))

    def _kernel_sum(self, matrix: torch.Tensor) -> Tuple[torch.Tensor, int]:
        samples = matrix.float()
        if self.row_normalize:
            samples = torch.nn.functional.normalize(samples, dim=-1)
        if self.sample_mode == "rows":
            samples = samples.flatten(0, 1)
            base_projection = samples @ self.rff_weight
        else:
            # This is exactly
            # samples.flatten(1) @ rff_weight.flatten(0, 1), while retaining
            # the [N, D] axes and their positional identity.
            base_projection = torch.einsum(
                "bnd,ndr->br", samples, self.rff_weight)
        bandwidth = self.bandwidth_sq.sqrt().clamp_min(1e-6)
        scale = math.sqrt(
            2.0 / (self.num_rff_features * len(self.bandwidth_multipliers)))
        features = [
            torch.cos(
                base_projection / (bandwidth * multiplier) + self.rff_phase
            ) * scale
            for multiplier in self.bandwidth_multipliers
        ]
        return torch.cat(features, dim=-1).sum(dim=0), int(samples.shape[0])

    def forward(
        self,
        student: torch.Tensor,
        teacher: torch.Tensor,
        *,
        stage: bool = True,
    ) -> Tuple[torch.Tensor, dict]:
        local_shape_ok = (
            student.ndim == 3
            and teacher.ndim == 3
            and student.shape == teacher.shape
            and student.device == teacher.device
            and min(student.shape[0], student.shape[1], student.shape[2]) > 0
        )
        if not self._all_ranks_true(local_shape_ok, student.device):
            raise ValueError(
                "invalid Wan KV shape/device on at least one rank; each local "
                "student/teacher pair must match [B, tokens, features]")

        # A mismatched feature width would initialize incompatible RFF maps.
        # Matrix samples additionally require a shared N because every token
        # position has its own slice of the position-sensitive random weight.
        if self._distributed():
            rff_sample_shape = (
                student.shape[-1:]
                if self.sample_mode == "rows"
                else student.shape[1:]
            )
            shape_bounds = torch.tensor(
                [*rff_sample_shape, *(-value for value in rff_sample_shape)],
                device=student.device,
                dtype=torch.long,
            )
            dist.all_reduce(shape_bounds, op=dist.ReduceOp.MIN)
            num_dimensions = len(rff_sample_shape)
            if any(
                int(shape_bounds[index].item())
                != -int(shape_bounds[index + num_dimensions].item())
                for index in range(num_dimensions)
            ):
                raise ValueError(
                    "Wan KV RFF sample shape differs across distributed ranks")

        rff_initialized = self.rff_weight.numel() > 0
        expected_rff_shape = (
            (student.shape[-1], self.num_rff_features)
            if self.sample_mode == "rows"
            else (
                student.shape[-2],
                student.shape[-1],
                self.num_rff_features,
            )
        )
        local_rff_shape_ok = (
            not rff_initialized
            or self.rff_weight.shape == expected_rff_shape
        )
        if not self._all_ranks_true(local_rff_shape_ok, student.device):
            raise ValueError(
                "RFF map shape is incompatible on at least one rank")
        if self._distributed():
            bandwidth_initialized = bool(
                torch.isfinite(self.bandwidth_sq).item())
            state_bounds = torch.tensor(
                [
                    int(rff_initialized),
                    -int(rff_initialized),
                    int(bandwidth_initialized),
                    -int(bandwidth_initialized),
                ],
                device=student.device,
                dtype=torch.int32,
            )
            dist.all_reduce(state_bounds, op=dist.ReduceOp.MIN)
            if (
                int(state_bounds[0].item()) != -int(state_bounds[1].item())
                or int(state_bounds[2].item()) != -int(state_bounds[3].item())
            ):
                raise RuntimeError(
                    "RFF/bandwidth initialization state differs across ranks")

        local_finite = bool(
            torch.isfinite(student).all().item()
            and torch.isfinite(teacher).all().item()
        )
        if not self._all_ranks_true(local_finite, student.device):
            raise FloatingPointError(
                "non-finite Wan KV feature on at least one rank")

        teacher = teacher.detach()
        self._initialize_rff(student, teacher)
        student_sum, student_count = self._kernel_sum(student)
        teacher_sum, teacher_count = self._kernel_sum(teacher)
        current_student, global_count = self._global_mean(
            student_sum, student_count, return_count=True)
        current_teacher = self._global_mean(teacher_sum, teacher_count)

        if self.student_ema.numel() == 0:
            current_coefficient = 1.0
            effective_student = current_student
            effective_teacher = current_teacher.detach()
        else:
            current_coefficient = 1.0 - self.beta
            # History has no graph. Consequently only current_student receives
            # gradients, with the mathematically correct (1-beta) coefficient.
            effective_student = (
                self.beta * self.student_ema.detach()
                + (1.0 - self.beta) * current_student)
            effective_teacher = (
                self.beta * self.teacher_ema.detach()
                + (1.0 - self.beta) * current_teacher.detach())

        loss = (effective_student - effective_teacher).square().sum()
        if stage:
            self._pending_student = current_student.detach()
            self._pending_teacher = current_teacher.detach()
        return loss, {
            "bandwidth": self.bandwidth_sq.detach().sqrt(),
            "history_updates": self.num_updates.detach().clone(),
            "current_coefficient": loss.new_tensor(current_coefficient),
            "local_samples": torch.tensor(
                student_count, device=loss.device, dtype=torch.long),
            "global_samples": global_count.detach().to(dtype=torch.long),
        }

    @torch.no_grad()
    def commit(self) -> None:
        """Commit staged kernel means after a successful optimizer step."""
        if self._pending_student is None or self._pending_teacher is None:
            return
        if not (
            torch.isfinite(self._pending_student).all()
            and torch.isfinite(self._pending_teacher).all()
        ):
            self.discard()
            raise FloatingPointError("refusing to commit non-finite KV kernel mean")
        if self.student_ema.numel() == 0:
            self.student_ema.resize_as_(self._pending_student).copy_(
                self._pending_student)
            self.teacher_ema.resize_as_(self._pending_teacher).copy_(
                self._pending_teacher)
        else:
            self.student_ema.mul_(self.beta).add_(
                self._pending_student, alpha=1.0 - self.beta)
            self.teacher_ema.mul_(self.beta).add_(
                self._pending_teacher, alpha=1.0 - self.beta)
        self.num_updates.add_(1)
        self.discard()

    def discard(self) -> None:
        """Drop staged state when backward/optimizer is skipped."""
        self._pending_student = None
        self._pending_teacher = None

    def get_extra_state(self):
        state = {
            "version": self._STATE_VERSION,
            "beta": self.beta,
            "num_rff_features": self.num_rff_features,
            "bandwidth_multipliers": self.bandwidth_multipliers,
            "seed": self.seed,
            "row_normalize": self.row_normalize,
        }
        # Preserve the exact legacy rows-mode payload.  Wan v3 checkpoints
        # contain this dict and must remain loadable without migration.
        if self.sample_mode != "rows":
            state["sample_mode"] = self.sample_mode
        return state

    def set_extra_state(self, state):
        expected = self.get_extra_state()
        if state != expected:
            raise ValueError(
                f"Wan KV MMD checkpoint/config mismatch: saved={state}, "
                f"current={expected}")

    def _load_from_state_dict(
        self, state_dict, prefix, local_metadata, strict, missing_keys,
        unexpected_keys, error_msgs,
    ):
        for name in (
            "rff_weight", "rff_phase", "student_ema", "teacher_ema",
        ):
            key = prefix + name
            if key in state_dict:
                getattr(self, name).resize_as_(state_dict[key])
        super()._load_from_state_dict(
            state_dict, prefix, local_metadata, strict, missing_keys,
            unexpected_keys, error_msgs)