misc

test/test_muon/test.py

@@ -2,7 +2,7 @@ import logging
 
 import torch
 import torch.distributed as dist
-from muon import Muon, get_default_muon_param_groups
+from optimizer.muon import Muon, get_default_muon_param_groups
 from torch.distributed.fsdp import FSDPModule, fully_shard
 from torch.distributed.tensor import DTensor
 from torch.distributed.tensor.placement_types import Replicate

torch-ext/optimizer/muon.py

@@ -701,10 +701,10 @@ class Muon(torch.optim.Optimizer):
                 new_scale = math.sqrt(threshold / v_ele)
                 if new_scale < scales_full[head_idx]:
                     scales_full[head_idx] = new_scale
-                    #logger.info(
-                    #    f"[{kind}] Head {head_idx} exceeded threshold "
-                    #    f"(value={v_ele:.4f}, threshold={threshold:.4f}) -> applying scale={new_scale:.4f}"
-                    #)
+                    logger.info(
+                        f"[{kind}] Head {head_idx} exceeded threshold "
+                        f"(value={v_ele:.4f}, threshold={threshold:.4f}) -> applying scale={new_scale:.4f}"
+                    )
                     scaling += 1
 
         return scales_full if scaling > 0 else None