Refactor model specific code to example_utils

Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com>

Author: ajrasane

examples/llm_ptq/example_utils.py

@@ -42,11 +42,21 @@
     snapshot_download = None
 
 import modelopt.torch.quantization as mtq
+from modelopt.torch.utils.dataset_utils import get_dataset_dataloader
 from modelopt.torch.utils.image_processor import (
     BaseImageProcessor,
     MllamaImageProcessor,
     Qwen3OmniImageProcessor,
 )
+from modelopt.torch.utils.video_dataset_utils import (
+    Qwen3OmniVideoProcessor,
+    get_supported_video_datasets,
+    get_video_dataset_dataloader,
+)
+from modelopt.torch.utils.vlm_dataset_utils import (
+    get_supported_vlm_datasets,
+    get_vlm_dataset_dataloader,
+)
 
 SPECULATIVE_MODEL_LIST = ["Eagle", "Medusa"]
 
@@ -244,9 +254,33 @@ def build_quant_cfg(
         quant_cfg["quant_cfg"]["*self_attn.q*"] = {"enable": False}
         quant_cfg["quant_cfg"]["*self_attn.kv*"] = {"enable": False}
 
+    if model_type == "qwen3omni":
+        print(
+            "Disabling quantization for conv layers, audio tower and visual encoder in Qwen3Omni model"
+        )
+        quant_cfg["quant_cfg"]["*conv*"] = {"enable": False}
+        quant_cfg["quant_cfg"]["*audio_tower*"] = {"enable": False}
+        quant_cfg["quant_cfg"]["*visual*"] = {"enable": False}
+
     return quant_cfg
 
 
+def get_generation_kwargs(model_type: str) -> dict[str, Any]:
+    """Get model-specific generation kwargs for calibration.
+
+    Args:
+        model_type: The model type string.
+
+    Returns:
+        Dictionary of generation kwargs for the model.
+    """
+    generation_kwargs = {}
+    if model_type == "qwen3omni":
+        generation_kwargs["return_audio"] = False
+        generation_kwargs["thinker_max_new_tokens"] = 1
+    return generation_kwargs
+
+
 def is_speculative(hf_config):
     """Check if the model architecture is a speculative model."""
     return hf_config.architectures and any(
@@ -617,3 +651,81 @@ def copy_custom_model_files(source_path: str, export_path: str, trust_remote_cod
         print(f"Successfully copied {len(copied_files)} custom model files to {export_path}")
     else:
         print("No custom model files found to copy")
+
+
+def get_qwen3omni_dataloader(
+    dataset_name: str | list[str] | None,
+    processor: Qwen3OmniImageProcessor | None,
+    tokenizer,
+    batch_size: int,
+    num_samples: int | list[int],
+    device: torch.device,
+    model_dtype: torch.dtype,
+    include_labels: bool = False,
+):
+    """Create a calibration dataloader for Qwen3Omni models.
+
+    Handles video, VLM, and text-only dataset configurations.
+
+    Args:
+        dataset_name: Name of the dataset(s) to use for calibration.
+        processor: The Qwen3OmniImageProcessor for multimodal inputs.
+        tokenizer: The tokenizer for text-only fallback.
+        batch_size: Batch size for the dataloader.
+        num_samples: Number of samples to use (int or list for multi-dataset).
+        device: Target device for tensors.
+        model_dtype: Model dtype for proper tensor conversion.
+        include_labels: Whether to include labels (for gradient-based auto_quantize).
+
+    Returns:
+        DataLoader for calibration.
+    """
+    if dataset_name is None:
+        dataset_name = ["cnn_dailymail", "nemotron-post-training-dataset-v2"]
+
+    if processor is not None:
+        if dataset_name in get_supported_video_datasets():
+            assert isinstance(dataset_name, str)
+            video_processor = Qwen3OmniVideoProcessor(
+                processor.tokenizer if hasattr(processor, "tokenizer") else processor,
+                device=device,
+                dtype=model_dtype,
+                use_audio_in_video=True,
+            )
+            calib_dataloader = get_video_dataset_dataloader(
+                dataset_name=dataset_name,
+                processor=video_processor,
+                batch_size=batch_size,
+                num_samples=num_samples if isinstance(num_samples, int) else num_samples[0],
+            )
+        elif dataset_name in get_supported_vlm_datasets():
+            assert isinstance(dataset_name, str)
+            assert isinstance(processor, Qwen3OmniImageProcessor), (
+                "The Qwen3OmniImageProcessor must be set."
+            )
+            # Set the dtype for proper tensor conversion in collate_function
+            processor.dtype = model_dtype
+            calib_dataloader = get_vlm_dataset_dataloader(
+                dataset_name=dataset_name,
+                processor=processor,
+                batch_size=batch_size,
+                num_samples=num_samples if isinstance(num_samples, int) else num_samples[0],
+            )
+        else:
+            raise ValueError(
+                f"Dataset '{dataset_name}' not supported for Qwen3Omni with processor. "
+                f"Supported video datasets: {get_supported_video_datasets()}, "
+                f"Supported VLM datasets: {get_supported_vlm_datasets()}"
+            )
+    else:
+        # Text-only fallback
+        calib_dataloader = get_dataset_dataloader(
+            dataset_name=dataset_name if isinstance(dataset_name, list) else [dataset_name],
+            tokenizer=tokenizer,
+            batch_size=batch_size,
+            num_samples=num_samples if isinstance(num_samples, list) else [num_samples],
+            device=device,
+            include_labels=include_labels,
+        )
+
+    return calib_dataloader

examples/llm_ptq/hf_ptq.py

@@ -14,9 +14,7 @@
 # limitations under the License.
 
 import argparse
-import io
 import random
-import sys
 import time
 import warnings
 from typing import Any
@@ -28,8 +26,10 @@
     build_quant_cfg,
     copy_custom_model_files,
     create_vlm_calibration_loop,
+    get_generation_kwargs,
     get_model,
     get_processor,
+    get_qwen3omni_dataloader,
     get_tokenizer,
     is_enc_dec,
     is_nemotron_vl,
@@ -72,15 +72,7 @@
 )
 from modelopt.torch.utils.memory_monitor import launch_memory_monitor
 from modelopt.torch.utils.speech_dataset_utils import get_speech_dataset_dataloader
-from modelopt.torch.utils.video_dataset_utils import (
-    Qwen3OmniVideoProcessor,
-    get_supported_video_datasets,
-    get_video_dataset_dataloader,
-)
-from modelopt.torch.utils.vlm_dataset_utils import (
-    get_supported_vlm_datasets,
-    get_vlm_dataset_dataloader,
-)
+from modelopt.torch.utils.vlm_dataset_utils import get_vlm_dataset_dataloader
 
 RAND_SEED = 1234
 
@@ -194,47 +186,20 @@ def make_calib_dataloader(
             num_samples=args.calib_size[0],
         )
     elif model_type == "qwen3omni":
-        dataset_name = args.dataset[0] if args.dataset else "cnn_dailymail"
-        # Check if using video dataset (e.g., finevideo)
-        if processor is not None:
-            if dataset_name in get_supported_video_datasets():
-                video_processor = Qwen3OmniVideoProcessor(
-                    processor.tokenizer if hasattr(processor, "tokenizer") else processor,
-                    device=device,
-                    dtype=language_model.dtype,
-                    use_audio_in_video=True,
-                )
-                calib_dataloader = get_video_dataset_dataloader(
-                    dataset_name=dataset_name,
-                    processor=video_processor,
-                    batch_size=args.batch_size,
-                    num_samples=args.calib_size[0],
-                )
-            elif dataset_name in get_supported_vlm_datasets():
-                assert isinstance(processor, Qwen3OmniImageProcessor), (
-                    "The Qwen3OmniImageProcessor must be set."
-                )
-                # Set the dtype for proper tensor conversion in collate_function
-                processor.dtype = language_model.dtype
-                calib_dataloader = get_vlm_dataset_dataloader(
-                    dataset_name=dataset_name,
-                    processor=processor,
-                    batch_size=args.batch_size,
-                    num_samples=args.calib_size[0],
-                )
-        else:
-            # Labels are only needed for gradient-based auto_quantize
-            include_labels = (
-                args.auto_quantize_bits is not None and args.auto_quantize_method == "gradient"
-            )
-            calib_dataloader = get_dataset_dataloader(
-                dataset_name=args.dataset,
-                tokenizer=tokenizer,
-                batch_size=args.batch_size,
-                num_samples=args.calib_size,
-                device=device,
-                include_labels=include_labels,
-            )
+        # Labels are only needed for gradient-based auto_quantize
+        include_labels = (
+            args.auto_quantize_bits is not None and args.auto_quantize_method == "gradient"
+        )
+        calib_dataloader = get_qwen3omni_dataloader(
+            dataset_name=args.dataset[0] if args.dataset else None,
+            processor=processor,
+            tokenizer=tokenizer,
+            batch_size=args.batch_size,
+            num_samples=args.calib_size[0] if processor else args.calib_size,
+            device=device,
+            model_dtype=language_model.dtype,
+            include_labels=include_labels,
+        )
     elif model_type == "whisper":
         assert processor is not None and isinstance(processor, WhisperProcessor), (
             "The AutoProcessor must be set."
@@ -566,17 +531,8 @@ def mono_quantize(
         quant_cfg["quant_cfg"]["*radio*"] = {"enable": False}
         quant_cfg["quant_cfg"]["*visual*"] = {"enable": False}
 
-    # For Qwen3Omni models, disable quantization of conv layers
-    generation_kwargs = {}
-    if model_type == "qwen3omni":
-        print(
-            "Disabling quantization for conv layers, audio tower and visual encoder in Qwen3Omni model"
-        )
-        quant_cfg["quant_cfg"]["*conv*"] = {"enable": False}
-        quant_cfg["quant_cfg"]["*audio_tower*"] = {"enable": False}
-        quant_cfg["quant_cfg"]["*visual*"] = {"enable": False}
-        generation_kwargs["return_audio"] = False
-        generation_kwargs["thinker_max_new_tokens"] = 1
+    # Get model-specific generation kwargs (e.g., for Qwen3Omni)
+    generation_kwargs = get_generation_kwargs(model_type)
 
     if not model_is_already_quantized or calibration_only:
         if model_type == "gptoss" and args.qformat == "nvfp4_mlp_only":
@@ -799,20 +755,7 @@ def post_quantize(
     """
 
     if args.verbose:
-        if args.quant_summary_path:
-            # Capture the summary output to a file
-            old_stdout = sys.stdout
-            sys.stdout = buffer = io.StringIO()
-            try:
-                mtq.print_quant_summary(full_model)
-            finally:
-                sys.stdout = old_stdout
-            summary = buffer.getvalue()
-            with open(args.quant_summary_path, "w") as f:
-                f.write(summary)
-            print(f"Quantization summary saved to {args.quant_summary_path}")
-        else:
-            mtq.print_quant_summary(full_model)
+        mtq.print_quant_summary(full_model, save_path=args.quant_summary_path)
 
     # Run some samples
     torch.cuda.empty_cache()

modelopt/torch/export/model_utils.py

@@ -17,45 +17,45 @@
 import torch.nn as nn
 
 MODEL_NAME_TO_TYPE = {
-    "ArcticForCausalLM": "llama",
-    "baichuan": "baichuan",
-    "Bart": "bart",
-    "Bloom": "bloom",
-    "ChatGLM": "chatglm",
-    "Dbrx": "dbrx",
-    "Deepseek": "deepseek",
-    "ExaoneForCausalLM": "exaone",
-    "FalconForCausalLM": "falcon",
-    "Gemma": "gemma",
-    "Gemma2": "gemma2",
-    "Gemma3": "gemma3",
-    "GLM": "glm",
     "GPT2": "gpt",
-    "GPTJ": "gptj",
-    "gptoss": "gptoss",
-    "InternLM2ForCausalLM": "internlm",
-    "Llama": "llama",
+    "Mllama": "mllama",
     "Llama4": "llama4",
+    "Llama": "llama",
     "Mistral": "llama",
-    "MixtralForCausalLM": "llama",
-    "Mllama": "mllama",
+    "GPTJ": "gptj",
+    "FalconForCausalLM": "falcon",
+    "RWForCausalLM": "falcon",
+    "baichuan": "baichuan",
     "MPT": "mpt",
-    "Nemotron": "gpt",
-    "phi": "phi",
-    "phi3": "phi3",
-    "phi3small": "phi3small",
-    "Phi4MMForCausalLM": "phi4mm",
-    "PhiMoEForCausalLM": "phi3",
+    "Bloom": "bloom",
+    "ChatGLM": "chatglm",
     "Qwen3Moe": "qwen3moe",
     "Qwen3Next": "qwen3next",
     "Qwen3OmniMoeForConditionalGeneration": "qwen3omni",
     "QWen": "qwen",
     "RecurrentGemma": "recurrentgemma",
-    "RWForCausalLM": "falcon",
+    "Gemma3": "gemma3",
+    "Gemma2": "gemma2",
+    "Gemma": "gemma",
+    "phi3small": "phi3small",
+    "phi3": "phi3",
+    "PhiMoEForCausalLM": "phi3",
+    "Phi4MMForCausalLM": "phi4mm",
+    "phi": "phi",
+    "TLGv4ForCausalLM": "phi",
+    "MixtralForCausalLM": "llama",
+    "ArcticForCausalLM": "llama",
     "StarCoder": "gpt",
+    "Dbrx": "dbrx",
     "T5": "t5",
-    "TLGv4ForCausalLM": "phi",
+    "Bart": "bart",
+    "GLM": "glm",
+    "InternLM2ForCausalLM": "internlm",
+    "ExaoneForCausalLM": "exaone",
+    "Nemotron": "gpt",
+    "Deepseek": "deepseek",
     "Whisper": "whisper",
+    "gptoss": "gptoss",
 }
 
 __doc__ = f"""Utility functions for model type detection and classification.

modelopt/torch/export/unified_export_hf.py

@@ -1003,10 +1003,12 @@ def export_hf_checkpoint(
         if hasattr(model, "generation_config") and model.generation_config is not None:
             gen_config = model.generation_config
             if not getattr(gen_config, "do_sample", True):
-                # Remove sampling-related params when do_sample is False
-                for attr in ["temperature", "top_p", "top_k"]:
-                    if hasattr(gen_config, attr):
-                        setattr(gen_config, attr, None)
+                # Enable sampling if sampling params are present
+                if any(
+                    getattr(gen_config, attr, None) is not None
+                    for attr in ["temperature", "top_p", "top_k"]
+                ):
+                    gen_config.do_sample = True
 
         # Save model
         model.save_pretrained(