tinygrad 在 tinybox 上运行 MLPerf BERT 训练基准数据准备、环境配置与完整执行指南【免费下载链接】tinygradYou like pytorch? You like micrograd? You love tinygrad! ❤️项目地址: https://gitcode.com/GitHub_Trending/tiny/tinygrad本篇指南围绕 examples/mlperf/training_submission_v4.1/tinycorp/benchmarks/bert/implementations/tinybox_green/README.md 展开讲解如何用 tinygrad 在 tinybox green6×RTX 4090与 tinybox red6×RX 7900 XTX两台硬件平台上完成 MLPerf Training v4.1 BERT掩码语言建模预训练基准的复现从环境安装、维基百科语料下载与预处理到一键运行run_and_time.sh并产出 MLPerf 合规日志。读完本文你将掌握该提交的完整执行链路以及背后model_train.py中 BERT 训练循环、LAMB 优化器、Beam Search 编译参数与日志系统的实现细节。1. 基准问题与提交概况该 README 描述的 Problem 是使用 BERTBidirectional Encoder Representations from Transformers进行 NLP 预训练属于 MLPerf Training v4.1 的 closed division、available on-premise 提交提交方为 tinycorp。目录examples/mlperf/training_submission_v4.1/tinycorp/下包含两个实现tinybox_green/NVIDIA 平台的 BERT 实现另有 resnet 实现。tinybox_red/AMD 平台的 BERT 实现额外带一个setup.sh硬件调优脚本。对应的硬件规格记录在 systems/tinybox_green.json 与 systems/tinybox_red.json 中两份配置高度对称配置项tinybox greentinybox redCPUAMD EPYC 7532 32-Core Processor1 路64 vCPU同左内存128GB8×16GB DDR4同左存储4TB RAID 阵列 1TB 启动盘NVMe SSD同左加速卡6× NVIDIA GeForce RTX 409024GB GDDR6X6× AMD Radeon RX 7900 XTX24GB GDDR6卡间/主机互连PCIe 4.0 x16同左操作系统Ubuntu 22.04.4同左软件栈Python 3.10.12CUDA 12.4Python 3.10.12ROCm 6.1.3tinygrad 版本commitb5546912e24e0a864b35924da4efa5d71cfe368b同左两平台均为单节点 6 卡、24GB 显存/卡因此训练超参GPUS6、BS66、EVAL_BS6完全一致仅底层驱动与 Beam Search 配置有差异见下文。2. 环境要求与软件安装README 明确要求从 master 分支安装 tinygrad 与 mlperf-logginggit clone https://github.com/tinygrad/tinygrad.git python3 -m pip install -e .[mlperf].[mlperf]是 pyproject.toml 中声明的可选依赖组用于拉取 mlperf-logging 等基准所需依赖。随后安装两个额外包pip install tqdm tensorflowtqdm训练进度条训练脚本直接from tqdm import tqdm见 examples/mlperf/model_train.py。tensorflowBERT 预训练数据预处理阶段生成 TFRecord 数据所需。2.1 tinybox_greenP2P 驱动README 要求按照 tinygrad 的 open-gpu-kernel-modules550.54.15-p2p 分支说明安装 P2P 驱动。该驱动是生产环境 tinybox green 的默认配置This is the default on production tinybox green。2.2 tinybox_red自定义 amdgpu 驱动tinybox red 需要禁用 cwsr并发 wave 调度相关特性并增大 MESMicro Engine Scheduler超时安装针对 Ubuntu 22.04 定制的 amdgpu 驱动v6.1.3。这也是 setup.sh 存在的背景。3. 数据准备下载与校验BERT 预训练使用维基百科语料。下载脚本位于 extra/datasets/wikipedia_download.pyBASEDIR/raid/datasets/wiki WIKI_TRAIN1 VERIFY_CHECKSUM1 python3 extra/datasets/wikipedia_download.pyBASEDIR数据存放根目录训练脚本默认BASEDIR/raid/datasets/wiki与此保持一致。WIKI_TRAIN1下载训练用维基转储。VERIFY_CHECKSUM1下载完成后校验校验和确保数据完整。4. 数据预处理预处理脚本为 extra/datasets/wikipedia.py它基于 MLCommons 官方训练仓库中 BERT 的create_pretraining_data.py修改而来用于把下载的维基文本切分成带掩码的预训练样本。脚本头部注释列出了可调环境变量MAX_SEQ_LENGTH最大序列长度训练时固定为 512。MAX_PREDICTIONS_PER_SEQ每条序列最多掩码预测数训练时固定为 76。RANDOM_SEED随机种子。DUPE_FACTOR输入数据以不同掩码重复的次数。MASKED_LM_PROBtoken 掩码概率。SHORT_SEQ_PROB选取短于MAX_SEQ_LENGTH序列的概率。4.1 训练数据BASEDIR/raid/datasets/wiki NUM_WORKERS16 python3 extra/datasets/wikipedia.py pre-train all注意 README 的提示预处理线程数受可用内存限制128GB 内存下最多建议 16 线程NUM_WORKERS16。如果只想生成某个特定 topic编号 0~499BASEDIR/raid/datasets/wiki python3 extra/datasets/wikipedia.py pre-train 424.2 验证数据BASEDIR/raid/datasets/wiki python3 extra/datasets/wikipedia.py pre-eval5. 运行基准5.1 tinybox_green直接运行examples/mlperf/training_submission_v4.1/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh5.2 tinybox_red先做一次性硬件设置再运行examples/mlperf/training_submission_v4.1/tinycorp/benchmarks/bert/implementations/tinybox_red/setup.sh examples/mlperf/training_submission_v4.1/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.shsetup.sh通过 rocm-smi 将 GPU 调整到基准所需状态rocm-smi --setprofile compute rocm-smi --setmclk 3 rocm-smi --setperflevel high其中--setprofile compute切换到计算 profile--setmclk 3锁定显存时钟档位--setperflevel high将性能等级设为 high。脚本中还注释了一个可选的功耗上限操作将 6 张卡统一限制到 350W# echo 350000000 | sudo tee /sys/class/drm/card{1..6}/device/hwmon/hwmon*/power1_cap6. run_and_time.sh 逐段解析两份run_and_time.sh结构相同仅平台名与 Beam 参数不同以 green 版 run_and_time.sh 为例export PYTHONPATH. export MODELbert export SUBMISSION_PLATFORMtinybox_green export DEFAULT_FLOATHALF GPUS6 BS66 EVAL_BS6 export BEAM4 BEAM_UOPS_MAX2000 BEAM_UPCAST_MAX64 BEAM_LOCAL_MAX512 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_green_${DATETIME}_${SEED}.log # init BENCHMARK10 INITMLPERF1 python3 examples/mlperf/model_train.py | tee $LOGFILE # run PARALLEL0 RUNMLPERF1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE关键环境变量说明变量greenred作用MODELbertbert选择 model_train.py 中的train_bert()入口SUBMISSION_PLATFORMtinybox_greentinybox_red写入 mllog 的SUBMISSION_PLATFORM字段DEFAULT_FLOATHALFHALF默认浮点精度 fp16配合 loss scaling 训练GPUS66使用 6 张卡脚本据此构造CUDA:0..5/AMD:0..5BS/EVAL_BS66 / 666 / 6训练/验证 batch size单卡 BS11、EVAL_BS1 的 6 卡总和BEAM43Beam Search 编译搜索宽度BEAM_UOPS_MAX/BEAM_UPCAST_MAX/BEAM_LOCAL_MAX2000 / 64 / 512仅 green限制 Beam 搜索的 UOps 数、upcast 数与局部 buffer 大小用于控制编译开销IGNORE_JIT_FIRST_BEAM11JIT 首次捕获时忽略 Beam 搜索避免首步编译失真BASEDIR/raid/datasets/wiki同左预处理后数据的根目录LOGMLPERF11启用 mlperf-loggingmllog合规日志SEED$RANDOM$RANDOM每次运行随机种子决定初始化与数据顺序脚本分两阶段执行输出统一追加到bert_{green|red}_{MMDDHHMM}_{SEED}.log初始化阶段BENCHMARK10 INITMLPERF1—— 只跑 10 步 fake dataget_fake_data_bert生成随机输入来清空磁盘缓存diskcache_clear()、预热并触发 Beam 编译同时把INIT_START/INIT_STOP等初始化事件写入 mllog。正式运行阶段PARALLEL0 RUNMLPERF1—— 加载真实维基数据训练并在达到目标精度后记录RUN_STOP。7. 训练脚本中的 BERT 实现细节train_bert()examples/mlperf/model_train.py 中MODELbert时执行是本次基准的核心以下参数均可通过环境变量覆盖超参数默认值来自源码BS默认11 * len(GPUS)fp16/bf16 时即 6 卡共 66EVAL_BS默认1 * len(GPUS)。全局 batch sizeGBS BS * grad_acc其中GRADIENT_ACC_STEPS默认 1。学习率OPT_BASE_LEARNING_RATE 0.000175 * sqrt(GBS/96)随 batch 缩放。TRAIN_STEPS 3600000 // GBSNUM_WARMUP_STEPS 1POLY_POWER 1.0多项式衰减。优化器LAMBLAMBadamFalseweight_decay0.01beta10.9beta20.999epsilon1e-6bias 与 LayerNorm 参数进入独立的LAMB组weight_decay0两组通过OptimizerGroup联合调度LR 由PolynomialDecayWithWarmup管理。损失缩放fp16 下LOSS_SCALER 2**11backward 前乘、更新前除回并实现按 global norm 的梯度裁剪global_norm 1.0时归一化。评估MAX_EVAL_STEPS保证EVAL_BS * MAX_EVAL_STEPS 10000目标为TARGET 0.72的 masked LM accuracy。模型配置MAX_SEQUENCE_LENGTH512、max_predictions_per_seq76写入 mllog。训练/评估步骤使用TinyJitContext(TRAINING1/0)编译TinyJit Context(TRAINING1) def train_step_bert(input_ids, segment_ids, attention_mask, masked_positions, masked_lm_ids, masked_lm_weights, next_sentence_labels): for t in [input_ids, segment_ids, attention_mask, masked_positions, masked_lm_ids, masked_lm_weights, next_sentence_labels]: if len(GPUS) 1: t.shard_(GPUS, axis0) else: t.to_(GPUS[0]) ...多卡时输入张量按 batch 维度shard_(GPUS, axis0)切分到 6 卡权重通过p.to_(GPUS)复制到全部设备数据并行。评估步骤eval_step_bert同样把结果聚合回 CPU 后Tensor.realize。数据流RUNMLPERF1时用batch_load_train_bert(BS, seedseed)/batch_load_val_bert(EVAL_BS)examples/mlperf/dataloader.py读取预处理后的数据INITMLPERF阶段则用get_fake_data_bert重复产出假数据以完成编译预热。检查点与收敛每SAVE_CKPT_FREQ默认 1000步保存.safe检查点并只保留最近KEEP_CKPT_AMOUNT默认 5份当 masked LM accuracy 首次达到 0.72 时打印收敛耗时、写入RUN_STOP(statusSUCCESS)并保存bert-large.safe后提前终止。8. 日志与合规输出LOGMLPERF1时脚本使用mlperf_logging.mllog输出合规事件包括SUBMISSION_ORGtinycorp、SUBMISSION_PLATFORM、SUBMISSION_DIVISIONCLOSED、SUBMISSION_STATUSONPREM以及 GLOBAL_BATCH_SIZE、OPT_NAMELAMB、LARS/LAMB 相关超参、EVAL_ACCURACY、RUN_STOP 等标准键结果写入result_bert_{seed}.txt。同时终端输出被tee到bert_{green|red}_{MMDDHHMM}_{SEED}.log每一步会打印运行时间、Python 开销、取数开销、显存占用与 GFLOPS 等诊断信息便于对照BENCHMARK10阶段的预计总训练时长评估进度。9. 快速排查要点若预处理阶段内存不足先降低NUM_WORKERS128GB 上限为 16。若INITMLPERF阶段编译过慢可调小BEAMred 版即从 4 降到 3或调低BEAM_UOPS_MAX等上限IGNORE_JIT_FIRST_BEAM1可避免首次 JIT 捕获被 Beam 搜索拖慢。若评估时报EVAL_BS * MAX_EVAL_STEPS断言失败说明验证集覆盖不足需增大MAX_EVAL_STEPS。tinybox red 若出现 GPU 挂起/超时类问题需重新确认自定义 amdgpu 驱动的 cwsr 禁用与 MES 超时配置并重跑setup.sh。综上这套 v4.1 BERT 提交把数据→编译→训练→合规日志整条 MLPerf 流水线收敛到两个 shell 脚本与一个 model_train.py 入口中是研究 tinygrad 在多卡异构平台上复现行业基准的完整范本。【免费下载链接】tinygradYou like pytorch? You like micrograd? You love tinygrad! ❤️项目地址: https://gitcode.com/GitHub_Trending/tiny/tinygrad创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考