Add files using upload-large-folder tool

.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

chat_template.jinja

@@ -0,0 +1,86 @@
+{% set thinking_option = 'off' %}
+{{- '<role>SYSTEM</role>' }}
+{%- if messages[0].role == 'system' %}
+    {{- messages[0].content + '\n' }}
+{%- endif %}
+{%- if tools %}
+    {{- "# 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>\n" }}
+{%- endif %}
+{{- 'detailed thinking ' + thinking_option + '<|role_end|>' }}
+{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
+{%- for message in messages[::-1] %}
+    {%- set index = (messages|length - 1) - loop.index0 %}
+    {%- if ns.multi_step_tool and message.role == "user" and message.content is string and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}
+        {%- set ns.multi_step_tool = false %}
+        {%- set ns.last_query_index = index %}
+    {%- endif %}
+{%- endfor %}
+{%- for message in messages %}
+    {%- if message.content is string %}
+        {%- set content = message.content %}
+    {%- else %}
+        {%- set content = '' %}
+    {%- endif %}
+    {%- if message.role == "user" %}
+        {{- '<role>HUMAN</role>' + message.content + '<|role_end|>' }}
+    {%- elif message.role == "system" and not loop.first %}
+        {{- '<role>SYSTEM</role>' + message.content + '<|role_end|>' }}
+    {%- elif message.role == "assistant" %}
+        {%- set reasoning_content = '' %}
+        {%- if message.reasoning_content is string %}
+            {%- set reasoning_content = message.reasoning_content %}
+        {%- else %}
+            {%- if '</think>' in content %}
+                {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+                {%- set content = content.split('</think>')[-1].lstrip('\n') %}
+            {%- endif %}
+        {%- endif %}
+        {%- if loop.index0 > ns.last_query_index %}
+            {%- if reasoning_content %}
+                {{- '<role>ASSISTANT</role>' + '\n<think>\n' + reasoning_content.strip('\n') + '\n</think>\n\n' + content.lstrip('\n') }}
+            {%- else %}
+                {{- '<role>ASSISTANT</role>' + content }}
+            {%- endif %}
+        {%- else %}
+            {{- '<role>ASSISTANT</role>' + content }}
+        {%- endif %}
+        {%- if message.tool_calls %}
+            {%- for tool_call in message.tool_calls %}
+                {%- if (loop.first and content) or (not loop.first) %}
+                    {{- '\n' }}
+                {%- endif %}
+                {%- if tool_call.function %}
+                    {%- set tool_call = tool_call.function %}
+                {%- endif %}
+                {{- '<tool_call>\n{"name": "' }}
+                {{- tool_call.name }}
+                {{- '", "arguments": ' }}
+                {%- if tool_call.arguments is string %}
+                    {{- tool_call.arguments }}
+                {%- else %}
+                    {{- tool_call.arguments | tojson }}
+                {%- endif %}
+                {{- '}\n</tool_call>' }}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|role_end|>' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
+            {{- '<role>OBSERVATION</role>' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- content }}
+        {{- '\n</tool_response>' }}
+        {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
+            {{- '<|role_end|>' }}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<role>ASSISTANT</role>' }}
+{%- endif %}

config.json

@@ -0,0 +1,61 @@
+{
+  "_moe_implementation": "fused",
+  "architectures": [
+    "LLaDA2MoeModelLM"
+  ],
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoConfig": "configuration_llada2_moe.LLaDA2MoeConfig",
+    "AutoModelForCausalLM": "modeling_llada2_moe.LLaDA2MoeModelLM"
+  },
+  "dtype": "float32",
+  "embedding_dropout": 0.0,
+  "first_k_dense_replace": 1,
+  "fuse_cross_entropy": false,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 2048,
+  "initializer_range": 0.02,
+  "intermediate_size": 5120,
+  "max_position_embeddings": 16384,
+  "max_window_layers": 28,
+  "model_type": "llada2_moe",
+  "moe_intermediate_size": 512,
+  "moe_router_enable_expert_bias": true,
+  "n_group": 8,
+  "norm_head": false,
+  "norm_softmax": false,
+  "norm_topk_prob": true,
+  "num_attention_heads": 16,
+  "num_experts": 256,
+  "num_experts_per_tok": 8,
+  "num_hidden_layers": 20,
+  "num_key_value_heads": 4,
+  "num_shared_experts": 1,
+  "output_dropout": 0.0,
+  "output_router_logits": false,
+  "pad_token_id": 156892,
+  "partial_rotary_factor": 0.5,
+  "rms_norm_eps": 1e-06,
+  "rope_parameters": {
+    "partial_rotary_factor": 0.5,
+    "rope_theta": 600000,
+    "rope_type": "default"
+  },
+  "rope_theta": 600000,
+  "rotary_dim": 64,
+  "routed_scaling_factor": 2.5,
+  "router_dtype": "fp32",
+  "score_function": "sigmoid",
+  "sliding_window": 4096,
+  "tie_word_embeddings": false,
+  "topk_group": 4,
+  "transformers_version": "5.9.0",
+  "use_bias": false,
+  "use_cache": false,
+  "use_qkv_bias": false,
+  "use_rmsnorm": true,
+  "use_sliding_window": false,
+  "using_split_qkv_in_self_attention": false,
+  "vocab_size": 157184
+}

configuration_llada2_moe.py

@@ -0,0 +1,88 @@
+"""LLaDA2 MoE model configuration"""
+
+from transformers.configuration_utils import PretrainedConfig
+
+
+class LLaDA2MoeConfig(PretrainedConfig):
+    model_type = "llada2_moe"
+
+    def __init__(
+        self,
+        vocab_size=30592,
+        hidden_size=1024,
+        intermediate_size=None,
+        num_hidden_layers=24,
+        num_attention_heads=16,
+        num_key_value_heads=0,
+        hidden_act="silu",
+        use_qkv_bias=False,  # llada2 only
+        use_qk_norm=False,
+        use_bias=True,  # llada2 only
+        rms_norm_eps=1e-05,
+        norm_head=False,  # llada2 only
+        tie_word_embeddings=False,  # PretrainedConfig key, here change default value.
+        embedding_dropout=0.1,
+        attention_dropout=0.1,
+        output_dropout=0.1,
+        initializer_range=0.02,
+        max_position_embeddings=16384,
+        rope_theta=10000.0,
+        use_cache=True,
+        use_sliding_window=False,
+        sliding_window=4096,
+        max_window_layers=28,
+        rope_scaling=None,
+        pad_token_id=None,
+        num_experts=16,
+        num_shared_experts=0,
+        num_experts_per_tok=2,
+        n_group=8,
+        topk_group=4,
+        routed_scaling_factor=2.5,
+        moe_intermediate_size=None,
+        first_k_dense_replace=0,
+        head_dim=None,
+        output_router_logits=False,
+        partial_rotary_factor=0.5,
+        **kwargs,
+    ):
+        self.num_hidden_layers = num_hidden_layers
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.hidden_act = hidden_act
+        self.use_qkv_bias = use_qkv_bias
+        self.use_bias = use_bias
+        self.norm_head = norm_head
+        self.rms_norm_eps = rms_norm_eps
+        self.embedding_dropout = embedding_dropout
+        self.attention_dropout = attention_dropout
+        self.output_dropout = output_dropout
+        self.initializer_range = initializer_range
+        self.max_position_embeddings = max_position_embeddings
+        self.rope_theta = rope_theta
+        self.use_cache = use_cache
+        self.use_sliding_window = use_sliding_window
+        self.sliding_window = sliding_window
+        self.max_window_layers = max_window_layers
+        self.head_dim = head_dim or self.hidden_size // self.num_attention_heads
+        self.rope_scaling = rope_scaling
+
+        # MoE configs
+        self.num_experts = num_experts
+        self.num_shared_experts = num_shared_experts
+        self.num_experts_per_tok = num_experts_per_tok
+        self.n_group = n_group
+        self.topk_group = topk_group
+        self.moe_intermediate_size = moe_intermediate_size
+        self.first_k_dense_replace = first_k_dense_replace
+        self.output_router_logits = output_router_logits
+        self.routed_scaling_factor = routed_scaling_factor
+        self.partial_rotary_factor = partial_rotary_factor
+
+        super().__init__(pad_token_id=pad_token_id, tie_word_embeddings=tie_word_embeddings, **kwargs)
+
+
+__all__ = ["LLaDA2MoeConfig"]

model-00001-of-00009.safetensors

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

model-00002-of-00009.safetensors

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

model-00003-of-00009.safetensors

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

model-00004-of-00009.safetensors

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

model-00005-of-00009.safetensors

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

model-00006-of-00009.safetensors

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

model-00007-of-00009.safetensors

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

model-00008-of-00009.safetensors

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

model-00009-of-00009.safetensors

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

modeling_llada2_moe.py

@@ -0,0 +1,370 @@
+# coding=utf-8
+# Copyright 2025 Antgroup and The HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# 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.
+"""DMax LLaDA2 training layers; see NOTICE.md for upstream revision and adaptations.
+
+Explicit 4D attention masks are passed unchanged to SDPA (True means attend).
+"""
+import math
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.utils.checkpoint import checkpoint
+from transformers.activations import ACT2FN
+from transformers.modeling_outputs import MoeCausalLMOutputWithPast
+from transformers.modeling_utils import PreTrainedModel
+from .configuration_llada2_moe import LLaDA2MoeConfig
+
+class LLaDA2MoeRMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        """
+        LLaDA2MoeRMSNorm is equivalent to T5LayerNorm
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        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 rotate_half(x):
+    """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)
+
+
+# Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb
+def apply_rotary_pos_emb(q, k, cos, sin):
+    """Rotate the configured head dimensions at the supplied logical positions."""
+    # RoPE values are [batch, token, rotary_dim]; broadcast over attention heads.
+    cos = cos.unsqueeze(1)
+    sin = sin.unsqueeze(1)
+
+    # Keep half or full tensor for later concatenation
+    rotary_dim = cos.shape[-1]
+    q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]
+    k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:]
+
+    # Apply rotary embeddings on the first half or full tensor
+    q_embed = (q_rot * cos) + (rotate_half(q_rot) * sin)
+    k_embed = (k_rot * cos) + (rotate_half(k_rot) * sin)
+
+    # Concatenate back to full shape
+    q_embed = torch.cat([q_embed, q_pass], dim=-1)
+    k_embed = torch.cat([k_embed, k_pass], dim=-1)
+    return q_embed, k_embed
+
+
+class LLaDA2MoeMLP(nn.Module):
+    def __init__(self, config: LLaDA2MoeConfig, intermediate_size: int):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = 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):
+        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+
+
+class LLaDA2MoeGate(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.top_k = config.num_experts_per_tok
+        self.num_experts = config.num_experts
+
+        self.n_group = config.n_group
+        self.topk_group = config.topk_group
+
+        # topk selection algorithm
+        self.gating_dim = config.hidden_size
+        self.weight = nn.Parameter(torch.empty((self.num_experts, self.gating_dim)))
+        self.routed_scaling_factor = config.routed_scaling_factor
+
+        self.register_buffer("expert_bias", torch.zeros((self.num_experts)))
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        import torch.nn.init as init
+
+        init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+
+    def group_limited_topk(
+        self,
+        scores: torch.Tensor,
+    ):
+        num_tokens, _ = scores.size()
+        # Organize the experts into groups
+        group_scores = scores.view(num_tokens, self.n_group, -1).topk(2, dim=-1)[0].sum(dim=-1)
+        group_idx = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)[1]
+        group_mask = torch.zeros_like(group_scores)
+        group_mask.scatter_(1, group_idx, 1)
+
+        # Mask the experts based on selection groups
+        score_mask = (
+            group_mask.unsqueeze(-1)
+            .expand(num_tokens, self.n_group, self.num_experts // self.n_group)
+            .reshape(num_tokens, -1)
+        )
+
+        masked_scores = scores.masked_fill(~score_mask.bool(), float('-inf'))
+        probs, top_indices = torch.topk(masked_scores, k=self.top_k, dim=-1)
+
+        return probs, top_indices
+
+    def forward(self, hidden_states):
+        # compute gating score
+        hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
+        logits = F.linear(hidden_states.type(torch.float32), self.weight.type(torch.float32))
+
+        scores = torch.sigmoid(logits.float()).type_as(logits)
+
+        scores_for_routing = scores + self.expert_bias
+        _, topk_idx = self.group_limited_topk(scores_for_routing)
+
+        scores = torch.gather(scores, dim=1, index=topk_idx).type_as(logits)
+
+        topk_weight = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20) if self.top_k > 1 else scores
+        topk_weight = topk_weight * self.routed_scaling_factor
+
+        return topk_idx, topk_weight, logits
+
+
+class LLaDA2MoeRotaryEmbedding(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        scaling = config.rope_scaling or {}
+        if scaling.get("rope_type", scaling.get("type", "default")) != "default":
+            raise ValueError("yrDMax supports the released LLaDA2 default RoPE only")
+        self.rebuild(torch.device("cpu"))
+
+    def rebuild(self, device):
+        dim = int(self.config.head_dim * self.config.partial_rotary_factor)
+        inv_freq = 1.0 / (self.config.rope_theta ** (
+            torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+    @torch.no_grad()
+    def forward(self, x, position_ids):
+        inv = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
+        with torch.autocast(device_type=x.device.type, enabled=False):
+            freqs = (inv @ position_ids[:, None, :].float()).transpose(1, 2)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            return emb.cos().to(x.dtype), emb.sin().to(x.dtype)
+
+
+class LLaDA2MoeSparseMoeBlock(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.num_experts_per_tok = config.num_experts_per_tok
+        self.experts = nn.ModuleList([
+            LLaDA2MoeMLP(config, config.moe_intermediate_size)
+            for _ in range(config.num_experts)])
+        self.gate = LLaDA2MoeGate(config)
+        self.shared_experts = (LLaDA2MoeMLP(
+            config, config.moe_intermediate_size * config.num_shared_experts)
+            if config.num_shared_experts else None)
+
+    def forward(self, hidden_states):
+        # Same expert and combine equations as upstream _forward; autograd also
+        # works in eval mode, so temporary mode changes cannot freeze experts.
+        shape = hidden_states.shape
+        indices, weights, router_logits = self.gate(hidden_states)
+        flat = hidden_states.reshape(-1, shape[-1])
+        expanded = flat.repeat_interleave(self.num_experts_per_tok, dim=0)
+        result = torch.empty_like(expanded)
+        for i, expert in enumerate(self.experts):
+            selected = indices.reshape(-1) == i
+            # Unwrapped/DDP FP32 masters may produce BF16 expert activations
+            # under autocast. Indexed assignment does not promote dtypes.
+            result[selected] = expert(expanded[selected]).to(result.dtype)
+        result = (result.view(*weights.shape, -1) * weights.unsqueeze(-1)).sum(1)
+        result = result.to(hidden_states.dtype).view(shape)
+        if self.shared_experts is not None:
+            result = result + self.shared_experts(hidden_states)
+        return result
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    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)
+
+
+# Copied from transformers.models.llama.modeling_llama.LlamaAttention with Llama->LLaDA2Moe
+class LLaDA2MoeAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config: LLaDA2MoeConfig, layer_idx=None):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+
+        self.attention_dropout = config.attention_dropout
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = config.head_dim or self.hidden_size // self.num_heads
+        partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0
+        self.rope_dim = int(self.head_dim * partial_rotary_factor)
+        self.num_key_value_heads = config.num_key_value_heads
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.max_position_embeddings = config.max_position_embeddings
+        self.rope_theta = config.rope_theta
+        self.is_causal = False
+
+        self.query_key_value = nn.Linear(
+            self.hidden_size,
+            (self.num_heads + 2 * self.num_key_value_heads) * self.head_dim,
+            bias=config.use_qkv_bias,
+        )
+
+        self.query_layernorm = LLaDA2MoeRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.key_layernorm = LLaDA2MoeRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.dense = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.use_bias)
+
+    def forward(self, hidden_states, attention_mask, position_embeddings):
+        bsz, q_len, _ = hidden_states.shape
+        qkv = self.query_key_value(hidden_states).view(
+            bsz, q_len, self.num_heads + 2 * self.num_key_value_heads, self.head_dim)
+        query, key, value = qkv.split(
+            [self.num_heads, self.num_key_value_heads, self.num_key_value_heads], dim=-2)
+        query = self.query_layernorm(query.transpose(1, 2))
+        key = self.key_layernorm(key.transpose(1, 2))
+        value = value.transpose(1, 2)
+        query, key = apply_rotary_pos_emb(query, key, *position_embeddings)
+        key, value = repeat_kv(key, self.num_key_value_groups), repeat_kv(value, self.num_key_value_groups)
+        out = F.scaled_dot_product_attention(
+            query.contiguous(), key.contiguous(), value.contiguous(),
+            attn_mask=attention_mask,
+            dropout_p=self.attention_dropout if self.training else 0.0,
+            is_causal=False)
+        out = self.dense(out.transpose(1, 2).contiguous().reshape(bsz, q_len, -1))
+        return out
+
+
+class LLaDA2MoeDecoderLayer(nn.Module):
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.attention = LLaDA2MoeAttention(config, layer_idx)
+        self.mlp = (LLaDA2MoeSparseMoeBlock(config)
+                    if config.num_experts is not None and layer_idx >= config.first_k_dense_replace
+                    else LLaDA2MoeMLP(config, config.intermediate_size))
+        self.input_layernorm = LLaDA2MoeRMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.post_attention_layernorm = LLaDA2MoeRMSNorm(config.hidden_size, config.rms_norm_eps)
+
+    def forward(self, hidden_states, attention_mask, position_embeddings):
+        attention_output = self.attention(self.input_layernorm(hidden_states), attention_mask,
+                                          position_embeddings)
+        hidden_states = hidden_states + attention_output
+        hidden_states = hidden_states + self.mlp(self.post_attention_layernorm(hidden_states))
+        return hidden_states
+
+
+class LLaDA2MoeModel(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, config.pad_token_id)
+        self.layers = nn.ModuleList([LLaDA2MoeDecoderLayer(config, i) for i in range(config.num_hidden_layers)])
+        self.norm = LLaDA2MoeRMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.rotary_emb = LLaDA2MoeRotaryEmbedding(config)
+        self.gradient_checkpointing = False
+
+    def forward(self, input_ids, attention_mask, position_ids, inputs_embeds=None):
+        hidden = self.word_embeddings(input_ids) if inputs_embeds is None else inputs_embeds
+        positions = self.rotary_emb(hidden, position_ids)
+        for layer in self.layers:
+            if self.gradient_checkpointing and self.training:
+                hidden = checkpoint(layer, hidden, attention_mask, positions, use_reentrant=False)
+            else:
+                hidden = layer(hidden, attention_mask, positions)
+        return self.norm(hidden)
+
+
+class LLaDA2MoeModelLM(PreTrainedModel):
+    config_class = LLaDA2MoeConfig
+    base_model_prefix = "model"
+    _no_split_modules = ["LLaDA2MoeDecoderLayer"]
+    _supports_sdpa = True
+    supports_gradient_checkpointing = True
+    _tied_weights_keys = {"lm_head.weight": "model.word_embeddings.weight"}
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = LLaDA2MoeModel(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+
+    def _init_weights(self, module):
+        if isinstance(module, (nn.Linear, nn.Embedding)):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if getattr(module, "bias", None) is not None:
+                module.bias.data.zero_()
+            if getattr(module, "padding_idx", None) is not None:
+                module.weight.data[module.padding_idx].zero_()
+        elif isinstance(module, LLaDA2MoeGate):
+            module.reset_parameters()
+        elif isinstance(module, LLaDA2MoeRMSNorm):
+            module.weight.data.fill_(1.0)
+
+    def get_input_embeddings(self):
+        return self.model.word_embeddings
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None):
+        self.model.gradient_checkpointing = True
+
+    def forward(self, input_ids=None, attention_mask=None, position_ids=None,
+                inputs_embeds=None, logits_to_keep=None):
+        """Apply the explicit OPUT grid and project only the requested loss positions."""
+        if (input_ids is None) == (inputs_embeds is None):
+            raise ValueError("Supply exactly one of input_ids and inputs_embeds")
+        shape = input_ids.shape if input_ids is not None else inputs_embeds.shape[:2]
+        if attention_mask is None or attention_mask.ndim != 4:
+            raise ValueError("Supply an explicit 4D block attention mask")
+        if attention_mask.shape[-2:] != (shape[1], shape[1]):
+            raise ValueError("Attention mask does not match the input grid")
+        if position_ids is None:
+            raise ValueError("Supply original DMax position_ids explicitly")
+        hidden = self.model(input_ids, attention_mask, position_ids, inputs_embeds)
+        if logits_to_keep is not None:
+            # The student needs only candidate logits; the frozen teacher needs
+            # no vocabulary projection when its decoder hook collects features.
+            hidden = hidden[:, logits_to_keep if isinstance(logits_to_keep, slice) else slice(*logits_to_keep)]
+        return MoeCausalLMOutputWithPast(logits=self.lm_head(hidden))

tokenizer.json

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

tokenizer_config.json

@@ -0,0 +1,23 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "<|startoftext|>",
+  "clean_up_tokenization_spaces": false,
+  "cls_token": "[CLS]",
+  "eos_token": "<|endoftext|>",
+  "fast_tokenizer": true,
+  "gmask_token": "[gMASK]",
+  "is_local": false,
+  "local_files_only": true,
+  "mask_token": "<|mask|>",
+  "merges_file": null,
+  "model_max_length": 32768,
+  "model_specific_special_tokens": {
+    "gmask_token": "[gMASK]"
+  },
+  "pad_token": "<|endoftext|>",
+  "padding_side": "right",
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "extra_special_tokens": {
+    "gmask_token": "[gMASK]"
+  }
+}