| # mypy: allow-untyped-decorators |
| # mypy: allow-untyped-defs |
| from typing import cast, Dict, List, Optional, Tuple, TYPE_CHECKING, Union |
| |
| import torch |
| from torch import Tensor |
| |
| from .optimizer import ( |
| _disable_dynamo_if_unsupported, |
| _get_scalar_dtype, |
| _maximize_doc, |
| Optimizer, |
| ParamsT, |
| TensorListList, |
| ) |
| |
| |
| __all__ = ["Adafactor", "adafactor"] |
| |
| |
| class Adafactor(Optimizer): |
| def __init__( |
| self, |
| params: ParamsT, |
| lr: Union[float, Tensor] = 1e-2, |
| beta2_decay: float = -0.8, |
| eps: Tuple[Optional[float], float] = (None, 1e-3), |
| d: float = 1.0, |
| weight_decay: float = 0.0, |
| *, |
| foreach: Optional[bool] = None, |
| maximize: bool = False, |
| ): |
| if isinstance(lr, Tensor) and lr.numel() != 1: |
| raise ValueError("Tensor lr must be 1-element") |
| if not 0.0 <= lr: |
| raise ValueError(f"Learning rate should be >= 0 but is: {lr}") |
| if not 0.0 >= beta2_decay: |
| raise ValueError(f"beta2_decay should be <= 0 but is: {beta2_decay}") |
| if eps[0] is not None and not 0.0 <= eps[0]: |
| raise ValueError(f"epsilon1 should be >= 0 but is: {eps[0]}") |
| if not 0.0 <= eps[1]: |
| raise ValueError(f"epsilon2 should be >= 0 but is: {eps[1]}") |
| if not 1.0 <= d: |
| raise ValueError(f"Clipping threshold d should be >= 1 but is: {d}") |
| if not 0.0 <= weight_decay: |
| raise ValueError(f"weight_decay should be >= 0 but is: {weight_decay}") |
| defaults = dict( |
| lr=lr, |
| beta2_decay=beta2_decay, |
| eps=eps, |
| d=d, |
| weight_decay=weight_decay, |
| foreach=foreach, |
| maximize=maximize, |
| ) |
| super().__init__(params, defaults) |
| |
| def __setstate__(self, state): |
| super().__setstate__(state) |
| for group in self.param_groups: |
| group.setdefault("foreach", None) |
| for p in group["params"]: |
| p_state = self.state.get(p, []) |
| if len(p_state) != 0 and not torch.is_tensor(p_state["step"]): |
| step_val = float(p_state["step"]) |
| p_state["step"] = torch.tensor(step_val, dtype=_get_scalar_dtype()) |
| |
| def _init_group( |
| self, |
| group, |
| params_with_grad, |
| grads, |
| row_vars, |
| col_vars, |
| variances, |
| state_steps, |
| ): |
| for p in group["params"]: |
| if p.grad is None: |
| continue |
| if torch.is_complex(p): |
| raise RuntimeError("Adafactor does not support complex parameters") |
| if p.grad.is_sparse: |
| raise RuntimeError("Adafactor does not support sparse gradients") |
| |
| params_with_grad.append(p) |
| grads.append(p.grad) |
| |
| state = self.state[p] |
| |
| # State initialization |
| if len(state) == 0: |
| # note(crcrpar): Deliberately host `step` on CPU if both capturable and fused are off. |
| # This is because kernel launches are costly on CUDA and XLA. |
| state["step"] = torch.tensor(0.0, dtype=_get_scalar_dtype()) |
| |
| if p.grad.dim() > 1: |
| row_shape = list(p.grad.shape) |
| row_shape[-1] = 1 |
| # Row factor of variance, NOT the same shape as grads (will be reduced along last dim) |
| state["row_var"] = p.grad.new_zeros(row_shape) |
| |
| col_shape = list(p.grad.shape) |
| col_shape[-2] = 1 |
| # Col factor of variance, NOT the same shape as grads (will be reduced along penultimate dim) |
| state["col_var"] = p.grad.new_zeros(col_shape) |
| else: |
| state["variance"] = torch.zeros_like( |
| p.grad, memory_format=torch.preserve_format |
| ) |
| |
| row_vars.append(state.get("row_var", None)) |
| col_vars.append(state.get("col_var", None)) |
| variances.append(state.get("variance", None)) |
| state_steps.append(state["step"]) |
| return False # has_complex |
| |
| @torch.no_grad() |
| def step(self, closure=None): |
| r"""Perform a single optimization step. |
| |
| Args: |
| closure (Callable, optional): A closure that reevaluates the model |
| and returns the loss. |
| """ |
| self._cuda_graph_capture_health_check() |
| |
| loss = None |
| if closure is not None: |
| with torch.enable_grad(): |
| loss = closure() |
| |
| for group in self.param_groups: |
| params_with_grad: List[Tensor] = [] |
| grads: List[Tensor] = [] |
| row_vars: List[Optional[Tensor]] = [] |
| col_vars: List[Optional[Tensor]] = [] |
| variances: List[Optional[Tensor]] = [] |
| state_steps: List[Tensor] = [] |
| eps1, eps2 = group["eps"] |
| |
| has_complex = self._init_group( |
| group, |
| params_with_grad, |
| grads, |
| row_vars, |
| col_vars, |
| variances, |
| state_steps, |
| ) |
| |
| adafactor( |
| params_with_grad, |
| grads, |
| row_vars, |
| col_vars, |
| variances, |
| state_steps, |
| d=group["d"], |
| lr=group["lr"], |
| beta2_decay=group["beta2_decay"], |
| weight_decay=group["weight_decay"], |
| eps1=eps1, |
| eps2=eps2, |
| foreach=group["foreach"], |
| maximize=group["maximize"], |
| grad_scale=getattr(self, "grad_scale", None), |
| found_inf=getattr(self, "found_inf", None), |
| has_complex=has_complex, |
| ) |
| |
| return loss |
| |
| |
| Adafactor.__doc__ = ( |
| r"""Implements Adafactor algorithm. |
| |
| .. math:: |
| \begin{aligned} |
| &\rule{110mm}{0.4pt} \\ |
| &\textbf{input} : \gamma \text{(lr)}, \: \tau |
| \text{(}\beta_2\text{ decay)}, \: \theta_0 \text{(params)}, \: f(\theta) \text{(objective)}, \\ |
| &\hspace{15mm} \: \epsilon_1, \epsilon_2 \text{ (epsilons)}, \: d \text{(clipping threshold)}, \\ |
| &\hspace{15mm} \: \lambda \text{(weight decay)}, |
| \: \textit{maximize} \\ |
| &\textbf{initialize} : \: R_0 \leftarrow 0 \text{ (second moment row factor)}, \\ |
| &\hspace{23mm} \: C_0 \leftarrow 0 \text{ (second moment col factor)}, \\ |
| &\hspace{23mm} \: \widehat{V}_0 \leftarrow 0 \text{ (second moment for vectors)} \\[-1.ex] |
| &\rule{110mm}{0.4pt} \\ |
| &\textbf{for} \: t=1 \: \textbf{to} \: \ldots \: \textbf{do} \\ |
| |
| &\hspace{5mm}\textbf{if} \: \textit{maximize}: \\ |
| &\hspace{10mm}G_t \leftarrow -\nabla_{\theta} f_t (\theta_{t-1}) \\ |
| &\hspace{5mm}\textbf{else} \\ |
| &\hspace{10mm}G_t \leftarrow \nabla_{\theta} f_t (\theta_{t-1}) \\ |
| &\hspace{5mm}\widehat{\beta}_{2_t} \leftarrow 1 - t^{\tau} \\ |
| &\hspace{5mm}\rho_t \leftarrow min(lr, \frac{1}{\sqrt{t}}) \\ |
| &\hspace{5mm}\alpha_t \leftarrow max(\epsilon_2, |
| \text{RMS}(\theta_{t-1}))\rho_t \\ |
| &\hspace{5mm}\theta_t \leftarrow \theta_{t-1} - \gamma \lambda \theta_{t-1} \\ |
| &\hspace{5mm}\textbf{if} \: \text{dim}(G_t) > 1: \\ |
| &\hspace{10mm}R_t \leftarrow \widehat{\beta}_{2_t}R_{t-1}+ |
| (1-\widehat{\beta}_{2_t})(G_t \odot G_t) \cdot 1_m \\ |
| &\hspace{10mm}C_t \leftarrow \widehat{\beta}_{2_t}C_{t-1}+ |
| (1-\widehat{\beta}_{2_t}) 1^\top_n \cdot (G_t \odot G_t) \\ |
| &\hspace{10mm}\widehat{V}_t \leftarrow |
| \frac{R_t \cdot C_t}{max(1^\top_n \cdot R_t, \epsilon_1)} \\ |
| &\hspace{5mm}\textbf{else} \\ |
| &\hspace{10mm}\widehat{V}_t \leftarrow \widehat{\beta}_{2_t}\widehat{V}_{t-1}+ |
| (1-\widehat{\beta}_{2_t}) \cdot (G_t \odot G_t) \\ |
| &\hspace{5mm}U_t \leftarrow |
| \frac{G_t}{max(\sqrt{\widehat{V}_t}, \epsilon_1)} \\ |
| &\hspace{5mm}\widehat{U}_t \leftarrow \frac{U_t}{max(1, \frac{\text{RMS}(U_t)}{d})} \\ |
| &\hspace{5mm}\theta_t \leftarrow \theta_{t-1} - \alpha_t \widehat{U}_t \\ |
| |
| &\rule{110mm}{0.4pt} \\[-1.ex] |
| &\bf{return} \: \theta_t \\[-1.ex] |
| &\rule{110mm}{0.4pt} \\[-1.ex] |
| \end{aligned} |
| |
| For further details regarding the algorithm we refer to `Adafactor: Adaptive Learning Rates with Sublinear Memory Cost`_. |
| """ |
| + rf""" |
| Args: |
| params (iterable): iterable of parameters to optimize or dicts defining |
| parameter groups |
| lr (float, Tensor, optional): unlike other optimizers, Adafactor does not require a |
| learning rate, and Shazeer, Noam, and Mitchell Stern do not use lr at all. |
| Deviating from the paper, this implementation uses lr for applying weight |
| decay and as the maximum value for relative step size rho_t. Note that in |
| the paper, a constant of 0.01 is used as the maximum value for relative |
| step size, and so we set 0.01 as the default value. (default: 1e-2) |
| beta2_decay (float, optional): the decay rate of beta2. beta2 standardly refers |
| to the coefficient used for computing the running average of the gradient |
| squared. (default: -0.8) |
| eps (Tuple[float, float], optional): epsilon1 is the term added to the denominator |
| of the update calculation to improve numerical stability. This use of epsilon1 |
| deviates from the algorithm written in the paper! See note below for more details. |
| epsilon2 is the term used to avoid having too small a weight update when applying |
| parameter scaling. (default: (None, 1e-3)) |
| d (float, optional): the clipping threshold, used to avoid larger-than-desired |
| updates. |
| weight_decay (float, optional): weight decay coefficient (default: 1e-2) |
| foreach (bool, optional): whether foreach implementation of optimizer is used. Note |
| that the foreach implementation uses ~ sizeof(params) more peak memory than the |
| for-loop version due to the intermediates being a tensorlist vs just one tensor. |
| As Adafactor is commonly used when memory is prohibitive, Adafactor will default |
| to the slower single tensor for-loop implementation unless this flag is explicitly |
| True. This behavior is contrary to other optimizers, which will attempt defaulting |
| to foreach on CUDA for faster runtime. (default: None) |
| {_maximize_doc}""" |
| + r""" |
| .. Note:: |
| The implementation of Adafactor subtly differs from Shazeer, Noam, and Mitchell Stern |
| and implementations in some other frameworks with its use of learning rate and |
| :math:`\epsilon_1`. |
| |
| Regarding the learning rate hyperparameter: Shazeer, Noam, and Mitchell Stern do not |
| use lr at all, as the stated algorithm uses :math:`\rho_t` and update clipping to |
| affect the step size. |
| |
| This implementation allows `lr` to influence the maximum value for :math:`\rho_t`: |
| |
| .. math:: |
| \begin{aligned} |
| &\hspace{5mm}\rho_t \leftarrow min(lr, \frac{1}{\sqrt{t}}) |
| \end{aligned} |
| |
| This differs from Shazeer, Noam, and Mitchell Stern, who use a constant of 0.01 as |
| the maximum value of :math:`\rho_t` |
| |
| .. math:: |
| \begin{aligned} |
| &\hspace{5mm}\rho_t \leftarrow min(0.01, \frac{1}{\sqrt{t}}) |
| \end{aligned} |
| |
| Shazeer, Noam, and Mitchell Stern do not enforce an opinion on how weight decay should |
| be computed, and so we use the learning rate as a coefficient for decoupled weight |
| decay, similar to what is suggested in `Decoupled Weight Decay Regularization`_. |
| |
| Regarding the use of :math:`\epsilon_1`: The implementation attempts to replicate the |
| presumed intention of Shazeer, Noam, and Mitchell Stern to use :math:`\epsilon_1` as |
| a stabilizing term when the squared gradient becomes small. |
| |
| This stabilization can be written as |
| |
| .. math:: |
| \begin{aligned} |
| &\hspace{5mm}R_t \leftarrow \widehat{\beta}_{2_t}R_{t-1}+ |
| (1-\widehat{\beta}_{2_t})(G_t \odot G_t + 1_n \cdot 1^\top_m) \cdot 1_m \\ |
| &\hspace{5mm}C_t \leftarrow \widehat{\beta}_{2_t}C_{t-1}+ |
| (1-\widehat{\beta}_{2_t}) 1^\top_n \cdot (G_t \odot G_t + 1_n \cdot 1^\top_m) \\ |
| &\hspace{5mm}\widehat{V}_t \leftarrow |
| \frac{R_t \cdot C_t}{max(1^\top_n \cdot R_t, \epsilon_1)} \\ |
| &\hspace{5mm}U_t \leftarrow \frac{G_t}{max(\sqrt{\widehat{V}_t}, \epsilon_1)} \\ |
| \end{aligned} |
| |
| where the row and column factors of gradient squared :math:`R_t` and :math:`C_t` |
| are left alone, and we apply :math:`\epsilon_1` at the final calculation of |
| the variance estimate :math:`\widehat{V}_t` and for the update :math:`U_t`. |
| |
| This is in contrast to Shazeer, Noam, and Mitchell Stern and other frameworks which |
| apply :math:`\epsilon_1` to both row and column factors of the squared gradient, but |
| not in the calculations after: |
| |
| .. math:: |
| \begin{aligned} |
| &\hspace{5mm}R_t \leftarrow \widehat{\beta}_{2_t}R_{t-1}+ |
| (1-\widehat{\beta}_{2_t})(G_t \odot G_t + \epsilon_1 1_n \cdot 1^\top_m) \cdot 1_m \\ |
| &\hspace{5mm}C_t \leftarrow \widehat{\beta}_{2_t}C_{t-1}+ |
| (1-\widehat{\beta}_{2_t}) 1^\top_n \cdot (G_t \odot G_t + \epsilon_1 1_n \cdot 1^\top_m) \\ |
| &\hspace{5mm}\widehat{V}_t \leftarrow \frac{R_t \cdot C_t}{1^\top_n \cdot R_t} \\ |
| &\hspace{5mm}U_t \leftarrow \frac{G_t}{\sqrt{\widehat{V}_t}} \\ |
| \end{aligned} |
| |
| |
| .. _Adafactor\: Adaptive Learning Rates with Sublinear Memory Cost: |
| https://arxiv.org/pdf/1804.04235 |
| .. _Decoupled Weight Decay Regularization: |
| https://arxiv.org/abs/1711.05101 |
| """ |
| ) |
| |
| |
| def _single_tensor_adafactor( |
| params: List[Tensor], |
| grads: List[Tensor], |
| # If grad is 1-dimensional (aka a vector), there is no factorization necessary |
| # so row_var and col_var will be None while variance will be filled. |
| # Contrarily, for a grad with multiple dimensions, we will factor along the last |
| # 2 dimensions, and so row_var and col_var will be filled and variance will be None. |
| row_vars: List[Optional[Tensor]], |
| col_vars: List[Optional[Tensor]], |
| variances: List[Optional[Tensor]], |
| state_steps: List[Tensor], |
| grad_scale: Optional[Tensor], |
| found_inf: Optional[Tensor], |
| *, |
| d: float, |
| lr: Union[Tensor, float], |
| beta2_decay: float, |
| weight_decay: float, |
| eps1: Optional[float], |
| eps2: float, |
| maximize: bool, |
| has_complex: bool, |
| ): |
| assert ( |
| grad_scale is None and found_inf is None |
| ), "Grad scaling should occur outside of optimizer.step()" |
| |
| if torch.jit.is_scripting(): |
| # this assert is due to JIT being dumb and not realizing that the ops below |
| # have overloads to handle both float and Tensor lrs, so we just assert it's |
| # a float since most people using JIT are using floats |
| assert isinstance(lr, float) |
| |
| for i, param in enumerate(params): |
| grad = grads[i] if not maximize else -grads[i] |
| step_t = state_steps[i] |
| row_var = row_vars[i] |
| col_var = col_vars[i] |
| variance = variances[i] |
| if eps1 is None: |
| eps1 = torch.finfo(param.dtype).eps |
| |
| # update step |
| step_t += 1 |
| step_float = step_t.item() |
| |
| one_minus_beta2_t = step_float**beta2_decay |
| rho_t = min(lr, 1 / (step_float**0.5)) |
| alpha = max(eps2, param.norm(2).item() / (param.numel() ** 0.5)) * rho_t |
| |
| # Perform stepweight decay |
| if weight_decay != 0: |
| param.mul_(1 - lr * weight_decay) |
| |
| if grad.dim() > 1: |
| assert ( |
| row_var is not None and col_var is not None |
| ), "row_var and col_var should be defined when grad is multidimensional" |
| # same as (g * g).mean(dim=-1) w/o materializing an intermediate size g |
| row_mean = ( |
| torch.norm(grad, dim=-1, keepdim=True).square_().div_(grad.size(-1)) |
| ) |
| row_var.lerp_(row_mean, one_minus_beta2_t) |
| # same as (g * g).mean(dim=-2) w/o materializing an intermediate size g |
| col_mean = ( |
| torch.norm(grad, dim=-2, keepdim=True).square_().div_(grad.size(-2)) |
| ) |
| col_var.lerp_(col_mean, one_minus_beta2_t) |
| var_estimate = row_var @ col_var |
| var_estimate.div_(row_var.mean(dim=-2, keepdim=True).clamp_(min=eps1)) |
| else: |
| assert ( |
| variance is not None |
| ), "variance should be defined when grad is a vector" |
| grad_squared = grad * grad |
| variance.lerp_(grad_squared, one_minus_beta2_t) |
| # avoid writing into variance during update |
| var_estimate = variance.clone() |
| |
| # square the eps1 as we sqrt after to keep eps1's magnitude |
| update = var_estimate.clamp_(min=eps1 * eps1).rsqrt_() |
| update.mul_(grad) |
| denom = max(1.0, update.norm(2).item() / ((update.numel() ** 0.5) * d)) |
| param.add_(update, alpha=-alpha / denom) |
| |
| |
| def _group_tensors_by_device_dtype_and_is_multidim( |
| tensorlists: TensorListList, |
| ) -> Dict[ |
| Tuple[Optional[torch.device], Optional[torch.dtype], bool], |
| List[List[Optional[Tensor]]], |
| ]: |
| """Groups tensors by device, dtype, AND multidimensionality -- whether the tensor |
| has multiple dims or just one dim (is a vector). This allows the foreach impl of |
| Adafactor to assume that every group of params will either be factored or not.""" |
| grouped_tensors = Optimizer._group_tensors_by_device_and_dtype(tensorlists) |
| ultra_grouped_tensors: Dict[ |
| Tuple[Optional[torch.device], Optional[torch.dtype], bool], |
| List[List[Optional[Tensor]]], |
| ] = {} |
| for (device, dtype), (tensorlists, _) in grouped_tensors.items(): |
| matrix_key = (device, dtype, True) |
| vector_key = (device, dtype, False) |
| |
| # assumes grad is the second tensorlist |
| for j, tensor in enumerate(tensorlists[1]): |
| assert tensor is not None, "grad should not be None" |
| if tensor.dim() > 1: |
| if matrix_key not in ultra_grouped_tensors: |
| ultra_grouped_tensors[matrix_key] = [[] for _ in tensorlists] |
| for i in range(len(tensorlists)): |
| ultra_grouped_tensors[matrix_key][i].append(tensorlists[i][j]) |
| else: |
| if vector_key not in ultra_grouped_tensors: |
| ultra_grouped_tensors[vector_key] = [[] for _ in tensorlists] |
| for i in range(len(tensorlists)): |
| ultra_grouped_tensors[vector_key][i].append(tensorlists[i][j]) |
| return ultra_grouped_tensors |
| |
| |
| def _multi_tensor_adafactor( |
| params: List[Tensor], |
| grads: List[Tensor], |
| # If grad is 1-dimensional (aka a vector), there is no factorization necessary |
| # so row_var and col_var will be None while variance will be filled. |
| # Contrarily, for a grad with multiple dimensions, we will factor along the last |
| # 2 dimensions, and so row_var and col_var will be filled and variance will be None. |
| row_vars: List[Optional[Tensor]], |
| col_vars: List[Optional[Tensor]], |
| variances: List[Optional[Tensor]], |
| state_steps: List[Tensor], |
| grad_scale: Optional[Tensor], |
| found_inf: Optional[Tensor], |
| *, |
| d: float, |
| lr: Union[Tensor, float], |
| beta2_decay: float, |
| weight_decay: float, |
| eps1: Optional[float], |
| eps2: float, |
| maximize: bool, |
| has_complex: bool, |
| ): |
| if len(params) == 0: |
| return |
| |
| assert ( |
| grad_scale is None and found_inf is None |
| ), "Grad scaling should occur outside of optimizer.step()" |
| |
| grouped_tensors = _group_tensors_by_device_dtype_and_is_multidim( |
| [params, grads, row_vars, col_vars, variances, state_steps] # type: ignore[list-item] |
| ) |
| for (_, dtype, is_multidim), ( |
| ( |
| device_params_, |
| device_grads_, |
| device_row_vars_, |
| device_col_vars_, |
| device_variances_, |
| device_state_steps_, |
| ) |
| ) in grouped_tensors.items(): |
| device_params = cast(List[Tensor], device_params_) |
| device_grads = cast(List[Tensor], device_grads_) |
| device_state_steps = cast(List[Tensor], device_state_steps_) |
| if eps1 is None: |
| assert ( |
| dtype is not None |
| ), "dtype is needed to compute eps1 when eps1 is unset" |
| eps1 = torch.finfo(dtype).eps |
| |
| if TYPE_CHECKING: |
| assert device_state_steps[0] is not None |
| |
| if maximize: |
| device_grads = torch._foreach_neg(device_grads) # type: ignore[assignment] |
| |
| # Update steps |
| # If steps are on CPU, foreach will fall back to the slow path, which is a for-loop calling t.add(1) over |
| # and over. 1 will then be wrapped into a Tensor over and over again, which is slower than if we just |
| # wrapped it once now. The alpha is required to assure we go to the right overload. |
| if not torch._utils.is_compiling() and device_state_steps[0].is_cpu: |
| torch._foreach_add_( |
| device_state_steps, torch.tensor(1.0, device="cpu"), alpha=1.0 |
| ) |
| else: |
| torch._foreach_add_(device_state_steps, 1.0) |
| |
| one_minus_beta2_ts = [] |
| beta2_ts = [] |
| rho_ts = [] |
| for s in device_state_steps: |
| one_minus_beta2_ts.append(s.item() ** beta2_decay) |
| beta2_ts.append(1 - s.item() ** beta2_decay) |
| rho_ts.append(min(lr, 1 / (s.item() ** 0.5))) |
| |
| alphas = [ |
| max(eps2, p.norm(2).item() / (p.numel() ** 0.5)) * r |
| for p, r in zip(device_params, rho_ts) |
| ] |
| |
| # Perform stepweight decay |
| if weight_decay != 0: |
| torch._foreach_mul_(device_params, 1 - lr * weight_decay) |
| |
| if is_multidim: |
| device_row_vars = cast(List[Tensor], device_row_vars_) |
| device_col_vars = cast(List[Tensor], device_col_vars_) |
| assert ( |
| device_row_vars[0] is not None and device_col_vars[0] is not None |
| ), "row_var and col_var should be defined when grad is multidimensional" |
| # same as (g * g).mean(dim=-1) w/o materializing an intermediate size g |
| row_means = [ |
| torch.norm(grad, dim=-1, keepdim=True) for grad in device_grads |
| ] |
| torch._foreach_mul_(row_means, row_means) |
| torch._foreach_div_(row_means, [grad.size(-1) for grad in device_grads]) |
| torch._foreach_mul_(device_row_vars, beta2_ts) |
| torch._foreach_mul_(row_means, one_minus_beta2_ts) |
| torch._foreach_add_(device_row_vars, row_means) |
| del row_means |
| |
| # same as (g * g).mean(dim=-2) w/o materializing an intermediate size g |
| col_means = [ |
| torch.norm(grad, dim=-2, keepdim=True) for grad in device_grads |
| ] |
| torch._foreach_mul_(col_means, col_means) |
| torch._foreach_div_(col_means, [grad.size(-2) for grad in device_grads]) |
| torch._foreach_mul_(device_col_vars, beta2_ts) |
| torch._foreach_mul_(col_means, one_minus_beta2_ts) |
| torch._foreach_add_(device_col_vars, col_means) |
| del col_means |
| |
| var_estimates = [ |
| row_var @ col_var |
| for row_var, col_var in zip(device_row_vars, device_col_vars) |
| ] |
| row_var_means = [ |
| row_var.mean(dim=-2, keepdim=True) for row_var in device_row_vars |
| ] |
| torch._foreach_clamp_min_(row_var_means, eps1) |
| torch._foreach_div_(var_estimates, row_var_means) |
| del row_var_means |
| else: |
| device_variances = cast(List[Tensor], device_variances_) |
| assert ( |
| device_variances[0] is not None |
| ), "variance should be defined when grad is a vector" |
| |
| grads_squared = torch._foreach_mul(device_grads, device_grads) |
| torch._foreach_mul_(device_variances, beta2_ts) |
| torch._foreach_mul_(grads_squared, one_minus_beta2_ts) |
| torch._foreach_add_(device_variances, grads_squared) |
| del grads_squared |
| |
| # avoid writing into variance during update |
| var_estimates = [v.clone() for v in device_variances] |
| |
| # square the eps1 as we sqrt after to keep eps1's magnitude |
| torch._foreach_clamp_min_(var_estimates, eps1 * eps1) |
| torch._foreach_sqrt_(var_estimates) |
| torch._foreach_reciprocal_(var_estimates) |
| torch._foreach_mul_(var_estimates, device_grads) |
| updates = var_estimates |
| |
| alphas = [ |
| -a / (max(1.0, update.norm(2).item() / ((update.numel() ** 0.5) * d))) |
| for a, update in zip(alphas, updates) |
| ] |
| torch._foreach_mul_(updates, alphas) |
| torch._foreach_add_(device_params, updates) |
| |
| |
| @_disable_dynamo_if_unsupported(single_tensor_fn=_single_tensor_adafactor) |
| def adafactor( |
| params: List[Tensor], |
| grads: List[Tensor], |
| row_vars: List[Optional[Tensor]], |
| col_vars: List[Optional[Tensor]], |
| variances: List[Optional[Tensor]], |
| state_steps: List[Tensor], |
| # kwonly args with defaults are not supported by functions compiled with torchscript issue #70627 |
| # setting this as kwarg for now as functional API is compiled by torch/distributed/optim |
| foreach: Optional[bool] = None, |
| grad_scale: Optional[Tensor] = None, |
| found_inf: Optional[Tensor] = None, |
| has_complex: bool = False, |
| *, |
| d: float, |
| lr: Union[float, Tensor], |
| beta2_decay: float, |
| weight_decay: float, |
| eps1: float, |
| eps2: float, |
| maximize: bool, |
| ): |
| r"""Functional API that performs Adafactor algorithm computation. |
| |
| See :class:`~torch.optim.Adafactor` for details. |
| """ |
| if not torch._utils.is_compiling() and not all( |
| isinstance(t, torch.Tensor) for t in state_steps |
| ): |
| raise RuntimeError( |
| "`state_steps` argument must contain a list of singleton tensors" |
| ) |
| |
| if foreach: |
| func = _multi_tensor_adafactor |
| else: |
| func = _single_tensor_adafactor |
| |
| func( |
| params, |
| grads, |
| row_vars, |
| col_vars, |
| variances, |
| state_steps, |
| d=d, |
| lr=lr, |
| beta2_decay=beta2_decay, |
| weight_decay=weight_decay, |
| eps1=eps1, |
| eps2=eps2, |
| maximize=maximize, |
| grad_scale=grad_scale, |
| found_inf=found_inf, |
| has_complex=has_complex, |
| ) |
