-Muon(; [opt, eta, mu, lambda, fallback])</code></pre><p>Muon - MomentUm Orthogonalized by Newton-schulz (https://github.com/KellerJordan/Muon)</p><p>Muon internally runs standard SGD-momentum, and then performs an orthogonalization post-processing step, in which each 2D parameter's update is replaced with the nearest orthogonal matrix using Newton-Schulz iteration.</p><p><strong>Parameters</strong></p><ul><li>Fallback optimizer (<code>opt</code>): Optimizer to use for 1D parameters or when the <code>fallback</code> function returns true</li><li>Learning rate (<code>η == eta</code>): Amount by which gradients are discounted before updating the weights</li><li>Momentum (<code>μ == mu</code>): Controls the acceleration of gradient descent in the prominent direction</li><li>Weight decay (<code>λ == lambda</code>): Controls the strength of <span>$L_2$</span> regularisation.</li><li>Fallback function (<code>fallback</code>): Function to control when, in addition to 1D arrays, the fallback optimizer should be used. Will be passed the parameter array and must return a boolean.</li></ul><p>Note: Works best with large batch sizes and may not be suitable for fine-tuning. In nanoGPT speedrun experiments, Muon is used for the internal layer >2D weights, and AdamW is used for the 1D weights, embeddings, and heads.</p><p><code>Optimisers.adjust!(optimiser_state, η::Real)</code> will adjust the fallback optimizer's <code>eta</code> to <code>η * (opt.eta / eta)</code>, and Muon's <code>eta</code> to <code>η</code>, preserving their ratio, but <code>Optimisers.adjust!(optimiser, eta = η)</code> will only adjust Muon's learning rate (allowing you to adjust the fallback optimizer's learning rate separately).</p></div><a class="docs-sourcelink" target="_blank" href="https://github.com/MurrellGroup/CannotWaitForTheseOptimisers.jl/blob/02e1cb427814523cfcda8d99040b0e5d0de95710/src/rules.jl#L4-L25">source</a></section></article><article class="docstring"><header><a class="docstring-article-toggle-button fa-solid fa-chevron-down" href="javascript:;" title="Collapse docstring"></a><a class="docstring-binding" id="CannotWaitForTheseOptimisers.NormGrowthCap" href="#CannotWaitForTheseOptimisers.NormGrowthCap"><code>CannotWaitForTheseOptimisers.NormGrowthCap</code></a> — <span class="docstring-category">Type</span><span class="is-flex-grow-1 docstring-article-toggle-button" title="Collapse docstring"></span></header><section><div><pre><code class="language-julia hljs">NormGrowthCap(τ = 1.01; ϵ = 1e-8, lb = 1e-7, throw = true, scale = true)</code></pre><p>Gradient norm growth limiter. <code>τ</code> controls the maximum that the gradient norm can grow from one step to the next, such that if <code>||dx||/||dx_prev|| > τ</code> & <code>||dx|| > lb</code>, then <code>dx = dx * τ*||dx_prev||/(||dx||+ϵ)</code> Inspired by <a href="https://arxiv.org/abs/2410.01623">Chen et al.</a> and used with Apollo in <a href="https://arxiv.org/abs/2412.05270">Zhu et al.</a>, but with Optimisers.jl this will apply per-tensor instead of per-model. This implementation also introduces <code>lb</code> as a hard minimum on the gradient norm threshold, and never rescales grads below this, preventing a tensor from getting "trapped" near zero. This can be a fixed min, or scaled by the square root of the number of parameters in the tensor (with <code>scale = true</code>).</p></div><a class="docs-sourcelink" target="_blank" href="https://github.com/MurrellGroup/CannotWaitForTheseOptimisers.jl/blob/02e1cb427814523cfcda8d99040b0e5d0de95710/src/rules.jl#L77-L86">source</a></section></article></article><nav class="docs-footer"><p class="footer-message">Powered by <a href="https://github.com/JuliaDocs/Documenter.jl">Documenter.jl</a> and the <a href="https://julialang.org/">Julia Programming Language</a>.</p></nav></div><div class="modal" id="documenter-settings"><div class="modal-background"></div><div class="modal-card"><header class="modal-card-head"><p class="modal-card-title">Settings</p><button class="delete"></button></header><section class="modal-card-body"><p><label class="label">Theme</label><div class="select"><select id="documenter-themepicker"><option value="auto">Automatic (OS)</option><option value="documenter-light">documenter-light</option><option value="documenter-dark">documenter-dark</option><option value="catppuccin-latte">catppuccin-latte</option><option value="catppuccin-frappe">catppuccin-frappe</option><option value="catppuccin-macchiato">catppuccin-macchiato</option><option value="catppuccin-mocha">catppuccin-mocha</option></select></div></p><hr/><p>This document was generated with <a href="https://github.com/JuliaDocs/Documenter.jl">Documenter.jl</a> version 1.8.0 on <span class="colophon-date" title="Monday 23 December 2024 19:57">Monday 23 December 2024</span>. Using Julia version 1.11.2.</p></section><footer class="modal-card-foot"></footer></div></div></div></body></html>
0 commit comments