# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file. import sys import time import warnings from collections.abc import Iterator from pathlib import Path from pprint import pprint from typing import Any, Literal import lightning as L import torch import torch._dynamo.config import torch._inductor.config import torch.nn.functional as F from lightning.fabric.plugins import BitsandbytesPrecision from litgpt.config import Config from litgpt.constants import _BITANDBYTES_AVAILABLE_NOT_EQUAL_0_42_0 from litgpt.generate.base import multinomial_num_samples_1, next_token, sample_top_p from litgpt.model import GPT from litgpt.prompts import PromptStyle, has_prompt_style, load_prompt_style from litgpt.tokenizer import Tokenizer from litgpt.utils import ( check_file_size_on_cpu_and_warn, check_valid_checkpoint_dir, extend_checkpoint_dir, get_default_supported_precision, load_checkpoint, ) def sample( logits: torch.Tensor, temperature: float = 1.0, top_k: int | None = None, top_p: float = 1.0, apply_softmax: bool = True, ) -> torch.Tensor: if top_p < 0.0 or top_p > 1.0: raise ValueError(f"top_p must be in [0, 1], got {top_p}") logits = logits[0, -1] # optionally crop the logits to only the top k options if top_k is not None: v, i = torch.topk(logits, min(top_k, logits.size(-1))) # do not use `torch.where` as in nanogpt because it will repeat top-k collisions fill_value = float("-inf") if apply_softmax else float(0) logits = torch.full_like(logits, fill_value).scatter_(-1, i, v) # optionally scale the logits and sample from a probability distribution if temperature > 0.0 and top_p > 0.0: if temperature > 0.0: logits = logits / temperature # optionally crop the logits to smallest set of logits with a cumulative probability above top_p if top_p < 1.0: logits = sample_top_p(logits, top_p) probs = F.softmax(logits, dim=-1) if apply_softmax else logits return multinomial_num_samples_1(probs), probs return torch.argmax(logits, dim=-1, keepdim=True), F.softmax(logits, dim=-1) def speculative_decoding( draft_model: GPT, target_model: GPT, token: torch.Tensor, input_pos: torch.Tensor, input_pos_maxp1: int, speculative_k: int, **sample_kwargs: dict[str, Any], ) -> torch.Tensor: """Performs speculative decoding using a draft and a target model. This implements the speculative decoding algorithm from "Fast Inference from Transformers via Speculative Decoding" (https://arxiv.org/pdf/2211.17192). The core idea is to: 1. Use a faster draft model to predict multiple tokens ahead 2. Verify those predictions with the slower but more accurate target model 3. Accept tokens where the target model agrees with high probability 4. Reject and resample tokens where there is a disagreement This allows leveraging a smaller/faster model to speed up generation while maintaining the quality of the larger target model. Args: draft_model: Smaller/faster model used for initial token predictions target_model: Larger/slower model used for verification token: Current input token tensor of shape [1] input_pos: Position index of the token tensor for KV-cache input_pos_maxp1: Maximum position + 1 for managing KV-cache buffer speculative_k: Number of tokens to speculatively generate at once sample_kwargs: Additional sampling parameters (temperature, top_k, top_p) Returns: torch.Tensor: Generated tokens that were either accepted from draft model or resampled from target model """ if speculative_k < 1: raise ValueError(f"speculative_k must be >= 1, got {speculative_k}") # Step 1: Generate candidate tokens using draft model # The draft model autoregressively generates k tokens, keeping track of probabilities draft_input_pos = input_pos.clone() draft_input_pos_maxp1 = input_pos_maxp1 draft_tokens, draft_probs = [], [] draft_token = token for idx in range(speculative_k): logits = draft_model( idx=draft_token.unsqueeze(0), input_pos=draft_input_pos, input_pos_maxp1=draft_input_pos_maxp1 ) draft_token, draft_prob = sample(logits, **sample_kwargs) draft_input_pos.add_(1) draft_input_pos_maxp1 += 1 draft_tokens.append(draft_token) draft_probs.append(draft_prob) draft_tokens = torch.cat(draft_tokens) # Step 2: Get target model predictions for comparison # Feed both original token and draft tokens to get target probabilities candidate_tokens = torch.cat((token, draft_tokens)) candidate_input_pos = input_pos + torch.arange(0, speculative_k + 1, device=input_pos.device) candidate_input_pos_maxp1 = input_pos_maxp1 + speculative_k target_logits = target_model( idx=candidate_tokens.unsqueeze(0), input_pos=candidate_input_pos, input_pos_maxp1=candidate_input_pos_maxp1 ) # Step 3: Convert target logits to probabilities using same sampling params target_probs = [] for target_logit in target_logits.split(1, dim=1): _, target_prob = sample(target_logit, **sample_kwargs) target_probs.append(target_prob) # Step 4: Accept/reject draft tokens based on probability comparison # Using rejection sampling: keep token if target_prob >= draft_prob. # Otherwise reject with probability 1 - target_prob / draft_prob. # If rejected, sample from an adjusted distribution: norm(max(0, target_prob_distribution - draft_prob_distribution) instead. accepted_tokens = [] for idx in range(len(draft_tokens)): draft_token = draft_tokens[idx].unsqueeze(0) draft_prob = draft_probs[idx][draft_token] target_prob = target_probs[idx][draft_token] # Accept the draft token if the target model is "confident" in it if target_prob >= draft_prob: accepted_tokens.append(draft_token) continue # If not accepted, probabilistically reject it discard_prob = 1 - target_prob / draft_prob should_discard_token = torch.rand(1, device=discard_prob.device) <= discard_prob if not should_discard_token: accepted_tokens.append(draft_token) continue # On rejection: sample new token from adjusted distribution # p'(x) = normalize(max(0, p_target(x) - p_draft(x))) adjusted_distribution = target_probs[idx] - draft_probs[idx] adjusted_distribution = torch.clamp(adjusted_distribution, 0.0) adjusted_distribution = adjusted_distribution / adjusted_distribution.sum() new_token, _ = sample(adjusted_distribution[None, None, ...], apply_softmax=False, **sample_kwargs) return torch.cat((*accepted_tokens, new_token)) # If all draft tokens were accepted: # 1. Update draft model's key-value cache # 2. Sample one more token from target model draft_model(idx=draft_token.unsqueeze(0), input_pos=draft_input_pos, input_pos_maxp1=draft_input_pos_maxp1) new_token, _ = sample(target_logits, **sample_kwargs) return torch.cat((*accepted_tokens, new_token)) @torch.inference_mode() def generate( draft_model: GPT, target_model: GPT, prompt: torch.Tensor, max_returned_tokens: int, *, temperature: float = 1.0, top_k: int | None = None, top_p: float = 1.0, stop_tokens: tuple[list[int], ...] = (), include_prompt: bool = True, speculative_k: int, ) -> Iterator[torch.Tensor]: """Generates tokens using speculative decoding with a draft and a target model. This function implements token generation using speculative decoding, where a faster draft model makes initial token predictions that are verified by a slower but more accurate target model. Args: draft_model: Smaller/faster model used for initial token predictions target_model: Larger/more accurate model used to verify draft predictions prompt: Input tensor of token ids to generate from, shape [sequence_length] max_returned_tokens: Maximum total tokens (prompt + generated) to return temperature: Sampling temperature (higher = more random, lower = more deterministic) top_k: If set, only sample from the top k most likely next tokens top_p: If <1.0, only sample from tokens whose cumulative probability exceeds top_p stop_tokens: List of token sequences that will stop generation if produced include_prompt: Whether to include prompt tokens in the returned sequence speculative_k: Number of tokens to speculatively generate at each step Returns: - tokens: Tensor of generated token ids - acceptance_rate: Ratio of accepted draft model predictions This implements an optimized decoding process: 1. Both models process the initial prompt 2. Draft model speculatively generates k tokens ahead 3. Target model verifies the draft predictions 4. Accepted tokens are kept, rejected ones trigger resampling 5. Process repeats until max tokens or stop sequence reached """ prompt_size = prompt.size(0) device = prompt.device assert max_returned_tokens > prompt_size, ( f"Not enough space for {prompt_size} prompt tokens in a context length of {max_returned_tokens}." ) if draft_model.max_seq_length < max_returned_tokens - 1: raise NotImplementedError( f"max_seq_length {draft_model.max_seq_length} needs to be >= {max_returned_tokens - 1}" ) if target_model.max_seq_length < max_returned_tokens - 1: raise NotImplementedError( f"max_seq_length {target_model.max_seq_length} needs to be >= {max_returned_tokens - 1}" ) # Step 1: Prefill draft and target models with the prompt. input_pos = torch.arange(0, prompt_size, device=device, dtype=torch.int64) # We want to skip if ThunderModules are involved, either directly or wrapped in LightningModule etc. input_pos_maxp1 = ( prompt_size if all(m.__class__.__name__ != "ThunderModule" for m in target_model.modules()) else None ) next_token( draft_model, input_pos, prompt.view(1, -1), input_pos_maxp1=input_pos_maxp1, temperature=temperature, top_k=top_k, top_p=top_p, ) token = next_token( target_model, input_pos, prompt.view(1, -1), input_pos_maxp1=input_pos_maxp1, temperature=temperature, top_k=top_k, top_p=top_p, ) # Update position trackers after prompt input_pos = torch.tensor([prompt_size], device=device, dtype=torch.int64) input_pos_maxp1 += 1 # Step 2: Main generation loop. tokens = [] total_generated, total_accepted = 0, 0 # Track acceptance statistics while input_pos < max_returned_tokens - 1: # Calculate speculative tokens to generate _speculative_k = min(speculative_k, (max_returned_tokens - input_pos - 1).item()) # Get new tokens via speculative decoding new_tokens = speculative_decoding( draft_model=draft_model, target_model=target_model, token=token, input_pos=input_pos, input_pos_maxp1=input_pos_maxp1, speculative_k=_speculative_k, temperature=temperature, top_k=top_k, top_p=top_p, ) # Update statistics accepted_tokens_len = len(new_tokens) total_generated += _speculative_k total_accepted += accepted_tokens_len - 1 # accepted +1 sampled from a target model # Process tokens and check for stop condition should_break = False for new_token in new_tokens: if new_token in stop_tokens: should_break = True break tokens.append(new_token) if should_break: break # Update positions for next iteration input_pos.add_(accepted_tokens_len) input_pos_maxp1 += accepted_tokens_len token = new_tokens[-1].unsqueeze(0) # Finalize generated sequence tokens = torch.stack(tokens) if include_prompt: tokens = torch.cat([prompt, tokens]) acceptance_rate = total_accepted / total_generated if total_generated > 0 else 0.0 return tokens, acceptance_rate def setup_model(config: Config, max_returned_tokens: int, fabric: L.Fabric) -> GPT: """Helper function to setup a model with common configuration.""" with fabric.init_module(empty_init=True): model = GPT(config) with fabric.init_tensor(): # set the max_seq_length to limit the memory usage to what we need model.max_seq_length = max_returned_tokens # enable the kv cache model.set_kv_cache(batch_size=1) model.eval() return fabric.setup_module(model) def load_model(checkpoint_dir: Path, fabric: L.Fabric) -> tuple[Config, Path]: """Helper function to validate and load model configuration.""" check_valid_checkpoint_dir(checkpoint_dir) config = Config.from_file(checkpoint_dir / "model_config.yaml") checkpoint_path = checkpoint_dir / "lit_model.pth" check_file_size_on_cpu_and_warn(checkpoint_path, fabric.device) return config, checkpoint_path @torch.inference_mode() def main( draft_model_checkpoint_dir: Path, target_model_checkpoint_dir: Path, prompt: str = "What food do llamas eat?", *, sys_prompt: str | None = None, num_samples: int = 1, max_new_tokens: int = 50, speculative_k: int = 3, top_k: int | None = 50, top_p: float = 1.0, temperature: float = 0.8, quantize: Literal["bnb.nf4", "bnb.nf4-dq", "bnb.fp4", "bnb.fp4-dq", "bnb.int8"] | None = None, precision: str | None = None, compile: bool = False, ) -> None: """Default generation option. Generates text samples based on pre-trained models and a tokenizer. Args: draft_model: Smaller/faster model used for initial token predictions target_model: Larger/more accurate model used to verify draft predictions prompt: The prompt string to use for generating the samples. sys_prompt: The system prompt to use for generating the samples. num_samples: The number of text samples to generate. max_new_tokens: The number of generation steps to take. speculative_k: Number of tokens to speculatively generate at each step top_k: The number of top most probable tokens to consider in the sampling process. top_p: If specified, it represents the cumulative probability threshold to consider in the sampling process. In top-p sampling, the next token is sampled from the highest probability tokens whose cumulative probability exceeds the threshold `top_p`. When specified, it must be `0 <= top_p <= 1`. Here, `top_p=0` is equivalent to sampling the most probable token, while `top_p=1` samples from the whole distribution. It can be used in conjunction with `top_k` and `temperature` with the following order of application: 1. `top_k` sampling 2. `temperature` scaling 3. `top_p` sampling For more details, see https://arxiv.org/abs/1904.09751 or https://huyenchip.com/2024/01/16/sampling.html#top_p temperature: A value controlling the randomness of the sampling process. Higher values result in more random samples. quantize: Whether to quantize the model and using which method: - bnb.nf4, bnb.nf4-dq, bnb.fp4, bnb.fp4-dq: 4-bit quantization from bitsandbytes - bnb.int8: 8-bit quantization from bitsandbytes for more details, see https://github.com/Lightning-AI/litgpt/blob/main/tutorials/quantize.md precision: Indicates the Fabric precision setting to use. compile: Whether to compile the model. """ draft_model_checkpoint_dir = extend_checkpoint_dir(draft_model_checkpoint_dir) target_model_checkpoint_dir = extend_checkpoint_dir(target_model_checkpoint_dir) pprint(locals()) # Setup Fabric precision = precision or get_default_supported_precision(training=False) plugins = None if quantize is not None and quantize.startswith("bnb."): if "mixed" in precision: raise ValueError("Quantization and mixed precision is not supported.") if _BITANDBYTES_AVAILABLE_NOT_EQUAL_0_42_0: warnings.warn( "LitGPT only supports bitsandbytes v0.42.0. This may result in errors when using quantization." ) dtype = {"16-true": torch.float16, "bf16-true": torch.bfloat16, "32-true": torch.float32}[precision] plugins = BitsandbytesPrecision(quantize[4:], dtype) precision = None fabric = L.Fabric(devices=1, precision=precision, plugins=plugins) # Load model configs and checkpoints draft_config, draft_checkpoint_path = load_model(draft_model_checkpoint_dir, fabric) target_config, target_checkpoint_path = load_model(target_model_checkpoint_dir, fabric) # Setup tokenizer and validate draft_tokenizer = Tokenizer(draft_model_checkpoint_dir) target_tokenizer = Tokenizer(target_model_checkpoint_dir) if draft_tokenizer.vocab_size != target_tokenizer.vocab_size: raise ValueError("Draft and target models have different vocab sizes.") tokenizer = target_tokenizer # Setup prompt prompt_style = ( load_prompt_style(target_model_checkpoint_dir) if has_prompt_style(target_model_checkpoint_dir) else PromptStyle.from_config(target_config) ) prompt = prompt_style.apply(prompt, sys_prompt=sys_prompt) encoded = tokenizer.encode(prompt, device=fabric.device) prompt_length = encoded.size(0) max_returned_tokens = prompt_length + max_new_tokens # Initialize models fabric.print(f"Loading draft model {str(draft_checkpoint_path)!r} with {draft_config.__dict__}", file=sys.stderr) fabric.print(f"Loading target model {str(target_checkpoint_path)!r} with {target_config.__dict__}", file=sys.stderr) t0 = time.perf_counter() draft_model = setup_model(draft_config, max_returned_tokens, fabric) target_model = setup_model(target_config, max_returned_tokens, fabric) fabric.print(f"Time to instantiate models: {time.perf_counter() - t0:.02f} seconds.", file=sys.stderr) # Setup compilation if needed if compile: torch._dynamo.config.automatic_dynamic_shapes = True torch._inductor.config.triton.unique_kernel_names = True torch._inductor.config.coordinate_descent_tuning = True global next_token next_token = torch.compile(next_token, mode="reduce-overhead") # Load model weights t0 = time.perf_counter() load_checkpoint(fabric, draft_model, draft_checkpoint_path) load_checkpoint(fabric, target_model, target_checkpoint_path) fabric.print(f"Time to load the models weights: {time.perf_counter() - t0:.02f} seconds.", file=sys.stderr) # Generate samples L.seed_everything(1234) for i in range(num_samples): t0 = time.perf_counter() y, acceptance_rate = generate( draft_model, target_model, encoded, max_returned_tokens, temperature=temperature, top_k=top_k, top_p=top_p, stop_tokens=([tokenizer.eos_id] if tokenizer.eos_id is not None else []), speculative_k=speculative_k, ) t = time.perf_counter() - t0 # Reset KV cache for model in (draft_model, target_model): for block in model.transformer.h: block.attn.kv_cache.reset_parameters() # Print results fabric.print(tokenizer.decode(y)) tokens_generated = y.size(0) - prompt_length print(f"Acceptance rate: {acceptance_rate * 100:.2f}%") fabric.print( f"Time for inference {i + 1}: {t:.02f} sec total, {tokens_generated / t:.02f} tokens/sec", file=sys.stderr ) if fabric.device.type == "cuda": fabric.print(f"Memory used: {torch.cuda.max_memory_allocated() / 1e9:.02f} GB", file=sys.stderr)