optimizer-model-consistency-pretrained / code /optimizer /activation_double_preconditioned_adam.py

Upload original AdamW and Muon pretrained checkpoints (step 31000)

ad93bc0 verified 1 day ago

5.69 kB

# activation_gradient_preconditioned_adam.py

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

class ActivationDoublePreconditionedAdam(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.999, 0.95), weight_decay=0.1, eps=1e-6):
        # 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)
        super(ActivationDoublePreconditionedAdam, 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']
                beta3_precond = beta2_adam
            else:
                beta1, beta2_adam, beta3_precond = group['betas']
            weight_decay = group['weight_decay']
            eps = group['eps']
            
            for p in group['params']:
                if p.grad is None:
                    continue
                # This optimizer needs attributes from both forward and backward hooks
                if not hasattr(p, 'activation_second_moment') or not hasattr(p, 'output_gradient_second_moment'):
                    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)
                    # L_t: Gradient preconditioner
                    state['grad_preconditioner'] = torch.zeros((p.shape[0], p.shape[0]), device=p.device)
                    # R_t: Activation preconditioner
                    state['activation_preconditioner'] = torch.zeros((p.shape[1], p.shape[1]), device=p.device)
                    # V'_t: Adam's 2nd moment buffer in the *projected* space
                    state['adam_v_prime'] = torch.zeros_like(p, memory_format=torch.preserve_format)

                state['step'] += 1
                step = state['step']
                
                momentum_buffer = state['momentum_buffer']
                grad_preconditioner = state['grad_preconditioner']
                activation_preconditioner = state['activation_preconditioner']
                adam_v_prime = state['adam_v_prime']

                xxt = p.activation_second_moment
                gygyt = p.output_gradient_second_moment

                # Step 5: Update momentum
                momentum_buffer.mul_(beta1).add_(grad, alpha=1 - beta1)

                # Step 6 & 7: Update preconditioners
                grad_preconditioner.mul_(beta3_precond).add_(gygyt, alpha=1 - beta3_precond)
                activation_preconditioner.mul_(beta3_precond).add_(xxt, alpha=1 - beta3_precond)

                # Step 8 & 9: Eigendecomposition of preconditioners
                # we use float32 for eigendecomposition to ensure numerical stability
                eps_L = torch.eye(grad_preconditioner.size(0), device=p.device) * eps
                # use svd instead of eigh for better numerical stability
                U, S_L, _ = torch.linalg.svd(grad_preconditioner.float() + eps_L)
                
                eps_R = torch.eye(activation_preconditioner.size(0), device=p.device) * eps
                # S_R, V = torch.linalg.eigh(activation_preconditioner.float() + eps_R)
                # use svd instead of eigh for better numerical stability
                V, S_R, _ = torch.linalg.svd(activation_preconditioner.float() + eps_R)

                # Step 10: Project momentum and gradient
                M_prime = torch.mm(U.t(), torch.mm(momentum_buffer, V))
                G_prime = torch.mm(U.t(), torch.mm(grad, V))

                # +++ ADDED: Store G_prime in state for external logging +++
                # will cost more memory, only use it for logging
                state['G_prime'] = G_prime

                # Step 11: Update Adam's second moment in projected space
                adam_v_prime.mul_(beta2_adam).addcmul_(G_prime, G_prime, value=1 - beta2_adam)

                # Bias Correction for Adam
                bias_correction1 = 1 - beta1 ** step
                bias_correction2 = 1 - beta2_adam ** step
                
                M_prime_hat = M_prime / bias_correction1
                adam_v_prime_hat = adam_v_prime / bias_correction2

                # Step 12: Run Adam in the preconditioned space
                denom = adam_v_prime_hat.sqrt().add_(eps)
                Q_t = M_prime_hat / denom

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

                # Step 13: Project back to do the update
                update = torch.mm(U, torch.mm(Q_t, V.t()))
                p.add_(update, alpha=-lr)

                # Cleanup hook attributes
                delattr(p, 'activation_second_moment')
                delattr(p, 'output_gradient_second_moment')
        return loss

	# activation_gradient_preconditioned_adam.py

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

	class ActivationDoublePreconditionedAdam(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.999, 0.95), weight_decay=0.1, eps=1e-6):
	# 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)
	super(ActivationDoublePreconditionedAdam, 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']
	beta3_precond = beta2_adam
	else:
	beta1, beta2_adam, beta3_precond = group['betas']
	weight_decay = group['weight_decay']
	eps = group['eps']

	for p in group['params']:
	if p.grad is None:
	continue
	# This optimizer needs attributes from both forward and backward hooks
	if not hasattr(p, 'activation_second_moment') or not hasattr(p, 'output_gradient_second_moment'):
	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)
	# L_t: Gradient preconditioner
	state['grad_preconditioner'] = torch.zeros((p.shape[0], p.shape[0]), device=p.device)
	# R_t: Activation preconditioner
	state['activation_preconditioner'] = torch.zeros((p.shape[1], p.shape[1]), device=p.device)
	# V'_t: Adam's 2nd moment buffer in the projected space
	state['adam_v_prime'] = torch.zeros_like(p, memory_format=torch.preserve_format)

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

	momentum_buffer = state['momentum_buffer']
	grad_preconditioner = state['grad_preconditioner']
	activation_preconditioner = state['activation_preconditioner']
	adam_v_prime = state['adam_v_prime']

	xxt = p.activation_second_moment
	gygyt = p.output_gradient_second_moment

	# Step 5: Update momentum
	momentum_buffer.mul_(beta1).add_(grad, alpha=1 - beta1)

	# Step 6 & 7: Update preconditioners
	grad_preconditioner.mul_(beta3_precond).add_(gygyt, alpha=1 - beta3_precond)
	activation_preconditioner.mul_(beta3_precond).add_(xxt, alpha=1 - beta3_precond)

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

	eps_R = torch.eye(activation_preconditioner.size(0), device=p.device) * eps
	# S_R, V = torch.linalg.eigh(activation_preconditioner.float() + eps_R)
	# use svd instead of eigh for better numerical stability
	V, S_R, _ = torch.linalg.svd(activation_preconditioner.float() + eps_R)

	# Step 10: Project momentum and gradient
	M_prime = torch.mm(U.t(), torch.mm(momentum_buffer, V))
	G_prime = torch.mm(U.t(), torch.mm(grad, V))

	# +++ ADDED: Store G_prime in state for external logging +++
	# will cost more memory, only use it for logging
	state['G_prime'] = G_prime

	# Step 11: Update Adam's second moment in projected space
	adam_v_prime.mul_(beta2_adam).addcmul_(G_prime, G_prime, value=1 - beta2_adam)

	# Bias Correction for Adam
	bias_correction1 = 1 - beta1 ** step
	bias_correction2 = 1 - beta2_adam ** step

	M_prime_hat = M_prime / bias_correction1
	adam_v_prime_hat = adam_v_prime / bias_correction2

	# Step 12: Run Adam in the preconditioned space
	denom = adam_v_prime_hat.sqrt().add_(eps)
	Q_t = M_prime_hat / denom

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

	# Step 13: Project back to do the update
	update = torch.mm(U, torch.mm(Q_t, V.t()))
	p.add_(update, alpha=-lr)

	# Cleanup hook attributes
	delattr(p, 'activation_second_moment')
	delattr(p, 'output_gradient_second_moment')
	return loss