基于 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.pydef 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.pyexamples/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.75975.9%作为达标线训练最多进行 37 个 epochgetenv(EPOCHS, 37)当验证准确率达到目标后提前停止计时。2. 环境要求安装 tinygrad 与 mlperf-logging原文档要求从 master 分支安装 tinygrad 与 mlperf-logginggit 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 redMLPerf 提交以 system 为单位。v4.1 提交包含两套系统systems 目录项目tinybox greentinybox_green.jsontinybox redtinybox_red.jsonCPUAMD EPYC 753232 核 / 64 线程AMD EPYC 753232 核 / 64 线程内存128GB8×16GB DDR4128GB8×16GB DDR4加速卡6× NVIDIA GeForce RTX 409024GB GDDR6X6× AMD Radeon RX 7900 XTX24GB GDDR6互连PCIe 4.0 x16PCIe 4.0 x16软件栈CUDA 12.4ROCm 6.1.3操作系统Ubuntu 22.04.4Ubuntu 22.04.4存储4TB RAID 阵列 1TB 系统盘NVMe SSD4TB RAID 阵列 1TB 系统盘NVMe SSDtinygrad 版本commit b5546912e24e0a864b35924da4efa5d71cfe368b同左两套系统均为 6 卡、24GB 显存、PCIe 4.0 x16 直连的架构因此共享大部分运行脚本仅在驱动级设置上分叉。3.1 tinybox green安装 p2p 驱动green 使用 NVIDIA GPU需要安装支持 P2Ppeer-to-peer通信的驱动。原文档指出应按照 tinygrad 维护的 open-gpu-kernel-modules 的550.54.15-p2p分支 README 安装该驱动并说明这是生产环境 tinybox green 的默认配置即无需额外手工操作。3.2 tinybox red禁用 cwsrred 使用 AMD GPU需要禁用 AMD 驱动中的cwsrCompute Wave Save/Restore计算波前保存/恢复机制。这也是生产环境 tinybox red 的默认配置若在新机器上手动配置步骤如下sudo vi /etc/modprobe.d/amdgpu.conf cat EOF /etc/modprobe.d/amdgpu.conf options amdgpu cwsr_enable0 EOF sudo update-initramfs -u sudo reboot # validate sudo cat /sys/module/amdgpu/parameters/cwsr_enable # 0关键点通过内核模块参数cwsr_enable0关闭 cwsr写入/etc/modprobe.d/amdgpu.conf使其在每次启动时生效修改后需update-initramfs -u重建 initramfs 并重启验证步骤不可省略重启后检查/sys/module/amdgpu/parameters/cwsr_enable输出必须为0否则设置未生效。关闭 cwsr 的动机在于性能cwsr 会在上下文切换时保存/恢复计算波前状态禁用后可减少 GPU 调度的额外开销对长时训练吞吐有明显收益。4. 数据准备下载并校验 ImageNet原文档给出的数据下载命令为IMGNET_TRAIN1 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_TRAIN1并预留足够的磁盘空间验证集 7GB 训练集 138GB 解压空间仓库系统清单显示 tinybox 标配 4TB RAID 阵列正是为此设计。数据会被组织在extra/datasets/imagenet/目录下训练与验证文件列表分别由 extra/datasets/imagenet.py 的get_train_files()/get_val_files()提供model_train.py 据此计算每个 epoch 的步数。5. 一次性设置仅 tinybox redtinybox 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 MODELresnet export SUBMISSION_PLATFORMtinybox_green export DEFAULT_FLOATHALF GPUS6 BS1536 EVAL_BS192 export LAZYCACHE0 RESET_STEP0 export TRAIN_BEAM4 IGNORE_JIT_FIRST_BEAM1 BEAM_UOPS_MAX1500 BEAM_UPCAST_MAX64 BEAM_LOCAL_MAX1024 BEAM_MIN_PROGRESS10 BEAM_PADTO0 # pip install -e .[mlperf] export LOGMLPERF1 export SEED$RANDOM DATETIME$(date %m%d%H%M) LOGFILEresnet_green_${DATETIME}_${SEED}.log # init BENCHMARK10 INITMLPERF1 python3 examples/mlperf/model_train.py | tee $LOGFILE # run PARALLEL0 RUNMLPERF1 EVAL_START_EPOCH3 EVAL_FREQ4 python3 examples/mlperf/model_train.py | tee -a $LOGFILE脚本分为init初始化与run正式计时两个阶段最终都进入同一个入口 examples/mlperf/model_train.py。下面按阶段与变量逐一说明。6.1 基础配置变量变量值含义源码位置MODELresnet选择基准模型model_spec.py 遍历入口SUBMISSION_PLATFORMtinybox_green/tinybox_red提交平台名写入 mllogmodel_train.pyDEFAULT_FLOATHALF全局默认浮点类型为 fp16halfmodel_train.py 记录到 configGPUS6使用 6 张卡训练model_train.pyBS1536训练 batch size全局即每卡 256model_train.pyEVAL_BS192验证 batch sizemodel_train.pyLAZYCACHE0关闭 lazy 调度缓存保证编译路径可复现tinygrad 全局环境变量RESET_STEP0不在验证前重置 train_step 内存见下model_train.pyBS1536的设计与学习率联动base_lr默认取7.2 * (BS/1536)model_train.py即按 1536 作为基准 batch size 线性缩放学习率BS1536时恰好等于基准值 7.2。6.2 阶段一init初始化与编译预热BENCHMARK10 INITMLPERF1 python3 examples/mlperf/model_train.py | tee $LOGFILEINITMLPERF1进入初始化模式。训练循环会先用假数据fake_data_get跑前几步用于触发 Beam Search 内核编译、JIT 预热与内存布局定型但不产生真实梯度统计model_train.pyBENCHMARK10只跑 10 个训练步即进入基准输出逻辑。此时打印中位单步耗时median_step_time预估总训练时长estimated_total_minutes单个 epoch 的全局计算量global_ops与内存量global_memmodel_train.pyINITMLPERF模式下若启用了 BeamTRAIN_BEAM/EVAL_BEAM完成首个 epoch 的编译后会提前break为正式阶段热好全部内核缓存输出同时通过tee写入以时间戳与随机种子命名的日志文件如resnet_green_07121534_12345.log供后续审计。6.3 阶段二run正式计时运行PARALLEL0 RUNMLPERF1 EVAL_START_EPOCH3 EVAL_FREQ4 python3 examples/mlperf/model_train.py | tee -a $LOGFILERUNMLPERF1进入正式提交模式mllog 会记录完整的 MLPerf 合规事件见第 7 节EVAL_START_EPOCH3从第 3 个 epoch 之后才开始做验证model_train.py避免训练初期频繁验证拖慢进度EVAL_FREQ4之后每 4 个 epoch 验证一次同时源码内置兜底逻辑epoch ≥ 33 后每个 epoch 都验证以便准确率在 [33, 37] 窗口内达标时立即停止计时model_train.pyPARALLEL0关闭数据加载与计算的重叠并行该提交采用串行流水配合STORE_COOKIES机制管理数据预取RESET_STEP0与 v4.1 之前版本不同正式阶段不在验证前调用train_step.reset()释放编译缓存从而缩短验证间隙耗时这是经过调优的提交配置。6.4 配套开发脚本dev_run.sh 与 dev_beam.sh同目录下还有两个开发辅助脚本供调参而非正式提交使用dev_run.sh在 run_and_time.sh 基础上追加WANDB1并固定EVAL_START_EPOCH3 EVAL_FREQ4将每步的 loss、Top-1 准确率、学习率、步耗时、GFLOPS 等指标同步到 Weights Biases 可视化model_train.pydev_beam.sh设置BENCHMARK10 DEBUG2以调试级日志快速跑 10 步专门用于验证 Beam Search 编译是否成功、是否有性能回退。7. 深入源码训练循环与 MLPerf 合规日志7.1 优化器与学习率调度ResNet-50 提交采用经典的两组优化器组合model_train.pyoptimizer LARS(parameters, base_lr, momentum.9, weight_decaydecay) # decay 默认 2e-4 optimizer_skip SGD(skip_list, base_lr, momentum.9, weight_decay0.0, classicTrue) optimizer_group OptimizerGroup(optimizer, optimizer_skip) scheduler PolynomialDecayWithWarmup(optimizer, initial_lrbase_lr, end_lr1e-4, power?, warmuplr_warmup_epochs * steps_in_train_epoch)主体参数使用LARSLayer-wise Adaptive Rate Scalingmomentum 0.9、weight decay 默认 2e-4DECAYmodel_train.pyBN 相关参数running_mean/running_var 等走独立的SGDclassic 模式、零 weight decay避免 LARS 对 BN 统计量造成干扰两组优化器共享PolynomialDecayWithWarmup调度器实现见 examples/mlperf/lr_schedulers.py初始学习率base_lr、终止学习率 1e-4、warmup 2 个 epoch每步梯度会除以loss_scalerfp16 下默认 256.0见 model_train.py以配合半精度训练DEFAULT_FLOATHALF正是为了让这一路径生效非同步 BNUnsyncedBatchNorm替代标准BatchNorm配合x.shard_(GPUS, axis0)实现 6 卡数据并行model_train.py默认关闭SYNCBN。7.2 验证与达标判定验证循环以EVAL_START_EPOCH/EVAL_FREQ为节奏执行model_train.py计算验证集 Top-1 准确率当achieved标志置真准确率 ≥TARGET0.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 合规事件LOGMLPERF1时model_train.py脚本会初始化MLLOGGER并向 MLPerf 标准 mllog 写入合规事件包括SUBMISSION_PLATFORM平台名、SUBMISSION_BENCHMARK RESNETSEED随机种子、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_DECAYEPOCH_START/EPOCH_STOP、EVAL_START/EVAL_STOP等训练生命周期事件[model_train.py](https://link.gitcode.com/i/ee91a85509b8e2af9c0d102f2304bd1d#L118-L136, L262-L269)。SEED$RANDOM每轮提交随机化种子保证多轮结果独立同时种子参与数据打乱seed*epochs emodel_train.py使每个 epoch 的 shuffle 序列可复现。8. 可调参数速查表以下变量均通过环境变量注入默认值取自 model_train.py 源码供在 dev 脚本中探索变量默认值说明EPOCHS37最大训练 epoch 数BS104 * GPUS训练 batch size提交用 1536EVAL_BSBS验证 batch size提交用 192LR7.2 * (BS/1536)基准学习率随 BS 线性缩放WARMUP_EPOCHS2学习率 warmup epoch 数DECAY2e-4LARS weight decayLOSS_SCALERfp16 下 256.0梯度损失缩放TARGET0.759达标 Top-1 准确率EVAL_START_EPOCH0开始验证的 epoch提交用 3EVAL_FREQ1验证频率提交用 4TRAIN_BEAM/EVAL_BEAMBEAM.value训练/验证使用的 Beam Search 档位提交用 4TRAIN_BEAM相关BEAM_*见各脚本Beam 搜索约束UOPS/UPCAST/LOCAL 上限、最小收益、PADTOSYNCBN关是否启用同步 BN提交关闭RESUME空从 checkpoint 恢复训练WANDB/WANDB_RESUME关是否上报 wandb 及恢复 run idBENCHMARK0跑 N 步后输出基准信息ASSERT_MIN_STEP_TIME空单步耗时断言性能回归保护STORE_COOKIES1数据预取 cookie 保留数量8.1 Beam Search 编译参数TRAIN_BEAM4是 tinygrad 特有的编译器搜索机制调度器会穷举多种 kernel 优化方案UOPS 组合、向量化宽度 upcast、本地内存布局等并选取最快者代价是编译期显著变长。配套约束参数的作用BEAM_UOPS_MAX单个 kernel 允许的最大 UOPS 数green 1500 / red 2000BEAM_UPCAST_MAX向量化/上转型最大宽度green 64 / red 96BEAM_LOCAL_MAX本地内存local memory最大占用均为 1024BEAM_MIN_PROGRESS搜索剪枝的最小收益阈值green 10 / red 5收益低于该值即停止加深BEAM_PADTO0不对 kernel 做 padding 对齐IGNORE_JIT_FIRST_BEAM1跳过 JIT 首次编译时的 Beam 搜索避免初始化阶段过慢。red 的阈值比 green 更激进UOPS 上限 2000、UPCAST 96、MIN_PROGRESS 5这是因为 RX 7900 XTX 与 RTX 4090 的指令集与缓存特征不同需要不同的搜索空间参见 red 的 run_and_time.sh。9. 端到端复现清单综合原文档与上述源码分析在一台全新 tinybox 上复现该提交的完整步骤为安装git clonetinygrad 后执行python3 -m pip install -e .[mlperf]需 Python 3.10green 配 CUDA 12.4red 配 ROCm 6.1.3驱动层green 安装 p2p 版 NVIDIA 驱动出厂默认red 写入/etc/modprobe.d/amdgpu.conf禁用 cwsr 并重启验证/sys/module/amdgpu/parameters/cwsr_enable 0数据IMGNET_TRAIN1 python3 extra/datasets/imagenet_download.py确认extra/datasets/imagenet/下 train/val 目录结构完整一次性调优red执行 setup.shcompute 档位、显存时钟、high 性能等级、350W 功率上限开发验证先跑dev_beam.sh10 步 DEBUG2确认编译通过再跑dev_run.shWANDB1观察 loss 与准确率收敛正式提交执行 run_and_time.shred 版同理init 阶段完成编译预热后自动进入计时阶段日志写入resnet_{green,red}_MMDDHHMM_SEED.log审计核对日志中 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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 基于 tinygrad 复现 MLPerf 训练提交：ResNet-50 在 tinybox green/red 上的完整部署与调优指南
基于 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.pydef 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.pyexamples/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.75975.9%作为达标线训练最多进行 37 个 epochgetenv(EPOCHS, 37)当验证准确率达到目标后提前停止计时。2. 环境要求安装 tinygrad 与 mlperf-logging原文档要求从 master 分支安装 tinygrad 与 mlperf-logginggit 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 redMLPerf 提交以 system 为单位。v4.1 提交包含两套系统systems 目录项目tinybox greentinybox_green.jsontinybox redtinybox_red.jsonCPUAMD EPYC 753232 核 / 64 线程AMD EPYC 753232 核 / 64 线程内存128GB8×16GB DDR4128GB8×16GB DDR4加速卡6× NVIDIA GeForce RTX 409024GB GDDR6X6× AMD Radeon RX 7900 XTX24GB GDDR6互连PCIe 4.0 x16PCIe 4.0 x16软件栈CUDA 12.4ROCm 6.1.3操作系统Ubuntu 22.04.4Ubuntu 22.04.4存储4TB RAID 阵列 1TB 系统盘NVMe SSD4TB RAID 阵列 1TB 系统盘NVMe SSDtinygrad 版本commit b5546912e24e0a864b35924da4efa5d71cfe368b同左两套系统均为 6 卡、24GB 显存、PCIe 4.0 x16 直连的架构因此共享大部分运行脚本仅在驱动级设置上分叉。3.1 tinybox green安装 p2p 驱动green 使用 NVIDIA GPU需要安装支持 P2Ppeer-to-peer通信的驱动。原文档指出应按照 tinygrad 维护的 open-gpu-kernel-modules 的550.54.15-p2p分支 README 安装该驱动并说明这是生产环境 tinybox green 的默认配置即无需额外手工操作。3.2 tinybox red禁用 cwsrred 使用 AMD GPU需要禁用 AMD 驱动中的cwsrCompute Wave Save/Restore计算波前保存/恢复机制。这也是生产环境 tinybox red 的默认配置若在新机器上手动配置步骤如下sudo vi /etc/modprobe.d/amdgpu.conf cat EOF /etc/modprobe.d/amdgpu.conf options amdgpu cwsr_enable0 EOF sudo update-initramfs -u sudo reboot # validate sudo cat /sys/module/amdgpu/parameters/cwsr_enable # 0关键点通过内核模块参数cwsr_enable0关闭 cwsr写入/etc/modprobe.d/amdgpu.conf使其在每次启动时生效修改后需update-initramfs -u重建 initramfs 并重启验证步骤不可省略重启后检查/sys/module/amdgpu/parameters/cwsr_enable输出必须为0否则设置未生效。关闭 cwsr 的动机在于性能cwsr 会在上下文切换时保存/恢复计算波前状态禁用后可减少 GPU 调度的额外开销对长时训练吞吐有明显收益。4. 数据准备下载并校验 ImageNet原文档给出的数据下载命令为IMGNET_TRAIN1 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_TRAIN1并预留足够的磁盘空间验证集 7GB 训练集 138GB 解压空间仓库系统清单显示 tinybox 标配 4TB RAID 阵列正是为此设计。数据会被组织在extra/datasets/imagenet/目录下训练与验证文件列表分别由 extra/datasets/imagenet.py 的get_train_files()/get_val_files()提供model_train.py 据此计算每个 epoch 的步数。5. 一次性设置仅 tinybox redtinybox 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 MODELresnet export SUBMISSION_PLATFORMtinybox_green export DEFAULT_FLOATHALF GPUS6 BS1536 EVAL_BS192 export LAZYCACHE0 RESET_STEP0 export TRAIN_BEAM4 IGNORE_JIT_FIRST_BEAM1 BEAM_UOPS_MAX1500 BEAM_UPCAST_MAX64 BEAM_LOCAL_MAX1024 BEAM_MIN_PROGRESS10 BEAM_PADTO0 # pip install -e .[mlperf] export LOGMLPERF1 export SEED$RANDOM DATETIME$(date %m%d%H%M) LOGFILEresnet_green_${DATETIME}_${SEED}.log # init BENCHMARK10 INITMLPERF1 python3 examples/mlperf/model_train.py | tee $LOGFILE # run PARALLEL0 RUNMLPERF1 EVAL_START_EPOCH3 EVAL_FREQ4 python3 examples/mlperf/model_train.py | tee -a $LOGFILE脚本分为init初始化与run正式计时两个阶段最终都进入同一个入口 examples/mlperf/model_train.py。下面按阶段与变量逐一说明。6.1 基础配置变量变量值含义源码位置MODELresnet选择基准模型model_spec.py 遍历入口SUBMISSION_PLATFORMtinybox_green/tinybox_red提交平台名写入 mllogmodel_train.pyDEFAULT_FLOATHALF全局默认浮点类型为 fp16halfmodel_train.py 记录到 configGPUS6使用 6 张卡训练model_train.pyBS1536训练 batch size全局即每卡 256model_train.pyEVAL_BS192验证 batch sizemodel_train.pyLAZYCACHE0关闭 lazy 调度缓存保证编译路径可复现tinygrad 全局环境变量RESET_STEP0不在验证前重置 train_step 内存见下model_train.pyBS1536的设计与学习率联动base_lr默认取7.2 * (BS/1536)model_train.py即按 1536 作为基准 batch size 线性缩放学习率BS1536时恰好等于基准值 7.2。6.2 阶段一init初始化与编译预热BENCHMARK10 INITMLPERF1 python3 examples/mlperf/model_train.py | tee $LOGFILEINITMLPERF1进入初始化模式。训练循环会先用假数据fake_data_get跑前几步用于触发 Beam Search 内核编译、JIT 预热与内存布局定型但不产生真实梯度统计model_train.pyBENCHMARK10只跑 10 个训练步即进入基准输出逻辑。此时打印中位单步耗时median_step_time预估总训练时长estimated_total_minutes单个 epoch 的全局计算量global_ops与内存量global_memmodel_train.pyINITMLPERF模式下若启用了 BeamTRAIN_BEAM/EVAL_BEAM完成首个 epoch 的编译后会提前break为正式阶段热好全部内核缓存输出同时通过tee写入以时间戳与随机种子命名的日志文件如resnet_green_07121534_12345.log供后续审计。6.3 阶段二run正式计时运行PARALLEL0 RUNMLPERF1 EVAL_START_EPOCH3 EVAL_FREQ4 python3 examples/mlperf/model_train.py | tee -a $LOGFILERUNMLPERF1进入正式提交模式mllog 会记录完整的 MLPerf 合规事件见第 7 节EVAL_START_EPOCH3从第 3 个 epoch 之后才开始做验证model_train.py避免训练初期频繁验证拖慢进度EVAL_FREQ4之后每 4 个 epoch 验证一次同时源码内置兜底逻辑epoch ≥ 33 后每个 epoch 都验证以便准确率在 [33, 37] 窗口内达标时立即停止计时model_train.pyPARALLEL0关闭数据加载与计算的重叠并行该提交采用串行流水配合STORE_COOKIES机制管理数据预取RESET_STEP0与 v4.1 之前版本不同正式阶段不在验证前调用train_step.reset()释放编译缓存从而缩短验证间隙耗时这是经过调优的提交配置。6.4 配套开发脚本dev_run.sh 与 dev_beam.sh同目录下还有两个开发辅助脚本供调参而非正式提交使用dev_run.sh在 run_and_time.sh 基础上追加WANDB1并固定EVAL_START_EPOCH3 EVAL_FREQ4将每步的 loss、Top-1 准确率、学习率、步耗时、GFLOPS 等指标同步到 Weights Biases 可视化model_train.pydev_beam.sh设置BENCHMARK10 DEBUG2以调试级日志快速跑 10 步专门用于验证 Beam Search 编译是否成功、是否有性能回退。7. 深入源码训练循环与 MLPerf 合规日志7.1 优化器与学习率调度ResNet-50 提交采用经典的两组优化器组合model_train.pyoptimizer LARS(parameters, base_lr, momentum.9, weight_decaydecay) # decay 默认 2e-4 optimizer_skip SGD(skip_list, base_lr, momentum.9, weight_decay0.0, classicTrue) optimizer_group OptimizerGroup(optimizer, optimizer_skip) scheduler PolynomialDecayWithWarmup(optimizer, initial_lrbase_lr, end_lr1e-4, power?, warmuplr_warmup_epochs * steps_in_train_epoch)主体参数使用LARSLayer-wise Adaptive Rate Scalingmomentum 0.9、weight decay 默认 2e-4DECAYmodel_train.pyBN 相关参数running_mean/running_var 等走独立的SGDclassic 模式、零 weight decay避免 LARS 对 BN 统计量造成干扰两组优化器共享PolynomialDecayWithWarmup调度器实现见 examples/mlperf/lr_schedulers.py初始学习率base_lr、终止学习率 1e-4、warmup 2 个 epoch每步梯度会除以loss_scalerfp16 下默认 256.0见 model_train.py以配合半精度训练DEFAULT_FLOATHALF正是为了让这一路径生效非同步 BNUnsyncedBatchNorm替代标准BatchNorm配合x.shard_(GPUS, axis0)实现 6 卡数据并行model_train.py默认关闭SYNCBN。7.2 验证与达标判定验证循环以EVAL_START_EPOCH/EVAL_FREQ为节奏执行model_train.py计算验证集 Top-1 准确率当achieved标志置真准确率 ≥TARGET0.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 合规事件LOGMLPERF1时model_train.py脚本会初始化MLLOGGER并向 MLPerf 标准 mllog 写入合规事件包括SUBMISSION_PLATFORM平台名、SUBMISSION_BENCHMARK RESNETSEED随机种子、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_DECAYEPOCH_START/EPOCH_STOP、EVAL_START/EVAL_STOP等训练生命周期事件[model_train.py](https://link.gitcode.com/i/ee91a85509b8e2af9c0d102f2304bd1d#L118-L136, L262-L269)。SEED$RANDOM每轮提交随机化种子保证多轮结果独立同时种子参与数据打乱seed*epochs emodel_train.py使每个 epoch 的 shuffle 序列可复现。8. 可调参数速查表以下变量均通过环境变量注入默认值取自 model_train.py 源码供在 dev 脚本中探索变量默认值说明EPOCHS37最大训练 epoch 数BS104 * GPUS训练 batch size提交用 1536EVAL_BSBS验证 batch size提交用 192LR7.2 * (BS/1536)基准学习率随 BS 线性缩放WARMUP_EPOCHS2学习率 warmup epoch 数DECAY2e-4LARS weight decayLOSS_SCALERfp16 下 256.0梯度损失缩放TARGET0.759达标 Top-1 准确率EVAL_START_EPOCH0开始验证的 epoch提交用 3EVAL_FREQ1验证频率提交用 4TRAIN_BEAM/EVAL_BEAMBEAM.value训练/验证使用的 Beam Search 档位提交用 4TRAIN_BEAM相关BEAM_*见各脚本Beam 搜索约束UOPS/UPCAST/LOCAL 上限、最小收益、PADTOSYNCBN关是否启用同步 BN提交关闭RESUME空从 checkpoint 恢复训练WANDB/WANDB_RESUME关是否上报 wandb 及恢复 run idBENCHMARK0跑 N 步后输出基准信息ASSERT_MIN_STEP_TIME空单步耗时断言性能回归保护STORE_COOKIES1数据预取 cookie 保留数量8.1 Beam Search 编译参数TRAIN_BEAM4是 tinygrad 特有的编译器搜索机制调度器会穷举多种 kernel 优化方案UOPS 组合、向量化宽度 upcast、本地内存布局等并选取最快者代价是编译期显著变长。配套约束参数的作用BEAM_UOPS_MAX单个 kernel 允许的最大 UOPS 数green 1500 / red 2000BEAM_UPCAST_MAX向量化/上转型最大宽度green 64 / red 96BEAM_LOCAL_MAX本地内存local memory最大占用均为 1024BEAM_MIN_PROGRESS搜索剪枝的最小收益阈值green 10 / red 5收益低于该值即停止加深BEAM_PADTO0不对 kernel 做 padding 对齐IGNORE_JIT_FIRST_BEAM1跳过 JIT 首次编译时的 Beam 搜索避免初始化阶段过慢。red 的阈值比 green 更激进UOPS 上限 2000、UPCAST 96、MIN_PROGRESS 5这是因为 RX 7900 XTX 与 RTX 4090 的指令集与缓存特征不同需要不同的搜索空间参见 red 的 run_and_time.sh。9. 端到端复现清单综合原文档与上述源码分析在一台全新 tinybox 上复现该提交的完整步骤为安装git clonetinygrad 后执行python3 -m pip install -e .[mlperf]需 Python 3.10green 配 CUDA 12.4red 配 ROCm 6.1.3驱动层green 安装 p2p 版 NVIDIA 驱动出厂默认red 写入/etc/modprobe.d/amdgpu.conf禁用 cwsr 并重启验证/sys/module/amdgpu/parameters/cwsr_enable 0数据IMGNET_TRAIN1 python3 extra/datasets/imagenet_download.py确认extra/datasets/imagenet/下 train/val 目录结构完整一次性调优red执行 setup.shcompute 档位、显存时钟、high 性能等级、350W 功率上限开发验证先跑dev_beam.sh10 步 DEBUG2确认编译通过再跑dev_run.shWANDB1观察 loss 与准确率收敛正式提交执行 run_and_time.shred 版同理init 阶段完成编译预热后自动进入计时阶段日志写入resnet_{green,red}_MMDDHHMM_SEED.log审计核对日志中 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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考