feat: Add NPU export (TFLite, QNN, CoreML)

optimization/npu_export.py

@@ -0,0 +1,450 @@
+"""
+NPU Export Module for MiniMind Max2
+Export to TFLite, QNN (Qualcomm), and other NPU formats.
+"""
+
+from dataclasses import dataclass
+from typing import List, Optional, Dict, Any, Tuple, Union
+from pathlib import Path
+import torch
+import torch.nn as nn
+import json
+
+
+@dataclass
+class NPUExportConfig:
+    """Configuration for NPU export."""
+    # Target platforms
+    target_platform: str = "tflite"  # tflite, qnn, coreml, nnapi
+
+    # Quantization
+    quantization: str = "int8"  # float16, int8, int4
+    calibration_samples: int = 100
+
+    # Optimization
+    optimize_for_inference: bool = True
+    enable_xnnpack: bool = True  # TFLite XNNPACK delegate
+
+    # Model settings
+    max_sequence_length: int = 2048
+    batch_size: int = 1
+
+    # QNN specific
+    qnn_target: str = "gpu"  # cpu, gpu, dsp, htp
+
+    # Output
+    include_metadata: bool = True
+
+
+class TFLiteExporter:
+    """Export MiniMind models to TensorFlow Lite format."""
+
+    def __init__(self, config: NPUExportConfig):
+        self.config = config
+
+    def export(
+        self,
+        model: nn.Module,
+        output_path: str,
+        sample_input: Optional[torch.Tensor] = None,
+    ) -> str:
+        """
+        Export model to TFLite format.
+
+        Args:
+            model: PyTorch model to export
+            output_path: Path for output .tflite file
+            sample_input: Sample input for tracing
+
+        Returns:
+            Path to exported model
+        """
+        try:
+            import tensorflow as tf
+        except ImportError:
+            print("TensorFlow not installed. Install with: pip install tensorflow")
+            return self._export_via_onnx(model, output_path, sample_input)
+
+        model.eval()
+
+        # Get model config
+        if hasattr(model, 'config'):
+            vocab_size = model.config.vocab_size
+            hidden_size = model.config.hidden_size
+        else:
+            vocab_size = 102400
+            hidden_size = 1024
+
+        # Create sample input if not provided
+        if sample_input is None:
+            sample_input = torch.randint(
+                0, vocab_size,
+                (self.config.batch_size, self.config.max_sequence_length),
+            )
+
+        # Export via ONNX as intermediate
+        onnx_path = output_path.replace('.tflite', '.onnx')
+        self._export_to_onnx(model, onnx_path, sample_input)
+
+        # Convert ONNX to TFLite
+        try:
+            import onnx
+            from onnx_tf.backend import prepare
+
+            # Load ONNX model
+            onnx_model = onnx.load(onnx_path)
+            tf_rep = prepare(onnx_model)
+
+            # Save as SavedModel
+            saved_model_path = output_path.replace('.tflite', '_saved_model')
+            tf_rep.export_graph(saved_model_path)
+
+            # Convert to TFLite
+            converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_path)
+
+            # Quantization settings
+            if self.config.quantization == "int8":
+                converter.optimizations = [tf.lite.Optimize.DEFAULT]
+                converter.target_spec.supported_types = [tf.int8]
+            elif self.config.quantization == "float16":
+                converter.optimizations = [tf.lite.Optimize.DEFAULT]
+                converter.target_spec.supported_types = [tf.float16]
+
+            # Enable optimizations
+            if self.config.optimize_for_inference:
+                converter.optimizations = [tf.lite.Optimize.DEFAULT]
+
+            tflite_model = converter.convert()
+
+            # Save
+            with open(output_path, 'wb') as f:
+                f.write(tflite_model)
+
+            print(f"Exported TFLite model to: {output_path}")
+            return output_path
+
+        except Exception as e:
+            print(f"TFLite conversion failed: {e}")
+            return onnx_path
+
+    def _export_to_onnx(
+        self,
+        model: nn.Module,
+        output_path: str,
+        sample_input: torch.Tensor,
+    ) -> str:
+        """Export to ONNX as intermediate format."""
+        torch.onnx.export(
+            model,
+            sample_input,
+            output_path,
+            export_params=True,
+            opset_version=14,
+            do_constant_folding=True,
+            input_names=['input_ids'],
+            output_names=['logits'],
+            dynamic_axes={
+                'input_ids': {0: 'batch_size', 1: 'sequence_length'},
+                'logits': {0: 'batch_size', 1: 'sequence_length'},
+            },
+        )
+        return output_path
+
+    def _export_via_onnx(
+        self,
+        model: nn.Module,
+        output_path: str,
+        sample_input: torch.Tensor,
+    ) -> str:
+        """Fallback: export to ONNX only."""
+        onnx_path = output_path.replace('.tflite', '.onnx')
+        return self._export_to_onnx(model, onnx_path, sample_input)
+
+
+class QNNExporter:
+    """Export MiniMind models to Qualcomm QNN format."""
+
+    def __init__(self, config: NPUExportConfig):
+        self.config = config
+
+    def export(
+        self,
+        model: nn.Module,
+        output_path: str,
+        sample_input: Optional[torch.Tensor] = None,
+    ) -> Dict[str, str]:
+        """
+        Export model to QNN format for Qualcomm NPUs.
+
+        Returns:
+            Dictionary with paths to exported files
+        """
+        model.eval()
+
+        # Get model config
+        if hasattr(model, 'config'):
+            vocab_size = model.config.vocab_size
+        else:
+            vocab_size = 102400
+
+        if sample_input is None:
+            sample_input = torch.randint(
+                0, vocab_size,
+                (self.config.batch_size, self.config.max_sequence_length),
+            )
+
+        output_dir = Path(output_path).parent
+        output_dir.mkdir(parents=True, exist_ok=True)
+
+        # Step 1: Export to ONNX
+        onnx_path = str(output_dir / "model.onnx")
+        torch.onnx.export(
+            model,
+            sample_input,
+            onnx_path,
+            export_params=True,
+            opset_version=14,
+            do_constant_folding=True,
+            input_names=['input_ids'],
+            output_names=['logits'],
+        )
+
+        outputs = {"onnx": onnx_path}
+
+        # Step 2: Generate QNN conversion script
+        qnn_script = self._generate_qnn_script(onnx_path, output_path)
+        script_path = str(output_dir / "convert_to_qnn.sh")
+        with open(script_path, 'w') as f:
+            f.write(qnn_script)
+
+        outputs["conversion_script"] = script_path
+
+        # Step 3: Generate model config for QNN
+        config_path = str(output_dir / "qnn_config.json")
+        qnn_config = {
+            "model_name": "minimind_max2",
+            "input_tensors": [{
+                "name": "input_ids",
+                "dims": [self.config.batch_size, self.config.max_sequence_length],
+                "data_type": "int32"
+            }],
+            "output_tensors": [{
+                "name": "logits",
+                "data_type": "float32"
+            }],
+            "backend": self.config.qnn_target,
+            "quantization": self.config.quantization,
+        }
+        with open(config_path, 'w') as f:
+            json.dump(qnn_config, f, indent=2)
+
+        outputs["config"] = config_path
+
+        print(f"QNN export prepared. Run {script_path} with QNN SDK installed.")
+        return outputs
+
+    def _generate_qnn_script(self, onnx_path: str, output_path: str) -> str:
+        """Generate shell script for QNN conversion."""
+        return f'''#!/bin/bash
+# QNN Conversion Script for MiniMind Max2
+# Requires Qualcomm QNN SDK
+
+# Check QNN SDK
+if [ -z "$QNN_SDK_ROOT" ]; then
+    echo "Error: QNN_SDK_ROOT not set. Please install Qualcomm QNN SDK."
+    exit 1
+fi
+
+# Convert ONNX to QNN
+$QNN_SDK_ROOT/bin/x86_64-linux-clang/qnn-onnx-converter \\
+    --input_network {onnx_path} \\
+    --output_path {output_path}.cpp
+
+# Compile model library
+$QNN_SDK_ROOT/bin/x86_64-linux-clang/qnn-model-lib-generator \\
+    -c {output_path}.cpp \\
+    -b {output_path}.bin \\
+    -t {self.config.qnn_target}
+
+echo "QNN model exported to {output_path}.bin"
+'''
+
+
+class CoreMLExporter:
+    """Export MiniMind models to Apple Core ML format."""
+
+    def __init__(self, config: NPUExportConfig):
+        self.config = config
+
+    def export(
+        self,
+        model: nn.Module,
+        output_path: str,
+        sample_input: Optional[torch.Tensor] = None,
+    ) -> str:
+        """Export model to Core ML format for Apple Neural Engine."""
+        try:
+            import coremltools as ct
+        except ImportError:
+            print("coremltools not installed. Install with: pip install coremltools")
+            return ""
+
+        model.eval()
+
+        # Get model config
+        if hasattr(model, 'config'):
+            vocab_size = model.config.vocab_size
+        else:
+            vocab_size = 102400
+
+        if sample_input is None:
+            sample_input = torch.randint(
+                0, vocab_size,
+                (self.config.batch_size, self.config.max_sequence_length),
+            )
+
+        # Trace model
+        traced = torch.jit.trace(model, sample_input)
+
+        # Convert to Core ML
+        mlmodel = ct.convert(
+            traced,
+            inputs=[ct.TensorType(
+                name="input_ids",
+                shape=sample_input.shape,
+                dtype=int,
+            )],
+            compute_units=ct.ComputeUnit.ALL,  # Use Neural Engine when available
+        )
+
+        # Quantization
+        if self.config.quantization == "float16":
+            mlmodel = ct.models.neural_network.quantization_utils.quantize_weights(
+                mlmodel, nbits=16
+            )
+        elif self.config.quantization == "int8":
+            mlmodel = ct.models.neural_network.quantization_utils.quantize_weights(
+                mlmodel, nbits=8
+            )
+
+        # Save
+        mlmodel.save(output_path)
+        print(f"Core ML model exported to: {output_path}")
+        return output_path
+
+
+class NPUExporter:
+    """Unified NPU export interface."""
+
+    def __init__(self, config: Optional[NPUExportConfig] = None):
+        self.config = config or NPUExportConfig()
+
+        self.exporters = {
+            "tflite": TFLiteExporter(self.config),
+            "qnn": QNNExporter(self.config),
+            "coreml": CoreMLExporter(self.config),
+        }
+
+    def export(
+        self,
+        model: nn.Module,
+        output_path: str,
+        target_platform: Optional[str] = None,
+        sample_input: Optional[torch.Tensor] = None,
+    ) -> Union[str, Dict[str, str]]:
+        """
+        Export model to specified NPU format.
+
+        Args:
+            model: PyTorch model
+            output_path: Output file path
+            target_platform: Target platform (tflite, qnn, coreml)
+            sample_input: Sample input for tracing
+
+        Returns:
+            Path(s) to exported model(s)
+        """
+        platform = target_platform or self.config.target_platform
+
+        if platform not in self.exporters:
+            raise ValueError(f"Unknown platform: {platform}. Supported: {list(self.exporters.keys())}")
+
+        exporter = self.exporters[platform]
+        return exporter.export(model, output_path, sample_input)
+
+    def export_all(
+        self,
+        model: nn.Module,
+        output_dir: str,
+        sample_input: Optional[torch.Tensor] = None,
+    ) -> Dict[str, Any]:
+        """Export to all supported formats."""
+        output_dir = Path(output_dir)
+        output_dir.mkdir(parents=True, exist_ok=True)
+
+        results = {}
+
+        for platform, exporter in self.exporters.items():
+            try:
+                if platform == "tflite":
+                    path = str(output_dir / "model.tflite")
+                elif platform == "qnn":
+                    path = str(output_dir / "qnn" / "model")
+                elif platform == "coreml":
+                    path = str(output_dir / "model.mlpackage")
+                else:
+                    continue
+
+                result = exporter.export(model, path, sample_input)
+                results[platform] = {"success": True, "path": result}
+            except Exception as e:
+                results[platform] = {"success": False, "error": str(e)}
+
+        return results
+
+
+def export_for_mobile(
+    model: nn.Module,
+    output_dir: str,
+    platforms: Optional[List[str]] = None,
+    config: Optional[NPUExportConfig] = None,
+) -> Dict[str, Any]:
+    """
+    High-level function to export model for mobile devices.
+
+    Args:
+        model: PyTorch model
+        output_dir: Output directory
+        platforms: List of target platforms (default: all)
+        config: Export configuration
+
+    Returns:
+        Dictionary with export results for each platform
+    """
+    config = config or NPUExportConfig()
+    exporter = NPUExporter(config)
+
+    if platforms is None:
+        return exporter.export_all(model, output_dir)
+
+    results = {}
+    output_dir = Path(output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    for platform in platforms:
+        try:
+            if platform == "tflite":
+                path = str(output_dir / "model.tflite")
+            elif platform == "qnn":
+                path = str(output_dir / "qnn" / "model")
+            elif platform == "coreml":
+                path = str(output_dir / "model.mlpackage")
+            else:
+                continue
+
+            result = exporter.export(model, path, target_platform=platform)
+            results[platform] = {"success": True, "path": result}
+        except Exception as e:
+            results[platform] = {"success": False, "error": str(e)}
+
+    return results