Laplace 3B登陆AWS Trainium2
AWS Machine Learning Blog称，i14与AWS通过TorchNeuron Native将30亿参数的Laplace 3B适配至AWS Trainium2，并保留Gated Linear Attention架构和Muon优化器。在单个trn2.48xlarge实例上，该实现训练速度约为每秒12万个token，并在1280个并发序列下达到每秒11544个输出token。博客称其循环状态无论上下文长度如何都保持每序列29.5 MB，同时提醒固定状态压缩仍需针对具体工作负载评估检索质量。
Laplace 3B展示了AWS Trainium2支持非标准循环架构和恒定单序列推理状态的能力，为从业者评估线性注意力部署提供了具体案例。
Most large language models (LLMs) use transformer architectures with softmax attention. During inference, their key-value (KV) cache grows with the conversation or document context. This can reduce concurrency, increase serving cost, and force agentic applications to discard useful history.
i14 designed Laplace 3B as a linear attention alternative. Its recurrent state stays at 29.5 MB per sequence, independent of context length. Working with AWS, i14 used native PyTorch support through TorchNeuron Native to bring this custom architecture to AWS Trainium2. This preserved its Gated Linear Attention (GLA) design and Muon optimizer.
On one trn2.48xlarge instance, Laplace 3B trained at about 120,000 tokens per second. The model served 11,544 output tokens per second across 1,280 concurrent sequences. This post explains the model design, the changes i14 made for Trainium2, and the training and inference results.
About i14
i14 is an AI research and technology company that develops model architectures, components, and runtimes. Its work draws on AI, computational cognitive science, and research into how biological systems learn over time. i14 applies these ideas to models that can retain long-range context without allowing serving state to grow with every token.
#i14#Laplace 3B#AWS Trainium2#TorchNeuron Native#AWS Neuron SDK#Gated Linear Attention#Muon optimizer#PyTorch