Enum AdjustLrFn 

pub enum AdjustLrFn {
    Original,
    MatchRmsAdamW,
}

Learning rate adjustment method for Muon optimizer.

Muon adjusts the learning rate based on parameter shape to maintain consistent RMS across rectangular matrices.

References

• Original: Muon: An optimizer for hidden layers
• Moonshot: Muon is Scalable for LLM Training

Variants

Original

Keller Jordan’s original method: lr * sqrt(max(1, A/B))

This scales the learning rate based on the aspect ratio of the weight matrix, ensuring that tall matrices (more rows than columns) get proportionally larger learning rates.

Example

For a [1024, 512] matrix: lr * sqrt(1024/512) = lr * 1.414

MatchRmsAdamW

Moonshot’s method: lr * 0.2 * sqrt(max(A, B))

This method is designed to match AdamW’s RMS, allowing Muon to directly reuse learning rates and weight decay values tuned for AdamW without retuning.

Example

For a [1024, 512] matrix: lr * 0.2 * sqrt(1024) = lr * 6.4

Trait Implementations

impl Clone for AdjustLrFn

fn clone(&self) -> AdjustLrFn

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for AdjustLrFn

fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), Error>

Formats the value using the given formatter.

impl Default for AdjustLrFn

fn default() -> AdjustLrFn

Returns the “default value” for a type.

impl<'de> Deserialize<'de> for AdjustLrFn

fn deserialize<__D>( __deserializer: __D, ) -> Result<AdjustLrFn, <__D as Deserializer<'de>>::Error>

where __D: Deserializer<'de>,

Deserialize this value from the given Serde deserializer.

impl PartialEq for AdjustLrFn

fn eq(&self, other: &AdjustLrFn) -> bool

Tests for self and other values to be equal, and is used by ==.

fn ne(&self, other: &Rhs) -> bool

Tests for !=. The default implementation is almost always sufficient, and should not be overridden without very good reason.

impl Serialize for AdjustLrFn

fn serialize<__S>( &self, __serializer: __S, ) -> Result<<__S as Serializer>::Ok, <__S as Serializer>::Error>

where __S: Serializer,

Serialize this value into the given Serde serializer.

impl Copy for AdjustLrFn

impl Eq for AdjustLrFn

impl StructuralPartialEq for AdjustLrFn