企业官网建设流程全解析

tinygrad 在 MLPerf BERT 训练基准中的实战指南：数据流水线、环境配置与多平台运行

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

本指南以 MLPerf Training v5.0 提交中 BERT（NLP 预训练）基准的 tinygrad 实现为主线，完整讲解从依赖安装、Wikipedia 原始数据下载与校验、MLM/NSP 预训练样本预处理，到 tinybox_green / tinybox_red / tinybox_8xMI300X 三种硬件平台上运行基准的完整流程。读完本文，你将掌握 MLPerf BERT 提交的全部操作步骤，以及驱动这些脚本的底层源码机制与关键环境变量。

一、基准背景：用 tinygrad 复现 MLPerf BERT 训练

MLPerf Training 是衡量机器学习训练性能的行业标准基准，其中 BERT（Bidirectional 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两份并列实现，内容基本一致）。

整套流程分为两大阶段：

1. 数据阶段：下载 2020 年 Wikipedia 数据集，做 checksum 校验，然后完成 BERT 标准的 tokenize、Masked LM 采样、序列截断、padding 等预处理，生成可直接供训练读取的 pickle 特征文件；
2. 训练阶段：通过 run_and_time.sh 依次完成"预热/初始化 + 正式计时训练"两轮运行，同时输出符合 MLPerf logging 规范的日志。

二、环境准备：tinygrad 与依赖安装

基准运行要求从mlperf_training_v5.0分支安装 tinygrad，并同时安装mlperf-logging与数据集工具链。依赖声明位于仓库根目录的 pyproject.toml，其中[mlperf]extra 通过注释形式指向mlperf-logging @ git+https://github.com/mlperf/logging.git@5.0.0-rc3，需要在安装前取消注释启用：

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

随后安装数据集下载与处理所需的辅助库：

pip install gdown numpy tqdm tensorflow

其中：

• gdown：负责从 Google Drive 下载 Wikipedia 数据集（下载逻辑见 wikipedia_download.py）；
• numpy/tqdm：预处理与进度展示依赖，预处理脚本大量使用tqdm与tqdm.contrib.concurrent.process_map实现多进程并行（见 wikipedia.py）；
• tensorflow：BERT 参考实现（model.ckpt-28252检查点）与校验文件来自 MLPerf 官方 TensorFlow 流程，预处理后对齐其数据格式。

tinybox_green 的专用说明：p2p 驱动

对于tinybox_green平台，README 特别注明需要按官方 p2p（peer-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.py

3.1 下载命令

在仓库根目录执行：

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

三个环境变量的作用（见 wikipedia_download.py）：

3.2 脚本实际下载的内容

脚本通过gdrive_download逐个拉取以下文件（均来自 MLPerf 官方 BERT 数据集页）：

• bert_config.json：BERT-Base 模型配置；
• vocab.txt：WordPiece 词表；
• model.ckpt-28252.data-00000-of-00001/model.ckpt-28252.index/model.ckpt-28252.meta：官方参考检查点三件套，并自动生成checkpoint索引文件（见 wikipedia_download.py）；
• bert_reference_results_text_md5.txt：训练语料 500 个分片的 MD5 清单（仅WIKI_TRAIN=1时）；
• results_text.tar.gz：约 500 个分片的原始文本压缩包（仅WIKI_TRAIN=1时）。

下载完成后，脚本用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_WORKERS=16 python3 extra/datasets/wikipedia.py pre-train all

pre-train all会通过process_map以NUM_WORKERS（默认取min(os.cpu_count(), 32)）个 worker 并行处理 0～499 号分片，每个分片的结果写入BASEDIR/train/{part}.pkl（见 process_part）。已存在的分片会被跳过，因此可以断点续跑。

4.2 生成单个分片（调试用）

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

第二个参数取 0～499 之间的分片号，仅处理results4/part-00042-of-00500，适合小规模验证流水线正确性。

4.3 生成验证集

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

pre-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_tokenize）；
• NSP 配对：create_instances_from_document把文档切成约MAX_SEQ_LENGTH - 3长度的块，随机决定 A/B 段划分，并以 50% 概率从其他文档随机拼接 B 段构成负样本is_random_next（见 wikipedia.py）；
• Masked LM：候选 token 打乱后按MASKED_LM_PROB（默认 0.15）比例采样，其中 80% 替换为[MASK]、10% 保持不变、10% 随机替换为词表中任意词（见 create_masked_lm_predictions）；
• 特征化：instance_to_features把序列 padding 到MAX_SEQ_LENGTH=512、掩码位置 padding 到MAX_PREDICTIONS_PER_SEQ=76，并统一转为 int32/float32 的 numpy 数组（见 wikipedia.py）。

4.5 预处理环境变量一览

五、运行基准：三个 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.sh

5.1 run_and_time.sh 的运行机制

以 tinybox_green 的 run_and_time.sh 为例，脚本分"init"与"run"两个阶段，均调用examples/mlperf/model_train.py：

# init：2 层 BERT 的快速预热，同时完成 MLPerf 日志初始化 BENCHMARK=10 INITMLPERF=1 BERT_LAYERS=2 python3 examples/mlperf/model_train.py | tee $LOGFILE # run：正式计时训练（PARALLEL=0 关闭数据并行以外的并行加载） PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE

两个阶段对应的执行分支由 model_train.py 中的INITMLPERF/RUNMLPERF控制：

• init 阶段：INITMLPERF=1时，脚本用MLLOGGER依次上报SUBMISSION_PLATFORM、SUBMISSION_BENCHMARK等 MLPerf 标准事件（mllog_constants），并用BENCHMARK=10指定只跑 10 步、取中位数作为步时基准（sorted(step_times)[BENCHMARK // 2]）；
• run 阶段：RUNMLPERF=1时执行完整训练循环，并在关键节点持续上报 mlperf 日志，日志文件以bert_green_${DATETIME}_${SEED}.log命名，SEED=$RANDOM保证每次运行的随机性可独立复现。

5.2 核心训练/调度环境变量

5.3 平台差异与注意事项

• tinybox_green（6×RTX 4090）：依赖 p2p 内核驱动完成多卡通信；dev_run.sh额外开启WANDB=1便于实验监控，dev_beam.sh通过DEBUG=2输出 beam 搜索日志（BEAM_LOG_SURPASS_MAX=1），两者均不参与正式计时；
• tinybox_red（6×AMD GPU）：脚本开头执行sudo rmmod amdgpu/sudo modprobe amdgpu做驱动重置，并设置HCQDEV_WAIT_TIMEOUT_MS=100000防止等待超时挂起；注释标明 AMD 的 AM 驱动曾导致 NaN，故回退使用标准amdgpu驱动；
• tinybox_8xMI300X（8×AMD MI300X）：为对齐 MLPerf 官方参考结果，显式设置 LAMB 优化器超参OPT_BASE_LEARNING_RATE=0.0011、OPT_LAMB_BETA_1=0.60466、OPT_LAMB_BETA_2=0.85437、DECAY=0.1及TRAIN_STEPS=3900，batch size 提升到 1024；init 阶段需DEBUG=2防止挂起。

六、提交硬件配置参考：tinybox_green 系统清单

README 对应的提交系统配置记录在 tinybox_green.json（submitter: tinycorp，division: closed），其硬件基线可作为复现环境的最低参考：

• 主机：1 节点，AMD EPYC 7532（32 核 / 64 vCPU），128GB DDR4（8×16GB），NVMe SSD（4TB RAID + 1TB 启动盘），Ubuntu 22.04.4，Python 3.10.12，CUDA 12.4；
• 加速器：6× NVIDIA GeForce RTX 4090（24GB GDDR6X），PCIe 4.0 x16 互联，风冷。

七、故障排查与实操建议

• 预处理内存不足：将NUM_WORKERS调低（128GB 内存建议 ≤16），或改用pre-train <分片号>分批处理；
• 数据校验失败：VERIFY_CHECKSUM=1会直接以异常中断，重新执行下载脚本即可（已下载文件会被gdown跳过，仅重下缺失项）；
• 训练日志缺失：确认LOGMLPERF=1与SUBMISSION_PLATFORM已设置，否则MLLOGGER不会初始化（见 model_train.py）；
• AMD 平台异常：参照 tinybox_red 的做法，在运行前重置amdgpu驱动并调大HCQDEV_WAIT_TIMEOUT_MS；
• 快速验证：先用BERT_LAYERS=2+BENCHMARK=10的 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