训练nGPT

归一化的Transformer（nGPT）通过将模型参数向量和激活向量约束到单位超球面来实现超球面表示学习。在本文中，我们描述了 nGPT 的实用训练方法，并在现代混合 Mamba-2——Transformer Mixture-of-Experts (MoE) 模型上对其进行了评估。该配方介绍了 logits 梯度预处理、对数学习率衰减、GatedAdamW、角度更新控制和可选的探索机制。与使用 AdamW 训练的相同混合 MoE 架构的非归一化模型相比，30B 总参数 nGPT 模型使用大约一半的训练词元达到相同的验证损失。该配方适用于所考虑的模型，其中包含多达 30B 个总参数。