企业官网建设流程全解析

tinygrad 在 tinybox green 上复现 MLPerf ResNet-50 图像分类训练基准

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

导读

本文以 examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/resnet/implementations/tinybox_green/README.md 为骨架，完整讲解如何在 tinygrad 仓库中复现 MLPerf 训练基准中的 ResNet-50 图像分类任务：从环境依赖安装、ImageNet 数据集下载与验证，到使用run_and_time.sh一键跑出带 mlperf-logging 合规日志的完整训练。文中同时结合仓库源码（examples/mlperf/model_train.py、examples/mlperf/dataloader.py、extra/datasets/imagenet_download.py）与各提交脚本，逐项拆解硬件前置条件、数据管线、模型与优化器配置、Beam Search 编译参数以及 MLPerf 日志协议，帮助你既会“跑通基准”，也能理解每一步背后的实现原理。

一、问题定义：在 tinygrad 上训练 ResNet-50 完成 ImageNet 分类

该 MLPerf 提交的“问题（Problem）”是使用 ResNet-50 卷积神经网络（CNN）完成 ImageNet 图像分类任务，即给定一张 224×224 的 RGB 图像，输出其在 1000 个类别上的概率分布。这也是 MLPerf Training v6.0 Closed Division 中 resnet 基准的标准设定。

从源码可以确认，模型本体定义在 extra/models/resnet.py（ResNet50类），实际训练入口位于 examples/mlperf/model_train.py 的train_resnet()函数，它负责：

• 读取全部超参数（epochs、batch size、学习率、损失缩放等）；
• 构建 ResNet-50 并把权重按设备切分（shard_）到多张 GPU；
• 组织 LARS + SGD 双优化器组与多项式衰减学习率调度器；
• 用TinyJit把训练步与评估步编译成 JIT 内核；
• 运行 epoch 循环，并在达到目标精度（默认 Top-1 0.759）时停止并保存模型。

关于环境与硬件前提的说明

该 README 面向 tinygrad 自有的 tinybox 硬件平台（tinybox green / tinybox red）。仓库中对应的系统描述文件 examples/mlperf/training_submission_v6.0/tinycorp/systems/tinybox_green.json 记录了 tinybox green 的硬件配置：单节点 1 颗 AMD EPYC 7532（32 核 / 64 线程）、128GB 内存、6 张 NVIDIA GeForce RTX 4090（24GB GDDR6X，PCIe 4.0 x16）、4TB RAID + 1TB 启动盘、Ubuntu 22.04.4 + CUDA 12.4 + Python 3.10.12。tinybox red 则为 AMD 平台（提交脚本中以DEV=AMD区分）。需要说明的是，这些具体硬件规格是提交时的机器配置，本文介绍的是通用复现流程，在不同 GPU 环境下需要按显存与算力调整 batch size、beam 参数等。

二、Requirements：安装 tinygrad 与 mlperf-logging

README 首先要求从 master 分支安装 tinygrad 与 mlperf-logging：

git clone https://github.com/tinygrad/tinygrad.git python3 -m pip install -e ".[mlperf]"

第二条命令以可编辑（-e）模式安装当前仓库，并带入mlperf可选依赖。该可选依赖定义在仓库根目录的 pyproject.toml 中，它负责拉取mlperf-logging包——正是model_train.py里from mlperf_logging import mllog的依赖来源。安装完成后，mllog会把训练事件按 MLPerf 规定的 key（如run_start、epoch_start、eval_accuracy、run_stop）写入日志文件，用于提交合规性校验。

平台差异前置条件

README 对不同 tinybox 型号给出了额外的驱动/内核级前置条件：

• tinybox_green：需要安装 p2p 驱动（即 tinygrad 的 open-gpu-kernel-modules 550.54.15-p2p 分支），生产环境 tinybox green 默认已启用该驱动，无需手动处理。
• tinybox_red（AMD）：需要禁用 compute wave stack rewind（cwsr），这也是生产环境默认配置。若需手动设置，执行以下命令：

sudo vi /etc/modprobe.d/amdgpu.conf cat <<EOF > /etc/modprobe.d/amdgpu.conf options amdgpu cwsr_enable=0 EOF sudo update-initramfs -u sudo reboot # validate sudo cat /sys/module/amdgpu/parameters/cwsr_enable #= 0

从结构上看，这份 README 同时覆盖 green（NVIDIA）与 red（AMD）两个平台，本仓库目录tinybox_green/下的提交脚本统一使用DEV=NV，而tinybox_red/下的脚本使用DEV=AMD。

三、数据准备：下载并校验 ImageNet

下载命令

README 给出如下下载指令：

IMGNET_TRAIN=1 python3 extra/datasets/imagenet_download.py

从 extra/datasets/imagenet_download.py 的入口代码看，该脚本实际做了这些事（imagenet_download.py 第 39-51 行）：

1. 创建imagenet/、imagenet/val/、imagenet/train/目录；
2. 下载类别索引文件imagenet_class_index.json与验证集 synset 标签文件；
3. 下载验证集压缩包ILSVRC2012_img_val.tar（约 7GB）并解压、整理验证集目录结构；
4. 仅当设置IMGNET_TRAIN=1时，额外下载训练集ILSVRC2012_img_train.tar（约 138GB）、解压并按类别目录整理（imagenet_prepare_train）。

也就是说，IMGNET_TRAIN=1是“同时准备训练集+验证集”的开关；如果不设置该变量，脚本只准备验证集。训练集约 138GB、验证集约 7GB，请预留足够磁盘空间。

数据集在训练管线中的消费方式

下载完成后，训练与评估数据通过 examples/mlperf/dataloader.py 的batch_load_resnet()（dataloader.py 第 72 行）以多进程队列方式加载：它读取get_train_files()/get_val_files()的文件列表，按 batch 打包送入共享内存队列，model_train.py侧则用data_get()取出并做shard(GPUS, axis=0)切分。dataloader 中的preprocess_train负责训练增强，padding 批次则用训练均值[123.68, 116.78, 103.94]填充——这与model_train.py中normalize使用的input_mean保持一致（model_train.py 第 155 行）。注意 mlperf 参考实现不做input_std归一化，tinygrad 实现也遵循这一点（源码中有对应注释）。

四、一次性环境设置（tinybox_red）

README 还提到 tinybox_red 需要执行一次性的性能调优脚本：

examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/setup.sh

注：README 原文指向 v4.0 目录；本仓库中 v6.0 目录下对应的设置脚本位于 examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/setup.sh，内容为 AMD 平台的 ROCm 电源/频率设置（rocm-smi --setprofile compute、--setmclk 3、--setperflevel high，并将功率上限设为 350W）。

该步骤仅针对 AMD 平台（tinybox_red）。NVIDIA 平台（tinybox_green）无需此步骤，这也是 tinybox_green 目录下没有setup.sh的原因。

五、运行基准：run_and_time.sh 全流程拆解

5.1 运行命令

README 给出的运行命令是：

examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/run_and_time.sh

本仓库 v6.0 目录下，tinybox_green 对应的实际脚本为 examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/resnet/implementations/tinybox_green/run_and_time.sh。MLPerf 的run_and_time.sh是提交方必须提供的“一次命令跑完全程”的入口脚本，它负责：初始化（清缓存、启动计时）→ 正式训练（含阶段性评估）→ 达标即停止 → 全程输出合规日志。

5.2 脚本内容与关键环境变量

以 tinybox_green 的run_and_time.sh为例（run_and_time.sh）：

#!/bin/bash set -e # Exit on any error set -o pipefail # Make pipeline fail if any command fails export PYTHONPATH="." DEV=NV export MODEL="resnet" export SUBMISSION_PLATFORM="tinybox_green" export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192 export RESET_STEP=0 export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=10 BEAM_PADTO=0 # pip install -e ".[mlperf]" export LOGMLPERF=${LOGMLPERF:-1} export SEED=$RANDOM DATETIME=$(date "+%m%d%H%M") LOGFILE="resnet_green_${DATETIME}_${SEED}.log" # init BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE # run PARALLEL=0 RUNMLPERF=1 EVAL_START_EPOCH=3 EVAL_FREQ=4 python3 examples/mlperf/model_train.py | tee -a $LOGFILE

脚本分为两段：

1. 初始化段（init）：BENCHMARK=10 INITMLPERF=1运行一次，用假数据（fake_data_get）跑 10 步，完成 JIT 内核编译预热、清空磁盘缓存并发出init_start/init_stop事件，同时估算总训练时长。
2. 正式运行段（run）：RUNMLPERF=1启动真实训练，从 epoch 3 开始每 4 个 epoch 评估一次（EVAL_START_EPOCH=3 EVAL_FREQ=4），达到目标精度即run_stop。

对比 tinybox_red 版本（run_and_time.sh），差异仅在于DEV=AMD、SUBMISSION_PLATFORM="tinybox_red"，以及初始化前多了一条sleep 5 && sudo rmmod amdgpu || true（卸载并重新加载 amdgpu 模块以重置驱动状态），且 beam 搜索参数不同（BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=96 BEAM_MIN_PROGRESS=5）。

5.3 核心环境变量速查表

结合run_and_time.sh与 examples/mlperf/model_train.py 的getenv读取逻辑，整理出与 ResNet-50 提交直接相关的关键变量：

训练超参数在 model_train.py 第 76-87 行 中被读取并写入config字典，同时通过MLLOGGER.event记录到日志中。

5.4 训练循环：优化器、学习率与 JIT 步

train_resnet()中的核心训练设计（model_train.py 第 99-180 行）：

• 双优化器组：主参数（卷积/全连接权重）使用LARS，而 BN 相关参数、bias 与downsample.1层权重使用SGD(classic=True)且无权重衰减，两者组成OptimizerGroup同步更新。这与 MLPerf 参考实现的“LARS + 无衰减 skip 参数”约定一致。
• 学习率调度：使用PolynomialDecayWithWarmup（定义在 examples/mlperf/lr_schedulers.py），从base_lr多项式衰减到1e-4，前WARMUP_EPOCHS个 epoch 线性 warmup，训练与 skip 两组共享相同的调度参数。
• 损失与精度：交叉熵采用label_smoothing=0.1，loss * loss_scaler反向传播后再除以loss_scaler归一化梯度；同时统计 Top-1 精度用于达标判断。
• JIT 编译：train_step/eval_step均以@TinyJit+@Context(TRAINING=1/0)装饰，训练与评估使用不同的 beam 配置。
• BN 策略：默认（未设置SYNCBN）将BatchNorm替换为UnsyncedBatchNorm（来自 examples/hlb_cifar10.py），即每个设备独立统计 BN，并将running_mean/running_var沿设备轴切分，避免跨设备同步通信。

5.5 Beam Search 编译参数解读

提交脚本中的TRAIN_BEAM=4 BEAM_UOPS_MAX=...等参数是 tinygrad 编译器（linearizer）在 beam search 阶段的搜索上限配置：

这些参数通过 tinygrad 的BEAM全局配置（见 tinygrad/helpers.py）进入编译流程，model_train.py中TRAIN_BEAM = getenv("TRAIN_BEAM", BEAM.value)在训练/评估阶段动态切换BEAM.value。由于 NVIDIA 与 AMD 后端指令集不同，两套提交选择了不同的搜索上限——这也解释了为什么 green（NV）与 red（AMD）脚本的 beam 参数存在差异。

六、开发辅助脚本：dev_run.sh 与 dev_beam.sh

tinybox_green/目录下还提供两个开发调试脚本：

• dev_run.sh：与正式脚本共享相同的模型/精度/beam 配置，但固定EVAL_START_EPOCH=3 EVAL_FREQ=4，并开启WANDB=1 PARALLEL=0以便在开发阶段用 Weights & Biases 监控训练曲线；不设置LOGMLPERF，因此不产生提交日志。
• dev_beam.sh：仅设BENCHMARK=10 DEBUG=2，用假数据快速跑 10 步，用于在完整训练前验证 beam 编译产物与性能预估（model_train.py会在第 10 步打印“Estimated training time”与每 epoch 的 global_ops / global_mem）。

两者与run_and_time.sh的差异集中体现在INITMLPERF/RUNMLPERF/LOGMLPERF/WANDB/BENCHMARK这几个开关上，体现了“调试（dev）→ 提交（run_and_time）”的工作流分层。

七、日志、结果与 checkpoint

7.1 MLPerf 日志产出

run_and_time.sh会把输出tee到形如resnet_green_MMDDHHMM_SEED.log的本地文件，同时model_train.py在LOGMLPERF=1时通过mllog额外写入result_resnet_<seed>.txt（model_train.py 第 33 行）。日志中记录的 MLPerf 事件包括：submission_org=tinycorp、submission_platform、division=closed、status=onprem、cache_clear、init_start/init_stop、run_start/run_stop、epoch_start/epoch_stop、eval_accuracy以及全部超参数事件（global_batch_size、opt_name=lars、LARS 的 base/end LR、poly power、warmup epochs、momentum、weight decay 等，见 model_train.py 第 117-136 行）。

7.2 达标停止与模型保存

训练默认目标为 Top-1 精度0.759（TARGET可覆盖）。从 epoch 33 开始每个 epoch 都会评估（即使不满足EVAL_FREQ间隔），一旦total_top_1 >= target：

• 记录run_stop事件（status=success）；
• 将模型权重以.safe格式保存到./ckpts/resnet50_<seed>.safe；
• 停止训练循环（见 model_train.py 第 330-339 行）。

7.3 断点续训

若需从中断处继续，设置RESUME=/path/to/ckpt.safe即可：model_train.py用load_training_state恢复模型、优化器与调度器状态，并根据scheduler.epoch_counter推算出起始 epoch（model_train.py 第 139-143 行）。配合CKPT=1在每次评估时自动保存含优化器状态的完整训练 checkpoint，可支撑长训场景。

八、常见问题与调参建议

• 显存不足（OOM）：提交配置BS=1536是 6×RTX 4090 的组合，单卡显存较小的环境可降低BS，同时按LR = 7.2 * (BS/1536)的线性缩放规则同步下调学习率，保持训练稳定性。
• beam 搜索过慢：dev_beam.sh用BENCHMARK=10先验证编译耗时；若搜索时间过长，可降低TRAIN_BEAM或收紧BEAM_UOPS_MAX/BEAM_UPCAST_MAX，但会损失部分内核质量。
• 评估阶段显存压力：默认RESET_STEP=1会在评估前释放训练 JIT 的显存；若希望提升吞吐可设RESET_STEP=0（提交脚本即如此），但需确认显存余量足够。
• 数据加载成为瓶颈：每步日志会打印fetch data时间，若明显偏高，说明磁盘/解压 IO 受限，应将数据集放在高速 NVMe 上（tinybox green 为 4TB RAID 阵列）。
• 随机种子：正式提交用SEED=$RANDOM保证每次运行独立；复现实验可固定SEED获得确定性结果。

九、结语

通过tinybox_green的这份 README 与其配套脚本，可以清晰看到 tinygrad 参与 MLPerf Training 基准的完整工程链路：环境安装 → 数据准备 → 平台调优 → 双阶段 run_and_time 提交 → mlperf-logging 合规输出 → 达标 checkpoint。对于想要在自有 NVIDIA 多卡环境复现 ResNet-50 训练基准的开发者，只需按 run_and_time.sh 调整GPUS/BS与 beam 参数即可复用整套流程；对底层机制感兴趣的读者，可继续深入 examples/mlperf/model_train.py、examples/mlperf/dataloader.py 与 extra/datasets/imagenet_download.py 三个核心文件，理解每一个环境变量在训练管线中的真实作用。

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