Upload model

README.md

@@ -1 +1,4 @@
+---
+{}
+---
 Coming soon.

config.json

@@ -0,0 +1,50 @@
+{
+  "adaptive_mask_rate": false,
+  "architectures": [
+    "DiffEncoderModel"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "attn_implementation": "sdpa",
+  "auto_map": {
+    "AutoConfig": "configuration_nvrdiff.NVRDiffConfig",
+    "AutoModel": "modeling_nvrdiff.DiffEncoderModel"
+  },
+  "block_size": 32,
+  "diff_loss_weight": 1,
+  "disable_qk_norm": false,
+  "dlm_arch": "encoder",
+  "dlm_paradigm": "bidirectional",
+  "dlm_type": "dream",
+  "enforce_mask": false,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 2560,
+  "initializer_range": 0.02,
+  "intermediate_size": 9728,
+  "intl_mask": false,
+  "mask_token_id": 151662,
+  "max_position_embeddings": 32768,
+  "max_window_layers": 28,
+  "model_type": "qwen3",
+  "multi_sampling": null,
+  "num_ar_layers": 0,
+  "num_attention_heads": 32,
+  "num_diffusion_layers": 0,
+  "num_hidden_layers": 36,
+  "num_key_value_heads": 8,
+  "prefix_ratio": 0.8,
+  "random_length_prob": 0,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 1000000,
+  "seq_length": 1024,
+  "sliding_window": null,
+  "tie_word_embeddings": true,
+  "tok_mask_half_life_ratio": null,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.52.2",
+  "use_cache": false,
+  "use_sliding_window": false,
+  "vocab_size": 151936
+}

configuration_nvrdiff.py

@@ -0,0 +1,250 @@
+# coding=utf-8
+# Copyright 2024 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Qwen3 model configuration"""
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.modeling_rope_utils import rope_config_validation
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class NVRDiffConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`Qwen3Model`]. It is used to instantiate a
+    Qwen3 model according to the specified arguments, defining the model architecture. Instantiating a configuration
+    with the defaults will yield a similar configuration to that of
+    Qwen3-8B [Qwen/Qwen3-8B](https://huggingface.co/Qwen/Qwen3-8B).
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+
+    Args:
+        vocab_size (`int`, *optional*, defaults to 151936):
+            Vocabulary size of the Qwen3 model. Defines the number of different tokens that can be represented by the
+            `inputs_ids` passed when calling [`Qwen3Model`]
+        hidden_size (`int`, *optional*, defaults to 4096):
+            Dimension of the hidden representations.
+        intermediate_size (`int`, *optional*, defaults to 22016):
+            Dimension of the MLP representations.
+        num_hidden_layers (`int`, *optional*, defaults to 32):
+            Number of hidden layers in the Transformer encoder.
+        num_attention_heads (`int`, *optional*, defaults to 32):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        num_key_value_heads (`int`, *optional*, defaults to 32):
+            This is the number of key_value heads that should be used to implement Grouped Query Attention. If
+            `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
+            `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
+            converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
+            by meanpooling all the original heads within that group. For more details checkout [this
+            paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to `32`.
+        head_dim (`int`, *optional*, defaults to 128):
+            The attention head dimension.
+        hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
+            The non-linear activation function (function or string) in the decoder.
+        max_position_embeddings (`int`, *optional*, defaults to 32768):
+            The maximum sequence length that this model might ever be used with.
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        rms_norm_eps (`float`, *optional*, defaults to 1e-06):
+            The epsilon used by the rms normalization layers.
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether or not the model should return the last key/values attentions (not used by all models). Only
+            relevant if `config.is_decoder=True`.
+        tie_word_embeddings (`bool`, *optional*, defaults to `False`):
+            Whether the model's input and output word embeddings should be tied.
+        rope_theta (`float`, *optional*, defaults to 10000.0):
+            The base period of the RoPE embeddings.
+        rope_scaling (`Dict`, *optional*):
+            Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type
+            and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value
+            accordingly.
+            Expected contents:
+                `rope_type` (`str`):
+                    The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope',
+                    'llama3'], with 'default' being the original RoPE implementation.
+                `factor` (`float`, *optional*):
+                    Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In
+                    most scaling types, a `factor` of x will enable the model to handle sequences of length x *
+                    original maximum pre-trained length.
+                `original_max_position_embeddings` (`int`, *optional*):
+                    Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during
+                    pretraining.
+                `attention_factor` (`float`, *optional*):
+                    Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention
+                    computation. If unspecified, it defaults to value recommended by the implementation, using the
+                    `factor` field to infer the suggested value.
+                `beta_fast` (`float`, *optional*):
+                    Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear
+                    ramp function. If unspecified, it defaults to 32.
+                `beta_slow` (`float`, *optional*):
+                    Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear
+                    ramp function. If unspecified, it defaults to 1.
+                `short_factor` (`List[float]`, *optional*):
+                    Only used with 'longrope'. The scaling factor to be applied to short contexts (<
+                    `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
+                    size divided by the number of attention heads divided by 2
+                `long_factor` (`List[float]`, *optional*):
+                    Only used with 'longrope'. The scaling factor to be applied to long contexts (<
+                    `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
+                    size divided by the number of attention heads divided by 2
+                `low_freq_factor` (`float`, *optional*):
+                    Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE
+                `high_freq_factor` (`float`, *optional*):
+                    Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE
+        attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`):
+            Whether to use a bias in the query, key, value and output projection layers during self-attention.
+        use_sliding_window (`bool`, *optional*, defaults to `False`):
+            Whether to use sliding window attention.
+        sliding_window (`int`, *optional*, defaults to 4096):
+            Sliding window attention (SWA) window size. If not specified, will default to `4096`.
+        max_window_layers (`int`, *optional*, defaults to 28):
+            The number of layers that use SWA (Sliding Window Attention). The bottom layers use SWA while the top use full attention.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+
+    ```python
+    >>> from transformers import Qwen3Model, Qwen3Config
+
+    >>> # Initializing a Qwen3 style configuration
+    >>> configuration = Qwen3Config()
+
+    >>> # Initializing a model from the Qwen3-8B style configuration
+    >>> model = Qwen3Model(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "qwen3"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    # Default tensor parallel plan for base model `Qwen3`
+    base_model_tp_plan = {
+        "layers.*.self_attn.q_proj": "colwise",
+        "layers.*.self_attn.k_proj": "colwise",
+        "layers.*.self_attn.v_proj": "colwise",
+        "layers.*.self_attn.o_proj": "rowwise",
+        "layers.*.mlp.gate_proj": "colwise",
+        "layers.*.mlp.up_proj": "colwise",
+        "layers.*.mlp.down_proj": "rowwise",
+    }
+    base_model_pp_plan = {
+        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+        "norm": (["hidden_states"], ["hidden_states"]),
+    }
+
+    def __init__(
+        self,
+        vocab_size=151936,
+        hidden_size=4096,
+        intermediate_size=22016,
+        num_hidden_layers=32,
+        num_attention_heads=32,
+        num_key_value_heads=32,
+        head_dim=128,
+        hidden_act="silu",
+        max_position_embeddings=32768,
+        initializer_range=0.02,
+        rms_norm_eps=1e-6,
+        use_cache=True,
+        tie_word_embeddings=False,
+        rope_theta=10000.0,
+        rope_scaling=None,
+        attention_bias=False,
+        use_sliding_window=False,
+        sliding_window=4096,
+        max_window_layers=28,
+        attention_dropout=0.0,
+        attn_implementation="sdpa",
+        seq_length=1024,
+        mask_token_id=-1,
+        dlm_type='llada',
+        random_length_prob=None,
+        num_ar_layers=4,
+        num_diffusion_layers=4,
+        diff_loss_weight=1,
+        enforce_mask=False,
+        prefix_ratio=0.8,
+        dlm_paradigm='bidirectional',
+        dlm_arch='encoder',
+        block_size=32,
+        disable_qk_norm=False,
+        intl_mask=False,
+        tok_mask_half_life_ratio=None,
+        adaptive_mask_rate=False,
+        multi_sampling=None,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.use_sliding_window = use_sliding_window
+        self.sliding_window = sliding_window  # we check `use_sliding_window` in the modeling code
+        self.max_window_layers = max_window_layers
+
+        # for backward compatibility
+        if num_key_value_heads is None:
+            num_key_value_heads = num_attention_heads
+
+        self.num_key_value_heads = num_key_value_heads
+        self.head_dim = head_dim
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+        self.attention_bias = attention_bias
+        self.attention_dropout = attention_dropout
+        # Validate the correctness of rotary position embeddings parameters
+        # BC: if there is a 'type' field, move it to 'rope_type'.
+        if self.rope_scaling is not None and "type" in self.rope_scaling:
+            self.rope_scaling["rope_type"] = self.rope_scaling["type"]
+        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
+        self.random_length_prob = random_length_prob
+        self.num_ar_layers = num_ar_layers
+        self.num_diffusion_layers = num_diffusion_layers
+        self.diff_loss_weight = diff_loss_weight
+        self.enforce_mask = enforce_mask
+        self.prefix_ratio = prefix_ratio
+        self.dlm_paradigm = dlm_paradigm
+        self.dlm_arch = dlm_arch
+        self.block_size = block_size
+        self.disable_qk_norm = disable_qk_norm
+        self.intl_mask = intl_mask
+        self.tok_mask_half_life_ratio = tok_mask_half_life_ratio
+        self.adaptive_mask_rate = adaptive_mask_rate
+        self.multi_sampling = multi_sampling
+        
+        super().__init__(
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+
+__all__ = ["Qwen3Config"]

generation_config.json

@@ -0,0 +1,5 @@
+{
+  "_from_model_config": true,
+  "transformers_version": "4.52.2",
+  "use_cache": false
+}

model-00001-of-00002.safetensors

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

model-00002-of-00002.safetensors

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

model.safetensors.index.json

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

modeling_nvrdiff.py

@@ -0,0 +1,534 @@
+import copy
+from typing import Callable, Optional, Tuple, Union
+import random
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from torchtitan.models.qwen3.modeling_qwen3 import Qwen3Config, Qwen3Model, Qwen3PreTrainedModel, Qwen3Attention, apply_rotary_pos_emb, repeat_kv
+
+from torch.nn.attention.flex_attention import flex_attention, create_block_mask
+
+from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
+
+from transformers.processing_utils import Unpack
+
+from transformers.cache_utils import Cache, DynamicCache
+
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+
+from transformers.generation import GenerationMixin
+
+import math
+
+# @torch.compile(dynamic=True, mode="reduce-overhead")
+# @torch.compile(mode="default")
+# @torch.compile(fullgraph=True, mode="reduce-overhead", dynamic=False)
+@torch.compile(fullgraph=True, mode="max-autotune-no-cudagraphs", dynamic=False)
+def fused_flex_attention(q, k, v, block_mask=None):
+    return flex_attention(q, k, v, block_mask=block_mask)
+
+# with reference to https://github.com/pytorch-labs/attention-gym/blob/main/examples/flex_attn.ipynb
+class Qwen3FlexAttention(Qwen3Attention):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+ 
+        self.max_seq_length = self.config.seq_length
+        self.prefix_len_orig = int(self.config.seq_length * self.config.prefix_ratio)
+        self.block_size_orig = self.config.block_size
+
+        if self.config.dlm_paradigm == 'bidirectional':
+            self.bidirectional_mask = self.compute_block_mask(mode='bidirectional')
+        elif self.config.dlm_paradigm == 'prefix_bidirectional':
+            self.prefix_bidirectional_mask = self.compute_block_mask(mode='prefix_bidirectional', prefix_len=self.prefix_len_orig)
+        elif self.config.dlm_paradigm == 'efficient_block_diff':
+            self.efficient_block_diff_mask = self.compute_block_mask(mode='efficient_block_diff', block_size=self.block_size_orig)
+        elif self.config.dlm_paradigm == 'block_diff':
+            self.block_diff_mask = self.compute_block_mask(mode='block_diff', block_size=self.block_size_orig)
+        else:
+            raise ValueError(f"Unknown attention mode: {self.config.dlm_paradigm}")
+
+        self.prefix_len = self.prefix_len_orig
+        self.block_size = self.block_size_orig
+        self.mode = 'bidirectional'
+
+        import torch._dynamo.config as dcfg
+        dcfg.cache_size_limit = 512
+
+
+    def set_attention_mode(self, mode, prefix_len=None, block_size=None):
+        self.mode = mode
+        self.prefix_len = prefix_len
+        self.block_size = block_size
+
+
+    def compute_block_mask(self, mode, prefix_len=None, q_len=None, block_size=None):
+
+        def bidirectional_mask(b, h, q, kv): 
+            return (q >= kv) | (q < kv)
+        
+        def prefix_bidirectional_mask(prefix_len, b, h, q, kv):
+            return (kv <= prefix_len) | (q >= prefix_len)
+
+        def efficient_block_diff_mask(block_size, b, h, q, kv):
+            return (q // block_size) >= (kv // block_size)
+
+        def block_diff_mask(block_size, b, h, q_idx, kv_idx, n):
+            """
+            Constructs the specialized block diffusion attention mask for training
+            composed of three masks:
+            - **Block Diagonal Mask (M_BD)**: Self-attention within noised blocks
+            - **Offset Block Causal Mask (M_OBC)**: Cross-attention for conditional context
+            - **Block Causal Mask (M_BC)**: Attention to update x0
+
+            Args:
+                b, h: Batch and head indices (ignored for mask logic).
+                q_idx, kv_idx: Query and Key indices.
+                seq_len: Total sequence length.
+                block_size: Defines the block structure.
+
+            Returns:
+                A boolean attention mask.
+            """
+
+            # Indicate whether token belongs to xt or x0
+            x0_flag_q = (q_idx >= n)
+            x0_flag_kv = (kv_idx >= n)
+
+            # Compute block indices
+            block_q = torch.where(x0_flag_q == 1,
+                                    (q_idx - n) // block_size,
+                                    q_idx // block_size)
+            block_kv = torch.where(x0_flag_kv == 1,
+                                    (kv_idx - n) // block_size,
+                                    kv_idx // block_size)
+
+            # **1. Block Diagonal Mask (M_BD) **
+            block_diagonal = (block_q == block_kv) & (x0_flag_q == x0_flag_kv)
+
+            # **2. Offset Block-Causal Mask (M_OBC) **
+            offset_block_causal = (
+                (block_q > block_kv)
+                & (x0_flag_kv == 1)
+                & (x0_flag_q == 0)
+            )
+
+            # **3. Block-Causal Mask (M_BC) **
+            block_causal = (block_q >= block_kv) & (x0_flag_kv == 1) & (x0_flag_q == 1)
+
+            # **4. Combine Masks **
+            return block_diagonal | offset_block_causal | block_causal
+
+        if mode == 'bidirectional':
+            attn_mask = bidirectional_mask
+        elif mode == 'prefix_bidirectional':
+            assert prefix_len is not None
+            attn_mask = lambda b, h, q, kv: prefix_bidirectional_mask(prefix_len, b, h, q, kv)
+        elif mode == 'efficient_block_diff':
+            assert block_size is not None
+            attn_mask = lambda b, h, q, kv: efficient_block_diff_mask(block_size, b, h, q, kv)
+        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)
+        else:
+            raise ValueError(f"Unknown attention mode: {mode}")
+
+        if q_len is not None:
+            Q_LEN = q_len
+        else:
+            if mode == '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
+        )
+
+        return block_mask
+
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: Tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[torch.Tensor],
+        past_key_value: Optional[Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        is_training: bool = True,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        bsz, q_len, _ = hidden_states.size()
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+        key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+        cos, sin = position_embeddings
+
+        if self.mode == 'block_diff' and is_training:
+            # Split query and key states in half along sequence length dimension
+            q1, q2 = query_states.chunk(2, dim=2)
+            k1, k2 = key_states.chunk(2, dim=2)
+            
+            # Apply RoPE independently to each half
+            q1, k1 = apply_rotary_pos_emb(q1, k1, cos, sin)
+            q2, k2 = apply_rotary_pos_emb(q2, k2, cos, sin)
+            
+            # Recombine the halves
+            query_states = torch.cat([q1, q2], dim=2)
+            key_states = torch.cat([k1, k2], dim=2)
+        else:
+            query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        if past_key_value is not None:
+            # sin and cos are specific to RoPE models; cache_position needed for the static cache
+            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+        key_states = repeat_kv(key_states, self.num_key_value_groups)
+        value_states = repeat_kv(value_states, self.num_key_value_groups)
+
+        if self.mode == 'bidirectional':
+            if q_len != self.bidirectional_mask.shape[-2]:
+                block_mask = self.compute_block_mask(mode='bidirectional', prefix_len=self.prefix_len, q_len=q_len)
+            else:
+                block_mask = self.bidirectional_mask
+
+        elif self.mode == 'prefix_bidirectional':
+            if self.prefix_len != self.prefix_len_orig or q_len != self.prefix_bidirectional_mask.shape[-2]:
+                block_mask = self.compute_block_mask(mode='prefix_bidirectional', prefix_len=self.prefix_len, q_len=q_len)
+
+                # print('create new block mask length for:',self.prefix_len)
+                # print(f"Block mask shape: {block_mask.shape}")
+                # print("Block mask pattern:")
+                # print(block_mask)
+            else:
+                block_mask = self.prefix_bidirectional_mask
+        elif self.mode == 'efficient_block_diff':
+            if self.block_size != self.block_size_orig or q_len != self.efficient_block_diff_mask.shape[-2]:
+                block_mask = self.compute_block_mask(mode='efficient_block_diff', block_size=self.block_size, q_len=q_len)
+            else:
+                block_mask = self.efficient_block_diff_mask
+        elif self.mode == 'block_diff':
+            if 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
+        else:
+            raise ValueError(f"Unknown attention mode: {self.mode}")
+
+        attn_output = fused_flex_attention(query_states, key_states, value_states, block_mask=block_mask)
+        attn_output = attn_output.transpose(1, 2).reshape(*input_shape, -1).contiguous()
+
+        attn_output = self.o_proj(attn_output)
+
+        return attn_output, None
+
+
+def gumbel_topk(log_w: torch.Tensor, k: int) -> torch.Tensor:
+    """Return a Bool mask of length len(log_w) with exactly k True."""
+    g = -torch.log(-torch.log(torch.rand_like(log_w) + 1e-9) + 1e-9)
+    topk = torch.topk(log_w + g, k).indices
+    mask = torch.zeros_like(log_w, dtype=torch.bool)
+    mask[topk] = True
+    return mask
+            
+
+class DiffEncoderModel(Qwen3PreTrainedModel, GenerationMixin):
+    """
+    A single model with:
+      - a bidirectional encoder + diffusion‐LM head over A
+      - a causal decoder + LM head over B, conditioned on F_A
+    """
+
+    def __init__(self, config: Qwen3Config):
+        super().__init__(config)
+
+        self.mask_token_id = config.mask_token_id
+
+        diffusion_config = copy.deepcopy(config)
+        diffusion_config.diffusion_lm = True
+
+        if config.dlm_paradigm in ['prefix_bidirectional', 'efficient_block_diff', 'block_diff']:
+            diffusion_config.attn_class = Qwen3FlexAttention
+        elif config.dlm_paradigm in ['bidirectional', 'autoregressive']:
+            diffusion_config.attn_class = Qwen3Attention
+
+            if config.dlm_paradigm == 'autoregressive':
+                diffusion_config.diffusion_lm = False
+        else:
+            raise ValueError(f"Unsupported DLM paradigm: {config.dlm_paradigm}")
+        
+        self.encoder = Qwen3Model(diffusion_config)
+        self.diffusion_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.vocab_size = config.vocab_size
+
+        self.post_init()
+
+
+    def forward_process(self, input_ids, eps=1e-3, block_size=None, loss_mask=None):
+        b, l = input_ids.shape
+        device = input_ids.device
+
+        t = torch.rand(b, device=device)
+
+        p_mask = (1 - eps) * t + eps  # shape: (b,)
+        p_mask = p_mask[:, None].expand(-1, l)  # shape: (b, l)
+
+        masked_indices = torch.rand((b, l), device=device) < p_mask
+
+        if loss_mask is not None:
+            masked_indices[loss_mask == 0] = 0
+
+        noisy_batch = torch.where(masked_indices, self.mask_token_id, input_ids)        
+
+        return noisy_batch, masked_indices, p_mask
+
+
+    def forward_process_exp(
+        self,
+        input_ids: torch.Tensor,
+        eps: float = 1e-3,
+        block_size: int | None = None,
+        half_life_ratio: float = 0.25, # λ = ln 2 / (half_life_ratio·L)
+        loss_mask: Optional[torch.Tensor] = None,
+    ):
+        """
+        Two-stage corruption with optional per-block sampling.
+
+        • Stage 1:  m ~ U(eps, 1)   →   k = round(m · len)  (exact budget).
+        • Stage 2:  sample exactly k positions with weights
+                    w_i(m) = exp[ λ · (1−m) · i ]   (late-heavy when m→0,
+                                                     uniform when m→1).
+
+          If `block_size` is given, the procedure is run *independently*
+          inside each contiguous block of that length (last block may be shorter).
+          When block_size is provided, m is sampled per-block and p_mask is per-block.
+
+        Args
+        ----
+        input_ids : (B, L)  LongTensor
+        eps       : minimum corruption ratio
+        block_size: if not None, operate block-wise with per-block m sampling
+        half_life_ratio : controls steepness when m→0
+        """
+        B, L = input_ids.shape
+        device = input_ids.device
+        dtype  = torch.float32
+
+        masked_indices = torch.zeros((B, L), dtype=torch.bool, device=device)
+        p_mask = torch.zeros((B, L), dtype=dtype, device=device)
+
+        # ---------- Stage 1 & 2: whole-sentence or block-wise -------------------
+        for b in range(B):
+            if block_size is None:
+                # ---------- Per-batch sampling (original behavior) ----------
+                m = eps + (1.0 - eps) * torch.rand(1, device=device).item()   # scalar
+                k_tot = int(round(m * L))
+                k_tot = max(1, min(k_tot, L))  # clamp to [1, L]
+                
+                # Fill p_mask for this batch
+                p_mask[b, :] = m
+                
+                slope = 1.0 - m          # ∈ [0,1]; 0 ⇒ uniform, 1 ⇒ late-heavy
+                
+                # ------- single pool over the whole sentence -------------
+                lam_base = math.log(2.0) / (half_life_ratio * L) # base decay rate (λ when slope=1)
+
+                pos   = torch.arange(L, device=device, dtype=dtype)
+                log_w = (lam_base * slope * pos).clone()
+
+                masked_indices[b] = gumbel_topk(log_w, k_tot)
+
+            else:
+                # ---------- Per-block sampling ----------
+                num_blocks = math.ceil(L / block_size)
+                lam_base = math.log(2.0) / (half_life_ratio * block_size) # base decay rate (λ when slope=1)
+
+                for blk in range(num_blocks):
+                    start = blk * block_size
+                    end   = min((blk + 1) * block_size, L)
+                    blk_len = end - start
+
+                    # Sample m per block
+                    m_blk = eps + (1.0 - eps) * torch.rand(1, device=device).item()
+                    
+                    # Fill p_mask for this block
+                    p_mask[b, start:end] = m_blk
+                    
+                    # per-block budget
+                    k_blk = int(round(m_blk * blk_len))
+                    k_blk = max(0, min(k_blk, blk_len))
+                    if k_blk == 0:
+                        continue
+
+                    slope = 1.0 - m_blk          # ∈ [0,1]; 0 ⇒ uniform, 1 ⇒ late-heavy
+
+                    pos   = torch.arange(blk_len, device=device, dtype=dtype)
+                    log_w = lam_base * slope * pos
+
+                    blk_mask = gumbel_topk(log_w, k_blk)
+                    masked_indices[b, start:end] = blk_mask
+
+        if loss_mask is not None:
+            masked_indices[loss_mask == 0] = 0            
+
+        noisy_batch = torch.where(masked_indices, self.mask_token_id, input_ids)
+        return noisy_batch, masked_indices, p_mask
+    
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor]   = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        labels: Optional[torch.LongTensor]       = None,
+        split_len: Optional[int]                 = None,
+        past_key_values: Optional[Cache]         = None,
+        block_size: Optional[int]                = None,
+        block_diff_ppl: bool                     = False,
+        eps: float                               = 1e-3,
+        is_teacher: bool                        = False,
+        masked_indices: Optional[torch.Tensor]   = None,
+        p_mask: Optional[torch.Tensor]           = None,
+        loss_mask: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> CausalLMOutputWithPast:
+
+        batch_size, seq_len = input_ids.shape
+
+        if self.config.dlm_paradigm == 'bidirectional':
+            if labels is not None and torch.rand(1) < self.config.random_length_prob:
+                random_length = torch.randint(2, input_ids.shape[1] + 1, (1,))
+                input_ids = input_ids[:, :random_length]
+                labels = labels[:, :random_length]
+                
+                if attention_mask is not None:
+                    attention_mask = attention_mask[:, :random_length]
+                if position_ids is not None:
+                    position_ids = position_ids[:, :random_length]
+
+        elif self.config.dlm_paradigm == 'prefix_bidirectional':
+            if labels is not None and split_len is None:
+                if torch.rand(1) < self.config.random_length_prob:
+                    split_len = torch.randint(1, seq_len//64, (1,)).item() * 64  ## [64, seq_len] divisible by 64
+                else:
+                    split_len = int(seq_len * self.config.prefix_ratio)
+
+        elif self.config.dlm_paradigm == 'efficient_block_diff':
+            if labels is not None and block_size is None:
+                if torch.rand(1) < self.config.random_length_prob:
+                    block_size = torch.randint(1, 8, (1,)).item() * 4  ## [4, 32] divisible by 4
+                else:
+                    block_size = self.config.block_size
+
+        elif self.config.dlm_paradigm == 'block_diff':
+            if labels is not None and block_size is None:
+                if torch.rand(1) < self.config.random_length_prob:
+                    block_size = torch.randint(1, 8, (1,)).item() * 4  ## [4, 32] divisible by 4
+                else:
+                    block_size = self.config.block_size
+
+        if labels is not None and self.config.dlm_paradigm != 'autoregressive':
+            if masked_indices is not None:
+                assert p_mask is not None
+
+                if loss_mask is not None:
+                    masked_indices[loss_mask == 0] = 0
+
+                noisy_inputs = torch.where(masked_indices, self.mask_token_id, input_ids)
+
+            else:
+                if self.config.tok_mask_half_life_ratio is not None:
+                    noisy_inputs, masked_indices, p_mask = self.forward_process_exp(input_ids, eps=eps, block_size=block_size, half_life_ratio=self.config.tok_mask_half_life_ratio, loss_mask=loss_mask)
+                else:
+                    noisy_inputs, masked_indices, p_mask = self.forward_process(input_ids, eps=eps, block_size=block_size, loss_mask=loss_mask)
+
+        else:
+            noisy_inputs = input_ids
+            masked_indices = None
+            p_mask = None
+
+        if self.config.dlm_paradigm in ['prefix_bidirectional', 'efficient_block_diff', 'block_diff']:
+            for layer in self.encoder.layers:
+                if hasattr(layer.self_attn, 'set_attention_mode'):
+                    layer.self_attn.set_attention_mode(self.config.dlm_paradigm, prefix_len=split_len, block_size=block_size)
+
+        input_ids_len = noisy_inputs.shape[1]
+        if labels is not None and self.config.dlm_paradigm == 'block_diff':
+            if position_ids is None:
+                position_ids = torch.arange(input_ids_len, device=noisy_inputs.device).unsqueeze(0)
+            noisy_inputs = torch.cat([noisy_inputs, input_ids], dim=1)
+
+        if block_diff_ppl:
+            if position_ids is None:
+                position_ids = torch.arange(input_ids_len // 2, device=noisy_inputs.device).unsqueeze(0)
+
+        enc_out  = self.encoder(
+            past_key_values=past_key_values,
+            input_ids=noisy_inputs,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            is_training=(labels is not None) or (block_diff_ppl),
+            **kwargs,
+        )
+
+        logits = self.diffusion_head(enc_out.last_hidden_state)  # (batch, len_B, vocab)
+
+        if labels is not None and self.config.dlm_paradigm == 'block_diff':
+            logits = logits[:, :input_ids_len]
+
+        loss = None
+        if labels is not None:  
+            if self.config.dlm_paradigm == 'autoregressive':
+                shift_logits = logits[..., :-1, :].contiguous()
+                shift_labels = labels[..., 1:].contiguous()
+                
+                if loss_mask is None:
+                    loss_fct = CrossEntropyLoss()
+                    shift_logits = shift_logits.view(-1, shift_logits.size(-1))
+                    shift_labels = shift_labels.view(-1)
+                    loss = loss_fct(shift_logits, shift_labels)
+
+                else:
+                    loss_mask = loss_mask[..., 1:].contiguous()
+
+                    loss_fct = CrossEntropyLoss(reduction='none')
+                    shift_logits = shift_logits.view(-1, shift_logits.size(-1))
+                    shift_labels = shift_labels.view(-1)
+                    shift_labels = shift_labels.to(shift_logits.device)
+                    
+                    token_losses = loss_fct(shift_logits, shift_labels)
+                                    
+                    loss = token_losses[loss_mask].sum() / loss_mask.sum()
+
+            else:
+                # Handle DREAM vs LLADA style losses
+                if hasattr(self.config, 'dlm_type') and self.config.dlm_type == 'dream':
+                    logits = logits[..., :-1, :].contiguous()
+                    labels = labels[..., 1:].contiguous()
+                    masked_indices = masked_indices[:, 1:]
+                    p_mask = p_mask[:, 1:]
+
+                # Calculate token-wise cross entropy loss for masked positions in B
+                token_loss = torch.nn.functional.cross_entropy(
+                    logits[masked_indices], 
+                    labels[masked_indices], 
+                    reduction='none'
+                ) / p_mask[masked_indices]
+                
+                loss = token_loss.sum() / masked_indices.sum()
+
+        return CausalLMOutputWithPast(
+            loss=loss if not is_teacher else logits,
+            logits=logits,
+            past_key_values=enc_out.past_key_values,
+            hidden_states=None,
+            attentions=None,
+        )
+