企业官网建设流程全解析

HRM-Text环境搭建指南：Docker镜像一键部署与H100双节点集群配置详解

【免费下载链接】HRM-TextHRM-Text is a 1B text generation model based on the HRM architecture, strengthened by task completion and latent space reasoning.项目地址: https://gitcode.com/gh_mirrors/hr/HRM-Text

HRM-Text 是一个基于 HRM 架构的 10 亿参数文本生成模型，本文带你用 Docker 镜像一键完成 HRM-Text 环境搭建，并详解 H100 双节点训练集群的完整配置流程。只需按步骤操作，几十美元级的算力预算就能从零预训练一个基础模型。

💡 提示：由于注意力计算路径依赖 FlashAttention 3 的 Hopper 专用内核，训练需要 80GB 的 H100 显卡。

一、先搞清楚资源需求：为什么是 H100？

HRM-Text 基于 PyTorch FSDP2、FlashAttention 3 与 PrefixLM 序列打包构建，官方给出两档参考训练配置：

*按每 H100 小时 $2 估算。

HRM-Text 1B 模型在 16 块 GPU 上用 1.9 天训练完成，MATH、DROP、ARC-C、MMLU 等基准成绩超过训练开销高几十到几万倍的模型——这正是环境搭建值得认真做的理由。

二、Docker 镜像一键部署（推荐方式）

官方发布的sapientai/hrm-text:latest镜像已内置全部依赖：Python 3.13、CUDA 12.8.1、cu128 版 PyTorch、FlashAttention 3，甚至为多机训练预装好了 InfiniBand/RDMA 驱动，完整清单见 docker/Dockerfile。

先确保宿主机装有 NVIDIA Container Toolkit（让 Docker 能识别 GPU），然后在仓库根目录执行：

docker run --gpus all --ipc=host --network=host -it \ -v "$PWD":/workspace \ sapientai/hrm-text:latest

各参数作用一目了然：

• --gpus all：把全部 NVIDIA GPU 暴露给容器
• --ipc=host：共享宿主进程间通信空间，避免 NCCL 共享内存不足
• --network=host：使用宿主网络，多节点 NCCL 通信的必选项
• -v "$PWD":/workspace：把代码目录挂载为容器内工作目录

备选：从源码手动安装环境

不用 Docker 时，需先自行安装 PyTorch、CUDA 与 FlashAttention 3（测试版本以 docker/Dockerfile 为准），再安装 Python 依赖：

pip install -r requirements.txt

核心依赖列表（含flash_attn_3、vllm、wandb、hydra-core等）维护在 requirements.txt。

三、H100 双节点集群配置详解

1. 共享存储与训练数据准备

双节点训练时，让每个节点挂载路径完全相同的共享工作区，避免各进程间代码与数据出现版本漂移，推荐目录布局：

/shared/ |-- HRM-Text/ |--- checkpoints/ |-- data_io/

训练数据由配套的data_io管线清洗并分词后，在每个训练节点上独立做分层采样（采样快速且确定性，各节点结果一致）：

cd <DATA_IO_PATH> python sample_tokenized.py epochs=4 output_path=/dev/shm/sampled

⚠️ 采样条目的epochs数必须与训练配置保持一致，默认 4 个 epoch。

2. 验证多机通信（NCCL）

在启动长任务前，至少确认torchrun能在目标节点间完成初始化；如果集群提供nccl-tests，建议跑一轮机内与机间带宽测试——不要等 46 小时后才发现问题出在网络上。

3. 配置训练日志

HRM-Text 通过 Weights & Biases 记录训练指标，启动前执行wandb login登录；无头环境可直接wandb login <API_KEY>。

4. 双节点启动预训练

在每个节点上分别运行（XL 规模，共 16 卡）：

OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 \ torchrun \ --nproc_per_node=8 \ --nnodes=2 \ --node_rank=<NODE_RANK> \ --master_addr=<MASTER_ADDR> \ --master_port=<MASTER_PORT> \ pretrain.py

如果是单节点 8×H100 跑 L 规模，一条命令即可：

OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 \ torchrun --nproc_per_node=8 pretrain.py \ arch/size@arch=L lr=2.5e-4 global_batch_size=172032

入口脚本 pretrain.py 负责 FSDP2 包装、优化器、学习率调度与分布式 checkpoint 保存；XL 规模的网络结构定义在 config/arch/size/XL.yaml。

5. 检查点保存注意事项

Checkpoint 每个 epoch 保存一次，默认输出到checkpoints/。多节点模式下每个节点只保存自己的分片，务必挂载共享存储，否则评估和导出会缺分片。

四、训练后：一键评估与模型导出

评估会自动加载最新 epoch 的 checkpoint：

python -m evaluation.main ckpt_path="checkpoints/..."

• 需要 1 块 80GB GPU；显存不足时追加generation_config.batch_size=16
• 只跑指定基准：追加run_only=[MATH,DROP,ARC,MMLU]

评估完成后，可导出为 HuggingFace Transformers 格式方便接入生态：

python -m conversion.convert_to_hf \ --ckpt_path "checkpoints/..." \ --out_dir "<OUTPUT_PATH>"

转换逻辑位于 conversion/convert_to_hf.py，默认使用 checkpoint 中的 EMA 权重。

五、常见问题速查

按照以上步骤，你就能在 Docker 容器里一键拉起完整环境，并让 16 块 H100 组成的双节点集群稳定跑起 HRM-Text 预训练。🚀

【免费下载链接】HRM-TextHRM-Text is a 1B text generation model based on the HRM architecture, strengthened by task completion and latent space reasoning.项目地址: https://gitcode.com/gh_mirrors/hr/HRM-Text