technical_notes.md

SSMoELM Technical Notes

2026-05-27

アーキテクチャ確定

モデル仕様

• 47.04M total params / 25.80M active per token
• d_model=768, L=6, n_heads=12 (kv_heads=3), head_dim=64
• MoE: 8 routed experts + 1 shared, top-2 routing, d_ff=256
• vocab=8192, context=2048, tie_word_embeddings=True

パラメータ内訳

Component	Params	Active
Embedding	6.29M	6.29M
Attention (全6層)	8.85M	8.85M
Shared Expert (全6層)	3.54M	3.54M
Routed Experts (top-2/8)	28.31M	7.08M
Norm 等	0.05M	0.05M
合計	47.04M	25.80M

ストレージ（sb3）

Component	Precision	Size
Embedding	4-bit	3.00 MB
Attention L1+L6	4-bit	1.41 MB
Attention L2-5	1-bit	0.70 MB
Shared Expert	4-bit	1.69 MB
Routed Experts	1-bit	3.38 MB
Router+Norm	bf16	0.09 MB
合計		10.27 MB binary / ~9.04 MB sb3

---

学習方針

QAT (BitNet b1.58 スタイル)

• master weights を bf16 で保持
• forward pass で layer-wise 精度に量子化（STE）
• bf16 勾配で master weights を更新
• 推論時に master weights を破棄

STE 実装

# 1-bit: binary * row_scale, gradient は w をそのまま通す
return w + mx.stop_gradient(quantized - w)

選定理由: Bonsai (PTQ) は 8B モデルの冗長性を活用する手法。 47M params / Chinchilla 付近では冗長性が不足し PTQ 品質が壊滅的になるリスクがある。 Bonsai スケール別の性能低下: 8B→-11%, 4B→-19%, 1.7B→-26%（スケール縮小で加速）。

---

訓練環境

• M2 Mac MLX (mlx==0.31.2)
• Python 3.12
• uv でパッケージ管理

---

データセット

Pretraining: 1B tokens（当初 4B 予定→ M2 Mac 16日超のため 1B に変更、約 4.2日）

• FineWeb-Edu sample-10BT: 60% (2.4B tokens)
• FineWeb sample-10BT: 40% (1.6B tokens)
• 保存形式: jsonl.gz, 100M tokens/shard × 40 shards, 6.3 GB
• シャッフル: 訓練時に shard 単位でランダム順 + shard 内シャッフル（方式 B）

Tokenizer 訓練データ

• 全 40 shard から均一サンプリング (25M tokens/shard × 40 = 1B tokens)
• 目的: 時代的偏りを排除

---

Tokenizer

仕様

• BPE, byte-level, vocab=8192
• ライブラリ: HuggingFace tokenizers
• 保存先: tokenizer/tokenizer.json

特殊トークン

Token	ID	用途
<bos>	0	系列開始
<eos>	1	会話全体の終了
<pad>	2	バッチパディング
<|system|>	3	system ターン開始
<|user|>	4	user ターン開始
<|assistant|>	5	assistant ターン開始
<|eot|>	6	ターン終了 (end of turn)

Chat template

<bos><|system|>
{system}<|eot|>
<|user|>
{user}<|eot|>
<|assistant|>
{response}<|eot|><eos>

効率測定 (Wikipedia EN + UltraChat 各5M chars)

Tokenizer	chars/tok (wiki)	chars/tok (chat)	tok/word (wiki)
SSMoELM (8,192)	3.667	3.925	1.738
GPT-2 (50,257)	4.570	4.670	1.394

• GPT-2 比 +19〜25% 多くトークンを消費（vocab 縮小の代償）
• 会話テキストの方が効率が良い（訓練データに会話を含めた効果）
• context 2048 tokens ≈ 会話で ~1300 words（許容範囲）

---

Ablation 計画（スキップ・2026-05-27）

当初 A/B/C の 3 run ablation を予定していたが、下記理由でスキップして Run C 構成で直接本番学習へ。

• Q/K 4-bit は安全側の設定（1-bit Q/K より精度劣化リスクが低い）
• MoE > Dense は文献上ほぼ確実（routing overhead が Dense より有利）
• ablation 3 run × 50M = 150M tokens ≈ 7.5時間 の節約

採用構成: attn-4bit-qk（L2-5 の Q/K を 4-bit、V/O は 1-bit のまま）

ストレージ影響: binary +1.06 MB → sb3 推定 ~9.97 MB（10 MB ギリギリ）

本番学習設定（2026-05-27 開始）

run          = attn-4bit-qk
tokens       = 1B
batch_size   = 2,  grad_accum = 16
eff_batch    = 65,536 tokens/step
total_steps  ≈ 15,259
warmup_frac  = 5% ≈ 763 steps
lr_max = 3e-4 → lr_min = 1e-5 (cosine)
router_lr    = 1e-4
optimizer    = AdamW (β=0.9/0.95, wd=0.1, clip=1.0)
save_every   = 1526 steps ≈ 100M tokens
速度見込み   ≈ 2800 tok/s → 約 4.2 日

---

ファイル構成

SSMoELM/
  specification.md          仕様書
  technical_notes.md        ← 本ファイル
  pyproject.toml
  data/
    prepare_dataset.py      4B tokens 収集（完了）
    make_tokenizer_sample.py 均一サンプリング（完了）
    raw/                    shard_000〜039.jsonl.gz (6.3 GB)
    tokenizer_sample/       sample.jsonl.gz (1.4 GB)
  tokenizer/
    train_tokenizer.py      訓練スクリプト（完了）
    eval_tokenizer.py       効率測定（完了）
    tokenizer.json          訓練済み tokenizer
  train/
    config.py               ModelConfig / DenseConfig
    quantize.py             QAT (1-bit, 4-bit, STE)
    model.py                SSMoELM / Dense (MLX)
    data.py                 (未実装)
    train.py                (未実装)
  tasks/
    ablation.md             Ablation 計画

---

未解決の設計リスク（Ablation で検証予定）

1. 1-bit Attention (L2-5): Q/K が precision-sensitive かどうか未検証
2. MoE vs Dense at 47M: routing overhead がペイするか未検証
3. d_ff=256 (0.33×d_model): 通常の 2.67×〜4× より大幅に小さい

---

---

2026-06-03

Pretraining 完了

• 採用チェックポイント: step 013734（900M tokens）— 800M・900M・1B比較で900Mが最良
• 学習スループット: ~3,000 tok/s（Flash Attention + QAT cache）
• 実際のパックサイズ: 12.1 MB（理論 11.49 MB、safetensors形式）

ベンチマーク結果（0-shot, 500 samples）

Task	Metric	Score	Random
HellaSwag	acc_norm	33.4%	25%
LAMBADA	acc	13.8%	N/A
PIQA	acc_norm	53.2%	50%
WinoGrande	acc	49.6%	50%
ARC-Easy	acc_norm	35.0%	25%
ARC-Challenge	acc_norm	21.0%	25%
BoolQ	acc	36.2%	50%
MMLU	acc avg (57 tasks, ≤500/task)	23.4%	25%

---

推論コード（PyTorch）

• ファイル: inference_pytorch.py（スタンドアロン、MLX依存なし）
• dtype: fp32（CPU/MPS）— MPS の bf16 は PyTorch で未対応
• メモリ: ディスク 12.1 MB → ランタイム ~188 MB（fp32展開）
• 正確な推論には packed safetensors が必須（raw bf16 master weights は QAT適用前の値のため不正確、LAMBADA ppl: 606 vs 6366）
• PyTorch/MLX スコア差: ±1%以内（サンプリング誤差範囲）

---

MoEルーティング統計（136トークン）

Layer	CV	特徴
0	0.32	E4/E5偏り
1	0.15	最も均一
2	0.22	E1偏り
3	0.42	最も不均一、E6突出（22.8%）
4	0.25	E5偏り
5	0.25	E0偏り

エキスパート崩壊なし。aux loss + z-loss が効果的。

---

SFT計画

• データセット: databricks/databricks-dolly-15k（CC BY-SA 3.0、商用利用可）
• 選定理由: 商用OK、人手アノテーション15K件、短い入出力、英語のみ
• 除外候補:
  • LIMA: CC BY-NC-SA（非商用のみ）
  • Alpaca: OpenAI ToS制約あり

---

TurboWarp 速度推定

Barynsor-1 実測 (21M dense, 8 tok/s) から逆算: ~151M MACs/sec

SSMoELM active compute per token:

• Attention: ~402M MACs/token
• Shared Expert: ~9.4M MACs/token
• Routed (top-2): ~18.9M MACs/token
• 推定 ~7.8 tok/s（Barynsor-1 の 151M MACs/sec 基準）

---

2026-06-08

公開モデル

• Base: hf.co/brulee-1/SSMoELM-Base
• Instruct: hf.co/brulee-1/SSMoELM-it

hf CLI で取得確認済み:

hf models info brulee-1/SSMoELM-Base
hf models info brulee-1/SSMoELM-it
hf download brulee-1/SSMoELM-Base README.md --local-dir /tmp/ssmoelm-hf/base
hf download brulee-1/SSMoELM-it README.md --local-dir /tmp/ssmoelm-hf/it

公開先のモデルカードは checkpoints/README.md（Base）と checkpoints/README_it.md（Instruct）を元に管理する。 HuggingFace 上では packed safetensors の uint8 配列を直接読むため、約12M params / 8-bit のように表示される場合があるが、実際は 47.04M params の 1-bit + 4-bit mixed precision MoE。

---

SFT 完了

• Base checkpoint: checkpoints/attn-4bit-qk_step013734.safetensors
• Final SFT checkpoint: checkpoints/sft_final.safetensors
• Packed Instruct weights:
  • checkpoints/ssmoelm_it_packed.safetensors (~12 MB)
  • checkpoints/ssmoelm_it_packed.npz (~11 MB)
• SFT log: checkpoints/sft2_log.csv
• 最終ログ行: step 177050, loss 1.592862, 2636.6 tok/s

SFT 実装は train/sft.py。assistant tokens only loss を使い、user/prompt 部分と特殊トークンの prefix は loss mask で無視する。

SFT データセット

Dataset	用途	備考
databricks/databricks-dolly-15k	instruction following	CC BY-SA 3.0
OpenAssistant/oasst1 EN	dialogue / assistant response	Apache 2.0
trivia_qa	QA factuality	CLI 既定に含む
natural_questions	short-answer QA	streaming
Open-Orca/FLAN	instruction mixture	streaming, default limit 70K

checkpoints/README_it.md の公開モデルカードでは Dolly-15k + oasst1 EN を主要 SFT データとして記載している。追加 QA/FLAN 系は train/sft.py の CLI 既定データセットに含まれるため、再現性メモとして本ファイルに残す。

---

Instruct 評価結果

SSMoELM-it は SFT 後も base の知識ベンチを概ね維持し、BoolQ / ARC-Challenge / PIQA が改善。

Task	Base	Instruct	Δ
HellaSwag acc_norm	33.4%	33.2%	-0.2%
LAMBADA acc	13.8%	14.8%	+1.0%
PIQA acc_norm	53.2%	55.4%	+2.2%
WinoGrande acc	49.6%	49.6%	0.0%
ARC-Easy acc_norm	35.0%	35.2%	+0.2%
ARC-Challenge acc_norm	21.0%	24.0%	+3.0%
BoolQ acc	36.2%	44.4%	+8.2%
MMLU avg	23.4%	23.2%	-0.2%

---

現在のファイル構成

SSMoELM/
  specification.md              最終仕様
  technical_notes.md            実装・実験ログ
  pyproject.toml                uv project dependencies
  data/
    prepare_dataset.py          FineWeb/FineWeb-Edu shard 作成
    make_tokenizer_sample.py    tokenizer 訓練サンプル作成
    raw/                        shard_000〜039.jsonl.gz（git ignore）
    tokenizer_sample/           sample.jsonl.gz（git ignore）
  tokenizer/
    train_tokenizer.py          tokenizer 訓練
    eval_tokenizer.py           tokenizer 効率測定
    tokenizer.json              訓練済み tokenizer
  train/
    config.py                   ModelConfig / DenseConfig
    quantize.py                 QAT 1-bit/4-bit + STE
    model.py                    SSMoELM / Dense MLX model
    data.py                     pretraining data loader
    train.py                    pretraining script
    lm_eval_adapter.py          lm-eval adapter
    sft.py                      supervised fine-tuning script
  inference.py                  packed PyTorch inference / chat entrypoint
  inference_packed.py           packed inference implementation
  inference_pytorch.py          PyTorch reference inference
  chat.py                       interactive chat helper
  export_packed.py              pack/export weights
  eval_benchmark.py             MLX evaluation
  eval_benchmark_pt.py          PyTorch evaluation
  eval_mmlu.py                  MMLU evaluation
  analyze_routing.py            MoE routing analysis
  checkpoints/
    README.md                   Base model card
    README_it.md                Instruct model card
    ssmoelm_it_packed.safetensors  Instruct packed weights（git ignore）
    ssmoelm_it_packed.npz          Instruct packed weights（git ignore）

---

TurboWarp / goboscript MVP

goboscript で開ける .sb3 artifact を生成開始。

• 生成スクリプト: scripts/generate_turbowarp_demo.py
• goboscript source: turbowarp/ssmoelm_mvp/main.gs, stage.gs
• 生成済み sb3: turbowarp/ssmoelm_mvp/ssmoelm_mvp.sb3
• 検証テスト: tests/test_goboscript_export.py

現時点の sb3 は UI/input loop + deterministic response shell + quantized dot-product kernel + packed uint8 decoder を含む。

実装済み Scratch kernels

Kernel	goboscript proc	入力	出力
1-bit dot	dot1_kernel	activation list, bit list, scale	dot1_result
4-bit dot	dot4_kernel	activation list, signed nibble list, scale	dot4_result
packed 1-bit decode	get_packed1_bit	uint8 list, 0-based bit index	packed1_bit
packed signed int4 decode	get_packed4_signed	uint8 list, 0-based nibble index	packed4_value
packed 1-bit dot	dot1_packed_kernel	activation list, packed uint8 bit list, scale	dot1_packed_result
packed 4-bit dot	dot4_packed_kernel	activation list, packed uint8 nibble list, scale	dot4_packed_result
self test	kernel_self_test, packed_decode_self_test, packed_kernel_self_test	toy vectors / bytes	kernel_status, decode_status, packed_kernel_status

Self-test toy values:

• 1-bit dot: raw dot -1.5, scale 0.25 → -0.375
• 4-bit dot: raw dot 9, scale 0.125 → 1.125
• packed 1-bit decode: byte 13 (0b00001101, LSB-first) → bits 1,0,1,1
• packed int4 decode: bytes 195 (0xC3) and 47 (0x2F), low-nibble-first → signed nibbles 3,-4,-1,2
• packed 1-bit / 4-bit dot は decoder 経由で pre-expanded dot と同じ toy result（-0.375, 1.125）を返す。

Build command:

/Users/scratchbrulee/.cargo/bin/goboscript build turbowarp/ssmoelm_mvp

Build result:

Emitted 361 blocks
ssmoelm_mvp.sb3 size: 7348 bytes
main sprite: 292 blocks
lists: packed1_bytes, k1_x, k1_w_bits, packed4_bytes, k4_w_nibbles, k4_x

Next step is 1-token projection verification: compare a tiny Scratch packed matmul/projection fixture against the Python reference, then extend it toward attention and MoE expert rows.

---

Decoding sweep（2026-06-08）

SSMoELM-it の生成設定を小規模 prompt sweep と手動確認で比較。

temperature	top_k	repetition_penalty	傾向
0.0	1	1.0	最も文法的だが同じ句を繰り返しやすい
0.0	1	1.15	反復は減るが bullet/list 系で引用符反復が残る
0.0	1	1.3	低品質語の混入を抑えつつ反復も比較的少ないため採用
0.55	20	1.15	多様だが固有名詞・断片語の混入が増える
0.7	30	1.15	ヒューリスティックスコアは高いが意味崩壊が多い

採用値:

temperature = 0.0
top_k = 1
top_p = 0.9
repetition_penalty = 1.3

理由: 47M / 8k vocab の小型モデルでは sampling の多様性より greedy + 強め repetition penalty の方が破綻しにくい。top_p は temperature=0.0 では実質未使用だが、sampling に切り替える場合の互換既定として 0.9 を維持。

---

Shared expert only ablation（2026-06-08）

PyTorch lm-eval adapter に --shared-only を追加し、routed experts を無効化して shared expert だけで短い 0-shot ベンチを実行。

Command:

/etc/profiles/per-user/scratchbrulee/bin/uv run python eval_benchmark_pt.py \
  --ckpt checkpoints/ssmoelm_it_packed.safetensors \
  --tasks hellaswag,lambada_openai \
  --limit 100

/etc/profiles/per-user/scratchbrulee/bin/uv run python eval_benchmark_pt.py \
  --ckpt checkpoints/ssmoelm_it_packed.safetensors \
  --tasks hellaswag,lambada_openai \
  --limit 100 \
  --shared-only

Mode	HellaSwag acc_norm	HellaSwag acc	LAMBADA acc	LAMBADA ppl
Full MoE	40.0%	35.0%	11.0%	732.20
Shared only	30.0%	31.0%	7.0%	2742.00

結論: shared expert だけでも HellaSwag は random を少し超えるが、LAMBADA perplexity が大きく悪化。routed experts は品質に効いている。