feat: Add capabilities/speculative.py

capabilities/speculative.py

@@ -0,0 +1,439 @@
+"""
+Speculative Decoding Module for MiniMind Max2
+Use small draft model to accelerate large model inference.
+"""
+
+from dataclasses import dataclass
+from typing import List, Optional, Dict, Any, Tuple
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import time
+
+
+@dataclass
+class SpeculativeConfig:
+    """Configuration for speculative decoding."""
+    # Speculation settings
+    num_speculative_tokens: int = 5  # Number of tokens to speculate
+    max_speculation_length: int = 8
+
+    # Acceptance settings
+    acceptance_method: str = "rejection"  # rejection, nucleus
+    temperature: float = 1.0
+    top_p: float = 0.95
+
+    # Performance tuning
+    adaptive_speculation: bool = True  # Adjust speculation based on acceptance rate
+    min_speculative_tokens: int = 2
+    max_speculative_tokens: int = 10
+    target_acceptance_rate: float = 0.8
+
+
+class DraftModel:
+    """
+    Wrapper for draft model in speculative decoding.
+    Typically a smaller, faster model (e.g., max2-nano for max2-pro).
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+        tokenizer = None,
+        device: str = "cuda",
+    ):
+        self.model = model
+        self.tokenizer = tokenizer
+        self.device = device
+        self.model.eval()
+
+    @torch.no_grad()
+    def speculate(
+        self,
+        input_ids: torch.Tensor,
+        num_tokens: int = 5,
+        temperature: float = 1.0,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Generate speculative tokens.
+
+        Args:
+            input_ids: Current input sequence [batch, seq_len]
+            num_tokens: Number of tokens to speculate
+            temperature: Sampling temperature
+
+        Returns:
+            Tuple of (speculated_tokens, speculated_probs)
+        """
+        batch_size = input_ids.shape[0]
+        speculated_tokens = []
+        speculated_probs = []
+
+        current_ids = input_ids
+
+        for _ in range(num_tokens):
+            # Forward pass
+            _, logits, _, _ = self.model(current_ids)
+            next_logits = logits[:, -1, :] / temperature
+
+            # Sample
+            probs = F.softmax(next_logits, dim=-1)
+            next_token = torch.multinomial(probs, num_samples=1)
+
+            # Get probability of selected token
+            token_prob = probs.gather(1, next_token)
+
+            speculated_tokens.append(next_token)
+            speculated_probs.append(token_prob)
+
+            # Append to sequence
+            current_ids = torch.cat([current_ids, next_token], dim=1)
+
+        # Stack results
+        speculated_tokens = torch.cat(speculated_tokens, dim=1)  # [batch, num_tokens]
+        speculated_probs = torch.cat(speculated_probs, dim=1)  # [batch, num_tokens]
+
+        return speculated_tokens, speculated_probs
+
+
+class SpeculativeDecoder:
+    """
+    Speculative decoding for accelerated generation.
+    Uses a small draft model to propose tokens, verified by target model.
+    """
+
+    def __init__(
+        self,
+        target_model: nn.Module,
+        draft_model: nn.Module,
+        config: Optional[SpeculativeConfig] = None,
+        device: str = "cuda",
+    ):
+        self.target = target_model
+        self.draft = DraftModel(draft_model, device=device)
+        self.config = config or SpeculativeConfig()
+        self.device = device
+
+        # Statistics
+        self.total_generated = 0
+        self.total_accepted = 0
+        self.speculation_lengths = []
+
+    def _rejection_sampling(
+        self,
+        draft_probs: torch.Tensor,
+        target_probs: torch.Tensor,
+        draft_tokens: torch.Tensor,
+    ) -> Tuple[torch.Tensor, int]:
+        """
+        Rejection sampling for token acceptance.
+
+        Returns:
+            Tuple of (accepted_mask, num_accepted)
+        """
+        batch_size, num_tokens = draft_tokens.shape
+
+        # Compute acceptance probability: min(1, target_p / draft_p)
+        acceptance_probs = torch.min(
+            torch.ones_like(draft_probs),
+            target_probs / (draft_probs + 1e-10),
+        )
+
+        # Sample uniform for rejection test
+        uniform = torch.rand_like(acceptance_probs)
+        accepted = uniform < acceptance_probs
+
+        # Find first rejection point
+        accepted_mask = torch.cumprod(accepted.float(), dim=1).bool()
+        num_accepted = accepted_mask.sum(dim=1).min().item()
+
+        return accepted_mask, num_accepted
+
+    @torch.no_grad()
+    def generate_step(
+        self,
+        input_ids: torch.Tensor,
+        num_speculative: Optional[int] = None,
+    ) -> Tuple[torch.Tensor, Dict[str, Any]]:
+        """
+        Single speculative generation step.
+
+        Args:
+            input_ids: Current sequence [batch, seq_len]
+            num_speculative: Number of tokens to speculate (uses config if None)
+
+        Returns:
+            New tokens and statistics
+        """
+        num_spec = num_speculative or self.config.num_speculative_tokens
+
+        # Phase 1: Draft model speculation
+        draft_tokens, draft_probs = self.draft.speculate(
+            input_ids,
+            num_tokens=num_spec,
+            temperature=self.config.temperature,
+        )
+
+        # Phase 2: Target model verification (single forward pass)
+        spec_input = torch.cat([input_ids, draft_tokens], dim=1)
+        _, target_logits, _, _ = self.target(spec_input)
+
+        # Get target probabilities for draft tokens
+        target_probs = F.softmax(target_logits[:, -num_spec-1:-1, :] / self.config.temperature, dim=-1)
+        target_probs_selected = target_probs.gather(2, draft_tokens.unsqueeze(-1)).squeeze(-1)
+
+        # Phase 3: Rejection sampling
+        accepted_mask, num_accepted = self._rejection_sampling(
+            draft_probs,
+            target_probs_selected,
+            draft_tokens,
+        )
+
+        # Accept verified tokens
+        if num_accepted > 0:
+            new_tokens = draft_tokens[:, :num_accepted]
+        else:
+            new_tokens = torch.empty(input_ids.shape[0], 0, dtype=torch.long, device=self.device)
+
+        # Sample one more token from target if not all accepted
+        if num_accepted < num_spec:
+            # Resample from target distribution at rejection point
+            next_logits = target_logits[:, input_ids.shape[1] + num_accepted - 1, :]
+            next_probs = F.softmax(next_logits / self.config.temperature, dim=-1)
+            bonus_token = torch.multinomial(next_probs, num_samples=1)
+            new_tokens = torch.cat([new_tokens, bonus_token], dim=1)
+
+        # Statistics
+        self.total_generated += new_tokens.shape[1]
+        self.total_accepted += num_accepted
+        self.speculation_lengths.append(num_spec)
+
+        stats = {
+            "num_speculated": num_spec,
+            "num_accepted": num_accepted,
+            "num_generated": new_tokens.shape[1],
+            "acceptance_rate": num_accepted / num_spec if num_spec > 0 else 0,
+        }
+
+        return new_tokens, stats
+
+    @torch.no_grad()
+    def generate(
+        self,
+        input_ids: torch.Tensor,
+        max_new_tokens: int = 100,
+        eos_token_id: Optional[int] = None,
+    ) -> Tuple[torch.Tensor, Dict[str, Any]]:
+        """
+        Full speculative generation.
+
+        Args:
+            input_ids: Initial input [batch, seq_len]
+            max_new_tokens: Maximum tokens to generate
+            eos_token_id: EOS token to stop generation
+
+        Returns:
+            Generated sequence and statistics
+        """
+        self.target.eval()
+
+        generated = input_ids.clone()
+        total_stats = {
+            "steps": 0,
+            "tokens_generated": 0,
+            "acceptance_rates": [],
+        }
+
+        start_time = time.time()
+        num_speculative = self.config.num_speculative_tokens
+
+        while total_stats["tokens_generated"] < max_new_tokens:
+            # Speculative step
+            new_tokens, step_stats = self.generate_step(generated, num_speculative)
+
+            if new_tokens.shape[1] == 0:
+                break
+
+            generated = torch.cat([generated, new_tokens], dim=1)
+
+            # Update stats
+            total_stats["steps"] += 1
+            total_stats["tokens_generated"] += new_tokens.shape[1]
+            total_stats["acceptance_rates"].append(step_stats["acceptance_rate"])
+
+            # Check for EOS
+            if eos_token_id is not None and (new_tokens == eos_token_id).any():
+                break
+
+            # Adaptive speculation
+            if self.config.adaptive_speculation:
+                avg_acceptance = sum(total_stats["acceptance_rates"][-5:]) / min(5, len(total_stats["acceptance_rates"]))
+                if avg_acceptance > self.config.target_acceptance_rate:
+                    num_speculative = min(num_speculative + 1, self.config.max_speculative_tokens)
+                elif avg_acceptance < self.config.target_acceptance_rate - 0.1:
+                    num_speculative = max(num_speculative - 1, self.config.min_speculative_tokens)
+
+        end_time = time.time()
+
+        total_stats["time_seconds"] = end_time - start_time
+        total_stats["tokens_per_second"] = total_stats["tokens_generated"] / total_stats["time_seconds"]
+        total_stats["avg_acceptance_rate"] = sum(total_stats["acceptance_rates"]) / max(1, len(total_stats["acceptance_rates"]))
+        total_stats["avg_tokens_per_step"] = total_stats["tokens_generated"] / max(1, total_stats["steps"])
+
+        return generated, total_stats
+
+    def get_statistics(self) -> Dict[str, float]:
+        """Get overall statistics."""
+        return {
+            "total_generated": self.total_generated,
+            "total_accepted": self.total_accepted,
+            "overall_acceptance_rate": self.total_accepted / max(1, self.total_generated),
+            "avg_speculation_length": sum(self.speculation_lengths) / max(1, len(self.speculation_lengths)),
+        }
+
+    def reset_statistics(self):
+        """Reset statistics counters."""
+        self.total_generated = 0
+        self.total_accepted = 0
+        self.speculation_lengths = []
+
+
+class TreeSpeculativeDecoder(SpeculativeDecoder):
+    """
+    Tree-based speculative decoding for higher acceptance rates.
+    Generates multiple speculation branches.
+    """
+
+    def __init__(
+        self,
+        target_model: nn.Module,
+        draft_model: nn.Module,
+        num_branches: int = 3,
+        config: Optional[SpeculativeConfig] = None,
+        device: str = "cuda",
+    ):
+        super().__init__(target_model, draft_model, config, device)
+        self.num_branches = num_branches
+
+    @torch.no_grad()
+    def generate_tree(
+        self,
+        input_ids: torch.Tensor,
+        depth: int = 3,
+    ) -> List[Tuple[torch.Tensor, torch.Tensor]]:
+        """
+        Generate tree of speculative tokens.
+
+        Returns:
+            List of (tokens, probs) tuples for each branch
+        """
+        branches = []
+
+        # Generate multiple branches from draft model
+        for _ in range(self.num_branches):
+            tokens, probs = self.draft.speculate(
+                input_ids,
+                num_tokens=depth,
+                temperature=self.config.temperature,
+            )
+            branches.append((tokens, probs))
+
+        return branches
+
+    @torch.no_grad()
+    def generate_step(
+        self,
+        input_ids: torch.Tensor,
+        num_speculative: Optional[int] = None,
+    ) -> Tuple[torch.Tensor, Dict[str, Any]]:
+        """Tree-based speculative step."""
+        num_spec = num_speculative or self.config.num_speculative_tokens
+
+        # Generate tree of speculations
+        branches = self.generate_tree(input_ids, num_spec)
+
+        best_tokens = None
+        best_accepted = 0
+
+        # Verify each branch and pick best
+        for draft_tokens, draft_probs in branches:
+            spec_input = torch.cat([input_ids, draft_tokens], dim=1)
+            _, target_logits, _, _ = self.target(spec_input)
+
+            target_probs = F.softmax(
+                target_logits[:, -num_spec-1:-1, :] / self.config.temperature, dim=-1
+            )
+            target_probs_selected = target_probs.gather(2, draft_tokens.unsqueeze(-1)).squeeze(-1)
+
+            _, num_accepted = self._rejection_sampling(
+                draft_probs,
+                target_probs_selected,
+                draft_tokens,
+            )
+
+            if num_accepted > best_accepted:
+                best_accepted = num_accepted
+                best_tokens = draft_tokens[:, :num_accepted]
+
+        if best_tokens is None or best_tokens.shape[1] == 0:
+            # Fallback: sample from target
+            _, logits, _, _ = self.target(input_ids)
+            probs = F.softmax(logits[:, -1, :] / self.config.temperature, dim=-1)
+            best_tokens = torch.multinomial(probs, num_samples=1)
+            best_accepted = 0
+
+        stats = {
+            "num_speculated": num_spec * self.num_branches,
+            "num_accepted": best_accepted,
+            "num_generated": best_tokens.shape[1],
+            "acceptance_rate": best_accepted / num_spec if num_spec > 0 else 0,
+            "num_branches": self.num_branches,
+        }
+
+        return best_tokens, stats
+
+
+def benchmark_speculative_decoding(
+    target_model: nn.Module,
+    draft_model: nn.Module,
+    input_ids: torch.Tensor,
+    num_tokens: int = 100,
+    device: str = "cuda",
+) -> Dict[str, Any]:
+    """
+    Benchmark speculative decoding vs standard generation.
+    """
+    import time
+
+    # Standard generation
+    target_model.eval()
+    start = time.time()
+    with torch.no_grad():
+        standard_output = target_model.generate(
+            input_ids,
+            max_new_tokens=num_tokens,
+        )
+    standard_time = time.time() - start
+
+    # Speculative generation
+    decoder = SpeculativeDecoder(target_model, draft_model, device=device)
+    start = time.time()
+    spec_output, spec_stats = decoder.generate(
+        input_ids,
+        max_new_tokens=num_tokens,
+    )
+    spec_time = time.time() - start
+
+    return {
+        "standard": {
+            "time": standard_time,
+            "tokens_per_second": num_tokens / standard_time,
+        },
+        "speculative": {
+            "time": spec_time,
+            "tokens_per_second": spec_stats["tokens_per_second"],
+            "acceptance_rate": spec_stats["avg_acceptance_rate"],
+            "avg_tokens_per_step": spec_stats["avg_tokens_per_step"],
+        },
+        "speedup": standard_time / spec_time,
+    }