Add nadamuon trained dwee / dpwee vit weights. Add comment to muon impl

Author: rwightman

timm/models/vision_transformer.py

@@ -2703,12 +2703,18 @@ def _cfg(url: str = '', **kwargs) -> Dict[str, Any]:
     'vit_wee_patch16_reg1_gap_256.sbb_in1k': _cfg(
         hf_hub_id='timm/',
         input_size=(3, 256, 256), crop_pct=0.95),
+    'vit_dwee_patch16_reg1_gap_256.sbb_nadamuon_in1k': _cfg(
+        hf_hub_id='timm/',
+        input_size=(3, 256, 256), crop_pct=0.95),
     'vit_dwee_patch16_reg1_gap_256.sbb_in1k': _cfg(
         hf_hub_id='timm/',
         input_size=(3, 256, 256), crop_pct=0.95),
     'vit_pwee_patch16_reg1_gap_256.sbb_in1k': _cfg(
         hf_hub_id='timm/',
         input_size=(3, 256, 256), crop_pct=0.95),
+    'vit_dpwee_patch16_reg1_gap_256.sbb_nadamuon_in1k': _cfg(
+        hf_hub_id='timm/',
+        input_size=(3, 256, 256), crop_pct=0.95),
     'vit_dpwee_patch16_reg1_gap_256.sbb_in1k': _cfg(
         hf_hub_id='timm/',
         input_size=(3, 256, 256), crop_pct=0.95),

timm/optim/muon.py

@@ -602,6 +602,7 @@ def _single_tensor_adamuon(
         # RMS-aligned rescaling: normalize by update norm, then scale by shape factor
         # Used by AdaMuon paper approach (match_rms_adamw), not by μP approach (rms_to_rms)
         if use_rms_norm:
+            # eq(8) in AdaMuon paper, 0.2 / RMS(update) = 0.2 * sqrt(ndim) / frob(update)
             update_norm = update_adaptive.norm().add_(eps)
             update_adaptive = update_adaptive / update_norm