toolkit/optimizers/automagic.py

from typing import List
import torch
from toolkit.optimizers.optimizer_utils import Auto8bitTensor, copy_stochastic, stochastic_grad_accummulation
from optimum.quanto import QBytesTensor
import random


class Automagic(torch.optim.Optimizer):
    def __init__(
        self,
        params,
        lr=1e-6, # lr is start lr
        min_lr=1e-7,
        max_lr=1e-3,
        lr_bump=1e-6, # amount to bump the lr when adjusting
        eps=(1e-30, 1e-3),
        clip_threshold=1.0,
        beta2=0.999,
        weight_decay=0.0,
        do_paramiter_swapping=False,
        paramiter_swapping_factor=0.1,
    ):
        self.lr = lr
        if self.lr > 1e-3:
            print(f"Warning! Start lr is very high: {self.lr}. Forcing to 1e-6. this does not work like prodigy")
            self.lr = 1e-6
        self.min_lr = min_lr
        self.max_lr = max_lr
        self.lr_bump = lr_bump

        defaults = {
            "lr": lr,
            "eps": eps,
            "clip_threshold": clip_threshold,
            "beta2": beta2,
            "weight_decay": weight_decay,
        }
        super().__init__(params, defaults)

        self.base_lrs: List[float] = [
            lr for group in self.param_groups
        ]

        self.is_stochastic_rounding_accumulation = False

        # setup stochastic grad accum hooks
        for group in self.param_groups:
            for param in group['params']:
                if param.requires_grad and param.dtype != torch.float32:
                    self.is_stochastic_rounding_accumulation = True
                    param.register_post_accumulate_grad_hook(
                        stochastic_grad_accummulation
                    )

        self.do_paramiter_swapping = do_paramiter_swapping
        self.paramiter_swapping_factor = paramiter_swapping_factor
        self._total_paramiter_size = 0
        # count total paramiters
        for group in self.param_groups:
            for param in group['params']:
                self._total_paramiter_size += torch.numel(param)
        # pretty print total paramiters with comma seperation
        print(f"Total training paramiters: {self._total_paramiter_size:,}")

        # needs to be enabled to count paramiters
        if self.do_paramiter_swapping:
            self.enable_paramiter_swapping(self.paramiter_swapping_factor)

    def enable_paramiter_swapping(self, paramiter_swapping_factor=0.1):
        self.do_paramiter_swapping = True
        self.paramiter_swapping_factor = paramiter_swapping_factor
        # call it an initial time
        self.swap_paramiters()

    def swap_paramiters(self):
        all_params = []
        # deactivate all paramiters
        for group in self.param_groups:
            for param in group['params']:
                param.requires_grad_(False)
                # remove any grad
                param.grad = None
                all_params.append(param)
        # shuffle all paramiters
        random.shuffle(all_params)

        # keep activating paramiters until we are going to go over the target paramiters
        target_paramiters = int(
            self._total_paramiter_size * self.paramiter_swapping_factor)
        total_paramiters = 0
        for param in all_params:
            total_paramiters += torch.numel(param)
            if total_paramiters >= target_paramiters:
                break
            else:
                param.requires_grad_(True)

    @staticmethod
    def _get_lr(param_group, param_state):
        if 'avg_lr' in param_state:
            lr = param_state["avg_lr"]
        else:
            lr = 0.0
        return lr

    def _get_group_lr(self, group):
        group_lrs = []
        for p in group["params"]:
            group_lrs.append(self._get_lr(group, self.state[p]))
        # return avg
        if len(group_lrs) == 0:
            return self.lr
        return sum(group_lrs) / len(group_lrs)

    @staticmethod
    def _rms(tensor):
        return tensor.norm(2) / (tensor.numel() ** 0.5)

    @staticmethod
    def _approx_sq_grad(exp_avg_sq_row, exp_avg_sq_col):
        r_factor = (exp_avg_sq_row / exp_avg_sq_row.mean(dim=-
                    1, keepdim=True)).rsqrt_().unsqueeze(-1)
        c_factor = exp_avg_sq_col.unsqueeze(-2).rsqrt()
        return torch.mul(r_factor, c_factor)

    def step_hook(self):
        if not self.is_stochastic_rounding_accumulation:
            return
        # copy over stochastically rounded grads
        for group in self.param_groups:
            for param in group['params']:
                if param.requires_grad and hasattr(param, "_accum_grad"):
                    param.grad = param._accum_grad
                    del param._accum_grad

    # automagic manages its own lr
    def get_learning_rates(self):

        lrs = [
            self._get_group_lr(group)
            for group in self.param_groups
        ]
        if len(lrs) == 0:
            lrs = self.base_lrs  # if called before stepping
        return lrs

    def get_avg_learning_rate(self):
        lrs = self.get_learning_rates()
        return sum(lrs) / len(lrs)

    @torch.no_grad()
    def step(self, closure=None):
        """
        Performs a single optimization step

        Arguments:
            closure (callable, optional): A closure that reevaluates the model
                and returns the loss.
        """
        self.step_hook()
        loss = None
        if closure is not None:
            loss = closure()

        for group in self.param_groups:
            for p in group["params"]:
                if p.grad is None or not p.requires_grad:
                    continue

                grad = p.grad
                if grad.dtype != torch.float32:
                    grad = grad.to(torch.float32)
                if grad.is_sparse:
                    raise RuntimeError(
                        "Automagic does not support sparse gradients.")

                state = self.state[p]
                grad_shape = grad.shape

                factored = len(grad_shape) >= 2
                # State Initialization
                if len(state) == 0:
                    self.initialize_state(p)
                else:
                    # Check if exp_avg_sq_row and exp_avg_sq_col exist for factored case
                    if factored:
                        if "exp_avg_sq_row" not in state or "exp_avg_sq_col" not in state:
                            state["exp_avg_sq_row"] = torch.zeros(p.shape[:-1]).to(grad)
                            state["exp_avg_sq_col"] = torch.zeros(p.shape[:-2] + p.shape[-1:]).to(grad)
                        else:
                            state["exp_avg_sq_row"] = state["exp_avg_sq_row"].to(grad)
                            state["exp_avg_sq_col"] = state["exp_avg_sq_col"].to(grad)
                    # Check if exp_avg_sq exists for non-factored case
                    else:
                        if "exp_avg_sq" not in state:
                            state["exp_avg_sq"] = torch.zeros_like(grad)
                        else:
                            state["exp_avg_sq"] = state["exp_avg_sq"].to(grad)

                p_data_fp32 = p

                if isinstance(p_data_fp32, QBytesTensor):
                    p_data_fp32 = p_data_fp32.dequantize()
                if p.dtype != torch.float32:
                    p_data_fp32 = p_data_fp32.clone().float()

                # Initialize step if it doesn't exist
                if "step" not in state:
                    state["step"] = 0
                state["step"] += 1
                state["RMS"] = self._rms(p_data_fp32)

                # Use fixed beta2 from group instead of decay_rate calculation
                beta2 = group["beta2"]
                eps = group["eps"]
                if isinstance(eps, tuple) or isinstance(eps, list):
                    eps = eps[0]
                update = (grad**2) + eps
                if factored:
                    exp_avg_sq_row = state["exp_avg_sq_row"]
                    exp_avg_sq_col = state["exp_avg_sq_col"]

                    exp_avg_sq_row.mul_(beta2).add_(
                        update.mean(dim=-1), alpha=(1.0 - beta2))
                    exp_avg_sq_col.mul_(beta2).add_(
                        update.mean(dim=-2), alpha=(1.0 - beta2))

                    # Approximation of exponential moving average of square of gradient
                    update = self._approx_sq_grad(
                        exp_avg_sq_row, exp_avg_sq_col)
                    update.mul_(grad)
                else:
                    exp_avg_sq = state["exp_avg_sq"]

                    exp_avg_sq.mul_(beta2).add_(update, alpha=(1.0 - beta2))
                    update = exp_avg_sq.rsqrt().mul_(grad)

                update.div_(
                    (self._rms(update) / group["clip_threshold"]).clamp_(min=1.0))

                # Ensure state is properly initialized
                if 'last_polarity' not in state or 'lr_mask' not in state:
                    self.initialize_state(p)
                
                # Get signs of current last update and updates
                last_polarity = state['last_polarity']
                current_polarity = (update > 0).to(torch.bool)
                sign_agreement = torch.where(
                    last_polarity == current_polarity, 1, -1)
                state['last_polarity'] = current_polarity

                lr_mask = state['lr_mask'].to(torch.float32)

                # Update learning rate mask based on sign agreement
                new_lr = torch.where(
                    sign_agreement > 0,
                    lr_mask + self.lr_bump,  # Increase lr
                    lr_mask - self.lr_bump  # Decrease lr
                )

                # Clip learning rates to bounds
                new_lr = torch.clamp(
                    new_lr,
                    min=self.min_lr,
                    max=self.max_lr
                )

                # Apply the learning rate mask to the update
                update.mul_(new_lr)

                state['lr_mask'] = Auto8bitTensor(new_lr)
                state['avg_lr'] = torch.mean(new_lr)

                if group["weight_decay"] != 0:
                    # Apply weight decay with per-parameter learning rates
                    # Instead of using add_ with a tensor alpha (which isn't supported),
                    # we'll use element-wise multiplication to apply the weight decay
                    weight_decay_update = p_data_fp32 * (-group["weight_decay"]) * new_lr
                    p_data_fp32.add_(weight_decay_update)

                p_data_fp32.add_(-update)

                if p.dtype != torch.float32:
                    # apply stochastic rounding
                    copy_stochastic(p, p_data_fp32)

        return loss
    
    def initialize_state(self, p):
        state = self.state[p]
        state["step"] = 0

        # store the lr mask
        if 'lr_mask' not in state:
            state['lr_mask'] = Auto8bitTensor(torch.ones(
                p.shape).to(p.device, dtype=torch.float32) * self.lr
            )
        state['avg_lr'] = torch.mean(
            state['lr_mask'].to(torch.float32))
        if 'last_polarity' not in state:
            state['last_polarity'] = torch.zeros(
                p.shape, dtype=torch.bool, device=p.device)
        
        factored = len(p.shape) >= 2
        if factored:
            state["exp_avg_sq_row"] = torch.zeros(
                p.shape[:-1]).to(p)
            state["exp_avg_sq_col"] = torch.zeros(
                p.shape[:-2] + p.shape[-1:]).to(p)
        else:
            state["exp_avg_sq"] = torch.zeros_like(p)

        state["RMS"] = 0
    
    # override the state_dict to save the lr_mask
    def state_dict(self, *args, **kwargs):
        orig_state_dict = super().state_dict(*args, **kwargs)
        # convert the state to quantized tensor to scale and quantized
        new_sace_state = {}
        for p, state in orig_state_dict['state'].items():
            save_state = {k: v for k, v in state.items() if k != 'lr_mask'}
            
            # Check if lr_mask exists in the state before trying to access it
            if 'lr_mask' in state:
                save_state['lr_mask'] = state['lr_mask'].state_dict()
            
            new_sace_state[p] = save_state
            
        orig_state_dict['state'] = new_sace_state
        
        return orig_state_dict
    
    def load_state_dict(self, state_dict, strict=True):
        # Validate that the state_dict is from an Automagic optimizer
        is_valid_automagic_state = False
        
        # Check if state_dict has the expected structure
        if 'state' in state_dict and isinstance(state_dict['state'], dict):
            # Check if at least one state entry has an lr_mask, which is specific to Automagic
            for param_id, param_state in state_dict['state'].items():
                if isinstance(param_state, dict) and 'lr_mask' in param_state:
                    is_valid_automagic_state = True
                    break
        
        if not is_valid_automagic_state:
            return
        
        # First, call the parent class's load_state_dict to load the basic optimizer state
        # We'll handle the lr_mask separately
        state_dict_copy = {
            'state': {},
            'param_groups': state_dict['param_groups']
        }
        
        # Copy all state entries except lr_mask
        for param_id, param_state in state_dict['state'].items():
            state_dict_copy['state'][param_id] = {
                k: v for k, v in param_state.items() if k != 'lr_mask'
            }
        
        # Call parent class load_state_dict with the modified state dict
        super().load_state_dict(state_dict_copy)
        
        # Now handle the lr_mask separately
        # We need to map the saved parameters to the current parameters
        # This is tricky because the parameter IDs might be different
        
        # Get all current parameters that require gradients
        current_params = []
        for group in self.param_groups:
            for p in group['params']:
                if p.requires_grad:
                    current_params.append(p)
        
        # If the number of parameters doesn't match, we can't reliably map them
        if len(current_params) != len(state_dict['param_groups'][0]['params']):
            print(f"WARNING: Number of parameters doesn't match between saved state ({len(state_dict['param_groups'][0]['params'])}) "
                  f"and current model ({len(current_params)}). Learning rate masks may not be correctly loaded.")
        
        # Map parameters by their position in the param_groups
        # This assumes the order of parameters is preserved between saving and loading
        saved_param_ids = list(state_dict['state'].keys())
        
        for i, current_param in enumerate(current_params):
            if i >= len(saved_param_ids):
                break
                
            saved_param_id = saved_param_ids[i]
            saved_state = state_dict['state'][saved_param_id]
            
            # Skip if this saved state doesn't have an lr_mask
            if 'lr_mask' not in saved_state:
                continue
                
            # Initialize the state for this parameter if it doesn't exist
            if current_param not in self.state:
                self.initialize_state(current_param)
                
            # Get the current state for this parameter
            current_state = self.state[current_param]
            
            # Load the lr_mask from the saved state
            saved_lr_mask = saved_state['lr_mask']
            
            # Reconstruct the Auto8bitTensor from its state dict
            try:
                # Make sure the shapes match
                if 'quantized' in saved_lr_mask and saved_lr_mask['quantized'].shape == current_param.shape:
                    current_state['lr_mask'] = Auto8bitTensor(saved_lr_mask)
                else:
                    print(f"WARNING: Shape mismatch for parameter {i}. "
                          f"Expected {current_param.shape}, got {saved_lr_mask['quantized'].shape if 'quantized' in saved_lr_mask else 'unknown'}. "
                          f"Initializing new lr_mask.")
                    # Initialize a new lr_mask
                    current_state['lr_mask'] = Auto8bitTensor(torch.ones(
                        current_param.shape).to(current_param.device, dtype=torch.float32) * self.lr
                    )
            except Exception as e:
                print(f"ERROR: Failed to load lr_mask for parameter {i}: {e}")
                # Initialize a new lr_mask
                current_state['lr_mask'] = Auto8bitTensor(torch.ones(
                    current_param.shape).to(current_param.device, dtype=torch.float32) * self.lr
                )