ããã«ã¡ã¯ãã€ãããŒã·ã§ã³ã»ã³ã¿ãŒã®å è€ã§ããæ®æ®µã¯ã³ã³ãã¥ãŒã¿ããžã§ã³ã®æè¡éçºãAI/æ©æ¢°åŠç¿ïŒML: Machine LearningïŒã·ã¹ãã ã®æ€èšŒã«åãçµãã§ããŸããäžæ¹ã§ãå
Œå 1 ã§å€§èŠæš¡èšèªã¢ãã«ïŒLLM: Large Language ModelïŒã«ã€ããŠèª¿æ»ãè¡ãªã£ãŠãããç¹ã«LLMã®æšè«ãåŠç¿ã®é«éåã«é¢å¿ãæã£ãŠããŸãã ä»åã¯ãå°ããªèšèªã¢ãã«ã«ããå
èªã¿ã掻çšããŠLLMã®æç« çæãé«éåããææ³ïŒAssisted Generation 2 , Speculative Sampling 3 ãªã©ãšåŒã°ããŠããŸãïŒã«ã€ããŠã玹ä»ããŸãã LLMã®æšè«ã¯èšç®ã³ã¹ããé«ããæç« çæã®é
ãã課é¡ãšããŠããæããããŠããŸããç¹ã«æ¥æ¬èªã¯ããŒã¯ã³ãããã®æåæ°ãå°ãªããChatGPTã®ãããªã¹ããªãŒã åºåã§ãããªãçæãé
ãæããããšæããŸãã ããã«å¯ŸããŠãããããäœåã«ã¡ã¢ãªãå©çšããŠãå
ã
ã®LLMãšå
šãåããã®ãããé«éã«åºåã§ããAssisted Generationãææ¡ãããŠããŸããAIéçºã»æ©æ¢°åŠç¿ã®ããã®ãã©ãããã©ãŒã ãéå¶ãã Hugging Face ãå
¬éããŠããæ©æ¢°åŠç¿ã©ã€ãã©ãªã§ããTransformersã§ãå®è£
ãããŠããŸãããã®èšäºã§ã¯ãã®ææ³ã«ã€ããŠãŸãšããããçºå±çãªææ³ãšãšãã«Hugging Faceã®ã¢ãã«ã§å®éšããŸããã LLMã®æç« çæã¡ã«ããºã ã¢ãã«ã®æšè«é床ãšAssisted Generation generate()ã®é«éå sample()ã®é«éå ããŒã¯ãã€ã¶ã®äºææ§ æŽãªãé«éåïŒ Token Tree Verification LLMã®æç« çæã¡ã«ããºã â¡ Causal Language Modelã«ãããAttention Maskã«ã€ã㊠Token Tree Verificationãçšããgenerate() å
è«æã®å®è£
ãããã« LLMã®æç« çæã¡ã«ããºã æç« çæã«äœ¿ãããLLMã®å€ãã¯Causal Language Modelã§äœãããŠããŸããCausal Language Modelãšã¯åèªåïŒå³å¯ã«ã¯ããŒã¯ã³åïŒãäžããããæã«ã®ã¡ã«ç¶ãåèªïŒããŒã¯ã³ïŒãäºæž¬ããã¢ãã«ã§ãäžããããããŒã¯ã³å ã«å¯Ÿããå ã®æ¬¡ã«æ¥ãã¹ãããŒã¯ã³ ã®ååž ãåºåããŸãã 以äžã®äŸã§ã¯ãå
¥åæ"The man worked as a"ã®åããŒã¯ã³ã«å¯ŸããŠã facebook/opt-125m ã¢ãã«ãçšããŠæãæ¬¡ã«æ¥ãããªããŒã¯ã³ãäºæž¬ããçµæãå³ç€ºããŠããŸããäžåºŠã®æšè«ã§å
šããŒã¯ã³ã䞊åã«äºæž¬ã§ããŠããããšããCausal Language Modelã§ã¯åŸãã®å
¥åããŒã¯ã³ãåã®äºæž¬ããŒã¯ã³ã«åœ±é¿ããªãããšãä»å玹ä»ããæè¡ã®ãã€ã³ãã§ãã ãããçšãããš1åã®æšè«ã§æ°ããããŒã¯ã³ã1ã€ïŒããã§ã¯securityïŒçæã§ãããããäœåºŠãç¹°ãè¿ãããšã§æç« ãçæããŸãã ãããAutoregressive GenerationãšåŒã³ãTransformersã§ã¯ generate() ãšãã颿°ãæäŸãããŠããŸããç¹ã«ãã®äŸã®ããã«æã確çã®é«ããã®ã次ã®ããŒã¯ã³ã«éžæããæ¹æ³ãGreedy DecodingãšåŒã³ãŸãã ïŒçµå±æçµããŒã¯ã³ã®äºæž¬ "security" ãã䜿ãéãç¡ãããã«èŠããŸãããæ®ãã®ããŒã¯ã³ã®äºæž¬ãåŸè¿°ã®ææ³ã§æŽ»çšã§ããŸããïŒ import torch import transformers from transformers import AutoModelForCausalLM, AutoTokenizer print (torch.__version__) # 2.1.0+cu121 print (transformers.__version__) # 4.34.1 model = AutoModelForCausalLM.from_pretrained( "facebook/opt-125m" ).cuda() tokenizer = AutoTokenizer.from_pretrained( "facebook/opt-125m" , use_fast= False ) prompt = "The man worked as a" input_ids = tokenizer(prompt, return_tensors= "pt" ).input_ids.cuda() print (tokenizer.decode(model.generate(input_ids, max_new_tokens= 10 )[ 0 ], skip_special_tokens= True )) # The man worked as a security guard at a hotel in the city of K ãŸããLLMãåºåãã確çã«æ²¿ã£ãŠåºåããµã³ããªã³ã° sample() ããããšã§ã倿§æ§ã®ããæç« ãçæããããšãã§ããŸãã print (tokenizer.decode(model.generate(input_ids, max_new_tokens= 10 , do_sample= True )[ 0 ], skip_special_tokens= True )) print (tokenizer.decode(model.generate(input_ids, max_new_tokens= 10 , do_sample= True )[ 0 ], skip_special_tokens= True )) # The man worked as a clerk to a bakery. He worked for a bakery # The man worked as a janitor on a cleaning staff in our house. ã¢ãã«ã®æšè«é床ãšAssisted Generation æè¿å
¬éãããŠããLLMã¯10å~1000åèŠæš¡ã®ãã©ã¡ãŒã¿ãæã£ãŠãããæšè«ã«éåžžã«æéããããããšãç¥ãããŠããŸãã A100 GPUã§ facebook/opt ã¢ãã«ãåäœãããã©ã³ãã ãªæç« ( C4ããŒã¿ã»ãã ã®åé éšãå©çš)ã®ç¶ããæšè«ãããå Žåã¯ä»¥äžã®å³ã®ããã«ãªãã13Bã¢ãã«ã¯æè»œéã®125Mã¢ãã«ã®30åã®çææéãããããŸãã13B, 125Mã¯ã¢ãã«ã®ãã©ã¡ãŒã¿æ°ã衚ããŠããããããã13 billionã125 millionã瀺ããŠããŸãã ããã§1åã®æšè«ã§åäœçœ®ã®æ¬¡ããŒã¯ã³ã䞊åã«äºæž¬ã§ãããšããæ§è³ªãå©çšããå°ããªã¢ãã«ã§ããã€ããå
èªã¿ãããŠããLLMã§ãããæ€èšŒããAssisted Generationãšããææ³ã§é«éã«æç« ãçæã§ããŸãã以äžã®ç¯ã§ã¯ãã®ææ³ã説æããŸãã generate()ã®é«éå ãŸãå°ããªã¢ãã«ïŒä»¥éã§ã¯ãã©ããã¢ãã«ãšåŒã³ãŸãïŒã§æ°ããåèªãè€æ°çæããŸãããã®æäœ¿çšããã¢ãã«ã¯LLMããã軜éã§ãã€åãããŒã¯ã³IDã䜿ã£ãŠããå¿
èŠããããŸãã以äžã®äŸã§ã¯1.3Bã¢ãã«ã«å¯ŸããåãããŒã¯ã³IDã䜿ã£ãŠããŠãã軜éãª125Mã¢ãã«ããã©ããã¢ãã«ã«ããŠããŸãã from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig import torch assist = AutoModelForCausalLM.from_pretrained( "facebook/opt-125m" ).cuda() model = AutoModelForCausalLM.from_pretrained( "facebook/opt-1.3B" , torch_dtype=torch.float16).cuda() tokenizer = AutoTokenizer.from_pretrained( "facebook/opt-1.3B" , use_fast= False ) prompt = "The man worked as a" input_ids = tokenizer(prompt, return_tensors= "pt" ).input_ids.cuda() initial_len = input_ids.size( 1 ) for i in range ( 10 ): with torch.no_grad(): y = assist(input_ids) next_id = y.logits[:, - 1 , :].argmax(- 1 , keepdims= True ) input_ids = torch.cat([input_ids, next_id], dim=- 1 ) # å
¥åæç« + [ãã©ããã¢ãã«ãçæããæç« ] print (prompt + f "[{tokenizer.decode(input_ids[0,initial_len:], skip_special_tokens=True)}]" ) # The man worked as a[ security guard at a hotel in the city of K] 次ã«çæããæç« "The man worked as a security guard at a hotel..." ãLLMã«å
¥ããåäœçœ®ã®æ¬¡ããŒã¯ã³ãäºæž¬ããŸãã with torch.no_grad(): y = model(input_ids) next_ids = y.logits.argmax(- 1 ) for i in range ( 10 ): input_words = tokenizer.decode(input_ids[ 0 , :initial_len+i], skip_special_tokens= True ) next_words = tokenizer.decode(next_ids[ 0 , initial_len+i- 1 :initial_len+i], skip_special_tokens= True ) print (f "model[{i}]: {input_words}[{next_words}]" ) if next_ids[ 0 , initial_len+i- 1 ] != input_ids[ 0 , initial_len+i]: assist_input = tokenizer.decode(input_ids[ 0 , :initial_len+i], skip_special_tokens= True ) assist_word = tokenizer.decode(input_ids[ 0 , initial_len+i:initial_len+i+ 1 ], skip_special_tokens= True ) print (f "assistant: {assist_input}[{assist_word}]" ) break """ model[0]: The man worked as a[ security] model[1]: The man worked as a security[ guard] model[2]: The man worked as a security guard[ at] model[3]: The man worked as a security guard at[ the] assistant: The man worked as a security guard at[ a] """ LLMã®äºæž¬ããŒã¯ã³ã "security guard at" ãŸã§ã¯ãã©ããã¢ãã«ã®åºåãšäžèŽããŠããããšãããããŸããããã§ãã®äžèŽããåãšãã®æ¬¡ã®LLMã®äºæž¬ "the" ãåããã "security guard at the" ãAssisted Generationã®åºåãšããŠäžæ°ã«æ¡çšããŠããŸããŸãã Greedy Decodingã®æ§è³ªäžãLLMã§åæããã³ãããã1ã€ãã€ããŒã¯ã³ãçæããŠããåããŠäºæž¬ãããããã®"the"ãŸã§äžèŽããããšã以äžã®ããã«ç¢ºãããããŸãã prompt = "The man worked as a" input_ids = tokenizer(prompt, return_tensors= "pt" ).input_ids.cuda() initial_len = input_ids.size( 1 ) for i in range ( 10 ): with torch.no_grad(): model_inputs = model.prepare_inputs_for_generation(input_ids) y = model(**model_inputs) next_id = y.logits[:, - 1 , :].argmax(- 1 , keepdims= True ) input_ids = torch.cat([input_ids, next_id], dim=- 1 ) print (tokenizer.decode(input_ids[ 0 ], skip_special_tokens= True )) # The man worked as a security guard at the University of California, Berkeley, # "The man worked as a security guard at the" ã model[3] ãšäžèŽããŠãã ã€ãŸã1åã®LLMæšè«ãš10åã®ãã©ããã¢ãã«æšè«ã§4ããŒã¯ã³çæã§ããŸããããããç¹°ãè¿ãããšã§å·šå€§ãªã¢ãã«ã®æšè«åæ°ãç¯çŽããªãããã®ã¢ãã«ã®åºåãå®å
šã«åçŸã§ããŸããããã¯Transformersã§ã assistant_model ãåŒæ°ã«æå®ããããšã§å©çšã§ããŸãã facebook/opt ã®ããŸããŸãªãã©ã¡ãŒã¿æ°ã®ã¢ãã«ã«å¯Ÿã125Mã¢ãã«( facebook/opt-125m )ãçšããŠAssisted Generationãè¡ã£ãéã®å®è¡æéã¯ä»¥äžã®ããã«ãªããŸãã ã¢ãã«èŠæš¡ã®å·®ã倧ãããã°å€§ããã»ã©Assisted Generationã®æ©æµãåããããããšãåãããŸãã sample()ã®é«éå ãµã³ããªã³ã°ã«ãã©ããã¢ãã«ã掻çšããéã¯ãå
èªã¿ã§åŸããã確çååžãçšããŠæ£åŽãµã³ããªã³ã°ãè¡ããLLMã®åºåãã確çååžãåçŸããŸããæ£åŽãµã³ããªã³ã°ãšã¯ã確çååžãéåžžã«è€éã§ããããã®ãµã³ããªã³ã°ãå°é£ã§ããæã«äœ¿ãããææ³ã§ãç®çã®ç¢ºçååž ã«å¯ŸããŠãµã³ããªã³ã°ãããã確çååž ãçšæããŠæ¬¡ã®ããã«è¡ããŸãã (åæºå) ä»»æã® ã«å¯Ÿã ãæãç«ã€ãããªå®æ° ãçšæãããïŒçšæã§ããªãå Žåã¯æ£åŽãµã³ããªã³ã°ã䜿ããªããããªãã¡ ãå®çŸ©ãããŠãã ã§ã¯åžžã« ãå®çŸ©ãããŠããå¿
èŠããããããã«ç¢ºçæ¯ ã¯æçã§ããå¿
èŠããããïŒ ãããµã³ããªã³ã°ãã äžæ§ä¹±æ°ããã®ãµã³ããªã³ã° ] ã ãæºããã° ã ããã®ãµã³ããªã³ã°çµæãšããŠæ¡çšãããã§ãªããã°æ£åŽã 2. ããããçŽãã ããã«ãã ããçŽæ¥ãµã³ããªã³ã°ããããšãªãããã®ç¢ºçååžã«åŸã ãçæã§ããŸãããã®ææ³ã®åŒ±ç¹ã®1ã€ã«ãçšæãã ãããŸãã«å€§ãããšæ£åŽã®å²åãå¢ããµã³ããªã³ã°ã®å¹çãèœã¡ãããšãæããããŸãã ãã ãäžè¬çã«äœ¿ãããŠããæ£åŽãµã³ããªã³ã°ã®åæãšã¯ç°ãªããLLMã®ç¢ºçååžããã®ãµã³ããªã³ã°ãå
èªã¿ãçµãã床ã«èšç®ã§ãããããæ¬¡ã®ãããªå·¥å€«ããªãããŠããŸãã å
¥åããŒã¯ã³ ã«å¯Ÿãããã©ããã¢ãã«ã§ åå
ãŸã§ã®ç¢ºçååž ãå
èªã¿ãã å
èªã¿ããããŒã¯ã³ãçšããŠLLMã®ç¢ºçååž ãèšç®ããïŒäžåºŠã®æšè«ã§å¯èœïŒ ã«å¯ŸããŠãäžæ§ä¹±æ° ] ã ãæºããã°æ¡çšãããã§ãªããã°æ£åŽãã åããŠæ£åŽããã ã®ããŒã¯ã³ã ãæ£èŠåããååžãããµã³ããªã³ã°ãã ã« ãå ã, 1.ã«æ»ã ãã®ææ³ãçšããããšã«ãã£ãŠã1ã«ãŒãããã1åèªãå¿
ãçæã§ãããŸãæ£åŽãµã³ããªã³ã°ã®ä»®å®ã§ãã ã®å®çŸ©åãç¢ºçæ¯ã®äžçãªã©ãæ°ã«ããå¿
èŠããªããªããŸãã å®è£
ã®ãœãŒã¹ã³ãŒãã¯ãã¡ã from typing import Dict import torch from transformers import TemperatureLogitsWarper, LogitsProcessorList from transformers.generation.stopping_criteria import StoppingCriteriaList class AssistMixin : def draft ( self, eos_token_id: int , input_ids: torch.LongTensor, max_assistant_tokens: int , do_sample: bool , ) -> torch.LongTensor: draft_ids = input_ids self.cache[ "assistant_prob_list" ] = [] for idx in range (max_assistant_tokens): if "assistant_past_key_values" in self.cache: prev_seq_len = self.cache[ "assistant_past_key_values" ][ 0 ][ 0 ].shape[- 2 ] # `new_token_len` can be 1 or 2 (next token in assistant + last token picked by the larger model) new_token_len = draft_ids.shape[ 1 ] - prev_seq_len assist_inputs = draft_ids[:, -new_token_len:] assist_attn = torch.ones_like(draft_ids) assistant_model_outputs = self.assist_model( assist_inputs, attention_mask=assist_attn, past_key_values=self.cache[ "assistant_past_key_values" ]) else : assistant_model_outputs = self.assist_model(draft_ids) self.cache[ "assistant_past_key_values" ] = assistant_model_outputs.past_key_values assist_new_logits = assistant_model_outputs.logits[:, - 1 , :] assist_new_logits = self.logits_processor(draft_ids, assist_new_logits) assist_new_logits = self.logits_warper(draft_ids, assist_new_logits) if do_sample: assist_new_probs = assist_new_logits.softmax(- 1 ) self.cache[ "assistant_prob_list" ].append(assist_new_probs) new_token = torch.multinomial(assist_new_probs, num_samples= 1 ).squeeze( 1 ) else : new_token = assist_new_logits.argmax(- 1 ) draft_ids = torch.cat((draft_ids, new_token[:, None ]), dim=- 1 ) if new_token[ 0 ] == eos_token_id: break if do_sample: self.cache[ "assistant_prob_list" ] = torch.stack(self.cache[ "assistant_prob_list" ], dim= 1 ) return draft_ids def verify ( self, eos_token_id: int , input_ids: torch.LongTensor, candidate_input_ids: torch.LongTensor, max_len: int , do_sample: bool , ) -> torch.LongTensor: candidate_length = candidate_input_ids.shape[ 1 ] - input_ids.shape[ 1 ] cur_len = input_ids.shape[ 1 ] if "past_key_values" in self.cache: model_attn = torch.ones_like(candidate_input_ids) model_input_ids = candidate_input_ids[:, -candidate_length- 1 :] outputs = self.model( model_input_ids, attention_mask=model_attn, past_key_values=self.cache[ "past_key_values" ], ) else : outputs = self.model(candidate_input_ids) logits = outputs.logits for i in range (candidate_length): logits[:, i, :] = self.logits_processor(candidate_input_ids[:, :cur_len + i], logits[:, i, :]) for i in range (candidate_length): logits[:, i, :] = self.logits_warper(candidate_input_ids[:, :cur_len + i], logits[:, i, :]) speculative_ids = candidate_input_ids[:, -candidate_length:] if do_sample: probs = logits[:, -candidate_length- 1 :, :].float().softmax(- 1 ) speculative_probs = self.cache[ "assistant_prob_list" ].gather(dim=- 1 , index=speculative_ids[:,:, None ]).squeeze(- 1 ) speculative_actual_probs = probs[:, :- 1 , :].gather(dim=- 1 , index=speculative_ids[:,:, None ]).squeeze(- 1 ) resample_probs = probs.clone() resample_probs[:, :- 1 , :] = torch.clamp(resample_probs[:, :- 1 , :] - self.cache[ "assistant_prob_list" ].float(), min = 0 ) resample_probs /= resample_probs.sum(- 1 , keepdim= True ) acceptance_thresholds = speculative_actual_probs / speculative_probs unif = torch.rand_like(acceptance_thresholds) n_matches = ((~(unif <= acceptance_thresholds)).cumsum(dim=- 1 ) < 1 ).sum() else : selected_tokens = logits[:, -candidate_length- 1 :, :].argmax(- 1 ) n_matches = ((~(speculative_ids == selected_tokens[:,:- 1 ])).cumsum(dim=- 1 ) < 1 ).sum().cpu().item() n_matches = min (max_len - cur_len, n_matches) self.cache[ "matches" ].append(n_matches) self.cache[ "past_key_values" ] = outputs.past_key_values input_ids = torch.cat((input_ids, speculative_ids[:, :n_matches]), dim=- 1 ) if input_ids[ 0 , - 1 ] == eos_token_id or input_ids.shape[ 1 ] == max_len: # if EOS or max_len, STOP return input_ids # add one more token if do_sample: return torch.cat((input_ids, torch.multinomial(resample_probs[:, n_matches, :], num_samples= 1 )), dim=- 1 ) else : return torch.cat((input_ids, selected_tokens[:, n_matches:n_matches+ 1 ]), dim=- 1 ) def crop_cache (self, assist_input_ids, large_input_ids): # Discard past key values relative to unused assistant tokens self.cache[ "past_key_values" ] = tuple ([( kv[ 0 ][:, :, :large_input_ids.shape[ 1 ]- 1 , :], kv[ 1 ][:, :, :large_input_ids.shape[ 1 ]- 1 , :], ) for kv in self.cache[ "past_key_values" ]]) self.cache[ "assistant_past_key_values" ] = tuple ([( kv[ 0 ][:, :, :assist_input_ids.shape[ 1 ]- 1 , :], kv[ 1 ][:, :, :assist_input_ids.shape[ 1 ]- 1 , :], ) for kv in self.cache[ "assistant_past_key_values" ]]) class SpecSampler (AssistMixin): def __init__ (self, tokenizer, large_model, assist_model): self.tokenizer = tokenizer self.model = large_model self.assist_model = assist_model @torch.no_grad() def generate (self, input_ids: torch.LongTensor, max_new_len: int , temperature: float ): max_len = input_ids.shape[ 1 ] + max_new_len self.cache = {} self.cache[ "matches" ] = [] self.max_assistant_tokens = 5 self.logits_processor = LogitsProcessorList() self.logits_warper = LogitsProcessorList([TemperatureLogitsWarper(temperature)]) while True : draft_ids = self.draft(self.tokenizer.eos_token_id, input_ids, max_assistant_tokens=self.max_assistant_tokens, do_sample= True ) new_input_ids = self.verify(self.tokenizer.eos_token_id, input_ids, draft_ids, max_len=max_len, do_sample= True ) n_matches = new_input_ids.shape[ 1 ] - input_ids.shape[ 1 ] - 1 if n_matches == self.max_assistant_tokens: self.max_assistant_tokens += 2 else : self.max_assistant_tokens = max ( 1 , self.max_assistant_tokens - 1 ) input_ids = new_input_ids self.crop_cache(input_ids[:, :- 1 ], input_ids) if input_ids.shape[ 1 ] >= max_len or input_ids[ 0 , - 1 ] == self.tokenizer.eos_token_id: break return input_ids ãã¡ããTransformersã§ã¯ assistant_model ãåŒæ°ã«æå®ããããšã§å©çšã§ããŸãããããŒãžã§ã³4.34.1æç¹ã§ã¯è¥å¹²å¹çãæªãå®è£
ã«ãªã£ãŠããŸãã ãšããã®ãTransformerså®è£
ã§ã¯å
èªã¿ã®ç¢ºçååžã®Softmax枩床ãåžžã«0ã«ããŠããã®ãšç䟡ãªãã®ã«ãªã£ãŠãããLLMã®åºåãã確çååžããããªãé¢ããŠããŸã£ãŠããŸãããããäžã§æãããœãŒã¹ã³ãŒãã®ããã«LLMåŽã§å©çšããæž©åºŠãšåäžã®ãã®ãå©çšããããã«å®è£
ãçŽããšã以äžã®å³ã®ããã«ã»ãšãã©ã®æž©åºŠãã©ã¡ãŒã¿ã§å¹çåã§ããŸããç¹ã«é«ã枩床ã«ãããŠããTransformerså®è£
ã§ã¯ãã©ããã¢ãã«ãªãã§ã®çæãããé
ããªã£ãŠããã®ã«å¯Ÿããç§ã®å®è£
ã§ã¯éããªã£ãŠããããšãããããŸãã assistã¯ç¬èªå®è£
ãassist(hf)ã¯Transformerså®è£
ã6.7B(float16)ã¯ãã©ããã¢ãã«ç¡ãã®é床ã125mã¯ãã©ããã¢ãã«ã®ã¿ã§çæããæã®é床 ããŒã¯ãã€ã¶ã®äºææ§ ããŒã¯ãã€ã¶(Tokenizer)ã¯å
¥åæç« ãèšèªã¢ãã«ãæ±ããããã«ããŒã¯ã³IDã®åã«å€æããã¢ãžã¥ãŒã«ã§ãããã¢ãã«ã«ãã£ãŠãã°ãã°ç°ãªãããŒã¯ãã€ã¶ã䜿ãããŠããŸãã ãããAssisted Generationã䜿ãã«ããããLLMãšãã©ããã¢ãã«ã®ããŒã¯ãã€ã¶ã¯åããã®ã䜿ãå¿
èŠããããŸããããã¯å
éšã§ã¯ããŒã¯ã³IDãçšããŠã¢ãã«éãããåãããŠãããIDãå
±éã§ããã°ããŒã¯ã³åããã®ãŸãŸæž¡ããŠå¹çãè¯ãããã§ããããã«å¯Ÿãã軜éçã¢ãã«ãååšããªããªã©ã§ã©ãããŠãç°ãªãããŒã¯ãã€ã¶ã䜿ãããå Žåã¯2ã€ã®ããŒã¯ãã€ã¶ã§ããŒã¯ã³åãçžäºã«å€æããããšã§Assisted Generationãæ£ããåããããšãã§ããŸãããã ãããŸãé«éåã¯æåŸ
ã§ããªããããããŸããã以äžã®äŸã§ã¯llm-jp-13Bãšopen-calm-small (160Mãã©ã¡ãŒã¿)ãçšããŠçæããŠããŸãããçæé床ã«ããŸãå·®ãçãŸããŸããã§ããã å®è£
ã®ãœãŒã¹ã³ãŒãã¯ãã¡ã import torch import numpy as np import transformers from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig from transformers import TemperatureLogitsWarper, LogitsProcessorList import time class JpCALM (AssistMixin): def __init__ (self, large_tokenizer, assist_tokenizer, large_model, assist_model): self.large_tokenizer = large_tokenizer self.assist_tokenizer = assist_tokenizer self.model = large_model self.assist_model = assist_model def generate (self, text: str , max_len: int ): assist_input_ids = self.assist_tokenizer.encode(text, add_special_tokens= False , return_tensors= "pt" ).to(self.assist_model.device) large_input_ids = self.large_tokenizer.encode(text, add_special_tokens= False , return_tensors= "pt" ).to(self.model.device) self.cache = {} self.cache[ "matches" ] = [] self.logits_processor = LogitsProcessorList() self.logits_warper = LogitsProcessorList() while True : # 5ã€å
èªã¿ assist_draft_ids = self.draft(self.assist_tokenizer.eos_token_id, assist_input_ids, max_assistant_tokens= 5 , do_sample= False ) # å
èªã¿ããassistantã®tokenã倿 candidate_words = self.assist_tokenizer.decode(assist_draft_ids[ 0 , assist_input_ids.shape[ 1 ]:], skip_special_tokens= True ) large_candidate_ids = self.large_tokenizer.encode(candidate_words, add_special_tokens= False , return_tensors= "pt" ).to(self.model.device)[:, 1 :] large_draft_ids = torch.cat((large_input_ids, large_candidate_ids), dim= 1 ) # æ€èšŒ large_next_input_ids = self.verify(self.large_tokenizer.eos_token_id, large_input_ids, large_draft_ids, max_len, do_sample= False ) # attentionã®ãã£ãã·ã¥ãæŽç self.crop_cache(assist_input_ids, large_next_input_ids) # æ€èšŒããLLMã®tokenã倿 selected_tokens = large_next_input_ids[:, large_input_ids.shape[ 1 ]:] large_input_ids = large_next_input_ids large_input_words = self.large_tokenizer.decode(large_input_ids[ 0 ], skip_special_tokens= True ) valid_words = self.large_tokenizer.decode(selected_tokens[ 0 ], skip_special_tokens= True ) assist_valid_ids = self.assist_tokenizer.encode(valid_words, add_special_tokens= False , return_tensors= "pt" ).to(self.assist_model.device) assist_input_ids = torch.cat((assist_input_ids, assist_valid_ids.long()), dim=- 1 ) if large_input_ids.shape[ 1 ] >= max_len or large_input_ids[ 0 , - 1 ] == self.large_tokenizer.eos_token_id: break return large_input_words large_tokenizer = AutoTokenizer.from_pretrained( "llm-jp/llm-jp-13b-v1.0" ) assist_tokenizer = AutoTokenizer.from_pretrained( "cyberagent/open-calm-small" ) large_model = AutoModelForCausalLM.from_pretrained( "llm-jp/llm-jp-13b-v1.0" , device_map= "auto" , torch_dtype=torch.float16, load_in_8bit= True ) assist_model = AutoModelForCausalLM.from_pretrained( "cyberagent/open-calm-small" , device_map= "auto" , torch_dtype=torch.float16) assisted_jp = JpCALM(large_tokenizer, assist_tokenizer, large_model, assist_model) print ( "=========assisted===========" ) tic = time.perf_counter() with torch.no_grad(): out1 = assisted_jp.generate( "èªç¶èšèªåŠçãšã¯äœã" , 128 ) tac = time.perf_counter() print (out1) print (tac - tic, "sec" ) matches = np.array(assisted_jp.cache[ "matches" ]) print ( "generated tokens per cycle:" , matches.mean() + 1 ) print ( "acceptance rate:" , (matches > 0 ).sum() / len (matches)) print ( "=========baseline===========" ) tic = time.perf_counter() with torch.no_grad(): tokenized_input = assisted_jp.large_tokenizer.encode( "èªç¶èšèªåŠçãšã¯äœã" , add_special_tokens= False , return_tensors= "pt" ).to(assisted_jp.model.device) output = assisted_jp.model.generate( tokenized_input, max_length= 128 , )[ 0 ] out2 = assisted_jp.large_tokenizer.decode(output) tac = time.perf_counter() print (out2) print (tac - tic, "sec" ) assert out1 == out2 """ =========assisted=========== èªç¶èšèªåŠçãšã¯äœã (å²©æ³¢æ°æž) | è¥¿å£ é |æ¬ | é販 | AmazonKindleã¹ãã¢ã§ã¯ã èªç¶èšèªåŠçãšã¯äœã (å²©æ³¢æ°æž)ããKindleç¡æã¢ããªãŸãã¯Kindleé»åæžç±ãªãŒããŒã§ä»ãããèªã¿ããã ããŸããKindleé»åæžç±ãªãŒããŒã® 詳现ã¯ãã¡ãèªç¶èšèªåŠçãšã¯äœã (å²©æ³¢æ°æž) ãã«ãŒãã«å
¥ããŸããèªç¶èšèªåŠçãšã¯äœã (å²©æ³¢æ°æž) æ°æž â 2016/11/21ï¿¥ 886 ï¿¥ 13.10530449775979 sec generated tokens per cycle: 1.7910447761194028 acceptance rate: 0.23880597014925373 =========baseline=========== èªç¶èšèªåŠçãšã¯äœã (å²©æ³¢æ°æž) | è¥¿å£ é |æ¬ | é販 | AmazonKindleã¹ãã¢ã§ã¯ã èªç¶èšèªåŠçãšã¯äœã (å²©æ³¢æ°æž)ããKindleç¡æã¢ããªãŸãã¯Kindleé»åæžç±ãªãŒããŒã§ä»ãããèªã¿ããã ããŸããKindleé»åæžç±ãªãŒããŒã® 詳现ã¯ãã¡ãèªç¶èšèªåŠçãšã¯äœã (å²©æ³¢æ°æž) ãã«ãŒãã«å
¥ããŸããèªç¶èšèªåŠçãšã¯äœã (å²©æ³¢æ°æž) æ°æž â 2016/11/21ï¿¥ 886 ï¿¥ 14.313117458950728 sec """ æŽãªãé«éåïŒ Token Tree Verification Assisted Generationã®é床ã¯å
èªã¿ãã©ãã ãåœãããã«ããã£ãŠããŸãããããé£ç¶ããŠåœãŠããã確çã¯ææ°é¢æ°çã«æžå°ããŠããããããã©ããã¢ãã«ãLLMã®åºåååžãããŸãåçŸãããªãã®ç²ŸåºŠã§å
èªã¿ãåœãŠãªããšçæé床ãäŒžã³æ©ãã§ããŸããŸããäŸãã°åç¯ã§æããllm-jp-13Bãšopen-calm-smallã®çµã¿åããã§ã¯å¹³åã§1.79åããçæã§ããŠããŸããã§ãããããã«å¯Ÿãããã©ããã¢ãã«ãå€ããã«å
èªã¿ã®æ£è§£çãåäžããã工倫ãšããŠToken Tree Verification 4 ãšããææ³ã掻çšã§ããŸãã LLMã®æç« çæã¡ã«ããºã â¡ Causal Language Modelã«ãããAttention Maskã«ã€ã㊠æç« çæã«äœ¿ãããŠããCausal Language Modelã§ã¯åŸãã®å
¥åããŒã¯ã³ãåã®äºæž¬ããŒã¯ã³ã«åœ±é¿ãäžããªããšãã話ããããŸãããããã¯Transformerã¢ãã«å
ã§ããŒã¯ã³å士ã圱é¿ãåãSelf-Attentionéšåã«ãã¹ã¯ãæœãããšã§å®çŸããŠããŸãã ãã®ãã¹ã¯ãæ¹é ããããšã§ããè€éãªå¶çŽã課ãããšãã§ããŸãã äŸãã° "The quick brown fox jumps over" ãš "The quick brown fox runs around" ãåæã«LLMã«å
¥åãããå Žåã "The quick brown fox jumps over runs around" ãš1è¡ã«ãŸãšããã®ã¡ã"jumps"ãš"runs", "over"ãš"around"ãªã©ãšãã£ãããŒã¯ã³éã®attentionããã¹ã¯ãåãé€ãããšã§ãããããã®æç« ãç¬ç«ã«æšè«ããæãšåãçµæãåŸãããšãã§ããŸãããã®ææ³ã¯ãåæ¹ã§äžèŽããŠããããŒã¯ã³åãé·ããã°é·ãã»ã©ãè€æ°æããããã§ãŸãšããããã¡ã¢ãªéãèšç®éçãªé¢ã§æå©ã«ãªããŸãã ãããäžè¬åãããšãåå²ããå
èªã¿ããŒã¯ã³æšã1åã®LLMæšè«ã§æ€èšŒããããšãã§ããŸãã å
è«æ ã§ã¯è€æ°ã®ãã©ããã¢ãã«ããããããå
èªã¿ãçæãããããã1ã€ã®ããŒã¯ã³æšã«éçŽããŠããŸããããä»åã¯1ã€ã®ãã©ããã¢ãã«ããè€æ°ã®å
èªã¿ãçæããŠã¿ãŸãã Token Tree Verificationãçšããgenerate() facebook/opt ã¢ãã«ã®äŸã«æ»ã£ãŠãå
èªã¿ãã©ããããåœãã£ãŠããã®ããå¯èŠåããŠã¿ãŸãã6.7B, 13Bã«å¯ŸããŠ125Mã®ã¢ãã«ããã©ããã¢ãã«ãšããŠå©çšããLLMãåºåããããŒã¯ã³ããã©ããã¢ãã«ã«ãããŠäœçªç®ãŸã§ã®åè£ã«å
¥ã£ãŠãããã以äžã®å³ã«ç€ºããŸããã ãããŸã§ã®ææ³ã§ã¯ãã©ããã¢ãã«ãæ¡çšããã®ã¯äžçªç¢ºçã®é«ãããŒã¯ã³ãããªãã¡top-1ã®ããŒã¯ã³ãªã®ã§çŽ70%ã®å²åã§å
èªã¿ã«æåããŠããããšãããããŸããããã«top-2ãŸã§æ¡çšãããšãããšæ£è§£çã¯80%ã«äžæããŸããå
èªã¿ã«è€æ°éžæè¢ãçšæããããšã«ã¯å¹æãããããã§ãã ã§ã¯ã©ã®ããã«éžæè¢ãçšæããã°è¯ãã§ãããããããŸãåœãããªãããªå
èªã¿ããŒã¯ã³ã«å¯ŸããŠã¯ããã®ããã«å
ãèªãããšã«ããå©çã¯å°ãªããã§ããããã§ãã©ããã¢ãã«ã®åºå確çããã®ãŸãŸå
èªã¿ãåœãã確çãšèŠåããŠããŸããå
èªã¿ãåœããæ°ã®æåŸ
å€ãæå€§ã«ãªãããã«æšãçæããŸãã å®è£
ã®ãœãŒã¹ã³ãŒãã¯ãã¡ã from dataclasses import dataclass, field import heapq from typing import Dict, List import torch import torch.nn.functional as F from transformers.generation.logits_process import LogitsProcessorList from transformers.generation.stopping_criteria import StoppingCriteriaList @ dataclass (order= True ) class Node : nll: float # 確çã®negative logãå°ããé ã§å¹
åªå
æ¢çŽ¢ãè¡ã gain: int tokens: torch.LongTensor = field(compare= False ) attention_mask: torch.LongTensor = field(compare= False ) all_draft_tokens: torch.LongTensor = field(compare= False ) class AssistTreeMixin : def draft ( self, eos_token_id: int , input_ids: torch.LongTensor, max_assistant_tokens: int , ) -> torch.LongTensor: self.cache[ "assistant_prob_list" ] = [] seq_len = input_ids.shape[ 1 ] max_len = seq_len + max_assistant_tokens pq = [] device = self.assist_model.device heapq.heappush(pq, Node( 0 , 0 , torch.LongTensor([]).to(device), torch.ones_like(input_ids), torch.LongTensor([]).to(device))) draft_ids = input_ids draft_masks = [] draft_nodes = [] for idx in range (max_assistant_tokens+ 1 ): top = heapq.heappop(pq) if len (top.tokens) > 0 : draft_ids = torch.cat((draft_ids, top.tokens[:, None ]), dim=- 1 ) draft_len = draft_ids.shape[ 1 ] attention_mask = F.pad(top.attention_mask, ( 0 , draft_len - top.attention_mask.shape[ 1 ])) attention_mask[:, - 1 ] = 1 if idx > 0 : # ããŒã¯ã³æšã«æ¡çš draft_masks.append(F.pad(attention_mask, ( 0 , max_len - attention_mask.shape[ 1 ]))) draft_nodes.append(top) if "assistant_past_key_values" in self.cache: prev_seq_len = self.cache[ "assistant_past_key_values" ][ 0 ][ 0 ].shape[- 2 ] new_token_len = draft_ids.shape[ 1 ] - prev_seq_len assist_inputs = draft_ids[:, -new_token_len:] assistant_model_outputs = self.assist_model( assist_inputs, attention_mask=attention_mask, past_key_values=self.cache[ "assistant_past_key_values" ]) else : assistant_model_outputs = self.assist_model(draft_ids, attention_mask=attention_mask) self.cache[ "assistant_past_key_values" ] = assistant_model_outputs.past_key_values assist_new_logits = assistant_model_outputs.logits[:, - 1 , :] # (batch, vocab) assist_new_logits = self.logits_processor(draft_ids, assist_new_logits) assist_new_logits = self.logits_warper(draft_ids, assist_new_logits) assist_new_logprobs = F.log_softmax(assist_new_logits, dim=- 1 ) # (batch, vocab) # èšç®éç¯çŽã®ããã«æ¢çŽ¢æ°ã5ã€ã«å¶é assist_new_topk = torch.topk(assist_new_logprobs, k= 5 , dim=- 1 ) # (batch, k) for k in range ( 5 ): new_token = assist_new_topk.indices[:, k] new_nll = -assist_new_topk.values[ 0 , k] if new_token != eos_token_id: heapq.heappush(pq, Node(top.nll + new_nll, top.gain+ 1 , new_token, attention_mask, torch.cat((top.all_draft_tokens, new_token[:])))) return draft_ids, torch.cat(draft_masks, dim= 0 ), draft_nodes def verify ( self, eos_token_id: int , input_ids: torch.LongTensor, candidate_input_ids: torch.LongTensor, draft_masks: torch.LongTensor, draft_nodes: List[Node], max_len: int , ) -> torch.LongTensor: candidate_length = candidate_input_ids.shape[ 1 ] - input_ids.shape[ 1 ] tgt_len = candidate_input_ids.shape[ 1 ] def make_tree_mask (_attention_mask, _input_shape, inputs_embeds, past_key_values_length): # Causal Maskãäžæžããã tree_mask = torch.tril(torch.ones(tgt_len, tgt_len)) tree_mask[-candidate_length:, -candidate_length:] = draft_masks[:, input_ids.shape[ 1 ]:] tree_mask = torch.full((tgt_len, tgt_len), torch.finfo(inputs_embeds.dtype).min).masked_fill(tree_mask > 0 , 0 ) if past_key_values_length > 0 : tree_mask = torch.cat((torch.zeros(tgt_len-past_key_values_length, past_key_values_length), tree_mask[past_key_values_length:, past_key_values_length:]), dim=- 1 ) return tree_mask[ None , None ].to(inputs_embeds.dtype).to(inputs_embeds.device) self.model.model.decoder._prepare_decoder_attention_mask = make_tree_mask cur_len = input_ids.shape[ 1 ] if "past_key_values" in self.cache: prev_seq_len = self.cache[ "past_key_values" ][ 0 ][ 0 ].shape[- 2 ] new_token_len = candidate_input_ids.shape[ 1 ] - prev_seq_len model_attn = torch.ones_like(candidate_input_ids) model_input_ids = candidate_input_ids[:, -new_token_len:] outputs = self.model( model_input_ids, attention_mask=model_attn, past_key_values=self.cache[ "past_key_values" ], ) else : outputs = self.model(candidate_input_ids) logits = outputs.logits for i in range (candidate_length): logits[:, i, :] = self.logits_processor(candidate_input_ids[:, :cur_len + i], logits[:, i, :]) for i in range (candidate_length): logits[:, i, :] = self.logits_warper(candidate_input_ids[:, :cur_len + i], logits[:, i, :]) speculative_ids = candidate_input_ids[:, -candidate_length:] selected_tokens = logits[:, -candidate_length- 1 :, :].argmax(- 1 ) best_sele = torch.LongTensor([]).to(input_ids.device) best_draft_mask = torch.cat((torch.ones(input_ids.shape[ 1 ]), torch.zeros(candidate_length))) n_matches = - 1 longest_tokens = 0 for i, node in enumerate (draft_nodes): selected_tokens_i = torch.cat((selected_tokens[ 0 , 0 : 1 ], selected_tokens[ 0 , 1 :][draft_masks[i,-candidate_length:]> 0 ])) streak = (~(node.all_draft_tokens == selected_tokens_i[:- 1 ])).cumsum( 0 ) < 1 n_matches_i = streak.sum().cpu().item() longest_tokens = max (longest_tokens, len (node.all_draft_tokens)) if n_matches_i > n_matches: n_matches = n_matches_i best_sele = selected_tokens_i best_draft_mask = draft_masks[i] self.cache[ "best" ] = longest_tokens == n_matches self.cache[ "mask_to_cache" ] = best_draft_mask > 0 n_matches = min (max_len - cur_len, n_matches) self.cache[ "matches" ].append(n_matches) self.cache[ "past_key_values" ] = outputs.past_key_values verified = torch.cat((input_ids, best_sele[ None , :n_matches]), dim=- 1 ) if verified[ 0 , - 1 ] == eos_token_id or verified.shape[ 1 ] == max_len: return verified # add one more token verified = torch.cat((verified, best_sele[ None , n_matches:n_matches+ 1 ]), dim=- 1 ) return verified def crop_cache (self, input_ids): # Discard past key values relative to unused assistant tokens mask = self.cache[ "mask_to_cache" ] length = input_ids.shape[ 1 ] - 2 self.cache[ "past_key_values" ] = tuple ([( kv[ 0 ][:, :, mask, :][:,:,:length,:], kv[ 1 ][:, :, mask, :][:,:,:length,:], ) for kv in self.cache[ "past_key_values" ]]) self.cache[ "assistant_past_key_values" ] = tuple ([( kv[ 0 ][:, :, mask, :][:,:,:length,:], kv[ 1 ][:, :, mask, :][:,:,:length,:], ) for kv in self.cache[ "assistant_past_key_values" ]]) class SpecDecoder (AssistTreeMixin): def __init__ (self, tokenizer, large_model, assist_model): self.tokenizer = tokenizer self.model = large_model self.assist_model = assist_model @torch.no_grad() def generate (self, input_ids: torch.LongTensor, max_new_len: int , only_draft= False ): max_len = input_ids.shape[ 1 ] + max_new_len self.cache = {} self.cache[ "matches" ] = [] self.cache[ "first_assistant_prob" ] = [] self.cache[ "verified" ] = [] self.max_assistant_tokens = 5 while True : draft_ids, draft_masks, draft_nodes = self.draft(self.tokenizer.eos_token_id, input_ids, max_assistant_tokens=self.max_assistant_tokens) self.cache[ "draft" ] = (draft_ids, draft_masks, draft_nodes) if only_draft: break new_input_ids = self.verify(self.tokenizer.eos_token_id, input_ids, draft_ids, draft_masks, draft_nodes, max_len=max_len) n_matches = new_input_ids.shape[ 1 ] - input_ids.shape[ 1 ] - 1 if self.cache[ "best" ]: self.max_assistant_tokens += 2 else : self.max_assistant_tokens = max ( 1 , self.max_assistant_tokens - 1 ) self.crop_cache(new_input_ids) input_ids = new_input_ids if input_ids.shape[ 1 ] >= max_len or input_ids[ 0 , - 1 ] == self.tokenizer.eos_token_id: break return input_ids facebook/opt-125m ããã©ããã¢ãã«ãšããŠã"The man worked as a"ã®ç¶ãã5ã€å
èªã¿ãããšæ¬¡ã®ããã«ãªããŸããã ç¢å°ã®äžèº«ã¯é·ç§»ç¢ºçã§ãã"worked as a"ã®æ¬¡ã®åèªã¯èªä¿¡ããªããã®ã®ã"security"ãšæ¥ãã"guard", "truck"ãšæ¥ãã"driver"ãšããã®ã¯ãããªãã«èªä¿¡ãããããã§ãã ãã®ããã«ã"security guard"ãšå
èªã¿ããã ãã§ã¯ãããããšãªãæã«"truck driver"ãšããå
èªã¿ãå«ããããšãã§ãããšããã®ããã®ããŒã¯ã³æšã掻çšããå©ç¹ã§ãã ãããçšããŠAssisted Generationãã©ãã ãéããªããå®éšããŸããã ...æ®å¿µãªããéãã¯ãªããŸããã§ãããããŒã¯ã³æšã®çæã¢ã«ãŽãªãºã ãå
èªã¿æ°ãªã©ã®ãã¥ãŒãã³ã°ãå¿
èŠããã§ãã å
è«æã®å®è£
å
è«æ ã§ã¯è€æ°ã®ãã©ããã¢ãã«ã䜿ã£ãŠå
èªã¿ãããããšãæ³å®ããŠãããè€æ°ã®GPUã«ãã©ããã¢ãã«ã忣ããããªã©ã®å¹çåææ³ã«ã€ããŠãè§ŠããŠããŸããããã«å ããŠããã®ç« ã§ã¯Greedy Decodingã«ããè§ŠããŸããã§ãããããµã³ããªã³ã°ãè¡ã£ãéã®é床ã«ã€ããŠãå®éšãããŠããŸãããŸããå
¬åŒå®è£
㯠FlexFlow ã§å©çšã§ããŸãã ãããã« æ¬èšäºã§ã¯ãAssisted GenerationãšããLLMã®æšè«é«éåææ³ã«ã€ããŠã玹ä»ããŸããã NTT Com ã§ã¯ãå€§èŠæš¡èšèªã¢ãã«ããã³çæAIãæŽ»çšãããããã¯ãã»ãœãªã¥ãŒã·ã§ã³éçºã瀟å
DXã«ææŠãããœãããŠã§ã¢ãšã³ãžãã¢ãåéããŠããŸãïŒè©³çްã¯ä»¥äžã®ãªã³ã¯ãã芧ãã ããã hrmos.co ã€ãããŒã·ã§ã³ã»ã³ã¿ãŒã§ã¯ãæ¬äººã®æå¿ã«å¿ããŠè€æ°ã®ãããžã§ã¯ããžåç»ã§ããå
Œåå¶åºŠãçšæãããŠããŸãã詳ãã㯠ã€ãããŒã·ã§ã³ã»ã³ã¿ãŒ ãã¯ãããžãŒéšé 玹ä»ããã ãã芧ãã ããã ↩ https://huggingface.co/blog/assisted-generation ↩ https://arxiv.org/abs/2302.01318 ↩ https://arxiv.org/abs/2305.09781v2 ↩