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