fix(megatron): narrow the output_layer tiedness fallback to expected exceptions

Only two exceptions are anticipated when reading Megatron-LM's
untie_embeddings_and_output_weights: ImportError when megatron.training is
absent, and AssertionError when get_args() runs before initialize_megatron.
Anything else now propagates rather than silently downgrading to "tied", which
is the path that exports an unquantized output_layer.

Adds coverage for the uninitialized case, which had none.

Signed-off-by: James Shen <yueshen@nvidia.com>

Author: yueshen2016

modelopt/torch/quantization/plugins/megatron.py

@@ -309,7 +309,8 @@ def _output_layer_untied(config) -> bool:
         from megatron.training import get_args as _mlm_get_args
 
         return bool(getattr(_mlm_get_args(), "untie_embeddings_and_output_weights", False))
-    except Exception as e:
+    except (ImportError, AssertionError) as e:
+        # ImportError: no megatron.training. AssertionError: get_args() before initialize_megatron.
         # Warn once per config rather than on every save and every load.
         if not getattr(config, "_modelopt_warned_output_layer_untied", False):
             warn_rank_0(

tests/gpu_megatron/torch/quantization/plugins/test_megatron.py

@@ -1862,6 +1862,28 @@ class _Config:
     assert warn.call_count == 1
 
 
+def test_output_layer_untied_warns_when_args_uninitialized():
+    """Megatron-LM importable but not initialized: treated as tied, warned once."""
+
+    class _Config:
+        pass
+
+    def _uninitialized():
+        raise AssertionError("args is not initialized.")
+
+    fake_training = types.ModuleType("megatron.training")
+    fake_training.get_args = _uninitialized
+
+    config = _Config()
+    with (
+        patch.dict(sys.modules, {"megatron.training": fake_training}),
+        patch("modelopt.torch.quantization.plugins.megatron.warn_rank_0") as warn,
+    ):
+        assert _output_layer_untied(config) is False
+        assert _output_layer_untied(config) is False
+    assert warn.call_count == 1
+
+
 def test_output_layer_untied_not_stamped_onto_teacher_config():
     """A distillation teacher keeps its own tiedness; the student's answer must not leak in."""