[AI21] Jamba-1.5: Hybrid Transformer-Mamba Models at Scale
AI21 Labs Jamba Team
Статья: https://arxiv.org/abs/2408.12570
Пост: https://www.ai21.com/blog/announcing-jamba-model-family
Модели: https://huggingface.co/collections/ai21labs/jamba-15-66c44befa474a917fcf55251

Малозамеченным прошёл релиз моделей Jamba-1.5, отскейленных версий мартовской Jamba (https://www.tgoop.com/gonzo_ML/2492).

Напомним, что Jamba — это гибрид SSM (https://www.tgoop.com/gonzo_ML/1424) и трансформера, точнее Mamba (https://www.tgoop.com/gonzo_ML/2148) + MoE (Mixture-of-Experts, про это мы писали много, можно начать отсюда https://www.tgoop.com/gonzo_ML/472) + трансформерные слои.

В оригинале блок Jamba состоял из 8 слоёв, из них каждый второй MoE, всего четыре штуки; три слоя Mamba, и один трансформерный. Малое количество трансформерных слоёв позволяло уменьшить размер KV-кеша (получается в 8 раз меньше обычного трансформера с таким же количеством слоёв).

Оригинальная Jamba содержала 52B параметров, из которых активны в каждый момент были 12B (потому что MoE).

Благодаря более скромному memory footprint, модель позволяла использовать контекст размером 140k на одном GPU A100-80 Gb, намного больше, чем влезало у Llama-2 70B или Mixtral 8x7B. Полный размер контекста модели был 256k токенов. Это также позволяло использовать более крупные батчи, так что итоговый throughput начиная с размера батча 4 был выше упомянутых конкурентов.

По качеству оригинальная Jamba показала себя достойно в сравнении с Llama-2 13B-70B, Gemma 7B и Mixtral.

Это была базовая модель, никакого alignment или instruction tuning. Доступна под Apache 2.0

Теперь в августе вышло обновление, Jamba-1.5, включающая две модели:
* Jamba-1.5-Mini: 12B/52B active/total params (как оригинальная Jamba)
* Jamba-1.5-Large: 94B/398B active/total params

Пробовали блоки Mamba-2 (https://www.tgoop.com/gonzo_ML/2718), но они оказались не лучше и в архитектуре оставили Mamba-1.

Для эффективного инференса разработали новую квантизацию ExpertsInt8, когда веса MoE и MLP квантуются в INT8, а перед вычислением приводятся к BF16, чтобы использовать быстрые BF16 кернелы. Это всё происходит внутри vLLM в fused_moe кернеле. На H100 latency ExpertsInt8 соответствует FP8, а на A100, где нет FP8, намного превосходит GPTQ.

В обучение добавили Activation Loss, так как некоторые активации вырастали до 4e6, что вроде ничему не мешало, но на всякий случай.

Throughput и latency у Jamba хороши по сравнению с конкурентами (Llama 3.1 8B, Mixtral-8x7B, Mistral Nemo 12B для Mini; Llama 3.1 70B, Mistral Large 2, Llama 3.1 405B для Large), особенно на большом размере контекста.

Обучалось на каком-то внутреннем датасете в три фазы. В pre-train по сравнению с предыдущей Jamba добавили мультиязычные данные с фокусом на English, Spanish, French, Portueguse, Italian, Dutch, German, Arabic, Hebrew. Затем был mid-training с фокусом на длинных документах. Затем post-training с SFT на качественных разговорных данных, скилл-специфичных и с длинным контекстом. Как я понял, отдельного preference tuning типа PPO/DPO не было, обошлись качественной синтетикой, фильтрацией и SFT.

Модель обучена с function calling. Я рад, что эта тема (https://www.tgoop.com/gonzo_ML/2821) развивается.

Итоговые модели сравнимы с соразмерными конкурентами из линеек Llama-3.1, Gemma-2, Mistral-Large-2.

Отдельно проверили способности на задачах с большим контекстом через бенчмарк RULER (https://arxiv.org/abs/2404.06654) с 8 вариантами needle-in-a-haystack задач. Заявляют, что они единственные, кто поддерживает эффективный контекст в 256k, остальные хоть и заявляют большие длины, но лажают. На ∞BENCH тоже хороши.

Короче, выглядит хорошо. Кажется, это первая реально большая нетрансформерная (ну почти) модель. Лицензия у новой модели правда изменилась с Apache 2.0 на Jamba Open Model License, которая personal, revocable, и не разрешает коммерческое использование, если вы зарабатываете больше $50M в год (problems nice to have).

В общем, интересное развитие, ждём больше нетрансформерных SSM (и не только) моделей. У NVIDIA был тоже гибрид Mamba-2-Hybrid (https://arxiv.org/abs/2406.07887), и есть ещё StripedHyena (свёртки Hyena + attention, https://www.together.ai/blog/stripedhyena-7b), но последние две были маленькие, 7-8B.

У кого охлаждение интереса к GenAI, а у кого и сид раунды на миллиард:

https://www.reuters.com/technology/artificial-intelligence/openai-co-founder-sutskevers-new-safety-focused-ai-startup-ssi-raises-1-billion-2024-09-04/

И в продолжение банкета. Не могу это развидеть

https://www.facebook.com/share/v/NeWF1pCNzJb8fW3f/

Here are my slides from today's talk at Datafest Yerevan.

The talk was about non-transformer architectures, e.g., good old MLPs, CNNs, RNNs, and brand-new SSMs. It may be too dense with too many model names, but I think it may be useful as a reference for further exploration.

https://docs.google.com/presentation/d/19jpt6sSScUb1yKnlO3a47SsMRIL7UmqQZKkuADyI7nM/edit#slide=id.g2f6fb83b821_0_15

Интересная статистика сбоев в распределенной системе для обучения большой LLM. Из работы про Llama 3 (https://ai.meta.com/research/publications/the-llama-3-herd-of-models/)

Ризонинг подвезли!

https://openai.com/index/learning-to-reason-with-llms/