#!/usr/bin/env python3 """Project Pyrite Phase 2: MoE-Stream Engine Runner. Runs DeepSeek-Coder-V2-Lite on Apple Silicon Metal using NVMe On-Demand Expert Streaming. Resident RAM usage is capped at ~2.5 GiB by offloading all 64 switch_mlp routed experts to NVMe SSD. """ import gc import json import os import sys import time from pathlib import Path # Add current dir to path sys.path.insert(0, str(Path(__file__).parent)) import mlx.core as mx import mlx.nn as nn from mlx_lm import stream_generate from mlx_lm.models.deepseek_v2 import Model, ModelArgs from mlx_lm.tokenizer_utils import TokenizerWrapper from transformers import AutoTokenizer from expert_registry import ModelSafetensorsIndex from moe_stream_layer import StreamingMoELayer def create_streaming_model(model_path: Path, max_cache_experts_per_layer: int = 8): """Instantiate and wire the MoE-Stream model architecture.""" model_path = Path(model_path) print("=" * 70) print(f" Initializing MoE-Stream Engine: {model_path.name}") print("=" * 70) # 1. Load Config with open(model_path / "config.json") as f: config = json.load(f) model_args = ModelArgs.from_dict(config) print(f"Model: {model_args.model_type} ({model_args.num_hidden_layers} layers, {config.get('n_routed_experts', 64)} experts)") # 2. Build Model Safetensors Registry for zero-copy streaming print("Building zero-copy NVMe safetensors registry...") registry = ModelSafetensorsIndex(model_path) registry.open() # 3. Load ONLY Resident Weights into RAM (skip switch_mlp) print("Loading resident weights into RAM (Filtering out switch_mlp experts)...") resident_weights = {} skipped_tensors = 0 for weight_name in registry.weight_map.keys(): if "switch_mlp" in weight_name: skipped_tensors += 1 continue resident_weights[weight_name] = registry.get_tensor(weight_name) print(f"Resident tensors loaded: {len(resident_weights):,} | Offloaded expert tensors: {skipped_tensors:,} (Zero RAM!)") # 4. Instantiate MLX Model (with lightweight MLP to avoid dummy 64-expert RAM allocation) print("Instantiating base model without dummy expert allocation...") saved_n_routed = model_args.n_routed_experts model_args.n_routed_experts = None # Temporarily disable MoE allocation model = Model(model_args) model_args.n_routed_experts = saved_n_routed # Restore # Quantize resident linear layers auto-detecting group_size and bits from safetensors print("Quantizing resident attention, shared_experts, and embeddings...") def auto_class_predicate(p, m): if "switch_mlp" in p: return False if not hasattr(m, "to_quantized"): return False w_key = f"{p}.weight" s_key = f"{p}.scales" if w_key in resident_weights and s_key in resident_weights: w = resident_weights[w_key] s = resident_weights[s_key] in_dim = m.weight.shape[1] if hasattr(m, "weight") else (m.input_dims if hasattr(m, "input_dims") else None) if in_dim is None: return False bits = (w.shape[-1] * 32) // in_dim group_size = in_dim // s.shape[-1] return {"bits": bits, "group_size": group_size, "mode": "affine"} return False nn.quantize(model, class_predicate=auto_class_predicate) # Load resident weights into base model print("Loading resident weights into base model...") model.load_weights(list(resident_weights.items()), strict=False) del resident_weights gc.collect() # 5. Intercept Layers 1 to 26 and inject StreamingMoELayer print(f"Injecting StreamingMoELayers with LRU cache ({max_cache_experts_per_layer} experts/layer)...") for layer_idx in range(1, model_args.num_hidden_layers): l_prefix = f"model.layers.{layer_idx}.mlp" # Resident Gate gate_w = registry.get_tensor(f"{l_prefix}.gate.weight") # Resident Shared Experts shared = { "gate_w": registry.get_tensor(f"{l_prefix}.shared_experts.gate_proj.weight"), "gate_s": registry.get_tensor(f"{l_prefix}.shared_experts.gate_proj.scales"), "gate_b": registry.get_tensor(f"{l_prefix}.shared_experts.gate_proj.biases"), "up_w": registry.get_tensor(f"{l_prefix}.shared_experts.up_proj.weight"), "up_s": registry.get_tensor(f"{l_prefix}.shared_experts.up_proj.scales"), "up_b": registry.get_tensor(f"{l_prefix}.shared_experts.up_proj.biases"), "down_w": registry.get_tensor(f"{l_prefix}.shared_experts.down_proj.weight"), "down_s": registry.get_tensor(f"{l_prefix}.shared_experts.down_proj.scales"), "down_b": registry.get_tensor(f"{l_prefix}.shared_experts.down_proj.biases"), } stream_moe = StreamingMoELayer( layer_idx=layer_idx, gate_weight=gate_w, shared_experts=shared, index_reader=registry, top_k=config.get("num_experts_per_tok", 6), routed_scaling_factor=config.get("routed_scaling_factor", 1.0), topk_method=config.get("topk_method", "greedy"), n_group=config.get("n_group", 1), topk_group=config.get("topk_group", 1), max_cache_experts=max_cache_experts_per_layer, ) model.layers[layer_idx].mlp = stream_moe print("MoE-Stream Engine successfully mounted on Apple Metal!") return model, registry def main(): model_dir = Path("/Volumes/NVMe2TB/AI-LLM/LLMjp/Pyrite/pyrite-mlx/results/models/deepseek_coder_v2_lite_pyrite_final_mlx") # 1. Load Tokenizer print("Loading tokenizer...") raw_tok = AutoTokenizer.from_pretrained(str(model_dir), trust_remote_code=True) tokenizer = TokenizerWrapper(raw_tok) # 2. Build MoE-Stream Model model, registry = create_streaming_model(model_dir, max_cache_experts_per_layer=8) # 3. Run Inference Generation prompt = "def fibonacci(n: int) -> int:\n \"\"\"Return the nth Fibonacci number with recursion.\"\"\"\n" print("\n" + "=" * 70) print(f" Prompt:\n{prompt.strip()}") print("=" * 70) print("Streaming generation started:\n") response_gen = stream_generate( model, tokenizer, prompt=prompt, max_tokens=64, ) last_resp = None for resp in response_gen: print(resp.text, end="", flush=True) last_resp = resp print("\n\n" + "=" * 70) if last_resp is not None: print(f" Prompt TPS: {last_resp.prompt_tps:.2f} tok/s ({last_resp.prompt_tokens} tokens)") print(f" Generation Speed: {last_resp.generation_tps:.2f} tok/s ({last_resp.generation_tokens} tokens)") print(f" Peak Memory: {last_resp.peak_memory:.2f} GB") print("=" * 70) registry.close() if __name__ == "__main__": main()