# 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 実装**
```python
# 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 で取得確認済み:

```bash
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% |

---

### 現在のファイル構成

```text
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:

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

Build result:

```text
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 | ヒューリスティックスコアは高いが意味崩壊が多い |

採用値:

```text
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:

```bash
/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 は品質に効いている。