| """Shared export path for the Qwen LLM recipes: pinned HF snapshot -> ExecuTorch checkpoint |
| -> export_llm (Vulkan partitioner, which WebGPUPartitioner wraps unchanged) -> .pte. |
| |
| Graph contract (matches the internal WebGPU LLM artifacts): forward(tokens[1, <=prefill_max] |
| int64, input_pos[1] int64) -> logits[1, vocab] float32 for the last position, KV cache inside the |
| program (fp32 buffers, max_context_length positions), 4-bit groupwise weight-only linears |
| (et_vk.linear_q4gsw) and 4-bit groupwise embedding (et_vk.embedding_q4gsw). |
| """ |
|
|
| import importlib |
| import json |
| import os |
| import shlex |
| import shutil |
| import subprocess |
| import sys |
| from pathlib import Path |
|
|
| import _common |
|
|
| SHIMS = Path(__file__).resolve().parent / "_shims" |
|
|
| TOKENIZER_FILES = [ |
| "tokenizer.json", |
| "tokenizer_config.json", |
| "vocab.json", |
| "merges.txt", |
| "generation_config.json", |
| "config.json", |
| "LICENSE", |
| ] |
|
|
|
|
| def export( |
| *, name, upstream_repo, revision, model_class, params, converter, max_context, args |
| ): |
| from huggingface_hub import snapshot_download |
|
|
| out = Path(args.out).resolve() |
| work = out / "_work" |
| work.mkdir(parents=True, exist_ok=True) |
| snap = Path( |
| snapshot_download( |
| upstream_repo, |
| revision=revision, |
| allow_patterns=["*.safetensors", "*.json", "merges.txt", "LICENSE"], |
| ) |
| ) |
|
|
| sys.path.insert(0, str(SHIMS)) |
| env = {**os.environ, "PYTHONPATH": os.pathsep.join(filter(None, [str(SHIMS), os.environ.get("PYTHONPATH")]))} |
| ckpt = work / "checkpoint.pth" |
| if not ckpt.exists(): |
| importlib.import_module(converter).convert_weights(str(snap), str(ckpt)) |
| params_path = work / "params.json" |
| params_path.write_text(json.dumps(params, indent=2) + "\n") |
|
|
| output_name = f"{name}.pte" |
| overrides = [ |
| f"base.model_class={model_class}", |
| f"base.checkpoint={ckpt}", |
| f"base.params={params_path}", |
| "model.use_kv_cache=true", |
| "model.use_sdpa_with_kv_cache=true", |
| "model.enable_dynamic_shape=true", |
| "model.dtype_override=fp32", |
| f"export.max_seq_length={args.prefill_max + 1}", |
| f"export.max_context_length={max_context}", |
| f"export.output_dir={out}", |
| f"export.output_name={out / output_name}", |
| "quantization.qmode=4w", |
| f"quantization.group_size={args.group_size}", |
| f"quantization.embedding_quantize='4,{args.group_size}'", |
| "backend.vulkan.enabled=true", |
| ] |
| cmd = [ |
| sys.executable, |
| "-m", |
| "executorch.extension.llm.export.export_llm", |
| *overrides, |
| ] |
| print("+", " ".join(shlex.quote(c) for c in cmd), flush=True) |
| subprocess.run(cmd, check=True, env=env) |
|
|
| pte = out / output_name |
| tok_dir = out / "tokenizer" |
| tok_dir.mkdir(exist_ok=True) |
| tokenizer_files = {} |
| for f in TOKENIZER_FILES: |
| if (snap / f).exists(): |
| shutil.copyfile(snap / f, tok_dir / f) |
| tokenizer_files[f] = _common.sha256_file(tok_dir / f) |
| shutil.rmtree(work) |
|
|
| data = pte.read_bytes() |
| recipe = { |
| "model": name, |
| "upstream": { |
| "model": upstream_repo, |
| "revision": revision, |
| "official": True, |
| "safetensors_sha256": { |
| p.name: _common.sha256_file(p) |
| for p in sorted(snap.glob("*.safetensors")) |
| }, |
| }, |
| "versions": _common.versions(), |
| "export": { |
| "api": "python -m executorch.extension.llm.export.export_llm (hydra overrides below)", |
| "overrides": [_common.scrub(o, out).replace(str(work.resolve()), "<work>") for o in overrides], |
| "checkpoint_conversion": f"{converter}.convert_weights(<snapshot>, checkpoint.pth)", |
| "import_shim": "recipes/_shims/torchtune (re-exports executorch.examples.models.checkpoint.get_mapped_key; torchtune itself is not installed)", |
| "params": params, |
| "partitioner": "VulkanPartitioner via backend.vulkan.enabled (WebGPUPartitioner is a pass-through wrapper of it; delegate id VulkanBackend)", |
| "inputs": [ |
| { |
| "name": "tokens", |
| "shape": [1, f"1..{args.prefill_max}"], |
| "dtype": "int64", |
| }, |
| {"name": "input_pos", "shape": [1], "dtype": "int64"}, |
| ], |
| "outputs": [ |
| { |
| "name": "logits", |
| "shape": [1, params["vocab_size"]], |
| "dtype": "float32", |
| "note": "last position only", |
| } |
| ], |
| "prefill_max": args.prefill_max, |
| "max_seq_len_metadata": args.prefill_max + 1, |
| "max_context_length": max_context, |
| "dtype": "float32 activations", |
| "quantization": f"4-bit groupwise weight-only linears (torchao IntxWeightOnlyConfig int4, group {args.group_size}) and 4-bit groupwise embedding (group {args.group_size}); output head quantized too", |
| "kv_cache": "fp32 buffers in the program; fp16 KV is a runtime option of the WebGPU runtime (et_set_kv_f16), not an export flag", |
| }, |
| "files": { |
| pte.name: {"bytes": len(data), "sha256": _common.sha256_file(pte)}, |
| **{f"tokenizer/{k}": {"sha256": v} for k, v in tokenizer_files.items()}, |
| }, |
| "graph": _common.delegate_summary(data), |
| "command": _common.command_line(out), |
| } |
| (out / "recipe.json").write_text(json.dumps(recipe, indent=2) + "\n") |
| print(json.dumps({"files": recipe["files"], "graph": recipe["graph"]}, indent=2)) |
|
|
|
|
| def add_args(ap, max_context): |
| ap.add_argument("--out", required=True, help="output directory") |
| ap.add_argument( |
| "--prefill-max", type=int, default=512, help="max tokens per forward call" |
| ) |
| ap.add_argument( |
| "--max-context", type=int, default=max_context, help="KV cache length" |
| ) |
| ap.add_argument("--group-size", type=int, default=64) |
|
|
