diff --git a/docs/float16.md b/docs/float16.md new file mode 100644 index 0000000..6fb63aa --- /dev/null +++ b/docs/float16.md @@ -0,0 +1,20 @@ +# Float16 training + +## Motivation + +Pure float16 training is tricky due to the limited dynamic range of float16 numbers. The largest representable number is 65504 (approximately $2^{16}$), and the smallest positive normal number is $6.103\times 10^{-5}$ (approximately $2^{-14}$). This is in stark contrast with float32 and bfloat16 which both share the same dynamic range that is able to represent numbers as large as $10^{38}$ and as small as $10^-{38}$. The most immediate problem with pure float16 training is that the loss function is usually around 1 in magnitude, so the gradients ($\partial L/\partial w_i$) end up on the order of $10^{-3}$ to $10^{-7}$. As a result, many gradients underflow or become [subnormal](https://en.wikipedia.org/wiki/Subnormal_number) and lose precision. Usually, the model simply fails to converge. + +When a machine learning practitioner encounters this problem, the solution suggested to them is [gradient scaling](https://docs.pytorch.org/docs/stable/amp.html#gradient-scaling) which scales up the loss (and therefore the gradients) before backpropagation, and then un-scales these gradients inside the optimizer right before using them. However, the optimizer state usually still uses the same low-bit dtype (i.e., float16) to store values that are derived from the gradients and which have similarly tiny magnitudes. Even worse, some state values are proportional to the square of the gradients (Adam's exponential moving average of the gradient variance for example). Squaring a gradient with magnitude $10^{-5}$ results in a variance value of $10^{-10}$, laughably far below float16's limited dynamic range. + +The next obvious solution is to scale up the loss (and gradients) and then simply not downscale them inside the optimizer. This certainly prevents the loss of precision when unscaling gradients. It works well with scaling factors that bring the magnitude of the gradients to be around 1, and it is an improvement over not scaling the gradients at all. It may be tempting to scale the gradients such that they cover the full dynamic range of float16 values. But consider a (scaled) gradient value of 1024. Inside Adam, the corresponding variance is its square $(2^{10})^2 = 2^{20}$, far above the maximum representable float16 number (65504), thus causing an overflow. + +What we want is to ensure that all of the state variables take full advantage of float16's limited dynamic range. The different state variables must then be scaled by different amounts. Each state variable would introduce a new scaling factor to tune, and each type of optimizer (SGD, RMSProp, Adam, etc) would have its own set of new hyperparameters. This is burdensome to the user. Is there a way to unify these scaling parameters? + +Yes - we can assume that the gradient is within a certain range, and automatically choose state variable scaling factors that do not overflow if the gradient is within the range. That is, **if** + +- The magnitudes of the scaled gradients are less than or equal to `gradient_max`. +- The magnitudes of the scaled old optimizer state variables are less than or equal to `optimizer_state_max`. + +**then** + +- The magnitude of the new scaled optimizer state variables will be less than or equal to `optimizer_state_max`. diff --git a/optimi/__init__.py b/optimi/__init__.py index 9276e1c..bbaef94 100644 --- a/optimi/__init__.py +++ b/optimi/__init__.py @@ -4,6 +4,7 @@ from .adamw import AdamW, adamw from .adan import Adan, adan from .gradientrelease import prepare_for_gradient_release, remove_gradient_release +from .gradientscaling import GradScalerBackport, ScalingInfo from .lion import Lion, lion from .radam import RAdam, radam from .ranger import Ranger, ranger diff --git a/optimi/adam.py b/optimi/adam.py index 7ed6b56..16765af 100644 --- a/optimi/adam.py +++ b/optimi/adam.py @@ -15,19 +15,22 @@ # lion-pytorch - MIT License - Copyright (c) 2023 Phil Wang - https://github.com/lucidrains/lion-pytorch from collections.abc import Callable, Iterable +from functools import partial +import math from typing import Any import torch from torch import Tensor from torch.utils._foreach_utils import _group_tensors_by_device_and_dtype -from optimi.optimizer import OptimiOptimizer -from optimi.utils import HAS_TRITON, _default_to_triton, _device_guard, _get_triton_block_size, debias_beta +from .gradientscaling import find_maximum_upscale, OptimizerScalingHelper, ScalingInfo +from .optimizer import KAHAN_DTYPES, OptimiOptimizerWithGradientScaling +from .utils import HAS_TRITON, _default_to_triton, _device_guard, _get_triton_block_size, debias_beta __all__ = ["Adam", "adam"] -class Adam(OptimiOptimizer): +class Adam(OptimiOptimizerWithGradientScaling): """Adam optimizer. Optionally with decoupled weight decay (AdamW). Args: @@ -51,6 +54,12 @@ class Adam(OptimiOptimizer): gradient_release: Fuses optimizer step and zero_grad as part of the parameter's backward pass. Requires model hooks created with `register_gradient_release`. Incompatible with closure (default: False) + scale_down_hysteresis: When the optimizer scaling factor is updated downwards, then further + divide the scaling factor by this amount. (default: None) + scale_up_hysteresis: When the optimizer scaling factor is updated upwards, then further + divide the scaling factor by this amount. Upwards rescaling must be triggered manually + by setting group["auto_rescale_on_next_iteration"] to True before an optimizer + iteration. (default: None) """ def __init__( @@ -67,6 +76,8 @@ def __init__( foreach: bool | None = None, triton: bool | None = None, gradient_release: bool = False, + scale_down_hysteresis: float | None = None, + scale_up_hysteresis: float | None = None, ): if not 0.0 <= betas[0] < 1.0: raise ValueError(f"Invalid beta1 parameter: {betas[0]=}") @@ -89,6 +100,7 @@ def __init__( triton=triton, gradient_release=gradient_release, setup=False, + **self._gradient_scaling_init_arguments(scale_down_hysteresis=scale_down_hysteresis, scale_up_hysteresis=scale_up_hysteresis), ) super().__init__(params, defaults) @@ -97,7 +109,7 @@ def _init_state(self, group: dict[str, Any], state: dict[Tensor, Any], param: Te state["exp_avg"] = torch.zeros_like(param, memory_format=torch.preserve_format) state["exp_avg_sq"] = torch.zeros_like(param, memory_format=torch.preserve_format) - if (group["kahan_sum"] or group["kahan_sum"] is None) and param.dtype in [torch.float16, torch.bfloat16]: + if (group["kahan_sum"] or group["kahan_sum"] is None) and param.dtype in KAHAN_DTYPES: state["kahan_comp"] = torch.zeros_like(param, memory_format=torch.preserve_format) group["kahan_sum"] = True elif group["triton"]: @@ -133,6 +145,7 @@ def _init_group( if not group["setup"]: group["setup"] = True + self._gradient_scaling_setup_group(group) group["step"] = torch.tensor(0, dtype=torch.int32) if group["triton"] is None and group["foreach"] is None: @@ -154,6 +167,9 @@ def step(self, closure: Callable | None = None, param: Tensor | None = None): with torch.enable_grad(): loss = closure() + if not self._gradient_scaling_pre_step(): + return loss # skip because found_inf + if param is None: for group in self.param_groups: params, grads, exp_avgs, exp_avg_sqs, kahan_comps = [], [], [], [], [] @@ -170,6 +186,7 @@ def step(self, closure: Callable | None = None, param: Tensor | None = None): beta2=group["beta2"], weight_decay=group["weight_decay"], eps=group["eps"], + scaling=self._gradient_scaling_group_info(group), step=group["step"], decouple_wd=group["decouple_wd"], decouple_lr=group["decouple_lr"], @@ -196,6 +213,7 @@ def step(self, closure: Callable | None = None, param: Tensor | None = None): beta2=group["beta2"], weight_decay=group["weight_decay"], eps=group["eps"], + scaling=self._gradient_scaling_group_info(group), step=state["step"], decouple_wd=group["decouple_wd"], decouple_lr=group["decouple_lr"], @@ -231,6 +249,7 @@ def adam( triton: bool = False, gradient_release: bool = False, optimizer_accumulation: bool = False, + scaling: ScalingInfo | None = None, ): """Functional API to apply an Adam or AdamW optimization step. @@ -256,7 +275,22 @@ def adam( triton: Enables the faster Triton implementation gradient_release: Fuses optimizer step as part of the parameter's backward pass optimizer_accumulation: Accumulate gradients into state during gradient release step + scaling: Information used to implement gradient scaling """ + + grad_to_opt_scale = None + rescale_factor = None + rescale_factor_sqr = None + + if scaling: + scaler = OptimizerScalingHelper(scaling_info=scaling) + scaler.get_max_optimizer_state_scale_up_factor = partial(_adam_max_scale_up, exp_avgs=exp_avgs, exp_avg_sqs=exp_avg_sqs) + scaler.update() + grad_to_opt_scale = scaler.gradient_to_optimizer_scale + rescale_factor = scaler.rescale_factor + if rescale_factor is not None: + rescale_factor_sqr = rescale_factor * rescale_factor + # calculate debiased beta hat & complement terms step.add_(1) step_int = step.item() @@ -268,13 +302,32 @@ def adam( # calculate decoupled weight decay or fully decoupled weight decay if weight_decay != 0: if decouple_lr: - weight_decay = 1 - (lr / max_lr) * weight_decay + weight_decay = (lr / max_lr) * weight_decay elif decouple_wd: - weight_decay = 1 - lr * weight_decay + weight_decay = lr * weight_decay if kahan_comps is None: kahan_comps = [None] * len(params) + if grad_to_opt_scale is not None: + if foreach: + raise ValueError(f"Gradient scaling and {foreach=} cannot be used together") + + # Scale the variables to prevent both overflow and underflow. + beta1_comp *= grad_to_opt_scale + beta2_comp *= grad_to_opt_scale * grad_to_opt_scale / scaler.optimizer_max + eps = eps * scaler.optimizer_scale + scale_rms = scaler.optimizer_max**0.5 + if weight_decay != 0 and not (decouple_wd or decouple_lr): + # This type of weight decay adds a grad*weight_decay term directly onto the gradients. This can + # overflow the gradients or the variables used to store their moving averages. There is no safe way + # to implement it without tracking or limiting the magnitudes of the parameters themselves. + + # weight_decay *= scaling.gradient_scale + raise ValueError(f"Gradient scaling and L2 weight decay {decouple_wd or decouple_lr=} cannot be used together") + else: + scale_rms = None + if gradient_release: if triton: func = _single_param_triton_adam @@ -302,13 +355,27 @@ def adam( beta2_hat=beta2_hat, beta2_comp=beta2_comp, weight_decay=weight_decay, + weight_decay_comp=1 - weight_decay, eps=eps, decouple_wd=(decouple_wd or decouple_lr), kahan_sum=kahan_sum, update_parameters=(not optimizer_accumulation), + scale_rms=scale_rms, + rescale_factor=rescale_factor, + rescale_factor_sqr=rescale_factor_sqr, ) +def _adam_max_scale_up(exp_avgs: list[Tensor], exp_avg_sqs: list[Tensor], dtype_to_max_value: dict) -> list[Tensor]: + # enqueue as much work as possible on the GPU(s) + m1 = find_maximum_upscale(exp_avgs, target_value_per_dtype=dtype_to_max_value, out_device=None) + m2 = find_maximum_upscale(exp_avg_sqs, target_value_per_dtype=dtype_to_max_value, out_device=None) + m1 += (x.sqrt_() for x in m2) + + # this returns a list of tensors that may be on the GPU + return m1 + + def _single_adam( params: list[Tensor], grads: list[Tensor], @@ -317,6 +384,7 @@ def _single_adam( kahan_comps: list[Tensor | None], *, lr: float, + beta1_hat: float, beta1_comp: float, beta2_hat: float, beta2_comp: float, @@ -325,6 +393,9 @@ def _single_adam( decouple_wd: bool, kahan_sum: bool = False, update_parameters: bool = True, + scale_rms: float | None = None, + rescale_factor: float | None = None, + rescale_factor_sqr: float | None = None, **kwargs, ): for i, param in enumerate(params): @@ -340,6 +411,7 @@ def _single_adam( exp_avg_sq=exp_avg_sq, kahan_comp=kahan_comp, lr=lr, + beta1_hat=beta1_hat, beta1_comp=beta1_comp, beta2_hat=beta2_hat, beta2_comp=beta2_comp, @@ -348,6 +420,9 @@ def _single_adam( decouple_wd=decouple_wd, kahan_sum=kahan_sum, update_parameters=update_parameters, + scale_rms=scale_rms, + rescale_factor=rescale_factor, + rescale_factor_sqr=rescale_factor_sqr, ) @@ -359,6 +434,7 @@ def _single_param_adam( kahan_comp: Tensor | None, *, lr: float, + beta1_hat: float, beta1_comp: float, beta2_hat: float, beta2_comp: float, @@ -367,23 +443,38 @@ def _single_param_adam( decouple_wd: bool, kahan_sum: bool = False, update_parameters: bool = True, + scale_rms: float | None = None, + rescale_factor: float | None = None, + rescale_factor_sqr: float | None = None, **kwargs, ): + if rescale_factor is not None: + # in-place rescale + exp_avg.mul_(rescale_factor) + exp_avg_sq.mul_(rescale_factor_sqr) + # decoupled weight decay, fully decoupled weight decay, or L2 weight decay - if weight_decay != 0 and update_parameters: - if decouple_wd: - param.mul_(weight_decay) - else: - grad.add_(param, alpha=weight_decay) + if weight_decay != 0 and update_parameters and not decouple_wd: + grad.add_(param, alpha=weight_decay) # update gradient moving averages with debiased betas - exp_avg.lerp_(grad, weight=beta1_comp) + exp_avg.mul_(beta1_hat).add_(grad, alpha=beta1_comp) exp_avg_sq.mul_(beta2_hat).addcmul_(grad, grad, value=beta2_comp) if update_parameters: - if kahan_sum and param.dtype in [torch.float16, torch.bfloat16]: + denom = exp_avg_sq.sqrt() + if scale_rms is not None: + denom.mul_(scale_rms) + denom.add_(eps) + + if kahan_sum and param.dtype in KAHAN_DTYPES: + if weight_decay != 0 and decouple_wd: + # Apply decoupled weight decay directly. + kahan_comp.sub_(kahan_comp, alpha=weight_decay) + kahan_comp.sub_(param.detach(), alpha=weight_decay) + # Adam step - kahan_comp.addcdiv_(exp_avg, exp_avg_sq.sqrt().add_(eps), value=-lr) + kahan_comp.addcdiv_(exp_avg, denom, value=-lr) # update weights with kahan compensation using grad as temp buffer grad.copy_(param.detach()) @@ -393,7 +484,9 @@ def _single_param_adam( kahan_comp.add_(grad.sub_(param)) else: # Adam step - param.addcdiv_(exp_avg, exp_avg_sq.sqrt().add_(eps), value=-lr) + if weight_decay != 0 and decouple_wd: + param.mul_(1 - weight_decay) + param.addcdiv_(exp_avg, denom, value=-lr) def _foreach_adam( @@ -408,6 +501,7 @@ def _foreach_adam( beta2_hat: float, beta2_comp: float, weight_decay: float, + weight_decay_comp: float, eps: float, decouple_wd: bool, kahan_sum: bool = False, @@ -418,14 +512,11 @@ def _foreach_adam( (dev_params, dev_grads, dev_exp_avgs, dev_exp_avg_sqs, dev_kahan_comps), _, ) in grouped_tensors.items(): - do_kahan_sum = kahan_sum and dtype in [torch.float16, torch.bfloat16] + do_kahan_sum = kahan_sum and dtype in KAHAN_DTYPES # decoupled weight decay, fully decoupled weight decay, or L2 weight decay - if weight_decay != 0: - if decouple_wd: - torch._foreach_mul_(dev_params, scalar=weight_decay) - else: - torch._foreach_add_(dev_grads, dev_params, alpha=weight_decay) + if weight_decay != 0 and not decouple_wd: + torch._foreach_add_(dev_grads, dev_params, alpha=weight_decay) # update gradient moving averages with debiased betas torch._foreach_lerp_(dev_exp_avgs, dev_grads, weight=beta1_comp) @@ -440,6 +531,8 @@ def _foreach_adam( if do_kahan_sum: # Adam step torch._foreach_addcdiv_(dev_kahan_comps, dev_exp_avgs, dev_grads, value=-lr) + if weight_decay != 0 and decouple_wd: + torch._foreach_add_(dev_kahan_comps, dev_params, alpha=weight_decay) # update weights with kahan compensation using dev_grads as temp buffer torch._foreach_copy_(dev_grads, dev_params) @@ -449,6 +542,9 @@ def _foreach_adam( torch._foreach_sub_(dev_grads, dev_params, alpha=1) torch._foreach_add_(dev_kahan_comps, dev_grads, alpha=1) else: + if weight_decay != 0 and decouple_wd: + torch._foreach_mul_(dev_params, scalar=weight_decay_comp) + # Adam step torch._foreach_addcdiv_(dev_params, dev_exp_avgs, dev_grads, value=-lr) @@ -471,6 +567,11 @@ def _adam_kernel( beta2_comp, weight_decay, eps, + scale_rms, # used with do_scale + rescale_factor, # used with do_rescale + rescale_factor_sqr, # used with do_rescale + do_scale: tl.constexpr, + do_rescale: tl.constexpr, do_weight_decay: tl.constexpr, kahan_sum: tl.constexpr, decouple_wd: tl.constexpr, @@ -491,24 +592,34 @@ def _adam_kernel( exp_avg = tl.load(exp_avg_ptr + offsets, mask=mask).to(tl.float32) exp_avg_sq = tl.load(exp_avg_sq_ptr + offsets, mask=mask).to(tl.float32) - # decoupled weight decay, fully decoupled weight decay, or L2 weight decay - if do_weight_decay and update_parameters: - if decouple_wd: - param = tl.cast(param * weight_decay, param.dtype) - else: - grad = grad + param.to(tl.float32) * weight_decay + if do_rescale: + # in-place rescale + exp_avg *= rescale_factor + exp_avg_sq *= rescale_factor_sqr + + # L2 weight decay + if do_weight_decay and update_parameters and not decouple_wd: + grad = grad + param.to(tl.float32) * weight_decay # update gradient moving averages: exp_avg = tl.fma(exp_avg, beta1_hat, beta1_comp * grad) exp_avg_sq = tl.fma(exp_avg_sq, beta2_hat, beta2_comp * grad * grad) if update_parameters: + rms = tl.sqrt(exp_avg_sq) + if do_scale: + rms *= scale_rms + if kahan_sum: # load kahan compensation, casting to fp32 kahan_comp = tl.load(kahan_ptr + offsets, mask=mask).to(tl.float32) # AdamW step, using the kahan comp instead of param - kahan_comp = kahan_comp - (lr * exp_avg / (tl.sqrt(exp_avg_sq) + eps)) + kahan_comp -= lr * exp_avg / (rms + eps) + + # apply decoupled weight decay or fully decoupled weight decay + if do_weight_decay and decouple_wd: + kahan_comp -= weight_decay * param # update weights with downcasted kahan update prev_param = param @@ -521,7 +632,10 @@ def _adam_kernel( tl.store(kahan_ptr + offsets, tl.cast(kahan_comp, param.dtype), mask=mask) else: # Standard AdamW step, optionally downcasting to param.dtype from fp32 intermediates - param = param + tl.cast((-lr * exp_avg / (tl.sqrt(exp_avg_sq) + eps)), param.dtype) + subtrahend = lr * exp_avg / (rms + eps) + if do_weight_decay and decouple_wd: + subtrahend += weight_decay * param + param = param - tl.cast(subtrahend, param.dtype) # Store updated parameters tl.store(param_ptr + offsets, param, mask=mask) @@ -546,8 +660,14 @@ def _triton_adam( eps: float, decouple_wd: bool, kahan_sum: bool, + scale_rms: float | None, + rescale_factor: float | None, + rescale_factor_sqr: float | None, **kwargs, ): + do_scale = scale_rms is not None + do_rescale = rescale_factor is not None + for i, param in enumerate(params): grad = grads[i] exp_avg = exp_avgs[i] @@ -574,9 +694,14 @@ def _triton_adam( beta2_comp=beta2_comp, weight_decay=weight_decay, eps=eps, + scale_rms=scale_rms, + do_scale=do_scale, + rescale_factor=rescale_factor, + rescale_factor_sqr=rescale_factor_sqr, + do_rescale=do_rescale, do_weight_decay=weight_decay != 0.0, decouple_wd=decouple_wd, - kahan_sum=kahan_sum and param.dtype in [torch.float16, torch.bfloat16], + kahan_sum=kahan_sum and param.dtype in KAHAN_DTYPES, update_parameters=True, n_elements=n_elements, BLOCK_SIZE=block_size, @@ -596,11 +721,17 @@ def _single_param_triton_adam( beta2_comp: float, weight_decay: float, eps: float, + scale_rms: float | None, + rescale_factor: float | None, + rescale_factor_sqr: float | None, decouple_wd: bool, kahan_sum: bool, update_parameters: bool, **kwargs, ): + do_scale = scale_rms is not None + do_rescale = rescale_factor is not None + n_elements = param.numel() block_size = _get_triton_block_size(n_elements) grid = (triton.cdiv(n_elements, block_size),) @@ -621,9 +752,14 @@ def _single_param_triton_adam( beta2_comp=beta2_comp, weight_decay=weight_decay, eps=eps, + scale_rms=scale_rms, + do_scale=do_scale, + rescale_factor=rescale_factor, + rescale_factor_sqr=rescale_factor_sqr, + do_rescale=do_rescale, do_weight_decay=weight_decay != 0.0, decouple_wd=decouple_wd, - kahan_sum=kahan_sum and param.dtype in [torch.float16, torch.bfloat16], + kahan_sum=kahan_sum and param.dtype in KAHAN_DTYPES, update_parameters=update_parameters, n_elements=n_elements, BLOCK_SIZE=block_size, diff --git a/optimi/adamw.py b/optimi/adamw.py index 900d146..574d723 100644 --- a/optimi/adamw.py +++ b/optimi/adamw.py @@ -14,7 +14,8 @@ from torch import Tensor -from optimi import Adam, adam +from .adam import Adam, adam +from .gradientscaling import ScalingInfo __all__ = ["AdamW", "adamw"] @@ -58,6 +59,8 @@ def __init__( foreach: bool | None = None, triton: bool | None = None, gradient_release: bool = False, + scale_down_hysteresis: float | None = None, + scale_up_hysteresis: float | None = None, ): super().__init__( params=params, @@ -72,6 +75,8 @@ def __init__( foreach=foreach, triton=triton, gradient_release=gradient_release, + scale_down_hysteresis=scale_down_hysteresis, + scale_up_hysteresis=scale_up_hysteresis, ) @@ -95,6 +100,7 @@ def adamw( triton: bool = False, gradient_release: bool = False, optimizer_accumulation: bool = False, + scaling: ScalingInfo | None = None, ): """Functional API to apply an AdamW optimization step. @@ -140,4 +146,5 @@ def adamw( triton=triton, gradient_release=gradient_release, optimizer_accumulation=optimizer_accumulation, + scaling=scaling, ) diff --git a/optimi/adan.py b/optimi/adan.py index 2b37c94..16ee945 100644 --- a/optimi/adan.py +++ b/optimi/adan.py @@ -18,19 +18,21 @@ # lion-pytorch - MIT License - Copyright (c) 2023 Phil Wang - https://github.com/lucidrains/lion-pytorch from collections.abc import Callable, Iterable +from functools import partial from typing import Any import torch from torch import Tensor from torch.utils._foreach_utils import _group_tensors_by_device_and_dtype -from optimi.optimizer import OptimiOptimizer -from optimi.utils import HAS_TRITON, _default_to_triton, _device_guard, _get_triton_block_size, debias_beta +from .gradientscaling import find_maximum_upscale, OptimizerScalingHelper, ScalingInfo +from .optimizer import KAHAN_DTYPES, OptimiOptimizerWithGradientScaling +from .utils import HAS_TRITON, _default_to_triton, _device_guard, _get_triton_block_size, debias_beta __all__ = ["Adan", "adan"] -class Adan(OptimiOptimizer): +class Adan(OptimiOptimizerWithGradientScaling): """Adan Optimizer: Adaptive Nesterov Momentum Algorithm. Args: @@ -57,6 +59,12 @@ class Adan(OptimiOptimizer): gradient_release: Fuses optimizer step and zero_grad as part of the parameter's backward pass. Requires model hooks created with `register_gradient_release`. Incompatible with closure (default: False) + scale_down_hysteresis: When the optimizer scaling factor is updated downwards, then further + divide the scaling factor by this amount. (default: None) + scale_up_hysteresis: When the optimizer scaling factor is updated upwards, then further + divide the scaling factor by this amount. Upwards rescaling must be triggered manually + by setting group["auto_rescale_on_next_iteration"] to True before an optimizer + iteration. (default: None) """ def __init__( @@ -73,6 +81,8 @@ def __init__( foreach: bool | None = None, triton: bool | None = None, gradient_release: bool = False, + scale_down_hysteresis: float | None = None, + scale_up_hysteresis: float | None = None, ): if not 0.0 <= betas[0] < 1.0: raise ValueError(f"Invalid beta1 parameter: {betas[0]=}") @@ -98,17 +108,20 @@ def __init__( triton=triton, gradient_release=gradient_release, setup=False, + **self._gradient_scaling_init_arguments(scale_down_hysteresis=scale_down_hysteresis, scale_up_hysteresis=scale_up_hysteresis), ) super().__init__(params, defaults) - def _init_state(self, group: dict[str, Any], state: dict[Tensor, Any], param: Tensor, gradient_release: bool = False): + def _init_state(self, group: dict[str, Any], state: dict[Tensor, Any], param: Tensor, gradient_release: bool = False, scaling_pre=None): if "kahan_comp" not in state: state["exp_avg"] = torch.zeros_like(param, memory_format=torch.preserve_format) state["exp_avg_diff"] = torch.zeros_like(param, memory_format=torch.preserve_format) state["exp_avg_sq"] = torch.zeros_like(param, memory_format=torch.preserve_format) state["prev_grad"] = param.grad.clone().mul_(-1) + if scaling_pre is not None: + group["optimizer_state_scale"].fill_(scaling_pre["gradient_scale"]) - if (group["kahan_sum"] or group["kahan_sum"] is None) and param.dtype in [torch.float16, torch.bfloat16]: + if (group["kahan_sum"] or group["kahan_sum"] is None) and param.dtype in KAHAN_DTYPES: state["kahan_comp"] = torch.zeros_like(param, memory_format=torch.preserve_format) group["kahan_sum"] = True elif group["triton"]: @@ -148,6 +161,7 @@ def _init_group( if not group["setup"]: group["setup"] = True + self._gradient_scaling_setup_group(group) group["step"] = torch.tensor(0, dtype=torch.int32) if group["triton"] is None and group["foreach"] is None: @@ -172,6 +186,9 @@ def step(self, closure: Callable | None = None, param: Tensor | None = None): with torch.enable_grad(): loss = closure() + if not self._gradient_scaling_pre_step(): + return loss # skip because found_inf + if param is None: for group in self.param_groups: params, grads, exp_avgs, exp_avg_diffs, exp_avg_sqs = [], [], [], [], [] @@ -200,6 +217,7 @@ def step(self, closure: Callable | None = None, param: Tensor | None = None): triton=group["triton"], gradient_release=False, optimizer_accumulation=False, + scaling=self._gradient_scaling_group_info(group), ) else: state = self.state[param] @@ -228,6 +246,7 @@ def step(self, closure: Callable | None = None, param: Tensor | None = None): triton=group["triton"], gradient_release=True, optimizer_accumulation=self._optimizer_accumulation, + scaling=self._gradient_scaling_group_info(group), ) return loss @@ -257,6 +276,7 @@ def adan( triton: bool = False, gradient_release: bool = False, optimizer_accumulation: bool = False, + scaling: ScalingInfo | None = None, ): """Functional API to apply a Adan optimization step. @@ -286,6 +306,7 @@ def adan( gradient_release: Fuses optimizer step as part of the parameter's backward pass optimizer_accumulation: Accumulate gradients into state during gradient release step """ + # calculate debiased beta hat & complement terms step.add_(1) step_int = step.item() @@ -308,6 +329,35 @@ def adan( else: weight_decay = 1 + weight_decay + if scaling: + if not triton: + raise ValueError(f"Gradient scaling and {triton=} must be used together") + + # We need an extra margin of safety because of the `exp_avg_diffs` state variable. This variable could end up containing a value + # with twice the magnitude of the input gradient. + scaler = OptimizerScalingHelper(scaling_info=scaling, scale_opt_upper_limit_extra_factor=1 / 3) + scaler.get_max_optimizer_state_scale_up_factor = partial( + _adan_max_scale_up, exp_avgs=exp_avgs, exp_avg_diffs=exp_avg_diffs, exp_avg_sqs=exp_avg_sqs, prev_grads=prev_grads + ) + scaler.update() + grad_to_opt_scale = scaler.gradient_to_optimizer_scale + rescale_factor = scaler.rescale_factor + if rescale_factor is not None: + rescale_factor_sqr = rescale_factor * rescale_factor + else: + rescale_factor_sqr = None + + # Scale the variables to prevent both overflow and underflow. + scale_grad = grad_to_opt_scale + beta3_comp *= 1 / scaler.optimizer_max + eps = eps * scaler.optimizer_scale + scale_rms = scaler.optimizer_max**0.5 + else: + scale_rms = None + scale_grad = None + rescale_factor = None + rescale_factor_sqr = None + if kahan_comps is None: kahan_comps = [None] * len(params) @@ -343,6 +393,10 @@ def adan( beta3_hat=beta3_hat, beta3_comp=beta3_comp, eps=eps, + scale_grad=scale_grad, + scale_rms=scale_rms, + rescale_factor=rescale_factor, + rescale_factor_sqr=rescale_factor_sqr, weight_decay=weight_decay, adam_wd=adam_wd, kahan_sum=kahan_sum, @@ -350,6 +404,19 @@ def adan( ) +def _adan_max_scale_up( + exp_avgs: list[Tensor], exp_avg_diffs: list[Tensor], exp_avg_sqs: list[Tensor], prev_grads: list[Tensor], dtype_to_max_value: dict +) -> list[Tensor]: + # state variables that have the same magnitudes as the gradients + m1 = find_maximum_upscale(exp_avgs + exp_avg_diffs + prev_grads, target_value_per_dtype=dtype_to_max_value, out_device=None) + m2 = find_maximum_upscale(exp_avg_sqs, target_value_per_dtype=dtype_to_max_value, out_device=None) + + m1 += (x.sqrt_() for x in m2) + + # this returns a list of tensors that may be on the GPU + return m1 + + def _single_adan( params: list[Tensor], grads: list[Tensor], @@ -448,7 +515,7 @@ def _single_param_adan( if adam_wd and weight_decay != 0: param.mul_(weight_decay) - if kahan_sum and param.dtype in [torch.float16, torch.bfloat16]: + if kahan_sum and param.dtype in KAHAN_DTYPES: # Adan step kahan_comp.addcdiv_(exp_avg, denom, value=-lr) kahan_comp.addcdiv_(exp_avg_diff, denom, value=-lr * beta2) @@ -495,7 +562,7 @@ def _foreach_adan( (dev_params, dev_grads, dev_exp_avgs, dev_exp_avg_sqs, dev_exp_avg_diffs, dev_prev_grads, dev_kahan_comps), _, ) in grouped_tensors.items(): - do_kahan_sum = kahan_sum and dtype in [torch.float16, torch.bfloat16] + do_kahan_sum = kahan_sum and dtype in KAHAN_DTYPES # difference between current & previous gradients, prev_grad is negated in last step torch._foreach_add_(dev_prev_grads, dev_grads) @@ -569,6 +636,12 @@ def _adan_kernel( beta3_comp, eps, weight_decay, + scale_grad, # used with do_scale + scale_rms, # used with do_scale + rescale_factor, # used with do_rescale + rescale_factor_sqr, # used with do_rescale + do_scale: tl.constexpr, + do_rescale: tl.constexpr, do_weight_decay: tl.constexpr, adam_wd: tl.constexpr, kahan_sum: tl.constexpr, @@ -591,6 +664,16 @@ def _adan_kernel( exp_avg_sq = tl.load(exp_avg_sq_ptr + offsets, mask=mask).to(tl.float32) prev_grad = tl.load(prev_grad_ptr + offsets, mask=mask).to(tl.float32) + if do_rescale: + # in-place rescale + exp_avg *= rescale_factor + exp_avg_diff *= rescale_factor + exp_avg_sq *= rescale_factor_sqr + prev_grad *= rescale_factor + + if do_scale: + grad *= scale_grad + # difference between current & previous gradients, prev_grad is negated in last step prev_grad = prev_grad + grad @@ -608,18 +691,23 @@ def _adan_kernel( prev_grad_next = -grad if update_parameters: - # Adam-style weight decay - if do_weight_decay and adam_wd: - param = tl.cast(param * weight_decay, param.dtype) + rms = tl.sqrt(exp_avg_sq) + if do_scale: + rms *= scale_rms # calculate η_k and Adan update - scale = -lr / (tl.sqrt(exp_avg_sq) + eps) + scale = -lr / (rms + eps) update = tl.fma(beta2, exp_avg_diff, exp_avg) if kahan_sum: # load kahan compensation, casting to fp32 kahan_comp = tl.load(kahan_ptr + offsets, mask=mask).to(tl.float32) + # Adam-style weight decay + if do_weight_decay and adam_wd: + kahan_comp *= weight_decay + kahan_comp -= param * (1.0 - weight_decay) + # Adan step kahan_comp = kahan_comp + scale * update @@ -630,15 +718,24 @@ def _adan_kernel( # save error back to kahan compensation for next iteration kahan_comp = kahan_comp + prev_param.to(tl.float32) - param.to(tl.float32) + # Adan-style weight decay + if do_weight_decay and (not adam_wd): + kahan_comp /= weight_decay + kahan_comp -= (1.0 - 1 / weight_decay) * param + # store kahan compensation tl.store(kahan_ptr + offsets, tl.cast(kahan_comp, param.dtype), mask=mask) else: + # Adam-style weight decay + if do_weight_decay and adam_wd: + param = tl.cast(param * weight_decay, param.dtype) + # Adan step param = tl.cast(tl.fma(scale, update, param), param.dtype) - # Adan-style weight decay - if do_weight_decay and (not adam_wd): - param = tl.cast(param / weight_decay, param.dtype) + # Adan-style weight decay + if do_weight_decay and (not adam_wd): + param = tl.cast(param / weight_decay, param.dtype) # store updated parameter tl.store(param_ptr + offsets, param, mask=mask) @@ -671,10 +768,16 @@ def _triton_adan( adam_wd: bool, kahan_sum: bool, update_parameters: bool = True, + scale_grad: float | None = None, + scale_rms: float | None = None, + rescale_factor: float | None = None, + rescale_factor_sqr: float | None = None, **kwargs, ) -> None: """Apply a fused Adan step to a list of tensors using the Triton kernel.""" do_weight_decay = weight_decay != 0.0 + do_scale = scale_rms is not None + do_rescale = rescale_factor is not None for i, param in enumerate(params): n_elements = param.numel() @@ -701,10 +804,16 @@ def _triton_adan( beta3_hat=beta3_hat, beta3_comp=beta3_comp, eps=eps, + scale_grad=scale_grad, + scale_rms=scale_rms, + do_scale=do_scale, + rescale_factor=rescale_factor, + rescale_factor_sqr=rescale_factor_sqr, + do_rescale=do_rescale, weight_decay=weight_decay, do_weight_decay=do_weight_decay, adam_wd=adam_wd, - kahan_sum=kahan_sum and param.dtype in [torch.float16, torch.bfloat16], + kahan_sum=kahan_sum and param.dtype in KAHAN_DTYPES, update_parameters=update_parameters, n_elements=n_elements, BLOCK_SIZE=block_size, @@ -732,9 +841,16 @@ def _single_param_triton_adan( adam_wd: bool, kahan_sum: bool, update_parameters: bool, + scale_grad: float | None = None, + scale_rms: float | None = None, + rescale_factor: float | None = None, + rescale_factor_sqr: float | None = None, ) -> None: """Fused Adan step for a single parameter tensor (used with gradient release).""" do_weight_decay = weight_decay != 0.0 + do_scale = scale_rms is not None + do_rescale = rescale_factor is not None + n_elements = param.numel() block_size = _get_triton_block_size(n_elements) @@ -760,10 +876,16 @@ def _single_param_triton_adan( beta3_hat=beta3_hat, beta3_comp=beta3_comp, eps=eps, + scale_grad=scale_grad, + scale_rms=scale_rms, + do_scale=do_scale, + rescale_factor=rescale_factor, + rescale_factor_sqr=rescale_factor_sqr, + do_rescale=do_rescale, weight_decay=weight_decay, do_weight_decay=do_weight_decay, adam_wd=adam_wd, - kahan_sum=kahan_sum and param.dtype in [torch.float16, torch.bfloat16], + kahan_sum=kahan_sum and param.dtype in KAHAN_DTYPES, update_parameters=update_parameters, n_elements=n_elements, BLOCK_SIZE=block_size, diff --git a/optimi/gradientscaling.py b/optimi/gradientscaling.py new file mode 100644 index 0000000..5c94163 --- /dev/null +++ b/optimi/gradientscaling.py @@ -0,0 +1,293 @@ +from collections import defaultdict +import dataclasses + +import torch +from torch import Tensor, optim + + +@dataclasses.dataclass(slots=True) +class ScalingInfo: + """ + Args: + gradient_scale: Current gradient scaling factor + gradient_max_magnitude: Largest magnitude among all gradients (optional) + optimizer_scale: Current optimizer scaling factor + scale_down_hysteresis: When the optimizer scaling factor is updated downwards, then further divide the + resulting scaling factor by this amount (larger than or equal to 1). + scale_up_hysteresis: If not None, then the next step will automatically rescale the optimizer state to + improve its coverage of the available dynamic range. This operation necessarily incurs a CPU-GPU + sync, so you should not do it too often. If not None, then it must be a number greater than or equal + to 1. + """ + + gradient_scale: float + gradient_max_magnitude: float | None + optimizer_scale: Tensor # on CPU! + scale_down_hysteresis: float + scale_up_hysteresis: float | None + + def with_(self, **kw): + return dataclasses.replace(self, **kw) + + +@dataclasses.dataclass +class MagnitudeLimits: + gradient_max: dict[torch.dtype, float] + optimizer_state_max: dict[torch.dtype, float] + + +DEFAULT_MAGNITUDE_LIMITS = MagnitudeLimits(gradient_max={torch.float16: 65536.0}, optimizer_state_max={torch.float16: 65504.0 * (7 / 8)}) + + +def group_tensors_by_device_and_dtype(tensors: list[Tensor]) -> dict[tuple[torch.device, torch.dtype], list[Tensor]]: + # lazy import to avoid affecting users who do not use gradient scaling + from torch.utils._foreach_utils import _group_tensors_by_device_and_dtype + + if not tensors: + return {} + return {k: v for k, [[v], _] in _group_tensors_by_device_and_dtype([tensors]).items()} + + +@torch.no_grad() +def find_maximum_magnitude_single_device_and_dtype(tensors: list[Tensor], full: bool = False): + lst = [] + for t in tensors: + # Tensor.aminmax() propagates nans + m = t.aminmax() + lst.append(m.min) + lst.append(m.max) + + if full: + a = torch.stack(lst).abs().reshape(-1, 2).amax(dim=0) + m = a.amax() + return m, torch.isfinite(a) + else: + a = torch.stack(lst).abs().amax() + return a + + +@torch.no_grad() +def find_maximum_upscale( + tensors: list[Tensor], target_value_per_dtype: dict[torch.dtype, float], out_device="cpu" +) -> Tensor | list[Tensor]: + """ + What is the greatest value `s` such that `(tensor * s).abs() <= target_value[tensor.dtype]` for every + tensor in tensors? + + If `out_device` is None, then return a list of tensors belonging to different devices and dtypes. You are + then responsible for computing the minimum across all of them. + """ + results = [ + target_value / find_maximum_magnitude_single_device_and_dtype(ts) + for (ts_device, ts_dtype), ts in group_tensors_by_device_and_dtype(tensors).items() + if (target_value := target_value_per_dtype.get(ts_dtype)) is not None + ] + if out_device is None: + return results + else: + if results: + return torch.stack([s.to(device=out_device) for s in results]).amin() + else: + return torch.full((), 65536.0, dtype=torch.float32, device=out_device) + + +class _OptimizerScalingHelper: + # used to define the types of non-dataclass fields + gradient_max: float + optimizer_max: float + optimizer_scale: float + gradient_to_optimizer_scale: float + rescale_factor: float | None + + +@dataclasses.dataclass +class OptimizerScalingHelper(_OptimizerScalingHelper): + """ + Args: + scaling_info: Scaling information. This includes the gradient and optimizer scaling factors, and hysteresis + factors to prevent overly frequent optimizer rescaling. + + Attributes: + + gradient_max: The maximum allowed input gradient value. This is set to be above the maximum representable + float16 value, or provided by the gradient scaler. + optimizer_max: The maximum allowed value inside the optimizer state. This is set just below the maximum + representable float16 value. Input gradients that are at most `gradient_max` must not be able to + produce an optimizer state that exceeds `optimizer_max`. + optimizer_scale: The current scaling factor of the optimizer. + gradient_to_optimizer_scale: The scaling factor by which to multiply a gradient value to obtain a value + compatible with the current optimizer scaling. + rescale_factor: If not None, then you MUST rescale the optimizer state by this factor on the next + iteration. + """ + + scaling_info: ScalingInfo + scale_opt_upper_limit_extra_factor: float = 1.0 + magnitude_limits = DEFAULT_MAGNITUDE_LIMITS + + def get_max_optimizer_state_scale_up_factor(self, dtype_to_max_value: dict[torch.dtype, float]) -> list[Tensor]: + raise NotImplementedError + + def update(self): + mag = self.magnitude_limits + info = self.scaling_info + + # hard coded for now, sorry! + gradient_max = mag.gradient_max[torch.float16] + opt_state_max = mag.optimizer_state_max[torch.float16] + + if info.gradient_max_magnitude is not None: + # Leave a tiny margin of safety. + gradient_max = min(gradient_max, info.gradient_max_magnitude * 1.0625) + + self.gradient_max = gradient_max + self.optimizer_max = opt_state_max + + # Compute upper limit on optimizer scale given the current gradient scaling factor (which may + # have changed since the last iteration). + scale_opt_upper_limit = self.scale_opt_upper_limit_extra_factor * opt_state_max / gradient_max * info.gradient_scale + scale_opt = info.optimizer_scale.item() + if scale_opt > scale_opt_upper_limit: + # The current optimizer scale is larger than the upper allowed limit, so we must + # scale the optimizer state downwards. This can happen when the gradients grow which results in + # a shrinking `gradient_scale`. + new_scale_opt = scale_opt_upper_limit / info.scale_down_hysteresis + elif info.scale_up_hysteresis is not None: + # User requested a periodic upscale, so we must figure out the maximum upscaling factor. + + # This returns a list of tensors that may or may not be on the CPU. We must take the minimum of + # all of them. + max_scale_up_opt_tensors = self.get_max_optimizer_state_scale_up_factor( + dtype_to_max_value={k: v / info.scale_up_hysteresis for k, v in mag.optimizer_state_max.items()} + ) + max_scale_up_opt = 65536.0 + if max_scale_up_opt_tensors: + value = torch.stack([x.cpu() for x in max_scale_up_opt_tensors]).amin() + if value.isfinite(): + max_scale_up_opt = value.item() + + new_scale_opt = min(scale_opt * max_scale_up_opt, scale_opt_upper_limit) + if new_scale_opt <= scale_opt: + # this would not be an upscale, cancel it + new_scale_opt = None + else: + new_scale_opt = None + + if new_scale_opt is not None: + # A rescale is necessary. + info.optimizer_scale.fill_(new_scale_opt) + self.rescale_factor = new_scale_opt / scale_opt + scale_opt = new_scale_opt + else: + self.rescale_factor = None + + self.optimizer_scale = scale_opt + self.gradient_to_optimizer_scale = scale_opt / info.gradient_scale + + +class GradScalerBackport: + """ + Backport of the newer PyTorch GradScaler. + + TODO: implement a better scaling policy + """ + + def __init__( + self, + device: str = "cpu", + init_scale: float = 65536.0, + *, + growth_factor: float = 2.0 ** (1 / 50), + backoff_factor: float = 0.5 ** (1 / 2), + enabled: bool = True, + gradient_max_magnitude: float = torch.finfo(torch.float16).max, + ): + self.device = device + self.growth_factor = growth_factor + self.backoff_factor = backoff_factor + self.enabled = enabled + self.gradient_max_magnitude = gradient_max_magnitude + self._found_inf = torch.full((), 0, dtype=torch.int32, device=device) + self._scale = torch.full((), init_scale, dtype=torch.float32, device=device) + self._gradient_max = torch.full((), 0.0, dtype=torch.float32, device=device) + self._performed_step = False + + def scale(self, loss: Tensor | list[Tensor]): + if not self.enabled: + return loss + + if single := isinstance(loss, Tensor): + loss = (loss,) + + loss = [x * self._scale.to(x.device) for x in loss] + + return loss[0] if single else loss + + def step(self, optimizer: optim.Optimizer, *args, **kwargs): + if not self.enabled: + return optimizer.step(*args, **kwargs) + + self._performed_step = True + assert getattr(optimizer, "_step_supports_amp_scaling", False) + + self._update_found_inf(optimizer) + optimizer.found_inf = self._found_inf + optimizer.grad_scale = self._scale + optimizer.grad_maximum_magnitude = self._gradient_max + try: + return optimizer.step(*args, **kwargs) + finally: + del optimizer.found_inf, optimizer.grad_scale, optimizer.grad_maximum_magnitude + + def _update_found_inf(self, optimizer): + # adapted from torch.amp.GradScaler + + # https://stackoverflow.com/questions/5029934/defaultdict-of-defaultdict + # Google says mypy struggles with defaultdicts type annotations. + per_device_and_dtype_grads: dict[tuple[torch.device, torch.dtype], list[Tensor]] = defaultdict(list) + with torch.no_grad(): + for group in optimizer.param_groups: + for param in group["params"]: + assert isinstance(param, torch.Tensor) + if param.grad is None: + continue + if param.grad.is_sparse: + # is_coalesced() == False means the sparse grad has values with duplicate indices. + # coalesce() deduplicates indices and adds all values that have the same index. + # For scaled fp16 values, there's a good chance coalescing will cause overflow, + # so we should check the coalesced _values(). + if param.grad.dtype is torch.float16: + param.grad = param.grad.coalesce() + to_unscale = param.grad._values() + else: + to_unscale = param.grad + + # TODO: is there a way to split by device and dtype without appending in the inner loop? + per_device_and_dtype_grads[to_unscale.device, to_unscale.dtype].append(to_unscale) + + maxs = [] + for k, grads in per_device_and_dtype_grads.items(): + m = find_maximum_magnitude_single_device_and_dtype(grads) + maxs.append(m) + + device = self.device + max_magnitude = torch.stack([m.to(device) for m in maxs]).amax() + self._gradient_max.fill_(torch.maximum(self._gradient_max, max_magnitude)) + self._found_inf |= ~max_magnitude.isfinite() | (max_magnitude > self.gradient_max_magnitude) + + def get_scale(self): + return self._scale.cpu().item() + + def update(self, new_scale=None): + if not self._performed_step: + self._found_inf.fill_(0) + # nothing can be done + return + + # CUDA-friendly, no CPU-GPU sync required if all values are on the same device. + rescale = torch.where(self._found_inf != 0, self.backoff_factor, self.growth_factor) + + self._scale *= rescale + self._found_inf.fill_(0) + self._gradient_max.fill_(0.0) + self._performed_step = False diff --git a/optimi/optimizer.py b/optimi/optimizer.py index 5a1c277..19b7a53 100644 --- a/optimi/optimizer.py +++ b/optimi/optimizer.py @@ -6,7 +6,11 @@ from torch import Tensor from torch.optim.optimizer import Optimizer -from optimi.utils import HAS_TRITON, MIN_TORCH_2_1, MIN_TORCH_2_6 +from .gradientscaling import ScalingInfo +from .utils import HAS_TRITON, MIN_TORCH_2_1, MIN_TORCH_2_6 + + +KAHAN_DTYPES = {torch.float16, torch.bfloat16} class OptimiOptimizer(Optimizer): @@ -112,3 +116,67 @@ def zero_grad(self, set_to_none: bool = True, param: Tensor | None = None): param.grad.detach_() else: param.grad.requires_grad_(False) + + +class OptimiOptimizerWithGradientScaling(OptimiOptimizer): + _step_supports_amp_scaling = True # see torch.amp.GradScaler.step() + + def _gradient_scaling_init_arguments(self, scale_down_hysteresis, scale_up_hysteresis): + if (scale_down_hysteresis is None) != (scale_up_hysteresis is None): + raise ValueError(f"Both must be set ({scale_down_hysteresis=} and {scale_up_hysteresis=})") + + if scale_down_hysteresis is None: + return dict(scale_up_hysteresis=None, scale_down_hysteresis=None) + + if scale_down_hysteresis < 1.0: + raise ValueError(f"{scale_down_hysteresis=} must be >= 1") + if scale_up_hysteresis < 1.0: + raise ValueError(f"{scale_up_hysteresis=} must be >= 1") + + return dict( + scale_up_hysteresis=scale_up_hysteresis, + scale_down_hysteresis=scale_down_hysteresis, + auto_rescale_on_next_iteration=False, + ) + + def _gradient_scaling_setup_group(self, group): + """This runs exactly once on setup BEFORE the first iteration of the optimizer.""" + if group["scale_up_hysteresis"] is not None: + info = self._gradient_scaling_pre_step_info + group["optimizer_state_scale"] = torch.tensor(info["gradient_scale"], dtype=torch.float32) + + def _gradient_scaling_pre_step(self) -> bool: + """ + Returns False if gradient scaling is enabled but the gradient scaler found infs, and is therefore asking us to skip the current + iteration to avoid poisoning the optimizer state. Returns True otherwise. + """ + grad_scale = getattr(self, "grad_scale", None) + if grad_scale is None: + return True + + found_inf = getattr(self, "found_inf", None) + if hasattr(found_inf, "item"): + found_inf = found_inf.item() + if found_inf: + self.skipped_last_step = True + return False + self.skipped_last_step = False + + grad_scale = grad_scale.item() + grad_max = getattr(self, "grad_maximum_magnitude", None) + if grad_max is not None: + grad_max = max(grad_max.item(), 1e-3) + + self._gradient_scaling_pre_step_info = dict(gradient_scale=grad_scale, gradient_max_magnitude=grad_max) + return True + + def _gradient_scaling_group_info(self, group): + if (info := getattr(self, "_gradient_scaling_pre_step_info", None)) is None: + return None + + return ScalingInfo( + optimizer_scale=group["optimizer_state_scale"], + scale_down_hysteresis=group["scale_down_hysteresis"], + scale_up_hysteresis=(group["scale_up_hysteresis"] if group["auto_rescale_on_next_iteration"] else None), + **info, + ) diff --git a/optimi/stableadamw.py b/optimi/stableadamw.py index 0cb68f3..9644a85 100644 --- a/optimi/stableadamw.py +++ b/optimi/stableadamw.py @@ -15,14 +15,17 @@ # lion-pytorch - MIT License - Copyright (c) 2023 Phil Wang - https://github.com/lucidrains/lion-pytorch from collections.abc import Callable, Iterable +from functools import partial from typing import Any import torch from torch import Tensor from torch.utils._foreach_utils import _group_tensors_by_device_and_dtype -from optimi.optimizer import OptimiOptimizer -from optimi.utils import ( +from .gradientscaling import OptimizerScalingHelper, ScalingInfo +from .adam import _adam_max_scale_up +from .optimizer import KAHAN_DTYPES, OptimiOptimizerWithGradientScaling +from .utils import ( HAS_TRITON, TORCH_TO_TRITON_DTYPE, _default_to_triton, @@ -35,7 +38,7 @@ # this is required as Optimizer.load_state_dict casts the state to the param's dtype -def _restore_triton_scratch_state(optim: OptimiOptimizer): +def _restore_triton_scratch_state(optim: OptimiOptimizerWithGradientScaling): "Restore or create scratch to fp32 after potentially cast to low precision by load_state_dict." for group in optim.param_groups: if group["triton"]: @@ -47,7 +50,7 @@ def _restore_triton_scratch_state(optim: OptimiOptimizer): state["mean_square"] = torch.zeros(1, dtype=torch.float32, device=p.device) -class StableAdamW(OptimiOptimizer): +class StableAdamW(OptimiOptimizerWithGradientScaling): """StableAdamW optimizer. An AdamW-Adafactor hybrid with learning rate update clipping. Args: @@ -71,6 +74,12 @@ class StableAdamW(OptimiOptimizer): gradient_release: Fuses optimizer step and zero_grad as part of the parameter's backward pass. Requires model hooks created with `register_gradient_release`. Incompatible with closure (default: False) + scale_down_hysteresis: When the optimizer scaling factor is updated downwards, then further + divide the scaling factor by this amount. (default: None) + scale_up_hysteresis: When the optimizer scaling factor is updated upwards, then further + divide the scaling factor by this amount. Upwards rescaling must be triggered manually + by setting group["auto_rescale_on_next_iteration"] to True before an optimizer + iteration. (default: None) """ def __init__( @@ -86,6 +95,8 @@ def __init__( foreach: bool | None = None, triton: bool | None = None, gradient_release: bool = False, + scale_down_hysteresis: float | None = None, + scale_up_hysteresis: float | None = None, ): if not 0.0 <= betas[0] < 1.0: raise ValueError(f"Invalid beta1 parameter: {betas[0]=}") @@ -107,6 +118,7 @@ def __init__( triton=triton, gradient_release=gradient_release, setup=False, + **self._gradient_scaling_init_arguments(scale_down_hysteresis=scale_down_hysteresis, scale_up_hysteresis=scale_up_hysteresis), ) super().__init__(params, defaults) @@ -118,7 +130,7 @@ def _init_state(self, group: dict[str, Any], state: dict[Tensor, Any], param: Te state["exp_avg_sq"] = torch.zeros_like(param, memory_format=torch.preserve_format) state["eps_sq"] = torch.tensor(group["eps"] ** 2, dtype=param.dtype, device=param.device) - if (group["kahan_sum"] or group["kahan_sum"] is None) and param.dtype in [torch.float16, torch.bfloat16]: + if (group["kahan_sum"] or group["kahan_sum"] is None) and param.dtype in KAHAN_DTYPES: state["kahan_comp"] = torch.zeros_like(param, memory_format=torch.preserve_format) group["kahan_sum"] = True elif group["triton"]: @@ -145,6 +157,7 @@ def _init_group( ): if not group["setup"]: group["setup"] = True + self._gradient_scaling_setup_group(group) group["step"] = torch.tensor(0, dtype=torch.int32) if group["triton"] is None and group["foreach"] is None: @@ -184,6 +197,9 @@ def step(self, closure: Callable | None = None, param: Tensor | None = None): with torch.enable_grad(): loss = closure() + if not self._gradient_scaling_pre_step(): + return loss # skip because found_inf + if param is None: for group in self.param_groups: params, grads, exp_avgs, exp_avg_sqs, eps_sqs, kahan_comps, mean_squares = [], [], [], [], [], [], [] @@ -219,6 +235,7 @@ def step(self, closure: Callable | None = None, param: Tensor | None = None): gradient_release=False, optimizer_accumulation=False, mean_squares=mean_squares, + scaling=self._gradient_scaling_group_info(group), ) else: state = self.state[param] @@ -247,6 +264,7 @@ def step(self, closure: Callable | None = None, param: Tensor | None = None): gradient_release=True, optimizer_accumulation=self._optimizer_accumulation, mean_squares=state["mean_square"], + scaling=self._gradient_scaling_group_info(group), ) else: stableadamw( @@ -269,6 +287,7 @@ def step(self, closure: Callable | None = None, param: Tensor | None = None): triton=False, gradient_release=True, optimizer_accumulation=self._optimizer_accumulation, + scaling=self._gradient_scaling_group_info(group), ) return loss @@ -296,6 +315,7 @@ def stableadamw( gradient_release: bool = False, optimizer_accumulation: bool = False, mean_squares: list[Tensor] | None = None, + scaling: ScalingInfo | None = None, ): """Functional API to apply a StableAdamW optimization step. @@ -322,6 +342,7 @@ def stableadamw( gradient_release: Fuses optimizer step as part of the parameter's backward pass optimizer_accumulation: Accumulate gradients into state during gradient release step mean_squares: RMS calculation scratch tensor for triton kernel + scaling: Gradient scaling information """ # calculate debiased beta hat & complement terms step.add_(1) @@ -331,6 +352,36 @@ def stableadamw( beta2_hat = debias_beta(beta2, step_int) beta2_comp = 1 - beta2_hat + if scaling: + if not triton: + raise ValueError(f"Gradient scaling and {triton=} must be used together") + + scaler = OptimizerScalingHelper(scaling_info=scaling) + + # The state variables are the same as for adam + scaler.get_max_optimizer_state_scale_up_factor = partial(_adam_max_scale_up, exp_avgs=exp_avgs, exp_avg_sqs=exp_avg_sqs) + scaler.update() + grad_to_opt_scale = scaler.gradient_to_optimizer_scale + rescale_factor = scaler.rescale_factor + if rescale_factor is not None: + rescale_factor_sqr = rescale_factor * rescale_factor + else: + rescale_factor_sqr = None + + # Scale the variables to prevent both overflow and underflow. + beta1_comp *= grad_to_opt_scale + beta2_comp *= grad_to_opt_scale * grad_to_opt_scale / scaler.optimizer_max + eps = eps * scaler.optimizer_scale + eps_sq = eps * eps / scaler.optimizer_max + scale_rms = scaler.optimizer_max**0.5 + scale_grad_for_rms = grad_to_opt_scale / scale_rms + else: + rescale_factor = None + rescale_factor_sqr = None + eps_sq = eps * eps + scale_rms = None + scale_grad_for_rms = None + if kahan_comps is None: kahan_comps = [None] * len(params) @@ -363,11 +414,16 @@ def stableadamw( beta2_comp=beta2_comp, weight_decay=weight_decay, eps=eps, + eps_sq_float=eps_sq, decouple_lr=decouple_lr, max_lr=max_lr, kahan_sum=kahan_sum, update_parameters=(not optimizer_accumulation), mean_squares=mean_squares, + scale_rms=scale_rms, + scale_grad_for_rms=scale_grad_for_rms, + rescale_factor=rescale_factor, + rescale_factor_sqr=rescale_factor_sqr, ) @@ -434,6 +490,7 @@ def _single_param_stableadamw( eps: float, decouple_lr: bool, max_lr: float | None, + scale_rms: float | None = None, kahan_sum: bool = False, update_parameters: bool = True, **kwargs, @@ -443,6 +500,11 @@ def _single_param_stableadamw( exp_avg_sq.mul_(beta2_hat).addcmul_(grad, grad, value=beta2_comp) if update_parameters: + denom = exp_avg_sq.sqrt() + if scale_rms is not None: + denom.mul_(scale_rms) + denom.add_(eps) + # compute per tensor RMS stabilization term rms = grad.pow(2).div_(exp_avg_sq.maximum(eps_sq)).mean().sqrt() @@ -457,7 +519,7 @@ def _single_param_stableadamw( weight_decay = 1 - lr * weight_decay param.mul_(weight_decay) - if kahan_sum and param.dtype in [torch.float16, torch.bfloat16]: + if kahan_sum and param.dtype in KAHAN_DTYPES: # Adam step kahan_comp.addcdiv_(exp_avg, exp_avg_sq.sqrt().add_(eps), value=-lr) @@ -493,7 +555,7 @@ def _foreach_stableadamw( ): grouped_tensors = _group_tensors_by_device_and_dtype([params, grads, exp_avgs, exp_avg_sqs, eps_sqs, kahan_comps]) for (_, dtype), ((dev_params, dev_grads, dev_exp_avgs, dev_exp_avg_sqs, dev_eps_sqs, dev_kahan_comps), _) in grouped_tensors.items(): - do_kahan_sum = kahan_sum and dtype in [torch.float16, torch.bfloat16] + do_kahan_sum = kahan_sum and dtype in KAHAN_DTYPES # update gradient moving averages with debiased betas torch._foreach_lerp_(dev_exp_avgs, dev_grads, weight=beta1_comp) @@ -555,11 +617,17 @@ def _stableadamw_exp_avg_kernel( exp_avg_sq_ptr, mean_square_ptr, eps, + eps_sq, beta1_hat, beta1_comp, beta2_hat, beta2_comp, + scale_grad_for_rms, # used with do_scale + rescale_factor, # used with do_rescale + rescale_factor_sqr, # used with do_rescale n_elements, + do_rescale: tl.constexpr, + do_scale: tl.constexpr, update_parameters: tl.constexpr, BLOCK_SIZE: tl.constexpr, param_dtype: tl.constexpr = tl.float32, @@ -575,13 +643,22 @@ def _stableadamw_exp_avg_kernel( exp_avg = tl.load(exp_avg_ptr + offsets, mask=mask).to(tl.float32) exp_avg_sq = tl.load(exp_avg_sq_ptr + offsets, mask=mask).to(tl.float32) + if do_rescale: + # in-place rescale + exp_avg *= rescale_factor + exp_avg_sq *= rescale_factor_sqr + exp_avg = tl.fma(exp_avg, beta1_hat, beta1_comp * grad) exp_avg_sq = tl.fma(exp_avg_sq, beta2_hat, beta2_comp * grad * grad) # partial calculation of per-element stabilisation term if update_parameters: - square = tl.where(mask, (grad * grad) / tl.maximum(exp_avg_sq, eps * eps), 0.0) + # We don't need to worry about underflow or underflow here because all the variables below are float32. + square = tl.where(mask, (grad * grad) / tl.maximum(exp_avg_sq, eps_sq), 0.0) block_sum = tl.sum(square, axis=0, dtype=tl.float32) / n_elements + if do_scale: + block_sum *= scale_grad_for_rms * scale_grad_for_rms + # in testing, this atomic_add was faster then storing the results in # a temporary buffer then summing in PyTorch or another Triton kernel tl.atomic_add(mean_square_ptr, block_sum) @@ -601,6 +678,8 @@ def _stableadamw_update_kernel( weight_decay, eps, max_lr, + scale_rms, # used with do_scale + do_scale: tl.constexpr, do_weight_decay: tl.constexpr, kahan_sum: tl.constexpr, decouple_lr: tl.constexpr, @@ -621,22 +700,32 @@ def _stableadamw_update_kernel( # RMS stabilized learning rate mean_square = tl.load(mean_square_ptr) - lr = lr / tl.maximum(1.0, tl.sqrt(mean_square)) + lr = lr / tl.maximum(1.0, tl.sqrt(mean_square)) # TODO: move this out of the kernel + + rms = tl.sqrt(exp_avg_sq) + if do_scale: + rms *= scale_rms - # decoupled weight decay or fully decoupled weight decay if do_weight_decay: if decouple_lr: - weight_decay = 1.0 - (lr / max_lr) * weight_decay + weight_decay = (lr / max_lr) * weight_decay else: - weight_decay = 1.0 - lr * weight_decay - param = tl.cast(param * weight_decay, param.dtype) + weight_decay = lr * weight_decay + + # AdamW update + delta = -(lr * exp_avg / (rms + eps)) if kahan_sum: # load kahan compensation, casting to fp32 kahan_comp = tl.load(kahan_ptr + offsets, mask=mask).to(tl.float32) + # decoupled weight decay or fully decoupled weight decay + if do_weight_decay: + kahan_comp *= 1.0 - weight_decay + kahan_comp -= param * weight_decay + # AdamW step, using the kahan comp instead of param - kahan_comp = kahan_comp - (lr * exp_avg / (tl.sqrt(exp_avg_sq) + eps)) + kahan_comp = kahan_comp + delta # update weights with downcasted kahan update prev_param = param @@ -648,8 +737,11 @@ def _stableadamw_update_kernel( # store kahan compensation tl.store(kahan_ptr + offsets, tl.cast(kahan_comp, param.dtype), mask=mask) else: + if do_weight_decay: + delta -= param * weight_decay + # Standard AdamW step, optionally downcasting to param.dtype from fp32 intermediates - param = param + tl.cast((-lr * exp_avg / (tl.sqrt(exp_avg_sq) + eps)), param.dtype) + param = param + tl.cast(delta, param.dtype) # Store updated parameters tl.store(param_ptr + offsets, param, mask=mask) @@ -670,11 +762,18 @@ def _triton_stableadamw( beta2_comp: float, weight_decay: float, eps: float, + eps_sq_float: float, decouple_lr: bool, max_lr: float | None = None, kahan_sum: bool = False, + scale_rms: float | None = None, + scale_grad_for_rms: float | None = None, + rescale_factor: float | None = None, + rescale_factor_sqr: float | None = None, **kwargs, ): + do_scale = scale_rms is not None + for i, param in enumerate(params): grad = grads[i] exp_avg = exp_avgs[i] @@ -695,11 +794,17 @@ def _triton_stableadamw( exp_avg_sq_ptr=exp_avg_sq, mean_square_ptr=mean_square, eps=eps, + eps_sq=eps_sq_float, beta1_hat=beta1_hat, beta1_comp=beta1_comp, beta2_hat=beta2_hat, beta2_comp=beta2_comp, update_parameters=True, + do_scale=do_scale, + do_rescale=rescale_factor is not None, + rescale_factor=rescale_factor, + rescale_factor_sqr=rescale_factor_sqr, + scale_grad_for_rms=scale_grad_for_rms, n_elements=n_elements, BLOCK_SIZE=block_size, param_dtype=TORCH_TO_TRITON_DTYPE[param.dtype], @@ -715,8 +820,10 @@ def _triton_stableadamw( weight_decay=weight_decay, eps=eps, max_lr=max_lr, + do_scale=do_scale, + scale_rms=scale_rms, do_weight_decay=weight_decay != 0.0, - kahan_sum=kahan_sum and param.dtype in [torch.float16, torch.bfloat16], + kahan_sum=kahan_sum and param.dtype in KAHAN_DTYPES, decouple_lr=decouple_lr, n_elements=n_elements, BLOCK_SIZE=block_size, @@ -740,12 +847,19 @@ def _single_param_triton_stableadamw( beta2_comp: float, weight_decay: float, eps: float, + eps_sq_float: float, decouple_lr: bool, max_lr: float | None = None, kahan_sum: bool = False, update_parameters: bool = True, + scale_rms: float | None = None, + scale_grad_for_rms: float | None = None, + rescale_factor: float | None = None, + rescale_factor_sqr: float | None = None, **kwargs, ): + do_scale = scale_rms is not None + n_elements = param.numel() block_size = _get_triton_block_size(n_elements) @@ -760,11 +874,17 @@ def _single_param_triton_stableadamw( exp_avg_sq_ptr=exp_avg_sq, mean_square_ptr=mean_squares, eps=eps, + eps_sq=eps_sq_float, beta1_hat=beta1_hat, beta1_comp=beta1_comp, beta2_hat=beta2_hat, beta2_comp=beta2_comp, update_parameters=update_parameters, + do_scale=do_scale, + do_rescale=rescale_factor is not None, + rescale_factor=rescale_factor, + rescale_factor_sqr=rescale_factor_sqr, + scale_grad_for_rms=scale_grad_for_rms, n_elements=n_elements, BLOCK_SIZE=block_size, param_dtype=TORCH_TO_TRITON_DTYPE[param.dtype], @@ -782,9 +902,11 @@ def _single_param_triton_stableadamw( eps=eps, max_lr=max_lr, do_weight_decay=weight_decay != 0.0, - kahan_sum=kahan_sum and param.dtype in [torch.float16, torch.bfloat16], + kahan_sum=kahan_sum and param.dtype in KAHAN_DTYPES, decouple_lr=decouple_lr, n_elements=n_elements, + do_scale=do_scale, + scale_rms=scale_rms, BLOCK_SIZE=block_size, ) # reset mean_square scratch for next iteration diff --git a/tests/config.py b/tests/config.py index 4064237..745c466 100644 --- a/tests/config.py +++ b/tests/config.py @@ -319,10 +319,10 @@ def default_variants(base: OptTest) -> list[OptTest]: replace( base_test, name=f"{base.optimizer_name}_decoupled_lr", - optimi_params=base.optimi_params.with_(weight_decay=1e-5, decouple_lr=True), + optimi_params=base.optimi_params.with_(weight_decay=3e-5, decouple_lr=True), reference_class=ref_cls, reference_params=(base.reference_params or base.optimi_params).with_( - weight_decay=1e-5 if base.fully_decoupled_reference else 0.01, + weight_decay=3e-5 if base.fully_decoupled_reference else 0.03, decouple_lr=True, ), ) diff --git a/tests/conftest.py b/tests/conftest.py index f2a0008..af44944 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -19,6 +19,7 @@ def pytest_configure(config): # Register dtype marks config.addinivalue_line("markers", "float32: mark test to run with float32 dtype") config.addinivalue_line("markers", "bfloat16: mark test to run with bfloat16 dtype") + config.addinivalue_line("markers", "float16: mark test to run with float16 dtype") # Register backend marks config.addinivalue_line("markers", "torch: mark test to run with torch backend") diff --git a/tests/opt_adam.py b/tests/opt_adam.py index e662184..1533e5e 100644 --- a/tests/opt_adam.py +++ b/tests/opt_adam.py @@ -11,7 +11,7 @@ @dataclass class AdamParams(BaseParams): betas: tuple[float, float] = (0.9, 0.99) - eps: float = 1e-6 + eps: float = 1e-5 # Provide BASE so the framework generates base/l2/decoupled variants as applicable. diff --git a/tests/opt_adamw.py b/tests/opt_adamw.py index 9ea41d3..7f03226 100644 --- a/tests/opt_adamw.py +++ b/tests/opt_adamw.py @@ -12,7 +12,7 @@ @dataclass class AdamWParams(BaseParams): betas: tuple[float, float] = (0.9, 0.99) - eps: float = 1e-6 + eps: float = 1e-5 # Provide BASE with fully_decoupled_reference so decoupled_lr uses DecoupledAdamW diff --git a/tests/opt_adan.py b/tests/opt_adan.py index 524d5c7..9c80d55 100644 --- a/tests/opt_adan.py +++ b/tests/opt_adan.py @@ -7,13 +7,13 @@ import torch from tests import reference -from .config import BaseParams, DeviceType, OptTest, OptTestType +from .config import BaseParams, DeviceType, OptTest, OptTestType, Tolerance, with_updated_spec @dataclass class AdanParams(BaseParams): betas: tuple[float, float, float] = (0.98, 0.92, 0.99) - eps: float = 1e-6 + eps: float = 1e-5 weight_decouple: bool = False # For adam_wd variant (maps to no_prox in reference) adam_wd: bool = False # For optimi optimizer @@ -25,6 +25,12 @@ def to_reference_kwargs(self, reference_class: type) -> dict[str, Any]: return kwargs +_custom_iterations = {(OptTestType.normal, DeviceType.gpu, torch.bfloat16): 20} +_spec = with_updated_spec( + spec=None, + test_type=OptTestType.normal, + tolerances_override={torch.float16: Tolerance(atol=1.4e-4, rtol=1.4e-3, max_error_rate=0.01)}, +) TESTS = [ OptTest( name="adan_base", @@ -32,7 +38,8 @@ def to_reference_kwargs(self, reference_class: type) -> dict[str, Any]: optimi_params=AdanParams(), reference_class=reference.Adan, reference_params=AdanParams(), - custom_iterations={(OptTestType.normal, DeviceType.gpu, torch.bfloat16): 20}, + custom_iterations=_custom_iterations, + spec=_spec, ), OptTest( name="adan_weight_decay", @@ -40,7 +47,8 @@ def to_reference_kwargs(self, reference_class: type) -> dict[str, Any]: optimi_params=AdanParams(weight_decay=2e-2), reference_class=reference.Adan, reference_params=AdanParams(weight_decay=2e-2), - custom_iterations={(OptTestType.normal, DeviceType.gpu, torch.bfloat16): 20}, + custom_iterations=_custom_iterations, + spec=_spec, ), OptTest( name="adan_adam_wd", @@ -48,7 +56,8 @@ def to_reference_kwargs(self, reference_class: type) -> dict[str, Any]: optimi_params=AdanParams(weight_decay=2e-2, adam_wd=True), reference_class=reference.Adan, reference_params=AdanParams(weight_decay=2e-2, weight_decouple=True), - custom_iterations={(OptTestType.normal, DeviceType.gpu, torch.bfloat16): 20}, + custom_iterations=_custom_iterations, + spec=_spec, ), OptTest( name="adan_decoupled_lr", @@ -56,6 +65,7 @@ def to_reference_kwargs(self, reference_class: type) -> dict[str, Any]: optimi_params=AdanParams(weight_decay=2e-5, decouple_lr=True), reference_class=reference.Adan, reference_params=AdanParams(weight_decay=2e-2), - custom_iterations={(OptTestType.normal, DeviceType.gpu, torch.bfloat16): 20}, + custom_iterations=_custom_iterations, + spec=_spec, ), ] diff --git a/tests/opt_stableadamw.py b/tests/opt_stableadamw.py index 69f1507..a719c08 100644 --- a/tests/opt_stableadamw.py +++ b/tests/opt_stableadamw.py @@ -11,7 +11,7 @@ @dataclass class StableAdamWParams(BaseParams): betas: tuple[float, float] = (0.9, 0.99) - eps: float = 1e-6 + eps: float = 1e-5 BASE = OptTest( diff --git a/tests/runner.py b/tests/runner.py index f0cecba..2ffc255 100644 --- a/tests/runner.py +++ b/tests/runner.py @@ -4,7 +4,7 @@ import random import torch -from optimi import prepare_for_gradient_release, remove_gradient_release +from optimi import prepare_for_gradient_release, remove_gradient_release, GradScalerBackport from torch import Tensor from .config import Backend, DeviceType, OptTest, OptTestType @@ -134,6 +134,15 @@ def run_test( optimi_kwargs = opttest.to_optimi_kwargs(backend) reference_class = opttest.reference_class + # Pure float16 training requires special care. Scale up the gradients and downscale the + # variances (exponential moving average of gradient squared in case of AdamW). + if dtype == torch.float16: + enable_grad_scaler = True + optimi_kwargs["scale_up_hysteresis"] = 1.2 + optimi_kwargs["scale_down_hysteresis"] = 1.2 + else: + enable_grad_scaler = False + reference_optimizer = None optimi_optimizer = None torch_optimizers: dict[torch.nn.Parameter, torch.optim.Optimizer] | None = None @@ -165,7 +174,19 @@ def optimizer_hook(parameter) -> None: prepare_for_gradient_release(m2, optimi_optimizer) gradient_accumulation_steps = accumulation_spec.gradient_accumulation_steps + grad_scaler = GradScalerBackport(init_scale=16384.0, enabled=enable_grad_scaler) + max_skipped_iteration_fraction = 0.25 + skipped_iterations = 0 + if enable_grad_scaler: + # Increase the number of iteration to make up for the skipped iterations due to the gradient scaler + # finding nans or infs inside the gradients. + iterations = round(iterations / (1 - max_skipped_iteration_fraction)) + for i in range(iterations): + if enable_grad_scaler: + for pg in optimi_optimizer.param_groups: + pg["auto_rescale_on_next_iteration"] = i % 3 == 2 + input1 = torch.randn(batch_size, dim1, device=device, dtype=dtype) if test_type == OptTestType.normal: input2 = input1.detach().clone() @@ -200,15 +221,21 @@ def optimizer_hook(parameter) -> None: loss3 = torch.nn.functional.mse_loss(output3, target3) if output3 is not None else None loss1.backward() - loss2.backward() + losses = [loss2] if loss3 is not None: - loss3.backward() + losses.append(loss3) + for loss in grad_scaler.scale(losses): + loss.backward() if test_type == OptTestType.normal: - reference_optimizer.step() - optimi_optimizer.step() + grad_scaler.step(optimi_optimizer) + if getattr(optimi_optimizer, "skipped_last_step", False): + skipped_iterations += 1 + else: + reference_optimizer.step() reference_optimizer.zero_grad() optimi_optimizer.zero_grad() + grad_scaler.update() elif test_type == OptTestType.gradient_release: reference_optimizer.step() reference_optimizer.zero_grad() @@ -218,8 +245,10 @@ def optimizer_hook(parameter) -> None: if test_type in (OptTestType.gradient_release, OptTestType.accumulation): if random.random() < 0.5: - optimi_optimizer.step() + grad_scaler.step(optimi_optimizer) + assert not getattr(optimi_optimizer, "skipped_last_step", False) optimi_optimizer.zero_grad() + grad_scaler.update() if test_type == OptTestType.normal: assert_most_approx_close( @@ -306,6 +335,8 @@ def optimizer_hook(parameter) -> None: name="PyTorch-Optimi: ", ) + assert skipped_iterations < iterations * max_skipped_iteration_fraction + if test_type == OptTestType.accumulation: assert_most_approx_close( m1.fc1.weight, diff --git a/tests/test_optimizers.py b/tests/test_optimizers.py index 5b2e65f..74f119b 100644 --- a/tests/test_optimizers.py +++ b/tests/test_optimizers.py @@ -2,6 +2,7 @@ import torch from _pytest.mark.structures import ParameterSet +import optimi from .config import Backend, DeviceType, OptTest, OptTestType, discover_tests from .runner import run_test @@ -13,6 +14,7 @@ DTYPE_PARAMS = [ pytest.param(torch.float32, marks=pytest.mark.float32, id="float32"), pytest.param(torch.bfloat16, marks=pytest.mark.bfloat16, id="bfloat16"), + pytest.param(torch.float16, marks=pytest.mark.float16, id="float16"), ] BACKEND_PARAMS = [ pytest.param(Backend.torch, marks=pytest.mark.torch, id=Backend.torch.value), @@ -80,6 +82,18 @@ def _should_skip(test_type: OptTestType, opttest: OptTest, device_type: DeviceTy if test_type != OptTestType.normal and backend == Backend.foreach: return True + # 9. Use float16 only on supported optimizers + if dtype == torch.float16 and not ( + issubclass(opttest.optimi_class, optimi.Adam) + and backend != Backend.foreach + and (opttest.optimi_params.decouple_lr or opttest.optimi_params.decouple_wd or opttest.optimi_params.weight_decay == 0) + or issubclass(opttest.optimi_class, optimi.StableAdamW) + and backend == Backend.triton + or issubclass(opttest.optimi_class, optimi.Adan) + and backend == Backend.triton + ): + return True + return False