Barbados / bundle /source /models /adam_atan2.py

Release Barbados looped MoE research prototype

d5147ba verified 15 days ago

4.14 kB

from typing import Tuple, Optional

import torch
from torch import Tensor
from torch.optim.optimizer import Optimizer, ParamsT


class AdamATan2(Optimizer):
    def __init__(
        self,
        params: ParamsT,
        # Optimizer parameters
        lr: float | Tensor = 1e-3,
        betas: Tuple[float, float] = (0.9, 0.95),
        weight_decay: float = 0.1,
        # Extra features
        ema: Optional[float] = None,
    ):
        # Initialize the Adam-atan2 optimizer
        if isinstance(lr, Tensor):
            if lr.numel() != 1:
                raise ValueError("Tensor lr must be 1-element")
        if not 0.0 <= lr:
            raise ValueError(f"Invalid learning rate: {lr}")
        if not 0.0 <= betas[0] < 1.0:
            raise ValueError(f"Invalid beta parameter at index 0: {betas[0]}")
        if not 0.0 <= betas[1] < 1.0:
            raise ValueError(f"Invalid beta parameter at index 1: {betas[1]}")
        if not 0.0 <= weight_decay:
            raise ValueError(f"Invalid weight_decay value: {weight_decay}")

        defaults = {
            "lr": lr,
            "betas": betas,
            "weight_decay": weight_decay,
            "ema": ema
        }
        super().__init__(params, defaults)
        # Initialize state
        self._init_state()

    @torch.no_grad()
    def _init_state(self):
        for group in self.param_groups:
            for p in group["params"]:
                state = self.state[p]

                # Step counter
                state["step"] = torch.tensor(0.0, dtype=torch.get_default_dtype())

                # Momentum
                if group["betas"][0] > 0:
                    state["exp_avg"] = torch.zeros_like(p)
                state["exp_avg_sq"] = torch.zeros_like(p)

                # Extra features
                if group["ema"] is not None:
                    state["param_ema"] = torch.empty_like(p).copy_(p)

    @torch.no_grad()
    def step(self, closure=None):  # pyright: ignore[reportIncompatibleMethodOverride]
        """Perform a single optimization step."""
        assert closure is None, "Closure is not supported"

        for group in self.param_groups:
            for param in group["params"]:
                if param.grad is None:
                    continue
                
                state = self.state[param]
                grad = param.grad

                # Weight decay update
                if group["weight_decay"] != 0:
                    param.mul_(1 - group["lr"] * group["weight_decay"])

                # Momentums
                if "exp_avg" in state:
                    state["exp_avg"].lerp_(grad, 1 - group["betas"][0])
                state["exp_avg_sq"].mul_(group["betas"][1]).addcmul_(grad, grad, value=1 - group["betas"][1])
                
                state["step"] += 1
                bias_correction1 = 1 - group["betas"][0] ** state["step"]
                bias_correction2 = 1 - group["betas"][1] ** state["step"]
                step_size = group["lr"] / bias_correction1
                bias_correction2_sqrt = bias_correction2.sqrt()

                denom = state["exp_avg_sq"].sqrt() / bias_correction2_sqrt
                # AdamW-atan2
                if "exp_avg" in state:
                    param.add_(torch.atan2(state["exp_avg"], denom), alpha=-step_size)  # pyright: ignore[reportArgumentType]
                else:
                    param.add_(torch.atan2(grad, denom), alpha=-group["lr"])  # pyright: ignore[reportArgumentType]

                # [Extra features] EMA
                if "param_ema" in state:
                    state["param_ema"].lerp_(param, 1 - group["ema"])

    @torch.no_grad()
    def swap_ema(self):
        """Swap param buffer and EMA buffer for evaluation. Remember to swap back after evaluation."""
        for group in self.param_groups:
            for param in group["params"]:
                state = self.state[param]
                if "param_ema" in state:
                    temp = torch.empty_like(param).copy_(param)
                    param.copy_(state["param_ema"])
                    state["param_ema"].copy_(temp)

	from typing import Tuple, Optional

	import torch
	from torch import Tensor
	from torch.optim.optimizer import Optimizer, ParamsT


	class AdamATan2(Optimizer):
	def __init__(
	self,
	params: ParamsT,
	# Optimizer parameters
	lr: float \| Tensor = 1e-3,
	betas: Tuple[float, float] = (0.9, 0.95),
	weight_decay: float = 0.1,
	# Extra features
	ema: Optional[float] = None,
	):
	# Initialize the Adam-atan2 optimizer
	if isinstance(lr, Tensor):
	if lr.numel() != 1:
	raise ValueError("Tensor lr must be 1-element")
	if not 0.0 <= lr:
	raise ValueError(f"Invalid learning rate: {lr}")
	if not 0.0 <= betas[0] < 1.0:
	raise ValueError(f"Invalid beta parameter at index 0: {betas[0]}")
	if not 0.0 <= betas[1] < 1.0:
	raise ValueError(f"Invalid beta parameter at index 1: {betas[1]}")
	if not 0.0 <= weight_decay:
	raise ValueError(f"Invalid weight_decay value: {weight_decay}")

	defaults = {
	"lr": lr,
	"betas": betas,
	"weight_decay": weight_decay,
	"ema": ema
	}
	super().__init__(params, defaults)
	# Initialize state
	self._init_state()

	@torch.no_grad()
	def _init_state(self):
	for group in self.param_groups:
	for p in group["params"]:
	state = self.state[p]

	# Step counter
	state["step"] = torch.tensor(0.0, dtype=torch.get_default_dtype())

	# Momentum
	if group["betas"][0] > 0:
	state["exp_avg"] = torch.zeros_like(p)
	state["exp_avg_sq"] = torch.zeros_like(p)

	# Extra features
	if group["ema"] is not None:
	state["param_ema"] = torch.empty_like(p).copy_(p)

	@torch.no_grad()
	def step(self, closure=None): # pyright: ignore[reportIncompatibleMethodOverride]
	"""Perform a single optimization step."""
	assert closure is None, "Closure is not supported"

	for group in self.param_groups:
	for param in group["params"]:
	if param.grad is None:
	continue

	state = self.state[param]
	grad = param.grad

	# Weight decay update
	if group["weight_decay"] != 0:
	param.mul_(1 - group["lr"] * group["weight_decay"])

	# Momentums
	if "exp_avg" in state:
	state["exp_avg"].lerp_(grad, 1 - group["betas"][0])
	state["exp_avg_sq"].mul_(group["betas"][1]).addcmul_(grad, grad, value=1 - group["betas"][1])

	state["step"] += 1
	bias_correction1 = 1 - group["betas"][0] ** state["step"]
	bias_correction2 = 1 - group["betas"][1] ** state["step"]
	step_size = group["lr"] / bias_correction1
	bias_correction2_sqrt = bias_correction2.sqrt()

	denom = state["exp_avg_sq"].sqrt() / bias_correction2_sqrt
	# AdamW-atan2
	if "exp_avg" in state:
	param.add_(torch.atan2(state["exp_avg"], denom), alpha=-step_size) # pyright: ignore[reportArgumentType]
	else:
	param.add_(torch.atan2(grad, denom), alpha=-group["lr"]) # pyright: ignore[reportArgumentType]

	# [Extra features] EMA
	if "param_ema" in state:
	state["param_ema"].lerp_(param, 1 - group["ema"])

	@torch.no_grad()
	def swap_ema(self):
	"""Swap param buffer and EMA buffer for evaluation. Remember to swap back after evaluation."""
	for group in self.param_groups:
	for param in group["params"]:
	state = self.state[param]
	if "param_ema" in state:
	temp = torch.empty_like(param).copy_(param)
	param.copy_(state["param_ema"])
	state["param_ema"].copy_(temp)