作者昇腾实战派知识地图https://blog.csdn.net/Lumos_Lovegood/article/details/161601003背景概述单细胞 RNA 测序scRNA-seq技术的爆发式发展使人类细胞图谱等综合数据库的规模已膨胀至数千万个细胞。如何从海量且高度异质的单细胞数据中高效提取生物学知识——包括细胞类型识别、批次效应消除、基因扰动预测和调控网络推断——成为计算生物学的核心挑战。当前主流方法大多针对特定任务开发专用模型导致数据孤岛和任务孤岛的局限。scGPTsingle-cell Generative Pre-trained Transformer由多伦多大学 Bo Wang 实验室于 2024 年发表于 Nature Methods是一种专为单细胞多组学设计的基础模型。scGPT 建立了语言与细胞生物学之间的平行联系——正如文本由单词组成细胞由基因定义。模型将基因视为 Token将细胞视为文本在 CELLxGENE 中超过3300 万个人类正常单细胞的 RNA 测序数据上完成生成式预训练通过迁移学习支持多种下游任务在相应基准上达到 SOTA。本文介绍 scGPT 的昇腾 Ascend NPU 适配版本——对微调与推理流程进行 NPU 适配使用昇腾原生 FlashAttention2 加速注意力计算为大规模单细胞数据分析提供高效的硬件加速方案。模型介绍scGPT 概述scGPT 的核心设计理念是将自然语言处理中的生成式预训练范式引入单细胞组学领域。模型在 51 个器官/组织、超过 3300 万个正常人类细胞上完成预训练学习到基因和细胞的通用嵌入表示。支持的下游任务细胞类型注释对未知细胞进行分类标注多批次整合消除技术批次效应保留生物学变异多组学整合融合 RNA-seq 和 ATAC-seq 等多模态数据扰动响应预测预测基因编辑后的细胞状态变化基因网络推断从注意力权重提取基因调控网络整体架构scGPT 基于 Transformer 架构但针对单细胞数据的非序列化特征进行了专门设计。输入编码单细胞数据的输入整合了三个维度的信息编码类型说明维度Gene Token Embedding基因名称的离散编码ntoken × d_modelExpression Value Encoding基因表达值的连续/离散编码d_modelCondition Token批次、模态、扰动条件等元信息d_model表达值编码支持三种模式continuous连续值编码器MLP 映射category分箱离散化后的嵌入编码scaling归一化缩放核心 Transformer堆叠的 Transformer Encoder 层默认 12 层多头注意力默认 8 头支持 Flash Attention 加速前馈网络隐藏维度 d_hid支持 Pre-Norm 和 Post-Norm 两种归一化方案掩蔽注意力机制Masked Attention由于基因表达数据不具有自然语言的严格语序scGPT 设计了专门的注意力掩码策略已知基因可互相关注双向注意力被掩蔽基因只能关注已知基因单向约束这使得模型在生成式预训练中能同时学习基因间的相互依赖关系多任务输出头输出头功能场景ExprDecoder基因表达值重建预训练、批次整合ClsDecoder细胞分类细胞类型注释MVCDecoder掩码值预测自监督学习AdversarialDiscriminator对抗批次判别批次整合细胞嵌入策略cls使用 [CLS] token 的输出作为细胞级表示avg-pool对所有基因 token 的输出取平均w-pool加权池化预训练与微调预训练数据CELLxGENE 3300 万人类正常单细胞51 个器官/组织任务生成式预训练——随机掩蔽部分基因表达值模型预测被掩蔽的值损失函数掩蔽 MSE 损失 弹性细胞相似度ECS正则化微调冻结预训练层仅训练任务特定的输出头逐步解冻全部层进行端到端微调支持 Domain-Specific BatchNormDSBN处理多批次数据缩放定律研究发现 scGPT 的性能与预训练数据规模呈正相关——当预训练数据从 3 万扩展到 3300 万个细胞时下游任务表现持续提升与 NLP 领域的缩放定律高度一致。昇腾 NPU 适配版本迁移动机原始 scGPT 依赖 CUDA Flash Attention 进行高效注意力计算。为在昇腾 Ascend NPU 上运行本项目使用torch_npu原生的npu_fusion_attention替代 CUDA Flash Attention实现等效的加速效果。仓库结构scGPT_npu/ ├── README.md ├── requirements.txt ├── pyproject.toml ├── examples/ │ ├── finetune_integration.py # 批次整合微调脚本NPU 适配 │ ├── inference.py # 推理脚本NPU 适配 │ ├── save/ │ │ └── scGPT_human/ # 预训练权重 │ │ ├── best_model.pt │ │ ├── vocab.json │ │ └── args.json │ └── data/ # 示例数据 ├── scgpt/ │ ├── model/ │ │ ├── model.py # TransformerModel 主模型 │ │ ├── dsbn.py # Domain-Specific BatchNorm │ │ ├── generation_model.py # 生成式模型 │ │ └── multiomic_model.py # 多组学模型 │ ├── utils/ │ │ └── flash_attention.py # ⭐ NPU Flash Attention 实现 │ ├── tasks/ │ │ ├── cell_emb.py # 细胞嵌入提取 │ │ └── grn.py # 基因调控网络推断 │ ├── tokenizer/ │ │ └── gene_tokenizer.py # 基因词表编码 │ ├── preprocess.py # 数据预处理 │ ├── loss.py # 损失函数 │ ├── trainer.py # 训练器 │ └── data_collator.py # 数据整理器 ├── data/ │ ├── pbmc3k.h5ad # 示例数据 │ └── cellxgene/ # 大规模数据构建脚本 └── tests/核心迁移改动1. Flash Attention NPU 实现scgpt/utils/flash_attention.py使用昇腾原生的npu_fusion_attention替代 CUDA flash-attnfromtorch_npuimportnpu_fusion_attention支持特性FlashAttention2下右对齐因果掩码默认或左上对齐因果掩码环境变量控制NPU_FA2_SPARSE_MODE2左上对齐或3下右对齐注意力掩码缓存复用预计算的掩码减少内存分配2. 推理与微调入口适配所有入口脚本添加 NPU 自动迁移importtorch_npufromtorch_npu.contribimporttransfer_to_npu3. NPU 性能优化配置exportCPU_AFFINITY_CONF1# CPU 亲和性绑定exportTASK_QUEUE_ENABLE2# 任务队列优化exportASCEND_RT_VISIBLE_DEVICES0# 指定 NPU 卡版本信息软件版本HDK25.5.0CANN8.3.RC1Python3.11PyTorch2.1.0torch_npu2.1.0环境配置创建 Conda 环境conda create-nscgptpython3.11-yconda activate scgpt克隆代码gitlfsinstallgitclone https://atomgit.com/AI4Science/scGPT.gitcdscGPT安装依赖exportPIP_INDEX_URLhttps://repo.huaweicloud.com/repository/pypi/simplecdscGPTpipinstall-e.pipinstall-rrequirements.txt pipinstalltorchtext0.15.2torchdata0.7.1 --no-deps主要依赖包括torch2.1.0、torch_npu2.1.0scanpy1.10.3单细胞数据分析核心库scvi-tools1.2.1单细胞变分推断工具scib1.1.5单细胞整合基准评估anndata0.10.8单细胞数据结构einops0.8.1张量操作工具验证安装python3-cimport torch;import torch_npu; a torch.randn(3, 4).npu(); print(a a);NPU 运行环境配置exportASCEND_RT_VISIBLE_DEVICES0exportCPU_AFFINITY_CONF1exportTASK_QUEUE_ENABLE2exportWANDB_MODEofflineexportWANDB_OFFLINEtrue模型权重预训练权重已包含在仓库中通过 Git LFSexamples/save/scGPT_human/ ├── best_model.pt # 预训练模型权重 ├── vocab.json # 基因词表约 60,000 个基因 └── args.json # 模型超参数配置模型参数配置来自 args.jsond_model: 512隐藏维度nhead: 8注意力头数nlayers: 12Transformer 层数d_hid: 512前馈网络维度微调批次整合微调cdexamples/ python finetune_integration.py该脚本完成加载 PBMC 10K 数据集含多批次使用 scGPT_human 预训练权重初始化以 40% mask ratio 进行掩蔽表达值预测结合对抗批次判别器DAB消除批次效应使用 scib 评估整合质量AvgBIO 等指标主要超参数mask_ratio: 0.4epochs: 30n_bins: 51表达值分箱数learning_rate: 1e-4batch_size: 64推理细胞嵌入提取cdexamples/ python inference.py推理流程加载 Kim2020 肺部单细胞数据h5ad 格式筛选高变异基因Top 3000 HVG使用预训练模型提取细胞嵌入输出推理耗时统计embed_adatascg.tasks.embed_data(adata,model_dir,gene_colgene_col,batch_size64,)Docker 镜像已预置 Conda 运行环境、模型权重、scGPT 源码开箱即用sudodockerpull swr.cn-north-4.myhuaweicloud.com/ascend_ai4s/scgpt:v1镜像内容项目源码路径/root/scGPT/已配置好 Conda 虚拟环境已内置模型权重文件无需额外下载迁移适配要点Flash Attention 替换原实现NPU 实现CUDAflash_attn库torch_npu.npu_fusion_attentionflash_attn_func(q, k, v)npu_fusion_attention(q, k, v, head_num, ...)CUDA causal maskNPU sparse_mode 控制2左上3下右自动设备迁移importtorch_npufromtorch_npu.contribimporttransfer_to_nputransfer_to_npu自动将所有 CUDA 调用重定向到 NPU无需手动修改设备代码。性能优化CPU_AFFINITY_CONF1启用 CPU 亲和性绑定减少跨 NUMA 访问TASK_QUEUE_ENABLE2启用任务队列优化提升 NPU 利用率Flash Attention 注意力掩码缓存避免重复创建大尺寸掩码张量已知限制当前版本聚焦于微调与推理不包含完整的 3300 万细胞预训练流程Flash Attention 的npu_fusion_attention要求输入为 float16/bfloat16大规模数据集100 万细胞的微调可能需要多卡并行scvi-tools的部分功能仍在 CPU 上运行参考文献Cui, H., Wang, C., Maan, H., Pang, K., Luo, F., Duan, N., Wang, B. scGPT: toward building a foundation model for single-cell multi-omics using generative AI. Nature Methods, 21(8), 1470-1480 (2024). https://www.nature.com/articles/s41592-024-02201-0上游代码仓库https://github.com/bowang-lab/scGPT昇腾 NPU 适配版https://atomgit.com/AI4Science/scGPT