Merge branch 'feature/puzzletron' into jrausch/distillation-consolidation

Author: kevalmorabia97

examples/puzzletron/evaluation/hf_deployable_anymodel.py

@@ -331,7 +331,6 @@ def get_triton_input(self):
             Tensor(name="top_p", shape=(-1,), dtype=np.single, optional=True),
             Tensor(name="temperature", shape=(-1,), dtype=np.single, optional=True),
             Tensor(name="random_seed", shape=(-1,), dtype=np.int_, optional=True),
-            Tensor(name="max_length", shape=(-1,), dtype=np.int_, optional=True),
             Tensor(name="output_logits", shape=(-1,), dtype=np.bool_, optional=True),
             Tensor(name="output_scores", shape=(-1,), dtype=np.bool_, optional=True),
         )

modelopt/torch/puzzletron/__init__.py

@@ -27,7 +27,6 @@
     puzzletron_nas_plugin,
     replacement_library,
     scoring,
-    sewing_kit,
     subblock_stats,
     tools,
     utils,

modelopt/torch/puzzletron/activation_scoring/score_pruning_activations.py

@@ -22,6 +22,8 @@
 
 from ..tools.logger import mprint
 
+__all__ = ["launch_score_activations"]
+
 
 def has_checkpoint_support(activation_hooks_kwargs: dict) -> bool:
     """Determine if the activation hook method has proper checkpoint support implemented.

modelopt/torch/puzzletron/anymodel/models/gpt_oss/__init__.py

@@ -18,5 +18,5 @@
 
 """GPT-OSS model support for AnyModel."""
 
-from .gpt_oss_converter import GptOssConverter
-from .gpt_oss_model_descriptor import GptOssModelDescriptor
+from .gpt_oss_converter import *
+from .gpt_oss_model_descriptor import *

modelopt/torch/puzzletron/anymodel/models/gpt_oss/gpt_oss_converter.py

@@ -23,6 +23,8 @@
 from ....block_config import AttentionConfig, BlockConfig, FFNConfig, MoEConfig
 from ...converter import Converter, ConverterFactory
 
+__all__ = ["GptOssConverter"]
+
 
 @ConverterFactory.register_decorator("gpt_oss")
 class GptOssConverter(Converter):

modelopt/torch/puzzletron/anymodel/models/gpt_oss/gpt_oss_model_descriptor.py

@@ -37,6 +37,8 @@
 from ...model_descriptor import ModelDescriptor, ModelDescriptorFactory
 from ...puzzformer.no_op import MatchingZeros, Same, return_tuple_of_size
 
+__all__ = ["GptOssModelDescriptor", "GptOssExpertRemovalLayerDescriptor"]
+
 
 @ModelDescriptorFactory.register_decorator("gpt_oss")
 class GptOssModelDescriptor(ModelDescriptor):

modelopt/torch/puzzletron/anymodel/models/gpt_oss/gpt_oss_pruned_to_mxfp4.py

@@ -35,6 +35,8 @@
 from tqdm import tqdm
 from transformers.integrations.mxfp4 import convert_moe_packed_tensors
 
+__all__ = []
+
 
 def deduce_experts_for_layer(
     layer: int,

modelopt/torch/puzzletron/anymodel/models/llama/__init__.py

@@ -13,5 +13,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .llama_converter import LlamaConverter
-from .llama_model_descriptor import LlamaModelDescriptor
+from .llama_converter import *
+from .llama_model_descriptor import *

modelopt/torch/puzzletron/anymodel/models/llama/llama_converter.py

@@ -23,6 +23,8 @@
 from ....block_config import AttentionConfig, BlockConfig, FFNConfig
 from ...converter import Converter, ConverterFactory
 
+__all__ = ["LlamaConverter"]
+
 
 @ConverterFactory.register_decorator("llama")
 class LlamaConverter(Converter):

modelopt/torch/puzzletron/anymodel/models/llama/llama_model_descriptor.py

@@ -32,6 +32,12 @@
 from ...model_descriptor import ModelDescriptor, ModelDescriptorFactory
 from ...puzzformer.no_op import MatchingZeros, Same, return_tuple_of_size
 
+__all__ = [
+    "LlamaModelDescriptor",
+    "LlamaFFNIntermediateLayerDescriptor",
+    "LlamaKVHeadsLayerDescriptor",
+]
+
 
 @ModelDescriptorFactory.register_decorator("llama")
 class LlamaModelDescriptor(ModelDescriptor):