Top1MoE-0.6B

OLMoE-1B-7B-0924 weights with top-1 routing hard-wired. 64 experts per layer, exactly one expert active per token. Total parameters 6.92B, active parameters per token 0.58B (top-8 upstream: 1.28B), hence "0.6B".

This repository ships its own model code (model_type = "top1moe", classes Top1MoeConfig, Top1MoeForCausalLM, ...). Load it with trust_remote_code=True. It requires transformers >= 5.0 (the code is derived from transformers 5.3.0 internals and raises a clear ImportError on 4.x).

Usage

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

repo = "JIRU-W/top1moe-0.6b"
tok = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)  # tokenizer itself is the stock OLMoE one
model = AutoModelForCausalLM.from_pretrained(
    repo, trust_remote_code=True, dtype=torch.bfloat16, device_map="cuda")

inputs = tok("The capital of France is", return_tensors="pt").to(model.device)
out = model.generate(**inputs, max_new_tokens=16, do_sample=False)
print(tok.decode(out[0], skip_special_tokens=True))

AutoConfig / AutoModel / AutoModelForCausalLM are all mapped through auto_map in config.json. After one trust_remote_code=True load, transformers has registered the classes in the current process, so model.save_pretrained(...) and model.push_to_hub(...) keep the custom code.

What "top-1 hard-wired" means

Three independent enforcement points, so the model cannot silently become top-k > 1:

1. configuration_top1moe.py: FORCED_TOP_K = 1; Top1MoeConfig.__setattr__ coerces num_experts_per_tok to 1 on every path (constructor, config.json, from_pretrained(..., num_experts_per_tok=8), later assignment), logging a warning.
2. modeling_top1moe.py: Top1MoeRouter.top_k = FORCED_TOP_K (never reads the config); the routing op is torch.topk(softmax(logits), 1).
3. Top1MoeForCausalLM.num_experts_per_tok = FORCED_TOP_K (used by the load-balancing auxiliary loss).

norm_topk_prob stays false as upstream: the single expert's output is scaled by its softmax probability. (Setting it to true, i.e. weight 1.0, was measured to be much worse: NLL 11.3 vs 4.5 on a short probe.)

Important: these weights were trained with top-8

The weights are byte-for-byte the upstream OLMoE-1B-7B-0924 checkpoint, only re-saved in the fused expert layout (experts.gate_up_proj of shape [64, 2048, 2048], experts.down_proj of shape [64, 2048, 1024]). Running them with one expert per token degrades them heavily:

probe (two short texts, bf16)	NLL
upstream OLMoE-1B-7B-0924, top-8	0.997
this model, top-1	4.498

Greedy generation is not coherent out of the box. Treat this checkpoint as a starting point for top-1 continued training / routing experiments, not as a ready-to-use 0.6B language model.

Validation

• Logits are bit-identical (max abs diff 0.0) to transformers' built-in OlmoeForCausalLM loaded from the upstream checkpoint with num_experts_per_tok=1.
• from_pretrained reports 0 missing / 0 unexpected / 0 mismatched keys.
• All 16 routers emit index tensors of shape (tokens, 1).

Files

file	content
configuration_top1moe.py	Top1MoeConfig, FORCED_TOP_K
modeling_top1moe.py	Top1MoeForCausalLM, Top1MoeModel, router / experts / attention modules
config.json	architecture hyper-parameters (OLMoE-1B-7B), num_experts_per_tok: 1, auto_map
model-0000*-of-00003.safetensors, model.safetensors.index.json	bf16 weights, fused expert layout
tokenizer.json, tokenizer_config.json, special_tokens_map.json	upstream OLMoE tokenizer
generation_config.json	upstream generation defaults

Provenance

• Weights: allenai/OLMoE-1B-7B-0924, commit 6d84c48581ece794365f2b8e9cfb043c68ade9c5, Apache-2.0.
• Code: derived from transformers==5.3.0 models/olmoe/{configuration,modeling}_olmoe.py (Apache-2.0), with the class renames and the top-1 enforcement described above. Attention, RoPE, norms, expert MLPs and the loss are unchanged.

中文说明

这是 OLMoE-1B-7B-0924 的原权重，把 top-k 硬编码成 top-1（每 token 只激活 1 个 expert，激活参数 0.58B）。 代码是自带的（model_type="top1moe"，类名 Top1Moe*），用 trust_remote_code=True 加载，需要 transformers>=5.0。权重是按 top-8 训的，直接 top-1 推理退化明显（NLL 4.5 对 1.0），适合当 top-1 续训或 路由实验的起点。强制 top-1 的三个点和验证结果见上文。