pyrite-stream / moe_stream_runner.py

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

cab2ef1 verified about 20 hours ago

6.97 kB

#!/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()

	#!/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()