Package: dualdec
Dual DecodingDual Decoding is a new decoding algorithm that has much faster generation speed compared with Speculative Decoding and Lookahead Decoding, without any training cost.Benchmark ResultThe generation speed tested on HumanEval and MBPP are presented here. The shadowed parts are the extra improvement caused by warm start.Here are more experimental results on GSM8K, CNNDM and WMT16. Dual with "†" represents Dual Decoding with warm start, and Dual without "†" represents Dual Decoding without warm start.TaskAlgorithmYi 34b/6b tok/sYi 34b/6b speed-upLlama-2 70b/7b tok/sLlama-2 70b/7b speed-upGSM8kGreedy15.331.008.961.00GSM8kSpeculative16.991.1116.861.88GSM8kLookahead25.141.6413.771.54GSM8kDual26.411.7222.182.48GSM8kDual†28.231.8424.032.68CNN/DMGreedy14.621.008.121.00CNN/DMSpeculative17.821.2212.771.57CNN/DMLookahead18.771.289.471.17CNN/DMDual21.241.4514.621.80CNN/DMDual†22.651.5514.671.81WMT16Greedy14.781.009.521.00WMT16Speculative17.481.1814.721.55WMT16Lookahead17.981.2214.651.54WMT16Dual19.751.3419.112.01WMT16Dual†19.941.3519.272.02InstallInstall from pippipinstalldualdecInstall from source codeFirst, clone this repo asdualdecThen, run the following commandcddualdec
pipinstall-e.Quick StartAn ExampleHere is an example of using Dual Decoding. Note: current version ofdualdeccan only support Llama model inference. Please importLlamaForCausalLMfromdualdec.modelsinstead oftransformersThe following code shows an example of comparing the output of Dual Decoding with Greedy decoding. Ifdualdecis installed correctly, the following code should outputTrue.importtorchfromdualdecimportdualdecfromtransformersimportAutoTokenizerfromdualdec.modelsimportLlamaForCausalLMwindow_size=20guess_set_size=20lookahead_level=7gamma=12small_model=LlamaForCausalLM.from_pretrained("yourpath",torch_dtype=torch.float16,device_map='cuda')target_model=LlamaForCausalLM.from_pretrained("yourpath",torch_dtype=torch.float16,device_map='cuda')tokenizer=AutoTokenizer.from_pretrained("yourpath")prompt="Please summarize the following paragraph. Officers searched properties in the Waterfront Park and Colonsay View areas of the city on Wednesday. Detectives said three firearms, ammunition and a five-figure sum of money were recovered. A 26-year-old man who was arrested and charged appeared at Edinburgh Sheriff Court on Thursday. Summary: "input_ids=tokenizer(prompt,return_tensors='pt').to('cuda')['input_ids']dualdec_output=dualdec(input_ids,small_model,target_model,max_len=64,gamma=gamma,window_size=window_size,guess_set_size=guess_set_size,lookahead_level=lookahead_level)std_output=target_model.generate(input_ids,do_sample=False,min_length=64,max_length=64)print(dualdec_output[:,:64].equal(std_output[:,:64]))Using Dual DecodingFirst, prepare two huggingface transformers models (the target model and draft model). Load the model withLlamaForCausalLMfromdualdec.models.Then import the generation functiondualdecfromdualdec.Call the function by the following parameters:@torch.no_grad()defdualdec(prefix:torch.Tensor,approx_model:torch.nn.Module,target_model:torch.nn.Module,ngram_cache:CacheEngine=None,max_len:int=512,gamma:int=4,window_size=20,guess_set_size=20,lookahead_level=7,eos_token_id=2,topk=3)->torch.Tensor:"""Performs dual decoding with an approximate model and a target model to generate a sequence of tokens.Args:prefix (torch.Tensor): The initial sequence of tokens to start the generation from.approx_model (torch.nn.Module): The approximate model used for initial token generation. The model should support huggingface transformers model methods.target_model (torch.nn.Module): The target model used for refining the generated tokens. The model should support huggingface transformers model methods.ngram_cache (CacheEngine, optional): A cache engine for storing and retrieving n-gram predictions. Defaults to None, in which case a new cache engine is created.max_len (int, optional): The maximum length of the generated sequence. Defaults to 512.gamma (int, optional): The lookahead parameter for generation. Defaults to 4.window_size (int, optional): The window size used for n-gram generation. Defaults to 20. Currently, must be equal to guess_set_size.guess_set_size (int, optional): The size of the guess set for n-gram retrieving. Defaults to 20. Currently, must be equal to window_size.lookahead_level (int, optional): The level of lookahead decoding. Defaults to 7.eos_token_id (int, optional): The token id representing the end-of-sequence token. Defaults to 2. Should be given by tokenizer.eos_token_id.topk (int, optional): The k of verifying k candidate suffixes.Returns:torch.Tensor: The generated sequence of tokens, including the initial prefix and any additional tokens generated by the function."""CitationPlease cite ourpaperif you find our work valuable.@article{dualdecoding,
  title={Dual Decoding: Speculative Decoding with Large Model Enhanced Drafting},
  author={Weilin Zhao, Yuxiang Huang, Xu Han, Chaojun Xiao, Zhiyuan Liu, Maosong Sun},
  journal={arXiv preprint arXiv:TODO},
  year={2024}
}