feat(trainer): Nesterov momentum for Muon by default

packages/prime-rl-configs/src/prime_rl/configs/trainer.py

@@ -433,6 +433,9 @@ class MuonConfig(BaseOptimizerConfig):
     mu: float = Field(0.95, ge=0)
     """Momentum factor for the Muon algorithm."""
 
+    nesterov: bool = True
+    """Nesterov momentum for the Muon updates."""
+
     betas1: float = Field(0.9, ge=0)
     """β1 for the AdamW/Lion sub-optimizer used on non-Muon params."""
 

src/prime_rl/trainer/optim/__init__.py

@@ -224,6 +224,7 @@ def muon_enabled(n, p):
         matrix_partitions=matrix_partitions,
         lr=lr,
         mu=config.mu,
+        nesterov=config.nesterov,
         betas=(config.betas1, config.betas2),
         weight_decay=config.weight_decay,
         adjust_lr="rms_norm",