Конвертация Chroma Context-1 в MXFP4 для домашней 4090
Позавчера Chroma выкатила https://trychroma.com/research/context-1 по своей модельке https://huggingface.co/chromadb/context-1, прочёл я его сегодня и был крайне впечатлен.
Это 20B параметров, MoE-архитектура на базе https://huggingface.co/openai/gpt-oss-20b, модель натренирована на агентный поиск, делает декомпозицию сложных запросов на подзапросы, выполнняет итеративный поиск по корпусу документов, и самое интересное - self-editing context, когда модель сама решает какие из найденных документов оставить, а какие выкинуть чтобы не засорять контекстное окно. На бенчмарках показывает результаты сопоставимые с топовыми LLM при скорости как у gpt-oss-20b.
Естественно, захотелось потрогать руками, но... модель выложена в BF16, а это 39 ГБ только веса, ну а моя 4090 хоть и прокачана до 48 ГБ VRAM, но 39 ГБ чистых весов + KV-кеш + накладные расходы CUDA в неё не влезут, нужно минимум ~45-50 ГБ чтобы хоть что-то сгенерировать с минимальным контекстом, но нужен не минимальный, а хотя бы 30-60к токенов, агенты же.
Что делать? Смотрю на оригинальную gpt-oss-20b от OpenAI, она в формате MXFP4, это https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf, хитрый 4-битный формат, где каждый вес хранится как E2M1 (2 бита экспонента, 1 бит мантиссы), а группа из 32 весов делит общий 8-битный масштаб (E8M0), нативно поддерживается в vLLM через Marlin-ядра и должно отлично работать на RTX 50xx серии, в таком формате модель занимает ~14 ГБ, на мою 4090 влезает с запасом.
Сформировал спеки для кодового агента в Cursor, говорю мол вот тебе веса модели gpt-oss-20b как референс, вот https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf по методам сжатия MXFP4, вот веса модели context-1 в BF16 которую надо конвертировать, вот docker-compose.yaml с vLLM 0.18.0 на которой агент должен проверять модель.
Задача написать скрипт конвертации, конвертировать модель, запустить в vLLM, убедиться по HTTP что приходят адекватные ответы.
Агент написал https://huggingface.co/evilfreelancer/context-1-mxfp4/blob/main/convert_mxfp4.py, который читает safetensors, для MoE-экспертов (gate_up_proj и down_proj) транспонирует веса в layout который ожидает vLLM, квантует и пакует всё в MXFP4, по аналогии с gpt-oss-20b, а слои внимания, роутера и эмбеддингов копируются как есть в BF16. В результате 39 ГБ превращаются в ~14 ГБ.
Запустил на vLLM, всё работает, reasoning что-то там рефлексирует, ответы вроде адекватные, но при тестировании tool calling (function calling) выяснилась забавная штука - модель при вызове инструментов валится с ошибкой:
openai_harmony. HarmonyError: Unexpected token 12606 while expecting start token 200006
Покопавшись в конфигах gpt-oss-20b обнаружил, что в https://huggingface.co/chromadb/context-1/blob/main/generation_config.json модели отсутствует токен <|call|> (id 200012) в списке eos_token_id, из-за этого модель не останавливает генерацию после вызова tool call и harmony ломается. По сути баг в vLLM, но фикс элементарный, добавить 200012 в eos_token_id в generation_config.json модели, после этого и обычный чат и tool calling заработали как надо.
Создал https://huggingface.co/chromadb/context-1/discussions/3 с фиксом в трекере оригинальной модели.
Ну а квантованную в MXFP4 версию загрузил на HuggingFace: https://huggingface.co/evilfreelancer/context-1-mxfp4
PS. Поднял её на домашней видяшке и сделал доступной через мою апишку, пример как заюзать через openai-клиент добавил в комментарии.
