Upload original AdamW and Muon pretrained checkpoints (step 31000)

ad93bc0 verified 1 day ago

5.36 kB

# activation_gradient_preconditioned_adam.py

import torch
import math
from torch.optim.optimizer import Optimizer

class ASGO(Optimizer):
    """
    Implements Activation-Gradient-Preconditioned Adam.
    This optimizer uses second-order statistics from both input activations (X_t)
    and output gradients (G_Yt) to form two preconditioners. This is closely
    related to the Shampoo optimizer.
    """
    def __init__(self, params, lr=1e-3, betas=(0.9, 0.8), weight_decay=0.1, eps=1e-10, scale=0.2, use_grafting=False, use_nesterov=False):
        # betas[0]: momentum (β1)
        # betas[1]: adam second moment (β2 in pseudocode, for V')
        # betas[2]: preconditioners (β3 in pseudocode, for L and R)
        # if not (isinstance(betas, (list, tuple)) and len(betas) == 3):
        #     raise ValueError("Betas must be a tuple/list of 3 elements.")
        
        defaults = dict(lr=lr, betas=betas, weight_decay=weight_decay, eps=eps, scale=scale, use_grafting=use_grafting, use_nesterov=use_nesterov)
        super(ASGO, self).__init__(params, defaults)

    @torch.no_grad()
    def step(self, closure=None):
        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()

        for group in self.param_groups:
            lr = group['lr']
            if len(group['betas']) == 2:
                beta1, beta2_adam = group['betas']
                beta2_precond = beta2_adam
            else:
                beta1, beta2_adam, beta2_precond = group['betas']
            weight_decay = group['weight_decay']
            eps = group['eps']
            scale = group['scale']
            use_grafting = group['use_grafting']
            use_nesterov = group['use_nesterov']
            
            for p in group['params']:
                if p.grad is None:
                    continue

                grad = p.grad
                state = self.state[p]
                
                if len(state) == 0:
                    state['step'] = 0
                    state['momentum_buffer'] = torch.zeros_like(p, memory_format=torch.preserve_format)
                    # R_t: Right preconditioner
                    state['right_preconditioner'] = torch.zeros((p.shape[-1], p.shape[-1]), device=p.device)
                    # V'_t: Adam's 2nd moment buffer for grafting
                    if use_grafting:
                        # beta2_adam = beta2_precond
                        state['exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)

                state['step'] += 1
                step = state['step']
                
                momentum_buffer = state['momentum_buffer']
                right_preconditioner = state['right_preconditioner']
                
                if use_grafting:
                    exp_avg_sq = state['exp_avg_sq']
                    exp_avg_sq.mul_(beta2_adam).addcmul_(grad, grad, value=1 - beta2_adam)


                
                GTG = torch.mm(grad.t(), grad)

                # Step 5: Update momentum
                momentum_buffer.mul_(beta1).add_(grad, alpha=1 - beta1)
                if use_nesterov:
                    Gt = torch.lerp(grad, momentum_buffer, beta1)
                else:
                    Gt = momentum_buffer

                # Step 6 & 7: Update preconditioners
                # left_preconditioner.mul_(beta2_precond).add_(GGT, alpha=1 - beta2_precond)
                right_preconditioner.mul_(beta2_precond).add_(GTG, alpha=1 - beta2_precond)

                # Step 8 & 9: Eigendecomposition of preconditioners
                # we use float32 for eigendecomposition to ensure numerical stability
                # eps_L = torch.eye(left_preconditioner.size(0), device=p.device) * eps
                # S_L, U = torch.linalg.eigh(left_preconditioner.float() + eps_L)
                
                eps_R = torch.eye(right_preconditioner.size(0), device=p.device) * eps
                
                # S_R, V = torch.linalg.eigh(right_preconditioner.float() + eps_R)
                # update_direction = Gt @ V @ torch.diag(S_R.rsqrt()) @ V.t()
                
                U, S_R, Vh = torch.linalg.svd(right_preconditioner.float() + eps_R)
                update_direction = Gt @ U @ torch.diag(S_R.rsqrt()) @ Vh

                
                # normalize the update
                update_direction = update_direction / (torch.linalg.norm(update_direction) + eps)

                if use_grafting:
                    bias_correction1 = 1 - beta1 ** step
                    bias_correction2 = 1 - beta2_adam ** step
                        
                    denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(eps)
                    # The magnitude is calculated based on the effective gradient
                    adam_update_magnitude = torch.linalg.norm((Gt / bias_correction1) / denom)
                    alpha = -lr * adam_update_magnitude
                
                else:
                    # Default scaling if not grafting
                    alpha = -lr * scale * math.sqrt(p.numel())

                
                if weight_decay > 0.0:
                    p.add_(p, alpha=-lr * weight_decay)

                
                p.add_(update_direction, alpha=alpha)

                
        return loss

	# activation_gradient_preconditioned_adam.py

	import torch
	import math
	from torch.optim.optimizer import Optimizer

	class ASGO(Optimizer):
	"""
	Implements Activation-Gradient-Preconditioned Adam.
	This optimizer uses second-order statistics from both input activations (X_t)
	and output gradients (G_Yt) to form two preconditioners. This is closely
	related to the Shampoo optimizer.
	"""
	def __init__(self, params, lr=1e-3, betas=(0.9, 0.8), weight_decay=0.1, eps=1e-10, scale=0.2, use_grafting=False, use_nesterov=False):
	# betas[0]: momentum (β1)
	# betas[1]: adam second moment (β2 in pseudocode, for V')
	# betas[2]: preconditioners (β3 in pseudocode, for L and R)
	# if not (isinstance(betas, (list, tuple)) and len(betas) == 3):
	# raise ValueError("Betas must be a tuple/list of 3 elements.")

	defaults = dict(lr=lr, betas=betas, weight_decay=weight_decay, eps=eps, scale=scale, use_grafting=use_grafting, use_nesterov=use_nesterov)
	super(ASGO, self).__init__(params, defaults)

	@torch.no_grad()
	def step(self, closure=None):
	loss = None
	if closure is not None:
	with torch.enable_grad():
	loss = closure()

	for group in self.param_groups:
	lr = group['lr']
	if len(group['betas']) == 2:
	beta1, beta2_adam = group['betas']
	beta2_precond = beta2_adam
	else:
	beta1, beta2_adam, beta2_precond = group['betas']
	weight_decay = group['weight_decay']
	eps = group['eps']
	scale = group['scale']
	use_grafting = group['use_grafting']
	use_nesterov = group['use_nesterov']

	for p in group['params']:
	if p.grad is None:
	continue

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

	if len(state) == 0:
	state['step'] = 0
	state['momentum_buffer'] = torch.zeros_like(p, memory_format=torch.preserve_format)
	# R_t: Right preconditioner
	state['right_preconditioner'] = torch.zeros((p.shape[-1], p.shape[-1]), device=p.device)
	# V'_t: Adam's 2nd moment buffer for grafting
	if use_grafting:
	# beta2_adam = beta2_precond
	state['exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)

	state['step'] += 1
	step = state['step']

	momentum_buffer = state['momentum_buffer']
	right_preconditioner = state['right_preconditioner']

	if use_grafting:
	exp_avg_sq = state['exp_avg_sq']
	exp_avg_sq.mul_(beta2_adam).addcmul_(grad, grad, value=1 - beta2_adam)



	GTG = torch.mm(grad.t(), grad)

	# Step 5: Update momentum
	momentum_buffer.mul_(beta1).add_(grad, alpha=1 - beta1)
	if use_nesterov:
	Gt = torch.lerp(grad, momentum_buffer, beta1)
	else:
	Gt = momentum_buffer

	# Step 6 & 7: Update preconditioners
	# left_preconditioner.mul_(beta2_precond).add_(GGT, alpha=1 - beta2_precond)
	right_preconditioner.mul_(beta2_precond).add_(GTG, alpha=1 - beta2_precond)

	# Step 8 & 9: Eigendecomposition of preconditioners
	# we use float32 for eigendecomposition to ensure numerical stability
	# eps_L = torch.eye(left_preconditioner.size(0), device=p.device) * eps
	# S_L, U = torch.linalg.eigh(left_preconditioner.float() + eps_L)

	eps_R = torch.eye(right_preconditioner.size(0), device=p.device) * eps

	# S_R, V = torch.linalg.eigh(right_preconditioner.float() + eps_R)
	# update_direction = Gt @ V @ torch.diag(S_R.rsqrt()) @ V.t()

	U, S_R, Vh = torch.linalg.svd(right_preconditioner.float() + eps_R)
	update_direction = Gt @ U @ torch.diag(S_R.rsqrt()) @ Vh


	# normalize the update
	update_direction = update_direction / (torch.linalg.norm(update_direction) + eps)

	if use_grafting:
	bias_correction1 = 1 - beta1 ** step
	bias_correction2 = 1 - beta2_adam ** step

	denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(eps)
	# The magnitude is calculated based on the effective gradient
	adam_update_magnitude = torch.linalg.norm((Gt / bias_correction1) / denom)
	alpha = -lr * adam_update_magnitude

	else:
	# Default scaling if not grafting
	alpha = -lr * scale * math.sqrt(p.numel())


	if weight_decay > 0.0:
	p.add_(p, alpha=-lr * weight_decay)


	p.add_(update_direction, alpha=alpha)


	return loss