企业官网建设流程全解析

基于 tinygrad 复现 MLPerf 训练提交：ResNet-50 在 tinybox green/red 上的完整部署与调优指南

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

本文档是 tinygrad 仓库内 MLPerf Training v4.1 提交目录中 ResNet-50 实现（关联 README）的深度展开。它完整继承原文档的环境要求、数据准备、平台特化配置与基准运行步骤，并进一步结合仓库中的运行脚本与 model_train.py 源码，逐一拆解每个环境变量的含义、默认值与对训练的影响。读完本文，你将掌握如何在一台 6 卡 tinybox 机器上从零准备 ImageNet 数据、完成驱动与内核层面的特化设置、运行一次符合 MLPerf 规范的 ResNet-50 训练提交，并理解日志合规（mllog）与 Beam Search 编译优化在其中扮演的角色。

1. 问题定义：用 ResNet-50 做图像分类

该提交所解决的问题与 MLPerf Training 的经典基准一致：使用 ResNet-50 卷积神经网络在 ImageNet-2012 数据集上完成图像分类训练。原文档将其归类为 "1. Problem"，仓库中的规格定义位于 examples/mlperf/model_spec.py：

def spec_resnet(): # Resnet50-v1.5 from extra.models.resnet import ResNet50 mdl = ResNet50() img = Tensor.randn(1, 3, 224, 224) test_model(mdl, img)

• 模型为 ResNet50-v1.5，输入张量为(1, 3, 224, 224)，即单张 224×224 的 RGB 图像；
• 实际模型实现在 extra/models/resnet.py；
• examples/mlperf/README 指出，本仓库提交的 5 个模型（Resnet50-v1.5、Retinanet、3D UNET、RNNT、BERT-large）同时覆盖训练与推理基准，ResNet-50 的计算量约为 8.2 GOPS/input。

训练质量目标在 model_train.py 中硬编码：

target, achieved = getenv("TARGET", 0.759), False

即默认以Top-1 准确率 0.759（75.9%）作为达标线，训练最多进行 37 个 epoch（getenv("EPOCHS", 37)），当验证准确率达到目标后提前停止计时。

2. 环境要求：安装 tinygrad 与 mlperf-logging

原文档要求从 master 分支安装 tinygrad 与 mlperf-logging：

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

其中.[mlperf]是 tinygrad 的 extras 安装选项，负责拉取 mlperf-logging、tqdm、wandb 等 MLPerf 训练所需的依赖。仓库根目录的 pyproject.toml 中定义了该 extras 的具体依赖清单，安装时请确保 Python 版本与系统环境匹配——参考提交记录，v4.1 提交使用的软件栈为Python 3.10.12（详见下文系统清单）。

安装完成后，可以在仓库根目录执行规格冒烟测试确认模型可运行：

python3 -m examples.mlperf.model_spec

该入口会遍历MODEL环境变量指定的模型列表（默认resnet,retinanet,unet3d,rnnt,bert,mrcnn），逐个打印模型的 GOPS 与耗时（见 model_spec.py），可作为提交前的最小验证。

3. 平台特化设置：tinybox green 与 tinybox red

MLPerf 提交以 "system" 为单位。v4.1 提交包含两套系统（systems 目录）：

两套系统均为 6 卡、24GB 显存、PCIe 4.0 x16 直连的架构，因此共享大部分运行脚本，仅在驱动级设置上分叉。

3.1 tinybox green：安装 p2p 驱动

green 使用 NVIDIA GPU，需要安装支持 P2P（peer-to-peer）通信的驱动。原文档指出应按照 tinygrad 维护的 open-gpu-kernel-modules 的550.54.15-p2p分支 README 安装该驱动，并说明这是生产环境 tinybox green 的默认配置，即无需额外手工操作。

3.2 tinybox red：禁用 cwsr

red 使用 AMD GPU，需要禁用 AMD 驱动中的cwsr（Compute Wave Save/Restore，计算波前保存/恢复）机制。这也是生产环境 tinybox red 的默认配置；若在新机器上手动配置，步骤如下：

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

关键点：

• 通过内核模块参数cwsr_enable=0关闭 cwsr，写入/etc/modprobe.d/amdgpu.conf使其在每次启动时生效；
• 修改后需update-initramfs -u重建 initramfs 并重启；
• 验证步骤不可省略：重启后检查/sys/module/amdgpu/parameters/cwsr_enable输出必须为0，否则设置未生效。

关闭 cwsr 的动机在于性能：cwsr 会在上下文切换时保存/恢复计算波前状态，禁用后可减少 GPU 调度的额外开销，对长时训练吞吐有明显收益。

4. 数据准备：下载并校验 ImageNet

原文档给出的数据下载命令为：

IMGNET_TRAIN=1 python3 extra/datasets/imagenet_download.py

该命令由 extra/datasets/imagenet_download.py 实现，其行为如下：

• 无条件下载验证集相关文件：
  • imagenet_class_index.json（类别索引）；
  • imagenet_2012_validation_synset_labels.txt（验证集标签）；
  • ILSVRC2012_img_val.tar（约 7GB 验证集）；
• 解压验证集后，依据标签文件将图片整理进按 synset 命名的子目录（imagenet_prepare_val），并删除临时标签文件；
• 仅当设置了IMGNET_TRAIN环境变量时，才继续下载并解压ILSVRC2012_img_train.tar（约 138GB 训练集），并按train/<synset>/目录结构拆包（imagenet_prepare_train，源码位置）。

因此：

• 只做推理或快速验证：直接运行python3 extra/datasets/imagenet_download.py（跳过 train）；
• 做完整训练提交：必须带上IMGNET_TRAIN=1，并预留足够的磁盘空间（验证集 7GB + 训练集 138GB + 解压空间，仓库系统清单显示 tinybox 标配 4TB RAID 阵列正是为此设计）。

数据会被组织在extra/datasets/imagenet/目录下，训练与验证文件列表分别由 extra/datasets/imagenet.py 的get_train_files()/get_val_files()提供，model_train.py 据此计算每个 epoch 的步数。

5. 一次性设置（仅 tinybox red）

tinybox green 无需额外设置（p2p 驱动为出厂默认），而 red 除了禁用 cwsr，还需要通过setup.sh做一次性的 ROCm 运行参数调优：

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

注意：v4.1 的 red 目录下也有同名脚本，v4.0 与 v4.1 内容一致。脚本内容（setup.sh）逐行含义如下：

rocm-smi --setprofile compute # 将 GPU 设置为 compute 计算档（而非图形档） rocm-smi --setmclk 3 # 将显存时钟锁定到较高的性能档位 rocm-smi --setperflevel high # 性能等级设为 high，避免省电策略降频 # power cap to 350W echo "350000000" | sudo tee /sys/class/drm/card{1..6}/device/hwmon/hwmon*/power1_cap

• 前三行通过rocm-smi将 6 张 GPU 统一设置为计算模式、锁定显存时钟、提高性能等级，保证长时训练时钟稳定；
• 最后一行以350W为功率上限（power1_cap单位是微瓦，350000000 µW = 350W），对card1到card6六张卡统一限功率，既控制功耗又提升多卡协同的稳定性；
• 该脚本需要sudo权限，属于"一次性"设置：重启后部分参数（如功率上限）可能恢复默认值，需在提交运行前再次执行。

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

原文档给出的运行命令为：

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

即仓库约定由run_and_time.sh负责"运行并计时"。green 与 red 各有一份，两者结构相同、Beam 参数略有差异。以 green 的 v4.1 版本为例（run_and_time.sh），完整内容如下：

#!/bin/bash export PYTHONPATH="." export MODEL="resnet" export SUBMISSION_PLATFORM="tinybox_green" export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192 export LAZYCACHE=0 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=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

脚本分为init（初始化）与run（正式计时）两个阶段，最终都进入同一个入口 examples/mlperf/model_train.py。下面按阶段与变量逐一说明。

6.1 基础配置变量

BS=1536的设计与学习率联动：base_lr默认取7.2 * (BS/1536)（model_train.py），即按 1536 作为基准 batch size 线性缩放学习率，BS=1536时恰好等于基准值 7.2。

6.2 阶段一：init（初始化与编译预热）

BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE

• INITMLPERF=1：进入"初始化"模式。训练循环会先用假数据（fake_data_get）跑前几步，用于触发 Beam Search 内核编译、JIT 预热与内存布局定型，但不产生真实梯度统计（model_train.py）；
• BENCHMARK=10：只跑 10 个训练步即进入基准输出逻辑。此时打印：
  • 中位单步耗时（median_step_time）；
  • 预估总训练时长estimated_total_minutes；
  • 单个 epoch 的全局计算量（global_ops）与内存量（global_mem）（model_train.py）；
• INITMLPERF模式下若启用了 Beam（TRAIN_BEAM/EVAL_BEAM），完成首个 epoch 的编译后会提前break，为正式阶段"热好"全部内核缓存；
• 输出同时通过tee写入以时间戳与随机种子命名的日志文件（如resnet_green_07121534_12345.log），供后续审计。

6.3 阶段二：run（正式计时运行）

PARALLEL=0 RUNMLPERF=1 EVAL_START_EPOCH=3 EVAL_FREQ=4 python3 examples/mlperf/model_train.py | tee -a $LOGFILE

• RUNMLPERF=1：进入正式提交模式，mllog 会记录完整的 MLPerf 合规事件（见第 7 节）；
• EVAL_START_EPOCH=3：从第 3 个 epoch 之后才开始做验证（model_train.py），避免训练初期频繁验证拖慢进度；
• EVAL_FREQ=4：之后每 4 个 epoch 验证一次；同时源码内置兜底逻辑：epoch ≥ 33 后每个 epoch 都验证，以便准确率在 [33, 37] 窗口内达标时立即停止计时（model_train.py）；
• PARALLEL=0：关闭数据加载与计算的重叠并行（该提交采用串行流水，配合STORE_COOKIES机制管理数据预取）；
• RESET_STEP=0：与 v4.1 之前版本不同，正式阶段不在验证前调用train_step.reset()释放编译缓存，从而缩短验证间隙耗时；这是经过调优的提交配置。

6.4 配套开发脚本：dev_run.sh 与 dev_beam.sh

同目录下还有两个开发辅助脚本，供调参而非正式提交使用：

• dev_run.sh：在 run_and_time.sh 基础上追加WANDB=1并固定EVAL_START_EPOCH=3 EVAL_FREQ=4，将每步的 loss、Top-1 准确率、学习率、步耗时、GFLOPS 等指标同步到 Weights & Biases 可视化（model_train.py）；
• dev_beam.sh：设置BENCHMARK=10 DEBUG=2，以调试级日志快速跑 10 步，专门用于验证 Beam Search 编译是否成功、是否有性能回退。

7. 深入源码：训练循环与 MLPerf 合规日志

7.1 优化器与学习率调度

ResNet-50 提交采用经典的两组优化器组合（model_train.py）：

optimizer = LARS(parameters, base_lr, momentum=.9, weight_decay=decay) # decay 默认 2e-4 optimizer_skip = SGD(skip_list, base_lr, momentum=.9, weight_decay=0.0, classic=True) optimizer_group = OptimizerGroup(optimizer, optimizer_skip) scheduler = PolynomialDecayWithWarmup(optimizer, initial_lr=base_lr, end_lr=1e-4, power=?, warmup=lr_warmup_epochs * steps_in_train_epoch)

• 主体参数使用LARS（Layer-wise Adaptive Rate Scaling），momentum 0.9、weight decay 默认 2e-4（DECAY，model_train.py）；
• BN 相关参数（running_mean/running_var 等）走独立的SGD（classic 模式、零 weight decay），避免 LARS 对 BN 统计量造成干扰；
• 两组优化器共享PolynomialDecayWithWarmup调度器（实现见 examples/mlperf/lr_schedulers.py），初始学习率base_lr、终止学习率 1e-4、warmup 2 个 epoch；
• 每步梯度会除以loss_scaler（fp16 下默认 256.0，见 model_train.py）以配合半精度训练，DEFAULT_FLOAT=HALF正是为了让这一路径生效；
• 非同步 BN（UnsyncedBatchNorm）替代标准BatchNorm，配合x.shard_(GPUS, axis=0)实现 6 卡数据并行（model_train.py），默认关闭SYNCBN。

7.2 验证与达标判定

验证循环以EVAL_START_EPOCH/EVAL_FREQ为节奏执行（model_train.py），计算验证集 Top-1 准确率；当achieved标志置真（准确率 ≥TARGET=0.759）时结束训练。训练中每步也会打印top_1_acc（训练集上的 Top-1），便于实时观察收敛趋势：

123 45.67 ms run, 3.21 ms python, 2.10 ms fetch data, 40.36 ms GPU, 2.35 loss, 58.32 acc, 7.200000 LR, 8.90 GB used, 180234.12 GFLOPS

（该输出格式定义于 model_train.py。）

7.3 mllog 合规事件

LOGMLPERF=1时（model_train.py），脚本会初始化MLLOGGER并向 MLPerf 标准 mllog 写入合规事件，包括：

• SUBMISSION_PLATFORM（平台名）、SUBMISSION_BENCHMARK（= RESNET）；
• SEED（随机种子）、TRAIN_SAMPLES/EVAL_SAMPLES（样本数）；
• GLOBAL_BATCH_SIZE、LARS_OPT_BASE_LEARNING_RATE、LARS_OPT_END_LR、LARS_OPT_LR_DECAY_POLY_POWER、LARS_OPT_LEARNING_RATE_WARMUP_EPOCHS、LARS_OPT_MOMENTUM、LARS_OPT_WEIGHT_DECAY；
• EPOCH_START/EPOCH_STOP、EVAL_START/EVAL_STOP等训练生命周期事件（[model_train.py](https://link.gitcode.com/i/ee91a85509b8e2af9c0d102f2304bd1d#L118-L136, L262-L269)）。

SEED=$RANDOM每轮提交随机化种子，保证多轮结果独立；同时种子参与数据打乱（seed*epochs + e，model_train.py），使每个 epoch 的 shuffle 序列可复现。

8. 可调参数速查表

以下变量均通过环境变量注入（默认值取自 model_train.py 源码），供在 dev 脚本中探索：

8.1 Beam Search 编译参数

TRAIN_BEAM=4是 tinygrad 特有的编译器搜索机制：调度器会穷举多种 kernel 优化方案（UOPS 组合、向量化宽度 upcast、本地内存布局等）并选取最快者，代价是编译期显著变长。配套约束参数的作用：

• BEAM_UOPS_MAX：单个 kernel 允许的最大 UOPS 数（green 1500 / red 2000）；
• BEAM_UPCAST_MAX：向量化/上转型最大宽度（green 64 / red 96）；
• BEAM_LOCAL_MAX：本地内存（local memory）最大占用（均为 1024）；
• BEAM_MIN_PROGRESS：搜索剪枝的最小收益阈值（green 10 / red 5），收益低于该值即停止加深；
• BEAM_PADTO=0：不对 kernel 做 padding 对齐；
• IGNORE_JIT_FIRST_BEAM=1：跳过 JIT 首次编译时的 Beam 搜索，避免初始化阶段过慢。

red 的阈值比 green 更激进（UOPS 上限 2000、UPCAST 96、MIN_PROGRESS 5），这是因为 RX 7900 XTX 与 RTX 4090 的指令集与缓存特征不同，需要不同的搜索空间，参见 red 的 run_and_time.sh。

9. 端到端复现清单

综合原文档与上述源码分析，在一台全新 tinybox 上复现该提交的完整步骤为：

1. 安装：git clonetinygrad 后执行python3 -m pip install -e ".[mlperf]"（需 Python 3.10+，green 配 CUDA 12.4，red 配 ROCm 6.1.3）；
2. 驱动层：green 安装 p2p 版 NVIDIA 驱动（出厂默认）；red 写入/etc/modprobe.d/amdgpu.conf禁用 cwsr 并重启验证/sys/module/amdgpu/parameters/cwsr_enable = 0；
3. 数据：IMGNET_TRAIN=1 python3 extra/datasets/imagenet_download.py，确认extra/datasets/imagenet/下 train/val 目录结构完整；
4. 一次性调优（red）：执行 setup.sh（compute 档位、显存时钟、high 性能等级、350W 功率上限）；
5. 开发验证：先跑dev_beam.sh（10 步 + DEBUG=2）确认编译通过；再跑dev_run.sh（WANDB=1）观察 loss 与准确率收敛；
6. 正式提交：执行 run_and_time.sh（red 版同理），init 阶段完成编译预热后自动进入计时阶段，日志写入resnet_{green,red}_MMDDHHMM_SEED.log；
7. 审计：核对日志中 mllog 事件的完整性（平台、种子、batch size、LARS 超参、epoch/eval 生命周期），确认 Top-1 达到 0.759 目标。

10. 小结

这份 MLPerf Training v4.1 提交展示了 tinygrad 在真实多卡训练基准上的完整落地路径：从驱动级特化（NVIDIA P2P / AMD cwsr 关闭 + ROCm 调优）、大规模数据管线（138GB ImageNet 下载与按 synset 重组）、半精度 LARS 训练（fp16 + loss scaling + 非同步 BN 数据并行），到用 Beam Search 编译器搜索为每张卡定制最快 kernel，并以 mllog 标准格式记录全部合规信息。与 README 中的核心步骤配合 model_train.py 与各 run_and_time.sh 脚本，你可以在自己的 tinybox green/red 上完整复现这套可审计、可复现、可调优的 ResNet-50 训练提交流程。

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