Buckets:
| """Optimizer-step semantics for fixed-budget OPD logical updates.""" | |
| from __future__ import annotations | |
| import math | |
| from dataclasses import dataclass, field | |
| from collections.abc import Iterable | |
| from typing import Callable | |
| import torch | |
| from shared.training import assert_finite_training_state | |
| class OptimizerUpdateResult: | |
| successful: bool | |
| zero_gradient: bool | |
| grad_norm: float | |
| gradient_statistics: dict[str, float] = field(default_factory=dict) | |
| def apply_optimizer_update( | |
| *, | |
| losses: Iterable[torch.Tensor], | |
| parameters: list[torch.nn.Parameter], | |
| model, | |
| optimizer, | |
| scheduler, | |
| ema_teacher, | |
| max_grad_norm: float, | |
| synchronize_gradients: Callable[[list[torch.nn.Parameter]], None] | None = None, | |
| synchronize_finite: Callable[[bool], bool] | None = None, | |
| gradient_diagnostics=None, | |
| ) -> OptimizerUpdateResult: | |
| """Apply one logical budget step, including the accepted=0 no-op case.""" | |
| if max_grad_norm <= 0 or not parameters: | |
| raise ValueError("optimizer update requires parameters and positive clipping") | |
| optimizer.zero_grad(set_to_none=True) | |
| finite_loss = True | |
| for loss in losses: | |
| if loss.numel() != 1: | |
| raise ValueError("each OPD/replay loss term must be scalar") | |
| finite_loss = finite_loss and bool(torch.isfinite(loss.detach())) | |
| loss.backward() | |
| if synchronize_gradients is not None: | |
| synchronize_gradients(parameters) | |
| if synchronize_finite is not None: | |
| finite_loss = synchronize_finite(finite_loss) | |
| gradient_statistics = gradient_diagnostics() if gradient_diagnostics is not None else {} | |
| gradients = [parameter.grad for parameter in parameters if parameter.grad is not None] | |
| zero_gradient = not gradients | |
| grad_norm = torch.nn.utils.clip_grad_norm_(parameters, max_grad_norm) | |
| grad_norm_value = float(grad_norm.detach().float().cpu()) | |
| successful = bool(finite_loss and math.isfinite(grad_norm_value)) | |
| if successful: | |
| optimizer.step() | |
| assert_finite_training_state(parameters, optimizer) | |
| scheduler.step() | |
| if ema_teacher is not None: | |
| ema_teacher.update(model) | |
| if not ema_teacher.is_finite(): | |
| raise FloatingPointError("EMA teacher became non-finite") | |
| optimizer.zero_grad(set_to_none=True) | |
| return OptimizerUpdateResult(successful, zero_gradient, grad_norm_value, gradient_statistics) | |
| __all__ = ["OptimizerUpdateResult", "apply_optimizer_update"] | |
Xet Storage Details
- Size:
- 2.52 kB
- Xet hash:
- 53959d67ea11fa1778eefb95df869b676a46435ae7c04a2f791717225d6052b1
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.