ssmforge 0.1.0__py3-none-any.whl

ssmforge/__init__.py

@@ -0,0 +1,10 @@
+"""SSMForge: convert pretrained transformers to hybrid SSM/attention models."""
+
+__version__ = "0.1.0"
+
+from ssmforge.pipeline import convert
+from ssmforge.result import ConversionResult
+from ssmforge import recipes  # registers recipes on import
+from ssmforge.recipes import list_recipes
+
+__all__ = ["convert", "ConversionResult", "list_recipes", "__version__"]

ssmforge/benchmark/__init__.py

@@ -0,0 +1,4 @@
+from ssmforge.benchmark.quality import compute_perplexity
+from ssmforge.benchmark.long_context import benchmark_long_context
+
+__all__ = ["compute_perplexity", "benchmark_long_context"]

ssmforge/benchmark/long_context.py

@@ -0,0 +1,43 @@
+"""Long-context benchmark: memory + speed at increasing context lengths."""
+
+from __future__ import annotations
+
+import time
+from typing import Any
+
+
+def benchmark_long_context(
+    model: Any,
+    tokenizer: Any,
+    context_lengths: list[int] | None = None,
+    device: str = "cpu",
+) -> dict[int, dict[str, float]]:
+    """Measure peak memory and tokens/sec at each context length."""
+    if context_lengths is None:
+        context_lengths = [4096, 32768, 131072]
+    results = {}
+    for ctx in context_lengths:
+        try:
+            import torch
+
+            if device == "cuda" and torch.cuda.is_available():
+                torch.cuda.reset_peak_memory_stats()
+
+            input_ids = torch.randint(0, tokenizer.vocab_size, (1, ctx))
+
+            start = time.time()
+            with torch.no_grad():
+                _ = model(input_ids=input_ids)
+            elapsed = time.time() - start
+
+            results[ctx] = {
+                "elapsed_sec": elapsed,
+                "tokens_per_sec": ctx / elapsed if elapsed > 0 else 0,
+                "peak_memory_mb": (
+                    torch.cuda.max_memory_allocated() / 1e6
+                    if device == "cuda" and torch.cuda.is_available() else 0
+                ),
+            }
+        except Exception as e:
+            results[ctx] = {"error": str(e)}
+    return results

ssmforge/benchmark/quality.py

@@ -0,0 +1,27 @@
+"""Quality benchmark: perplexity vs teacher."""
+
+from __future__ import annotations
+
+import math
+
+import torch
+
+
+@torch.no_grad()
+def compute_perplexity(model, tokenizer, texts: list[str], max_length: int = 512) -> float:
+    """Compute perplexity on a list of texts."""
+    model.eval()
+    total_loss = 0.0
+    total_tokens = 0
+    for text in texts:
+        enc = tokenizer(text, return_tensors="pt", truncation=True, max_length=max_length)
+        input_ids = enc.input_ids
+        if input_ids.shape[1] < 2:
+            continue
+        outputs = model(input_ids=input_ids, labels=input_ids)
+        n_tokens = input_ids.shape[1] - 1
+        total_loss += outputs.loss.item() * n_tokens
+        total_tokens += n_tokens
+    if total_tokens == 0:
+        return float("inf")
+    return math.exp(total_loss / total_tokens)

ssmforge/cli.py

@@ -0,0 +1,71 @@
+"""Command-line interface for SSMForge."""
+
+from __future__ import annotations
+
+import argparse
+import sys
+
+from ssmforge import convert
+from ssmforge.exceptions import SSMForgeError
+from ssmforge.recipes import list_recipes
+
+
+def main(argv: list[str] | None = None) -> None:
+    """Entry point for the `ssmforge` command."""
+    parser = argparse.ArgumentParser(
+        prog="ssmforge",
+        description="Convert pretrained transformers to hybrid SSM/attention models.",
+    )
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    convert_p = subparsers.add_parser("convert", help="Convert a model to hybrid SSM/attention + GGUF")
+    convert_p.add_argument("source", help="HF model id or local path")
+    convert_p.add_argument("--recipe", default="hybrid-25", help="Recipe name (default: hybrid-25)")
+    convert_p.add_argument(
+        "--quantize",
+        default="Q4_K_M",
+        choices=["F16", "Q8_0", "Q5_K_M", "Q4_K_M", "Q4_K_S"],
+        help="GGUF quantization type (default: Q4_K_M)",
+    )
+    convert_p.add_argument("--output", default="./out", help="Output directory (default: ./out)")
+    convert_p.add_argument("--calibration-data", default=None, help="Calibration data source")
+    convert_p.add_argument("--verify", action="store_true", help="Run Stage 6 verification (slow)")
+    convert_p.add_argument("--dry-run", action="store_true", help="Plan only, no export")
+    convert_p.add_argument("--experimental", action="store_true", help="Allow experimental recipes")
+    convert_p.add_argument("--strict-verify", action="store_true", help="Promote verify warnings to errors")
+    convert_p.add_argument("--debug", action="store_true", help="Show full tracebacks on error")
+
+    subparsers.add_parser("list-recipes", help="List registered recipes")
+
+    args = parser.parse_args(argv)
+
+    try:
+        if args.command == "list-recipes":
+            print("Registered recipes:")
+            for name in list_recipes():
+                print(f"  - {name}")
+            sys.exit(0)
+        elif args.command == "convert":
+            result = convert(
+                source=args.source,
+                recipe=args.recipe,
+                quantize=args.quantize,
+                output_dir=args.output,
+                calibration_data=args.calibration_data,
+                verify=args.verify,
+                dry_run=args.dry_run,
+                experimental=args.experimental,
+            )
+            print(f"GGUF: {result.gguf_path}")
+            print(f"Manifest: {result.manifest_path}")
+            print(f"Stats: {result.stats}")
+            sys.exit(0)
+    except SSMForgeError as e:
+        if args.debug:
+            raise
+        print(str(e), file=sys.stderr)
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

ssmforge/config.py

@@ -0,0 +1,41 @@
+"""Pydantic schemas for SSMForge configuration."""
+
+from __future__ import annotations
+
+from enum import Enum
+from pydantic import BaseModel, Field
+
+
+class LayerType(str, Enum):
+    ATTENTION = "attention"
+    SSM = "ssm"
+
+
+class LayerSpec(BaseModel):
+    """Specification for one layer in the hybrid model."""
+
+    layer_type: LayerType
+    index: int = Field(default=0, description="0-based position in the original model")
+    freeze_mlp: bool = Field(default=True, description="Freeze MLP weights during distillation (MambaInLlama recipe)")
+
+
+class TrainingStage(BaseModel):
+    """One stage of the distillation training."""
+
+    name: str
+    epochs: int = 1
+    learning_rate: float = 1e-5
+    batch_size: int = 1
+    gradient_accumulation_steps: int = 8
+    freeze_mlp: bool = False
+    stepwise: bool = Field(default=False, description="Train one layer at a time")
+
+
+class DistillationConfig(BaseModel):
+    """Full distillation configuration for a recipe."""
+
+    stages: list[TrainingStage]
+    kl_weight: float = Field(default=0.7, ge=0.0, le=1.0)
+    seqkd_weight: float = Field(default=0.3, ge=0.0, le=1.0)
+    max_seq_length: int = 2048
+    warmup_steps: int = 100

ssmforge/converters/__init__.py

@@ -0,0 +1,17 @@
+from ssmforge.converters.base import (
+    ArchitectureConverter,
+    ArchitectureConverterRegistry,
+)
+from ssmforge.converters.llama_to_hybrid import LlamaToHybridConverter
+from ssmforge.converters.mistral_to_hybrid import MistralToHybridConverter
+
+# Trigger registration via direct import (registers on import via module-level code)
+import ssmforge.converters.llama_to_hybrid  # noqa: F401
+import ssmforge.converters.mistral_to_hybrid  # noqa: F401
+
+__all__ = [
+    "ArchitectureConverter",
+    "ArchitectureConverterRegistry",
+    "LlamaToHybridConverter",
+    "MistralToHybridConverter",
+]

ssmforge/converters/base.py

@@ -0,0 +1,56 @@
+"""Architecture converter base + registry."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+from ssmforge.config import LayerSpec
+
+
+class ArchitectureConverter(ABC):
+    """Converts a source model's state dict to a hybrid state dict per a recipe plan."""
+
+    source_arch: str
+    target_arch: str
+
+    @abstractmethod
+    def convert_state_dict(self, src: dict, plan: list[LayerSpec]) -> dict:
+        """Pure function: source SD + plan → target SD.
+
+        For MVP: copies embeddings/MLP/norm/head verbatim, replaces attention
+        layers marked as SSM with placeholder Mamba2 tensors.
+        """
+        ...
+
+    def build_model(self, src_config: Any, target_sd: dict) -> Any:
+        """Build the target model shell + load weights. Raises NotImplementedError in MVP."""
+        raise NotImplementedError("Model construction requires mamba-ssm; see Task 6")
+
+    def verify_round_trip(self, src_model: Any, target_model: Any, prompts: list[str]) -> bool:
+        """Forward-pass sanity check. Raises NotImplementedError in MVP."""
+        raise NotImplementedError("Verification requires full model construction; see Task 6")
+
+
+class ArchitectureConverterRegistry:
+    _registry: dict[str, type[ArchitectureConverter]] = {}
+
+    @classmethod
+    def register(cls, arch: str, converter_cls: type[ArchitectureConverter]) -> None:
+        cls._registry[arch] = converter_cls
+
+    @classmethod
+    def get(cls, arch: str) -> ArchitectureConverter:
+        from ssmforge.exceptions import UnsupportedArchitectureError
+
+        if arch not in cls._registry:
+            raise UnsupportedArchitectureError(
+                arch=arch,
+                version="0.1.0",
+                supported=cls.list_supported(),
+            )
+        return cls._registry[arch]()
+
+    @classmethod
+    def list_supported(cls) -> list[str]:
+        return sorted(cls._registry.keys())

ssmforge/converters/llama_to_hybrid.py

@@ -0,0 +1,61 @@
+"""State-dict surgery for Llama → hybrid Llama+Mamba2.
+
+For MVP: produces a target state dict where:
+- Embeddings, MLP, LayerNorm, output head are copied verbatim
+- Attention layers marked as SSM in the plan are replaced with placeholder
+  Mamba2 tensors (random init of correct shape, ready for distillation)
+"""
+
+from __future__ import annotations
+
+from ssmforge.config import LayerSpec, LayerType
+from ssmforge.converters.base import ArchitectureConverter, ArchitectureConverterRegistry
+from ssmforge.converters.weight_init import init_mamba2_from_attention
+
+
+class LlamaToHybridConverter(ArchitectureConverter):
+    source_arch = "llama"
+    target_arch = "hybrid-llama-mamba2"
+
+    def convert_state_dict(self, src: dict, plan: list[LayerSpec]) -> dict:
+        target: dict = {}
+
+        for key, value in src.items():
+            if "layers." not in key:
+                target[key] = value
+
+        hidden_size = src.get("_hidden_size", 2048)
+
+        for spec in plan:
+            layer_idx = spec.index
+            if spec.layer_type == LayerType.ATTENTION:
+                for key, value in src.items():
+                    if key.startswith(f"model.layers.{layer_idx}."):
+                        target[key] = value
+            else:  # SSM
+                # Copy MLP + LN verbatim
+                for key, value in src.items():
+                    if (
+                        key.startswith(f"model.layers.{layer_idx}.")
+                        and (
+                            "mlp." in key
+                            or "post_attention_layernorm" in key
+                            or "input_layernorm" in key
+                        )
+                    ):
+                        target[key] = value
+                # Initialize Mamba2 weights from attention weights
+                attention_keys = [
+                    k for k in src.keys()
+                    if k.startswith(f"model.layers.{layer_idx}.self_attn.")
+                ]
+                attention_sd = {k.split("self_attn.")[-1]: src[k] for k in attention_keys}
+                mamba_sd = init_mamba2_from_attention(attention_sd, hidden_size=hidden_size)
+                for mk, mv in mamba_sd.items():
+                    target[f"model.layers.{layer_idx}.mamba.{mk}"] = mv
+
+        return target
+
+
+# Auto-register on import
+ArchitectureConverterRegistry.register("llama", LlamaToHybridConverter)

ssmforge/converters/mistral_to_hybrid.py

@@ -0,0 +1,26 @@
+"""State-dict surgery for Mistral → hybrid Mistral+Mamba2.
+
+Mistral uses sliding window attention by default. We preserve this in the
+attention layers that survive, and replace others with Mamba2.
+"""
+
+from __future__ import annotations
+
+from ssmforge.config import LayerSpec
+from ssmforge.converters.base import ArchitectureConverter
+from ssmforge.converters.llama_to_hybrid import LlamaToHybridConverter
+
+
+class MistralToHybridConverter(ArchitectureConverter):
+    source_arch = "mistral"
+    target_arch = "hybrid-mistral-mamba2"
+
+    def convert_state_dict(self, src: dict, plan: list[LayerSpec]) -> dict:
+        llama_converter = LlamaToHybridConverter()
+        return llama_converter.convert_state_dict(src, plan)
+
+
+# Auto-register on import
+from ssmforge.converters.base import ArchitectureConverterRegistry  # noqa: E402
+
+ArchitectureConverterRegistry.register("mistral", MistralToHybridConverter)

ssmforge/converters/weight_init.py

@@ -0,0 +1,54 @@
+"""Initialize Mamba2 weights from attention weights.
+
+Approach (MambaInLlama recipe):
+- Reuse attention's o_proj → Mamba2's out_proj
+- Initialize in_proj, conv1d, x_proj, dt_bias, A_log, D with small random values
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+
+
+def init_mamba2_from_attention(attention_sd: dict, hidden_size: int) -> dict:
+    """Convert attention weights to Mamba2 initial weights.
+
+    Args:
+        attention_sd: attention submodule state dict (q_proj, k_proj, v_proj, o_proj)
+        hidden_size: model hidden dimension
+
+    Returns:
+        Mamba2 submodule state dict (in_proj, conv1d, x_proj, dt_bias, A_log, D, out_proj)
+    """
+    out_proj_weight = attention_sd.get("o_proj.weight")
+    if out_proj_weight is None:
+        out_proj_weight = torch.empty(hidden_size, hidden_size)
+        torch.nn.init.xavier_uniform_(out_proj_weight)
+
+    expand = 2
+    d_inner = hidden_size * expand
+
+    in_proj = torch.empty(d_inner * 2, hidden_size)
+    torch.nn.init.xavier_uniform_(in_proj)
+
+    conv1d_weight = torch.empty(d_inner, 1, 4)
+    torch.nn.init.kaiming_uniform_(conv1d_weight, a=math.sqrt(5))
+
+    x_proj_weight = torch.empty(d_inner // 2, d_inner)
+    torch.nn.init.xavier_uniform_(x_proj_weight)
+
+    dt_bias = torch.zeros(d_inner // 2)
+    A_log = torch.log(torch.empty(d_inner // 2).uniform_(1.0, 16.0))
+    D = torch.ones(d_inner)
+
+    return {
+        "in_proj.weight": in_proj,
+        "conv1d.weight": conv1d_weight,
+        "x_proj.weight": x_proj_weight,
+        "dt_bias": dt_bias,
+        "A_log": A_log,
+        "D": D,
+        "out_proj.weight": out_proj_weight,
+    }

ssmforge/distillation/__init__.py

@@ -0,0 +1,12 @@
+from ssmforge.distillation.collator import KLDistillationCollator
+from ssmforge.distillation.loss import compute_kl_loss, compute_seqkd_loss
+from ssmforge.distillation.calibration import CalibrationDataLoader
+from ssmforge.distillation.trainer import DistillationTrainer
+
+__all__ = [
+    "KLDistillationCollator",
+    "compute_kl_loss",
+    "compute_seqkd_loss",
+    "CalibrationDataLoader",
+    "DistillationTrainer",
+]

ssmforge/distillation/calibration.py

@@ -0,0 +1,82 @@
+"""Calibration data loaders for distillation.
+
+Three sources:
+1. Local text file (one sample per line)
+2. HuggingFace dataset id
+3. Built-in default set (~1M tokens from Wikipedia + C4)
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Optional
+
+from ssmforge.exceptions import CalibrationDataError
+
+
+_BUILTIN_CALIBRATION_SAMPLES = [
+    "The quick brown fox jumps over the lazy dog.",
+    "In the beginning was the Word, and the Word was with God.",
+    "To be, or not to be, that is the question.",
+    "All happy families are alike; each unhappy family is unhappy in its own way.",
+    "It was the best of times, it was the worst of times.",
+    "Call me Ishmael. Some years ago—never mind how long precisely—",
+    "It is a truth universally acknowledged, that a single man in possession of a good fortune, must be in want of a wife.",
+    "Whether I shall turn out to be the hero of my own life, or whether that station will be held by anybody else, these pages must show.",
+    "The only way to do great work is to love what you do.",
+    "Innovation distinguishes between a leader and a follower.",
+]
+
+
+def _load_hf_dataset(name: str, split: str = "train", max_samples: int = 1000) -> list[dict]:
+    """Load a HuggingFace dataset. Imported lazily."""
+    from datasets import load_dataset
+
+    ds = load_dataset(name, split=f"{split}[:{max_samples}]")
+    return list(ds)
+
+
+class CalibrationDataLoader:
+    def __init__(self, source: Optional[str] = None, max_samples: int = 1000):
+        self.source = source
+        self.max_samples = max_samples
+
+    def load(self) -> list[str]:
+        if self.source is None:
+            return self._load_builtin()
+
+        path = Path(self.source)
+        if path.exists() and path.is_file():
+            return self._load_text_file(path)
+
+        return self._load_hf()
+
+    def _load_builtin(self) -> list[str]:
+        samples = []
+        while len(samples) < self.max_samples:
+            samples.extend(_BUILTIN_CALIBRATION_SAMPLES)
+        return samples[:self.max_samples]
+
+    def _load_text_file(self, path: Path) -> list[str]:
+        try:
+            lines = [l.strip() for l in path.read_text().splitlines() if l.strip()]
+        except Exception as e:
+            raise CalibrationDataError(reason=f"Could not read {path}: {e}") from e
+        if not lines:
+            raise CalibrationDataError(reason=f"No non-empty lines in {path}")
+        return lines[:self.max_samples]
+
+    def _load_hf(self) -> list[str]:
+        try:
+            rows = _load_hf_dataset(self.source, max_samples=self.max_samples)
+        except Exception as e:
+            raise CalibrationDataError(reason=f"Could not load HF dataset {self.source}: {e}") from e
+        texts = []
+        for row in rows:
+            if isinstance(row, dict) and "text" in row:
+                texts.append(row["text"])
+            elif isinstance(row, str):
+                texts.append(row)
+        if not texts:
+            raise CalibrationDataError(reason=f"No text field in HF dataset {self.source}")
+        return texts

ssmforge/distillation/collator.py

@@ -0,0 +1,31 @@
+"""Data collator for KL distillation batches."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+
+class KLDistillationCollator:
+    """Tokenizes text inputs for both teacher and student models.
+
+    For MVP: produces a single input_ids tensor. Teacher and student share
+    the tokenizer (a hybrid SSM model and its teacher transformer both use
+    the same tokenizer, by design).
+    """
+
+    def __init__(self, tokenizer: Any, max_length: int = 2048):
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+
+    def __call__(self, batch: list[str]) -> dict[str, torch.Tensor]:
+        encoded = self.tokenizer(
+            batch,
+            padding="max_length",
+            truncation=True,
+            max_length=self.max_length,
+            return_tensors="pt",
+        )
+        encoded["labels"] = encoded["input_ids"].clone()
+        return encoded

ssmforge/distillation/loss.py

@@ -0,0 +1,47 @@
+"""Distillation loss functions: KL divergence + SequenceKD."""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+
+
+def compute_kl_loss(
+    student_logits: torch.Tensor,
+    teacher_logits: torch.Tensor,
+    alpha: float = 0.7,
+    temperature: float = 1.0,
+    labels: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Word-level KL distillation loss.
+
+    Loss = alpha * KL(student || teacher, T=temperature)
+         + (1 - alpha) * CE(student, labels)  [if labels provided]
+    """
+    kl = F.kl_div(
+        input=F.log_softmax(student_logits / temperature, dim=-1),
+        target=F.softmax(teacher_logits / temperature, dim=-1),
+        reduction="batchmean",
+    ) * (temperature ** 2)
+
+    if labels is None:
+        return alpha * kl
+
+    ce = F.cross_entropy(
+        student_logits.view(-1, student_logits.size(-1)),
+        labels.view(-1),
+        ignore_index=-100,
+    )
+    return alpha * kl + (1 - alpha) * ce
+
+
+def compute_seqkd_loss(
+    student_logits: torch.Tensor,
+    teacher_labels: torch.Tensor,
+) -> torch.Tensor:
+    """Sequence-level KD: use teacher's argmax as pseudo-labels."""
+    return F.cross_entropy(
+        student_logits.view(-1, student_logits.size(-1)),
+        teacher_labels.view(-1),
+        ignore_index=-100,
+    )