За пределами KV-кэша: обучение и обслуживание линейной модели внимания Laplace 3B от i14 на AWS Trainium2 с использованием нативного PyTorch
aws ·
В статье подробно описывается процесс обучения и развертывания модели линейного внимания Laplace 3B от i14 на инфраструктуре AWS Trainium2 с использованием нативного PyTorch через TorchNeuron Native. В отличие от традиционных трансформеров, использующих растущий ключ-значения (KV) кэш, Laplace 3B применяет затворенное линейное внимание (Gated Linear Attention), что позволяет поддерживать фиксированное рекуррентное состояние объемом 29,5 МБ на последовательность независимо от длины контекста. Реализация потребовала трех основных инженерных адаптаций: переписывания ядра GLA с использованием стандартных операторов PyTorch, оптимизации распределенной коммуникации для оптимизатора Muon и статического разворачивания обратного прохода для компилятора. Обучение на одном экземпляре trn2.48xlarge достигло скорости около 120 000 токенов в секунду при стабильной сходимости функции потерь. Производительность инференса составила 11 544 выходных токена в секунду для 1 280 одновременных последовательностей благодаря компиляции шага декодирования для снижения накладных расходов на диспетчеризацию операторов. Эти результаты демонстрируют, что пользовательские архитектуры могут обеспечивать высокую эффективность на оборудовании AWS, предлагая предсказуемое использование памяти для приложений с длинным контекстом.
aws