biomed.omics.bl.sm.ma-ted-458m_flagos

## 模型介绍 MAMMAL 由 IBM Research 提出，面向药物研发全流程中「蛋白质、小分子、单细胞基因表达」等多模态生物医学数据难以统一建模的问题。它用一套**任务提示语法（task prompt syntax）**，把分类、回归与生成任务统一表达为 token 与标量的动态组合，从而在同一模型内支持单域任务与跨域实体（如蛋白—小分子相互作用）。 **MAMMAL -- Molecular Aligned Multi-Modal Architecture and Language** <https://arxiv.org/abs/2410.22367> 模型在超过 20 亿个生物样本上完成预训练，覆盖蛋白质序列、小分子结构与单细胞基因数据，在药物发现管线上的多个任务取得当时最优结果。 --- ## 模型描述 MAMMAL 的骨干是 **T5 风格的编码器-解码器**（项目 `mammal/keys.py` 自述为 *"t5 multitask generic implementation"*）。本仓库适配的 `biomed.omics.bl.sm.ma-ted-458m` 规模如下： | 项目 | 数值 | | :--- | :--- | | 参数量 | **458,007,202**（加载后实测，与标称 458M 精确吻合） | | 编码器 / 解码器层数 | 12 / 12 | | `d_model` | 768（实测 `hidden_states` 形状 `(1, 295, 768)`） | | 词表大小 | 100001（实测 `logits` 形状 `(1, 3, 100001)`） | | 权重精度 | fp32 | | 模态 | 蛋白质（AA）、小分子（SMILES/BPE）、单细胞基因表达、细胞属性 | > 层数 12/12 由算子统计交叉验证得出：单次前向共 217 次 `aten::mm`，恰为 > `encoder 12 层 × (q/k/v/o + FF 3 个) = 84` 与 `decoder 12 层 × (self 4 + cross 4 + FF 3) = 132` > 再加 `lm_head` 1 次之和。 任务的切换完全靠 **prompt 语法**完成，无需改动模型结构。例如蛋白—蛋白结合亲和力分类： ```text <@TOKENIZER-TYPE=AA><BINDING_AFFINITY_CLASS><SENTINEL_ID_0><MOLECULAR_ENTITY> <MOLECULAR_ENTITY_GENERAL_PROTEIN><SEQUENCE_NATURAL_START>{蛋白A}<SEQUENCE_NATURAL_END> <MOLECULAR_ENTITY><MOLECULAR_ENTITY_GENERAL_PROTEIN><SEQUENCE_NATURAL_START>{蛋白B}<SEQUENCE_NATURAL_END><EOS> ``` 其中 `<@TOKENIZER-TYPE=AA>` 是 fuse `ModularTokenizerOp` 要求的**分词器提示前缀**，它决定了后续序列用哪个分词器编码——这是使用本模型时最容易出错的地方：任何送入分词器的字符串都必须带该前缀，标签也不例外。 ```text 带提示前缀的 prompt ──► ModularTokenizerOp ──► T5 Encoder ──► T5 Decoder ──► 类别 / 标量 / 序列 ``` --- ## 适用场景 | 场景 | 说明 | | :--- | :--- | | **药物研发多任务推理** | 结合亲和力分类、药物-靶点相互作用（DTI）、致癌性、蛋白溶解度、单细胞类型注释等；通过替换 prompt 中的任务 token 切换 | | **跨域建模** | 蛋白质 + 小分子 + 基因表达混合输入，支持单域与跨域实体组合 | | **本地快速验证** | 使用官方示例（calmodulin / calcineurin 结合亲和力）检查加载、分词、推理链路是否正常 | | **ModelScope / OneCode 运行** | 作为独立模型包下载后直接安装依赖并运行脚本 | | **算子性能分析** | 使用 `prof.py` + `run_mammal.py` 采集单样本前向的算子调用与耗时分布，用于国产芯片算子覆盖与适配分析 | | **多卡训练** | 通过 `torchrun` 启动多进程训练（本仓库聚焦推理适配，未包含训练脚本） | --- ## 使用说明 ### 1. OneCode 使用 可通过 OneCode 在线环境体验智能化一键式 AI4S 编程： [点击体验智能化一键式 AI4S 编程](https://www.onescience.ai/) <!-- 待确认：OneCode 体验页的具体 URL，请按平台实际地址替换 --> ### 2. 手动安装使用 以下命令均假定在模型项目根目录执行。 #### 硬件要求 - 推理可在 CPU 或 DCU 上运行；本仓库的适配验证在 DCU 容器的 CPU 模式下完成。 - DCU 用户需要预先安装 DTK，建议使用 DTK 25.04.2 以上版本或与当前集群匹配的 OneScience 推荐版本。 - 单样本前向（`1×295` token 输入、458M 规模、fp32）在 CPU 上约需 **3.8 秒**，模型加载约 30 秒，内存占用约 3 GB。 - 训练建议使用 DCU 或 GPU，多卡训练使用 NCCL 后端，请确保设备驱动、通信库和 PyTorch 版本匹配。 #### 下载模型包 ```bash modelscope download --model FrontierScientificModels/biomed.omics.bl.sm.ma-ted-458m --local_dir ./mammal cd mammal ``` 若使用 Hugging Face 权重，可直接下载 `ibm/biomed.omics.bl.sm.ma-ted-458m` 并放到 `mammal-458m/` 目录下： ```bash export HF_ENDPOINT=https://hf-mirror.com hf download ibm/biomed.omics.bl.sm.ma-ted-458m --local-dir ./mammal-458m # 或（不受 CLI 改名影响） python3 -c "from huggingface_hub import snapshot_download; \ print(snapshot_download('ibm/biomed.omics.bl.sm.ma-ted-458m', local_dir='./mammal-458m'))" ``` 所需文件如下： | 文件 | 大小 | 说明 | | :--- | ---: | :--- | | `model.safetensors` | **1,832,066,472** | 模型权重（fp32） | | `config.json` | 2.8 KB | 模型配置 | | `tokenizer/config.yaml` | 967 B | fuse 分词器总配置 | | `tokenizer/t5_tokenizer_AA_special.json` | 70 KB | AA（蛋白质）分词器 | | `tokenizer/bpe_tokenizer_trained_on_chembl_zinc_*.json` | 186 KB | 小分子 BPE 分词器 | | `tokenizer/gene_tokenizer.json` | 2.6 MB | 基因分词器 | | `tokenizer/cell_attributes_tokenizer.json` | 93 KB | 单细胞属性分词器 | > 完整性校验：`458,007,202` 参数 × 4 字节 = `1,832,028,808`，与实测 > `1,832,066,472` 相差 37,664 字节（safetensors 头部与额外张量），属正常范围。 > 若本地文件字节数与上表不符，说明传输被截断，请重新下载。 #### 安装运行环境 **DCU 环境** ```bash # 请先激活 DTK 及 CONDA conda create -n onescience310 python=3.10 -y conda activate onescience310 pip install onescience[bio-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai ``` **GPU 环境** 暂未提供。药物研发领域（含生命科学）当前仅提供 DCU 版本，GPU 版本后续按需补充。 #### 安装模型依赖（关键：一律 `--no-deps`） 上游 `pyproject.toml` 声明了 `tensorflow>=2.17`，但经源码核查，`mammal/` 包内**没有任何 tensorflow 引用**，该声明属于过度依赖。在部分 DCU 容器中 tensorflow 的 `libtensorflow_cc.so` 存在符号缺失（如 `ncclCommRegister`），一旦被安装链带入就会导致 import 失败。因此**必须逐项手动安装**： ```bash # 1) fuse-med-ml（分词器 ModularTokenizerOp 的提供者） pip download fuse-med-ml==0.4.1 --no-deps -d ./wheels pip install --no-deps ./wheels/fuse_med_ml-0.4.1-py3-none-any.whl # 2) fuse 运行时依赖链（缺什么补什么，不要全量安装 fuse 的 34 个依赖） pip install --no-deps hdf5plugin h5py pip install deepdiff orderly_set paramiko psutil wget termcolor click # 3) mammal 本体 pip install --no-deps git+https://github.com/BiomedSciAI/biomed-multi-alignment.git pip install omegaconf peft # 4) 验证 python3 -c "from mammal.model import Mammal; print('mammal ok')" python3 -c "from fuse.data.tokenizers.modular_tokenizer.op import ModularTokenizerOp; print('tokenizer ok')" ``` 若 `mammal ok` 因缺少某个模块而失败，按报错末尾的 `No module named 'xxx'` 单独补装即可，**不要**改用不带 `--no-deps` 的安装命令。 本机另一条可行路径：`mammal/__init__.py` 是空文件，因此把 `mammal/` 目录直接置于 Python 路径下即可 import，无需 pip 安装。 #### 运行前必须设置的环境变量 在 DCU 容器中，torch 自带的厂商设备后端扩展可能加载失败，且 CUDA 探测接口在无 GPU 环境下会触发段错误，因此运行任何脚本前需设置： ```bash export TORCH_DEVICE_BACKEND_AUTOLOAD=0 # 关闭厂商设备后端自动加载 export CUDA_VISIBLE_DEVICES="" # 屏蔽 CUDA 探测（CPU 推理 / 算子采集时必需） export USE_TF=NO USE_JAX=NO # 可选保险：让 transformers 跳过 TF/JAX 分支 ``` #### 快速验证 使用官方示例（calmodulin 与 calcineurin 的结合亲和力分类）检查完整链路： ```bash export TORCH_DEVICE_BACKEND_AUTOLOAD=0 python3 test_mammal.py --quick --model-path ./mammal-458m ``` 该脚本依次检查运行环境、模型加载与参数量、分词器可用性、官方任务推理，并输出逐项 PASS / FAIL 汇总。 官方示例输出： ```text 任务 : 结合亲和力分类（BINDING_AFFINITY_CLASS） 输入 : calmodulin 与 calcineurin 两条蛋白序列（按预训练语法包装） tokenize : input_ids (295,) 输出 : generated_output = '<SENTINEL_ID_0><1><EOS>' ``` 实测结果：`exit=0`，**PASS 12 / FAIL 0**，模型加载 28.5 秒，参数量 458,007,202。 #### 算子性能分析 采集单样本前向的算子调用与耗时分布： ```bash export TORCH_DEVICE_BACKEND_AUTOLOAD=0 # 推理：仅前向（teacher forcing，一次完整 encoder + decoder） CUDA_VISIBLE_DEVICES="" python3 prof.py run_mammal.py # 训练：前向 + 反向 CUDA_VISIBLE_DEVICES="" MAMMAL_PROF_MODE=train python3 prof.py run_mammal.py ``` 运行后在当前目录生成 `op_profiles/<run_id>/`，其中包含 `top_operators.md`（TOP-N 算子耗时表）与 `operator_profile.json`（完整统计）。采集范围是「第一个 `nn.Module` 前向到脚本结束」，因此 `run_mammal.py` 只执行一个样本的一次前向。 实测输入规格（用于归档命名）： ```text model_dtypes : ['torch.float32'] input_dtype : torch.int64 input_ids.shape : (1, 295) hidden_states.shape : (1, 295, 768) logits.shape : (1, 3, 100001) 单次前向耗时 : 3814.0 ms ``` 对应产物目录名：`biomed.omics.bl.sm.ma-ted-458m_推理_1_float32_1x295x768`。 ```bash cd op_profiles mv run_mammal__<run_id> biomed.omics.bl.sm.ma-ted-458m_推理_1_float32_1x295x768 tar -czf "biomed.omics.bl.sm.ma-ted-458m_推理_1_float32_1x295x768.tar.gz" \ "biomed.omics.bl.sm.ma-ted-458m_推理_1_float32_1x295x768" ``` **实测算子分布**（单次前向，共 90 种算子、6640 次调用，算子 Self CPU 合计 3723.4 ms）： | 排名 | 算子 | 调用数 | Self CPU (ms) | 占比 | | ---: | :--- | ---: | ---: | ---: | | 1 | `aten::mm` | 217 | 2795.8 | **75.1%** | | 2 | `aten::mul` | 247 | 343.7 | 9.2% | | 3 | `aten::_scaled_dot_product_flash_attention_for_cpu` | 36 | 188.0 | 5.0% | | 4 | `aten::add` | 176 | 103.8 | 2.8% | | 5 | `aten::copy_` | 258 | 98.5 | 2.6% | | 6 | `aten::sum` | 62 | 81.5 | 2.2% | | 7 | 其余（`pow` / `tanh` / `cat` / `softmax` 等） | — | 112.1 | 3.1% | 矩阵乘法以 75% 的占比构成绝对瓶颈，其余为逐元素运算与归约，符合 Transformer 类模型的典型形态，可作为国产芯片算子覆盖与优化的优先目标。 > 说明：本次采集在 DCU 容器的 CPU 模式下完成（报告中 `Accelerator: none`）。 > 若需获取加速卡上的真实算子耗时，可在 flagOS 环境（`torch_fl` 可用，设备名 `flagos`）下重跑同一命令。 #### 训练权重 本仓库聚焦推理适配，不包含训练脚本。推理所用权重来自 Hugging Face `ibm/biomed.omics.bl.sm.ma-ted-458m`，下载后置于 `mammal-458m/` 目录即可。微调相关的完整流程请参考上游仓库中的 `mammal/examples/` 各任务目录。 #### 本仓库脚本一览 | 文件 | 作用 | | :--- | :--- | | `test_mammal.py` | 完整自测：环境、模型加载与参数量、分词器、官方任务推理 | | `run_mammal.py` | 算子采集用运行脚本：单样本、单次前向（可选反向） | | `prof.py` / `torch_op_profiler.py` | 算子采集入口与实现（官方提供） | | `mammal/` | 上游模型包（MAMMAL 实现，可直接置于 Python 路径下使用） | --- ## OneScience 官方信息 | 平台 | OneScience 主仓库 | Skills 仓库 | | :--- | :--- | :--- | | Gitee | <https://gitee.com/onescience-ai/onescience> | <https://gitee.com/onescience-ai/oneskills> | | GitHub | <https://github.com/onescience-ai/OneScience> | <https://github.com/onescience-ai/oneskills> | --- ## 引用与许可证 **许可证** 本仓库为 **MAMMAL** 原始工作的 OneScience 适配复现版本，参考 Hugging Face 上的官方权重与 GitHub 上的官方实现完成： - 源地址（权重）：<https://huggingface.co/ibm/biomed.omics.bl.sm.ma-ted-458m> - 源地址（代码）：<https://github.com/BiomedSciAI/biomed-multi-alignment> - 论文：<https://arxiv.org/abs/2410.22367> 原始仓库采用 **Apache License 2.0**（<https://www.apache.org/licenses/LICENSE-2.0>），本仓库沿用该许可，并在其中附加了面向 DCU 环境的适配代码与使用说明。使用本仓库即视为同时接受上游许可条款。 **引用** 如果您使用了本仓库的内容，请同时引用原始论文： ```bibtex @misc{shoshan2024mammalmolecularaligned, title = {MAMMAL -- Molecular Aligned Multi-Modal Architecture and Language}, author = {Yoel Shoshan and Moshiko Raboh and Michal Ozery-Flato and Vadim Ratner and Alex Golts and Jeffrey K. Weber and Ella Barkan and Simona Rabinovici-Cohen and Sagi Polaczek and Ido Amos and Ben Shapira and Liam Hazan and Matan Ninio and Sivan Ravid and Michael M. Danziger and Joseph A. Morrone and Parthasarathy Suryanarayanan and Michal Rosen-Zvi and Efrat Hexter}, year = {2024}, eprint = {2410.22367}, archivePrefix = {arXiv}, primaryClass = {q-bio.QM}, url = {https://arxiv.org/abs/2410.22367} } ```