Refactor puzzletron imports: relative imports, public API, logger fix

- Convert all internal absolute imports to relative imports across puzzletron
- Add `__all__` to each submodule; rewrite `__init__.py` files to use `from .X import *`
- Rename `puzzletron.py` → `entrypoint.py`; move `puzzletron_nas_plugin` from `nas/plugins/` to top-level
- Standardize external callers to `import modelopt.torch.puzzletron as mtpz`
- Fix `tools/logger.py`: remove global torch logger hijack (`torch_logger.handlers = ...`) and `setLoggerClass(DistributedLogger)` that caused verbose torch._dynamo output and test hangs

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>

Author: kevalmorabia97

examples/llm_eval/lm_eval_hf.py

@@ -52,11 +52,7 @@
 from modelopt.torch.sparsity.attention_sparsity.conversion import is_attn_sparsified
 
 try:
-    import modelopt.torch.puzzletron.anymodel.models  # noqa: F401
-    from modelopt.torch.puzzletron.anymodel.model_descriptor.model_descriptor_factory import (
-        resolve_descriptor_from_pretrained,
-    )
-    from modelopt.torch.puzzletron.anymodel.puzzformer import deci_x_patcher
+    import modelopt.torch.puzzletron as mtpz
 
     _ANYMODEL_AVAILABLE = True
 except ImportError:
@@ -68,12 +64,12 @@ def _anymodel_patcher_context(pretrained, trust_remote_code=False):
     if not _ANYMODEL_AVAILABLE or not pretrained:
         return contextlib.nullcontext()
     try:
-        descriptor = resolve_descriptor_from_pretrained(
+        descriptor = mtpz.resolve_descriptor_from_pretrained(
             pretrained, trust_remote_code=trust_remote_code
         )
     except (ValueError, AttributeError):
         return contextlib.nullcontext()
-    return deci_x_patcher(model_descriptor=descriptor)
+    return mtpz.deci_x_patcher(model_descriptor=descriptor)
 
 
 def create_from_arg_obj(cls: type[T], arg_dict: dict, additional_config: dict | None = None) -> T:

examples/puzzletron/evaluation/hf_deployable_anymodel.py

@@ -31,10 +31,7 @@
 from peft import PeftModel
 from transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer
 
-from modelopt.torch.puzzletron.anymodel.model_descriptor.model_descriptor_factory import (
-    resolve_descriptor_from_pretrained,
-)
-from modelopt.torch.puzzletron.anymodel.puzzformer import deci_x_patcher
+import modelopt.torch.puzzletron as mtpz
 
 try:
     from pytriton.decorators import batch
@@ -148,11 +145,11 @@ def _load(
             # See: modelopt/torch/puzzletron/anymodel/puzzformer/utils.py
             # =========================================================================
 
-            descriptor = resolve_descriptor_from_pretrained(
+            descriptor = mtpz.resolve_descriptor_from_pretrained(
                 self.hf_model_id_path, trust_remote_code=hf_kwargs.get("trust_remote_code", False)
             )
 
-            with deci_x_patcher(model_descriptor=descriptor):
+            with mtpz.deci_x_patcher(model_descriptor=descriptor):
                 self.model = AutoModelForCausalLM.from_pretrained(
                     self.hf_model_id_path,
                     torch_dtype=torch_dtype,

examples/puzzletron/main.py

@@ -37,15 +37,8 @@
 from pathlib import Path
 
 import modelopt.torch.nas as mtn
-import modelopt.torch.puzzletron.mip.mip_and_realize_models as mip_and_realize_models
-import modelopt.torch.puzzletron.mip.sweep as sweep
+import modelopt.torch.puzzletron as mtpz
 import modelopt.torch.utils.distributed as dist
-from modelopt.torch.puzzletron.nas.plugins.puzzletron_nas_plugin import PuzzletronModel
-from modelopt.torch.puzzletron.tools.hydra_utils import (
-    initialize_hydra_config_for_dir,
-    register_hydra_resolvers,
-)
-from modelopt.torch.puzzletron.tools.logger import mprint
 
 
 def parse_args():
@@ -74,26 +67,26 @@ def run_full_puzzletron(hydra_config_path: str):
     Args:
         config_path: Path to the YAML configuration file
     """
-    mprint("Puzzletron Progress 1/8: starting puzzletron pipeline")
+    mtpz.tools.mprint("Puzzletron Progress 1/8: starting puzzletron pipeline")
     dist.setup(timeout=timedelta(minutes=10))
 
     # Register Hydra custom resolvers (needed for config resolution)
-    register_hydra_resolvers()
+    mtpz.tools.register_hydra_resolvers()
 
     hydra_config_path = Path(hydra_config_path).resolve()
     hydra_config_dir = str(hydra_config_path.parent)
     hydra_config_name = hydra_config_path.stem
 
     # Load hydra config
-    hydra_cfg = initialize_hydra_config_for_dir(
+    hydra_cfg = mtpz.tools.initialize_hydra_config_for_dir(
         config_dir=hydra_config_dir,
         config_name=hydra_config_name,
         overrides=[],
     )
 
     # Convert model (convert from HF to DeciLM, score pruning activations,
     # prune the model and save pruned checkpoints)
-    input_model = PuzzletronModel()
+    input_model = mtpz.puzzletron_nas_plugin.PuzzletronModel()
     converted_model = mtn.convert(
         input_model,
         mode=[
@@ -120,7 +113,7 @@ def run_full_puzzletron(hydra_config_path: str):
     )
 
     dist.cleanup()
-    mprint("Puzzletron Progress 8/8: puzzletron pipeline completed (multi-gpu)")
+    mtpz.tools.mprint("Puzzletron Progress 8/8: puzzletron pipeline completed (multi-gpu)")
 
 
 def run_mip_only(hydra_config_path: str):
@@ -135,33 +128,33 @@ def run_mip_only(hydra_config_path: str):
     dist.setup(timeout=timedelta(minutes=10))
 
     # Register Hydra custom resolvers (needed for config resolution)
-    register_hydra_resolvers()
+    mtpz.tools.register_hydra_resolvers()
 
     hydra_config_path = Path(hydra_config_path).resolve()
     hydra_config_dir = str(hydra_config_path.parent)
     hydra_config_name = hydra_config_path.stem
 
     # Load hydra config
-    hydra_cfg = initialize_hydra_config_for_dir(
+    hydra_cfg = mtpz.tools.initialize_hydra_config_for_dir(
         config_dir=hydra_config_dir,
         config_name=hydra_config_name,
         overrides=[],
     )
 
     # Check if sweep mode is enabled
     if hasattr(hydra_cfg.mip, "sweep") and hydra_cfg.mip.sweep.get("enabled", False):
-        mprint(
+        mtpz.tools.mprint(
             "Puzzletron Progress 7/8: running MIP sweep for multiple compression rates (multi-gpu)"
         )
-        sweep.run_mip_sweep(hydra_cfg)
+        mtpz.mip.run_mip_sweep(hydra_cfg)
     else:
         # mip_and_realize_models (distributed processing)
         # TODO: How to make it part of mnt.search() api, similarly to run_full_puzzletron() API
-        mprint("Puzzletron Progress 7/8: running MIP and realizing models (multi-gpu)")
-        mip_and_realize_models.launch_mip_and_realize_model(hydra_cfg)
+        mtpz.tools.mprint("Puzzletron Progress 7/8: running MIP and realizing models (multi-gpu)")
+        mtpz.mip.launch_mip_and_realize_model(hydra_cfg)
 
     dist.cleanup()
-    mprint("Puzzletron Progress 8/8: puzzletron pipeline completed (multi-gpu)")
+    mtpz.tools.mprint("Puzzletron Progress 8/8: puzzletron pipeline completed (multi-gpu)")
 
 
 def main():

examples/puzzletron/mbridge_distillation/distill_hf.py

@@ -46,7 +46,7 @@
 from megatron.core.distributed import DistributedDataParallelConfig
 
 # Import to register heterogeneous bridges (side effect)
-import modelopt.torch.puzzletron.export.mbridge  # noqa: F401
+import modelopt.torch.puzzletron as mtpz  # noqa: F401
 import modelopt.torch.utils.distributed as dist
 from modelopt.torch.utils import print_rank_0
 

modelopt/torch/puzzletron/__init__.py

@@ -13,3 +13,26 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# NOTE: Some modules also trigger factory registration as side effect
+from . import (
+    activation_scoring,
+    anymodel,
+    block_config,
+    build_library_and_stats,
+    dataset,
+    entrypoint,
+    export,
+    mip,
+    pruning,
+    puzzletron_nas_plugin,
+    replacement_library,
+    scoring,
+    sewing_kit,
+    subblock_stats,
+    tools,
+    utils,
+)
+
+# Import functions from important modules to top-level
+from .anymodel import *
+from .entrypoint import *

modelopt/torch/puzzletron/activation_scoring/__init__.py

@@ -0,0 +1,17 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .activation_hooks import *
+from .score_pruning_activations import *

modelopt/torch/puzzletron/activation_scoring/activation_hooks/utils.py

@@ -22,8 +22,9 @@
 import torch
 
 from modelopt.torch.prune.importance_hooks.base_hooks import ForwardHook as ActivationsHook
-from modelopt.torch.puzzletron.tools.logger import aprint
-from modelopt.torch.puzzletron.utils.dummy_modules import DummyBlock, DummyModule
+
+from ...tools.logger import aprint
+from ...utils.dummy_modules import DummyBlock, DummyModule
 
 
 def register_activation_hooks(

modelopt/torch/puzzletron/activation_scoring/score_pruning_activations.py

@@ -19,8 +19,8 @@
 from omegaconf import DictConfig
 
 import modelopt.torch.utils.distributed as dist
-from modelopt.torch.puzzletron.tools.logger import mprint
-from modelopt.torch.puzzletron.tools.validate_model import validate_model
+
+from ..tools.logger import mprint
 
 
 def has_checkpoint_support(activation_hooks_kwargs: dict) -> bool:
@@ -127,10 +127,9 @@ def should_skip_scoring_completely(cfg: DictConfig) -> bool:
     return is_completed
 
 
-# Old progress tracking removed - checkpoint manager handles all progress tracking
-
-
 def launch_score_activations(cfg: DictConfig):
+    from ..tools.validate_model import validate_model
+
     # Check if we should skip scoring entirely (only if 100% complete)
     if should_skip_scoring_completely(cfg):
         return

modelopt/torch/puzzletron/anymodel/__init__.py

@@ -37,28 +37,7 @@
     - (more to come: qwen2, mistral_small, etc.)
 """
 
-# Import models to trigger factory registration
-from modelopt.torch.puzzletron.anymodel import models  # noqa: F401
-from modelopt.torch.puzzletron.anymodel.converter import Converter, ConverterFactory, convert_model
-from modelopt.torch.puzzletron.anymodel.model_descriptor import (
-    ModelDescriptor,
-    ModelDescriptorFactory,
-)
-from modelopt.torch.puzzletron.anymodel.puzzformer import (
-    MatchingZeros,
-    Same,
-    deci_x_patcher,
-    return_tuple_of_size,
-)
-
-__all__ = [
-    "Converter",
-    "ConverterFactory",
-    "ModelDescriptor",
-    "ModelDescriptorFactory",
-    "deci_x_patcher",
-    "MatchingZeros",
-    "Same",
-    "return_tuple_of_size",
-    "convert_model",
-]
+from . import models  # trigger factory registration
+from .converter import *
+from .model_descriptor import *
+from .puzzformer import *

modelopt/torch/puzzletron/anymodel/converter/convert_any_model.py

@@ -18,9 +18,9 @@
 
 from pathlib import Path
 
-from modelopt.torch.puzzletron.anymodel.converter.converter import Converter
-from modelopt.torch.puzzletron.anymodel.converter.converter_factory import ConverterFactory
-from modelopt.torch.puzzletron.anymodel.model_descriptor import ModelDescriptorFactory
+from ..model_descriptor import ModelDescriptorFactory
+from .converter import Converter
+from .converter_factory import ConverterFactory
 
 __all__ = ["convert_model"]