economize

Composable efficiency techniques for Hugging Face causal language models.

economize attaches techniques to an unmodified transformers model. The base model keeps its module tree and its forward. Drafters and other added components bring their own parameters and act only while decoding; weight transforms (pruning, low rank, quantization) rewrite the base weights in place, once, when they attach. Any set of registered techniques can be combined in one call, and every architecture adapter, technique, verifier, cache and checkpoint format is added by registration.

Version 0.3 ships cross-layer jump drafting, static depth pruning, Wanda / SparseGPT pruning, SVD and SVD-LLM low rank, GPTQ / AWQ / round-to-nearest quantization, TEAL activation sparsity and draft-only quantization, on 44 model types (dense, MoE, and hybrid attention plus Mamba / linear-attention models), with batched decoding of ragged prompts, draft trees and lossless speculative sampling.

Install

pip install economize                       # torch>=2.4, transformers>=5.0, safetensors, huggingface_hub
pip install "economize[all]"                # every optional integration below except vLLM and SGLang

From source:

git clone https://github.com/Thiggel/economize.git
pip install -e "./economize[dev]"

Extra	Installs	For
test	pytest	the test suite
peft	peft	PEFT-wrapped base models and LoRA training of the base
quant	torchao, compressed-tensors	quantized kernels, compressed-tensors export
kernels	kernels	Hugging Face kernels hub kernels (economize.integrations.kernels)
train	accelerate, datasets, trl	drafter training and online RL rollouts
eval	lm-eval[hf]	economize-lm-eval, lm-evaluation-harness with economize's decoding loop
bench	datasets	prompt suites of economize.bench (Spec-Bench, GSM8K, HumanEval)
vllm, sglang	vllm, sglang	serving backends (economize.backends; drafters are written by economize.backends.export_drafter)

Installing the package registers the vLLM plugin (vllm.general_plugins), the compressed-tensors quantization backend and the economize-lm-eval command through entry points.

Quick start

import torch
import economize
from transformers import AutoModelForCausalLM, AutoTokenizer

base = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B", dtype=torch.bfloat16, device_map="cuda")
model = economize.load(base, "path/to/drafter")            # a trained drafter: a folder or a Hub repo id
# or in one step, loading the base model named in the checkpoint:
model = economize.from_pretrained("path/to/drafter", dtype=torch.bfloat16, device_map="cuda")

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B")
ids = tokenizer("The water cycle describes", return_tensors="pt").input_ids.cuda()
out = model.generate(ids, max_new_tokens=128)             # lossless: the same tokens as base.generate

model.generate accepts the usual transformers arguments (logits processors such as repetition_penalty, eos_token_id, return_dict_in_generate, ...) and routes them through economize's decoding loop. model.decode returns the tokens together with a log of every model call and the acceptance of each cycle:

result = model.decode(ids[0], 128, decoding=economize.DecodeConfig(cycle_length=4))
result.cost_summary()      # exact passes, draft steps, layer applications relative to plain decoding
result.mean_accepted       # accepted drafts per cycle

Jump drafting

A draft step for the next token predicts its hidden state after layer m from the current token's state after layer u and the next token's embedding, runs only layers m+1..r and reads out logits. Layers 1..m (and r+1..L) are skipped. One cycle of N = cycle_length tokens runs the unmodified model once, in parallel, over every token not yet computed exactly, which also yields one exact token, followed by up to N-1 draft steps.

config = economize.JumpConfig(
    target_layer=12,            # m
    source_layer=18,            # u
    exit_layer=None,            # r (default: the last layer; below it a readout adapter is added)
    adaptation="frozen",        # or "lora" (LoRA active only on the jump path) or "full" (trainable layer copies)
)
model = economize.apply(base, config)
list(model.technique_parameters())   # the added parameters only; the base model is untouched

DecodeConfig chooses how drafts are used:

Option	Meaning
cycle_length	Tokens per cycle N (1 is ordinary decoding).
lossless=True	Speculative decoding: drafts are verified by the exact pass, the output equals greedy decoding.
lossless=False	Every draft is kept and the exact pass only refreshes the cache (lossy, cheaper).
min_confidence	Stop drafting in a cycle once the drafter's top probability falls below this value.
verifier	A registered verifier by name: "greedy", "rejection" (lossless speculative sampling at T > 0, chains and trees), "typical" (Medusa's typical acceptance, lossy) or "accept_all".
verifier_options	Keyword options of the verifier, e.g. {"threshold": 0.09, "alpha": 0.3} for "typical".
cuda_graphs	Replay the exact passes and the draft steps from CUDA graphs (see below).
sync_interval	Ordinary decoding reads tokens back and checks stop conditions every this many steps (default 8; the output is the same as reading every step).

Other conditioning variants are part of JumpConfig: several source layers (extra_source_layers) combined by a registered mixer ("concat", "fused_linear", "layer_attention"), the predictor ("mlp", "linear", its width and depth, or "flow" / "flow_linear": that mean followed by a few Euler steps of a rectified flow, trained by flow matching, with flow_steps, flow_sigma, flow_width), and ablations (use_state, use_token, token_residual).

Several prompts of different lengths decode as one batch, each row accepting its own number of drafts per cycle. With verifier="rejection" and a sampling selector the output follows the model's own sampling distribution (speculative sampling):

results = model.decode_batch([ids_a, ids_b, ids_c], 128)          # one DecodeResult per prompt
sampled = model.decode_batch(
    [ids_a, ids_b], 128,
    decoding=economize.DecodeConfig(cycle_length=4, verifier="rejection"),
    selector=economize.TopPSampler(temperature=0.7, top_p=0.9, seed=0, device="cuda"),
)

Each row may also have its own selector, stop condition and token limit (selectors=, stops=, and a list for max_new_tokens). Rows that share a selector object are drafted and verified together.

At batch size 1 a draft step and a short exact pass are bound by kernel launches. With DecodeConfig(cuda_graphs=True) the decoder records its exact chain passes and the jump drafter's steps as CUDA graphs on a static KV cache and replays them. The model keeps one decoder per configuration, so the graphs recorded in the first generation serve the following ones; model.refresh() drops them after the weights or the techniques change. Draft trees, dynamic-depth policies, KV-cache techniques and models with recurrent layers decode eagerly whatever the flag says.

Combining techniques

economize.apply takes any number of configurations and attaches them in a fixed order, whatever the order of the arguments: weight transforms (by their order: depth pruning, pruning, low rank, quantization), activation hooks, KV-cache wrappers, dynamic depth, the drafter, then the rest. Techniques that need data calibrate on the model as it is at that point, so quantization sees the pruned weights.

from economize.techniques.prune import PruneConfig
from economize.techniques.quant import QuantConfig
from economize.techniques.sparsity import TealConfig

model = economize.apply(
    base,
    PruneConfig(method="wanda", sparsity="2:4"),
    QuantConfig(format="int4", method="gptq", group_size=128),
    TealConfig(sparsity=0.4),
    economize.JumpConfig(target_layer=12, source_layer=18),
    calibration=calibration_ids,          # LongTensor [N, S], an iterable of batches, or a callable
)

See docs/weights.md for every weight and activation technique and its options.

Training

economize.train trains what the techniques add (drafters, readout adapters, LoRA on a drafter's path, early-exit heads) and, where a technique needs it, the base model:

from economize.train import fit

model = economize.apply(base, economize.JumpConfig(target_layer=16, source_layer=17, exit_layer=17))
fit(model, "data/qwen/train.bin", steps=2000, output_dir="runs/jump_m16_u17")   # uint32 token file
model = economize.load(base, "runs/jump_m16_u17/final")

fit picks the objective registered for the attached drafter (jump_rollout for jump, the objective of tskip/train.py). An existing transformers.Trainer pipeline uses EffTrainer with the same objectives:

from economize.train import EffTrainer, JumpRollout, economize_training_args

objective = JumpRollout(horizon=16)
args = economize_training_args(objective, "runs/jump_trainer", max_steps=2000, per_device_train_batch_size=8)
EffTrainer(model, objective, args, train_dataset=dataset).train()

pip install "economize[train]" adds accelerate, datasets and trl (fit itself needs only torch). docs/training.md lists the objectives, the data formats, data-parallel runs under accelerate launch and how to register an objective.

Ecosystem

Library	Entry point	Details
transformers	model.generate(...), custom_generate Hub recipe	logits processors, stopping criteria, return_dict_in_generate
Hugging Face Hub	model.save_pretrained, push_to_hub, economize.from_pretrained	only the technique weights are uploaded
PEFT	economize.apply(peft_model, ...)	the adapter finds the base model inside a PeftModel
kernels	economize.integrations.kernels.kernelize(model), op(name)	hub kernels for the base model and every added module
lm-evaluation-harness	economize-lm-eval run --model economize --model_args pretrained=...,economize=...	generation runs economize's decoding loop
TRL	economize.integrations.trl.attach_rollouts	in-process GRPO / RLOO rollouts with the drafter refreshed after each step
vLLM	economize.backends.vllm	drafter export (export_drafter(model, dir, backend="vllm")), plugin entry point, jump drafting in the V2 model runner (tested with vLLM 0.30)
SGLang	economize.backends.sglang	interface only (JumpDraftBackend); its registered exporter raises UnsupportedBackendError, and nothing is patched in SGLang yet

docs/ecosystem.md describes each integration and its tests.

Saving and loading

model.save_pretrained("my_drafter")       # economize_config.json + technique weights in safetensors
model = economize.load(base, "my_drafter")

Only the technique weights are saved. Checkpoints of the research code (tskip, a jump.pt file or its run directory) are detected and converted on load. Further formats are added through the FORMATS registry.

Extending economize

Everything that varies is a registry, and new entries are added by registration, never by editing economize:

Registry	Module	Entry-point group
TECHNIQUES	economize.core	economize.techniques
ARCHITECTURES, LAYER_KINDS	economize.arch	economize.architectures, economize.layer_kinds
FORMATS	economize.core	economize.formats
VERIFIERS, POLICIES	economize.decoding	economize.verifiers, economize.draft_policies
CACHES, MASKS, MASK_FORMATS	economize.runtime	economize.caches, economize.masks, economize.mask_formats
PREDICTORS, MIXERS	economize.techniques.jump	economize.jump_predictors, economize.jump_mixers
SKIP_SEARCHES, SKIP_METRICS	economize.techniques.layer_skip	economize.skip_searches, economize.skip_metrics
MEDUSA_TREES	economize.techniques.medusa	economize.medusa_trees
MTP_FAMILIES	economize.techniques.vendor_mtp	economize.mtp_families
CONFIDENCE_MEASURES, FLEXI_ROUTERS	economize.techniques.dynamic_depth	economize.calm_confidence, economize.flexi_routers
PRUNE_METHODS, DEPTH_CRITERIA	economize.techniques.prune	economize.prune_methods, economize.depth_criteria
QUANT_METHODS, QUANT_BACKENDS, QUANT_FORMATS, AWQ_SCALE_GROUPS, TORCHAO_TENSORS	economize.techniques.quant	economize.quant_methods, economize.quant_backends, economize.quant_formats, economize.awq_scale_groups, economize.torchao_tensors
FACTORIZERS	economize.techniques.lowrank	economize.factorizers
TEAL_BACKENDS	economize.techniques.sparsity	economize.teal_backends
EVICTION_METHODS	economize.techniques.kv_evict	economize.kv_eviction
KERNEL_OPS	economize.integrations.kernels	economize.kernel_ops
OBJECTIVES, DIVERGENCES, SCHEDULES, STATE_LOSSES	economize.train	economize.objectives, economize.divergences, economize.schedules, economize.state_losses
PROMPT_SOURCES, HARDWARE, SHAPE_READERS	economize.bench	economize.prompt_sources, economize.hardware, economize.shape_readers
EXPORTERS	economize.backends	economize.exporters
VLLM_DRAFTERS	economize.backends.vllm	economize.vllm_drafters
CT_SCHEMES	economize.integrations.compressed_tensors	economize.ct_schemes

A new predictor, for example:

from economize.techniques.jump import PREDICTORS, Predictor

@PREDICTORS.register("residual_linear")
class ResidualLinear(Predictor):
    def __init__(self, hidden_size, config):
        super().__init__(hidden_size, config)
        self.proj = torch.nn.Linear(2 * hidden_size, hidden_size)

    def forward(self, source, token_embedding):
        return token_embedding.float() + self.proj(torch.cat([source, token_embedding], -1).float())

model = economize.apply(base, economize.JumpConfig(target_layer=12, source_layer=18, predictor="residual_linear"))

An installed package can contribute the same entry without an import, through the entry-point group in its pyproject.toml:

[project.entry-points."economize.jump_predictors"]
residual_linear = "my_pkg.predictors:ResidualLinear"

A new technique subclasses economize.Technique with a frozen TechniqueConfig dataclass and is registered with @economize.register_technique. A drafting technique declares Capability.DRAFTER and returns an object that follows the economize.decoding.Drafter protocol. A new model family registers an ArchAdapter under its model_type; most decoder-only families share the built-in StandardDecoderAdapter. Family facts that techniques and the cost model need are set on the adapter: layer_layout = "llama" gives AWQ the Llama scale groups, shape_fields corrects the cost model's config reading, and module_path(module) names modules for exporters. A drafter is written for a serving backend by a DrafterExporter registered with @economize.backends.register_exporter("<backend>", "<technique>"); a new backend is one such class. docs/ARCHITECTURE.md describes every contract.

Scope of version 0.2

• Techniques: jump drafting (frozen, LoRA, full; one or several source layers mixed by concatenation, a fused linear map or layer attention; MLP or linear predictors), depth_prune, prune, lowrank, quant, teal, draft_quant.
• Decoding: batched ragged prompts with per-row selectors, stops and limits, chains and draft trees, greedy verification, lossless speculative sampling on chains and trees, typical acceptance, lossy drafting, fixed or confidence-based draft length, CUDA graphs for chain decoding, a log of every model call for cost accounting.
• Training: economize.train (fit, EffTrainer, registered objectives).
• Checkpoints record a fingerprint of the base weights; loading onto different weights warns (economize.TeacherMismatchWarning, off with check_teacher=False).
• Not yet: flex-attention masks for draft trees (trees use sdpa or eager masks), draft trees on models with recurrent layers, and the vLLM / SGLang draft step.

Tests

cd TJEPA-Eff && source env.sh
PYTHONPATH=economize/src:economize/tests:. $PY -m pytest economize/tests            # unit, technique, integration, oracle (CPU)
PYTHONPATH=economize/src:economize/tests:. $PY -m pytest economize/tests -m slow    # trained checkpoints on Qwen2.5-0.5B
PYTHONPATH=economize/src:economize/tests:. $PY -m pytest economize/tests -m gpu     # batched decoding on a CUDA device

The lint (uvx ruff check src tests hub) and type (uvx mypy --python-executable <python> src/economize) gates and the design rules a change must follow are in docs/CONTRIBUTING.md.

The unit tests check that the exact path equals the model's own logits on every supported family, that batched lossless decoding equals greedy decoding of each prompt alone (chains, trees and jump drafters, with partial acceptance and hybrid recurrent layers), that rejection sampling outputs the target distribution (chi-square tests with a negative control), the recipe's ordering and rollback, and the checkpoint round trips. The oracle tests check that economize produces the same tokens, the same log of model calls and the same acceptance as the research decoder tskip.decode.Decoder, on random drafters of a tiny model and on three trained drafters of Qwen2.5-0.5B.

License

Apache License 2.0, see LICENSE.