best_engine_ai_helper.streaming module

streaming — run a model that does NOT fit in memory, with no airllm.

A transformer never needs all of its layers resident at once. Decoding reads layer 1, then layer 2, then layer 3; only the layer being computed has to be in memory. Keep one decoder layer resident, read the rest off disk as they are reached, and the memory a model needs stops scaling with its parameter count and starts scaling with its widest single layer.

This module is the runtime that does that. best_engine_ai_helper.offload is the part that decides whether to, sizes it and prices it; this is the part that actually runs it.

Why it is written here rather than imported

It used to be airllm. That dependency was removed, and the reasons were measured rather than assumed:

1. It did not work on Apple Silicon at all. airllm’s AutoModel.from_pretrained dispatches on architecture on every platform except macOS, where it hard-returns a hand-written Llama transformer in MLX whose projections are all bias=False. Qwen, Gemma and Mistral carry attention biases; Llama 3.2 ties its embeddings and has no lm_head tensor to load; older exports persist rotary_emb.inv_freq, which that class has no parameter for. Each of those is a different crash, and the platform this suite is developed on hit all three.

2. It wrote a second copy of the checkpoint to disk. Its splitter materialises one file per layer before anything runs, so streaming a 70B model needed ~294 GB free on top of the download. That is what made almost every row of airllm plan read runs: NO on a normal laptop — the models fit in memory-per-layer terms and failed on scratch space.

3. It carried MLX, bitsandbytes and a dozen per-family subclasses for architectures nobody here asked for.

What this does instead

The technique is the same one airllm’s generic (non-macOS) path uses, and it is the right one: build the model on the meta device so it costs nothing, then hang a forward pre-hook on every decoder layer that materialises that layer’s weights just before it runs, and a post-hook that returns it to meta afterwards. transformers owns the forward pass, so the architecture is whatever transformers says it is — which is why this works for Qwen and Gemma without a line of per-family code.

Three deliberate differences from that design, all of them things this project needs and that one did not do:

• No pre-split, and therefore no scratch space. safetensors is a memory-mapped format with a tensor-level index, so one layer’s weights can be read straight out of the original shards. Nothing is copied, nothing is written, and the disk requirement for streaming drops from “the checkpoint again” to zero.

• Apple Silicon is torch on MPS, not MLX. torch has supported Metal for years; routing through a second framework is what cost airllm every non-Llama architecture on this platform.

• The decoder stack is discovered, not named. airllm keeps a per-family table of module paths (model.layers here, transformer.h there). This looks for the ModuleList whose length matches the config’s layer count, which is a property every causal LM has and no table has to be kept current.

What it does NOT do

Said plainly, because a runtime that pretends is worse than one that refuses:

• safetensors only. A checkpoint published only as pytorch_model.bin is refused by name. Pickle has no tensor-level index, so “read one layer” would mean reading the whole file per layer.

• No per-expert MoE streaming. A sparse mixture-of-experts layer is materialised whole, so the resident figure for one is its full expert set rather than the handful a token routes to.

• No weight quantisation on the fly. The checkpoint is run at the dtype it ships in (cast to the compute dtype). Quantise before streaming if the layer itself is too wide.

Author

Warith Harchaoui <warithmetics@deraison.ai>

best_engine_ai_helper.streaming.DOWNLOAD_PATTERNS = ('*.safetensors', '*.safetensors.index.json', '*.json', '*.model', '*.txt')

weights, the config that builds the model, and the tokenizer. Nothing else – no ONNX exports, no duplicate .bin mirrors of tensors already present as safetensors, no TensorFlow copies.

Type:

What a snapshot download needs

class best_engine_ai_helper.streaming.StreamedModel(model, *, device=None, dtype=None, prefetching=True, hf_token=None, trust_remote_code=False, max_seq_len=None)

Bases: object

A causal LM whose decoder layers live on disk until the moment they run.

Built on the meta device, so constructing one costs no memory whatever the parameter count. Each decoder layer is materialised by a forward pre-hook and returned to meta by a post-hook, so the resident footprint is everything-outside-the-decoder-stack (embeddings, final norm, head) plus one layer.

Parameters:

• model (str) – Local path or Hugging Face repository id.

• device (str or None) – torch device; None picks default_device().

• dtype (torch.dtype or None) – Compute dtype; None picks default_dtype().

• prefetching (bool) – Read the next layer on a worker thread while the current one computes. Disk is the binding constraint here, so this is the main speed lever that costs nothing but one layer of extra host memory.

• hf_token (str or None) – Token for a gated repository.

• trust_remote_code (bool) – Passed to transformers for architectures that ship their own modelling code. Off by default: it executes code from the repository.

• max_seq_len (int | None)

model

The underlying model, with hooks installed.

Type:

transformers.PreTrainedModel

tokenizer

Type:

transformers.PreTrainedTokenizerBase

device

Type:

str

layers

The streamed decoder layers, in execution order.

Type:

list

Examples

>>> StreamedModel("meta-llama/Llama-3.3-70B-Instruct")

close()

Stop the prefetch thread. Idempotent.

Return type:

None

forward(*args, **kwargs)

Delegate to transformers.

Parameters:

• args (Any)

• kwargs (Any)

Return type:

Any

generate(*args, **kwargs)

Delegate to transformers, which owns the forward pass.

Parameters:

• args (Any)

• kwargs (Any)

Return type:

Any

tokenizer: Any | None

best_engine_ai_helper.streaming.default_device()

The best torch device on this machine for streamed decoding.

MPS is named explicitly rather than left to a cuda-or-cpu guess: that guess is what sent Apple Silicon down a separate framework and cost it every architecture that was not Llama.

Returns:

'cuda', 'mps' or 'cpu'.

Return type:

str

Examples

>>> default_device() in ("cuda", "mps", "cpu")
True

best_engine_ai_helper.streaming.default_dtype(device)

The compute dtype to materialise streamed weights in.

Half precision everywhere there is an accelerator, because the whole point is to halve what crosses the disk-to-memory boundary per layer. CPU stays in float32: half precision on CPU is emulated and slower, and a machine streaming on CPU is already paying enough.

Parameters:

device (str) – Output of default_device().

Return type:

torch.dtype

Examples

>>> import torch
>>> default_dtype("cpu") is torch.float32
True
>>> default_dtype("mps") is torch.float16
True

best_engine_ai_helper.streaming.resolve_checkpoint(model, *, hf_token=None)

The local directory holding a checkpoint, downloading it if needed.

Parameters:

• model (str) – A local path, or a Hugging Face repository id.

• hf_token (str or None) – Token for a gated repository.

Returns:

Directory containing the safetensors shards and the config.

Return type:

pathlib.Path

Examples

>>> resolve_checkpoint("meta-llama/Llama-3.3-70B-Instruct")

best_engine_ai_helper.streaming.tensor_index(checkpoint)

Map every tensor name in a checkpoint to the file that holds it.

This is what makes the pre-split unnecessary. A sharded checkpoint already ships model.safetensors.index.json, which is exactly the “which file holds this tensor” table a splitter would otherwise rebuild by copying every byte.

Parameters:

checkpoint (pathlib.Path) – Directory from resolve_checkpoint().

Returns:

Tensor name -> the .safetensors file containing it.

Return type:

dict[str, pathlib.Path]

Raises:

FileNotFoundError – When the checkpoint ships no safetensors at all, naming the reason.

Examples

>>> tensor_index(Path("/models/llama"))