Step 3.5 Flash [hf] — опенсорсная MoE-модель на 196 миллиардов параметров, из которых 11 миллиардов активные.При таком размере модель демонстрирует конкурентные результаты в сравнении с более крупными опенсорсными и проприетарными моделями. Разберём, как устроена Step 3.5 Flash.
Архитектура
Одна из главных целей, которую ставили перед архитектурой разработчики — создать модель с низкой задержкой (latency) для использования в агентских сценариях. Добиваются этого с помощью гибридного механизма внимания, sparse MoE и Multi Token Prediction (MTP). В каждом MoE-слое — 288 routed-экспертов и один shared-эксперт, при этом для каждого токена активируются восемь routed-экспертов. Модель содержит 45 слоёв. По сравнению с современными открытыми MoE-моделями (DeepSeek-V3.2, Qwen-3.5, GLM-5), здесь больше мелких экспертов, более высокая sparsity и меньшее число слоёв, что снижает вычислительные затраты на один шаг инференса.
Авторы используют схему 3:1 — последовательных слоя внимания используют механизм скользящего окна (sliding-window attention, SWA), а каждый четвёртый — классический (full attention, FA). Есть Grouped Query Attention с восемью KV-головами (GQA-8), что даёт эффективное распределение KV-кэша на стандартных нодах с восемью GPU и тензорный параллелизм на восемь частей. Это ведёт к увеличению утилизации памяти. Attention становится memory-bound, освободившиеся вычислительные ресурсы идут на MTP — speculative drafting и verification.
Изначально метод чередования слоёв (3 SWA на 1 FA) уступал классической FA-архитектуре по качеству на бенчмарках. Авторы смогли исправить эту проблему, увеличив число query-голов с 64 до 96. Вторая модификация — Head-Wise Gated Attention, который в SWA снижает влияние шумовых активаций в случаях, когда релевантный контекст отсутствует внутри локального окна.
В модели используется стандартный Sparse MoE с fine-grained-сегментацией, при которой отдельные подпространства скрытого представления маршрутизируются независимо. Для лучшего распределения токенов по экспертам применяется механизм EP-Group Balanced MoE Routing, который оптимизирует равномерное распределение токенов по экспертам и по GPU.
Используют три MTP-головы. В течение обучения учится только одна, а остальные две — лишь на последней стадии посттрейнинга. Добавляют также position-dependent loss reweighting, чтобы снизить влияние далёких токенов на лосс и не переучиваться на их предсказание.
Нестабильности в обучении
Авторы описывают три фактора нестабильности, с которыми боролись во время претрейна: спайки в лоссе, «мёртвые» эксперты и «взрывы» экспертов. Первые выявили в процессе обучения — обнаружили накопление ошибки сложения при вычислении полярной декомпозиции в оптимизаторе Muon. Смена типа данных с bfloat16 на float16 исправила проблему.
Даже при хорошем роутинге эксперты могут либо перестать учиться, либо генерировать «вредные» активации. Shared-эксперт может «давить» отдельных экспертов, не позволяя им учиться; плохо обученный эксперт способен «игнорироваться» впоследствии; при идеальном роутинге эксперты «соревнуются» даже за неподходящие им токены — всё это приводит к коллапсу отдельных экспертов. Таким образом, статистика роутинга — не показатель здоровья обучения: необходимо отслеживать нормы выходных активаций, весов, динамику изменений, распределение по экспертам. И использовать эти метрики для диагностики и стабилизации обучения.
Существует обратная проблема — «взрывы» активаций. Они происходят, когда эксперт, который обрёл узкую специализацию, даёт высокие активации важным для себя биграммам, особенно частотным. При использовании pre-norm ничто не ограничивает абсолютные значения при добавлении к residual. Если через гейт SwiGLU проходят активации, рост становится ещё сильнее. Muon при обучении может усиливать подобные паттерны, приводя к положительной обратной связи между активациями и обновлениями весов. Для борьбы с этим авторы использовали клиппинг на выходные активации и на веса экспертов.
Разбор подготовил
Душный NLP
