Upload folder using huggingface_hub

.gitattributes

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

chat_template.jinja

@@ -0,0 +1,397 @@
+{#-
+  In addition to the normal inputs of `messages` and `tools`, this template also accepts the
+  following kwargs:
+  - "builtin_tools": A list, can contain "browser" and/or "python".
+  - "model_identity": A string that optionally describes the model identity.
+  - "reasoning_effort": A string that describes the reasoning effort, defaults to "medium".
+ #}
+
+{#- Tool Definition Rendering ============================================== #}
+{%- macro render_typescript_type(param_spec, required_params, is_nullable=false) -%}
+    {%- if param_spec.type == "array" -%}
+        {%- if param_spec['items'] -%}
+            {%- if param_spec['items']['type'] == "string" -%}
+                {{- "string[]" }}
+            {%- elif param_spec['items']['type'] == "number" -%}
+                {{- "number[]" }}
+            {%- elif param_spec['items']['type'] == "integer" -%}
+                {{- "number[]" }}
+            {%- elif param_spec['items']['type'] == "boolean" -%}
+                {{- "boolean[]" }}
+            {%- else -%}
+                {%- set inner_type = render_typescript_type(param_spec['items'], required_params) -%}
+                {%- if inner_type == "object | object" or inner_type|length > 50 -%}
+                    {{- "any[]" }}
+                {%- else -%}
+                    {{- inner_type + "[]" }}
+                {%- endif -%}
+            {%- endif -%}
+            {%- if param_spec.nullable -%}
+                {{- " | null" }}
+            {%- endif -%}
+        {%- else -%}
+            {{- "any[]" }}
+            {%- if param_spec.nullable -%}
+                {{- " | null" }}
+            {%- endif -%}
+        {%- endif -%}
+    {%- elif param_spec.type is defined and param_spec.type is iterable and param_spec.type is not string and param_spec.type is not mapping and param_spec.type[0] is defined -%}
+        {#- Handle array of types like ["object", "object"] from Union[dict, list] #}
+        {%- if param_spec.type | length > 1 -%}
+            {{- param_spec.type | join(" | ") }}
+        {%- else -%}
+            {{- param_spec.type[0] }}
+        {%- endif -%}
+    {%- elif param_spec.oneOf -%}
+        {#- Handle oneOf schemas - check for complex unions and fallback to any #}
+        {%- set has_object_variants = false -%}
+        {%- for variant in param_spec.oneOf -%}
+            {%- if variant.type == "object" -%}
+                {%- set has_object_variants = true -%}
+            {%- endif -%}
+        {%- endfor -%}
+        {%- if has_object_variants and param_spec.oneOf|length > 1 -%}
+            {{- "any" }}
+        {%- else -%}
+            {%- for variant in param_spec.oneOf -%}
+                {{- render_typescript_type(variant, required_params) -}}
+                {%- if variant.description %}
+                    {{- "// " + variant.description }}
+                {%- endif -%}
+                {%- if variant.default is defined %}
+                    {{ "// default: " + variant.default|tojson }}
+                {%- endif -%}
+                {%- if not loop.last %}
+                    {{- " | " }}
+                {% endif -%}
+            {%- endfor -%}
+        {%- endif -%}
+    {%- elif param_spec.type == "string" -%}
+        {%- if param_spec.enum -%}
+            {{- '"' + param_spec.enum|join('" | "') + '"' -}}
+        {%- else -%}
+            {{- "string" }}
+            {%- if param_spec.nullable %}
+                {{- " | null" }}
+            {%- endif -%}
+        {%- endif -%}
+    {%- elif param_spec.type == "number" -%}
+        {{- "number" }}
+    {%- elif param_spec.type == "integer" -%}
+        {{- "number" }}
+    {%- elif param_spec.type == "boolean" -%}
+        {{- "boolean" }}
+
+    {%- elif param_spec.type == "object" -%}
+        {%- if param_spec.properties -%}
+            {{- "{
+" }}
+            {%- for prop_name, prop_spec in param_spec.properties.items() -%}
+                {{- prop_name -}}
+                {%- if prop_name not in (param_spec.required or []) -%}
+                    {{- "?" }}
+                {%- endif -%}
+                {{- ": " }}
+                {{ render_typescript_type(prop_spec, param_spec.required or []) }}
+                {%- if not loop.last -%}
+                    {{-", " }}
+                {%- endif -%}
+            {%- endfor -%}
+            {{- "}" }}
+        {%- else -%}
+            {{- "object" }}
+        {%- endif -%}
+    {%- else -%}
+        {{- "any" }}
+    {%- endif -%}
+{%- endmacro -%}
+
+{%- macro render_tool_namespace(namespace_name, tools) -%}
+    {{- "## " + namespace_name + "
+
+" }}
+    {{- "namespace " + namespace_name + " {
+
+" }}
+    {%- for tool in tools %}
+        {%- set tool = tool.function %}
+        {{- "// " + tool.description + "
+" }}
+        {{- "type "+ tool.name + " = " }}
+        {%- if tool.parameters and tool.parameters.properties %}
+            {{- "(_: {
+" }}
+            {%- for param_name, param_spec in tool.parameters.properties.items() %}
+                {%- if param_spec.description %}
+                    {{- "// " + param_spec.description + "
+" }}
+                {%- endif %}
+                {{- param_name }}
+                {%- if param_name not in (tool.parameters.required or []) -%}
+                    {{- "?" }}
+                {%- endif -%}
+                {{- ": " }}
+                {{- render_typescript_type(param_spec, tool.parameters.required or []) }}
+                {%- if param_spec.default is defined -%}
+                    {%- if param_spec.enum %}
+                        {{- ", // default: " + param_spec.default }}
+                    {%- elif param_spec.oneOf %}
+                        {{- "// default: " + param_spec.default }}
+                    {%- else %}
+                        {{- ", // default: " + param_spec.default|tojson }}
+                    {%- endif -%}
+                {%- endif -%}
+                {%- if not loop.last %}
+                    {{- ",
+" }}
+                {%- else %}
+                    {{- "
+" }}
+                {%- endif -%}
+            {%- endfor %}
+            {{- "}) => any;
+
+" }}
+        {%- else -%}
+            {{- "() => any;
+
+" }}
+        {%- endif -%}
+    {%- endfor %}
+    {{- "} // namespace " + namespace_name }}
+{%- endmacro -%}
+
+{%- macro render_builtin_tools(browser_tool, python_tool) -%}
+    {%- if browser_tool %}
+        {{- "## browser
+
+" }}
+        {{- "// Tool for browsing.
+" }}
+        {{- "// The `cursor` appears in brackets before each browsing display: `[{cursor}]`.
+" }}
+        {{- "// Cite information from the tool using the following format:
+" }}
+        {{- "// `【{cursor}†L{line_start}(-L{line_end})?】`, for example: `【6†L9-L11】` or `【8†L3】`.
+" }}
+        {{- "// Do not quote more than 10 words directly from the tool output.
+" }}
+        {{- "// sources=web (default: web)
+" }}
+        {{- "namespace browser {
+
+" }}
+        {{- "// Searches for information related to `query` and displays `topn` results.
+" }}
+        {{- "type search = (_: {
+" }}
+        {{- "query: string,
+" }}
+        {{- "topn?: number, // default: 10
+" }}
+        {{- "source?: string,
+" }}
+        {{- "}) => any;
+
+" }}
+        {{- "// Opens the link `id` from the page indicated by `cursor` starting at line number `loc`, showing `num_lines` lines.
+" }}
+        {{- "// Valid link ids are displayed with the formatting: `【{id}†.*】`.
+" }}
+        {{- "// If `cursor` is not provided, the most recent page is implied.
+" }}
+        {{- "// If `id` is a string, it is treated as a fully qualified URL associated with `source`.
+" }}
+        {{- "// If `loc` is not provided, the viewport will be positioned at the beginning of the document or centered on the most relevant passage, if available.
+" }}
+        {{- "// Use this function without `id` to scroll to a new location of an opened page.
+" }}
+        {{- "type open = (_: {
+" }}
+        {{- "id?: number | string, // default: -1
+" }}
+        {{- "cursor?: number, // default: -1
+" }}
+        {{- "loc?: number, // default: -1
+" }}
+        {{- "num_lines?: number, // default: -1
+" }}
+        {{- "view_source?: boolean, // default: false
+" }}
+        {{- "source?: string,
+" }}
+        {{- "}) => any;
+
+" }}
+        {{- "// Finds exact matches of `pattern` in the current page, or the page given by `cursor`.
+" }}
+        {{- "type find = (_: {
+" }}
+        {{- "pattern: string,
+" }}
+        {{- "cursor?: number, // default: -1
+" }}
+        {{- "}) => any;
+
+" }}
+        {{- "} // namespace browser
+
+" }}
+    {%- endif -%}
+
+    {%- if python_tool %}
+        {{- "## python
+
+" }}
+        {{- "Use this tool to execute Python code in your chain of thought. The code will not be shown to the user. This tool should be used for internal reasoning, but not for code that is intended to be visible to the user (e.g. when creating plots, tables, or files).
+
+" }}
+        {{- "When you send a message containing Python code to python, it will be executed in a stateful Jupyter notebook environment. python will respond with the output of the execution or time out after 120.0 seconds. The drive at '/mnt/data' can be used to save and persist user files. Internet access for this session is UNKNOWN. Depends on the cluster.
+
+" }}
+    {%- endif -%}
+{%- endmacro -%}
+
+{#- System Message Construction ============================================ #}
+{%- macro build_system_message() -%}
+    {%- if model_identity is not defined %}
+        {%- set model_identity = "You are ChatGPT, a large language model trained by OpenAI." %}
+    {%- endif %}
+    {{- model_identity + "
+" }}
+    {{- "Knowledge cutoff: 2024-06
+" }}
+    {{- "Current date: " + strftime_now("%Y-%m-%d") + "
+
+" }}
+    {%- if reasoning_effort is not defined %}
+        {%- set reasoning_effort = "medium" %}
+    {%- endif %}
+    {{- "Reasoning: " + reasoning_effort + "
+
+" }}
+    {%- if builtin_tools %}
+        {{- "# Tools
+
+" }}
+        {%- set available_builtin_tools = namespace(browser=false, python=false) %}
+        {%- for tool in builtin_tools %}
+            {%- if tool == "browser" %}
+                {%- set available_builtin_tools.browser = true %}
+            {%- elif tool == "python" %}
+                {%- set available_builtin_tools.python = true %}
+            {%- endif %}
+        {%- endfor %}
+        {{- render_builtin_tools(available_builtin_tools.browser, available_builtin_tools.python) }}
+    {%- endif -%}
+    {{- "# Valid channels: analysis, commentary, final. Channel must be included for every message." }}
+    {%- if tools -%}
+        {{- "
+Calls to these tools must go to the commentary channel: 'functions'." }}
+    {%- endif -%}
+{%- endmacro -%}
+
+{#- Main Template Logic ================================================= #}
+{#- Set defaults #}
+
+{#- Render system message #}
+{{- "<|start|>system<|message|>" }}
+{{- build_system_message() }}
+{{- "<|end|>" }}
+
+{#- Extract developer message #}
+{%- if messages[0].role == "developer" or messages[0].role == "system" %}
+    {%- set developer_message = messages[0].content %}
+    {%- set loop_messages = messages[1:] %}
+{%- else %}
+    {%- set developer_message = "" %}
+    {%- set loop_messages = messages %}
+{%- endif %}
+
+{#- Render developer message #}
+{%- if developer_message or tools %}
+    {{- "<|start|>developer<|message|>" }}
+    {%- if developer_message %}
+        {{- "# Instructions
+
+" }}
+        {{- developer_message }}
+    {%- endif %}
+    {%- if tools -%}
+        {{- "
+
+" }}
+        {{- "# Tools
+
+" }}
+        {{- render_tool_namespace("functions", tools) }}
+    {%- endif -%}
+    {{- "<|end|>" }}
+{%- endif %}
+
+{#- Render messages #}
+{%- set last_tool_call = namespace(name=none) %}
+{%- for message in loop_messages -%}
+    {#- At this point only assistant/user/tool messages should remain #}
+    {%- if message.role == 'assistant' -%}
+        {#- Checks to ensure the messages are being passed in the format we expect #}
+        {%- if "content" in message %}
+            {%- if "<|channel|>analysis<|message|>" in message.content or "<|channel|>final<|message|>" in message.content %}
+                {{- raise_exception("You have passed a message containing <|channel|> tags in the content field. Instead of doing this, you should pass analysis messages (the string between '<|message|>' and '<|end|>') in the 'thinking' field, and final messages (the string between '<|message|>' and '<|end|>') in the 'content' field.") }}
+            {%- endif %}
+        {%- endif %}
+        {%- if "thinking" in message %}
+            {%- if "<|channel|>analysis<|message|>" in message.thinking or "<|channel|>final<|message|>" in message.thinking %}
+                {{- raise_exception("You have passed a message containing <|channel|> tags in the thinking field. Instead of doing this, you should pass analysis messages (the string between '<|message|>' and '<|end|>') in the 'thinking' field, and final messages (the string between '<|message|>' and '<|end|>') in the 'content' field.") }}
+            {%- endif %}
+        {%- endif %}
+        {%- if "tool_calls" in message %}
+            {#- We assume max 1 tool call per message, and so we infer the tool call name #}
+            {#- in "tool" messages from the most recent assistant tool call name #}
+            {%- set tool_call = message.tool_calls[0] %}
+            {%- if tool_call.function %}
+                {%- set tool_call = tool_call.function %}
+            {%- endif %}
+            {%- if message.content and message.thinking %}
+                {{- raise_exception("Cannot pass both content and thinking in an assistant message with tool calls! Put the analysis message in one or the other, but not both.") }}
+            {%- elif message.content %}
+                {{- "<|start|>assistant<|channel|>analysis<|message|>" + message.content + "<|end|>" }}
+            {%- elif message.thinking %}
+                {{- "<|start|>assistant<|channel|>analysis<|message|>" + message.thinking + "<|end|>" }}
+            {%- endif %}
+            {{- "<|start|>assistant to=" }}
+            {{- "functions." + tool_call.name + "<|channel|>commentary " }}
+            {{- (tool_call.content_type if tool_call.content_type is defined else "json") + "<|message|>" }}
+            {{- tool_call.arguments|tojson }}
+            {{- "<|call|>" }}
+            {%- set last_tool_call.name = tool_call.name %}
+        {%- elif loop.last and not add_generation_prompt %}
+            {#- Only render the CoT if the final turn is an assistant turn and add_generation_prompt is false #}
+            {#- This is a situation that should only occur in training, never in inference. #}
+            {%- if "thinking" in message %}
+                {{- "<|start|>assistant<|channel|>analysis<|message|>" + message.thinking + "<|end|>" }}
+            {%- endif %}
+            {#- <|return|> indicates the end of generation, but <|end|> does not #}
+            {#- <|return|> should never be an input to the model, but we include it as the final token #}
+            {#- when training, so the model learns to emit it. #}
+            {{- "<|start|>assistant<|channel|>final<|message|>" + message.content + "<|return|>" }}
+        {%- else %}
+            {#- CoT is dropped during all previous turns, so we never render it for inference #}
+            {{- "<|start|>assistant<|channel|>final<|message|>" + message.content + "<|end|>" }}
+            {%- set last_tool_call.name = none %}
+        {%- endif %}
+    {%- elif message.role == 'tool' -%}
+        {%- if last_tool_call.name is none %}
+            {{- raise_exception("Message has tool role, but there was no previous assistant message with a tool call!") }}
+        {%- endif %}
+        {{- "<|start|>functions." + last_tool_call.name }}
+        {{- " to=assistant<|channel|>commentary<|message|>" + message.content|tojson + "<|end|>" }}
+    {%- elif message.role == 'user' -%}
+        {{- "<|start|>user<|message|>" + message.content + "<|end|>" }}
+    {%- endif -%}
+{%- endfor -%}
+
+{#- Generation prompt #}
+{%- if add_generation_prompt -%}
+<|start|>assistant
+{%- endif -%}

config.json

@@ -0,0 +1,638 @@
+{
+  "architectures": [
+    "GptOssForCausalLM"
+  ],
+  "attention_bias": true,
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoModel": "modeling_gpt_oss.GptOssModel",
+    "AutoModelForCausalLM": "modeling_gpt_oss.GptOssForCausalLM"
+  },
+  "eos_token_id": 200002,
+  "experts_per_token": 4,
+  "head_dim": 64,
+  "hidden_act": "silu",
+  "hidden_size": 2880,
+  "initial_context_length": 4096,
+  "initializer_range": 0.02,
+  "intermediate_size": 2880,
+  "layer_types": [
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "full_attention"
+  ],
+  "max_position_embeddings": 131072,
+  "model_type": "gpt_oss",
+  "num_attention_heads": 64,
+  "num_experts_per_tok": 4,
+  "num_hidden_layers": 36,
+  "num_key_value_heads": 8,
+  "num_local_experts": 128,
+  "output_router_logits": false,
+  "pad_token_id": 199999,
+  "quantization_config": {
+    "autoround_version": "0.6.1.dev",
+    "batch_size": 4,
+    "bits": 4,
+    "data_type": "int",
+    "extra_config": {
+      "model.layers.0.mlp.router.router": {
+        "bits": 16
+      },
+      "model.layers.0.self_attn.k_proj": {
+        "bits": 16
+      },
+      "model.layers.0.self_attn.o_proj": {
+        "bits": 16
+      },
+      "model.layers.0.self_attn.q_proj": {
+        "bits": 16
+      },
+      "model.layers.0.self_attn.v_proj": {
+        "bits": 16
+      },
+      "model.layers.1.mlp.router.router": {
+        "bits": 16
+      },
+      "model.layers.1.self_attn.k_proj": {
+        "bits": 16
+      },
+      "model.layers.1.self_attn.o_proj": {
+        "bits": 16
+      },
+      "model.layers.1.self_attn.q_proj": {
+        "bits": 16
+      },
+      "model.layers.1.self_attn.v_proj": {
+        "bits": 16
+      },
+      "model.layers.10.mlp.router.router": {
+        "bits": 16
+      },
+      "model.layers.10.self_attn.k_proj": {
+        "bits": 16
+      },
+      "model.layers.10.self_attn.o_proj": {
+        "bits": 16
+      },
+      "model.layers.10.self_attn.q_proj": {
+        "bits": 16
+      },
+      "model.layers.10.self_attn.v_proj": {
+        "bits": 16
+      },
+      "model.layers.11.mlp.router.router": {
+        "bits": 16
+      },
+      "model.layers.11.self_attn.k_proj": {
+        "bits": 16
+      },
+      "model.layers.11.self_attn.o_proj": {
+        "bits": 16
+      },
+      "model.layers.11.self_attn.q_proj": {
+        "bits": 16
+      },
+      "model.layers.11.self_attn.v_proj": {
+        "bits": 16
+      },
+      "model.layers.12.mlp.router.router": {
+        "bits": 16
+      },
+      "model.layers.12.self_attn.k_proj": {
+        "bits": 16
+      },
+      "model.layers.12.self_attn.o_proj": {
+        "bits": 16
+      },
+      "model.layers.12.self_attn.q_proj": {
+        "bits": 16
+      },
+      "model.layers.12.self_attn.v_proj": {
+        "bits": 16
+      },
+      "model.layers.13.mlp.router.router": {
+        "bits": 16
+      },
+      "model.layers.13.self_attn.k_proj": {
+        "bits": 16
+      },
+      "model.layers.13.self_attn.o_proj": {
+        "bits": 16
+      },
+      "model.layers.13.self_attn.q_proj": {
+        "bits": 16
+      },
+      "model.layers.13.self_attn.v_proj": {
+        "bits": 16
+      },
+      "model.layers.14.mlp.router.router": {
+        "bits": 16
+      },
+      "model.layers.14.self_attn.k_proj": {
+        "bits": 16
+      },
+      "model.layers.14.self_attn.o_proj": {
+        "bits": 16
+      },
+      "model.layers.14.self_attn.q_proj": {
+        "bits": 16
+      },
+      "model.layers.14.self_attn.v_proj": {
+        "bits": 16
+      },
+      "model.layers.15.mlp.router.router": {
+        "bits": 16
+      },
+      "model.layers.15.self_attn.k_proj": {
+        "bits": 16
+      },
+      "model.layers.15.self_attn.o_proj": {
+        "bits": 16
+      },
+      "model.layers.15.self_attn.q_proj": {
+        "bits": 16
+      },
+      "model.layers.15.self_attn.v_proj": {
+        "bits": 16
+      },
+      "model.layers.16.mlp.router.router": {
+        "bits": 16
+      },
+      "model.layers.16.self_attn.k_proj": {
+        "bits": 16
+      },
+      "model.layers.16.self_attn.o_proj": {
+        "bits": 16
+      },
+      "model.layers.16.self_attn.q_proj": {
+        "bits": 16
+      },
+      "model.layers.16.self_attn.v_proj": {
+        "bits": 16
+      },
+      "model.layers.17.mlp.router.router": {
+        "bits": 16
+      },
+      "model.layers.17.self_attn.k_proj": {
+        "bits": 16
+      },
+      "model.layers.17.self_attn.o_proj": {
+        "bits": 16
+      },
+      "model.layers.17.self_attn.q_proj": {
+        "bits": 16
+      },
+      "model.layers.17.self_attn.v_proj": {
+        "bits": 16
+      },
+      "model.layers.18.mlp.router.router": {
+        "bits": 16
+      },
+      "model.layers.18.self_attn.k_proj": {
+        "bits": 16
+      },
+      "model.layers.18.self_attn.o_proj": {
+        "bits": 16
+      },
+      "model.layers.18.self_attn.q_proj": {
+        "bits": 16
+      },
+      "model.layers.18.self_attn.v_proj": {
+        "bits": 16
+      },
+      "model.layers.19.mlp.router.router": {
+        "bits": 16
+      },
+      "model.layers.19.self_attn.k_proj": {
+        "bits": 16
+      },
+      "model.layers.19.self_attn.o_proj": {
+        "bits": 16
+      },
+      "model.layers.19.self_attn.q_proj": {
+        "bits": 16
+      },
+      "model.layers.19.self_attn.v_proj": {
+        "bits": 16
+      },
+      "model.layers.2.mlp.router.router": {
+        "bits": 16
+      },
+      "model.layers.2.self_attn.k_proj": {
+        "bits": 16
+      },
+      "model.layers.2.self_attn.o_proj": {
+        "bits": 16
+      },
+      "model.layers.2.self_attn.q_proj": {
+        "bits": 16
+      },
+      "model.layers.2.self_attn.v_proj": {
+        "bits": 16
+      },
+      "model.layers.20.mlp.router.router": {
+        "bits": 16
+      },
+      "model.layers.20.self_attn.k_proj": {
+        "bits": 16
+      },
+      "model.layers.20.self_attn.o_proj": {
+        "bits": 16
+      },
+      "model.layers.20.self_attn.q_proj": {
+        "bits": 16
+      },
+      "model.layers.20.self_attn.v_proj": {
+        "bits": 16
+      },
+      "model.layers.21.mlp.router.router": {
+        "bits": 16
+      },
+      "model.layers.21.self_attn.k_proj": {
+        "bits": 16
+      },
+      "model.layers.21.self_attn.o_proj": {
+        "bits": 16
+      },
+      "model.layers.21.self_attn.q_proj": {
+        "bits": 16
+      },
+      "model.layers.21.self_attn.v_proj": {
+        "bits": 16
+      },
+      "model.layers.22.mlp.router.router": {
+        "bits": 16
+      },
+      "model.layers.22.self_attn.k_proj": {
+        "bits": 16
+      },
+      "model.layers.22.self_attn.o_proj": {
+        "bits": 16
+      },
+      "model.layers.22.self_attn.q_proj": {
+        "bits": 16
+      },
+      "model.layers.22.self_attn.v_proj": {
+        "bits": 16
+      },
+      "model.layers.23.mlp.router.router": {
+        "bits": 16
+      },
+      "model.layers.23.self_attn.k_proj": {
+        "bits": 16
+      },
+      "model.layers.23.self_attn.o_proj": {
+        "bits": 16
+      },
+      "model.layers.23.self_attn.q_proj": {
+        "bits": 16
+      },
+      "model.layers.23.self_attn.v_proj": {
+        "bits": 16
+      },
+      "model.layers.24.mlp.router.router": {
+        "bits": 16
+      },
+      "model.layers.24.self_attn.k_proj": {
+        "bits": 16
+      },
+      "model.layers.24.self_attn.o_proj": {
+        "bits": 16
+      },
+      "model.layers.24.self_attn.q_proj": {
+        "bits": 16
+      },
+      "model.layers.24.self_attn.v_proj": {
+        "bits": 16
+      },
+      "model.layers.25.mlp.router.router": {
+        "bits": 16
+      },
+      "model.layers.25.self_attn.k_proj": {
+        "bits": 16
+      },
+      "model.layers.25.self_attn.o_proj": {
+        "bits": 16
+      },
+      "model.layers.25.self_attn.q_proj": {
+        "bits": 16
+      },
+      "model.layers.25.self_attn.v_proj": {
+        "bits": 16
+      },
+      "model.layers.26.mlp.router.router": {
+        "bits": 16
+      },
+      "model.layers.26.self_attn.k_proj": {
+        "bits": 16
+      },
+      "model.layers.26.self_attn.o_proj": {
+        "bits": 16
+      },
+      "model.layers.26.self_attn.q_proj": {
+        "bits": 16
+      },
+      "model.layers.26.self_attn.v_proj": {
+        "bits": 16
+      },
+      "model.layers.27.mlp.router.router": {
+        "bits": 16
+      },
+      "model.layers.27.self_attn.k_proj": {
+        "bits": 16
+      },
+      "model.layers.27.self_attn.o_proj": {
+        "bits": 16
+      },
+      "model.layers.27.self_attn.q_proj": {
+        "bits": 16
+      },
+      "model.layers.27.self_attn.v_proj": {
+        "bits": 16
+      },
+      "model.layers.28.mlp.router.router": {
+        "bits": 16
+      },
+      "model.layers.28.self_attn.k_proj": {
+        "bits": 16
+      },
+      "model.layers.28.self_attn.o_proj": {
+        "bits": 16
+      },
+      "model.layers.28.self_attn.q_proj": {
+        "bits": 16
+      },
+      "model.layers.28.self_attn.v_proj": {
+        "bits": 16
+      },
+      "model.layers.29.mlp.router.router": {
+        "bits": 16
+      },
+      "model.layers.29.self_attn.k_proj": {
+        "bits": 16
+      },
+      "model.layers.29.self_attn.o_proj": {
+        "bits": 16
+      },
+      "model.layers.29.self_attn.q_proj": {
+        "bits": 16
+      },
+      "model.layers.29.self_attn.v_proj": {
+        "bits": 16
+      },
+      "model.layers.3.mlp.router.router": {
+        "bits": 16
+      },
+      "model.layers.3.self_attn.k_proj": {
+        "bits": 16
+      },
+      "model.layers.3.self_attn.o_proj": {
+        "bits": 16
+      },
+      "model.layers.3.self_attn.q_proj": {
+        "bits": 16
+      },
+      "model.layers.3.self_attn.v_proj": {
+        "bits": 16
+      },
+      "model.layers.30.mlp.router.router": {
+        "bits": 16
+      },
+      "model.layers.30.self_attn.k_proj": {
+        "bits": 16
+      },
+      "model.layers.30.self_attn.o_proj": {
+        "bits": 16
+      },
+      "model.layers.30.self_attn.q_proj": {
+        "bits": 16
+      },
+      "model.layers.30.self_attn.v_proj": {
+        "bits": 16
+      },
+      "model.layers.31.mlp.router.router": {
+        "bits": 16
+      },
+      "model.layers.31.self_attn.k_proj": {
+        "bits": 16
+      },
+      "model.layers.31.self_attn.o_proj": {
+        "bits": 16
+      },
+      "model.layers.31.self_attn.q_proj": {
+        "bits": 16
+      },
+      "model.layers.31.self_attn.v_proj": {
+        "bits": 16
+      },
+      "model.layers.32.mlp.router.router": {
+        "bits": 16
+      },
+      "model.layers.32.self_attn.k_proj": {
+        "bits": 16
+      },
+      "model.layers.32.self_attn.o_proj": {
+        "bits": 16
+      },
+      "model.layers.32.self_attn.q_proj": {
+        "bits": 16
+      },
+      "model.layers.32.self_attn.v_proj": {
+        "bits": 16
+      },
+      "model.layers.33.mlp.router.router": {
+        "bits": 16
+      },
+      "model.layers.33.self_attn.k_proj": {
+        "bits": 16
+      },
+      "model.layers.33.self_attn.o_proj": {
+        "bits": 16
+      },
+      "model.layers.33.self_attn.q_proj": {
+        "bits": 16
+      },
+      "model.layers.33.self_attn.v_proj": {
+        "bits": 16
+      },
+      "model.layers.34.mlp.router.router": {
+        "bits": 16
+      },
+      "model.layers.34.self_attn.k_proj": {
+        "bits": 16
+      },
+      "model.layers.34.self_attn.o_proj": {
+        "bits": 16
+      },
+      "model.layers.34.self_attn.q_proj": {
+        "bits": 16
+      },
+      "model.layers.34.self_attn.v_proj": {
+        "bits": 16
+      },
+      "model.layers.35.mlp.router.router": {
+        "bits": 16
+      },
+      "model.layers.35.self_attn.k_proj": {
+        "bits": 16
+      },
+      "model.layers.35.self_attn.o_proj": {
+        "bits": 16
+      },
+      "model.layers.35.self_attn.q_proj": {
+        "bits": 16
+      },
+      "model.layers.35.self_attn.v_proj": {
+        "bits": 16
+      },
+      "model.layers.4.mlp.router.router": {
+        "bits": 16
+      },
+      "model.layers.4.self_attn.k_proj": {
+        "bits": 16
+      },
+      "model.layers.4.self_attn.o_proj": {
+        "bits": 16
+      },
+      "model.layers.4.self_attn.q_proj": {
+        "bits": 16
+      },
+      "model.layers.4.self_attn.v_proj": {
+        "bits": 16
+      },
+      "model.layers.5.mlp.router.router": {
+        "bits": 16
+      },
+      "model.layers.5.self_attn.k_proj": {
+        "bits": 16
+      },
+      "model.layers.5.self_attn.o_proj": {
+        "bits": 16
+      },
+      "model.layers.5.self_attn.q_proj": {
+        "bits": 16
+      },
+      "model.layers.5.self_attn.v_proj": {
+        "bits": 16
+      },
+      "model.layers.6.mlp.router.router": {
+        "bits": 16
+      },
+      "model.layers.6.self_attn.k_proj": {
+        "bits": 16
+      },
+      "model.layers.6.self_attn.o_proj": {
+        "bits": 16
+      },
+      "model.layers.6.self_attn.q_proj": {
+        "bits": 16
+      },
+      "model.layers.6.self_attn.v_proj": {
+        "bits": 16
+      },
+      "model.layers.7.mlp.router.router": {
+        "bits": 16
+      },
+      "model.layers.7.self_attn.k_proj": {
+        "bits": 16
+      },
+      "model.layers.7.self_attn.o_proj": {
+        "bits": 16
+      },
+      "model.layers.7.self_attn.q_proj": {
+        "bits": 16
+      },
+      "model.layers.7.self_attn.v_proj": {
+        "bits": 16
+      },
+      "model.layers.8.mlp.router.router": {
+        "bits": 16
+      },
+      "model.layers.8.self_attn.k_proj": {
+        "bits": 16
+      },
+      "model.layers.8.self_attn.o_proj": {
+        "bits": 16
+      },
+      "model.layers.8.self_attn.q_proj": {
+        "bits": 16
+      },
+      "model.layers.8.self_attn.v_proj": {
+        "bits": 16
+      },
+      "model.layers.9.mlp.router.router": {
+        "bits": 16
+      },
+      "model.layers.9.self_attn.k_proj": {
+        "bits": 16
+      },
+      "model.layers.9.self_attn.o_proj": {
+        "bits": 16
+      },
+      "model.layers.9.self_attn.q_proj": {
+        "bits": 16
+      },
+      "model.layers.9.self_attn.v_proj": {
+        "bits": 16
+      }
+    },
+    "group_size": 64,
+    "low_gpu_mem_usage": true,
+    "nsamples": 512,
+    "packing_format": "auto_round:auto_gptq",
+    "quant_method": "auto-round",
+    "sym": true
+  },
+  "rms_norm_eps": 1e-05,
+  "rope_scaling": {
+    "beta_fast": 32.0,
+    "beta_slow": 1.0,
+    "factor": 32.0,
+    "original_max_position_embeddings": 4096,
+    "rope_type": "yarn",
+    "truncate": false
+  },
+  "rope_theta": 150000,
+  "router_aux_loss_coef": 0.9,
+  "sliding_window": 128,
+  "swiglu_limit": 7.0,
+  "tie_word_embeddings": false,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.55.0",
+  "use_cache": true,
+  "vocab_size": 201088
+}

generation_config.json

@@ -0,0 +1,6 @@
+{
+  "_from_model_config": true,
+  "eos_token_id": 200002,
+  "pad_token_id": 199999,
+  "transformers_version": "4.55.0"
+}

model-00001-of-00014.safetensors

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

model-00002-of-00014.safetensors

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

model-00003-of-00014.safetensors

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

model-00004-of-00014.safetensors

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

model-00005-of-00014.safetensors

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

model-00006-of-00014.safetensors

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

model-00007-of-00014.safetensors

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

model-00008-of-00014.safetensors

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

model-00009-of-00014.safetensors

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

model-00010-of-00014.safetensors

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

model-00011-of-00014.safetensors

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

model-00012-of-00014.safetensors

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

model-00013-of-00014.safetensors

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

model-00014-of-00014.safetensors

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

modeling_gpt_oss.py

@@ -0,0 +1,893 @@
+import copy
+from typing import Callable, Optional, Union
+
+import torch
+from accelerate import init_empty_weights
+from torch import nn
+from torch.nn import functional as F
+
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.generation import GenerationMixin
+from transformers.integrations.hub_kernels import use_kernel_forward_from_hub
+from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask
+from transformers.modeling_layers import GradientCheckpointingLayer
+from transformers.modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast
+from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from transformers.processing_utils import Unpack
+from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple
+from transformers.utils.generic import OutputRecorder, check_model_inputs
+
+from transformers.configuration_utils import PretrainedConfig, layer_type_validation
+from transformers.modeling_rope_utils import rope_config_validation
+
+class GptOssConfig(PretrainedConfig):
+    r"""
+    This will yield a configuration to that of the BERT
+    [google-bert/bert-base-uncased](https://huggingface.co/google-bert/bert-base-uncased) architecture.
+
+    """
+
+    model_type = "gpt_oss"
+    base_model_pp_plan = {
+        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+        "norm": (["hidden_states"], ["hidden_states"]),
+    }
+    base_model_tp_plan = {
+        "layers.*.self_attn.q_proj": "colwise",
+        "layers.*.self_attn.k_proj": "colwise",
+        "layers.*.self_attn.v_proj": "colwise",
+        "layers.*.self_attn.o_proj": "rowwise",
+        "layers.*.self_attn.sinks": "local_rowwise",
+        "layers.*.mlp.experts": "gather",
+        "layers.*.mlp.router": "ep_router",
+        "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
+        "layers.*.mlp.experts.gate_up_proj_bias": "grouped_gemm",
+        "layers.*.mlp.experts.down_proj": "grouped_gemm",
+        "layers.*.mlp.experts.down_proj_bias": "grouped_gemm",
+    }
+
+    def __init__(
+        self,
+        num_hidden_layers: int = 36,
+        num_local_experts: int = 128,
+        vocab_size: int = 201088,
+        hidden_size: int = 2880,
+        intermediate_size: int = 2880,
+        head_dim: int = 64,
+        num_attention_heads: int = 64,
+        num_key_value_heads: int = 8,
+        sliding_window: int = 128,
+        rope_theta: float = 150000.0,
+        tie_word_embeddings=False,
+        hidden_act: str = "silu",
+        initializer_range: float = 0.02,
+        max_position_embeddings=131072,
+        rms_norm_eps: float = 1e-5,
+        rope_scaling={"rope_type": "yarn", "factor": 32.0, "beta_fast": 32.0, "beta_slow": 1.0, "truncate": False},
+        attention_dropout: float = 0.0,
+        num_experts_per_tok=4,
+        router_aux_loss_coef: float = 0.9,
+        output_router_logits=False,
+        use_cache=True,
+        layer_types=None,
+        **kwargs,
+    ):
+        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_local_experts = num_local_experts
+        self.sliding_window = sliding_window
+        self.num_experts_per_tok = num_experts_per_tok
+        # for backward compatibility
+        if num_key_value_heads is None:
+            num_key_value_heads = num_attention_heads
+
+        self.num_key_value_heads = num_key_value_heads
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+        self.attention_dropout = attention_dropout
+        self.head_dim = head_dim if head_dim is not None else self.hidden_size // self.num_attention_heads
+        self.layer_types = layer_types
+        if self.layer_types is None:
+            self.layer_types = [
+                "sliding_attention" if bool((i + 1) % 2) else "full_attention" for i in range(self.num_hidden_layers)
+            ]
+        layer_type_validation(self.layer_types)
+
+        # Validate the correctness of rotary position embeddings parameters
+        # BC: if there is a 'type' field, copy it it to 'rope_type'.
+        if self.rope_scaling is not None and "type" in self.rope_scaling:
+            self.rope_scaling["rope_type"] = self.rope_scaling["type"]
+        rope_config_validation(self)
+
+        self.attention_bias = True
+        self.max_position_embeddings = max_position_embeddings
+        self.router_aux_loss_coef = router_aux_loss_coef
+        self.output_router_logits = output_router_logits
+        self.use_cache = use_cache
+        super().__init__(
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+@use_kernel_forward_from_hub("RMSNorm")
+class GptOssRMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        """
+        GptOssRMSNorm 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)  # main diff with Llama
+
+    def extra_repr(self):
+        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+#
+# class GptOssExperts(nn.Module):
+#     def __init__(self, config):
+#         super().__init__()
+#         self.intermediate_size = config.intermediate_size
+#         self.num_experts = config.num_local_experts
+#         self.hidden_size = config.hidden_size
+#         self.expert_dim = self.intermediate_size
+#         self.gate_up_proj = nn.Parameter(torch.empty(self.num_experts, self.hidden_size, 2 * self.expert_dim))
+#         self.gate_up_proj_bias = nn.Parameter(torch.empty(self.num_experts, 2 * self.expert_dim))
+#         self.down_proj = nn.Parameter(torch.empty((self.num_experts, self.expert_dim, self.hidden_size)))
+#         self.down_proj_bias = nn.Parameter(torch.empty(self.num_experts, self.hidden_size))
+#         self.alpha = 1.702
+#         self.limit = 7.0
+#
+#
+#
+#     def forward(self, hidden_states: torch.Tensor, router_indices=None, routing_weights=None) -> torch.Tensor:
+#         """
+#         When training is is more efficient to just loop over the experts and compute the output for each expert
+#         as otherwise the memory would explode.
+#
+#         For inference we can sacrifice some memory and compute the output for all experts at once. By repeating the inputs.
+#
+#         Args:
+#             hidden_states (torch.Tensor): (batch_size, seq_len, hidden_size)
+#             selected_experts (torch.Tensor): (batch_size * token_num, top_k)
+#             routing_weights (torch.Tensor): (batch_size * token_num, num_experts)
+#         Returns:
+#             torch.Tensor
+#         """
+#         batch_size = hidden_states.shape[0]
+#         hidden_states = hidden_states.reshape(-1, self.hidden_size)  # (num_tokens, hidden_size)
+#         num_experts = routing_weights.shape[1]
+#         if self.training:
+#             next_states = torch.zeros_like(hidden_states, dtype=hidden_states.dtype, device=hidden_states.device)
+#             with torch.no_grad():
+#                 expert_mask = torch.nn.functional.one_hot(router_indices, num_classes=num_experts)
+#                 expert_mask = expert_mask.permute(2, 1, 0)
+#                 # we sum on the top_k and on the sequence lenght to get which experts
+#                 # are hit this time around
+#                 expert_hitted = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
+#             for expert_idx in expert_hitted[:]:
+#                 with torch.no_grad():
+#                     _, token_idx = torch.where(expert_mask[expert_idx[0]])
+#                 current_state = hidden_states[token_idx]
+#                 gate_up = current_state @ self.gate_up_proj[expert_idx] + self.gate_up_proj_bias[expert_idx]
+#                 gate, up = gate_up[..., ::2], gate_up[..., 1::2]
+#                 gate = gate.clamp(min=None, max=self.limit)
+#                 up = up.clamp(min=-self.limit, max=self.limit)
+#                 glu = gate * torch.sigmoid(gate * self.alpha)
+#                 gated_output = (up + 1) * glu
+#                 out = gated_output @ self.down_proj[expert_idx] + self.down_proj_bias[expert_idx]
+#                 weighted_output = out[0] * routing_weights[token_idx, expert_idx, None]
+#                 next_states.index_add_(0, token_idx, weighted_output.to(hidden_states.dtype))
+#             next_states = next_states.view(batch_size, -1, self.hidden_size)
+#         else:
+#             hidden_states = hidden_states.repeat(num_experts, 1)
+#             hidden_states = hidden_states.view(num_experts, -1, self.hidden_size)
+#             gate_up = torch.bmm(hidden_states, self.gate_up_proj) + self.gate_up_proj_bias[..., None, :]
+#             gate, up = gate_up[..., ::2], gate_up[..., 1::2]
+#             gate = gate.clamp(min=None, max=self.limit)
+#             up = up.clamp(min=-self.limit, max=self.limit)
+#             glu = gate * torch.sigmoid(gate * self.alpha)
+#             next_states = torch.bmm(((up + 1) * glu), self.down_proj)
+#             next_states = next_states + self.down_proj_bias[..., None, :]
+#             next_states = next_states.view(num_experts, batch_size, -1, self.hidden_size)
+#             next_states = next_states * routing_weights.transpose(0, 1).view(num_experts, batch_size, -1)[..., None]
+#             next_states = next_states.sum(dim=0)
+#         return next_states
+
+class GptOssExperts(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.intermediate_size = config.intermediate_size
+        self.num_experts = config.num_local_experts
+        self.hidden_size = config.hidden_size
+        self.expert_dim = self.intermediate_size
+
+        # 使用nn.Linear替代手动矩阵乘法
+        self.gate_up_projs = nn.ModuleList([
+            nn.Linear(self.hidden_size, 2 * self.expert_dim)
+            for _ in range(self.num_experts)
+        ])
+
+        self.down_projs = nn.ModuleList([
+            nn.Linear(self.expert_dim, self.hidden_size)
+            for _ in range(self.num_experts)
+        ])
+
+        self.alpha = 1.702
+        self.limit = 7.0
+
+    def forward(self, hidden_states: torch.Tensor, router_indices=None, routing_weights=None) -> torch.Tensor:
+        batch_size = hidden_states.shape[0]
+        hidden_states = hidden_states.reshape(-1, self.hidden_size)  # (num_tokens, hidden_size)
+        num_experts = routing_weights.shape[1]
+
+        if self.training:
+            next_states = torch.zeros_like(hidden_states, dtype=hidden_states.dtype, device=hidden_states.device)
+
+            with torch.no_grad():
+                expert_mask = torch.nn.functional.one_hot(router_indices, num_classes=num_experts)
+                expert_mask = expert_mask.permute(2, 1, 0)
+                expert_hitted = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
+
+            for expert_idx in expert_hitted[:]:
+                with torch.no_grad():
+                    _, token_idx = torch.where(expert_mask[expert_idx[0]])
+
+                current_state = hidden_states[token_idx]
+
+                # 使用Linear层替代手动矩阵乘法
+                gate_up = self.gate_up_projs[expert_idx](current_state)
+                gate, up = gate_up[..., ::2], gate_up[..., 1::2]
+                gate = gate.clamp(min=None, max=self.limit)
+                up = up.clamp(min=-self.limit, max=self.limit)
+
+                glu = gate * torch.sigmoid(gate * self.alpha)
+                gated_output = (up + 1) * glu
+
+                # 使用Linear层替代手动矩阵乘法
+                out = self.down_projs[expert_idx](gated_output)
+
+                weighted_output = out[0] * routing_weights[token_idx, expert_idx, None]
+                next_states.index_add_(0, token_idx, weighted_output.to(hidden_states.dtype))
+
+            next_states = next_states.view(batch_size, -1, self.hidden_size)
+        else:
+            hidden_states = hidden_states.repeat(num_experts, 1)
+            hidden_states = hidden_states.view(num_experts, -1, self.hidden_size)
+
+            # 批量处理所有专家
+            gate_up = torch.stack([proj(hidden_states[i]) for i, proj in enumerate(self.gate_up_projs)])
+            gate, up = gate_up[..., ::2], gate_up[..., 1::2]
+            gate = gate.clamp(min=None, max=self.limit)
+            up = up.clamp(min=-self.limit, max=self.limit)
+
+            glu = gate * torch.sigmoid(gate * self.alpha)
+            next_states = torch.stack([proj((up[i] + 1) * glu[i]) for i, proj in enumerate(self.down_projs)])
+
+            next_states = next_states.view(num_experts, batch_size, -1, self.hidden_size)
+            next_states = next_states * routing_weights.transpose(0, 1).view(num_experts, batch_size, -1)[..., None]
+            next_states = next_states.sum(dim=0)
+
+        return next_states
+
+# class GptOssTopKRouter(nn.Module):
+#     def __init__(self, config):
+#         super().__init__()
+#         self.top_k = config.num_experts_per_tok
+#         self.num_experts = config.num_local_experts
+#         self.hidden_dim = config.hidden_size
+#         self.weight = nn.Parameter(torch.empty(self.num_experts, self.hidden_dim))
+#         self.bias = nn.Parameter(torch.empty(self.num_experts))
+#
+#     def forward(self, hidden_states):
+#         hidden_states = hidden_states.reshape(-1, self.hidden_dim)
+#         router_logits = F.linear(hidden_states, self.weight, self.bias)  # (seq_len, num_experts)
+#         router_top_value, router_indices = torch.topk(router_logits, self.top_k, dim=-1)  # (seq_len, top_k)
+#         router_top_value = torch.nn.functional.softmax(router_top_value, dim=1, dtype=router_top_value.dtype)
+#         router_scores = torch.zeros_like(router_logits).scatter_(1, router_indices, router_top_value)
+#         return router_scores, router_indices
+
+
+class GptOssTopKRouter(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.top_k = config.num_experts_per_tok
+        self.num_experts = config.num_local_experts
+        self.hidden_dim = config.hidden_size
+
+        # 使用nn.Linear替代手动参数
+        self.router = nn.Linear(self.hidden_dim, self.num_experts)
+
+    def forward(self, hidden_states):
+        # 展平输入 (batch_size * seq_len, hidden_dim)
+        hidden_states = hidden_states.reshape(-1, self.hidden_dim)
+        router_logits = self.router(hidden_states)  # (num_tokens, num_experts)
+
+        router_top_value, router_indices = torch.topk(
+            router_logits,
+            self.top_k,
+            dim=-1
+        )  # (num_tokens, top_k)
+
+        router_top_value = F.softmax(router_top_value, dim=-1, dtype=router_top_value.dtype)
+
+        router_scores = torch.zeros_like(router_logits).scatter_(
+            dim=1,
+            index=router_indices,
+            src=router_top_value
+        )
+
+        return router_scores, router_indices
+
+
+
+
+@use_kernel_forward_from_hub("MegaBlocksMoeMLP")
+class GptOssMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.router = GptOssTopKRouter(config)
+        self.experts = GptOssExperts(config)
+
+    def forward(self, hidden_states):
+        router_scores, router_indices = self.router(hidden_states)  # (num_experts, seq_len)
+        routed_out = self.experts(hidden_states, router_indices=router_indices, routing_weights=router_scores)
+        return routed_out, router_scores
+
+
+class GptOssRotaryEmbedding(nn.Module):
+    def __init__(self, config: GptOssConfig, device=None):
+        super().__init__()
+        # BC: "rope_type" was originally "type"
+        if hasattr(config, "rope_scaling") and isinstance(config.rope_scaling, dict):
+            self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
+        else:
+            self.rope_type = "default"
+        self.max_seq_len_cached = config.max_position_embeddings
+        self.original_max_seq_len = config.max_position_embeddings
+
+        self.config = config
+        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+
+        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.original_inv_freq = self.inv_freq
+
+    @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 torch.autocast(device_type=device_type, enabled=False):  # Force float32
+            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+            emb = freqs
+            cos = emb.cos() * self.attention_scaling
+            sin = emb.sin() * self.attention_scaling
+
+        return cos.to(x.dtype), sin.to(x.dtype)
+
+
+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 _apply_rotary_emb(
+    x: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+) -> torch.Tensor:
+    first_half, second_half = torch.chunk(x, 2, dim=-1)
+    first_ = first_half * cos - second_half * sin
+    second_ = second_half * cos + first_half * sin
+    return torch.cat((first_, second_), dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    q_embed = _apply_rotary_emb(q, cos, sin)
+    k_embed = _apply_rotary_emb(k, cos, sin)
+    return q_embed, k_embed
+
+
+def eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: Optional[torch.Tensor],
+    scaling: float,
+    dropout: float = 0.0,
+    **kwargs,
+):
+    key_states = repeat_kv(key, module.num_key_value_groups)
+    value_states = repeat_kv(value, module.num_key_value_groups)
+    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+    if attention_mask is not None:
+        causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
+        attn_weights = attn_weights + causal_mask
+
+    sinks = module.sinks.reshape(1, -1, 1, 1).expand(query.shape[0], -1, query.shape[-2], -1)
+    combined_logits = torch.cat([attn_weights, sinks], dim=-1)
+
+    # This was not in the original implementation and slightly affect results; it prevents overflow in BF16/FP16
+    # when training with bsz>1 we clamp max values.
+
+    combined_logits = combined_logits - combined_logits.max(dim=-1, keepdim=True).values
+    probs = F.softmax(combined_logits, dim=-1, dtype=combined_logits.dtype)
+    scores = probs[..., :-1]  # we drop the sink here
+    attn_weights = nn.functional.dropout(scores, 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
+
+
+class GptOssAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config: GptOssConfig, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+        self.scaling = self.head_dim**-0.5
+        self.attention_dropout = config.attention_dropout
+        self.is_causal = True
+        self.q_proj = nn.Linear(
+            config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.k_proj = nn.Linear(
+            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.v_proj = nn.Linear(
+            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.o_proj = nn.Linear(
+            config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
+        )
+        self.sliding_window = config.sliding_window if config.layer_types[layer_idx] == "sliding_attention" else None
+        self.sinks = nn.Parameter(torch.empty(config.num_attention_heads))
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[torch.Tensor],
+        past_key_value: Optional[Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        if past_key_value is not None:
+            cache_kwargs = {"cache_position": cache_position}
+            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+        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,
+            sliding_window=self.sliding_window,
+            s_aux=self.sinks,  # diff with Llama
+            **kwargs,
+        )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+
+class GptOssDecoderLayer(GradientCheckpointingLayer):
+    def __init__(self, config: GptOssConfig, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.self_attn = GptOssAttention(config=config, layer_idx=layer_idx)
+        self.mlp = GptOssMLP(config)
+        self.input_layernorm = GptOssRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = GptOssRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.attention_type = config.layer_types[layer_idx]
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        use_cache: Optional[bool] = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,  # necessary, but kept here for BC
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple[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_value=past_key_value,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            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)  # diff with llama: router scores
+        hidden_states = residual + hidden_states
+        return hidden_states
+
+
+@auto_docstring
+class GptOssPreTrainedModel(PreTrainedModel):
+    config: GptOssConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["GptOssDecoderLayer"]
+    _skip_keys_device_placement = ["past_key_values"]
+    _supports_flash_attn = True
+    _supports_sdpa = False
+    _supports_flex_attn = True
+
+    _can_compile_fullgraph = True
+    _supports_attention_backend = True
+    _can_record_outputs = {
+        "router_logits": OutputRecorder(GptOssTopKRouter, index=0),
+        "hidden_states": GptOssDecoderLayer,
+        "attentions": GptOssAttention,
+    }
+    _keep_in_fp32_modules = ["post_attention_layernorm", "input_layernorm", "norm"]
+    _supports_flash_attention = False
+    _supports_flex_attention = False
+
+    def _init_weights(self, module):
+        std = self.config.initializer_range
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Parameter):
+            module.data.normal_(mean=0.0, std=std)
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+        elif isinstance(module, GptOssRMSNorm):
+            module.weight.data.fill_(1.0)
+        # elif isinstance(module, GptOssExperts):##too slow
+        #     for gate_up_proj in module.gate_up_projs:
+        #         gate_up_proj.weight.normal_(mean=0.0, std=std)
+        #         gate_up_proj.bias.data.zero_()
+        #     for down_proj in module.down_projs:
+        #         down_proj.weight.data.normal_(mean=0.0, std=std)
+        #         down_proj.bias.data.zero_()
+        elif isinstance(module, GptOssAttention):
+            module.sinks.data.normal_(mean=0.0, std=std)
+        # elif isinstance(module, GptOssTopKRouter):
+        #     module.weight.data.normal_(mean=0.0, std=std)
+        #     module.bias.data.normal_(mean=0.0, std=std)
+
+
+@auto_docstring
+class GptOssModel(GptOssPreTrainedModel):
+    _no_split_modules = ["GptOssDecoderLayer"]
+
+    def __init__(self, config: GptOssConfig):
+        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(
+            [GptOssDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self.norm = GptOssRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = GptOssRotaryEmbedding(config=config)
+        self.gradient_checkpointing = False
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @check_model_inputs
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[list[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> MoeModelOutputWithPast:
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache()
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        if cache_position is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(
+                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+            )
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        # It may already have been prepared by e.g. `generate`
+        if not isinstance(causal_mask_mapping := attention_mask, dict):
+            mask_kwargs = {
+                "config": self.config,
+                "input_embeds": inputs_embeds,
+                "attention_mask": attention_mask,
+                "cache_position": cache_position,
+                "past_key_values": past_key_values,
+            }
+            causal_mask_mapping = {
+                "full_attention": create_causal_mask(**mask_kwargs),
+                "sliding_attention": create_sliding_window_causal_mask(**mask_kwargs),
+            }
+
+        hidden_states = inputs_embeds
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        for decoder_layer in self.layers:
+            hidden_states = decoder_layer(
+                hidden_states,
+                attention_mask=causal_mask_mapping[decoder_layer.attention_type],
+                position_ids=position_ids,
+                past_key_value=past_key_values,
+                use_cache=use_cache,
+                cache_position=cache_position,
+                position_embeddings=position_embeddings,
+                **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: Union[torch.Tensor, tuple[torch.Tensor], None],
+    num_experts: Optional[int] = None,
+    top_k=2,
+    attention_mask: Optional[torch.Tensor] = None,
+) -> Union[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 GptOssForCausalLM(GptOssPreTrainedModel, GenerationMixin):
+    _tied_weights_keys = ["lm_head.weight"]
+    _tp_plan = {"lm_head": "colwise_rep"}
+    _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = GptOssModel(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_local_experts
+        self.num_experts_per_tok = config.num_experts_per_tok
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def get_decoder(self):
+        return self.model
+
+    @can_return_tuple
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_router_logits: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        logits_to_keep: Union[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]`.
+
+        Example:
+
+        ```python
+        >>> from transformers import AutoTokenizer, GptOssForCausalLM
+
+        >>> model = GptOssForCausalLM.from_pretrained("mistralai/GptOss-8x7B-v0.1")
+        >>> tokenizer = AutoTokenizer.from_pretrained("mistralai/GptOss-8x7B-v0.1")
+
+        >>> prompt = "Hey, are you conscious? Can you talk to me?"
+        >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+        >>> # Generate
+        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+        "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
+        ```"""
+
+        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,
+            cache_position=cache_position,
+            **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__ = ["GptOssForCausalLM", "GptOssModel", "GptOssPreTrainedModel"]

quantization_config.json

@@ -0,0 +1,554 @@
+{
+  "bits": 4,
+  "group_size": 64,
+  "sym": true,
+  "data_type": "int",
+  "batch_size": 4,
+  "nsamples": 512,
+  "low_gpu_mem_usage": true,
+  "autoround_version": "0.6.1.dev",
+  "quant_method": "auto-round",
+  "packing_format": "auto_round:auto_gptq",
+  "extra_config": {
+    "model.layers.0.self_attn.q_proj": {
+      "bits": 16
+    },
+    "model.layers.0.self_attn.k_proj": {
+      "bits": 16
+    },
+    "model.layers.0.self_attn.v_proj": {
+      "bits": 16
+    },
+    "model.layers.0.self_attn.o_proj": {
+      "bits": 16
+    },
+    "model.layers.0.mlp.router.router": {
+      "bits": 16
+    },
+    "model.layers.1.self_attn.q_proj": {
+      "bits": 16
+    },
+    "model.layers.1.self_attn.k_proj": {
+      "bits": 16
+    },
+    "model.layers.1.self_attn.v_proj": {
+      "bits": 16
+    },
+    "model.layers.1.self_attn.o_proj": {
+      "bits": 16
+    },
+    "model.layers.1.mlp.router.router": {
+      "bits": 16
+    },
+    "model.layers.2.self_attn.q_proj": {
+      "bits": 16
+    },
+    "model.layers.2.self_attn.k_proj": {
+      "bits": 16
+    },
+    "model.layers.2.self_attn.v_proj": {
+      "bits": 16
+    },
+    "model.layers.2.self_attn.o_proj": {
+      "bits": 16
+    },
+    "model.layers.2.mlp.router.router": {
+      "bits": 16
+    },
+    "model.layers.3.self_attn.q_proj": {
+      "bits": 16
+    },
+    "model.layers.3.self_attn.k_proj": {
+      "bits": 16
+    },
+    "model.layers.3.self_attn.v_proj": {
+      "bits": 16
+    },
+    "model.layers.3.self_attn.o_proj": {
+      "bits": 16
+    },
+    "model.layers.3.mlp.router.router": {
+      "bits": 16
+    },
+    "model.layers.4.self_attn.q_proj": {
+      "bits": 16
+    },
+    "model.layers.4.self_attn.k_proj": {
+      "bits": 16
+    },
+    "model.layers.4.self_attn.v_proj": {
+      "bits": 16
+    },
+    "model.layers.4.self_attn.o_proj": {
+      "bits": 16
+    },
+    "model.layers.4.mlp.router.router": {
+      "bits": 16
+    },
+    "model.layers.5.self_attn.q_proj": {
+      "bits": 16
+    },
+    "model.layers.5.self_attn.k_proj": {
+      "bits": 16
+    },
+    "model.layers.5.self_attn.v_proj": {
+      "bits": 16
+    },
+    "model.layers.5.self_attn.o_proj": {
+      "bits": 16
+    },
+    "model.layers.5.mlp.router.router": {
+      "bits": 16
+    },
+    "model.layers.6.self_attn.q_proj": {
+      "bits": 16
+    },
+    "model.layers.6.self_attn.k_proj": {
+      "bits": 16
+    },
+    "model.layers.6.self_attn.v_proj": {
+      "bits": 16
+    },
+    "model.layers.6.self_attn.o_proj": {
+      "bits": 16
+    },
+    "model.layers.6.mlp.router.router": {
+      "bits": 16
+    },
+    "model.layers.7.self_attn.q_proj": {
+      "bits": 16
+    },
+    "model.layers.7.self_attn.k_proj": {
+      "bits": 16
+    },
+    "model.layers.7.self_attn.v_proj": {
+      "bits": 16
+    },
+    "model.layers.7.self_attn.o_proj": {
+      "bits": 16
+    },
+    "model.layers.7.mlp.router.router": {
+      "bits": 16
+    },
+    "model.layers.8.self_attn.q_proj": {
+      "bits": 16
+    },
+    "model.layers.8.self_attn.k_proj": {
+      "bits": 16
+    },
+    "model.layers.8.self_attn.v_proj": {
+      "bits": 16
+    },
+    "model.layers.8.self_attn.o_proj": {
+      "bits": 16
+    },
+    "model.layers.8.mlp.router.router": {
+      "bits": 16
+    },
+    "model.layers.9.self_attn.q_proj": {
+      "bits": 16
+    },
+    "model.layers.9.self_attn.k_proj": {
+      "bits": 16
+    },
+    "model.layers.9.self_attn.v_proj": {
+      "bits": 16
+    },
+    "model.layers.9.self_attn.o_proj": {
+      "bits": 16
+    },
+    "model.layers.9.mlp.router.router": {
+      "bits": 16
+    },
+    "model.layers.10.self_attn.q_proj": {
+      "bits": 16
+    },
+    "model.layers.10.self_attn.k_proj": {
+      "bits": 16
+    },
+    "model.layers.10.self_attn.v_proj": {
+      "bits": 16
+    },
+    "model.layers.10.self_attn.o_proj": {
+      "bits": 16
+    },
+    "model.layers.10.mlp.router.router": {
+      "bits": 16
+    },
+    "model.layers.11.self_attn.q_proj": {
+      "bits": 16
+    },
+    "model.layers.11.self_attn.k_proj": {
+      "bits": 16
+    },
+    "model.layers.11.self_attn.v_proj": {
+      "bits": 16
+    },
+    "model.layers.11.self_attn.o_proj": {
+      "bits": 16
+    },
+    "model.layers.11.mlp.router.router": {
+      "bits": 16
+    },
+    "model.layers.12.self_attn.q_proj": {
+      "bits": 16
+    },
+    "model.layers.12.self_attn.k_proj": {
+      "bits": 16
+    },
+    "model.layers.12.self_attn.v_proj": {
+      "bits": 16
+    },
+    "model.layers.12.self_attn.o_proj": {
+      "bits": 16
+    },
+    "model.layers.12.mlp.router.router": {
+      "bits": 16
+    },
+    "model.layers.13.self_attn.q_proj": {
+      "bits": 16
+    },
+    "model.layers.13.self_attn.k_proj": {
+      "bits": 16
+    },
+    "model.layers.13.self_attn.v_proj": {
+      "bits": 16
+    },
+    "model.layers.13.self_attn.o_proj": {
+      "bits": 16
+    },
+    "model.layers.13.mlp.router.router": {
+      "bits": 16
+    },
+    "model.layers.14.self_attn.q_proj": {
+      "bits": 16
+    },
+    "model.layers.14.self_attn.k_proj": {
+      "bits": 16
+    },
+    "model.layers.14.self_attn.v_proj": {
+      "bits": 16
+    },
+    "model.layers.14.self_attn.o_proj": {
+      "bits": 16
+    },
+    "model.layers.14.mlp.router.router": {
+      "bits": 16
+    },
+    "model.layers.15.self_attn.q_proj": {
+      "bits": 16
+    },
+    "model.layers.15.self_attn.k_proj": {
+      "bits": 16
+    },
+    "model.layers.15.self_attn.v_proj": {
+      "bits": 16
+    },
+    "model.layers.15.self_attn.o_proj": {
+      "bits": 16
+    },
+    "model.layers.15.mlp.router.router": {
+      "bits": 16
+    },
+    "model.layers.16.self_attn.q_proj": {
+      "bits": 16
+    },
+    "model.layers.16.self_attn.k_proj": {
+      "bits": 16
+    },
+    "model.layers.16.self_attn.v_proj": {
+      "bits": 16
+    },
+    "model.layers.16.self_attn.o_proj": {
+      "bits": 16
+    },
+    "model.layers.16.mlp.router.router": {
+      "bits": 16
+    },
+    "model.layers.17.self_attn.q_proj": {
+      "bits": 16
+    },
+    "model.layers.17.self_attn.k_proj": {
+      "bits": 16
+    },
+    "model.layers.17.self_attn.v_proj": {
+      "bits": 16
+    },
+    "model.layers.17.self_attn.o_proj": {
+      "bits": 16
+    },
+    "model.layers.17.mlp.router.router": {
+      "bits": 16
+    },
+    "model.layers.18.self_attn.q_proj": {
+      "bits": 16
+    },
+    "model.layers.18.self_attn.k_proj": {
+      "bits": 16
+    },
+    "model.layers.18.self_attn.v_proj": {
+      "bits": 16
+    },
+    "model.layers.18.self_attn.o_proj": {
+      "bits": 16
+    },
+    "model.layers.18.mlp.router.router": {
+      "bits": 16
+    },
+    "model.layers.19.self_attn.q_proj": {
+      "bits": 16
+    },
+    "model.layers.19.self_attn.k_proj": {
+      "bits": 16
+    },
+    "model.layers.19.self_attn.v_proj": {
+      "bits": 16
+    },
+    "model.layers.19.self_attn.o_proj": {
+      "bits": 16
+    },
+    "model.layers.19.mlp.router.router": {
+      "bits": 16
+    },
+    "model.layers.20.self_attn.q_proj": {
+      "bits": 16
+    },
+    "model.layers.20.self_attn.k_proj": {
+      "bits": 16
+    },
+    "model.layers.20.self_attn.v_proj": {
+      "bits": 16
+    },
+    "model.layers.20.self_attn.o_proj": {
+      "bits": 16
+    },
+    "model.layers.20.mlp.router.router": {
+      "bits": 16
+    },
+    "model.layers.21.self_attn.q_proj": {
+      "bits": 16
+    },
+    "model.layers.21.self_attn.k_proj": {
+      "bits": 16
+    },
+    "model.layers.21.self_attn.v_proj": {
+      "bits": 16
+    },
+    "model.layers.21.self_attn.o_proj": {
+      "bits": 16
+    },
+    "model.layers.21.mlp.router.router": {
+      "bits": 16
+    },
+    "model.layers.22.self_attn.q_proj": {
+      "bits": 16
+    },
+    "model.layers.22.self_attn.k_proj": {
+      "bits": 16
+    },
+    "model.layers.22.self_attn.v_proj": {
+      "bits": 16
+    },
+    "model.layers.22.self_attn.o_proj": {
+      "bits": 16
+    },
+    "model.layers.22.mlp.router.router": {
+      "bits": 16
+    },
+    "model.layers.23.self_attn.q_proj": {
+      "bits": 16
+    },
+    "model.layers.23.self_attn.k_proj": {
+      "bits": 16
+    },
+    "model.layers.23.self_attn.v_proj": {
+      "bits": 16
+    },
+    "model.layers.23.self_attn.o_proj": {
+      "bits": 16
+    },
+    "model.layers.23.mlp.router.router": {
+      "bits": 16
+    },
+    "model.layers.24.self_attn.q_proj": {
+      "bits": 16
+    },
+    "model.layers.24.self_attn.k_proj": {
+      "bits": 16
+    },
+    "model.layers.24.self_attn.v_proj": {
+      "bits": 16
+    },
+    "model.layers.24.self_attn.o_proj": {
+      "bits": 16
+    },
+    "model.layers.24.mlp.router.router": {
+      "bits": 16
+    },
+    "model.layers.25.self_attn.q_proj": {
+      "bits": 16
+    },
+    "model.layers.25.self_attn.k_proj": {
+      "bits": 16
+    },
+    "model.layers.25.self_attn.v_proj": {
+      "bits": 16
+    },
+    "model.layers.25.self_attn.o_proj": {
+      "bits": 16
+    },
+    "model.layers.25.mlp.router.router": {
+      "bits": 16
+    },
+    "model.layers.26.self_attn.q_proj": {
+      "bits": 16
+    },
+    "model.layers.26.self_attn.k_proj": {
+      "bits": 16
+    },
+    "model.layers.26.self_attn.v_proj": {
+      "bits": 16
+    },
+    "model.layers.26.self_attn.o_proj": {
+      "bits": 16
+    },
+    "model.layers.26.mlp.router.router": {
+      "bits": 16
+    },
+    "model.layers.27.self_attn.q_proj": {
+      "bits": 16
+    },
+    "model.layers.27.self_attn.k_proj": {
+      "bits": 16
+    },
+    "model.layers.27.self_attn.v_proj": {
+      "bits": 16
+    },
+    "model.layers.27.self_attn.o_proj": {
+      "bits": 16
+    },
+    "model.layers.27.mlp.router.router": {
+      "bits": 16
+    },
+    "model.layers.28.self_attn.q_proj": {
+      "bits": 16
+    },
+    "model.layers.28.self_attn.k_proj": {
+      "bits": 16
+    },
+    "model.layers.28.self_attn.v_proj": {
+      "bits": 16
+    },
+    "model.layers.28.self_attn.o_proj": {
+      "bits": 16
+    },
+    "model.layers.28.mlp.router.router": {
+      "bits": 16
+    },
+    "model.layers.29.self_attn.q_proj": {
+      "bits": 16
+    },
+    "model.layers.29.self_attn.k_proj": {
+      "bits": 16
+    },
+    "model.layers.29.self_attn.v_proj": {
+      "bits": 16
+    },
+    "model.layers.29.self_attn.o_proj": {
+      "bits": 16
+    },
+    "model.layers.29.mlp.router.router": {
+      "bits": 16
+    },
+    "model.layers.30.self_attn.q_proj": {
+      "bits": 16
+    },
+    "model.layers.30.self_attn.k_proj": {
+      "bits": 16
+    },
+    "model.layers.30.self_attn.v_proj": {
+      "bits": 16
+    },
+    "model.layers.30.self_attn.o_proj": {
+      "bits": 16
+    },
+    "model.layers.30.mlp.router.router": {
+      "bits": 16
+    },
+    "model.layers.31.self_attn.q_proj": {
+      "bits": 16
+    },
+    "model.layers.31.self_attn.k_proj": {
+      "bits": 16
+    },
+    "model.layers.31.self_attn.v_proj": {
+      "bits": 16
+    },
+    "model.layers.31.self_attn.o_proj": {
+      "bits": 16
+    },
+    "model.layers.31.mlp.router.router": {
+      "bits": 16
+    },
+    "model.layers.32.self_attn.q_proj": {
+      "bits": 16
+    },
+    "model.layers.32.self_attn.k_proj": {
+      "bits": 16
+    },
+    "model.layers.32.self_attn.v_proj": {
+      "bits": 16
+    },
+    "model.layers.32.self_attn.o_proj": {
+      "bits": 16
+    },
+    "model.layers.32.mlp.router.router": {
+      "bits": 16
+    },
+    "model.layers.33.self_attn.q_proj": {
+      "bits": 16
+    },
+    "model.layers.33.self_attn.k_proj": {
+      "bits": 16
+    },
+    "model.layers.33.self_attn.v_proj": {
+      "bits": 16
+    },
+    "model.layers.33.self_attn.o_proj": {
+      "bits": 16
+    },
+    "model.layers.33.mlp.router.router": {
+      "bits": 16
+    },
+    "model.layers.34.self_attn.q_proj": {
+      "bits": 16
+    },
+    "model.layers.34.self_attn.k_proj": {
+      "bits": 16
+    },
+    "model.layers.34.self_attn.v_proj": {
+      "bits": 16
+    },
+    "model.layers.34.self_attn.o_proj": {
+      "bits": 16
+    },
+    "model.layers.34.mlp.router.router": {
+      "bits": 16
+    },
+    "model.layers.35.self_attn.q_proj": {
+      "bits": 16
+    },
+    "model.layers.35.self_attn.k_proj": {
+      "bits": 16
+    },
+    "model.layers.35.self_attn.v_proj": {
+      "bits": 16
+    },
+    "model.layers.35.self_attn.o_proj": {
+      "bits": 16
+    },
+    "model.layers.35.mlp.router.router": {
+      "bits": 16
+    }
+  }
+}

special_tokens_map.json

@@ -0,0 +1,23 @@
+{
+  "bos_token": {
+    "content": "<|startoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|return|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer.json

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

tokenizer_config.json

@@ -0,0 +1,183 @@
+{
+  "added_tokens_decoder": {
+    "199998": {
+      "content": "<|startoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "199999": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "200000": {
+      "content": "<|reserved_200000|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "200001": {
+      "content": "<|reserved_200001|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "200002": {
+      "content": "<|return|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "200003": {
+      "content": "<|constrain|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "200004": {
+      "content": "<|reserved_200004|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "200005": {
+      "content": "<|channel|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "200006": {
+      "content": "<|start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "200007": {
+      "content": "<|end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "200008": {
+      "content": "<|message|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "200009": {
+      "content": "<|reserved_200009|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "200010": {
+      "content": "<|reserved_200010|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "200011": {
+      "content": "<|reserved_200011|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "200012": {
+      "content": "<|call|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "200013": {
+      "content": "<|reserved_200013|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "200014": {
+      "content": "<|reserved_200014|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "200015": {
+      "content": "<|reserved_200015|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "200016": {
+      "content": "<|reserved_200016|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "200017": {
+      "content": "<|reserved_200017|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "200018": {
+      "content": "<|endofprompt|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<|startoftext|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|return|>",
+  "extra_special_tokens": {},
+  "model_input_names": [
+    "input_ids",
+    "attention_mask"
+  ],
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<|endoftext|>",
+  "tokenizer_class": "PreTrainedTokenizerFast"
+}