Add info-level logging for param group classification (Muon vs AdamW)

Log each parameter's skip/expert status and effective ndim in
default_is_muon(), and summarize Muon/AdamW param lists in
get_default_muon_param_groups().

[skip-build]

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

torch-ext/optimizer/core.py

@@ -1,3 +1,4 @@
+import logging
 import math
 from dataclasses import dataclass
 
@@ -12,6 +13,8 @@ from torch.distributed.tensor import DTensor
 # expert_keys, QK layer parsing) works regardless of wrapper nesting.
 _WRAPPER_PARTS = frozenset({"_orig_mod", "_checkpoint_wrapped_module"})
 
+logger = logging.getLogger(__name__)
+
 
 def normalize_fqn(name: str) -> str:
     """Strip torch.compile / checkpoint wrapper components from a parameter FQN."""
@@ -92,11 +95,18 @@ def default_is_muon(name, x, expert_keys=None):
     parts = normalize_fqn(name).split(".")
     skip_keys = ["embed_tokens", "lm_head", "tok_embeddings", "output"]
     if any(key in parts for key in skip_keys):
+        logger.info("[is_muon] %s: skip (matched skip_key), ndim=%d", name,
+                    x.ndim)
         return False
     effective_ndim = x.ndim
-    if expert_keys and any(key in parts for key in expert_keys):
+    is_expert = expert_keys and any(key in parts for key in expert_keys)
+    if is_expert:
         effective_ndim -= 1
-    return effective_ndim >= 2
+    result = effective_ndim >= 2
+    logger.info("[is_muon] %s: ndim=%d, expert=%s, effective_ndim=%d → %s",
+                name, x.ndim, is_expert, effective_ndim,
+                "Muon" if result else "AdamW")
+    return result
 
 
 def get_default_muon_param_groups(model, is_muon_func=None, expert_keys=None):
@@ -104,7 +114,7 @@ def get_default_muon_param_groups(model, is_muon_func=None, expert_keys=None):
         is_muon_func = lambda n, x: default_is_muon(n, x, expert_keys)
 
     muon_params, muon_names = [], []
-    non_muon_params = []
+    non_muon_params, non_muon_names = [], []
 
     for n, p in model.named_parameters():
         if not p.requires_grad:
@@ -114,6 +124,13 @@ def get_default_muon_param_groups(model, is_muon_func=None, expert_keys=None):
             muon_names.append(n)
         else:
             non_muon_params.append(p)
+            non_muon_names.append(n)
+
+    logger.info("[param_groups] expert_keys=%s", expert_keys)
+    logger.info("[param_groups] Muon params (%d): %s", len(muon_names),
+                muon_names)
+    logger.info("[param_groups] AdamW params (%d): %s", len(non_muon_names),
+                non_muon_names)
 
     return [
         {