src/minimal_sd.py

#!/usr/bin/env python3
"""Minimal speculative decoding implementation for toy claims C4 and C5."""
import torch
import numpy as np
from transformers import AutoModelForCausalLM, AutoTokenizer

def sample_from_logits(logits, temperature=1.0):
    """Sample a token from logits using temperature."""
    if temperature == 0:
        return logits.argmax(dim=-1)
    probs = torch.softmax(logits / temperature, dim=-1)
    return torch.multinomial(probs, num_samples=1).squeeze(-1)

def speculative_decode(target_model, draft_model, tokenizer, prompt, max_draft_len=4, max_new_tokens=50, device='cpu', draft_temp=1.2):
    """
    Minimal speculative decoding.
    Draft model uses temperature sampling; target uses greedy.
    """
    target_model.to(device)
    draft_model.to(device)
    target_model.eval()
    draft_model.eval()
    
    input_ids = tokenizer.encode(prompt, return_tensors='pt').to(device)
    init_len = input_ids.shape[1]
    acceptance_lengths = []
    n_draft_tokens = 0
    n_target_calls = 0
    
    with torch.no_grad():
        while input_ids.shape[1] < init_len + max_new_tokens:
            current_len = input_ids.shape[1]
            
            # Draft model generates k tokens with temperature sampling
            draft_ids = input_ids.clone()
            draft_tokens = []
            for _ in range(max_draft_len):
                draft_logits = draft_model(draft_ids).logits[:, -1, :]
                next_token = sample_from_logits(draft_logits, temperature=draft_temp)
                draft_tokens.append(next_token.item())
                draft_ids = torch.cat([draft_ids, next_token.unsqueeze(0)], dim=1)
            
            n_draft_tokens += len(draft_tokens)
            
            # Target model verifies all draft tokens in one forward pass (greedy)
            draft_tensor = torch.tensor([draft_tokens], device=device)
            verify_ids = torch.cat([input_ids, draft_tensor], dim=1)
            target_logits = target_model(verify_ids).logits
            n_target_calls += 1
            
            # Verify each draft token greedily
            accepted = 0
            for i, draft_tok in enumerate(draft_tokens):
                logit_pos = current_len + i - 1
                if logit_pos >= target_logits.shape[1]:
                    break
                target_tok = target_logits[:, logit_pos, :].argmax(dim=-1).item()
                if target_tok == draft_tok:
                    accepted += 1
                    input_ids = torch.cat([input_ids, torch.tensor([[draft_tok]], device=device)], dim=1)
                else:
                    input_ids = torch.cat([input_ids, torch.tensor([[target_tok]], device=device)], dim=1)
                    break
            else:
                # All draft tokens accepted; append bonus token from target
                if input_ids.shape[1] < init_len + max_new_tokens:
                    bonus_tok = target_logits[:, -1, :].argmax(dim=-1)
                    input_ids = torch.cat([input_ids, bonus_tok.unsqueeze(0)], dim=1)
            
            acceptance_lengths.append(accepted)
            
            if input_ids.shape[1] >= init_len + max_new_tokens:
                break
    
    generated_text = tokenizer.decode(input_ids[0], skip_special_tokens=True)
    return generated_text, acceptance_lengths, n_draft_tokens, n_target_calls

def generate_random_prompt(tokenizer, length=20, seed=None):
    """Generate a prompt with random tokens."""
    if seed is not None:
        torch.manual_seed(seed)
    vocab_size = tokenizer.vocab_size
    random_ids = torch.randint(0, vocab_size, (1, length))
    return tokenizer.decode(random_ids[0], skip_special_tokens=True)

def main():
    print("Loading models...")
    model_name = "gpt2"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    target_model = AutoModelForCausalLM.from_pretrained(model_name)
    draft_model = AutoModelForCausalLM.from_pretrained(model_name)
    
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f"Using device: {device}")
    
    # Claim 4: Compare random tokens vs real prompts
    print("\n" + "="*60)
    print("CLAIM 4: Random tokens vs real prompts")
    print("="*60)
    
    real_prompts = [
        "The quick brown fox jumps over the lazy dog.",
        "In machine learning, speculative decoding accelerates inference by",
        "The capital of France is Paris, and it is known for",
        "To solve this equation, we first need to",
        "The weather today is sunny with a chance of"
    ]
    
    results_c4 = {'real': [], 'random': []}
    
    for prompt in real_prompts:
        _, al_real, _, _ = speculative_decode(target_model, draft_model, tokenizer, prompt, max_draft_len=4, max_new_tokens=20, device=device, draft_temp=1.2)
        avg_al = np.mean(al_real) if al_real else 0
        results_c4['real'].append(avg_al)
        print(f"Real prompt: avg AL = {avg_al:.2f} (details: {al_real})")
    
    for seed in range(5):
        random_prompt = generate_random_prompt(tokenizer, length=15, seed=seed)
        _, al_rand, _, _ = speculative_decode(target_model, draft_model, tokenizer, random_prompt, max_draft_len=4, max_new_tokens=20, device=device, draft_temp=1.2)
        avg_al = np.mean(al_rand) if al_rand else 0
        results_c4['random'].append(avg_al)
        print(f"Random prompt: avg AL = {avg_al:.2f} (details: {al_rand})")
    
    real_avg = np.mean(results_c4['real'])
    rand_avg = np.mean(results_c4['random'])
    print(f"\nReal prompts avg AL: {real_avg:.2f}")
    print(f"Random prompts avg AL: {rand_avg:.2f}")
    print(f"Random > Real: {rand_avg > real_avg}")
    
    # Claim 5: Optimal draft length
    print("\n" + "="*60)
    print("CLAIM 5: Optimal draft length")
    print("="*60)
    
    prompt = "The weather today is"
    
    results_c5 = {}
    for draft_len in [1, 2, 3, 4]:
        _, al, n_draft, n_target = speculative_decode(target_model, draft_model, tokenizer, prompt, max_draft_len=draft_len, max_new_tokens=20, device=device, draft_temp=1.2)
        avg_al = np.mean(al) if al else 0
        total_generated = sum(al) + len(al)
        speedup = total_generated / n_target if n_target > 0 else 0
        results_c5[draft_len] = {'avg_al': avg_al, 'speedup': speedup, 'n_target': n_target, 'al_list': al}
        print(f"Draft length {draft_len}: avg AL = {avg_al:.2f}, approx speedup = {speedup:.2f}x, n_target={n_target}, ALs={al}")
    
    optimal_dl = max(results_c5, key=lambda k: results_c5[k]['speedup'])
    print(f"\nOptimal draft length (by speedup): {optimal_dl}")
    
    # Save results
    import json
    with open('outputs/claims_c4_c5.json', 'w') as f:
        json.dump({
            'claim4': {
                'real_avg_al': float(real_avg),
                'random_avg_al': float(rand_avg),
                'random_gt_real': bool(rand_avg > real_avg),
                'real_details': [float(x) for x in results_c4['real']],
                'random_details': [float(x) for x in results_c4['random']]
            },
            'claim5': {
                str(k): {kk: (float(vv) if isinstance(vv, (int, float, np.floating)) else vv) for kk, vv in v.items()} 
                for k, v in results_c5.items()
            },
            'optimal_draft_length': int(optimal_dl)
        }, f, indent=2)
    print("\nResults saved to outputs/claims_c4_c5.json")

if __name__ == "__main__":
    main()