Grug 67B A2B Datakit SFT 262K — Special-token learning rate

Research artifact. This model has not been properly tested or evaluated and is not validated for production use. The 262K context length is a configuration value; long-context performance remains untested.

This repository contains a BF16 Grug mixture-of-experts checkpoint after supervised fine-tuning (SFT) on the agentic and reasoning mixture specified below. Long-context pretraining replay supplied 20% of scheduled token positions.

Training data

The run scheduled 67,108,864,000 token positions: 1,000 updates × 256 packed sequences × 262,144 tokens. The fixed mixture assigned 53,687,091,200 positions (80%) to SFT and 13,421,772,800 positions (20%) to long-context pretraining replay. It selected 203 SFT leaf sources. Sources smaller than one sample in a 64,000-sample mixture block were placed in a pooled component so they remained sampleable.

The SFT allocation by source family was:

training_data_mix.json is the exact machine-readable specification. It records all 203 SFT leaf allocations, the excluded-source list, the replay recipe, and the 356 top-level weights saved in the finished W&B run. The checked in allocation and the W&B configuration agree for every unpooled SFT source; the pooled SFT weight also agrees exactly.

Model and provenance

The starting point was an unpublished base checkpoint at step 157,000. Training resumed its optimizer state without warmup. The peak base learning rate was 5e-5; updates to seven special-token rows in both input embeddings and the LM head were scaled by sqrt(32), giving an effective peak learning rate of 2.828427e-4 for those rows. Router weights remained trainable. Router biases and deferred load-balancing bias updates stayed frozen.

Seven LM-head rows were initialized from single-token anchors: 128006 from 3560 ( role), 128007 from 1984 ( message), 128002 from 1781 ( think), 128003 from 4320 ( answer), 128005 from 5507 ( tool), 128011 from 842 ( end), and 128009 from EOS (128001). Only those LM-head optimizer moments were reset. Input embeddings and their optimizer moments were preserved at initialization and remained trainable. BOS and EOS were preserved; no new tokens were added. The SFT data loader enforced a one-epoch limit.

Training run: grug-67b-sft-20260920-special-token-lr-frozen-bias-1000-cpfix.

Seven attention layers use full causal attention; the other 19 use the sliding window.

Chat and generation settings

Omitting enable_thinking defaults to thinking. For a vLLM chat request, select nonthinking mode with:

{"chat_template_kwargs": {"enable_thinking": false}}

For local prompt formatting:

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained(
    "open-athena/Grug-67B-A2B-Datakit-SFT-262K-2026.09.21",
)
prompt_ids = tokenizer.apply_chat_template(
    [{"role": "user", "content": "What is 17 times 23?"}],
    tokenize=True, return_dict=False,
    add_generation_prompt=True, enable_thinking=False,
)

The exported generation configuration specifies both <|end_of_text|> (128001) and <|eot_id|> (128009) as stop IDs. The tokenizer's EOS remains 128001. Preserve both generation stop IDs when configuring another runtime.

Thinking output uses <|start_think|> (128002) and <|end_think|> (128003). Decode with skip_special_tokens=False when these boundaries are needed. Tools use JSON inside <tool_call>...</tool_call>; input arguments may be objects or JSON strings. Tool-call delimiters remain ordinary text, not newly added special tokens.

The inference template adds the thinking default to the training template. Export checks cover default and explicit thinking modes, assistant text beside tool calls, and equivalent formatting for string messages and messages represented as lists of text parts.

Serving

This architecture requires a runtime with Grug MoE support. No serving runtime has been validated for this checkpoint. Stock Transformers model loading has not been validated.

Evaluation and limitations

This model is a research artifact and has not been properly tested or evaluated. No inference smoke test or benchmark result is established for this checkpoint. Training loss and export-format checks do not establish generation quality, safety, or numerical equivalence to the native checkpoint.

License

The model materials are released under the OpenMDW License Agreement, version 1.1.

We release our tools to the public with the intent that they are used safely, ethically, and responsibly. They are meant to contribute to the public's understanding of science, and aid in its advancement. By using this Marin model, you agree to comply with all applicable law.