Upload model

config.json

@@ -22,7 +22,6 @@
   "dlm_paradigm": "bidirectional",
   "dlm_type": "llada",
   "dp_varying_mask_ratio": false,
-  "dtype": "float32",
   "enforce_mask": false,
   "eos_token_id": 2,
   "global_loss_avg": false,
@@ -57,12 +56,25 @@
     "rope_type": "yarn",
     "type": "yarn"
   },
+  "rope_scaling": {
+    "beta_fast": 32.0,
+    "beta_slow": 1.0,
+    "factor": 16.0,
+    "llama_4_scaling_beta": 0.1,
+    "mscale": 1.0,
+    "mscale_all_dim": 1.0,
+    "original_max_position_embeddings": 16384,
+    "rope_theta": 1000000.0,
+    "rope_type": "yarn",
+    "type": "yarn"
+  },
   "rope_theta": 1000000.0,
   "seq_length": 8192,
   "sliding_window": null,
   "tie_word_embeddings": false,
   "tok_mask_half_life_ratio": null,
-  "transformers_version": "5.0.0rc1",
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.55.4",
   "use_cache": false,
   "vocab_size": 131072
 }

configuration_ministral_dlm.py

@@ -155,6 +155,7 @@ class MinistralDLMConfig(PretrainedConfig):
         tie_word_embeddings=False,
         rope_theta=1000000.0,
         rope_parameters=None,
+        rope_scaling=None,
         attention_bias=False,
         attention_dropout=0.0,
         mlp_bias=False,
@@ -204,6 +205,7 @@ class MinistralDLMConfig(PretrainedConfig):
         self.use_cache = use_cache
         self.rope_theta = rope_theta
         self.rope_parameters = rope_parameters
+        self.rope_scaling = rope_scaling
         self.attention_bias = attention_bias
         self.attention_dropout = attention_dropout
         self.mlp_bias = mlp_bias

generation_config.json

@@ -2,6 +2,6 @@
   "_from_model_config": true,
   "bos_token_id": 1,
   "eos_token_id": 2,
-  "transformers_version": "5.0.0rc1",
+  "transformers_version": "4.55.4",
   "use_cache": false
 }

model-00001-of-00004.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:59f86ca799393141bddf2d14132ac34374fea32b3cf87bf4bf1ed1839be30d1f
+size 4999767504

model-00002-of-00004.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fbfe58b42e25e9b11b3c1f34d59de5ae84133285bc2c3dc24ffc9d2280bf1c1b
+size 4999802928

model-00003-of-00004.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ade91c2da2eb811769316e84a2bd8034f2c48323e4249cf0ebc5f58a4ca703cf
+size 4915916376

model-00004-of-00004.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4dbf0f5cba01173593d5a4cd658bf64bd37eda9e33b418493576eee072013636
+size 2063657632

model.safetensors.index.json

@@ -0,0 +1,317 @@
+{
+  "metadata": {
+    "total_parameters": 8489553920,
+    "total_size": 16979107840
+  },
+  "weight_map": {
+    "diffusion_head.weight": "model-00004-of-00004.safetensors",
+    "encoder.embed_tokens.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.0.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.0.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.0.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.0.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.0.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.0.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.0.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.0.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.0.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.1.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.1.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.1.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.1.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.1.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.1.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.1.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.1.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.1.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.10.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.10.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.10.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.10.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.10.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.10.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.10.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.10.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.10.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.11.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.11.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.11.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.11.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.11.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.11.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.11.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.11.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.11.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.12.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.12.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.12.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.12.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.12.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.12.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.12.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.12.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.12.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.13.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.13.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.13.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.13.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.13.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.13.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.13.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.13.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.13.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.14.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.14.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.14.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.14.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.14.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.14.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.14.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.14.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.14.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.15.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.15.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.15.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.15.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.15.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.15.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.15.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.15.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.15.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.16.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.16.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.16.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.16.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.16.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.16.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.16.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.16.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.16.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.17.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.17.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.17.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.17.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.17.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.17.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.17.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.17.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.17.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.18.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.18.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.18.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.18.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.18.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.18.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.18.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.18.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.18.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.19.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.19.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.19.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.19.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.19.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.19.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.19.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.19.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.19.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.2.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.2.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.2.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.2.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.2.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.2.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.2.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.2.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.2.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.20.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.20.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.20.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.20.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.20.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.20.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.20.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.20.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.20.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.21.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.21.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.21.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.21.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.21.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.21.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.21.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.21.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.21.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.22.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.22.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.22.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.22.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.22.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.22.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.22.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.22.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.22.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.23.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.23.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.23.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.23.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.23.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.23.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.23.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.23.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.23.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.24.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.24.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.24.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.24.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.24.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.24.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.24.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.24.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.24.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.25.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.25.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.25.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.25.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.25.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.25.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.25.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.25.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.25.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.26.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.26.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.26.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.26.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.26.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.26.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.26.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.26.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.26.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.27.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.27.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.27.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.27.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.27.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.27.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.27.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.27.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.27.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.28.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.28.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.28.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.28.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.28.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.28.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.28.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.28.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.28.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.29.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.29.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.29.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.29.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.29.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.29.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.29.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.29.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.29.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.3.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.3.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.3.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.3.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.3.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.3.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.3.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.3.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.3.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.30.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.30.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.30.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.30.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.30.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.30.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.30.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.30.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.30.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.31.input_layernorm.weight": "model-00004-of-00004.safetensors",
+    "encoder.layers.31.mlp.down_proj.weight": "model-00004-of-00004.safetensors",
+    "encoder.layers.31.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.31.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.31.post_attention_layernorm.weight": "model-00004-of-00004.safetensors",
+    "encoder.layers.31.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.31.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.31.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.31.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "encoder.layers.32.input_layernorm.weight": "model-00004-of-00004.safetensors",
+    "encoder.layers.32.mlp.down_proj.weight": "model-00004-of-00004.safetensors",
+    "encoder.layers.32.mlp.gate_proj.weight": "model-00004-of-00004.safetensors",
+    "encoder.layers.32.mlp.up_proj.weight": "model-00004-of-00004.safetensors",
+    "encoder.layers.32.post_attention_layernorm.weight": "model-00004-of-00004.safetensors",
+    "encoder.layers.32.self_attn.k_proj.weight": "model-00004-of-00004.safetensors",
+    "encoder.layers.32.self_attn.o_proj.weight": "model-00004-of-00004.safetensors",
+    "encoder.layers.32.self_attn.q_proj.weight": "model-00004-of-00004.safetensors",
+    "encoder.layers.32.self_attn.v_proj.weight": "model-00004-of-00004.safetensors",
+    "encoder.layers.33.input_layernorm.weight": "model-00004-of-00004.safetensors",
+    "encoder.layers.33.mlp.down_proj.weight": "model-00004-of-00004.safetensors",
+    "encoder.layers.33.mlp.gate_proj.weight": "model-00004-of-00004.safetensors",
+    "encoder.layers.33.mlp.up_proj.weight": "model-00004-of-00004.safetensors",
+    "encoder.layers.33.post_attention_layernorm.weight": "model-00004-of-00004.safetensors",
+    "encoder.layers.33.self_attn.k_proj.weight": "model-00004-of-00004.safetensors",
+    "encoder.layers.33.self_attn.o_proj.weight": "model-00004-of-00004.safetensors",
+    "encoder.layers.33.self_attn.q_proj.weight": "model-00004-of-00004.safetensors",
+    "encoder.layers.33.self_attn.v_proj.weight": "model-00004-of-00004.safetensors",
+    "encoder.layers.4.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.4.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.4.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.4.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.4.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.4.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.4.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.4.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.4.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.5.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.5.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.5.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.5.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.5.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.5.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.5.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.5.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.5.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.6.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.6.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.6.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.6.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.6.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.6.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.6.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.6.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.6.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.7.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.7.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.7.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.7.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.7.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.7.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.7.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.7.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.7.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.8.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.8.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.8.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.8.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.8.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.8.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.8.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.8.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.8.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "encoder.layers.9.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.9.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.9.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.9.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.9.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.9.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.9.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.9.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.layers.9.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "encoder.norm.weight": "model-00004-of-00004.safetensors"
+  }
+}

modeling_ministral.py

@@ -9,7 +9,8 @@ from transformers.utils.generic import check_model_inputs
 from transformers.activations import ACT2FN
 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, 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
 from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
 from transformers.modeling_layers import (
@@ -23,7 +24,7 @@ from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_u
 from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
 from transformers.processing_utils import Unpack
 from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple
-from transformers.utils.generic import maybe_autocast
+# from transformers.utils.generic import maybe_autocast
 from .configuration_ministral_dlm import MinistralDLMConfig
 
 
@@ -33,8 +34,7 @@ def rotate_half(x):
     x2 = x[..., x.shape[-1] // 2 :]
     return torch.cat((-x2, x1), dim=-1)
 
-
-@use_kernel_func_from_hub("rotary_pos_emb")
+# @use_kernel_func_from_hub("rotary_pos_emb")
 def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
     """Applies Rotary Position Embedding to the query and key tensors.
 
@@ -105,7 +105,7 @@ def _get_llama_4_attn_scale(positions_ids: torch.Tensor, beta: float, max_positi
     return scaling.unsqueeze(-1)
 
 
-@use_kernelized_func(apply_rotary_pos_emb)
+# @use_kernelized_func(apply_rotary_pos_emb)
 class Ministral3Attention(nn.Module):
     """Multi-headed attention from 'Attention Is All You Need' paper"""
 
@@ -356,12 +356,13 @@ class Ministral3RotaryEmbedding(nn.Module):
         inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
         position_ids_expanded = position_ids[:, None, :].float()
 
-        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
-        with maybe_autocast(device_type=device_type, enabled=False):  # Force float32
-            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
-            emb = torch.cat((freqs, freqs), dim=-1)
-            cos = emb.cos() * self.attention_scaling
-            sin = emb.sin() * self.attention_scaling
+        # device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+        # with maybe_autocast(device_type=device_type, enabled=False):  # Force float32
+
+        freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        cos = emb.cos() * self.attention_scaling
+        sin = emb.sin() * self.attention_scaling
 
         return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
 
@@ -404,7 +405,8 @@ class Ministral3Model(Ministral3PreTrainedModel):
             inputs_embeds = self.embed_tokens(input_ids)
 
         if use_cache and past_key_values is None:
-            past_key_values = DynamicCache(config=self.config)
+            # past_key_values = DynamicCache(config=self.config)
+            past_key_values = DynamicCache()
 
         if cache_position is None:
             past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0

modeling_ministral_dlm.py

@@ -1,5 +1,4 @@
 import copy
-from dataclasses import dataclass
 from typing import Callable, Optional, Tuple, Union
 import random
 import os
@@ -11,7 +10,6 @@ import torch
 import torch.nn.functional as F
 from torch import nn
 from transformers.modeling_outputs import CausalLMOutputWithPast, BaseModelOutput
-from transformers.utils import ModelOutput
 
 from torch.nn.attention.flex_attention import flex_attention, create_block_mask
 
@@ -31,17 +29,6 @@ from .chat_utils import generate_with_prefix_cache_block_diff
 from .modeling_ministral import Ministral3Model, Ministral3PreTrainedModel, Ministral3Attention, apply_rotary_pos_emb, repeat_kv, _get_llama_4_attn_scale
 from .configuration_ministral_dlm import MinistralDLMConfig
 
-
-@dataclass
-class MinistralDiffOutputWithPast(ModelOutput):
-    loss: torch.FloatTensor | None = None
-    logits: torch.FloatTensor | None = None
-    causal_logits: torch.FloatTensor | None = None
-    past_key_values: Cache | None = None
-    hidden_states: tuple[torch.FloatTensor, ...] | None = None
-    attentions: tuple[torch.FloatTensor, ...] | None = None
-
-
 # @torch.compile(dynamic=True, mode="reduce-overhead")
 # @torch.compile(mode="default")
 # @torch.compile(fullgraph=True, mode="reduce-overhead", dynamic=False)
@@ -492,7 +479,6 @@ class MinistralDiffEncoderModel(Ministral3PreTrainedModel, GenerationMixin):
         loss_mask: Optional[torch.Tensor] = None,
         ce_loss_weight: float = 1.0,
         output_last_hidden_states_only: bool = False,
-        skip_loss: bool = False,
         **kwargs,
     ) -> CausalLMOutputWithPast:
 
@@ -579,7 +565,7 @@ class MinistralDiffEncoderModel(Ministral3PreTrainedModel, GenerationMixin):
             logits = logits[:, :input_ids_len]
 
         loss = None
-        if labels is not None and not skip_loss:
+        if labels is not None:  
             if self.config.dlm_paradigm == 'autoregressive':
                 shift_logits = logits[..., :-1, :].contiguous()
                 shift_labels = labels[..., 1:].contiguous()
@@ -716,10 +702,9 @@ class MinistralDiffEncoderModel(Ministral3PreTrainedModel, GenerationMixin):
                     else: 
                         loss = (loss, num_mask_tokens)
 
-        return MinistralDiffOutputWithPast(
+        return CausalLMOutputWithPast(
             loss=loss if not is_teacher else logits,
             logits=logits,
-            causal_logits=causal_logits,
             past_key_values=enc_out.past_key_values,
             hidden_states=None,
             attentions=None,