Update modeling_cogvlm.py

modeling_cogvlm.py

@@ -117,7 +117,8 @@ def attention_fn(
         attention_mask: "torch.tensor(B, H, L, HD)",
         *,
         scaling_attention_score: bool = True,
-        attention_dropout: nn.Module = None
+        attention_dropout: nn.Module = None,
+        print_values: bool = False,
 ):
     attention_mask_bool = (attention_mask == 0)
     is_low_triangle = (attention_mask_bool == torch.ones_like(attention_mask_bool, dtype=torch.float).tril()).all()
@@ -126,6 +127,10 @@ def attention_fn(
         warnings.warn("It's recommended to use torch2.0 or higher.")
     if int(torch.__version__.split('.')[0]) >= 2 and scaling_attention_score and (is_full or is_low_triangle):
         dropout_p = 0. if attention_dropout is None or not attention_dropout.training else attention_dropout.p
+
+        if print_values:
+            print("Is_causal:", not is_full)
+
         return torch.nn.functional.scaled_dot_product_attention(
             query_layer, key_layer, value_layer,
             attn_mask=None,
@@ -225,6 +230,7 @@ class VisionExpertAttention(nn.Module):
             past_key_value: Optional[Tuple[torch.Tensor]] = None,
             output_attentions: bool = False,
             use_cache: bool = False,
+            print_values: bool = False,
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
         bsz, q_len, _ = hidden_states.size()
         vision_token_mask, language_token_mask = get_expert_mask(token_type_ids)
@@ -240,6 +246,36 @@ class VisionExpertAttention(nn.Module):
         key_states = self._transpose_for_scores(key_states)  # B, H, L, HD
         value_states = self._transpose_for_scores(value_states)  # B, H, L, HD
 
+        # if print_values:
+
+            # torch.save(query_states, "query_states.pt")
+            # torch.save(key_states, "key_states.pt")
+            # torch.save(value_states, "value_states.pt")
+
+            # from huggingface_hub import HfApi
+
+            # api = HfApi()
+            # api.upload_file(
+            #     path_or_fileobj="query_states.pt",
+            #     path_in_repo="query_states.pt",
+            #     repo_id="nielsr/test-cogvlm",
+            #     repo_type="dataset",
+            # )
+            # api = HfApi()
+            # api.upload_file(
+            #     path_or_fileobj="key_states.pt",
+            #     path_in_repo="key_states.pt",
+            #     repo_id="nielsr/test-cogvlm",
+            #     repo_type="dataset",
+            # )
+            # api = HfApi()
+            # api.upload_file(
+            #     path_or_fileobj="value_states.pt",
+            #     path_in_repo="value_states.pt",
+            #     repo_id="nielsr/test-cogvlm",
+            #     repo_type="dataset",
+            # )
+
         kv_seq_len = key_states.shape[-2]
         if past_key_value is not None:
             kv_seq_len += past_key_value[0].shape[-2]
@@ -252,9 +288,31 @@ class VisionExpertAttention(nn.Module):
 
         past_key_value = (key_states, value_states) if use_cache else None
 
+        if print_values:
+            print("Shape of query_states:", query_states.shape)
+            print("Last values of query_states:", query_states[0,0,-3:,-3:])
+            print("Mean of query_states:", query_states.mean())
+
+            print("Shape of key_states:", key_states.shape)
+            print("Last values of key_states:", key_states[0,0,-3:,-3:])
+            print("Mean of key_states:", key_states.mean())
+
+            print("Shape of value_states:", value_states.shape)
+            print("First values of value_states:", value_states[0,0,-3:,-3:])
+            print("Mean of value_states:", value_states.mean())
+
+            print("Shape of the attention_mask:", attention_mask.shape)
+            print("Mean of the attention_mask:", attention_mask.float().mean())
+            print("Is_full:", (attention_mask > 0).all())
+
         context_layer = attention_fn(
             query_layer=query_states, key_layer=key_states, value_layer=value_states, attention_mask=attention_mask,
-            scaling_attention_score=True, attention_dropout=None)
+            scaling_attention_score=True, attention_dropout=None, print_values=print_values)
+        
+        if print_values:
+            print("Shape of context_layer:", context_layer.shape)
+            print("First values of context_layer:", context_layer[0,0,:3,:3])
+        
         if context_layer.size() != (bsz, self.num_heads, q_len, self.head_dim):
             raise ValueError(
                 f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
@@ -290,11 +348,18 @@ class CogVLMDecoderLayer(nn.Module):
             past_key_value: Optional[Tuple[torch.Tensor]] = None,
             output_attentions: Optional[bool] = False,
             use_cache: Optional[bool] = False,
+            print_values = False,
     ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
         residual = hidden_states
 
+        # if print_values:
+        #     print("Hidden states before RMS norm:", hidden_states[0, :3, :3])
+
         hidden_states = self.input_layernorm(hidden_states)
 
+        # if print_values:
+        #     print("Hidden states after RMS norm, before self attention:", hidden_states[0,:3,:3])
+
         # Self Attention
         hidden_states, self_attn_weights, present_key_value = self.self_attn(
             hidden_states=hidden_states,
@@ -304,7 +369,12 @@ class CogVLMDecoderLayer(nn.Module):
             past_key_value=past_key_value,
             output_attentions=output_attentions,
             use_cache=use_cache,
+            print_values=print_values,
         )
+
+        # if print_values:
+        #     print("Hidden states after self attention:", hidden_states[0,:3,:3])
+
         hidden_states = residual + hidden_states
 
         # Fully Connected
@@ -413,6 +483,7 @@ class CogVLMModel(CogVLMPreTrainedModel):
             output_attentions: Optional[bool] = None,
             output_hidden_states: Optional[bool] = None,
             return_dict: Optional[bool] = None,
+            step: int = None,
     ) -> Union[Tuple, BaseModelOutputWithPast]:
         """take care of image_encode, token_type_ids, position_ids and (attention_mask = None is fine)"""
 
@@ -425,10 +496,44 @@ class CogVLMModel(CogVLMPreTrainedModel):
                 assert token_type_ids is not None, f"multi-modality requires `token_type_ids`!"
                 assert len(input_ids) == len(images), f"{len(input_ids)} {len(images)}"
                 inputs_embeds = self.embed_tokens(input_ids)
+
                 images_features = self.encode_images(images)
                 images_features = rearrange(images_features, 'b n d -> (b n) d')
                 images_features = images_features.to(dtype=inputs_embeds.dtype, device=inputs_embeds.device)
+
+                # from huggingface_hub import HfApi
+
+                # torch.save(images_features, "images_features.pt")
+                # torch.save(inputs_embeds, "inputs_embeds.pt")
+                # torch.save(token_type_ids, "token_type_ids.pt")
+                
+                # api = HfApi()
+                # api.upload_file(
+                #     path_or_fileobj="images_features.pt",
+                #     path_in_repo="images_features.pt",
+                #     repo_id="nielsr/test-cogvlm",
+                #     repo_type="dataset",
+                # )
+                # api.upload_file(
+                #     path_or_fileobj="inputs_embeds.pt",
+                #     path_in_repo="inputs_embeds.pt",
+                #     repo_id="nielsr/test-cogvlm",
+                #     repo_type="dataset",
+                # )
+                # api.upload_file(
+                #     path_or_fileobj="token_type_ids.pt",
+                #     path_in_repo="token_type_ids.pt",
+                #     repo_id="nielsr/test-cogvlm",
+                #     repo_type="dataset",
+                # )
+
+                # print("First values of text embeddings:", inputs_embeds[0, :3, :3])
+                # print("First values of images_features:", images_features[0, :3])
+
                 inputs_embeds = inputs_embeds.index_put([token_type_ids == VISION_TOKEN_TYPE], images_features)
+
+                # print("First values of inputs_embeds after index_put:", inputs_embeds[0, :3, :3])
+
             else:  # single-modality
                 if token_type_ids is None:
                     token_type_ids = torch.ones_like(input_ids, dtype=torch.long, device=input_ids.device) * LANGUAGE_TOKEN_TYPE
@@ -450,6 +555,7 @@ class CogVLMModel(CogVLMPreTrainedModel):
             output_attentions=output_attentions,
             output_hidden_states=output_hidden_states,
             return_dict=return_dict,
+            step=step,
         )
 
     def llm_forward(
@@ -464,6 +570,7 @@ class CogVLMModel(CogVLMPreTrainedModel):
             output_attentions: Optional[bool] = None,
             output_hidden_states: Optional[bool] = None,
             return_dict: Optional[bool] = None,
+            step: int = None,
     ) -> Union[Tuple, BaseModelOutputWithPast]:
         """largely copy from llama forward and adapt for cogvlm with `token_type_ids`"""
         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
@@ -513,6 +620,48 @@ class CogVLMModel(CogVLMPreTrainedModel):
 
         hidden_states = inputs_embeds
 
+        if step == 1:
+            torch.save(hidden_states, "hidden_states_step_1.pt")
+            torch.save(attention_mask, "attention_mask_step_1.pt")
+            torch.save(token_type_ids, "token_type_ids_step_1.pt")
+            torch.save(position_ids, "position_ids_step_1.pt")
+            torch.save(past_key_values, "past_key_value_step_1.pt")
+            
+            from huggingface_hub import HfApi
+
+            api = HfApi()
+
+            api.upload_file(
+                path_or_fileobj="hidden_states_step_1.pt",
+                path_in_repo="hidden_states_step_1.pt",
+                repo_id="nielsr/test-cogvlm",
+                repo_type="dataset",
+            )
+            api.upload_file(
+                path_or_fileobj="attention_mask_step_1.pt",
+                path_in_repo="attention_mask_step_1.pt",
+                repo_id="nielsr/test-cogvlm",
+                repo_type="dataset",
+            )
+            api.upload_file(
+                path_or_fileobj="token_type_ids_step_1.pt",
+                path_in_repo="token_type_ids_step_1.pt",
+                repo_id="nielsr/test-cogvlm",
+                repo_type="dataset",
+            )
+            api.upload_file(
+                path_or_fileobj="position_ids_step_1.pt",
+                path_in_repo="position_ids_step_1.pt",
+                repo_id="nielsr/test-cogvlm",
+                repo_type="dataset",
+            )
+            api.upload_file(
+                path_or_fileobj="past_key_value_step_1.pt",
+                path_in_repo="past_key_value_step_1.pt",
+                repo_id="nielsr/test-cogvlm",
+                repo_type="dataset",
+            )
+
         # decoder layers
         all_hidden_states = () if output_hidden_states else None
         all_self_attns = () if output_attentions else None
@@ -531,9 +680,21 @@ class CogVLMModel(CogVLMPreTrainedModel):
                 past_key_value=past_key_value,
                 output_attentions=output_attentions,
                 use_cache=use_cache,
+                print_values=idx==0 and step==1,
             )
             hidden_states = layer_outputs[0]
 
+            # if idx == 0:
+            #     torch.save(hidden_states, "hidden_states_after_layer_0.pt")
+                
+            #     api = HfApi()
+            #     api.upload_file(
+            #         path_or_fileobj="hidden_states_after_layer_0.pt",
+            #         path_in_repo="hidden_states_after_layer_0.pt",
+            #         repo_id="nielsr/test-cogvlm",
+            #         repo_type="dataset",
+            #     )
+
             if use_cache:
                 next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)
 
@@ -542,6 +703,10 @@ class CogVLMModel(CogVLMPreTrainedModel):
 
         hidden_states = self.norm(hidden_states)
 
+        if step == 1:
+            print("Shape of hidden states:", hidden_states.shape)
+            print("First values of hidden states:", hidden_states[0,:3,:3])
+
         # add hidden states from the last decoder layer
         if output_hidden_states:
             all_hidden_states += (hidden_states,)
@@ -649,6 +814,7 @@ class CogVLMForCausalLM(CogVLMPreTrainedModel):
             output_hidden_states: Optional[bool] = None,
             return_dict: Optional[bool] = None,
             labels: Optional[torch.LongTensor] = None,
+            step: int = None,
     ) -> Union[Tuple, CausalLMOutputWithPast]:
         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
         output_hidden_states = (
@@ -669,6 +835,7 @@ class CogVLMForCausalLM(CogVLMPreTrainedModel):
             output_attentions=output_attentions,
             output_hidden_states=output_hidden_states,
             return_dict=return_dict,
+            step=step,
         )
 
         hidden_states = outputs[0]
@@ -745,6 +912,7 @@ class CogVLMForCausalLM(CogVLMPreTrainedModel):
             model_kwargs: Dict[str, Any],
             is_encoder_decoder: bool = False,
             standardize_cache_format: bool = False,
+            model_inputs: Optional[Dict[str, Any]] = None,
     ) -> Dict[str, Any]:
         # update past_key_values
         model_kwargs["past_key_values"] = self._extract_past_from_model_output(

visual.py

@@ -31,7 +31,7 @@ class Attention(nn.Module):
         self.dense = nn.Linear(config.hidden_size, config.hidden_size)
         self.output_dropout = torch.nn.Dropout(config.dropout_prob)
 
-    def forward(self, x: "tensor(B, L, D)") -> "tensor(B, L, D)":
+    def forward(self, x: "tensor(B, L, D)", print_values=False) -> "tensor(B, L, D)":
         B, L, _ = x.shape
         qkv = self.query_key_value(x)
         qkv = qkv.reshape(B, L, 3, self.num_heads, -1).permute(2, 0, 1, 3, 4)  # 3, B, L, H, D
@@ -40,6 +40,7 @@ class Attention(nn.Module):
         out = xops.memory_efficient_attention(
             q, k, v, scale=self.scale,
         )
+
         output = self.dense(out.view(B, L, -1))
         output = self.output_dropout(output)
         return output
@@ -74,9 +75,18 @@ class TransformerLayer(nn.Module):
         self.mlp = MLP(config)
         self.post_attention_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
 
-    def forward(self, hidden_states):
+    def forward(self, hidden_states, print_values=False):
         attention_input = hidden_states
-        attention_output = self.input_layernorm(self.attention(attention_input))
+
+        if print_values:
+            print("Hidden states before attention:", attention_input[0, :3, :3])
+
+        attention_output = self.attention(attention_input, print_values=print_values)
+
+        if print_values:
+            print("Hidden states after attention:", attention_output[0, :3, :3])
+
+        attention_output = self.input_layernorm(attention_output)
         hidden_states = attention_input + attention_output
         mlp_input = hidden_states
         mlp_output = self.post_attention_layernorm(self.mlp(mlp_input))
@@ -90,8 +100,36 @@ class Transformer(nn.Module):
         self.layers = nn.ModuleList([TransformerLayer(config) for _ in range(config.num_hidden_layers)])
 
     def forward(self, hidden_states):
-        for layer_module in self.layers:
-            hidden_states = layer_module(hidden_states)
+
+        print("Shape of hidden states before CLIP:", hidden_states.shape)
+        # torch.save(hidden_states, "hidden_states_before_clip.pt")
+
+        # from huggingface_hub import HfApi
+
+        # api = HfApi()
+        # api.upload_file(
+        #     path_or_fileobj="hidden_states_before_clip.pt",
+        #     path_in_repo="hidden_states_before_clip.pt",
+        #     repo_id="nielsr/test-cogvlm",
+        #     repo_type="dataset",
+        # )
+
+        for idx, layer_module in enumerate(self.layers):
+            hidden_states = layer_module(hidden_states, print_values=idx==0)
+
+        print("Shape of hidden states after CLIP:", hidden_states.shape)
+        # torch.save(hidden_states, "hidden_states_after_clip.pt")
+
+        # from huggingface_hub import HfApi
+
+        # api = HfApi()
+        # api.upload_file(
+        #     path_or_fileobj="hidden_states_after_clip.pt",
+        #     path_in_repo="hidden_states_after_clip.pt",
+        #     repo_id="nielsr/test-cogvlm",
+        #     repo_type="dataset",
+        # )
+
         return hidden_states