OpenRLHF

本文在单卡昇腾 NPU 上安装 OpenRLHF，并对 Qwen/Qwen2.5-0.5B-Instruct 完成 4 步 LoRA 监督微调。

前置条件

硬件

Atlas 800T / 900 A2 训练系列（Ascend 910B），单卡。

软件

类别	要求
CANN	toolkit 与驱动已安装，并能 `source set_env.sh`
Python	3.12
PyTorch	`torch==2.11.0` 与 `torch_npu==2.11.0`
OpenRLHF	从 GitHub 克隆当前正式 Release，见第 6 节
模型	Qwen/Qwen2.5-0.5B-Instruct

配套机器：Atlas 900 A2 PODc（Ascend 910B4）。配套镜像：swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:9.1.0-910b-ubuntu22.04-py3.12。上游安装说明见 OpenRLHF Quick Start — Installation。

1. 加载 CANN 环境

source /usr/local/Ascend/ascend-toolkit/set_env.sh
export PATH=/usr/local/sbin:/usr/local/bin:$PATH

2. 检查环境

确认 NPU 在线；找不到 npu-smi 时参阅快速安装昇腾环境。

npu-smi info

检查 Python 版本：

python --version

输出结果如下：

Python 3.12.xxx

3. 安装 PyTorch NPU 栈

从 PyTorch CPU 轮子和华为昇腾索引安装与 CANN 9.1.0 匹配的 torch==2.11.0 / torch_npu==2.11.0，并一并装上 numpy、pyyaml。版本行出现 2.11.0+cpu 属正常，npu_available 须为 True。

python -m pip install --extra-index-url https://download.pytorch.org/whl/cpu torch==2.11.0
python -m pip install --extra-index-url https://repo.huaweicloud.com/ascend/repos/pypi \
  torch_npu==2.11.0 numpy pyyaml
python -c "import numpy, yaml, torch, torch_npu; print('torch', torch.__version__); print('torch_npu', torch_npu.__version__); print('npu_available', torch.npu.is_available())"

输出结果如下：

...torch 2.11.0...
torch_npu 2.11.0
npu_available True

4. 让 OpenRLHF 走 NPU

OpenRLHF 通过 torch.cuda 选设备。下面把 transfer_to_npu 写入 sitecustomize.py，新的 Python 进程会把 torch.cuda 映射到 NPU。

工作目录为 /root/openrlhf-qs：

mkdir -p /root/openrlhf-qs

保存为 /root/openrlhf-qs/enable_npu.py：

from pathlib import Path
import site

path = Path(site.getsitepackages()[0]) / "sitecustomize.py"
path.write_text("from torch_npu.contrib import transfer_to_npu\n")
print(path)

python /root/openrlhf-qs/enable_npu.py
python -c "import torch; print(torch.zeros(1, device='cuda').device)"

输出结果如下：

...sitecustomize.py
npu:0

5. 补一份 flash_attn 占位包

OpenRLHF 启动训练时会立刻 import flash_attn，与是否启用 Flash Attention 无关。昇腾上没有 CUDA 版这个包，下面写入一份只满足 import 的占位包。第 8 节用 --ds.attn_implementation eager，不要加 --ds.packing_samples。

保存为 /root/openrlhf-qs/write_flash_attn_stub.py：

from pathlib import Path

root = Path("/root/openrlhf-qs/flash_attn")
(root / "utils").mkdir(parents=True, exist_ok=True)
(root / "__init__.py").write_text("")
(root / "utils" / "__init__.py").write_text("")
(root / "bert_padding.py").write_text(
    "from einops import rearrange\n"
    "\n"
    "def index_first_axis(*args, **kwargs):\n"
    "    raise RuntimeError('packing needs CUDA flash_attn')\n"
    "\n"
    "def unpad_input(*args, **kwargs):\n"
    "    raise RuntimeError('packing needs CUDA flash_attn')\n"
    "\n"
    "def pad_input(*args, **kwargs):\n"
    "    raise RuntimeError('packing needs CUDA flash_attn')\n"
)
(root / "utils" / "distributed.py").write_text(
    "def all_gather(*args, **kwargs):\n"
    "    raise RuntimeError('ring attention needs CUDA flash_attn')\n"
)
print(root)

python /root/openrlhf-qs/write_flash_attn_stub.py

输出结果如下：

/root/openrlhf-qs/flash_attn

6. 克隆 OpenRLHF 并安装依赖

把 <ref> 换成 Releases 里当前正式 tag，克隆后把源码目录加入 PYTHONPATH。

mkdir -p /root/openrlhf-qs
if [ ! -d /root/openrlhf-qs/OpenRLHF/.git ]; then
  GIT_TERMINAL_PROMPT=0 GIT_HTTP_VERSION=HTTP/1.1 git clone --depth 1 --branch "<ref>" \
    https://github.com/OpenRLHF/OpenRLHF.git /root/openrlhf-qs/OpenRLHF
fi
export PYTHONPATH=/root/openrlhf-qs/OpenRLHF:/root/openrlhf-qs:${PYTHONPATH:-}
python -c "import openrlhf; print(openrlhf.__file__)"

输出结果如下：

/root/openrlhf-qs/OpenRLHF/openrlhf/__init__.py

安装 SFT 路径需要的包：

python -m pip install \
  accelerate aiohttp datasets 'deepspeed==0.19.5' einops 'grpcio>=1.74.0' \
  'huggingface_hub>=1.0.0' jsonlines loralib optimum 'optree>=0.15.0' \
  packaging peft pylatexenc tensorboard torchdata torchmetrics tqdm \
  'transformers==5.15.0' transformers_stream_generator
python -c "import deepspeed, transformers; print('deepspeed', deepspeed.__version__); print('transformers', transformers.__version__)"

输出结果如下：

...deepspeed 0.19.5
transformers 5.15.0

7. 准备模型和数据

从 Hugging Face Hub 下载底座模型，并链到工作目录：

mkdir -p /root/openrlhf-qs
rm -f /root/openrlhf-qs/model
ln -s "$(python -c 'from huggingface_hub import snapshot_download; print(snapshot_download("Qwen/Qwen2.5-0.5B-Instruct"))' | grep '^/' | tail -n 1)" /root/openrlhf-qs/model
ls /root/openrlhf-qs/model/config.json

输出结果如下：

/root/openrlhf-qs/model/config.json

训练数据如下，保存为 /root/openrlhf-qs/tiny_sft.jsonl：

{"question": "Translate to English: 你好", "response": "Hello"}
{"question": "Name a color.", "response": "Blue"}
{"question": "Add 2 and 3.", "response": "5"}
{"question": "Say yes or no.", "response": "Yes"}

8. 在 NPU 上训练

下面用 4 条样本做 1 个 epoch 的 LoRA，确认能在 NPU 上跑完。日志里的 npu::npu_format_cast 和 device='npu' 表示这一次跑在 NPU 上。

export PYTHONPATH=/root/openrlhf-qs/OpenRLHF:/root/openrlhf-qs:${PYTHONPATH:-}
deepspeed --module openrlhf.cli.train_sft \
  --data.max_len 256 \
  --data.dataset /root/openrlhf-qs/tiny_sft.jsonl \
  --data.input_key question \
  --data.output_key response \
  --train.batch_size 1 \
  --train.micro_batch_size 1 \
  --data.max_samples 4 \
  --model.model_name_or_path /root/openrlhf-qs/model \
  --ckpt.output_dir /root/openrlhf-qs/ckpt \
  --ckpt.save_steps -1 \
  --logger.logging_steps 1 \
  --eval.steps -1 \
  --ds.zero_stage 2 \
  --ds.adam_offload \
  --train.max_epochs 1 \
  --ds.param_dtype bf16 \
  --ds.attn_implementation eager \
  --adam.lr 5e-6 \
  --ds.lora.rank 8 \
  --model.gradient_checkpointing_enable \
  --data.dataloader_num_workers 0 \
  2>&1

输出结果如下：

...Setting ASCEND_RT_VISIBLE_DEVICES=0...
...npu::npu_format_cast...
...device='npu'...
...Train step of epoch 0: 100%...
...exits successfully.