Upload model

chat_utils.py

@@ -133,7 +133,7 @@ def generate_with_prefix_cache_block_diff(
                 layer.self_attn.diffusion_lm=False
 
     # Compute KV cache for the prompt initially
-    output = model(prompt, use_cache=True, use_causal_mask=causal_context)
+    output = model(prompt, use_cache=True)
     past_key_values = output.past_key_values
 
     if causal_context:
@@ -230,8 +230,7 @@ def generate_with_prefix_cache_block_diff(
         output = model(
             x_accum[:, block_slice],
             past_key_values=past_key_values,
-            use_cache=True,
-            use_causal_mask=causal_context
+            use_cache=True
         )
         past_key_values = output.past_key_values
 

config.json

@@ -69,7 +69,6 @@
     "type": "yarn"
   },
   "rope_theta": 1000000.0,
-  "seq_length": 8192,
   "sliding_window": null,
   "tie_word_embeddings": false,
   "tok_mask_half_life_ratio": null,

configuration_ministral_dlm.py

@@ -70,8 +70,6 @@ class MinistralDLMConfig(PretrainedConfig):
             Whether to use a bias in up_proj, down_proj and gate_proj layers.
         sliding_window (`int`, *optional*, defaults to None):
             Sliding window attention size.
-        seq_length (`int`, *optional*, defaults to 8192):
-            Sequence length for training.
         mask_token_id (`int`, *optional*, defaults to -1):
             Token ID for masking in diffusion.
         dlm_type (`str`, *optional*, defaults to 'llada'):
@@ -161,7 +159,6 @@ class MinistralDLMConfig(PretrainedConfig):
         mlp_bias=False,
         sliding_window=None,
         attn_implementation="sdpa",
-        seq_length=8192,
         mask_token_id=-1,
         dlm_type='llada',
         random_length_prob=None,
@@ -214,7 +211,6 @@ class MinistralDLMConfig(PretrainedConfig):
         rope_config_validation(self)
         
         self.attn_implementation = attn_implementation
-        self.seq_length = seq_length
         
         self.mask_token_id = mask_token_id
         self.dlm_type = dlm_type

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:581e534c77fd49b8ab1d234f65bf08b88b03bd4dd10e397285a3441260957a8d
+oid sha256:73af2cd1c982f85bac01c7da43765deb3f2deced76eb93dbd2a6a968ff531349
 size 16979144720

modeling_ministral.py

@@ -11,7 +11,7 @@ from transformers.cache_utils import Cache, DynamicCache
 from transformers.generation import GenerationMixin
 # from transformers.integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func
 from transformers.integrations import use_kernel_forward_from_hub
-from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask, ALL_MASK_ATTENTION_FUNCTIONS, sdpa_mask_older_torch
+from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask
 from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
 from transformers.modeling_layers import (
     GenericForQuestionAnswering,
@@ -27,7 +27,6 @@ from transformers.utils import TransformersKwargs, auto_docstring, can_return_tu
 # from transformers.utils.generic import maybe_autocast
 from .configuration_ministral_dlm import MinistralDLMConfig
 
-#ALL_MASK_ATTENTION_FUNCTIONS._global_mapping['sdpa'] = sdpa_mask_older_torch
 
 def rotate_half(x):
     """Rotates half the hidden dims of the input."""
@@ -418,9 +417,10 @@ class Ministral3Model(Ministral3PreTrainedModel):
         if position_ids is None:
             position_ids = cache_position.unsqueeze(0)
 
-        #if self.training:
-        #    causal_mask = None
-        if kwargs.get("use_causal_mask", False):
+        if self.training:
+            causal_mask = None
+        
+        else:
             mask_function = create_causal_mask if self.config.sliding_window is None else create_sliding_window_causal_mask
             causal_mask = mask_function(
                 config=self.config,
@@ -430,8 +430,6 @@ class Ministral3Model(Ministral3PreTrainedModel):
                 past_key_values=past_key_values,
                 position_ids=position_ids,
             )
-        else:
-            causal_mask = None
 
         hidden_states = inputs_embeds
         position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids)

modeling_ministral_dlm.py

@@ -54,7 +54,6 @@ class MinistralFlexAttention(Ministral3Attention):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
  
-        self.max_seq_length = self.config.seq_length
         self.block_size_orig = self.config.block_size
 
         if self.config.dlm_paradigm == 'bidirectional':
@@ -62,9 +61,9 @@ class MinistralFlexAttention(Ministral3Attention):
         elif self.config.dlm_paradigm == 'autoregressive':
             self.autoregressive_mask = self.compute_block_mask(mode='autoregressive')
         elif self.config.dlm_paradigm == 'block_diff':
-            self.block_diff_mask = self.compute_block_mask(mode='block_diff', block_size=self.block_size_orig)
+            self.block_diff_mask = None
         elif self.config.dlm_paradigm == 'sbd_block_diff':
-            self.sbd_block_diff_mask = self.compute_block_mask(mode='sbd_block_diff', block_size=self.block_size_orig)
+            self.sbd_block_diff_mask = None
         else:
             raise ValueError(f"Unknown attention mode: {self.config.dlm_paradigm}")
 
@@ -79,7 +78,7 @@ class MinistralFlexAttention(Ministral3Attention):
         self.mode = mode
         self.block_size = block_size
 
-    def compute_block_mask(self, mode, q_len=None, block_size=None):
+    def compute_block_mask(self, mode, q_len, block_size=None):
 
         def bidirectional_mask(b, h, q, kv): 
             return (q >= kv) | (q < kv)
@@ -166,23 +165,15 @@ class MinistralFlexAttention(Ministral3Attention):
             attn_mask = autoregressive_mask
         elif mode == 'block_diff':
             assert block_size is not None
-            attn_mask = lambda b, h, q, kv: block_diff_mask(block_size, b, h, q, kv, self.max_seq_length)
+            attn_mask = lambda b, h, q, kv: block_diff_mask(block_size, b, h, q, kv, q_len//2)
         elif mode == 'sbd_block_diff':
             assert block_size is not None
-            attn_mask = lambda b, h, q, kv: sbd_block_diff_mask(block_size, b, h, q, kv, self.max_seq_length)
+            attn_mask = lambda b, h, q, kv: sbd_block_diff_mask(block_size, b, h, q, kv, q_len//2)
         else:
             raise ValueError(f"Unknown attention mode: {mode}")
 
-        if q_len is not None:
-            Q_LEN = q_len
-        else:
-            if mode in ['block_diff', 'sbd_block_diff']:
-                Q_LEN = self.max_seq_length * 2
-            else:
-                Q_LEN = self.max_seq_length
-
         block_mask = create_block_mask(
-            attn_mask, B=None, H=None, Q_LEN=Q_LEN, KV_LEN=Q_LEN
+            attn_mask, B=None, H=None, Q_LEN=q_len, KV_LEN=q_len
         )
 
         return block_mask
@@ -238,24 +229,24 @@ class MinistralFlexAttention(Ministral3Attention):
         value_states = repeat_kv(value_states, self.num_key_value_groups)
 
         if self.mode == 'bidirectional':
-            if q_len != self.bidirectional_mask.shape[-2]:
+            if self.bidirectional_mask is None or q_len != self.bidirectional_mask.shape[-2]:
                 block_mask = self.compute_block_mask(mode='bidirectional', q_len=q_len)
             else:
                 block_mask = self.bidirectional_mask
 
         elif self.mode == 'autoregressive':
-            if q_len != self.autoregressive_mask.shape[-2]:
+            if self.autoregressive_mask is None or q_len != self.autoregressive_mask.shape[-2]:
                 block_mask = self.compute_block_mask(mode='autoregressive', q_len=q_len)
             else:
                 block_mask = self.autoregressive_mask
 
         elif self.mode == 'block_diff':
-            if self.block_size != self.block_size_orig or q_len != self.block_diff_mask.shape[-2]:
+            if self.block_diff_mask is None or self.block_size != self.block_size_orig or q_len != self.block_diff_mask.shape[-2]:
                 block_mask = self.compute_block_mask(mode='block_diff', block_size=self.block_size, q_len=q_len)
             else:
                 block_mask = self.block_diff_mask
         elif self.mode == 'sbd_block_diff':
-            if self.block_size != self.block_size_orig or q_len != self.sbd_block_diff_mask.shape[-2]:
+            if self.sbd_block_diff_mask is None or self.block_size != self.block_size_orig or q_len != self.sbd_block_diff_mask.shape[-2]:
                 block_mask = self.compute_block_mask(mode='sbd_block_diff', block_size=self.block_size, q_len=q_len)
             else:
                 block_mask = self.sbd_block_diff_mask