Normalize parameter FQNs to handle torch.compile / checkpoint wrappers

torch.compile wraps modules as OptimizedModule, inserting _orig_mod into
parameter FQNs. Activation checkpointing similarly inserts
_checkpoint_wrapped_module. These wrapper components break name-based
matching for skip_keys, expert_keys, and QK layer parsing.

Add normalize_fqn() that strips these wrapper components, and apply it
in default_is_muon(), _expand_expert_params(), and parse_qk_layer().

[skip-build]

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

test/test_normalize_fqn.py

@@ -0,0 +1,52 @@
+"""Unit tests for FQN normalization (no GPU / distributed required)."""
+
+import pytest
+from optimizer.core import normalize_fqn
+from optimizer.qk_clip import parse_qk_layer
+
+
+class TestNormalizeFqn:
+
+    def test_passthrough(self):
+        assert normalize_fqn("model.layers.3.attn.q_proj.weight") == \
+            "model.layers.3.attn.q_proj.weight"
+
+    def test_strip_orig_mod(self):
+        assert normalize_fqn("model._orig_mod.layers.3.attn.q_proj.weight") == \
+            "model.layers.3.attn.q_proj.weight"
+
+    def test_strip_checkpoint_wrapped(self):
+        name = "model.layers.0._checkpoint_wrapped_module.moe.experts.w1.weight"
+        assert normalize_fqn(name) == \
+            "model.layers.0.moe.experts.w1.weight"
+
+    def test_strip_both(self):
+        name = "model._orig_mod.layers.0._checkpoint_wrapped_module.attn.q_proj.weight"
+        assert normalize_fqn(name) == \
+            "model.layers.0.attn.q_proj.weight"
+
+    def test_strip_nested_orig_mod(self):
+        name = "_orig_mod._orig_mod.layers.0.mlp.gate_proj.weight"
+        assert normalize_fqn(name) == \
+            "layers.0.mlp.gate_proj.weight"
+
+
+class TestParseQkLayerWithWrappers:
+
+    def test_plain_name(self):
+        assert parse_qk_layer("model.layers.3.attn.q_proj.weight") == ("q_proj", 3)
+
+    def test_orig_mod(self):
+        assert parse_qk_layer("model._orig_mod.layers.3.attn.wq.weight") == ("wq", 3)
+
+    def test_checkpoint_wrapped(self):
+        name = "model.layers.5._checkpoint_wrapped_module.self_attn.k_proj.weight"
+        assert parse_qk_layer(name) == ("k_proj", 5)
+
+    def test_both_wrappers(self):
+        name = "_orig_mod.model._checkpoint_wrapped_module.layers.7.attn.wk.weight"
+        assert parse_qk_layer(name) == ("wk", 7)
+
+    def test_non_qk_still_none(self):
+        name = "model._orig_mod.layers.2.attn.v_proj.weight"
+        assert parse_qk_layer(name) == (None, -1)

torch-ext/optimizer/core.py

@@ -7,6 +7,18 @@ from torch.distributed import ProcessGroup
 from torch.distributed.tensor import DTensor
 
 
+# torch.compile wraps modules as OptimizedModule, inserting "_orig_mod" into
+# parameter FQNs.  Activation checkpointing similarly inserts
+# "_checkpoint_wrapped_module".  Strip these so name-based matching (skip_keys,
+# expert_keys, QK layer parsing) works regardless of wrapper nesting.
+_WRAPPER_PARTS = frozenset({"_orig_mod", "_checkpoint_wrapped_module"})
+
+
+def normalize_fqn(name: str) -> str:
+    """Strip torch.compile / checkpoint wrapper components from a parameter FQN."""
+    return ".".join(p for p in name.split(".") if p not in _WRAPPER_PARTS)
+
+
 @dataclass
 class _muon_state:
     worker_rank: int
@@ -78,6 +90,7 @@ def adjust_lr_for_muon(lr, param_shape):
 
 
 def default_is_muon(name, x, expert_keys=None):
+    name = normalize_fqn(name)
     skip_keys = ["embed_tokens", "lm_head", "tok_embeddings", "output"]
     if any(key in name for key in skip_keys):
         return False

torch-ext/optimizer/muon.py

@@ -11,7 +11,8 @@ from torch.profiler import record_function
 from .adamw import step_adamw
 from .async_utils import run_pipeline
 from .core import (_muon_state, adjust_lr_for_muon,
-                   get_default_muon_param_groups, update_g, update_p)
+                   get_default_muon_param_groups, normalize_fqn, update_g,
+                   update_p)
 from .distributed.utils import (_is_shard, construct_shard_mesh,
                                 get_slices_of_dtensor)
 from .newton_schulz import (COMM_DTYPE, DEFAULT_CHUNK_SIZE_RATIO,
@@ -45,7 +46,8 @@ def _expand_expert_params(names, params, expert_keys):
     expanded_params = []
 
     for n, p in zip(names, params):
-        is_expert = expert_keys and any(key in n for key in expert_keys)
+        is_expert = expert_keys and any(
+            key in normalize_fqn(n) for key in expert_keys)
         is_dtensor = isinstance(p.data, DTensor)
 
         if not is_expert:

torch-ext/optimizer/qk_clip.py

@@ -5,6 +5,8 @@ from dataclasses import dataclass
 import torch
 from torch.distributed.tensor import DTensor
 
+from .core import normalize_fqn
+
 logger = logging.getLogger(__name__)
 
 
@@ -23,7 +25,7 @@ def parse_qk_layer(name: str) -> tuple[str | None, int]:
         'model.7.attn.k_proj.weight'  -> ('k_proj', 7)
         'model.4.attn.v_proj.weight'  -> (None, -1)
     """
-    parts = name.split('.')
+    parts = normalize_fqn(name).split('.')
     if len(parts) < 3:
         return None, -1