Add print statements

modeling_cogvlm.py

@@ -225,6 +225,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 +241,34 @@ 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
 
+        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]
@@ -308,6 +337,7 @@ 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: