HRM(Hierarchical Reasoning Model)部署训练记录

前两天朋友刷到HRM这个27M模型的文章，想让我试着部署训练一下。此文用于记录部署过程
前期准备
克隆仓库
sapientinc/HRM
安装CUDA
我的CUDA是已经安装好的12.8版本，安装过程不再赘述
安装torch
torch版本如下 Version: 2.7.1+cu128
pip install torch torchvision torchaudio -f https://mirrors.aliyun.com/pytorch-wheels/cu128/
这里贴出的链接为CUDA12.8版本的阿里云镜像torch安装
安装Flash Attention
这里由于我的系统是windows所以我使用flash-attention-for-windows
github地址
Pip install flash_attn-2.8.2+cu128torch2.7.1cxx11abiFALSEfullbackward-cp311-cp311-win_amd64.whl
这里cu128指CUDA12.8
torch版本为2.7.1
python版本3.11
安装依赖
pip install -r requirements.txt
注册并创建wandb key
由于项目使用wandb记录数据，所以此处需要注册wandb并且在训练时要保持网络畅通
pip install wandb
安装好后使用key登录
wandb login
安装triton
虽然原文没有提到，但是我在部署过程中发现需要，故此添加安装
需要注意triton版本、torch版本以及CUDA版本需要对应

同样我这里使用的windows版
github链接
此处我使用的为3.3版本 Version: 3.3.1.post19
pip install -U "triton-windows