tinygrad 复现 MLPerf BERT 预训练tinybox_red 六卡 AMD 平台的提交环境搭建与基准运行指南【免费下载链接】tinygradYou like pytorch? You like micrograd? You love tinygrad! ❤️项目地址: https://gitcode.com/GitHub_Trending/tiny/tinygrad导读本文以 tinygrad 仓库中 examples/mlperf/training_submission_v5.1/tinycorp/benchmarks/bert/implementations/tinybox_red/README.md 为骨架完整讲解如何在 tinybox_red6 × AMD Radeon RX 7900 XTX这一硬件平台上从零复现 MLPerf Training v5.1 的 BERT 大规模预训练基准包括环境安装、Wikipedia 数据集的下载与校验、训练/验证样本预处理以及通过run_and_time.sh一键运行整个基准并输出合规日志的完整流程。读完本文你将掌握该提交目录中所有脚本与环境变量的含义并能结合 tinygrad 训练入口源码理解每一步背后的实现机制。1. 问题定义用 BERT 完成 NLP 预训练基准该提交目录对应的 benchmark 是BERTBidirectional Encoder Representations from Transformers任务类型为 NLP 大规模预训练。README 第一部分明确写道This problem uses BERT for NLP.在 MLPerf 的语境下这个基准的目标是使用 BERT-large 架构在 Wikipedia 语料上完成掩码语言建模Masked LM与下一句预测NSP两个预训练目标并在验证集上达到指定的 masked LM 准确率目标。从训练入口 examples/mlperf/model_train.py 的train_bert()函数可以看到tinygrad 实现的默认收敛目标为target, achieved getenv(TARGET, 0.72), False即当验证集 masked LM 准确率达到 0.72 时判定训练收敛并停止计时examples/mlperf/model_train.py。模型默认使用 BERT-large 配置24 层、hidden size 1024、16 注意力头、vocab 30522该配置由 examples/mlperf/helpers.py 中的get_mlperf_bert_config()定义并可通过BERT_SIZE、BERT_LAYERS等环境变量缩小规模用于调试。2. 提交平台tinybox_red 硬件规格tinybox_red 是 tinycorptinygrad 项目团队的 AMD 平台机器其完整硬件清单记录在系统描述文件 examples/mlperf/training_submission_v5.1/tinycorp/systems/tinybox_red.json 中核心参数如下项目配置节点数1主机 CPUAMD EPYC 753232 核 / 64 线程主机内存128GB8 × 16GB DDR4存储4TB RAID 阵列 1TB 系统盘NVMe SSD加速卡6 × AMD Radeon RX 7900 XTX24GB GDDR6PCIe 4.0 x16操作系统Ubuntu 22.04.4Python 3.10.12软件栈tinygradmlperf_training_v5.0 分支这些硬件参数与后续脚本中的关键配置一一对应GPUS6对应 6 张 7900 XTXBASEDIR/raid/datasets/wiki对应 4TB RAID 存储阵列README 中“128GB RAM 下预处理最多建议 16 线程”的说明也来自host_memory_capacity。3. 环境准备安装 tinygrad 与依赖README 的 Requirements 部分给出了标准安装流程。需要说明的是mlperf这个可选依赖在 pyproject.toml 中默认被注释掉因此在安装前需要先取消注释# mlperf [mlperf-logging githttps://github.com/mlperf/logging.git5.0.0-rc3]并确保处于mlperf_training_v5.0分支README 中实际写作 v5.0与 v5.1 提交目录配套使用git clone https://gitcode.com/GitHub_Trending/tiny/tinygrad python3 -m pip install -e .[mlperf]随后安装数据与预处理阶段所需的辅助包pip install gdown numpy tqdm tensorflow各依赖的用途分别为gdown从 Google Drive 下载 MLPerf BERT 官方 Wikipedia 数据集见 extra/datasets/wikipedia_download.py 文件头注释# pip install gdownnumpy / tqdm预处理脚本 extra/datasets/wikipedia.py 的数值与进度显示依赖tensorflow加载官方 BERT checkpoint 的依赖在RUNMLPERF模式下model.load_from_pretrained(init_ckpt)会读取官方model.ckpt-28252预训练权重见 examples/mlperf/model_train.py。说明README 中还提到 tinybox_greenNVIDIA 平台需要按对应文档安装 p2p 驱动并注明这是生产环境 tinybox green 的默认配置该驱动安装仅适用于 green 平台tinybox_red 无需此步骤。4. 数据准备下载、校验与预处理4.1 下载原始数据含校验README 给出的第一步是下载 2020 版 Wikipedia 原始语料MLPerf BERT 官方使用的数据集并启用校验BASEDIR/raid/datasets/wiki WIKI_TRAIN1 VERIFY_CHECKSUM1 python3 extra/datasets/wikipedia_download.py对应源码 extra/datasets/wikipedia_download.py 会依次完成下载bert_config.json、vocab.txt、model.ckpt-28252data/index/meta 三个文件及checkpoint索引文件这些是官方预训练 checkpoint 与分词词表当WIKI_TRAIN1时额外下载bert_reference_results_text_md5.txt官方 md5 清单和results_text.tar.gz约 500 个分片的原始文本并解压到results4/目录当VERIFY_CHECKSUM1时调用verify_checksum()对results4/下每个文件逐块计算 MD5 并与官方清单比对不匹配会抛出ValueError见 extra/datasets/wikipedia_download.py。三个环境变量的含义如下环境变量默认值作用BASEDIRextra/datasets/wiki数据根目录README 生产环境固定为/raid/datasets/wikiWIKI_TRAIN0是否下载约 500 分片的大型训练语料results_text.tar.gzVERIFY_CHECKSUM0解压后是否逐文件校验 MD54.2 预处理训练与验证数据下载完成后进入预处理阶段。README 特别强调预处理使用的线程数受可用内存限制128GB RAM 下最多建议 16 线程。训练数据全量预处理500 个分片全部生成BASEDIR/raid/datasets/wiki NUM_WORKERS16 python3 extra/datasets/wikipedia.py pre-train all生成单个指定分片分片编号为 0499BASEDIR/raid/datasets/wiki python3 extra/datasets/wikipedia.py pre-train 42验证数据预处理生成 10000 条验证样本BASEDIR/raid/datasets/wiki python3 extra/datasets/wikipedia.py pre-eval预处理脚本 extra/datasets/wikipedia.py 是 MLPerf 官方create_pretraining_data.py的改写版本其内部流程见该文件 L384-L398 的__main__入口与get_features_from_part()为分词Tokenizer实现完整的 BasicTokenizer空白切分、中文逐字切分、Unicode 规范化、去除重音、标点拆分与 WordpieceTokenizer词表来自BASEDIR/vocab.txt实例构造create_instances_from_document按MAX_SEQ_LENGTH512截断以 10% 概率生成短序列SHORT_SEQ_PROB构造[CLS] tokens_a [SEP] tokens_b [SEP]格式并随机决定tokens_b是否来自其他文档NSP 任务掩码create_masked_lm_predictions按 15% 概率挑选 token其中 80% 替换为[MASK]、10% 保持原词、10% 随机替换对应官方 BERT 掩码策略特征化instance_to_features填充到固定长度并输出input_ids / input_mask / segment_ids / masked_lm_positions / masked_lm_ids / masked_lm_weights / next_sentence_labels七个 numpy 数组以.pkl形式写入BASEDIR/train/{part}.pkl或BASEDIR/eval.pkl。预处理阶段还支持以下环境变量均在源码中以getenv读取环境变量默认值作用NUM_WORKERSmin(os.cpu_count(), 32)并行处理分片的进程数README 建议 128GB 内存下取 16MAX_SEQ_LENGTH512最大序列长度MAX_PREDICTIONS_PER_SEQ76每序列最大掩码预测数与模型训练日志一致MASKED_LM_PROB0.15token 掩码概率DUPE_FACTOR10同一文档用不同随机掩码重复生成的次数SHORT_SEQ_PROB0.1生成短序列的概率RANDOM_SEED12345预处理随机种子训练阶段加载train/*.pkl、验证阶段加载eval.pkl的逻辑分别由 wikipedia.py 中的get_wiki_train_files()带diskcache缓存和 examples/mlperf/dataloader.py 中的batch_load_train_bert/batch_load_val_bert完成。5. 运行基准run_and_time.sh 逐行解读数据就绪后在仓库根目录执行 README 指定的脚本即可运行完整基准examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh注README 中引用的是training_submission_v5.0路径当前 v5.1 提交目录下存在内容一致的脚本 examples/mlperf/training_submission_v5.1/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh两者均可执行。该脚本是基准的“大脑”其完整内容如下我们逐段拆解#!/bin/bash set -e # Exit on any error set -o pipefail # Make pipeline fail if any command fails export PYTHONPATH. DEVAMD export MODELbert export SUBMISSION_PLATFORMtinybox_red export DEFAULT_FLOATHALF SUM_DTYPEHALF GPUS6 BS90 EVAL_BS90 export IGNORE_OOB1 export BEAM5 BEAM_UOPS_MAX8000 BEAM_UPCAST_MAX256 BEAM_LOCAL_MAX1024 BEAM_MIN_PROGRESS5 export IGNORE_JIT_FIRST_BEAM1 export BASEDIR/raid/datasets/wiki # pip install -e .[mlperf] export LOGMLPERF1 export SEED$RANDOM DATETIME$(date %m%d%H%M) LOGFILEbert_red_${DATETIME}_${SEED}.log export HCQDEV_WAIT_TIMEOUT_MS100000 # prevents hang? # init sleep 5 sudo rmmod amdgpu || true BENCHMARK10 INITMLPERF1 BERT_LAYERS2 python3 examples/mlperf/model_train.py | tee $LOGFILE # run PARALLEL0 RUNMLPERF1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE5.1 平台与模型开关环境变量取值含义PYTHONPATH.使仓库根目录可被 Python 导入DEVAMD选择 AMD 后端设备tinygrad 多后端架构中的 Device.DEFAULTMODELbert指定训练模型入口脚本据此分发到train_bert()SUBMISSION_PLATFORMtinybox_red写入 MLPerf 合规日志的提交平台名DEFAULT_FLOAT/SUM_DTYPEHALF默认浮点精度与规约累加精度为 FP16半精度训练GPUS6使用 6 张 GPU对应[f{Device.DEFAULT}:{i} for i in range(getenv(GPUS, 1))]见 examples/mlperf/model_train.pyBS/EVAL_BS90训练/验证批次大小IGNORE_OOB1忽略越界out-of-bounds校验避免长序列样本触发索引断言BASEDIR/raid/datasets/wiki指向第 4 节准备好的数据目录5.2 Beam Search 内核搜索参数tinygrad 的编译优化依赖 beam search 在巨大的调度空间中搜索更优内核run_and_time.sh为 tinybox_red 调优了一组参数环境变量取值含义BEAM5beam search 宽度每步保留的候选内核数BEAM_UOPS_MAX8000单个内核 uop 数量上限BEAM_UPCAST_MAX256upcast向量化提升上限BEAM_LOCAL_MAX1024local 维度上限BEAM_MIN_PROGRESS5beam 搜索最小进度阈值IGNORE_JIT_FIRST_BEAM1忽略 JIT 首次编译时的 beam 结果避免把编译噪声计入候选这些变量由 tinygrad/helpers.py 中的BEAM等全局配置读取最终作用于代码生成阶段的调度搜索tinygrad/codegen 目录。对比同仓库其他平台脚本可知tinybox_greenNVIDIA使用BEAM8 BEAM_UOPS_MAX10000tinybox_8xMI300X 使用BEAM3 BEAM_UOPS_MAX6000可见这是一组与 GPU 架构强相关的调优参数不可随意跨平台套用。5.3 日志与随机种子export LOGMLPERF1 export SEED$RANDOM DATETIME$(date %m%d%H%M) LOGFILEbert_red_${DATETIME}_${SEED}.logLOGMLPERF1启用mlperf_logging.mllog合规日志训练过程中会按 MLPerf 规范记录SUBMISSION_ORG / SUBMISSION_PLATFORM / GLOBAL_BATCH_SIZE / OPT_NAME / RUN_STOP等关键事件见 examples/mlperf/model_train.py 与 L1056-L1077并生成result_bert_{seed}.logSEED$RANDOM每次运行使用随机种子保证提交的多次运行可复现性统计日志文件以bert_red_{月日时分}_{种子}.log命名两阶段输出通过tee追加写入同一文件。5.4 AMD 设备预处理export HCQDEV_WAIT_TIMEOUT_MS100000 # prevents hang? sleep 5 sudo rmmod amdgpu || trueHCQDEV_WAIT_TIMEOUT_MS用于调高 AMD HCQ 设备等待超时注释表明用于防止挂起rmmod amdgpu用于在启动前重置 AMDGPU 内核模块状态|| true保证模块未加载时不会因set -e中断脚本。5.5 两阶段执行init 与 run脚本主体分两次调用训练入口这是 MLPerf 时间合规的关键设计BENCHMARK10 INITMLPERF1 BERT_LAYERS2 python3 examples/mlperf/model_train.py | tee $LOGFILE PARALLEL0 RUNMLPERF1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE第一阶段initINITMLPERF1配合BENCHMARK10只跑 10 步使用伪造数据get_fake_data_bert见 examples/mlperf/helpers.py的训练与验证同时BERT_LAYERS2将模型缩到 2 层用于完成 JIT 捕获、内核 beam 搜索预热与 MLPerfINIT_START → INIT_STOP事件记录。init 阶段不计入基准时钟第二阶段runRUNMLPERF1加载真实数据与官方 checkpoint 开始正式训练PARALLEL0关闭并行数据加载相关路径此时MLLOGGER记录RUN_START / EPOCH_START / EVAL_ACCURACY / RUN_STOP(SUCCESS)当验证集 masked LM 准确率达到TARGET0.72时停止计时并输出收敛耗时见 examples/mlperf/model_train.py。训练循环本身由TinyJit编译的train_step_bert驱动examples/mlperf/model_train.py内部完成LAMB 优化器梯度更新、loss_scalerFP16 下为2**11缩放与反缩放、全局梯度范数计算与梯度裁剪global_norm 1.0时按范数缩放。超参数方面max_lr 0.000175 * sqrt(GBS/96)train_steps 3600000 // GBS优化器使用 LAMBadamFalse并配合PolynomialDecayWithWarmup多项式衰减学习率调度examples/mlperf/model_train.py。6. 开发辅助脚本dev_run.sh 与 dev_beam.sh除正式基准脚本外同一目录还提供了两个开发用途脚本dev_run.sh与run_and_time.sh相同的平台/beam 配置但以WANDB1 PARALLEL0 RUNMLPERF1直接进入训练适合配合 wandb 进行实验跟踪与调参dev_beam.sh面向内核优化的调试脚本额外设置BEAM_LOG_SURPASS_MAX1记录 beam 搜索中超越最优的日志、RESET_STEP1每次 eval 前重置 JIT 释放显存、DEBUG2输出调试信息并以BENCHMARK10 BERT_LAYERS2的小规模配置运行用于在完整训练前验证 beam 参数组合。7. 同提交下的平台对照v5.1 提交目录 examples/mlperf/training_submission_v5.1/tinycorp/benchmarks/bert/implementations/ 下共包含三个运行平台方便对比 tinybox_red 的配置差异平台后端GPU 数量BS/EVAL_BSBEAM特性tinybox_redDEVAMD690 / 905本文主体RX 7900 XTX含HCQDEV_WAIT_TIMEOUT_MS防挂起配置与rmmod amdgpu重置tinybox_greenDEVNV690 / 908NVIDIA 平台见 run_and_time.shtinybox_8xMI300XDEVAMD81024 / 10243MI300X 平台显存更大故 BS 大幅提升并显式调优OPT_BASE_LEARNING_RATE / OPT_LAMB_BETA_1 / OPT_LAMB_BETA_2 / DECAY / TRAIN_STEPS见 run_and_time.sh三者的数据准备流程第 4 节完全一致差异集中在设备后端、batch size 与 beam 参数这正体现了 tinygrad 同一套训练代码examples/mlperf/model_train.py 的train_bert()通过环境变量驱动即可适配不同硬件的能力。8. 总结tinybox_red 上的 MLPerf BERT 基准复现可以归纳为一条清晰的链路安装pip install -e .[mlperf]gdown/numpy/tqdm/tensorflow数据wikipedia_download.py下载 MD5 校验→wikipedia.py pre-train all / pre-train N / pre-eval分词、掩码、NSP 实例化与特征化运行run_and_time.sh以“initbeam 预热 时钟外初始化→ run真实训练 合规日志”两阶段驱动 examples/mlperf/model_train.py在 6 张 RX 7900 XTX 上以 FP16 LAMB 梯度裁剪完成 BERT-large 预训练达标masked LM acc ≥ 0.72后停止并记录合规日志。整个流程中GPUS/BS/EVAL_BS/DEFAULT_FLOAT决定训练配置BEAM系列变量决定内核搜索强度LOGMLPERF/INITMLPERF/RUNMLPERF决定基准合规语义——理解这些环境变量的作用均可对照 examples/mlperf/model_train.py 与 extra/datasets/wikipedia.py 源码中的getenv读取点你就掌握了在 tinybox_red 上复现并扩展该基准的完整能力。【免费下载链接】tinygradYou like pytorch? You like micrograd? You love tinygrad! ❤️项目地址: https://gitcode.com/GitHub_Trending/tiny/tinygrad创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考