AirLLM выпустил три версии подряд: v3.1.0 (29 Jul 01:08) с Kimi K3 на 2.8T параметров, v3.2.0 (18 Aug 23:15) с Qwen3.8-27B и v3.3.0 (28 Aug 02:11) с Qwen3.8-Flash-Next. Заявленные пики VRAM при генерации — 3.72 GB, 3.33 GB и 5.95 GB соответственно. Одновременно сменился публичный API: вместо класса под конкретную архитектуру — фабрика, которая сама определяет архитектуру по конфигу и подбирает подходящий класс модели, а поверх стриминга появился LoRA-трейнинг.
Ниже — что именно изменилось в механике загрузки весов, откуда взялись требования к окружению и что это значит для тех, кто запускает такие модели на потребительском железе.
Что изменилось в API
Прямой класс для Llama-моделей заменён фабрикой: в примере явно указано from airllm import AutoModel #<----- instead of AirLLMLlama2, дальше модель создаётся тем же привычным способом, что и раньше. Фабрика определяет архитектуру модели и по словарю ARCH_OVERRIDESтаблица соответствия «имя архитектуры из конфига модели — класс реализации в AirLLM» выбирает класс; если совпадения нет, берётся базовая реализация. Так ChatGLM, QWen, Baichuan и InternLM получают собственные подклассы, а всё остальное — базовую реализацию.
Параметр compression принимает '4bit', а для 8-битного блочного квантования — '8bit'; по умолчанию сжатия нет.
Для дообучения добавлены два потоковых LoRA-тренера. Первый предназначен для Qwen3.5/Qwen3.8 dense VL (только текст) и описан как тот же API для Qwen/Qwen3.8-27B. Второй — для Qwen3.8-Flash-Next (125B MoE + 51B PLE); его шаг обучения выполняет один SFTsupervised fine-tuning, дообучение на парах «вход — эталонный ответ»-шаг и возвращает loss.
Как устроена послойная загрузка
Основной механизм прежний: модель не помещается в память целиком, а слои подкачиваются по одному. move_layer_to_device принимает state_dict слоя, восстанавливает обычные весовые модули, распаковывает state_dict и для каждого параметра решает, куда его положить. Часть весов остаётся на CPU на всё время работы процесса — например, вложенная таблица n-gram модели Flash-Next, которая лежит под слоем декодера, но на GPU попадать не должна. Такие параметры регистрируются отдельно и при переносе слоя на карту не трогаются. Если есть квантователькомпонент, преобразующий веса в квантованный формат при загрузке, параметр восстанавливается через create_quantized_param; иначе значение приводится к нужной форме и размещается на running_device — либо дословно, либо с приведением к running_dtype. Функция возвращает список перемещённых имён, чтобы потом было что выгружать.
Сама модель строится на meta-устройствевиртуальном устройстве PyTorch, на котором тензоры существуют без выделения памяти. _instantiate_on_meta создаёт модель transformers внутри init_empty_weights(include_buffers=False), пробуя фабрики из _auto_model_classes и проверяя соответствие именам слоёв; init_model сначала запрашивает реализацию внимания sdpa, при ошибке откатывается на eager. После выполнения слоя _post_hook освобождает память: при наличии квантователя, экспертного стриминга или cpu-resident параметров он возвращает на meta только те тензоры, что сам разместил, иначе выгружает весь модуль на meta-устройство.
Стриминг экспертов MoE
Разреженный MoE-слой содержит сотни экспертов, но каждый токен маршрутизируется лишь к нескольким. Для Kimi K3 эксперты одного слоя в развёрнутом виде занимают ~55GB, а токен трогает ~1GB. Поэтому вместо целого слоя стримятся отдельные эксперты.
Стриминг включается в _setup_expert_streaming только при трёх условиях: в layer_names_dict есть expert_prefix, хранилище весов — персистеробъект, который отвечает за чтение и раскладку весов чекпоинта формата safetensors, сжатие отключено. Ключи тензоров делятся по маркеру '.{expert_prefix}.' на «прочие» и по номеру эксперта (номер обязан быть цифрой). На каждый модуль эксперта вешаются pre- и post-хуки. Pre-hook по паре (индекс слоя, индекс эксперта) загружает ключи только этого эксперта и переносит их на устройство, запоминая перемещённое; post-hook возвращает эти параметры на meta и очищает список.
Отдельный жизненный цикл нужен потому, что при активном _expert_streaming слой после выполнения освобождает только те тензоры, которые разместил сам, а не выгружает модуль целиком на meta — иначе он вытеснил бы экспертов, которые управляют своей памятью сами, а заодно и cpu-resident потомков вроде n-gram-таблицы Flash-Next.
mmap-эмбеддинги и резидентные на CPU модули
Часть весов не должна попадать на GPU вообще. AirLLM собирает имена таких модулей: сначала из списка cpu_resident в layer_names_dict, затем — если задан cpu_resident_marker — из стеблей файлов *.mmap.json и *.safetensors в каталоге чекпоинта, оканчивающихся на этот маркер.
Для каждого такого имени AirLLM проверяет наличие <name>.mmap.json. Если файл есть, n-gram-таблица (большая таблица n-грамм, например ~51B у Flash-Next) открывается через open_ngram_mmap_table и подключается к модулю так, что её содержимое читается с диска по мере обращения, а не копируется в RAM; имя веса при этом добавляется в _cpu_resident_params. Без mmap-файла модуль грузится через load_merged_ngram_embedding, а если файла нет вовсе, модуль просто пропускается. move_layer_to_device затем пропускает такие параметры, даже если они ещё лежат в шарде декодер-слоя.
Сжатие весов
_decompress_state_dict распаковывает упакованные веса (weight_packed/weight_scale) в обычные weight. Причина: compressed-tensorsформат хранения весов, при котором они лежат в упакованном виде вместе с масштабами не имеет квантованных вычислительных ядер — готовых операций, которые считали бы прямо по упакованным данным, — поэтому перед первым forward он разворачивает всю модель хуком, после чего каждый квантованный модуль ждёт обычный weight. Распаковка идёт на GPU уже после передачи упакованных байтов — для MXFP4 это в 4 раза меньше данных по PCIe, чем передавать развёрнутый вес.
Для модулей, чей forward сам умеет потреблять упакованные данные, распаковка не нужна: их оставляют упакованными. Для остальных упакованные параметры подменяются на обычный weight, а модуль помечается статусом COMPRESSED, чтобы forward пропустил fake-quantize.
Отличие compression от обычной квантизации авторы объясняют так: квантизация обычно требует квантовать и веса, и активации, чтобы реально ускорить работу, а это усложняет сохранение точности и борьбу с выбросами. Здесь узкое место — загрузка с диска, поэтому квантуются только веса, что проще для точности.
LoRA поверх стриминга
Трейнинг реализован через torch.autograd.Function _StreamedModule: forward выполняется под no_grad, вход и выход держатся на CPU, а в backward сохранённое скрытое состояние переносится на устройство и модуль пересчитывается уже с градиентом. После backward слой выгружается, а предыдущий префетчится. _prefetch_stream через executor заранее загружает слой на CPU, если префетчинг включён и индекс в допустимых границах.
У Flash-Next есть особенности. _make_masks строит маски из token-width embeds (hidden), а не из 4-stream hyper state, и возвращает словарь с qwen_sparse_attention и linear_attention. Если маски построить не удалось, система сообщает об этом и переходит на путь is_causal слоя. _embed_and_context повторяет hidden столько раз, сколько задано в конфигурации модели параметром hc_count (по умолчанию 1): это число гипер-состояний, которые модель ожидает на входе, и повторение приводит hidden к нужной кратности. _run_decoder_layer берёт маску qwen_sparse_attention и передаёт conv_mask и ple_input_ids.
Требования к окружению
Для Kimi K3 flash-attn объявлен обязательным: код модели K3 перезаписывает любую запрошенную реализацию внимания. Нужна сборка torch под CUDA 12 — готового колеса flash-attn под CUDA 13 пока нет. Нужен transformers 4.56.x, так как удалённый код K3 не загружается на 5.x.
Для Flash-Next требуется in-tree qwen4_exp из main-ветки transformers. Достаточно машины с 64GB RAM, на диске планируется ~360GB либо можно передать delete_original=True, чтобы освободить оригиналы после разбиения. Установка: pip install airllm compressed-tensors flash-attn, transformers — из репозитория.
Отдельная функция restore_relocated_transformers_symbols лечит несовместимость с transformers 5.0: та перенесла два вспомогательных символа из transformers.utils.generic в transformers.utils.output_capturing, а удалённый код моделей всё ещё импортирует их из старого места и потому не импортируется вовсе. Функция повторно экспортирует перенесённые символы под старыми именами там, где их не хватает, и вызывается в __init__ базового класса до загрузки модели.
Что это меняет на практике
Замеры для K3 сделаны на одной RTX 6000 Ada (48GB) против полного чекпоинта 1.56TB: пик VRAM при генерации 3.72 GB, после инициализации — 0.83 GB. Инициализация занимает 900 s, генерация — 292 s/token и ограничена диском. Для Qwen3.8-27B в заголовке релиза указано 3.33GB, обучение — ~2GB при seq 512. Flash-Next стримится на одной карте в 5.95GB VRAM (RTX 4090), обучается под 6GB на RTX 3060 Ti; таблица n-gram ~51B при этом отображается в файл на хосте, а слои декодера стримятся.
При миграции на K3 ломается выбор attention-реализации: AirLLM проталкивает выбор вниз по вложенным суб-конфигам, потому что мультимодальные обёртки держат декодер под суб-конфигом (у K3 это text_config), а transformers записывает запрос только в тот конфиг, который ему передали. Без этого суб-модель читает незаданное значение и уходит на путь flash-attention, падающий на машине без flash-attn.
Для MoE transformers по умолчанию использует grouped_mm; на 4090 это Python-фолбэк, чей autograd всегда выделяет полноразмерный dW для упакованных тензоров экспертов (~3.4GB), даже если веса заморожены. Eager F.linear на маршрутизированных срезах этого не делает, поэтому LoRA-класс при наличии set_experts_implementation переключает экспертов на eager.
Ограничения и открытые вопросы
Авторы прямо называют цену подхода: при инференсе исходная модель сначала разбирается на слои и сохраняется на диск, поэтому нужно достаточно места в каталоге кэша huggingface. Для Flash-Next это ~360GB диска и 64GB RAM. Зависимость от main-ветки transformers — тоже признанное ограничение: qwen4_exp пока не в стабильном релизе.
В README описаны две типовые ошибки. 401 Client Error возникает на gated-моделях — им нужен токен HuggingFace API, который передаётся через hf_token и при загрузке конфигурации модели уходит в HuggingFace как токен доступа. ValueError про padding-токен возникает, когда у токенизатора его нет; помогает либо явное задание padding-токена, либо отключение padding — пример в README использует padding=False вместе с return_attention_mask=False, truncation=True и max_length=MAX_LENGTH.
