Add training and inference source code

2ea2835 verified 29 days ago

21.1 kB

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

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