FlagEmbedding项目中的大模型预训练数据处理内存优化实践

问题背景

在使用FlagEmbedding项目进行大模型预训练数据准备时，用户遇到了内存不足的问题。具体表现为在运行main.pretrain_data脚本处理2B规模的预训练数据时，任务因内存不足被系统终止。该用户的机器配置为250GB内存，但在处理CommonCrawl数据集时仍然出现内存耗尽的情况。

技术分析

这种现象在大规模数据处理中较为常见，主要原因在于：

1. CommonCrawl数据集规模庞大，单次加载全部数据会消耗大量内存
2. 数据处理过程中可能产生大量中间变量
3. 并行处理(num_proc=16)虽然提高了速度，但也增加了内存压力

解决方案

方案一：移除CommonCrawl数据集

最简单的解决方案是直接移除CommonCrawl数据集，这种方法虽然可行，但会损失约52.2%的训练数据，可能影响模型最终性能。

方案二：分阶段处理数据（推荐）

更专业的做法是将数据处理分为两个阶段：

1. 第一阶段处理CommonCrawl数据

   • 创建专用配置文件slimpajama-part1.json
   • 仅包含CommonCrawl数据集配置
   • 运行单独的数据处理命令

2. 第二阶段处理其他数据集

   • 创建另一个配置文件slimpajama-part2.json
   • 包含C4、GitHub、Book、ArXiv、Wiki、StackExchange等数据集
   • 运行单独的数据处理命令

在后续预训练阶段，通过--train_data参数同时指定两个阶段生成的数据目录即可。

技术细节与注意事项

1. 配置文件设计：配置文件中需要明确定义各数据集的混合比例(mixture)和平均token数量(num_tokens_avg)

2. Llama-3的特殊性：值得注意的是，使用打包数据(packed data)预训练Llama-3可能导致性能下降，这与Llama-3修改了注意力掩码机制有关，它禁止了跨文档边界的自注意力计算。

3. 内存管理：对于超大模型(如Llama-3)，建议采用文档级别的数据处理方式而非打包方式，这能更好地控制内存使用。

最佳实践建议

1. 对于内存受限的环境，始终采用分阶段数据处理策略
2. 监控内存使用情况，适当调整num_proc参数
3. 针对不同模型架构选择合适的数据处理方式
4. 在处理前评估各数据集的内存需求，合理规划处理顺序

通过这种分阶段处理的方法，可以在有限的内存资源下完成大规模预训练数据的准备工作，同时保证数据的完整性和模型训练的效果。