Fi7ddlmZddlZddlmZd d dZd d dZGddej ZGdd ej Zy)) annotationsN)optimct|jdk(sJ|j}||j|zz}|j d|j dkDr |j }dD].\}}}||j z}||z||z|zz}||z||zz}0|j d|j dkDr |j }|S)a"Compute the zeroth power / orthogonalization of matrix G using Newton-Schulz iteration. This function implements a quintic Newton-Schulz iteration to compute an approximate orthogonalization of the input matrix G. The iteration coefficients are optimized to maximize convergence slope at zero, producing a result similar to UV^T from SVD, where USV^T = G, but with relaxed convergence guarantees that empirically work well for optimization purposes. Args: G (torch.Tensor): Input 2D tensor/matrix to orthogonalize. eps (float, optional): Small epsilon value added to norm for numerical stability. Default: 1e-7. Returns: (torch.Tensor): Orthogonalized matrix with same shape as input G. Examples: >>> G = torch.randn(128, 64) >>> G_ortho = zeropower_via_newtonschulz5(G) >>> print(G_ortho.shape) torch.Size([128, 64]) Notes: - Uses bfloat16 precision for computation. - Performs exactly 5 Newton-Schulz iteration steps with fixed coefficients. - Automatically transposes for efficiency when rows > columns. - Output approximates US'V^T where S' has diagonal entries ~ Uniform(0.5, 1.5). - Does not produce exact UV^T but works well empirically for neural network optimization. r)guV @ggn@@rrrr)lenshapebfloat16normsizeT)GepsXabcABs X/var/www/pod-logistic/pod-ai/venv/lib/python3.12/site-packages/ultralytics/optim/muon.pyzeropower_via_newtonschulz5r s8 qww<1   ACAvvay166!9 CC 1a G EAEAI  EAEM  vvay166!9 CC Hc,|j|d|z |r|j||n|}|jdk(r|jt |d}t |}|t d|jd|jdz dzz}|S)aXCompute Muon optimizer update with momentum and orthogonalization. This function applies momentum to the gradient, optionally uses Nesterov acceleration, and then orthogonalizes the update using Newton-Schulz iterations. For convolutional filters (4D tensors), it reshapes before orthogonalization and scales the final update based on parameter dimensions. Args: grad (torch.Tensor): Gradient tensor to update. Can be 2D or 4D (for conv filters). momentum (torch.Tensor): Momentum buffer tensor, modified in-place via lerp. beta (float, optional): Momentum coefficient for exponential moving average. Default: 0.95. nesterov (bool, optional): Whether to use Nesterov momentum acceleration. Default: True. Returns: (torch.Tensor): Orthogonalized update tensor with same shape as input grad. For 4D inputs, returns reshaped result matching original dimensions. Examples: >>> grad = torch.randn(64, 128) >>> momentum = torch.zeros_like(grad) >>> update = muon_update(grad, momentum, beta=0.95, nesterov=True) >>> print(update.shape) torch.Size([64, 128]) Notes: - Momentum buffer is updated in-place: momentum = beta * momentum + (1-beta) * grad. - With Nesterov: update = beta * momentum + (1-beta) * grad. - Without Nesterov: update = momentum. - 4D tensors (conv filters) are reshaped to 2D as (out_channels, in_channels*height*width) for orthogonalization. - Final update is scaled by sqrt(max(1, dim[-2] / dim[-1])) to account for parameter dimensions. r?)lerp_lerpndimviewr rmaxr )gradmomentumbetanesterovupdates r muon_updater);s> NN4T"*2TYYx &F {{aS["- ( 0F c!TYYr]TYYr]23s::F Mrc~eZdZdZ d dfd Zej ddZxZS)MuSGDa"Hybrid optimizer combining Muon and SGD updates for neural network training. This optimizer implements a combination of Muon (a momentum-based optimizer with orthogonalization via Newton-Schulz iterations) and standard SGD with momentum. It allows different parameter groups to use either the hybrid Muon+SGD approach or pure SGD. Args: params (Iterable): Parameters to optimize or dicts defining parameter groups. muon (float, optional): Weight factor for Muon updates in hybrid mode. Default: 0.5. sgd (float, optional): Weight factor for SGD updates in hybrid mode. Default: 0.5. Attributes: muon (float): Scaling factor applied to Muon learning rate. sgd (float): Scaling factor applied to SGD learning rate in hybrid mode. Examples: >>> param_groups = [ ... { ... "params": model.conv_params, ... "lr": 0.02, ... "use_muon": True, ... "momentum": 0.95, ... "nesterov": True, ... "weight_decay": 0.01, ... }, ... { ... "params": model.other_params, ... "lr": 0.01, ... "use_muon": False, ... "momentum": 0.9, ... "nesterov": False, ... "weight_decay": 0, ... }, ... ] >>> optimizer = MuSGD(param_groups, muon=0.5, sgd=0.5) >>> loss = model(data) >>> loss.backward() >>> optimizer.step() Notes: - Parameter groups with 'use_muon': True will receive both Muon and SGD updates. - Parameter groups with 'use_muon': False will receive only SGD updates. - The Muon update uses orthogonalization which works best for 2D+ parameter tensors. c bt|||||} t | || ||_||_y)aNInitialize MuSGD optimizer with hybrid Muon and SGD capabilities. Args: params (Iterable): Iterable of parameters to optimize or dicts defining parameter groups. lr (float): Learning rate. momentum (float): Momentum factor for SGD. weight_decay (float): Weight decay (L2 penalty). nesterov (bool): Whether to use Nesterov momentum. use_muon (bool): Whether to enable Muon updates. muon (float): Scaling factor for Muon component. sgd (float): Scaling factor for SGD component. )lrr% weight_decayr'use_muonN)dictsuper__init__muonsgd) selfparamsr-r%r.r'r/r3r4defaults __class__s rr2zMuSGD.__init__s=.%   * rc d}|$tj5|}ddd|jD]4}|drS|dD]H}|d}|j|j}|j|}t |dk(r0tj ||d<tj ||d<t||d|d|d  }|j|j|j||jz |d dk7r|j||d  }|dj|dj||d r|j|d|d n|d} |j| ||jz K]|dD]}|d}|j|j}|d dk7r|j||d  }|j|}t |dk(rtj ||d<|dj|dj||d r|j|d|d n|d}|j|| 7|S#1swYPxYw) a#Perform a single optimization step. Applies either hybrid Muon+SGD updates or pure SGD updates depending on the 'use_muon' flag in each parameter group. For Muon-enabled groups, parameters receive both an orthogonalized Muon update and a standard SGD momentum update. Args: closure (Callable, optional): A closure that reevaluates the model and returns the loss. Default: None. Returns: (torch.Tensor | None): The loss value if closure is provided, otherwise None. Notes: - Parameters with None gradients are skipped. - Muon updates use Newton-Schulz orthogonalization and work best on 2D+ tensors. - Weight decay is applied only to the SGD component in hybrid mode. Nr/r6r-rmomentum_buffermomentum_buffer_SGDr%r')r&r'alphar.)torch enable_grad param_groupsr$stater zeros_liker)add_reshaper r3addmul_r4) r5closurelossgrouppr-r$rAr( sgd_updates rstepz MuSGD.steps(  ""$ !y !&&. .EZ x?AtBvv~ 66D JJqME5zQ383C3CA3F/07<7G7G7J34(e$56U:=NY^_iYjFFF6>>!''2BN:KFL^,1#xx~1FxG/055eJ6GHMMdS!,'1?4x.AtBvv~ 66D^,1#xx~1FxG JJqME5zQ383C3CA3F/0+,11% 2CDII$O!,'8!9zARS"#45 FF6"F-!.=. .^ e ! !s I//I9)gMbP?rMFFrr)r-floatr%rNr.rNr'boolr/rOr3rNr4rNN __name__ __module__ __qualname____doc__r2r>no_gradrL __classcell__r8s@rr+r+cs+`!           DU]]_GGrr+cVeZdZdZddfd Zej ddZxZS)MuonaSMuon optimizer for usage in non-distributed settings. This optimizer implements the Muon algorithm, which combines momentum-based updates with orthogonalization via Newton-Schulz iterations. It applies weight decay and learning rate scaling to parameter updates. Args: params (iterable): Iterable of parameters to optimize or dicts defining parameter groups. lr (float, optional): Learning rate. Default: 0.02. weight_decay (float, optional): Weight decay (L2 penalty) coefficient. Default: 0. momentum (float, optional): Momentum coefficient for exponential moving average. Default: 0.95. Attributes: param_groups (list): List of parameter groups with their optimization settings. state (dict): Dictionary containing optimizer state for each parameter. Examples: >>> model = YourModel() >>> optimizer = Muon(model.parameters(), lr=0.02, weight_decay=0.01, momentum=0.95) >>> loss = model(data) >>> loss.backward() >>> optimizer.step() Notes: - Designed for non-distributed training environments. - Uses Muon updates with orthogonalization for all parameters. - Weight decay is applied multiplicatively before parameter update. - Parameters with None gradients are assigned zero gradients for synchronization. cBt|||}t| ||y)a}Initialize Muon optimizer with orthogonalization-based updates. Args: params (Iterable): Iterable of parameters to optimize or dicts defining parameter groups. lr (float): Learning rate. weight_decay (float): Weight decay factor applied multiplicatively. momentum (float): Momentum factor for gradient accumulation. )r-r.r%N)r0r1r2)r5r6r-r.r%r7r8s rr2z Muon.__init__s"2L8L *rc0d}|$tj5|}ddd|jD]}|dD]}|jtj||_|j |}t |dk(rtj||d<t|j|d|d}|jd|d|d zz |j|j|j|d |S#1swYxYw) aPerform a single optimization step. Applies Muon updates to all parameters, incorporating momentum and orthogonalization. Weight decay is applied multiplicatively before the parameter update. Args: closure (Callable[[], torch.Tensor] | None, optional): A closure that reevaluates the model and returns the loss. Default: None. Returns: (torch.Tensor | None): The loss value if closure is provided, otherwise None. Examples: >>> optimizer = Muon(model.parameters()) >>> loss = model(inputs) >>> loss.backward() >>> optimizer.step() Notes: - Parameters with None gradients are assigned zero gradients for synchronization. - Weight decay is applied as: p *= (1 - lr * weight_decay). - Muon update uses Newton-Schulz orthogonalization and works best on 2D+ tensors. Nr6rr:r%)r&rr-r.r<) r>r?r@r$rBrAr r)rFrCrDr )r5rGrHrIrJrAr(s rrLz Muon.step(s2  ""$ !y !&& DE8_ D66>"--a0AF 1 u:?/4/?/?/BE+,$QVVU3D-EER\L]^q5;~)>>>?v~~agg.uT{lC D D  ! !s D  D)g{Gz?rffffff?)r-rNr.rNr%rNrPrQrXs@rrZrZs): +U]]_))rrZ)gHz>)r torch.TensorrrNreturnr^)r]T) r$r^r%r^r&rNr'rOr_r^) __future__rr>rrr) Optimizerr+rZrrrcsA# / d%PXEOOXvT5??Tr