makeMoE

makeMoE 项目介绍

项目背景

makeMoE 是一个从零开始实现的稀疏专家混合语言模型，受到 Andrei Karpathy 的 makemore 项目的启发，并在其基础上进行了改进。此项目的目标是使用稀疏专家混合架构来生成类似莎士比亚风格的文本。该项目的代码使用 PyTorch 编写，并强调可读性和可扩展性。

项目的独特之处

主要创新点

• 稀疏专家混合架构：与传统的单一前馈神经网络不同，makeMoE 采用了稀疏专家混合的方法。这种架构使得模型能够在计算资源相对有限的情况下, 实现更为复杂的推理。

• 门控机制：实现了 top-k 门控和噪声 top-k 门控。门控机制的创新使模型在选择激活专家时更为灵活。

• 初始化选项：项目中使用了 Kaiming He 初始化，但用户可以根据需求更换成其他方法（例如 Xavier Glorot 初始化）以进行尝试。

• 专家容量：最近的更新（2024年3月18日）中增加了专家容量功能，提升了训练效率。

相较于 makemore 的保持不变之处

• 数据集和语言模型任务：继续使用 Andrei Karpathy 原本选择的数据集和任务，即生成类似莎士比亚的文本。

• 因果自注意力机制实现：沿用了 makemore 项目中的实现方式。

• 训练循环与推理逻辑：基础的训练和推理流程保持不变，确保了原有架构的经典性和稳定性。

参考文献

此项目广泛参考了以下文献：

• Andrew Karpathy 的 makemore 项目。
• 关于稀疏门控专家混合层的论文《Outrageously Large Neural Networks》。
• 专家混合的研究论文《Mixtral of experts》。

使用与上手指南

项目的主要实现被集中在单一的 makeMoE.py 文件中。如欲深入了解整体模型架构和各组件如何协调工作，推荐从 makMoE_from_Scratch.ipynb 开始学习。为了提高训练效率和理解专家容量，建议查阅 makeMoE_from_Scratch_with_Expert_Capacity.ipynb。

如果希望对源码进行调整和优化，可以从 makeMoE_Concise.ipynb 出发，这是一个更为精简和易于改动的版本。

makeMoE 的开发在 Databricks 完成，使用单个 A100 GPU 进行计算。项目支持在云服务商提供的更大规模 GPU 集群上扩展。此外，虽然项目中使用了 MLFlow 来记录和追踪各类指标，但这一部分是可选的，用户可以根据需要自行 pip 安装。

项目展望

makeMoE 强调了代码的可读性和可扩展性，期待更多开发者能够在此基础上进行创新和改进。开发者们可以自由地 hack、理解、改进并创造属于自己的版本，希望大家在这一项目中找到乐趣，欢迎大家的反馈和建议！