pyrite-stream / SPECIFICATION.md

Release Project Pyrite Phase 2: MoE-Stream Engine with Mixtral-8x22B & DeepSeek support

cab2ef1 verified about 13 hours ago

5.81 kB

# Project Pyrite Phase 2: MoE-Stream Engine 設計仕様書
## 〜 16GB Mac でフル精度 (BF16) MoE を動かす NVMe-RAM 分離ストリーミング 〜

**著者**: 中津川 昌弘 (Masahiro NAKATSUGAWA)  
**対象モデル**: `deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct` (15.7B / 2.36B active, 27 layers, 64 experts)

---

### 1. 理論的ブレイクスルーの数学的証明

従来の AirLLM 等の「全層オフロード」が遅かった理由は、Dense モデルでは毎トークンで全パラメータ（70GB）を読まねばならなかったからです。
しかし、**MoE アーキテクチャの本質は「スパース性（疎結合性）」** にあります。

DeepSeek-Coder-V2-Lite のテンソル構造を実機解析した結果、以下の数値が判明しました：

1. **常駐メモリ (Resident Working Set in RAM)**:
   - Embeddings (`embed_tokens`) + LM Head (`lm_head`)
   - 27層分の Attention 機構 (`self_attn`: MLA)
   - 27層分の LayerNorm
   - 27層分の Router Gate (`mlp.gate.weight`: [64, 2048])
   - 27層分の 共有エキスパート (`mlp.shared_experts`: 2 experts)
   👉 **常駐必須サイズ: わずか約 2.5 GiB** (16GB Mac のメモリに余裕で収まる！)

2. **動的ストリーミング (On-Demand Experts from NVMe SSD)**:
   - 各層の `switch_mlp` には 64 個のエキスパートが存在（Tensor shape: `[64, 1408, ...]`）。
   - 1 回のトークン生成でルーターが選ぶのは **Top-6 エキスパートのみ**（残り 58 個は完全に休眠）。
   - **1 エキスパートあたりのサイズ (BF16時)**: 約 15.8 MB
   - **1 レイヤーあたりの必要転送量**: $6 \times 15.8\text{ MB} = \mathbf{94.8\text{ MB}}$
   - **全 27 レイヤーを通じた必要転送量**: $26 \times 94.8\text{ MB} = \mathbf{2.46\text{ GB}}$

3. **転送レイテンシと速度予測**:
   - Mac の NVMe SSD 連続読み出し速度: **約 4,000 MB/s 〜 5,000 MB/s (4〜5 GB/s)**
   - 1 トークンあたりの SSD 転送時間: $2.46\text{ GB} \div 4.5\text{ GB/s} \approx \mathbf{0.54\text{ 秒}}$
   - **LRU キャッシュ（頻出エキスパートの RAM 保持）導入時**:
     プログラミングコード生成では同じエキスパートが約 40〜60% 再利用されるため、実効転送量は **約 1.0〜1.2 GB** まで半減。
     👉 **期待される実効生成速度: 2〜5 tokens/sec（フル精度 BF16、メモリ消費 5GB 以内）**
   - **8-bit / 4-bit 量子化エキスパート併用時**:
     転送量がさらに 1/2〜1/4 に激減し、**10〜20 tokens/sec（実用爆速レベル）** に到達可能！
---

### 2. システム・アーキテクチャ図

```
┌────────────────────────────────────────────────────────┐
│               NVMe SSD Storage (2TB)                   │
│  Layer 1〜26: Switch MLP Experts 0〜63 (BF16 / 31.4GB) │
└───────────────────────────┬────────────────────────────┘
                            │ On-Demand Streaming (Direct I/O / mmap)
                            ▼ (Top-6 Experts only: ~95MB/layer)
┌────────────────────────────────────────────────────────┐
│            RAM / Unified Memory (16GB Mac)             │
│                                                        │
│  [常駐ゾーン (~2.5 GiB)]                                │
│   ・Embedding & LM Head                                │
│   ・Multi-Head Latent Attention (MLA)                  │
│   ・LayerNorm & Router Gates                           │
│   ・Shared Experts (2 experts)                         │
│                                                        │
│  [動的エキスパート・キャッシュ枠 (~2.5 GiB)]           │
│   ・Active Top-6 Experts Buffer (Ring Buffer)          │
│   ・LRU Recent Experts Cache                           │
└───────────────────────────┬────────────────────────────┘
                            │ Zero-Copy
                            ▼
┌────────────────────────────────────────────────────────┐
│            Apple Silicon Metal 4 GPU Execution         │
└────────────────────────────────────────────────────────┘
```

---

### 3. 実装計画

1. **Phase 1: エキスパート・スプリッター（分割保存ツール）**
   `split_moe_experts.py`:
   31GB の巨大 safetensors を、「常駐ベース（Base）」と「各層・各エキスパートの独立ブロック（Expert Pack）」に分割・高速インデックス化。
2. **Phase 2: ストリーミング MoE レイヤー（MLX カスタムモジュール）**
   `StreamingMoE`:
   ルーターゲートの出力（Top-K）を監視し、ヒットしたエキスパートのみを SSD から `mmap` または事前割り当てバッファに非同期でロードして FFN 計算を行う。
3. **Phase 3: LRU キャッシュ & パイプライン・プレフェッチ**
   レイヤー $L$ の計算中に、レイヤー $L+1$ のルーターを先読みして SSD からプリフェッチし、転送待ち時間を計算時間で隠蔽（Overlap）。

	# Project Pyrite Phase 2: MoE-Stream Engine 設計仕様書
	## 〜 16GB Mac でフル精度 (BF16) MoE を動かす NVMe-RAM 分離ストリーミング〜

	著者: 中津川昌弘 (Masahiro NAKATSUGAWA)
	対象モデル: `deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct` (15.7B / 2.36B active, 27 layers, 64 experts)

	---

	### 1. 理論的ブレイクスルーの数学的証明

	従来の AirLLM 等の「全層オフロード」が遅かった理由は、Dense モデルでは毎トークンで全パラメータ（70GB）を読まねばならなかったからです。
	しかし、MoE アーキテクチャの本質は「スパース性（疎結合性）」にあります。

	DeepSeek-Coder-V2-Lite のテンソル構造を実機解析した結果、以下の数値が判明しました：

	1. 常駐メモリ (Resident Working Set in RAM):
	- Embeddings (`embed_tokens`) + LM Head (`lm_head`)
	- 27層分の Attention 機構 (`self_attn`: MLA)
	- 27層分の LayerNorm
	- 27層分の Router Gate (`mlp.gate.weight`: [64, 2048])
	- 27層分の共有エキスパート (`mlp.shared_experts`: 2 experts)
	👉 常駐必須サイズ: わずか約 2.5 GiB (16GB Mac のメモリに余裕で収まる！)

	2. 動的ストリーミング (On-Demand Experts from NVMe SSD):
	- 各層の `switch_mlp` には 64 個のエキスパートが存在（Tensor shape: `[64, 1408, ...]`）。
	- 1 回のトークン生成でルーターが選ぶのは Top-6 エキスパートのみ（残り 58 個は完全に休眠）。
	- 1 エキスパートあたりのサイズ (BF16時): 約 15.8 MB
	- 1 レイヤーあたりの必要転送量: $6 \times 15.8\text{ MB} = \mathbf{94.8\text{ MB}}$
	- 全 27 レイヤーを通じた必要転送量: $26 \times 94.8\text{ MB} = \mathbf{2.46\text{ GB}}$

	3. 転送レイテンシと速度予測:
	- Mac の NVMe SSD 連続読み出し速度: 約 4,000 MB/s 〜 5,000 MB/s (4〜5 GB/s)
	- 1 トークンあたりの SSD 転送時間: $2.46\text{ GB} \div 4.5\text{ GB/s} \approx \mathbf{0.54\text{ 秒}}$
	- LRU キャッシュ（頻出エキスパートの RAM 保持）導入時:
	プログラミングコード生成では同じエキスパートが約 40〜60% 再利用されるため、実効転送量は約 1.0〜1.2 GB まで半減。
	👉 期待される実効生成速度: 2〜5 tokens/sec（フル精度 BF16、メモリ消費 5GB 以内）
	- 8-bit / 4-bit 量子化エキスパート併用時:
	転送量がさらに 1/2〜1/4 に激減し、10〜20 tokens/sec（実用爆速レベル）に到達可能！

	---

	### 2. システム・アーキテクチャ図

	```
	┌────────────────────────────────────────────────────────┐
	│ NVMe SSD Storage (2TB) │
	│ Layer 1〜26: Switch MLP Experts 0〜63 (BF16 / 31.4GB) │
	└───────────────────────────┬────────────────────────────┘
	│ On-Demand Streaming (Direct I/O / mmap)
	▼ (Top-6 Experts only: ~95MB/layer)
	┌────────────────────────────────────────────────────────┐
	│ RAM / Unified Memory (16GB Mac) │
	│ │
	│ [常駐ゾーン (~2.5 GiB)] │
	│ ・Embedding & LM Head │
	│ ・Multi-Head Latent Attention (MLA) │
	│ ・LayerNorm & Router Gates │
	│ ・Shared Experts (2 experts) │
	│ │
	│ [動的エキスパート・キャッシュ枠 (~2.5 GiB)] │
	│ ・Active Top-6 Experts Buffer (Ring Buffer) │
	│ ・LRU Recent Experts Cache │
	└───────────────────────────┬────────────────────────────┘
	│ Zero-Copy
	▼
	┌────────────────────────────────────────────────────────┐
	│ Apple Silicon Metal 4 GPU Execution │
	└────────────────────────────────────────────────────────┘
	```

	---

	### 3. 実装計画

	1. Phase 1: エキスパート・スプリッター（分割保存ツール）
	`split_moe_experts.py`:
	31GB の巨大 safetensors を、「常駐ベース（Base）」と「各層・各エキスパートの独立ブロック（Expert Pack）」に分割・高速インデックス化。
	2. Phase 2: ストリーミング MoE レイヤー（MLX カスタムモジュール）
	`StreamingMoE`:
	ルーターゲートの出力（Top-K）を監視し、ヒットしたエキスパートのみを SSD から `mmap` または事前割り当てバッファに非同期でロードして FFN 計算を行う。
	3. Phase 3: LRU キャッシュ & パイプライン・プレフェッチ
	レイヤー $L$ の計算中に、レイヤー $L+1$ のルーターを先読みして SSD からプリフェッチし、転送待ち時間を計算時間で隠蔽（Overlap）。