Add files using upload-large-folder tool

3cd1076 verified 3 months ago

3.82 kB

#!/usr/bin/env python
# -*- coding: utf-8 -*-
import torch
import torch.nn as nn
import math
import os
import torch.distributed as dist
import torch.nn as nn
from torch import Tensor
def zeropower_via_newtonschulz5(G: Tensor, steps: int) -> Tensor:
    """
    Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a
    quintic iteration whose coefficients are selected to maximize the slope at zero. For the purpose
    of minimizing steps, it turns out to be empirically effective to keep increasing the slope at
    zero even beyond the point where the iteration no longer converges all the way to one everywhere
    on the interval. This iteration therefore does not produce UV^T but rather something like US'V^T
    where S' is diagonal with S_{ii}' ~ Uniform(0.5, 1.5), which turns out not to hurt model
    performance at all relative to UV^T, where USV^T = G is the SVD.
    """
    assert G.ndim >= 2 # batched Muon implementation by @scottjmaddox, and put into practice in the record by @YouJiacheng
    a, b, c = (3.4445, -4.7750,  2.0315)
    X = G.bfloat16()
    if G.size(-2) > G.size(-1):
        X = X.mT

    # Ensure spectral norm is at most 1
    X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7)
    # Perform the NS iterations
    for _ in range(steps):
        A = X @ X.mT
        B = b * A + c * A @ A # quintic computation strategy adapted from suggestion by @jxbz, @leloykun, and @YouJiacheng
        X = a * X + B @ X
    
    if G.size(-2) > G.size(-1):
        X = X.mT
    return X
class Muon(torch.optim.Optimizer):
    """
    Adam optimizer with orthogonalization step.
    """
    def __init__(self, params, lr=0.001, betas=(0.9, 0.999), eps=1e-8, weight_decay=0, ns_steps=5):
        defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, ns_steps=ns_steps)
        super().__init__(params, defaults)

    @torch.no_grad()
    def step(self, closure=None):
        """
        Performs a single optimization step.
        Args:
            closure (callable, optional): A closure that reevaluates the model
                and returns the loss.
        """
        loss = None
        if closure is not None:
            loss = closure()

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

                # Initialize state
                if len(state) == 0:
                    state['step'] = 0
                    state['exp_avg'] = torch.zeros_like(p)
                    state['exp_avg_sq'] = torch.zeros_like(p)

                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
                beta1, beta2 = group['betas']

                state['step'] += 1
                bias_correction1 = 1 - beta1 ** state['step']
                bias_correction2 = 1 - beta2 ** state['step']

                # Update momentum and squared gradient
                exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)

                # Compute the update
                denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(group['eps'])
                step_size = group['lr'] / bias_correction1

                # Orthogonalize the update
                update = exp_avg / denom
                if update.ndim >= 2:
                    update = zeropower_via_newtonschulz5(update, steps=group['ns_steps'])

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

                # Apply weight decay
                if group['weight_decay'] != 0:
                    p.add_(p, alpha=-group['lr'] * group['weight_decay'])

        return loss

	#!/usr/bin/env python
	# -- coding: utf-8 --
	import torch
	import torch.nn as nn
	import math
	import os
	import torch.distributed as dist
	import torch.nn as nn
	from torch import Tensor
	def zeropower_via_newtonschulz5(G: Tensor, steps: int) -> Tensor:
	"""
	Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a
	quintic iteration whose coefficients are selected to maximize the slope at zero. For the purpose
	of minimizing steps, it turns out to be empirically effective to keep increasing the slope at
	zero even beyond the point where the iteration no longer converges all the way to one everywhere
	on the interval. This iteration therefore does not produce UV^T but rather something like US'V^T
	where S' is diagonal with S_{ii}' ~ Uniform(0.5, 1.5), which turns out not to hurt model
	performance at all relative to UV^T, where USV^T = G is the SVD.
	"""
	assert G.ndim >= 2 # batched Muon implementation by @scottjmaddox, and put into practice in the record by @YouJiacheng
	a, b, c = (3.4445, -4.7750, 2.0315)
	X = G.bfloat16()
	if G.size(-2) > G.size(-1):
	X = X.mT

	# Ensure spectral norm is at most 1
	X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7)
	# Perform the NS iterations
	for _ in range(steps):
	A = X @ X.mT
	B = b * A + c * A @ A # quintic computation strategy adapted from suggestion by @jxbz, @leloykun, and @YouJiacheng
	X = a * X + B @ X

	if G.size(-2) > G.size(-1):
	X = X.mT
	return X
	class Muon(torch.optim.Optimizer):
	"""
	Adam optimizer with orthogonalization step.
	"""
	def __init__(self, params, lr=0.001, betas=(0.9, 0.999), eps=1e-8, weight_decay=0, ns_steps=5):
	defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, ns_steps=ns_steps)
	super().__init__(params, defaults)

	@torch.no_grad()
	def step(self, closure=None):
	"""
	Performs a single optimization step.

	Args:
	closure (callable, optional): A closure that reevaluates the model
	and returns the loss.
	"""
	loss = None
	if closure is not None:
	loss = closure()

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

	# Initialize state
	if len(state) == 0:
	state['step'] = 0
	state['exp_avg'] = torch.zeros_like(p)
	state['exp_avg_sq'] = torch.zeros_like(p)

	exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
	beta1, beta2 = group['betas']

	state['step'] += 1
	bias_correction1 = 1 - beta1 ** state['step']
	bias_correction2 = 1 - beta2 ** state['step']

	# Update momentum and squared gradient
	exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
	exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)

	# Compute the update
	denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(group['eps'])
	step_size = group['lr'] / bias_correction1

	# Orthogonalize the update
	update = exp_avg / denom
	if update.ndim >= 2:
	update = zeropower_via_newtonschulz5(update, steps=group['ns_steps'])

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

	# Apply weight decay
	if group['weight_decay'] != 0:
	p.add_(p, alpha=-group['lr'] * group['weight_decay'])

	return loss