[Bugfix] BioMedGPT fp16 (PharMolix#28)

.gitignore

@@ -141,6 +141,7 @@ cython_debug/
 /.vscode
 /assets/*
 /ckpts/**/*
+/datasets/**/*
 !/ckpts/**/.placeholder
 !/datasets/**/.placeholder
 

configs/encoders/multimodal/biomedgptv.json

@@ -36,7 +36,7 @@
             "use_float16": true
         },
         "llm": {
-            "model_name_or_path": "./ckpts/text_ckpts/biomed-llama-7b",
+            "model_name_or_path": "./ckpts/text_ckpts/biomedgpt-lm-7b",
             "use_float16": true
         }
     }

examples/biomedgpt_inference.ipynb

@@ -25,6 +25,13 @@
     "path"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "BioMedGPT-10B is composed of a language model (BioMedGPT-LM-7B), a molecule encoder (GraphMVP), a protein encoder (ESM2-3B) and two modality adaptors. For inference, install `config.json`, `special_token_map.json`, `tokenizer_config.json`, `tokenizer.json`, `tokenizer.model` for BioMedGPT-LM-7B [here](https://huggingface.co/PharMolix/BioMedGPT-LM-7B) and put them under `OpenBioMed/ckpts/text_ckpts/biomedgpt-lm-7b`. Install `config.json`, `tokenizer_config.json`, `vocab.txt` for ESM2-3B [here](https://huggingface.co/facebook/esm2_t36_3B_UR50D) and put them under `OpenBioMed/ckpts/text_ckpts/biomedgpt-lm-7b`. Approximately 20GB GPU Memory is required to load BioMedGPT-10B with fp16."
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 2,
@@ -58,8 +65,7 @@
     "fix_path_in_config(config, path)\n",
     "print(\"Config: \", config)\n",
     "\n",
-    "# ~20G GPU Memory is required to load BioMedGPT-10B\n",
-    "device = torch.device(\"cuda:1\")\n",
+    "device = torch.device(\"cuda:0\")\n",
     "config[\"network\"][\"device\"] = device\n",
     "model = BioMedGPTV(config[\"network\"])\n",
     "ckpt = torch.load(\"../ckpts/fusion_ckpts/biomedgpt_10b.pth\")\n",

open_biomed/models/multimodal/biomedgpt/biomedgpt.py

@@ -99,9 +99,9 @@ def __init__(self, config):
 
         # load protein structure encoder
         self.prot_structure_config = EsmConfig.from_json_file(os.path.join(config["protein"]["model_name_or_path"], "config.json"))
-        if config["protein"]["use_float16"]:
-            self.prot_structure_config.torch_dtype = "float16" 
         self.prot_structure_encoder = EsmModel(self.prot_structure_config)
+        if config["protein"]["use_float16"]:
+            self.prot_structure_encoder = self.prot_structure_encoder.half()
         if config["protein"]["lora"]:
             from peft import get_peft_model, LoraConfig, TaskType
             logger.info("applying lora to protein structure encoder")
@@ -122,9 +122,9 @@ def __init__(self, config):
         self.llm_tokenizer.add_special_tokens({'unk_token': '<unk>'})
         logger.info("loading llm")
         self.llm_config = LlamaConfig.from_json_file(os.path.join(config["llm"]["model_name_or_path"], "config.json"))
-        if config["llm"]["use_float16"]:
-            self.llm_config.torch_dtype = "float16"
         self.llm = LlamaForCausalLM(self.llm_config)
+        if config["llm"]["use_float16"]:
+            self.llm = self.llm.half()
         for name, param in self.llm.named_parameters():
             param.requires_grad = False
         #self.llm = LlamaForCausalLM.from_pretrained(config["llm"]["ckpt"], torch_dtype=torch.float16)

requirements.txt

@@ -1,9 +1,9 @@
 torch>=1.13.0
 torch_geometric
-transformers
+transformers>=4.28.1
 SentencePiece
 rdkit
 pandas
 scanpy
-eniops
+einops
 local_attention