BMTrain

收录理由

BMTrain 面向那些需要预训练或微调数百亿参数规模模型、却不想把代码重写成分布式训练框架的团队。它把 ZeRO-2/3 风格的分片和张量并行带进 PyTorch，你仍然写普通的 nn.Module 代码，通信和计算与数据传输的重叠由工具在底层处理。上手时只需替换几个基类、包装一下 transformer 模块，然后用你熟悉的 torchrun 命令启动即可。这种轻量迁移路径对研究组和小团队很有吸引力，他们想训练大模型，但不想为此引入一个笨重的框架。项目自带 GPT-2 示例，并且能接入 OpenPrompt 和 OpenDelta，所以从零开始的预训练和微调流程都能覆盖。