tinygrad 在 tinybox 8xMI300X 上运行 MLPerf BERT 训练基准环境搭建、数据管线与 8 卡调优实战【免费下载链接】tinygradYou like pytorch? You like micrograd? You love tinygrad! ❤️项目地址: https://gitcode.com/GitHub_Trending/tiny/tinygrad本篇文章以 tinygrad 仓库中 BERT 基准说明文档 为骨架完整讲解如何在 8x AMD Instinct MI300X 节点tinybox 8xMI300X上复现 MLPerf BERT 大规模预训练从分支安装、Wikipedia 数据下载与预处理到 run_and_time.sh 一键运行与逐项环境变量解析并下沉到 model_train.py 的训练循环与 LAMB 优化器实现帮助读者理解 tinygrad 如何用一套 Python 代码完成封闭分区的 BERT 训练提交。1. 问题定义用 BERT 做 NLP 的 MLPerf 训练基准MLPerf 训练基准Training benchmark要求参与者在规定的数据、模型与精度约束下训练出达到目标精度的模型并记录耗时。本目录提交的是BERTBidirectional Encoder Representations from Transformers用于自然语言处理NLP的掩码语言模型与下一句预测预训练任务对应官方 BERT-large 规模的训练负载。tinygrad 以tinycorp作为 submitter、closed分区、Available on-premise状态参与提交模型在纯 tinygrad 框架内实现训练框架版本为mlperf_training_v5.0分支。整体提交说明可参见 examples/mlperf/README其中列举了 ResNet50-v1.5、RetinaNet、3D UNet、RNNT、BERT-large 五个目标模型本 README 聚焦其中的 BERT。2. 目标硬件与软件栈tinybox 8xMI300X运行该基准的目标节点是 tinybox 8xMI300X其完整硬件规格记录在 tinybox_8xMI300X.json 中组件规格主机处理器2x AMD EPYC 9354每颗 32 核 / 64 vCPU主机内存2304GB24x 96GB DDR5存储NVMe SSD3x 4TB RAID 阵列加速器8x AMD Instinct MI300X192GB HBM3加速器互联PCIe 5.0 x16操作系统Ubuntu 24.04.1 LTS运行时Python 3.10.16ROCm 3.0.094441cb框架tinygrad分支 mlperf_training_v5.0可以看到8 张 MI300X 各带 192GB HBM3配合 2.3TB 主机内存为 BERT-large 这类大规模预训练提供了充裕的显存与内存预算。基准脚本中BS1024、EVAL_BS1024的大批量设置正是建立在这一显存规模之上。注意当前仓库中的提交目录为training_submission_v5.1README 内文引用的脚本路径写作training_submission_v5.0实际可执行脚本位于examples/mlperf/training_submission_v5.1/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/下本文一律使用当前仓库中真实存在的路径。3. 环境准备安装 tinygrad、mlperf-logging 与依赖3.1 安装 tinygradmlperf 分支基准要求在mlperf_training_v5.0分支上安装 tinygrad并启用[mlperf]扩展该扩展会引入 mlperf-logging即代码中的“uncomment mlperf from setup.py”git clone https://github.com/tinygrad/tinygrad.git python3 -m pip install -e .[mlperf]pip install -e以可编辑模式安装便于在基准代码与 tinygrad 本体之间同步迭代。日志记录依赖mlperf_logging.mllog包在 model_train.py 中通过from mlperf_logging import mllog导入用于生成符合 MLPerf 规范的result_bert_seed.log事件流。3.2 安装数据集与辅助依赖pip install gdown numpy tqdm tensorflowgdown从 Google Drive 下载 Wikipedia 原始语料包numpy预处理与张量搬运tqdm训练/预处理进度条训练循环中from tqdm import tqdm显式导入见 model_train.pytensorflowBERT 训练数据加载采用 TFRecord 风格的 interleave 打乱见下文 dataloader。3.3 关于 tinybox_green 的 p2p 驱动README 特别注明在 tinybox_green 上运行需要按 NVIDIA 开源内核模块的 p2p 分支安装 P2P 直连驱动且该驱动是生产版 tinybox green 的默认配置。本主题的 tinybox_8xMI300X 是 AMD 平台由DEVAMD环境变量选择 AMD 后端因此不需要该 NVIDIA 驱动——这属于文档针对不同硬件变体的差异化说明。4. 数据准备下载并校验 Wikipedia 原始语料4.1 下载原始数据BERT 预训练数据来自 Wikipedia dump。执行 wikipedia_download.pyBASEDIR/raid/datasets/wiki WIKI_TRAIN1 VERIFY_CHECKSUM1 python3 extra/datasets/wikipedia_download.pyBASEDIR数据集根目录/raid/datasets/wiki是 tinybox 的约定路径应与后续所有预处理、训练命令保持一致WIKI_TRAIN1下载训练用 Wikipedia dumpVERIFY_CHECKSUM1下载完成后校验校验和确保语料完整。4.2 预处理训练数据原始 dump 需要经过 tokenize、masking 等处理生成 BERT 预训练样本这一步由 wikipedia.py 完成BASEDIR/raid/datasets/wiki NUM_WORKERS16 python3 extra/datasets/wikipedia.py pre-train allpre-train all表示预处理全部 500 个主题topic 编号 0~499。也可以只生成某一个主题做快速验证BASEDIR/raid/datasets/wiki python3 extra/datasets/wikipedia.py pre-train 42内存限制提示README 明确指出预处理线程数受可用内存限制——128GB 内存的机器建议最多 16 线程NUM_WORKERS16线程数过高会导致内存耗尽。tinybox 8xMI300X 拥有 2304GB 主机内存可安全采用 16 线程。4.3 预处理验证数据BASEDIR/raid/datasets/wiki python3 extra/datasets/wikipedia.py pre-eval验证集会被序列化为eval.pkl训练时由 dataloader.py 的batch_load_val_bert循环读取BASEDIR/eval.pkl并按EVAL_BS分批、越界回绕。4.4 预处理参数环境变量说明预处理脚本头部的注释wikipedia.py列出了可调环境变量用于控制 BERT 预训练样本的生成方式与 MLPerf 参考实现create_pretraining_data.py一致环境变量作用MAX_SEQ_LENGTH最大序列长度训练配置为 512MAX_PREDICTIONS_PER_SEQ每条序列最多掩码预测数训练配置为 76RANDOM_SEED随机种子保证样本可复现DUPE_FACTOR同一输入用不同掩码重复生成的次数MASKED_LM_PROB掩码概率SHORT_SEQ_PROB采样短于MAX_SEQ_LENGTH序列的概率5. 一键运行基准run_and_time.shREADME 给出的运行方式即执行提交脚本examples/mlperf/training_submission_v5.1/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh同一目录下的 tinybox_green、tinybox_red 变体脚本位于 implementations/tinybox_green/ 与 implementations/tinybox_red/。5.1 脚本全貌该脚本在仓库中的完整内容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_8xMI300X export DEFAULT_FLOATHALF GPUS8 BS1024 EVAL_BS1024 # 超参数参照 MLPerf v3.1 参考结果 export OPT_BASE_LEARNING_RATE0.0011 OPT_LAMB_BETA_10.60466 OPT_LAMB_BETA_20.85437 DECAY0.1 export TRAIN_STEPS3900 export IGNORE_OOB1 export BEAM3 BEAM_UOPS_MAX6000 BEAM_UPCAST_MAX256 BEAM_LOCAL_MAX1024 BEAM_MIN_PROGRESS5 export IGNORE_JIT_FIRST_BEAM1 FREE_INTERMEDIATE0 export BASEDIR/raid/datasets/wiki # pip install -e .[mlperf] export LOGMLPERF1 export SEED$RANDOM DATETIME$(date %m%d%H%M) LOGFILEbert_8xMI300x_${DATETIME}_${SEED}.log BENCHMARK10 INITMLPERF1 BERT_LAYERS2 python3 examples/mlperf/model_train.py | tee $LOGFILE # run PARALLEL0 RUNMLPERF1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE5.2 运行模式与后端PYTHONPATH.从仓库根目录导入examples/mlperf与tinygrad包DEVAMD选择 AMD 后端tinygrad 通过Device.DEFAULT解析训练循环中构造[f{Device.DEFAULT}:{i} for i in range(GPUS)]model_train.pyMODELbertmodel_train.py依据该变量分派到train_bert()SUBMISSION_PLATFORMtinybox_8xMI300X写入 MLPerf 日志的提交平台字段DEFAULT_FLOATHALF默认浮点精度为 FP16配合LOSS_SCALER2^11见 model_train.py做混合精度训练GPUS88 卡数据并行权重与梯度通过p.to_(GPUS)/t.shard_(GPUS, axis0)分布在 8 个设备上BS1024、EVAL_BS1024训练/评估全局批量。5.3 优化器与学习率超参export OPT_BASE_LEARNING_RATE0.0011 OPT_LAMB_BETA_10.60466 OPT_LAMB_BETA_20.85437 DECAY0.1 export TRAIN_STEPS3900OPT_BASE_LEARNING_RATE0.0011LAMB 优化器初始学习率注释注明参照 MLPerf v3.1 官方参考实现Quanta D54U-3U 的 BERT 结果OPT_LAMB_BETA_1 / OPT_LAMB_BETA_2LAMB 的动量与二阶矩衰减系数取值与参考实现一致DECAY0.1LAMB 权重衰减系数TRAIN_STEPS3900总训练步数全局批量 1024单轮覆盖约 400 万样本。这些值在 model_train.py 中通过getenv读取并会通过mllog以OPT_BASE_LR、OPT_LAMB_BETA_1、OPT_LAMB_BETA_2、OPT_LAMB_WEIGHT_DECAY等键记录model_train.py。5.4 编译调优参数Beam Search 与 JITexport BEAM3 BEAM_UOPS_MAX6000 BEAM_UPCAST_MAX256 BEAM_LOCAL_MAX1024 BEAM_MIN_PROGRESS5 export IGNORE_JIT_FIRST_BEAM1 FREE_INTERMEDIATE0这是 tinygrad 特有的“编译搜索”参数直接影响 kernel 生成质量与显存占用BEAM3beam search 宽度tinygrad 会对每个算子尝试多种 schedule 组合选择局部最优 kernel对应BEAM.value训练循环中通过BEAM.value TRAIN_BEAM切换BEAM_UOPS_MAX6000beam 搜索中允许的 uop 上限约束搜索空间BEAM_UPCAST_MAX256upcast向量化/线程展开上限BEAM_LOCAL_MAX1024local size 上限BEAM_MIN_PROGRESS5beam 剪枝的最小进展阈值IGNORE_JIT_FIRST_BEAM1忽略 JIT 捕获阶段的首次 beam 结果首次编译不计入最终 kernel 选择FREE_INTERMEDIATE0默认不释放 JIT 捕获的中间张量释放中间显存的功能在 model_train.py 的 eval 分支中可选项启用注释提示该特性在 tiny green 上长时间训练后存在挂起问题。IGNORE_OOB1表示关闭越界检查out-of-bounds validation减少运行时校验开销BASEDIR/raid/datasets/wiki指向已预处理的数据目录。5.5 两阶段执行INITMLPERF 与 RUNMLPERF脚本将一次正式提交拆成两个进程BENCHMARK10 INITMLPERF1 BERT_LAYERS2 python3 examples/mlperf/model_train.py | tee $LOGFILE PARALLEL0 RUNMLPERF1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE第一阶段初始化INITMLPERF1 BENCHMARK10 BERT_LAYERS2用 2 层小 BERT 跑 10 步基准完成 MLPerf 日志初始化清空 disk cache、写入 submission 元数据、INIT_START/INIT_STOP事件并打印预估总训练时间第二阶段正式运行RUNMLPERF1 PARALLEL0加载预训练 checkpoint 与真实数据执行完整 3900 步训练直到 masked LM 精度达到TARGET默认 0.72或跑满步数期间写RUN_START、EPOCH_*、EVAL_*、RUN_STOP等标准事件。日志同时以tee写入bert_8xMI300x_MMDDHHMM_SEED.logSEED$RANDOM保证每次提交独立随机种子。LOGMLPERF1开启mlperf_logging输出生成result_bert_seed.logmodel_train.py。6. 训练循环源码解析model_train.py 中的 BERT 路径6.1 超参数解析train_bert() 通过getenv解析全部超参除脚本已导出的外还有一批带默认值的可选项参数默认值说明BSFP16 下11*GPUS训练批量脚本覆盖为 1024GRADIENT_ACC_STEPS1梯度累积步数当前仅支持 1有断言GBSBS*grad_acc全局批量EVAL_BS1*GPUS评估批量脚本覆盖为 1024NUM_WARMUP_STEPS1学习率 warmup 步数MAX_EVAL_STEPSceil(10000/EVAL_BS)保证评估覆盖至少 10000 样本EVAL_STEP_FREQ按公式计算评估频率SAVE_CKPT_FREQ/KEEP_CKPT_AMOUNT1000 / 5检查点保存频率与保留份数LOSS_SCALERFP16 下2^11混合精度损失缩放DECAY/EPSILON/POLY_POWER0.01 / 1e-6 / 1.0LAMB 权重衰减、epsilon、多项式衰减幂次TARGET0.72masked LM 精度达标线同时有断言10000 EVAL_BS * MAX_EVAL_STEPS确保评估覆盖完整验证集。6.2 模型、优化器与学习率调度模型由 helpers.py 的get_mlperf_bert_model()构造替换LinearBert、EmbeddingBert、LayerNormBert初始化器后实例化BertForPretrainingDISABLE_DROPOUT1可关闭 dropout 以提升确定性参数按“是否含 bias/LayerNorm”分组带权重衰减组与无权重衰减组各用一个LAMB优化器adamFalse走 LAMB 原生路径合并为OptimizerGroupmodel_train.py学习率用PolynomialDecayWithWarmup幂次POLY_POWER1.0即线性衰减到 0两组调度器合并为LRSchedulerGroup训练步函数train_step_bert以TinyJitContext(TRAINING1)编译为图数据张量按axis0shard 到 8 卡前向得到lm_logits与seq_relationship_logits计算 loss 后乘loss_scaler反传再除以loss_scaler还原梯度model_train.py。6.3 梯度裁剪与收敛判定训练步内实现了全局梯度范数裁剪对所有参数梯度求平方和开根号得global_norm若global_norm 1.0则按grad/global_norm缩放model_train.py这是 BERT 大模型稳定训练的关键环节。训练循环每步打印 step time拆分为 python / fetch data / device 三段耗时、loss、LR、显存占用与 GFLOPS按EVAL_STEP_FREQ触发评估评估步eval_step_bert在Context(TRAINING0)下计算 masked LM / next sentence 的 accuracy 与 loss。当avg_lm_acc TARGET时记录RUN_STOP(statusSUCCESS)并保存bert-large.safecheckpointmodel_train.py打印参考收敛点的总耗时与样本数。6.4 数据加载与检查点训练数据batch_load_train_bertdataloader.py基于预处理生成的多个 train 文件用 TF 风格的随机 shuffle 构造InterleavedDatasetcycle length 取 CPU 核数与文件数的较小值验证数据batch_load_val_bert循环读取BASEDIR/eval.pkl非正式运行时无RUNMLPERF用get_fake_data_berthelpers.py生成全零假数据用于编译预热与基准测速检查点每SAVE_CKPT_FREQ步通过safe_save(get_training_state(...))保存超出的旧检查点自动清理支持RESUMEckpt从中断步恢复训练model_train.py。7. 开发调试脚本dev_run.sh 与 dev_beam.sh同一目录下还提供了两个面向开发者的脚本dev_beam.sh以BENCHMARK10 BERT_LAYERS2运行 beam search 编译基准用于快速验证 kernel 编译配置与预估训练时长不加载真实数据dev_run.sh开启WANDB1的实验运行版本可在 wandb 中跟踪 loss、LR、step time、显存与 GFLOPS 曲线适合超参调优阶段。两者的后端、批量、LAMB 超参与正式脚本保持一致DEVAMD GPUS8 BS1024 EVAL_BS1024及相同 beam 参数区别仅在于是否走 MLPerf 日志流程。8. 结果产物与验证要点一次完整运行结束后工作目录会产生bert_8xMI300x_MMDDHHMM_SEED.logtee 记录的完整控制台输出含每步耗时、loss、LR、GFLOPS 与评估精度result_bert_seed.logmlperf_logging.mllog生成的标准事件流供 MLPerf 官方工具校验submission org/platform/division、超参、eval accuracy、run stop 状态等ckpts/目录启用CKPT1时bert-large.safe或带时间戳的中间检查点可用 safe 格式 的safe_load加载。验证要点TARGET0.72的 masked LM 精度是否达成、RUN_STOP事件状态是否为 SUCCESS、训练步数是否在TRAIN_STEPS3900内完成以及日志中 step time 与显存曲线是否符合 8xMI300X 的预期水平。9. 小结本目录 README 篇幅精炼但串联起了 MLPerf BERT 提交的完整链路分支安装pip install -e .[mlperf]→ Wikipedia 下载校验wikipedia_download.py→ 训练/验证预处理wikipedia.py pre-train / pre-eval→run_and_time.sh两阶段执行INITMLPERF 预热 RUNMLPERF 正式训练。结合 model_train.py 与 dataloader.py 的源码可以看到tinygrad 用约两千行 Python 完成了 LAMB 优化、梯度裁剪、FP16 混合精度、8 卡数据并行、beam 编译搜索与 MLPerf 日志记录为复现 BERT 大规模预训练提供了一份完整、可审计的参考实现。【免费下载链接】tinygradYou like pytorch? You like micrograd? You love tinygrad! ❤️项目地址: https://gitcode.com/GitHub_Trending/tiny/tinygrad创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考