convert_embeds_mixed.py

"""Mixed-precision CoreML convert for Qwen3-ASR LLM (input_embeds variant).

Attention layers (q/k/v/o_proj) → LUT-8 (8-bit palettize).
MLP layers (gate/up/down_proj) → LUT-4 (4-bit palettize).
Everything else (norms, lm_head, embed) → kept as fp16.

Compute precision = fp32 to avoid Qwen3-ASR RMSNorm/attention NaN.
"""
from __future__ import annotations
import argparse
import os
import sys
import re
from pathlib import Path

sys.path.insert(0, "/tmp/Anemll")
import numpy as np
import torch
import torch.nn as nn
import coremltools as ct
import coremltools.optimize as cto


def patch_qwen_for_inputs_embeds():
    from anemll.models import qwen_model as qm

    orig_model_forward = qm.QwenModel.forward

    def model_forward_or_embeds(
        self, input_ids, causal_mask, position_ids, current_pos, IN_PREFILL: bool = False,
    ):
        if input_ids.dtype in (torch.float16, torch.float32, torch.bfloat16):
            hidden_states = input_ids
            if IN_PREFILL:
                rotary_emb = self.get_rotary_embedding_prefill(position_ids)
            else:
                rotary_emb = self.get_rotary_embeddings_s(current_pos)
            hidden_states = self.process_layers(
                hidden_states, position_ids, causal_mask,
                current_pos, rotary_emb, start_layer=0, end_layer=None,
                IN_PREFILL=IN_PREFILL,
            )
            hidden_states = self.norm(hidden_states)
            return hidden_states
        return orig_model_forward(self, input_ids, causal_mask, position_ids,
                                  current_pos, IN_PREFILL=IN_PREFILL)

    qm.QwenModel.forward = model_forward_or_embeds

    orig_causal_forward = qm.QwenForCausalLM.forward

    def causal_forward_or_embeds(
        self, input_ids, update_mask, position_ids, causal_mask, current_pos,
        IN_PREFILL: bool = False,
    ):
        if input_ids.dtype in (torch.float16, torch.float32, torch.bfloat16):
            hidden_states = self.model(
                input_ids, causal_mask, position_ids, current_pos,
                IN_PREFILL=IN_PREFILL,
            )
            if not IN_PREFILL and current_pos is not None:
                seq_len = hidden_states.shape[1]
                if seq_len == 1:
                    pos_tensor = torch.tensor([0], device=hidden_states.device, dtype=torch.long)
                else:
                    if isinstance(current_pos, torch.Tensor):
                        pos_tensor = current_pos if current_pos.dim() > 0 else current_pos.unsqueeze(0)
                    else:
                        pos_tensor = torch.tensor([current_pos], device=hidden_states.device, dtype=torch.long)
                hidden_states = torch.index_select(hidden_states, dim=1, index=pos_tensor)
            hidden_states = hidden_states.permute(0, 2, 1).unsqueeze(2).to(qm.MODEL_DTYPE)
            return tuple(
                getattr(self, f"lm_head16_{k}")(hidden_states).squeeze(2).transpose(1, 2)
                for k in range(1, 17)
            )
        return orig_causal_forward(
            self, input_ids, update_mask, position_ids, causal_mask, current_pos,
            IN_PREFILL=IN_PREFILL,
        )

    qm.QwenForCausalLM.forward = causal_forward_or_embeds
    print("[patch] QwenModel + QwenForCausalLM accept inputs_embeds")


def select_attn_layer(op):
    """Return True if op is in a self_attn projection (q/k/v/o_proj)."""
    n = op.name.lower()
    return ("self_attn" in n and any(p in n for p in ("q_proj", "k_proj", "v_proj", "o_proj")))


def select_mlp_layer(op):
    """Return True if op is in an MLP projection (gate/up/down_proj)."""
    n = op.name.lower()
    return "mlp" in n and any(p in n for p in ("gate_proj", "up_proj", "down_proj"))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True, type=Path)
    ap.add_argument("--output", required=True, type=Path)
    ap.add_argument("--attn-bits", type=int, default=8)
    ap.add_argument("--mlp-bits", type=int, default=4)
    ap.add_argument("--group-size", type=int, default=8)
    ap.add_argument("--context-length", type=int, default=512)
    args = ap.parse_args()

    patch_qwen_for_inputs_embeds()

    from anemll.models.qwen_model import (
        QwenForCausalLM, QwenConfig, MODEL_DTYPE, TEST_DEVICE,
    )
    from anemll.ane_converter import qwen_converter as qc
    import anemll.models.qwen_model as qm
    qm.ENABLE_COREML = True

    import json
    cfg = json.load(open(args.model / "config.json"))
    cfg["context_length"] = args.context_length
    cfg["state_length"] = args.context_length
    config = QwenConfig(**cfg)
    model = QwenForCausalLM(config, enable_coreml=True)
    model.load_pretrained_weights(str(args.model))
    model.eval()
    for p in model.parameters():
        p.requires_grad = False
    print(f"Model loaded: hidden={config.hidden_size}, layers={config.num_hidden_layers}")

    class WrapperEmbeds(torch.nn.Module):
        def __init__(self, model):
            super().__init__()
            self.model = model

        def forward(self, inputs_embeds, position_ids, causal_mask, current_pos, update_mask):
            return self.model(
                input_ids=inputs_embeds, update_mask=update_mask,
                position_ids=position_ids, causal_mask=causal_mask,
                current_pos=current_pos, IN_PREFILL=False,
            )

    wrapper = WrapperEmbeds(model).eval()

    sample_inputs_embeds = torch.zeros((1, 1, config.hidden_size), dtype=torch.float16, device=TEST_DEVICE)
    sample_position_ids = torch.zeros((1,), dtype=torch.int32, device=TEST_DEVICE)
    sample_causal_mask = torch.zeros((1, 1, 1, args.context_length), dtype=torch.float16, device=TEST_DEVICE)
    sample_current_pos = torch.zeros((1,), dtype=torch.int32, device=TEST_DEVICE)
    sample_update_mask = torch.zeros((1, 1, args.context_length, 1), dtype=torch.float16, device=TEST_DEVICE)

    print("Tracing ...")
    traced = torch.jit.trace(
        wrapper,
        (sample_inputs_embeds, sample_position_ids, sample_causal_mask,
         sample_current_pos, sample_update_mask),
    )
    print("Converting (fp32 compute, no palettize yet) ...")

    states = qc.QwenConverter.GetTransformerStates(model, prefix="model.model.")

    mlmodel = ct.convert(
        traced,
        inputs=[
            ct.TensorType(name="inputs_embeds", shape=sample_inputs_embeds.shape, dtype=np.float16),
            ct.TensorType(name="position_ids",  shape=sample_position_ids.shape,  dtype=np.int32),
            ct.TensorType(name="causal_mask",   shape=sample_causal_mask.shape,   dtype=np.float16),
            ct.TensorType(name="current_pos",   shape=sample_current_pos.shape,   dtype=np.int32),
            ct.TensorType(name="update_mask",   shape=sample_update_mask.shape,   dtype=np.float16),
        ],
        outputs=[ct.TensorType(name=f"logits{i+1}", dtype=np.float16) for i in range(16)],
        states=states,
        minimum_deployment_target=ct.target.iOS18,
        compute_precision=ct.precision.FLOAT32,
        compute_units=ct.ComputeUnit.CPU_AND_NE,
        convert_to="mlprogram",
        skip_model_load=True,
    )

    # Walk the MIL program to enumerate const-weight ops; classify by name.
    prog = mlmodel._mil_program
    fn = prog.functions["main"]
    attn_op_names, mlp_op_names = [], []
    for op in fn.operations:
        if op.op_type != "const":
            continue
        n = op.name.lower()
        # Skip tiny constants (norms, biases, indices); only target large weight matrices.
        try:
            arr = op.val.val
            if hasattr(arr, "shape") and arr.ndim >= 2 and arr.size >= 64 * 64:
                pass
            else:
                continue
        except Exception:
            continue
        if ("self_attn" in n or "self.attn" in n) and any(p in n for p in ("q_proj", "k_proj", "v_proj", "o_proj")):
            attn_op_names.append(op.name)
        elif ("mlp" in n) and any(p in n for p in ("gate_proj", "up_proj", "down_proj")):
            mlp_op_names.append(op.name)
    print(f"Found {len(attn_op_names)} attention weight ops and {len(mlp_op_names)} MLP weight ops")
    if not attn_op_names or not mlp_op_names:
        print("WARN: matched zero ops — falling back to global LUT-4")
        cfg = cto.coreml.OpPalettizerConfig(
            nbits=args.mlp_bits, mode="kmeans",
            granularity="per_grouped_channel", group_size=args.group_size,
        )
        mlmodel = cto.coreml.palettize_weights(
            mlmodel, cto.coreml.OptimizationConfig(global_config=cfg),
        )
    else:
        cfg_attn = cto.coreml.OpPalettizerConfig(
            nbits=args.attn_bits, mode="kmeans",
            granularity="per_grouped_channel", group_size=args.group_size,
        )
        cfg_mlp = cto.coreml.OpPalettizerConfig(
            nbits=args.mlp_bits, mode="kmeans",
            granularity="per_grouped_channel", group_size=args.group_size,
        )
        op_name_configs = {**{n: cfg_attn for n in attn_op_names},
                           **{n: cfg_mlp for n in mlp_op_names}}
        pal_cfg = cto.coreml.OptimizationConfig(op_name_configs=op_name_configs)
        print(f"Mixed palettize: {len(attn_op_names)} ops @ LUT-{args.attn_bits}, {len(mlp_op_names)} ops @ LUT-{args.mlp_bits}, rest fp16")
        mlmodel = cto.coreml.palettize_weights(mlmodel, pal_cfg)

    args.output.parent.mkdir(parents=True, exist_ok=True)
    mlmodel.save(str(args.output))
    print(f"Saved: {args.output}")


if __name__ == "__main__":
    main()