tinygrad 在 MLPerf BERT 训练基准中的实战指南数据流水线、环境配置与多平台运行【免费下载链接】tinygradYou like pytorch? You like micrograd? You love tinygrad! ❤️项目地址: https://gitcode.com/GitHub_Trending/tiny/tinygrad本指南以 MLPerf Training v5.0 提交中 BERTNLP 预训练基准的 tinygrad 实现为主线完整讲解从依赖安装、Wikipedia 原始数据下载与校验、MLM/NSP 预训练样本预处理到 tinybox_green / tinybox_red / tinybox_8xMI300X 三种硬件平台上运行基准的完整流程。读完本文你将掌握 MLPerf BERT 提交的全部操作步骤以及驱动这些脚本的底层源码机制与关键环境变量。一、基准背景用 tinygrad 复现 MLPerf BERT 训练MLPerf Training 是衡量机器学习训练性能的行业标准基准其中 BERTBidirectional Encoder Representations from Transformers任务要求在大规模 Wikipedia 语料上完成掩码语言模型Masked LM与下一句预测NSP的预训练。tinygrad 在examples/mlperf/training_submission_v5.0/下维护了完整的 closed 分卷提交其 BERT 实现位于examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/核心入口是 model_train.py而本指南对应的文档正是该提交在tinybox_green实现目录下的 README另有tinybox_red、tinybox_8xMI300X两份并列实现内容基本一致。整套流程分为两大阶段数据阶段下载 2020 年 Wikipedia 数据集做 checksum 校验然后完成 BERT 标准的 tokenize、Masked LM 采样、序列截断、padding 等预处理生成可直接供训练读取的 pickle 特征文件训练阶段通过 run_and_time.sh 依次完成预热/初始化 正式计时训练两轮运行同时输出符合 MLPerf logging 规范的日志。二、环境准备tinygrad 与依赖安装基准运行要求从mlperf_training_v5.0分支安装 tinygrad并同时安装mlperf-logging与数据集工具链。依赖声明位于仓库根目录的 pyproject.toml其中[mlperf]extra 通过注释形式指向mlperf-logging githttps://github.com/mlperf/logging.git5.0.0-rc3需要在安装前取消注释启用git clone tinygrad 仓库地址 python3 -m pip install -e .[mlperf]随后安装数据集下载与处理所需的辅助库pip install gdown numpy tqdm tensorflow其中gdown负责从 Google Drive 下载 Wikipedia 数据集下载逻辑见 wikipedia_download.pynumpy/tqdm预处理与进度展示依赖预处理脚本大量使用tqdm与tqdm.contrib.concurrent.process_map实现多进程并行见 wikipedia.pytensorflowBERT 参考实现model.ckpt-28252检查点与校验文件来自 MLPerf 官方 TensorFlow 流程预处理后对齐其数据格式。tinybox_green 的专用说明p2p 驱动对于tinybox_green平台README 特别注明需要按官方 p2ppeer-to-peer分支的说明安装 GPU 内核驱动对应550.54.15-p2p版本并指出该驱动在生产型 tinybox green 上是默认配置。p2p 能力让多卡间的显存直访与 allreduce 通信可以绕过主机内存拷贝这对 6 卡并行的 BERT 训练至关重要——tinygrad 的多卡 allreduce 实现位于tinygrad/schedule/allreduce.py多卡调度逻辑见tinygrad/schedule/multi.py。三、数据下载与校验wikipedia_download.py3.1 下载命令在仓库根目录执行BASEDIR/raid/datasets/wiki WIKI_TRAIN1 VERIFY_CHECKSUM1 python3 extra/datasets/wikipedia_download.py三个环境变量的作用见 wikipedia_download.py环境变量作用BASEDIR数据集根目录默认指向extra/datasets/wiki所有文件配置、词表、检查点、压缩包、解压结果都落在此目录下WIKI_TRAIN1额外下载训练语料results_text.tar.gz及其 md5 校验文件bert_reference_results_text_md5.txt并解压VERIFY_CHECKSUM1解压后对results4目录逐一做 MD5 校验任一文件不匹配即抛出ValueError见 verify_checksum3.2 脚本实际下载的内容脚本通过gdrive_download逐个拉取以下文件均来自 MLPerf 官方 BERT 数据集页bert_config.jsonBERT-Base 模型配置vocab.txtWordPiece 词表model.ckpt-28252.data-00000-of-00001/model.ckpt-28252.index/model.ckpt-28252.meta官方参考检查点三件套并自动生成checkpoint索引文件见 wikipedia_download.pybert_reference_results_text_md5.txt训练语料 500 个分片的 MD5 清单仅WIKI_TRAIN1时results_text.tar.gz约 500 个分片的原始文本压缩包仅WIKI_TRAIN1时。下载完成后脚本用tarfile流式解压带 tqdm 进度条并os.remove删除压缩包以释放磁盘空间见 wikipedia_uncompress_and_extract。四、预训练数据预处理wikipedia.py预处理脚本 wikipedia.py 是对 MLPerf 官方create_pretraining_data.py的改写实现负责把原始 Wikipedia 文本转换成 BERT 训练所需的input_ids / input_mask / segment_ids / masked_lm_positions / masked_lm_ids / masked_lm_weights / next_sentence_labels七元组特征并以 pickle 落盘。注意README 明确指出预处理线程数受可用内存限制——128GB 内存的机器上NUM_WORKERS建议上限为 16。4.1 生成训练集全量 500 个分片BASEDIR/raid/datasets/wiki NUM_WORKERS16 python3 extra/datasets/wikipedia.py pre-train allpre-train all会通过process_map以NUM_WORKERS默认取min(os.cpu_count(), 32)个 worker 并行处理 0499 号分片每个分片的结果写入BASEDIR/train/{part}.pkl见 process_part。已存在的分片会被跳过因此可以断点续跑。4.2 生成单个分片调试用BASEDIR/raid/datasets/wiki python3 extra/datasets/wikipedia.py pre-train 42第二个参数取 0499 之间的分片号仅处理results4/part-00042-of-00500适合小规模验证流水线正确性。4.3 生成验证集BASEDIR/raid/datasets/wiki python3 extra/datasets/wikipedia.py pre-evalpre-eval从results4/eval.txt生成样本后按均匀间隔抽取出10000 条验证样本pick_ratio len(instances) / 10000写入BASEDIR/eval.pkl见 get_features_from_part。4.4 预处理核心流程源码级解读单条样本的生成链路为文本清洗 → BasicTokenizer → WordPiece 分词 → 文档切块与 NSP 配对 → Masked LM 采样 → padding 成定长特征文本清洗_clean_text剔除控制字符Unicode 类别C、非法码点_tokenize_chinese_chars将 CJK 字符两侧加空格便于切分_is_chinese_char覆盖了0x4E00–0x9FFF、扩展区及兼容区等全部中文字符区间见 wikipedia.py分词BasicTokenizer 处理小写化、去重音符号NFD 归一化后丢弃Mn类别、按标点切分WordPiece 采用从最长子串贪婪匹配的策略长度超过 200 字符或无法匹配的 token 一律映射为[UNK]见 _wordpiece_tokenizeNSP 配对create_instances_from_document把文档切成约MAX_SEQ_LENGTH - 3长度的块随机决定 A/B 段划分并以 50% 概率从其他文档随机拼接 B 段构成负样本is_random_next见 wikipedia.pyMasked LM候选 token 打乱后按MASKED_LM_PROB默认 0.15比例采样其中 80% 替换为[MASK]、10% 保持不变、10% 随机替换为词表中任意词见 create_masked_lm_predictions特征化instance_to_features把序列 padding 到MAX_SEQ_LENGTH512、掩码位置 padding 到MAX_PREDICTIONS_PER_SEQ76并统一转为 int32/float32 的 numpy 数组见 wikipedia.py。4.5 预处理环境变量一览环境变量默认值说明BASEDIRextra/datasets/wiki输入文本与输出 pickle 的根目录MAX_SEQ_LENGTH512最大序列长度控制 padding 与内存MAX_PREDICTIONS_PER_SEQ76每序列最多掩码预测数MASKED_LM_PROB0.15掩码概率SHORT_SEQ_PROB0.1采样短序列长度 2512的概率DUPE_FACTOR10同一文本以不同掩码重复生成的次数RANDOM_SEED12345随机种子保证可复现NUM_WORKERSmin(os.cpu_count(), 32)预处理并行进程数受内存限制五、运行基准三个 tinybox 平台预处理完成后即可运行基准。README 提供了三个平台各自的入口脚本它们在examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/下运行方式一致examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh5.1 run_and_time.sh 的运行机制以 tinybox_green 的 run_and_time.sh 为例脚本分init与run两个阶段均调用examples/mlperf/model_train.py# init2 层 BERT 的快速预热同时完成 MLPerf 日志初始化 BENCHMARK10 INITMLPERF1 BERT_LAYERS2 python3 examples/mlperf/model_train.py | tee $LOGFILE # run正式计时训练PARALLEL0 关闭数据并行以外的并行加载 PARALLEL0 RUNMLPERF1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE两个阶段对应的执行分支由 model_train.py 中的INITMLPERF/RUNMLPERF控制init 阶段INITMLPERF1时脚本用MLLOGGER依次上报SUBMISSION_PLATFORM、SUBMISSION_BENCHMARK等 MLPerf 标准事件mllog_constants并用BENCHMARK10指定只跑 10 步、取中位数作为步时基准sorted(step_times)[BENCHMARK // 2]run 阶段RUNMLPERF1时执行完整训练循环并在关键节点持续上报 mlperf 日志日志文件以bert_green_${DATETIME}_${SEED}.log命名SEED$RANDOM保证每次运行的随机性可独立复现。5.2 核心训练/调度环境变量环境变量tinybox_greentinybox_redtinybox_8xMI300X说明DEVNVAMDAMD后端设备NVCUDAAMDROCm/HIPDEFAULT_FLOATHALFHALF—默认默认浮点精度HALF 即 FP16 混合精度SUM_DTYPEHALFHALF—归约累加精度GPUS668参与训练的 GPU 数BS/EVAL_BS96961024训练/验证 batch sizeFUSE_ARANGE11—融合 arange 算子减少 kernel 数BEAM853调度器 beam search 宽度BEAM_UOPS_MAX1000080006000beam 搜索的 uop 上限BEAM_UPCAST_MAX256256256upcast 上限BEAM_LOCAL_MAX102410241024局部搜索上限BEAM_MIN_PROGRESS555beam 最小进步阈值IGNORE_JIT_FIRST_BEAM111忽略 JIT 首轮 beamBASEDIR/raid/datasets/wiki同左同左数据集位置LOGMLPERF111启用 MLPerf 日志上报TRAIN_STEPS——39008xMI300X 的固定训练步数5.3 平台差异与注意事项tinybox_green6×RTX 4090依赖 p2p 内核驱动完成多卡通信dev_run.sh额外开启WANDB1便于实验监控dev_beam.sh通过DEBUG2输出 beam 搜索日志BEAM_LOG_SURPASS_MAX1两者均不参与正式计时tinybox_red6×AMD GPU脚本开头执行sudo rmmod amdgpu/sudo modprobe amdgpu做驱动重置并设置HCQDEV_WAIT_TIMEOUT_MS100000防止等待超时挂起注释标明 AMD 的 AM 驱动曾导致 NaN故回退使用标准amdgpu驱动tinybox_8xMI300X8×AMD MI300X为对齐 MLPerf 官方参考结果显式设置 LAMB 优化器超参OPT_BASE_LEARNING_RATE0.0011、OPT_LAMB_BETA_10.60466、OPT_LAMB_BETA_20.85437、DECAY0.1及TRAIN_STEPS3900batch size 提升到 1024init 阶段需DEBUG2防止挂起。六、提交硬件配置参考tinybox_green 系统清单README 对应的提交系统配置记录在 tinybox_green.jsonsubmitter: tinycorpdivision: closed其硬件基线可作为复现环境的最低参考主机1 节点AMD EPYC 753232 核 / 64 vCPU128GB DDR48×16GBNVMe SSD4TB RAID 1TB 启动盘Ubuntu 22.04.4Python 3.10.12CUDA 12.4加速器6× NVIDIA GeForce RTX 409024GB GDDR6XPCIe 4.0 x16 互联风冷。七、故障排查与实操建议预处理内存不足将NUM_WORKERS调低128GB 内存建议 ≤16或改用pre-train 分片号分批处理数据校验失败VERIFY_CHECKSUM1会直接以异常中断重新执行下载脚本即可已下载文件会被gdown跳过仅重下缺失项训练日志缺失确认LOGMLPERF1与SUBMISSION_PLATFORM已设置否则MLLOGGER不会初始化见 model_train.pyAMD 平台异常参照 tinybox_red 的做法在运行前重置amdgpu驱动并调大HCQDEV_WAIT_TIMEOUT_MS快速验证先用BERT_LAYERS2BENCHMARK10的 init 阶段验证整条流水线再启动完整训练。至此你已经可以从零复现 tinygrad 在 MLPerf Training v5.0 上的 BERT 基准安装依赖 → 下载并校验 Wikipedia 数据 → 多进程预处理 → 在目标 tinybox 平台上执行 init/run 两阶段计时训练并产出符合规范的 MLPerf 日志。本提交的全部脚本含 resnet、retinanet、stable_diffusion 等其他基准均可在 training_submission_v5.0 目录下继续探索。【免费下载链接】tinygradYou like pytorch? You like micrograd? You love tinygrad! ❤️项目地址: https://gitcode.com/GitHub_Trending/tiny/tinygrad创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考