Blame for 14B/code/trainer/feature_gradient_replay.py

"""Separate distribution estimation from generator backpropagation (Sec. 3.3).

The distribution batch B supplies all cross-video interactions. Only the smaller
configured differentiation batch b needs feature gradients and rollout replay.
This module caches features and returns their cotangents (dL/d feature); the
trainer later replays the encoders, VAE, and generator to obtain dL/d parameters.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Mapping, Sequence

import torch
import torch.distributed as dist

from losses.multi_encoder_mmd import (
    MultiEncoderMMD,
)


def _detach(value: Any) -> Any:
    if isinstance(value, torch.Tensor):
        return value.detach()
    if isinstance(value, Mapping):
        return {key: _detach(item) for key, item in value.items()}
    if isinstance(value, tuple):
        return tuple(_detach(item) for item in value)
    if isinstance(value, list):
        return [_detach(item) for item in value]
    return value


def _detach_to_cpu(value: Any) -> Any:
    if isinstance(value, torch.Tensor):
        return value.detach().to(device="cpu", non_blocking=False)
    if isinstance(value, Mapping):
        return {key: _detach_to_cpu(item) for key, item in value.items()}
    if isinstance(value, tuple):
        return tuple(_detach_to_cpu(item) for item in value)
    if isinstance(value, list):
        return [_detach_to_cpu(item) for item in value]
    return value


@dataclass(frozen=True)
class FeatureGradientReplayResult:
    """Detached loss, per-microbatch feature gradients, and replay bookkeeping.

    ``local_cotangents`` are the derivatives of the global distribution loss.
    ``fsdp_cotangents`` additionally compensate for FSDP's gradient averaging.
    Field names stay compatible with existing training logs and checkpoints.
    """

    loss: torch.Tensor
    optimized_loss: torch.Tensor
    metadata: Mapping[str, Any]
    local_cotangents: tuple[Mapping[str, torch.Tensor], ...]
    fsdp_cotangents: tuple[Mapping[str, torch.Tensor], ...]
    replay_batch_indices: tuple[int, ...]
    replay_sample_ids: tuple[tuple[str, ...], ...]
    local_samples: int
    global_samples: int
    gradient_local_samples: int
    gradient_global_samples: int
    current_teacher_local_samples: int
    current_teacher_global_samples: int
    world_size: int


class FeatureGradientReplay:
    """Cache B generated feature sets and differentiate the selected b samples.

    B and b denote global sample counts; constructor counts are per rank. Each
    encoder keeps its own [tokens, channels] layout while sharing sample order.
    This implementation selects a per-rank suffix. All B samples still define
    the distribution loss, including interactions with unselected samples.
    """

    def __init__(
        self,
        loss_module: MultiEncoderMMD,
        *,
        target_local_samples: int,
        gradient_local_samples: int,
        loss_scale: float = 1.0,
        require_sample_id_alignment: bool = True,
        full_student_partial_gradient: bool = True,
    ) -> None:
        if not isinstance(loss_module, MultiEncoderMMD):
            raise TypeError(
                "loss_module must be MultiRepresentationVideoNystromMMD")
        if not 0 < int(gradient_local_samples) < int(target_local_samples):
            raise ValueError(
                "dual GradCache requires a strict positive gradient suffix")
        if not full_student_partial_gradient:
            raise ValueError(
                "dual GradCache requires full_student_partial_gradient=true")
        self.loss_module = loss_module
        self.target_local_samples = int(target_local_samples)
        self.gradient_local_samples = int(gradient_local_samples)
        self.loss_scale = float(loss_scale)
        self.require_sample_id_alignment = bool(require_sample_id_alignment)
        self.full_student_partial_gradient = True
        self.partial_gradient_enabled = True
        self.current_teacher_attraction_enabled = bool(
            loss_module.current_teacher_attraction_enabled)
        self._feature_batches: list[dict[str, torch.Tensor]] = []
        self._teacher_batches: list[dict[str, torch.Tensor]] = []
        self._batch_sizes: list[int] = []
        self._gradient_batch_indices: list[int] = []
        self._gradient_batch_sizes: list[int] = []
        self._gradient_sample_ids: list[tuple[str, ...]] = []
        self._selection_started = False
        self._device: torch.device | None = None
        self._result: FeatureGradientReplayResult | None = None

    @staticmethod
    def _world_size() -> int:
        return dist.get_world_size() if dist.is_initialized() else 1

    @staticmethod
    def _rank() -> int:
        return dist.get_rank() if dist.is_initialized() else 0

    @property
    def local_samples(self) -> int:
        return sum(self._batch_sizes)

    @property
    def selected_local_samples(self) -> int:
        return sum(self._gradient_batch_sizes)

    @property
    def finalized(self) -> bool:
        return self._result is not None

    def _validate_mapping(
        self,
        value: Mapping[str, torch.Tensor],
        label: str,
        *,
        expected_batch: int | None = None,
    ) -> tuple[int, torch.device]:
        if not isinstance(value, Mapping):
            raise TypeError(f"{label} must be a branch tensor mapping")
        if set(value) != set(self.loss_module.branch_names):
            raise ValueError(
                f"{label} keys must be {self.loss_module.branch_names}")
        batch_size = None
        device = None
        for name in self.loss_module.branch_names:
            tensor = value[name]
            if (
                not isinstance(tensor, torch.Tensor)
                or tensor.ndim != 3
                or not tensor.is_floating_point()
                or min(tensor.shape) <= 0
            ):
                raise ValueError(f"{label}[{name!r}] must be floating [B,N,D]")
            if tuple(tensor.shape[1:]) != self.loss_module.sample_shapes[name]:
                raise ValueError(
                    f"{label}[{name!r}] geometry {tuple(tensor.shape[1:])} "
                    f"!= {self.loss_module.sample_shapes[name]}")
            if not bool(torch.isfinite(tensor).all().item()):
                raise FloatingPointError(f"{label}[{name!r}] is non-finite")
            if batch_size is None:
                batch_size = int(tensor.shape[0])
                device = tensor.device
            elif int(tensor.shape[0]) != batch_size or tensor.device != device:
                raise ValueError(
                    f"all {label} branches must share batch size and device")
        if expected_batch is not None and batch_size != expected_batch:
            raise ValueError(
                f"{label} batch size {batch_size} != {expected_batch}")
        return int(batch_size), device

    def add(
        self,
        generated_features: Mapping[str, torch.Tensor],
        current_teacher_features: Mapping[str, torch.Tensor] | None = None,
        *,
        gradient_selected: bool = True,
        sample_ids: Sequence[str] | None = None,
        current_teacher_sample_ids: Sequence[str] | None = None,
    ) -> int:
        """Store one rollout's features without retaining its generator graph."""
        if self.finalized:
            raise RuntimeError("cannot add after finalize")
        batch_size, device = self._validate_mapping(
            generated_features, "generated_features")
        if self._device is None:
            self._device = device
        elif device != self._device:
            raise ValueError("all microbatches must share a device")
        if self.local_samples + batch_size > self.target_local_samples:
            raise ValueError("dual GradCache exceeds target_local_samples")
        if gradient_selected:
            self._selection_started = True
            if self.selected_local_samples + batch_size > self.gradient_local_samples:
                raise ValueError("dual GradCache exceeds gradient suffix")
        elif self._selection_started:
            raise ValueError("gradient-selected microbatches must form a suffix")

        normalized_ids = None
        if sample_ids is not None:
            normalized_ids = tuple(str(value) for value in sample_ids)
            if len(normalized_ids) != batch_size or any(not value for value in normalized_ids):
                raise ValueError("sample_ids must contain one non-empty ID per row")
        if (
            not self.current_teacher_attraction_enabled
            and current_teacher_sample_ids is not None
        ):
            raise ValueError(
                "current_teacher_sample_ids were provided while current-teacher "
                "attraction is disabled")
        if (
            gradient_selected
            and self.require_sample_id_alignment
            and self.current_teacher_attraction_enabled
        ):
            if normalized_ids is None:
                raise ValueError("selected dual samples require sample_ids")
            teacher_ids = (
                tuple(str(value) for value in current_teacher_sample_ids)
                if current_teacher_sample_ids is not None else None
            )
            if teacher_ids != normalized_ids:
                raise ValueError("student and teacher sample IDs/order differ")

        if gradient_selected and self.current_teacher_attraction_enabled:
            if current_teacher_features is None:
                raise ValueError("selected dual samples require teacher features")
            self._validate_mapping(
                current_teacher_features,
                "current_teacher_features",
                expected_batch=batch_size,
            )
        elif current_teacher_features is not None:
            raise ValueError(
                "current teacher features were provided when this microbatch "
                "must not carry teacher payloads")

        # Collection stage (Algorithm 1): keep only feature values and sample
        # order. Replay state is stored by the trainer, outside this cache.
        staged = {
            name: generated_features[name].detach().float().clone()
            for name in self.loss_module.branch_names
        }
        self._feature_batches.append(staged)
        index = len(self._feature_batches) - 1
        if gradient_selected:
            self._gradient_batch_indices.append(index)
            self._gradient_batch_sizes.append(batch_size)
            self._gradient_sample_ids.append(normalized_ids)
            if self.current_teacher_attraction_enabled:
                self._teacher_batches.append({
                    name: current_teacher_features[name].detach().float().clone()
                    for name in self.loss_module.branch_names
                })
        self._batch_sizes.append(batch_size)
        return index

    def _gather(self, value: torch.Tensor) -> list[torch.Tensor]:
        if self._world_size() == 1:
            return [value.detach()]
        gathered = [torch.empty_like(value) for _ in range(self._world_size())]
        dist.all_gather(gathered, value.detach().contiguous())
        return gathered

    def _validate_ready(self) -> None:
        local = torch.tensor(
            [self.local_samples, self.selected_local_samples],
            device=self._device,
            dtype=torch.long,
        )
        expected = (self.target_local_samples, self.gradient_local_samples)
        if self._world_size() == 1:
            observed = tuple(int(value) for value in local.tolist())
            if observed != expected:
                raise RuntimeError(
                    f"incomplete dual GradCache: {observed} != {expected}")
            return
        gathered = [torch.empty_like(local) for _ in range(self._world_size())]
        dist.all_gather(gathered, local)
        observed = [tuple(int(value) for value in item.tolist()) for item in gathered]
        if any(item != expected for item in observed):
            raise RuntimeError(
                f"incomplete distributed dual GradCache: {observed} != {expected}")

    def finalize(self) -> FeatureGradientReplayResult:
        """Evaluate the global objective and cache dL/d feature for replay."""
        if self._result is not None:
            return self._result
        self._validate_ready()
        world_size = self._world_size()
        global_distribution_features = {}
        global_differentiation_features = {}
        global_reference_features = (
            {} if self.current_teacher_attraction_enabled else None)
        local_leaves = {}
        # Map each rank's selected suffix into the concatenated B-sample batch.
        # These indices preserve full-batch interactions while only b rows have
        # an autograd path. This function does not randomly resample the subset.
        suffix_start = self.target_local_samples - self.gradient_local_samples
        gradient_indices = torch.cat([
            torch.arange(
                rank * self.target_local_samples + suffix_start,
                (rank + 1) * self.target_local_samples,
                device=self._device,
                dtype=torch.long,
            )
            for rank in range(world_size)
        ])
        for name in self.loss_module.branch_names:
            local_distribution_features = torch.cat([
                batch[name] for batch in self._feature_batches
            ], dim=0)
            local_differentiation_features = torch.cat([
                self._feature_batches[index][name]
                for index in self._gradient_batch_indices
            ], dim=0)
            distribution_parts = self._gather(local_distribution_features)
            differentiation_parts = self._gather(local_differentiation_features)
            # Only this rank's selected features are differentiable leaves.
            # Remote features contribute values; their owning rank gets their
            # derivatives. No generator or encoder graph is retained here.
            leaf = local_differentiation_features.detach().requires_grad_(True)
            differentiation_parts[self._rank()] = leaf
            local_leaves[name] = leaf
            global_distribution_features[name] = torch.cat(distribution_parts, dim=0).detach()
            global_differentiation_features[name] = torch.cat(differentiation_parts, dim=0)
            if self.current_teacher_attraction_enabled:
                local_reference_features = torch.cat([
                    batch[name] for batch in self._teacher_batches
                ], dim=0)
                reference_parts = self._gather(local_reference_features)
                global_reference_features[name] = torch.cat(
                    reference_parts, dim=0).detach()

        # Resolve all cross-video interactions once. The loss combines the
        # persistent Nystrom summary with current reference features (Eq. 9),
        # then aggregates the encoder objectives (Eqs. 11 and 13).
        raw_loss, metadata = self.loss_module(
            gradient_students=global_differentiation_features,
            student_reference=global_distribution_features,
            current_teacher=global_reference_features,
            gradient_student_indices=gradient_indices,
        )
        optimized = raw_loss * self.loss_scale
        gradients = torch.autograd.grad(
            optimized,
            tuple(local_leaves[name] for name in self.loss_module.branch_names),
            retain_graph=False,
            create_graph=False,
        )
        # Restore the collection order so each feature gradient is replayed
        # through the exact sample and encoder that produced its feature value.
        split_by_branch = {
            name: tuple(part.contiguous() for part in gradient.detach().split(
                self._gradient_batch_sizes, dim=0))
            for name, gradient in zip(self.loss_module.branch_names, gradients)
        }
        local_parts = tuple({
            name: split_by_branch[name][index]
            for name in self.loss_module.branch_names
        } for index in range(len(self._gradient_batch_sizes)))
        # FSDP averages parameter gradients across ranks, so cancel that factor
        # here. This world_size multiplier is separate from the B/b selective
        # differentiation correction (Eq. 12), applied later by the trainer.
        fsdp_parts = tuple({
            name: (value * float(world_size)).contiguous()
            for name, value in part.items()
        } for part in local_parts)
        global_samples = self.target_local_samples * world_size
        gradient_global = self.gradient_local_samples * world_size
        kss_denominator = global_samples * global_samples
        # Keep established metadata names for checkpoint/log compatibility.
        detached_metadata = dict(_detach(metadata))
        teacher_local = (
            self.gradient_local_samples
            if self.current_teacher_attraction_enabled else 0)
        teacher_global = teacher_local * world_size
        detached_metadata.update({
            "raw_loss": raw_loss.detach(),
            "optimized_loss": optimized.detach(),
            "loss_scale": raw_loss.new_tensor(self.loss_scale),
            "gradcache_local_samples": raw_loss.new_tensor(
                self.target_local_samples, dtype=torch.long),
            "gradcache_global_samples": raw_loss.new_tensor(
                global_samples, dtype=torch.long),
            "gradcache_forward_local_samples": raw_loss.new_tensor(
                self.target_local_samples, dtype=torch.long),
            "gradcache_forward_global_samples": raw_loss.new_tensor(
                global_samples, dtype=torch.long),
            "gradcache_gradient_local_samples": raw_loss.new_tensor(
                self.gradient_local_samples, dtype=torch.long),
            "gradcache_gradient_global_samples": raw_loss.new_tensor(
                gradient_global, dtype=torch.long),
            "gradcache_current_teacher_local_samples": raw_loss.new_tensor(
                teacher_local, dtype=torch.long),
            "gradcache_current_teacher_global_samples": raw_loss.new_tensor(
                teacher_global, dtype=torch.long),
            "gradcache_kss_denominator": raw_loss.new_tensor(
                kss_denominator, dtype=torch.long),
            "gradcache_partial_gradient_enabled": raw_loss.new_ones(
                (), dtype=torch.long),
            "gradcache_full_student_partial_gradient": raw_loss.new_ones(
                (), dtype=torch.long),
            "gradcache_world_size": raw_loss.new_tensor(
                world_size, dtype=torch.long),
            "gradcache_fsdp_cotangent_scale": raw_loss.new_tensor(
                float(world_size)),
            "gradcache_branch_names": self.loss_module.branch_names,
        })
        self._result = FeatureGradientReplayResult(
            loss=raw_loss.detach(),
            optimized_loss=optimized.detach(),
            metadata=detached_metadata,
            local_cotangents=local_parts,
            fsdp_cotangents=fsdp_parts,
            replay_batch_indices=tuple(self._gradient_batch_indices),
            replay_sample_ids=tuple(self._gradient_sample_ids),
            local_samples=self.target_local_samples,
            global_samples=global_samples,
            gradient_local_samples=self.gradient_local_samples,
            gradient_global_samples=gradient_global,
            current_teacher_local_samples=teacher_local,
            current_teacher_global_samples=teacher_global,
            world_size=world_size,
        )
        return self._result

    def offload_finalized_to_cpu(
        self,
    ) -> FeatureGradientReplayResult:
        """Move finalized feature/cotangent storage to CPU transactionally."""
        if self._result is None:
            raise RuntimeError(
                "cannot offload a multi-representation cache before finalize")
        self._feature_batches = [
            _detach_to_cpu(value) for value in self._feature_batches
        ]
        self._teacher_batches = [
            _detach_to_cpu(value) for value in self._teacher_batches
        ]
        result = self._result
        self._result = FeatureGradientReplayResult(
            loss=_detach_to_cpu(result.loss),
            optimized_loss=_detach_to_cpu(result.optimized_loss),
            metadata=_detach_to_cpu(result.metadata),
            local_cotangents=_detach_to_cpu(result.local_cotangents),
            fsdp_cotangents=_detach_to_cpu(result.fsdp_cotangents),
            replay_batch_indices=result.replay_batch_indices,
            replay_sample_ids=result.replay_sample_ids,
            local_samples=result.local_samples,
            global_samples=result.global_samples,
            gradient_local_samples=result.gradient_local_samples,
            gradient_global_samples=result.gradient_global_samples,
            current_teacher_local_samples=result.current_teacher_local_samples,
            current_teacher_global_samples=result.current_teacher_global_samples,
            world_size=result.world_size,
        )
        return self._result


    def _clear(self) -> None:
        self._feature_batches.clear()
        self._teacher_batches.clear()
        self._batch_sizes.clear()
        self._gradient_batch_indices.clear()
        self._gradient_batch_sizes.clear()
        self._gradient_sample_ids.clear()
        self._selection_started = False
        self._device = None
        self._result = None

    def commit(self) -> None:
        """Release cached values after the corresponding optimizer update."""
        if self._result is None:
            raise RuntimeError("cannot commit before finalize")
        self._clear()

    def discard(self) -> None:
        """Release an abandoned batch without marking an optimizer update."""
        self._clear()


# Preserve old imports and serialized class references during the naming cleanup.
MultiRepresentationFreshVideoGradCache = FeatureGradientReplay
MultiRepresentationFreshVideoGradCacheResult = FeatureGradientReplayResult

__all__ = [
    "FeatureGradientReplay",
    "FeatureGradientReplayResult",
    "MultiRepresentationFreshVideoGradCache",
    "MultiRepresentationFreshVideoGradCacheResult",
]