Upload original AdamW and Muon pretrained checkpoints (step 31000)

ad93bc0 verified 1 day ago

4.91 kB

"""
Implementation of normalized steepest descent algorithm. We additionally apply multipliers to align the RMS norm update with Adam.
Note that when p=2 is chosen, the algorithm is the same as normalized SGD.
"""

import torch
import math

class NormalizedSteepestDescent(torch.optim.Optimizer):
    """
    Implements a normalized version of Stochastic Gradient Descent (with momentum).
    The update for each parameter is normalized by its Frobenius norm before
    being applied. This is a form of layer-wise adaptivity.
    """
    def __init__(self, params, norm_p=2, lr=1e-3, momentum=0.9, weight_decay=0.1, scale=0.2, eps=1e-12, use_nesterov=False, use_grafting=False):
        # scale is used for aligning the update RMS with Adam
        if lr < 0.0:
            raise ValueError(f"Invalid learning rate: {lr}")
        if momentum < 0.0:
            raise ValueError(f"Invalid momentum value: {momentum}")
        if weight_decay < 0.0:
            raise ValueError(f"Invalid weight_decay value: {weight_decay}")
        # use_grafting not implemented yet
        defaults = dict(norm_p=norm_p, lr=lr, momentum=momentum, weight_decay=weight_decay, scale=scale, eps=eps, use_nesterov=use_nesterov)
        super(NormalizedSteepestDescent, self).__init__(params, defaults)

    def __setstate__(self, state):
        super(NormalizedSteepestDescent, 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:
            norm_p = group['norm_p']
            lr = group['lr']
            momentum = group['momentum']
            weight_decay = group['weight_decay']
            eps = group['eps']
            scale = group['scale']
            use_nesterov = group['use_nesterov']

            for p in group['params']:
                if p.grad is None:
                    continue
                
                # Get the gradient
                d_p = p.grad

                # Apply momentum
                if momentum != 0:
                    param_state = self.state[p]
                    if 'momentum_buffer' not in param_state:
                        buf = param_state['momentum_buffer'] = torch.clone(d_p).detach()
                    else:
                        buf = param_state['momentum_buffer']
                        buf.mul_(momentum).add_(d_p, alpha=1-momentum) # m_t = beta * m_{t-1} + (1-beta) * g_t
                    
                    # This is the update tensor (m_t)
                    if use_nesterov:
                        # use nesterov momentum, following Muon and Lion
                        Gt = torch.lerp(d_p, buf, momentum)
                    else:
                        Gt = buf
                else:
                    # If no momentum, the gradient is the update tensor
                    Gt = d_p

                mn = p.numel() # parameter count

                if norm_p <= 1 + eps:
                    # in this case, only the largest coordiantes have update
                    g_max_abs = torch.max(Gt.abs())
                    mask = Gt.abs() >= g_max_abs - eps
                    # Count how many elements have the max value (k in the formula)
                    k = torch.sum(mask).float()
                    # The update direction is sign(G_t) where |G_t| is max, and 0 otherwise.
                    update_direction = torch.zeros_like(Gt)
                    update_direction[mask] = torch.sign(Gt[mask])
                    
                    # Calculate the final update magnitude based on the limit
                    alpha = -lr * scale * torch.sqrt(mn / (k + eps))
                
                elif norm_p == float('inf'):
                    # in this case, we apply SignSGD
                    update_direction = torch.sign(Gt)
                    alpha = -lr * scale 
                
                else:
                    norm_q = 1 / (1 - 1 / norm_p)
                    assert norm_q > 0

                    # Calculate the update direction: |G_t|^(q-1) * sign(G_t)
                    update_direction = torch.pow(Gt.abs(), norm_q - 1) * torch.sign(Gt)
                    
                    update_norm = torch.linalg.norm(update_direction)

                    update_direction = update_direction / (update_norm + eps)
                    
                    # Calculate the pre-factor
                    alpha = -lr * scale * math.sqrt(mn)
                    
                    
                    

                # Apply weight decay
                if weight_decay > 0.0:
                    p.mul_(1.0 - lr * weight_decay)
                
                # Apply the update
                p.add_(update_direction, alpha=alpha)

        return loss

	"""
	Implementation of normalized steepest descent algorithm. We additionally apply multipliers to align the RMS norm update with Adam.
	Note that when p=2 is chosen, the algorithm is the same as normalized SGD.
	"""

	import torch
	import math

	class NormalizedSteepestDescent(torch.optim.Optimizer):
	"""
	Implements a normalized version of Stochastic Gradient Descent (with momentum).

	The update for each parameter is normalized by its Frobenius norm before
	being applied. This is a form of layer-wise adaptivity.
	"""
	def __init__(self, params, norm_p=2, lr=1e-3, momentum=0.9, weight_decay=0.1, scale=0.2, eps=1e-12, use_nesterov=False, use_grafting=False):
	# scale is used for aligning the update RMS with Adam
	if lr < 0.0:
	raise ValueError(f"Invalid learning rate: {lr}")
	if momentum < 0.0:
	raise ValueError(f"Invalid momentum value: {momentum}")
	if weight_decay < 0.0:
	raise ValueError(f"Invalid weight_decay value: {weight_decay}")
	# use_grafting not implemented yet
	defaults = dict(norm_p=norm_p, lr=lr, momentum=momentum, weight_decay=weight_decay, scale=scale, eps=eps, use_nesterov=use_nesterov)
	super(NormalizedSteepestDescent, self).__init__(params, defaults)

	def __setstate__(self, state):
	super(NormalizedSteepestDescent, 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:
	norm_p = group['norm_p']
	lr = group['lr']
	momentum = group['momentum']
	weight_decay = group['weight_decay']
	eps = group['eps']
	scale = group['scale']
	use_nesterov = group['use_nesterov']

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

	# Get the gradient
	d_p = p.grad

	# Apply momentum
	if momentum != 0:
	param_state = self.state[p]
	if 'momentum_buffer' not in param_state:
	buf = param_state['momentum_buffer'] = torch.clone(d_p).detach()
	else:
	buf = param_state['momentum_buffer']
	buf.mul_(momentum).add_(d_p, alpha=1-momentum) # m_t = beta * m_{t-1} + (1-beta) * g_t

	# This is the update tensor (m_t)
	if use_nesterov:
	# use nesterov momentum, following Muon and Lion
	Gt = torch.lerp(d_p, buf, momentum)
	else:
	Gt = buf
	else:
	# If no momentum, the gradient is the update tensor
	Gt = d_p

	mn = p.numel() # parameter count

	if norm_p <= 1 + eps:
	# in this case, only the largest coordiantes have update
	g_max_abs = torch.max(Gt.abs())
	mask = Gt.abs() >= g_max_abs - eps
	# Count how many elements have the max value (k in the formula)
	k = torch.sum(mask).float()
	# The update direction is sign(G_t) where \|G_t\| is max, and 0 otherwise.
	update_direction = torch.zeros_like(Gt)
	update_direction[mask] = torch.sign(Gt[mask])

	# Calculate the final update magnitude based on the limit
	alpha = -lr * scale * torch.sqrt(mn / (k + eps))

	elif norm_p == float('inf'):
	# in this case, we apply SignSGD
	update_direction = torch.sign(Gt)
	alpha = -lr * scale

	else:
	norm_q = 1 / (1 - 1 / norm_p)
	assert norm_q > 0

	# Calculate the update direction: \|G_t\|^(q-1) * sign(G_t)
	update_direction = torch.pow(Gt.abs(), norm_q - 1) * torch.sign(Gt)

	update_norm = torch.linalg.norm(update_direction)

	update_direction = update_direction / (update_norm + eps)

	# Calculate the pre-factor
	alpha = -lr * scale * math.sqrt(mn)




	# Apply weight decay
	if weight_decay > 0.0:
	p.mul_(1.0 - lr * weight_decay)

	# Apply the update
	p.add_(update_direction, alpha=alpha)

	return loss