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