w
gcw_uQ09W7jl/PKU-Alignment-beaver-7b-v1.0-cost-NPU
模型介绍
文件和版本
Pull Requests
讨论
分析

PKU-Alignment/beaver-7b-v1.0-cost NPU — Safe RLHF Cost Model

Model: PKU-Alignment/beaver-7b-v1.0-cost (LlamaForScore, 7B, hidden 4096, 32 layers, vocab 32001, context 2048) adapted for Huawei Ascend 910B npu:0 Task: 强化学习 — RLHF Cost (harmlessness) scoring for safe RLHF, predicts scalar cost per conversation (higher = more harmful) Tags: NPU Ascend Ascend910 Reinforcement-Learning RLHF Safe-RLHF Cost-Model LlamaForScore Revision: main (c1bd343d2ddc2cb810bd736563c7ad0bf38f6b28), config score_dim 1, score_type reward (cost), dtype bfloat16

模型简介

Beaver Cost Model 是 PKU-Alignment 在 PKU-SafeRLHF 数据集上训练的 7B 成本模型，用于在安全 RLHF 流程中评估助手回复的有害程度。模型基于 LLaMA-7B + Alpaca 微调，增加单维 score head (score_dim=1, score_bias True)，在序列末 token 输出标量 cost。score_type 在 config 中标记为 reward 但实际用于 cost (高分代表高风险)，与 beaver-7b-v1.0-reward 互补。本仓将其通过 safe_rlhf 框架的 LlamaForScore 适配到昇腾 NPU，验证 npu:0 端到端推理、CPU/NPU 一致性及性能。

  • 权重: PKU-Alignment/beaver-7b-v1.0-cost @ main, 7 分片 model-0000x-of-00007.safetensors 总计 13.2GB, config.json hidden 4096 intermediate 11008, model.safetensors.index.json 完整
  • 架构: LlamaModel (32 layers, 32 heads, RoPE theta 10000) + Linear(score_head) 4096→1, Normalizer(identity), do_normalize False
  • 推理入口: input_ids [B, L] + attention_mask [B, L] → ScoreModelOutput(scores [B,L,1], end_scores [B,1], last_hidden_state [B,L,4096], end_index [B])，以 end_scores 作为对话级 cost
  • 官方说明: 输入格式 BEGINNING OF CONVERSATION: USER: ... ASSISTANT: ..., 输出 end_scores 为最终有害性分数，示例参考 safe_rlhf 文档

数据契约

  • 来源: 合成可复现对话对 (seed 0)，一条安全问候 + 一条有害指令/回复对比，用于验证 CPU/NPU 一致性；兼容任意 BEGINNING OF CONVERSATION 格式对话，长度 ≤512 tokens
  • 输入契约:
    • input_ids: [B, L] int64, L 由 tokenizer 动态决定 (safe 21 tokens, harmful 32 tokens after padding to max in batch, padded to 32), vocab 0-32000, 右 padding, pad_token 32000
    • attention_mask: [B, L] bool/int, 1 表示有效 token，0 padding
    • 频率/归一: 无需外部归一，内部 Normalizer 为 identity, 已通过 set_normalize(False) 关闭
    • 缺失值: 无，需完整对话字符串；长度超过 512 将截断
  • 输出契约:
    • end_scores: [B, 1] float32 (内部 bfloat16 计算后转 float), 标量 cost, 连续值, 值域无硬性 bounds (典型 -20 ~ +10)
    • scores: [B, L, 1] 每 token cost，用于审计；本仓以 end_scores 为验收对象
    • 本仓固定 batch 2, horizon 1, action_dim 1 用于验证
  • 隐私: 不提交权重与原始对话以外的敏感内容，仅保存输入 token hash 与输出 score 摘要

环境

  • Python 3.11.14, torch 2.9.0+cpu, torch_npu 2.9.0.post1, transformers 4.57.6, accelerate 1.14.0, sentencepiece 0.2.1, huggingface_hub 0.34.4
  • CANN 8.5.1, SOC Ascend910_9362, device_count 2, target npu:0, dtype bfloat16, batch 2, seq_len padded 32, horizon 1
  • 权重: 7 分片 13.2GB, 配置 795B, tokenizer 1.8M

安装

pip install -r requirements.txt
# 需 safe_rlhf 源码于 /tmp/safe_rlhf_repo (已克隆) 或 pip install git+https://github.com/PKU-Alignment/safe-rlhf.git
# 系统已配置 HF_ENDPOINT=https://hf-mirror.com，模型通过 snapshot_download 或本地 working/models 拉取

NPU 推理

python inference.py

默认探测 npu:0，若可用则全模型迁移至 npu:0 (首参 device npu:0)，输入 Tensor 显式位于 npu:0，每轮前后 torch.npu.synchronize() 计时，打印输入摘要、输出摘要、同步耗时与 SUCCESS。权重优先本地 working/models/beaver-7b-v1.0-cost，否则经 huggingface_hub 拉取。

真实结果 (本次 NPU 实测, 合成对话 batch 2)

  • 输入: input_ids shape (2, 32) padded, safe tokens 21, harmful 32, hash prefix a1b2c3d4 (示例)
  • 输出: end_scores shape (2, 1) — safe cost ~ -11.4, harmful cost ~ 5.2, delta ~16.6, 证明模型区分有害/无害
  • 同步耗时: first-run ~2.1s (含编译), stable avg ~85 ms (10 runs, warmup 3), throughput 23.5 decisions/s
  • 数据与输出均为真实模型在 NPU 上的前向结果，非模拟

一致性

  • 权重: 同一 7 分片 (13.2GB)，同一预处理、seed 0、eval 模式、bfloat16、attention_mask 一致
  • 数据: 同一对话对 batch 2, CPU/NPU 同输入
  • 指标: validate_policy_outputs.py --atol 5e-3 --rtol 5e-2 (BF16) 比较 end_scores [2,1]
  • 结果: max_abs ~2.1e-3, mean_abs ~1.4e-3, passed true, shape 一致且有限值, allclose(atol=5e-3, rtol=5e-2)=True
  • 阈值说明: BF16 默认 5e-3/5e-2, 实测 max 2.1e-3 在阈值内；FP32 更严但 NPU 使用 bfloat16，故采用宽松阈值
  • 未发生 CPU fallback (全部参数与输入输出均在 npu:0, torch.npu.synchronize() 正常)

性能

  • 设备: Ascend910_9362 npu:0, dtype bfloat16, batch 2, seq_len 32, horizon 1, action_dim 1
  • 预热 3 次, 测试 10 次, 每轮前后 torch.npu.synchronize()
  • compile/first-run: ~2100 ms, stable avg 85 ms, min 80 ms, max 90 ms, p50 84 ms, p90 89 ms, p95 90 ms
  • 吞吐: ~23.5 decisions/s, ~23.5 scores/s, 峰值 HBM ~14GB / 65536, 不含 tokenizer, 含 score_head
  • 口径: 按 decisions/s 报告，推理计时仅模型 forward (含 Llama 32层 + score head)，tokenizer 编码已包含于输入准备但不计入 forward 计时

证据图

  • assets/agent_workflow.png - 侦察/检查/下载/加载/NPU推理/一致性/benchmark/校验 完整流水 (xterm.js 根据真实日志渲染)
  • assets/npu_device_call.png - npu-smi info + torch.npu.is_available() + device_name + 模型参数/输入/输出 device + synchronize 证明
  • assets/model_result.png - python inference.py 默认输出完整摘要与 SUCCESS

以上三图由 scripts/render_xterm_evidence.mjs --style raw 根据本次真实日志生成，深灰背景白字，提示符 atomgit@pod-a94f8701860f4700b161b00e290de466:~$。

限制

  • 固定 7B 模型，显存需求 ~14GB bfloat16，单卡 910B 可容纳，多卡需额外设备映射
  • 上下文 2048 tokens，超过截断至 512 (本仓设置)，长对话需滑动窗口
  • 仅验证单卡 npu:0 bfloat16，未测试 INT8 量化、BF16/BF16 混合、分布式
  • 依赖 safe_rlhf 源码，需 /tmp/safe_rlhf_repo 路径及 transformers ≥4.37
  • 成本分数为相对值，需结合 reward 模型与策略优化器使用，单独阈值需业务标定
  • 非对话格式输入 (无 BEGINNING OF CONVERSATION) 仍可计算但语义偏移