InferZO

Inference-native zeroth-order optimization with a vLLM execution backend and Hugging Face Trainer integration. This release includes low-rank and dense parameter perturbations, synchronous and asynchronous execution, co-serving, Countdown evolution-strategies experiments, and their processed results.

See the experiment guide for the mapping from each experiment to its code and the evidence index for results. Dependency versions are pinned in uv.lock.

Layout

Path	What it is
zo_vllm/	The library: ZO directions (LoRA-factorized and full-parameter), estimators, update transforms, the LoRA update bank, the vLLM engine wrapper and worker-side scoring, the async (serving) runtime, task adapters, experiment infrastructure (job queue, manifests, timing).
zo_trainer/	Hugging Face Trainer integration used by the runners (probe loop, callbacks, metrics, checkpointing).
zo_serve/	The serving-time training API mounted into the vLLM OpenAI server (co-serving and async ES).
phase3/inference_native.py	Job-group definitions: every paper experiment is a named group of jobs (model, task, arm, seed, steps, ...).
phase3/runners/hf_trainer_speed_migration.py	The training/timing runner for every inference-native arm.
phase3/runners/launch_inference_native.py	Builds a group's jobs, writes them to a queue directory and starts one queue worker per GPU (zo_vllm.experiment.runners.job_queue_worker -> backend_job -> runner).
phase3/runners/* (other)	Stand-alone experiment scripts (forward bridge microbenchmark, mixed-batch stress, padding/packing, CUDA-graph control, analytic state traffic, task padding statistics) and result collectors/analyzers.
phase7/scripts/run_throughput_parity.sh	Co-serving experiment driver (one vLLM server; foreground load plus LoRA-prefill and ZO background).
es_countdown/	Evolution-strategies Countdown experiments (generation-reward ZO) and the step-level consistency check against the ES-at-Scale code.
baselines/	Reference-alignment oracle and goldens (LoZO, MeZO, DPZero, ZO-Bench sign/momentum, ES-at-Scale), the transformers-compatibility shim and bootstrap script for the released-LoZO baseline on Qwen3/Llama. The tests use the shipped JSON goldens; the generate_*_golden.py scripts are provided for reference (the ZO-Bench generator is not included). The ES-at-Scale generator verifies the released reference-source manifest.
docs/	Design notes on the execution graph; some modules they reference are not part of this release.
third_party_patches/	Our changes to vLLM, LoZO and ES-at-Scale as patches against public upstream commits, with the upstream licenses.
tests/, e2e_tests/	Unit tests (CPU) and GPU end-to-end tests (-m gpu / -m e2e, not run by default).

phase3/inference_native.py defines more job groups than the paper uses; the released groups are listed in docs/experiments.md. The library registers further ZO direction providers and update rules at import time; the paper uses LoZO-style low-rank directions, dense Gaussian directions, and the ZO-SGD, ZO-sign and DPZero update rules.

Setup

Requirements: Linux x86_64, NVIDIA GPU with a CUDA 12.9 or 13.0 driver, uv 0.11.15, git, tmux (the queue launcher starts its workers in tmux windows).

git clone https://github.com/playeriv65/InferZO.git
cd InferZO
# 1. vLLM fork = upstream vLLM at a fixed commit + our patch
git clone https://github.com/vllm-project/vllm third_party/vllm
git -C third_party/vllm checkout 0867497368f390212a3f9684e2e05f698f8d1149
git -C third_party/vllm apply "$PWD/third_party_patches/vllm-zo.patch"

# 2. Released LoZO baseline = upstream LoZO + our patch (timing hooks, scopes,
#    Qwen3/Llama loading, per-step traces)
git clone https://github.com/optsuite/LOZO third_party/LOZO
git -C third_party/LOZO checkout 5d1ade5bf41846471d33ba526f877829b9a19856
git -C third_party/LOZO apply "$PWD/third_party_patches/lozo-zo.patch"

# 3. Main environment (Python 3.12, torch 2.11.0, transformers 5.8, editable
#    vLLM built from the precompiled wheel of the same upstream commit)
scripts/sync-env.sh cu130          # or cu129 on drivers without CUDA 13
scripts/run-in-env.sh python -c "import zo_vllm, zo_trainer, zo_serve, vllm"

# 4. Released-LoZO interpreter (Python 3.10, transformers 4.28.1, OPT models)
uv venv --python 3.10 third_party/LOZO/large_models/.venv
uv pip install --python third_party/LOZO/large_models/.venv/bin/python \
    -r phase3/presets/lozo_paper_requirements_lock.txt
#    Qwen3/Llama baselines need a transformers-4.51 interpreter built by
#    baselines/bootstrap_lozo_qwen3.sh (it verifies itself against the 4.28 one
#    on an OPT job); point ZO_LOZO_BASELINE_PYTHON_QWEN3 / _LLAMA at it.

.env.example lists every environment variable the code reads (artifact root, Hugging Face cache, optional W&B, baseline interpreters); W&B is off unless enabled explicitly. scripts/run-in-env.sh refuses to start when an interpreter or compile cache sits on a declared network mount (ZO_NETWORK_MOUNTS, see scripts/local_disk_guard.py).

For the ES experiments install the helper package into the same environment:

uv pip install --python .venv/bin/python --no-deps -e es_countdown
git clone https://github.com/VsonicV/es-at-scale es_countdown/third_party/es-at-scale
git -C es_countdown/third_party/es-at-scale checkout 574a9d134da1ffce2a8bb812019899e5c96b588a
git -C es_countdown/third_party/es-at-scale apply "$PWD/third_party_patches/es-at-scale-zo.patch"
python es_countdown/es/official_consistency/reference_source.py es_countdown/third_party/es-at-scale

The Countdown train/eval data are read from that checkout (es_countdown/third_party/es-at-scale/datasets/...). The released consistency config and golden generator verify the public upstream revision, patch digest, and complete reference-file manifest (third_party_patches/es-at-scale-source.json). No local patch commit or config edit is required. Git metadata, .venv, and Python bytecode are excluded; changed or additional source files are rejected. The consistency check also needs the ES-at-Scale reference interpreter (vLLM 0.11.0; ES_AT_SCALE_PYTHON=<venv>/bin/python); see es_countdown/es/official_consistency/launch.py.

Running experiments

Queued job groups (training, timing, component chains, speedups, appendix):

# dry run: build the group's jobs and print them
scripts/run-in-env.sh python phase3/runners/launch_inference_native.py \
    --group official_lozo_speedup_all2d --gpus 0 --dry-run
# launch: one worker per listed GPU drains the queue; results land under
# phase3/results/<run id>/jobs/<job id>/ (make it a link to large storage)
scripts/run-in-env.sh python phase3/runners/launch_inference_native.py \
    --group official_lozo_speedup_all2d --gpus 0 --reuse-engine

--model, --variant and --steps restrict or shorten a group. Every job writes result.json, a manifest (code revision, environment, command line) and per-step timing/metric files. Timing values in the paper are tail means (steps 101-200 of a 200-step run unless stated otherwise).

Stand-alone scripts are run directly, e.g.

scripts/run-in-env.sh python phase3/runners/run_bridge_microbench.py --help
scripts/run-in-env.sh python phase3/runners/task_padding_stats.py --help
scripts/run-in-env.sh python phase3/runners/analytic_state_traffic.py --help
scripts/run-in-env.sh python -m phase3.runners.run_mixed_batch_stress --help
GPU=0 RANK=8 phase7/scripts/run_throughput_parity.sh     # co-serving
scripts/run-in-env.sh python es_countdown/es/run_countdown.py \
    --config es_countdown/configs/es/countdown_qwen25_1p5b_full_parameter.json

backend_job can also dispatch to a phase-4 SST-2 runner and a legacy vllm_zo_task runner, and serving_load has timeline/acceptance analysis modes; those scripts are not in this release.

Tests

scripts/run-in-env.sh python -m pytest            # CPU suite (default markers)
scripts/run-in-env.sh python -m pytest -c es_countdown/pyproject.toml \
    es_countdown/tests es_countdown/es/tests

Several tests read the patched third_party/vllm and third_party/LOZO sources, so apply the patches first. tests/test_phase3_inference_native.py:: test_source_provenance_captures_current_tree expects a git checkout, and test_strict_state_transition_plumbs_trace_to_both_runners expects the released-LoZO interpreter of step 4. Two ES tests (es/tests/test_official_consistency_cpu_probes.py) need ES_AT_SCALE_PYTHON and the reconstructed reference checkout described above.

Release history

The earlier implementation remains available at the legacy-public-2026-05-27 tag. Its source, dependency lock and experiment notes are preserved together. The current tree contains the maintained implementation and the released experiment evidence; runtime outputs, model weights and datasets are external.