Instructions to use DreamFoundries/Laguna-S-2.1-6bit with libraries, inference providers, notebooks, and local apps. Follow these links to get started.

Libraries

How to use DreamFoundries/Laguna-S-2.1-6bit with MLX:

# Make sure mlx-lm is installed
# pip install --upgrade mlx-lm

# Generate text with mlx-lm
from mlx_lm import load, generate

model, tokenizer = load("DreamFoundries/Laguna-S-2.1-6bit")

prompt = "Write a story about Einstein"
messages = [{"role": "user", "content": prompt}]
prompt = tokenizer.apply_chat_template(
    messages, add_generation_prompt=True
)

text = generate(model, tokenizer, prompt=prompt, verbose=True)

Notebooks
Google Colab
Kaggle
Local Apps Settings
LM Studio

How to use DreamFoundries/Laguna-S-2.1-6bit with Pi:

Start the MLX server

# Install MLX LM:
uv tool install mlx-lm
# Start a local OpenAI-compatible server:
mlx_lm.server --model "DreamFoundries/Laguna-S-2.1-6bit"

Configure the model in Pi

# Install Pi:
npm install -g @mariozechner/pi-coding-agent
# Add to ~/.pi/agent/models.json:
{
  "providers": {
    "mlx-lm": {
      "baseUrl": "http://localhost:8080/v1",
      "api": "openai-completions",
      "apiKey": "none",
      "models": [
        {
          "id": "DreamFoundries/Laguna-S-2.1-6bit"
        }
      ]
    }
  }
}

Run Pi

# Start Pi in your project directory:
pi

OpenClaw new

How to use DreamFoundries/Laguna-S-2.1-6bit with OpenClaw:

Start the MLX server

# Install MLX LM:
uv tool install mlx-lm
# Start a local OpenAI-compatible server:
mlx_lm.server --model "DreamFoundries/Laguna-S-2.1-6bit"

Configure OpenClaw

# Install OpenClaw:
npm install -g openclaw@latest
# Register the local server and set it as the default model:
openclaw onboard --non-interactive --mode local \
  --auth-choice custom-api-key \
  --custom-base-url http://127.0.0.1:8080/v1 \
  --custom-model-id "DreamFoundries/Laguna-S-2.1-6bit" \
  --custom-provider-id mlx-lm \
  --custom-compatibility openai \
  --custom-text-input \
  --accept-risk \
  --skip-health

Run OpenClaw

openclaw agent --local --agent main --message "Hello from Hugging Face"

MLX LM

How to use DreamFoundries/Laguna-S-2.1-6bit with MLX LM:

Generate or start a chat session

# Install MLX LM
uv tool install mlx-lm
# Interactive chat REPL
mlx_lm.chat --model "DreamFoundries/Laguna-S-2.1-6bit"

Run an OpenAI-compatible server

# Install MLX LM
uv tool install mlx-lm
# Start the server
mlx_lm.server --model "DreamFoundries/Laguna-S-2.1-6bit"
# Calling the OpenAI-compatible server with curl
curl -X POST "http://localhost:8000/v1/chat/completions" \
   -H "Content-Type: application/json" \
   --data '{
     "model": "DreamFoundries/Laguna-S-2.1-6bit",
     "messages": [
       {"role": "user", "content": "Hello"}
     ]
   }'

Hermes Agent

How to use DreamFoundries/Laguna-S-2.1-6bit with Hermes Agent:

Start the MLX server

# Install MLX LM:
uv tool install mlx-lm
# Start a local OpenAI-compatible server:
mlx_lm.server --model "DreamFoundries/Laguna-S-2.1-6bit"

Configure Hermes

# Install Hermes:
curl -fsSL https://hermes-agent.nousresearch.com/install.sh | bash
hermes setup
# Point Hermes at the local server:
hermes config set model.provider custom
hermes config set model.base_url http://127.0.0.1:8080/v1
hermes config set model.default DreamFoundries/Laguna-S-2.1-6bit

Run Hermes

hermes

DreamFoundries commited on 17 days ago

Commit

8f1ee35

verified ·

1 Parent(s): bde4d98

Upload MLX metadata

Browse files

Files changed (10) hide show

LICENSE.md +49 -0
README.md +34 -0
chat_template.jinja +94 -0
config.json +272 -0
configuration_laguna.py +249 -0
generation_config.json +24 -0
model.safetensors.index.json +0 -0
modeling_laguna.py +886 -0
tokenizer.json +0 -0
tokenizer_config.json +16 -0

LICENSE.md ADDED Viewed

	@@ -0,0 +1,49 @@

+OpenMDW License Agreement, version 1.1 (OpenMDW-1.1)
+By exercising rights granted to you under this agreement, you accept and agree
+to its terms.
+As used in this agreement, "Model Materials" means the materials provided to
+you under this agreement, consisting of: (1) one or more machine learning
+models (including architecture and parameters); and (2) all related artifacts
+(including associated data, documentation and software) that are provided to
+you hereunder.
+Subject to your compliance with this agreement, permission is hereby granted,
+free of charge, to deal in the Model Materials without restriction, including
+under all copyright, patent, database, and trade secret rights included or
+embodied therein.
+If you distribute any portion of the Model Materials, you shall retain in your
+distribution (1) a copy of this agreement, and (2) all copyright notices and
+other notices of origin included in the Model Materials that are applicable to
+your distribution.
+If you file, maintain, or voluntarily participate in a lawsuit against any
+person or entity asserting that the Model Materials directly or indirectly
+infringe any patent or copyright, then all rights and grants made to you
+hereunder are terminated, unless that lawsuit was in response to a
+corresponding lawsuit first brought against you.
+This agreement does not impose any restrictions or obligations with respect to
+any use, modification, or sharing of any outputs generated by using the Model
+Materials.
+THE MODEL MATERIALS ARE PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
+OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE, TITLE, NONINFRINGEMENT, ACCURACY, OR THE
+ABSENCE OF LATENT OR OTHER DEFECTS OR ERRORS, WHETHER OR NOT DISCOVERABLE, ALL
+TO THE GREATEST EXTENT PERMISSIBLE UNDER APPLICABLE LAW.
+YOU ARE SOLELY RESPONSIBLE FOR (1) CLEARING RIGHTS OF OTHER PERSONS THAT MAY
+APPLY TO THE MODEL MATERIALS OR ANY USE THEREOF, INCLUDING WITHOUT LIMITATION
+ANY PERSON'S COPYRIGHTS OR OTHER RIGHTS INCLUDED OR EMBODIED IN THE MODEL
+MATERIALS; (2) OBTAINING ANY NECESSARY CONSENTS, PERMISSIONS OR OTHER RIGHTS
+REQUIRED FOR ANY USE OF THE MODEL MATERIALS; OR (3) PERFORMING ANY DUE
+DILIGENCE OR UNDERTAKING ANY OTHER INVESTIGATIONS INTO THE MODEL MATERIALS OR
+ANYTHING INCORPORATED OR EMBODIED THEREIN.
+IN NO EVENT SHALL THE PROVIDERS OF THE MODEL MATERIALS BE LIABLE FOR ANY CLAIM,
+DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
+OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE MODEL MATERIALS, THE
+USE THEREOF OR OTHER DEALINGS THEREIN.

README.md ADDED Viewed

	@@ -0,0 +1,34 @@

+---
+library_name: mlx
+base_model: poolside/Laguna-S-2.1
+license: openmdw-1.1
+pipeline_tag: text-generation
+tags:
+- mlx
+- mlx-lm
+- laguna
+- quantized
+---
+# Laguna-S-2.1-6bit
+This repository contains an MLX conversion of [poolside/Laguna-S-2.1](https://huggingface.co/poolside/Laguna-S-2.1).
+## Conversion Details
+- Original model: `poolside/Laguna-S-2.1`
+- Model family: Laguna S 2.1
+- Model size: 118B total parameters, approximately 8B activated parameters per token
+- Conversion: MLX-LM conversion using the Laguna architecture implementation
+- Quantization: 6-bit affine quantization
+- Quantization policy: Uniform 6-bit affine quantization
+- Group size: 64
+- Published MLX package size: 89.01 GiB
+## Benchmarks
+No comparative benchmarks are available yet. This repository does not currently provide quality, speed, memory, or benchmark comparisons against the original weights or other quantizations.
+## License
+This is a converted and quantized derivative of the original checkpoint. It retains the upstream [OpenMDW-1.1 license](https://huggingface.co/poolside/Laguna-S-2.1/blob/main/LICENSE.md) and applicable notices of origin.

chat_template.jinja ADDED Viewed

	@@ -0,0 +1,94 @@

+{#- Iteration on laguna_glm_thinking_v8/chat_template.jinja -#}
+{#- No formatting instructions -#}
+{{- "〈|EOS|〉" -}}
+{%- set enable_thinking = enable_thinking | default(true) -%}
+{%- set add_generation_prompt = add_generation_prompt | default(false) -%}
+{%- set preserve_thinking = preserve_thinking | default(false) -%}
+{#- ───── header (system message) ───── -#}
+{#- A caller-supplied system message with empty content opts out of the default below, producing no <system> block — used to train without a system message. -#}
+{%- set system_message = "You are a helpful, conversationally-fluent assistant made by Poolside. You are here to be helpful to users through natural language conversations." -%}
+{%- if messages and messages[0].role == "system" -%}
+  {%- set system_message = messages[0].content -%}
+  {%- set messages = messages[1:] -%}
+{%- endif -%}
+{%- set has_sys = system_message and system_message.strip() -%}
+{%- if has_sys or tools or enable_thinking -%}
+  {{- "<system>" -}}
+  {%- if has_sys -%}
+    {{- system_message.rstrip() -}}
+    {%- if tools -%}{{- "\n\n" -}}{%- endif -%}
+  {%- endif -%}
+  {%- if tools -%}
+    {{- "### Tools\n\n" -}}
+    {{- "You may call functions to assist with the user query.\n" -}}
+    {{- "All available function signatures are listed below:\n" -}}
+    {{- "<available_tools>\n" -}}
+    {%- for tool in tools -%}
+      {{- (tool | tojson) ~ "\n" -}}
+    {%- endfor -%}
+    {{- "</available_tools>" -}}
+  {%- endif -%}
+  {{- "</system>\n" -}}
+{%- endif -%}
+{#- ───── main loop ───── -#}
+{%- for message in messages -%}
+  {%- set content = message.content if message.content is string else "" -%}
+  {%- if message.role == "user" -%}
+    {{- "<user>" + content + "</user>\n" -}}
+  {%- elif message.role == "assistant" -%}
+    {%- generation -%}
+      {{- "<assistant>" -}}
+      {#- Extract reasoning content from message.reasoning (vLLM field name) or message.reasoning_content -#}
+      {%- set reasoning_content = '' -%}
+      {%- if message.reasoning is string -%}
+        {%- set reasoning_content = message.reasoning -%}
+      {%- elif message.reasoning_content is string -%}
+        {%- set reasoning_content = message.reasoning_content -%}
+      {%- endif -%}
+      {#- Display reasoning content for all messages if enable_thinking -#}
+      {%- if enable_thinking or preserve_thinking -%}
+        {{- '<think>' + reasoning_content + '</think>' -}}
+      {%- else -%}
+        {{- '</think>' -}}
+      {%- endif -%}
+      {#- Display main content (trailing newline only when no tool_calls follow) -#}
+      {%- if content -%}
+        {{- content -}}
+      {%- endif -%}
+      {%- if message.tool_calls -%}
+        {%- for tool_call in message.tool_calls -%}
+          {%- set function_data = tool_call.function -%}
+          {{- '<tool_call>' + function_data.name -}}
+          {%- set _args = function_data.arguments -%}
+          {%- for k, v in _args.items() -%}
+            {{- "<arg_key>" ~ k ~ "</arg_key>" -}}
+            {{- "<arg_value>" -}}{{- v | tojson(ensure_ascii=False) if v is not string else v -}}{{- "</arg_value>" -}}
+          {%- endfor -%}
+          {{- "</tool_call>" -}}
+        {%- endfor -%}
+      {%- endif -%}
+      {{- "</assistant>\n" -}}
+    {%- endgeneration -%}
+  {%- elif message.role == "tool" -%}
+    {{- "<tool_response>" + content + "</tool_response>\n" -}}
+  {%- elif message.role == "system" -%}
+    {#- Render additional system messages (the first one, if any, is handled separately in the header and was sliced off above) -#}
+    {{- "<system>" + content + "</system>\n" -}}
+  {%- endif -%}
+{%- endfor -%}
+{#- ───── generation prompt ───── -#}
+{%- if add_generation_prompt -%}
+  {{- "<assistant>" -}}
+  {#- ───── Include reasoning mode directive ───── -#}
+  {%- if enable_thinking -%}
+    {{- '<think>' -}}
+  {%- else -%}
+    {{- '</think>' -}}
+  {%- endif -%}
+{%- endif -%}

config.json ADDED Viewed

	@@ -0,0 +1,272 @@

+{
+    "architectures": [
+        "LagunaForCausalLM"
+    ],
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "auto_map": {
+        "AutoConfig": "configuration_laguna.LagunaConfig",
+        "AutoModelForCausalLM": "modeling_laguna.LagunaForCausalLM"
+    },
+    "bos_token_id": 2,
+    "decoder_sparse_step": 1,
+    "eos_token_id": [
+        2,
+        24
+    ],
+    "gating": "per-head",
+    "gating_types": [
+        "per_head",
+        "per_head",
+        "per_head",
+        "per_head",
+        "per_head",
+        "per_head",
+        "per_head",
+        "per_head",
+        "per_head",
+        "per_head",
+        "per_head",
+        "per_head",
+        "per_head",
+        "per_head",
+        "per_head",
+        "per_head",
+        "per_head",
+        "per_head",
+        "per_head",
+        "per_head",
+        "per_head",
+        "per_head",
+        "per_head",
+        "per_head",
+        "per_head",
+        "per_head",
+        "per_head",
+        "per_head",
+        "per_head",
+        "per_head",
+        "per_head",
+        "per_head",
+        "per_head",
+        "per_head",
+        "per_head",
+        "per_head",
+        "per_head",
+        "per_head",
+        "per_head",
+        "per_head",
+        "per_head",
+        "per_head",
+        "per_head",
+        "per_head",
+        "per_head",
+        "per_head",
+        "per_head",
+        "per_head"
+    ],
+    "head_dim": 128,
+    "hidden_size": 3072,
+    "intermediate_size": 12288,
+    "layer_types": [
+        "full_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "full_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "full_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "full_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "full_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "full_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "full_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "full_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "full_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "full_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "full_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "full_attention",
+        "sliding_attention",
+        "sliding_attention",
+        "sliding_attention"
+    ],
+    "max_position_embeddings": 1048576,
+    "mlp_layer_types": [
+        "dense",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse",
+        "sparse"
+    ],
+    "mlp_only_layers": [
+        0
+    ],
+    "model_type": "laguna",
+    "moe_apply_router_weight_on_input": false,
+    "moe_intermediate_size": 1024,
+    "moe_routed_scaling_factor": 2.5,
+    "moe_router_logit_softcapping": 0.0,
+    "norm_topk_prob": true,
+    "num_attention_heads": 48,
+    "num_attention_heads_per_layer": [
+        48,
+        72,
+        72,
+        72,
+        48,
+        72,
+        72,
+        72,
+        48,
+        72,
+        72,
+        72,
+        48,
+        72,
+        72,
+        72,
+        48,
+        72,
+        72,
+        72,
+        48,
+        72,
+        72,
+        72,
+        48,
+        72,
+        72,
+        72,
+        48,
+        72,
+        72,
+        72,
+        48,
+        72,
+        72,
+        72,
+        48,
+        72,
+        72,
+        72,
+        48,
+        72,
+        72,
+        72,
+        48,
+        72,
+        72,
+        72
+    ],
+    "num_experts": 256,
+    "num_experts_per_tok": 10,
+    "num_hidden_layers": 48,
+    "num_key_value_heads": 8,
+    "pad_token_id": 9,
+    "quantization": {
+        "group_size": 64,
+        "bits": 6,
+        "mode": "affine"
+    },
+    "quantization_config": {
+        "group_size": 64,
+        "bits": 6,
+        "mode": "affine"
+    },
+    "rms_norm_eps": 1e-06,
+    "rope_parameters": {
+        "full_attention": {
+            "rope_theta": 500000.0,
+            "rope_type": "yarn",
+            "factor": 128.0,
+            "original_max_position_embeddings": 8192,
+            "beta_slow": 1.0,
+            "beta_fast": 32.0,
+            "attention_factor": 1.4852030263919618,
+            "partial_rotary_factor": 0.5
+        },
+        "sliding_attention": {
+            "rope_type": "default",
+            "rope_theta": 10000.0,
+            "partial_rotary_factor": 1.0
+        }
+    },
+    "router_aux_loss_coef": 0.0,
+    "shared_expert_intermediate_size": 1024,
+    "sliding_window": 512,
+    "tie_word_embeddings": false,
+    "torch_dtype": "bfloat16",
+    "use_cache": true,
+    "vocab_size": 100352
+}

configuration_laguna.py ADDED Viewed

	@@ -0,0 +1,249 @@

+# ruff: noqa
+# Copyright 2025 Poolside 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.
+from transformers.configuration_utils import PreTrainedConfig
+from transformers.modeling_rope_utils import RopeParameters
+from transformers.utils.import_utils import is_causal_conv1d_available, is_flash_linear_attention_available
+class LagunaConfig(PreTrainedConfig):
+    r"""
+    Configuration class for Laguna model.
+    Laguna is Poolside's MoE architecture with:
+    - Attention output gating (softplus gate)
+    - Sigmoid routing instead of softmax
+    - No QKV bias
+    - Explicit head_dim parameter
+    Args:
+        head_dim (`int`, *optional*, defaults to 128):
+            Dimension of attention heads. Laguna uses explicit head_dim rather than
+            computing it from hidden_size // num_attention_heads.
+        qkv_bias (`bool`, *optional*, defaults to `False`):
+            Whether to add bias to QKV projections. Laguna uses no QKV bias.
+        attention_bias (`bool`, *optional*, defaults to `False`):
+            Whether to add bias to attention output projection. Laguna uses no attention bias.
+        gating (`bool` or `str`, *optional*, defaults to `True`):
+            Attention output gating mode. When ``True`` or ``"per-element"`` a g_proj
+            linear layer with output size ``num_attention_heads * head_dim`` is added
+            and ``attn_output = attn_output * softplus(g_proj(x))``. When ``"per-head"``
+            g_proj has output size ``num_attention_heads`` and the gate broadcasts across
+            ``head_dim``. When ``False`` no gating is applied.
+        partial_rotary_factor (`float`, *optional*):
+            Fraction of head_dim to apply rotary embeddings to. When set, this value is
+            injected into ``rope_parameters`` (and ``swa_rope_parameters``) if not already
+            specified there. When ``None`` the default behaviour of the rope implementation
+            is used (typically full rotary).
+        num_attention_heads_per_layer (`list[int]`, *optional*):
+            Optional per-layer override for ``num_attention_heads``. When provided the list
+            length must equal ``num_hidden_layers`` and each entry is the head count used by
+            that layer. When ``None`` every layer uses ``num_attention_heads``.
+        vocab_size (`int`, *optional*, defaults to 100352):
+            Vocabulary size of the Laguna model.
+        hidden_size (`int`, *optional*, defaults to 2048):
+            Dimension of the hidden representations.
+        intermediate_size (`int`, *optional*, defaults to 8192):
+            Dimension of the MLP representations for dense layers.
+        num_hidden_layers (`int`, *optional*, defaults to 48):
+            Number of hidden layers in the Transformer.
+        num_attention_heads (`int`, *optional*, defaults to 32):
+            Number of attention heads.
+        num_key_value_heads (`int`, *optional*, defaults to 8):
+            Number of key-value heads for GQA.
+        max_position_embeddings (`int`, *optional*, defaults to 4096):
+            Maximum sequence length.
+        rms_norm_eps (`float`, *optional*, defaults to 1e-6):
+            Epsilon for RMSNorm layers.
+        sliding_window (`int`, *optional*):
+            Sliding window attention size. Used by layers whose type in ``layer_types``
+            is ``"sliding_attention"``. When ``None``, all layers use full attention.
+        layer_types (`list[str]`, *optional*):
+            Per-layer attention type. Each element should be ``"sliding_attention"`` or
+            ``"full_attention"``. Length must equal ``num_hidden_layers``. When ``None``,
+            all layers default to global attention.
+        swa_attention_sink_enabled (`bool`, *optional*, defaults to `False`):
+            Whether to enable learnable attention sinks on sliding-window attention layers.
+            When enabled, a per-head bias parameter is added that allows the model to attend
+            to position 0 even when it falls outside the sliding window.
+        swa_rope_parameters (`RopeParameters`, *optional*):
+            Separate RoPE configuration for sliding-window attention layers. When ``None``,
+            SWA layers use the same RoPE as global attention layers.
+        num_experts (`int`, *optional*, defaults to 256):
+            Number of routed experts.
+        num_experts_per_tok (`int`, *optional*, defaults to 16):
+            Number of experts selected per token (top-k).
+        moe_intermediate_size (`int`, *optional*, defaults to 1024):
+            Intermediate size of routed experts.
+        shared_expert_intermediate_size (`int`, *optional*, defaults to 1024):
+            Intermediate size of the shared expert.
+        norm_topk_prob (`bool`, *optional*, defaults to `True`):
+            Whether to normalize top-k routing probabilities.
+        decoder_sparse_step (`int`, *optional*, defaults to 1):
+            Frequency of MoE layers (1 = every layer is MoE after mlp_only_layers).
+        mlp_only_layers (`list[int]`, *optional*, defaults to `[0]`):
+            Layer indices that use dense MLP instead of MoE.
+        router_aux_loss_coef (`float`, *optional*, defaults to 0.001):
+            Auxiliary loss coefficient for load balancing.
+        moe_routed_scaling_factor (`float`, *optional*, defaults to 1.0):
+            Scalar multiplier applied to the routed-expert output before combining with the
+            shared-expert output.
+        moe_apply_router_weight_on_input (`bool`, *optional*, defaults to `False`):
+            When ``True`` the top-k routing weights are multiplied into each expert's input
+            rather than its output. Matches the numerical form used by the trained checkpoint.
+        moe_router_logit_softcapping (`float`, *optional*, defaults to 0.0):
+            Optional soft-capping value ``c`` applied to router logits as
+            ``x = tanh(x / c) * c`` before sigmoid + top-k. Disabled when ``0``.
+        rope_parameters (`RopeParameters`, *optional*):
+            RoPE configuration. Defaults to rope_theta=500000.0.
+    """
+    model_type = "laguna"
+    keys_to_ignore_at_inference = ["past_key_values"]
+    # PreTrainedConfig in transformers v5 no longer auto-declares these; subclasses
+    # opt in by providing class-level annotations with defaults.
+    pad_token_id: int | None = None
+    bos_token_id: int | None = None
+    eos_token_id: int | list[int] | None = None
+    base_model_tp_plan = {
+        "layers.*.self_attn.q_proj": "colwise",
+        "layers.*.self_attn.k_proj": "colwise",
+        "layers.*.self_attn.v_proj": "colwise",
+        "layers.*.self_attn.g_proj": "colwise",  # Laguna-specific gating projection
+        "layers.*.self_attn.o_proj": "rowwise",
+        "layers.*.mlp.gate_proj": "colwise",
+        "layers.*.mlp.up_proj": "colwise",
+        "layers.*.mlp.down_proj": "rowwise",
+    }
+    base_model_pp_plan = {
+        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+        "norm": (["hidden_states"], ["hidden_states"]),
+    }
+    def __init__(
+        self,
+        vocab_size: int = 100352,
+        hidden_size: int = 2048,
+        intermediate_size: int = 8192,
+        num_hidden_layers: int = 48,
+        num_attention_heads: int = 32,
+        num_key_value_heads: int = 8,
+        head_dim: int = 128,
+        qkv_bias: bool = False,
+        attention_bias: bool = False,
+        gating: bool | str = True,
+        hidden_act: str = "silu",
+        max_position_embeddings: int = 4096,
+        initializer_range: float = 0.02,
+        rms_norm_eps: float = 1e-6,
+        use_cache: bool = True,
+        tie_word_embeddings: bool = False,
+        rope_parameters: RopeParameters | dict[str, RopeParameters] | None = None,
+        partial_rotary_factor: float | None = None,
+        attention_dropout: float = 0.0,
+        sliding_window: int | None = None,
+        layer_types: list[str] | None = None,
+        num_attention_heads_per_layer: list[int] | None = None,
+        swa_attention_sink_enabled: bool = False,
+        swa_rope_parameters: RopeParameters | None = None,
+        num_experts: int = 256,
+        num_experts_per_tok: int = 16,
+        moe_intermediate_size: int = 1024,
+        shared_expert_intermediate_size: int = 1024,
+        norm_topk_prob: bool = True,
+        decoder_sparse_step: int = 1,
+        mlp_only_layers: list[int] | None = None,
+        router_aux_loss_coef: float = 0.001,
+        moe_routed_scaling_factor: float = 1.0,
+        moe_apply_router_weight_on_input: bool = False,
+        moe_router_logit_softcapping: float = 0.0,
+        output_router_logits: bool = False,
+        **kwargs,
+    ):
+        # Default mlp_only_layers: first layer is dense (moe_first_k_dense_replace=1)
+        if mlp_only_layers is None:
+            mlp_only_layers = [0]
+        # Default layer_types: all layers use full attention (Laguna-M). Laguna-XS
+        # ships an explicit list with a mix of "full_attention" and "sliding_attention".
+        # Downstream mask builders (``create_masks_for_generate``) iterate
+        # ``layer_types``, so it must be a list — not left as ``None``.
+        if layer_types is None:
+            layer_types = ["full_attention"] * num_hidden_layers
+        # Default rope_parameters with Laguna's theta
+        if rope_parameters is None:
+            rope_parameters = {"rope_type": "default", "rope_theta": 500000.0}
+        # config.json stores SWA rope nested in rope_parameters["sliding_attention"]
+        # and carries no top-level swa_rope_parameters. Derive it here, else the
+        # sliding-window layers silently reuse the full-attention rope.
+        if swa_rope_parameters is None and isinstance(rope_parameters, dict):
+            swa_rope_parameters = rope_parameters.get("sliding_attention")
+        # If ``partial_rotary_factor`` is set at the top level, inject it into any
+        # rope dict that does not already carry one so the rotary embedding picks
+        # it up consistently for both full-attention and SWA layers.
+        if partial_rotary_factor is not None:
+            if isinstance(rope_parameters, dict) and "partial_rotary_factor" not in rope_parameters:
+                rope_parameters = {**rope_parameters, "partial_rotary_factor": partial_rotary_factor}
+            if isinstance(swa_rope_parameters, dict) and "partial_rotary_factor" not in swa_rope_parameters:
+                swa_rope_parameters = {
+                    **swa_rope_parameters,
+                    "partial_rotary_factor": partial_rotary_factor,
+                }
+        self.vocab_size = vocab_size
+        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.num_key_value_heads = num_key_value_heads
+        self.head_dim = head_dim
+        self.qkv_bias = qkv_bias
+        self.attention_bias = attention_bias
+        self.gating = gating
+        self.hidden_act = hidden_act
+        self.max_position_embeddings = max_position_embeddings
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.rope_parameters = rope_parameters
+        self.partial_rotary_factor = partial_rotary_factor
+        self.attention_dropout = attention_dropout
+        # Sliding window attention arguments
+        self.sliding_window = sliding_window
+        self.layer_types = layer_types
+        self.num_attention_heads_per_layer = num_attention_heads_per_layer
+        self.swa_attention_sink_enabled = swa_attention_sink_enabled
+        self.swa_rope_parameters = swa_rope_parameters
+        # MoE arguments
+        self.num_experts = num_experts
+        self.num_experts_per_tok = num_experts_per_tok
+        self.moe_intermediate_size = moe_intermediate_size
+        self.shared_expert_intermediate_size = shared_expert_intermediate_size
+        self.norm_topk_prob = norm_topk_prob
+        self.decoder_sparse_step = decoder_sparse_step
+        self.mlp_only_layers = mlp_only_layers
+        self.router_aux_loss_coef = router_aux_loss_coef
+        self.moe_routed_scaling_factor = moe_routed_scaling_factor
+        self.moe_apply_router_weight_on_input = moe_apply_router_weight_on_input
+        self.moe_router_logit_softcapping = moe_router_logit_softcapping
+        self.output_router_logits = output_router_logits
+        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
+__all__ = ["LagunaConfig"]

generation_config.json ADDED Viewed

	@@ -0,0 +1,24 @@

+{
+  "bos_token_id": 2,
+  "do_sample": true,
+  "eos_token_id": [
+    2,
+    24
+  ],
+  "pad_token_id": 9,
+  "temperature": 1.0,
+  "top_p": 1.0,
+  "min_p": 0.0,
+  "top_k": 20,
+  "tool_call_parser": "poolside_v1",
+  "reasoning_parser": "poolside_v1",
+  "default_chat_template_kwargs": {
+    "enable_thinking": true
+  },
+  "speculative_config": {
+    "method": "dflash",
+    "source": "huggingface",
+    "model": "poolside/Laguna-S-2.1-DFlash",
+    "num_speculative_tokens": 15
+  }
+}

model.safetensors.index.json ADDED Viewed

The diff for this file is too large to render. See raw diff

modeling_laguna.py ADDED Viewed

	@@ -0,0 +1,886 @@

+# ruff: noqa
+# Copyright 2025 Poolside 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.
+from collections.abc import Callable
+import torch
+import torch.nn.functional as F
+from torch import nn
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache
+from transformers.integrations import use_experts_implementation, use_kernelized_func
+from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
+from transformers.modeling_layers import GradientCheckpointingLayer
+from transformers.modeling_outputs import MoeModelOutputWithPast
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+from transformers.processing_utils import Unpack
+from transformers.utils import auto_docstring, can_return_tuple, is_grouped_mm_available
+from transformers.utils.generic import TransformersKwargs, merge_with_config_defaults
+from transformers.utils.output_capturing import OutputRecorder, capture_outputs
+from transformers.cache_utils import DynamicCache
+from transformers.generation import GenerationMixin
+from transformers.integrations import use_kernel_forward_from_hub
+from transformers.masking_utils import create_causal_mask
+from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils.generic import maybe_autocast
+from .configuration_laguna import LagunaConfig
+from transformers import initialization as init
+from transformers.masking_utils import create_sliding_window_causal_mask
+from transformers.modeling_outputs import MoeCausalLMOutputWithPast
+from transformers.utils.import_utils import is_causal_conv1d_available, is_flash_linear_attention_available
+@use_kernel_forward_from_hub("RMSNorm")
+class LagunaRMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps: float = 1e-6) -> None:
+        """
+        LagunaRMSNorm is equivalent to T5LayerNorm
+        """
+        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):
+        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+class LagunaRotaryEmbedding(nn.Module):
+    inv_freq: torch.Tensor  # fix linting for `register_buffer`
+    def __init__(self, config: LagunaConfig, device=None):
+        super().__init__()
+        self.max_seq_len_cached = config.max_position_embeddings
+        self.original_max_seq_len = config.max_position_embeddings
+        self.config = config
+        self.rope_type = self.config.rope_parameters["rope_type"]
+        rope_init_fn: Callable = self.compute_default_rope_parameters
+        if self.rope_type != "default":
+            rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+        inv_freq, self.attention_scaling = rope_init_fn(self.config, device)
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False)
+    @staticmethod
+    def compute_default_rope_parameters(config, device=None, seq_len=None) -> tuple["torch.Tensor", float]:
+        """
+        Computes the inverse frequencies according to the original RoPE implementation
+        Args:
+            config ([`~transformers.PreTrainedConfig`]):
+                The model configuration.
+            device (`torch.device`):
+                The device to use for initialization of the inverse frequencies.
+            seq_len (`int`, *optional*):
+                The current sequence length. Unused for this type of RoPE.
+        Returns:
+            Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
+            post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
+        """
+        base = config.rope_parameters["rope_theta"]
+        head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+        partial = config.rope_parameters.get("partial_rotary_factor", 1.0)
+        dim = int(head_dim * partial)
+        inv_freq = 1.0 / (
+            base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
+        )
+        return inv_freq, 1.0
+    @torch.no_grad()
+    @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)
+    def forward(self, x, position_ids):
+        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
+        position_ids_expanded = position_ids[:, None, :].float()
+        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+        with maybe_autocast(device_type=device_type, enabled=False):  # Force float32
+            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cos = emb.cos() * self.attention_scaling
+            sin = emb.sin() * self.attention_scaling
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+class LagunaMLP(nn.Module):
+    def __init__(self, config, intermediate_size=None):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size if intermediate_size is None else 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):
+        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+        return down_proj
+class LagunaTopKRouter(nn.Module):
+    """Laguna MoE router using sigmoid scoring (not softmax).
+    Supports optional router-logit soft-capping and auxiliary-loss-free load
+    balancing (arXiv:2408.15664): the per-expert bias ``e_score_correction_bias``
+    is added to selection scores but the returned routing weights remain unbiased.
+    The bias lives on the router so accelerate's per-module hooks can co-locate it
+    with the gate — moving it to the experts module would cross a hook boundary
+    and leave the bias on meta under ``device_map="auto"`` / CPU-offload.
+    """
+    def __init__(self, config):
+        super().__init__()
+        self.top_k = config.num_experts_per_tok
+        self.num_experts = config.num_experts
+        self.norm_topk_prob = config.norm_topk_prob
+        self.hidden_dim = config.hidden_size
+        self.weight = nn.Parameter(torch.zeros(self.num_experts, self.hidden_dim))
+        # Zero-initialised so inference on checkpoints that don't ship the bias
+        # is a no-op. ``_checkpoint_conversion_mapping`` below remaps the
+        # ``mlp.experts.e_score_correction_bias`` key from vLLM-trained
+        # checkpoints onto this attribute.
+        self.e_score_correction_bias = nn.Parameter(torch.zeros(config.num_experts), requires_grad=False)
+        self.router_logit_softcapping = float(getattr(config, "moe_router_logit_softcapping", 0.0) or 0.0)
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        hidden_states = hidden_states.reshape(-1, self.hidden_dim)
+        router_logits = F.linear(hidden_states, self.weight).float()
+        if self.router_logit_softcapping > 0.0:
+            router_logits = torch.tanh(router_logits / self.router_logit_softcapping) * self.router_logit_softcapping
+        routing_scores = torch.sigmoid(router_logits)
+        scores_for_selection = routing_scores + self.e_score_correction_bias.to(routing_scores.dtype)
+        _, selected_experts = torch.topk(scores_for_selection, self.top_k, dim=-1)
+        routing_weights = routing_scores.gather(-1, selected_experts)
+        if self.norm_topk_prob:
+            routing_weights = routing_weights / routing_weights.sum(dim=-1, keepdim=True)
+        routing_weights = routing_weights.to(hidden_states.dtype)
+        return router_logits, routing_weights, selected_experts
+@use_experts_implementation
+class LagunaExperts(nn.Module):
+    """Fused expert weights as 3D tensors for batched execution."""
+    def __init__(self, config):
+        super().__init__()
+        self.num_experts = config.num_experts
+        self.hidden_dim = config.hidden_size
+        self.intermediate_dim = config.moe_intermediate_size
+        self.gate_up_proj = nn.Parameter(torch.empty(self.num_experts, 2 * self.intermediate_dim, self.hidden_dim))
+        self.down_proj = nn.Parameter(torch.empty(self.num_experts, self.hidden_dim, self.intermediate_dim))
+        self.act_fn = ACT2FN[config.hidden_act]
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        top_k_index: torch.Tensor,
+        top_k_weights: torch.Tensor,
+    ) -> torch.Tensor:
+        final_hidden_states = torch.zeros_like(hidden_states)
+        with torch.no_grad():
+            expert_mask = F.one_hot(top_k_index, num_classes=self.num_experts)
+            expert_mask = expert_mask.permute(2, 1, 0)
+            expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
+        for expert_idx in expert_hit:
+            expert_idx = expert_idx[0]
+            if expert_idx == self.num_experts:
+                continue
+            top_k_pos, token_idx = torch.where(expert_mask[expert_idx])
+            current_state = hidden_states[token_idx]
+            gate, up = F.linear(current_state, self.gate_up_proj[expert_idx]).chunk(2, dim=-1)
+            current_hidden_states = self.act_fn(gate) * up
+            current_hidden_states = F.linear(current_hidden_states, self.down_proj[expert_idx])
+            current_hidden_states = current_hidden_states * top_k_weights[token_idx, top_k_pos, None]
+            final_hidden_states.index_add_(0, token_idx, current_hidden_states.to(final_hidden_states.dtype))
+        return final_hidden_states
+class LagunaSparseMoeBlock(nn.Module):
+    """Laguna MoE block using sigmoid router, fused expert tensors, and a shared expert."""
+    def __init__(self, config):
+        super().__init__()
+        self.num_experts = config.num_experts
+        self.routed_scaling_factor = float(getattr(config, "moe_routed_scaling_factor", 1.0))
+        # ``moe_apply_router_weight_on_input=True`` would require scaling each expert's
+        # input (rather than its output) by the routing weight. Supporting it cleanly
+        # alongside the fused experts kernels (``grouped_mm`` / ``batched_mm``) is future
+        # work; for now we fail loudly so a checkpoint that needs it can't silently
+        # diverge from its numerical form.
+        if getattr(config, "moe_apply_router_weight_on_input", False):
+            raise NotImplementedError(
+                "moe_apply_router_weight_on_input=True is not yet supported in the "
+                "transformers implementation of Laguna."
+            )
+        self.gate = LagunaTopKRouter(config)
+        self.experts = LagunaExperts(config)
+        self.shared_expert = LagunaMLP(config, intermediate_size=config.shared_expert_intermediate_size)
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        batch_size, sequence_length, hidden_dim = hidden_states.shape
+        hidden_states = hidden_states.view(-1, hidden_dim)
+        shared_expert_output = self.shared_expert(hidden_states)
+        _, routing_weights, selected_experts = self.gate(hidden_states)
+        expert_output = self.experts(hidden_states, selected_experts, routing_weights)
+        if self.routed_scaling_factor != 1.0:
+            expert_output = expert_output * self.routed_scaling_factor
+        expert_output = expert_output + shared_expert_output
+        expert_output = expert_output.reshape(batch_size, sequence_length, hidden_dim)
+        return expert_output
+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)
+# Adapted from transformers.models.glm.modular_glm.apply_rotary_pos_emb
+def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
+    """Applies Rotary Position Embedding to the query and key tensors.
+    Removes the interleaving of cos and sin from GLM
+    Args:
+        q (`torch.Tensor`): The query tensor.
+        k (`torch.Tensor`): The key tensor.
+        cos (`torch.Tensor`): The cosine part of the rotary embedding.
+        sin (`torch.Tensor`): The sine part of the rotary embedding.
+        unsqueeze_dim (`int`, *optional*, defaults to 1):
+            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+    Returns:
+        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+    """
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    # 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
+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)
+def eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: torch.Tensor | None,
+    scaling: float,
+    dropout: float = 0.0,
+    **kwargs: Unpack[TransformersKwargs],
+):
+    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:
+        attn_weights = attn_weights + attention_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
+# Laguna attention is identical to Qwen2MoE attention except:
+# - No QKV bias
+# - Explicit head_dim from config
+# - Output gating: attn_output = attn_output * softplus(g_proj(hidden_states)) (optional)
+# - Per-layer sliding window attention with optional attention sinks
+@use_kernelized_func(apply_rotary_pos_emb)
+class LagunaAttention(nn.Module):
+    def __init__(self, config: LagunaConfig, layer_idx: int, num_heads: int | None = None):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.head_dim = config.head_dim
+        # Allow the caller (decoder layer) to supply a per-layer head count; fall back
+        # to config.num_attention_heads when not provided.
+        self.num_heads = num_heads if num_heads is not None else config.num_attention_heads
+        self.num_key_value_groups = self.num_heads // config.num_key_value_heads
+        self.scaling = self.head_dim**-0.5
+        self.attention_dropout = config.attention_dropout
+        self.is_causal = True
+        # Per-layer sliding window (follows Gemma2/Cohere2 convention)
+        layer_types = getattr(config, "layer_types", None)
+        if layer_types is not None:
+            self.is_sliding = layer_types[layer_idx] == "sliding_attention"
+            self.sliding_window = config.sliding_window if self.is_sliding else None
+        else:
+            self.is_sliding = False
+            self.sliding_window = None
+        # Laguna: no QKV bias, explicit head_dim
+        self.q_proj = nn.Linear(config.hidden_size, self.num_heads * config.head_dim, bias=False)
+        self.k_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * config.head_dim, bias=False)
+        self.v_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * config.head_dim, bias=False)
+        self.o_proj = nn.Linear(self.num_heads * config.head_dim, config.hidden_size, bias=False)
+        # Laguna-specific: optional gating projection.
+        # ``gating`` may be:
+        #   - True / "per-element": one gate per (head, head_dim) channel
+        #   - "per-head":           one gate per head, broadcast across head_dim
+        #   - False:                no gating
+        gating = getattr(config, "gating", True)
+        self.gating = bool(gating)
+        self.gate_per_head = gating == "per-head"
+        if self.gating:
+            g_out = self.num_heads if self.gate_per_head else self.num_heads * config.head_dim
+            self.g_proj = nn.Linear(config.hidden_size, g_out, bias=False)
+        # Attention sinks (learnable per-head bias for SWA layers)
+        if self.is_sliding and getattr(config, "swa_attention_sink_enabled", False):
+            self.sink = nn.Parameter(torch.zeros(self.num_heads))
+        # QK normalization (RMSNorm applied per-head after reshape, before RoPE)
+        self.q_norm = LagunaRMSNorm(config.head_dim, eps=config.rms_norm_eps)
+        self.k_norm = LagunaRMSNorm(config.head_dim, eps=config.rms_norm_eps)
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: torch.Tensor | None,
+        past_key_values: Cache | None = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+        query_states = query_states.view(hidden_shape).transpose(1, 2)
+        key_states = key_states.view(hidden_shape).transpose(1, 2)
+        value_states = value_states.view(hidden_shape).transpose(1, 2)
+        # QK normalization (applied per-head before RoPE)
+        query_states = self.q_norm(query_states)
+        key_states = self.k_norm(key_states)
+        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:
+            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+        # ``attention_mask`` here is already the correct mask for this layer type —
+        # ``LagunaModel.forward`` builds separate full-attention and sliding-attention
+        # masks (using ``create_causal_mask`` / ``create_sliding_window_causal_mask``)
+        # and the decoder layer passes the right one in.
+        attention_interface: Callable = eager_attention_forward
+        if self.config._attn_implementation != "eager":
+            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,
+            **kwargs,
+        )
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        # Laguna-specific: apply gating BEFORE o_proj (optional)
+        if self.gating:
+            gate = F.softplus(self.g_proj(hidden_states).float()).to(attn_output.dtype)
+            if self.gate_per_head:
+                # gate: [..., num_heads]; broadcast across head_dim
+                attn_shape = attn_output.shape
+                attn_output = (
+                    attn_output.view(*attn_shape[:-1], self.num_heads, self.head_dim) * gate.unsqueeze(-1)
+                ).view(attn_shape)
+            else:
+                attn_output = attn_output * gate
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+class LagunaDecoderLayer(GradientCheckpointingLayer):
+    """Laguna decoder layer with gated attention and sigmoid-routed MoE."""
+    def __init__(self, config: LagunaConfig, layer_idx: int):
+        super().__init__()
+        per_layer_heads = getattr(config, "num_attention_heads_per_layer", None)
+        layer_num_heads = per_layer_heads[layer_idx] if per_layer_heads is not None else config.num_attention_heads
+        # Layer type drives mask and position-embedding dispatch in ``LagunaModel.forward``.
+        layer_types = getattr(config, "layer_types", None)
+        self.attention_type = layer_types[layer_idx] if layer_types is not None else "full_attention"
+        self.self_attn = LagunaAttention(config, layer_idx, num_heads=layer_num_heads)
+        # Use MoE or dense MLP based on layer configuration
+        if (layer_idx not in config.mlp_only_layers) and (
+            config.num_experts > 0 and (layer_idx + 1) % config.decoder_sparse_step == 0
+        ):
+            self.mlp = LagunaSparseMoeBlock(config)
+        else:
+            self.mlp = LagunaMLP(config, intermediate_size=config.intermediate_size)
+        self.input_layernorm = LagunaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = LagunaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.hidden_size = config.hidden_size
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        use_cache: bool | None = False,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        # Self Attention
+        hidden_states, _ = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )
+        hidden_states = residual + hidden_states
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+        return hidden_states
+@auto_docstring
+class LagunaPreTrainedModel(PreTrainedModel):
+    config: LagunaConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["LagunaDecoderLayer"]
+    _skip_keys_device_placement = ["past_key_values"]
+    _supports_flash_attn = True
+    _supports_sdpa = True
+    _supports_flex_attn = True
+    _can_compile_fullgraph = (
+        is_grouped_mm_available()
+    )  # https://huggingface.co/docs/transformers/experts_interface#torchcompile
+    _supports_attention_backend = True
+    _can_record_outputs = {
+        "router_logits": OutputRecorder(LagunaTopKRouter, index=0),
+        "hidden_states": LagunaDecoderLayer,
+        "attentions": LagunaAttention,
+    }
+    # vLLM-trained Laguna checkpoints store the aux-loss-free routing bias on the
+    # experts module (``mlp.experts.e_score_correction_bias``). In this impl the
+    # bias lives on the router to stay co-located with its consumer across
+    # accelerate's per-module hooks, so remap the legacy key on load.
+    _checkpoint_conversion_mapping = {
+        r"^(.*)\.mlp\.experts\.e_score_correction_bias$": r"\1.mlp.gate.e_score_correction_bias",
+    }
+    @torch.no_grad()
+    def _init_weights(self, module):
+        super()._init_weights(module)
+        std = self.config.initializer_range
+        if isinstance(module, LagunaExperts):
+            init.normal_(module.gate_up_proj, mean=0.0, std=std)
+            init.normal_(module.down_proj, mean=0.0, std=std)
+        elif isinstance(module, LagunaTopKRouter):
+            init.normal_(module.weight, mean=0.0, std=std)
+        # Bare ``nn.Parameter``s that are not covered by the parent's generic
+        # Linear/Embedding/norm handling need their own rules so that the
+        # __init__ and from_pretrained(state_dict={}) paths produce identical
+        # weights under a fixed seed.
+        if isinstance(module, LagunaTopKRouter):
+            torch.nn.init.zeros_(module.e_score_correction_bias)
+        if isinstance(module, LagunaAttention) and hasattr(module, "sink"):
+            torch.nn.init.zeros_(module.sink)
+class LagunaModel(LagunaPreTrainedModel):
+    def __init__(self, config: LagunaConfig):
+        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(
+            [LagunaDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self.norm = LagunaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        # ``LagunaRotaryEmbedding`` inherits ``Qwen2MoeRotaryEmbedding``'s flat-shape
+        # contract — it reads ``config.rope_parameters["rope_type"]`` at the outer
+        # level. Laguna stores rope nested by layer type (``{"full_attention": {...},
+        # ...}``), so pass a config clone with the full-attention sub-dict flattened.
+        rp = getattr(config, "rope_parameters", None)
+        if isinstance(rp, dict) and isinstance(rp.get("full_attention"), dict):
+            import copy
+            full_config = copy.deepcopy(config)
+            full_config.rope_parameters = dict(rp["full_attention"])
+            self.rotary_emb = LagunaRotaryEmbedding(config=full_config)
+        else:
+            self.rotary_emb = LagunaRotaryEmbedding(config=config)
+        # Separate RoPE for sliding-window attention layers (when configured).
+        # Be careful with ``partial_rotary_factor`` — ``PreTrainedConfig.standardize_rope_params``
+        # unconditionally overwrites ``rope_parameters["partial_rotary_factor"]`` with
+        # ``self.partial_rotary_factor``, so we must align the top-level field on the
+        # cloned config to the SWA value, otherwise the global partial factor silently
+        # clobbers the SWA one.
+        if getattr(config, "swa_rope_parameters", None) is not None:
+            import copy
+            swa_config = copy.deepcopy(config)
+            swa_config.rope_parameters = dict(config.swa_rope_parameters)
+            swa_partial = swa_config.rope_parameters.get("partial_rotary_factor")
+            swa_config.partial_rotary_factor = swa_partial
+            self.swa_rotary_emb = LagunaRotaryEmbedding(config=swa_config)
+        else:
+            self.swa_rotary_emb = None
+        self.gradient_checkpointing = False
+        # Initialize weights and apply final processing
+        self.post_init()
+    @merge_with_config_defaults
+    @capture_outputs
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        use_cache: bool | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> MoeModelOutputWithPast:
+        from transformers.cache_utils import DynamicCache
+        from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask
+        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 position_ids is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            position_ids = (
+                torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
+            ).unsqueeze(0)
+        # Build one mask per layer-type so each layer can be dispatched with the right
+        # attention pattern (follows the afmoe / cohere2 v5 convention).
+        layer_types = getattr(self.config, "layer_types", None)
+        has_swa = layer_types is not None and "sliding_attention" in layer_types
+        if not isinstance(causal_mask_mapping := attention_mask, dict):
+            mask_kwargs = {
+                "config": self.config,
+                "inputs_embeds": inputs_embeds,
+                "attention_mask": attention_mask,
+                "past_key_values": past_key_values,
+                "position_ids": position_ids,
+            }
+            causal_mask_mapping = {"full_attention": create_causal_mask(**mask_kwargs)}
+            if has_swa:
+                causal_mask_mapping["sliding_attention"] = create_sliding_window_causal_mask(**mask_kwargs)
+        hidden_states = inputs_embeds
+        global_pe = self.rotary_emb(hidden_states, position_ids)
+        # Per-layer-type position embeddings: Laguna optionally uses a different rope for
+        # sliding layers (``swa_rope_parameters``). When absent, SWA layers share the
+        # global rope.
+        if has_swa:
+            swa_pe = self.swa_rotary_emb(hidden_states, position_ids) if self.swa_rotary_emb is not None else global_pe
+            position_embeddings_mapping = {"full_attention": global_pe, "sliding_attention": swa_pe}
+        else:
+            position_embeddings_mapping = None
+        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+            layer_attn_mask = causal_mask_mapping[decoder_layer.attention_type]
+            layer_pos_emb = (
+                position_embeddings_mapping[decoder_layer.attention_type]
+                if position_embeddings_mapping is not None
+                else global_pe
+            )
+            hidden_states = decoder_layer(
+                hidden_states,
+                attention_mask=layer_attn_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+                position_embeddings=layer_pos_emb,
+                **kwargs,
+            )
+        hidden_states = self.norm(hidden_states)
+        return MoeModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values,
+        )
+def load_balancing_loss_func(
+    gate_logits: torch.Tensor | tuple[torch.Tensor] | None,
+    num_experts: int | None = None,
+    top_k=2,
+    attention_mask: torch.Tensor | None = None,
+) -> torch.Tensor | int:
+    r"""
+    Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch.
+    See Switch Transformer (https://huggingface.co/papers/2101.03961) for more details. This function implements the loss
+    function presented in equations (4) - (6) of the paper. It aims at penalizing cases where the routing between
+    experts is too unbalanced.
+    Args:
+        gate_logits:
+            Logits from the `gate`, should be a tuple of model.config.num_hidden_layers tensors of
+            shape [batch_size X sequence_length, num_experts].
+        num_experts:
+            Number of experts
+        top_k:
+            The number of experts to route per-token, can be also interpreted as the `top-k` routing
+            parameter.
+        attention_mask (`torch.Tensor`, *optional*):
+            The attention_mask used in forward function
+            shape [batch_size X sequence_length] if not None.
+    Returns:
+        The auxiliary loss.
+    """
+    if gate_logits is None or not isinstance(gate_logits, tuple):
+        return 0
+    if isinstance(gate_logits, tuple):
+        compute_device = gate_logits[0].device
+        concatenated_gate_logits = torch.cat([layer_gate.to(compute_device) for layer_gate in gate_logits], dim=0)
+    routing_weights = torch.nn.functional.softmax(concatenated_gate_logits, dim=-1)
+    _, selected_experts = torch.topk(routing_weights, top_k, dim=-1)
+    expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts)
+    if attention_mask is None:
+        # Compute the percentage of tokens routed to each experts
+        tokens_per_expert = torch.mean(expert_mask.float(), dim=0)
+        # Compute the average probability of routing to these experts
+        router_prob_per_expert = torch.mean(routing_weights, dim=0)
+    else:
+        batch_size, sequence_length = attention_mask.shape
+        num_hidden_layers = concatenated_gate_logits.shape[0] // (batch_size * sequence_length)
+        # Compute the mask that masks all padding tokens as 0 with the same shape of expert_mask
+        expert_attention_mask = (
+            attention_mask[None, :, :, None, None]
+            .expand((num_hidden_layers, batch_size, sequence_length, top_k, num_experts))
+            .reshape(-1, top_k, num_experts)
+            .to(compute_device)
+        )
+        # Compute the percentage of tokens routed to each experts
+        tokens_per_expert = torch.sum(expert_mask.float() * expert_attention_mask, dim=0) / torch.sum(
+            expert_attention_mask, dim=0
+        )
+        # Compute the mask that masks all padding tokens as 0 with the same shape of tokens_per_expert
+        router_per_expert_attention_mask = (
+            attention_mask[None, :, :, None]
+            .expand((num_hidden_layers, batch_size, sequence_length, num_experts))
+            .reshape(-1, num_experts)
+            .to(compute_device)
+        )
+        # Compute the average probability of routing to these experts
+        router_prob_per_expert = torch.sum(routing_weights * router_per_expert_attention_mask, dim=0) / torch.sum(
+            router_per_expert_attention_mask, dim=0
+        )
+    overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(0))
+    return overall_loss * num_experts
+@auto_docstring
+class LagunaForCausalLM(LagunaPreTrainedModel, GenerationMixin):
+    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+    _tp_plan = {"lm_head": "colwise_gather_output"}
+    _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = LagunaModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.router_aux_loss_coef = config.router_aux_loss_coef
+        self.num_experts = config.num_experts
+        self.num_experts_per_tok = config.num_experts_per_tok
+        # Initialize weights and apply final processing
+        self.post_init()
+    @can_return_tuple
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        output_router_logits: bool | None = None,
+        logits_to_keep: int | torch.Tensor = 0,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> MoeCausalLMOutputWithPast:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+        """
+        output_router_logits = (
+            output_router_logits if output_router_logits is not None else self.config.output_router_logits
+        )
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs: MoeModelOutputWithPast = 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_router_logits=output_router_logits,
+            **kwargs,
+        )
+        hidden_states = outputs.last_hidden_state
+        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+        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, :])
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(logits, labels, self.vocab_size, **kwargs)
+        aux_loss = None
+        if output_router_logits:
+            aux_loss = load_balancing_loss_func(
+                outputs.router_logits,
+                self.num_experts,
+                self.num_experts_per_tok,
+                attention_mask,
+            )
+            if labels is not None:
+                loss += self.router_aux_loss_coef * aux_loss.to(loss.device)  # make sure to reside in the same device
+        return MoeCausalLMOutputWithPast(
+            loss=loss,
+            aux_loss=aux_loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            router_logits=outputs.router_logits,
+        )
+__all__ = ["LagunaForCausalLM", "LagunaModel", "LagunaPreTrainedModel"]
+# --- Added: register the native Laguna checkpoint-conversion for trust_remote_code loads.
+# transformers >=5.12 skips checkpoint-conversion mappings for custom (remote) code
+# unless explicitly registered, which broke loading the shipped per-expert MoE weights.
+try:
+    from transformers.conversion_mapping import (
+        get_checkpoint_conversion_mapping as _lg_get,
+        register_checkpoint_conversion_mapping as _lg_reg,
+        USER_REGISTERED_MAPPINGS as _lg_user,
+    )
+    if "laguna" not in _lg_user:
+        _lg_m = _lg_get("laguna")
+        if _lg_m is not None:
+            _lg_reg("laguna", _lg_m, overwrite=True)
+except Exception:
+    pass

tokenizer.json ADDED Viewed

The diff for this file is too large to render. See raw diff

tokenizer_config.json ADDED Viewed

	@@ -0,0 +1,16 @@

+{
+  "backend": "tokenizers",
+  "bos_token": "〈|EOS|〉",
+  "clean_up_tokenization_spaces": false,
+  "cls_token": "〈|CLS|〉",
+  "eos_token": "〈|EOS|〉",
+  "is_local": true,
+  "local_files_only": false,
+  "mask_token": "〈|MASK|〉",
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "〈|PAD|〉",
+  "sep_token": "〈|SEP|〉",
+  "tokenizer_class": "TokenizersBackend",
+  "tool_parser_type": "glm47",
+  "unk_token": "〈|UNK|〉"
+}