trash-panda

Run models bigger than your GPU: stream weights layer by layer across GPU, RAM and disk.

Status: early (Phase 1). The API may change. The package name is temporary.

Why

Loading a big checkpoint the usual way needs two full copies in RAM, and a model that does not fit in VRAM simply fails. trash-panda builds the model without allocating weights (meta device), reads the checkpoint tensor by tensor, and keeps each block of the model wherever it fits best:

Only one block at a time is expanded to the compute dtype, so peak VRAM is roughly one block plus your activations, whatever the model size.

It works with any floating-point checkpoint: fp32, fp16 and bf16 are streamed as they are (no quantization needed), and fp8 (with or without scale tensors) is kept compact so more fits.

Install

Install PyTorch first (the build that matches your GPU, see https://pytorch.org), then:

pip install dumpsterpanda

The package is imported as trash_panda:

import trash_panda

To try the latest development version straight from GitHub instead:

pip install git+https://github.com/yla382/trashpanda.git

Why is it called dumpsterpanda on PyPI? 🦝😢

The project is called trash-panda, but the name trashpanda was already taken on PyPI, and PyPI wouldn't let trash-panda sit that close to it either. So the raccoon moved one dumpster over. The import is still trash_panda.

Use

import torch
from trash_panda import plan, stream_load, install_offload

with torch.device("meta"):                      # the structure only, no weights in memory
    model = MyModel(config).eval()

p = plan(model, "model.safetensors")            # what fits where, with warnings
print(p)
print(stream_load(model, "model.safetensors", plan=p))   # lazy load straight into the right tiers
install_offload(model)                          # hooks that stream blocks during forward

out = model(x)

Automatic or manual

plan() decides from the checkpoint's sizes and the free memory on your machine. Everything it decides can be overridden:

from trash_panda import Plan, Tier, plan

p = plan(model, path, resident_blocks=12, max_vram="9GB", min_free_ram="4GB", compute_dtype="bf16")
p.assign(Tier.DISK, indices=[27])               # edit it afterwards

p = Plan.manual(resident=["blocks.0.", "blocks.1."], streamed=["blocks.2.", "blocks.3."])   # or write it by hand

Or skip plans entirely and give prefixes directly (the low-level API):

stream_load(model, path, offload_prefixes=["blocks."], resident_prefixes=["blocks.0."], compute_dtype="bf16")
offload_blocks(model.blocks, device="cuda", compute_dtype=torch.bfloat16)

Fitting the GPU to your workload

How many blocks can stay on the GPU depends on how much memory your activations need, which grows with image size and cannot be known before the model runs. By default the planner keeps a fixed headroom_vram="1.5GB" free. To use a measured value instead, run your real call once and read the peak:

from trash_panda import install_offload, measure_activation_vram, plan, stream_load

p = plan(model, path)
stream_load(model, path, plan=p)
install_offload(model)

need = measure_activation_vram(lambda: run_my_generation(), p)   # include everything that runs, such as a VAE decode
print(f"activations need {need / 1e9:.2f} GB")

p = plan(model, path, activation_vram="2.6GB")                   # next time: fill the rest of the GPU

The value is per model and per input size, so measure again if you change the resolution.

Releasing a model

When a model is done, release() gives its memory back, so the next stage can be planned and loaded with it. Pass the handle that install_offload() (or offload_blocks()) returned:

from trash_panda import install_offload, plan, release, stream_load

p = plan(encoder, "encoder.safetensors")
stream_load(encoder, "encoder.safetensors", plan=p)
handle = install_offload(encoder)
embeddings = encoder(tokens)

release(encoder, handle)                         # its RAM and VRAM are free again

p = plan(unet, "unet.safetensors")               # this plan sees the memory the encoder gave back

release() removes the hooks, replaces every weight with an empty tensor on the meta device, and empties the GPU cache and PyTorch's cache of pinned memory. The model cannot be run afterwards. handle.remove() alone does not do this: it restores the blocks, but the weights stay in memory.

Storage dtype

By default weights are kept exactly as stored in the file. Optionally, storage_dtype="bf16" casts waiting weights down, and storage_dtype="fp8" quantizes large 2-D weights to fp8 on load (lossy; small tensors such as norms are left alone).

Benchmarks

Measurements on one machine, and a script to reproduce them with SDXL, are in BENCHMARK.md and examples/. They are single runs on a Windows desktop whose free memory changes from run to run, so they are indicative only, and they do not show that trash-panda is faster than the other tools.

What it does not do (yet)

• coordinating memory between several models that are loaded at the same time (for example a text encoder and a diffusion model). Each model is planned on its own from the memory that is free when you plan it, so plan them one after the other, or free the first model before loading the next
• one-call cleanup after a model is done (a release(model) helper and a pipeline stage() helper that remove the hooks, drop the weights, empty the GPU cache and report memory before and after). For now, free memory yourself: handle.remove(), del model, gc.collect(), torch.cuda.empty_cache(); for the strictest cleanup (system RAM in particular), run heavy stages in their own process
• full overlap of reading and computing for disk blocks: blocks waiting in RAM and on disk are prefetched through pinned buffers by a background thread, two blocks ahead by default, but the drive's speed sets a floor that no amount of prefetching removes (see "Known issues")
• Macs and other non-NVIDIA GPUs: the prefetching (a second CUDA stream, events and pinned buffers) is CUDA-only. On a Mac the library should fall back to the CPU, with plain copies and no prefetching, but this has not been tested on a Mac yet, and the Apple GPU (mps) has not been tried at all. A Mac's CPU and GPU also share one memory pool, so splitting a model between "GPU" and "RAM" does not free anything there
• int8 / 4-bit formats (fp8 and ordinary float dtypes only)
• LoRA/LoKr adapters, multi-GPU, model definitions

Known issues

• Automatic plans change from run to run. plan() uses the RAM that is free at the moment you call it, and free RAM varies. On a 16 GB machine the same script gave a RAM budget of 6.2 GB in one run and 8.0 GB in another, which moved several blocks between RAM and disk
• Load other models before you plan. If a second model (a text encoder, a VAE) is loaded after plan(), the memory the plan counted on is taken. In RAM, Windows pages the weights out and the first step is very slow. On the GPU, too many blocks stay resident and the overflow spills into shared memory, which can make every step several times slower. Create the other models first so the plan sees the memory that is really left
• The default RAM reserve is a rule of thumb. min_free_ram="auto" keeps 8% of the machine's RAM free (between 1 and 4 GB, so about 1.3 GB on a 16 GB machine). If your other programs use a lot of memory, or very little, set min_free_ram yourself, for example "500MB" or 0
• The default VRAM headroom is a fixed guess. headroom_vram="1.5GB" does not know how much memory your activations need (it grows with image size), and measuring it is a manual step for now (see "Fitting the GPU to your workload"). If too many blocks are resident, Windows silently spills GPU memory into shared system memory and every step becomes several times slower, with no error. Pass a measured activation_vram, lower resident_blocks, or set "CUDA - Sysmem Fallback Policy" to "Prefer No Sysmem Fallback" in the NVIDIA Control Panel so it fails loudly
• Disk blocks are limited by the speed of the drive. Blocks on disk are read ahead of time (two blocks ahead by default, lookahead=), which hid about a third of the read time in testing, but a read takes much longer than one block's compute on a slow drive. On the test machine's SATA SSD (about 400 MB/s measured, with 1, 2, 4 or 8 reader threads alike) a plan with 11 blocks (4.1 GB) on disk took about 16 s per step, against about 9.3 s with every streamed block in RAM. A deeper look-ahead (4) was no better and uses more VRAM. These are single runs
• The first use of a RAM-tier block may be slow. The RAM tier is mapped from the checkpoint file, so Windows may read a block from disk the first time it is used. In one profiled run a few blocks made the main thread wait 1 to 3 s each, and a text encoder that runs once copied its blocks at about disk speed. The cause is not confirmed yet (see the roadmap)
• Tested on one setup so far: a 12B diffusion model (fp8 with scales) on Windows with an RTX 3060 (12 GB) and 16 GB of RAM. Other models, GPUs and operating systems are not covered yet. The unit tests run on the CPU, and examples/sdxl_benchmark.py has a path for machines without CUDA (the Apple GPU or the CPU) that is untested

Roadmap

Roughly in the order they are planned. Nothing here is promised or scheduled.

Speed

• read the RAM tier into memory ahead of use. In the profile of an all-RAM plan, 21 of the 140 streamed blocks were late by about 1 s each, which is roughly the time to read one block from the drive, and each of the slowest was late in only 1 of 7 passes. That suggests nearly all the waiting is in the first pass, when the file-backed RAM tier is read from disk on first use. If so, steady steps may be much faster than the average the progress bar shows (about 6 s against 9.3 s) and the first step much slower. This is an inference, not a measurement: the profile now lists the time and the waiting of every pass (TRASH_PANDA_PROFILE=1), which will confirm or reject it. If it holds, a background thread would touch the RAM-tier pages once when the hooks are installed, keeping the tier file-backed
• prefetch models that run a single pass (a text encoder, a VAE) from their first block: the prefetcher learns the block order during the first pass, so a one-pass model is not prefetched at all. Hand it the order of the blocks when the hooks are installed, and let the order seen at run time correct it
• the number of pinned buffers as a tested option (TRASH_PANDA_PINNED_BUFFERS, 3 or 4 instead of 2), if the profile's "worker waited for a free pinned buffer" line shows the worker is held up by them
• overlap the stages of a pipeline: while the text encoder runs, load as much of the diffusion model as there is spare VRAM and RAM for, then release the encoder and start sampling with what is already loaded while the rest keeps loading; the same for the VAE during sampling. It needs the staged planning above, a plan that predicts the memory that will be free after a stage, and blocks that can change from streamed to resident while running. Preloading would also compete with the encoder for the drive when the encoder reads from disk. An idea, not a design
• a memory-mapped disk tier as an experiment: map the checkpoint once and let the operating system keep recently used blocks cached, instead of an explicit split between RAM and disk. To be measured against explicit read-ahead, because a least-recently-used cache can thrash when the model is read in a cycle
• cheaper expansion: reuse preallocated buffers and fuse the fp8 to bf16 conversion and scaling, instead of creating new tensors for every weight on every pass
• a cost model that chooses between GPU, RAM and disk from measured disk and PCIe speed

Planner In this order, because the silent slowdowns seen so far (GPU memory spilling into shared memory at about 11.6 to 11.8 GB, a headroom that was too thin, an overflow after sampling at 2K) all come from the planner knowing too little:

• warn when a plan fills almost all of the GPU, since the overflow is silent and slow
• measure the activation size automatically (and remember it per model and input size) instead of asking you to run measure_activation_vram once, and base the default headroom on it instead of the fixed 1.5 GB guess
• plan a pipeline in stages (encode, denoise, decode) so memory used by an earlier stage can be released. release() does the releasing by hand now. In single runs, releasing the text encoder before the diffusion model was planned gave the planner several GB more RAM, and releasing the diffusion model before the VAE was built stopped the decode of a 2048 x 1024 image from spilling into shared GPU memory (watched by eye, not logged). What is left is a helper that runs the stages in order, and a smaller step such as release_gpu() / restore_gpu() on the handle, which moves the resident blocks off the GPU and back while keeping the hooks. The cost of putting them back for each image, and the RAM they need meanwhile, have not been measured
• count the pinned buffers at the size PyTorch gives them: it appears to round each pinned allocation up to the next power of two (two buffers of 0.7 GB took 2.15 GB), so the RAM the planner counts for them can be too low
• take memory held by other loaded models into account, or at least warn about it
• a KV-cache budget for language models, from the context length you intend to use

Formats

• 4-bit and 8-bit storage (for example NF4 with bitsandbytes, or our own implementation), to fit more of the model in RAM and on the GPU and halve the traffic
• fp8 checkpoints with block-wise scales (the weight_scale_inv layout used on the Hugging Face Hub)
• GGUF models, possibly through llama-cpp-python as a separate backend

Compatibility

• from_pretrained-style helper for transformers models: build on the meta device with real buffers, find the weight files (Hub or local), translate old checkpoint names, plan, load and install the hooks in one call
• ignore buffers that are not in the checkpoint (rotary tables, position ids) instead of failing
• make model.dtype, model.device and model.to(...) behave on offloaded models, so generate() and pipeline() work
• diffusers pipelines: build the large components with stream_load and hand them to the pipeline class, so from_pretrained does not load everything into RAM first
• LoRA and other adapters on offloaded blocks
• compare speed and memory against the offloading built into diffusers and accelerate

Larger ideas

• compute the RAM tier on the CPU instead of copying it, as llama.cpp does, for token-by-token generation
• batching for language models, so one weight copy serves many sequences
• multi-GPU support

Reliability and project

• investigate sudden exits: rarely, and only when RAM is already high (about 12 GB in use on a 16 GB machine) before a large model is loaded, the process ends with no Python error, usually while the weights are being loaded. It is probably the operating system ending the process or a crash in native code, but this has not been confirmed. Next time, record the exit code and run with PYTHONFAULTHANDLER=1, and check whether it also happens from a clean memory state
• find out where a sudden exit goes wrong: when a run ends without a Python error, check the exit code right afterwards ($LASTEXITCODE in PowerShell) and run with PYTHONFAULTHANDLER=1. A large negative code such as -1073741819 is a crash in native code, while a normal code points to the process being ended by the system. Also print the memory in use and the available commit (RAM plus page file) before loading, so a failed run shows what the machine had left
• a stage() helper that runs a pipeline's stages in turn and releases each one (release() itself is done)
• tests on tiny transformers and diffusers models on the CPU, in CI
• tests on more models, GPUs and operating systems (Linux in particular), with benchmarks published here. A first test on a Mac is planned: the unit tests and examples/sdxl_benchmark.py --size 512 --steps 10 in the baseline and trash-panda modes, to see whether it installs and runs on the CPU or the Apple GPU. Real Mac support, if there is demand, would be a separate piece of work (mps, unified memory, fp8 on the Apple GPU)
• a documentation site (GitHub Pages) with guides and an API reference
• a final project name, and a contributor agreement before outside code is accepted

Development

pip install pytest psutil
pytest

Tests run on the CPU with a tiny fake model.

License

AGPL-3.0-only. See LICENSE.