/j7ddlmZddlZddlmZddd ZdddZGddejZGddejZdS)) annotationsN)optimHz>G torch.Tensorepsfloatreturnct|jdksJ|}|||zz}|d|dkr|j}dD])\}}}||jz}||z||z|zz}||z||zz}*|d|dkr|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)rrXabcABs [/home/longshao/multi-rider-rag/.venv/lib/python3.11/site-packages/ultralytics/optim/muon.pyzeropower_via_newtonschulz5r s8 qw<<1     ACAvvayy166!99 C  1a G EAEAI  EAEMvvayy166!99 C Hffffff?Tgradmomentumbetanesterovboolcd||d|z |r|||n|}|jdkr#|t |d}t |}|t d|d|dz 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_lerpndimviewrrmaxr)rr r!r"updates r muon_updater/;s> NN4T"""*2 @TYYx & & &F {aS[["-- ( 0 0F c!TYYr]]TYYr]]233s::F MrcfeZdZdZ ddfd ZejddZxZS)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. MbP?Fr(lrr r weight_decayr"r#use_muonmuonsgdc t|||||} t|| ||_||_dS)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. )r4r r5r"r6N)dictsuper__init__r7r8) selfparamsr4r r5r"r6r7r8defaults __class__s rr<zMuSGD.__init__sV.%     *** rNc d}|5tj5|}dddn #1swxYwY|jD]}|dr|dD]u}|d}|j|j}|j|}t |dkr.tj||d<tj||d<t||d|d|d  }|| |j ||j z |d dkr| ||d  }|d |d||d r#| |d|d n|d} || ||jz w|dD]}|d}|j|j}|d dkr| ||d  }|j|}t |dkrtj||d<|d |d||d r#| |d|d n|d}||| |S) 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. Nr6r>r4rmomentum_buffermomentum_buffer_SGDr r")r!r"alphar5)torch enable_grad param_groupsrstater zeros_liker/add_reshaperr7addmul_r8) r=closurelossgrouppr4rrIr. sgd_updates rstepz MuSGD.stepsK(  "$$ ! !wyy ! ! ! ! ! ! ! ! ! ! ! ! ! ! !&. .. .EZ , .x??AtBv~ 6D JqME5zzQ383CA3F3F/07<7G7J7J34(e$56U:=NY^_iYjFFF6>>!'22BN:KFLLL^,11#xx~1FxGG/055eJ6GHHMMdSSS!,:'>>>1?4x..AtBv~ 6D^,11#xx~1FxGG JqME5zzQ383CA3F3F/0+,11% 2CDDII$OOO!,6'8!9zARSSS"#45 FF6"F----!." /33)r2r3r3FFr(r()r4r r r r5r r"r#r6r#r7r r8r N __name__ __module__ __qualname____doc__r<rFno_gradrT __classcell__r@s@rr1r1cs++`!       DU]__GGG_GGGGGrr1cXeZdZdZd d fd Zejdd ZxZS)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. {Gz?rrr4r r5r cpt|||}t||dS)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. )r4r5r N)r:r;r<)r=r>r4r5r r?r@s rr<z Muon.__init__s92L8LLL *****rNcPd}|5tj5|}dddn #1swxYwY|jD]}|dD]}|jtj||_|j|}t |dkrtj||d<t|j|d|d}|d|d|d zz | | |j |d ڌ|S) 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. Nr>rrBr )r!r r4r5rD) rFrGrHrrJrIrr/rNrKrLr)r=rOrPrQrRrIr.s rrTz Muon.step(ss2  "$$ ! !wyy ! ! ! ! ! ! ! ! ! ! ! ! ! ! !& D DE8_ D D6>"-a00AF 1 u::??/4/?/B/BE+,$QVU3D-EER\L]^^^q5;~)>>>???v~~ag..uT{lCCCC D rU)rarr)r4r r5r r r rVrWr^s@rr`r`st: + + + + + + +U]__)))_)))))rr`)r)rrrr r r)rT) rrr rr!r r"r#r r) __future__rrFrrr/ Optimizerr1r`rrrgs#""""" / / / / / d%%%%%PXXXXXEOXXXvTTTTT5?TTTTTr