feat: re-apply dynamic MoE expansion, multi-head CoT, plugin architecture, improved MoD

src/architecture.py

@@ -80,10 +80,22 @@ class MorphConfig:
     lora_dropout: float = 0.05
     # MoE
     num_experts: int = 4
+    max_experts: int = 64
     expert_hidden: int = 512
     moe_top_k: int = 2
     moe_aux_weight: float = 0.01
     moe_prune_threshold: float = 0.02
+    moe_expand_threshold: float = 0.15
+    # MoD - Mixture of Depths
+    use_mod: bool = True
+    mod_hidden: int = 128
+    mod_dropout: float = 0.1
+    mod_keep_prob: float = 0.8
+    mod_temperature: float = 1.0
+    mod_temperature_anneal: float = 0.995
+    # multi-head CoT reasoning
+    num_cot_heads: int = 4
+    cot_hidden: int = 256
     # memory
     memory_size: int = 1024
     memory_dim: int = 768
@@ -122,6 +134,8 @@ class MorphConfig:
     draft_layers: int = 2
     # RoPE scaling for extended context
     rope_scaling: Optional[dict] = None
+    # plugin architecture
+    plugin_dir: Optional[str] = None
     # training
     max_seq_len: int = 2048
     # quantization
@@ -387,6 +401,40 @@ class MultiStepReasoner(nn.Module):
         return out, scratch[-1]
 
 
+class MultiHeadCoT(nn.Module):
+    """Multi-head chain-of-thought reasoning: generates N parallel reasoning paths
+    and fuses them for higher accuracy on complex tasks."""
+
+    def __init__(self, config: MorphConfig, hidden_dim: int):
+        super().__init__()
+        self.num_heads = config.num_cot_heads
+        cot_dim = config.cot_hidden
+        self.heads = nn.ModuleList([
+            nn.Sequential(
+                nn.Linear(hidden_dim, cot_dim),
+                nn.GELU(),
+                nn.LayerNorm(cot_dim),
+                nn.Linear(cot_dim, hidden_dim),
+            ) for _ in range(self.num_heads)
+        ])
+        self.fusion = nn.Sequential(
+            nn.Linear(hidden_dim * (self.num_heads + 1), hidden_dim),
+            nn.GELU(),
+            nn.LayerNorm(hidden_dim),
+            nn.Linear(hidden_dim, hidden_dim),
+        )
+        nn.init.zeros_(self.fusion[-1].weight)
+        nn.init.zeros_(self.fusion[-1].bias)
+
+    def forward(self, hidden: torch.Tensor) -> torch.Tensor:
+        B, T, H = hidden.shape
+        paths = [hidden]
+        for head in self.heads:
+            paths.append(head(hidden))
+        fused = self.fusion(torch.cat(paths, dim=-1))
+        return hidden + fused  # Residual connection
+
+
 # ---------------------------------------------------------------------------
 # Code structure awareness
 # ---------------------------------------------------------------------------
@@ -539,27 +587,32 @@ class QuantizedMemoryModule(MemoryModule):
 
 class MixtureOfDepths(nn.Module):
     """MoD: per-token gating to dynamically skip transformer layers.
-    Uses a lightweight router to predict keep_prob per token, reducing
-    compute by ~30-50% with minimal accuracy loss.
+    Uses a lightweight router with temperature annealing for adaptive layer skipping,
+    reducing compute by ~30-50% with minimal accuracy loss.
     """
 
-    def __init__(self, hidden_dim: int, mod_hidden: int, keep_prob: float = 0.8, dropout: float = 0.1):
+    def __init__(self, hidden_dim: int, mod_hidden: int, keep_prob: float = 0.8, 
+                 dropout: float = 0.1, temperature: float = 1.0, temperature_anneal: float = 0.995):
         super().__init__()
         self.keep_prob = keep_prob
+        self.temperature = temperature
+        self.temperature_anneal = temperature_anneal
         self.router = nn.Sequential(
             nn.Linear(hidden_dim, mod_hidden),
             nn.GELU(),
             nn.LayerNorm(mod_hidden),
             nn.Linear(mod_hidden, 1),
-            nn.Sigmoid(),
         )
         self.dropout = nn.Dropout(dropout)
 
     def forward(self, hidden: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
         B, T, H = hidden.shape
-        probs = self.router(hidden.detach())  # (B, T, 1)
+        logits = self.router(hidden.detach())  # (B, T, 1)
+        probs = torch.sigmoid(logits / self.temperature)
         mask = torch.bernoulli(probs * 0.5 + self.keep_prob * 0.5).expand_as(hidden)
         mask = self.dropout(mask)
+        if self.training:
+            self.temperature = max(0.1, self.temperature * self.temperature_anneal)
         return hidden * mask, probs
 
 
@@ -610,20 +663,23 @@ class MemoryEfficientAttention(nn.Module):
 
 
 class DynamicMoEBlock(nn.Module):
-    """Sparse MoE with dynamic expert pruning and load-balancing."""
+    """Sparse MoE with dynamic expert expansion, pruning, and load-balancing."""
 
     def __init__(self, hidden_dim: int, num_experts: int, expert_hidden: int,
-                 top_k: int, prune_threshold: float = 0.02):
+                 top_k: int, prune_threshold: float = 0.02, expand_threshold: float = 0.15, max_experts: int = 64):
         super().__init__()
         self.num_experts = num_experts
+        self.max_experts = max_experts
         self.top_k = top_k
         self.prune_threshold = prune_threshold
+        self.expand_threshold = expand_threshold
         self.gate = nn.Linear(hidden_dim, num_experts, bias=False)
         self.experts = nn.ModuleList([
             Expert(hidden_dim, expert_hidden) for _ in range(num_experts)
         ])
         self.expert_usage = torch.zeros(num_experts)
         self._pruned = set()
+        self._expansion_count = 0
 
     def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
         B, T, H = x.shape
@@ -643,7 +699,8 @@ class DynamicMoEBlock(nn.Module):
             sel = routing[:, i] > 0
             if sel.any():
                 out[sel] += routing[sel, i].unsqueeze(-1) * expert(flat[sel])
-            self.expert_usage[i] += sel.sum().item()
+            if i < len(self.expert_usage):
+                self.expert_usage[i] += sel.sum().item()
 
         f_i = routing.mean(0)
         P_i = probs.mean(0)
@@ -651,17 +708,40 @@ class DynamicMoEBlock(nn.Module):
 
         return out.view(B, T, H), aux
 
-    def prune_experts(self):
-        """Prune experts with usage below threshold (called periodically during training)."""
+    def prune_and_expand_experts(self):
+        """Dynamically prune underused experts and clone overused ones."""
         total = self.expert_usage.sum()
         if total == 0:
+            self.expert_usage.zero_()
             return
+        
         usage_ratios = self.expert_usage / total
-        for i, ratio in enumerate(usage_ratios):
-            if ratio < self.prune_threshold and len(self._pruned) < self.num_experts - 1:
+        active_experts = [i for i in range(len(self.experts)) if i not in self._pruned]
+        
+        for i in active_experts:
+            if usage_ratios[i] < self.prune_threshold and len(self._pruned) < len(self.experts) - 1:
                 self._pruned.add(i)
-                print(f"Pruned expert {i} (usage ratio {ratio:.4f})")
+                print(f"Pruned expert {i} (usage {usage_ratios[i]:.4f})")
+        
+        if len(self.experts) < self.max_experts:
+            avg_usage = usage_ratios[active_experts].mean().item()
+            for i in active_experts:
+                if usage_ratios[i] > self.expand_threshold and len(self.experts) < self.max_experts:
+                    new_expert = Expert(
+                        self.experts[i].in_proj.in_features,
+                        self.experts[i].in_proj.out_features
+                    )
+                    new_expert.load_state_dict(self.experts[i].state_dict())
+                    with torch.no_grad():
+                        for param in new_expert.parameters():
+                            param.add_(torch.randn_like(param) * 0.01)
+                    self.experts.append(new_expert)
+                    self.expert_usage = torch.cat([self.expert_usage, torch.zeros(1)])
+                    self._expansion_count += 1
+                    print(f"Expanded expert {i} -> new expert {len(self.experts)-1}")
+        
         self.expert_usage.zero_()
+        print(f"Active experts: {len(self.experts) - len(self._pruned)}/{len(self.experts)}")
 
 
 class MultimodalFusion(nn.Module):
@@ -954,14 +1034,22 @@ class MorphModel(nn.Module):
         self.verifier = VerifierHead(hidden_dim)
         self.skill_module = SkillTokenModule(self.cfg, hidden_dim)
         self.depth_module = DepthEmbeddings(self.cfg, hidden_dim)
-        self.moe_block = DynamicMoEBlock(hidden_dim, self.cfg.num_experts, self.cfg.expert_hidden, self.cfg.moe_top_k, self.cfg.moe_prune_threshold)
+        self.moe_block = DynamicMoEBlock(
+            hidden_dim, self.cfg.num_experts, self.cfg.expert_hidden,
+            self.cfg.moe_top_k, self.cfg.moe_prune_threshold,
+            self.cfg.moe_expand_threshold, self.cfg.max_experts
+        )
         self.memory = QuantizedMemoryModule(self.cfg.memory_size, self.cfg.memory_dim, hidden_dim, self.cfg.memory_quantize, self.cfg.memory_quant_bits)
-        self.mod = MixtureOfDepths(hidden_dim, self.cfg.mod_hidden, self.cfg.mod_keep_prob, self.cfg.mod_dropout)
+        self.mod = MixtureOfDepths(
+            hidden_dim, self.cfg.mod_hidden, self.cfg.mod_keep_prob, 
+            self.cfg.mod_dropout, self.cfg.mod_temperature, self.cfg.mod_temperature_anneal
+        )
         self.multimodal_fusion = MultimodalFusion(self.cfg, hidden_dim)
         self.tool_use = ToolUseModule(self.cfg, hidden_dim)
         self.document_module = DocumentModule(self.cfg, hidden_dim)
         self.video_module = VideoModule(self.cfg, hidden_dim)
         self.code_sandbox = CodeSandbox(self.cfg.sandbox_timeout, self.cfg.sandbox_max_memory)
+        self.cot_reasoner = MultiHeadCoT(self.cfg, hidden_dim)
 
         # cast novel components to the base model's compute dtype
         self._dtype = self.base_model_raw.model.embed_tokens.weight.dtype
@@ -969,13 +1057,51 @@ class MorphModel(nn.Module):
             self.coordinator, self.reasoner, self.code_bias, self.scratchpad,
             self.verifier, self.skill_module, self.depth_module, self.moe_block,
             self.memory, self.mod, self.multimodal_fusion, self.tool_use,
-            self.document_module, self.video_module,
+            self.document_module, self.video_module, self.code_sandbox, self.cot_reasoner,
         ):
             mod.to(self._dtype)
 
         self.vocab_size = vocab_size
         self.base_model = None
         self._skill_lora_modules: Dict[str, nn.Module] = {}
+        self._plugins: Dict[str, nn.Module] = {}
+
+        # Load plugins from plugin_dir if specified
+        if self.cfg.plugin_dir:
+            self.load_plugins(self.cfg.plugin_dir)
+
+    def load_plugins(self, plugin_dir: str):
+        """Load custom capability plugins from a directory."""
+        import os
+        import importlib.util
+        plugin_path = Path(plugin_dir)
+        if not plugin_path.exists():
+            print(f"Plugin directory not found: {plugin_dir}")
+            return
+        
+        for file in plugin_path.glob("*.py"):
+            if file.name.startswith("_"):
+                continue
+            try:
+                spec = importlib.util.spec_from_file_location(file.stem, file)
+                mod = importlib.util.module_from_spec(spec)
+                spec.loader.exec_module(mod)
+                for attr_name in dir(mod):
+                    attr = getattr(mod, attr_name)
+                    if isinstance(attr, type) and issubclass(attr, nn.Module) and attr is not nn.Module:
+                        plugin_name = getattr(attr, 'plugin_name', attr_name)
+                        plugin_instance = attr(self.cfg, hidden_dim=self.base_model_raw.config.hidden_size)
+                        setattr(self, f"plugin_{plugin_name}", plugin_instance)
+                        self._plugins[plugin_name] = plugin_instance
+                        plugin_instance.to(self._dtype)
+                        print(f"Loaded plugin: {plugin_name} from {file.name}")
+            except Exception as e:
+                print(f"Failed to load plugin {file.name}: {e}")
+        self._plugins: Dict[str, nn.Module] = {}
+
+        # Load plugins from plugin_dir if specified
+        if self.cfg.plugin_dir:
+            self.load_plugins(self.cfg.plugin_dir)
 
     # ---- gradient-checkpointing passthrough (Trainer calls these on the top model) ----
 
@@ -1086,6 +1212,9 @@ class MorphModel(nn.Module):
         # ---- System-2 reasoning loop ----
         reasoned, scratch = self.reasoner(base_hidden, steps)
 
+        # ---- multi-head CoT reasoning ----
+        reasoned = self.cot_reasoner(reasoned)
+
         # ---- subsystem gates ----
         g_think, g_code, g_mem, g_scratch = gates[:, 0], gates[:, 1], gates[:, 2], gates[:, 3]
         thresh = self.cfg.adaptive_threshold
@@ -1335,9 +1464,13 @@ class MorphModel(nn.Module):
             "coordinator", "reasoner", "code_bias", "scratchpad", "verifier",
             "skill_module", "depth_module", "moe_block", "memory",
             "mod", "multimodal_fusion", "tool_use", "document_module", "video_module",
+            "code_sandbox", "cot_reasoner",
         ):
             for k, v in getattr(self, name).state_dict().items():
                 sd[f"{name}.{k}"] = v
+        for plugin_name, plugin in self._plugins.items():
+            for k, v in plugin.state_dict().items():
+                sd[f"plugin_{plugin_name}.{k}"] = v
         if self.base_model is not None:
             try:
                 from peft import get_peft_model_state_dict
@@ -1351,10 +1484,16 @@ class MorphModel(nn.Module):
             "coordinator", "reasoner", "code_bias", "scratchpad", "verifier",
             "skill_module", "depth_module", "moe_block", "memory",
             "mod", "multimodal_fusion", "tool_use", "document_module", "video_module",
+            "code_sandbox", "cot_reasoner",
         ):
             sub = {k[len(name) + 1:]: v for k, v in sd.items() if k.startswith(name + ".")}
             if sub:
                 getattr(self, name).load_state_dict(sub)
+        for plugin_name in self._plugins:
+            prefix = f"plugin_{plugin_name}."
+            sub = {k[len(prefix):]: v for k, v in sd.items() if k.startswith(prefix)}
+            if sub:
+                self._plugins[plugin_name].load_state_dict(sub)
         if self.base_model is not None:
             peft_sd = {k: v for k, v in sd.items() if k.startswith("base_model")}
             if peft_sd:
@@ -1381,6 +1520,9 @@ class MorphModel(nn.Module):
                 "tool_use": self.tool_use.state_dict(),
                 "document_module": self.document_module.state_dict(),
                 "video_module": self.video_module.state_dict(),
+                "code_sandbox": self.code_sandbox.state_dict(),
+                "cot_reasoner": self.cot_reasoner.state_dict(),
+                **{f"plugin_{k}": v.state_dict() for k, v in self._plugins.items()},
                 "config": self.cfg,
             },
             f"{path}/morph_components.pt",
@@ -1401,9 +1543,14 @@ class MorphModel(nn.Module):
                 "coordinator", "reasoner", "code_bias", "scratchpad", "verifier",
                 "skill_module", "depth_module", "moe_block", "memory",
                 "mod", "multimodal_fusion", "tool_use", "document_module", "video_module",
+                "code_sandbox", "cot_reasoner",
             ):
                 if name in ckpt:
                     getattr(self, name).load_state_dict(ckpt[name])
+            for plugin_name in self._plugins:
+                key = f"plugin_{plugin_name}"
+                if key in ckpt:
+                    self._plugins[plugin_name].load_state_dict(ckpt[key])
             self.novel_trained = True
         else:
             trainer_ckpt = self._find_trainer_checkpoint(path)

src/runtime.py

@@ -101,12 +101,13 @@ class MorphRuntime:
     Audio (ASR + TTS), Tools, Documents, Video, Multimodal fusion.
     """
 
-    def __init__(self, model_path: str, use_4bit: bool = True, use_cpu: bool = False):
+    def __init__(self, model_path: str, use_4bit: bool = True, use_cpu: bool = False, plugin_dir: Optional[str] = None):
         self.model_path = Path(model_path)
         self.use_4bit = use_4bit and not use_cpu
         self.use_cpu = use_cpu
         self.skills: Dict[str, Skill] = {}
         self.active_skill: Optional[str] = None
+        self.plugin_dir = plugin_dir
 
         # v5 pipeline layers
         self.fsm = RuntimeFSM()
@@ -278,6 +279,8 @@ class MorphRuntime:
 
     def _load_model(self):
         config = MorphConfig()
+        if self.plugin_dir:
+            config.plugin_dir = self.plugin_dir
         self.model = MorphModel(config)
 
         # base model was loaded bf16 by MorphModel; load the trained adapter
@@ -294,7 +297,21 @@ class MorphRuntime:
             self.device = torch.device("cpu")
 
         self.model = self.model.to(self.device)
+        
+        # Enable extended context if configured
+        if hasattr(self.model.base_model_raw, 'config') and hasattr(self.model.base_model_raw.config, 'rope_scaling'):
+            if self.model.base_model_raw.config.rope_scaling:
+                print(f"Extended context enabled: {self.model.base_model_raw.config.rope_scaling}")
+        
         print(f"Model loaded on {self.device}")
+        if self.model._plugins:
+            print(f"Loaded plugins: {list(self.model._plugins.keys())}")
+
+    def expand_moe_experts(self):
+        """Dynamically expand MoE experts based on usage patterns."""
+        if hasattr(self.model, 'moe_block') and hasattr(self.model.moe_block, 'prune_and_expand_experts'):
+            self.model.moe_block.prune_and_expand_experts()
+            print(f"MoE expanded. Total experts: {len(self.model.moe_block.experts)}")
 
     def chat(self, prompt: str, max_tokens: int = 512, temperature: float = 0.7,
              skill: Optional[str] = None, image_path: Optional[str] = None,