1 分•作者: yousef_g•6 天前
Hi HN，我构建这个项目是因为我想看看，我是否能在独立研究人员可用的硬件上预训练大词汇量LLM（比如拥有26.2万个token的Gemma）。
标准的精确交叉熵在16GB GPU上，遇到这种规模时会立刻耗尽内存（OOM）。
为了绕过这个问题，我实现了MAXIS损失函数。它使用“幽灵Logit”来数学模拟未采样token的缺失概率质量，而不是实例化完整的26.2万宽的矩阵。
在16GB VRAM显卡（T4）上的基准测试结果：
* 与Triton优化的Liger Kernel相比，损失层速度快17.5倍。
* 目标计算中VRAM减少约39%。
包含RandNLA注意力机制，该机制使用因果克罗内克素描来保持内存随着序列长度的增长而平稳。
我在代码库中包含了带有正式数学公式的技术报告。我非常欢迎大家对分区函数模拟或素描方法的任何技术反馈。