MiniOneRec

MiniOneRec 是生成式推荐框架。本文在单张昇腾 NPU 上依次完成 SFT 训练、约束解码评估与推荐导向 RL 三个阶段。

前置条件

硬件

Atlas 800T / 900 A2 训练系列,芯片为 Ascend 910B,需支持 bf16。本文示例为单卡。

软件

类别

要求

CANN

toolkit 与驱动固件已安装,并可 source set_env.sh,9.1.0 或以上

Python

3.10 或以上

git

用于克隆上游仓库

模型

Qwen/Qwen2.5-0.5B(base 版)

本文验证环境

项目

内容

机器

Atlas 900 A2,Ascend 910B,单卡

镜像

swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:9.1.0-910b-ubuntu22.04-py3.12

CANN

9.1.0

Python

3.12

torch / torch_npu

2.9.0 / 2.9.0.post2

阅读本文前,请先按 快速安装昇腾环境 准备好 CANN 与驱动。

1. 加载 CANN 环境

每次新开终端后执行:

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

2. 检查环境是否就绪

source /usr/local/Ascend/ascend-toolkit/set_env.sh
test -n "$ASCEND_HOME_PATH"
command -v npu-smi
npu-smi info
python --version

确认 npu-smi info 输出设备列表且 Health 为 OK。若提示找不到 npu-smi,按 快速安装昇腾环境 检查驱动与设备挂载。

3. 获取代码

克隆 MiniOneRec 仓库:

git clone --depth 1 --branch main https://github.com/AkaliKong/MiniOneRec.git
cd MiniOneRec

备注

本文使用 MiniOneRec main 分支的最新提交。

4. 安装依赖

安装以下依赖(与 CANN 9.1.0 配套):

source /usr/local/Ascend/ascend-toolkit/set_env.sh
pip install torch==2.9.0 torch-npu==2.9.0.post2
pip install transformers==4.57.1 accelerate==1.10.1 datasets==4.2.0 \
    tokenizers==0.22.1 fire==0.7.1 pandas==2.2.2 tqdm bitsandbytes==0.48.1 \
    trl==0.24.0 scikit-learn==1.5.2 wandb==0.22.2

5. 确认运行环境

执行以下命令,确认 NPU 与训练设备就绪:

source /usr/local/Ascend/ascend-toolkit/set_env.sh
cd MiniOneRec
python -c "
import npu_support
import torch
from transformers import TrainingArguments
print('IS_NPU:', npu_support.IS_NPU)
print('cuda mapped:', torch.cuda.is_available())
print('evaluate.py device:', 'cuda' if torch.cuda.is_available() else 'cpu')
print('trainer device:', TrainingArguments(output_dir='/tmp/check-npu').device)
print('tensor landed on:', torch.zeros(1, device='cuda').device.type)
"

输出结果如下:

IS_NPU: True
cuda mapped: True
evaluate.py device: cuda
trainer device: npu:0
tensor landed on: npu

检查克隆目录状态:

cd MiniOneRec
test -z "$(git status --porcelain)" && echo "source untouched"

输出结果如下:

source untouched

6. 下载基座模型与预处理数据

基座模型取 Qwen2.5 的 base 版(Instruct 版会使约束解码失效)。模型与预处理数据从 HuggingFace 下载:

cd MiniOneRec
python -c "from huggingface_hub import snapshot_download; snapshot_download(repo_id='Qwen/Qwen2.5-0.5B', local_dir='./Qwen2.5-0.5B')"
python -c "from huggingface_hub import snapshot_download; snapshot_download(repo_id='kkknight/MiniOneRec', allow_patterns=['Amazon/index/Industrial_and_Scientific.index.json','Amazon/index/Industrial_and_Scientific.item.json','Amazon/train/Industrial_and_Scientific_5_2016-10-2018-11.csv','Amazon/valid/Industrial_and_Scientific_5_2016-10-2018-11.csv','Amazon/test/Industrial_and_Scientific_5_2016-10-2018-11.csv','Amazon/info/Industrial_and_Scientific_5_2016-10-2018-11.txt'], local_dir='data')"

校验资产落位:

cd MiniOneRec
test -f Qwen2.5-0.5B/config.json && echo "base model ok"
test -f data/Amazon/index/Industrial_and_Scientific.index.json && echo "sid index ok"
test -f data/Amazon/index/Industrial_and_Scientific.item.json && echo "item meta ok"
test -f data/Amazon/train/Industrial_and_Scientific_5_2016-10-2018-11.csv && echo "train csv ok"
test -f data/Amazon/valid/Industrial_and_Scientific_5_2016-10-2018-11.csv && echo "valid csv ok"
test -f data/Amazon/info/Industrial_and_Scientific_5_2016-10-2018-11.txt && echo "info txt ok"

输出结果如下:

base model ok
sid index ok
item meta ok
train csv ok
valid csv ok
info txt ok

7. SFT 训练

在 0 号卡上执行 1 个 epoch:

source /usr/local/Ascend/ascend-toolkit/set_env.sh
export ASCEND_RT_VISIBLE_DEVICES=0
export WANDB_MODE=disabled
cd MiniOneRec
python sft.py \
    --base_model ./Qwen2.5-0.5B \
    --train_file ./data/Amazon/train/Industrial_and_Scientific_5_2016-10-2018-11.csv \
    --eval_file ./data/Amazon/valid/Industrial_and_Scientific_5_2016-10-2018-11.csv \
    --output_dir ./sft_out \
    --batch_size 8 \
    --micro_batch_size 2 \
    --num_epochs 1 \
    --cutoff_len 256 \
    --sample 600 \
    --category Industrial_and_Scientific \
    --sid_index_path ./data/Amazon/index/Industrial_and_Scientific.index.json \
    --item_meta_path ./data/Amazon/index/Industrial_and_Scientific.item.json \
    --freeze_LLM False \
    --train_from_scratch False

训练完成后校验产物:

cd MiniOneRec
test -f sft_out/final_checkpoint/model.safetensors && echo "model saved"
test -f sft_out/final_checkpoint/config.json && echo "config saved"
test -f sft_out/final_checkpoint/tokenizer_config.json && echo "tokenizer saved"
grep -o '"model_type": "[a-z0-9_.-]*"' sft_out/final_checkpoint/config.json

输出结果如下:

model saved
config saved
tokenizer saved
"model_type": "qwen2"

8. 约束解码评估

用上一节的 final_checkpoint 在 test 子集上做约束解码评估并计算 HR/NDCG/CC。CC 必须为 0。

切出 test 子集:

cd MiniOneRec
head -n 200 data/Amazon/test/Industrial_and_Scientific_5_2016-10-2018-11.csv > test_small.csv
wc -l test_small.csv

运行评估:

source /usr/local/Ascend/ascend-toolkit/set_env.sh
export ASCEND_RT_VISIBLE_DEVICES=0
cd MiniOneRec
python evaluate.py \
    --base_model ./sft_out/final_checkpoint \
    --info_file ./data/Amazon/info/Industrial_and_Scientific_5_2016-10-2018-11.txt \
    --test_data_path test_small.csv \
    --result_json_data ./eval_result.json \
    --category Industrial_and_Scientific \
    --num_beams 20 --batch_size 2

输出结果如下:

industrial and scientific items
...

校验预测产物:

cd MiniOneRec
python -c "import json; d=json.load(open('eval_result.json')); print(len(d)); print('predict' in d[0])"

输出结果如下:

199
True

计算 HR/NDCG/CC:

cd MiniOneRec
python calc.py \
    --path ./eval_result.json \
    --item_path ./data/Amazon/info/Industrial_and_Scientific_5_2016-10-2018-11.txt

输出结果如下:

20
[1, 3, 5, 10, 20]
NDCG:...
HR...
0

9. 推荐导向 RL

在 0 号卡上执行 1 个 epoch。--sample 40 用于限制数据规模,请勿省略;train_batch_size 必须能被 num_generations 整除。

source /usr/local/Ascend/ascend-toolkit/set_env.sh
export ASCEND_RT_VISIBLE_DEVICES=0
cd MiniOneRec
python rl.py \
    --model_path ./sft_out/final_checkpoint \
    --train_file ./data/Amazon/train/Industrial_and_Scientific_5_2016-10-2018-11.csv \
    --eval_file ./data/Amazon/valid/Industrial_and_Scientific_5_2016-10-2018-11.csv \
    --info_file ./data/Amazon/info/Industrial_and_Scientific_5_2016-10-2018-11.txt \
    --category Industrial_and_Scientific \
    --train_batch_size 4 \
    --eval_batch_size 4 \
    --num_generations 2 \
    --num_train_epochs 1 \
    --gradient_accumulation_steps 1 \
    --learning_rate 1e-5 \
    --beta 1e-3 \
    --beam_search True \
    --reward_type ranking \
    --sync_ref_model True \
    --test_during_training False \
    --eval_step 1.0 \
    --sample 40 \
    --output_dir ./rl_out \
    --sid_index_path ./data/Amazon/index/Industrial_and_Scientific.index.json \
    --item_meta_path ./data/Amazon/index/Industrial_and_Scientific.item.json

输出结果如下:

Industrial_and_Scientific
...

日志中的 reward、reward_std、loss、grad_norm、kl 可能为 0,属正常现象。

校验产物:

cd MiniOneRec
test -f rl_out/final_checkpoint/model.safetensors && echo "rl model saved"
test -f rl_out/final_checkpoint/config.json && echo "rl config saved"
test -f rl_out/final_checkpoint/tokenizer_config.json && echo "rl tokenizer saved"

输出结果如下:

rl model saved
rl config saved
rl tokenizer saved

相关链接