Большинство больших языковых моделей используют трансформеры с softmax-вниманием, где размер кеша ключ-значение растет с длиной контекста, что увеличивает нагрузку на память и снижает параллелизм. Laplace 3B применяет Gated Linear Attention (GLA), который фиксирует размер состояния в 29,5 МБ независимо от длины последовательности, что позволяет эффективно работать с длинными контекстами.
Для адаптации модели к AWS Trainium2 команда переписала ядро внимания GLA на стандартные PyTorch-операторы, заменив GPU-оптимизированные Triton-кернелы. Использование TorchNeuron Native дало прирост производительности. Также оптимизировали коммуникации и компиляцию обратного прохода для соответствия аппаратным ограничениям.
AWS Trainium2 — специализированные ML-ускорители с управляемой программно памятью и высокоскоростным межчиповым соединением. На одном экземпляре trn2.48xlarge Laplace 3B обучалась со скоростью около 120 тысяч токенов в секунду и обслуживала 1280 параллельных последовательностей с пропускной способностью более 11 тысяч токенов в секунду.
Такой подход позволяет создавать модели с длительной памятью и предсказуемыми требованиями к ресурсам, что важно для приложений с анализом больших документов, агентных систем и сервисов с постоянным состоянием.
i14 сочетает идеи из когнитивных наук и ИИ, создавая архитектуры, которые сохраняют контекст без экспоненциального роста вычислительных затрат, что открывает новые возможности для масштабируемых и экономичных решений.
Laplace 3B с линейным вниманием и оптимизация под AWS Trainium2 позволяют компаниям запускать масштабируемые языковые модели с длительной памятью и меньшими затратами на инфраструктуру. Это снижает стоимость и повышает эффективность AI-сервисов, особенно в задачах с большими контекстами и длительным взаимодействием.
Материал подготовлен редакцией VOSHOD по открытому источнику. Мы сокращаем и адаптируем текст, чтобы выделить практический смысл.