作者：昇腾实战派
知识地图：【昇腾实战派】综合指导

本文记录 NVIDIA 药物发现通用分子生成模型 GenMol 迁移适配到华为昇腾（Ascend）NPU 平台的完整过程，涵盖模型原理、整体架构、环境搭建、推理代码解析、运行结果与常见问题，力求给出一份可复现的实操指南。

一、GenMol 介绍

GenMol 是 NVIDIA 与 KAIST 联合推出的药物发现"通才"生成模型，论文《GenMol: A Drug Discovery Generalist with Discrete Diffusion》发表于 ICML 2025，官方代码仓：https://github.com/NVIDIA-BioNeMo/genmol 。

传统分子生成模型多为"一个任务一个模型"：linker design、R-group exploration、lead optimization、hit generation 各自依赖独立模型甚至 RL 框架。而真实药物研发 pipeline 充满"片段替换、方向探索、结构跳跃、属性优化"交织的需求。GenMol 的目标是用一个统一的生成框架覆盖药物发现全链路任务，其三大基石：

1. SAFE 分子表征（片段序列，而非原子序列）：以 BRICS 片段为词汇，用片段序列表示整个分子，且片段顺序不影响分子结构（天然非自回归）——这恰好对应药化师"以片段为语义单元"的思维方式；
2. 掩码离散扩散（Masked Discrete Diffusion, MDLM）：主体是双向 BERT Transformer，训练时在扩散时间 t 上随机 mask SAFE 序列并学习恢复。相比 GPT 的从左到右逐 token 生成，扩散模型全局并行、顺序无关、错误可被后续步骤修正——与 SAFE"片段顺序只是编码伪像"的特性天然契合；
3. 片段重掩码（Fragment Remasking）与置信度采样：以完整片段（而非单原子）为探索单元做结构搜索；采样时优先填充模型最有把握的位置，并行解码多个 token。这是其在 PMO benchmark 上超越 REINVENT 等 RL 框架的核心原因。

能力覆盖（单卡推理即可完成）：

• De Novo 生成：无约束自由生成全新分子
• 片段约束生成：linker design（两步/一步）、motif extension、scaffold decoration、superstructure generation
• 目标导向 hit 生成（PMO benchmark，AUC 18.2 全场第一）
• 目标导向 lead 优化（86.7% 成功率）

2025 年 10 月官方发布 GenMol V2，引入尖括号扩展 SAFE 语法以区分片段内/片段间连接点，在一步 linker design 上有效性从 16.7% 提升到 81.8%。本文实践基于 V1 权重（genmol_v1，随 NGC 发布）。

二、整体架构

2.1 整体数据流

GenMol 推理由四部分组成：SAFE 编码 → 构造掩码输入 → 迭代去噪解码 → SAFE 解码回 SMILES。

SMILES / 片段输入
    │  SAFEConverter / BracketSAFEConverter      ← 片段化（BRICS 风格）+ 连接点标记
    ▼
SAFE 片段序列（token 序列，词表 ~1880）
    │  tokenizer + _insert_mask                  ← 片段间插入 [MASK]（长度先验采样自 len.pk）
    ▼
x = [BOS] · frag · [MASK]×N · [EOS]
    │
    ▼
┌─────────────────────────────────────────────────┐
│  BertForMaskedLM（BERT-base：12 层 / 768 维 /    │
│  12 头 / max_len 256 / vocab ~1880）             │
│  每步：logits = model(x, attention_mask)         │
│  （可选 MCG 引导：logits = w·logits              │
│       + (1-w)·logits_poor，见 2.4）              │
└─────────────────────────────────────────────────┘
    │  step_confidence：softmax 采样 + Gumbel 扰动
    │  + 置信度 top-k 选位并行填充
    ▼
完整 SAFE 序列 ──► safe_to_smiles 解码 ──► SMILES 分子

对应源码：src/genmol/sampler.py（Sampler 类）与 src/genmol/model.py（GenMol LightningModule）。

2.2 骨干网络：一个标准的 BERT-base

GenMol 的骨干直接使用 HuggingFace BertForMaskedLM（configs/base.yaml）：

配置	值
num_hidden_layers	12
hidden_size	768
num_attention_heads	12
intermediate_size	3072
max_position_embeddings	256
vocab_size	~1880（含 < > 两个 V2 扩展 token）

没有位置编码改造、没有额外解码器——生成能力完全来自离散扩散框架与训练策略，这也是它易于跨硬件迁移的原因（纯标准算子）。

2.3 离散扩散：MDLM 的训练与采样

模型训练采用 bionemo.moco 的 MDLM（Masked Diffusion Language Model）插值器：

• 前向过程：按 LogLinearExp 噪声调度，在时间 t∈[0,1] 对 token 以概率 1-t^β 做 [MASK] 替换，t=1 时全掩码（DiscreteMaskedPrior）；
• 反向过程：模型在 x_t 上预测 x_0 的分布（subs parameterization），训练损失为加权交叉熵；
• 采样：从全掩码序列出发，迭代 num_steps 步，每步"预测 → 选位置 → 填充"。

置信度采样（step_confidence，bionemo/moco/interpolants/continuous_time/discrete/mdlm.py）是并行解码的关键：

probs = torch.softmax(log_p_x0 / logit_temperature, dim=-1)
preds = torch.distributions.Categorical(probs=probs).sample()
confidence = probs.gather(-1, preds.unsqueeze(-1)).squeeze(-1)

# Gumbel 噪声随进程衰减：前期探索，后期收敛
ratio = curr_step / (num_steps - 1)
gumbel_noise = gumbel_sample * randomness * (1 - ratio)
confidence = (torch.log(confidence) + gumbel_noise) / confidence_temperature

mask = xt == self.mask_index            # 只在未填充位置里选
confidence[~mask] = -torch.inf
confidence_threshold, _ = torch.topk(confidence, k=num_tokens_unmask, dim=-1)
to_replace = confidence >= confidence_threshold
xt[to_replace] = preds[to_replace]      # 一步并行填充多个位置

即：每步从 softmax 采样候选 token，用"概率 + 递减 Gumbel 噪声"作为置信度，只解锁置信度最高的若干位置并行填充——模型最有把握的部分（通常是骨架核心）先固定，不确定的边缘留到后续步骤修复。

2.4 分子上下文引导（MCG）

sampler.generate() 中的 gamma/w 参数实现了论文的 Molecular Context Guidance：

# sampler.py generate()
if gamma and w:
    x_poor = x.clone()
    # 将给定片段（context）随机 mask 掉 gamma 比例
    num_mask_poor = int(context_tokens.sum() * gamma)
    mask_idx_poor = random.sample(context_token_ids, num_mask_poor)
    x_poor[:, mask_idx_poor] = self.model.mask_index
    logits_poor = self.model(x_poor, attention_mask=attention_mask)
    logits = w * logits + (1 - w) * logits_poor   # 引导强度 = w

直觉：logits 看得见完整片段上下文，logits_poor 只见残缺上下文，二者差值放大了"片段上下文对生成的影响"，等效于 classifier-free guidance 的分子版本。

2.5 掩码长度先验与片段重掩码

• _insert_mask：插入 [MASK] 的长度不是固定值，而是从训练集长度分布（data/len.pk）随机采样，让生成分子的尺寸符合真实药物分布；
• 片段重掩码（mask_modification / remask）：随机选择一个完整片段（以 . 分隔符定位）整段替换为 5~15 个 [MASK] 让模型重新生成——片段级、大步长的化学空间跳跃，用于 PMO hit 生成与 lead 优化任务。

2.6 片段约束生成的任务路由

frag.py 中四类任务共用 Sampler，靠输入构造与后处理区分：

任务	方法	输入构造	后处理
linker_design	fragment_linking	两侧片段分别生成（V1 两步式）	mix_sequences 重组 + 子结构过滤
motif_extension	fragment_completion	motif + 掩码	子结构过滤
scaffold_decoration	fragment_completion	scaffold + 掩码	子结构过滤
superstructure	fragment_completion	随机选一个 attach point（depth=3）	子结构过滤

三、实验环境

组件	版本
硬件	Ascend 910（64G HBM）x16，x86_64 Docker 容器
操作系统	openEuler 24.03 LTS-SP2
HDK	25.5.1
CANN	8.5.1
Python	3.10（conda）
torch / torch_npu	2.12.0 / 2.12.0
transformers	4.52.4
safe-mol / rdkit	0.1.13 / rdkit-pypi 2022.9.5（必须，见 7.1）
bionemo-moco	0.0.2.1
pytdc	0.4.1（评估指标）

注：官方适配仓标注验证组合为 CANN 8.3.RC1 + torch 2.6.0（aarch64）。本次实践在 x86_64 + CANN 8.5.1 + torch 2.12.0 上验证通过，并对 torch 2.6.0 与 2.12.0 做了指标一致性对比（见 6.3）。

四、环境搭建

4.1 拉取模型代码

git clone https://atomgit.com/AI4Science/Genmol.git
cd Genmol

4.2 创建 conda 环境

conda create -n Genmol python=3.10 -y
conda activate Genmol
# conda 的 libstdc++ 较新时需前插，避免与系统库冲突（报 CXXABI_1.3.15 not found 时）
export LD_LIBRARY_PATH=$CONDA_PREFIX/lib:$LD_LIBRARY_PATH

4.3 安装依赖

# 安装 PyTorch 框架与 torch_npu 插件（版本需与本机 CANN 配套）
pip install torch==2.12.0 torch-npu==2.12.0

# 安装项目基础依赖包
# 注意：requirements.txt 中 torch==2.6.0、torch-npu==2.6.0 两行需先注释，避免覆盖上方已装版本
pip install -r env/requirements.txt
pip install pyyaml decorator attrs psutil scipy setuptools==65.5.1

# rdkit 修复（重要）：依赖安装会同时装入 rdkit（新版）与 rdkit-pypi（旧版），两者写入同一
# site-packages/rdkit/ 目录造成混装损坏；且 rdkit 新版与 torch_npu 共存时 import 段错误。
# 统一改用 rdkit-pypi==2022.9.5：
pip uninstall -y rdkit rdkit-pypi
rm -rf $CONDA_PREFIX/lib/python3.10/site-packages/rdkit*
pip install rdkit-pypi==2022.9.5

pip install -e .

4.4 验证 PyTorch 与 torch_npu

source /usr/local/Ascend/ascend-toolkit/set_env.sh
python3 -c "import torch;import torch_npu; a = torch.randn(3, 4).npu(); print(a + a);"

输出 device='npu:0' 张量即成功。报错时排查顺序：set_env.sh 是否已 source → pyyaml/decorator 等运行时依赖是否安装 → CANN 与 torch_npu 版本是否匹配。

4.5 修复 safe 包（transformers 4.52 兼容性）

chmod +x fix_safe.sh
./fix_safe.sh

该脚本将 safe 包 __init__.py 精简为仅导出 SAFEConverter / decode / encode，绕开 ImportError: cannot import name '_CONFIG_FOR_DOC' from 'transformers.models.gpt2.modeling_gpt2'（safe-mol 的 trainer 模块与新版 transformers 不兼容，推理只需 converter）。

4.6 系统库与权重

# ImportError: libXrender.so.1 时安装（rdkit 绘图依赖）
dnf install -y libXrender libXext libSM

# 权重：需 NGC 账号，或使用直链（匿名可下，~1.33GB）
ngc registry resource download-version "nvidia/clara/genmol_v1:1.0"
mv ./genmol_v1_v1.0/model.ckpt ./

tokenizer（datamol-io/safe-gpt）运行时自动从 HuggingFace 拉取，网络受限时配置 export HF_ENDPOINT=https://hf-mirror.com。

五、推理代码解析

5.1 推理入口（scripts/exps/denovo.py）

from genmol.sampler import Sampler
from tdc import Oracle, Evaluator

sampler = Sampler('model.ckpt')
# 1000 个分子，softmax_temp=0.5（低温度=高质量），randomness=0.5（低噪声=保守）
samples = sampler.de_novo_generation(num_samples, softmax_temp=0.5, randomness=0.5)

# 评估：有效性 / 唯一性 / 多样性 / 质量（QED>=0.6 且 SA<=4）
df = pd.DataFrame({'smiles': samples,
                   'qed': oracle_qed(samples),
                   'sa': oracle_sa(samples)})

5.2 Sampler 初始化与 NPU 迁移

# sampler.py 模块头部 —— NPU 迁移的全部秘密就这两行
import torch_npu
from torch_npu.contrib import transfer_to_npu   # cuda -> npu 自动重映射

class Sampler:
    def __init__(self, path):
        self.model = load_model_from_path(path)  # Lightning checkpoint 加载 + EMA 权重切换
        self.slicer = Slicer()
        self.dot_index = self.model.tokenizer('.')['input_ids'][1]   # 片段分隔符
        self.pad_index = self.model.tokenizer.pad_token_id
        self.mdlm = self.model.mdlm
        self.mdlm.to_device(self.model.device)

迁移要点：源码在模块级 import torch_npu + transfer_to_npu，把 torch.cuda.* 全量重映射为 torch.npu.*，因此脚本内 device='npu'、.to(device) 无需任何业务代码改动。这是典型的"侵入式一行迁移"风格，对 LightningModule 加载 checkpoint（内部 map_location 走 device）同样生效。

5.3 生成主循环

@torch.no_grad()
def generate(self, x, softmax_temp=1.2, randomness=2, fix=True, gamma=0, w=2):
    num_steps = max(self.mdlm.get_num_steps_confidence(x), 2)   # 步数=最大掩码数
    attention_mask = x != self.pad_index

    for i in range(num_steps):
        logits = self.model(x, attention_mask)
        # ... MCG 引导（gamma/w，见 2.4）...
        x = self.mdlm.step_confidence(logits, x, i, num_steps,
                                       softmax_temp, randomness)
    # batch_decode -> safe_to_smiles -> 取最大连通片段
    samples = self.model.tokenizer.batch_decode(x, skip_special_tokens=True)
    samples = [safe_to_smiles(s, fix=fix) for s in samples]
    samples = [sorted(s.split('.'), key=len)[-1] for s in samples if s]
    return samples

注意最后的 sorted(s.split('.'), key=len)[-1]：解码结果按 .（SAFE 片段分隔符）拆分后只保留最大片段，等价于天然过滤掉断裂的小分子副产物——这也是 validity 指标高企的机制之一。

六、运行推理与结果展示

6.1 执行命令

source /usr/local/Ascend/ascend-toolkit/set_env.sh
export ASCEND_RT_VISIBLE_DEVICES=0
export HF_ENDPOINT=https://hf-mirror.com
export CPU_AFFINITY_CONF=1      # 流水优化（可选）
export TASK_QUEUE_ENABLE=2

python scripts/exps/denovo.py   # De Novo 生成，~36s
python scripts/exps/frag.py     # 片段约束生成（4 任务），~151s

首次运行会联网下载 tdc 评估数据（Harvard Dataverse，网络较慢，见 7.2）。

6.2 De Novo 生成结果（NPU，1000 分子）

Time:           35.80 sec
Validity:       0.997
Uniqueness:     0.999
Diversity:      0.816
Quality:        0.857

前 10 个高质量分子示例（QED ≥ 0.6 且 SA ≤ 4）：

                                         smiles      qed       sa
             COc1c(C(=O)NCC2CCCCC2)ccc(F)c1F    0.921     2.069
CCc1ccc(CNC(=O)C(=O)NCc2ccc(C[NH+]3CCCC3)cc2)cc1  0.635   3.039
            CC[C@H](C)CNC(=O)N1CCCN(c2ccccc2)CC1  0.925     2.345

6.3 片段约束生成结果（NPU，4 任务 × 100 样本）

任务	Validity	Uniqueness	Quality	Diversity	Distance
linker_design	1.000	0.816	0.205	0.540	0.558
motif_extension	0.824	0.781	0.282	0.625	0.683
scaffold_decoration	0.963	0.814	0.298	0.588	0.651
superstructure_generation	0.989	0.858	0.375	0.589	0.749

与官方 A100 结果对比（官方 README，GenMol V1）：

指标	NPU 910（本次）	官方 A100	备注
De Novo Validity	0.997	1.000	随机采样波动
De Novo Uniqueness	0.999	0.997
De Novo Quality	0.857	0.846
De Novo Diversity	0.816	0.818
Linker Validity / Quality	1.000 / 0.205	1.000 / 0.219
Motif ext. Validity / Quality	0.824 / 0.282	0.829 / 0.301
Scaffold deco. Validity / Quality	0.963 / 0.298	0.966 / 0.318
Superstructure Validity / Quality	0.989 / 0.375	0.975 / 0.348

结论：NPU 结果与官方 A100 基线在随机波动范围内一致（脚本未固定随机种子，且 superstructure 任务的 attach point 每次随机选择），数值精度无漂移。

另外验证了 torch 2.6.0 → 2.12.0 升级前后的指标一致性（denovo 32.4s→35.8s，frag 149.5s→150.6s，各指标差异均在采样波动内），说明 GenMol 纯 BERT 算子面在昇腾后向兼容性良好。

6.4 性能数据（单卡 910）

阶段	耗时
De Novo 生成 1000 分子	35.80 s
片段约束生成（4 任务 × 100 样本 × 多片段）	150.57 s

吞吐参考：单卡约 28 molecules/s（de novo，含 512 token 级序列的 ~40 步迭代解码），满足虚拟筛选前构建 idea pool 的实用需求。

七、常见问题与告警说明

7.1 rdkit 与 torch_npu 共存段错误（本文最大的坑）

现象：import torch_npu 后再 import rdkit 的 Chem 模块（或 safe、datamol 等依赖链），进程直接 Segmentation fault（exit 139），无任何 traceback。纯 torch（TORCH_DEVICE_BACKEND_AUTOLOAD=0）+ rdkit 无问题。

根因：两层问题叠加。其一，pip 依赖解析同时装入了 rdkit 2026.x（datamol 依赖）与 rdkit-pypi 2022.9.5（safe-mol 依赖链），两个发行版往同一个 site-packages/rdkit/ 目录写文件，混装损坏；其二，清理后单装 rdkit 2026.x 仍段错误——其新版工具链编译的 C++ 扩展与 CANN 加载的库存在 ABI 冲突，初始化 libRDKitFileParsers 时崩溃。

解法：统一使用旧版发行版（贴近官方验证环境）：

pip uninstall -y rdkit rdkit-pypi
rm -rf $CONDA_PREFIX/lib/python3.10/site-packages/rdkit*
pip install rdkit-pypi==2022.9.5

导入顺序也有讲究：safe（rdkit）先于 torch_npu 导入时部分路径可绕开，但用 rdkit-pypi 2022.9.5 后无顺序要求。

7.2 tdc 评估数据下载慢 / pickle data was truncated

现象：denovo.py 生成完成后在 oracle_sa(samples) 处崩溃，_pickle.UnpicklingError: pickle data was truncated。

根因：pytdc 的 SA Score 需要下载 fpscores.pkl（~9MB，Harvard Dataverse 源，实测仅 ~20 kiB/s），下载中断会留下截断的缓存文件 oracle/fpscores.pkl。

解法：删除损坏文件重试，或手动下载完整文件放到 oracle/ 目录。官方 GitHub 也给出了替代源：https://github.com/rdkit/rdkit/tree/master/Contrib/SA_Score/fpscores.pkl.gz

7.3 libXrender 与 _CONFIG_FOR_DOC（README 已预判的两个坑）

• ImportError: libXrender.so.1：rdkit 绘图链的 X11 依赖缺失，dnf install -y libXrender libXext libSM；
• ImportError: cannot import name '_CONFIG_FOR_DOC' from 'transformers.models.gpt2.modeling_gpt2'：safe-mol 内部 trainer 模块与 transformers 4.52 不兼容，执行仓内 fix_safe.sh 精简 safe 包导入。

7.4 aarch64 与 x86_64 的 wheel 差异

官方适配仓 README 的 torch 安装命令是 aarch64 wheel（manylinux_2_28_aarch64），x86_64 环境需改用 pip install torch==2.12.0 torch-npu==2.12.0（华为云 PyPI 源有全架构包），torch_npu 版本需与 CANN 配套。

7.5 运行期告警（非报错，不影响结果）

• In npu autocast, but the target dtype is not supported：NPU autocast 仅支持 fp16/bf16，此处被禁用。GenMol 以 fp32 推理，无影响；
• transfer_to_npu 的 fp64→fp32 替换：NPU 不支持 double，DoubleTensor 被替换为 FloatTensor。GenMol 全程 fp32 张量，无影响；与 GPU fp64 实现做逐位对比时需注意；
• jit script 禁用：GenMol 未使用 jit，无影响。

八、实践小结与使用建议

1. 迁移成本低：GenMol 骨干是标准 BertForMaskedLM，无自定义算子，迁移面收敛在"导入 torch_npu + transfer_to_npu + 设备字符串"，主要工作量在环境依赖治理（rdkit 版本冲突是最大障碍）；
2. 指标可信：NPU 与官方 A100 的 8 项核心指标全部对齐，且跨 torch 版本（2.6.0/2.12.0）一致，说明数值路径无平台差异；
3. 使用提示（来自社区实践经验）：输入片段建议先用 RDKit MolToSmiles 统一为芳香式 SMILES（凯库勒式可能不被识别导致零输出）；scaffold 拓展位点过多时容易生成大环，需控制 attach point 数量；
4. 延伸方向：PMO hit 生成与 lead 优化任务（scripts/exps/pmo、scripts/exps/lead）依赖 fragment remasking 迭代搜索，单任务耗时 2~4 小时（A100），可作为 NPU 长稳性与吞吐的进一步验证场景。

参考资料

• GenMol 官方仓库：https://github.com/NVIDIA-BioNeMo/genmol
• 论文：Lee et al., GenMol: A Drug Discovery Generalist with Discrete Diffusion, ICML 2025（arXiv:2501.06158）
• 昇腾适配仓：https://atomgit.com/AI4Science/Genmol
• 权重：NGC nvidia/clara/genmol_v1（https://catalog.ngc.nvidia.com/orgs/nvidia/teams/clara/resources/genmol_v1）
• SAFE 表征：Atkinson et al., Breaking BAD: Fragmenting and Synthesizing Molecules with Fragments and Attributes, 2023
• MDLM：Sahoo et al., Simple and Effective Masked Diffusion Language Models, 2024