Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

added_tokens.json

@@ -0,0 +1,24 @@
+{
+  "</tool_call>": 151658,
+  "<tool_call>": 151657,
+  "<|box_end|>": 151649,
+  "<|box_start|>": 151648,
+  "<|endoftext|>": 151643,
+  "<|file_sep|>": 151664,
+  "<|fim_middle|>": 151660,
+  "<|fim_pad|>": 151662,
+  "<|fim_prefix|>": 151659,
+  "<|fim_suffix|>": 151661,
+  "<|im_end|>": 151645,
+  "<|im_start|>": 151644,
+  "<|image_pad|>": 151655,
+  "<|object_ref_end|>": 151647,
+  "<|object_ref_start|>": 151646,
+  "<|quad_end|>": 151651,
+  "<|quad_start|>": 151650,
+  "<|repo_name|>": 151663,
+  "<|video_pad|>": 151656,
+  "<|vision_end|>": 151653,
+  "<|vision_pad|>": 151654,
+  "<|vision_start|>": 151652
+}

chat_template.jinja

@@ -0,0 +1,54 @@
+{%- if tools %}
+    {{- '<|im_start|>system\n' }}
+    {%- if messages[0]['role'] == 'system' %}
+        {{- messages[0]['content'] }}
+    {%- else %}
+        {{- 'You are a helpful assistant.' }}
+    {%- endif %}
+    {{- "\n\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
+{%- else %}
+    {%- if messages[0]['role'] == 'system' %}
+        {{- '<|im_start|>system\n' + messages[0]['content'] + '<|im_end|>\n' }}
+    {%- else %}
+        {{- '<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+{%- for message in messages %}
+    {%- if (message.role == "user") or (message.role == "system" and not loop.first) or (message.role == "assistant" and not message.tool_calls) %}
+        {{- '<|im_start|>' + message.role + '\n' + message.content + '<|im_end|>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {{- '<|im_start|>' + message.role }}
+        {%- if message.content %}
+            {{- '\n' + message.content }}
+        {%- endif %}
+        {%- for tool_call in message.tool_calls %}
+            {%- if tool_call.function is defined %}
+                {%- set tool_call = tool_call.function %}
+            {%- endif %}
+            {{- '\n<tool_call>\n{"name": "' }}
+            {{- tool_call.name }}
+            {{- '", "arguments": ' }}
+            {{- tool_call.arguments | tojson }}
+            {{- '}\n</tool_call>' }}
+        {%- endfor %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != "tool") %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- message.content }}
+        {{- '\n</tool_response>' }}
+        {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+{%- endif %}

config.json

@@ -0,0 +1,103 @@
+{
+  "architectures": [
+    "WeDLMForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoConfig": "configuration_wedlm.WeDLMConfig",
+    "AutoModel": "modeling_wedlm.WeDLMForCausalLM",
+    "AutoModelForCausalLM": "modeling_wedlm.WeDLMForCausalLM"
+  },
+  "dtype": "bfloat16",
+  "eos_token_id": 151645,
+  "flow_block_size": 64,
+  "flow_discretization": "argmax",
+  "flow_inference_steps": 8,
+  "flow_init_sigma": 1.0,
+  "flow_loss_weight": 1.0,
+  "flow_temperature": 1.0,
+  "flow_time_embedding_dim": 256,
+  "flow_time_embedding_max_period": 10000,
+  "flow_time_scale": 1000.0,
+  "flow_top_k": 0,
+  "flow_top_p": 1.0,
+  "flow_train_block_size": 64,
+  "flow_train_mask_ratio": 0.15,
+  "flow_train_min_target_tokens": 1,
+  "flow_train_strategy": "suffix_block",
+  "flow_type": "rectified_flow",
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 4096,
+  "initializer_range": 0.02,
+  "intermediate_size": 12288,
+  "layer_types": [
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention"
+  ],
+  "mask_token_id": null,
+  "max_position_embeddings": 16384,
+  "max_window_layers": 48,
+  "model_type": "wedlm",
+  "num_attention_heads": 32,
+  "num_hidden_layers": 48,
+  "num_key_value_heads": 8,
+  "pad_token_id": 151643,
+  "qk_norm": true,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 1000000.0,
+  "sliding_window": null,
+  "tie_word_embeddings": false,
+  "transformers_version": "4.57.1",
+  "use_cache": false,
+  "use_sliding_window": false,
+  "vocab_size": 151936
+}

configuration_wedlm.py

@@ -0,0 +1,248 @@
+# coding=utf-8
+# Copyright 2024 The WeDLM team 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.
+"""WeDLM 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 WeDLMConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`WeDLMModel`]. It is used to instantiate a
+    WeDLM model according to the specified arguments, defining the model architecture.
+
+    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 WeDLM model. Defines the number of different tokens that can be represented by the
+            `inputs_ids` passed when calling [`WeDLMModel`]
+        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 decoder.
+        num_attention_heads (`int`, *optional*, defaults to 32):
+            Number of attention heads for each attention layer in the Transformer decoder.
+        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.
+        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.
+        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 using full attention.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+        attention_bias (`bool`, *optional*, defaults to `True`):
+            Whether to use bias in QKV projections. Set to `True` for Qwen2.5 compatibility,
+            `False` for Qwen3 compatibility.
+        qk_norm (`bool`, *optional*, defaults to `False`):
+            Whether to use QK normalization. Set to `True` for Qwen3 compatibility.
+        head_dim (`int`, *optional*):
+            The dimension of each attention head. If not specified, defaults to hidden_size // num_attention_heads.
+
+        mask_token_id (`int`, *optional*):
+            Token id used by upstream discrete diffusion-style block decoding as the MASK placeholder.
+            (Kept for backward compatibility; not used by Flow-Matching decoding.)
+
+        ----------------------------------------------------------------------
+        Flow Matching / Rectified Flow (primary training & decoding)
+        ----------------------------------------------------------------------
+        flow_type (`str`, *optional*, defaults to `"rectified_flow"`):
+            Flow family. Currently `"rectified_flow"` is implemented (straight-line probability paths).
+        flow_time_embedding_dim (`int`, *optional*, defaults to 256):
+            Dimensionality of the sinusoidal timestep embedding before projection.
+        flow_time_embedding_max_period (`int`, *optional*, defaults to 10000):
+            Max period for sinusoidal timestep embeddings.
+        flow_time_scale (`float`, *optional*, defaults to 1000.0):
+            Scalar multiplier applied to normalized timesteps (t in [0,1]) before sinusoidal embedding.
+        flow_init_sigma (`float`, *optional*, defaults to 1.0):
+            Standard deviation of the Gaussian base distribution used for flow sampling/training in embedding space.
+        flow_loss_weight (`float`, *optional*, defaults to 1.0):
+            Weight for Flow Matching loss.
+
+        flow_train_strategy (`str`, *optional*, defaults to `"suffix_block"`):
+            Training target selection strategy when `flow_target_mask` is not provided to `forward`.
+            - `"suffix_block"`: use a suffix contiguous block as flow targets (aligned with block decoding).
+            - `"random"`: random Bernoulli masking by `flow_train_mask_ratio`.
+        flow_train_block_size (`int`, *optional*, defaults to 64):
+            Target block size for `"suffix_block"` strategy.
+        flow_train_min_target_tokens (`int`, *optional*, defaults to 1):
+            Minimum number of flow targets per sample (after filtering pads / ignored labels).
+        flow_train_mask_ratio (`float`, *optional*, defaults to 0.15):
+            Target masking ratio for `"random"` strategy.
+
+        flow_inference_steps (`int`, *optional*, defaults to 8):
+            Default number of Euler steps per block for `generate_wedlm`.
+        flow_block_size (`int`, *optional*, defaults to 64):
+            Default block size (new tokens per block) for `generate_wedlm`.
+        flow_discretization (`str`, *optional*, defaults to `"argmax"`):
+            Discretization method from final embeddings to token IDs:
+            - `"argmax"`: argmax over vocabulary logits.
+            - `"sample"`: temperature + (top-k/top-p) sampling from vocabulary logits.
+        flow_temperature (`float`, *optional*, defaults to 1.0):
+            Temperature for `"sample"` discretization.
+        flow_top_p (`float`, *optional*, defaults to 1.0):
+            Top-p (nucleus) sampling for `"sample"` discretization.
+        flow_top_k (`int`, *optional*, defaults to 0):
+            Top-k sampling for `"sample"` discretization. 0 disables top-k filtering.
+    """
+
+    model_type = "wedlm"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    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,
+        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,
+        use_sliding_window=False,
+        sliding_window=4096,
+        max_window_layers=28,
+        attention_dropout=0.0,
+        attention_bias=True,
+        qk_norm=False,
+        head_dim=None,
+        mask_token_id=None,
+        # Flow Matching / Rectified Flow (primary)
+        flow_type="rectified_flow",
+        flow_time_embedding_dim=256,
+        flow_time_embedding_max_period=10000,
+        flow_time_scale=1000.0,
+        flow_init_sigma=1.0,
+        flow_loss_weight=1.0,
+        flow_train_strategy="suffix_block",
+        flow_train_block_size=64,
+        flow_train_min_target_tokens=1,
+        flow_train_mask_ratio=0.15,
+        flow_inference_steps=8,
+        flow_block_size=64,
+        flow_discretization="argmax",
+        flow_temperature=1.0,
+        flow_top_p=1.0,
+        flow_top_k=0,
+        **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 if self.use_sliding_window else None
+        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.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_dropout = attention_dropout
+        self.attention_bias = attention_bias
+        self.qk_norm = qk_norm
+        self.mask_token_id = mask_token_id
+
+        if head_dim is None:
+            self.head_dim = hidden_size // num_attention_heads
+        else:
+            self.head_dim = head_dim
+
+        # ----------------------------
+        # Flow Matching configuration
+        # ----------------------------
+        self.flow_type = flow_type
+        self.flow_time_embedding_dim = flow_time_embedding_dim
+        self.flow_time_embedding_max_period = flow_time_embedding_max_period
+        self.flow_time_scale = flow_time_scale
+        self.flow_init_sigma = flow_init_sigma
+        self.flow_loss_weight = flow_loss_weight
+
+        self.flow_train_strategy = flow_train_strategy
+        self.flow_train_block_size = flow_train_block_size
+        self.flow_train_min_target_tokens = flow_train_min_target_tokens
+        self.flow_train_mask_ratio = flow_train_mask_ratio
+
+        self.flow_inference_steps = flow_inference_steps
+        self.flow_block_size = flow_block_size
+        self.flow_discretization = flow_discretization
+        self.flow_temperature = flow_temperature
+        self.flow_top_p = flow_top_p
+        self.flow_top_k = flow_top_k
+
+        # 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)
+
+        # Generate layer_types based on sliding window configuration
+        self.layer_types = [
+            "sliding_attention"
+            if self.sliding_window is not None and i >= self.max_window_layers
+            else "full_attention"
+            for i in range(self.num_hidden_layers)
+        ]
+
+        super().__init__(
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+
+__all__ = ["WeDLMConfig"]

generation_config.json

@@ -0,0 +1,10 @@
+{
+  "_from_model_config": true,
+  "eos_token_id": [
+    151645,
+    151643
+  ],
+  "pad_token_id": 151643,
+  "transformers_version": "4.57.1",
+  "use_cache": false
+}

model-00001-of-00005.safetensors

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

model-00002-of-00005.safetensors

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

model-00003-of-00005.safetensors

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

model-00004-of-00005.safetensors

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

model-00005-of-00005.safetensors

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

model.safetensors.index.json

@@ -0,0 +1,544 @@
+{
+  "metadata": {
+    "total_parameters": 10540703744,
+    "total_size": 21081407488
+  },
+  "weight_map": {
+    "flow_head.weight": "model-00005-of-00005.safetensors",
+    "flow_time_embed.linear_1.bias": "model-00005-of-00005.safetensors",
+    "flow_time_embed.linear_1.weight": "model-00005-of-00005.safetensors",
+    "flow_time_embed.linear_2.bias": "model-00005-of-00005.safetensors",
+    "flow_time_embed.linear_2.weight": "model-00005-of-00005.safetensors",
+    "lm_head.weight": "model-00005-of-00005.safetensors",
+    "model.embed_tokens.weight": "model-00001-of-00005.safetensors",
+    "model.layers.0.input_layernorm.weight": "model-00001-of-00005.safetensors",
+    "model.layers.0.mlp.down_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.0.mlp.gate_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.0.mlp.up_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.0.post_attention_layernorm.weight": "model-00001-of-00005.safetensors",
+    "model.layers.0.self_attn.k_norm.weight": "model-00001-of-00005.safetensors",
+    "model.layers.0.self_attn.k_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.0.self_attn.o_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.0.self_attn.q_norm.weight": "model-00001-of-00005.safetensors",
+    "model.layers.0.self_attn.q_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.0.self_attn.v_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.1.input_layernorm.weight": "model-00001-of-00005.safetensors",
+    "model.layers.1.mlp.down_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.1.mlp.gate_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.1.mlp.up_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.1.post_attention_layernorm.weight": "model-00001-of-00005.safetensors",
+    "model.layers.1.self_attn.k_norm.weight": "model-00001-of-00005.safetensors",
+    "model.layers.1.self_attn.k_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.1.self_attn.o_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.1.self_attn.q_norm.weight": "model-00001-of-00005.safetensors",
+    "model.layers.1.self_attn.q_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.1.self_attn.v_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.10.input_layernorm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.10.mlp.down_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.10.mlp.gate_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.10.mlp.up_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.10.post_attention_layernorm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.10.self_attn.k_norm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.10.self_attn.k_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.10.self_attn.o_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.10.self_attn.q_norm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.10.self_attn.q_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.10.self_attn.v_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.11.input_layernorm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.11.mlp.down_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.11.mlp.gate_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.11.mlp.up_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.11.post_attention_layernorm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.11.self_attn.k_norm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.11.self_attn.k_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.11.self_attn.o_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.11.self_attn.q_norm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.11.self_attn.q_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.11.self_attn.v_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.12.input_layernorm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.12.mlp.down_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.12.mlp.gate_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.12.mlp.up_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.12.post_attention_layernorm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.12.self_attn.k_norm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.12.self_attn.k_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.12.self_attn.o_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.12.self_attn.q_norm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.12.self_attn.q_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.12.self_attn.v_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.13.input_layernorm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.13.mlp.down_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.13.mlp.gate_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.13.mlp.up_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.13.post_attention_layernorm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.13.self_attn.k_norm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.13.self_attn.k_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.13.self_attn.o_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.13.self_attn.q_norm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.13.self_attn.q_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.13.self_attn.v_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.14.input_layernorm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.14.mlp.down_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.14.mlp.gate_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.14.mlp.up_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.14.post_attention_layernorm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.14.self_attn.k_norm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.14.self_attn.k_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.14.self_attn.o_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.14.self_attn.q_norm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.14.self_attn.q_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.14.self_attn.v_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.15.input_layernorm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.15.mlp.down_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.15.mlp.gate_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.15.mlp.up_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.15.post_attention_layernorm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.15.self_attn.k_norm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.15.self_attn.k_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.15.self_attn.o_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.15.self_attn.q_norm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.15.self_attn.q_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.15.self_attn.v_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.16.input_layernorm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.16.mlp.down_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.16.mlp.gate_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.16.mlp.up_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.16.post_attention_layernorm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.16.self_attn.k_norm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.16.self_attn.k_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.16.self_attn.o_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.16.self_attn.q_norm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.16.self_attn.q_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.16.self_attn.v_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.17.input_layernorm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.17.mlp.down_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.17.mlp.gate_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.17.mlp.up_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.17.post_attention_layernorm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.17.self_attn.k_norm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.17.self_attn.k_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.17.self_attn.o_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.17.self_attn.q_norm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.17.self_attn.q_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.17.self_attn.v_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.18.input_layernorm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.18.mlp.down_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.18.mlp.gate_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.18.mlp.up_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.18.post_attention_layernorm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.18.self_attn.k_norm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.18.self_attn.k_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.18.self_attn.o_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.18.self_attn.q_norm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.18.self_attn.q_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.18.self_attn.v_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.19.input_layernorm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.19.mlp.down_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.19.mlp.gate_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.19.mlp.up_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.19.post_attention_layernorm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.19.self_attn.k_norm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.19.self_attn.k_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.19.self_attn.o_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.19.self_attn.q_norm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.19.self_attn.q_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.19.self_attn.v_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.2.input_layernorm.weight": "model-00001-of-00005.safetensors",
+    "model.layers.2.mlp.down_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.2.mlp.gate_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.2.mlp.up_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.2.post_attention_layernorm.weight": "model-00001-of-00005.safetensors",
+    "model.layers.2.self_attn.k_norm.weight": "model-00001-of-00005.safetensors",
+    "model.layers.2.self_attn.k_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.2.self_attn.o_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.2.self_attn.q_norm.weight": "model-00001-of-00005.safetensors",
+    "model.layers.2.self_attn.q_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.2.self_attn.v_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.20.input_layernorm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.20.mlp.down_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.20.mlp.gate_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.20.mlp.up_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.20.post_attention_layernorm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.20.self_attn.k_norm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.20.self_attn.k_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.20.self_attn.o_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.20.self_attn.q_norm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.20.self_attn.q_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.20.self_attn.v_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.21.input_layernorm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.21.mlp.down_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.21.mlp.gate_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.21.mlp.up_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.21.post_attention_layernorm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.21.self_attn.k_norm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.21.self_attn.k_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.21.self_attn.o_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.21.self_attn.q_norm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.21.self_attn.q_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.21.self_attn.v_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.22.input_layernorm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.22.mlp.down_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.22.mlp.gate_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.22.mlp.up_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.22.post_attention_layernorm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.22.self_attn.k_norm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.22.self_attn.k_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.22.self_attn.o_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.22.self_attn.q_norm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.22.self_attn.q_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.22.self_attn.v_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.23.input_layernorm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.23.mlp.down_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.23.mlp.gate_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.23.mlp.up_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.23.post_attention_layernorm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.23.self_attn.k_norm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.23.self_attn.k_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.23.self_attn.o_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.23.self_attn.q_norm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.23.self_attn.q_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.23.self_attn.v_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.24.input_layernorm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.24.mlp.down_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.24.mlp.gate_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.24.mlp.up_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.24.post_attention_layernorm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.24.self_attn.k_norm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.24.self_attn.k_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.24.self_attn.o_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.24.self_attn.q_norm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.24.self_attn.q_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.24.self_attn.v_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.25.input_layernorm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.25.mlp.down_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.25.mlp.gate_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.25.mlp.up_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.25.post_attention_layernorm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.25.self_attn.k_norm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.25.self_attn.k_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.25.self_attn.o_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.25.self_attn.q_norm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.25.self_attn.q_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.25.self_attn.v_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.26.input_layernorm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.26.mlp.down_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.26.mlp.gate_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.26.mlp.up_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.26.post_attention_layernorm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.26.self_attn.k_norm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.26.self_attn.k_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.26.self_attn.o_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.26.self_attn.q_norm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.26.self_attn.q_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.26.self_attn.v_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.27.input_layernorm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.27.mlp.down_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.27.mlp.gate_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.27.mlp.up_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.27.post_attention_layernorm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.27.self_attn.k_norm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.27.self_attn.k_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.27.self_attn.o_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.27.self_attn.q_norm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.27.self_attn.q_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.27.self_attn.v_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.28.input_layernorm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.28.mlp.down_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.28.mlp.gate_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.28.mlp.up_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.28.post_attention_layernorm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.28.self_attn.k_norm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.28.self_attn.k_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.28.self_attn.o_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.28.self_attn.q_norm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.28.self_attn.q_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.28.self_attn.v_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.29.input_layernorm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.29.mlp.down_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.29.mlp.gate_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.29.mlp.up_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.29.post_attention_layernorm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.29.self_attn.k_norm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.29.self_attn.k_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.29.self_attn.o_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.29.self_attn.q_norm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.29.self_attn.q_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.29.self_attn.v_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.3.input_layernorm.weight": "model-00001-of-00005.safetensors",
+    "model.layers.3.mlp.down_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.3.mlp.gate_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.3.mlp.up_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.3.post_attention_layernorm.weight": "model-00001-of-00005.safetensors",
+    "model.layers.3.self_attn.k_norm.weight": "model-00001-of-00005.safetensors",
+    "model.layers.3.self_attn.k_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.3.self_attn.o_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.3.self_attn.q_norm.weight": "model-00001-of-00005.safetensors",
+    "model.layers.3.self_attn.q_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.3.self_attn.v_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.30.input_layernorm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.30.mlp.down_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.30.mlp.gate_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.30.mlp.up_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.30.post_attention_layernorm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.30.self_attn.k_norm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.30.self_attn.k_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.30.self_attn.o_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.30.self_attn.q_norm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.30.self_attn.q_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.30.self_attn.v_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.31.input_layernorm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.31.mlp.down_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.31.mlp.gate_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.31.mlp.up_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.31.post_attention_layernorm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.31.self_attn.k_norm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.31.self_attn.k_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.31.self_attn.o_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.31.self_attn.q_norm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.31.self_attn.q_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.31.self_attn.v_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.32.input_layernorm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.32.mlp.down_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.32.mlp.gate_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.32.mlp.up_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.32.post_attention_layernorm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.32.self_attn.k_norm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.32.self_attn.k_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.32.self_attn.o_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.32.self_attn.q_norm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.32.self_attn.q_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.32.self_attn.v_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.33.input_layernorm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.33.mlp.down_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.33.mlp.gate_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.33.mlp.up_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.33.post_attention_layernorm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.33.self_attn.k_norm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.33.self_attn.k_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.33.self_attn.o_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.33.self_attn.q_norm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.33.self_attn.q_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.33.self_attn.v_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.34.input_layernorm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.34.mlp.down_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.34.mlp.gate_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.34.mlp.up_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.34.post_attention_layernorm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.34.self_attn.k_norm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.34.self_attn.k_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.34.self_attn.o_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.34.self_attn.q_norm.weight": "model-00003-of-00005.safetensors",
+    "model.layers.34.self_attn.q_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.34.self_attn.v_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.35.input_layernorm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.35.mlp.down_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.35.mlp.gate_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.35.mlp.up_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.35.post_attention_layernorm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.35.self_attn.k_norm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.35.self_attn.k_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.35.self_attn.o_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.35.self_attn.q_norm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.35.self_attn.q_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.35.self_attn.v_proj.weight": "model-00003-of-00005.safetensors",
+    "model.layers.36.input_layernorm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.36.mlp.down_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.36.mlp.gate_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.36.mlp.up_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.36.post_attention_layernorm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.36.self_attn.k_norm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.36.self_attn.k_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.36.self_attn.o_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.36.self_attn.q_norm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.36.self_attn.q_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.36.self_attn.v_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.37.input_layernorm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.37.mlp.down_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.37.mlp.gate_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.37.mlp.up_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.37.post_attention_layernorm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.37.self_attn.k_norm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.37.self_attn.k_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.37.self_attn.o_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.37.self_attn.q_norm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.37.self_attn.q_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.37.self_attn.v_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.38.input_layernorm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.38.mlp.down_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.38.mlp.gate_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.38.mlp.up_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.38.post_attention_layernorm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.38.self_attn.k_norm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.38.self_attn.k_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.38.self_attn.o_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.38.self_attn.q_norm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.38.self_attn.q_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.38.self_attn.v_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.39.input_layernorm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.39.mlp.down_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.39.mlp.gate_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.39.mlp.up_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.39.post_attention_layernorm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.39.self_attn.k_norm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.39.self_attn.k_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.39.self_attn.o_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.39.self_attn.q_norm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.39.self_attn.q_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.39.self_attn.v_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.4.input_layernorm.weight": "model-00001-of-00005.safetensors",
+    "model.layers.4.mlp.down_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.4.mlp.gate_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.4.mlp.up_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.4.post_attention_layernorm.weight": "model-00001-of-00005.safetensors",
+    "model.layers.4.self_attn.k_norm.weight": "model-00001-of-00005.safetensors",
+    "model.layers.4.self_attn.k_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.4.self_attn.o_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.4.self_attn.q_norm.weight": "model-00001-of-00005.safetensors",
+    "model.layers.4.self_attn.q_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.4.self_attn.v_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.40.input_layernorm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.40.mlp.down_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.40.mlp.gate_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.40.mlp.up_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.40.post_attention_layernorm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.40.self_attn.k_norm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.40.self_attn.k_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.40.self_attn.o_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.40.self_attn.q_norm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.40.self_attn.q_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.40.self_attn.v_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.41.input_layernorm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.41.mlp.down_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.41.mlp.gate_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.41.mlp.up_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.41.post_attention_layernorm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.41.self_attn.k_norm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.41.self_attn.k_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.41.self_attn.o_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.41.self_attn.q_norm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.41.self_attn.q_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.41.self_attn.v_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.42.input_layernorm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.42.mlp.down_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.42.mlp.gate_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.42.mlp.up_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.42.post_attention_layernorm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.42.self_attn.k_norm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.42.self_attn.k_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.42.self_attn.o_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.42.self_attn.q_norm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.42.self_attn.q_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.42.self_attn.v_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.43.input_layernorm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.43.mlp.down_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.43.mlp.gate_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.43.mlp.up_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.43.post_attention_layernorm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.43.self_attn.k_norm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.43.self_attn.k_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.43.self_attn.o_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.43.self_attn.q_norm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.43.self_attn.q_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.43.self_attn.v_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.44.input_layernorm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.44.mlp.down_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.44.mlp.gate_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.44.mlp.up_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.44.post_attention_layernorm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.44.self_attn.k_norm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.44.self_attn.k_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.44.self_attn.o_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.44.self_attn.q_norm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.44.self_attn.q_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.44.self_attn.v_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.45.input_layernorm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.45.mlp.down_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.45.mlp.gate_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.45.mlp.up_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.45.post_attention_layernorm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.45.self_attn.k_norm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.45.self_attn.k_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.45.self_attn.o_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.45.self_attn.q_norm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.45.self_attn.q_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.45.self_attn.v_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.46.input_layernorm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.46.mlp.down_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.46.mlp.gate_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.46.mlp.up_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.46.post_attention_layernorm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.46.self_attn.k_norm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.46.self_attn.k_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.46.self_attn.o_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.46.self_attn.q_norm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.46.self_attn.q_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.46.self_attn.v_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.47.input_layernorm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.47.mlp.down_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.47.mlp.gate_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.47.mlp.up_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.47.post_attention_layernorm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.47.self_attn.k_norm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.47.self_attn.k_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.47.self_attn.o_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.47.self_attn.q_norm.weight": "model-00004-of-00005.safetensors",
+    "model.layers.47.self_attn.q_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.47.self_attn.v_proj.weight": "model-00004-of-00005.safetensors",
+    "model.layers.5.input_layernorm.weight": "model-00001-of-00005.safetensors",
+    "model.layers.5.mlp.down_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.5.mlp.gate_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.5.mlp.up_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.5.post_attention_layernorm.weight": "model-00001-of-00005.safetensors",
+    "model.layers.5.self_attn.k_norm.weight": "model-00001-of-00005.safetensors",
+    "model.layers.5.self_attn.k_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.5.self_attn.o_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.5.self_attn.q_norm.weight": "model-00001-of-00005.safetensors",
+    "model.layers.5.self_attn.q_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.5.self_attn.v_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.6.input_layernorm.weight": "model-00001-of-00005.safetensors",
+    "model.layers.6.mlp.down_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.6.mlp.gate_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.6.mlp.up_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.6.post_attention_layernorm.weight": "model-00001-of-00005.safetensors",
+    "model.layers.6.self_attn.k_norm.weight": "model-00001-of-00005.safetensors",
+    "model.layers.6.self_attn.k_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.6.self_attn.o_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.6.self_attn.q_norm.weight": "model-00001-of-00005.safetensors",
+    "model.layers.6.self_attn.q_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.6.self_attn.v_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.7.input_layernorm.weight": "model-00001-of-00005.safetensors",
+    "model.layers.7.mlp.down_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.7.mlp.gate_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.7.mlp.up_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.7.post_attention_layernorm.weight": "model-00001-of-00005.safetensors",
+    "model.layers.7.self_attn.k_norm.weight": "model-00001-of-00005.safetensors",
+    "model.layers.7.self_attn.k_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.7.self_attn.o_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.7.self_attn.q_norm.weight": "model-00001-of-00005.safetensors",
+    "model.layers.7.self_attn.q_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.7.self_attn.v_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.8.input_layernorm.weight": "model-00001-of-00005.safetensors",
+    "model.layers.8.mlp.down_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.8.mlp.gate_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.8.mlp.up_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.8.post_attention_layernorm.weight": "model-00001-of-00005.safetensors",
+    "model.layers.8.self_attn.k_norm.weight": "model-00001-of-00005.safetensors",
+    "model.layers.8.self_attn.k_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.8.self_attn.o_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.8.self_attn.q_norm.weight": "model-00001-of-00005.safetensors",
+    "model.layers.8.self_attn.q_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.8.self_attn.v_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.9.input_layernorm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.9.mlp.down_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.9.mlp.gate_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.9.mlp.up_proj.weight": "model-00002-of-00005.safetensors",
+    "model.layers.9.post_attention_layernorm.weight": "model-00002-of-00005.safetensors",
+    "model.layers.9.self_attn.k_norm.weight": "model-00001-of-00005.safetensors",
+    "model.layers.9.self_attn.k_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.9.self_attn.o_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.9.self_attn.q_norm.weight": "model-00001-of-00005.safetensors",
+    "model.layers.9.self_attn.q_proj.weight": "model-00001-of-00005.safetensors",
+    "model.layers.9.self_attn.v_proj.weight": "model-00001-of-00005.safetensors",
+    "model.norm.weight": "model-00004-of-00005.safetensors"
+  }
+}

modeling_wedlm.py

@@ -0,0 +1,1413 @@
+# coding=utf-8
+# Copyright 2024 The WeDLM team 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.
+"""PyTorch WeDLM model."""
+
+from typing import Optional, Tuple, Union, Dict, List, Callable
+
+import math
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from transformers import PreTrainedModel, GenerationMixin
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers.processing_utils import Unpack
+from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple
+from transformers.utils.generic import check_model_inputs
+from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask
+from transformers.modeling_layers import GradientCheckpointingLayer
+from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+# Import attention-related utilities
+try:
+    from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
+except ImportError:
+    FlashAttentionKwargs = dict
+
+from .configuration_wedlm import WeDLMConfig
+
+import logging
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)
+
+
+# ============================================================================
+# Flow Matching / Rectified Flow helpers
+# ============================================================================
+
+class WeDLMFlowTimeEmbedding(nn.Module):
+    """Sinusoidal timestep embedding + MLP, used to condition Flow Matching / Rectified Flow.
+
+    The module is intentionally lightweight and conditions the velocity field on continuous timesteps.
+    Timesteps are assumed to be normalized to [0, 1] (float). Internally, a configurable scale is applied
+    before sinusoidal features are computed.
+    """
+
+    def __init__(self, config: WeDLMConfig):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.time_embed_dim = int(getattr(config, "flow_time_embedding_dim", 256))
+        self.max_period = int(getattr(config, "flow_time_embedding_max_period", 10000))
+        self.time_scale = float(getattr(config, "flow_time_scale", 1000.0))
+
+        self.linear_1 = nn.Linear(self.time_embed_dim, self.hidden_size)
+        self.act = nn.SiLU()
+        self.linear_2 = nn.Linear(self.hidden_size, self.hidden_size)
+
+    @staticmethod
+    def _sinusoidal_embedding(timesteps: torch.Tensor, dim: int, max_period: int) -> torch.Tensor:
+        """Create sinusoidal timestep embeddings.
+
+        timesteps: (batch,) float tensor.
+        Returns: (batch, dim) float tensor.
+        """
+        if timesteps.ndim != 1:
+            timesteps = timesteps.view(-1)
+        half = dim // 2
+        device = timesteps.device
+        dtype = torch.float32
+
+        freqs = torch.exp(
+            -math.log(max_period) * torch.arange(0, half, device=device, dtype=dtype) / max(half, 1)
+        )
+        args = timesteps.to(dtype)[:, None] * freqs[None]
+        emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2 == 1:
+            emb = torch.cat([emb, torch.zeros((emb.shape[0], 1), device=device, dtype=dtype)], dim=-1)
+        return emb
+
+    def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
+        # timesteps expected in [0,1]; scale to a more typical diffusion timestep range.
+        t = timesteps.to(dtype=torch.float32)
+        if t.ndim == 0:
+            t = t[None]
+        if t.ndim != 1:
+            t = t.view(-1)
+        t = t.clamp(0.0, 1.0) * self.time_scale
+
+        emb = self._sinusoidal_embedding(t, self.time_embed_dim, self.max_period)
+
+        # NOTE: In bf16/fp16 training (and especially under PEFT/LoRA), the wrapped Linear's base weights
+        # can be low-precision (e.g. bfloat16) while the sinusoidal features are float32.
+        # Torch's F.linear requires the input and weight dtypes to match, so we explicitly cast here
+        # to the *base* layer's weight dtype.
+        base_linear_1 = getattr(self.linear_1, "base_layer", self.linear_1)
+        w1 = getattr(base_linear_1, "weight", None)
+        if w1 is not None:
+            emb = emb.to(dtype=w1.dtype)
+        emb = self.linear_1(emb)
+        emb = self.act(emb)
+
+        base_linear_2 = getattr(self.linear_2, "base_layer", self.linear_2)
+        w2 = getattr(base_linear_2, "weight", None)
+        if w2 is not None:
+            emb = emb.to(dtype=w2.dtype)
+        emb = self.linear_2(emb)
+        return emb
+
+
+# ============================================================================
+# ============================================================================
+# Core Components (self-contained, no Qwen2 dependency)
+# ============================================================================
+
+class WeDLMMLP(nn.Module):
+    """WeDLM MLP module with SwiGLU activation."""
+    
+    def __init__(self, config: WeDLMConfig):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+        return down_proj
+
+
+class WeDLMRMSNorm(nn.Module):
+    """WeDLM RMSNorm, equivalent to T5LayerNorm."""
+    
+    def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+    def extra_repr(self) -> str:
+        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+class WeDLMRotaryEmbedding(nn.Module):
+    """WeDLM Rotary Position Embedding."""
+    
+    def __init__(self, config: WeDLMConfig, device=None):
+        super().__init__()
+        # Determine rope_type from config
+        if hasattr(config, "rope_scaling") and isinstance(config.rope_scaling, dict):
+            self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type", "default"))
+        else:
+            self.rope_type = "default"
+        
+        self.max_seq_len_cached = config.max_position_embeddings
+        self.original_max_seq_len = config.max_position_embeddings
+        self.config = config
+        
+        # Get initialization function
+        if self.rope_type == "default":
+            inv_freq, self.attention_scaling = self._compute_default_rope_parameters(config, device)
+        else:
+            rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+            inv_freq, self.attention_scaling = rope_init_fn(config, device)
+        
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.original_inv_freq = self.inv_freq
+
+    @staticmethod
+    def _compute_default_rope_parameters(
+        config: WeDLMConfig,
+        device: Optional[torch.device] = None,
+    ) -> Tuple[torch.Tensor, float]:
+        """
+        Computes the inverse frequencies for default RoPE.
+        
+        Args:
+            config: Model configuration
+            device: Device to place the tensors on
+            
+        Returns:
+            Tuple of (inv_freq tensor, attention_scaling factor)
+        """
+        base = config.rope_theta
+        dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+        
+        # Compute the inverse frequencies
+        inv_freq = 1.0 / (
+            base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
+        )
+        attention_factor = 1.0
+        return inv_freq, attention_factor
+
+    @torch.no_grad()
+    def forward(self, x: torch.Tensor, position_ids: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Compute rotary position embeddings.
+        
+        Args:
+            x: Input tensor, used for dtype and device
+            position_ids: Position indices
+            
+        Returns:
+            Tuple of (cos, sin) tensors
+        """
+        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"
+        
+        # Force float32 computation for numerical stability
+        with torch.amp.autocast(device_type=device_type, enabled=False):
+            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)
+
+
+# ============================================================================
+# Attention Utilities
+# ============================================================================
+
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(
+    q: torch.Tensor, 
+    k: torch.Tensor, 
+    cos: torch.Tensor, 
+    sin: torch.Tensor, 
+    position_ids: Optional[torch.Tensor] = None, 
+    unsqueeze_dim: int = 1
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Applies Rotary Position Embedding to the query and key tensors."""
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    Repeats key/value heads to match the number of query heads (for GQA).
+    
+    Equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep).
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+def eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: Optional[torch.Tensor],
+    scaling: float,
+    dropout: float = 0.0,
+    **kwargs,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Eager (standard) attention implementation."""
+    key_states = repeat_kv(key, module.num_key_value_groups)
+    value_states = repeat_kv(value, module.num_key_value_groups)
+
+    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+
+    if attention_mask is not None:
+        causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
+        attn_weights = attn_weights + causal_mask
+
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+    attn_output = torch.matmul(attn_weights, value_states)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+
+    return attn_output, attn_weights
+
+
+# ============================================================================
+# Attention Layer
+# ============================================================================
+
+class WeDLMAttention(nn.Module):
+    """
+    WeDLM Attention module.
+    
+    Supports both:
+    - Qwen2.5 style: with QKV bias, no QK Norm
+    - Qwen3 style: configurable QKV bias, with QK Norm
+    """
+
+    def __init__(self, config: WeDLMConfig, layer_idx: int):
+        super().__init__()
+        self.layer_type = config.layer_types[layer_idx] if hasattr(config, "layer_types") else None
+        self.config = config
+        self.layer_idx = layer_idx
+        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+        self.scaling = self.head_dim ** -0.5
+        self.attention_dropout = config.attention_dropout
+        self.is_causal = True
+        
+        # Support configurable attention_bias (Qwen2.5: True, Qwen3: False by default)
+        attention_bias = getattr(config, "attention_bias", True)
+        
+        self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=attention_bias)
+        self.k_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=attention_bias)
+        self.v_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=attention_bias)
+        self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False)
+        
+        # Support optional QK Norm (Qwen3 feature)
+        self.qk_norm = getattr(config, "qk_norm", False)
+        if self.qk_norm:
+            self.q_norm = WeDLMRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+            self.k_norm = WeDLMRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        
+        self.sliding_window = config.sliding_window if self.layer_type == "sliding_attention" else None
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: Tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[torch.Tensor],
+        past_key_values: Optional[Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        if self.qk_norm:
+            # Qwen3 style: apply norm after projection, before transpose
+            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)
+        else:
+            # Qwen2 style: no norm
+            query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+            key_states = 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
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        if past_key_values 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_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+        # Select attention implementation
+        attention_interface: Callable = eager_attention_forward
+        if self.config._attn_implementation != "eager" and self.config._attn_implementation in ALL_ATTENTION_FUNCTIONS:
+            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            sliding_window=self.sliding_window,
+            **kwargs,
+        )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+
+# ============================================================================
+# Decoder Layer
+# ============================================================================
+
+class WeDLMDecoderLayer(GradientCheckpointingLayer):
+    """WeDLM Decoder Layer with pre-norm architecture."""
+    
+    def __init__(self, config: WeDLMConfig, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+
+        self.self_attn = WeDLMAttention(config=config, layer_idx=layer_idx)
+        self.mlp = WeDLMMLP(config)
+        self.input_layernorm = WeDLMRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = WeDLMRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.attention_type = config.layer_types[layer_idx]
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        output_attentions: Optional[bool] = False,
+        use_cache: Optional[bool] = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
+        """
+        Args:
+            hidden_states: Input tensor of shape `(batch, seq_len, embed_dim)`
+            attention_mask: Attention mask of size `(batch, sequence_length)`
+            position_ids: Position indices
+            past_key_values: Cached past key and value projection states
+            output_attentions: Whether to return attention weights
+            use_cache: Whether to use KV cache
+            cache_position: Position in the cache
+            position_embeddings: Tuple of (cos, sin) for rotary embeddings
+        """
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        
+        # Self Attention
+        hidden_states, self_attn_weights = self.self_attn(
+            hidden_states=hidden_states,
+            position_embeddings=position_embeddings,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            cache_position=cache_position,
+            **kwargs,
+        )
+        hidden_states = residual + hidden_states
+
+        # Feed Forward
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        outputs = (hidden_states,)
+
+        if output_attentions:
+            outputs += (self_attn_weights,)
+
+        return outputs
+
+
+# ============================================================================
+# Model Classes
+# ============================================================================
+
+@auto_docstring
+class WeDLMPreTrainedModel(PreTrainedModel):
+    """Base class for WeDLM models."""
+    
+    config_class = WeDLMConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["WeDLMDecoderLayer"]
+    _skip_keys_device_placement = ["past_key_values"]
+    _supports_flash_attn = True
+    _supports_sdpa = True
+    _supports_flex_attn = True
+    _can_compile_fullgraph = True
+    _supports_attention_backend = True
+    _can_record_outputs = {
+        "hidden_states": WeDLMDecoderLayer,
+        "attentions": WeDLMAttention,
+    }
+
+
+@auto_docstring
+class WeDLMModel(WeDLMPreTrainedModel):
+    """
+    WeDLM base model outputting raw hidden states.
+    """
+    
+    def __init__(self, config: WeDLMConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList(
+            [WeDLMDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self.norm = WeDLMRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = WeDLMRotaryEmbedding(config=config)
+        self.gradient_checkpointing = False
+        self.has_sliding_layers = "sliding_attention" in self.config.layer_types
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.embed_tokens = value
+
+    @check_model_inputs
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache(config=self.config)
+
+        if cache_position is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(
+                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+            )
+
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        # Prepare attention masks
+        if not isinstance(causal_mask_mapping := attention_mask, dict):
+            mask_kwargs = {
+                "config": self.config,
+                "input_embeds": inputs_embeds,
+                "attention_mask": attention_mask,
+                "cache_position": cache_position,
+                "past_key_values": past_key_values,
+                "position_ids": position_ids,
+            }
+            causal_mask_mapping = {
+                "full_attention": create_causal_mask(**mask_kwargs),
+            }
+            if self.has_sliding_layers:
+                causal_mask_mapping["sliding_attention"] = create_sliding_window_causal_mask(**mask_kwargs)
+
+        hidden_states = inputs_embeds
+
+        # Create position embeddings to be shared across the decoder layers
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        # Decoder layers
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+
+        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            layer_outputs = decoder_layer(
+                hidden_states,
+                attention_mask=causal_mask_mapping[decoder_layer.attention_type],
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                output_attentions=output_attentions,
+                use_cache=use_cache,
+                cache_position=cache_position,
+                position_embeddings=position_embeddings,
+                **kwargs,
+            )
+
+            hidden_states = layer_outputs[0]
+
+            if output_attentions:
+                all_self_attns += (layer_outputs[1],)
+
+        hidden_states = self.norm(hidden_states)
+
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        if not return_dict:
+            return tuple(v for v in [hidden_states, past_key_values, all_hidden_states, all_self_attns] if v is not None)
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values if use_cache else None,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+        )
+
+
+@auto_docstring
+class WeDLMForCausalLM(WeDLMPreTrainedModel, GenerationMixin):
+    """
+    WeDLM Model for Flow-Matching language modeling (Rectified Flow in token-embedding space).
+
+    - Training (`labels` provided): optimizes a Flow Matching objective on a selected subset of token positions.
+      Large-vocabulary projection (`lm_head`) is skipped by default during training for lower cost.
+    - Inference (no `labels`): behaves like a standard causal LM (returns logits).
+    - Fast decoding: use `generate_wedlm` (Flow-Matching block decoding).
+    """
+    _tied_weights_keys = ["lm_head.weight"]
+
+    def __init__(self, config: WeDLMConfig):
+        super().__init__(config)
+        self.model = WeDLMModel(config)
+        self.vocab_size = config.vocab_size
+
+        # Token discretization head (used for inference / evaluation / discretization at the end of each flow block)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        # Flow Matching modules
+        self.flow_time_embed = WeDLMFlowTimeEmbedding(config)
+        self.flow_head = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    # ----------------------------
+    # Embedding plumbing
+    # ----------------------------
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def get_decoder(self):
+        return self.model
+
+    # ----------------------------
+    # Flow Matching utilities
+    # ----------------------------
+    def _select_flow_targets(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor],
+        labels: Optional[torch.LongTensor],
+        flow_target_mask: Optional[torch.BoolTensor],
+    ) -> torch.BoolTensor:
+        """
+        Determine which token positions participate in Flow Matching loss.
+
+        Priority:
+        1) `flow_target_mask` argument
+        2) `config.flow_train_strategy`
+        """
+        bsz, seq_len = input_ids.shape
+        device = input_ids.device
+
+        if attention_mask is not None:
+            valid = attention_mask.to(dtype=torch.bool, device=device)
+        else:
+            # Best-effort fallback: treat non-pad as valid.
+            pad_id = getattr(self.config, "pad_token_id", None)
+            if pad_id is None:
+                valid = torch.ones((bsz, seq_len), dtype=torch.bool, device=device)
+            else:
+                valid = input_ids.ne(pad_id)
+
+        if labels is not None:
+            valid = valid & labels.ne(-100)
+
+        if flow_target_mask is not None:
+            target = flow_target_mask.to(dtype=torch.bool, device=device)
+            return target & valid
+
+        strategy = str(getattr(self.config, "flow_train_strategy", "suffix_block")).lower()
+        min_targets = int(getattr(self.config, "flow_train_min_target_tokens", 1))
+
+        if strategy == "random":
+            ratio = float(getattr(self.config, "flow_train_mask_ratio", 0.15))
+            # Sample only from valid positions; guarantee at least `min_targets` if possible.
+            rand = torch.rand((bsz, seq_len), device=device)
+            target = (rand < ratio) & valid
+            if min_targets > 0:
+                for b in range(bsz):
+                    if valid[b].any() and target[b].sum().item() < min_targets:
+                        valid_idx = valid[b].nonzero(as_tuple=True)[0]
+                        # Select the last positions (deterministic tie-break) to fill up.
+                        need = min(min_targets - target[b].sum().item(), valid_idx.numel())
+                        if need > 0:
+                            target[b, valid_idx[-need:]] = True
+            return target
+
+        # Default: suffix_block
+        block_size = int(getattr(self.config, "flow_train_block_size", 64))
+        target = torch.zeros((bsz, seq_len), dtype=torch.bool, device=device)
+        for b in range(bsz):
+            valid_idx = valid[b].nonzero(as_tuple=True)[0]
+            if valid_idx.numel() == 0:
+                continue
+            # Do not target the very first valid token by default (no context); if needed, user can pass flow_target_mask.
+            if valid_idx.numel() == 1:
+                continue
+            # Suffix contiguous block among valid positions.
+            k = min(block_size, valid_idx.numel() - 1)
+            k = max(k, min_targets)
+            k = min(k, valid_idx.numel() - 1)  # ensure at least one context token remains
+            if k <= 0:
+                continue
+            target_pos = valid_idx[-k:]
+            target[b, target_pos] = True
+        return target
+
+    def _normalize_timesteps(
+        self,
+        timesteps: Optional[torch.FloatTensor],
+        target_mask: torch.BoolTensor,
+    ) -> torch.FloatTensor:
+        """
+        Returns per-token normalized timesteps t in [0, 1], shape (bsz, seq_len).
+        """
+        device = target_mask.device
+        bsz, seq_len = target_mask.shape
+
+        if timesteps is None:
+            t = torch.rand((bsz, seq_len), device=device, dtype=torch.float32)
+            # Non-target positions: t=1.0 (data endpoint) so that any accidental use is benign.
+            t = torch.where(target_mask, t, torch.ones_like(t))
+            return t
+
+        t_in = timesteps.to(device=device, dtype=torch.float32)
+        if t_in.ndim == 0:
+            t = t_in.view(1, 1).expand(bsz, seq_len)
+        elif t_in.ndim == 1:
+            if t_in.shape[0] == 1 and bsz > 1:
+                t = t_in.view(1, 1).expand(bsz, seq_len)
+            elif t_in.shape[0] == bsz:
+                t = t_in.view(bsz, 1).expand(bsz, seq_len)
+            else:
+                raise ValueError(
+                    f"flow_timesteps must be scalar, shape (bsz,), or shape (bsz, seq_len); got {tuple(t_in.shape)}"
+                )
+        elif t_in.ndim == 2:
+            if t_in.shape != (bsz, seq_len):
+                raise ValueError(
+                    f"flow_timesteps must have shape (bsz, seq_len) == {(bsz, seq_len)}; got {tuple(t_in.shape)}"
+                )
+            t = t_in
+        else:
+            raise ValueError(
+                f"flow_timesteps must be scalar, 1D, or 2D; got ndim={t_in.ndim} with shape {tuple(t_in.shape)}"
+            )
+
+        # Clamp into [0,1] for numerical safety.
+        t = torch.clamp(t, 0.0, 1.0)
+        t = torch.where(target_mask, t, torch.ones_like(t))
+        return t
+
+    def _build_flow_inputs(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor],
+        labels: Optional[torch.LongTensor],
+        flow_target_mask: Optional[torch.BoolTensor],
+        flow_timesteps: Optional[torch.FloatTensor],
+        flow_noise: Optional[torch.FloatTensor],
+    ) -> Tuple[torch.FloatTensor, torch.BoolTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:
+        """
+        Prepare (inputs_embeds, target_mask, t, noise, clean_embeds) for Flow Matching training.
+        """
+        clean_embeds = self.model.embed_tokens(input_ids)
+        bsz, seq_len, hidden = clean_embeds.shape
+        device = clean_embeds.device
+
+        target_mask = self._select_flow_targets(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            labels=labels,
+            flow_target_mask=flow_target_mask,
+        )
+
+        t = self._normalize_timesteps(flow_timesteps, target_mask=target_mask)
+
+        sigma = float(getattr(self.config, "flow_init_sigma", 1.0))
+        if flow_noise is None:
+            noise = torch.randn_like(clean_embeds, dtype=torch.float32) * sigma
+        else:
+            noise = flow_noise.to(device=device, dtype=torch.float32)
+            if noise.shape != clean_embeds.shape:
+                raise ValueError(
+                    f"flow_noise must have the same shape as token embeddings {tuple(clean_embeds.shape)}; got {tuple(noise.shape)}"
+                )
+
+        # Rectified Flow path: X_t = (1 - t) * X_0 + t * X_1
+        # Here X_0 is noise, X_1 is data (token embeddings).
+        t_exp = t.unsqueeze(-1)
+        x_t = (1.0 - t_exp) * noise + t_exp * clean_embeds.to(dtype=torch.float32)
+
+        # Time conditioning is provided by adding a learned timestep embedding to the *input embeddings*
+        # for flow-target positions (so the Transformer can use t).
+        inputs_embeds = clean_embeds.to(dtype=torch.float32)
+        if target_mask.any():
+            # Compute time embedding only for targets to reduce overhead.
+            t_flat = t[target_mask].reshape(-1)
+            time_cond = self.flow_time_embed(t_flat).to(dtype=inputs_embeds.dtype)
+            inputs_embeds = inputs_embeds.clone()
+            inputs_embeds[target_mask] = x_t[target_mask] + time_cond
+        else:
+            inputs_embeds = inputs_embeds.clone()
+
+        return inputs_embeds.to(dtype=clean_embeds.dtype), target_mask, t, noise, clean_embeds
+
+    # ----------------------------
+    # Decoding: Flow Matching block decoding
+    # ----------------------------
+    def _top_k_top_p_filtering(
+        self,
+        logits: torch.Tensor,
+        top_k: int = 0,
+        top_p: float = 1.0,
+        filter_value: float = -float("inf"),
+    ) -> torch.Tensor:
+        """Apply top-k and/or nucleus (top-p) filtering to logits."""
+        if top_k is not None and top_k > 0:
+            top_k = min(top_k, logits.size(-1))
+            indices_to_remove = logits < torch.topk(logits, top_k, dim=-1).values[..., -1, None]
+            logits = logits.masked_fill(indices_to_remove, filter_value)
+
+        if top_p is not None and 0.0 < top_p < 1.0:
+            sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
+            cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+
+            sorted_indices_to_remove = cumulative_probs > top_p
+            sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+            sorted_indices_to_remove[..., 0] = False
+
+            indices_to_remove = sorted_indices_to_remove.scatter(-1, sorted_indices, sorted_indices_to_remove)
+            logits = logits.masked_fill(indices_to_remove, filter_value)
+
+        return logits
+
+    def _sample_from_logits(
+        self,
+        logits: torch.Tensor,
+        temperature: float = 1.0,
+        top_p: float = 1.0,
+        top_k: int = 0,
+    ) -> torch.Tensor:
+        """Sample token IDs from logits with temperature + (top-k, top-p) filtering."""
+        if temperature is None or temperature <= 0:
+            temperature = 1.0
+        logits = logits / float(temperature)
+        logits = self._top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p)
+        probs = F.softmax(logits, dim=-1)
+        return torch.multinomial(probs, num_samples=1).squeeze(-1)
+
+    @torch.no_grad()
+    def generate_wedlm(
+        self,
+        input_ids: torch.LongTensor,
+        max_new_tokens: int,
+        block_size: Optional[int] = None,
+        num_steps: Optional[int] = None,
+        flow_init_sigma: Optional[float] = None,
+        discretization: Optional[str] = None,
+        temperature: Optional[float] = None,
+        top_p: Optional[float] = None,
+        top_k: Optional[int] = None,
+        pad_token_id: Optional[int] = None,
+        eos_token_id: Optional[int] = None,
+        return_stats: bool = True,
+        **kwargs,
+    ) -> Union[torch.LongTensor, Dict]:
+        """
+        Flow-Matching block decoding.
+
+        Generates `block_size` tokens per block using `num_steps` Euler steps, and runs the vocabulary projection
+        (`lm_head`) only once per block (final discretization).
+        """
+        device = input_ids.device
+        if pad_token_id is None:
+            pad_token_id = self.config.pad_token_id
+        if eos_token_id is None:
+            eos_token_id = getattr(self.config, "eos_token_id", None)
+
+        if block_size is None:
+            block_size = int(getattr(self.config, "flow_block_size", 64))
+        if num_steps is None:
+            num_steps = int(getattr(self.config, "flow_inference_steps", 8))
+        if flow_init_sigma is None:
+            flow_init_sigma = float(getattr(self.config, "flow_init_sigma", 1.0))
+
+        if discretization is None:
+            discretization = str(getattr(self.config, "flow_discretization", "argmax")).lower()
+        if temperature is None:
+            temperature = float(getattr(self.config, "flow_temperature", 1.0))
+        if top_p is None:
+            top_p = float(getattr(self.config, "flow_top_p", 1.0))
+        if top_k is None:
+            top_k = int(getattr(self.config, "flow_top_k", 0))
+
+        batch_size = input_ids.shape[0]
+        all_generated: List[torch.Tensor] = []
+        all_sample_stats: List[Dict] = []
+
+        num_blocks = (max_new_tokens + block_size - 1) // block_size
+
+        for batch_idx in range(batch_size):
+            sample_ids = input_ids[batch_idx]
+            if pad_token_id is not None:
+                pad_mask = sample_ids.ne(pad_token_id)
+                if pad_mask.any():
+                    valid_length = int(pad_mask.sum().item())
+                    prefix_ids = sample_ids[:valid_length]
+                else:
+                    prefix_ids = sample_ids
+            else:
+                prefix_ids = sample_ids
+
+            prefix_ids = prefix_ids.clone()
+            prefix_length = prefix_ids.shape[0]
+
+            sample_stats = {
+                "input_length": prefix_length,
+                "num_blocks": num_blocks,
+                "block_size": block_size,
+                "num_steps": num_steps,
+                "sigma": float(flow_init_sigma),
+                "generated_tokens": 0,
+                "blocks": [],
+            }
+
+            current_ids = prefix_ids
+
+            for block_idx in range(num_blocks):
+                remaining = max_new_tokens - block_idx * block_size
+                cur_block = min(block_size, remaining)
+                if cur_block <= 0:
+                    break
+
+                # State variable: current embedding estimates for the block (initialized from Gaussian noise)
+                x = torch.randn((cur_block, self.config.hidden_size), device=device, dtype=torch.float32) * float(flow_init_sigma)
+
+                dt = 1.0 / float(num_steps)
+                # Euler integration from t=0 (noise) to t=1 (data)
+                for step in range(num_steps):
+                    t = float(step) / float(num_steps)
+                    # Create per-token timesteps matching training behavior (all same t for this step)
+                    t_tensor = torch.full((cur_block,), t, device=device, dtype=torch.float32)
+                    # Build embeddings in *model dtype* to avoid dtype mismatch in Linear layers when AMP is off.
+                    context_embeds = self.model.embed_tokens(current_ids.unsqueeze(0))
+                    ctx_dtype = context_embeds.dtype
+                    # Per-token time conditioning (cur_block,) -> (cur_block, H)
+                    t_cond = self.flow_time_embed(t_tensor).to(dtype=ctx_dtype)  # (cur_block, H)
+
+                    # Context tokens are discrete; block tokens are continuous (x) + time conditioning.
+                    block_embeds = x.to(dtype=ctx_dtype) + t_cond  # (cur_block, H)
+                    block_embeds = block_embeds.view(1, cur_block, -1)  # (1, cur_block, H)
+                    inputs_embeds = torch.cat([context_embeds, block_embeds], dim=1)
+
+                    seq_len = inputs_embeds.shape[1]
+                    attention_mask = torch.ones((1, seq_len), dtype=torch.long, device=device)
+                    position_ids = torch.arange(seq_len, dtype=torch.long, device=device).unsqueeze(0)
+
+                    outputs = self.model(
+                        input_ids=None,
+                        inputs_embeds=inputs_embeds,
+                        attention_mask=attention_mask,
+                        position_ids=position_ids,
+                        use_cache=False,
+                        return_dict=True,
+                    )
+                    h_new = outputs.last_hidden_state[:, -cur_block:, :]  # (1, cur_block, H)
+
+                    # Predict velocity in a dtype-compatible way, then accumulate in fp32.
+                    h_in = h_new
+                    base_flow = getattr(self.flow_head, "base_layer", self.flow_head)
+                    w_flow = getattr(base_flow, "weight", None)
+                    if w_flow is not None:
+                        h_in = h_in.to(dtype=w_flow.dtype)
+                    v = self.flow_head(h_in).to(dtype=torch.float32).squeeze(0)  # (cur_block, H)
+                    x = x + dt * v
+
+                # Discretize final embeddings into token IDs
+                base_lm = getattr(self.lm_head, "base_layer", self.lm_head)
+                w_lm = getattr(base_lm, "weight", None)
+                x_lm = x
+                if w_lm is not None:
+                    x_lm = x_lm.to(dtype=w_lm.dtype)
+                logits = self.lm_head(x_lm).to(dtype=torch.float32)  # (cur_block, vocab)
+                if discretization == "sample":
+                    next_ids = self._sample_from_logits(logits, temperature=temperature, top_p=top_p, top_k=top_k)
+                else:
+                    next_ids = torch.argmax(logits, dim=-1)
+
+                # Optional early stop on EOS within the block
+                if eos_token_id is not None:
+                    eos_positions = (next_ids == eos_token_id).nonzero(as_tuple=True)[0]
+                    if eos_positions.numel() > 0:
+                        cut = int(eos_positions[0].item()) + 1
+                        next_ids = next_ids[:cut]
+
+                current_ids = torch.cat([current_ids, next_ids.to(dtype=torch.long)], dim=0)
+                sample_stats["generated_tokens"] += int(next_ids.numel())
+                sample_stats["blocks"].append(
+                    {
+                        "block_idx": block_idx,
+                        "target_block_size": int(cur_block),
+                        "actual_block_tokens": int(next_ids.numel()),
+                    }
+                )
+
+                if eos_token_id is not None and next_ids.numel() > 0 and next_ids[-1].item() == int(eos_token_id):
+                    break
+
+            sample_stats["output_length"] = int(current_ids.numel())
+            all_generated.append(current_ids)
+            all_sample_stats.append(sample_stats)
+
+        # Pad to max length
+        max_len = max(seq.numel() for seq in all_generated) if all_generated else 0
+        padded = []
+        for seq in all_generated:
+            if seq.numel() < max_len:
+                pad = torch.full(
+                    (max_len - seq.numel(),),
+                    int(pad_token_id) if pad_token_id is not None else 0,
+                    dtype=torch.long,
+                    device=device,
+                )
+                seq = torch.cat([seq, pad], dim=0)
+            padded.append(seq)
+        sequences = torch.stack(padded, dim=0) if padded else torch.empty((0, 0), dtype=torch.long, device=device)
+
+        if not return_stats:
+            return sequences
+
+        total_steps = int(num_steps) * int(num_blocks) * int(batch_size)
+        return {
+            "sequences": sequences,
+            "stats": {
+                "batch_size": int(batch_size),
+                "max_new_tokens": int(max_new_tokens),
+                "block_size": int(block_size),
+                "num_steps": int(num_steps),
+                "discretization": discretization,
+                "temperature": float(temperature),
+                "top_p": float(top_p),
+                "top_k": int(top_k),
+                "total_flow_evals": total_steps,
+                "per_sample_stats": all_sample_stats,
+            },
+        }
+
+    # ----------------------------
+    # Generate (override to use Flow Matching by default)
+    # ----------------------------
+    @torch.no_grad()
+    def generate(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        generation_config=None,
+        max_new_tokens: Optional[int] = None,
+        max_length: Optional[int] = None,
+        do_sample: Optional[bool] = None,
+        temperature: Optional[float] = None,
+        top_p: Optional[float] = None,
+        top_k: Optional[int] = None,
+        pad_token_id: Optional[int] = None,
+        eos_token_id: Optional[int] = None,
+        use_flow_matching: bool = True,
+        block_size: Optional[int] = None,
+        num_steps: Optional[int] = None,
+        flow_init_sigma: Optional[float] = None,
+        streamer=None,
+        **kwargs,
+    ):
+        """
+        Override generate() to use Flow Matching decoding by default.
+
+        Set `use_flow_matching=False` to fall back to standard AR generation.
+        """
+        # Fall back to standard AR generation if requested
+        if not use_flow_matching:
+            return super().generate(
+                input_ids=input_ids,
+                generation_config=generation_config,
+                max_new_tokens=max_new_tokens,
+                max_length=max_length,
+                do_sample=do_sample,
+                temperature=temperature,
+                top_p=top_p,
+                top_k=top_k,
+                pad_token_id=pad_token_id,
+                eos_token_id=eos_token_id,
+                streamer=streamer,
+                **kwargs,
+            )
+
+        # Extract parameters from generation_config if provided
+        if generation_config is not None:
+            if max_new_tokens is None:
+                max_new_tokens = getattr(generation_config, "max_new_tokens", None)
+            if max_length is None:
+                max_length = getattr(generation_config, "max_length", None)
+            if do_sample is None:
+                do_sample = getattr(generation_config, "do_sample", None)
+            if temperature is None:
+                temperature = getattr(generation_config, "temperature", None)
+            if top_p is None:
+                top_p = getattr(generation_config, "top_p", None)
+            if top_k is None:
+                top_k = getattr(generation_config, "top_k", None)
+            if pad_token_id is None:
+                pad_token_id = getattr(generation_config, "pad_token_id", None)
+            if eos_token_id is None:
+                eos_token_id = getattr(generation_config, "eos_token_id", None)
+
+        # Determine max_new_tokens
+        if max_new_tokens is None:
+            if max_length is not None and input_ids is not None:
+                max_new_tokens = max_length - input_ids.shape[1]
+            else:
+                max_new_tokens = 256  # Default
+
+        max_new_tokens = max(1, max_new_tokens)
+
+        # Map do_sample to discretization
+        discretization = "sample" if do_sample else "argmax"
+
+        # Use config defaults for flow parameters
+        if block_size is None:
+            block_size = getattr(self.config, "flow_block_size", 64)
+        if num_steps is None:
+            num_steps = getattr(self.config, "flow_inference_steps", 8)
+        if flow_init_sigma is None:
+            flow_init_sigma = getattr(self.config, "flow_init_sigma", 1.0)
+        if temperature is None:
+            temperature = getattr(self.config, "flow_temperature", 1.0)
+        if top_p is None:
+            top_p = getattr(self.config, "flow_top_p", 1.0)
+        if top_k is None:
+            top_k = getattr(self.config, "flow_top_k", 0)
+        if pad_token_id is None:
+            pad_token_id = self.config.pad_token_id
+        if eos_token_id is None:
+            eos_token_id = self.config.eos_token_id
+
+        # Call Flow Matching generation
+        result = self.generate_wedlm(
+            input_ids=input_ids,
+            max_new_tokens=max_new_tokens,
+            block_size=block_size,
+            num_steps=num_steps,
+            flow_init_sigma=flow_init_sigma,
+            discretization=discretization,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            pad_token_id=pad_token_id,
+            eos_token_id=eos_token_id,
+            return_stats=False,
+        )
+
+        # Handle streamer if provided (basic support)
+        if streamer is not None:
+            for token_id in result[0, input_ids.shape[1]:]:
+                streamer.put(token_id.unsqueeze(0).unsqueeze(0))
+            streamer.end()
+
+        return result
+
+    # ----------------------------
+    # Forward
+    # ----------------------------
+    @can_return_tuple
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        # Flow Matching controls (training-time)
+        flow_target_mask: Optional[torch.BoolTensor] = None,
+        flow_timesteps: Optional[torch.FloatTensor] = None,
+        flow_noise: Optional[torch.FloatTensor] = None,
+        return_logits: bool = False,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        """
+        When `labels` is provided: computes Flow Matching loss.
+        Otherwise: returns logits like a standard causal LM.
+
+        Args:
+            flow_target_mask (`torch.BoolTensor`, *optional*):
+                Boolean mask indicating which positions are flow targets.
+            flow_timesteps (`torch.FloatTensor`, *optional*):
+                Timesteps for flow matching, in range [0, 1].
+            flow_noise (`torch.FloatTensor`, *optional*):
+                Noise tensor for flow matching interpolation.
+            return_logits (`bool`, defaults to `False`):
+                If True, also compute vocabulary logits during Flow-Matching training.
+        """
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # ------------------------------------------------------------
+        # Flow Matching training path (primary)
+        # ------------------------------------------------------------
+        if labels is not None:
+            if input_ids is None:
+                raise ValueError("Flow-Matching training requires input_ids (token IDs) when labels is provided.")
+
+            if inputs_embeds is not None:
+                raise ValueError("Do not pass inputs_embeds when training with labels; Flow-Matching builds embeds internally.")
+
+            inputs_embeds, target_mask, t, noise, clean_embeds = self._build_flow_inputs(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                labels=labels,
+                flow_target_mask=flow_target_mask,
+                flow_timesteps=flow_timesteps,
+                flow_noise=flow_noise,
+            )
+
+            outputs = self.model(
+                input_ids=None,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                inputs_embeds=inputs_embeds,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=True,
+                cache_position=cache_position,
+                **kwargs,
+            )
+
+            hidden_states = outputs.last_hidden_state
+            loss = None
+            logits = None
+
+            if target_mask.any():
+                # Target velocity: d/dt X_t = X_1 - X_0 (data - noise) on straight-line path.
+                v_target = (clean_embeds.to(dtype=torch.float32) - noise.to(dtype=torch.float32))[target_mask]
+                # PEFT/LoRA can wrap `flow_head` and keep base weights in bf16/fp16.
+                # Ensure dtype alignment for the Linear matmul.
+                hs = hidden_states[target_mask]
+                base_flow = getattr(self.flow_head, "base_layer", self.flow_head)
+                w_flow = getattr(base_flow, "weight", None)
+                if w_flow is not None:
+                    hs = hs.to(dtype=w_flow.dtype)
+                v_pred = self.flow_head(hs).to(dtype=torch.float32)
+                flow_loss = F.mse_loss(v_pred, v_target, reduction="mean")
+
+                w = float(getattr(self.config, "flow_loss_weight", 1.0))
+                loss = w * flow_loss
+            else:
+                # No valid targets -> zero loss (avoid NaNs).
+                loss = hidden_states.new_tensor(0.0)
+
+            if return_logits:
+                slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+                logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+            if not return_dict:
+                output = (logits,) + (outputs.past_key_values, outputs.hidden_states, outputs.attentions)
+                return (loss,) + output
+
+            return CausalLMOutputWithPast(
+                loss=loss,
+                logits=logits,
+                past_key_values=outputs.past_key_values,
+                hidden_states=outputs.hidden_states,
+                attentions=outputs.attentions,
+            )
+
+        # ------------------------------------------------------------
+        # Standard causal LM path (no labels): logits for evaluation / AR generation
+        # ------------------------------------------------------------
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=True,
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+        hidden_states = outputs.last_hidden_state
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+        if not return_dict:
+            output = (logits,) + (outputs.past_key_values, outputs.hidden_states, outputs.attentions)
+            return output
+
+        return CausalLMOutputWithPast(
+            loss=None,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        past_key_values=None,
+        attention_mask=None,
+        inputs_embeds=None,
+        cache_position=None,
+        position_ids=None,
+        use_cache=True,
+        **kwargs,
+    ):
+        if past_key_values is not None:
+            if inputs_embeds is not None:
+                input_ids = input_ids[:, -cache_position.shape[0] :]
+            elif input_ids.shape[1] != cache_position.shape[0]:
+                input_ids = input_ids[:, cache_position]
+
+        if attention_mask is not None and position_ids is None:
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            if past_key_values:
+                position_ids = position_ids[:, -input_ids.shape[1] :]
+
+        if inputs_embeds is not None and cache_position is not None and cache_position[0] == 0:
+            model_inputs = {"inputs_embeds": inputs_embeds, "input_ids": None}
+        else:
+            model_inputs = {"input_ids": input_ids.clone(memory_format=torch.contiguous_format), "inputs_embeds": None}
+
+        if isinstance(past_key_values, DynamicCache) and attention_mask is not None and attention_mask.ndim == 2:
+            model_inputs["cache_position"] = cache_position
+            model_inputs["past_key_values"] = past_key_values
+            model_inputs["use_cache"] = use_cache
+            model_inputs["position_ids"] = position_ids
+            model_inputs["attention_mask"] = attention_mask
+            return model_inputs
+
+        model_inputs.update(
+            {
+                "position_ids": position_ids,
+                "cache_position": cache_position,
+                "past_key_values": past_key_values,
+                "use_cache": use_cache,
+                "attention_mask": attention_mask,
+            }
+        )
+        return model_inputs
+__all__ = [
+    "WeDLMConfig",
+    "WeDLMPreTrainedModel", 
+    "WeDLMModel",
+    "WeDLMForCausalLM",
+]

rng_state.pth

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

scheduler.pt

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

special_tokens_map.json

@@ -0,0 +1,31 @@
+{
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "eos_token": {
+    "content": "<|im_end|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer.json

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

tokenizer_config.json

@@ -0,0 +1,208 @@
+{
+  "add_bos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "151643": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151644": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151645": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151646": {
+      "content": "<|object_ref_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151647": {
+      "content": "<|object_ref_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151648": {
+      "content": "<|box_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151649": {
+      "content": "<|box_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151650": {
+      "content": "<|quad_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151651": {
+      "content": "<|quad_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151652": {
+      "content": "<|vision_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151653": {
+      "content": "<|vision_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151654": {
+      "content": "<|vision_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151655": {
+      "content": "<|image_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151656": {
+      "content": "<|video_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151657": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151658": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151659": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151660": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151661": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151662": {
+      "content": "<|fim_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151663": {
+      "content": "<|repo_name|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151664": {
+      "content": "<|file_sep|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "extra_special_tokens": {},
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "padding_side": "right",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}

training_args.bin

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