File size: 6,036 Bytes
f2b116f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 | """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)
|