DeepPlant-GEP：植物基因表达预测模型昇腾 NPU 服务化部署

简介

DeepPlant-GEP（Addaoud/DeepPlant-GEP）是一个面向植物调控基因组学研究的基因表达预测模型。它读取一段 DNA 序列，输出 2935 个基因表达相关信号（track）的预测值。其网络结构参照 Enformer 的设计思路：先用一维卷积塔在 DNA 序列上提取局部调控特征，再送入 Transformer encoder 建模长程依赖，最后通过注意力池化聚合为一个全局表征，经多层感知机回归到 2935 维的表达预测。

本仓库将该模型适配到昇腾 910B NPU，通过 vLLM-Ascend 提供 OpenAI 兼容的服务化推理（/pooling 接口，回归任务），客户端提交 DNA 碱基 id 序列即可拿到完整的 2935 维预测向量。

模型关键规格：

• 参数量约 248.9M，权重为 float32，单文件约 995MB
• 卷积塔：Conv1d(4→512) 骨干 + 两级残差卷积（512→2048）
• Transformer：6 层 pre-norm encoder，embed_dim=2048，8 头，FFN 2048
• 输出头：注意力池化 + 2048→4096→4096→2935 回归
• 输入：400 bp DNA 序列（碱基 A/C/G/T/N 编码为 0/1/2/3/4）

验证环境

本适配实际运行并验证的环境如下：

权重从 GitCode HF 镜像拉取（GitCode 拉取成功后不再回退到其他源）。该模型体积约 1GB，单卡 910B 可完整承载，服务进程 HBM 占用约 4.4GB。

服务启动

环境准备

pip install vllm==0.18.0 vllm-ascend==0.18.0 torch torch-npu transformers safetensors
source /usr/local/Ascend/ascend-toolkit/set_env.sh

插件注册与启动

DeepPlant-GEP 的自定义架构不在 vLLM 内置 registry 中，且 /vllm-workspace 为只读挂载。本适配把模型实现为外部插件（deepplant_plugin/），启动时在进程内完成两件事：把 DeepPlantGEPConfig 注册进 vLLM 的配置 registry（解析自定义 model_type=deepplant_gep），再把 DeepPlantGEP 注册进模型 registry。使用仓库内的启动脚本：

cd /tmp/deepplant-work/adapt
export ASCEND_RT_VISIBLE_DEVICES=1   # 指定 NPU 逻辑设备
export VLLM_ENABLE_V1_MULTIPROCESSING=0
export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1
nohup python3 run_deepplant_serve.py /tmp/deepplant-work/serve_model \
  --port 8055 \
  --dtype float32 \
  --max-model-len 400 \
  --enforce-eager \
  --max-num-seqs 1 \
  --gpu-memory-utilization 0.2 \
  --served-model-name DeepPlant-GEP \
  --trust-remote-code \
  > serve.log 2>&1 &

参数说明：

• --enforce-eager：关闭图捕获与 torch.compile，走 eager 执行，规避本环境下 GE/AOE 编译初始化失败的问题；
• --max-num-seqs 1：保证每个引擎步骤只有一条完整 DNA 序列（模型内部按固定 100 位置自适应池化，不依赖 vLLM 的序列打包）；
• VLLM_ALLOW_LONG_MAX_MODEL_LEN=1：模型 config 的 max_seq_length=100，但输入需要 400 bp；因模型用自适应池化忽略 vLLM position 语义，安全放行；
• HF_HOME/TRANSFORMERS_CACHE 建议指向本地可写目录（共享 NFS 配额有限）。

服务就绪检查

curl -s http://127.0.0.1:8055/v1/models

返回 DeepPlant-GEP 即表示服务已就绪。

Smoke 验证

服务启动后，通过 /pooling 接口做一次快速冒烟验证。以一段 400 bp 随机 DNA 为例（也可直接跑 inference.py）：

python3 inference.py --mode vllm --url http://127.0.0.1:8055

预期输出：

输入 DNA 长度: 400
[vllm] 服务推理耗时: 18 ms
[vllm] 输出维度: 2935
[vllm] 输出统计: min=-1.36 max=2.10 mean=0.41

也可手工 POST：

curl -s http://127.0.0.1:8055/pooling \
  -H "Content-Type: application/json" \
  -d '{"model":"DeepPlant-GEP","input":[0,1,2,3,0,1,2,3,4,0,1,2,3,4,0,1,2,3,4,0],"task":"classify","encoding_format":"float"}'

性能参考

在昇腾 910B 上实测的推理性能（400 bp 输入、batch=1）：

NPU 占用实测（推理突发期间 npu-smi info）：

| 6     Ascend910           | OK            | 201.7       44                0    / 0             |
| 1     13                  | 0000:0B:00.0  | 63          0    / 0          4474 / 65536         |

AICore 峰值约 63%，HBM 占用 4474MB / 65536MB。

精度测评

以同一段 400 bp DNA 序列分别通过 vLLM-Ascend 服务与直接 torch_npu 推理得到预测向量，对比结果：

服务化推理与本地 NPU 推理在 float32 下逐元素完全一致，说明 vLLM 插件路径没有引入任何精度损失。

典型输入输出样例（随机 400bp，seed=2026）：

输入 DNA: AGACTTTCAAAGATATGCTGGGTAGAGGTCGAGGTTATTATTTGTTACCAATTCTCATTG...
输出维度: 2935
输出范围: [-1.3618, 2.1036]
输出均值: 0.4079
Top-5 高表达 track: [2019, 831, 652, 561, 560]

不同 DNA 序列（随机 / GC-rich / AT-rich）得到的预测分布存在明显区分，模型具备序列敏感的表达预测能力。

注意事项

1. 自定义架构：DeepPlant-GEP 使用 PyTorchModelHubMixin 而非标准 transformers 架构，vLLM 无法识别。必须以插件方式注册 deepplant_gep 配置与 DeepPlantGEP 架构，否则启动即报 Unknown model type。
2. 注意力池化参数：checkpoint 中 attention_pool.to_attn_logits 是 [2048, 2048] 裸张量（无 bias），实现时必须用 nn.Parameter 而非 nn.Linear，否则权重加载缺一个 bias 键。
3. 位置编码：模型自带 [1, 100, 2048] 学习位置编码，卷积塔输出经自适应池化固定到 100 个位置后直接相加，与 vLLM 传入的 positions 无关，因此 VLLM_ALLOW_LONG_MAX_MODEL_LEN=1 是安全的。
4. NPU 设备选择：共享环境下请用 ASCEND_RT_VISIBLE_DEVICES 指定空闲 NPU，避免与既有推理任务争抢显存，也不要杀掉其他窗口的进程。
5. 磁盘配额：共享 NFS /opt/atomgit 常被其他任务占满配额，临时文件、HF 缓存请放到 /tmp。
6. ACLGraph：本环境 CANN 的 GE/AOE 初始化受限，启用 eager 模式可稳定服务化；如生产环境可用，可移除 --enforce-eager 以启用 ACLGraph 图捕获加速。

目录结构

DeepPlant-GEP/
|-- inference.py          # 推理脚本（vllm 服务调用 / direct 本地 NPU）
|-- README.md             # 本文档
|-- requirements.txt      # 依赖清单
|-- assets/
|   `-- README.md         # 观测资料说明