"""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()