optimize.py

# -------------------------------------------------------------------------
# Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved.
# Portions of this file consist of AI generated content.
# --------------------------------------------------------------------------
# SPDX-License-Identifier: MIT
# --------------------------------------------------------------------------
"""End-to-end optimization pipeline for Qwen3.5-35B-A3B MoE VLM.

Exports three sub-models (vision encoder, text embedding, text decoder),
applies graph optimizations and INT4 quantization via Olive passes.

Usage:
    python optimize.py --config-dir cpu_and_mobile --device cpu
    python optimize.py --config-dir cpu_and_mobile --device cpu --skip-export
"""
import argparse
import json
import logging
from pathlib import Path

logging.getLogger("onnxscript").setLevel(logging.WARNING)
logging.getLogger("onnx_ir").setLevel(logging.WARNING)

MODELS_DIR = "models"


def _read_model_name(config_dir: str) -> str:
    """Read the HuggingFace model name from the Olive text config."""
    text_cfg = json.loads((Path(config_dir) / "text.json").read_text())
    return text_cfg["input_model"]["model_path"]


def export_models(config_dir: str):
    """Run Olive for all 3 sub-models (text, embedding, vision)."""
    from olive import run

    config_path = Path(config_dir)
    print(f"=== Running Olive pipelines (configs from {config_path}) ===")
    for config in ("text.json", "embedding.json", "vision.json"):
        print(f"  Running {config}...")
        run(str(config_path / config))
    print()


def update_genai_config(output_dir: str, model_name: str, device: str = "cpu", context_length: int = 4096):
    """Patch genai_config.json with embedding/vision sections and processor_config.

    Reads token IDs, vision parameters, and preprocessor settings from the
    HuggingFace model config rather than hardcoding them.
    """
    from transformers import AutoConfig, GenerationConfig
    from huggingface_hub import hf_hub_download

    hf_config = AutoConfig.from_pretrained(model_name)
    gen_config = GenerationConfig.from_pretrained(model_name)
    vc = hf_config.vision_config

    config_path = Path(output_dir) / "genai_config.json"
    with open(config_path) as f:
        config = json.load(f)

    if device == "gpu":
        provider_options = [
            {"cuda": {"enable_cuda_graph": "1", "enable_skip_layer_norm_strict_mode": "1"}}
        ]
        vision_provider_options = [
            {"cuda": {"enable_cuda_graph": "0", "enable_skip_layer_norm_strict_mode": "1"}}
        ]
    else:
        provider_options = []
        vision_provider_options = []

    vision_session_options = {"log_id": "onnxruntime-genai", "provider_options": vision_provider_options}

    config["model"]["decoder"]["filename"] = "text.onnx"

    config["model"]["embedding"] = {
        "filename": "embedding.onnx",
        "inputs": {"input_ids": "input_ids", "image_features": "image_features"},
        "outputs": {"inputs_embeds": "inputs_embeds"},
        "session_options": vision_session_options,
    }

    config["model"]["vision"] = {
        "filename": "vision.onnx",
        "config_filename": "processor_config.json",
        "spatial_merge_size": vc.spatial_merge_size,
        "tokens_per_second": 2.0,
        "patch_size": vc.patch_size,
        "inputs": {"pixel_values": "pixel_values", "image_grid_thw": "image_grid_thw"},
        "outputs": {"image_features": "image_features"},
        "session_options": vision_session_options,
    }

    config["model"]["bos_token_id"] = gen_config.bos_token_id
    config["model"]["context_length"] = context_length
    config["model"]["eos_token_id"] = gen_config.eos_token_id
    config["model"]["pad_token_id"] = gen_config.pad_token_id
    config["model"]["image_token_id"] = hf_config.image_token_id
    config["model"]["video_token_id"] = hf_config.video_token_id
    config["model"]["vision_start_token_id"] = hf_config.vision_start_token_id

    config["search"]["max_length"] = context_length
    if gen_config.top_k is not None:
        config["search"]["top_k"] = gen_config.top_k
    if gen_config.top_p is not None and config["search"].get("top_p") is None:
        config["search"]["top_p"] = gen_config.top_p

    with open(config_path, "w") as f:
        json.dump(config, f, indent=4)
    print(f"  Updated {config_path}")

    pp_path = hf_hub_download(model_name, "preprocessor_config.json")
    with open(pp_path) as f:
        pp = json.load(f)
    pp_size = pp.get("size", {})

    processor_config = {
        "processor": {
            "name": "qwen2_5_image_processor",
            "transforms": [
                {"operation": {"name": "decode_image", "type": "DecodeImage", "attrs": {"color_space": "RGB"}}},
                {"operation": {"name": "convert_to_rgb", "type": "ConvertRGB"}},
                {"operation": {"name": "resize", "type": "Resize", "attrs": {
                    "width": 960, "height": 672, "smart_resize": 1,
                    "min_pixels": pp_size.get("shortest_edge", 65536),
                    "max_pixels": pp_size.get("longest_edge", 16777216),
                    "patch_size": vc.patch_size,
                    "merge_size": vc.spatial_merge_size,
                }}},
                {"operation": {"name": "rescale", "type": "Rescale", "attrs": {
                    "rescale_factor": 1.0 / 255,
                }}},
                {"operation": {"name": "normalize", "type": "Normalize", "attrs": {
                    "mean": pp.get("image_mean", [0.5, 0.5, 0.5]),
                    "std": pp.get("image_std", [0.5, 0.5, 0.5]),
                    "qwen2_5_vl": 1,
                }}},
                {"operation": {"name": "patch_image", "type": "PatchImage", "attrs": {
                    "patch_size": vc.patch_size,
                    "temporal_patch_size": vc.temporal_patch_size,
                    "merge_size": vc.spatial_merge_size,
                }}},
            ],
        }
    }

    processor_path = Path(output_dir) / "processor_config.json"
    with open(processor_path, "w") as f:
        json.dump(processor_config, f, indent=2)
    print(f"  Created {processor_path}")


def fix_tokenizer(output_dir: str = MODELS_DIR):
    """Fix tokenizer.json for C++ std::regex compatibility."""
    tk_path = Path(output_dir) / "tokenizer.json"
    if not tk_path.exists():
        return
    tk = json.loads(tk_path.read_text(encoding="utf-8"))
    pt = tk.get("pre_tokenizer", {})
    if pt.get("type") == "Sequence":
        pt["pretokenizers"] = [s for s in pt["pretokenizers"] if s.get("type") == "ByteLevel"]
        for s in pt["pretokenizers"]:
            s["use_regex"] = True
    tk_path.write_text(json.dumps(tk, ensure_ascii=False), encoding="utf-8")

    tc_path = Path(output_dir) / "tokenizer_config.json"
    if tc_path.exists():
        tc = json.loads(tc_path.read_text(encoding="utf-8"))
        tc["tokenizer_class"] = "Qwen2Tokenizer"
        tc_path.write_text(json.dumps(tc, indent=2, ensure_ascii=False), encoding="utf-8")
    print(f"  Fixed tokenizer for C++ std::regex compatibility")


def main():
    parser = argparse.ArgumentParser(description="Optimize Qwen3.5-35B-A3B MoE VLM")
    parser.add_argument("--device", choices=["gpu", "cpu"], default="cpu")
    parser.add_argument("--config-dir", default="cpu_and_mobile")
    parser.add_argument("--skip-export", action="store_true")
    parser.add_argument("--models-dir", default=None)
    parser.add_argument("--context-length", type=int, default=4096)
    args = parser.parse_args()

    models_dir = args.models_dir or str(Path(args.config_dir) / MODELS_DIR)
    model_name = _read_model_name(args.config_dir)

    if not args.skip_export:
        export_models(args.config_dir)

    print("=== Generating configs ===")
    update_genai_config(
        output_dir=models_dir,
        model_name=model_name,
        device=args.device,
        context_length=args.context_length,
    )
    fix_tokenizer(output_dir=models_dir)
    print()
    print("Done.")


if __name__ == "__main__":
    main()