add: sparse_moe architecture source

architecture/sparse_moe/__init__.py

@@ -0,0 +1,113 @@
+from .config import (
+    ExpertConfig,
+    RouterConfig,
+    TrainingConfig,
+    DataConfig,
+    BenchmarkConfig,
+    ProjectConfig,
+    load_preset,
+)
+from .experts import LoRAAdapter, TinyExpert, create_experts
+from .routing import LinearRouter, RoutingOutput
+from .layers import SparseMoELayer
+from .injection import inject_moe_layers
+from .trainer import MoETrainer
+from .evaluation import (
+    evaluate_model,
+    evaluate_per_domain,
+    parameter_summary,
+    print_domain_report,
+    print_report,
+)
+from .analysis import compute_specialization_matrix, print_specialization_report
+from .datasets import (
+    DictDataset,
+    load_domain_texts,
+    split_and_build_loaders,
+    build_mixed_loader,
+    tokenize_and_batch,
+)
+from .utils import build_labels, set_seed, to_serializable, parse_seed_list
+from .baselines import (
+    prepare_dense_lora_model,
+    train_dense_lora_model,
+    train_dense_lora_baseline,
+)
+from .reporting import (
+    save_dashboard,
+    generate_qualitative_samples,
+    build_before_after_rows,
+    print_before_after_report,
+    run_multi_seed,
+)
+from .visualization import (
+    ColorPalette,
+    BasePlot,
+    DomainScorePlot,
+    TrainingConvergencePlot,
+    ExpertRoutingHeatmap,
+)
+
+__all__ = [
+    # Config
+    "ExpertConfig",
+    "RouterConfig",
+    "TrainingConfig",
+    "DataConfig",
+    "BenchmarkConfig",
+
+    "ProjectConfig",
+    "load_preset",
+    # Experts
+    "LoRAAdapter",
+    "TinyExpert",
+    "create_experts",
+    # Router
+    "LinearRouter",
+    "RoutingOutput",
+    # MoE layer
+    "SparseMoELayer",
+    # Injection
+    "inject_moe_layers",
+    # Training
+    "MoETrainer",
+    # Evaluation
+    "evaluate_model",
+    "evaluate_per_domain",
+
+    "parameter_summary",
+    "print_domain_report",
+    "print_report",
+    # Analysis
+    "compute_specialization_matrix",
+    "print_specialization_report",
+    # Data
+    "DictDataset",
+    "load_domain_texts",
+    "split_and_build_loaders",
+    "build_mixed_loader",
+    "tokenize_and_batch",
+    # Utils
+    "build_labels",
+    "set_seed",
+    "to_serializable",
+    "parse_seed_list",
+    # Dense baseline
+    "prepare_dense_lora_model",
+    "train_dense_lora_model",
+    "train_dense_lora_baseline",
+    # Reporting
+    "save_dashboard",
+    "generate_qualitative_samples",
+    "build_before_after_rows",
+    "print_before_after_report",
+    "run_multi_seed",
+    # Visualization
+    "ColorPalette",
+    "BasePlot",
+    "DomainScorePlot",
+    "TrainingConvergencePlot",
+    "ExpertRoutingHeatmap",
+]
+
+__version__ = "0.1.0"

architecture/sparse_moe/analysis.py

@@ -0,0 +1,151 @@
+from __future__ import annotations
+
+from typing import Dict, Optional
+
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+
+from .evaluation import _named_moe_layers
+from .utils import build_labels
+
+
+# Core: expert–domain affinity matrix
+@torch.no_grad()
+def compute_specialization_matrix(model: nn.Module, domain_loaders: Dict[str, DataLoader], device: Optional[torch.device] = None) -> Dict[str, object]:
+    device = device or next(model.parameters()).device
+    was_training = model.training
+    model.eval()
+
+    try:
+        moe_layers = _named_moe_layers(model)
+        if not moe_layers:
+            raise ValueError("No SparseMoELayer found in model")
+        domains = list(domain_loaders.keys())
+        num_domains = len(domains)
+
+        # Accumulate weighted expert fractions per domain, keeping layers separate.
+        layer_affinity = {
+            layer_name: torch.zeros(layer.num_experts, num_domains)
+            for layer_name, layer in moe_layers
+        }
+        layer_weight_per_domain = {
+            layer_name: torch.zeros(num_domains) for layer_name, _ in moe_layers
+        }
+
+        for d_idx, domain in enumerate(domains):
+            loader = domain_loaders[domain]
+            for batch in loader:
+                batch = {k: v.to(device) for k, v in batch.items()}
+                labels = build_labels(batch)
+                _ = model(**batch, labels=labels)
+
+                for layer_name, module in moe_layers:
+                    stats = getattr(module, "_last_stats", None)
+                    if stats is None:
+                        continue
+                    fracs = stats.get("expert_fractions", [])
+                    if len(fracs) != module.num_experts:
+                        continue
+
+                    num_tokens = float(stats.get("num_tokens", 1))
+                    frac_t = torch.tensor(fracs, dtype=torch.float32)
+
+                    layer_affinity[layer_name][:, d_idx] += frac_t * num_tokens
+                    layer_weight_per_domain[layer_name][d_idx] += num_tokens
+
+        # Normalise and flatten rows as unique layer-expert slots.
+        affinity_rows = []
+        expert_labels = []
+        for layer_name, module in moe_layers:
+            layer_matrix = layer_affinity[layer_name]
+            for d_idx in range(num_domains):
+                weight = layer_weight_per_domain[layer_name][d_idx]
+                if weight > 0:
+                    layer_matrix[:, d_idx] /= weight
+            affinity_rows.append(layer_matrix)
+            for expert_idx in range(module.num_experts):
+                expert_labels.append(f"{layer_name}:E{expert_idx}")
+
+        affinity = torch.cat(affinity_rows, dim=0) if affinity_rows else torch.zeros(0, num_domains)
+
+        # Specialisation
+        row_max = affinity.max(dim=1).values
+        row_mean = affinity.mean(dim=1).clamp(min=1e-9)
+        specialisation = row_max / row_mean
+
+        # JS divergence
+        col_sums = affinity.sum(dim=0, keepdim=True).clamp(min=1e-9)
+        col_norm = affinity / col_sums
+        divergence = _pairwise_js(col_norm)
+
+        return {
+            "affinity": affinity,
+            "layer_affinity": layer_affinity,
+            "domains": domains,
+            "expert_labels": expert_labels,
+            "specialisation": specialisation,
+            "divergence": divergence,
+        }
+    finally:
+        model.train(was_training)
+
+
+# Helpers
+def _kl_divergence(P: torch.Tensor, q: torch.Tensor) -> torch.Tensor:
+    P_safe = P.clamp(min=1e-9)
+    q_safe = q.clamp(min=1e-9)
+    return (P_safe * (P_safe / q_safe).log()).sum()
+
+
+def _pairwise_js(col_norm: torch.Tensor) -> torch.Tensor:
+    D = col_norm.shape[1]
+    js = torch.zeros(D, D)
+
+    for i in range(D):
+        for j in range(i + 1, D):
+            p = col_norm[:, i]
+            q = col_norm[:, j]
+            m = 0.5 * (p + q)
+            jsd = 0.5 * _kl_divergence(p, m) + 0.5 * _kl_divergence(q, m)
+            js[i, j] = js[j, i] = jsd.item()
+
+    return js
+
+
+# Pretty-print
+def print_specialization_report(result: Dict[str, object]):
+    affinity = result["affinity"]
+    domains = result["domains"]
+    expert_labels = result.get("expert_labels", [f"E{i}" for i in range(affinity.shape[0])])
+    spec = result["specialisation"]
+    div = result["divergence"]
+
+    E, D = affinity.shape
+
+    print("\n" + "=" * 62)
+    print("  Expert Specialisation Report")
+    print("=" * 62)
+
+    # Affinity matrix
+    header = "  Layer-Expert                 " + "".join(f"  {d:>8s}" for d in domains) + "  Spec"
+    print(header)
+    print("  " + "-" * (len(header) - 2))
+    for e in range(E):
+        row = f"  {expert_labels[e]:<28s}"
+        for d in range(D):
+            row += f"  {affinity[e, d]:>8.3f}"
+        row += f"  {spec[e]:>5.2f}"
+        print(row)
+
+    # Divergence
+    print(f"\n  JS Divergence (domain routing distances):")
+    div_header = "          " + "".join(f"  {d:>8s}" for d in domains)
+    print(div_header)
+    for i, d_name in enumerate(domains):
+        row = f"  {d_name:<8s}"
+        for j in range(D):
+            row += f"  {div[i, j]:>8.4f}"
+        print(row)
+
+    print("=" * 62 + "\n")

architecture/sparse_moe/baselines.py

@@ -0,0 +1,290 @@
+from __future__ import annotations
+
+import copy
+import math
+from typing import Dict, List, Optional
+
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+from transformers import AutoModelForCausalLM
+
+from .config import ExpertConfig, RouterConfig, TrainingConfig
+from .evaluation import evaluate_per_domain
+from .experts import LoRAAdapter
+from .injection import _is_ffn
+from .trainer import _cosine_with_warmup, _build_amp
+from .utils import build_labels
+
+
+
+# LoRA injection (dense — no routing, no MoE)
+def _inject_dense_lora(model: nn.Module, expert_config: ExpertConfig, router_config: RouterConfig) -> nn.Module:
+    model_type = getattr(model.config, "model_type", "qwen2")
+
+    # Scale rank to match MoE
+    scaled_rank = expert_config.lora_rank * router_config.num_experts
+
+    # Freeze all base params
+    for p in model.parameters():
+        p.requires_grad = False
+
+    replaced = 0
+    for name, module in list(model.named_modules()):
+        if not _is_ffn(module, model_type):
+            continue
+
+        # Get input dimension (d_model) for the LoRA adapter
+        d_model = getattr(model.config, "hidden_size", None)
+        if d_model is None:
+            # Fallback heuristic if config lacks hidden_size
+            for child in module.children():
+                if hasattr(child, "in_features"):
+                    d_model = child.in_features
+                    break
+                elif hasattr(child, "weight"):
+                    d_model = child.weight.size(-1)
+                    break
+            
+        if d_model is None:
+            print(f"[WARNING] Could not determine d_model for {name}, skipping.")
+            continue
+
+        # Create a single LoRA adapter with scaled rank
+        mod_config = copy.copy(expert_config)
+        mod_config.lora_rank = min(scaled_rank, 64, d_model // 2)
+        mod_config.lora_alpha = expert_config.lora_alpha * (mod_config.lora_rank / expert_config.lora_rank)
+        
+        # Resolve the precise device and dtype of the layer being replaced
+        try:
+            layer_device = next(module.parameters()).device
+            layer_dtype = next(module.parameters()).dtype
+        except StopIteration:
+            layer_device = next(model.parameters()).device
+            layer_dtype = next(model.parameters()).dtype
+
+        lora = LoRAAdapter(
+            d_model=d_model,
+            config=mod_config
+        ).to(device=layer_device, dtype=layer_dtype)
+
+        # Wrap FFN + LoRA
+        wrapper = _DenseLoRAWrapper(module, lora)
+
+        # Replace in parent
+        parts = name.split(".")
+        parent = model
+        for part in parts[:-1]:
+            parent = getattr(parent, part)
+        setattr(parent, parts[-1], wrapper)
+        replaced += 1
+
+    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    if replaced > 0:
+        print(
+            f"Dense LoRA: injected {replaced} adapters, "
+            f"rank={mod_config.lora_rank}, "
+            f"trainable={trainable:,} params"
+        )
+    else:
+        raise ValueError(
+            "Dense LoRA: no FFN layers found; model was not modified. "
+            "Check architecture support."
+        )
+    return model
+
+
+class _DenseLoRAWrapper(nn.Module):
+
+    def __init__(self, base_ffn: nn.Module, lora: LoRAAdapter):
+        super().__init__()
+        self.base_ffn = base_ffn
+        self.lora = lora
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # saves memory
+        with torch.no_grad():
+            base_out = self.base_ffn(x)
+        lora_out = self.lora(x)
+        return base_out + lora_out.to(base_out.dtype)
+
+
+# Training loop (simplified — no routing, no aux loss)
+def _train_dense(model: nn.Module, train_loader: DataLoader, val_loader: Optional[DataLoader], config: TrainingConfig, device: torch.device) -> Dict[str, List[float]]:
+    model.train()
+
+    if device.type == "cuda":
+        torch.set_float32_matmul_precision("high")
+        torch.backends.cudnn.benchmark = True
+
+    trainable = [p for p in model.parameters() if p.requires_grad]
+    if not trainable:
+        raise ValueError(
+            "Dense LoRA baseline has no trainable parameters. "
+            "Check FFN detection for this architecture."
+        )
+    dense_lr = config.lr / 2.0
+    optimizer = torch.optim.AdamW(trainable, lr=dense_lr, weight_decay=0.01)
+
+    steps_per_epoch = max(
+        math.ceil(len(train_loader) / config.gradient_accumulation_steps),
+        1,
+    )
+    total_steps = steps_per_epoch * config.num_epochs
+    scheduler = torch.optim.lr_scheduler.LambdaLR(
+        optimizer,
+        lr_lambda=lambda step: _cosine_with_warmup(
+            step, config.warmup_steps, total_steps
+        ),
+    )
+
+    use_amp, amp_dtype, scaler = _build_amp(device, config.fp16)
+
+    history = {"train_loss": [], "val_loss": [], "lr": [], "epoch_boundaries": []}
+    accum = config.gradient_accumulation_steps
+    global_step = 0
+    best_val_loss = float("inf")
+    patience = 1
+    patience_counter = 0
+    best_state = None
+
+    for epoch in range(1, config.num_epochs + 1):
+        print(f"  [Dense LoRA] epoch {epoch}/{config.num_epochs}")
+        history["epoch_boundaries"].append(global_step)
+        optimizer.zero_grad(set_to_none=True)
+
+        epoch_loss = 0.0
+        epoch_steps = 0
+
+        for step, batch in enumerate(train_loader):
+            batch = {k: v.to(device, non_blocking=True) for k, v in batch.items()}
+            labels = build_labels(batch)
+
+            with torch.amp.autocast(device.type, dtype=amp_dtype, enabled=use_amp):
+                outputs = model(**batch, labels=labels)
+                loss = outputs.loss / accum
+
+            scaler.scale(loss).backward()
+
+            if (step + 1) % accum == 0 or (step + 1) == len(train_loader):
+                scaler.unscale_(optimizer)
+                nn.utils.clip_grad_norm_(trainable, config.max_grad_norm)
+                scaler.step(optimizer)
+                scaler.update()
+                scheduler.step()
+                optimizer.zero_grad(set_to_none=True)
+                global_step += 1
+
+                raw_loss = loss.item() * accum
+                current_lr = scheduler.get_last_lr()[0]
+                history["train_loss"].append(raw_loss)
+                history["lr"].append(current_lr)
+                epoch_loss += raw_loss
+                epoch_steps += 1
+
+                if global_step % config.log_every_steps == 0:
+                    print(
+                        f"  [Dense LoRA] step {global_step:>5d} | "
+                        f"loss={raw_loss:.4f} | "
+                        f"lr={current_lr:.2e} | "
+                        f"avg={epoch_loss / max(epoch_steps, 1):.4f}"
+                    )
+
+        if val_loader is None:
+            continue
+
+        val_loss = _evaluate_dense(model, val_loader, device, use_amp, amp_dtype)
+        history["val_loss"].append(val_loss)
+        print(f"  [Dense LoRA] val_loss={val_loss:.4f}")
+
+        if val_loss < best_val_loss:
+            best_val_loss = val_loss
+            patience_counter = 0
+            best_state = {
+                name: param.detach().clone()
+                for name, param in model.named_parameters()
+                if param.requires_grad
+            }
+        else:
+            patience_counter += 1
+            if patience_counter >= patience:
+                print("  [Dense LoRA] early stopping triggered; restoring best weights.")
+                if best_state is not None:
+                    for name, param in model.named_parameters():
+                        if name in best_state:
+                            param.data.copy_(best_state[name])
+                break
+
+    print(
+        f"Dense LoRA training complete ({global_step} steps across up to {config.num_epochs} epochs)."
+    )
+    return history
+
+
+@torch.no_grad()
+def _evaluate_dense(model: nn.Module, val_loader: DataLoader, device: torch.device, use_amp: bool, amp_dtype: torch.dtype) -> float:
+    model.eval()
+    total_loss, total_tokens = 0.0, 0
+    for batch in val_loader:
+        batch = {k: v.to(device, non_blocking=True) for k, v in batch.items()}
+        labels = build_labels(batch)
+        with torch.amp.autocast(device.type, dtype=amp_dtype, enabled=use_amp):
+            outputs = model(**batch, labels=labels)
+        valid_tokens = int((labels[:, 1:] != -100).sum().item())
+        total_loss += outputs.loss.item() * max(valid_tokens, 1)
+        total_tokens += valid_tokens
+    model.train()
+    return total_loss / max(total_tokens, 1)
+
+
+# Public API
+def prepare_dense_lora_model(base_model_name: str, expert_config: ExpertConfig, router_config: RouterConfig, device: torch.device) -> nn.Module:
+    model = AutoModelForCausalLM.from_pretrained(base_model_name).to(device)
+    model = _inject_dense_lora(model, expert_config, router_config)
+
+    try:
+        model.gradient_checkpointing_enable(
+            gradient_checkpointing_kwargs={"use_reentrant": False}
+        )
+    except TypeError:
+        model.gradient_checkpointing_enable()
+
+    return model
+
+
+def train_dense_lora_model(model: nn.Module, train_loader: DataLoader, val_loader: Optional[DataLoader], training_config: TrainingConfig, device: torch.device) -> Dict[str, List[float]]:
+    return _train_dense(model, train_loader, val_loader, training_config, device)
+
+
+def train_dense_lora_baseline(base_model_name: str, train_loader: DataLoader, eval_loaders: Dict[str, DataLoader], training_config: TrainingConfig, expert_config: ExpertConfig, router_config: RouterConfig, device: torch.device, val_loader: Optional[DataLoader] = None) -> Dict[str, Dict]:
+    print("═══ Dense LoRA Baseline ═══")
+
+    # Load fresh model (don't contaminate the MoE model)
+    model = prepare_dense_lora_model(
+        base_model_name,
+        expert_config,
+        router_config,
+        device,
+    )
+
+    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+    dense_config = copy.deepcopy(training_config)
+
+    # Train with the same epoch schedule as the main MoE path.
+    history = train_dense_lora_model(model, train_loader, val_loader, dense_config, device)
+
+    # Evaluate per-domain
+    print("Evaluating Dense LoRA per-domain...")
+    per_domain = evaluate_per_domain(model, eval_loaders, device)
+
+    # Cleanup
+    del model
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+    return {
+        "per_domain": per_domain,
+        "history": history,
+        "trainable_params": trainable,
+    }

architecture/sparse_moe/config.py

@@ -0,0 +1,332 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field, fields
+from typing import Any, Dict, List, Optional
+
+import torch
+
+
+# Expert configuration
+@dataclass
+class ExpertConfig:
+
+    expert_type: str = "lora"
+
+    # ── LoRA settings ──
+    lora_rank: int = 16
+    lora_alpha: float = 32.0
+    lora_dropout: float = 0.05
+
+    # ── TinyExpert settings ──
+    tiny_intermediate: Optional[int] = None
+    tiny_activation: str = "silu"
+    tiny_dropout: float = 0.1
+
+    def __post_init__(self):
+        if self.expert_type not in ("lora", "tiny"):
+            raise ValueError(
+                f"expert_type must be 'lora' or 'tiny', got {self.expert_type!r}"
+            )
+        if self.lora_rank <= 0:
+            raise ValueError(f"lora_rank must be positive, got {self.lora_rank}")
+        if self.lora_alpha <= 0:
+            raise ValueError(f"lora_alpha must be positive, got {self.lora_alpha}")
+        if not (0 <= self.lora_dropout < 1):
+            raise ValueError(f"lora_dropout must be in [0, 1), got {self.lora_dropout}")
+        if self.tiny_intermediate is not None and self.tiny_intermediate <= 0:
+            raise ValueError(
+                f"tiny_intermediate must be positive, got {self.tiny_intermediate}"
+            )
+        if self.tiny_activation not in ("gelu", "silu", "relu"):
+            raise ValueError(
+                f"tiny_activation must be gelu/silu/relu, got {self.tiny_activation!r}"
+            )
+        if not (0 <= self.tiny_dropout < 1):
+            raise ValueError(
+                f"tiny_dropout must be in [0, 1), got {self.tiny_dropout}"
+            )
+
+
+# Router configuration
+@dataclass
+class RouterConfig:
+
+    num_experts: int = 8
+    top_k: int = 2
+    noise_std: float = 0.1
+
+    z_loss_coeff: float = 1e-4
+
+    loss_type: str = "switch"
+
+    entropy_reg: float = 0.01
+
+    def __post_init__(self):
+        if self.num_experts <= 0:
+            raise ValueError(f"num_experts must be positive, got {self.num_experts}")
+        if self.top_k <= 0:
+            raise ValueError(f"top_k must be positive, got {self.top_k}")
+        if self.top_k > self.num_experts:
+            raise ValueError(
+                f"top_k ({self.top_k}) cannot exceed num_experts ({self.num_experts})"
+            )
+        if self.noise_std < 0:
+            raise ValueError(f"noise_std must be >= 0, got {self.noise_std}")
+        if self.loss_type not in ("switch", "gshard"):
+            raise ValueError(
+                f"loss_type must be 'switch' or 'gshard', got {self.loss_type!r}"
+            )
+
+
+# Training configuration
+@dataclass
+class TrainingConfig:
+
+    num_epochs: int = 8
+    lr: float = 2e-4
+    batch_size: int = 8
+    gradient_accumulation_steps: int = 4
+
+    max_seq_len: int = 512
+    max_grad_norm: float = 1.0
+    warmup_steps: int = 100
+    fp16: bool = True
+
+    entropy_reg_start: float = 0.01
+
+    aux_loss_weight: float = 0.01
+
+    output_dir: str = "checkpoints"
+    save_every_steps: int = 500
+    log_every_steps: int = 20
+
+
+# Data configuration
+@dataclass
+class DataConfig:
+
+    samples_per_domain: int = 3000
+    eval_frac: float = 0.15
+    domains: List[str] = field(default_factory=lambda: ["code", "prose", "qa", "math"])
+
+
+# Benchmark configuration
+@dataclass
+class BenchmarkConfig:
+
+    num_warmup_runs: int = 3
+    num_timed_runs: int = 10
+    max_gen_tokens: int = 80
+    report_dir: str = "benchmark/report"
+
+
+# Top-level project configuration
+@dataclass
+class ProjectConfig:
+
+    base_model: str = "Qwen/Qwen2.5-3B"
+    seed: int = 42
+    device: str = "auto"
+
+    expert: ExpertConfig = field(default_factory=ExpertConfig)
+    routing: RouterConfig = field(default_factory=RouterConfig)
+    training: TrainingConfig = field(default_factory=TrainingConfig)
+    data: DataConfig = field(default_factory=DataConfig)
+    benchmark: BenchmarkConfig = field(default_factory=BenchmarkConfig)
+
+    # ── Demo / experiment flags ──
+    artifact_dir: str = "artifacts"
+    run_dense_baseline: bool = True
+    run_multi_seed: bool = False
+    multi_seed_list: str = "42,43,44"
+    multi_seed_epochs: int = 2
+
+    # ── Pipeline / Resume flags ──
+    resume: bool = False
+    stage: str = "stage_03_moe_training"
+    hub_repo_id: Optional[str] = None
+
+    def __post_init__(self):
+            self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        # Disable FP16 on CPU
+        if self.device == "cpu":
+            self.training.fp16 = False
+
+    @property
+    def resolved_device(self) -> torch.device:
+        return torch.device(self.device)
+
+    # YAML loading
+    @classmethod
+    def from_yaml(cls, path: str) -> "ProjectConfig":
+        try:
+            import yaml
+        except ImportError:
+            raise ImportError("PyYAML required for YAML config loading: pip install pyyaml")
+
+        with open(path, "r") as f:
+            raw = yaml.safe_load(f) or {}
+
+        return cls._from_dict(raw)
+
+    @classmethod
+    def _from_dict(cls, d: Dict[str, Any]) -> "ProjectConfig":
+        sub_configs = {
+            "expert": ExpertConfig,
+            "routing": RouterConfig,
+            "training": TrainingConfig,
+            "data": DataConfig,
+            "benchmark": BenchmarkConfig,
+        }
+
+        kwargs: Dict[str, Any] = {}
+        for key, value in d.items():
+            if key in sub_configs and isinstance(value, dict):
+                # Build nested dataclass with type coercion
+                kwargs[key] = _coerce_dataclass(sub_configs[key], value)
+            else:
+                kwargs[key] = value
+
+        return _coerce_dataclass(cls, kwargs)
+
+    def to_dict(self) -> Dict[str, Any]:
+        return _asdict_recursive(self)
+
+
+# Type coercion helper
+def _coerce_dataclass(cls, kwargs: Dict[str, Any]):
+    import logging
+    _log = logging.getLogger(__name__)
+
+    def _coerce_bool(value: Any) -> bool:
+        if isinstance(value, str):
+            lowered = value.strip().lower()
+            if lowered in {"1", "true", "yes", "on"}:
+                return True
+            if lowered in {"0", "false", "no", "off"}:
+                return False
+            raise ValueError(f"Cannot parse boolean value from {value!r}")
+        return bool(value)
+
+    _str_to_type = {"int": int, "float": float, "bool": _coerce_bool, "str": str}
+    field_types = {f.name: f.type for f in fields(cls)}
+    coerced = {}
+    for key, value in kwargs.items():
+        # Filter out unknown keys with a warning (catches YAML typos)
+        if key not in field_types:
+            _log.warning(
+                f"Ignoring unknown config key {key!r} for {cls.__name__}. "
+                f"Valid keys: {list(field_types.keys())}"
+            )
+            continue
+
+        type_str = field_types[key]
+        cast_fn = _str_to_type.get(type_str)
+
+        # Handle Optional[T] annotations (rendered as 'Optional[T]' strings)
+        if cast_fn is None and isinstance(type_str, str) and "Optional" in type_str:
+            for base_name, base_fn in _str_to_type.items():
+                if base_name in type_str:
+                    cast_fn = base_fn
+                    break
+
+        if cast_fn is not None and not isinstance(value, type(None)):
+            try:
+                value = cast_fn(value)
+            except (ValueError, TypeError) as e:
+                _log.warning(
+                    f"Could not coerce {key!r}={value!r} to {type_str}: {e}"
+                )
+        coerced[key] = value
+
+    return cls(**coerced)
+
+
+# Scale presets
+PRESETS: Dict[str, Dict[str, Any]] = {
+    # 'test' is maintained for the unit test suite gate.
+    "test": {
+        "base_model": "Qwen/Qwen2.5-3B",
+        "expert": {"lora_rank": 16, "lora_alpha": 32.0, "lora_dropout": 0.1},
+        "routing": {"num_experts": 16, "top_k": 2},
+        "training": {
+            "num_epochs": 1,
+            "lr": 5e-5,    
+            "batch_size": 2,
+            "max_seq_len": 512,
+            "gradient_accumulation_steps": 1,
+            "fp16": True,
+            "warmup_steps": 5,
+            "save_every_steps": 0,
+            "log_every_steps": 1,
+            "aux_loss_weight": 0.05,
+        },
+        "data": {"samples_per_domain": 50, "eval_frac": 0.1},
+        "benchmark": {"num_warmup_runs": 1, "num_timed_runs": 1, "max_gen_tokens": 10},
+        "run_dense_baseline": False,
+        "run_multi_seed": False,
+    },
+
+    "production": {
+        "base_model": "Qwen/Qwen2.5-3B",
+        "expert": {"lora_rank": 16, "lora_alpha": 32.0, "lora_dropout": 0.1},
+        "routing": {"num_experts": 16, "top_k": 2},
+        "training": {
+            "num_epochs": 4,
+            "lr": 5e-5,    # Conservative for 3B — see RESULTS_OVER_SESSIONS.txt
+            "batch_size": 24,
+            "max_seq_len": 4096,
+            "gradient_accumulation_steps": 2,
+            "fp16": True,
+            "warmup_steps": 350,
+            "save_every_steps": 500,
+            "log_every_steps": 25,
+            "aux_loss_weight": 0.05,
+        },
+        "data": {"samples_per_domain": 15000, "eval_frac": 0.1},
+        "benchmark": {"num_timed_runs": 20, "max_gen_tokens": 100},
+        "run_dense_baseline": True,
+        "run_multi_seed": True,
+        "multi_seed_epochs": 3,
+    },
+    "cloud_validation": {
+        "base_model": "Qwen/Qwen2.5-3B",
+        "expert": {"lora_rank": 16, "lora_alpha": 32.0, "lora_dropout": 0.1},
+        "routing": {"num_experts": 16, "top_k": 2},
+        "training": {
+            "num_epochs": 1,
+            "lr": 5e-5,    
+            "batch_size": 8,           
+            "max_seq_len": 2048,       
+            "gradient_accumulation_steps": 2, 
+            "fp16": True,
+            "warmup_steps": 50,
+            "save_every_steps": 100,
+            "log_every_steps": 5,
+            "aux_loss_weight": 0.05,
+        },
+        "data": {"samples_per_domain": 500, "eval_frac": 0.1},
+        "benchmark": {"num_timed_runs": 10, "max_gen_tokens": 80},
+        "run_dense_baseline": True,
+        "run_multi_seed": False,
+    },
+}
+
+
+def load_preset(name: str) -> ProjectConfig:
+    if name not in PRESETS:
+        raise ValueError(f"Unknown preset {name!r}, available: {list(PRESETS.keys())}")
+    return ProjectConfig._from_dict(PRESETS[name])
+
+
+# Helpers
+def _asdict_recursive(obj: Any) -> Any:
+    if hasattr(obj, "__dataclass_fields__"):
+        return {k: _asdict_recursive(v) for k, v in obj.__dict__.items()}
+    elif isinstance(obj, (list, tuple)):
+        return [_asdict_recursive(v) for v in obj]
+    elif isinstance(obj, dict):
+        return {k: _asdict_recursive(v) for k, v in obj.items()}
+    elif isinstance(obj, torch.device):
+        return str(obj)
+    return obj

architecture/sparse_moe/datasets.py

@@ -0,0 +1,228 @@
+from __future__ import annotations
+import sys
+import random
+from typing import Dict, List, Optional, Tuple
+
+import torch
+from torch.utils.data import DataLoader, Dataset
+
+
+
+# Workers: 0 on Windows (no fork), 4 on Linux/macOS
+_DATALOADER_WORKERS = 0 if sys.platform == "win32" else 4
+
+
+# Dataset class
+class DictDataset(Dataset):
+
+    def __init__(self, input_ids: torch.Tensor, attention_mask: torch.Tensor):
+        self.input_ids = input_ids          # (N, seq_len)
+        self.attention_mask = attention_mask  # (N, seq_len)
+
+    def __len__(self) -> int:
+        return self.input_ids.shape[0]
+
+    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
+        return {
+            "input_ids": self.input_ids[idx],          # (seq_len,)
+            "attention_mask": self.attention_mask[idx],  # (seq_len,)
+        }
+
+
+# Tokenisation helpers
+def tokenize_and_batch(texts: List[str], tokenizer, max_len: int, batch_size: int, shuffle: bool = True) -> DataLoader:
+    if not texts:
+        raise ValueError("tokenize_and_batch requires at least one text sample")
+
+    enc = tokenizer(
+        texts,
+        truncation=True,
+        padding="max_length",
+        max_length=max_len,
+        return_tensors="pt",
+    )
+    dataset = DictDataset(enc["input_ids"], enc["attention_mask"])
+    return DataLoader(dataset, batch_size=batch_size, shuffle=shuffle,
+                       num_workers=_DATALOADER_WORKERS, pin_memory=True)
+
+
+# Domain text loading
+def load_domain_texts(samples_per_domain: int = 3000, domains: Optional[List[str]] = None) -> Dict[str, List[str]]:
+    domains = domains or ["code", "prose", "qa", "math"]
+    n = samples_per_domain
+    result: Dict[str, List[str]] = {}
+
+    for domain in domains:
+        try:
+            if domain == "code":
+                result["code"] = _load_code(n)
+            elif domain == "prose":
+                result["prose"] = _load_prose(n)
+            elif domain == "qa":
+                result["qa"] = _load_qa(n)
+            elif domain == "math":
+                result["math"] = _load_math(n)
+            else:
+                print(f"[WARNING] Unknown domain {domain!r}, skipping")
+                continue
+            print(f"  {domain}: loaded {len(result[domain])} samples")
+        except Exception as e:
+            print(f"[WARNING] Failed to load {domain}: {e}. Using synthetic fallback.")
+            result[domain] = _synthetic_fallback(domain, n)
+
+    return result
+
+
+def _load_code(n: int) -> List[str]:
+    from datasets import load_dataset
+    try:
+        ds = load_dataset("code_search_net", "python", split="train")
+        texts = [ex["whole_func_string"] for ex in ds if ex.get("whole_func_string") and len(ex["whole_func_string"].strip()) > 50]
+        if len(texts) >= n:
+            return texts[:n]
+    except Exception as e:
+        print(f"[WARNING] CodeSearchNet failed: {e}, trying MBPP fallback")
+    # Fallback to MBPP (only ~970 samples)
+    try:
+        ds = load_dataset("mbpp", split="train+test+validation")
+        texts = [f"{ex['text']}\n{ex['code']}" for ex in ds if ex.get("code")]
+        if 0 < len(texts) < n:
+            print(f"[WARNING] Code domain: only {len(texts)} samples available (requested {n}). Duplicating.")
+            texts = (texts * (n // len(texts) + 1))[:n]
+        return texts[:n]
+    except Exception as e:
+        print(f"[WARNING] MBPP fallback failed: {e}. Using synthetic fallback.")
+        return _synthetic_fallback("code", n)
+
+
+def _load_prose(n: int) -> List[str]:
+    from datasets import load_dataset
+    ds = load_dataset("wikitext", "wikitext-2-raw-v1", split="train")
+    texts = [t for t in ds["text"] if len(t.strip()) > 50]
+    if len(texts) < n:
+        import warnings
+        warnings.warn(
+            f"WikiText-2 has only {len(texts)} usable paragraphs but {n} were "
+            f"requested. Duplicating to fill the pool. This may overfit prose.",
+            UserWarning, stacklevel=2,
+        )
+        # Cycle through available texts until we reach n
+        texts = (texts * (n // len(texts) + 1))[:n]
+    return texts[:n]
+
+
+def _load_qa(n: int) -> List[str]:
+    from datasets import load_dataset
+    ds = load_dataset("squad", split="train")
+    texts = [
+        f"Q: {x['question']}\nA: {x['answers']['text'][0]}"
+        for x in ds
+        if x.get("answers") and x["answers"].get("text")
+    ]
+    return texts[:n]
+
+
+def _load_math(n: int) -> List[str]:
+    from datasets import load_dataset
+    ds = load_dataset("gsm8k", "main", split="train")
+    texts = [f"{x['question']}\n{x['answer']}" for x in ds]
+    if len(texts) < n:
+        import warnings
+        warnings.warn(
+            f"GSM8K has only {len(texts)} problems but {n} were requested. "
+            f"Duplicating to fill the pool.",
+            UserWarning, stacklevel=2,
+        )
+        texts = (texts * (n // len(texts) + 1))[:n]
+    return texts[:n]
+
+
+def _synthetic_fallback(domain: str, n: int) -> List[str]:
+    templates = {
+        "code": [
+            "def fibonacci(n):\n    if n <= 1: return n\n    return fibonacci(n-1) + fibonacci(n-2)",
+            "def factorial(n):\n    result = 1\n    for i in range(2, n+1):\n        result *= i\n    return result",
+            "class Stack:\n    def __init__(self):\n        self.items = []\n    def push(self, x):\n        self.items.append(x)",
+        ],
+        "prose": [
+            "The history of science spans many centuries of intellectual endeavour.",
+            "Climate patterns have shifted dramatically over the past century.",
+            "Ancient civilisations developed complex mathematical systems.",
+        ],
+        "qa": [
+            "Question: What is machine learning?\nAnswer: Machine learning is a subset of AI.",
+            "Question: How does the internet work?\nAnswer: Through a network of interconnected servers.",
+            "Question: What causes earthquakes?\nAnswer: Tectonic plate movement.",
+        ],
+        "math": [
+            "Problem: If x + 3 = 7, solve for x.\nSolution: x = 4.",
+            "Problem: Calculate 15% of 200.\nSolution: 30.",
+            "Problem: Find the area of a rectangle 5×3.\nSolution: 15 square units.",
+        ],
+    }
+    pool = templates.get(domain, templates["prose"])
+    return [pool[i % len(pool)] for i in range(n)]
+
+
+# Train/eval split
+def split_and_build_loaders(domain_texts: Dict[str, List[str]], tokenizer, max_seq_len: int, batch_size: int, eval_frac: float = 0.15, seed: int = 42) -> Tuple[Dict[str, DataLoader], Dict[str, DataLoader]]:
+    rng = random.Random(seed)
+    train_loaders: Dict[str, DataLoader] = {}
+    eval_loaders: Dict[str, DataLoader] = {}
+
+    for domain, texts in domain_texts.items():
+        shuffled = list(texts)
+        if not shuffled:
+            raise ValueError(f"Domain {domain!r} has no samples to split")
+        rng.shuffle(shuffled)
+
+        if len(shuffled) == 1:
+            print("[WARNING] " + str(
+                f"Domain {domain!r} has only one sample; duplicating it into train and eval."
+            ))
+            train_texts = shuffled
+            eval_texts = shuffled
+        else:
+            split_idx = int(len(shuffled) * (1 - eval_frac))
+            split_idx = min(max(1, split_idx), len(shuffled) - 1)
+            train_texts = shuffled[:split_idx]
+            eval_texts = shuffled[split_idx:]
+
+        train_loaders[domain] = tokenize_and_batch(
+            train_texts, tokenizer, max_seq_len, batch_size, shuffle=True
+        )
+        eval_loaders[domain] = tokenize_and_batch(
+            eval_texts, tokenizer, max_seq_len, batch_size, shuffle=False
+        )
+
+        print(
+            f"  {domain}: {len(train_texts)} train / {len(eval_texts)} eval"
+        )
+
+    return train_loaders, eval_loaders
+
+
+def build_mixed_loader(domain_loaders: Dict[str, DataLoader], batch_size: int, seed: Optional[int] = None, shuffle: bool = True) -> DataLoader:
+    if not domain_loaders:
+        raise ValueError("build_mixed_loader requires at least one domain loader")
+
+    all_ids, all_masks = [], []
+    for loader in domain_loaders.values():
+        ds = loader.dataset
+        all_ids.append(ds.input_ids)
+        all_masks.append(ds.attention_mask)
+
+    combined_ids = torch.cat(all_ids, dim=0)      # (N_total, seq_len)
+    combined_masks = torch.cat(all_masks, dim=0)   # (N_total, seq_len)
+
+    dataset = DictDataset(combined_ids, combined_masks)
+
+    generator = torch.Generator()
+    if seed is not None:
+        generator.manual_seed(seed)
+
+    return DataLoader(
+        dataset, batch_size=batch_size, shuffle=shuffle, 
+        generator=generator if shuffle else None,
+        num_workers=_DATALOADER_WORKERS, pin_memory=True
+    )

architecture/sparse_moe/evaluation.py

@@ -0,0 +1,241 @@
+from __future__ import annotations
+
+import math
+from typing import Dict, Optional
+
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+
+from .layers import SparseMoELayer
+from .utils import build_labels
+
+
+def _named_moe_layers(model: nn.Module) -> list[tuple[str, SparseMoELayer]]:
+    layers = []
+    for idx, module in enumerate(model.modules()):
+        if isinstance(module, SparseMoELayer):
+            layer_name = getattr(module, "layer_name", f"moe_layer_{idx}")
+            layers.append((layer_name, module))
+    return layers
+
+
+
+# Core evaluation
+def evaluate_model(model: nn.Module, dataloader: DataLoader, device: Optional[torch.device] = None) -> Dict[str, object]:
+    device = device or next(model.parameters()).device
+    was_training = model.training
+    model.eval()
+    moe_layers = _named_moe_layers(model)
+
+    try:
+        total_loss, total_tokens, n_batches = 0.0, 0, 0
+        total_correct, total_top5, total_mrr = 0.0, 0.0, 0.0
+        entropy_weighted_sum, entropy_weight = 0.0, 0.0
+        layer_entropy_sum = {name: 0.0 for name, _ in moe_layers}
+        layer_entropy_weight = {name: 0.0 for name, _ in moe_layers}
+        layer_weighted_counts = {
+            name: torch.zeros(layer.num_experts, dtype=torch.float32)
+            for name, layer in moe_layers
+        }
+        layer_weight_totals = {name: 0.0 for name, _ in moe_layers}
+
+        for batch in dataloader:
+            batch = {k: v.to(device) for k, v in batch.items()}
+            labels = build_labels(batch)
+            outputs = model(**batch, labels=labels)
+
+            # Reconstruct sum using same count
+            shifted_labels = labels[:, 1:]
+            predictable_tokens = int((shifted_labels != -100).sum().item())
+
+            # Language-modelling loss
+            total_loss += outputs.loss.item() * max(predictable_tokens, 1)
+            total_tokens += predictable_tokens
+            n_batches += 1
+
+            # Accuracy metrics
+            if predictable_tokens > 0 and hasattr(outputs, "logits"):
+                logits = outputs.logits  # (B, S, V)
+                valid_mask = shifted_labels != -100
+
+                # Top-1 accuracy
+                preds = logits[:, :-1].argmax(dim=-1)
+                correct = (preds == shifted_labels) & valid_mask
+                total_correct += correct.sum().item()
+
+                # Top-5 accuracy
+                top5 = logits[:, :-1].topk(5, dim=-1).indices
+                hits = (top5 == shifted_labels.unsqueeze(-1)).any(dim=-1) & valid_mask
+                total_top5 += hits.sum().item()
+
+                # MRR
+                valid_logits = logits[:, :-1][valid_mask]  # (N_valid, V)
+                valid_labels_flat = shifted_labels[valid_mask]   # (N_valid,)
+                correct_logits = valid_logits.gather(1, valid_labels_flat.unsqueeze(-1))  # (N, 1)
+                
+                n_valid = valid_logits.size(0)
+                ranks = torch.zeros(n_valid, device=device, dtype=torch.long)
+                chunk_size = 2048
+                for i in range(0, n_valid, chunk_size):
+                    end = i + chunk_size
+                    ranks[i:end] = (valid_logits[i:end] >= correct_logits[i:end]).sum(dim=-1)
+                
+                total_mrr += (1.0 / ranks.float()).sum().item()
+
+            # Collect routing stats from MoE layers
+            for layer_name, module in moe_layers:
+                stats = getattr(module, "_last_stats", None)
+                if stats is None:
+                    continue
+
+                token_weight = float(stats.get("num_tokens", predictable_tokens))
+                entropy_value = float(stats.get("routing_entropy", 0.0))
+                layer_entropy_sum[layer_name] += entropy_value * token_weight
+                layer_entropy_weight[layer_name] += token_weight
+                entropy_weighted_sum += entropy_value * token_weight
+                entropy_weight += token_weight
+
+                fracs = stats.get("expert_fractions", [])
+                if len(fracs) != module.num_experts:
+                    continue
+                frac_t = torch.tensor(fracs, dtype=torch.float32)
+                assign_weight = float(stats.get("num_assignments", token_weight))
+                layer_weighted_counts[layer_name] += frac_t * assign_weight
+                layer_weight_totals[layer_name] += assign_weight
+
+        if n_batches == 0:
+            raise ValueError("Evaluation dataloader yielded no batches")
+        if total_tokens == 0:
+            raise ValueError("Evaluation dataloader contained no predictable tokens")
+
+        avg_loss = total_loss / max(total_tokens, 1)
+        ppl = math.exp(min(avg_loss, 100))  # cap to avoid overflow
+
+        # Normalise
+        utilisation = []
+        expert_labels = []
+        layer_routing_entropy: Dict[str, float] = {}
+        layer_utilisation: Dict[str, list[float]] = {}
+        layer_utilisation_std: Dict[str, float] = {}
+        util_stds = []
+
+        for layer_name, module in moe_layers:
+            weight_total = layer_weight_totals[layer_name]
+            if layer_entropy_weight[layer_name] > 0:
+                layer_routing_entropy[layer_name] = (
+                    layer_entropy_sum[layer_name] / layer_entropy_weight[layer_name]
+                )
+            else:
+                layer_routing_entropy[layer_name] = 0.0
+
+            if weight_total > 0:
+                layer_values = (layer_weighted_counts[layer_name] / weight_total).tolist()
+            else:
+                layer_values = [0.0] * module.num_experts
+            layer_utilisation[layer_name] = layer_values
+
+            std = torch.tensor(layer_values).std().item() if len(layer_values) > 1 else 0.0
+            layer_utilisation_std[layer_name] = std
+            util_stds.append(std)
+
+            for expert_idx, value in enumerate(layer_values):
+                expert_labels.append(f"{layer_name}:E{expert_idx}")
+                utilisation.append(value)
+
+        util_std = sum(util_stds) / len(util_stds) if util_stds else 0.0
+
+        return {
+            "perplexity": ppl,
+            "avg_loss": avg_loss,
+            "bits_per_token": avg_loss / math.log(2),  # nats -> bits
+            "token_acc": total_correct / max(total_tokens, 1),
+            "top5_acc": total_top5 / max(total_tokens, 1),
+            "mrr": total_mrr / max(total_tokens, 1),
+            "routing_entropy": entropy_weighted_sum / max(entropy_weight, 1.0),
+            "layer_routing_entropy": layer_routing_entropy,
+            "expert_utilisation": utilisation,
+            "expert_labels": expert_labels,
+            "layer_expert_utilisation": layer_utilisation,
+            "layer_utilisation_std": layer_utilisation_std,
+            "utilisation_std": util_std,
+            "num_batches": n_batches,
+        }
+    finally:
+        model.train(was_training)
+
+
+# Per-domain evaluation
+def evaluate_per_domain(model: nn.Module, domain_loaders: Dict[str, DataLoader], device: Optional[torch.device] = None) -> Dict[str, Dict[str, object]]:
+    results = {}
+    for domain, loader in domain_loaders.items():
+        results[domain] = evaluate_model(model, loader, device)
+    return results
+
+
+# Parameter summary
+def parameter_summary(model: nn.Module) -> Dict[str, int]:
+    total = sum(p.numel() for p in model.parameters())
+    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    frozen = total - trainable
+
+    router_params, expert_params = 0, 0
+    for name, p in model.named_parameters():
+        if not p.requires_grad:
+            continue
+        if "routing" in name or ".gate." in name:
+            router_params += p.numel()
+        else:
+            expert_params += p.numel()
+
+    return {
+        "total": total,
+        "trainable": trainable,
+        "frozen": frozen,
+        "router": router_params,
+        "experts": expert_params,
+        "trainable_pct": 100.0 * trainable / max(total, 1),
+    }
+
+
+# Pretty-print helpers
+def print_domain_report(domain_results: Dict[str, Dict[str, float]]):
+    print("\n" + "=" * 62)
+    print("  Per-Domain Evaluation")
+    print("=" * 62)
+    header = f"  {'Domain':<15s} {'PPL':>8s} {'Loss':>8s} {'Top-1 Acc':>10s} {'Top-5 Acc':>10s}"
+    print(header)
+    print("  " + "-" * (len(header) - 2))
+    for domain, m in domain_results.items():
+        top1 = f"{m.get('token_acc', 0.0) * 100:.1f}%"
+        top5 = f"{m.get('top5_acc', 0.0) * 100:.1f}%"
+        print(f"  {domain:<15s} {m['perplexity']:>8.2f} {m['avg_loss']:>8.4f} "
+              f"{top1:>10s} {top5:>10s}")
+    print("=" * 62 + "\n")
+
+
+def print_report(metrics: Dict[str, float], params: Dict[str, int]):
+    print("\n" + "=" * 55)
+    print("  Sparse MoE — Evaluation Report")
+    print("=" * 55)
+
+    print(f"\n  Perplexity          : {metrics['perplexity']:.2f}")
+    print(f"  Avg loss            : {metrics['avg_loss']:.4f}")
+    print(f"  Routing entropy     : {metrics['routing_entropy']:.4f}")
+    print(f"  Utilisation std     : {metrics['utilisation_std']:.4f}")
+
+    util = metrics.get("expert_utilisation", [])
+    labels = metrics.get("expert_labels", [])
+    if util and labels:
+        ranked = sorted(zip(labels, util), key=lambda item: item[1], reverse=True)
+        print(f"  Top layer-experts   :")
+        for label, value in ranked[: min(10, len(ranked))]:
+            print(f"    {label:<28s} {value:.3f}")
+
+    print(f"\n  Parameters:")
+    print(f"    Total             : {params['total']:>12,}")
+    print(f"    Trainable         : {params['trainable']:>12,}  ({params['trainable_pct']:.2f}%)")
+    print(f"      Router          : {params['router']:>12,}")
+    print(f"      Experts         : {params['experts']:>12,}")
+    print(f"    Frozen            : {params['frozen']:>12,}")
+    print("=" * 55 + "\n")

architecture/sparse_moe/experts.py

@@ -0,0 +1,118 @@
+from __future__ import annotations
+
+import math
+from typing import Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .config import ExpertConfig
+
+
+
+# LoRA Adapter Expert
+class LoRAAdapter(nn.Module):
+
+    def __init__(self, d_model: int, config: ExpertConfig):
+        super().__init__()
+        r = config.lora_rank
+        self.scaling = config.lora_alpha / r
+
+        # A: down-projection  (d_model → rank)
+        # B: up-projection    (rank → d_model)
+        self.lora_A = nn.Parameter(torch.empty(r, d_model))
+        self.lora_B = nn.Parameter(torch.empty(d_model, r))
+        self.dropout = nn.Dropout(config.lora_dropout)
+
+        nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
+        nn.init.normal_(self.lora_B, std=0.01)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        out = F.linear(self.dropout(x), self.lora_A)
+        out = F.linear(out, self.lora_B)
+        return out * self.scaling
+
+    def extra_repr(self) -> str:
+        r = self.lora_A.shape[0]
+        d = self.lora_A.shape[1]
+        return f"d_model={d}, rank={r}, scaling={self.scaling:.2f}"
+
+
+# Batched LoRA Experts
+class BatchedLoRAExperts(nn.Module):
+    def __init__(self, d_model: int, config: ExpertConfig, num_experts: int):
+        super().__init__()
+        self.num_experts = num_experts
+        r = config.lora_rank
+        self.scaling = config.lora_alpha / r
+
+        # A: (num_experts, r, d_model)
+        self.lora_A = nn.Parameter(torch.empty(num_experts, r, d_model))
+        # B: (num_experts, d_model, r)
+        self.lora_B = nn.Parameter(torch.empty(num_experts, d_model, r))
+        self.dropout = nn.Dropout(config.lora_dropout)
+
+        for i in range(num_experts):
+            nn.init.kaiming_uniform_(self.lora_A[i], a=math.sqrt(5))
+            nn.init.normal_(self.lora_B[i], std=0.01)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.dropout(x)
+        # (E, C, D) @ (E, D, r) -> (E, C, r)
+        out = torch.bmm(x, self.lora_A.transpose(1, 2))
+        
+        # (E, C, r) @ (E, r, d) -> (E, C, D)
+        out = torch.bmm(out, self.lora_B.transpose(1, 2))
+        
+        return out * self.scaling
+
+    def extra_repr(self) -> str:
+        r = self.lora_A.shape[1]
+        d = self.lora_A.shape[2]
+        return f"d_model={d}, rank={r}, num_experts={self.num_experts}, scaling={self.scaling:.2f}"
+
+
+# TinyExpert MLP
+_ACTIVATIONS = {"gelu": nn.GELU, "silu": nn.SiLU, "relu": nn.ReLU}
+
+
+class TinyExpert(nn.Module):
+
+    def __init__(self, d_model: int, config: ExpertConfig):
+        super().__init__()
+        intermediate = config.tiny_intermediate or d_model * 4
+
+        self.fc1 = nn.Linear(d_model, intermediate, bias=False)
+        self.fc2 = nn.Linear(intermediate, d_model, bias=False)
+        self.dropout = nn.Dropout(config.tiny_dropout)
+
+        act_cls = _ACTIVATIONS.get(config.tiny_activation, nn.SiLU)
+        self.activation = act_cls()
+
+        self._init_weights()
+
+    def _init_weights(self):
+        for linear in (self.fc1, self.fc2):
+            nn.init.kaiming_uniform_(linear.weight, a=math.sqrt(5))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h = self.activation(self.fc1(x))
+        h = self.dropout(h)
+        return self.fc2(h)
+
+    def extra_repr(self) -> str:
+        return (f"d_model={self.fc1.in_features}, "
+                f"intermediate={self.fc1.out_features}")
+
+
+# Factory
+def create_experts(d_model: int, config: ExpertConfig, num_experts: int) -> Union[BatchedLoRAExperts, nn.ModuleList]:
+    if config.expert_type == "lora":
+        return BatchedLoRAExperts(d_model, config, num_experts)
+    elif config.expert_type == "tiny":
+        experts = [TinyExpert(d_model, config) for _ in range(num_experts)]
+    else:
+        raise ValueError(f"Unknown expert_type: {config.expert_type!r}")
+
+    return nn.ModuleList(experts)

architecture/sparse_moe/injection.py

@@ -0,0 +1,129 @@
+from __future__ import annotations
+
+import inspect
+from functools import wraps
+from typing import Set
+
+import torch.nn as nn
+
+from .config import ExpertConfig, RouterConfig
+from .layers import SparseMoELayer
+
+# Maps model_type → set of FFN class names to replace
+_FFN_CLASS_NAMES: dict[str, Set[str]] = {
+    "gpt2":       {"GPT2MLP"},
+    "llama":      {"LlamaMLP"},
+    "mistral":    {"MistralMLP"},
+    "phi":        {"PhiMLP", "Phi3MLP", "PhiMoEMLP"},
+    "phi3":       {"PhiMLP", "Phi3MLP"},
+    "qwen2":      {"Qwen2MLP"},
+    "gemma":      {"GemmaMLP"},
+    "gemma2":     {"GemmaMLP"},
+}
+
+
+def _is_ffn(module: nn.Module, model_type: str) -> bool:
+    class_name = type(module).__name__
+
+    # Known architectures
+    if model_type in _FFN_CLASS_NAMES:
+        return class_name in _FFN_CLASS_NAMES[model_type]
+
+    # Heuristic fallback for unknown architectures
+    heuristic_names = {"MLP", "FFN", "FeedForward", "MoEMLP"}
+    is_match = any(h in class_name for h in heuristic_names)
+
+    if is_match:
+        print("[WARNING] " +
+            f"Unknown model_type={model_type!r}: using heuristic to detect FFN "
+            f"(matched class {class_name!r}). Verify this is correct."
+        )
+
+    return is_match
+
+
+def _set_module(root: nn.Module, dotted_name: str, replacement: nn.Module):
+    parts = dotted_name.split(".")
+    parent = root
+    for part in parts[:-1]:
+        parent = getattr(parent, part)
+    setattr(parent, parts[-1], replacement)
+
+
+# Main injection function
+def inject_moe_layers(model: nn.Module, expert_config: ExpertConfig, router_config: RouterConfig, capacity_factor: float = 1.25) -> nn.Module:
+    model_type = getattr(model.config, "model_type", "qwen2")
+    d_model = model.config.hidden_size
+    print(
+        f"Injecting MoE layers: model_type={model_type!r}, d_model={d_model}, "
+        f"experts={router_config.num_experts}, top_k={router_config.top_k}, "
+        f"type={expert_config.expert_type}"
+    )
+
+    for param in model.parameters():
+        param.requires_grad = False
+
+    replacements = []
+    for name, module in model.named_modules():
+        if _is_ffn(module, model_type):
+            replacements.append((name, module))
+
+    if not replacements:
+        raise ValueError(
+            f"No FFN layers found for model_type={model_type!r}. "
+            f"The model was not modified. Check architecture support before training."
+        )
+
+    for name, original_ffn in replacements:
+        base_ffn = original_ffn if expert_config.expert_type == "lora" else None
+
+        try:
+            layer_device = next(original_ffn.parameters()).device
+            layer_dtype = next(original_ffn.parameters()).dtype
+        except StopIteration:
+            layer_device = next(model.parameters()).device
+            layer_dtype = next(model.parameters()).dtype
+
+        moe_layer = SparseMoELayer(
+            d_model=d_model,
+            expert_config=expert_config,
+            router_config=router_config,
+            base_ffn=base_ffn,
+            capacity_factor=capacity_factor,
+        ).to(device=layer_device, dtype=layer_dtype)
+        moe_layer.layer_name = name
+        _set_module(model, name, moe_layer)
+
+    _install_attention_mask_context(model)
+
+    return model
+
+
+def _install_attention_mask_context(model: nn.Module):
+    moe_layers = [m for m in model.modules() if isinstance(m, SparseMoELayer)]
+    model._moe_layers = moe_layers
+
+    if getattr(model, "_moe_attention_context_installed", False):
+        return
+
+    original_forward = model.forward
+    forward_signature = inspect.signature(original_forward)
+
+    @wraps(original_forward)
+    def wrapped_forward(*args, **kwargs):
+        attention_mask = kwargs.get("attention_mask")
+        if attention_mask is None:
+            try:
+                bound = forward_signature.bind_partial(*args, **kwargs)
+                attention_mask = bound.arguments.get("attention_mask")
+            except TypeError as e:
+                print(f"[DEBUG] Could not bind arguments to trace attention_mask: {e}")
+                attention_mask = None
+
+        current_layers = getattr(model, "_moe_layers", [])
+        for layer in current_layers:
+            layer._current_attention_mask = attention_mask
+        return original_forward(*args, **kwargs)
+
+    model.forward = wrapped_forward
+    model._moe_attention_context_installed = True

architecture/sparse_moe/layers.py

@@ -0,0 +1,338 @@
+from __future__ import annotations
+
+import math
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+from .config import ExpertConfig, RouterConfig
+from .routing import LinearRouter, RoutingOutput
+from .experts import BatchedLoRAExperts, create_experts
+
+
+# Sparse MoE Layer
+class SparseMoELayer(nn.Module):
+
+    def __init__(
+        self,
+        d_model: int,
+        expert_config: ExpertConfig,
+        router_config: RouterConfig,
+        base_ffn: Optional[nn.Module] = None,
+        capacity_factor: float = 1.25,
+    ):
+        super().__init__()
+        self.d_model = d_model
+        self.num_experts = router_config.num_experts
+        self.top_k = router_config.top_k
+        self.expert_type = expert_config.expert_type
+        self.capacity_factor = capacity_factor
+        self.layer_name = "unassigned"
+
+        # Router
+        self.routing = LinearRouter(d_model, router_config)
+
+        # Experts
+        self.experts = create_experts(d_model, expert_config, router_config.num_experts)
+
+        # Base FFN (frozen, used only in LoRA mode)
+        self.base_ffn: Optional[nn.Module] = None
+        if expert_config.expert_type == "lora":
+            if base_ffn is None:
+                raise ValueError("base_ffn is required for LoRA mode")
+            self.base_ffn = base_ffn
+            # Freeze the base FFN
+            for p in self.base_ffn.parameters():
+                p.requires_grad = False
+
+        # Cache for trainer/evaluator to read
+        self._last_aux_loss: Optional[torch.Tensor] = None
+        self._last_stats: Optional[dict] = None
+        self._current_attention_mask: Optional[torch.Tensor] = None
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, S, D = x.shape
+
+        # ── Base FFN (LoRA mode) ──
+        if self.base_ffn is not None:
+            with torch.no_grad():
+                base_out = self.base_ffn(x)
+        else:
+            base_out = torch.zeros_like(x)   # (B, S, D)
+
+        attention_mask = self._resolve_attention_mask(B, S, x.device)
+
+        # ── Routing ──
+        routing = self.routing(x, attention_mask=attention_mask)
+
+        # ── Sparse expert dispatch ──
+        expert_out = self._sparse_forward(x, routing, B, S, D, attention_mask)
+
+        # ── Combine ──
+        output = base_out + expert_out.to(base_out.dtype)
+
+        # Cache for downstream
+        self._last_aux_loss = routing.aux_loss
+        self._last_stats = routing.stats
+
+        return output
+
+    def _resolve_attention_mask(self, batch_size: int, seq_len: int, device: torch.device) -> Optional[torch.Tensor]:
+        attention_mask = self._current_attention_mask
+        if attention_mask is None:
+            return None
+
+        if attention_mask.dim() > 2:
+            if attention_mask.dim() == 4:
+                # Handle 4D causal masks e.g. [B, 1, S, S] → take last row
+                attention_mask = attention_mask[:, 0, -1, :]
+            elif attention_mask.dim() == 3:
+                attention_mask = attention_mask[:, -1, :]
+            else:
+                attention_mask = attention_mask.reshape(batch_size, -1)
+        if attention_mask.shape[1] != seq_len:
+            if attention_mask.shape[1] < seq_len:
+                raise ValueError(
+                    f"attention_mask length {attention_mask.shape[1]} does not match "
+                    f"hidden state length {seq_len} for {self.layer_name}"
+                )
+            attention_mask = attention_mask[:, -seq_len:]
+        return attention_mask.to(device=device, dtype=torch.bool)
+
+    def _sparse_forward(
+        self,
+        x: torch.Tensor,
+        routing: RoutingOutput,
+        B: int, S: int, D: int,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if isinstance(self.experts, BatchedLoRAExperts):
+            return self._batched_sparse_forward(x, routing, B, S, D, attention_mask)
+        return self._loop_sparse_forward(x, routing, B, S, D, attention_mask)
+
+    def _loop_sparse_forward(self, x: torch.Tensor, routing: RoutingOutput, B: int, S: int, D: int, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+        flat_x = x.reshape(B * S, D)
+        combined = torch.zeros_like(flat_x)
+        flat_valid = self._flat_valid_mask(attention_mask, B, S, flat_x.device)
+        
+        # Explicit graph break to cleanly branch out if zero tokens
+        num_valid_int = int(flat_valid.sum().item())
+        if num_valid_int == 0:
+            return combined.reshape(B, S, D)
+
+        flat_indices = routing.top_k_indices.reshape(B * S, self.top_k)
+        flat_weights = routing.top_k_weights.reshape(B * S, self.top_k)
+
+        # Capacity: max tokens any single expert may handle
+        expected = math.ceil(num_valid_int * self.top_k / self.num_experts)
+        capacity = max(1, int(self.capacity_factor * expected))
+
+        # ── Per-expert dispatch ──
+        for e_idx, expert in enumerate(self.experts):
+            mask = (flat_indices == e_idx) & flat_valid.unsqueeze(-1)
+            token_positions, k_slots = torch.where(mask)
+
+            if token_positions.numel() == 0:
+                continue
+
+            if token_positions.numel() > capacity:
+                weights_for_sort = flat_weights[token_positions, k_slots]
+                keep = torch.topk(weights_for_sort, k=capacity, largest=True, sorted=False).indices
+                token_positions = token_positions[keep]
+                k_slots = k_slots[keep]
+
+            # Gather tokens for this expert
+            expert_input = flat_x[token_positions]               # (num_selected, D)
+
+            # Run expert
+            expert_output = expert(expert_input)                 # (num_selected, D)
+
+            weights = flat_weights[token_positions, k_slots]
+            weighted_output = expert_output * weights.unsqueeze(-1)
+
+            # Coalesce duplicate token indices first so the final scatter uses
+            # unique destinations only. This avoids CUDA backward bugs around
+            # index_add_ with repeated indices on top-k routed tokens.
+            unique_pos, reduced_output = self._coalesce_duplicate_dispatch(
+                token_positions,
+                weighted_output,
+                combined.dtype,
+            )
+            combined.index_add_(0, unique_pos, reduced_output)
+
+        return combined.reshape(B, S, D)  # (B, S, D)
+
+    def _batched_sparse_forward(self, x: torch.Tensor, routing: RoutingOutput, B: int, S: int, D: int, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+        if not getattr(torch.compiler, "is_compiling", lambda: False)():
+            return self._forward_dynamic(x, routing, B, S, D, attention_mask)
+        return self._forward_padded_loop(x, routing, B, S, D, attention_mask)
+
+    def _forward_dynamic(
+        self, x: torch.Tensor, routing: RoutingOutput, B: int, S: int, D: int,
+        attention_mask: Optional[torch.Tensor]
+    ) -> torch.Tensor:
+        N = B * S
+        flat_x = x.reshape(N, D)
+        combined = torch.zeros_like(flat_x)
+        flat_valid = self._flat_valid_mask(attention_mask, B, S, flat_x.device)
+
+        num_valid_int = int(flat_valid.sum().item())
+        if num_valid_int == 0:
+            return combined.reshape(B, S, D)
+
+        flat_indices = routing.top_k_indices.reshape(N, self.top_k)
+        flat_weights = routing.top_k_weights.reshape(N, self.top_k)
+
+        expected = math.ceil(num_valid_int * self.top_k / self.num_experts)
+        capacity = max(1, int(self.capacity_factor * expected))
+        E, C = self.num_experts, capacity
+
+        one_hot = torch.zeros(N * self.top_k, E, dtype=torch.long, device=flat_x.device)
+        one_hot.scatter_(1, flat_indices.reshape(-1, 1), 1)
+
+        cum_count = one_hot.cumsum(0)
+        slot = ((cum_count - 1) * one_hot).sum(-1)
+
+        valid_flat = flat_valid.unsqueeze(-1).expand(N, self.top_k).reshape(-1)
+        capacity_mask = valid_flat & (slot < C)
+
+        active = capacity_mask.nonzero(as_tuple=False).squeeze(-1)
+        if active.numel() == 0:
+            return combined.reshape(B, S, D)
+
+        active_token_idx  = (active // self.top_k)
+        active_expert_idx = flat_indices.reshape(-1)[active]
+        active_slot_idx   = slot[active]
+        active_weights    = flat_weights.reshape(-1)[active]
+
+        buf_idx = active_expert_idx * C + active_slot_idx
+
+        batched_input = flat_x.new_zeros(E * C, D)
+        batched_input.index_copy_(0, buf_idx, flat_x[active_token_idx])
+        batched_input = batched_input.reshape(E, C, D)
+
+        expert_out = self.experts(batched_input)
+        expert_out_flat = expert_out.reshape(E * C, D)
+
+        weighted_out = expert_out_flat[buf_idx] * active_weights.unsqueeze(-1)
+
+        unique_pos, reduced_out = self._coalesce_duplicate_dispatch(
+            active_token_idx, weighted_out, combined.dtype,
+        )
+        combined.index_add_(0, unique_pos, reduced_out)
+
+        return combined.reshape(B, S, D)
+
+    def _forward_padded_loop(
+        self, x: torch.Tensor, routing: RoutingOutput, B: int, S: int, D: int,
+        attention_mask: Optional[torch.Tensor]
+    ) -> torch.Tensor:
+        flat_x = x.reshape(B * S, D)
+        combined = torch.zeros_like(flat_x)
+        flat_valid = self._flat_valid_mask(attention_mask, B, S, flat_x.device)
+
+        num_valid_int = int(flat_valid.sum().item())
+        if num_valid_int == 0:
+            return combined.reshape(B, S, D)
+
+        flat_indices = routing.top_k_indices.reshape(B * S, self.top_k)
+        flat_weights = routing.top_k_weights.reshape(B * S, self.top_k)
+
+        expected = math.ceil(num_valid_int * self.top_k / self.num_experts)
+        capacity = max(1, int(self.capacity_factor * expected))
+
+        inputs, positions, weights, valids = [], [], [], []
+
+        for e_idx in range(self.num_experts):
+            mask = (flat_indices == e_idx) & flat_valid.unsqueeze(-1)
+            token_idx, k_slots = torch.where(mask)
+
+            if token_idx.numel() > capacity:
+                w_sort = flat_weights[token_idx, k_slots]
+                keep = torch.topk(w_sort, k=capacity, largest=True, sorted=False).indices
+                token_idx = token_idx[keep]
+                k_slots = k_slots[keep]
+
+            n_tokens = token_idx.numel()
+
+            pad_inputs = torch.zeros(capacity, D, dtype=flat_x.dtype, device=flat_x.device)
+            pad_pos = torch.zeros(capacity, dtype=torch.long, device=flat_x.device)
+            pad_weights = torch.zeros(capacity, dtype=flat_weights.dtype, device=flat_x.device)
+            pad_valid = torch.zeros(capacity, dtype=torch.bool, device=flat_x.device)
+
+            if n_tokens > 0:
+                pad_inputs[:n_tokens] = flat_x[token_idx]
+                pad_pos[:n_tokens] = token_idx
+                pad_weights[:n_tokens] = flat_weights[token_idx, k_slots]
+                pad_valid[:n_tokens] = True
+
+            inputs.append(pad_inputs)
+            positions.append(pad_pos)
+            weights.append(pad_weights)
+            valids.append(pad_valid)
+
+        batched_input = torch.stack(inputs)                    # (E, C, D)
+        batched_pos = torch.stack(positions).reshape(-1)       # (E * C,)
+        batched_weights = torch.stack(weights).reshape(-1, 1)  # (E * C, 1)
+        batched_valid = torch.stack(valids).reshape(-1)        # (E * C,)
+
+        expert_out = self.experts(batched_input)               # (E, C, D)
+        expert_out = expert_out.reshape(-1, D)                 # (E * C, D)
+
+        weighted_out = expert_out * batched_weights            # (E * C, D)
+
+        if batched_valid.any():
+            unique_pos, reduced_output = self._coalesce_duplicate_dispatch(
+                batched_pos[batched_valid],
+                weighted_out[batched_valid],
+                combined.dtype,
+            )
+            combined.index_add_(0, unique_pos, reduced_output)
+
+        return combined.reshape(B, S, D)
+
+    def _coalesce_duplicate_dispatch(self, token_positions: torch.Tensor, weighted_output: torch.Tensor, out_dtype: torch.dtype) -> tuple[torch.Tensor, torch.Tensor]:
+        if token_positions.numel() == 0:
+            empty = weighted_output.new_zeros((0, weighted_output.shape[-1]), dtype=out_dtype)
+            return token_positions, empty
+
+        weighted_output = weighted_output.to(out_dtype).contiguous()
+        if token_positions.numel() == 1:
+            return token_positions, weighted_output
+
+        order = torch.argsort(token_positions)
+        sorted_pos = token_positions[order]
+        sorted_out = weighted_output[order]
+
+        is_new = torch.ones_like(sorted_pos, dtype=torch.bool)
+        is_new[1:] = sorted_pos[1:] != sorted_pos[:-1]
+        start_idx = torch.nonzero(is_new, as_tuple=False).flatten()
+        end_idx = torch.empty_like(start_idx)
+        if start_idx.numel() > 1:
+            end_idx[:-1] = start_idx[1:] - 1
+        end_idx[-1] = sorted_pos.numel() - 1
+
+        cumsum = sorted_out.cumsum(dim=0)
+        reduced = cumsum[end_idx].clone()
+        if start_idx.numel() > 1:
+            reduced[1:] -= cumsum[start_idx[1:] - 1]
+
+        return sorted_pos[start_idx], reduced.to(out_dtype).contiguous()
+
+    def _flat_valid_mask(
+        self,
+        attention_mask: Optional[torch.Tensor],
+        batch_size: int,
+        seq_len: int,
+        device: torch.device,
+    ) -> torch.Tensor:
+        """Flatten a token-validity mask, defaulting to all tokens valid."""
+        if attention_mask is None:
+            return torch.ones(batch_size * seq_len, dtype=torch.bool, device=device)
+        return attention_mask.reshape(batch_size * seq_len).to(device=device, dtype=torch.bool)
+
+    def extra_repr(self) -> str:
+        mode = "lora+residual" if self.base_ffn is not None else "tiny+replace"
+        return (f"d_model={self.d_model}, experts={self.num_experts}, "
+                f"top_k={self.top_k}, mode={mode}, capacity={self.capacity_factor}")

architecture/sparse_moe/prompts.py

@@ -0,0 +1,72 @@
+DOMAIN_PROMPTS = {
+    "code": [
+        "def quicksort(arr):\n    if len(arr) <= 1:\n        return arr",
+        "class BinarySearchTree:\n    def __init__(self):\n        self.root = None",
+        "import torch\nimport torch.nn as nn\n\nclass TransformerBlock(nn.Module):",
+        "def solve_n_queens(n):\n    \"\"\"Return all distinct solutions to the n-queens puzzle.\"\"\"",
+        "async def fetch_url(url):\n    import aiohttp\n    async with aiohttp.ClientSession() as session:",
+        "def fibonacci_memo(n, memo={}):\n    \"\"\"Compute Fibonacci with memoization.\"\"\"",
+        "def deep_copy_tree(node):\n    \"\"\"Create a deep copy of a binary tree.\"\"\"",
+        "def detect_cycle_in_linked_list(head):\n    \"\"\"Floyd's cycle-finding algorithm.\"\"\"",
+        "def parse_log_file(filename):\n    \"\"\"Efficiently parse a large log file with regex.\"\"\"",
+        "def rotate_matrix_90_degrees(matrix):\n    \"\"\"Rotate an n x n matrix in-place.\"\"\"",
+        "def is_palindrome_string(s):\n    \"\"\"Check if string is a palindrome ignoring case/spacing.\"\"\"",
+        "def compute_levenshtein_distance(s1, s2):\n    \"\"\"Compute edit distance between two strings.\"\"\"",
+    ],
+    "reasoning": [
+        "If Mary is twice as old as John was when Mary was... wait. Assume Mary is 40. Then John is",
+        "The murderer must be either the butler or the gardener. Since the gardener was seen in town,",
+        "If all Bloops are Razzies and all Razzies are Lurgs, then it must be true that",
+        "A man has to get a fox, a chicken, and a sack of corn across a river in a small boat.",
+        "Complete the sequence and explain the logic: 2, 6, 12, 20, 30, ...",
+        "If you have two buckets, one that holds 3 liters and one that holds 5 liters, how do you get exactly 4 liters?",
+        "Three gods A, B, and C are called True, False, and Random. True always speaks truly...",
+        "A grandfather, two fathers, and two sons went to the movie theater together and bought one ticket each.",
+        "What has keys but no locks, space but no room, and allows you to enter but never leave?",
+        "Which is heavier: a ton of bricks or a ton of feathers? Explain your reasoning.",
+        "I speak without a mouth and hear without ears. I have no body, but I come alive with wind. What am I?",
+        "If a plane crashes on the border between the United States and Canada, where do they bury the survivors?",
+    ],
+    "prose": [
+        "The history of artificial intelligence began in the mid-20th century when",
+        "Deep beneath the ocean's surface, where sunlight cannot reach,",
+        "The neon lights of Neo-Tokyo flickered against the damp pavement as",
+        "The time traveler stepped out of the capsule and inhaled the pungent air of 1666 London.",
+        "To explain quantum entanglement, imagine two magical dice that always show the same number",
+        "Describe the sensation of waking up on a planet with two suns and purple grass.",
+        "Write a letter from a pioneer in the year 1850 describing their journey across the American West.",
+        "A dialogue between a cynical detective and a naive robot about the nature of justice.",
+        "Describe the architecture of a city built entirely inside a giant, hollowed-out tree.",
+        "Tell a short fairy tale about a weaver who could spin moonlight into cloth.",
+        "Explain the concept of 'Ikigai' and how it relates to modern work-life balance.",
+        "Describe the smells and sounds of a bustling Spice Market in ancient Marrakesh.",
+    ],
+    "qa": [
+        "Question: What is the difference between supervised and unsupervised learning?\nAnswer:",
+        "Question: How do transformers achieve parallelism compared to RNNs?\nAnswer:",
+        "Question: What is the main objective of a Generative Adversarial Network (GAN)?\nAnswer:",
+        "Question: Explain the vanishing gradient problem and how Residual Connections help.\nAnswer:",
+        "Question: What are the primary benefits of using BFloat16 over Float16 for LLM training?\nAnswer:",
+        "Question: Explain the difference between Model Parallelism and Data Parallelism in deep learning.\nAnswer:",
+        "Question: What is 'Temperature' in the context of LLM sampling and how does it effect output diversity?\nAnswer:",
+        "Question: What is the 'Attention Bottleneck' in large scale transformers?\nAnswer:",
+        "Question: Define 'Zero-Shot Learning' and provide an example of its application.\nAnswer:",
+        "Question: What are the trade-offs between dense models and Mixture-of-Experts (MoE) models?\nAnswer:",
+        "Question: Explain the concept of 'Reward Modeling' in RLHF.\nAnswer:",
+        "Question: What is the difference between a Softmax and a Gumbel-Softmax?\nAnswer:",
+    ],
+    "math": [
+        "Problem: Prove that the sum of the first n natural numbers equals n(n+1)/2.\nSolution:",
+        "Problem: Find the derivative of f(x) = x³ · sin(x).\nSolution:",
+        "Problem: Solve for x in the quadratic equation 2x² - 5x + 3 = 0.\nSolution:",
+        "Problem: Calculate the volume of a sphere with a radius of 4 units.\nSolution:",
+        "Problem: If a train travels 450 km in 5 hours, what is its average speed in km/h?\nSolution:",
+        "Problem: Find the integral of f(x) = e^(2x) from 0 to 1.\nSolution:",
+        "Problem: In a triangle with sides measuring 7, 24, and 25, identify if it is a right triangle.\nSolution:",
+        "Problem: Calculate the probability of flipping exactly 3 heads in 5 coin tosses.\nSolution:",
+        "Problem: Solve for x: log2(x + 3) = 4.\nSolution:",
+        "Problem: A rectangular garden has a perimeter of 100 meters and its length is twice its width. Find dimensions.\nSolution:",
+        "Problem: Calculate the dot product of vectors A = [2, 5, -1] and B = [3, -2, 4].\nSolution:",
+        "Problem: Find the local maximum of the function f(x) = -x² + 6x - 5.\nSolution:",
+    ],
+}

architecture/sparse_moe/reporting.py

@@ -0,0 +1,324 @@
+from __future__ import annotations
+
+import copy
+import json
+from pathlib import Path
+from typing import Dict, List, Optional
+
+import numpy as np
+import torch
+
+from .visualization import (
+    HAS_MATPLOTLIB,
+    DomainScorePlot,
+    TrainingConvergencePlot,
+    ExpertRoutingHeatmap,
+)
+
+from .config import ProjectConfig
+from .evaluation import evaluate_per_domain
+from .datasets import build_mixed_loader
+from .injection import inject_moe_layers
+from .trainer import MoETrainer
+from .utils import set_seed, to_serializable, parse_seed_list
+
+# Before / after comparison
+def build_before_after_rows(before: Dict[str, Dict], after: Dict[str, Dict], dense: Optional[Dict[str, Dict]] = None) -> List[Dict]:
+    rows = []
+    for domain in before:
+        before_ent = before[domain]["routing_entropy"]
+        after_ent = after[domain]["routing_entropy"]
+        row = {
+            "domain": domain,
+            "ppl_before": before[domain]["perplexity"],
+            "ppl_after": after[domain]["perplexity"],
+            "ppl_change": after[domain]["perplexity"] - before[domain]["perplexity"],
+            "entropy_before": before_ent if before_ent > 0 else None,
+            "entropy_after": after_ent,
+        }
+        if dense and domain in dense:
+            row["ppl_dense"] = dense[domain]["perplexity"]
+            row["moe_vs_dense"] = after[domain]["perplexity"] - dense[domain]["perplexity"]
+        rows.append(row)
+    return rows
+
+
+def print_before_after_report(rows: List[Dict], include_dense: bool):
+    print("\n" + "=" * 75)
+    print("  Pretrained vs Trained — Per-Domain Perplexity")
+    print("=" * 75)
+
+    if include_dense:
+        header = (f"  {'Domain':<10s} {'Before':>8s} {'After':>8s} {'Δ':>8s} "
+                  f"{'Dense':>8s} {'MoE−Dense':>10s}")
+    else:
+        header = f"  {'Domain':<10s} {'Before':>8s} {'After':>8s} {'Δ':>8s}"
+    print(header)
+    print("  " + "-" * (len(header) - 2))
+
+    for r in rows:
+        delta_sym = "↓" if r["ppl_change"] < 0 else "↑"
+        line = (f"  {r['domain']:<10s} {r['ppl_before']:>8.2f} {r['ppl_after']:>8.2f} "
+                f"{r['ppl_change']:>+7.2f}{delta_sym}")
+        if include_dense and "ppl_dense" in r:
+            diff_sym = "✓" if r["moe_vs_dense"] <= 0 else "✗"
+            line += f" {r['ppl_dense']:>8.2f} {r['moe_vs_dense']:>+9.2f}{diff_sym}"
+        print(line)
+    print("=" * 75 + "\n")
+
+
+# Dashboard artifacts
+def save_dashboard(cfg: ProjectConfig, rows: List[Dict], before_results: Dict, after_results: Dict, after_spec: Dict, history: Dict, params: Dict, dense_results: Optional[Dict] = None):
+    out = Path(cfg.artifact_dir)
+    out.mkdir(parents=True, exist_ok=True)
+
+    # ── CSV ──
+    import csv
+    csv_path = out / "before_after.csv"
+    if not rows:
+        print("[WARNING] No comparison rows to save — skipping CSV.")
+    else:
+        with open(csv_path, "w", newline="") as f:
+            writer = csv.DictWriter(f, fieldnames=rows[0].keys())
+            writer.writeheader()
+            writer.writerows([{k: round(v, 4) if isinstance(v, float) else v
+                               for k, v in r.items()} for r in rows])
+        print(f"  CSV saved to      : {csv_path}")
+
+    # ── JSON ──
+    json_path = out / "eval_results.json"
+    with open(json_path, "w") as f:
+        json.dump(to_serializable({
+            "before": before_results,
+            "after": after_results,
+            "dense": dense_results,
+            "params": params,
+            "after_spec": {k: v for k, v in after_spec.items() if k != "affinity"},
+        }), f, indent=2)
+    print(f"  JSON saved to     : {json_path}")
+
+    spec_path = out / "specialization.json"
+    with open(spec_path, "w") as f:
+        json.dump(to_serializable(after_spec), f, indent=2)
+    print(f"  Spec JSON saved to: {spec_path}")
+
+    hist_path = out / "training_history.json"
+    with open(hist_path, "w") as f:
+        json.dump(to_serializable(history), f, indent=2)
+    print(f"  History saved to  : {hist_path}")
+
+    # ── Plots ──
+    if HAS_MATPLOTLIB:
+        # Per-Domain Comparison
+        domain_plot = DomainScorePlot()
+        domain_plot.plot_comparison(rows, include_dense=dense_results is not None)
+        domain_plot.save(out / "plot_domain_scores.png")
+        domain_plot.close()
+        print(f"  Domain plot saved : {out / 'plot_domain_scores.png'}")
+        
+        # Training Convergence
+        if history and "train_loss" in history:
+            train_plot = TrainingConvergencePlot()
+            train_plot.plot_history(history)
+            train_plot.save(out / "plot_training_loss.png")
+            train_plot.close()
+            print(f"  Loss plot saved   : {out / 'plot_training_loss.png'}")
+            
+        # Expert Routing Heatmap
+        routing_plot = ExpertRoutingHeatmap()
+        routing_plot.plot_routing(after_spec)
+        routing_plot.save(out / "plot_expert_routing.png")
+        routing_plot.close()
+        print(f"  Heatmap saved     : {out / 'plot_expert_routing.png'}")
+    else:
+        print("[WARNING] Matplotlib not available. Skipping plot generation.")
+
+    # ── Markdown summary ──
+    _save_dashboard_md(out / "dashboard_summary.md", cfg, rows, params,
+                       after_spec, history, include_dense=dense_results is not None)
+
+
+def _save_dashboard_md(path: Path, cfg: ProjectConfig, rows: List[Dict], params: Dict, spec: Dict, history: Dict, include_dense: bool):
+    lines = [
+        "# Keiro — Experiment Summary\n",
+        f"**Base model**: `{cfg.base_model}`  ",
+        f"**Experts**: {cfg.routing.num_experts} × {cfg.expert.expert_type} "
+        f"(rank {cfg.expert.lora_rank})  ",
+        f"**Top-k**: {cfg.routing.top_k}  ",
+        f"**Epochs**: {cfg.training.num_epochs}  ",
+        f"**Seq length**: {cfg.training.max_seq_len}  ",
+        f"**Samples/domain**: {cfg.data.samples_per_domain}  ",
+        f"**LR**: {cfg.training.lr}  ",
+        f"**Batch size**: {cfg.training.batch_size} × {cfg.training.gradient_accumulation_steps} "
+        f"(eff={cfg.training.batch_size * cfg.training.gradient_accumulation_steps})  \n",
+    ]
+
+    train_loss = history.get("train_loss", [])
+    val_losses = history.get("val_loss", [])
+    if train_loss:
+        lines.append("## Training\n")
+        lines.append("| Metric | Value |")
+        lines.append("|--------|-------|")
+        lines.append(f"| Total steps | {len(train_loss)} |")
+        lines.append(f"| Final train loss | {train_loss[-1]:.4f} |")
+        if val_losses:
+            best_ep = val_losses.index(min(val_losses)) + 1
+            lines.append(f"| Best val loss | {min(val_losses):.4f} (epoch {best_ep}) |")
+            if len(val_losses) > 1 and val_losses[-1] > min(val_losses):
+                lines.append(f"| Early stopping | Yes (restored epoch {best_ep}) |")
+        lines.append("")
+
+    lines.append("## Parameter Efficiency\n")
+    lines.append("| Metric | Value |")
+    lines.append("|--------|-------|")
+    lines.append(f"| Total params | {params['total']:,} |")
+    lines.append(f"| Trainable | {params['trainable']:,} ({params['trainable_pct']:.2f}%) |")
+    lines.append(f"| Router | {params['router']:,} |")
+    lines.append(f"| Experts | {params['experts']:,} |\n")
+
+    lines.append("## Per-Domain Results\n")
+    if include_dense:
+        lines.append("| Domain | Before | After | delta | % Change | Dense | MoE−Dense |")
+        lines.append("|--------|--------|-------|---|----------|-------|-----------|")
+    else:
+        lines.append("| Domain | Before | After | delta | % Change |")
+        lines.append("|--------|--------|-------|---|----------|")
+
+    for r in rows:
+        sym = "↓" if r["ppl_change"] < 0 else "↑"
+        pct = (r["ppl_change"] / r["ppl_before"]) * 100 if r["ppl_before"] > 0 else 0
+        line = (f"| {r['domain']} | {r['ppl_before']:.2f} | {r['ppl_after']:.2f} | "
+                f"{r['ppl_change']:+.2f}{sym} | {pct:+.1f}% |")
+        if include_dense and "ppl_dense" in r:
+            line += f" {r['ppl_dense']:.2f} | {r['moe_vs_dense']:+.2f} |"
+        lines.append(line)
+    lines.append("")
+
+    div = spec.get("divergence")
+    if div is not None:
+        if isinstance(div, torch.Tensor):
+            div = div.numpy()
+        n = div.shape[0]
+        upper = [div[i, j] for i in range(n) for j in range(i + 1, n)]
+        if upper:
+            lines.append("## Routing Divergence\n")
+            lines.append(f"Mean inter-domain JS divergence: **{np.mean(upper):.4f}**  ")
+            lines.append("(Higher = more distinct routing per domain)\n")
+
+    lines.append("### Visualizations\n")
+    lines.append("![Domain Scores](plot_domain_scores.png)")
+    lines.append("![Training Loss](plot_training_loss.png)")
+    lines.append("![Expert Routing](plot_expert_routing.png)")
+    lines.append("")
+
+    path.write_text("\n".join(lines), encoding="utf-8")
+    print(f"  Summary saved to  : {path}")
+
+
+# Qualitative text generation
+from .prompts import DOMAIN_PROMPTS
+
+
+@torch.no_grad()
+def generate_qualitative_samples(model, tokenizer, device, max_new_tokens: int = 80) -> str:
+    model.eval()
+    lines = []
+
+    for domain, prompts in DOMAIN_PROMPTS.items():
+        lines.append(f"{'='*70}")
+        lines.append(f"  DOMAIN: {domain.upper()}")
+        lines.append(f"{'='*70}")
+
+        for i, prompt in enumerate(prompts, 1):
+            inputs = tokenizer(
+                prompt, return_tensors="pt", truncation=True, max_length=256
+            ).to(device)
+
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=max_new_tokens,
+                do_sample=False,
+                pad_token_id=tokenizer.eos_token_id,
+                repetition_penalty=1.2,
+            )
+
+            if "attention_mask" in inputs:
+                prompt_len = int(inputs["attention_mask"][0].sum().item())
+            else:
+                prompt_len = inputs["input_ids"].shape[1]
+            completion = tokenizer.decode(outputs[0, prompt_len:], skip_special_tokens=True)
+
+            lines.append(f"\n  Prompt {i}:")
+            lines.append(f"  {'─'*60}")
+            for line in prompt.split('\n'):
+                lines.append(f"    {line}")
+            lines.append("\n  Completion:")
+            lines.append(f"  {'─'*60}")
+            for line in completion.strip().split('\n'):
+                lines.append(f"    {line}")
+            lines.append("")
+
+    return "\n".join(lines)
+
+
+# Multi-seed experiment
+def run_multi_seed(cfg: ProjectConfig, train_loaders: Dict, eval_loaders: Dict, device: torch.device):
+    from transformers import AutoModelForCausalLM
+    seeds = parse_seed_list(cfg.multi_seed_list)
+    all_results = []
+
+    for seed in seeds:
+        print(f"\n{'='*50}")
+        print(f"Multi-seed run: seed={seed}")
+        print(f"{'='*50}")
+        set_seed(seed)
+
+        model = AutoModelForCausalLM.from_pretrained(cfg.base_model).to(device)
+        model = inject_moe_layers(model, cfg.expert, cfg.routing)
+        try:
+            model.gradient_checkpointing_enable(
+                gradient_checkpointing_kwargs={"use_reentrant": False}
+            )
+        except TypeError:
+            model.gradient_checkpointing_enable()
+
+        mixed_train = build_mixed_loader(train_loaders, cfg.training.batch_size, seed)
+        mixed_eval = build_mixed_loader(eval_loaders, cfg.training.batch_size, seed)
+
+        train_cfg = copy.copy(cfg.training)
+        train_cfg.num_epochs = cfg.multi_seed_epochs
+        train_cfg.save_every_steps = 0
+
+        trainer = MoETrainer(model, mixed_train, mixed_eval, train_cfg)
+        trainer.train()
+
+        results = evaluate_per_domain(model, eval_loaders, device)
+        all_results.append(results)
+
+        del model, trainer
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+    _print_multi_seed_summary(seeds, all_results)
+
+    out = Path(cfg.artifact_dir) / "multi_seed_results.json"
+    with open(out, "w") as f:
+        json.dump(to_serializable({"seeds": seeds, "results": all_results}), f, indent=2)
+    print(f"Multi-seed results saved: {out}")
+
+
+def _print_multi_seed_summary(seeds: List[int], all_results: List[Dict]):
+    domains = list(all_results[0].keys())
+
+    print(f"\n{'='*60}")
+    print(f"  Multi-Seed Summary ({len(seeds)} seeds: {seeds})")
+    print(f"{'='*60}")
+    print(f"  {'Domain':<10s} {'Mean PPL':>10s} {'Std PPL':>10s} {'Min':>8s} {'Max':>8s}")
+    print(f"  {'-'*48}")
+
+    for domain in domains:
+        ppls = [r[domain]["perplexity"] for r in all_results]
+        mean = sum(ppls) / len(ppls)
+        std = (sum((p - mean)**2 for p in ppls) / max(len(ppls) - 1, 1)) ** 0.5
+        print(f"  {domain:<10s} {mean:>10.2f} {std:>10.2f} {min(ppls):>8.2f} {max(ppls):>8.2f}")
+    print(f"{'='*60}\n")

architecture/sparse_moe/routing.py

@@ -0,0 +1,171 @@
+from __future__ import annotations
+
+from typing import Dict, NamedTuple, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .config import RouterConfig
+
+
+# Router output container
+class RoutingOutput(NamedTuple):
+
+    gate_probs: torch.Tensor
+    top_k_indices: torch.Tensor
+    top_k_weights: torch.Tensor
+    aux_loss: torch.Tensor
+    stats: Dict[str, float]
+
+
+# Linear Router
+class LinearRouter(nn.Module):
+
+    def __init__(self, d_model: int, config: RouterConfig):
+        super().__init__()
+        self.num_experts = config.num_experts
+        self.top_k = min(config.top_k, config.num_experts)
+        self.noise_std = config.noise_std
+        self.z_loss_coeff = config.z_loss_coeff
+        self.loss_type = config.loss_type
+        self.entropy_reg = config.entropy_reg
+
+        # Gating projection: logits = x @ W + b
+        self.gate = nn.Linear(d_model, config.num_experts, bias=True)
+        nn.init.xavier_uniform_(self.gate.weight)
+        nn.init.zeros_(self.gate.bias)
+
+    def forward(self, x: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> RoutingOutput:
+        B, S, D = x.shape
+        token_mask = self._normalise_attention_mask(attention_mask, B, S, x.device)
+
+        # ── Logits ──
+        clean_logits = self.gate(x)
+        logits = clean_logits
+
+        # Exploration noise (training only)
+        if self.training and self.noise_std > 0:
+            logits = logits + torch.randn_like(logits) * self.noise_std
+
+        # ── Probabilities ──
+        orig_dtype = logits.dtype
+        probs = F.softmax(logits.float(), dim=-1)
+
+        # ── Top-k selection ──
+        top_k_weights, top_k_indices = torch.topk(probs, self.top_k, dim=-1)
+
+        if self.top_k > 1:
+            top_k_weights = top_k_weights / (
+                top_k_weights.sum(dim=-1, keepdim=True) + 1e-9
+            )
+
+        # Sparse gate tensor (for downstream masking)
+        gate_probs = torch.zeros_like(probs)
+        gate_probs.scatter_(-1, top_k_indices, top_k_weights)
+
+        # Cast weights back to original dtype for downstream expert dispatch
+        top_k_weights = top_k_weights.to(orig_dtype)
+        gate_probs = gate_probs.to(orig_dtype)
+
+        # Auxiliary losses
+        aux_loss = self._load_balance_loss(probs, top_k_indices, token_mask)
+        aux_loss = aux_loss + self._z_loss(clean_logits, token_mask)
+        if self.entropy_reg > 0:
+            aux_loss = aux_loss + self._entropy_loss(probs, token_mask)
+
+        valid_mask = token_mask.unsqueeze(-1)
+        top_k_weights = top_k_weights.masked_fill(~valid_mask, 0.0)
+
+        # Stats (detached, for logging only)
+        with torch.no_grad():
+            f = self._expert_fractions(top_k_indices, token_mask)
+            # Exact entropy: use torch.where to avoid log(0)
+            log_probs = torch.where(probs > 0, probs.log(), torch.zeros_like(probs))
+            token_entropy = -(probs * log_probs).sum(dim=-1)
+            valid_tokens = int(token_mask.sum().item())
+            if valid_tokens > 0:
+                entropy = token_entropy[token_mask].mean()
+            else:
+                entropy = probs.new_zeros(())
+            stats = {
+                "routing_entropy": entropy.item(),
+                "expert_fractions": f.cpu().tolist(),
+                "aux_loss": aux_loss.item(),
+                "num_tokens": valid_tokens,
+                "num_assignments": valid_tokens * self.top_k,
+            }
+
+        gate_probs = gate_probs.masked_fill(~valid_mask, 0.0)
+        return RoutingOutput(gate_probs, top_k_indices, top_k_weights, aux_loss, stats)
+
+    # Loss helpers
+    def _normalise_attention_mask(self, attention_mask: Optional[torch.Tensor], batch_size: int, seq_len: int, device: torch.device) -> torch.Tensor:
+        if attention_mask is None:
+            return torch.ones(batch_size, seq_len, dtype=torch.bool, device=device)
+
+        if attention_mask.dim() > 2:
+            if attention_mask.dim() == 4:
+                # Handle 4D causal masks e.g. [B, 1, S, S]
+                attention_mask = attention_mask[:, 0, -1, :]
+            elif attention_mask.dim() == 3:
+                attention_mask = attention_mask[:, -1, :]
+            else:
+                attention_mask = attention_mask.reshape(batch_size, -1)
+        if attention_mask.shape[1] != seq_len:
+            if attention_mask.shape[1] < seq_len:
+                raise ValueError(
+                    f"attention_mask length {attention_mask.shape[1]} does not match "
+                    f"sequence length {seq_len}"
+                )
+            attention_mask = attention_mask[:, -seq_len:]
+        return attention_mask.to(device=device, dtype=torch.bool)
+
+    def _expert_fractions(self, top_k_indices: torch.Tensor, token_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+        flat = top_k_indices.reshape(-1)
+        if token_mask is not None:
+            assignment_mask = token_mask.reshape(-1, 1).expand(-1, self.top_k).reshape(-1)
+            flat = flat[assignment_mask]
+        if flat.numel() == 0:
+            return torch.zeros(self.num_experts, device=top_k_indices.device)
+        counts = torch.bincount(flat, minlength=self.num_experts).float()
+        return counts / counts.sum().clamp(min=1)
+
+    def _load_balance_loss(self, probs: torch.Tensor, top_k_indices: torch.Tensor, token_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+        f = self._expert_fractions(top_k_indices, token_mask)
+        P = self._mean_probs(probs, token_mask)
+        return self.num_experts * (f * P).sum()
+
+    def _mean_probs(self, probs: torch.Tensor, token_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+        if token_mask is None:
+            return probs.mean(dim=(0, 1))
+
+        flat_mask = token_mask.reshape(-1)
+        flat_probs = probs.reshape(-1, probs.shape[-1])
+        n_valid = flat_mask.sum()
+        
+        # Avoid graph breaks from dynamic control flow
+        sum_probs = (flat_probs * flat_mask.unsqueeze(-1)).sum(dim=0)
+        return sum_probs / n_valid.clamp(min=1)
+
+    def _entropy_loss(self, probs: torch.Tensor, token_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+        P_mean = self._mean_probs(probs, token_mask)
+        entropy = -(P_mean * (P_mean + 1e-9).log()).sum()
+        return -self.entropy_reg * entropy
+
+    def _z_loss(self, logits: torch.Tensor, token_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+        if token_mask is None:
+            return self.z_loss_coeff * (logits ** 2).mean()
+
+        flat_mask = token_mask.reshape(-1)
+        flat_logits = logits.reshape(-1, logits.shape[-1])
+        n_valid = flat_mask.sum()
+        
+        # Avoid graph breaks from dynamic control flow
+        sum_z = ((flat_logits ** 2) * flat_mask.unsqueeze(-1)).sum()
+        return self.z_loss_coeff * (sum_z / n_valid.clamp(min=1))
+
+    def extra_repr(self) -> str:
+        return (f"experts={self.num_experts}, top_k={self.top_k}, "
+                f"loss={self.loss_type}, noise={self.noise_std}, "
+                f"entropy_reg={self.entropy_reg}")

architecture/sparse_moe/stage_runner.py

@@ -0,0 +1,65 @@
+import os
+import glob
+import re
+from typing import Tuple, Optional
+
+def find_latest_checkpoint(checkpoint_dir: str) -> Tuple[Optional[int], Optional[str]]:
+    if not os.path.isdir(checkpoint_dir):
+        return None, None
+        
+    epoch_files = glob.glob(os.path.join(checkpoint_dir, "moe-epoch_*.pt"))
+    if not epoch_files:
+        return None, None
+            
+    best_epoch = -1
+    best_file = None
+    
+    for f in epoch_files:
+        m = re.search(r"moe-epoch_(\d+)\.pt", f)
+        if m:
+            ep = int(m.group(1))
+            if ep > best_epoch:
+                best_epoch = ep
+                best_file = f
+                
+    if best_epoch >= 0:
+        return best_epoch, best_file
+    return None, None
+
+def download_stage_from_hub(repo_id: str, remote_path: str, local_dir: str, token: str = None) -> bool:
+    from huggingface_hub import snapshot_download
+    token = token or os.environ.get("HF_TOKEN")
+    try:
+        snapshot_download(
+            repo_id=repo_id,
+            repo_type="dataset",
+            allow_patterns=f"{remote_path}/moe-*.pt",
+            local_dir=local_dir,
+            local_dir_use_symlinks=False,
+            token=token
+        )
+        return True
+    except Exception as e:
+        print(f"[StageRunner] Could not download {remote_path} from Hub: {e}")
+        return False
+
+def get_resume_epoch(checkpoint_dir: str, hub_repo_id: str, stage_name: str) -> Optional[int]:
+    epoch, local_path = find_latest_checkpoint(checkpoint_dir)
+    
+    if epoch is not None:
+        print(f"[StageRunner] Found local checkpoint at epoch {epoch}")
+        return epoch
+        
+    print(f"[StageRunner] No local checkpoint in {checkpoint_dir}. Checking Hub...")
+   
+    remote_path = f"{stage_name}" 
+    success = download_stage_from_hub(hub_repo_id, remote_path, ".", token=None)
+    
+    if success:
+        epoch, local_path = find_latest_checkpoint(checkpoint_dir)
+        if epoch is not None:
+            print(f"[StageRunner] Downloaded checkpoint from Hub and resuming from epoch {epoch}")
+            return epoch
+            
+    print("[StageRunner] No checkpoints found locally or on Hub. Starting from scratch.")
+    return None

architecture/sparse_moe/sys_profiler.py

@@ -0,0 +1,157 @@
+import os
+import sys
+import time
+import json
+import threading
+import subprocess
+from datetime import datetime
+from typing import Dict, List, Optional, Any
+
+import torch
+
+
+class SystemEnvironment:
+    
+    @staticmethod
+    def get_nvidia_smi_info() -> str:
+        try:
+            result = subprocess.run(
+                ["nvidia-smi"], 
+                capture_output=True, 
+                text=True, 
+                timeout=5
+            )
+            return result.stdout if result.returncode == 0 else "nvidia-smi returned an error."
+        except Exception as e:
+            return f"nvidia-smi not available or failed: {e}"
+
+    @staticmethod
+    def get_pytorch_info() -> Dict[str, Any]:
+        return {
+            "python_version": sys.version,
+            "torch_version": torch.__version__,
+            "cuda_available": torch.cuda.is_available(),
+            "cuda_version": torch.version.cuda if torch.cuda.is_available() else None,
+            "cudnn_version": torch.backends.cudnn.version() if torch.cuda.is_available() else None,
+            "alloc_conf": os.environ.get("PYTORCH_CUDA_ALLOC_CONF", "Not Set"),
+        }
+
+    @staticmethod
+    def generate_system_report() -> str:
+        report = []
+        report.append(f"Date: {datetime.utcnow().strftime('%a %b %d %H:%M:%S UTC %Y')}")
+        report.append(f"Sys Platform: {sys.platform}")
+        report.append("\n=== PyTorch Info ===")
+        pt_info = SystemEnvironment.get_pytorch_info()
+        for k, v in pt_info.items():
+            report.append(f"{k}: {v}")
+            
+        report.append("\n=== NVIDIA-SMI ===")
+        report.append(SystemEnvironment.get_nvidia_smi_info())
+        
+        return "\n".join(report)
+
+
+class ResourceTracker:
+    
+    def __init__(self, polling_interval: float = 0.5, use_cuda_native: bool = True):
+        self.polling_interval = polling_interval
+        self.use_cuda_native = use_cuda_native
+        
+        # State
+        self.is_running = False
+        self._thread: Optional[threading.Thread] = None
+        
+        # Data
+        self.timestamps: List[float] = []
+        self.gpu_utilization: List[float] = []
+        self.memory_allocated_mb: List[float] = []
+        self.memory_reserved_mb: List[float] = []
+        
+        self._start_time: float = 0.0
+
+    def _poll_metrics(self):
+        while self.is_running:
+            current_time = time.time() - self._start_time
+            self.timestamps.append(current_time)
+            
+            # 1. PyTorch CUDA Native memory (very exact per-process)
+            if torch.cuda.is_available():
+                alloc_mb = torch.cuda.memory_allocated() / (1024 ** 2)
+                resrv_mb = torch.cuda.memory_reserved() / (1024 ** 2)
+            else:
+                alloc_mb, resrv_mb = 0.0, 0.0
+                
+            self.memory_allocated_mb.append(alloc_mb)
+            self.memory_reserved_mb.append(resrv_mb)
+            
+            # 2. System Level GPU Utilization (via fast nvidia-smi call)
+            util = 0.0
+            if torch.cuda.is_available():
+                    # Query utilization
+                    res = subprocess.run(
+                        ["nvidia-smi", "--query-gpu=utilization.gpu", "--format=csv,noheader,nounits"],
+                        capture_output=True, text=True
+                    )
+                    if res.returncode == 0:
+                        lines = res.stdout.strip().split('\n')
+                        if lines and lines[0].isdigit():
+                            util = float(lines[0])
+                except Exception:
+                    pass
+            self.gpu_utilization.append(util)
+            
+            time.sleep(self.polling_interval)
+
+    def start(self):
+        if self.is_running:
+            return
+        
+        self.is_running = True
+        self.timestamps.clear()
+        self.gpu_utilization.clear()
+        self.memory_allocated_mb.clear()
+        self.memory_reserved_mb.clear()
+        
+        if torch.cuda.is_available():
+            torch.cuda.reset_peak_memory_stats()
+            
+        self._start_time = time.time()
+        self._thread = threading.Thread(target=self._poll_metrics, daemon=True)
+        self._thread.start()
+
+    def stop(self) -> Dict[str, List[float]]:
+        self.is_running = False
+        if self._thread is not None:
+            self._thread.join(timeout=self.polling_interval * 2)
+            
+        return self.get_results()
+        
+    def get_results(self) -> Dict[str, List[float]]:
+        return {
+            "timestamps_sec": list(self.timestamps),
+            "gpu_utilization_pct": list(self.gpu_utilization),
+            "memory_allocated_mb": list(self.memory_allocated_mb),
+            "memory_reserved_mb": list(self.memory_reserved_mb),
+        }
+        
+    def save_tabular_log(self, filepath: str, label: str = "Workload"):
+        with open(filepath, "w") as f:
+            f.write(f"=== Resource Tracking: {label} ===\n")
+            f.write(f"Polling Interval: {self.polling_interval}s\n")
+            f.write("-" * 65 + "\n")
+            f.write(f"{'Time(s)':>10} | {'GPU Util(%)':>15} | {'Allocated(MB)':>15} | {'Reserved(MB)':>15}\n")
+            f.write("-" * 65 + "\n")
+            
+            for t, u, a, r in zip(self.timestamps, self.gpu_utilization, self.memory_allocated_mb, self.memory_reserved_mb):
+                f.write(f"{t:>10.2f} | {u:>15.1f} | {a:>15.1f} | {r:>15.1f}\n")
+            
+            f.write("-" * 65 + "\n")
+            if self.memory_allocated_mb:
+                f.write(f"Peak Allocated: {max(self.memory_allocated_mb):.1f} MB\n")
+            if self.memory_reserved_mb:
+                f.write(f"Peak Reserved: {max(self.memory_reserved_mb):.1f} MB\n")
+
+if __name__ == "__main__":
+    # Smoke test structure
+    print(SystemEnvironment.generate_system_report())

architecture/sparse_moe/trainer.py

@@ -0,0 +1,362 @@
+from __future__ import annotations
+
+import math
+from pathlib import Path
+from typing import Dict, List, Optional
+
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+
+from .config import TrainingConfig
+from .layers import SparseMoELayer
+from .utils import build_labels
+
+
+
+# LR schedule
+def _cosine_with_warmup(step: int, warmup: int, total: int) -> float:
+    if step < warmup:
+        return step / max(warmup, 1)
+    progress = (step - warmup) / max(total - warmup, 1)
+    return 0.5 * (1.0 + math.cos(math.pi * progress))
+
+
+def _build_amp(device: torch.device, fp16: bool) -> tuple[bool, torch.dtype, torch.amp.GradScaler]:
+    use_amp = fp16 and device.type == "cuda"
+    amp_dtype = torch.float16
+    if use_amp and torch.cuda.is_bf16_supported():
+        amp_dtype = torch.bfloat16
+    scaler = torch.amp.GradScaler(device.type, enabled=use_amp and amp_dtype == torch.float16)
+    return use_amp, amp_dtype, scaler
+
+
+# Trainer
+class MoETrainer:
+
+    def __init__(
+        self,
+        model: nn.Module,
+        train_loader: DataLoader,
+        val_loader: Optional[DataLoader],
+        config: TrainingConfig,
+    ):
+        self.model = model
+        self.train_loader = train_loader
+        self.val_loader = val_loader
+        self.config = config
+        self.start_epoch = 1
+
+        # Device
+        self.device = next(model.parameters()).device
+
+        trainable = [p for p in model.parameters() if p.requires_grad]
+        if not trainable:
+            raise ValueError(
+                "MoETrainer received a model with no trainable parameters. "
+                "Check that MoE injection succeeded for this architecture."
+            )
+        self.optimizer = torch.optim.AdamW(trainable, lr=config.lr, weight_decay=0.01)
+
+        self.use_amp, self._amp_dtype, self.scaler = _build_amp(self.device, config.fp16)
+        if self.use_amp and self._amp_dtype == torch.bfloat16:
+            print("Using BF16 mixed precision (Ampere+/Blackwell detected)")
+
+        if self.device.type == "cuda":
+            torch.set_float32_matmul_precision("high")
+            torch.backends.cudnn.benchmark = True
+            print("Enabled TF32 matmuls + cuDNN benchmark")
+
+        self._moe_layers: list[SparseMoELayer] = [
+            m for m in model.modules() if isinstance(m, SparseMoELayer)
+        ]
+
+        # Step tracking
+        self.global_step = 0
+        self._total_steps = self._estimate_total_steps()
+
+        self.scheduler = torch.optim.lr_scheduler.LambdaLR(
+            self.optimizer,
+            lr_lambda=lambda step: _cosine_with_warmup(
+                step, config.warmup_steps, self._total_steps
+            ),
+        )
+
+        # No external logging — all metrics tracked in the returned history dict
+
+        print(
+            f"MoETrainer: {sum(p.numel() for p in trainable):,} trainable params, "
+            f"AMP={'on' if self.use_amp else 'off'}, "
+            f"effective_batch={config.batch_size * config.gradient_accumulation_steps}, "
+            f"total_steps≈{self._total_steps}"
+        )
+
+        # Initialise router entropy_reg to step 0 value
+        self._anneal_entropy()
+
+    # Public API
+    def train(self) -> Dict[str, List[float]]:
+        history: Dict[str, List[float]] = {
+            "train_loss": [], "val_loss": [], "aux_loss": [], "lr": [],
+            "epoch_boundaries": [],  # global step at start of each epoch
+        }
+
+        best_val_loss = float("inf")
+        patience = 1
+        patience_counter = 0
+        best_state = None
+
+        for epoch in range(self.start_epoch, self.config.num_epochs + 1):
+            print(f"═══ Epoch {epoch}/{self.config.num_epochs} ═══")
+            history["epoch_boundaries"].append(self.global_step)
+            self._train_epoch(history)
+            self.save_checkpoint(f"epoch_{epoch}")
+
+            if self.val_loader is not None:
+                val_loss = self.evaluate()
+                history["val_loss"].append(val_loss)
+                print(f"  val_loss={val_loss:.4f}")
+
+                if val_loss < best_val_loss:
+                    best_val_loss = val_loss
+                    patience_counter = 0
+                    # Save best trainable weights
+                    best_state = {
+                        name: param.clone()
+                        for name, param in self.model.named_parameters()
+                        if param.requires_grad
+                    }
+                    self.save_checkpoint("best")
+                    print(f"  New best val_loss={best_val_loss:.4f} (Saved moe-best.pt)")
+                else:
+                    patience_counter += 1
+                    print(
+                        f"  val_loss did not improve "
+                        f"(best={best_val_loss:.4f}, patience={patience_counter}/{patience})"
+                    )
+                    if patience_counter >= patience:
+                        print("  Early stopping triggered — restoring best weights.")
+                        break
+        
+        # After loop completion, ensure best state is restored 
+        if best_state is not None:
+            for name, param in self.model.named_parameters():
+                if name in best_state:
+                    param.data.copy_(best_state[name])
+
+        print("Training complete.")
+        return history
+
+    @torch.no_grad()
+    def evaluate(self) -> float:
+        if self.val_loader is None:
+            return float("nan")
+
+        self.model.eval()
+        total_loss, total_tokens = 0.0, 0
+
+        for batch in self.val_loader:
+            batch = {k: v.to(self.device, non_blocking=True) for k, v in batch.items()}
+            labels = build_labels(batch)
+
+            with torch.amp.autocast(self.device.type, dtype=self._amp_dtype, enabled=self.use_amp):
+                outputs = self.model(**batch, labels=labels)
+
+            valid_tokens = int((labels[:, 1:] != -100).sum().item())
+            total_loss += outputs.loss.item() * max(valid_tokens, 1)
+            total_tokens += valid_tokens
+
+        self.model.train()
+        return total_loss / max(total_tokens, 1)
+
+    # Training epoch
+    def _train_epoch(self, history: Dict[str, List[float]]):
+        self.model.train()
+        accum_steps = self.config.gradient_accumulation_steps
+        epoch_loss = 0.0
+        epoch_steps = 0
+
+        self.optimizer.zero_grad(set_to_none=True)
+
+        for step_in_epoch, batch in enumerate(self.train_loader):
+            batch = {k: v.to(self.device, non_blocking=True) for k, v in batch.items()}
+            labels = build_labels(batch)
+
+            # ── Forward pass (AMP context) ──
+            with torch.amp.autocast(self.device.type, dtype=self._amp_dtype, enabled=self.use_amp):
+                outputs = self.model(**batch, labels=labels)
+                lm_loss = outputs.loss / accum_steps
+
+                # Auxiliary load-balance loss
+                aux_loss = self._collect_aux_loss() / accum_steps
+                total_loss = lm_loss + self.config.aux_loss_weight * aux_loss
+
+            # ── Backward pass ──
+            self.scaler.scale(total_loss).backward()
+
+            # ── Optimizer step (every accum_steps micro-batches) ──
+            is_accum_boundary = (step_in_epoch + 1) % accum_steps == 0 or (
+                step_in_epoch + 1
+            ) == len(self.train_loader)
+
+            # Accumulate per-micro-batch stats (not yet logged — done at step boundary)
+            raw_lm_loss = lm_loss.item() * accum_steps
+            raw_aux_loss = aux_loss.item() * accum_steps
+            epoch_loss += raw_lm_loss
+            epoch_steps += 1
+
+            if is_accum_boundary:
+                self.scaler.unscale_(self.optimizer)
+                nn.utils.clip_grad_norm_(
+                    [p for p in self.model.parameters() if p.requires_grad],
+                    self.config.max_grad_norm,
+                )
+                
+                old_scale = self.scaler.get_scale()
+                self.scaler.step(self.optimizer)
+                self.scaler.update()
+                
+                self.scheduler.step()
+                    
+                self.optimizer.zero_grad(set_to_none=True)
+                self.global_step += 1
+
+                # ── Entropy annealing ──
+                self._anneal_entropy()
+
+                history["lr"].append(current_lr)
+
+                if self.global_step % self.config.log_every_steps == 0:
+                    avg = epoch_loss / max(epoch_steps, 1)
+                    print(
+                        f"  step {self.global_step:>5d} | "
+                        f"loss={raw_lm_loss:.4f} | "
+                        f"aux={raw_aux_loss:.4f} | "
+                        f"lr={current_lr:.2e} | "
+                        f"avg={avg:.4f}"
+                        + (f" | mem={torch.cuda.max_memory_allocated()/1e9:.1f}GB"
+                           if torch.cuda.is_available() else "")
+                    )
+
+                # ── Checkpoint (once per optimizer step) ──
+                if (
+                    self.config.save_every_steps > 0
+                    and self.global_step % self.config.save_every_steps == 0
+                ):
+                    self.save_checkpoint(f"step-{self.global_step}")
+
+    # Helpers
+    def _collect_aux_loss(self) -> torch.Tensor:
+        total = torch.tensor(0.0, device=self.device)
+        count = 0
+        for moe in self._moe_layers:
+            loss = getattr(moe, "_last_aux_loss", None)
+            if loss is not None:
+                total = total + loss
+                count += 1
+        return total / max(count, 1)
+
+    def _anneal_entropy(self):
+        if self._total_steps <= 0:
+            return
+        progress = min(self.global_step / self._total_steps, 1.0)
+        new_reg = self.config.entropy_reg_start * 0.5 * (
+            1.0 + math.cos(math.pi * progress)
+        )
+        for moe in self._moe_layers:
+            moe.routing.entropy_reg = new_reg
+
+    def _estimate_total_steps(self) -> int:
+        batches = len(self.train_loader)
+        accum = self.config.gradient_accumulation_steps
+        steps_per_epoch = max(math.ceil(batches / accum), 1)
+        return steps_per_epoch * self.config.num_epochs
+
+    # Checkpointing
+    def save_checkpoint(self, tag: str):
+        out_dir = Path(self.config.output_dir)
+        out_dir.mkdir(parents=True, exist_ok=True)
+
+        trainable_keys = {
+            k for k, p in self.model.named_parameters() if p.requires_grad
+        }
+        state = {
+            k: v for k, v in self.model.state_dict().items()
+            if k in trainable_keys
+        }
+
+        path = out_dir / f"moe-{tag}.pt"
+        torch.save({
+            "model_state": state,
+            "optimizer_state": self.optimizer.state_dict(),
+            "scheduler_state": self.scheduler.state_dict(),
+            "scaler_state": self.scaler.state_dict(),
+            "global_step": self.global_step,
+        }, path)
+        print(f"Checkpoint saved: {path}")
+
+    def load_checkpoint(self, path: str):
+        checkpoint = torch.load(path, map_location=self.device, weights_only=True)
+        state = checkpoint.get("model_state", checkpoint)
+
+        # Validate keys
+        trainable_keys = {
+            k for k, p in self.model.named_parameters() if p.requires_grad
+        }
+        incoming_keys = set(state.keys())
+
+        matched = trainable_keys & incoming_keys
+        missing = trainable_keys - incoming_keys
+        unexpected = incoming_keys - trainable_keys
+
+        incompatible = []
+        for key in matched:
+            checkpoint_shape = state[key].shape
+            model_shape = self.model.state_dict()[key].shape
+            if checkpoint_shape != model_shape:
+                incompatible.append(key)
+
+        if incompatible:
+            print(
+                f"[WARNING] Checkpoint {path} has {len(incompatible)} parameters with "
+                f"mismatching shapes (e.g. num_experts change). "
+                f"Skipping incompatible checkpoint — training will start from scratch."
+            )
+            return
+
+        if missing:
+            print(f"[WARNING] Missing keys in checkpoint: {missing}")
+        if unexpected:
+            print(f"[WARNING] Unexpected keys in checkpoint: {unexpected}")
+
+        current_state = self.model.state_dict()
+        for key in matched:
+            current_state[key].copy_(state[key])
+
+        # Restore optimizer + step
+        if "optimizer_state" in checkpoint:
+            try:
+                self.optimizer.load_state_dict(checkpoint["optimizer_state"])
+            except (ValueError, RuntimeError) as e:
+                print(f"[WARNING] Could not restore optimizer state: {e}. "
+                      f"Optimizer will be re-initialized.")
+        if "scheduler_state" in checkpoint:
+            self.scheduler.load_state_dict(checkpoint["scheduler_state"])
+        if "scaler_state" in checkpoint:
+            self.scaler.load_state_dict(checkpoint["scaler_state"])
+        if "global_step" in checkpoint:
+            self.global_step = checkpoint["global_step"]
+        # Keep router entropy annealing aligned with the restored step.
+        self._anneal_entropy()
+
+        print(f"Checkpoint loaded: {path} (step {self.global_step})")
+
+    def _sync_scheduler_from_global_step(self):
+        multiplier = _cosine_with_warmup(
+            self.global_step,
+            self.config.warmup_steps,
+            self._total_steps,
+        )
+        self.scheduler.last_epoch = self.global_step
+        for group, base_lr in zip(self.optimizer.param_groups, self.scheduler.base_lrs):
+            group["lr"] = base_lr * multiplier

architecture/sparse_moe/utils.py

@@ -0,0 +1,51 @@
+from __future__ import annotations
+
+import random
+from typing import Any, Dict, List
+
+import torch
+import numpy as np
+
+
+# Label building
+def build_labels(batch: Dict[str, torch.Tensor]) -> torch.Tensor:
+    labels = batch["input_ids"].clone()  # (batch_size, seq_len)
+
+    if "attention_mask" in batch:
+        labels[batch["attention_mask"] == 0] = -100
+
+    return labels  # (batch_size, seq_len)
+
+
+# Reproducibility
+def set_seed(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+        torch.backends.cudnn.benchmark = True
+
+
+# Serialisation
+def to_serializable(value: Any) -> Any:
+    if isinstance(value, torch.Tensor):
+        return value.detach().cpu().tolist()
+    if isinstance(value, np.ndarray):
+        return value.tolist()
+    if isinstance(value, np.integer):
+        return int(value)
+    if isinstance(value, np.floating):
+        return float(value)
+    if isinstance(value, dict):
+        return {k: to_serializable(v) for k, v in value.items()}
+    if isinstance(value, (list, tuple)):
+        return [to_serializable(v) for v in value]
+    if hasattr(value, "__dataclass_fields__"):
+        return {k: to_serializable(v) for k, v in value.__dict__.items()}
+    return value
+
+
+# Multi-seed helpers
+def parse_seed_list(raw: str) -> List[int]:
+    return [int(s.strip()) for s in raw.split(",") if s.strip()]

architecture/sparse_moe/visualization.py

@@ -0,0 +1,278 @@
+from __future__ import annotations
+
+import warnings
+from typing import List, Optional, Tuple, Union, Dict, Any
+from pathlib import Path
+import numpy as np
+
+try:
+    import matplotlib.pyplot as plt
+    import matplotlib.colors as mcolors
+    from matplotlib.gridspec import GridSpec
+    HAS_MATPLOTLIB = True
+except ImportError:
+    HAS_MATPLOTLIB = False
+    warnings.warn(
+        "Matplotlib not installed. Visualization features unavailable. "
+        "Install with: pip install matplotlib"
+    )
+
+def _check_matplotlib():
+    if not HAS_MATPLOTLIB:
+        raise RuntimeError(
+            "Matplotlib is required for visualization. "
+            "Install with: pip install matplotlib"
+        )
+
+
+class KeiroPalette:
+    
+    # Primary line colors
+    PRIMARY = {
+        "dense": "#E74C3C",     # Red - Dense baseline ("Before")
+        "moe": "#3498DB",       # Blue - Sparse MoE ("After")
+        "expert_avg": "#9B59B6",# Purple - Expert average
+        "memory": "#2ECC71",    # Green - Memory bounds
+    }
+    
+    # Generative expert spectrum for heatmap/routing
+    EXPERTS = [
+        "#3498DB", "#2980B9", "#1ABC9C", "#27AE60", 
+        "#F39C12", "#D35400", "#E74C3C", "#8E44AD"
+    ]
+    
+    BG_LIGHT = "#FAFAFA"
+    BG_DARK = "#1A1A2E"
+    GRID_LIGHT = "#E0E0E0"
+    GRID_DARK = "#2D2D44"
+
+
+class BasePlot:
+    
+    def __init__(self, figsize=(12, 8), dpi=150, theme="dark", title=None):
+        _check_matplotlib()
+        self.figsize = figsize
+        self.dpi = dpi
+        self.theme = theme
+        self.fig, self.ax = plt.subplots(figsize=figsize)
+        self._apply_theme()
+        if title:
+            self.ax.set_title(title, fontsize=16, fontweight='bold', pad=20)
+            
+    def _apply_theme(self):
+        bg = KeiroPalette.BG_DARK if self.theme == "dark" else KeiroPalette.BG_LIGHT
+        grid = KeiroPalette.GRID_DARK if self.theme == "dark" else KeiroPalette.GRID_LIGHT
+        fg = 'white' if self.theme == "dark" else 'black'
+        
+        self.fig.patch.set_facecolor(bg)
+        self.ax.set_facecolor(bg)
+        self.ax.tick_params(colors=fg)
+        self.ax.xaxis.label.set_color(fg)
+        self.ax.yaxis.label.set_color(fg)
+        self.ax.title.set_color(fg)
+        self.ax.grid(True, alpha=0.2, color=grid)
+        for spine in self.ax.spines.values():
+            spine.set_color(grid)
+
+    def save(self, filepath: Union[str, Path]):
+        filepath = Path(filepath)
+        filepath.parent.mkdir(parents=True, exist_ok=True)
+        self.fig.savefig(filepath, dpi=self.dpi, bbox_inches="tight", facecolor=self.fig.get_facecolor())
+    
+    def close(self):
+        plt.close(self.fig)
+
+
+class ResourceUtilizationPlot(BasePlot):
+    
+    def __init__(self, title="Resource Utilization (Before vs After)", **kwargs):
+        super().__init__(title=title, **kwargs)
+        self.ax.set_xlabel("Time (seconds)", fontsize=12)
+        self.ax2 = self.ax.twinx()
+        self.ax.set_ylabel("Memory Allocated (MB)", fontsize=12)
+        self.ax2.set_ylabel("GPU Utilization (%)", fontsize=12)
+        
+        if self.theme == "dark":
+            self.ax2.tick_params(colors='white')
+            self.ax2.yaxis.label.set_color('white')
+            for spine in self.ax2.spines.values():
+                spine.set_color(KeiroPalette.GRID_DARK)
+
+    def add_trace(self, time_sec: List[float], values: List[float], label: str, metric: str = "memory"):
+        color = KeiroPalette.PRIMARY["dense"] if "Before" in label or "Dense" in label else KeiroPalette.PRIMARY["moe"]
+        linestyle = "-" if metric == "memory" else "--"
+        axis = self.ax if metric == "memory" else self.ax2
+        
+        axis.plot(
+            time_sec, values, label=label,
+            color=color, linestyle=linestyle, linewidth=2.5, alpha=0.8
+        )
+        
+    def finalize(self):
+        lines1, labels1 = self.ax.get_legend_handles_labels()
+        lines2, labels2 = self.ax2.get_legend_handles_labels()
+        self.ax2.legend(lines1 + lines2, labels1 + labels2, loc="best", framealpha=0.8)
+        plt.tight_layout()
+
+
+class KeiroDashboard:
+    
+    def __init__(self, figsize=(18, 12), dpi=150, theme="dark"):
+        _check_matplotlib()
+        self.dpi = dpi
+        self.theme = theme
+        self.fig, self.axes = plt.subplots(2, 2, figsize=figsize)
+        self._apply_theme()
+        
+    def _apply_theme(self):
+        bg = KeiroPalette.BG_DARK if self.theme == "dark" else KeiroPalette.BG_LIGHT
+        grid = KeiroPalette.GRID_DARK if self.theme == "dark" else KeiroPalette.GRID_LIGHT
+        fg = 'white' if self.theme == "dark" else 'black'
+        self.fig.patch.set_facecolor(bg)
+        
+        for ax in self.axes.flat:
+            ax.set_facecolor(bg)
+            ax.tick_params(colors=fg)
+            ax.xaxis.label.set_color(fg)
+            ax.yaxis.label.set_color(fg)
+            ax.title.set_color(fg)
+            ax.grid(True, alpha=0.2, color=grid)
+            for spine in ax.spines.values():
+                spine.set_color(grid)
+
+    def plot_memory_scaling(self, ax_idx=(0,0), seq_lens=None, data_dict=None):
+        ax = self.axes[ax_idx]
+        ax.set_title("Peak Memory vs Sequence Length", fontsize=14, fontweight='bold')
+        ax.set_xlabel("Sequence Length")
+        ax.set_ylabel("Memory (MB)")
+        if seq_lens and data_dict:
+            for k, v in data_dict.items():
+                color = KeiroPalette.PRIMARY["dense"] if "Dense" in k else KeiroPalette.PRIMARY["moe"]
+                ax.plot(seq_lens, v, label=k, color=color, marker='o', linewidth=2)
+            ax.legend()
+            
+    def plot_throughput(self, ax_idx=(0,1), seq_lens=None, data_dict=None):
+        ax = self.axes[ax_idx]
+        ax.set_title("Inference Throughput (tokens/sec)", fontsize=14, fontweight='bold')
+        ax.set_xlabel("Sequence Length")
+        ax.set_ylabel("Throughput")
+        if seq_lens and data_dict:
+            for k, v in data_dict.items():
+                color = KeiroPalette.PRIMARY["dense"] if "Dense" in k else KeiroPalette.PRIMARY["moe"]
+                ax.plot(seq_lens, v, label=k, color=color, marker='s', linewidth=2)
+            ax.legend()
+
+    def plot_expert_load(self, ax_idx=(1,0), expert_distribution=None):
+        ax = self.axes[ax_idx]
+        ax.set_title("MoE Expert Load Balancing", fontsize=14, fontweight='bold')
+        ax.set_xlabel("Expert ID")
+        ax.set_ylabel("Tokens Assigned (%)")
+        if expert_distribution:
+            x = np.arange(len(expert_distribution))
+            colors = [KeiroPalette.EXPERTS[i % len(KeiroPalette.EXPERTS)] for i in x]
+            total = sum(expert_distribution)
+            pcts = [100.0 * c / total for c in expert_distribution] if total > 0 else expert_distribution
+            ax.bar(x, pcts, color=colors, alpha=0.8)
+            ax.set_xticks(x)
+            ax.set_xticklabels([f"E{i}" for i in x])
+            ax.axhline(100.0 / len(expert_distribution), color='gray', linestyle='--', label='Perfect Balance')
+            ax.legend()
+            
+    def plot_speedup(self, ax_idx=(1,1), seq_lens=None, base_time=None, moe_time=None):
+        ax = self.axes[ax_idx]
+        ax.set_title("MoE Speedup vs Dense", fontsize=14, fontweight='bold')
+        ax.set_xlabel("Sequence Length")
+        ax.set_ylabel("Speedup (x)")
+        ax.axhline(1.0, color='gray', linestyle='--', alpha=0.5)
+        if seq_lens and base_time and moe_time:
+            speedups = [b/m if m > 0 else 0 for b, m in zip(base_time, moe_time)]
+            ax.plot(seq_lens, speedups, color=KeiroPalette.PRIMARY["expert_avg"], marker='D', linewidth=2, label="Speedup")
+            ax.legend()
+
+    def save(self, filepath: Union[str, Path]):
+        filepath = Path(filepath)
+        filepath.parent.mkdir(parents=True, exist_ok=True)
+        plt.tight_layout()
+        self.fig.savefig(filepath, dpi=self.dpi, bbox_inches="tight", facecolor=self.fig.get_facecolor())
+    
+    def close(self):
+        plt.close(self.fig)
+
+class ColorPalette(KeiroPalette):
+    pass
+
+class DomainScorePlot(BasePlot):
+    def __init__(self, figsize=(10, 6), **kwargs):
+        super().__init__(figsize=figsize, title="Per-Domain Perplexity", **kwargs)
+
+    def plot_comparison(self, rows: List[Dict], include_dense: bool = False):
+        if not rows: return
+        domains = [r["domain"] for r in rows]
+        before = [r["ppl_before"] for r in rows]
+        after = [r["ppl_after"] for r in rows]
+        x = np.arange(len(domains))
+        width = 0.35 if not include_dense else 0.25
+        self.ax.bar(x - width/2, before, width, label='Before (Dense)', color=KeiroPalette.PRIMARY["dense"])
+        self.ax.bar(x + width/2, after, width, label='After (MoE)', color=KeiroPalette.PRIMARY["moe"])
+        if include_dense:
+            dense = [r.get("ppl_dense", 0) for r in rows]
+            self.ax.bar(x + 1.5*width, dense, width, label='Dense Baseline', color=KeiroPalette.PRIMARY["expert_avg"])
+        self.ax.set_xticks(x)
+        self.ax.set_xticklabels(domains, rotation=45, ha='right')
+        self.ax.set_ylabel("Perplexity (Lower is better)")
+        self.ax.legend()
+        self.fig.tight_layout()
+
+class TrainingConvergencePlot(BasePlot):
+    def __init__(self, figsize=(10, 6), **kwargs):
+        super().__init__(figsize=figsize, title="Training Convergence", **kwargs)
+
+    def plot_history(self, history: Dict):
+        train_loss = history.get("train_loss", [])
+        val_loss = history.get("val_loss", [])
+        if train_loss:
+            self.ax.plot(train_loss, label="Train Loss", color=KeiroPalette.PRIMARY["dense"])
+        if val_loss:
+            if len(val_loss) < len(train_loss):
+                x_val = np.linspace(0, len(train_loss)-1, len(val_loss))
+                self.ax.plot(x_val, val_loss, label="Val Loss", marker='o', color=KeiroPalette.PRIMARY["moe"])
+            else:
+                self.ax.plot(val_loss, label="Val Loss", color=KeiroPalette.PRIMARY["moe"])
+        self.ax.set_xlabel("Steps (or Epochs)")
+        self.ax.set_ylabel("Cross Entropy Loss")
+        self.ax.legend()
+
+class ExpertRoutingHeatmap(BasePlot):
+    def __init__(self, figsize=(12, 8), **kwargs):
+        super().__init__(figsize=figsize, title="Expert Routing by Domain", **kwargs)
+
+    def plot_routing(self, spec_dict: Dict):
+        affinity = spec_dict.get("affinity")
+        domains = spec_dict.get("domains")
+        labels = spec_dict.get("expert_labels")
+
+        if affinity is None or domains is None:
+            return
+            
+        # Convert torch tensor to numpy
+        if hasattr(affinity, "cpu"):
+            matrix = affinity.cpu().numpy()
+        else:
+            matrix = np.array(affinity)
+            
+        # Plotting
+        im = self.ax.imshow(matrix, aspect="auto", cmap="viridis")
+        self.ax.set_xticks(range(len(domains)))
+        self.ax.set_xticklabels(domains, rotation=45, ha='right')
+        
+        # Only show individual expert labels if there aren't too many
+        if labels and len(labels) <= 64:
+            self.ax.set_yticks(range(len(labels)))
+            self.ax.set_yticklabels(labels, fontsize=6)
+        else:
+            self.ax.set_ylabel(f"{len(labels)} Layer-Experts")
+            self.ax.set_yticks([]) # Hide Y-axis labels for readability if too dense
+
+        self.fig.colorbar(im, ax=self.ax, fraction=0.046, pad=0.04)
+        self.fig.tight_layout()
+