auto: sync run_external_moe_family_wikitext_compare.py

scripts/run_external_moe_family_wikitext_compare.py

@@ -24,6 +24,7 @@ from src.data.wikitext_bpe import _collect_wikitext_text, load_wikitext_bpe
 from src.model import dense_transformer_baseline_external as local_dense_mod
 from src.model import motif_moe_external as local_motif_mod
 from src.model import plain_moe_transformer_external as local_flat_mod
+from src.model import wikitext_motif_combined_ensemble_external as local_ensemble_mod
 
 
 DEFAULT_DOWNLOADS_DIR = Path(r"C:\Users\Kharki\Downloads\Telegram Desktop")
@@ -123,6 +124,19 @@ def make_optimizer(model: torch.nn.Module, d_model: int) -> torch.optim.Optimize
     )
 
 
+def infer_model_width(model: torch.nn.Module, fallback: int) -> int:
+    config = getattr(model, "config", None)
+    if config is not None and hasattr(config, "d_model"):
+        return int(config.d_model)
+    members = getattr(model, "members", None)
+    if members:
+        first_member = next(iter(members.values()))
+        member_config = getattr(first_member, "config", None)
+        if member_config is not None and hasattr(member_config, "d_model"):
+            return int(member_config.d_model)
+    return int(fallback)
+
+
 def _dense_param_count(cfg: Any) -> int:
     hidden = max(16, int(round(cfg.d_model * cfg.ffn_hidden_ratio)))
     per_layer = (
@@ -280,6 +294,7 @@ def run_compare(
     motif_source: str,
     motif_profile: str,
     baseline_source: str,
+    include_combined_ensemble: bool,
 ) -> dict[str, Any]:
     modules = load_modules(downloads_dir, motif_source=motif_source, baseline_source=baseline_source)
     motif_mod = modules["motif"]
@@ -362,6 +377,39 @@ def run_compare(
     dense_model = dense_mod.DenseTransformerLM(dense_cfg).to(device)
     dense_params = int(dense_mod.count_parameters(dense_model))
 
+    combined_model: torch.nn.Module | None = None
+    combined_params: int | None = None
+    combined_config: dict[str, Any] | None = None
+    combined_history: list[dict[str, float]] | None = None
+    if include_combined_ensemble:
+        ensemble_build = local_ensemble_mod.build_wikitext_motif_ensemble(
+            local_ensemble_mod.EnsembleBuildConfig(
+                scale="0.125x",
+                block_size=seq_len,
+                vocab_size=actual_vocab_size,
+                include_text_motif_moe=True,
+                gate_hidden_dim=128,
+                gate_mode="contextual",
+                freeze_members=False,
+                member_dropout_p=0.05,
+                gate_temperature=1.0,
+                gate_entropy_bonus_coef=0.0,
+            )
+        )
+        combined_model = ensemble_build.model.to(device)
+        combined_params = int(local_ensemble_mod.count_parameters(combined_model))
+        combined_config = {
+            "scale": "0.125x",
+            "block_size": seq_len,
+            "gate_mode": "contextual",
+            "gate_hidden_dim": 128,
+            "freeze_members": False,
+            "member_dropout_p": 0.05,
+            "gate_temperature": 1.0,
+            "gate_entropy_bonus_coef": 0.0,
+            "member_names": [spec.name for spec in ensemble_build.specs],
+        }
+
     motif_history = train_model(
         name="MotifMoE",
         model=motif_model,
@@ -398,6 +446,19 @@ def run_compare(
         eval_every=eval_every,
         eval_batches=eval_batches,
     )
+    if combined_model is not None:
+        combined_history = train_model(
+            name="CombinedEnsemble",
+            model=combined_model,
+            d_model=infer_model_width(combined_model, fallback=motif_cfg.d_model),
+            train_data=train_data,
+            val_data=val_data,
+            device=device,
+            steps=steps,
+            batch_size=batch_size,
+            eval_every=eval_every,
+            eval_batches=eval_batches,
+        )
 
     report = {
         "downloads_dir": str(downloads_dir),
@@ -414,6 +475,7 @@ def run_compare(
         "motif_source": motif_source,
         "motif_profile": motif_profile,
         "baseline_source": baseline_source,
+        "include_combined_ensemble": include_combined_ensemble,
         "models": {
             "motif_moe": {
                 "params": motif_params,
@@ -457,6 +519,13 @@ def run_compare(
             },
         },
     }
+    if combined_model is not None and combined_history is not None and combined_params is not None and combined_config is not None:
+        report["models"]["combined_ensemble"] = {
+            "params": combined_params,
+            "config": combined_config,
+            "history": combined_history,
+            "final": combined_history[-1],
+        }
     return report
 
 
@@ -475,6 +544,7 @@ def main() -> None:
     parser.add_argument("--motif-source", choices=("local", "external"), default="external")
     parser.add_argument("--motif-profile", choices=("scaled", "text-lm"), default="scaled")
     parser.add_argument("--baseline-source", choices=("local", "external"), default="external")
+    parser.add_argument("--include-combined-ensemble", action="store_true")
     parser.add_argument("--model-name", default=None)
     parser.add_argument("--model", default=None)
     args = parser.parse_args()
@@ -493,6 +563,7 @@ def main() -> None:
         motif_source=args.motif_source,
         motif_profile=args.motif_profile,
         baseline_source=args.baseline_source,
+        include_combined_ensemble=args.include_combined_ensemble,
     )
 
     suffix = f"_densematch-{report['dense_match_target']}" if report["dense_match_target"] != "none" else ""
@@ -506,7 +577,8 @@ def main() -> None:
             f"_{report['motif_source']}-{report['motif_profile']}"
             f"_baseline-{report['baseline_source']}"
         )
-    out_path = ARCHIVE_DIR / f"external_moe_family_wikitext_compare_vocab{report['vocab_size']}{suffix}{profile_suffix}.json"
+    ensemble_suffix = "_with-ensemble" if report["include_combined_ensemble"] else ""
+    out_path = ARCHIVE_DIR / f"external_moe_family_wikitext_compare_vocab{report['vocab_size']}{suffix}{profile_suffix}{ensemble_suffix}.json"
     out_path.write_text(json.dumps(report, indent=2, ensure_ascii=False), encoding="utf-8")
     ranking = sorted(
         (

src/model/wikitext_motif_combined_ensemble_external.py

@@ -0,0 +1,1088 @@
+from __future__ import annotations
+
+import argparse
+import json
+import math
+from dataclasses import asdict, dataclass, field
+from pathlib import Path
+from typing import Any, Dict, Iterable, List, Mapping, Optional, Sequence, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, Dataset
+
+from src.model.dense_transformer_baseline_external import DenseTransformerConfig, DenseTransformerLM
+from src.model.motif_moe_external import MotifMoEConfig, MotifMoETransformer, count_parameters, make_motif_moe_config
+
+
+# -----------------------------------------------------------------------------
+# Data utilities
+# -----------------------------------------------------------------------------
+
+
+class ByteTokenizer:
+    """Dependency-free byte tokenizer for quick WikiText experiments.
+
+    This is intentionally simple so the code runs without external tokenization
+    libraries. If you already have pretokenized ids, use `load_token_tensor` and
+    skip this tokenizer entirely.
+    """
+
+    PAD_ID = 256
+    EOS_ID = 257
+    VOCAB_SIZE = 258
+
+    def __init__(self) -> None:
+        self.pad_id = self.PAD_ID
+        self.eos_id = self.EOS_ID
+        self.vocab_size = self.VOCAB_SIZE
+
+    def encode(self, text: str, add_eos: bool = True) -> List[int]:
+        ids = list(text.encode("utf-8", errors="replace"))
+        if add_eos:
+            ids.append(self.eos_id)
+        return ids
+
+    def decode(self, ids: Sequence[int]) -> str:
+        buf = bytearray()
+        for idx in ids:
+            if idx < 256:
+                buf.append(int(idx))
+            elif idx == self.eos_id:
+                buf.extend(b"\n")
+        return buf.decode("utf-8", errors="replace")
+
+
+def load_token_tensor(path: str | Path) -> torch.Tensor:
+    path = Path(path)
+    suffix = path.suffix.lower()
+    if suffix == ".pt":
+        obj = torch.load(path, map_location="cpu")
+        if isinstance(obj, torch.Tensor):
+            return obj.long().flatten()
+        raise TypeError(f"Expected a Tensor in {path}, got {type(obj)!r}")
+
+    if suffix == ".npy":
+        try:
+            import numpy as np
+        except Exception as exc:  # pragma: no cover - optional dependency
+            raise ImportError("numpy is required to load .npy token arrays") from exc
+        arr = np.load(path)
+        return torch.from_numpy(arr).long().flatten()
+
+    if suffix in {".bin", ".uint16", ".u16"}:
+        try:
+            import numpy as np
+        except Exception as exc:  # pragma: no cover - optional dependency
+            raise ImportError("numpy is required to load .bin token arrays") from exc
+        arr = np.fromfile(path, dtype=np.uint16)
+        return torch.from_numpy(arr.astype("int64"))
+
+    raise ValueError(f"Unsupported token file extension for {path}")
+
+
+def text_file_to_tokens(path: str | Path, tokenizer: Optional[ByteTokenizer] = None) -> torch.Tensor:
+    tokenizer = tokenizer or ByteTokenizer()
+    text = Path(path).read_text(encoding="utf-8")
+    ids: List[int] = []
+    for line in text.splitlines():
+        ids.extend(tokenizer.encode(line, add_eos=True))
+    return torch.tensor(ids, dtype=torch.long)
+
+
+class TokenBlockDataset(Dataset[torch.Tensor]):
+    def __init__(self, tokens: torch.Tensor, block_size: int, stride: Optional[int] = None) -> None:
+        if tokens.ndim != 1:
+            raise ValueError("tokens must be a 1D tensor")
+        if block_size < 2:
+            raise ValueError("block_size must be >= 2")
+        self.tokens = tokens.long().contiguous()
+        self.block_size = int(block_size)
+        self.stride = int(stride if stride is not None else block_size)
+        if self.stride < 1:
+            raise ValueError("stride must be >= 1")
+        self.max_start = max(0, self.tokens.numel() - self.block_size)
+        self.num_examples = 1 + (self.max_start // self.stride) if self.tokens.numel() >= self.block_size else 0
+
+    def __len__(self) -> int:
+        return self.num_examples
+
+    def __getitem__(self, index: int) -> torch.Tensor:
+        start = index * self.stride
+        chunk = self.tokens[start : start + self.block_size]
+        if chunk.numel() != self.block_size:
+            raise IndexError(index)
+        return chunk
+
+
+# -----------------------------------------------------------------------------
+# Branch specs and builders
+# -----------------------------------------------------------------------------
+
+
+@dataclass
+class MotifBranchSpec:
+    name: str
+    model_type: str  # 'dense' or 'motif_moe'
+    attn_qk_ratio: float = 1.0
+    attn_v_ratio: float = 1.0
+
+    # Dense-only fields
+    ffn_hidden_ratio: float = 4.0
+    ffn_activation: str = "gelu"
+
+    # Motif-MoE-only fields
+    motif_families: Tuple[str, ...] = ()
+    experts_per_family: int = 2
+    motif_hidden_ratios: Dict[str, float] = field(default_factory=dict)
+    top_k_motifs: int = 2
+    top_k_experts: int = 1
+    expert_pool_scale: float = 1.0
+    router_jitter_noise: float = 0.01
+
+    def validate(self) -> None:
+        if self.model_type not in {"dense", "motif_moe"}:
+            raise ValueError(f"Unsupported model_type={self.model_type!r}")
+        if self.attn_qk_ratio <= 0 or self.attn_v_ratio <= 0:
+            raise ValueError("Attention ratios must be > 0")
+        if self.model_type == "dense" and self.ffn_hidden_ratio <= 0:
+            raise ValueError("ffn_hidden_ratio must be > 0 for dense branches")
+        if self.model_type == "motif_moe":
+            if not self.motif_families:
+                raise ValueError("motif_moe branch requires motif_families")
+            if self.experts_per_family < 1:
+                raise ValueError("experts_per_family must be >= 1")
+
+
+@dataclass
+class EnsembleBuildConfig:
+    scale: str = "0.125x"
+    block_size: int = 256
+    vocab_size: int = ByteTokenizer.VOCAB_SIZE
+    include_text_motif_moe: bool = True
+    gate_hidden_dim: int = 128
+    gate_mode: str = "contextual"  # contextual | static
+    freeze_members: bool = True
+    member_dropout_p: float = 0.0
+    gate_temperature: float = 1.0
+    gate_entropy_bonus_coef: float = 0.0
+
+
+@dataclass
+class TrainConfig:
+    batch_size: int = 8
+    lr: float = 3e-4
+    weight_decay: float = 0.01
+    epochs: int = 1
+    max_steps: Optional[int] = None
+    grad_clip: float = 1.0
+    device: str = "cpu"
+    log_every: int = 50
+
+
+@dataclass
+class EvalSummary:
+    split: str
+    lm_loss: float
+    perplexity: float
+    tokens: int
+
+
+def make_wikitext_branch_specs(include_text_motif_moe: bool = True) -> List[MotifBranchSpec]:
+    """Default motif-profile ensemble for WikiText-like natural text.
+
+    Profiles reflect the observed pattern that text often benefits from stronger
+    expand/select/memory capacity and does not require the widest compare space.
+    """
+    specs = [
+        MotifBranchSpec(
+            name="uniform_dense",
+            model_type="dense",
+            attn_qk_ratio=1.0,
+            attn_v_ratio=1.0,
+            ffn_hidden_ratio=4.0,
+            ffn_activation="gelu",
+        ),
+        MotifBranchSpec(
+            name="narrow_compare_dense",
+            model_type="dense",
+            attn_qk_ratio=0.75,
+            attn_v_ratio=1.0,
+            ffn_hidden_ratio=4.5,
+            ffn_activation="gelu",
+        ),
+        MotifBranchSpec(
+            name="wide_memory_dense",
+            model_type="dense",
+            attn_qk_ratio=0.75,
+            attn_v_ratio=1.0,
+            ffn_hidden_ratio=6.0,
+            ffn_activation="swiglu",
+        ),
+    ]
+    if include_text_motif_moe:
+        specs.append(
+            MotifBranchSpec(
+                name="text_motif_moe",
+                model_type="motif_moe",
+                attn_qk_ratio=0.75,
+                attn_v_ratio=1.0,
+                motif_families=("expand", "select", "memory"),
+                experts_per_family=2,
+                motif_hidden_ratios={
+                    "expand": 6.0,
+                    "select": 2.0,
+                    "memory": 4.0,
+                },
+                top_k_motifs=2,
+                top_k_experts=1,
+                expert_pool_scale=0.5,
+                router_jitter_noise=0.01,
+            )
+        )
+    for spec in specs:
+        spec.validate()
+    return specs
+
+
+def _clone_motif_cfg(cfg: MotifMoEConfig) -> MotifMoEConfig:
+    return MotifMoEConfig(**asdict(cfg))
+
+
+class DenseTransformerLMWithHidden(DenseTransformerLM):
+    """DenseTransformerLM that can optionally return final hidden states."""
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        labels: Optional[torch.Tensor] = None,
+        *,
+        return_hidden_states: bool = False,
+        return_router_stats: bool = False,
+    ) -> Dict[str, Any]:
+        del return_router_stats
+        x = self.emb_dropout(self.tok_emb(input_ids))
+        for block in self.blocks:
+            x = block(x)
+        hidden = self.ln_f(x)
+        logits = self.lm_head(hidden)
+
+        lm_loss: Optional[torch.Tensor] = None
+        loss: Optional[torch.Tensor] = None
+        if labels is not None:
+            shift_logits = logits[:, :-1, :].contiguous()
+            shift_labels = labels[:, 1:].contiguous()
+            lm_loss = F.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1),
+                ignore_index=-100,
+            )
+            loss = lm_loss
+
+        zero = logits.new_zeros(())
+        out: Dict[str, Any] = {
+            "logits": logits,
+            "loss": loss,
+            "lm_loss": lm_loss,
+            "router_aux_loss": zero,
+            "router_aux_loss_raw": zero,
+            "router_stats": None,
+        }
+        if return_hidden_states:
+            out["hidden_states"] = hidden
+        return out
+
+
+class MotifMoETransformerWithHidden(MotifMoETransformer):
+    """MotifMoETransformer that can optionally return final hidden states."""
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        labels: Optional[torch.Tensor] = None,
+        *,
+        return_hidden_states: bool = False,
+        return_router_stats: bool = False,
+    ) -> Dict[str, Any]:
+        out = super().forward(input_ids, labels=labels, return_router_stats=return_router_stats)
+        if return_hidden_states:
+            x = self.emb_dropout(self.tok_emb(input_ids))
+            aux_losses: List[torch.Tensor] = []
+            router_stats: List[Dict[str, Any]] = []
+            for block in self.blocks:
+                x, aux_loss, stats = block(x, return_router_stats=return_router_stats)
+                aux_losses.append(aux_loss)
+                if return_router_stats:
+                    router_stats.append(stats)
+            hidden = self.ln_f(x)
+            out["hidden_states"] = hidden
+            # Preserve the already-computed outputs, but keep router stats consistent.
+            if return_router_stats:
+                out["router_stats"] = router_stats
+        return out
+
+
+def build_branch_model(
+    spec: MotifBranchSpec,
+    *,
+    scale: str,
+    vocab_size: int,
+    block_size: int,
+) -> nn.Module:
+    spec.validate()
+    base_cfg = make_motif_moe_config(
+        scale=scale,
+        vocab_size=vocab_size,
+        max_seq_len=block_size,
+        expert_pool_scale=max(0.5, spec.expert_pool_scale),
+        top_k_motifs=max(1, spec.top_k_motifs),
+        top_k_experts=max(1, spec.top_k_experts),
+        attn_qk_ratio=spec.attn_qk_ratio,
+        attn_v_ratio=spec.attn_v_ratio,
+    )
+
+    if spec.model_type == "dense":
+        dense_cfg = DenseTransformerConfig(
+            vocab_size=base_cfg.vocab_size,
+            max_seq_len=base_cfg.max_seq_len,
+            tie_word_embeddings=base_cfg.tie_word_embeddings,
+            n_layers=base_cfg.n_layers,
+            d_model=base_cfg.d_model,
+            n_heads=base_cfg.n_heads,
+            attn_qk_ratio=spec.attn_qk_ratio,
+            attn_v_ratio=spec.attn_v_ratio,
+            ffn_hidden_ratio=spec.ffn_hidden_ratio,
+            ffn_activation=spec.ffn_activation,
+            rope_base=base_cfg.rope_base,
+            bias=base_cfg.bias,
+            norm_eps=base_cfg.norm_eps,
+            resid_dropout=base_cfg.resid_dropout,
+            attn_dropout=base_cfg.attn_dropout,
+            emb_dropout=base_cfg.emb_dropout,
+            ffn_dropout=base_cfg.expert_dropout,
+            initializer_range=base_cfg.initializer_range,
+        )
+        return DenseTransformerLMWithHidden(dense_cfg)
+
+    motif_cfg = _clone_motif_cfg(base_cfg)
+    motif_cfg.attn_qk_ratio = spec.attn_qk_ratio
+    motif_cfg.attn_v_ratio = spec.attn_v_ratio
+    motif_cfg.motif_families = tuple(spec.motif_families)
+    motif_cfg.experts_per_family = {family: spec.experts_per_family for family in motif_cfg.motif_families}
+    motif_cfg.motif_hidden_ratios = {
+        family: spec.motif_hidden_ratios.get(family, base_cfg.motif_hidden_ratios.get(family, 4.0))
+        for family in motif_cfg.motif_families
+    }
+    motif_cfg.top_k_motifs = min(spec.top_k_motifs, len(motif_cfg.motif_families))
+    motif_cfg.top_k_experts = spec.top_k_experts
+    motif_cfg.router_jitter_noise = spec.router_jitter_noise
+    motif_cfg.validate()
+    return MotifMoETransformerWithHidden(motif_cfg)
+
+
+# -----------------------------------------------------------------------------
+# Ensemble model
+# -----------------------------------------------------------------------------
+
+
+class ContextualLogitGate(nn.Module):
+    """Learns per-token convex weights over branch logits.
+
+    Features are extracted from each branch's logits:
+      - entropy
+      - max probability
+      - top-1 / top-2 margin
+      - logit standard deviation
+    """
+
+    def __init__(
+        self,
+        num_branches: int,
+        *,
+        hidden_dim: int = 128,
+        mode: str = "contextual",
+        temperature: float = 1.0,
+    ) -> None:
+        super().__init__()
+        if mode not in {"contextual", "static"}:
+            raise ValueError("mode must be 'contextual' or 'static'")
+        self.num_branches = num_branches
+        self.hidden_dim = hidden_dim
+        self.mode = mode
+        self.temperature = temperature
+        self.branch_bias = nn.Parameter(torch.zeros(num_branches))
+        if mode == "contextual":
+            in_dim = num_branches * 4
+            self.mlp = nn.Sequential(
+                nn.Linear(in_dim, hidden_dim),
+                nn.GELU(),
+                nn.Linear(hidden_dim, hidden_dim),
+                nn.GELU(),
+                nn.Linear(hidden_dim, num_branches),
+            )
+        else:
+            self.mlp = None
+
+    def _extract_features(self, logits: torch.Tensor) -> torch.Tensor:
+        # logits: [M, B, T, V]
+        stats: List[torch.Tensor] = []
+        for branch_logits in logits.detach():
+            log_probs = F.log_softmax(branch_logits.float(), dim=-1)
+            probs = log_probs.exp()
+            entropy = -(probs * log_probs).sum(dim=-1)
+            top2 = torch.topk(branch_logits.float(), k=min(2, branch_logits.size(-1)), dim=-1).values
+            if top2.size(-1) == 1:
+                margin = top2[..., 0]
+            else:
+                margin = top2[..., 0] - top2[..., 1]
+            max_prob = probs.max(dim=-1).values
+            spread = branch_logits.float().std(dim=-1)
+            stats.append(torch.stack([entropy, max_prob, margin, spread], dim=-1))
+        return torch.cat(stats, dim=-1)  # [B, T, 4M]
+
+    def forward(self, logits: torch.Tensor) -> torch.Tensor:
+        # logits: [M, B, T, V]
+        _, batch, seqlen, _ = logits.shape
+        if self.mode == "static":
+            weights = F.softmax(self.branch_bias / max(self.temperature, 1e-6), dim=-1)
+            return weights.view(1, 1, -1).expand(batch, seqlen, -1)
+        features = self._extract_features(logits)
+        gate_logits = self.mlp(features) + self.branch_bias
+        return F.softmax(gate_logits / max(self.temperature, 1e-6), dim=-1)
+
+
+class WikiTextMotifEnsembleLM(nn.Module):
+    """Ensemble over motif-profile branches with a learned convex logit combiner."""
+
+    def __init__(
+        self,
+        members: Mapping[str, nn.Module],
+        *,
+        gate_hidden_dim: int = 128,
+        gate_mode: str = "contextual",
+        freeze_members: bool = True,
+        member_dropout_p: float = 0.0,
+        gate_temperature: float = 1.0,
+        gate_entropy_bonus_coef: float = 0.0,
+    ) -> None:
+        super().__init__()
+        if not members:
+            raise ValueError("At least one ensemble member is required")
+        self.members = nn.ModuleDict(members)
+        self.branch_names = list(members.keys())
+        self.freeze_members = freeze_members
+        self.member_dropout_p = member_dropout_p
+        self.gate_entropy_bonus_coef = gate_entropy_bonus_coef
+        self.gate = ContextualLogitGate(
+            len(self.members),
+            hidden_dim=gate_hidden_dim,
+            mode=gate_mode,
+            temperature=gate_temperature,
+        )
+        if self.freeze_members:
+            self.freeze_all_members()
+
+    def freeze_all_members(self) -> None:
+        for model in self.members.values():
+            for param in model.parameters():
+                param.requires_grad_(False)
+            model.eval()
+
+    def unfreeze_all_members(self) -> None:
+        for model in self.members.values():
+            for param in model.parameters():
+                param.requires_grad_(True)
+
+    def _forward_member(self, model: nn.Module, input_ids: torch.Tensor, return_router_stats: bool) -> Dict[str, Any]:
+        if self.freeze_members:
+            with torch.no_grad():
+                return model(input_ids, labels=None, return_hidden_states=False, return_router_stats=return_router_stats)
+        return model(input_ids, labels=None, return_hidden_states=False, return_router_stats=return_router_stats)
+
+    def _apply_member_dropout(self, weights: torch.Tensor) -> torch.Tensor:
+        if not self.training or self.member_dropout_p <= 0:
+            return weights
+        keep = torch.rand_like(weights) > self.member_dropout_p
+        # Ensure at least one branch stays alive per token.
+        all_zero = ~keep.any(dim=-1, keepdim=True)
+        if torch.any(all_zero):
+            fallback = torch.zeros_like(keep)
+            fallback[..., 0] = True
+            keep = torch.where(all_zero, fallback, keep)
+        dropped = weights * keep.float()
+        norm = dropped.sum(dim=-1, keepdim=True).clamp_min(1e-8)
+        return dropped / norm
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        labels: Optional[torch.Tensor] = None,
+        *,
+        return_member_outputs: bool = False,
+        return_router_stats: bool = False,
+    ) -> Dict[str, Any]:
+        member_outputs: Dict[str, Dict[str, Any]] = {}
+        member_logits: List[torch.Tensor] = []
+        member_aux_losses: List[torch.Tensor] = []
+
+        for name, model in self.members.items():
+            out = self._forward_member(model, input_ids, return_router_stats=return_router_stats)
+            member_outputs[name] = out
+            member_logits.append(out["logits"])
+            if out.get("router_aux_loss") is not None:
+                member_aux_losses.append(out["router_aux_loss"])
+
+        logits_stack = torch.stack(member_logits, dim=0)  # [M, B, T, V]
+        gate_weights = self.gate(logits_stack)  # [B, T, M]
+        gate_weights = self._apply_member_dropout(gate_weights)
+        weights_for_sum = gate_weights.permute(2, 0, 1).unsqueeze(-1)
+        ensemble_logits = (logits_stack * weights_for_sum).sum(dim=0)
+
+        lm_loss: Optional[torch.Tensor] = None
+        loss: Optional[torch.Tensor] = None
+        if labels is not None:
+            shift_logits = ensemble_logits[:, :-1, :].contiguous()
+            shift_labels = labels[:, 1:].contiguous()
+            lm_loss = F.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1),
+                ignore_index=-100,
+            )
+            loss = lm_loss
+            if not self.freeze_members and member_aux_losses:
+                loss = loss + torch.stack(member_aux_losses).mean()
+            if self.gate_entropy_bonus_coef != 0.0:
+                entropy = -(gate_weights * (gate_weights.clamp_min(1e-8).log())).sum(dim=-1).mean()
+                loss = loss - self.gate_entropy_bonus_coef * entropy
+
+        out: Dict[str, Any] = {
+            "logits": ensemble_logits,
+            "loss": loss,
+            "lm_loss": lm_loss,
+            "router_aux_loss": torch.stack(member_aux_losses).mean() if member_aux_losses else ensemble_logits.new_zeros(()),
+            "router_aux_loss_raw": torch.stack(member_aux_losses).mean() if member_aux_losses else ensemble_logits.new_zeros(()),
+            "ensemble_weights": gate_weights,
+            "branch_names": self.branch_names,
+        }
+        if return_member_outputs:
+            out["member_outputs"] = member_outputs
+        if return_router_stats:
+            out["member_router_stats"] = {name: member_outputs[name].get("router_stats") for name in self.branch_names}
+        return out
+
+    @torch.no_grad()
+    def summarize_weights(self, gate_weights: torch.Tensor) -> Dict[str, float]:
+        mean_weights = gate_weights.mean(dim=(0, 1)).cpu().tolist()
+        return {name: float(weight) for name, weight in zip(self.branch_names, mean_weights)}
+
+
+@dataclass
+class BuiltEnsemble:
+    model: WikiTextMotifEnsembleLM
+    specs: List[MotifBranchSpec]
+
+
+def build_wikitext_motif_ensemble(config: EnsembleBuildConfig) -> BuiltEnsemble:
+    specs = make_wikitext_branch_specs(include_text_motif_moe=config.include_text_motif_moe)
+    members = {
+        spec.name: build_branch_model(spec, scale=config.scale, vocab_size=config.vocab_size, block_size=config.block_size)
+        for spec in specs
+    }
+    model = WikiTextMotifEnsembleLM(
+        members,
+        gate_hidden_dim=config.gate_hidden_dim,
+        gate_mode=config.gate_mode,
+        freeze_members=config.freeze_members,
+        member_dropout_p=config.member_dropout_p,
+        gate_temperature=config.gate_temperature,
+        gate_entropy_bonus_coef=config.gate_entropy_bonus_coef,
+    )
+    return BuiltEnsemble(model=model, specs=specs)
+
+
+# -----------------------------------------------------------------------------
+# Checkpoint helpers
+# -----------------------------------------------------------------------------
+
+
+@dataclass
+class BranchCheckpoint:
+    spec: Dict[str, Any]
+    state_dict: Dict[str, Any]
+
+
+@torch.no_grad()
+def save_branch_checkpoint(model: nn.Module, spec: MotifBranchSpec, path: str | Path) -> None:
+    payload = {
+        "spec": asdict(spec),
+        "state_dict": model.state_dict(),
+    }
+    torch.save(payload, Path(path))
+
+
+@torch.no_grad()
+def load_branch_checkpoint(path: str | Path, *, scale: str, vocab_size: int, block_size: int, device: str = "cpu") -> Tuple[MotifBranchSpec, nn.Module]:
+    payload = torch.load(Path(path), map_location=device)
+    spec_dict = payload["spec"]
+    if isinstance(spec_dict.get("motif_families"), list):
+        spec_dict["motif_families"] = tuple(spec_dict["motif_families"])
+    spec = MotifBranchSpec(**spec_dict)
+    model = build_branch_model(spec, scale=scale, vocab_size=vocab_size, block_size=block_size)
+    model.load_state_dict(payload["state_dict"])
+    return spec, model
+
+
+def save_ensemble_checkpoint(
+    ensemble: WikiTextMotifEnsembleLM,
+    specs: Sequence[MotifBranchSpec],
+    path: str | Path,
+) -> None:
+    payload = {
+        "specs": [asdict(spec) for spec in specs],
+        "gate_state_dict": ensemble.gate.state_dict(),
+        "branch_names": list(ensemble.branch_names),
+    }
+    torch.save(payload, Path(path))
+
+
+def load_ensemble_gate(ensemble: WikiTextMotifEnsembleLM, path: str | Path, device: str = "cpu") -> Dict[str, Any]:
+    payload = torch.load(Path(path), map_location=device)
+    ensemble.gate.load_state_dict(payload["gate_state_dict"])
+    return payload
+
+
+# -----------------------------------------------------------------------------
+# Training and evaluation
+# -----------------------------------------------------------------------------
+
+
+@dataclass
+class RunningAverage:
+    total: float = 0.0
+    count: int = 0
+
+    def update(self, value: float, n: int = 1) -> None:
+        self.total += value * n
+        self.count += n
+
+    @property
+    def mean(self) -> float:
+        return self.total / max(1, self.count)
+
+
+def move_batch(batch: torch.Tensor, device: torch.device) -> torch.Tensor:
+    return batch.to(device, non_blocking=True)
+
+
+def make_loader(tokens: torch.Tensor, *, block_size: int, batch_size: int, shuffle: bool) -> DataLoader[torch.Tensor]:
+    dataset = TokenBlockDataset(tokens, block_size=block_size, stride=block_size)
+    return DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, drop_last=False)
+
+
+def train_lm_model(
+    model: nn.Module,
+    train_tokens: torch.Tensor,
+    valid_tokens: Optional[torch.Tensor],
+    *,
+    block_size: int,
+    cfg: TrainConfig,
+) -> Dict[str, Any]:
+    device = torch.device(cfg.device)
+    model.to(device)
+    model.train()
+    loader = make_loader(train_tokens, block_size=block_size, batch_size=cfg.batch_size, shuffle=True)
+    optimizer = torch.optim.AdamW(model.parameters(), lr=cfg.lr, betas=(0.9, 0.95), weight_decay=cfg.weight_decay)
+
+    step = 0
+    train_loss = RunningAverage()
+    for epoch in range(cfg.epochs):
+        for batch in loader:
+            batch = move_batch(batch, device)
+            optimizer.zero_grad(set_to_none=True)
+            out = model(batch, labels=batch)
+            loss = out["loss"]
+            if loss is None:
+                raise RuntimeError("Model did not return a loss")
+            loss.backward()
+            if cfg.grad_clip > 0:
+                torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip)
+            optimizer.step()
+            step += 1
+            train_loss.update(float(loss.detach().cpu()), n=batch.size(0))
+            if cfg.log_every > 0 and step % cfg.log_every == 0:
+                print(f"[train] epoch={epoch+1} step={step} loss={train_loss.mean:.4f}")
+            if cfg.max_steps is not None and step >= cfg.max_steps:
+                break
+        if cfg.max_steps is not None and step >= cfg.max_steps:
+            break
+
+    result: Dict[str, Any] = {
+        "train_loss": train_loss.mean,
+        "train_steps": step,
+    }
+    if valid_tokens is not None:
+        result["valid"] = asdict(
+            evaluate_lm_model(model, valid_tokens, block_size=block_size, batch_size=cfg.batch_size, device=cfg.device)
+        )
+    return result
+
+
+@torch.no_grad()
+def evaluate_lm_model(
+    model: nn.Module,
+    tokens: torch.Tensor,
+    *,
+    block_size: int,
+    batch_size: int,
+    device: str,
+    split: str = "valid",
+) -> EvalSummary:
+    model.eval()
+    dev = torch.device(device)
+    model.to(dev)
+    loader = make_loader(tokens, block_size=block_size, batch_size=batch_size, shuffle=False)
+    loss_avg = RunningAverage()
+    token_count = 0
+    for batch in loader:
+        batch = move_batch(batch, dev)
+        out = model(batch, labels=batch)
+        lm_loss = out["lm_loss"]
+        if lm_loss is None:
+            raise RuntimeError("Model did not return lm_loss")
+        valid_tokens = int((batch[:, 1:] != -100).sum().item())
+        loss_avg.update(float(lm_loss.detach().cpu()), n=max(1, valid_tokens))
+        token_count += valid_tokens
+    lm_loss = loss_avg.mean
+    ppl = float(math.exp(min(20.0, lm_loss)))
+    return EvalSummary(split=split, lm_loss=lm_loss, perplexity=ppl, tokens=token_count)
+
+
+@torch.no_grad()
+def evaluate_members(
+    members: Mapping[str, nn.Module],
+    tokens: torch.Tensor,
+    *,
+    block_size: int,
+    batch_size: int,
+    device: str,
+) -> Dict[str, Dict[str, float]]:
+    results: Dict[str, Dict[str, float]] = {}
+    for name, model in members.items():
+        summary = evaluate_lm_model(model, tokens, block_size=block_size, batch_size=batch_size, device=device, split=name)
+        results[name] = asdict(summary)
+    return results
+
+
+# -----------------------------------------------------------------------------
+# CLI helpers
+# -----------------------------------------------------------------------------
+
+
+@dataclass
+class DataBundle:
+    train_tokens: torch.Tensor
+    valid_tokens: torch.Tensor
+    test_tokens: Optional[torch.Tensor]
+    vocab_size: int
+
+
+def load_wikitext_data(
+    *,
+    train_path: str,
+    valid_path: str,
+    test_path: Optional[str] = None,
+    tokenized: bool = False,
+) -> DataBundle:
+    if tokenized:
+        train_tokens = load_token_tensor(train_path)
+        valid_tokens = load_token_tensor(valid_path)
+        test_tokens = load_token_tensor(test_path) if test_path is not None else None
+        vocab_size = int(max(train_tokens.max().item(), valid_tokens.max().item(), test_tokens.max().item() if test_tokens is not None else 0) + 1)
+    else:
+        tokenizer = ByteTokenizer()
+        train_tokens = text_file_to_tokens(train_path, tokenizer)
+        valid_tokens = text_file_to_tokens(valid_path, tokenizer)
+        test_tokens = text_file_to_tokens(test_path, tokenizer) if test_path is not None else None
+        vocab_size = tokenizer.vocab_size
+    return DataBundle(train_tokens=train_tokens, valid_tokens=valid_tokens, test_tokens=test_tokens, vocab_size=vocab_size)
+
+
+def train_members_cli(args: argparse.Namespace) -> None:
+    data = load_wikitext_data(
+        train_path=args.train_path,
+        valid_path=args.valid_path,
+        test_path=args.test_path,
+        tokenized=args.tokenized,
+    )
+    out_dir = Path(args.output_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    members_dir = out_dir / "members"
+    members_dir.mkdir(parents=True, exist_ok=True)
+
+    specs = make_wikitext_branch_specs(include_text_motif_moe=not args.no_text_motif_moe)
+    train_cfg = TrainConfig(
+        batch_size=args.batch_size,
+        lr=args.lr,
+        weight_decay=args.weight_decay,
+        epochs=args.epochs,
+        max_steps=args.max_steps,
+        grad_clip=args.grad_clip,
+        device=args.device,
+        log_every=args.log_every,
+    )
+
+    summary: Dict[str, Any] = {"members": {}}
+    for spec in specs:
+        print(f"\n=== Training member: {spec.name} ===")
+        model = build_branch_model(spec, scale=args.scale, vocab_size=data.vocab_size, block_size=args.block_size)
+        print(f"params[{spec.name}]={count_parameters(model):,}")
+        result = train_lm_model(
+            model,
+            data.train_tokens,
+            data.valid_tokens,
+            block_size=args.block_size,
+            cfg=train_cfg,
+        )
+        ckpt_path = members_dir / f"{spec.name}.pt"
+        save_branch_checkpoint(model, spec, ckpt_path)
+        summary["members"][spec.name] = {
+            "checkpoint": str(ckpt_path),
+            "params": count_parameters(model),
+            "result": result,
+        }
+        print(json.dumps(summary["members"][spec.name], indent=2))
+
+    (out_dir / "member_training_summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8")
+
+
+
+def train_ensemble_cli(args: argparse.Namespace) -> None:
+    data = load_wikitext_data(
+        train_path=args.train_path,
+        valid_path=args.valid_path,
+        test_path=args.test_path,
+        tokenized=args.tokenized,
+    )
+    out_dir = Path(args.output_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    members_dir = out_dir / "members"
+
+    specs: List[MotifBranchSpec] = []
+    members: Dict[str, nn.Module] = {}
+    for ckpt_path in sorted(members_dir.glob("*.pt")):
+        spec, model = load_branch_checkpoint(
+            ckpt_path,
+            scale=args.scale,
+            vocab_size=data.vocab_size,
+            block_size=args.block_size,
+            device=args.device,
+        )
+        specs.append(spec)
+        members[spec.name] = model
+
+    if not members:
+        raise FileNotFoundError(f"No member checkpoints found in {members_dir}")
+
+    ensemble = WikiTextMotifEnsembleLM(
+        members,
+        gate_hidden_dim=args.gate_hidden_dim,
+        gate_mode=args.gate_mode,
+        freeze_members=not args.joint_finetune,
+        member_dropout_p=args.member_dropout_p,
+        gate_temperature=args.gate_temperature,
+        gate_entropy_bonus_coef=args.gate_entropy_bonus_coef,
+    )
+    device = torch.device(args.device)
+    ensemble.to(device)
+    if args.joint_finetune:
+        ensemble.unfreeze_all_members()
+        for model in ensemble.members.values():
+            model.train()
+
+    train_cfg = TrainConfig(
+        batch_size=args.batch_size,
+        lr=args.lr,
+        weight_decay=args.weight_decay,
+        epochs=args.epochs,
+        max_steps=args.max_steps,
+        grad_clip=args.grad_clip,
+        device=args.device,
+        log_every=args.log_every,
+    )
+    result = train_lm_model(
+        ensemble,
+        data.train_tokens,
+        data.valid_tokens,
+        block_size=args.block_size,
+        cfg=train_cfg,
+    )
+    gate_path = out_dir / "motif_ensemble_gate.pt"
+    save_ensemble_checkpoint(ensemble, specs, gate_path)
+
+    valid_summary = evaluate_lm_model(
+        ensemble,
+        data.valid_tokens,
+        block_size=args.block_size,
+        batch_size=args.batch_size,
+        device=args.device,
+        split="valid_ensemble",
+    )
+    member_summaries = evaluate_members(
+        ensemble.members,
+        data.valid_tokens,
+        block_size=args.block_size,
+        batch_size=args.batch_size,
+        device=args.device,
+    )
+
+    summary = {
+        "ensemble_params_trainable": sum(p.numel() for p in ensemble.parameters() if p.requires_grad),
+        "ensemble_result": result,
+        "valid_ensemble": asdict(valid_summary),
+        "valid_members": member_summaries,
+        "gate_checkpoint": str(gate_path),
+    }
+    (out_dir / "ensemble_training_summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8")
+    print(json.dumps(summary, indent=2))
+
+
+
+def evaluate_cli(args: argparse.Namespace) -> None:
+    data = load_wikitext_data(
+        train_path=args.train_path,
+        valid_path=args.valid_path,
+        test_path=args.test_path,
+        tokenized=args.tokenized,
+    )
+    out_dir = Path(args.output_dir)
+    members_dir = out_dir / "members"
+    specs: List[MotifBranchSpec] = []
+    members: Dict[str, nn.Module] = {}
+    for ckpt_path in sorted(members_dir.glob("*.pt")):
+        spec, model = load_branch_checkpoint(
+            ckpt_path,
+            scale=args.scale,
+            vocab_size=data.vocab_size,
+            block_size=args.block_size,
+            device=args.device,
+        )
+        specs.append(spec)
+        members[spec.name] = model
+
+    if not members:
+        raise FileNotFoundError(f"No member checkpoints found in {members_dir}")
+
+    ensemble = WikiTextMotifEnsembleLM(
+        members,
+        gate_hidden_dim=args.gate_hidden_dim,
+        gate_mode=args.gate_mode,
+        freeze_members=True,
+        member_dropout_p=0.0,
+        gate_temperature=args.gate_temperature,
+        gate_entropy_bonus_coef=0.0,
+    )
+    gate_path = out_dir / "motif_ensemble_gate.pt"
+    if gate_path.exists():
+        load_ensemble_gate(ensemble, gate_path, device=args.device)
+    else:
+        print(f"Warning: {gate_path} not found. Evaluating with randomly initialized gate.")
+
+    splits = [("valid", data.valid_tokens)]
+    if data.test_tokens is not None:
+        splits.append(("test", data.test_tokens))
+
+    report: Dict[str, Any] = {"members": {}, "ensemble": {}}
+    for split_name, split_tokens in splits:
+        report["ensemble"][split_name] = asdict(
+            evaluate_lm_model(
+                ensemble,
+                split_tokens,
+                block_size=args.block_size,
+                batch_size=args.batch_size,
+                device=args.device,
+                split=split_name,
+            )
+        )
+        report["members"][split_name] = evaluate_members(
+            members,
+            split_tokens,
+            block_size=args.block_size,
+            batch_size=args.batch_size,
+            device=args.device,
+        )
+
+    sample_batch = next(iter(make_loader(data.valid_tokens, block_size=args.block_size, batch_size=min(2, args.batch_size), shuffle=False)))
+    sample_batch = sample_batch.to(args.device)
+    ensemble = ensemble.to(args.device)
+    out = ensemble(sample_batch, labels=sample_batch)
+    report["mean_gate_weights"] = ensemble.summarize_weights(out["ensemble_weights"])
+
+    report_path = out_dir / "ensemble_eval_report.json"
+    report_path.write_text(json.dumps(report, indent=2), encoding="utf-8")
+    print(json.dumps(report, indent=2))
+
+
+# -----------------------------------------------------------------------------
+# CLI
+# -----------------------------------------------------------------------------
+
+
+def build_arg_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="WikiText motif-profile ensemble")
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    def add_common(sub: argparse.ArgumentParser) -> None:
+        sub.add_argument("--train-path", type=str, required=True)
+        sub.add_argument("--valid-path", type=str, required=True)
+        sub.add_argument("--test-path", type=str, default=None)
+        sub.add_argument("--tokenized", action="store_true", help="Interpret input paths as pretokenized tensors (.pt/.npy/.bin)")
+        sub.add_argument("--output-dir", type=str, required=True)
+        sub.add_argument("--scale", type=str, default="0.125x")
+        sub.add_argument("--block-size", type=int, default=256)
+        sub.add_argument("--batch-size", type=int, default=8)
+        sub.add_argument("--device", type=str, default="cpu")
+        sub.add_argument("--lr", type=float, default=3e-4)
+        sub.add_argument("--weight-decay", type=float, default=0.01)
+        sub.add_argument("--epochs", type=int, default=1)
+        sub.add_argument("--max-steps", type=int, default=None)
+        sub.add_argument("--grad-clip", type=float, default=1.0)
+        sub.add_argument("--log-every", type=int, default=50)
+
+    train_members = subparsers.add_parser("train-members", help="Train individual motif-profile members")
+    add_common(train_members)
+    train_members.add_argument("--no-text-motif-moe", action="store_true")
+
+    train_ensemble = subparsers.add_parser("train-ensemble", help="Train the ensemble gate over saved members")
+    add_common(train_ensemble)
+    train_ensemble.add_argument("--joint-finetune", action="store_true", help="Also fine-tune member models while training the gate")
+    train_ensemble.add_argument("--gate-hidden-dim", type=int, default=128)
+    train_ensemble.add_argument("--gate-mode", type=str, default="contextual", choices=["contextual", "static"])
+    train_ensemble.add_argument("--member-dropout-p", type=float, default=0.05)
+    train_ensemble.add_argument("--gate-temperature", type=float, default=1.0)
+    train_ensemble.add_argument("--gate-entropy-bonus-coef", type=float, default=0.0)
+
+    evaluate = subparsers.add_parser("evaluate", help="Evaluate trained members and the ensemble")
+    add_common(evaluate)
+    evaluate.add_argument("--gate-hidden-dim", type=int, default=128)
+    evaluate.add_argument("--gate-mode", type=str, default="contextual", choices=["contextual", "static"])
+    evaluate.add_argument("--gate-temperature", type=float, default=1.0)
+
+    return parser
+
+
+def main() -> None:
+    parser = build_arg_parser()
+    args = parser.parse_args()
+    if args.command == "train-members":
+        train_members_cli(args)
+    elif args.command == "train-ensemble":
+        train_ensemble_cli(args)
+    elif args.command == "evaluate":
+        evaluate_cli(args)
+    else:  # pragma: no cover
+        raise KeyError(args.command)
+
+
+if __name__ == "__main__":
+    main()