Upload original AdamW and Muon pretrained checkpoints (step 31000)

ad93bc0 verified 1 day ago

9.42 kB

"""
Implementation of a hybrid optimizer.
For high-dimensional parameters (tensors with more than 2 dimensions), this optimizer
applies steepest descent with respect to a specified matrix norm (induced 1, 2, or infinity norm).
The update rules are derived from the provided formulas.
For all other parameters (2D or 1D), it defaults to the Adam optimization algorithm.
"""

import torch
import math
from optimizer.numerical_methods import zeropower_via_newtonschulz5

class MatrixNormSteepestDescent(torch.optim.Optimizer):
    """
    Implements a hybrid optimization algorithm.
    It uses matrix norm steepest descent for tensors with dim > 2, 
    and Adam for all other parameters.
    """
    def __init__(self, params, norm_p=2, lr=1e-3, momentum=0.95, betas=(0.9, 0.999), weight_decay=0.1, scale=0.2, eps=1e-8, use_nesterov=False):
        """
        Initializes the optimizer.
        Args:
            params (iterable): Iterable of parameters to optimize or dicts defining parameter groups.
            norm_p (int or float): The matrix norm to use for high-dimensional parameters. 
                                   Supported values are 1, 2, and float('inf').
            lr (float, optional): Learning rate (default: 1e-3).
            betas (Tuple[float, float], optional): Coefficients used for computing running averages 
                                                   of gradient and its square in Adam (default: (0.9, 0.999)).
                                                   beta1 is also used as momentum for the matrix norm updates.
            weight_decay (float, optional): Weight decay (L2 penalty) (default: 0.1).
            eps (float, optional): Term added to the denominator to improve numerical stability (default: 1e-8).
            use_nesterov (bool, optional): Enables Nesterov momentum for matrix norm updates (default: False).
        """
        if lr < 0.0:
            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 weight_decay < 0.0:
            raise ValueError(f"Invalid weight_decay value: {weight_decay}")

        defaults = dict(norm_p=norm_p, lr=lr, momentum=momentum, betas=betas, weight_decay=weight_decay, scale=scale, eps=eps, use_nesterov=use_nesterov)
        super(MatrixNormSteepestDescent, self).__init__(params, defaults)

    def __setstate__(self, state):
        super(MatrixNormSteepestDescent, self).__setstate__(state)

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

        for group in self.param_groups:
            lr = group['lr']
            momentum = group['momentum']
            beta1, beta2 = group['betas']
            weight_decay = group['weight_decay']
            scale = group['scale']
            eps = group['eps']
            norm_p = group['norm_p']
            use_nesterov = group['use_nesterov']

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

                # Apply decoupled weight decay for all parameters
                if weight_decay != 0:
                    p.add_(p, alpha=-weight_decay * lr)
                
                param_state = self.state[p]

                # Check parameter dimension to decide update rule
                if p.dim() >= 2:
                    # --- Matrix Norm Steepest Descent for Tensors with dim >= 2 ---

                    # Momentum calculation
                    if momentum != 0:
                        if 'momentum_buffer' not in param_state:
                            buf = param_state['momentum_buffer'] = torch.clone(grad).detach()
                        else:
                            buf = param_state['momentum_buffer']
                            buf.mul_(momentum).add_(grad, alpha=1 - momentum)
                        
                        if use_nesterov:
                            Gt = torch.lerp(grad, buf, momentum)
                        else:
                            Gt = buf
                    else:
                        Gt = grad

                    # original_shape = Gt.shape
                    # # Matricize the tensor by reshaping: (d1, d_2*d_3*...)
                    # matrix_Gt = Gt.reshape(original_shape[0], -1)
                    
                    update_magnitude = 0
                    update_direction = torch.zeros_like(Gt)

                    if norm_p == 1:
                        # Based on the formula for the induced 1-norm update
                        col_max_abs = torch.max(torch.abs(Gt), dim=-2, keepdim=True)[0]
                        update_magnitude = torch.sum(col_max_abs)
                        
                        # Direction is ColumnMax(Gt)
                        mask = (torch.abs(Gt) >= col_max_abs - 1e-12)
                        update_direction[mask] = torch.sign(Gt[mask])
                        
                    elif norm_p == float('inf'):
                        # Based on the formula for the induced infinity-norm update
                        row_max_abs = torch.max(torch.abs(Gt), dim=-1, keepdim=True)[0]
                        update_magnitude = torch.sum(row_max_abs)
                        
                        # Direction is RowMax(Gt)
                        mask = (torch.abs(Gt) >= row_max_abs - 1e-12)
                        update_direction[mask] = torch.sign(Gt[mask])
                        
                    elif norm_p == 2:
                        # Implement Muon (similar to the https://arxiv.org/pdf/2502.16982 version since we have the normalization)
                        try:
                            update_direction = zeropower_via_newtonschulz5(G=Gt)

                        except torch.linalg.LinAlgError:
                            # If NewtonSchulz fails, we skip the update for this parameter to avoid crashing.
                            update_magnitude = 0
                            update_direction.fill_(0)
                    
                    elif norm_p == '2,inf':
                        # +++ NEW: Row-wise Normalized Gradient Descent +++
                        # Formula: G_i / ||G_i||_2 for each row i
                        
                        # 1. Calculate L2 norm for each row (dim=-1)
                        # Gt shape is likely (Out, In) or (B, Out, In) etc. 
                        # We assume the last dimension is the "input" dimension for summation
                        # row_norms = torch.norm(Gt, p=2, dim=-1, keepdim=True)

                        # !!!!!!! here fore embedding layers, we consider using col-wise normalization

                        # This is also sth interesting: column-wise normalization is 1,2 induced steepest descent
                        col_norms = torch.norm(Gt, p=2, dim=-1, keepdim=True)

                        # print(f'The shape of column norm: {col_norms.shape}')
                        
                        # 2. Normalize gradient rows
                        # update_direction = Gt / (row_norms + 1e-12)
                        update_direction = Gt / (col_norms + 1e-12)
                    
                    else:
                        raise ValueError(f'Invalid matrix norm choice for p.dim() >= 2. Only 1, 2, and inf are supported.')

                    # Combine magnitude and direction and reshape back to original tensor shape
                    # final_update = (update_direction * update_magnitude).reshape(original_shape)
                    update_direction = update_direction / (torch.linalg.norm(update_direction) + 1e-12)
                    alpha = -lr * scale * math.sqrt(p.numel())
                    p.add_(update_direction, alpha=alpha)
                    
                else:
                    # --- Adam Fallback for Params with dim == 1 ---

                    # Initialize Adam state
                    if 'step' not in param_state:
                        param_state['step'] = 0
                        param_state['exp_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format)
                        param_state['exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)

                    exp_avg = param_state['exp_avg']
                    exp_avg_sq = param_state['exp_avg_sq']
                    
                    param_state['step'] += 1
                    step = param_state['step']

                    # Adam update logic
                    exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
                    exp_avg_sq.mul_(beta2).addcmul_(grad, grad.conj(), value=1 - beta2)
                    
                    bias_correction1 = 1 - beta1 ** step
                    bias_correction2 = 1 - beta2 ** step
                    
                    step_size = lr / bias_correction1
                    
                    denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(eps)
                    
                    p.addcdiv_(exp_avg, denom, value=-step_size)

        return loss

	"""
	Implementation of a hybrid optimizer.

	For high-dimensional parameters (tensors with more than 2 dimensions), this optimizer
	applies steepest descent with respect to a specified matrix norm (induced 1, 2, or infinity norm).
	The update rules are derived from the provided formulas.

	For all other parameters (2D or 1D), it defaults to the Adam optimization algorithm.
	"""

	import torch
	import math
	from optimizer.numerical_methods import zeropower_via_newtonschulz5

	class MatrixNormSteepestDescent(torch.optim.Optimizer):
	"""
	Implements a hybrid optimization algorithm.

	It uses matrix norm steepest descent for tensors with dim > 2,
	and Adam for all other parameters.
	"""
	def __init__(self, params, norm_p=2, lr=1e-3, momentum=0.95, betas=(0.9, 0.999), weight_decay=0.1, scale=0.2, eps=1e-8, use_nesterov=False):
	"""
	Initializes the optimizer.

	Args:
	params (iterable): Iterable of parameters to optimize or dicts defining parameter groups.
	norm_p (int or float): The matrix norm to use for high-dimensional parameters.
	Supported values are 1, 2, and float('inf').
	lr (float, optional): Learning rate (default: 1e-3).
	betas (Tuple[float, float], optional): Coefficients used for computing running averages
	of gradient and its square in Adam (default: (0.9, 0.999)).
	beta1 is also used as momentum for the matrix norm updates.
	weight_decay (float, optional): Weight decay (L2 penalty) (default: 0.1).
	eps (float, optional): Term added to the denominator to improve numerical stability (default: 1e-8).
	use_nesterov (bool, optional): Enables Nesterov momentum for matrix norm updates (default: False).
	"""
	if lr < 0.0:
	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 weight_decay < 0.0:
	raise ValueError(f"Invalid weight_decay value: {weight_decay}")

	defaults = dict(norm_p=norm_p, lr=lr, momentum=momentum, betas=betas, weight_decay=weight_decay, scale=scale, eps=eps, use_nesterov=use_nesterov)
	super(MatrixNormSteepestDescent, self).__init__(params, defaults)

	def __setstate__(self, state):
	super(MatrixNormSteepestDescent, self).__setstate__(state)

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

	for group in self.param_groups:
	lr = group['lr']
	momentum = group['momentum']
	beta1, beta2 = group['betas']
	weight_decay = group['weight_decay']
	scale = group['scale']
	eps = group['eps']
	norm_p = group['norm_p']
	use_nesterov = group['use_nesterov']

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

	grad = p.grad

	# Apply decoupled weight decay for all parameters
	if weight_decay != 0:
	p.add_(p, alpha=-weight_decay * lr)

	param_state = self.state[p]

	# Check parameter dimension to decide update rule
	if p.dim() >= 2:
	# --- Matrix Norm Steepest Descent for Tensors with dim >= 2 ---

	# Momentum calculation
	if momentum != 0:
	if 'momentum_buffer' not in param_state:
	buf = param_state['momentum_buffer'] = torch.clone(grad).detach()
	else:
	buf = param_state['momentum_buffer']
	buf.mul_(momentum).add_(grad, alpha=1 - momentum)

	if use_nesterov:
	Gt = torch.lerp(grad, buf, momentum)
	else:
	Gt = buf
	else:
	Gt = grad

	# original_shape = Gt.shape
	# # Matricize the tensor by reshaping: (d1, d_2d_3...)
	# matrix_Gt = Gt.reshape(original_shape[0], -1)

	update_magnitude = 0
	update_direction = torch.zeros_like(Gt)

	if norm_p == 1:
	# Based on the formula for the induced 1-norm update
	col_max_abs = torch.max(torch.abs(Gt), dim=-2, keepdim=True)[0]
	update_magnitude = torch.sum(col_max_abs)

	# Direction is ColumnMax(Gt)
	mask = (torch.abs(Gt) >= col_max_abs - 1e-12)
	update_direction[mask] = torch.sign(Gt[mask])

	elif norm_p == float('inf'):
	# Based on the formula for the induced infinity-norm update
	row_max_abs = torch.max(torch.abs(Gt), dim=-1, keepdim=True)[0]
	update_magnitude = torch.sum(row_max_abs)

	# Direction is RowMax(Gt)
	mask = (torch.abs(Gt) >= row_max_abs - 1e-12)
	update_direction[mask] = torch.sign(Gt[mask])

	elif norm_p == 2:
	# Implement Muon (similar to the https://arxiv.org/pdf/2502.16982 version since we have the normalization)
	try:
	update_direction = zeropower_via_newtonschulz5(G=Gt)

	except torch.linalg.LinAlgError:
	# If NewtonSchulz fails, we skip the update for this parameter to avoid crashing.
	update_magnitude = 0
	update_direction.fill_(0)

	elif norm_p == '2,inf':
	# +++ NEW: Row-wise Normalized Gradient Descent +++
	# Formula: G_i / \|\|G_i\|\|_2 for each row i

	# 1. Calculate L2 norm for each row (dim=-1)
	# Gt shape is likely (Out, In) or (B, Out, In) etc.
	# We assume the last dimension is the "input" dimension for summation
	# row_norms = torch.norm(Gt, p=2, dim=-1, keepdim=True)

	# !!!!!!! here fore embedding layers, we consider using col-wise normalization

	# This is also sth interesting: column-wise normalization is 1,2 induced steepest descent
	col_norms = torch.norm(Gt, p=2, dim=-1, keepdim=True)

	# print(f'The shape of column norm: {col_norms.shape}')

	# 2. Normalize gradient rows
	# update_direction = Gt / (row_norms + 1e-12)
	update_direction = Gt / (col_norms + 1e-12)

	else:
	raise ValueError(f'Invalid matrix norm choice for p.dim() >= 2. Only 1, 2, and inf are supported.')

	# Combine magnitude and direction and reshape back to original tensor shape
	# final_update = (update_direction * update_magnitude).reshape(original_shape)
	update_direction = update_direction / (torch.linalg.norm(update_direction) + 1e-12)
	alpha = -lr * scale * math.sqrt(p.numel())
	p.add_(update_direction, alpha=alpha)

	else:
	# --- Adam Fallback for Params with dim == 1 ---

	# Initialize Adam state
	if 'step' not in param_state:
	param_state['step'] = 0
	param_state['exp_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format)
	param_state['exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)

	exp_avg = param_state['exp_avg']
	exp_avg_sq = param_state['exp_avg_sq']

	param_state['step'] += 1
	step = param_state['step']

	# Adam update logic
	exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
	exp_avg_sq.mul_(beta2).addcmul_(grad, grad.conj(), value=1 - beta2)

	bias_correction1 = 1 - beta1 ** step
	bias_correction2 = 1 - beta2 ** step

	step_size = lr / bias_correction1

	denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(eps)

	p.addcdiv_(exp_avg, denom, value=-step_size)

	return loss