Update modeling_dream.py

modeling_dream.py

@@ -23,6 +23,7 @@ import math
 from typing import List, Optional, Tuple, Union
 import os
 import torch
+import hashlib
 import torch.utils.checkpoint
 from torch import nn
 
@@ -47,6 +48,9 @@ from .generation_utils import DreamGenerationMixin, DreamGenerationConfig
 
 if is_flash_attn_2_available():
     from transformers.modeling_flash_attention_utils import _flash_attention_forward
+    
+def check_hash(X):
+    t = X.detach().cpu().contiguous().view(torch.uint16); print(hashlib.md5(t.numpy().tobytes()).hexdigest())
 
 
 logger = logging.get_logger(__name__)
@@ -360,7 +364,9 @@ class DreamSdpaAttention(DreamAttention):
         use_cache: bool = False,
         cache_position: Optional[torch.LongTensor] = None,
         position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,  # will become mandatory in v4.46
+        use_flex_attn: Optional[bool] = False,
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+
         if output_attentions:
             # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented.
             logger.warning_once(
@@ -378,14 +384,45 @@ class DreamSdpaAttention(DreamAttention):
 
         bsz, q_len, _ = hidden_states.size()
 
+        # Debug: Print all hidden_states[0] values
+        # with open("mabmcm_mmm.txt", "a") as f:
+        #     f.write(f"\n=== Layer {self.layer_idx} ===\n")
+        #     f.write(f"hidden_states[0] - all positions:\n")
+        #     for idx in range(len(hidden_states[0])):
+        #         f.write(f"  idx {idx}: {hidden_states[0][idx]}\n")
+
         query_states = self.q_proj(hidden_states)
         key_states = self.k_proj(hidden_states)
         value_states = self.v_proj(hidden_states)
 
+        # Debug: Print all QKV[0] values after projection (before view/transpose)
+        # with open("mabmcm_mmm.txt", "a") as f:
+        #     f.write(f"\nquery_states[0] (after proj) - all positions:\n")
+        #     for idx in range(len(query_states[0])):
+        #         f.write(f"  idx {idx}: {query_states[0][idx]}\n")
+        #     f.write(f"\nkey_states[0] (after proj) - all positions:\n")
+        #     for idx in range(len(key_states[0])):
+        #         f.write(f"  idx {idx}: {key_states[0][idx]}\n")
+        #     f.write(f"\nvalue_states[0] (after proj) - all positions:\n")
+        #     for idx in range(len(value_states[0])):
+        #         f.write(f"  idx {idx}: {value_states[0][idx]}\n")
+
         query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
         key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
         value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
 
+        # Debug: Print all QKV[0][0] values after view/transpose
+        # with open("mabmcm_mmm.txt", "a") as f:
+        #     f.write(f"\nquery_states[0][0] (after view/transpose) - all positions:\n")
+        #     for idx in range(len(query_states[0][0])):
+        #         f.write(f"  idx {idx}: {query_states[0][0][idx]}\n")
+        #     f.write(f"\nkey_states[0][0] (after view/transpose) - all positions:\n")
+        #     for idx in range(len(key_states[0][0])):
+        #         f.write(f"  idx {idx}: {key_states[0][0][idx]}\n")
+        #     f.write(f"\nvalue_states[0][0] (after view/transpose) - all positions:\n")
+        #     for idx in range(len(value_states[0][0])):
+        #         f.write(f"  idx {idx}: {value_states[0][0][idx]}\n")
+
         if position_embeddings is None:
             logger.warning_once(
                 "The attention layers in this model are transitioning from computing the RoPE embeddings internally "
@@ -398,6 +435,15 @@ class DreamSdpaAttention(DreamAttention):
             cos, sin = position_embeddings
         query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
 
+        # Debug: Print all QKV[0][0] values after positional embedding
+        # with open("mabmcm_mmm.txt", "a") as f:
+        #     f.write(f"\nquery_states[0][0] (after positional embedding) - all positions:\n")
+        #     for idx in range(len(query_states[0][0])):
+        #         f.write(f"  idx {idx}: {query_states[0][0][idx]}\n")
+        #     f.write(f"\nkey_states[0][0] (after positional embedding) - all positions:\n")
+        #     for idx in range(len(key_states[0][0])):
+        #         f.write(f"  idx {idx}: {key_states[0][0][idx]}\n")
+
         if past_key_value is not None:
             cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}  # Specific to RoPE models
             key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
@@ -405,6 +451,18 @@ class DreamSdpaAttention(DreamAttention):
         key_states = repeat_kv(key_states, self.num_key_value_groups)
         value_states = repeat_kv(value_states, self.num_key_value_groups)
 
+        # Debug: Print all QKV[0][0] values after grouping
+        # with open("mabmcm_mmm.txt", "a") as f:
+        #     f.write(f"\nquery_states[0][0] (after grouping) - all positions:\n")
+        #     for idx in range(len(query_states[0][0])):
+        #         f.write(f"  idx {idx}: {query_states[0][0][idx]}\n")
+        #     f.write(f"\nkey_states[0][0] (after grouping) - all positions:\n")
+        #     for idx in range(len(key_states[0][0])):
+        #         f.write(f"  idx {idx}: {key_states[0][0][idx]}\n")
+        #     f.write(f"\nvalue_states[0][0] (after grouping) - all positions:\n")
+        #     for idx in range(len(value_states[0][0])):
+        #         f.write(f"  idx {idx}: {value_states[0][0][idx]}\n")
+
         # causal_mask = attention_mask
         # if attention_mask is not None:  # no matter the length, we just slice it
         #     causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
@@ -420,7 +478,14 @@ class DreamSdpaAttention(DreamAttention):
         # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling.
         # The q_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case q_len == 1.
         # is_causal = True if causal_mask is None and q_len > 1 else False
+        if use_flex_attn:
+            # L = attention_mask.shape[0]
+            # attention_mask_inverted = 1 - attention_mask
+            # attention_mask = torch.cat([attention_mask, attention_mask_inverted], dim=1)
+            # attention_mask = torch.cat([attention_mask, torch.zeros(L, 2*L, dtype=attention_mask.dtype, device=attention_mask.device)], dim=0)
+            attention_mask = attention_mask.bool()
 
+        
         attn_output = torch.nn.functional.scaled_dot_product_attention(
             query_states,
             key_states,
@@ -430,9 +495,21 @@ class DreamSdpaAttention(DreamAttention):
             is_causal=False, # hard coded
         )
 
+        # Debug: Print all attn_output[0][0] values after attention
+        # with open("mabmcm_mmm.txt", "a") as f:
+        #     f.write(f"\nattn_output[0][0] (after attention) - all positions:\n")
+        #     for idx in range(len(attn_output[0][0])):
+        #         f.write(f"  idx {idx}: {attn_output[0][0][idx]}\n")
+
         attn_output = attn_output.transpose(1, 2).contiguous()
         attn_output = attn_output.view(bsz, q_len, self.hidden_size)
 
+        # Debug: Print all attn_output[0] values after view
+        # with open("mabmcm_mmm.txt", "a") as f:
+        #     f.write(f"\nattn_output[0] (after view) - all positions:\n")
+        #     for idx in range(len(attn_output[0])):
+        #         f.write(f"  idx {idx}: {attn_output[0][idx]}\n")
+
         attn_output = self.o_proj(attn_output)
 
         return attn_output, None, past_key_value
@@ -466,6 +543,7 @@ class DreamDecoderLayer(nn.Module):
         use_cache: Optional[bool] = False,
         cache_position: Optional[torch.LongTensor] = None,
         position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,  # will become mandatory in v4.46
+        use_flex_attn: Optional[bool] = False,
         **kwargs,
     ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
         """
@@ -489,9 +567,7 @@ class DreamDecoderLayer(nn.Module):
                 Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code
                 into the model
         """
-
         residual = hidden_states
-
         hidden_states = self.input_layernorm(hidden_states)
 
         # Self Attention
@@ -504,6 +580,7 @@ class DreamDecoderLayer(nn.Module):
             use_cache=use_cache,
             cache_position=cache_position,
             position_embeddings=position_embeddings,
+            use_flex_attn=use_flex_attn,
         )
         hidden_states = residual + hidden_states
 
@@ -642,7 +719,9 @@ class DreamBaseModel(DreamPreTrainedModel):
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
         cache_position: Optional[torch.LongTensor] = None,
+        use_flex_attn: Optional[bool]=None,
     ) -> Union[Tuple, BaseModelOutput]:
+        
         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
         output_hidden_states = (
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
@@ -660,7 +739,13 @@ class DreamBaseModel(DreamPreTrainedModel):
                     "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
                 )
                 use_cache = False
-
+        
+        # Remark: append an [MASK]*L suffix to the input_ids
+        # if use_flex_attn:
+            # mask_id = 151666
+            # L = input_ids.shape[1]
+            # input_ids = torch.cat([input_ids, torch.full((input_ids.shape[0], L), mask_id, dtype=input_ids.dtype, device=input_ids.device)], dim=1)
+        
         if inputs_embeds is None:
             inputs_embeds = self.embed_tokens(input_ids)
         
@@ -678,6 +763,9 @@ class DreamBaseModel(DreamPreTrainedModel):
 
         hidden_states = inputs_embeds
 
+        if use_flex_attn:
+            position_ids = torch.cat([position_ids[:, :16], torch.tensor([[11, 14, 10, 13, 15]], device=position_ids.device)], dim=1)
+        
         # create position embeddings to be shared across the decoder layers
         position_embeddings = self.rotary_emb(hidden_states, position_ids)
 
@@ -711,6 +799,7 @@ class DreamBaseModel(DreamPreTrainedModel):
                     use_cache=use_cache,
                     cache_position=cache_position,
                     position_embeddings=position_embeddings,
+                    use_flex_attn=use_flex_attn,
                 )
 
             hidden_states = layer_outputs[0]
@@ -782,8 +871,14 @@ class DreamModel(DreamGenerationMixin, DreamPreTrainedModel):
         return_dict: Optional[bool] = None,
         cache_position: Optional[torch.LongTensor] = None,
         num_logits_to_keep: int = 0,
+        use_flex_attn: bool = False,
         **loss_kwargs,
     ) -> Union[Tuple, MaskedLMOutput]:
+
+        if not use_flex_attn:
+            attention_mask = "full"
+        
+        # Remark: in our method, attention_mask should be an L*L matrix
         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
         output_hidden_states = (
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
@@ -802,6 +897,7 @@ class DreamModel(DreamGenerationMixin, DreamPreTrainedModel):
             output_hidden_states=output_hidden_states,
             return_dict=return_dict,
             cache_position=cache_position,
+            use_flex_attn=use_flex_attn,
         )
 
         hidden_states = outputs[0]