update model weights to bf16

config.json

@@ -13,8 +13,8 @@
     "sliding_window": 2047,
     "torch_dtype": "bfloat16"
   },
-  "torch_dtype": "float32",
-  "transformers_version": "4.41.1",
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.44.2",
   "vision_encoder_config": {
     "anyres_patch_sampling": true,
     "image_aspect_ratio": "anyres",

generation_config.json

@@ -1,7 +1,7 @@
 {
   "_from_model_config": true,
   "bos_token_id": 1,
-  "eos_token_id": 32000,
+  "eos_token_id": 32007,
   "pad_token_id": 32000,
-  "transformers_version": "4.41.1"
+  "transformers_version": "4.44.2"
 }

model-00001-of-00004.safetensors → model-00001-of-00002.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:2394080dd0bac25f59461eba060e937989ab2336edffb415db3a36d7f3fc371e
-size 4962660968
+oid sha256:21e0452442b5c189d7f0a1078d243a4ad705036e12703f25f81f0711ae478d70
+size 4972926984

model-00002-of-00004.safetensors → model-00002-of-00002.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d63cfd9676a01932ea6674b0103a88a3863c3c8028b66d6716d4bd36945b66ed
-size 4983112136
+oid sha256:9a5e8bd3fbe75d20605d6760268337170a544f04bc4dfac00c2cba65981d7deb
+size 3745680670

model-00003-of-00004.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:cf6be012e8534481a9ca4e8dd8f3e482de42af52cb86b23c463b1e55ab5a40a2
-size 4983112168

model-00004-of-00004.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:71e9607ebf4884b3913619693b547dd4ae9bb66d9e6de42e80bcb5357de36914
-size 2508236156

modeling_xgenmm.py

@@ -162,6 +162,7 @@ class XGenMMConfig(PretrainedConfig):
                 "task_specific_params": None,
                 "problem_type": None,
                 "model_type": "phi3",
+                "_attn_implementation": "flash_attention_2",
             }
             logger.info(
                 "text_config is None. Initializing the text config with default values (`Phi3Config`)."
@@ -1031,7 +1032,7 @@ class VLM(nn.Module):
             num_additional_embeddings=len(self.special_tokens),
             _weight=self.lang_model.get_input_embeddings().weight,
             pad_token_id=self.pad_token_id,
-        )
+        ).to(self.lang_model.dtype)
         if hasattr(input_embeds, "additional_embedding"):
             input_embeds.additional_embedding.weight.data.normal_(
                 mean=0.0,
@@ -1052,7 +1053,7 @@ class VLM(nn.Module):
                 if hasattr(self.lang_model.get_output_embeddings(), "bias")
                 else None
             ),
-        )
+        ).to(self.lang_model.dtype)
         if hasattr(out_embeds, "additional_fc"):
             out_embeds.additional_fc.weight.data.normal_(
                 mean=0.0,
@@ -1642,7 +1643,8 @@ class VLMWithLanguageStream(VLM):
             if has_labels:
                 new_label = labels[i].clone()
 
-            for img_num, img_idx in enumerate(image_token_idxs):
+            for img_num in range(len(image_token_idxs)):
+                img_idx = image_token_idxs[img_num]
                 # Get vision token attention mask for padded llava-style any resolution image tokens.
                 if self.image_aspect_ratio == "anyres":
                     num_vis_tokens = vision_tokens[i][img_num].shape[0]
@@ -1662,6 +1664,10 @@ class VLMWithLanguageStream(VLM):
                     vis_attention_mask = torch.ones(
                         num_vis_tokens, dtype=torch.long
                     ).to(attention_mask.device)
+                
+                # Offset the rest of image tokens with current num_vis_tokens
+                for j in range(img_num+1, len(image_token_idxs)):
+                    image_token_idxs[j] += (num_vis_tokens - 1)
 
                 new_embed = torch.cat(
                     (
@@ -2029,11 +2035,15 @@ class XGenMMModelForConditionalGeneration(PreTrainedModel):
 
         # vision encoder initialization
         vision_encoder = AutoModel.from_pretrained(
-            config.vision_encoder_config.model_name
+            config.vision_encoder_config.model_name,
+            torch_dtype=config.text_config.torch_dtype,
         ).vision_model
 
         # language model initialization
-        language_model = AutoModelForCausalLM.from_config(config.text_config)
+        language_model = AutoModelForCausalLM.from_config(
+            config.text_config,
+            torch_dtype=config.text_config.torch_dtype,
+        )
         check_embedding_fns(language_model)
         # Update _tied_weights_keys using the base model used.
         if language_model._tied_weights_keys is not None:
@@ -2052,7 +2062,7 @@ class XGenMMModelForConditionalGeneration(PreTrainedModel):
                 f"Warning: The language embedding dimension in the vision tokenizer config is different from the language model's embedding dimension. Overwriting the language embedding dimension in the vision tokenizer config to {overwrite}."
             )
 
-        vision_tokenizer = XGenMMVisionTokenizer(config.vision_tokenizer_config).model
+        vision_tokenizer = XGenMMVisionTokenizer(config.vision_tokenizer_config).model.to(language_model.dtype)
 
         self.vlm = XGenMMPerceiver(
             vision_encoder=vision_encoder,

setup.sh

@@ -2,6 +2,6 @@ pip install torch==2.2.1 torchvision==0.17.1 torchaudio==2.2.1 --index-url https
 pip install open_clip_torch==2.24.0
 pip install einops
 pip install einops-exts
-pip install transformers==4.41.1
+pip install transformers==4.44.2
 # optional
 pip install ipywidgets

special_tokens_map.json

@@ -6,13 +6,7 @@
     "rstrip": false,
     "single_word": false
   },
-  "eos_token": {
-    "content": "<|endoftext|>",
-    "lstrip": false,
-    "normalized": false,
-    "rstrip": false,
-    "single_word": false
-  },
+  "eos_token": "<|end|>",
   "pad_token": {
     "content": "<pad>",
     "lstrip": false,

tokenizer_config.json

@@ -1,6 +1,7 @@
 {
   "add_bos_token": false,
   "add_eos_token": false,
+  "add_prefix_space": null,
   "added_tokens_decoder": {
     "0": {
       "content": "<unk>",
@@ -126,7 +127,7 @@
   "bos_token": "<s>",
   "chat_template": "{% for message in messages %}{% if message['role'] == 'system' %}{{'<|system|>\n' + message['content'] + '<|end|>\n'}}{% elif message['role'] == 'user' %}{{'<|user|>\n' + message['content'] + '<|end|>\n'}}{% elif message['role'] == 'assistant' %}{{'<|assistant|>\n' + message['content'] + '<|end|>\n'}}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ '<|assistant|>\n' }}{% else %}{{ eos_token }}{% endif %}",
   "clean_up_tokenization_spaces": false,
-  "eos_token": "<|endoftext|>",
+  "eos_token": "<|end|>",
   "legacy": false,
   "model_max_length": 4096,
   "pad_token": "<pad>",