test: Cover compiled Megatron wrapper names

Author: vivekkalyan

tests/unit/test_megatron_dedicated.py

@@ -18,6 +18,7 @@
     MergedWeightTransferInitInfo,
 )
 from art.megatron.service import MegatronService
+from art.megatron.train import _unwrap_art_wrapper_name
 
 
 @pytest.mark.asyncio
@@ -132,6 +133,21 @@ async def fake_create_subprocess_shell(
     assert seen["start_new_session"] is True
 
 
+def test_unwrap_art_wrapper_name_strips_compiled_wrapper_segments() -> None:
+    assert (
+        _unwrap_art_wrapper_name(
+            "module.module.decoder.layers.0._orig_mod.self_attention.linear_proj.linear_proj.weight"
+        )
+        == "decoder.layers.0.self_attention.linear_proj.weight"
+    )
+    assert (
+        _unwrap_art_wrapper_name(
+            "module.module.decoder.layers.0._orig_mod.mlp.experts.linear_fc1.linear_fc1.weight7"
+        )
+        == "decoder.layers.0.mlp.experts.linear_fc1.weight7"
+    )
+
+
 @pytest.mark.asyncio
 async def test_megatron_service_start_openai_server_dedicated_starts_subprocess(
     tmp_path: Path,