| """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") |
|
|
| |
| |
| 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 = {} |
| |
| |
| |
| 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) |
| |
| |
| |
| 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() |
|
|
| |
| |
| |
| 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, |
| ) |
| |
| |
| 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_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 |
| |
| 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() |
|
|
|
|
| |
| MultiRepresentationFreshVideoGradCache = FeatureGradientReplay |
| MultiRepresentationFreshVideoGradCacheResult = FeatureGradientReplayResult |
|
|
| __all__ = [ |
| "FeatureGradientReplay", |
| "FeatureGradientReplayResult", |
| "MultiRepresentationFreshVideoGradCache", |
| "MultiRepresentationFreshVideoGradCacheResult", |
| ] |
|
|
