Upload original AdamW and Muon pretrained checkpoints (step 31000)

ad93bc0 verified 1 day ago

11.6 kB

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

class EigenBlockOptimizer(Optimizer):
    """
    Implements Eigen-Block Basis Optimizer.
    
    Features:
    - Hessian Approximation via Shampoo/SOAP style preconditioners (L, R).
    - Eigen-Block Rotation: Groups close eigenvalues and rotates within blocks 
      to maximize L1 sparsity (using Newton-Schulz).
    - Momentum: Applies momentum to gradients in the original space.
    - Dual Mode: Supports both SignSGD (default) and Adam updates in the optimal basis.
    """
    def __init__(self, params, lr=1e-3, betas=(0.9, 0.95), weight_decay=0.0, 
                 eigen_tol=0.1, eps=1e-8, use_adam=False):
        """
        Args:
            lr (float): Learning rate.
            betas (tuple): (beta1, beta2). 
                           beta1: Momentum coefficient.
                           beta2: Preconditioner EMA coefficient (and Adam 2nd moment if use_adam=True).
            eigen_tol (float): Relative tolerance for grouping eigenvalues. 
                               Groups lambdas where (lambda_i - lambda_start) < tol * max_lambda.
            use_adam (bool): If True, uses Adam updates in the projected space.
                             If False, uses SignSGD updates in the projected space.
        """
        defaults = dict(lr=lr, betas=betas, weight_decay=weight_decay, 
                        eigen_tol=eigen_tol, eps=eps, use_adam=use_adam)
        super(EigenBlockOptimizer, 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']
            beta1, beta2 = group['betas']
            weight_decay = group['weight_decay']
            eigen_tol = group['eigen_tol']
            eps = group['eps']
            use_adam = group['use_adam']
            
            for p in group['params']:
                if p.grad is None:
                    continue

                grad = p.grad
                if grad.dim() != 2:
                    # Fallback for 1D vectors (bias/layernorm): Standard Adam or SignSGD
                    state = self.state[p]
                    if len(state) == 0:
                        state['step'] = 0
                        state['exp_avg'] = torch.zeros_like(p)
                        if use_adam:
                            state['exp_avg_sq'] = torch.zeros_like(p)

                    state['step'] += 1
                    exp_avg = state['exp_avg']
                    
                    # Update momentum
                    exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
                    
                    if weight_decay > 0:
                        p.mul_(1 - lr * weight_decay)

                    if use_adam:
                        exp_avg_sq = state['exp_avg_sq']
                        exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
                        denom = exp_avg_sq.sqrt().add_(eps)
                        
                        # Bias correction
                        bias_correction1 = 1 - beta1 ** state['step']
                        bias_correction2 = 1 - beta2 ** state['step']
                        step_size = lr * math.sqrt(bias_correction2) / bias_correction1
                        
                        p.addcdiv_(exp_avg, denom, value=-step_size)
                    else:
                        # SignSGD with momentum
                        p.add_(exp_avg.sign(), alpha=-lr)
                    continue

                # --- Handle 2D Matrices (Weights) ---
                state = self.state[p]
                
                # Init state
                if len(state) == 0:
                    state['step'] = 0
                    state['momentum_buffer'] = torch.zeros_like(p) # 1st moment (M)
                    # Shampoo/SOAP style preconditioners
                    state['L'] = torch.zeros((grad.size(0), grad.size(0)), device=grad.device, dtype=grad.dtype)
                    state['R'] = torch.zeros((grad.size(1), grad.size(1)), device=grad.device, dtype=grad.dtype)
                    
                    if use_adam:
                        # Adam 2nd moment (V') in the PROJECTED space
                        # Note: We store this same size as p, assuming basis consistency over time
                        state['adam_v_prime'] = torch.zeros_like(p)

                state['step'] += 1
                
                L = state['L']
                R = state['R']
                momentum = state['momentum_buffer']
                
                # 1. Update Momentum (M_t) in Original Space
                momentum.mul_(beta1).add_(grad, alpha=1 - beta1)

                # 2. Update Preconditioners (Hessian Approximation)
                # Using EMA of GG^T and G^TG
                L.mul_(beta2).add_(grad @ grad.T, alpha=1 - beta2)
                R.mul_(beta2).add_(grad.T @ grad, alpha=1 - beta2)

                # 3. Eigen Decomposition
                # Add 1e-12 for numerical stability
                L_evals, U_L = torch.linalg.eigh(L + 1e-12 * torch.eye(L.size(0), device=L.device))
                R_evals, U_R = torch.linalg.eigh(R + 1e-12 * torch.eye(R.size(0), device=R.device))
                
                # 4. Rotate Momentum to Eigen-basis: M_tilde = U_L^T M U_R
                # We rotate M instead of G because we are doing Momentum-based updates
                M_tilde = U_L.T @ momentum @ U_R

                # 5. Compute Optimal Basis Q_L (Left)
                Q_L = U_L.clone()
                # Use max eigenvalue to scale tolerance
                tol_L = eigen_tol * (L_evals.abs().max().item() + 1e-6)
                boundaries_L = self._find_block_boundaries(L_evals, tol_L)
                # --- NEW: Record L Block Stats ---
                self._store_block_stats(state, 'L', boundaries_L, L_evals.shape[0])
                
                for start, end in boundaries_L:
                    if end - start <= 1:
                        continue
                    # Extract rows corresponding to this block from M_tilde
                    # We use M_tilde to determine the rotation N within the block
                    M_sub = M_tilde[start:end, :]
                    
                    # Formula: N = msign(A * sign(A)^T)
                    # Using M_sub aligns the basis with the momentum direction's sparsity
                    A = M_sub @ M_sub.sign().T
                    N_block = zeropower_via_newtonschulz5(A)
                    
                    # Update Q_L columns: Q_L_new = Q_L_old @ N_block
                    Q_L[:, start:end] = U_L[:, start:end] @ N_block

                # 6. Compute Optimal Basis Q_R (Right)
                Q_R = U_R.clone()
                tol_R = eigen_tol * (R_evals.abs().max().item() + 1e-6)
                boundaries_R = self._find_block_boundaries(R_evals, tol_R)
                # --- NEW: Record R Block Stats ---
                self._store_block_stats(state, 'R', boundaries_R, R_evals.shape[0])
                
                for start, end in boundaries_R:
                    if end - start <= 1:
                        continue
                    M_sub = M_tilde[:, start:end]
                    
                    # Formula: N = msign(sign(A)^T * A)
                    A = M_sub.sign().T @ M_sub
                    N_block = zeropower_via_newtonschulz5(A)
                    
                    Q_R[:, start:end] = U_R[:, start:end] @ N_block

                # 7. Project Momentum and Gradient into the FULL optimal basis Q
                # M_opt = Q_L^T M Q_R
                M_opt = Q_L.T @ momentum @ Q_R
                
                if use_adam:
                    # If using Adam, we also need the projected Gradient to update V'
                    # G_opt = Q_L^T G Q_R
                    G_opt = Q_L.T @ grad @ Q_R
                    
                    adam_v_prime = state['adam_v_prime']
                    # Update V'_t using projected gradient square
                    adam_v_prime.mul_(beta2).addcmul_(G_opt, G_opt, value=1 - beta2)
                    
                    # Bias Correction
                    bias_correction1 = 1 - beta1 ** state['step']
                    bias_correction2 = 1 - beta2 ** state['step']
                    
                    M_hat = M_opt / bias_correction1
                    V_hat = adam_v_prime / bias_correction2
                    
                    # Adam Update in Basis
                    Update_inner = M_hat / (V_hat.sqrt() + eps)
                else:
                    # SignSGD with Momentum in Basis
                    Update_inner = M_opt.sign()
                
                # 8. Rotate update back to original space
                # D_t = Q_L * Update_inner * Q_R^T
                Update_step = Q_L @ Update_inner @ Q_R.T
                
                if weight_decay > 0:
                    p.mul_(1 - lr * weight_decay)
                
                p.add_(Update_step, alpha=-lr)

        return loss

    def _find_block_boundaries(self, evals, tol):
        """
        Groups eigenvalues based on difference from the FIRST element in the block.
        Iterates through sorted eigenvalues. If evals[i] - evals[start] > tol,
        a new block is started at i.
        """
        # evals is expected to be sorted (torch.linalg.eigh returns ascending)
        boundaries = [0]
        n = len(evals)
        
        if n == 0: return []
        
        start_idx = 0
        for i in range(1, n):
            # Check difference against the anchor (start_idx)
            diff = evals[i] - evals[start_idx]
            
            if diff > tol:
                # Close current block, start new one at i
                boundaries.append(i)
                start_idx = i
                
        boundaries.append(n)
        
        # Convert to (start, end) tuples
        blocks = []
        for k in range(len(boundaries) - 1):
            blocks.append((boundaries[k], boundaries[k+1]))
            
        return blocks
    
    def _store_block_stats(self, state, name, blocks, total_dim):
        """
        计算 10%, 25%, 50%, 75%, 90%, 100% 处的特征值索引属于第几个 Block，并存储到 state 中。
        """
        if 'eigen_stats' not in state:
            state['eigen_stats'] = {}
        
        # 1-based block index (第几个格子)
        stats = {}
        percentiles = [0.10, 0.25, 0.50, 0.75, 0.90, 1.00]
        
        for p in percentiles:
            # 计算对应的 index (0-based)
            # 例如 100% 对应最后一个元素 (total_dim - 1)
            target_idx = int(p * total_dim) - 1
            target_idx = max(0, min(target_idx, total_dim - 1))
            
            # 查找该 index 属于哪个 block
            found_block_idx = 0 # default
            for i, (start, end) in enumerate(blocks):
                if start <= target_idx < end:
                    found_block_idx = i + 1 # 使用 1-based 计数，方便观察
                    break
            
            # 记录 metric，例如 'p10': 1
            stats[f'p{int(p*100)}'] = found_block_idx
            
        stats['total_blocks'] = len(blocks)
        stats['total_dim'] = total_dim
        
        state['eigen_stats'][name] = stats

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

	class EigenBlockOptimizer(Optimizer):
	"""
	Implements Eigen-Block Basis Optimizer.

	Features:
	- Hessian Approximation via Shampoo/SOAP style preconditioners (L, R).
	- Eigen-Block Rotation: Groups close eigenvalues and rotates within blocks
	to maximize L1 sparsity (using Newton-Schulz).
	- Momentum: Applies momentum to gradients in the original space.
	- Dual Mode: Supports both SignSGD (default) and Adam updates in the optimal basis.
	"""
	def __init__(self, params, lr=1e-3, betas=(0.9, 0.95), weight_decay=0.0,
	eigen_tol=0.1, eps=1e-8, use_adam=False):
	"""
	Args:
	lr (float): Learning rate.
	betas (tuple): (beta1, beta2).
	beta1: Momentum coefficient.
	beta2: Preconditioner EMA coefficient (and Adam 2nd moment if use_adam=True).
	eigen_tol (float): Relative tolerance for grouping eigenvalues.
	Groups lambdas where (lambda_i - lambda_start) < tol * max_lambda.
	use_adam (bool): If True, uses Adam updates in the projected space.
	If False, uses SignSGD updates in the projected space.
	"""
	defaults = dict(lr=lr, betas=betas, weight_decay=weight_decay,
	eigen_tol=eigen_tol, eps=eps, use_adam=use_adam)
	super(EigenBlockOptimizer, 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']
	beta1, beta2 = group['betas']
	weight_decay = group['weight_decay']
	eigen_tol = group['eigen_tol']
	eps = group['eps']
	use_adam = group['use_adam']

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

	grad = p.grad
	if grad.dim() != 2:
	# Fallback for 1D vectors (bias/layernorm): Standard Adam or SignSGD
	state = self.state[p]
	if len(state) == 0:
	state['step'] = 0
	state['exp_avg'] = torch.zeros_like(p)
	if use_adam:
	state['exp_avg_sq'] = torch.zeros_like(p)

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

	# Update momentum
	exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)

	if weight_decay > 0:
	p.mul_(1 - lr * weight_decay)

	if use_adam:
	exp_avg_sq = state['exp_avg_sq']
	exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
	denom = exp_avg_sq.sqrt().add_(eps)

	# Bias correction
	bias_correction1 = 1 - beta1 ** state['step']
	bias_correction2 = 1 - beta2 ** state['step']
	step_size = lr * math.sqrt(bias_correction2) / bias_correction1

	p.addcdiv_(exp_avg, denom, value=-step_size)
	else:
	# SignSGD with momentum
	p.add_(exp_avg.sign(), alpha=-lr)
	continue

	# --- Handle 2D Matrices (Weights) ---
	state = self.state[p]

	# Init state
	if len(state) == 0:
	state['step'] = 0
	state['momentum_buffer'] = torch.zeros_like(p) # 1st moment (M)
	# Shampoo/SOAP style preconditioners
	state['L'] = torch.zeros((grad.size(0), grad.size(0)), device=grad.device, dtype=grad.dtype)
	state['R'] = torch.zeros((grad.size(1), grad.size(1)), device=grad.device, dtype=grad.dtype)

	if use_adam:
	# Adam 2nd moment (V') in the PROJECTED space
	# Note: We store this same size as p, assuming basis consistency over time
	state['adam_v_prime'] = torch.zeros_like(p)

	state['step'] += 1

	L = state['L']
	R = state['R']
	momentum = state['momentum_buffer']

	# 1. Update Momentum (M_t) in Original Space
	momentum.mul_(beta1).add_(grad, alpha=1 - beta1)

	# 2. Update Preconditioners (Hessian Approximation)
	# Using EMA of GG^T and G^TG
	L.mul_(beta2).add_(grad @ grad.T, alpha=1 - beta2)
	R.mul_(beta2).add_(grad.T @ grad, alpha=1 - beta2)

	# 3. Eigen Decomposition
	# Add 1e-12 for numerical stability
	L_evals, U_L = torch.linalg.eigh(L + 1e-12 * torch.eye(L.size(0), device=L.device))
	R_evals, U_R = torch.linalg.eigh(R + 1e-12 * torch.eye(R.size(0), device=R.device))

	# 4. Rotate Momentum to Eigen-basis: M_tilde = U_L^T M U_R
	# We rotate M instead of G because we are doing Momentum-based updates
	M_tilde = U_L.T @ momentum @ U_R

	# 5. Compute Optimal Basis Q_L (Left)
	Q_L = U_L.clone()
	# Use max eigenvalue to scale tolerance
	tol_L = eigen_tol * (L_evals.abs().max().item() + 1e-6)
	boundaries_L = self._find_block_boundaries(L_evals, tol_L)
	# --- NEW: Record L Block Stats ---
	self._store_block_stats(state, 'L', boundaries_L, L_evals.shape[0])

	for start, end in boundaries_L:
	if end - start <= 1:
	continue
	# Extract rows corresponding to this block from M_tilde
	# We use M_tilde to determine the rotation N within the block
	M_sub = M_tilde[start:end, :]

	# Formula: N = msign(A * sign(A)^T)
	# Using M_sub aligns the basis with the momentum direction's sparsity
	A = M_sub @ M_sub.sign().T
	N_block = zeropower_via_newtonschulz5(A)

	# Update Q_L columns: Q_L_new = Q_L_old @ N_block
	Q_L[:, start:end] = U_L[:, start:end] @ N_block

	# 6. Compute Optimal Basis Q_R (Right)
	Q_R = U_R.clone()
	tol_R = eigen_tol * (R_evals.abs().max().item() + 1e-6)
	boundaries_R = self._find_block_boundaries(R_evals, tol_R)
	# --- NEW: Record R Block Stats ---
	self._store_block_stats(state, 'R', boundaries_R, R_evals.shape[0])

	for start, end in boundaries_R:
	if end - start <= 1:
	continue
	M_sub = M_tilde[:, start:end]

	# Formula: N = msign(sign(A)^T * A)
	A = M_sub.sign().T @ M_sub
	N_block = zeropower_via_newtonschulz5(A)

	Q_R[:, start:end] = U_R[:, start:end] @ N_block

	# 7. Project Momentum and Gradient into the FULL optimal basis Q
	# M_opt = Q_L^T M Q_R
	M_opt = Q_L.T @ momentum @ Q_R

	if use_adam:
	# If using Adam, we also need the projected Gradient to update V'
	# G_opt = Q_L^T G Q_R
	G_opt = Q_L.T @ grad @ Q_R

	adam_v_prime = state['adam_v_prime']
	# Update V'_t using projected gradient square
	adam_v_prime.mul_(beta2).addcmul_(G_opt, G_opt, value=1 - beta2)

	# Bias Correction
	bias_correction1 = 1 - beta1 ** state['step']
	bias_correction2 = 1 - beta2 ** state['step']

	M_hat = M_opt / bias_correction1
	V_hat = adam_v_prime / bias_correction2

	# Adam Update in Basis
	Update_inner = M_hat / (V_hat.sqrt() + eps)
	else:
	# SignSGD with Momentum in Basis
	Update_inner = M_opt.sign()

	# 8. Rotate update back to original space
	# D_t = Q_L * Update_inner * Q_R^T
	Update_step = Q_L @ Update_inner @ Q_R.T

	if weight_decay > 0:
	p.mul_(1 - lr * weight_decay)

	p.add_(Update_step, alpha=-lr)

	return loss

	def _find_block_boundaries(self, evals, tol):
	"""
	Groups eigenvalues based on difference from the FIRST element in the block.
	Iterates through sorted eigenvalues. If evals[i] - evals[start] > tol,
	a new block is started at i.
	"""
	# evals is expected to be sorted (torch.linalg.eigh returns ascending)
	boundaries = [0]
	n = len(evals)

	if n == 0: return []

	start_idx = 0
	for i in range(1, n):
	# Check difference against the anchor (start_idx)
	diff = evals[i] - evals[start_idx]

	if diff > tol:
	# Close current block, start new one at i
	boundaries.append(i)
	start_idx = i

	boundaries.append(n)

	# Convert to (start, end) tuples
	blocks = []
	for k in range(len(boundaries) - 1):
	blocks.append((boundaries[k], boundaries[k+1]))

	return blocks

	def _store_block_stats(self, state, name, blocks, total_dim):
	"""
	计算 10%, 25%, 50%, 75%, 90%, 100% 处的特征值索引属于第几个 Block，并存储到 state 中。
	"""
	if 'eigen_stats' not in state:
	state['eigen_stats'] = {}

	# 1-based block index (第几个格子)
	stats = {}
	percentiles = [0.10, 0.25, 0.50, 0.75, 0.90, 1.00]

	for p in percentiles:
	# 计算对应的 index (0-based)
	# 例如 100% 对应最后一个元素 (total_dim - 1)
	target_idx = int(p * total_dim) - 1
	target_idx = max(0, min(target_idx, total_dim - 1))

	# 查找该 index 属于哪个 block
	found_block_idx = 0 # default
	for i, (start, end) in enumerate(blocks):
	if start <= target_idx < end:
	found_block_idx = i + 1 # 使用 1-based 计数，方便观察
	break

	# 记录 metric，例如 'p10': 1
	stats[f'p{int(p*100)}'] = found_block_idx

	stats['total_blocks'] = len(blocks)
	stats['total_dim'] = total_dim

	state['eigen_stats'][name] = stats