| """Fresh-update gradient cache for video-level distribution losses. |
| |
| The cache deliberately has no state across optimizer updates. It collects a |
| fixed number of detached ``[B, N, D]`` feature matrices on every rank, builds |
| the exact global fresh batch in deterministic rank order, and differentiates |
| the frozen reference loss only with respect to the local rank's feature leaf. |
| |
| The two returned cotangents have distinct meanings: |
| |
| ``local_cotangents`` |
| The mathematical derivatives ``d L_global / d feature_local``. |
| |
| ``fsdp_cotangents`` |
| ``world_size * local_cotangents``. These are the cotangents that replay |
| should use when FSDP/DDP will average parameter gradients across ranks: |
| |
| ``mean_r(world_size * J_r.T @ dL/dF_r) = sum_r J_r.T @ dL/dF_r``. |
| |
| Collection is transactional. ``finalize`` stages a result but retains the |
| fresh features until the caller reports optimizer success with ``commit`` or |
| failure with ``discard``. Both terminal operations clear all update-local |
| state, so no generated sample can leak into the next optimizer update. |
| """ |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
| import math |
| from typing import Any, Mapping, Sequence |
|
|
| import torch |
| import torch.distributed as dist |
|
|
| from losses.video_nystrom_mmd import VideoNystromMMD |
| from losses.multilayer_video_nystrom_mmd import ( |
| MultiLayerVideoNystromMMD, |
| ) |
|
|
|
|
| def _detach_metadata(value: Any) -> Any: |
| if isinstance(value, torch.Tensor): |
| return value.detach() |
| if isinstance(value, Mapping): |
| return { |
| key: _detach_metadata(item) |
| for key, item in value.items() |
| } |
| if isinstance(value, tuple): |
| return tuple(_detach_metadata(item) for item in value) |
| if isinstance(value, list): |
| return [_detach_metadata(item) for item in value] |
| return value |
|
|
|
|
| @dataclass(frozen=True) |
| class FreshVideoGradCacheResult: |
| """Finalized loss and replay cotangents for one fresh optimizer update.""" |
|
|
| loss: torch.Tensor |
| optimized_loss: torch.Tensor |
| metadata: Mapping[str, Any] |
| local_cotangents: tuple[torch.Tensor, ...] |
| fsdp_cotangents: tuple[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 FreshVideoGradCache: |
| """Collect and differentiate one globally fresh video feature batch.""" |
|
|
| def __init__( |
| self, |
| loss_module: VideoNystromMMD | MultiLayerVideoNystromMMD, |
| *, |
| target_local_samples: int = 16, |
| gradient_local_samples: int | None = None, |
| loss_scale: float = 1.0, |
| require_sample_id_alignment: bool = False, |
| full_student_partial_gradient: bool = False, |
| ) -> None: |
| if not isinstance( |
| loss_module, |
| (VideoNystromMMD, MultiLayerVideoNystromMMD), |
| ): |
| raise TypeError( |
| "loss_module must be a VideoNystromMMD or " |
| "MultiLayerVideoNystromMMD instance") |
| if int(target_local_samples) <= 0: |
| raise ValueError("target_local_samples must be positive") |
| if gradient_local_samples is None: |
| gradient_local_samples = int(target_local_samples) |
| if not 0 < int(gradient_local_samples) <= int(target_local_samples): |
| raise ValueError( |
| "gradient_local_samples must be in [1, target_local_samples]") |
| if ( |
| not math.isfinite(float(loss_scale)) |
| or float(loss_scale) < 0.0 |
| ): |
| raise ValueError("loss_scale must be finite and non-negative") |
| if any(parameter.requires_grad for parameter in loss_module.parameters()): |
| raise ValueError( |
| "VideoNystromMMD reference parameters must be frozen") |
|
|
| self.loss_module = loss_module |
| self.current_teacher_attraction_enabled = bool(getattr( |
| loss_module, "current_teacher_attraction_enabled", False)) |
| self.target_local_samples = int(target_local_samples) |
| self.gradient_local_samples = int(gradient_local_samples) |
| self.partial_gradient_enabled = bool( |
| self.gradient_local_samples < self.target_local_samples) |
| self.full_student_partial_gradient = bool( |
| full_student_partial_gradient) |
| if ( |
| self.full_student_partial_gradient |
| and not self.partial_gradient_enabled |
| ): |
| raise ValueError( |
| "full_student_partial_gradient requires " |
| "gradient_local_samples < target_local_samples") |
| if ( |
| self.partial_gradient_enabled |
| and isinstance(loss_module, MultiLayerVideoNystromMMD) |
| ): |
| raise TypeError( |
| "partial-gradient GradCache currently requires a single " |
| "VideoNystromMMD layer") |
| self.loss_scale = float(loss_scale) |
| self.require_sample_id_alignment = bool(require_sample_id_alignment) |
| self._feature_batches: list[torch.Tensor] = [] |
| self._teacher_feature_batches: list[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._sample_shape: tuple[int, int] | None = None |
| self._device: torch.device | None = None |
| self._result: FreshVideoGradCacheResult | None = None |
|
|
| @staticmethod |
| def _world_size() -> int: |
| if dist.is_available() and dist.is_initialized(): |
| return dist.get_world_size() |
| return 1 |
|
|
| @staticmethod |
| def _rank() -> int: |
| if dist.is_available() and dist.is_initialized(): |
| return dist.get_rank() |
| return 0 |
|
|
| @staticmethod |
| def _all_ranks_true(local_value: bool, device: torch.device) -> bool: |
| flag = torch.tensor( |
| int(bool(local_value)), device=device, dtype=torch.int32) |
| if ( |
| dist.is_available() |
| and dist.is_initialized() |
| and dist.get_world_size() > 1 |
| ): |
| dist.all_reduce(flag, op=dist.ReduceOp.MIN) |
| return bool(flag.item()) |
|
|
| @property |
| def local_samples(self) -> int: |
| return sum(self._batch_sizes) |
|
|
| @property |
| def num_microbatches(self) -> int: |
| return len(self._feature_batches) |
|
|
| @property |
| def selected_local_samples(self) -> int: |
| return sum(self._gradient_batch_sizes) |
|
|
| @property |
| def ready(self) -> bool: |
| return ( |
| self.local_samples == self.target_local_samples |
| and self.selected_local_samples == self.gradient_local_samples |
| ) |
|
|
| @property |
| def finalized(self) -> bool: |
| return self._result is not None |
|
|
| def add( |
| self, |
| generated_features: torch.Tensor, |
| current_teacher_features: torch.Tensor | None = None, |
| *, |
| gradient_selected: bool | None = None, |
| sample_ids: Sequence[str] | None = None, |
| current_teacher_sample_ids: Sequence[str] | None = None, |
| ) -> int: |
| """Add one detached microbatch and return its zero-based cache index.""" |
| if self.finalized: |
| raise RuntimeError( |
| "cannot add features after finalize; commit or discard first") |
| if not isinstance(generated_features, torch.Tensor): |
| raise TypeError("generated_features must be a torch.Tensor") |
| if ( |
| generated_features.ndim != 3 |
| or min(generated_features.shape) <= 0 |
| or not generated_features.is_floating_point() |
| ): |
| raise ValueError( |
| "generated_features must be a non-empty floating [B, N, D] " |
| "tensor") |
| if not bool(torch.isfinite(generated_features).all().item()): |
| raise FloatingPointError( |
| "generated_features contains non-finite values") |
|
|
| if gradient_selected is None: |
| if self.partial_gradient_enabled: |
| raise TypeError( |
| "gradient_selected must be explicit when partial-gradient " |
| "GradCache is enabled") |
| gradient_selected = True |
| gradient_selected = bool(gradient_selected) |
| if self._selection_started and not gradient_selected: |
| raise ValueError( |
| "gradient-selected microbatches must form a contiguous suffix") |
| if gradient_selected: |
| self._selection_started = True |
|
|
| batch_size = int(generated_features.shape[0]) |
| normalized_sample_ids: tuple[str, ...] | None = None |
| normalized_teacher_ids: tuple[str, ...] | None = None |
| if sample_ids is not None: |
| normalized_sample_ids = tuple(str(value) for value in sample_ids) |
| if len(normalized_sample_ids) != batch_size or any( |
| not value for value in normalized_sample_ids |
| ): |
| raise ValueError( |
| "sample_ids must contain one non-empty ID per student") |
| if current_teacher_sample_ids is not None: |
| normalized_teacher_ids = tuple( |
| str(value) for value in current_teacher_sample_ids) |
| if len(normalized_teacher_ids) != batch_size or any( |
| not value for value in normalized_teacher_ids |
| ): |
| raise ValueError( |
| "current_teacher_sample_ids must contain one non-empty ID " |
| "per selected student") |
| if self.require_sample_id_alignment and normalized_sample_ids is None: |
| raise ValueError( |
| "sample_ids are required by this GradCache experiment") |
|
|
| if self.current_teacher_attraction_enabled and gradient_selected: |
| if not isinstance(current_teacher_features, torch.Tensor): |
| raise TypeError( |
| "current_teacher_features must be a torch.Tensor for each " |
| "selected student microbatch") |
| if ( |
| current_teacher_features.ndim != 3 |
| or min(current_teacher_features.shape) <= 0 |
| or not current_teacher_features.is_floating_point() |
| ): |
| raise ValueError( |
| "current_teacher_features must be a non-empty floating " |
| "[B, N, D] tensor") |
| if tuple(current_teacher_features.shape) != tuple( |
| generated_features.shape |
| ): |
| raise ValueError( |
| "current teacher features must exactly match selected " |
| "student [B,N,D]: " |
| f"{tuple(current_teacher_features.shape)} != " |
| f"{tuple(generated_features.shape)}") |
| if current_teacher_features.device != generated_features.device: |
| raise ValueError( |
| "generated and current teacher features must share a device") |
| if not bool(torch.isfinite( |
| current_teacher_features.detach() |
| ).all().item()): |
| raise FloatingPointError( |
| "current_teacher_features contains non-finite values") |
| if self.require_sample_id_alignment and normalized_teacher_ids is None: |
| raise ValueError( |
| "current_teacher_sample_ids are required for selected " |
| "students") |
| if ( |
| normalized_teacher_ids is not None |
| and normalized_sample_ids != normalized_teacher_ids |
| ): |
| raise ValueError( |
| "matched teacher sample IDs must exactly equal selected " |
| "student sample IDs in order") |
| elif current_teacher_features is not None: |
| raise ValueError( |
| "current_teacher_features may only be provided for a selected " |
| "microbatch while exact current-teacher attraction is enabled") |
| elif current_teacher_sample_ids is not None: |
| raise ValueError( |
| "current_teacher_sample_ids may only be provided for a " |
| "selected current-teacher microbatch") |
|
|
| if self.local_samples + batch_size > self.target_local_samples: |
| raise ValueError( |
| "fresh feature cache would exceed target_local_samples: " |
| f"{self.local_samples} + {batch_size} > " |
| f"{self.target_local_samples}") |
| if ( |
| gradient_selected |
| and self.selected_local_samples + batch_size |
| > self.gradient_local_samples |
| ): |
| raise ValueError( |
| "fresh feature cache would exceed gradient_local_samples: " |
| f"{self.selected_local_samples} + {batch_size} > " |
| f"{self.gradient_local_samples}") |
|
|
| sample_shape = ( |
| int(generated_features.shape[1]), |
| int(generated_features.shape[2]), |
| ) |
| if self._sample_shape is None: |
| self._sample_shape = sample_shape |
| self._device = generated_features.device |
| elif ( |
| sample_shape != self._sample_shape |
| or generated_features.device != self._device |
| ): |
| raise ValueError( |
| "all fresh feature batches on a rank must share [N, D] and " |
| "device; expected " |
| f"{self._sample_shape} on {self._device}, got " |
| f"{sample_shape} on {generated_features.device}") |
|
|
| staged = generated_features.detach().to(dtype=torch.float32).clone() |
| self._feature_batches.append(staged) |
| batch_index = len(self._feature_batches) - 1 |
| if gradient_selected: |
| self._gradient_batch_indices.append(batch_index) |
| self._gradient_batch_sizes.append(batch_size) |
| self._gradient_sample_ids.append( |
| normalized_sample_ids or tuple("" for _ in range(batch_size))) |
| if self.current_teacher_attraction_enabled and gradient_selected: |
| staged_teacher = current_teacher_features.detach().to( |
| dtype=torch.float32).clone() |
| self._teacher_feature_batches.append(staged_teacher) |
| self._batch_sizes.append(batch_size) |
| return batch_index |
|
|
| def _validate_distributed_geometry( |
| self, |
| local_features: torch.Tensor, |
| *, |
| expected_samples: int, |
| label: str, |
| ) -> None: |
| return None |
|
|
| def _validate_ready_consensus(self) -> None: |
| return None |
|
|
| def _gather_detached( |
| self, |
| local_features: torch.Tensor, |
| ) -> list[torch.Tensor]: |
| if self._world_size() == 1: |
| return [local_features.detach()] |
| gathered = [ |
| torch.empty_like(local_features) |
| for _ in range(self._world_size()) |
| ] |
| dist.all_gather(gathered, local_features.detach().contiguous()) |
| return gathered |
|
|
| def finalize(self) -> FreshVideoGradCacheResult: |
| """Compute the global loss and stage suffix-only replay cotangents.""" |
| if self._result is not None: |
| return self._result |
| self._validate_ready_consensus() |
|
|
| local_reference = torch.cat(self._feature_batches, dim=0) |
| self._validate_distributed_geometry( |
| local_reference, |
| expected_samples=self.target_local_samples, |
| label="student reference", |
| ) |
| gathered_reference = self._gather_detached(local_reference) |
|
|
| if self.partial_gradient_enabled: |
| selected_batches = [ |
| self._feature_batches[index] |
| for index in self._gradient_batch_indices |
| ] |
| local_gradient = torch.cat(selected_batches, dim=0) |
| self._validate_distributed_geometry( |
| local_gradient, |
| expected_samples=self.gradient_local_samples, |
| label="gradient students", |
| ) |
| gathered_gradient = self._gather_detached(local_gradient) |
| local_leaf = local_gradient.detach().requires_grad_(True) |
| gathered_gradient[self._rank()] = local_leaf |
| global_gradient = torch.cat(gathered_gradient, dim=0) |
| global_reference = torch.cat( |
| gathered_reference, dim=0).detach() |
| gradient_student_indices = None |
| if self.full_student_partial_gradient: |
| |
| |
| |
| |
| |
| suffix_start = ( |
| self.target_local_samples |
| - self.gradient_local_samples |
| ) |
| gradient_student_indices = torch.cat([ |
| torch.arange( |
| rank * self.target_local_samples + suffix_start, |
| (rank + 1) * self.target_local_samples, |
| device=global_reference.device, |
| dtype=torch.long, |
| ) |
| for rank in range(self._world_size()) |
| ]) |
| else: |
| local_leaf = local_reference.detach().requires_grad_(True) |
| gathered_reference[self._rank()] = local_leaf |
| global_gradient = torch.cat(gathered_reference, dim=0) |
| global_reference = None |
|
|
| global_teacher = None |
| if self.current_teacher_attraction_enabled: |
| if len(self._teacher_feature_batches) != len( |
| self._gradient_batch_indices |
| ): |
| raise RuntimeError( |
| "matched teacher and selected student GradCache batch " |
| "counts differ") |
| local_teacher = torch.cat( |
| self._teacher_feature_batches, dim=0) |
| if tuple(local_teacher.shape) != tuple( |
| local_gradient.shape |
| if self.partial_gradient_enabled else local_reference.shape |
| ): |
| raise RuntimeError( |
| "matched teacher and gradient student shapes differ") |
| self._validate_distributed_geometry( |
| local_teacher, |
| expected_samples=self.gradient_local_samples, |
| label="current teacher", |
| ) |
| gathered_teacher = self._gather_detached(local_teacher) |
| global_teacher = torch.cat( |
| gathered_teacher, dim=0).detach().to(dtype=torch.float32) |
|
|
| if self.partial_gradient_enabled: |
| if isinstance(self.loss_module, MultiLayerVideoNystromMMD): |
| raise TypeError( |
| "partial-gradient GradCache currently requires a single " |
| "VideoNystromMMD layer") |
| partial_loss_kwargs = { |
| "gradient_students": global_gradient, |
| "student_reference": global_reference, |
| "current_teacher": global_teacher, |
| } |
| if self.full_student_partial_gradient: |
| partial_loss_kwargs["gradient_student_indices"] = ( |
| gradient_student_indices) |
| raw_loss, metadata = self.loss_module(**partial_loss_kwargs) |
| elif self.current_teacher_attraction_enabled: |
| raw_loss, metadata = self.loss_module( |
| global_gradient, current_teacher=global_teacher) |
| else: |
| raw_loss, metadata = self.loss_module(global_gradient) |
| if not isinstance(raw_loss, torch.Tensor) or raw_loss.ndim != 0: |
| raise TypeError( |
| "VideoNystromMMD must return a scalar tensor loss") |
| if not self._all_ranks_true( |
| bool(torch.isfinite(raw_loss.detach()).item()), |
| raw_loss.device, |
| ): |
| raise FloatingPointError( |
| "VideoNystromMMD returned a non-finite loss on at least one " |
| "rank") |
| optimized_loss = raw_loss * self.loss_scale |
|
|
| (local_cotangent,) = torch.autograd.grad( |
| optimized_loss, |
| local_leaf, |
| retain_graph=False, |
| create_graph=False, |
| ) |
| local_cotangent = local_cotangent.detach() |
| if not self._all_ranks_true( |
| bool(torch.isfinite(local_cotangent).all().item()), |
| local_cotangent.device, |
| ): |
| raise FloatingPointError( |
| "VideoNystromMMD returned a non-finite feature cotangent on " |
| "at least one rank") |
|
|
| world_size = self._world_size() |
| local_parts = tuple( |
| value.contiguous() |
| for value in local_cotangent.split( |
| self._gradient_batch_sizes, dim=0) |
| ) |
| fsdp_parts = tuple( |
| (value * float(world_size)).contiguous() |
| for value in local_parts |
| ) |
| global_samples = self.target_local_samples * world_size |
| gradient_global_samples = self.gradient_local_samples * world_size |
| teacher_local_samples = ( |
| self.gradient_local_samples |
| if self.current_teacher_attraction_enabled else 0 |
| ) |
| teacher_global_samples = teacher_local_samples * world_size |
| kss_denominator = global_samples * ( |
| global_samples |
| if self.full_student_partial_gradient |
| else gradient_global_samples |
| ) |
| detached_metadata = dict(_detach_metadata(metadata)) |
| detached_metadata.update({ |
| "raw_loss": raw_loss.detach(), |
| "optimized_loss": optimized_loss.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_samples, dtype=torch.long), |
| "gradcache_kss_denominator": raw_loss.new_tensor( |
| kss_denominator, dtype=torch.long), |
| "gradcache_partial_gradient_enabled": raw_loss.new_tensor( |
| int(self.partial_gradient_enabled), dtype=torch.long), |
| "gradcache_full_student_partial_gradient": raw_loss.new_tensor( |
| int(self.full_student_partial_gradient), 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)), |
| }) |
| detached_metadata.update({ |
| "gradcache_current_teacher_local_samples": raw_loss.new_tensor( |
| teacher_local_samples, dtype=torch.long), |
| "gradcache_current_teacher_global_samples": raw_loss.new_tensor( |
| teacher_global_samples, dtype=torch.long), |
| }) |
| core_global_samples = detached_metadata.get("global_samples") |
| if ( |
| core_global_samples is not None |
| and int(torch.as_tensor(core_global_samples).item()) |
| != global_samples |
| ): |
| raise ValueError( |
| "VideoNystromMMD global_samples metadata disagrees with the " |
| f"fresh GradCache: {int(torch.as_tensor(core_global_samples))} " |
| f"!= {global_samples}") |
| core_gradient_samples = detached_metadata.get( |
| "gradient_student_samples") |
| if ( |
| self.partial_gradient_enabled |
| and core_gradient_samples is not None |
| and int(torch.as_tensor(core_gradient_samples).item()) |
| != gradient_global_samples |
| ): |
| raise ValueError( |
| "VideoNystromMMD gradient_student_samples metadata disagrees " |
| "with the fresh GradCache") |
| core_repulsion_denominator = detached_metadata.get( |
| "repulsion_denominator") |
| if ( |
| self.partial_gradient_enabled |
| and core_repulsion_denominator is not None |
| and int(torch.as_tensor(core_repulsion_denominator).item()) |
| != kss_denominator |
| ): |
| raise ValueError( |
| "VideoNystromMMD Kss denominator disagrees with the fresh " |
| f"GradCache: expected {kss_denominator}") |
|
|
| self._result = FreshVideoGradCacheResult( |
| loss=raw_loss.detach(), |
| optimized_loss=optimized_loss.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_samples, |
| current_teacher_local_samples=teacher_local_samples, |
| current_teacher_global_samples=teacher_global_samples, |
| world_size=world_size, |
| ) |
| return self._result |
|
|
| def _clear(self) -> None: |
| self._feature_batches.clear() |
| self._teacher_feature_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._sample_shape = None |
| self._device = None |
| self._result = None |
|
|
| def commit(self) -> None: |
| """Clear a finalized cache after a successful optimizer update.""" |
| if self._result is None: |
| raise RuntimeError( |
| "cannot commit a fresh feature cache before finalize") |
| self._clear() |
|
|
| def discard(self) -> None: |
| """Clear collected/finalized data after a skipped or failed update.""" |
| self._clear() |
|
|
