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) -> Self

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl Copy for AdjustLrFn

impl Debug for AdjustLrFn

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

Formats the value using the given formatter.

impl Default for AdjustLrFn

fn default() -> Self

Returns the “default value” for a type.

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

fn deserialize<__D>(__deserializer: __D) -> Result<Self, __D::Error>

where __D: Deserializer<'de>,

Deserialize this value from the given Serde deserializer.

impl Eq for AdjustLrFn

impl PartialEq for AdjustLrFn

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

Equality operator ==.

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

Inequality operator !=.

impl Serialize for AdjustLrFn

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

where __S: Serializer,

Serialize this value into the given Serde serializer.

impl StructuralPartialEq for AdjustLrFn