Upload original AdamW and Muon pretrained checkpoints (step 31000)

ad93bc0 verified 1 day ago

3.08 kB

"""
Implementation of normalized SGD algorithm. We additionally apply 0.2 * \sqrt{mn} to align the RMS norm update with Adam.
"""

import torch
import math

class NormalizedSGD(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, lr=1e-3, momentum=0.9, weight_decay=0.1, scale=0.2, eps=1e-12):
        # 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}")

        defaults = dict(lr=lr, momentum=momentum, weight_decay=weight_decay, scale=scale, eps=eps)
        super(NormalizedSGD, self).__init__(params, defaults)

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

            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)
                    update_vec = buf
                else:
                    # If no momentum, the gradient is the update tensor
                    update_vec = d_p

                # +++ CORE MODIFICATION: Normalize the update vector +++
                # Calculate the Frobenius norm of the update vector
                norm = torch.linalg.norm(update_vec)

                # Apply weight decay
                if weight_decay > 0.0:
                    p.mul_(1.0 - lr * weight_decay)
                
                # Normalize and apply the update
                # p_new = p_old - lr * (m_t / ||m_t||_F)
                scale_align = scale * math.sqrt(p.numel())
                p.add_(update_vec, alpha=-lr * scale_align / (norm + eps))

        return loss

	"""
	Implementation of normalized SGD algorithm. We additionally apply 0.2 * \sqrt{mn} to align the RMS norm update with Adam.
	"""

	import torch
	import math

	class NormalizedSGD(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, lr=1e-3, momentum=0.9, weight_decay=0.1, scale=0.2, eps=1e-12):
	# 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}")

	defaults = dict(lr=lr, momentum=momentum, weight_decay=weight_decay, scale=scale, eps=eps)
	super(NormalizedSGD, self).__init__(params, defaults)

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

	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)
	update_vec = buf
	else:
	# If no momentum, the gradient is the update tensor
	update_vec = d_p

	# +++ CORE MODIFICATION: Normalize the update vector +++
	# Calculate the Frobenius norm of the update vector
	norm = torch.linalg.norm(update_vec)

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

	# Normalize and apply the update
	# p_new = p_old - lr * (m_t / \|\|m_t\|\|_F)
	scale_align = scale * math.sqrt(p.numel())
	p.add_(update_vec, alpha=-lr * scale_align / (norm + eps))

	return loss