diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..ca729f3 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,79 @@ +name: CI + +# End-to-end reproducibility on free runners: lint, then prove the whole +# prepare -> train -> export -> C-verify -> perplexity pipeline works from a +# clean checkout, and that the ESP32 sketch still compiles. No board needed. +on: + push: + branches: [main] + pull_request: + +jobs: + lint: + name: Lint (Python) + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: astral-sh/setup-uv@v6 + - name: Syntax check all Python + run: uv run python -m compileall -q src data experiments + # Informational only: the repo has no ruff config yet, so findings are + # reported without blocking the pipeline. + - name: Ruff check (non-blocking) + continue-on-error: true + run: uvx ruff check src data experiments + + host-pipeline: + name: Reproduce train -> export -> verify on host + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: astral-sh/setup-uv@v6 + - name: Install deps + run: uv sync --frozen + - name: Prepare a small TinyStories slice (16MB) + run: uv run python data/prepare.py --vocab 32768 --max-bytes 16777216 + - name: Train a tiny ple checkpoint (smoke, CPU only) + run: | + uv run python src/train.py --arm ple --vocab 32768 --target-core 590000 \ + --batch-size 16 --seq-len 256 --steps 100 --tag ci + - name: Export model.bin + golden logits + run: uv run python src/export.py ple-ci-s0 + - name: Build host verifier + run: cc -O3 -o /tmp/esp32-llm-verify firmware/host_verify/verify.c -lm + - name: Verify C matches the PyTorch golden + run: | + out=$(/tmp/esp32-llm-verify firmware/model/model.bin firmware/model/golden.txt) + echo "$out" + echo "$out" | grep -q "PASS" || { echo "::error::C port diverged from golden"; exit 1; } + - name: Build + run host perplexity harness + run: | + cc -O3 -o /tmp/esp32-llm-ppl firmware/host_verify/ppl.c -lm + /tmp/esp32-llm-ppl firmware/model/model.bin data/val_v32768.bin 8 + - name: Upload model artifacts + uses: actions/upload-artifact@v4 + with: + name: ci-model + path: firmware/model/ + + firmware-compile: + name: Compile ESP32 sketch (no board needed) + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: astral-sh/setup-uv@v6 + - uses: arduino/setup-arduino-cli@v2 + - name: Prepare a small tokenizer (vocab 4096) + run: uv run python data/prepare.py --vocab 4096 --max-bytes 8388608 + - name: Generate vocab.h for the sketch + run: uv run python src/gen_assets.py --tok data/bpe4096.json + - name: Install ESP32 core (pinned, matching README) + run: arduino-cli core update-index && arduino-cli core install esp32:esp32@3.3.10 + - name: Install display libraries + run: arduino-cli lib install "Adafruit GFX Library" "Adafruit BusIO" "Adafruit SH110X" + - name: Compile sketch + run: | + arduino-cli compile \ + --fqbn 'esp32:esp32:esp32s3:UploadSpeed=921600,USBMode=hwcdc,CDCOnBoot=cdc,UploadMode=default,CPUFreq=240,FlashMode=qio,FlashSize=16M,PartitionScheme=custom,PSRAM=opi,DebugLevel=info' \ + --build-property compiler.optimization_flags=-O3 \ + firmware/esp32_llm diff --git a/.gitignore b/.gitignore index d77a05e..065df50 100644 --- a/.gitignore +++ b/.gitignore @@ -1,23 +1,32 @@ -# python +# ---- Python --------------------------------------------------------------- __pycache__/ .venv/ *.pyc -# data + tokenizers (regenerate with data/prepare.py) +# ---- Dataset + tokenizer (downloaded / trained by data/prepare.py) -------- +# The TinyStories slice is ~300MB and the BPE tokenizer + uint16 token bins are +# regenerable, so none of these are committed. Reproduce with: +# uv run python data/prepare.py data/*.bin data/*.txt data/*.json -# training outputs +# ---- Training checkpoints (created by src/train.py) ------------------------ +# Large and machine-specific; reproduce with src/train.py. runs/ -# generated firmware artifacts (regenerate with src/export.py, src/gen_assets.py) +# ---- Generated firmware artifacts ------------------------------------------ +# model.bin and golden.* are produced by src/export.py from a checkpoint in +# runs/. They are excluded to keep the repo small; the trained weights are also +# published separately so nobody has to train to reproduce the demo -- see +# tools/fetch_model.py and firmware/esp32_llm/README.md. vocab.h is produced by +# src/gen_assets.py from the tokenizer JSON above. firmware/model/model.bin firmware/model/golden.npz firmware/model/golden.txt firmware/esp32_llm/vocab.h -# raw video / recordings (demo footage — keep out of the repo) +# ---- Raw video / recordings (demo footage -- keep out of the repo) --------- video/ videos/ *.mp4 diff --git a/data/prepare.py b/data/prepare.py index af4a8e2..9f0f86c 100644 --- a/data/prepare.py +++ b/data/prepare.py @@ -63,11 +63,16 @@ def train_tokenizer(text): def main(): - global VOCAB_SIZE + global VOCAB_SIZE, SLICE_BYTES ap = argparse.ArgumentParser() ap.add_argument("--vocab", type=int, default=4096) + ap.add_argument("--max-bytes", type=int, default=SLICE_BYTES, + help="bytes of TinyStories to download/train on (default " + f"{SLICE_BYTES / 1e6:.0f}MB; pass a smaller value for a " + "quick smoke test or CI)") args = ap.parse_args() VOCAB_SIZE = args.vocab + SLICE_BYTES = args.max_bytes # vocab 4096 keeps the original train.bin/val.bin; others get suffixed names # so both datasets coexist and train.py can pick by --vocab. suffix = "" if VOCAB_SIZE == 4096 else f"_v{VOCAB_SIZE}" diff --git a/firmware/esp32_llm/README.md b/firmware/esp32_llm/README.md index 42024ba..65ea916 100644 --- a/firmware/esp32_llm/README.md +++ b/firmware/esp32_llm/README.md @@ -6,6 +6,17 @@ embedding/output head is staged in PSRAM at boot. ## Build and verify +Get the trained `model.bin` first. It is not committed (see `.gitignore`) -- once +the weights are released it can be fetched with `tools/fetch_model.py`, which +verifies the SHA-256 below: + +```bash +python tools/fetch_model.py # after the model is published (issue #7) +python tools/fetch_model.py --check-only +``` + +Alternatively, export it yourself from a trained checkpoint with `src/export.py`. + Export the group-128 ragged-int4 model and verify the portable C runtime first: ```bash diff --git a/src/gen_assets.py b/src/gen_assets.py index a0d1f61..d6ce7af 100644 --- a/src/gen_assets.py +++ b/src/gen_assets.py @@ -4,6 +4,7 @@ - prints the prompt token ids for a fixed demo prompt to paste into the sketch. """ +import argparse import os from tokenizers import Tokenizer @@ -15,7 +16,11 @@ def main(): - tok = Tokenizer.from_file(TOK) + ap = argparse.ArgumentParser() + ap.add_argument("--tok", default=TOK, + help="path to the BPE tokenizer JSON (default data/bpe32768.json)") + args = ap.parse_args() + tok = Tokenizer.from_file(args.tok) V = tok.get_vocab_size() # Raw bytes per token: decode single-id sequences. For the ASCII TinyStories diff --git a/tools/fetch_model.py b/tools/fetch_model.py new file mode 100644 index 0000000..cc6ea6d --- /dev/null +++ b/tools/fetch_model.py @@ -0,0 +1,93 @@ +#!/usr/bin/env python3 +"""Fetch the published model artifacts for the ESP32 firmware. + +The trained 28.9M-parameter checkpoint is not committed (see .gitignore) and +reproducing it requires training. The released weights are hosted at +MODEL_URL once published (tracked in issues #5 and #7); this script downloads +model.bin, verifies its SHA-256 against the hash documented in +firmware/esp32_llm/README.md, and writes it into firmware/model/ next to the +golden files the host verifier needs. + +Usage: + python tools/fetch_model.py [--url URL] [--sha HEX] + python tools/fetch_model.py --check-only + +Stdlib only, so it runs anywhere. +""" + +import argparse +import hashlib +import os +import shutil +import sys +import tempfile +import urllib.request + +HERE = os.path.dirname(os.path.abspath(__file__)) +OUT = os.path.join(HERE, "..", "firmware", "model") +MODEL_BIN = os.path.join(OUT, "model.bin") + +# Fill this in once the weights are released (issue #5). Keep in sync with the +# SHA-256 printed in firmware/esp32_llm/README.md. +MODEL_URL = "" + +# SHA-256 of the artifact used for the on-device measurements, from +# firmware/esp32_llm/README.md. +EXPECTED_SHA = "21067f5d78113f6c64a8720b05ff7e5c774dab0276797a522f81a6797253d97c" + + +def sha256_of(path, chunk=1 << 20): + h = hashlib.sha256() + with open(path, "rb") as f: + while True: + block = f.read(chunk) + if not block: + break + h.update(block) + return h.hexdigest() + + +def check_only(): + if not os.path.exists(MODEL_BIN): + print(f"not present: {MODEL_BIN}") + return 1 + got = sha256_of(MODEL_BIN) + ok = got == EXPECTED_SHA + print(f"{MODEL_BIN}") + print(f" expected {EXPECTED_SHA}") + print(f" got {got} {'OK' if ok else 'MISMATCH'}") + return 0 if ok else 1 + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--url", default=MODEL_URL, help="direct URL to model.bin") + ap.add_argument("--sha", default=EXPECTED_SHA, help="expected SHA-256") + ap.add_argument("--check-only", action="store_true", + help="verify an already-downloaded model.bin and exit") + args = ap.parse_args() + + if args.check_only: + return check_only() + + if not args.url: + print("MODEL_URL is not set yet -- the trained weights are not published.") + print("Track issue #7 (https://github.com/slvDev/esp32-ai/issues/7) and") + print("issue #5 for the release. Once available, set MODEL_URL or pass --url.") + return 1 + + os.makedirs(OUT, exist_ok=True) + tmp = os.path.join(tempfile.gettempdir(), "model.bin.download") + print(f"downloading {args.url}") + urllib.request.urlretrieve(args.url, tmp) + got = sha256_of(tmp) + if got != args.sha: + print(f"SHA-256 mismatch: expected {args.sha}, got {got}") + return 1 + shutil.move(tmp, MODEL_BIN) + print(f"verified + saved {MODEL_BIN} ({os.path.getsize(MODEL_BIN) / 1e6:.2f} MB)") + return 0 + + +if __name__ == "__main__": + sys.exit(main())