| |
| """Llama 3.2 1B Instruct -> ExecuTorch WebGPU (4-bit groupwise weights, KV cache in program). |
| |
| Same export path as the Qwen recipes (_llm.py): export_llm with the Vulkan partitioner (which |
| WebGPUPartitioner wraps unchanged), qmode=4w group 64, 4-bit embedding, fp32 activations and fp32 |
| KV buffers, prefill up to 512 tokens, 8960-token context. Program contract (same as the previous |
| artifact): forward(tokens[1, 1..512] int64, input_pos[1] int64) -> logits[1, 128256] float32 for |
| the last position. |
| |
| Weights: the Meta-format checkpoint `original/consolidated.00.pth` + `original/params.json` of |
| meta-llama/Llama-3.2-1B-Instruct at a pinned revision (Llama 3.2 Community License). The upstream |
| repo is gated. `--source auto` (default) downloads from it when the token has access, otherwise |
| from a public mirror. Either way every file is checked against the hashes Hugging Face publishes |
| for the *official* repo at the pinned revision (LFS sha256 for large files, git blob sha1 for small |
| ones; both are visible without access), so the bytes used are the official bytes. |
| """ |
|
|
| import argparse |
| import hashlib |
| import json |
| import os |
| import shlex |
| import shutil |
| import subprocess |
| import sys |
| from pathlib import Path |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parent)) |
| import _common |
| import _llm |
|
|
| OFFICIAL = { |
| "repo": "meta-llama/Llama-3.2-1B-Instruct", |
| "revision": "9213176726f574b556790deb65791e0c5aa438b6", |
| } |
| |
| MIRROR = { |
| "repo": "LIF1014/ptdbench-Llama-3.2-1B-Instruct", |
| "revision": "828e5f37c833a02415416260d5a93780fd41d10e", |
| } |
| |
| |
| FILES = { |
| "original/consolidated.00.pth": { |
| "lfs": "fc17d497df5e4175b3a8acb4f5865b26f7fc1b009b25bef814b95fde10e8a1f3" |
| }, |
| "original/params.json": {"blob": "9cd8dbdf2dc6f4d8abb60bdb5ce64f4bec2fdfd9"}, |
| "tokenizer.json": {"blob": "5cc5f00a5b203e90a27a3bd60d1ec393b07971e8"}, |
| "tokenizer_config.json": {"blob": "4ff488a165e900e5129cda7c20ab32d568d2a475"}, |
| "special_tokens_map.json": {"blob": "02ee80b6196926a5ad790a004d9efd6ab1ba6542"}, |
| "config.json": {"blob": "3e3aaf51a035cb5092d9f6827a0dc074657ba88c"}, |
| "generation_config.json": {"blob": "75ae08310d6d23df373ee2644b497192b3cce6d8"}, |
| "LICENSE.txt": {"blob": "085b47c1575cb889b7024030e60b78f54f0b8c9e"}, |
| "USE_POLICY.md": {"blob": "ac3c5f21b9779e3da0677d6d3c587778fe3a331e"}, |
| } |
| |
| TOKENIZER_FILES = [ |
| "tokenizer.json", |
| "tokenizer_config.json", |
| "special_tokens_map.json", |
| "config.json", |
| "generation_config.json", |
| "LICENSE.txt", |
| "USE_POLICY.md", |
| ] |
| NAME = "llama3_2_1b_instruct_q4gsw" |
| MODEL_CLASS = "llama3_2" |
|
|
|
|
| def git_blob_sha1(path): |
| data = Path(path).read_bytes() |
| return hashlib.sha1(b"blob %d\0" % len(data) + data).hexdigest() |
|
|
|
|
| def check(path, want): |
| if "lfs" in want: |
| got = _common.sha256_file(path) |
| return got == want["lfs"], {"sha256": got} |
| got = git_blob_sha1(path) |
| return got == want["blob"], {"git_blob_sha1": got} |
|
|
|
|
| def official_metadata(): |
| """Cross-check the pinned hashes against what the Hub reports now (no access needed).""" |
| from huggingface_hub import HfApi |
|
|
| info = HfApi().model_info( |
| OFFICIAL["repo"], revision=OFFICIAL["revision"], files_metadata=True |
| ) |
| live = {s.rfilename: s for s in info.siblings} |
| for f, want in FILES.items(): |
| s = live[f] |
| if "lfs" in want and (s.lfs is None or s.lfs.sha256 != want["lfs"]): |
| sys.exit(f"official {f}: Hub reports {s.lfs} != pinned {want}") |
| if "blob" in want and s.blob_id != want["blob"]: |
| sys.exit(f"official {f}: Hub reports blob {s.blob_id} != pinned {want}") |
| return info.sha |
|
|
|
|
| def fetch(source): |
| from huggingface_hub import hf_hub_download |
| from huggingface_hub.errors import GatedRepoError |
|
|
| def get(repo): |
| return { |
| f: Path(hf_hub_download(repo["repo"], f, revision=repo["revision"])) |
| for f in FILES |
| } |
|
|
| used = None |
| if source in ("auto", "official"): |
| try: |
| paths, used = get(OFFICIAL), OFFICIAL |
| except GatedRepoError as e: |
| if source == "official": |
| raise |
| print( |
| f"official repo not accessible ({str(e).splitlines()[-1][:120]}); using mirror", |
| flush=True, |
| ) |
| if used is None: |
| paths, used = get(MIRROR), MIRROR |
| verified = {} |
| for f, p in paths.items(): |
| ok, got = check(p, FILES[f]) |
| if not ok: |
| sys.exit( |
| f"{used['repo']}@{used['revision']}:{f} {got} does not match official {FILES[f]}" |
| ) |
| verified[f] = got |
| return paths, used, verified |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser(description=__doc__.splitlines()[0]) |
| _llm.add_args(ap, max_context=8960) |
| ap.add_argument( |
| "--embedding-group-size", |
| type=int, |
| default=None, |
| help="group size of the 4-bit embedding (default: --group-size; the previous artifact used 32)", |
| ) |
| ap.add_argument("--source", choices=["auto", "official", "mirror"], default="auto") |
| args = ap.parse_args() |
|
|
| out = Path(args.out).resolve() |
| work = out / "_work" |
| work.mkdir(parents=True, exist_ok=True) |
| official_sha = official_metadata() |
| paths, used, verified = fetch(args.source) |
| params = json.loads(paths["original/params.json"].read_text()) |
|
|
| sys.path.insert(0, str(_llm.SHIMS)) |
| env = { |
| **os.environ, |
| "PYTHONPATH": os.pathsep.join( |
| filter(None, [str(_llm.SHIMS), os.environ.get("PYTHONPATH")]) |
| ), |
| } |
| ckpt = work / "checkpoint.pth" |
| if ckpt.is_symlink() or ckpt.exists(): |
| ckpt.unlink() |
| ckpt.symlink_to(paths["original/consolidated.00.pth"]) |
| params_path = work / "params.json" |
| shutil.copyfile(paths["original/params.json"], params_path) |
|
|
| emb_group = args.embedding_group_size or args.group_size |
| 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={args.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,{emb_group}'", |
| "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: |
| shutil.copyfile(paths[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": OFFICIAL["repo"], |
| "revision": OFFICIAL["revision"], |
| "official": True, |
| "license": "Llama 3.2 Community License (tokenizer/LICENSE.txt, tokenizer/USE_POLICY.md); Built with Llama", |
| "files": {f: {**FILES[f], **verified[f]} for f in FILES}, |
| "hub_revision_checked": official_sha, |
| "downloaded_from": { |
| **used, |
| "note": "all files verified byte-identical to the official hashes above", |
| }, |
| }, |
| "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": "original/consolidated.00.pth (Meta format, bf16; no conversion), params = original/params.json", |
| "params": params, |
| "rope": "llama3 scaled RoPE, rope_scale_factor=32 (set by ExecuTorch for model_class llama3_2 when use_scaled_rope)", |
| "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": args.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 {emb_group}); output head quantized too", |
| "kv_cache": "fp32 buffers in the program; fp16 KV is a runtime option of the WebGPU runtime (enable_f16_kv_cache / 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)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
