⚙️ Laplace 3B от i14: линейное внимание без KV-кэша обучили на AWS Trainium2

Компания i14 совместно с AWS Generative AI Innovation Center адаптировала свою модель Laplace 3B с линейным вниманием под ускорители Trainium2, используя нативный бэкенд PyTorch - TorchNeuron Native.

В архитектуре Laplace 3B (24 слоя, 3,03 млрд параметров) softmax-внимание заменено на Gated Linear Attention. Вместо растущего KV-кэша модель хранит рекуррентное состояние фиксированного размера - 29,5 МБ на последовательность независимо от длины контекста. Это меняет сложность обучения с O(T²) на O(T), а инференс делает O(1) на токен.

Инженеры внесли три ключевых изменения: переписали Triton-ядро fused_chunk_simple_gla на стандартные операторы PyTorch (16 чанков по 128 токенов под 128 партиций буфера NeuronCore), переработали коммуникации Muon через два all_to_all в MuonFSDP2 (трафик на ранг упал в 24 раза, пропускная способность выросла в 1,63 раза) и статически развернули 16 шагов рекуррентного backward pass, чтобы уложиться в лимиты компилятора.

Результат на одном инстансе trn2.48xlarge: обучение около 120 000 токенов в секунду и сервинг 11 544 выходных токенов в секунду при 1 280 одновременных последовательностях. Подход интересен для长期 агентов и анализа больших документов, где нужна предсказуемая память при подаче запросов, но качество извлечения информации из сжатого состояния требует отдельной проверки под конкретные задачи.

🎨 qvabo.studio/tg · картинки и видео

#AI #Models #AWS #Trainium2 #PyTorch #LLM #LinearAttention