F ast D raft: How to Train Your Draft (2025)

Abstract: Speculative Decoding has gained popularity as an effective technique for accelerating the auto-regressive inference process of Large Language Models. However, Speculative Decoding entirely relies on the availability of efficient draft models, which are often lacking for many existing language models due to a stringent constraint of vocabulary compatibility. In this work we introduce FastDraft, a novel and efficient approach for pre-training and aligning a draft model to any large language model by incorporating efficient pre-training, followed by fine-tuning over synthetic datasets generated by the target model. We demonstrate FastDraft by training two highly parameter efficient drafts for the popular Phi-3-mini and Llama-3.1-8B models. Using FastDraft, we were able to produce a draft model with approximately 10 billion tokens on a single server with 8 Intel Gaudi 2 accelerators in under 24 hours. Our results show that the draft model achieves impressive results in key metrics of acceptance rate, block efficiency and up to 3x memory bound speed up when evaluated on code completion and up to 2x in summarization, text completion and instruction tasks. We validate our theoretical findings through benchmarking on the latest Intel Core Ultra, achieving a wall-clock time speedup of up to 2x, indicating a significant reduction in runtime. Due to its high quality, FastDraft unlocks large language models inference on AI-PC and other edge-devices.

## introduction the advent of transformer architectures has fundamentally reshaped the field of natural language processing (nlp). in recent years, transformerbased models have achieved remarkable success across a broad spectrum of natural language understanding and generation tasks @xcite * equal contribution @xcite . their exceptional performance, particularly in large language models (llms), has made them highly desirable for deployment in numerous applications, ranging from conversational systems to content generation and beyond. despite their outstanding performance, llms suffer from slow inference speed due to substantial memory bandwidth requirements and the sequential nature of auto-regressive generation (arg). the introduction of speculative decoding (sd) @xcite offers a promising solution for accelerating arg without sacrificing generation quality, making it a compelling approach for improving llm inference efficiency. sd utilizes a draft language model (lm) to generate a sequence of tokens autoregressively, while the target model validates the batched tokens in parallel. in certain applications, sd can achieve a 2-3x speedup in llm inference without compromising the generation quality of the target model. achieving significant speedup with sd requires a high-quality draft model that is both efficient and well-aligned with the target. to date, such draft models remain scarce, even for widely used open-source llms @xcite due to vocabulary incompatibility. to address this limitation, we propose fastdraft, a method for producing hardware (hw) efficient draft models that are orders of magnitude smaller than their corresponding target models. while extensive research has focused on training and data generation for high-quality llms @xcite , these frameworks are not necessarily applicable to training draft models for sd. llms are typically trained to generate helpful, high-quality responses, whereas draft models should be trained to generate sequences that are likely to be accepted by the target model. we explore previously unexamined aspects of draft model training for sd. our contributions include: (1) introducing a method for producing quality and highly efficient draft models with low resource requirements for any given target llm and demonstrate it by training and benchmarking a draft for phi-3-mini, (2) conducting extensive ablation studies on pre-training data size, pre-training for both code and natural language, target-draft alignment via knowledge distillation (kd) and hwaware draft architectures, and (3) demonstrating the scalability of fastdraft by training a draft model for llama-3.1-8b-instruct, achieving performance improvements comparable to those attained for phi-3mini. we demonstrate significant improvements in key metrics using fastdraft, leading to theoretical speedups of up to 3x as measured by the memory bound speedup (mbsu) metric. we further validate our theoretical findings by evaluating the wallclock time speedup achieved with our approach on the latest intel ® core ™ ultra. our results demonstrate an average 1.5x speedup for natural language tasks and an average 2x speedup for code completion tasks. according to our findings, the small size of the draft model and the limited amount of data required to produce a high-quality draft enabled us to successfully train and align a draft model to phi-3-mini, end-to-end, in under 24 hours using a single node with 8 intel ® gaudi ® 2 accelerators. ## related work the widespread adoption of llms in cloud and edge devices has driven a significant body of research focused on developing alternative strategies for arg to address the low inference efficiency of llms @xcite @xcite . however, many of these approaches compromise generation quality or require additional training data and architectural modifications. the introduction of speculative decoding as a lossless approach for accelerating llm inference @xcite has inspired a new wave of follow-up research. some studies propose using plug-in prediction heads as a drafting mechanism @xcite , while others focus on improving the serving latency of stand-alone draft models @xcite @xcite . in contrast, our work focuses on directly enhancing stand-alone draft model's capabilities through pre-training and fine-tuning. we provide more details on the comparison with these methods in appendix a. other works apply kd to draft models to improve alignment with the target model. these studies explore various divergence functions for the kd algorithm, rather than relying solely on the commonly used kullback-leibler divergence (kld). for instance, @xcite proposed total variation distance (tvd) based on corollary 3.6 in @xcite , which posits that minimizing tvd maximizes the token-level acceptance rate. @xcite further developed this approach with tvd++, fine-tuning their model, pre-trained on 600 billion tokens. another study, @xcite , focuses on enhancing the hw efficiency of draft models by extensively analyzing the tradeoff between time latency and acceptance rates, rather than relying solely on the latter. while their work emphasizes the approach for obtaining a draft model by pruning of the shelf models, our method introduces an efficient approach for pre-training and aligning a draft model. @xcite proposes training a compact draft model based on the target model, using a relatively limited dataset. the primary aim of the paper is to develop a draft architecture that effectively utilizes the target model's hidden representations and weights. in contrast, our paper centers on training and aligning any draft architecture that shares only the vocabulary with the target model. ## speculative decoding sd is a lossless decoding paradigm introduced by @xcite for accelerating arg with llms. it is inspired by speculative execution @xcite and aims to mitigate the inherent latency bottleneck caused by the sequential nature of arg @xcite . sd employs a draft lm to generate a block of γ candidate tokens. the llm, referred to as the target model, then processes these candidate tokens in parallel. the algorithm examines each token's probability distribution, calculated by both the target and draft models, to determine whether the token should be accepted or rejected. as a result, any lm can function as a draft model, provided it shares the same vocabulary as the target model. however, since an lm's vocabulary is fixed during pre-training, leveraging existing models as draft models is only feasible if they were pre-trained on the same vocabulary. a key metric for benchmarking arg inference performance is time per output token (tpot), which represents the expected latency for generating a single output token, excluding pre-filling latency. 1 in sd, tpot is a function of the draft model latency l d , the speculation block size γ, the corresponding expected block efficiency foot_1 τ γ , and the target model latency l γ t . the equation is expressed as: in comparison, for traditional arg, tpot is simply tpot @xmath0. the expected speedup is calculated as follows: from equation 2 , two key requirements emerge for sd to yield meaningful speedups. first, the latency for generating the speculated tokens block must be negligible compared to the target model's latency. second, the increase in target model latency for block size γ should be insignificant compared to the latency for block size 1. the latter condition generally holds for most popular llms with sufficiently small γ, while the former depends on the availability of draft models that meet the vocabulary constraint. the draft architecture imposes only one strict requirement, it must produce a probability distribution over the target's vocabulary. beyond this, the design is flexible. nevertheless, certain factors should be considered when selecting the draft's architecture, with latency being the primary concern, as discussed in section 3. then, the chosen draft is trained over a pre-training dataset of natural language for language modeling with the objective of predicting the next token. a common application of llms is code completion. while the pre-training dataset may include some code, unless the dataset is specifically focused on code, the draft's performance on tasks requiring an understanding of code is often suboptimal. to address this, @xcite proposed continued pre-training (cp), where training begins with a pre-trained model and is extended using a combination of code and natural language data. fastdraft adopts this approach for producing drafts for code completion tasks. ## experiments in this section we report descriptions, setup and results of our experiments. in the following experiments, we report acceptance rate (section 4.3.1) results using two sampling methods, greedy and multinomial sampling decoding with temperature @xmath1.6 with block sizes γ = 3 and γ = 5 unless stated otherwise. the hyper-parameters used in our draft pre-training experiments are detailed appendix b.3 unless stated otherwise. ## results & reproducibility combining the findings from the ablation studies presented in section 5, we outline our comprehensive pipeline for draft pre-training and fine-tuning to optimize performance on key metrics for speculative decoding and demonstrate it by producing a draft for phi-3-mini. to further illustrate reproducibility, we employ the same pipeline to produce a draft for llama-3.1-8b-instruct (dubey et al., 2024). we present a 50m draft for the phi-3-mini and a 150m draft for llama-3.1-8b. drafts' architecture details are presented in appendix b.1. considering both performance and resource efficiency, our findings from sections 5.2, 5.3 suggest pretraining over 5bt of fineweb text data and continued pre-training on a mixture of 5bt code data from the stack v2 and 2.5bt fineweb text data as a best practice. for fine-tuning, we conclude from sections 5.4, 5.5 that sequence-level kd yields significant improvements while token-level kd benefits are not definitive, therefore, for fastdraft we only utilize sequence-level kd. we construct an alignment dataset for both phi-3 and llama-3.1 drafts by combining a number of synthetic datasets we generated with the appropriate target model as described in section 5.1. ## conclusion & future work in this paper, we introduced fastdraft, a novel approach for training and evaluating draft models for speculative decoding. our results demonstrate that fastdraft facilitates the rapid training of high-quality, efficient draft models that are well-aligned with target models. we conducted a comprehensive ablation study to examine various aspects of draft training, including pre-training data composition and draft-target alignment. using our method, we successfully trained a highly efficient 50m-parameter draft model for phi-3-mini, achieving an acceptance rate of up to 67% and up to a 3x memory-bound speedup. we further demonstrated the scalability of fastdraft by training a draft for llama-3.1-8b-instruct, underscoring its effectiveness for larger models. we validated our theoretical findings by benchmarking our drafts on the latest ai-pc hw achieving a wall-clock time speedup of up to 2x. we hope our findings will inspire additional research on efficient draft training, particularly focusing on the development of resource-efficient draft architectures and hardwareaware designs.