marcelz/elephant-1-BF16-small-head

elephant-1 (original BF16 weights) on a 16 GB GPU

This repo runs the original, unmodified BF16 weights of baselabs/elephant-1 on a graphics card with 16 GB of memory. The original does not fit such a card as shipped. I made it fit with two changes, explained below: the model's head is cut down to the rows the answers use, and one lookup table stays in RAM.

On my RTX 5070 Ti a short request takes about 70 ms (median 72 ms) instead of 260 ms, and the answers stay the same (the largest difference in any probability is 0.000001).

The weights are not stored in this repo. install.sh downloads them from Base Labs (one fixed revision) and sets up everything else.

Independent work, not affiliated with or endorsed by Base Labs, Baseten or OneJev. Code in this repo: MIT. elephant-1: MIT (Base Labs).

Quick start

You need: Linux, an NVIDIA GPU with 16 GB (I tested an RTX 5070 Ti), NVIDIA driver 580 or newer, Python 3.12, git, about 30 GB of free disk space, about 16 GB of free RAM, and internet access.

pip install huggingface_hub==2.2.0 hf_xet==1.7.0      # only to download this repo
hf download marcelz/elephant-1-BF16-small-head --local-dir elephant-1-bf16
cd elephant-1-bf16
bash install.sh          # one command: Python environment, packages, the model (17.5 GiB), the engine
bash run.sh example.py   # loads the model, answers one request, shows time and memory

Output on my machine (your timings will differ):

BF16 model loaded in 2 s; GPU memory peaked at 12.9 GiB while loading
answer: billing | billing 99.9%, shipping 0.1%, technical 0.0%
speed: 20 requests, median 55 ms, slowest 57 ms
GPU memory: 13.0 GiB during a request (the card offers 15.5 GiB)
RAM: this process peaked at 15.3 GiB

The first load reads 17.5 GiB from disk (about 40 s on my NVMe drive); later loads are faster while the file is in the disk cache.

In your own code. Save this as my_script.py in this folder and run bash run.sh my_script.py from here (that uses the environment install.sh built):

import engine_slots                                           # engine_slots.py is in this folder
from qev.schema import SystemOneRequest

eng = engine_slots.load_engine("elephant-1")                  # the model folder made by install.sh: a local folder, not a hub id
resp, meta = eng.decide(SystemOneRequest.model_validate({
    "state": "I was charged twice for my order last week and nobody has replied.",
    "questions": {"route": {"type": "choice", "instructions": "Which team should handle this?",
                            "criteria": {"billing": "payments and refunds", "shipping": None, "technical": None}}}}))
print(resp.answers["route"])

Two warnings (about causal_conv1d and Triton's FP fusion) are expected and harmless; example.py hides them.

load_engine runs the questions of a request one after the other (fork_mode="sequential"); every number in this README was measured that way. fork_mode="auto" lets the engine pick batched mode instead.

Why the original does not fit a 16 GB card

elephant-1 has about 9 billion numbers. Three parts decide how much GPU memory it needs. A token is a word piece; the model knows 248,320 of them.

part	what it does	memory
embedding table	looks up 4,096 numbers for every token of the prompt	1.9 GiB
32 decoder layers	do the thinking and end with 4,096 numbers	12.9 GiB
head (lm_head)	turns those 4,096 numbers into one score for each of the 248,320 tokens	1.9 GiB, or 3.8 GiB in the 32-bit form the engine uses

The engine is OneJev's qev, a Python library that your program imports: it writes the prompt, runs the model through PyTorch and turns the scores into answers. It puts all three parts on the GPU: 12.9 + 1.9 + 3.8 = 18.6 GiB. A 16 GB card gives PyTorch about 15.5 GiB. Even without the 32-bit copy the BF16 weights alone take 16.7 GiB. So the original does not fit. (Calculated from the sizes; I did not run the unmodified loading.)

A workaround that keeps the original code is to compute the head on the CPU. It works and gives the same answers, but costs about 0.2 s every time the model runs, which is why this setup needed 260 ms per short request in my tests.

What I changed

elephant-1 does not write text. It picks an answer: the engine lists the options as A., B., C. and reads only the scores of those letters (it knows 256 labels: A to Z, then AA, AB and so on). So of the head's 248,320 rows the answers use at most 256, which is 0.1 %. The other rows only feed one diagnostic number (allowed_token_mass) that no answer depends on.

Two changes, both in engine_slots.py:

1. A small head. After loading, it copies those 256 rows (4 MiB) out of the model's own head into a small table on the GPU and drops the big one. Nothing is retrained and the weights file is not touched; the small head holds the same numbers.
2. The embedding table stays in RAM. It is only a lookup, so just the looked-up numbers travel to the GPU. The small head alone would not be enough: with the table on the GPU the model would still occupy 12.9 + 1.9 = 14.8 GiB and leave almost no room for the prompt.

Now the GPU holds the decoder layers (12.9 GiB) and a 4 MiB head.

┌────────────────────────────────────────────────────────────────────────────┐
│ prompt, written by the engine:                                             │
│   Options:  A. billing   B. shipping   C. technical                        │
│   Answer with one letter: A, B, C.                                         │
└────────────────────────────────────────────────────────────────────────────┘
                                      │
                                      ▼  tokenizer: text -> token ids
┌────────────────────────────────────────────────────────────────────────────┐
│ THE MODEL (one weights file, never modified): embedding, decoder, head     │
│ embedding table (lookup, CPU)  ->  32 decoder layers (GPU, 12.9 GiB)       │
│ -> hidden state: 4,096 numbers, which go into the head:                    │
└────────────────────────────────────────────────────────────────────────────┘
                                      │
                  ┌───────────────────┴────────────────────┐
                  ▼                                        ▼
┌───────────────────────────────────┐    ┌───────────────────────────────────┐
│ BEFORE: head, all rows used       │    │ AFTER: same head, 256 rows used   │
│ 248,320 rows x 4,096, fp32        │    │ 256 rows x 4,096, fp32            │
│ = 3.8 GiB: does not fit next      │    │ = 4 MiB: sits on the GPU,         │
│ to the decoder on a 16 GB GPU,    │    │ no CPU step. The rows are         │
│ so it runs on the CPU             │    │ copied out of the model's         │
│ (~0.2 s per pass)                 │    │ head when loading                 │
│ -> 248,320 scores; the engine     │    │ -> 256 scores; the engine         │
│ keeps A, B, C, drops the rest     │    │ reads A, B, C directly            │
└───────────────────────────────────┘    └───────────────────────────────────┘
                  │                                        │
                  └───────────────────┬────────────────────┘
                                      ▼
┌────────────────────────────────────────────────────────────────────────────┐
│ softmax over the scores of A, B, C  ->  probabilities  ->  answer          │
│ same answers both ways (largest difference in any probability: 1e-6)       │
│ short requests, median time: 260 ms before -> 72 ms after                  │
└────────────────────────────────────────────────────────────────────────────┘

Measured

RTX 5070 Ti (16 GB), PyTorch 2.15.0, my 420 synthetic router-style requests (629 questions), engine time per request (writing the prompt and running the model; loading not included), median (90th percentile):

setup	short requests (under 200 tokens, 390 of them)	1,000 to 4,000 tokens (17)	over 4,000 tokens (13)	GPU memory, typical / longest request
original engine, head computed on the CPU	260 ms (616)	949 ms	1,625 ms	13.0 / 14.6 GiB
this repo: small head	72 ms (152)	664 ms	1,255 ms	13.0 / 14.6 GiB

Same answers: over all 629 questions the largest difference in any option probability is 1e-6, and no top choice changes. I re-ran this loader on all 420 requests and every answer was bit-identical to my stored results.

Limits

• Long prompts are tight. GPU memory peaks at 13.0 GiB for a typical request and at 14.6 GiB for the longest prompt of my set (9,551 tokens) on a card that offers 15.5 GiB. In one earlier test pass that request ran out of memory once; it passed when repeated and in the final run. For prompts above about 4,000 tokens the FP8 version is safer (peak 8.8 GiB): marcelz/elephant-1-FP8-W8A8.
• RAM. Loading peaks at about 14 GiB of RAM and the process holds about 15 GiB while running (measured). With less than 16 GB of free RAM, use the FP8 version.
• allowed_token_mass is reported as 1.0 and means nothing. It is a diagnostic; nothing else in the response changes.
• The small head can only score the 256 labels, so it cannot generate text. The engine never generates text, so this loses nothing here.
• Pass a local folder as the model path. The engine reads the model's attention setting only from a local config.json; with a hub id it silently runs the wrong attention. install.sh and example.py already do this.
• Tested on one GPU model and on synthetic data written by an AI model (its gold labels may contain errors); not tested on the public Decision Index benchmark. Timings on other cards will differ.
• The engine does not ship an option for this. engine_slots.py wraps the loaded model; it is not part of OneJev.
• PyTorch 2.15.0 is the CUDA 13 build, hence driver 580 or newer. causal-conv1d is not installed (no wheel for this PyTorch); transformers uses its own PyTorch version of that step.

Files

file	what
install.sh	one-command setup: Python 3.12 environment in .venv, pinned packages, OneJev qev at commit 81ce62f with Base Labs' patch (in .runtime), the model from baselabs/elephant-1 (revision b3f8a93) in elephant-1
run.sh	runs a Python file with that environment (clears variables such as PYTHONPATH that some shells export and that can shadow the installed packages)
example.py	the example above, with speed and memory
engine_slots.py	the loader: moves the layers to the GPU, keeps the embedding table in RAM, installs the small head (SlotHead)

More: the FP8 version for 12 GB cards is marcelz/elephant-1-FP8-W8A8. Please cite the original model. elephant-1 is by Base Labs / Baseten (post-trained from Qwen3.5-9B); the engine is OneJev's qev (Apache-2.0).