v1

app.py

@@ -18,8 +18,8 @@ from transformers import TextIteratorStreamer
 from torchvision.transforms.functional import pil_to_tensor
 
 # flash attention
-# import subprocess
-# subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
+import subprocess
+subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
 # accel
 accel = Accelerator()

trol/arch_internlm2/modeling_internlm2.py

@@ -867,13 +867,15 @@ class InternLM2Model(InternLM2PreTrainedModel):
         self.norm = InternLM2RMSNorm(
             config.hidden_size, eps=config.rms_norm_eps)
 
-        self.trol_gating = nn.ModuleList([nn.Linear(self.config.hidden_size, 1)]*self.config.num_hidden_layers)
-        self.trol_function = lambda x, idx: 0.5*F.tanh(self.trol_gating[idx](x))+0.5
-
         self.gradient_checkpointing = False
         # Initialize weights and apply final processing
         self.post_init()
 
+    def initialize_trol_gating(self):
+        self.trol_gating = nn.ModuleList([nn.Linear(self.config.hidden_size, 1).cuda()]*self.config.num_hidden_layers)
+        self.trol_function = lambda x, idx: 0.5*F.tanh(self.trol_gating[idx](x))+0.5
+
+
     def get_input_embeddings(self):
         return self.tok_embeddings
 

trol/arch_phi3/modeling_phi3.py

@@ -1031,13 +1031,15 @@ class Phi3Model(Phi3PreTrainedModel):
         self._attn_implementation = "flash_attention_2"
         self.norm = Phi3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 
-        self.trol_gating = nn.ModuleList([nn.Linear(self.config.hidden_size, 1)]*self.config.num_hidden_layers)
-        self.trol_function = lambda x, idx: 0.5*F.tanh(self.trol_gating[idx](x))+0.5
-
         self.gradient_checkpointing = False
         # Initialize weights and apply final processing
         self.post_init()
 
+    def initialize_trol_gating(self):
+        self.trol_gating = nn.ModuleList([nn.Linear(self.config.hidden_size, 1).cuda()]*self.config.num_hidden_layers)
+        self.trol_function = lambda x, idx: 0.5*F.tanh(self.trol_gating[idx](x))+0.5
+
+
     def get_input_embeddings(self):
         return self.embed_tokens
 

trol/load_trol.py

@@ -81,11 +81,17 @@ def load_trol(link):
     # setting config
     setting_trol_config(trol, tok_trol, image_special_token)
 
-
     # trol gating load
     from huggingface_hub import hf_hub_download
     try:
+        trol.model.initialize_trol_gating()
         trol.model.trol_gating.load_state_dict(torch.load(hf_hub_download(repo_id=path, filename="trol_gating.pt")))
     except:
+        trol.language_model.model.initialize_trol_gating()
         trol.language_model.model.trol_gating.load_state_dict(torch.load(hf_hub_download(repo_id=path, filename="trol_gating.pt")))
+
+    # X -> float16 conversion 
+    for param in trol.parameters():
+        if 'float32' in str(param.dtype).lower() or 'float16' in str(param.dtype).lower():
+            param.data = param.data.to(torch.float16)
     return trol, tok_trol