Release Project Pyrite Phase 2: MoE-Stream Engine with Mixtral-8x22B & DeepSeek support
cab2ef1 verified | #!/usr/bin/env python3 | |
| """Integration Test for StreamingMoELayer with Quantized Weights on Apple Silicon MLX.""" | |
| import time | |
| from pathlib import Path | |
| import mlx.core as mx | |
| from expert_registry import SafetensorsIndex | |
| from moe_stream_layer import StreamingMoELayer | |
| def run_test(): | |
| model_dir = Path("/Volumes/NVMe2TB/AI-LLM/LLMjp/Pyrite/pyrite-mlx/results/models/deepseek_coder_v2_lite_pyrite_final_mlx") | |
| shard1 = model_dir / "model-00001-of-00002.safetensors" | |
| print("=" * 70) | |
| print(" Streaming MoE Layer Integration Test (Layer 1)") | |
| print("=" * 70) | |
| idx = SafetensorsIndex(shard1) | |
| idx.open() | |
| # 1. Load resident gate & shared expert weights | |
| print("Loading resident weights (Gate + Shared Experts)...") | |
| gate_w = idx.get_tensor("model.layers.1.mlp.gate.weight") | |
| shared = { | |
| "gate_w": idx.get_tensor("model.layers.1.mlp.shared_experts.gate_proj.weight"), | |
| "gate_s": idx.get_tensor("model.layers.1.mlp.shared_experts.gate_proj.scales"), | |
| "gate_b": idx.get_tensor("model.layers.1.mlp.shared_experts.gate_proj.biases"), | |
| "up_w": idx.get_tensor("model.layers.1.mlp.shared_experts.up_proj.weight"), | |
| "up_s": idx.get_tensor("model.layers.1.mlp.shared_experts.up_proj.scales"), | |
| "up_b": idx.get_tensor("model.layers.1.mlp.shared_experts.up_proj.biases"), | |
| "down_w": idx.get_tensor("model.layers.1.mlp.shared_experts.down_proj.weight"), | |
| "down_s": idx.get_tensor("model.layers.1.mlp.shared_experts.down_proj.scales"), | |
| "down_b": idx.get_tensor("model.layers.1.mlp.shared_experts.down_proj.biases"), | |
| } | |
| print(f"Gate weight shape: {gate_w.shape}") | |
| print(f"Shared experts: gate_w={shared['gate_w'].shape}, down_w={shared['down_w'].shape}") | |
| # 2. Instantiate StreamingMoELayer with LRU cache (max 8 experts in RAM) | |
| layer = StreamingMoELayer( | |
| layer_idx=1, | |
| gate_weight=gate_w, | |
| shared_experts=shared, | |
| index_reader=idx, | |
| top_k=6, | |
| max_cache_experts=8, | |
| ) | |
| # 3. Simulate forward passes with synthetic activations | |
| hidden_dim = 2048 | |
| seq_len = 4 | |
| x = mx.random.normal((1, seq_len, hidden_dim)).astype(mx.bfloat16) | |
| print("\n--- Token Step 1 (Cold Start: Streaming from NVMe SSD) ---") | |
| t0 = time.perf_counter() | |
| out1 = layer(x) | |
| mx.eval(out1) | |
| elapsed_cold = (time.perf_counter() - t0) * 1000 | |
| print(f"Step 1 Output shape: {out1.shape}") | |
| print(f"Cold Start Time: {elapsed_cold:.2f} ms") | |
| print(f"Current Cached Experts in RAM: {list(layer.cache.cache.keys())}") | |
| print("\n--- Token Step 2 (Subsequent Token: Testing Cache Hits) ---") | |
| t0 = time.perf_counter() | |
| out2 = layer(x) | |
| mx.eval(out2) | |
| elapsed_warm = (time.perf_counter() - t0) * 1000 | |
| print(f"Warm Step Time: {elapsed_warm:.2f} ms") | |
| print(f"Current Cached Experts in RAM: {list(layer.cache.cache.keys())}") | |
| idx.close() | |
| print("\n" + "=" * 70) | |
| print(" Test Passed! Dynamic streaming & LRU caching verified on Apple Metal! ") | |
| print("=" * 70) | |
| if __name__ == "__main__": | |
| run_test() | |