Laplace 3B: фиксированное состояние вместо KV-кэша на Trainium2

На инференсе обычного трансформера кэш ключей и значений растёт вместе с контекстом: меньше параллельных сессий, дороже сервинг, агентные сценарии режут полезную историю. i14 собрала Laplace 3B на линейном внимании — рекуррентное состояние остаётся 29,5 МБ на последовательность и не зависит от длины. На AWS Trainium2 модель подняли через нативный PyTorch (TorchNeuron Native), сохранив Gated Linear Attention и оптимизатор Muon.

На одном инстансе trn2.48xlarge обучение шло примерно со скоростью 120 000 токенов в секунду. На выдаче — 11 544 выходных токена в секунду при 1 280 параллельных последовательностях. Вес — 3,028 млрд параметров, 24 слоя декодера, контекст 2 048 токенов; датасет — 100 млрд токенов FineWeb-Edu.

Линейное внимание вместо softmax

Классическое внимание сравнивает токены попарно: на обучении это квадратично по длине, на выдаче память уходит в KV-кэш. Simple GLA в Laplace сжимает прошлое в матрицу состояния 96×192 на голову (16 голов), обновление рекуррентное, по длине последовательности обучение линейное. Забывание по головам задано геометрическим расписанием — от одного до 2 048 токенов.

Три правки под Trainium2

• Ядро GLA перенесли с GPU-Triton на операторы PyTorch: 16 чанков по 128 токенов под 128 партиций буфера NeuronCore — прирост пропускной способности около 14%.
• Коммуникацию в MuonFSDP2 перестроили на два all_to_all вместо gather/broadcast: трафик на ранг упал в 24 раза (до 0,45 ГБ), скорость обучения выросла с 73 378 до 119 934 токенов в секунду при той же кривой loss.
• Обратный проход по 16 шагам рекуррентности статически развернули под ограничения компилятора Trainium2.

Контекст сжимается в фиксированное состояние, а не хранится поключно — качество извлечения фактов для своих ворклоадов всё равно нужно замерять отдельно.

Источник: Beyond the KV cache: Training and serving i14’s Laplace 3B linear attention model on AWS Trainium2 with native PyTorch.