AirLLM

AirLLM is a Python library that runs arbitrarily large LLMs on a single consumer GPU by streaming model layers — and, for MoE models, individual experts — from disk through VRAM rather than keeping the full model resident. It achieves this without quantization, distillation, or pruning by default, though optional 4-bit/8-bit block-wise compression is available via bitsandbytes. The project was extracted from a broader repository called Anima in late 2023; the Python package is airllm and its source lives under air_llm/, while dormant Anima-era directories (rlhf/, anima_100k/, training/) remain in the repo as historical artifacts. Since v3.0.0, streaming is driven by PyTorch forward hooks registered on Transformers' own forward/generate, which is why most modern architectures work through the generic AirLLMBaseModel without per-family subclasses.

The Core runtime section is the heart of the docs — start with Overview and dispatch for the big picture, then AirLLMBaseModel, Layer streaming internals, On-disk splitting and persistence, and Compression for how weights move from disk to GPU and back. The Supported models section documents per-family subclasses — Qwen family, Mixtral and Mistral, ChatGLM, InternLM and Baichuan, and Kimi K3 — each covering the layer-name overrides and quirks needed for that architecture. The Installation and setup section covers Installation and dependencies, the macOS MLX backend for Apple Silicon, Upgrading notes for version bumps, and the Anima legacy background. The Testing section describes the splitter tests (Splitter tests), the Compression test, and the manual GPU streaming test harness; the Release process section documents how versions are cut and published to PyPI.

If you want to understand how AirLLM fits so much model into so little VRAM, read Overview and dispatch and then Layer streaming internals. If you are adding support for a new model architecture, read AirLLMBaseModel first and then one of the existing subclasses in Supported models — Qwen family or Kimi K3 are the most instructive examples. If you are installing AirLLM or debugging an import error, go straight to Installation and dependencies; macOS users should also read macOS MLX backend. If you are contributing changes or cutting a release, read Testing for how to exercise the splitter and streaming paths, and Release process for the tag-and-publish workflow.

Pages in this section:

• Overview and dispatch
• AirLLMBaseModel
• Layer streaming internals
• On-disk splitting and persistence
• Compression

AirLLM streams model layers from disk through VRAM rather than keeping the full model resident, enabling low VRAM usage — dispatch through AutoModel.get_module_class automatically selects the appropriate subclass based on model architecture. On non-macOS platforms, core entry points (AirLLMBaseModel, AutoModel, split_and_save_layers, NotEnoughSpaceException) are exported from air_llm/airllm/__init__.py; on macOS, only AirLLMLlamaMlx and AutoModel are available, with AutoModel.from_pretrained always returning the MLX backend.

AirLLM achieves low VRAM usage by streaming model layers — and, for MoE models, individual experts — from disk through VRAM rather than keeping the full model resident.[1]

On non-macOS platforms, air_llm/airllm/__init__.py exports AirLLMBaseModel, AutoModel, split_and_save_layers, and NotEnoughSpaceException as the core entry points.[2] On macOS (platform == 'darwin'), air_llm/airllm/__init__.py exports only AirLLMLlamaMlx and AutoModel; no PyTorch-based subclasses are imported.[2]

AutoModel in air_llm/airllm/auto_model.py is a dispatch-only factory: it cannot be instantiated directly and raises EnvironmentError if you try.[3] AutoModel.get_module_class is the entry point for dispatch: it accepts a model repo ID or local path and returns a (module, class_name) tuple for the appropriate AirLLM subclass.[4] To load gated models, pass an hf_token kwarg; AutoModel.get_module_class forwards it as the token= argument to AutoConfig.from_pretrained.[3] AutoModel.from_pretrained on macOS always returns an AirLLMLlamaMlx instance, bypassing the architecture-dispatch table entirely — see macOS MLX backend for details.[3]

AutoModel.get_module_class reads the model's config.architectures[0] field to select a class, defaulting to AirLLMBaseModel for any architecture not found in ARCH_OVERRIDES.[3] ARCH_OVERRIDES in auto_model.py maps six architecture strings to dedicated subclasses — ChatGLMModel and ChatGLMForConditionalGeneration → AirLLMChatGLM; QWenLMHeadModel → AirLLMQWen; BaichuanForCausalLM and BaiChuanForCausalLM → AirLLMBaichuan; InternLMForCausalLM → AirLLMInternLM; KimiK3ForConditionalGeneration → AirLLMKimiK3 — and every other standard *ForCausalLM falls through to AirLLMBaseModel without any code change.[3] The dispatch table tested in air_llm/tests/test_automodel.py maps representative repo IDs to concrete class names: garage-bAInd/Platypus2-7B → AirLLMLlama2, Qwen/Qwen-7B → AirLLMQWen, internlm/internlm-chat-7b → AirLLMInternLM, THUDM/chatglm3-6b-base → AirLLMChatGLM, baichuan-inc/Baichuan2-7B-Base → AirLLMBaichuan, mistralai/Mistral-7B-Instruct-v0.1 → AirLLMMistral, mistralai/Mixtral-8x7B-v0.1 → AirLLMMixtral.[4] AirLLMLlama2 in air_llm/airllm/airllm.py is a trivial subclass of AirLLMBaseModel with no overrides; it exists as a named entry point for backward compatibility.[5]

Sources

1. README.md
2. air_llm/airllm/__init__.py
3. air_llm/airllm/auto_model.py
4. air_llm/tests/test_automodel.py
5. air_llm/airllm/airllm.py

AirLLMBaseModel is the base class for AirLLM's model implementations; it handles layer-name mapping, model loading with HF authentication, disk-storage paths, layer sharding, and prefetch/compression settings. During initialization, AirLLMBaseModel applies smart fallbacks: SDPA attention first (falling back to eager), trust_remote_code only when needed, and the model's native dtype before float16, ensuring robustness across architectures. Layer sharding splits a model's weights into separate per-layer files on disk so that only one layer's weights reside in GPU memory at a time; this is the core technique allowing AirLLM to run large models on GPUs with limited VRAM.

The default layer-names dict in AirLLMBaseModel maps to standard Llama-style module paths: model.embed_tokens, model.layers, model.norm, and lm_head; subclasses override set_layer_names_dict for non-standard architectures.[1]

AirLLMBaseModel.__init__ accepts a hf_token parameter for passing a Hugging Face API token at initialization time, required for gated models such as meta-llama/Llama-2-7b-hf.[2] A layer_shards_saving_path keyword argument may be passed at initialization to specify an alternative directory for storing split per-layer shards, defaulting to a path next to the model cache.[2] AirLLMBaseModel.__init__ accepts a delete_original flag; when True, the original downloaded Hugging Face checkpoint is deleted after splitting, retaining only the per-layer shards to save disk space — see On-disk splitting and persistence for splitting details.[1] Prefetching — overlapping next-layer disk load with current-layer GPU compute — is enabled by default and can be disabled by passing prefetching=False.[2] The profiling_mode parameter (default False) can be set to True to emit per-layer time-consumption data during inference.[2] Compression support is initialized via a compression parameter ('4bit' or '8bit'); if bitsandbytes is not installed, an ImportError is raised immediately — see Compression for behavioral details.[1]

The runtime dtype defaults to the model's own config.torch_dtype (typically bfloat16 for modern models) rather than a hardcoded float16; float16 is used only as a last fallback when the config provides no dtype.[1]

AirLLMBaseModel tries trust_remote_code=False first when loading a model config, falling back to trust_remote_code=True only when Transformers does not recognize the architecture — this avoids breakage from vendored remote code (e.g., DeepSeek-V2's modeling_deepseek.py) that calls long-removed Transformers APIs.[1]

init_model in airllm_base.py attempts to build the model with attn_implementation='sdpa'; on ValueError or TypeError (some remote-code architectures don't support SDPA), it falls back to eager attention.[1] _propagate_attn_implementation walks nested PretrainedConfig sub-configs up to depth 2 and copies the chosen attention implementation into each; this is needed for multimodal wrappers like Kimi K3 whose text decoder lives under a text_config sub-config, where an unset value would otherwise fall through to a flash-attention path and fail on machines without flash-attn installed.[1]

Sources

1. air_llm/airllm/airllm_base.py
2. README.md

Layer streaming in AirLLM prefetches layers from disk into pinned host memory via a serialised thread pool, enabling one layer to load while another computes on GPU; oversized layers fall back to pageable memory and the safetensors random-access API minimises per-expert load costs. AirLLM manages GPU and host memory between layer transitions through aggressive cleanup (gc.collect(), malloc_trim(), torch.cuda.empty_cache()) and tracks resource usage via LayeredProfiler, which records minimum free GPU memory and aggregates timing across repeated operations.

The prefetch background thread pool uses exactly one worker (ThreadPoolExecutor(max_workers=1)), serialising disk loads while allowing one layer to load concurrently with the current layer's GPU compute.[1] Pinned (page-locked) host memory per prefetched layer is capped at 2 GB (max_pinned_layer_bytes = 2 * 1024 ** 3); layers larger than this threshold are loaded into ordinary pageable memory instead.[1] Prefetching is automatically disabled when compression is also enabled, because the two features are incompatible in the current implementation — a warning is printed and self.prefetching is set to False.[1]

load_layer_subset in air_llm/airllm/utils.py uses safetensors' random-access API to load only the specified keys from a shard, making per-expert streaming cost only the target expert's bytes rather than the full layer file.[2] layer_tensor_names in air_llm/airllm/utils.py lists tensor names in a layer shard using safetensors' metadata-only access, reading no tensor data.[2]

clean_memory() in air_llm/airllm/utils.py calls gc.collect(), attempts malloc_trim(0) via libc (silently skipped on non-Linux), and then torch.cuda.empty_cache() to reclaim both CPU and GPU memory between layer loads.[2]

LayeredProfiler in air_llm/airllm/profiler.py tracks the minimum free GPU memory ever seen across all add_profiling_time calls, initialising min_free_mem to 1 TiB as a sentinel so the first real reading always wins.[3] LayeredProfiler.print_profiling_time() reports the sum of all recorded times per item, not a per-call breakdown.[3] LayeredProfiler.clear_profiling_time() resets every item's time list to an empty list but preserves the item keys, so previously observed items are not forgotten.[3]

Sources

1. air_llm/airllm/airllm_base.py
2. air_llm/airllm/utils.py
3. air_llm/airllm/profiler.py

On first load, AirLLM decomposes the original model into per-layer shards on disk; subsequent runs stream directly from those shards. Sufficient disk space is required in the HuggingFace cache directory.[1] The air_llm/airllm/persist/ package exports only ModelPersister via its __init__.py, making it the sole public symbol of the persist subpackage.[2]

SafetensorModelPersister.persist_model in air_llm/airllm/persist/safetensor_model_persister.py writes the weight data first, then touches a .safetensors.done marker file to signal that the shard is complete and safe to load.[3] SafetensorModelPersister.model_persist_exist requires BOTH a .safetensors file and a .safetensors.done marker to exist before treating a layer shard as complete — a write-then-marker pattern that guards against partial writes.[3] MlxModelPersister.model_persist_exist in air_llm/airllm/persist/mlx_model_persister.py mirrors this contract: it checks for both a .mlx.npz weight file and a .mlx.done marker before treating the shard as complete.[4] MlxModelPersister.persist_model casts all tensors to float16 before saving them as NumPy .npz files — weights are not preserved in their original dtype on macOS.[4]

link_or_copy_file in air_llm/airllm/utils.py always resolves the source to its real path via os.path.realpath before linking, because HuggingFace cache files are stored as symlinks into a blob directory.[5] When creating a shard reference, link_or_copy_file tries a hard link first, then a symlink, then a full copy — hard links are preferred because they keep the data alive even if the original checkpoint file is later deleted, and cost no extra disk space.[5] This linking strategy has practical significance at scale: for Kimi K3's 1.56 TB checkpoint, a naive split would require 3.12 TB, but because K3's shards are pure single-module files, split layers are hard-linked to the originals instead of copied — see Kimi K3 for further detail.[6]

Sources

1. README.md
2. air_llm/airllm/persist/__init__.py
3. air_llm/airllm/persist/safetensor_model_persister.py
4. air_llm/airllm/persist/mlx_model_persister.py
5. air_llm/airllm/utils.py
6. github.com/lyogavin/airllm/releases/tag/v3.1.0

Block-wise quantization (4-bit or 8-bit) is enabled by passing compression='4bit' or compression='8bit' to AutoModel.from_pretrained; it requires bitsandbytes and airllm ≥ 2.0.0 — see Installation and dependencies for setup details.[1] AirLLM's block-wise compression quantizes weights only — not activations — because the bottleneck is disk loading rather than matrix-multiply throughput, which makes accuracy loss easier to control.[1]

split_and_save_layers in air_llm/airllm/utils.py raises an AssertionError if a compression argument is passed without bitsandbytes installed.[2] When compression is active, split_and_save_layers appends the compression type to the shard directory name (e.g., splitted_model.4bit), keeping compressed and uncompressed shards in separate directories.[2] check_space in air_llm/airllm/utils.py adjusts the estimated model size before the disk-space check: 4bit compression scales the raw byte count by 1/0.2813 (~3.55×), while 8bit halves it.[2]

Sources

1. README.md
2. air_llm/airllm/utils.py

Pages in this section:

• Mixtral and Mistral
• Qwen family
• ChatGLM
• InternLM and Baichuan
• Kimi K3

AirLLMMistral and AirLLMMixtral are minimal subclasses of AirLLMBaseModel that disable BetterTransformer and use default generation configs, deferring all other behavior to the base class. Both models avoid the optimum-based BetterTransformer in favor of Transformers' built-in sdpa acceleration. SDPA (Scaled Dot-Product Attention) is a fused attention kernel built into PyTorch that delivers acceleration comparable to BetterTransformer without requiring the external optimum library.

AirLLMMistral (in air_llm/airllm/airllm_mistral.py) is a minimal subclass of AirLLMBaseModel that only disables BetterTransformer and returns a bare GenerationConfig(); all other behavior is inherited from the base class — see AirLLMBaseModel for the shared interface.[1] AirLLMMixtral (in air_llm/airllm/airllm_mixtral.py) explicitly disables BetterTransformer by returning False from get_use_better_transformer(), reflecting the drop of the optimum-based BetterTransformer dependency in favour of built-in sdpa.[2] AirLLMMixtral returns a bare GenerationConfig() (all defaults) from get_generation_config(), relying on Transformers' own generation defaults without any Mixtral-specific overrides.[2]

Sources

1. air_llm/airllm/airllm_mistral.py
2. air_llm/airllm/airllm_mixtral.py

AirLLMQWen and AirLLMQWen2 are thin AirLLMBaseModel subclasses that implement AirLLM's streaming inference protocol for the Qwen family by mapping layer names, managing rotary embeddings, and packing KV caches in Qwen's layer_past convention.

AirLLMQWen (in air_llm/airllm/airllm_qwen.py) and AirLLMQWen2 (in air_llm/airllm/airllm_qwen2.py) are both thin subclasses of AirLLMBaseModel — see AirLLMBaseModel — that disable BetterTransformer/optimum by returning False from get_use_better_transformer().[1][2] AirLLMQwen3_5 (in air_llm/airllm/airllm_qwen3_5.py) is a thin subclass of AirLLMBaseModel that implements AirLLM's streaming inference protocol for the Qwen3.8-27B dense vision-language model — the first Qwen3.x VL variant in the codebase. The Qwen3.8-27B model (served via AirLLMQwen3_5) runs on an RTX 3090 GPU consuming 3.33 GB of GPU memory through AirLLM's layer-streaming inference. AirLLMQwen4Exp (in air_llm/airllm/airllm_qwen4_exp.py) is a thin subclass of AirLLMBaseModel that implements AirLLM's streaming inference protocol for the Qwen3.8-Flash-Next model. The Qwen3.8-Flash-Next model (served via AirLLMQwen4Exp) runs on consumer hardware consuming as little as 5.95 GB of GPU memory through AirLLM's layer-streaming inference.

AirLLMQWen maps its layer names to QWen's transformer block naming scheme via set_layer_names_dict: embed → transformer.wte, layer_prefix → transformer.h, norm → transformer.ln_f, lm_head → lm_head.[2]

AirLLMQWen.get_generation_config returns a bare GenerationConfig() with no special settings.[2]

AirLLMQWen.get_pos_emb_args computes rotary_emb._ntk_alpha_cached_list on the model's transformer and returns a rotary_pos_emb_list argument dict used when passing positional embeddings to each QWen layer during streaming.[2] AirLLMQWen.get_past_key_value_args packs the KV cache as {'layer_past': (k_cache, v_cache)}, matching QWen's layer_past parameter convention (distinct from ChatGLM's kv_cache convention, covered on ChatGLM).[2] AirLLMQWen.get_past_key_values_cache_seq_len reads the cached sequence length from axis 1 of the key tensor (past_key_values[0][0].shape[1]), unlike ChatGLM which reads from axis 0.[2]

AirLLMQWen.get_attention_mask_args always passes attention_mask=None to QWen layers, and get_position_ids_args returns an empty dict — QWen's attention and position handling is driven entirely by rotary embeddings, not explicit masks or position IDs.[2]

test_qwen3_8_split.py serves as a template for validating on-disk splitting behaviour for large models. airllm_qwen4_exp.py and test_qwen38_flash_next_split.py together serve as a reference pattern for extending AirLLM to new Qwen-lineage models via the layer-streaming and on-disk-split pipeline.

Sources

1. air_llm/airllm/airllm_qwen2.py
2. air_llm/airllm/airllm_qwen.py

AirLLMChatGLM is an adapter that maps AirLLM's streaming interface to ChatGLM's module structure and tensor layout, handling layer name routing, sequence positioning, rotary embeddings, and KV cache packing.

AirLLMChatGLM maps its layer names to ChatGLM's module structure via set_layer_names_dict, with keys: embed → transformer.embedding.word_embeddings, layer_prefix → transformer.encoder.layers, norm → transformer.encoder.final_layernorm, lm_head → transformer.output_layer, and an additional rotary_pos_emb key → transformer.rotary_pos_emb.[1]

ChatGLM uses a [seq_len, batch, heads, dim] tensor layout: AirLLMChatGLM.get_sequence_len reads sequence length from axis 0 (seq.shape[0]), and get_past_key_values_cache_seq_len likewise reads from axis 0 of the key tensor (past_key_values[0][0].shape[0]).[1]

AirLLMChatGLM.get_pos_emb_args computes rotary embeddings over the full configured seq_length, slices the result to the current sequence length, and transposes to [seq_len, 1, ...] layout before returning them under the rotary_pos_emb key.[1] AirLLMChatGLM.get_attention_mask_args always passes attention_mask=None, and get_position_ids_args returns an empty dict — ChatGLM relies on rotary embeddings and does not use explicit masks or position IDs during AirLLM layer streaming.[1]

AirLLMChatGLM.get_past_key_value_args packs the KV cache as {'kv_cache': (k_cache, v_cache)}, matching the kv_cache parameter name required by ChatGLM's layer forward signature.[1]

Sources

1. air_llm/airllm/airllm_chatglm.py

AirLLMInternLM and AirLLMBaichuan are model-specific subclasses of AirLLMBaseModel that configure InternLM and Baichuan models for streaming inference, with Baichuan using a vendored tokenizer to avoid upstream HuggingFace bugs. These classes inherit shared streaming logic from the base and override only generation config and tokenizer loading to match each model's requirements.

AirLLMInternLM, defined in air_llm/airllm/airllm_internlm.py, is a minimal subclass of AirLLMBaseModel — BetterTransformer is disabled and get_generation_config returns a bare GenerationConfig(), with all streaming logic inherited from the base class.[1]

AirLLMBaichuan imports BaichuanTokenizer from the vendored local module .tokenization_baichuan rather than from transformers, specifically to avoid a known upstream HuggingFace Hub tokenizer bug for Baichuan2.[2] AirLLMBaichuan.get_tokenizer loads that vendored BaichuanTokenizer with use_fast=False and trust_remote_code=True from the local model path, working around the bug tracked in the Baichuan2-7B-Base Hub discussion.[2]

Sources

1. air_llm/airllm/airllm_internlm.py
2. air_llm/airllm/airllm_baichuan.py

AirLLMKimiK3 adapts streaming inference for K3's decoder architecture by pointing layer references through a language_model prefix and streaming its 896-expert MoE layers per-expert rather than per-layer, keeping fixed modules like vision and projection in resident memory. K3's MXFP4 quantization and dense expert structure demand specific dependencies: flash-attn, CUDA 12, and transformers 4.56.x, with compressed-tensors support for packed weight expansion.

AirLLMKimiK3, implemented in air_llm/airllm/airllm_kimi_k3.py, overrides set_layer_names_dict to point embedding, layers, norm, and lm_head under the language_model prefix, because K3's decoder lives one level deeper than standard *ForCausalLM models.[1] Four modules — output_attn_res_norm, output_attn_res_proj, mm_projector, and vision_tower — are listed under the resident key, meaning they are loaded once and kept in memory rather than streamed, because together they are well under 1 GB; this also ensures they are split out during on-disk splitting even though they are not decoder layers.[1][2] K3 uses per-expert streaming (keyed by expert_prefix) rather than per-layer streaming: each layer holds 896 experts but a token needs only ~1 GB of them, so streaming by expert avoids loading the full ~55 GB layer.[1] Kimi K3 uses MXFP4 weights that cross PCIe packed and expand on the GPU, moving 4× less data than full-precision weights.[3]

Kimi K3 (2.8T) requires three steps beyond a standard pip install airllm: running pip install compressed-tensors flash-attn (K3's model code mandates flash attention regardless of what you request), using a CUDA 12 build of torch (no prebuilt flash-attn wheel exists for CUDA 13), and pinning transformers to 4.56.x (K3's remote code does not load on 5.x) — see Installation and dependencies for general dependency guidance.[4]

Sources

1. air_llm/airllm/airllm_kimi_k3.py
2. air_llm/tests/test_kimi_k3_split.py
3. github.com/lyogavin/airllm/releases/tag/v3.1.0
4. README.md

Pages in this section:

• macOS MLX backend
• Installation and dependencies
• Upgrading
• Anima legacy

The AirLLMLlamaMlx backend runs Llama-family models natively on Apple Silicon Macs using MLX instead of PyTorch, handling checkpoint sharding, RoPE tuning, and greedy/sampled decoding entirely in MLX. Setup requires both mlx and torch installed on Apple Silicon hardware; the backend auto-splits model checkpoints on first use and tracks RAM consumption across forward passes. MLX is Apple's open-source array framework optimized for Apple Silicon's unified memory architecture, enabling tensor operations to run on the GPU or Neural Engine without copying data across separate memory pools.

AirLLMLlamaMlx, defined in air_llm/airllm/airllm_llama_mlx.py, is the macOS/MLX backend for Llama-family models, importing mlx.core and mlx.nn in place of PyTorch for all forward computation.[1] macOS support requires both mlx and torch to be installed, and only Apple Silicon hardware is supported — Intel Macs are not.[2]

AirLLMLlamaMlx.__init__ calls find_or_create_local_splitted_path to split the checkpoint into per-layer shards on first use, sharing the same on-disk splitting step as the CUDA backend — see On-disk splitting and persistence.[1] AirLLMLlamaMlx defines a fixed layer-name mapping via set_layer_names_dict(): embed → model.embed_tokens, layer_prefix → model.layers, norm → model.norm, and lm_head → lm_head; subclasses can override this mapping to support other architectures.[1]

sanitize_config in air_llm/airllm/airllm_llama_mlx.py defaults rope_theta to 10000 when the key is absent from the config dict, matching the RoPE base used by the original LLaMA.[1] sanitize_config also defaults n_kv_heads to the full n_heads count when the key is absent, enabling GQA-unaware configs to function with standard MHA semantics.[1] get_model_args_from_config hardcodes rope_traditional=False, so the MLX backend always uses the non-traditional (non-GPT-J) RoPE rotation order regardless of what the model config specifies.[1]

AirLLMLlamaMlx uses psutil.virtual_memory() — not GPU memory — to measure available RAM, tracking consumed and peak-consumed memory when show_memory_util=True.[1] The sample function uses greedy decoding (mx.argmax) when temperature=0 and categorical sampling otherwise; temperature=0 is the default for both generate and model_generate.[1]

Sources

1. air_llm/airllm/airllm_llama_mlx.py
2. README.md

AirLLM declares a lean set of required dependencies—torch, transformers, accelerate, and tokenization/serialization tools—while bitsandbytes and compressed-tensors remain optional so a basic install avoids unnecessary overhead. Per-family model classes import defensively inside try/except blocks, ensuring a missing optional dependency for one model never breaks the whole package or prevents generic streaming.

The airllm package (v3.1.0) declares torch>=2.4, transformers>=4.49,<5.13, accelerate>=1.0, safetensors, huggingface-hub, scipy, sentencepiece, and tqdm as its mandatory install_requires in setup.py.[1] sentencepiece was added as a declared mandatory dependency after its absence caused import airllm to crash on clean installs, because the Baichuan tokenizer is imported eagerly at package load time.[1] bitsandbytes and compressed-tensors are intentionally absent from install_requires: bitsandbytes is needed only for compression mode, and compressed-tensors only for MXFP4 checkpoints (Kimi K3), so a plain pip install airllm produces a lean, known-good stack.[1]

In air_llm/airllm/__init__.py, per-family subclasses (AirLLMChatGLM, AirLLMBaichuan, etc.) are imported inside a try/except loop so that a missing optional dependency for one model family never prevents the package from importing; a warning names the unavailable class and notes that the generic streaming path still works.[2]

requirements.txt is a legacy file from the repository's earlier Anima training/RLHF era and does NOT reflect the install_requires of the published airllm package.[3] That file installs transformers from the Hugging Face GitHub HEAD rather than a PyPI release, pins accelerate to the v0.20.3 tag and peft to the v0.3.0 tag from their respective GitHub repositories.[3] Fixed-version entries in requirements.txt include bitsandbytes==0.39.0 (optional 4-bit/8-bit compression), einops==0.6.1, evaluate==0.4.0, scikit-learn==1.2.2, sentencepiece==0.1.99, and wandb==0.15.3.[3] The requirements.txt file is preserved for historical reference only; users installing AirLLM for inference should rely on setup.py via pip install airllm to obtain the correct, tested dependency set.

Sources

1. air_llm/setup.py
2. air_llm/airllm/__init__.py
3. requirements.txt

In v3.0.1, import airllm failed on a clean install with ModuleNotFoundError: No module named 'sentencepiece' because the eagerly-imported Baichuan tokenizer needs it; sentencepiece is now a declared dependency, so pip install airllm works out of the box.[1] Also introduced in v3.0.1, per-model-family imports in __init__ are now defensive: a missing optional dependency for one niche model family only warns instead of breaking the whole package, and the core AutoModel/AirLLMBaseModel path always loads.[1]

Sources

1. github.com/lyogavin/airllm/releases/tag/v3.0.1

Anima is AirLLM's predecessor, a 33B instruction-tuned model built by fine-tuning Guanaco with QLoRA on a single H100 for 10,000 steps—a sweet-spot between training cost and model quality. The trained adapter weights and merged model are published separately on HuggingFace; training can be reproduced locally or across multiple A100/H100 GPUs via Hugging Face Accelerate. QLoRA (Quantized Low-Rank Adaptation) is a fine-tuning technique that quantizes base model weights while training only small low-rank adapter layers, drastically reducing GPU memory usage and enabling large-model fine-tuning on limited hardware. A Peft adapter stores only the fine-tuned weight differences (low-rank deltas) rather than a full model copy, enabling compact distribution of trained modifications.

The Anima model — the precursor project in this repository — is based on QLoRA fine-tuning of the 33B Guanaco model (timdettmers/guanaco-33b), trained on a single H100 GPU for 10,000 steps.[1] The choice of 10,000 steps reflects a finding from the QLoRA paper that more training samples are not always better, and that 10,000 steps offers a relatively good ROI.[1] Two HuggingFace repositories serve the Anima model: lyogavin/Anima33B (Peft adapter weights only) and lyogavin/Anima33B-merged (full merged standalone model).[1] Anima 33B training can be reproduced by installing dependencies and running ./run_Amina_training.sh, tested on a single 80 GB H100 or dual 40 GB A100 setup.[1] Multi-GPU training is supported out-of-the-box via Hugging Face Accelerate, verified on 2×A100 40 GB, where the training script runs seamlessly.[1] lyogavin/Anima33B-merged folds the Peft adapter deltas back into the base model weights, enabling dependency-free inference without needing the adapter weights separately.

Sources

1. README_ja.md

Pages in this section:

• Splitter tests
• GPU streaming test harness
• Compression test

split_and_save_layers is AirLLM's checkpoint splitting routine; tests verify it preserves every tensor bit-for-bit while hard-linking single-module shards and materializing real files when multiple modules share a shard. The test suite checks that splitting handles both uniform module-per-shard layouts (hard-link eligible) and complex real-world checkpoints with shared shards, including out-of-order residents and packed 4-bit dtypes that must not be corrupted.

air_llm/tests/test_kimi_k3_split.py bypasses airllm/__init__.py at import time to avoid pulling in the MLX backend and full Transformers stack, making the splitter tests runnable on any platform without those heavy dependencies.[1] The test file pins SafetensorModelPersister explicitly so the splitter tests exercise the Linux/CUDA code path regardless of the platform running the tests.[1]

When a Kimi K3 checkpoint has exactly one module per shard, split_and_save_layers hard-links the shard file instead of copying it, so a 1.56 TB checkpoint does not need double the disk space; this is verified by inode equality in test_kimi_k3_split.py.[1] When multiple modules share a single shard — embed, norm, lm_head, and residual norms in K3 — split_and_save_layers must materialise each module as a separate real file rather than hard-linking; the test asserts the inode differs from the source shard and that no foreign tensors leak into the per-module file.[1] For a standard checkpoint where multiple modules share each shard, split_and_save_layers must likewise write real per-module files; TestStandardCheckpointStillSplits verifies that split-file inodes differ from all source-shard inodes.[1]

split_and_save_layers must be bit-for-bit lossless: every tensor in the original checkpoint must appear in the split output with identical dtype and values; test_kimi_k3_split.py verifies this for both the K3 and standard layouts.[1] Packed 4-bit (uint8) tensor dtypes must be preserved verbatim through the split path; test_kimi_k3_split.py asserts that MXFP4 weight_packed tensors remain torch.uint8 after splitting, since any dtype cast would silently corrupt K3 weights.[1] The K3 fake checkpoint in test_kimi_k3_split.py deliberately places the projector shard before the vision-tower shard — out of the order they appear in the resident list — to ensure split_and_save_layers handles out-of-order resident modules without losing tensors.[1]

Sources

1. air_llm/tests/test_kimi_k3_split.py

air_llm/tests/test_streaming_gpu.py is a manual GPU test harness for AirLLM layer-streaming inference; it is not an automated pytest suite and must be invoked explicitly.[1]

The harness can cap the visible VRAM for the test process via torch.cuda.set_per_process_memory_fraction, emulating a smaller GPU card without physically needing one.[1] When the --max-vram-gb value equals or exceeds the device's total memory, the harness skips capping and prints a warning instead of applying the fraction.[1]

The harness invokes AirLLM via AutoModel.from_pretrained with compression and delete_original forwarded from CLI flags, and calls model.generate with do_sample=False (greedy decoding) and return_dict_in_generate=True.[1] --compression accepts only None, "4bit", or "8bit", corresponding to the compression options exposed by AutoModel.from_pretrained — see Compression for details on those modes.[1] The default prompt is "The capital of France is" and the default number of new tokens is 12.[1]

Peak VRAM usage (in MB) is reported via torch.cuda.max_memory_allocated(), with the peak counter reset immediately before the model.generate() call so only inference is measured.[1]

When --compare is passed, the harness runs the full-load Transformers reference before capping VRAM so the complete model fits on the GPU; only then is cap_vram applied for the AirLLM run.[1] The reference full-load run loads the model with dtype=torch.float16 and is documented as only feasible for small models.[1] After the reference run completes, the harness explicitly deletes the reference model and calls torch.cuda.empty_cache() to free VRAM before the AirLLM run starts.[1] --compare mode asserts token-level output equality between the AirLLM streaming run and the full-load Transformers run, exiting with code 1 on mismatch.[1]

How to run the GPU streaming test for a tiny model and verify output matches a full-load reference:

python test_streaming_gpu.py --model TinyLlama/TinyLlama-1.1B-Chat-v1.0 --compare

[1]

How to emulate a 4 GB GPU card running a 7B model with air_llm/tests/test_streaming_gpu.py:

python test_streaming_gpu.py --model Qwen/Qwen2.5-7B-Instruct --max-vram-gb 4

[1]

The canonical usage of AirLLMLlama2 in air_llm/inference_example.py accepts either a Hugging Face repo ID or a local model path as the first argument:

model = AirLLMLlama2("garage-bAInd/Platypus2-70B-instruct")
# or
model = AirLLMLlama2("/home/ubuntu/.cache/huggingface/hub/...")

[2]

The end-to-end inference pattern in air_llm/inference_example.py tokenizes with return_attention_mask=False, passes input_ids to .cuda(), calls .generate() with use_cache=True and return_dict_in_generate=True, then decodes generation_output.sequences[0]:

input_tokens = model.tokenizer(input_text, return_tensors="pt",
    return_attention_mask=False, truncation=True,
    max_length=MAX_LENGTH, padding=True)
generation_output = model.generate(
    input_tokens['input_ids'].cuda(),
    max_new_tokens=2, use_cache=True, return_dict_in_generate=True)
output = model.tokenizer.decode(generation_output.sequences[0])

[2]

Sources

1. air_llm/tests/test_streaming_gpu.py
2. air_llm/inference_example.py

The compression test in air_llm/tests/test_compression.py verifies that compress_layer_state_dict / uncompress_layer_state_dict achieves lossless round-trip for None compression and RMSE < 0.1 for '4bit' and '8bit' modes on float16 tensors — see Compression for the underlying API and behavior.[1] Test tensors are randomly generated with shape (32, 128) in float16 dtype and placed on CUDA, so a CUDA-capable GPU is required to run the suite.[1]

Sources

1. air_llm/tests/test_compression.py

AirLLM's release process is triggered by a GitHub Release or manual workflow dispatch; it validates the version tag matches setup.py, copies the top-level README into the package, and publishes to PyPI using OIDC Trusted Publishing without stored credentials. The workflow gates publishing to a protected PyPI environment that enforces review and validation steps, ensuring only intentional, correctly-versioned releases reach users.

A release is triggered by publishing a GitHub Release; .github/workflows/release.yml also supports a manual workflow_dispatch with a dry_run boolean input (default true) that builds but skips publishing to PyPI.[1]

To cut a release, contributors must: (1) bump version= in air_llm/setup.py and commit, then (2) create a GitHub Release whose tag matches that version (e.g. "3.0.1" or "v3.0.1").[1] On a real GitHub Release, .github/workflows/release.yml enforces a version guard: the workflow extracts the version from air_llm/setup.py and compares it to the release tag (stripping a leading v), failing the build with an error if they disagree.[1]

The release workflow uses Python 3.11 for both building and packaging.[1] Before building, .github/workflows/release.yml copies the top-level README.md into air_llm/README.md so the distributed package always contains the current top-level readme.[1] Distribution metadata is validated with twine check dist/* before uploading, catching packaging problems before they reach PyPI.[1]

.github/workflows/release.yml publishes the airllm package to PyPI using PyPI Trusted Publishing (OIDC), so no API tokens need to be stored in the repository.[1] The publish job requires the id-token: write permission, which is mandatory for PyPI Trusted Publishing via OIDC.[1] Publishing is gated by the pypi GitHub environment and only runs on a real GitHub Release or an explicit non-dry-run manual dispatch, ensuring the environment's protections (e.g. required reviewers) are always applied.[1]

Sources

1. .github/workflows/release.yml