MLPerf 2.0 Bert在AWS Gaudi1上训练的数据集处理问题咨询

环境背景 #

• 运行环境：AWS Gaudi1虚拟机（742GB内存、96核），镜像为Deep Learning AMI Habana PyTorch 1.12.0 SynapseAI 1.6.0 (Ubuntu 20.04)
• 目标：运行MLPerf 2.0 Bert训练代码，处理Wiki数据集

问题1：逐个转换数据集part文件为tf_records的流程是否正确？ #

这个流程是正确且合理的。原脚本批量处理所有part文件时会一次性加载大量数据到内存，导致内存耗尽。逐个转换每个part文件可以将内存负载分散到单个文件的处理量上，避免内存溢出，生成的单个tf_record文件和批量生成的文件格式完全一致，不会影响后续训练或打包流程。

问题2：将tf_records转换为二进制文件的方式是否可行？生成的文件能否用于非打包训练方式？ #

这种方式不可行。你遇到的masked_lm_ids解析错误，是因为第三方脚本生成的二进制文件格式和Habana Gaudi训练代码预期的tf_record格式不匹配。Gaudi的Bert训练代码依赖原生tf_record的Example结构，自定义二进制文件的字段解析逻辑和训练代码不兼容，会导致数据解析失败。因此不能用这种方式替代原生tf_record进行非打包训练。

问题3：如何在pack_pretraining_data_tfrec.py脚本中限制max-files为10或25？ #

直接在运行脚本时通过命令行参数指定即可，命令格式如下：

python pack_pretraining_data_tfrec.py --max-files=10 --input-dir=/path/to/your/tf_records --output-dir=/path/to/output

但需要注意：原脚本的内存占用问题并非仅由max-files参数决定——它在生成策略文件前仍会尝试加载所有tf_record的元数据。如果内存依然不足，可以修改脚本逻辑：

• 遍历输入目录时，每次只读取max-files指定数量的tf_record文件
• 分批次生成对应的打包策略文件和打包数据，而不是一次性处理所有文件

其他针对Gaudi打包Wiki数据集的方法 #

• 分批次打包：手动将tf_record文件分成若干组（每组10-25个文件），对每组单独运行pack_pretraining_data_tfrec.py，生成多个打包后的数据集文件，训练时指定所有打包文件的路径即可。
• 优化脚本内存占用：修改pack_pretraining_data_tfrec.py，将加载所有tf_record元数据的逻辑改为流式读取，避免一次性加载全部数据到内存。例如，使用TensorFlow的tf.data.TFRecordDataset逐个读取文件并提取元数据，而不是一次性加载所有文件内容。
• 适配Habana工具链：检查Habana SynapseAI工具链中是否有针对MLPerf数据集的预处理优化脚本，部分版本会提供内存高效的数据集打包工具，适配Gaudi的硬件特性。

内容的提问来源于stack exchange，提问作者Rajesh Muthu