Add print statements

modeling_cogvlm.py

@@ -296,8 +296,8 @@ class CogVLMDecoderLayer(nn.Module):
 
         hidden_states = self.input_layernorm(hidden_states)
 
-        if print_values:
-            print("Hidden states before self attention:", hidden_states[0,:3,:3])
+        # if print_values:
+        #     print("Hidden states before self attention:", hidden_states[0,:3,:3])
 
         # Self Attention
         hidden_states, self_attn_weights, present_key_value = self.self_attn(
@@ -310,8 +310,8 @@ class CogVLMDecoderLayer(nn.Module):
             use_cache=use_cache,
         )
 
-        if print_values:
-            print("Hidden states after self attention:", hidden_states[0,:3,:3])
+        # if print_values:
+        #     print("Hidden states after self attention:", hidden_states[0,:3,:3])
 
         hidden_states = residual + hidden_states
 
@@ -464,12 +464,12 @@ class CogVLMModel(CogVLMPreTrainedModel):
                     repo_type="dataset",
                 )
 
-                print("First values of text embeddings:", inputs_embeds[0, :3, :3])
-                print("First values of images_features:", images_features[0, :3])
+                # 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])
+                # print("First values of inputs_embeds after index_put:", inputs_embeds[0, :3, :3])
 
             else:  # single-modality
                 if token_type_ids is None:
@@ -542,8 +542,6 @@ class CogVLMModel(CogVLMPreTrainedModel):
         else:
             position_ids = position_ids.view(-1, seq_length).long()
 
-        print("Input ids:", input_ids)
-
         if inputs_embeds is None:
             inputs_embeds = self.embed_tokens(input_ids)
         # embed positions
@@ -578,9 +576,9 @@ class CogVLMModel(CogVLMPreTrainedModel):
             if output_hidden_states:
                 all_hidden_states += (hidden_states,)
 
-            if idx in [0, 1, 2]:
-                print(f"Hidden states before layer {idx}", hidden_states[0,:3,:3])
-                print(f"Mean of hidden states before layer {idx}", hidden_states.mean())
+            # if idx in [0, 1, 2]:
+            #     print(f"Hidden states before layer {idx}", hidden_states[0,:3,:3])
+            #     print(f"Mean of hidden states before layer {idx}", hidden_states.mean())
 
             past_key_value = past_key_values[idx] if past_key_values is not None else None
             layer_outputs = decoder_layer(

visual.py

@@ -74,9 +74,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)
+
+        if print_values:
+            print("Hidden states before attention:", attention_input[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))
@@ -105,7 +114,7 @@ class Transformer(nn.Module):
         )
 
         for idx, layer_module in enumerate(self.layers):
-            hidden_states = layer_module(hidden_states)
+            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")