From cca0b2f3fe2c6722fa0e02233c50bc0b9d1d6988 Mon Sep 17 00:00:00 2001 From: Alex-Wengg Date: Sun, 2 Aug 2026 13:42:29 -0400 Subject: [PATCH 1/2] trial(translate/riva-4b-v2): stateful CoreML probe of Riva-Translate-4B-Instruct-v2 Feasibility probe answering whether a 4B Mistral-class translation LLM is viable on-device via CoreML. Conversion recipe follows the qwen3-asr stateful-KV-cache pattern, traced fully in fp16 so it converts on a 24GB host (~70s), with embed_tokens host-side and lm_head as a separate model. Results (M-series 24GB): fp16/GPU decodes at 11.5 tok/s with exact 14/14 greedy parity vs fp16 torch (logits corr 0.999999). int4 per-block-32 cuts the decoder 7.0GB -> 2.0GB, loads in 3.8s, prefills 3x faster, and still produces a correct translation (different but valid greedy path). Decode is dispatch-overhead-bound, not bandwidth-bound (int4 did not speed it up). CPU_AND_NE is strictly worse (209ms/tok, corr 0.785) - same conclusion as the Qwen3-0.6B LLM-on-ANE trial: autoregressive decode of this shape gains nothing from the ANE. Verdict: converts cleanly, runs fine on Mac, marginal on iOS; MLX remains the better host for 4B-class decode. mlpackages/npy artifacts gitignored. --- .../riva-translate-4b-v2/coreml/.gitignore | 2 + .../riva-translate-4b-v2/coreml/README.md | 77 +++++ .../coreml/convert_stateful_decoder.py | 324 ++++++++++++++++++ .../coreml/quantize_int4.py | 56 +++ .../riva-translate-4b-v2/coreml/run_coreml.py | 145 ++++++++ .../coreml/run_reference.py | 93 +++++ 6 files changed, 697 insertions(+) create mode 100644 models/translate/riva-translate-4b-v2/coreml/.gitignore create mode 100644 models/translate/riva-translate-4b-v2/coreml/README.md create mode 100644 models/translate/riva-translate-4b-v2/coreml/convert_stateful_decoder.py create mode 100644 models/translate/riva-translate-4b-v2/coreml/quantize_int4.py create mode 100644 models/translate/riva-translate-4b-v2/coreml/run_coreml.py create mode 100644 models/translate/riva-translate-4b-v2/coreml/run_reference.py diff --git a/models/translate/riva-translate-4b-v2/coreml/.gitignore b/models/translate/riva-translate-4b-v2/coreml/.gitignore new file mode 100644 index 0000000..fbd2db9 --- /dev/null +++ b/models/translate/riva-translate-4b-v2/coreml/.gitignore @@ -0,0 +1,2 @@ +out/ +reference.npz diff --git a/models/translate/riva-translate-4b-v2/coreml/README.md b/models/translate/riva-translate-4b-v2/coreml/README.md new file mode 100644 index 0000000..c9a73d6 --- /dev/null +++ b/models/translate/riva-translate-4b-v2/coreml/README.md @@ -0,0 +1,77 @@ +# Riva-Translate-4B-Instruct-v2 → CoreML feasibility probe + +Probe conversion of [nvidia/Riva-Translate-4B-Instruct-v2](https://huggingface.co/nvidia/Riva-Translate-4B-Instruct-v2) +(text-to-text NMT, 37 languages) to CoreML, to answer: does a 4B Mistral-class +translation LLM run acceptably on-device via CoreML? + +## Architecture + +Standard `MistralForCausalLM` (pruned/distilled from Mistral-NeMo-12B): +34 layers, hidden 3072, 32 Q / 8 KV heads (GQA×4), head_dim 128, +intermediate 8640, vocab 131072 (Tekken), tied embeddings, 8k context, +~4.2B params. + +## Conversion recipe + +Same stateful-KV-cache pattern as `models/stt/qwen3-asr-0.6b/coreml/convert_stateful_decoder.py`: + +- 34-layer decoder stack as ONE stateful model, 68 fp16 KV state buffers + (`max_seq_len` 1024), `RangeDim` query length → single model serves both + prefill and decode +- traced end-to-end in **fp16** (no fp32 master copy — converts on a 24GB host, + peak well under RAM; conversion takes ~70s) +- `embed_tokens` kept host-side (`embed_tokens_fp16.npy`, 768MB, mmap lookup) +- lm_head = final RMSNorm + tied 3072×131072 projection, separate model +- RoPE cos/sin computed host-side and passed as inputs; additive causal mask input + +Scripts: `convert_stateful_decoder.py` (convert), `quantize_int4.py` +(post-training linear int4, per-block 32), `run_reference.py` / +`run_coreml.py` (parity + latency harness). + +## Results (M-series, 24GB, macOS 26) + +| variant | decoder size | load | prefill (36 tok) | decode mean | greedy parity vs fp16 torch | +|---|---|---|---|---|---| +| fp16, CPU_AND_GPU | 7.0GB | 24.3s | 1.61s (22 tok/s) | 86.9ms/tok (11.5 tok/s, p50 63ms) | 14/14 tokens, logits corr 0.999999 | +| int4 pb32, CPU_AND_GPU | 2.0GB | 3.8s | 0.55s (65 tok/s) | 81.9ms/tok (12.2 tok/s, p50 70ms) | 0/14 (corr 0.959) — but output is a correct alternative translation | +| int4 pb32, CPU_AND_NE | 2.0GB | 56.7s | 7.06s (5.1 tok/s) | 209ms/tok (4.8 tok/s) | 0/14 (corr 0.785) — same correct text as int4 GPU | + +- fp16 en→de output: "In San Francisco ist es für diese Jahreszeit ungewöhnlich warm." +- int4 en→de output: "Das Wetter in San Francisco ist für diese Jahreszeit ungewöhnlich warm." + (greedy path flips to an equally valid phrasing; no degradation visible on this sample — + proper eval would need COMET on FLORES) + +## Findings + +1. **Conversion itself is a non-issue.** fp16 tracing + stateful states convert + cleanly in ~70s with coremltools 9.0; smoke + parity pass. The + vocab-131072 lm_head converts fine as a plain matmul (no topk in-graph, so + the known CoreML topk mod-131072 bug is not triggered). +2. **Decode is overhead-bound in this harness, not bandwidth-bound.** + int4 shrank weights 3.5× and sped prefill 3×, but decode stayed ~80ms/tok. + ~63ms p50 for a 2GB weight read ≈ 32GB/s effective — far below the chip's + bandwidth. Per-step cost is dominated by predict dispatch (2 model calls + per token from Python + RangeDim shape handling). A Swift host with + preallocated buffers would likely land in the 20–40ms/tok range, but that + is still 25–50 tok/s at best — MLX gets similar or better with far less + machinery. +3. **Quality survives int4 per-block-32** on the sampled prompt; keep lm_head + fp16 if logit fidelity matters (int4 head dropped first-step corr to 0.959). +4. **Memory**: int4 stack = 2.0GB decoder + 216MB head + 768MB embedding + (embedding could be int8/int4'd too → ~200–400MB). Total ≈ 2.5GB + KV cache + (68 × 1×8×1024×128 fp16 = 272MB at seq 1024). Fits Mac/iPad easily; + iPhone Pro-class only with the extended-memory entitlement and little else + resident — pairing it with an ASR stack on-device would be tight. +5. **ANE is strictly worse** (int4, CPU_AND_NE): 56.7s load (ANE compile), + prefill 7.1s (4× slower than GPU), decode 209ms/tok (2.5× slower), and + first-step logits corr drops to 0.785 — consistent with partial placement + plus CPU-fallback sync overhead. Same outcome class as the Qwen3-0.6B + LLM-on-ANE trial: autoregressive decode of this shape does not benefit + from the ANE. Output text was still the correct translation. + +## Verdict + +Converting works; running is "fine on Mac, marginal on iOS". For FluidAudio +purposes a 4B translation LLM remains better served by MLX on macOS; the +CoreML route offers no advantage for GPU-bound autoregressive decode and the +ANE does not change the picture for this shape of model. diff --git a/models/translate/riva-translate-4b-v2/coreml/convert_stateful_decoder.py b/models/translate/riva-translate-4b-v2/coreml/convert_stateful_decoder.py new file mode 100644 index 0000000..9a77514 --- /dev/null +++ b/models/translate/riva-translate-4b-v2/coreml/convert_stateful_decoder.py @@ -0,0 +1,324 @@ +"""Convert Riva-Translate-4B-Instruct-v2 decoder to a stateful CoreML model. + +Probe conversion for on-device feasibility. Architecture is standard Mistral +(MistralForCausalLM, 34 layers, hidden 3072, 32Q/8KV heads, head_dim 128, +GQA x4, vocab 131072, tied embeddings). Pattern follows +models/stt/qwen3-asr-0.6b/coreml/convert_stateful_decoder.py, with two +memory-driven differences for a 4.2B model on a 24GB host: + + - the model is loaded and traced in fp16 end to end (no fp32 master copy) + - embed_tokens and lm_head stay OUT of the decoder graph; embedding lookup + happens host-side and lm_head converts as a separate model + +Usage: + uv run convert_stateful_decoder.py --output-dir ./out + uv run convert_stateful_decoder.py --max-seq-len 1024 --skip-lm-head +""" + +# /// script +# requires-python = ">=3.10,<3.13" +# dependencies = [ +# "torch>=2.4", +# "transformers>=4.48", +# "coremltools>=8.0", +# "numpy<2", +# "safetensors", +# "huggingface_hub", +# ] +# /// + +import argparse +import gc +import math +import time +from pathlib import Path + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F + +# Riva-Translate-4B-Instruct-v2 architecture constants (from config.json) +NUM_LAYERS = 34 +NUM_Q_HEADS = 32 +NUM_KV_HEADS = 8 +HEAD_DIM = 128 +HIDDEN_SIZE = 3072 +INTERMEDIATE_SIZE = 8640 +VOCAB_SIZE = 131_072 +GQA_REPEAT = NUM_Q_HEADS // NUM_KV_HEADS # 4 + + +def rotate_half(x: torch.Tensor) -> torch.Tensor: + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """Expand KV heads for grouped query attention: [B, H_kv, S, D] -> [B, H_q, S, D].""" + if n_rep == 1: + return hidden_states + batch, num_kv_heads, slen, head_dim = hidden_states.shape + hidden_states = hidden_states[:, :, None, :, :].expand( + batch, num_kv_heads, n_rep, slen, head_dim + ) + return hidden_states.reshape(batch, num_kv_heads * n_rep, slen, head_dim) + + +class StatefulMistralDecoder(nn.Module): + """Mistral decoder stack with stateful KV cache for CoreML export. + + Runs entirely in fp16. Final RMSNorm is NOT applied here — it lives in + the lm_head model. + """ + + def __init__(self, layers: nn.ModuleList, max_seq_len: int): + super().__init__() + self.layers = layers + self.max_seq_len = max_seq_len + self.scale = 1.0 / math.sqrt(HEAD_DIM) + + for i in range(NUM_LAYERS): + self.register_buffer( + f"k_cache_{i}", + torch.zeros(1, NUM_KV_HEADS, max_seq_len, HEAD_DIM, dtype=torch.float16), + ) + self.register_buffer( + f"v_cache_{i}", + torch.zeros(1, NUM_KV_HEADS, max_seq_len, HEAD_DIM, dtype=torch.float16), + ) + + def forward( + self, + hidden_states: torch.Tensor, + position_cos: torch.Tensor, + position_sin: torch.Tensor, + attention_mask: torch.Tensor, + ) -> torch.Tensor: + """Args: + hidden_states: [1, Q, 3072] fp16 — input embeddings + position_cos: [1, Q, 128] fp16 — RoPE cosines for query positions + position_sin: [1, Q, 128] fp16 — RoPE sines + attention_mask: [1, 1, Q, end_step] fp16 (0=attend, -inf-ish=ignore) + Returns: + [1, Q, 3072] fp16 — decoder output (pre final norm) + """ + q_len = hidden_states.shape[1] + end_step = attention_mask.shape[-1] + past_kv_len = end_step - q_len + + cos = position_cos.unsqueeze(1) + sin = position_sin.unsqueeze(1) + + for i in range(NUM_LAYERS): + layer = self.layers[i] + k_cache = getattr(self, f"k_cache_{i}") + v_cache = getattr(self, f"v_cache_{i}") + + residual = hidden_states + hidden_states = layer.input_layernorm(hidden_states) + + attn = layer.self_attn + q = attn.q_proj(hidden_states) # [1, Q, 32*128=4096] + k = attn.k_proj(hidden_states) # [1, Q, 8*128=1024] + v = attn.v_proj(hidden_states) # [1, Q, 8*128=1024] + + q = q.view(1, q_len, NUM_Q_HEADS, HEAD_DIM).transpose(1, 2) + k = k.view(1, q_len, NUM_KV_HEADS, HEAD_DIM).transpose(1, 2) + v = v.view(1, q_len, NUM_KV_HEADS, HEAD_DIM).transpose(1, 2) + + q = (q * cos) + (rotate_half(q) * sin) + k = (k * cos) + (rotate_half(k) * sin) + + k_cache[:, :, past_kv_len:end_step, :] = k + v_cache[:, :, past_kv_len:end_step, :] = v + + k_full = k_cache[:, :, :end_step, :] + v_full = v_cache[:, :, :end_step, :] + + k_full = repeat_kv(k_full, GQA_REPEAT) # [1, 32, end_step, 128] + v_full = repeat_kv(v_full, GQA_REPEAT) + + attn_weights = torch.matmul(q, k_full.transpose(2, 3)) * self.scale + attn_weights = attn_weights + attention_mask + attn_weights = F.softmax(attn_weights, dim=-1) + attn_output = torch.matmul(attn_weights, v_full) # [1, 32, Q, 128] + + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.view(1, q_len, NUM_Q_HEADS * HEAD_DIM) + hidden_states = attn.o_proj(attn_output) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = layer.post_attention_layernorm(hidden_states) + + mlp = layer.mlp + gate = mlp.gate_proj(hidden_states) + up = mlp.up_proj(hidden_states) + hidden_states = mlp.down_proj(F.silu(gate) * up) + hidden_states = residual + hidden_states + + return hidden_states + + +class LmHead(nn.Module): + """Final RMSNorm + tied-embedding lm_head projection.""" + + def __init__(self, norm: nn.Module, weight: torch.Tensor): + super().__init__() + self.norm = norm + self.proj = nn.Linear(HIDDEN_SIZE, VOCAB_SIZE, bias=False) + with torch.no_grad(): + self.proj.weight.copy_(weight) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.proj(self.norm(hidden_states)) + + +def main(): + parser = argparse.ArgumentParser(description="Convert Riva-Translate-4B decoder to stateful CoreML") + parser.add_argument("--model-id", default="nvidia/Riva-Translate-4B-Instruct-v2") + parser.add_argument("--max-seq-len", type=int, default=1024, help="Max sequence length for KV cache") + parser.add_argument("--output-dir", default=".") + parser.add_argument("--skip-lm-head", action="store_true") + parser.add_argument("--skip-decoder", action="store_true") + args = parser.parse_args() + + MAX_SEQ_LEN = args.max_seq_len + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + import coremltools as ct + + print(f"torch {torch.__version__}, coremltools {ct.__version__}") + + # ---- Step 1: Load model in fp16 (no fp32 master copy — 24GB host) ---- + print(f"Loading {args.model_id} in fp16...") + t0 = time.time() + from transformers import AutoModelForCausalLM + + model = AutoModelForCausalLM.from_pretrained( + args.model_id, torch_dtype=torch.float16, low_cpu_mem_usage=True + ) + model.eval() + print(f"Loaded in {time.time() - t0:.1f}s") + + layers = model.model.layers + assert len(layers) == NUM_LAYERS + attn0 = layers[0].self_attn + assert attn0.q_proj.out_features == NUM_Q_HEADS * HEAD_DIM + assert attn0.k_proj.out_features == NUM_KV_HEADS * HEAD_DIM + assert not hasattr(attn0, "q_norm"), "unexpected QK norms for Mistral" + + # ---- Step 2: lm_head model (norm + tied embedding projection) ---- + if not args.skip_lm_head: + print("\nConverting lm_head (norm + 3072x131072 projection)...") + lm_head = LmHead(model.model.norm, model.model.embed_tokens.weight.detach()) + lm_head.eval().half() + ex = torch.randn(1, 1, HIDDEN_SIZE, dtype=torch.float16) + with torch.no_grad(): + traced_head = torch.jit.trace(lm_head, ex) + head_ml = ct.convert( + traced_head, + inputs=[ct.TensorType("hidden_states", shape=(1, 1, HIDDEN_SIZE), dtype=np.float16)], + outputs=[ct.TensorType("logits", dtype=np.float16)], + minimum_deployment_target=ct.target.macOS15, + compute_precision=ct.precision.FLOAT16, + compute_units=ct.ComputeUnit.CPU_AND_GPU, + ) + head_path = output_dir / "riva4b_lm_head.mlpackage" + head_ml.save(str(head_path)) + print(f"Saved {head_path}") + del lm_head, traced_head, head_ml + gc.collect() + + if args.skip_decoder: + return + + # ---- Step 3: Stateful decoder ---- + # Free everything except the layer stack before tracing. + embed_path = output_dir / "embed_tokens_fp16.npy" + if not embed_path.exists(): + np.save(embed_path, model.model.embed_tokens.weight.detach().numpy()) + print(f"Saved host-side embedding table to {embed_path}") + model.model.embed_tokens = None + model.lm_head = None + gc.collect() + + print(f"\nCreating stateful decoder (max_seq_len={MAX_SEQ_LEN})...") + stateful_model = StatefulMistralDecoder(layers, max_seq_len=MAX_SEQ_LEN) + stateful_model.eval() + + trace_q, trace_end = 1, 5 + hidden = torch.randn(1, trace_q, HIDDEN_SIZE, dtype=torch.float16) + cos_in = torch.randn(1, trace_q, HEAD_DIM, dtype=torch.float16) + sin_in = torch.randn(1, trace_q, HEAD_DIM, dtype=torch.float16) + mask = torch.zeros(1, 1, trace_q, trace_end, dtype=torch.float16) + + print("Tracing (fp16, CPU)...") + t0 = time.time() + with torch.no_grad(): + traced = torch.jit.trace(stateful_model, (hidden, cos_in, sin_in, mask)) + traced.eval() + print(f"Trace complete in {time.time() - t0:.1f}s") + + query_length = ct.RangeDim(lower_bound=1, upper_bound=MAX_SEQ_LEN, default=1) + end_step_dim = ct.RangeDim(lower_bound=1, upper_bound=MAX_SEQ_LEN, default=1) + + inputs = [ + ct.TensorType("hidden_states", shape=(1, query_length, HIDDEN_SIZE), dtype=np.float16), + ct.TensorType("position_cos", shape=(1, query_length, HEAD_DIM), dtype=np.float16), + ct.TensorType("position_sin", shape=(1, query_length, HEAD_DIM), dtype=np.float16), + ct.TensorType("attention_mask", shape=(1, 1, query_length, end_step_dim), dtype=np.float16), + ] + outputs = [ct.TensorType("output_hidden", dtype=np.float16)] + + states = [] + for i in range(NUM_LAYERS): + for kv in ("k", "v"): + states.append( + ct.StateType( + wrapped_type=ct.TensorType( + shape=(1, NUM_KV_HEADS, MAX_SEQ_LEN, HEAD_DIM), dtype=np.float16 + ), + name=f"{kv}_cache_{i}", + ) + ) + + print("Converting decoder to CoreML (this is the slow, memory-heavy step)...") + t0 = time.time() + mlmodel = ct.convert( + traced, + inputs=inputs, + outputs=outputs, + states=states, + minimum_deployment_target=ct.target.macOS15, + compute_precision=ct.precision.FLOAT16, + compute_units=ct.ComputeUnit.CPU_AND_GPU, + ) + print(f"CoreML conversion complete in {time.time() - t0:.1f}s") + + out_path = output_dir / "riva4b_decoder_stateful.mlpackage" + mlmodel.save(str(out_path)) + print(f"Saved {out_path}") + + # ---- Step 4: smoke test ---- + print("\nSmoke test (decode step Q=1)...") + state = mlmodel.make_state() + out = mlmodel.predict( + { + "hidden_states": np.random.randn(1, 1, HIDDEN_SIZE).astype(np.float16), + "position_cos": np.random.randn(1, 1, HEAD_DIM).astype(np.float16), + "position_sin": np.random.randn(1, 1, HEAD_DIM).astype(np.float16), + "attention_mask": np.zeros((1, 1, 1, 1), dtype=np.float16), + }, + state=state, + ) + arr = out["output_hidden"] + print(f" output shape {arr.shape}, range [{np.min(arr):.3f}, {np.max(arr):.3f}]") + print("Done.") + + +if __name__ == "__main__": + main() diff --git a/models/translate/riva-translate-4b-v2/coreml/quantize_int4.py b/models/translate/riva-translate-4b-v2/coreml/quantize_int4.py new file mode 100644 index 0000000..c2564d2 --- /dev/null +++ b/models/translate/riva-translate-4b-v2/coreml/quantize_int4.py @@ -0,0 +1,56 @@ +"""Post-training int4 (linear, per-block) quantization of the Riva-4B decoder. + +Data-free linear_quantize_weights is used instead of k-means palettization +because k-means over 4B params takes hours on CPU; per-block linear int4 is +minutes and typically sufficient to judge feasibility. + +Usage: + uv run quantize_int4.py --model-dir ./out +""" + +# /// script +# requires-python = ">=3.10,<3.13" +# dependencies = [ +# "coremltools>=8.0", +# "numpy<2", +# ] +# /// + +import argparse +import time +from pathlib import Path + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--model-dir", default="./out") + parser.add_argument("--block-size", type=int, default=32) + args = parser.parse_args() + + import coremltools as ct + import coremltools.optimize as cto + + model_dir = Path(args.model_dir) + + config = cto.coreml.OptimizationConfig( + global_config=cto.coreml.OpLinearQuantizerConfig( + mode="linear_symmetric", + dtype="int4", + granularity="per_block", + block_size=args.block_size, + ) + ) + + for name in ("riva4b_decoder_stateful", "riva4b_lm_head"): + src = model_dir / f"{name}.mlpackage" + dst = model_dir / f"{name}_int4.mlpackage" + print(f"Quantizing {src.name} (block_size={args.block_size})...") + t0 = time.time() + model = ct.models.MLModel(str(src), compute_units=ct.ComputeUnit.CPU_ONLY, skip_model_load=True) + quantized = cto.coreml.linear_quantize_weights(model, config) + quantized.save(str(dst)) + print(f" saved {dst.name} in {time.time() - t0:.0f}s") + + +if __name__ == "__main__": + main() diff --git a/models/translate/riva-translate-4b-v2/coreml/run_coreml.py b/models/translate/riva-translate-4b-v2/coreml/run_coreml.py new file mode 100644 index 0000000..a6df0ff --- /dev/null +++ b/models/translate/riva-translate-4b-v2/coreml/run_coreml.py @@ -0,0 +1,145 @@ +"""Run the converted Riva-4B CoreML pipeline and compare against reference.npz. + +Pipeline: host-side embedding lookup (numpy) -> stateful decoder (prefill with +Q=prompt_len, then Q=1 decode steps) -> lm_head -> greedy argmax. + +Also reports prefill latency and per-token decode latency. + +Usage: + uv run run_coreml.py --model-dir ./out --reference reference.npz +""" + +# /// script +# requires-python = ">=3.10,<3.13" +# dependencies = [ +# "coremltools>=8.0", +# "numpy<2", +# ] +# /// + +import argparse +import time +from pathlib import Path + +import numpy as np + +HIDDEN_SIZE = 3072 +HEAD_DIM = 128 +ROPE_THETA = 1_000_000.0 + + +def rope_tables(positions: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """cos/sin tables in concatenated-halves layout, [1, len(positions), 128].""" + inv_freq = 1.0 / (ROPE_THETA ** (np.arange(0, HEAD_DIM, 2, dtype=np.float64) / HEAD_DIM)) + freqs = np.outer(positions.astype(np.float64), inv_freq) # [Q, 64] + emb = np.concatenate([freqs, freqs], axis=-1) # [Q, 128] + return ( + np.cos(emb)[None].astype(np.float16), + np.sin(emb)[None].astype(np.float16), + ) + + +def causal_mask(q_len: int, end_step: int) -> np.ndarray: + past = end_step - q_len + mask = np.zeros((1, 1, q_len, end_step), dtype=np.float16) + for r in range(q_len): + mask[0, 0, r, past + r + 1 :] = -30_000.0 + return mask + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--model-dir", default=".") + parser.add_argument("--reference", default="reference.npz") + parser.add_argument("--max-new-tokens", type=int, default=32) + parser.add_argument("--compute-units", default="CPU_AND_GPU", + choices=["ALL", "CPU_AND_GPU", "CPU_ONLY", "CPU_AND_NE"]) + parser.add_argument("--suffix", default="", help="Model filename suffix, e.g. _int4") + args = parser.parse_args() + + import coremltools as ct + + model_dir = Path(args.model_dir) + ref = np.load(args.reference) + prompt_ids = ref["prompt_ids"] + ref_gen = ref["gen_ids"] + eos_id = int(ref["eos_token_id"]) + + cu = getattr(ct.ComputeUnit, args.compute_units) + print(f"Loading models (compute_units={args.compute_units})...") + t0 = time.time() + decoder = ct.models.MLModel( + str(model_dir / f"riva4b_decoder_stateful{args.suffix}.mlpackage"), compute_units=cu + ) + lm_head = ct.models.MLModel(str(model_dir / f"riva4b_lm_head{args.suffix}.mlpackage"), compute_units=cu) + embed = np.load(model_dir / "embed_tokens_fp16.npy", mmap_mode="r") + print(f"Loaded in {time.time() - t0:.1f}s") + + state = decoder.make_state() + + def run_decoder(token_ids: np.ndarray, past_len: int) -> np.ndarray: + q = len(token_ids) + positions = np.arange(past_len, past_len + q) + cos, sin = rope_tables(positions) + hidden = embed[token_ids][None].astype(np.float16) # [1, Q, 3072] + out = decoder.predict( + { + "hidden_states": hidden, + "position_cos": cos, + "position_sin": sin, + "attention_mask": causal_mask(q, past_len + q), + }, + state=state, + ) + return out["output_hidden"] # [1, Q, 3072] + + def logits_for(hidden_last: np.ndarray) -> np.ndarray: + out = lm_head.predict({"hidden_states": hidden_last.reshape(1, 1, HIDDEN_SIZE)}) + return out["logits"].reshape(-1) + + # ---- Prefill ---- + t0 = time.time() + hidden = run_decoder(prompt_ids, past_len=0) + prefill_s = time.time() - t0 + print(f"Prefill: {len(prompt_ids)} tokens in {prefill_s:.2f}s ({len(prompt_ids)/prefill_s:.1f} tok/s)") + + logits = logits_for(hidden[0, -1]) + + # First-step logits parity + ref_logits = ref["first_logits"] + top1_ref = int(np.argmax(ref_logits)) + top1_cml = int(np.argmax(logits)) + corr = np.corrcoef(ref_logits.astype(np.float64), logits.astype(np.float64))[0, 1] + print(f"First-step logits: corr={corr:.6f}, argmax ref={top1_ref} coreml={top1_cml} " + f"{'MATCH' if top1_ref == top1_cml else 'MISMATCH'}") + + # ---- Greedy decode ---- + gen = [] + past = len(prompt_ids) + tok = top1_cml + decode_times = [] + while len(gen) < args.max_new_tokens: + gen.append(tok) + if tok == eos_id: + break + t0 = time.time() + hidden = run_decoder(np.array([tok]), past_len=past) + logits = logits_for(hidden[0, -1]) + decode_times.append(time.time() - t0) + tok = int(np.argmax(logits)) + past += 1 + + dt = np.array(decode_times) + print(f"Decode: {len(dt)} steps, mean {dt.mean()*1000:.1f}ms/tok ({1.0/dt.mean():.1f} tok/s), " + f"p50 {np.percentile(dt,50)*1000:.1f}ms") + + gen = np.array(gen) + n = min(len(gen), len(ref_gen)) + match = int((gen[:n] == ref_gen[:n]).sum()) + print(f"\nToken parity vs reference: {match}/{n} match") + print(f" ref: {ref_gen[:n].tolist()}") + print(f" coreml: {gen[:n].tolist()}") + + +if __name__ == "__main__": + main() diff --git a/models/translate/riva-translate-4b-v2/coreml/run_reference.py b/models/translate/riva-translate-4b-v2/coreml/run_reference.py new file mode 100644 index 0000000..25e1b82 --- /dev/null +++ b/models/translate/riva-translate-4b-v2/coreml/run_reference.py @@ -0,0 +1,93 @@ +"""Greedy reference generation with HF transformers for parity checking. + +Dumps prompt token ids, generated token ids, and first-step logits to +reference.npz for run_coreml.py to compare against. + +Usage: + uv run run_reference.py --max-new-tokens 32 +""" + +# /// script +# requires-python = ">=3.10,<3.13" +# dependencies = [ +# "torch>=2.4", +# "transformers>=4.48", +# "numpy<2", +# "safetensors", +# "huggingface_hub", +# ] +# /// + +import argparse +from pathlib import Path + +import numpy as np +import torch + +SOURCE_TEXT = "The weather in San Francisco is unusually warm for this time of year." + + +def build_prompt(tokenizer, text: str, source_lang: str, target_lang: str) -> str: + messages = [ + { + "role": "user", + "content": f"Translate the following text from {source_lang} to {target_lang}: {text}", + } + ] + return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--model-id", default="nvidia/Riva-Translate-4B-Instruct-v2") + parser.add_argument("--max-new-tokens", type=int, default=32) + parser.add_argument("--source-lang", default="English") + parser.add_argument("--target-lang", default="German") + parser.add_argument("--output", default="reference.npz") + args = parser.parse_args() + + from transformers import AutoModelForCausalLM, AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained(args.model_id) + prompt = build_prompt(tokenizer, SOURCE_TEXT, args.source_lang, args.target_lang) + print(f"Prompt:\n{prompt!r}\n") + + input_ids = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).input_ids + print(f"Prompt tokens: {input_ids.shape[1]}") + + model = AutoModelForCausalLM.from_pretrained( + args.model_id, torch_dtype=torch.float16, low_cpu_mem_usage=True + ) + model.eval() + + with torch.no_grad(): + # First-step logits for numeric comparison + first = model(input_ids) + first_logits = first.logits[0, -1, :].float().numpy() + + out = model.generate( + input_ids, + max_new_tokens=args.max_new_tokens, + do_sample=False, + temperature=None, + top_p=None, + top_k=None, + ) + + gen_ids = out[0, input_ids.shape[1] :].numpy() + text = tokenizer.decode(gen_ids, skip_special_tokens=True) + print(f"Reference translation: {text!r}") + print(f"Generated ids: {gen_ids.tolist()}") + + np.savez( + Path(args.output), + prompt_ids=input_ids[0].numpy(), + gen_ids=gen_ids, + first_logits=first_logits, + eos_token_id=np.int64(tokenizer.eos_token_id), + ) + print(f"Saved {args.output}") + + +if __name__ == "__main__": + main() From 6d4348d6b51b3f43e05e8c459d1c2bd333ba5e01 Mon Sep 17 00:00:00 2001 From: Alex-Wengg Date: Sun, 2 Aug 2026 14:42:36 -0400 Subject: [PATCH 2/2] trial(translate/riva-4b-v2): profile + optimize CoreML decode; MLX baseline Optimization pass on the Riva-4B CoreML probe. Decoder-only profiling across quant formats shows CoreML's quantized GEMV kernels hit a ~40ms/tok floor (int4 moves 4x less data than fp16 but decodes only 1.5x faster; palettized LUT is 2.3x SLOWER than fp16). RangeDim shape churn measured free; host-side gaps between predicts measured near-free (busy-wait probe); alternating between two MLModel instances costs ~24ms/step, fixed by a --fused variant (decoder + final norm + tied lm_head in one stateful graph, last-position logits) which lifts prefill to 228 tok/s. Sustained-load bimodality observed (~41ms vs ~68ms regimes, GPU clock management). MLX 4-bit baseline on the same machine: 106.7 tok/s decode, 3.5GB peak, same correct translation - 4.4x CoreML's best steady-state (24 tok/s). The gap lives in the kernels, not the host or model structure, so the verdict stands quantitatively: 4B decode belongs on MLX; CoreML keeps prefill/ encoder workloads. Speculative decoding and chunked pipelines intentionally not pursued (draft-model machinery / adds the handoff cost it would need to avoid). --- .../riva-translate-4b-v2/coreml/.gitignore | 1 + .../riva-translate-4b-v2/coreml/README.md | 58 ++++++++- .../coreml/convert_stateful_decoder.py | 42 ++++++- .../coreml/profile_decode.py | 118 ++++++++++++++++++ .../coreml/quantize_int4.py | 3 +- .../coreml/quantize_variants.py | 78 ++++++++++++ .../riva-translate-4b-v2/coreml/run_coreml.py | 32 +++-- 7 files changed, 312 insertions(+), 20 deletions(-) create mode 100644 models/translate/riva-translate-4b-v2/coreml/profile_decode.py create mode 100644 models/translate/riva-translate-4b-v2/coreml/quantize_variants.py diff --git a/models/translate/riva-translate-4b-v2/coreml/.gitignore b/models/translate/riva-translate-4b-v2/coreml/.gitignore index fbd2db9..e7242db 100644 --- a/models/translate/riva-translate-4b-v2/coreml/.gitignore +++ b/models/translate/riva-translate-4b-v2/coreml/.gitignore @@ -1,2 +1,3 @@ out/ reference.npz +mlx-4bit/ diff --git a/models/translate/riva-translate-4b-v2/coreml/README.md b/models/translate/riva-translate-4b-v2/coreml/README.md index c9a73d6..bc0d75b 100644 --- a/models/translate/riva-translate-4b-v2/coreml/README.md +++ b/models/translate/riva-translate-4b-v2/coreml/README.md @@ -69,9 +69,59 @@ Scripts: `convert_stateful_decoder.py` (convert), `quantize_int4.py` LLM-on-ANE trial: autoregressive decode of this shape does not benefit from the ANE. Output text was still the correct translation. +## Optimization pass (profiling, quant variants, fused model, MLX baseline) + +Follow-up pass to find where decode time goes and whether CoreML can approach +bandwidth-optimal decode. Scripts: `profile_decode.py`, `quantize_variants.py`, +`convert_stateful_decoder.py --fused`. + +Decoder-only latency (Q=1, short context, steady loop, CPU_AND_GPU): + +| weights | size | ms/tok | effective BW | +|---|---|---|---| +| fp16 | 7.0GB | 58.8 | ~142GB/s | +| int8 per-channel | 3.5GB | 47.5 | 74GB/s | +| int4 per-block-32 | 2.0GB | 41.3 | 49GB/s | +| int4 per-channel | 1.8GB | 39.1 | 46GB/s | +| 4-bit palettized LUT | 1.8GB | 95.5 | (avoid — slowest) | + +Findings, in causal order: + +1. **RangeDim shape churn is a non-issue** — growing `end_step` per decode + step costs nothing vs constant shape. No need for static-shape decode + models or position-input scatter designs. +2. **CoreML quantized GEMV kernels hit a ~40ms floor** on this GPU: int4 + moves 4× less data than fp16 but only decodes 1.5× faster. Palettized LUT + dequant is pathological (2.3× slower than fp16). fp16 runs at ~52% of + chip bandwidth; int4 at ~18%. +3. **Alternating between two MLModels costs ~24ms/step** (decoder 41.1ms + + head 1.5ms alone, but 66.8ms interleaved). Fixed by the `--fused` variant + (34 layers + final norm + tied lm_head in one graph, last-position logits + output): prefill 228 tok/s, one predict per token. +4. **Host-side gaps between predicts are nearly free** (busy-wait probe: + +0.5→10ms gaps add only 0–3ms net), so a Swift host loop wouldn't beat + the Python harness by much — the floor is in the kernels, not the host. +5. **The machine is bimodal under sustained GPU load**: identical benchmarks + oscillate between ~41ms and ~68ms regimes (GPU clock management / + thermals). Steady-state best ≈ 24 tok/s; observed sustained ≈ 15 tok/s. +6. **MLX 4-bit baseline on the same machine: 106.7 tok/s decode** + (`out/mlx-4bit`, 4.5 bits/weight, 3.5GB peak memory, same correct + translation output). MLX's quantized GEMV runs near bandwidth-optimal — + ~4.4× faster than CoreML's best steady-state. CoreML wins only prefill + (228 vs ~39 tok/s on this 36-token prompt, where MLX is setup-dominated). + +Not pursued (and why): speculative decoding needs a draft model and host +machinery disproportionate to a probe; chunked multi-model pipelines add +handoffs (the thing that costs 24ms) without reducing sequential work; ANE +already shown strictly worse. + ## Verdict -Converting works; running is "fine on Mac, marginal on iOS". For FluidAudio -purposes a 4B translation LLM remains better served by MLX on macOS; the -CoreML route offers no advantage for GPU-bound autoregressive decode and the -ANE does not change the picture for this shape of model. +Converting works; running is "fine on Mac, marginal on iOS". The optimization +pass makes the conclusion quantitative: CoreML's quantized GEMV kernels cap +decode at ~24 tok/s while MLX hits 106.7 tok/s on the same silicon — a 4.4× +gap that no amount of model restructuring closes, because it lives inside the +kernels. For FluidAudio purposes a 4B translation LLM is an MLX workload on +macOS. CoreML remains the right tool for the encoder-heavy sub-1B models the +framework is built around, and its strong prefill suggests the hybrid worth +remembering: CoreML/ANE encoders + MLX decoder. diff --git a/models/translate/riva-translate-4b-v2/coreml/convert_stateful_decoder.py b/models/translate/riva-translate-4b-v2/coreml/convert_stateful_decoder.py index 9a77514..453b070 100644 --- a/models/translate/riva-translate-4b-v2/coreml/convert_stateful_decoder.py +++ b/models/translate/riva-translate-4b-v2/coreml/convert_stateful_decoder.py @@ -162,6 +162,27 @@ def forward( return hidden_states +class StatefulMistralDecoderFused(StatefulMistralDecoder): + """Decoder stack + final RMSNorm + tied lm_head in one graph. + + Emits logits for the LAST query position only ([1, 1, vocab]) so a single + model serves prefill (Q=N) and decode (Q=1) with one predict per token and + no cross-model handoff (measured ~24ms/step penalty for two-model decode). + """ + + def __init__(self, layers: nn.ModuleList, norm: nn.Module, head_weight: torch.Tensor, max_seq_len: int): + super().__init__(layers, max_seq_len) + self.norm = norm + self.head = nn.Linear(HIDDEN_SIZE, VOCAB_SIZE, bias=False) + with torch.no_grad(): + self.head.weight.copy_(head_weight) + + def forward(self, hidden_states, position_cos, position_sin, attention_mask): + hidden = super().forward(hidden_states, position_cos, position_sin, attention_mask) + last = self.norm(hidden[:, -1:, :]) + return self.head(last) + + class LmHead(nn.Module): """Final RMSNorm + tied-embedding lm_head projection.""" @@ -183,6 +204,8 @@ def main(): parser.add_argument("--output-dir", default=".") parser.add_argument("--skip-lm-head", action="store_true") parser.add_argument("--skip-decoder", action="store_true") + parser.add_argument("--fused", action="store_true", + help="Fuse final norm + lm_head into the decoder (single model per token)") args = parser.parse_args() MAX_SEQ_LEN = args.max_seq_len @@ -242,13 +265,19 @@ def main(): if not embed_path.exists(): np.save(embed_path, model.model.embed_tokens.weight.detach().numpy()) print(f"Saved host-side embedding table to {embed_path}") + head_weight = model.model.embed_tokens.weight.detach() if args.fused else None model.model.embed_tokens = None model.lm_head = None gc.collect() - print(f"\nCreating stateful decoder (max_seq_len={MAX_SEQ_LEN})...") - stateful_model = StatefulMistralDecoder(layers, max_seq_len=MAX_SEQ_LEN) - stateful_model.eval() + print(f"\nCreating stateful decoder (max_seq_len={MAX_SEQ_LEN}, fused={args.fused})...") + if args.fused: + stateful_model = StatefulMistralDecoderFused( + layers, model.model.norm, head_weight, max_seq_len=MAX_SEQ_LEN + ) + else: + stateful_model = StatefulMistralDecoder(layers, max_seq_len=MAX_SEQ_LEN) + stateful_model.eval().half() trace_q, trace_end = 1, 5 hidden = torch.randn(1, trace_q, HIDDEN_SIZE, dtype=torch.float16) @@ -272,7 +301,8 @@ def main(): ct.TensorType("position_sin", shape=(1, query_length, HEAD_DIM), dtype=np.float16), ct.TensorType("attention_mask", shape=(1, 1, query_length, end_step_dim), dtype=np.float16), ] - outputs = [ct.TensorType("output_hidden", dtype=np.float16)] + out_name = "logits" if args.fused else "output_hidden" + outputs = [ct.TensorType(out_name, dtype=np.float16)] states = [] for i in range(NUM_LAYERS): @@ -299,7 +329,7 @@ def main(): ) print(f"CoreML conversion complete in {time.time() - t0:.1f}s") - out_path = output_dir / "riva4b_decoder_stateful.mlpackage" + out_path = output_dir / ("riva4b_decoder_fused.mlpackage" if args.fused else "riva4b_decoder_stateful.mlpackage") mlmodel.save(str(out_path)) print(f"Saved {out_path}") @@ -315,7 +345,7 @@ def main(): }, state=state, ) - arr = out["output_hidden"] + arr = out[out_name] print(f" output shape {arr.shape}, range [{np.min(arr):.3f}, {np.max(arr):.3f}]") print("Done.") diff --git a/models/translate/riva-translate-4b-v2/coreml/profile_decode.py b/models/translate/riva-translate-4b-v2/coreml/profile_decode.py new file mode 100644 index 0000000..16e2302 --- /dev/null +++ b/models/translate/riva-translate-4b-v2/coreml/profile_decode.py @@ -0,0 +1,118 @@ +"""Profile where decode time goes in the Riva-4B CoreML pipeline. + +Hypotheses tested: + A. growing mask shape per step (real decode) -> per-step shape re-specialization + B. constant mask shape per step -> steady-state kernel reuse + C. lm_head call cost + D. embedding lookup cost (host-side numpy) + +Usage: + uv run profile_decode.py --model-dir ./out --suffix _int4 +""" + +# /// script +# requires-python = ">=3.10,<3.13" +# dependencies = [ +# "coremltools>=8.0", +# "numpy<2", +# ] +# /// + +import argparse +import time +from pathlib import Path + +import numpy as np + +HIDDEN_SIZE = 3072 +HEAD_DIM = 128 +STEPS = 24 + + +def bench(fn, n=STEPS, warmup=2): + for _ in range(warmup): + fn(0) + times = [] + for i in range(n): + t0 = time.perf_counter() + fn(i) + times.append(time.perf_counter() - t0) + t = np.array(times) * 1000 + return f"mean {t.mean():6.1f}ms p50 {np.percentile(t,50):6.1f}ms min {t.min():6.1f}ms max {t.max():6.1f}ms" + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--model-dir", default="./out") + parser.add_argument("--suffix", default="_int4") + parser.add_argument("--compute-units", default="CPU_AND_GPU") + args = parser.parse_args() + + import coremltools as ct + + model_dir = Path(args.model_dir) + cu = getattr(ct.ComputeUnit, args.compute_units) + print(f"Loading (suffix={args.suffix!r}, cu={args.compute_units})...") + t0 = time.time() + decoder = ct.models.MLModel( + str(model_dir / f"riva4b_decoder_stateful{args.suffix}.mlpackage"), compute_units=cu + ) + head_path = model_dir / f"riva4b_lm_head{args.suffix}.mlpackage" + lm_head = ct.models.MLModel(str(head_path), compute_units=cu) if head_path.exists() else None + embed = np.load(model_dir / "embed_tokens_fp16.npy", mmap_mode="r") + print(f"Loaded in {time.time() - t0:.1f}s\n") + + hidden1 = np.random.randn(1, 1, HIDDEN_SIZE).astype(np.float16) + cos1 = np.random.randn(1, 1, HEAD_DIM).astype(np.float16) + sin1 = np.random.randn(1, 1, HEAD_DIM).astype(np.float16) + + state = decoder.make_state() + + def dec_step(end_step): + mask = np.zeros((1, 1, 1, end_step), dtype=np.float16) + return decoder.predict( + {"hidden_states": hidden1, "position_cos": cos1, "position_sin": sin1, + "attention_mask": mask}, + state=state, + ) + + # A: growing shape (real decode pattern), starting at end_step=40 + print("A. decoder, GROWING end_step (40..): ", bench(lambda i: dec_step(40 + i))) + + # B: constant shapes + for es in (64, 512, 1024): + print(f"B. decoder, CONSTANT end_step={es:4d}: ", bench(lambda i, es=es: dec_step(es))) + + # C: lm_head alone + if lm_head is not None: + def head_step(i): + return lm_head.predict({"hidden_states": hidden1}) + print("C. lm_head alone: ", bench(head_step)) + + # D: embedding lookup alone + def emb_step(i): + return embed[[i % 1000]][None].astype(np.float16) + print("D. embed lookup (numpy mmap): ", bench(emb_step)) + + # E: interleaved decoder + head (real decode pattern with two models) + if lm_head is not None: + def interleaved(i): + out = dec_step(64) + h = out["output_hidden"].astype(np.float16).reshape(1, 1, HIDDEN_SIZE) + return lm_head.predict({"hidden_states": h}) + print("E. interleaved decoder+head: ", bench(interleaved)) + + # F: decoder loop with deliberate host-side gap between predicts. + # Quantifies the GPU idle/power-state penalty as a function of gap length. + for gap_ms in (0.5, 2, 5, 10): + def gapped(i, g=gap_ms): + t_end = time.perf_counter() + g / 1000.0 + while time.perf_counter() < t_end: # busy-wait, no sleep syscall + pass + return dec_step(64) + t = bench(gapped) + print(f"F. decoder w/ {gap_ms:4.1f}ms host gap: ", t) + + +if __name__ == "__main__": + main() diff --git a/models/translate/riva-translate-4b-v2/coreml/quantize_int4.py b/models/translate/riva-translate-4b-v2/coreml/quantize_int4.py index c2564d2..dcc76d8 100644 --- a/models/translate/riva-translate-4b-v2/coreml/quantize_int4.py +++ b/models/translate/riva-translate-4b-v2/coreml/quantize_int4.py @@ -25,6 +25,7 @@ def main(): parser = argparse.ArgumentParser() parser.add_argument("--model-dir", default="./out") parser.add_argument("--block-size", type=int, default=32) + parser.add_argument("--models", default="riva4b_decoder_stateful,riva4b_lm_head") args = parser.parse_args() import coremltools as ct @@ -41,7 +42,7 @@ def main(): ) ) - for name in ("riva4b_decoder_stateful", "riva4b_lm_head"): + for name in args.models.split(","): src = model_dir / f"{name}.mlpackage" dst = model_dir / f"{name}_int4.mlpackage" print(f"Quantizing {src.name} (block_size={args.block_size})...") diff --git a/models/translate/riva-translate-4b-v2/coreml/quantize_variants.py b/models/translate/riva-translate-4b-v2/coreml/quantize_variants.py new file mode 100644 index 0000000..f457f79 --- /dev/null +++ b/models/translate/riva-translate-4b-v2/coreml/quantize_variants.py @@ -0,0 +1,78 @@ +"""Produce alternative compressed decoder variants to find the fastest GPU dequant path. + +Variants: + int8 — linear symmetric, per-channel (classic fast path) + pal4 — 4-bit palettization LUT, per-grouped-channel(16), uniform mode + int4c — linear symmetric int4, per-channel (no per-block scales) + +Usage: + uv run quantize_variants.py --model-dir ./out --variants int8,pal4 +""" + +# /// script +# requires-python = ">=3.10,<3.13" +# dependencies = [ +# "coremltools>=8.0", +# "numpy<2", +# ] +# /// + +import argparse +import time +from pathlib import Path + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--model-dir", default="./out") + parser.add_argument("--variants", default="int8,pal4,int4c") + args = parser.parse_args() + + import coremltools as ct + import coremltools.optimize as cto + + model_dir = Path(args.model_dir) + src = model_dir / "riva4b_decoder_stateful.mlpackage" + + def cfg_int8(): + return cto.coreml.OptimizationConfig( + global_config=cto.coreml.OpLinearQuantizerConfig( + mode="linear_symmetric", dtype="int8", granularity="per_channel" + ) + ) + + def cfg_pal4(): + return cto.coreml.OptimizationConfig( + global_config=cto.coreml.OpPalettizerConfig( + mode="uniform", nbits=4, granularity="per_grouped_channel", group_size=16 + ) + ) + + def cfg_int4c(): + return cto.coreml.OptimizationConfig( + global_config=cto.coreml.OpLinearQuantizerConfig( + mode="linear_symmetric", dtype="int4", granularity="per_channel" + ) + ) + + makers = {"int8": cfg_int8, "pal4": cfg_pal4, "int4c": cfg_int4c} + + for name in args.variants.split(","): + name = name.strip() + dst = model_dir / f"riva4b_decoder_stateful_{name}.mlpackage" + if dst.exists(): + print(f"skip {name} (exists)") + continue + print(f"Quantizing variant {name}...") + t0 = time.time() + model = ct.models.MLModel(str(src), compute_units=ct.ComputeUnit.CPU_ONLY, skip_model_load=True) + if name == "pal4": + quantized = cto.coreml.palettize_weights(model, makers[name]()) + else: + quantized = cto.coreml.linear_quantize_weights(model, makers[name]()) + quantized.save(str(dst)) + print(f" saved {dst.name} in {time.time() - t0:.0f}s") + + +if __name__ == "__main__": + main() diff --git a/models/translate/riva-translate-4b-v2/coreml/run_coreml.py b/models/translate/riva-translate-4b-v2/coreml/run_coreml.py index a6df0ff..9892e78 100644 --- a/models/translate/riva-translate-4b-v2/coreml/run_coreml.py +++ b/models/translate/riva-translate-4b-v2/coreml/run_coreml.py @@ -54,7 +54,10 @@ def main(): parser.add_argument("--max-new-tokens", type=int, default=32) parser.add_argument("--compute-units", default="CPU_AND_GPU", choices=["ALL", "CPU_AND_GPU", "CPU_ONLY", "CPU_AND_NE"]) - parser.add_argument("--suffix", default="", help="Model filename suffix, e.g. _int4") + parser.add_argument("--suffix", default="", help="Decoder filename suffix, e.g. _int4") + parser.add_argument("--head-suffix", default=None, help="lm_head suffix (defaults to --suffix)") + parser.add_argument("--fused", action="store_true", + help="Use riva4b_decoder_fused (in-graph norm+lm_head, logits output)") args = parser.parse_args() import coremltools as ct @@ -68,10 +71,17 @@ def main(): cu = getattr(ct.ComputeUnit, args.compute_units) print(f"Loading models (compute_units={args.compute_units})...") t0 = time.time() + decoder_base = "riva4b_decoder_fused" if args.fused else "riva4b_decoder_stateful" decoder = ct.models.MLModel( - str(model_dir / f"riva4b_decoder_stateful{args.suffix}.mlpackage"), compute_units=cu + str(model_dir / f"{decoder_base}{args.suffix}.mlpackage"), compute_units=cu ) - lm_head = ct.models.MLModel(str(model_dir / f"riva4b_lm_head{args.suffix}.mlpackage"), compute_units=cu) + if args.fused: + lm_head = None + else: + head_suffix = args.suffix if args.head_suffix is None else args.head_suffix + lm_head = ct.models.MLModel( + str(model_dir / f"riva4b_lm_head{head_suffix}.mlpackage"), compute_units=cu + ) embed = np.load(model_dir / "embed_tokens_fp16.npy", mmap_mode="r") print(f"Loaded in {time.time() - t0:.1f}s") @@ -91,19 +101,23 @@ def run_decoder(token_ids: np.ndarray, past_len: int) -> np.ndarray: }, state=state, ) + if args.fused: + return out["logits"].reshape(-1) # last-position logits return out["output_hidden"] # [1, Q, 3072] - def logits_for(hidden_last: np.ndarray) -> np.ndarray: - out = lm_head.predict({"hidden_states": hidden_last.reshape(1, 1, HIDDEN_SIZE)}) + def logits_for(dec_out) -> np.ndarray: + if args.fused: + return dec_out + out = lm_head.predict({"hidden_states": dec_out[0, -1].reshape(1, 1, HIDDEN_SIZE)}) return out["logits"].reshape(-1) # ---- Prefill ---- t0 = time.time() - hidden = run_decoder(prompt_ids, past_len=0) + dec_out = run_decoder(prompt_ids, past_len=0) prefill_s = time.time() - t0 print(f"Prefill: {len(prompt_ids)} tokens in {prefill_s:.2f}s ({len(prompt_ids)/prefill_s:.1f} tok/s)") - logits = logits_for(hidden[0, -1]) + logits = logits_for(dec_out) # First-step logits parity ref_logits = ref["first_logits"] @@ -123,8 +137,8 @@ def logits_for(hidden_last: np.ndarray) -> np.ndarray: if tok == eos_id: break t0 = time.time() - hidden = run_decoder(np.array([tok]), past_len=past) - logits = logits_for(hidden[0, -1]) + dec_out = run_decoder(np.array([tok]), past_len=past) + logits = logits_for(dec_out) decode_times.append(time.time() - t0) tok = int(np.argmax(logits)) past += 1