HRM-Text

环境要求

• Python 环境（源码安装时）
• PyTorch、CUDA 与 FlashAttention 3（源码安装时，版本见 docker/Dockerfile）
• Docker 及 NVIDIA Container Toolkit（GPU 可见）
• 目标尺寸对应 GPU：L/0.6B 需 8×H100 单节点，XL/1B 需 16×H100 双节点
• Weights & Biases 账号用于训练日志记录

安装与启动步骤

1. 1克隆仓库

   把 HRM-Text 源码拉到本地，后续 Docker 命令需在仓库目录下执行。

   git clone https://github.com/sapientinc/HRM-Text.git
   cd HRM-Text

   复制

2. 2启动 Docker 容器

   用官方镜像启动，--gpus all 让容器可见 GPU，并挂载当前仓库到 /workspace。

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

   复制

3. 3源码安装依赖

   不用 Docker 时先自行装好 PyTorch、CUDA、FlashAttention 3，再装 Python 依赖。

   pip install -r requirements.txt

   复制

4. 4登录 W&B

   训练指标会记录到 Weights & Biases，启动训练前需先登录。

   wandb login

   复制

5. 5无界面登录 W&B

   无浏览器环境时，从 wandb.ai/authorize 获取 API Key 后直接登录。

   wandb login <API_KEY>

   复制

关键配置

如何确认成功

容器正常启动并进入镜像内环境；wandb login 成功后可开始按配置发起预训练。

常见问题

注意事项

• Docker 运行需 NVIDIA Container Toolkit 让容器可见 GPU。
• 多节点建议统一共享工作区布局，如 /shared/HRM-Text 与 /shared/data_io。
• 预训练成本较高，L 约 $800、XL 约 $1472。