Upload modeling_vila.py

modeling_vila.py

@@ -212,6 +212,7 @@ class VILAPretrainedModel(PreTrainedModel):
             self.vision_tower = self.vision_tower.cuda()
         # set device_map auto can autoamtically shard llm to different devices
         self.llm, self.tokenizer = self.init_llm(llm_cfg, config, device_map=device_map)
+        self.llm_model_embed_tokens = self.llm.model.embed_tokens
 
         # NOTE(ligeng): hard code to set padding_side to left
         self.tokenizer.padding_side = "left"
@@ -221,6 +222,12 @@ class VILAPretrainedModel(PreTrainedModel):
         self.post_config()
         self.is_loaded = True
 
+        self.llm_only_need_embed = kwargs.get("llm_only_need_embed", False)
+        if self.llm_only_need_embed:
+            print("We only need the embed_tokens in llm.")
+            del self.llm
+            self.llm = None
+
         assert (
             self.llm is not None or self.vision_tower is not None or self.mm_projector is not None
         ), "At least one of the components must be instantiated."
@@ -628,7 +635,7 @@ class VILAForCasualLM(VILAPretrainedModel):
                 self.encoders[name].end_tokens = None
 
         # Extract text and media embeddings
-        text_embeds = self.llm.model.embed_tokens(input_ids)
+        text_embeds = self.llm_model_embed_tokens(input_ids)
         if media is not None:
             media_embeds = self.__embed_media_tokens(media, media_config)
         else:
@@ -712,6 +719,7 @@ class VILAForCasualLM(VILAPretrainedModel):
                     dummy = torch.zeros(infos[0]["shape"], dtype=infos[0]["dtype"], device=self.device)
                     embeds["dummy"].extend(self.encoders[name]([dummy], media_config[name]))
                     continue
+            media[name] = [a.to(torch.bfloat16) for a in media[name]]
             embeds[name] = deque(self.encoders[name](media[name], media_config[name]))
         return embeds