Skip to content

MPSLab-ASU/HELM

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

48 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

HELM: Heterogeneous Execution for Large Models

CI Artifact package License: Apache 2.0 Python 3.10+

HELM is a compiler and runtime for running large language models on consumer hardware — machines with a single GPU and CPU RAM. It microbenchmarks your hardware, compiles the model into static per-device FX subgraphs, and executes heterogeneously across CPU and GPU with a paged KV cache that offloads cold pages to CPU RAM during long-context decode.


Reproducing the paper

Docker-first, one obvious path (no GPU needed for the first step):

git clone https://github.com/MPSLab-ASU/HELM && cd HELM

# 1. CPU-only functional check: builds the CPU image and runs the 426-test
#    suite + the canonical-results self-check (the image's default command).
#    Expect "412 passed, 14 skipped" here (the skips are GPU-gated and run
#    on CUDA hosts; see REPRODUCING.md).
docker build --build-arg TORCH_INDEX=cpu -t helm:cpu . && docker run --rm helm:cpu

# 2. GPU image (CUDA 12.8 wheels); add the vLLM/DeepSpeed baselines with :full.
docker build -t helm .
docker build --build-arg INSTALL_GPU_EXTRAS=1 -t helm:full .

# 3. One-command reproduction (native or in-container): each target runs
#    its experiments AND then renders the corresponding paper-vs-measured
#    comparison (PNG + markdown in results/figures/) automatically.
#    Run with no args for the full target list + time estimates:
bash reproduce.sh test        # pytest + results validator (CPU-only, ~5 min)
bash reproduce.sh smoke       # small end-to-end GPU run + correctness (~30 min)
bash reproduce.sh table2      # Table II grid -> renders Table II (hours)
bash reproduce.sh fig4        # also: fig5, fig6, table3 — each renders its figure
bash reproduce.sh all         # everything above in sequence + PASS/FAIL/SKIPPED summary
# Utilities:
bash reproduce.sh figures [run_dir]   # re-render only (no experiments)
bash reproduce.sh validate <run_dir> "<GPU>" <Model>   # compare a run to the paper

Docker runs that download models should mount the Hugging Face cache, e.g.:

docker run --rm --gpus all -v $HOME/.cache/huggingface:/root/.cache/huggingface \
    helm bash reproduce.sh smoke

Per-claim commands, expected outputs and tolerances: REPRODUCING.md.


vs. Existing Tools

vs. Accelerate — Accelerate uses device_map='auto' which fills GPU from layer 0 upward with no runtime cost model. HELM microbenchmarks your actual hardware (FLOPS at decode/prefill regimes, memory bandwidth, PCIe throughput) and finds the partition that minimises decode latency under a roofline model. It compiles to isolated FX subgraphs rather than attaching per-layer forward hooks, eliminating Python-level synchronisation overhead at every boundary.

vs. vLLM — vLLM requires all model weights to fit in GPU VRAM and OOMs on models larger than VRAM capacity. HELM targets exactly this regime: models that cannot fit on a single GPU but need to run on one anyway.

vs. FlexGen — FlexGen overlaps weight streaming with compute, which only helps at batch sizes in the hundreds. HELM keeps all weights resident in CPU RAM + GPU VRAM statically; the only cross-device traffic is a single activation tensor at the stage boundary, once per decode step.

vs. DeepSpeed ZeRO-Inference — DeepSpeed streams weights from CPU per layer during inference, incurring O(L) PCIe round-trips per generated token. HELM places weights statically and only transfers activations at the stage boundary — no per-layer weight movement during decode.


How It Works

Model (nn.Module)
      │
      ▼
 [1] FX Tracing          — Capture a single decode graph (seq_len=1); the same
      │                     compiled graph is reused for both prefill and decode
      ▼
 [2] HelmGraph IR         — Lift FX nodes into typed IR with shape, FLOP, and
      │                     byte-cost metadata (shape propagation + fallback estimator)
      ▼
 [3] HybridAnalyzer       — Annotate each node with param bytes, activation bytes,
      │                     and KV bytes per token; aggregate per transformer block
      ▼
 [4] PartitionUnitBuilder — Group nodes into coarse units: one embedding unit,
      │                     L transformer-block units, one output-projection unit
      ▼
 [5] DeviceProfiler       — Microbenchmark actual GPU/CPU FLOPS (decode GEMV regime
      │                     + prefill GEMM regime), DRAM bandwidth, and PCIe H2D/D2H;
      │                     results cached across compile() calls
      ▼
 [6] StrategySelector     — Score all O(L) feasible CPU→GPU split points under the
      │                     roofline cost model (separate projection + attention
      │                     rooflines; L3 cache hierarchy for short-context KV;
      │                     GPU flash-attention activation traffic = 0);
      │                     return the split minimising the configured objective
      ▼
 [7] StageFXBuilder       — Fragment the FX graph into per-device standalone
      │                     GraphModules at the chosen split boundary
      ▼
 [8] PipelineRuntime      — Prefill pass → autoregressive decode loop (batch ≥ 1);
      │                     both phases reuse the same compiled stage subgraphs
      ▼
 [9] StageRuntimeExecutor — Execute stages in sequence; transfer activation tensor
      │                     across device boundary; pre-position weights one submodule
      │                     at a time to avoid peak-RAM spikes during load
      ▼
[10] KVOffloadManager     — Paged KV cache with GPU watermark; evict oldest pages
                            to CPU pinned RAM; stream back page-by-page during decode
                            via async H2D + online-softmax streaming attention

StrategySelector evaluates three plan classes: (1) all-GPU, (2) all feasible CPU→GPU 2-stage splits, (3) all-CPU. The exhaustive search runs in O(L) time and takes <1 ms.

KVOffloadManager patches attention forward at the class level (no FX graph changes needed). Supports batch_size > 1: one KVCacheManager per batch item sharing a single KVAllocator pool. Eviction policy: oldest page (by start_token) first, protecting the active tail page.


Results (paper Table II excerpt — RTX 3090, 24 GB VRAM, batch=1, fp16)

Decode throughput (tok/s) / TTFT p50 (s) at output length 128:

Model HELM Accelerate DeepSpeed vLLM
Qwen3-4B 87.1 / 0.019 25.7 / 0.041 26.7 / 0.040 87.0 / 0.018
Qwen3-8B 50.3 / 0.028 25.8 / 0.042 26.5 / 0.040 50.3 / 0.027
Qwen3-14B 8.1 / 4.7 1.3 / 0.782 OOM
Qwen3-32B 1.8 / 11.7 0.2 / 5.3 OOM

Key observations (full 3-GPU × 10-model grid in results/paper/):

  • On models that fit in VRAM, HELM dispatches to its vLLM-resident fast path and matches vLLM's throughput (1.0×) — the compile-time partitioning adds no overhead when offloading is unnecessary.
  • On models that exceed VRAM (14B/32B on the 3090; everything above 4B on an 8 GB GPU), HELM is the fastest feasible backend: up to 10.4× higher decode throughput than the strongest feasible baseline, 6.2× geometric mean across the nine overflow configurations.
  • HELM's paged KV offload extends the maximum context length by a geometric mean of 5.9× and up to 256× (Qwen3-14B: 32K vs. Accelerate's 128), with token-identical greedy outputs with and without offload.

Every number above is pinned in machine-readable form under results/paper/ and re-derived in CI by results/validate_results.py.


Installation

Requirements: Linux, Python 3.10+, NVIDIA driver ≥ 525 (the pinned cu128 wheels bundle the CUDA runtime — no host CUDA toolkit needed), uv

git clone https://github.com/MPSLab-ASU/HELM.git
cd HELM
uv sync

CPU-only (unit tests, cost model, planner — no GPU required):

pip install torch==2.10.0 torchvision==0.25.0 --index-url https://download.pytorch.org/whl/cpu
pip install -e ".[dev]"
pytest tests/ -q

Or via Docker (see the Dockerfile header for GPU/CPU build variants; the default container command runs the test suite + results validation):

docker build -t helm .
docker run --rm -it --gpus all helm bash    # interactive shell

The AVX2+F16C CPU kernel is built automatically via a C++ extension. Requires GCC ≥ 9 and a CPU with AVX2 support (Intel Haswell+ / AMD Ryzen+).


Usage

Quick start — auto partition

uv run helm \
    --model Qwen/Qwen3-4B \
    --mode execute_stagewise \
    --compiler-plan auto \
    --max-new-tokens 64 \
    --kv-offload

Manual partition

Specify exactly which layers run on CPU and which on GPU. Ranges are inclusive.

uv run helm \
    --model Qwen/Qwen3-14B \
    --mode execute_stagewise \
    --compiler-plan manual \
    --compiler-cpu-layers 0:7 \
    --compiler-gpu-layers 8:47 \
    --max-new-tokens 128 \
    --kv-offload

Inspect the partition plan without executing

uv run helm \
    --model Qwen/Qwen3-8B \
    --mode plan \
    --compiler-plan auto \
    --print-plan

Paper Experiments

The full experiment suite runs with a single command:

bash experiments/run_paper_experiments.sh

Self-contained — no arguments required. It detects hardware, selects feasible backends, and runs 4 models × 4 backends over sections E/A/B/C by default (add SECTIONS=...,F for the slow max-decode probe behind Fig. 5, and NO_LM_EVAL=0 for section D):

Section What it measures
E — Feasibility Does the model load? Peak GPU/CPU memory at load time
F — Max decode length Longest output before OOM
A — Latency sweep TTFT, decode tok/s — p50/p95/p99 across 10 requests (5 with QUICK=1)
B — Throughput Aggregate tok/s at batch sizes 1/2/4/8
C — Ablations batch size, context length, AVX on/off, KV offload on/off, CPU threads
D — LM quality MMLU, HellaSwag, ARC-Easy via lm-evaluation-harness

Models: Qwen/Qwen3-4B, Qwen/Qwen3-8B, Qwen/Qwen3-14B, Qwen/Qwen3-32B

Crash recovery: each section writes a .done sentinel on completion. Re-running skips completed sections.

Environment overrides:

QUICK=1              # 5 requests, shorter sequences — smoke-test on any machine
NO_LM_EVAL=1         # skip Section D
SECTIONS=E,F,A       # run only specific sections
OUT_ROOT=/scratch/results

Results land in experiments/results/<timestamp>/ with a SUMMARY.md containing paper-ready tables.


CLI Reference

Flag Default Description
--model required HuggingFace model name or local path
--mode plan dry_run · import · units · plan · lower · execute_stagewise · baseline
--compiler-plan auto auto (profiled) or manual
--compiler-cpu-layers Layer range for CPU, e.g. 0:7
--compiler-gpu-layers Layer range for GPU, e.g. 8:47
--kv-offload off Enable paged KV cache with CPU offloading
--max-new-tokens 8 Tokens to generate
--dtype float16 float16 · bfloat16 · float32
--cpu-threads 6 CPU thread count for PyTorch ops
--print-plan off Print full partition plan JSON

Artifact Evaluation

This repository is packaged as an ACM artifact (ESWEEK CASES call for artifacts, badges Available / Functional-Reusable / Results Reproduced):

  • ARTIFACT.md — artifact appendix: check-list, hardware/ software requirements, claims↔experiments mapping, badge justification and the HotCRP/Zenodo submission checklist.
  • REPRODUCING.md — command-by-command reproduction of every table and figure, with expected outputs, tolerances and time estimates (kick-the-tires path: ~30 min).
  • results/paper/ — the paper's numbers in machine-readable form; results/validate_results.py self-checks them and compares fresh runs against Table II.
  • CI (.github/workflows/) — CPU-only test suite (426 tests) on every push, canonical-results validation, artifact-tarball build (make_package.sh) and a Docker image build.
  • Dockerfile — pinned CUDA-12.8 / CPU-only container environments; CITATION.cff / .zenodo.json — citation and archival (DOI) metadata.

License

Apache License 2.0

About

HELM is a compiler and runtime for deploying Large Language Models (LLMs) on heterogeneous hardware clusters. It orchestrates execution across GPUs, and CPUs.

Resources

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages