shibatch/tinygptossmoe3m

Tiny gpt-oss MoE 3M

This repository contains a tiny gpt-oss-style text-only Mixture-of-Experts causal language model for validation and debugging.

The model is intentionally small. It is not intended to be a high-quality text generation model. Its main purpose is to provide a compact checkpoint that exercises gpt-oss MoE text-model code paths in Hugging Face Transformers and in independent inference engines.

This checkpoint is useful for implementation testing because it includes sliding attention and full attention layers, grouped-query attention, YaRN-style RoPE configuration, untied input/output embeddings, and top-4 MoE routing with multiple local experts.

This is a synthetic tiny validation checkpoint. It is not an official OpenAI model and does not contain weights from the original gpt-oss checkpoints.

Repository weight profiles

This repository contains two representations of the same trained model:

Subfolder	Purpose	Precision
hf/	Hugging Face GptOssForCausalLM	FP32
original/	OpenAI original-format loader	official-profile MXFP4 + BF16

The original/ checkpoint is not a blanket 4-bit conversion. Only expert mlp1_weight and mlp2_weight tensors use MXFP4 .blocks and .scales. Embedding, attention, router, bias, norm, sink, and unembedding tensors remain BF16, matching the official gpt-oss precision profile.

Model purpose

This model is designed for:

• testing GptOssForCausalLM
• validating GptOssConfig
• testing gpt-oss-style MoE model loading
• checking model save/load behavior
• checking tokenizer save/load behavior
• exercising sliding attention layers
• exercising full attention layers
• exercising grouped-query attention
• exercising YaRN-style RoPE configuration parsing
• exercising untied input/output embedding paths
• exercising MoE expert parameters
• exercising top-4 expert routing
• providing a compact gpt-oss-style MoE checkpoint for inference-engine validation
• testing official-profile MXFP4 tensor loading and decoding
• comparing full-sequence and KV-cached logits

It is not designed for:

• high-quality story generation
• instruction following
• chat use
• benchmark comparison against production language models
• production deployment
• reproducing the behavior of the original gpt-oss models

Model architecture

The model uses GptOssForCausalLM with a small gpt-oss-style MoE configuration.

Representative configuration:

model_type: gpt_oss
vocab_size: 1024

hidden_size: 128
intermediate_size: 128

num_hidden_layers: 6
num_attention_heads: 4
num_key_value_heads: 1
head_dim: 32

sliding_window: 128
max_position_embeddings: 4096
initial_context_length: 1024

layer_types:
  - sliding_attention
  - full_attention
  - sliding_attention
  - full_attention
  - sliding_attention
  - full_attention

hidden_act: silu
tie_word_embeddings: false
attention_bias: true
attention_dropout: 0.0
rms_norm_eps: 1e-05
initializer_range: 0.02

rope_theta: 150000.0
rope_scaling:
  rope_type: yarn
  factor: 4.0
  original_max_position_embeddings: 1024
  beta_fast: 32.0
  beta_slow: 1.0
  truncate: false

num_local_experts: 8
num_experts_per_tok: 4
experts_per_token: 4
output_router_logits: false
router_aux_loss_coef: 0.0
swiglu_limit: 7.0

pad_token_id: 1000
bos_token_id: 1000
eos_token_id: 1001

The attention pattern is:

sFsFsF

where s means sliding_attention and F means full_attention.

This pattern was chosen for validation coverage. A full-attention-only model may be easier to train, but it would not exercise the sliding attention path.

MoE configuration

This model enables gpt-oss-style MoE blocks.

num_local_experts: 8
num_experts_per_tok: 4
experts_per_token: 4
intermediate_size: 128

The num_local_experts=8 and num_experts_per_tok=4 setting is intentional. It keeps the model small while still exercising top-4 routing without selecting all experts on every token.

This checkpoint is intended to cover:

router parameters
multiple local experts
top-4 expert selection
weighted expert combination
MoE FFN parameters
sliding and full attention interaction
GQA with num_key_value_heads = 1
untied input/output embeddings

Training data

The model was trained on TinyStories-style English story text.

The model uses a small custom byte-level BPE tokenizer. The tokenizer is designed for tiny validation models rather than production text quality. The configuration keeps the checkpoint compact and makes the model easier to train at very small scale.

Tokenizer

The reference training script uses a legacy byte-level BPE tokenizer setup:

RawTokenizer(BPE())
ByteLevel(add_prefix_space=False)
BpeTrainer(vocab_size=1000, min_frequency=2, special_tokens=[], initial_alphabet=ByteLevel.alphabet())
normalizer: None

Special tokens are added after BPE training:

<s>          -> expected id 1000
</s>         -> expected id 1001
<|im_start|> -> expected id 1002

The tokenizer uses:

bos_token: <s>
eos_token: </s>
pad_token: <s>

The model config uses vocab_size=1024 to leave a small reserved range above the learned base vocabulary and special tokens.

This tokenizer setup was chosen because it produced substantially better tiny-model training behavior than the earlier tokenizer configuration where special tokens were included directly in the BPE training step.

Training setup

Representative training settings:

num_epochs: 1
learning_rate: 2e-4
batch_size: 32
block_size: 256

device: auto
dtype: auto, resolved to float32 by the training script
grad_clip: 1.0
error_on_nonfinite_gradients: true

vocab_size: 1024
base_vocab_size: 1000
hidden_size: 128
intermediate_size: 128
num_hidden_layers: 6
num_attention_heads: 4
num_key_value_heads: 1
head_dim: 32
layer_pattern: sFsFsF
sliding_window: 128
max_position_embeddings: 4096
initial_context_length: 1024

num_local_experts: 8
num_experts_per_tok: 4
router_aux_loss_coef: 0.0
attention_bias: true
tie_word_embeddings: false

The final evaluation loss in the reference run was approximately:

Final loss: 1.4606

This loss should not be interpreted as a general language-model quality benchmark. The model is very small and includes gpt-oss-specific architectural paths primarily for validation coverage.

Example generation

Example output from the reference checkpoint:

Prompt: Once upon

Once upon a time, there was a little girl named Lily. She loved to play outside in the sun. One day, she saw a big, scary dog. The dog was barking and running around. Lily wanted to pet the dog, but she didn't want to leave.

Suddenly, a big dog came running towards her. Lily was scared and didn't know what to do. But then she remembered that she had a friend

Prompt: There was a little

There was a little girl named Lily who was very curious. She wanted to know what was inside the box. She asked her mom, "What is in the box?" Her mom said, "It's a box of candy. It's a special treat for you."

Lily was excited to open the box. She opened the box and found a big, shiny candy inside. She was so happy and said, "Thank you, Mommy!" Her mom smiled and

Prompt: One day

One day, a little girl named Lily went to the park with her mom. She saw a big slide and wanted to try it. She asked her mom if she could go on the slide. Her mom said yes, but only if she promised to be careful.

Lily was so happy that she ran up the slide and started to slide down. She slid down the slide and laughed as she slid down. She felt so happy and free.

The model can generate TinyStories-like text fragments, but repetition, template collapse, and weak long-form coherence are expected. This is normal for this checkpoint and is not considered a failure for its intended purpose.

Hugging Face FP32 usage

import torch
from transformers import PreTrainedTokenizerFast, GptOssForCausalLM

repo = "shibatch/tinygptossmoe3m"

tokenizer = PreTrainedTokenizerFast.from_pretrained(repo, subfolder="hf")
model = GptOssForCausalLM.from_pretrained(
    repo,
    subfolder="hf",
    torch_dtype=torch.float32,
)
model.eval()

prompt = "Once upon"
inputs = tokenizer(prompt, return_tensors="pt")

with torch.no_grad():
    output_ids = model.generate(
        **inputs,
        max_new_tokens=100,
        do_sample=False,
        pad_token_id=tokenizer.pad_token_id,
        eos_token_id=tokenizer.eos_token_id,
    )

print(tokenizer.decode(output_ids[0], skip_special_tokens=True))

Loading requirements

The hf/ checkpoint requires a Transformers version that supports gpt-oss models.

The following imports should work:

from transformers import GptOssForCausalLM, GptOssConfig

If these imports fail, update Transformers to a version with gpt-oss support.

Tokenizer loading note

This repository uses a custom byte-level BPE tokenizer saved as a PreTrainedTokenizerFast.

For this reason, examples use:

from transformers import PreTrainedTokenizerFast

instead of AutoTokenizer.

Using AutoTokenizer may fail in some environments if the tokenizer backend cannot be inferred automatically.

The expected tokenizer files include:

tokenizer.json
tokenizer_config.json
special_tokens_map.json

Intended validation coverage

This checkpoint is intended to validate support for:

GptOssConfig
GptOssForCausalLM
sliding_attention layers
full_attention layers
GQA with num_key_value_heads = 1
YaRN-style rope_scaling config parsing
attention_bias = true
untied input/output embeddings
gpt-oss RMSNorm behavior
gpt-oss MLP activation path
gpt-oss MoE expert parameters
num_local_experts = 8
num_experts_per_tok = 4
experts_per_token = 4
MoE expert dispatch
MoE expert output combination
generate()
save_pretrained()
from_pretrained()
official gpt-oss original-format loading
MXFP4 E2M1 expert blocks
E8M0 expert scales
full versus KV-cached logits
top-4 routing IDs and weights
selected expert projection outputs

Limitations

This is a tiny debug model. It should not be used as a general-purpose language model.

Known limitations:

• frequent phrase repetition
• TinyStories template collapse
• weak long-form coherence
• small vocabulary
• weak semantic consistency
• no instruction tuning
• no chat formatting
• no production use
• no compatibility claim with original gpt-oss model quality
• the tiny attention dimensions require the documented reference-attention adapter for OpenAI's Triton implementation

The checkpoint is primarily intended to make gpt-oss-style MoE text-model code paths easy to test without downloading a large model.

Why include MoE?

A dense tiny model is simpler to train, but it does not cover MoE-specific implementation paths.

This checkpoint intentionally includes:

num_local_experts = 8
num_experts_per_tok = 4

to exercise a more realistic top-4 MoE routing path than a minimal top_k=1 configuration.

Why not full attention only?

A full-attention-only tiny model may train more cleanly, but it would not cover gpt-oss-style sliding attention behavior.

This checkpoint uses:

sliding_attention
full_attention
sliding_attention
full_attention
sliding_attention
full_attention

to cover both attention implementations.

Notes on MXFP4 and quantization

The original/model.safetensors file follows the naming, shape, packing, and mixed-precision profile used by OpenAI's original-format gpt-oss checkpoints. There is no separate .mxfp4 file extension: MXFP4 tensors are represented by paired .blocks and .scales tensors inside SafeTensors.

block.N.mlp.mlp1_weight.blocks   uint8 E2M1, two values per byte
block.N.mlp.mlp1_weight.scales   uint8 E8M0
block.N.mlp.mlp2_weight.blocks   uint8 E2M1, two values per byte
block.N.mlp.mlp2_weight.scales   uint8 E8M0

Packing and conversion rules:

• 32 E2M1 values per scale block
• first/even value in the low nibble
• second/odd value in the high nibble
• E8M0 exponent bias 127
• OCP ROUND_DOWN scale selection
• round-to-nearest, ties-to-even E2M1 conversion
• block scaling along the final logical matrix dimension

The model contains 93 physical tensors: 69 BF16, 12 FP4 block tensors, and 12 UE8 scale tensors. Exact names, dtypes, shapes, byte sizes, and raw-tensor SHA-256 values are in original/tensor_inventory.json. A standalone packing example is in mxfp4_packing_test_vector.json.

The fixed model SHA-256 is:

30dac8e68ccf869d872beb382cebd0be8b168de761dfb38f1358d260d4b8d09d

Fixed revisions

OpenAI gpt-oss:
599476783c6f88508dab8577808b5ead5cbee8d2

Triton kernels:
9e1e203f64752cf99abf0e44286231c5d5df7e76

Official profile reference openai/gpt-oss-20b:
6cee5e81ee83917806bbde320786a8fb61efebee

Exact package versions are recorded in MXFP4_VERSION_LOCK.json and requirements-mxfp4.txt.

python3 -m venv .venv-mxfp4
source .venv-mxfp4/bin/activate
pip install -r requirements-mxfp4.txt

git clone https://github.com/openai/gpt-oss.git third_party/openai-gpt-oss
git -C third_party/openai-gpt-oss checkout --detach \
  599476783c6f88508dab8577808b5ead5cbee8d2

git clone https://github.com/triton-lang/triton.git third_party/triton
git -C third_party/triton checkout --detach \
  9e1e203f64752cf99abf0e44286231c5d5df7e76

pip install --no-deps -e third_party/triton/python/triton_kernels

Reproduce the conversion

After cloning the pinned OpenAI and Triton repositories and installing triton_kernels, run from this repository's root:

PYTHONPATH=third_party/openai-gpt-oss:third_party/triton/python/triton_kernels \
python tools/convert_tinygptoss_to_mxfp4.py \
  --source-dir hf \
  --output-dir rebuilt/original \
  --openai-repo third_party/openai-gpt-oss \
  --triton-repo third_party/triton

Fixed reference outputs

mxfp4_reference_outputs/ contains:

• full FP32 logits [1, 5, 1024]
• token-by-token cached logits [1, 5, 1024]
• top-4 routing IDs, weights, and router logits for all six layers
• selected expert projection and routing-weighted output for all six layers
• a 64-token greedy continuation

For the fixed prompt Once upon a time,, full and cached logits are byte-identical:

max_abs_diff: 0.0
mean_abs_diff: 0.0
all_token_argmax_equal: true

Regenerate them with:

CUDA_VISIBLE_DEVICES=0 \
PYTHONPATH=third_party/openai-gpt-oss:third_party/triton/python/triton_kernels \
python tools/generate_gptoss_mxfp4_reference.py \
  --checkpoint original \
  --tokenizer hf \
  --output-dir rebuilt/mxfp4_reference_outputs \
  --openai-repo third_party/openai-gpt-oss \
  --triton-repo third_party/triton \
  --prompt 'Once upon a time,' \
  --max-new-tokens 64 \
  --context 256 \
  --device cuda:0

OpenAI's Triton attention code assumes dimensions used by the 20B/120B models. For this tiny model's one KV head, four GQA groups, and 32-wide head, the harness transposes the two grouping axes and calls the pinned official attention_ref. OpenAI's model source is not patched; its loader, KV cache, Transformer, routing, and Triton MXFP4 MoE paths remain in use. The adapter is fully recorded in mxfp4_reference_outputs/reference_manifest.json.

Verify all repository payload hashes, all 93 tensor hashes, and MXFP4 shapes:

python tools/verify_mxfp4.py

Suggested repository name

Suggested Hugging Face repository name:

shibatch/tinygptossmoe3m

Citation

This is a synthetic tiny validation checkpoint derived from gpt-oss-compatible MoE text architecture settings. It is intended for debugging and implementation testing.