企业官网建设流程全解析

tinygrad 复现 MLPerf Training v5.1 BERT 基准：tinybox_green 平台全流程实战指南

【免费下载链接】tinygradYou like pytorch? You like micrograd? You love tinygrad! ❤️项目地址: https://gitcode.com/GitHub_Trending/tiny/tinygrad

本文以 tinygrad 仓库中 MLPerf Training v5.1 提交（training_submission_v5.1/tinycorp/benchmarks/bert/implementations/tinybox_green/README.md）为核心，完整讲解在 tinybox_green 上从零搭建 BERT（NLP 预训练）基准测试的依赖安装、Wikipedia 数据集下载与预处理、基准运行与日志输出的全过程。读完本文，你将掌握该提交的完整复现步骤、run_and_time.sh中每个关键环境变量的作用，以及底层训练脚本 model_train.py 中 LAMB 优化器、warmup 调度与梯度裁剪的实现细节。

一、基准任务概览：用 tinygrad 训练 BERT

该 README 是 tinygrad 为 MLPerf Training 提交（v5.1 目录下）中BERT 基准（Problem: BERT for NLP）提供的运行说明。MLPerf BERT 任务本质是BERT 预训练（masked language modeling + next sentence prediction）：模型需要从 Wikipedia 语料中学习语言表征，最终以 masked LM 与 next-sentence 两个任务的准确率作为质量指标。

在 model_train.py 的train_bert()中，默认质量目标为TARGET=0.72（对应masked_lm_accuracy + seq_relationship_accuracy的合成指标，源码见 L1006），训练默认步数为TRAIN_STEPS = 3600000 // GBS（约 360 万样本等效训练），使用LAMB 优化器、PolynomialDecayWithWarmup 学习率调度和全局梯度范数裁剪（global_norm > 1.0时缩放到 1.0，见 L1122-L1128）。整个训练循环被@TinyJit编译，保证多 GPU 数据并行下的单步开销可控。

二、环境准备：安装 tinygrad 与 mlperf-logging

README 明确要求从mlperf_training_v5.0分支安装 tinygrad，并启用mlperf可选依赖（其中包含 mlperf-logging）：

git clone <tinygrad 仓库地址> python3 -m pip install -e ".[mlperf]"

• -e表示可编辑安装，便于在源码上直接改动调试；
• .[mlperf]触发 pyproject.toml 中定义的 mlperf 扩展依赖，安装后model_train.py才能from mlperf_logging import mllog（见 model_train.py）；
• 安装前需按 README 提示取消 setup.py 中 mlperf 依赖的注释，确保mlperf_logging包被正确引入。

此外还需安装数据处理相关的 Python 包：

pip install gdown numpy tqdm tensorflow

tinybox_green 硬件前提：P2P 驱动

README 单独强调了 tinybox_green 平台需要安装P2P 内核驱动（支持 GPU 间 peer-to-peer 直接通信，避免数据绕经主机内存），并注明这是生产环境 tinybox green 的默认配置。若在非默认环境复现，需按官方驱动说明安装后再运行基准，否则多卡（GPUS=6）数据并行训练中的张量分片（t.shard_(GPUS, axis=0)）性能会明显受损。

三、数据集准备：下载、校验与预处理

3.1 下载原始数据

使用 tinygrad 自带的下载脚本，从 Google Drive 拉取 2020 Wikipedia 语料（与 MLCommons 参考实现使用的数据集一致）：

BASEDIR="/raid/datasets/wiki" WIKI_TRAIN=1 VERIFY_CHECKSUM=1 python3 extra/datasets/wikipedia_download.py

该命令的关键环境变量：

• BASEDIR：数据集根目录（README 统一约定为/raid/datasets/wiki）；
• WIKI_TRAIN=1：除 BERT 配置、vocab 与初始 checkpoint 外，额外下载results_text.tar.gz训练语料（见 wikipedia_download.py）；
• VERIFY_CHECKSUM=1：解压后对照bert_reference_results_text_md5.txt对每个文件做MD5 校验，任一文件不匹配即抛出ValueError（见 wikipedia_download.py）。

3.2 预处理训练数据

BASEDIR="/raid/datasets/wiki" NUM_WORKERS=16 python3 extra/datasets/wikipedia.py pre-train all

• pre-train all：对所有 500 个主题执行预训练样本生成；
• NUM_WORKERS=16：并行预处理线程数。README 特别给出内存约束：预处理的线程数受可用内存限制，128GB RAM 环境下最多建议 16 线程。

如需只生成某个主题（取值范围 0 到 499）：

BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-train 42

3.3 预处理验证数据

BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-eval

该步骤生成与官方验证集对应的预训练特征，供 model_train.py 中的batch_load_val_bert使用。

预处理脚本本身还支持一组可调环境变量（见 wikipedia.py）：

脚本内置了完整的 BERT tokenization 流程（whitespace / wordpiece / 中文分字 / 标点切分 / 去除控制字符等，见 wikipedia.py），是对 MLCommonscreate_pretraining_data.py的 tinygrad 化移植。

四、运行基准：run_and_time.sh 逐行拆解

README 给出的运行方式为执行实现目录下的启动脚本。注意：README 正文书写的是 v5.0 路径，当前仓库 v5.1 提交中实际文件位于 tinybox_green/run_and_time.sh，直接执行：

bash examples/mlperf/training_submission_v5.1/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh

脚本会经历两个阶段：init（初始化/预编译）与run（正式计时运行），并将输出同时写入时间戳命名的日志文件：

# init BENCHMARK=10 INITMLPERF=1 BERT_LAYERS=2 python3 examples/mlperf/model_train.py | tee $LOGFILE # run PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE

4.1 关键环境变量一览（tinybox_green）

以下变量全部来自 tinybox_green/run_and_time.sh，与 model_train.py 中的getenv读取一一对应：

4.2 初始化与正式运行两阶段的意义

• init 阶段（INITMLPERF=1）：以BENCHMARK=10只跑 10 步、使用BERT_LAYERS=2的浅层模型与假数据完成设备预热与 kernel 编译缓存，同时向 MLPerf 日志写入SUBMISSION_*、CACHE_CLEAR等元数据，并触发INIT_START/INIT_STOP计时事件（见 model_train.py）。
• run 阶段（RUNMLPERF=1）：PARALLEL=0关闭多进程，加载真实数据（batch_load_train_bert/batch_load_val_bert），开始正式训练与周期评估，日志事件涵盖RUN_START、EPOCH_START/STOP、EVAL_START/STOP、EVAL_ACCURACY，达标（TARGET=0.72）后写入RUN_STOP(status=SUCCESS)。

训练步内部由@TinyJit @Context(TRAINING=1)的train_step_bert完成：输入张量shard_/to_到各 GPU → 前向得到lm_logits与seq_relationship_logits→ 计算 loss 并乘loss_scaler反传 → 梯度除以 scaler → 计算全局 L2 范数并裁剪到 1.0 → LAMB 更新 → 调度器步进（见 model_train.py）。

4.3 训练超参数（由 run_and_time.sh 或环境变量覆盖）

train_bert()中一组可覆盖的默认超参数（model_train.py）：

LAMB 优化器把参数分为两组：带权重衰减的普通参数，以及bias/LayerNorm相关的无衰减参数（parameters_no_wd），两者共享同一组学习率调度（见 model_train.py），这与 MLPerf 官方 BERT 收敛配方一致。

五、多平台对照：tinybox_red 与 tinybox_8xMI300X

README 同时列出了另外两个实现目录的运行方式，当前仓库 v5.1 中的对应脚本为：

• tinybox_red/run_and_time.sh（AMD 平台）
• tinybox_8xMI300X/run_and_time.sh（8 卡 MI300X）

三份脚本共享INITMLPERF+RUNMLPERF两阶段结构，差异主要在硬件相关配置：

几点说明：

• tinybox_red 在 init 前先sleep 5 && sudo rmmod amdgpu || true重置 AMD 驱动，并通过HCQDEV_WAIT_TIMEOUT_MS=100000放宽硬件队列等待超时，用于规避偶发挂起；
• tinybox_8xMI300X 的 LAMB 超参数与TRAIN_STEPS=3900参考了 MLCommons 公开结果中的收敛配方（脚本注释给出了出处），说明大 batch（1024）下需要配套调整学习率、动量与衰减才能稳定收敛；
• 三份脚本均开启LOGMLPERF=1，正式提交时日志由mllog写入result_bert_<seed>.log（见 model_train.py）。

此外，仓库还提供单卡开发脚本 tinybox_1xMI300X/dev_beam.sh：以GPUS=1 BS=128、BENCHMARK=10 BERT_LAYERS=2 DEBUG=2快速验证单卡性能，适合没有整机时的本地调优。

六、运行输出与日志解读

正式运行中，tqdm.write每步输出一行结构化信息（model_train.py）：

step ms run / ms python / ms fetch data / ms device / loss / LR / GB used / GFLOPS

例如：10 123.45 ms run, 12.30 ms python, 5.20 ms fetch data, 105.95 ms NV:0*6, 7.52 loss, 0.000175 LR, 11.20 GB used, 1482.35 GFLOPS

• ms run：单步总耗时；
• ms python：Python 侧开销，越小说明 JIT 编译越充分；
• ms fetch data：数据管线预取耗时；
• ms device：设备执行耗时（NV:0*6表示 6 卡并行）；
• GFLOPS：基于GlobalCounters.global_ops统计的实际计算吞吐。

BENCHMARK=10触发后，脚本会输出预估总训练时长与每 epoch 的总算力/访存统计，方便在正式提交前快速评估是否达到目标性能窗口。

七、总结与复现检查清单

按照 README 完整复现 tinybox_green 上的 MLPerf BERT 基准，可归纳为四步：

1. 安装：从mlperf_training_v5.0分支pip install -e ".[mlperf]"，并补齐gdown numpy tqdm tensorflow；
2. 数据：WIKI_TRAIN=1 VERIFY_CHECKSUM=1下载并校验，wikipedia.py pre-train all（128GB 内存下NUM_WORKERS=16）与pre-eval预处理；
3. 硬件：确认 tinybox_green 已安装 P2P 驱动（生产默认）；
4. 运行：执行 tinybox_green/run_and_time.sh，先INITMLPERF预热，再RUNMLPERF正式计时并输出result_bert_<seed>.log。

整套流程展示了 tinygrad 将 MLPerf 官方 BERT 预训练配方完整落地到自有硬件栈的能力：从数据管线的 tokenization 复刻（wikipedia.py），到训练核心的 LAMB + 多项式衰减 + 梯度裁剪（model_train.py），再到多平台（NV / AMD / 8xMI300X）的 beam 编译优化参数调优，均可直接对照本指南与仓库源码逐项核验。

【免费下载链接】tinygradYou like pytorch? You like micrograd? You love tinygrad! ❤️项目地址: https://gitcode.com/GitHub_Trending/tiny/tinygrad