o Zhs @sddlZddlmZddlZddlmZddlmZer/ddlm Z m Z m Z ddl m Z mZGdd d ejjZGd d d ejjZGd d d ejjZ    dde deejdeejdeedeef ddZGdddejjZGdddejjZGdddejjZdS)N)Optional)GenerationConfig)is_torch_available) HybridCachePreTrainedModel StaticCache)is_torch_greater_or_equal"is_torch_greater_or_equal_than_2_3cseZdZdZ  d"dededeffdd Zd ejd ejd ejfd d Z    d#d e ejd e ejde e de e d ej jf ddZ e      d$dej jdedede dedededed efd d!ZZS)%%TorchExportableModuleForDecoderOnlyLMa A recipe module designed to make a `PreTrainedModel` exportable with `torch.export`, specifically for decoder-only LM with cache. This module ensures that the exported model is compatible with further lowering and execution in `ExecuTorch`. modelmax_batch_size max_cache_lencstt|jdr|jjdurtdt|jds'tdt||_ d S|jj dkr6t ||||_ d Std|jj d) a Initializes the exportable module with `HybridCache`. Args: model (`PreTrainedModel`): The pretrained model to wrap. max_batch_size (int): Maximum batch size for the cache. max_cache_len (int): Maximum sequence length for the cache. Raises: ValueError: If the model is configured with a unsupported cache implementation. use_cacheFz5The model must have caching enabled to be performant.cache_implementationzXUsing `StaticCache` for export as `cache_implementation` is not specified in the config.hybridz"Unsupported cache implementation: z". Please use `hybrid` or `static`.N) super__init__hasattrconfigr ValueErrorlogginginfo$TorchExportableModuleWithStaticCacherr$TorchExportableModuleWithHybridCache)selfrrr __class__S/var/www/auris/lib/python3.10/site-packages/transformers/integrations/executorch.pyr#s    z.TorchExportableModuleForDecoderOnlyLM.__init__ input_idscache_positionreturncCs|j||S) Forward pass of the module, which is compatible with the ExecuTorch llm runner. Args: input_ids (`torch.Tensor`): Tensor representing current input token id to the module. cache_position (`torch.Tensor`): Tensor representing current input position in the cache. Returns: torch.Tensor: Logits output from the model. )rforward)rr"r#r r r!r&Gsz-TorchExportableModuleForDecoderOnlyLM.forwardNdynamic_shapesstrictcCsj|dur|n tjdggtjd}|dur|ntjdgtjd}tjj|j||fi||dur1|dSddS)a Export the wrapped module using `torch.export`. Args: input_ids (`Optional[torch.Tensor]`): Tensor representing current input token id to the module. If not provided, a default tensor will be used. cache_position (`Optional[torch.Tensor]`): Tensor representing current input position in the cache. If not provided, a default tensor will be used. dynamic_shapes (`Optional[dict]`): Dynamic shapes to use for export if specified. strict(`Optional[bool]`): Flag to instruct `torch.export` to use `torchdynamo`. Nr dtyperTargskwargsr'r()torchtensorlongexportr)rr"r#r'r(example_input_idsexample_cache_positionr r r!r1Xs  z,TorchExportableModuleForDecoderOnlyLM.exportF?2cpuexported_programpromptmax_new_tokens do_sample temperaturetop_ktop_pdevicec Cs|} ||ddj|} | } d} t| jdD]!} | dd| | df}tj| gtj|d}| ||}| d7} qt|D]}| ddddf}tj| gtj|d}| ||}|r|dkrg||}n|}|dkr|t ||ddk}t d ||<|d krtj |d d \}}tj tj |dd dd }||k}|dddf|dddf<d|d<|d||}t d ||<tj |dd }tj|dd}n|jdd d}|dkr|d}tj| |gdd } | d7} ||jkrnqA|j| dd dS)a Generate a sequence of tokens using an exported program. Args: exported_program (`torch.export.ExportedProgram`): The exported model being used for generate. tokenizer: The tokenizer to use. prompt (str): The input prompt. max_new_tokens (int): Maximum number of new tokens to generate. do_sample (bool): Whether to use sampling or greedy decoding. temperature (float): The temperature for sampling. top_k (int): The number of highest probability tokens to keep for top-k sampling. top_p (float): The cumulative probability for nucleus sampling. device (str): The device to use. Returns: str: The generated text. pt)Zreturn_tensorsrr Nr*r?).rBNz-infr5T)Z descendingdim.).r)Z num_samples)rDZkeepdimr)Zskip_special_tokens)moduler"toclonerangeshaper.r/r0ZtopkfloatsortZcumsumZsoftmaxZscatterZ multinomialargmaxrDZsqueezecatitem eos_token_iddecode)r8Z tokenizerr9r:r;r<r=r>r?Zexported_moduler" generated_idsZ curr_positioniZcurr_input_idsZcurr_cache_position_outputslogitsZindices_to_removeZ sorted_logitsZsorted_indicesZcumulative_probsZsorted_indices_to_removeZprobsZ next_token_idr r r!generatewsN          z.TorchExportableModuleForDecoderOnlyLM.generater r NNNN)r4Fr5r6r5r7)__name__ __module__ __qualname____doc__rintrr.Tensorr&rdictboolr1ExportedProgram staticmethodstrrJrV __classcell__r r rr!r sv $    r c sbeZdZdZdeffdd ZdejdejfddZe d ej j d ejd e d ejfd dZ ZS)ra A recipe module designed to make a `PreTrainedModel` exportable with `torch.export`, specifically for decoder-only LM to `StaticCache`. This module ensures that the exported model is compatible with further lowering and execution in `ExecuTorch`. Note: This class is specifically designed to support export process using `torch.export` in a way that ensures the model can be further lowered and run efficiently in `ExecuTorch`. rcs$t|jdurtd|jjstd|jjdkr td||_t|jj|jjj j |jjj j |jjj j |jj d|_tt|jjD] }|jd||jj|dd |jd ||jj|dd qEtd d |jjjD|_|jrttj|jj |jj tjd }|jd|dd dSdS)a Initializes the wrapper module with the pretrained model. Args: model (`PreTrainedModel`): The pretrained model to wrap. The model must have caching enabled and use a 'static' caching implementation. Raises: AssertionError: If the pretrained model does not have caching enabled or if it does not use a 'static' caching implementation in `model.generation_config`. NzkThe model must have a generation config to be exported with static caching. Please set `generation_config`.zvThe model must have caching enabled to be exported with static caching. Please set `generation_config.use_cache=True`.staticzThe model must use a 'static' caching implementation to be exported with static caching. Please set `generation_config.cache_implementation='static'`.rrrr?r* key_cache_F persistent value_cache_css|]}d|vVqdS)ZCausalLMNr ).0archr r r! sz@TorchExportableModuleWithStaticCache.__init__..r)mask)rrgeneration_configAssertionErrorrrrrr cache_config batch_sizerr?r* static_cacherHlen key_cacheregister_buffer value_cacheanyZ architectures is_causalr.Ztrilonesr`)rrrRZ causal_maskrr r!rsF      z-TorchExportableModuleWithStaticCache.__init__r"r#c CsR|j\}}|jr|j|d|fnd}|d}|j}|j|||||dd}|jS)a Forward pass of the module, which is compatible with the ExecuTorch runtime. Args: input_ids (`torch.Tensor`): Tensor representing current input token id to the module. cache_position (`torch.Tensor`): Tensor representing current input position in the cache. Returns: torch.Tensor: Logits output from the model. This forward adapter serves two primary purposes: 1. **Making the Model `torch.export`-Compatible**: The adapter hides unsupported objects, such as the `Cache`, from the graph inputs and outputs, enabling the model to be exportable using `torch.export` without encountering issues. 2. **Ensuring Compatibility with `ExecuTorch` runtime**: The adapter matches the model's forward signature with that in `executorch/extension/llm/runner`, ensuring that the exported model can be executed in `ExecuTorch` out-of-the-box. NrT)r"attention_mask position_idsr#past_key_valuesr)rIryrn unsqueezersrrU) rr"r#rSZseqlenZ attn_maskr|r}Zoutsr r r!r&'s  z,TorchExportableModuleWithStaticCache.forwardr8prompt_token_idsr:r$c Csb|jd}||}|D]\}}|dr"|jd}t||}nq g}tt||D]'} |j|dd| | dftj| gtj dd} | |d|  q,tj | dddddfdd  } | | t ||kr|jtj| ggtj dtjt |gtj dd} tj | dddddfdd  } | | t ||ksqtj|gtj dS) a Generate a sequence of tokens using an exported program. This util function is designed to test exported models by simulating the generation process. It processes the input prompt tokens sequentially (no parallel prefill). This generate function is not intended to replace the original `generate` method, and the support for leveraging the original `generate` is potentially planned! Args: exported_program (`torch.export.ExportedProgram`): The exported program generated via `torch.export`. prompt_token_ids (`torch.Tensor`): Tensor representing the input prompt token IDs. max_new_tokens (`int`): Maximum number of new tokens to generate. Note that the total generation length is limited by both `max_new_tokens` and the model's cache size. Returns: torch.Tensor: A tensor containing the generated sequence of token IDs, including the original prompt tokens. rBrurNr r))r"r#rrC)rIZ named_buffers startswithminrHrEr&r.r/r0appendrNrLrt) r8rr:Zprompt_token_lenZmax_generation_lengthZ buffer_namebufferrZresponse_tokensZ input_posresult current_tokenr r r!rVKs4    $  $  z-TorchExportableModuleWithStaticCache.generate)rYrZr[r\rrr.r^r&rbr1rar]rVrdr r rr!rs 8$rcsPeZdZdZ  ddededeffdd Zd ejd ejd ejfd d Z Z S)ra A recipe module designed to make a `PreTrainedModel` exportable with `torch.export`, specifically for decoder-only LM to `HybridCache`. This module ensures that the exported model is compatible with further lowering and execution in `ExecuTorch`. r r rrrcst||_|jjjstdt|jjdr|jjjdkr#tdt|jj|||jj |jj d|_ t t |j jD] }|jd||j j|dd|jd ||j j|ddq;d S) a Initializes the exportable module with `HybridCache`. Args: model (`PreTrainedModel`): The pretrained model to wrap. max_batch_size (int): Maximum batch size for the cache. max_cache_len (int): Maximum sequence length for the cache. Raises: AssertionError: If the model doesn't have the expected configuration for HybridCache. zModel must have caching enabledrrz,Model must use 'hybrid' cache implementationrfrgFrhrjN)rrrrrrprrrr?r*cacherHrtrurvrw)rrrrrRrr r!rs&     z-TorchExportableModuleWithHybridCache.__init__r"r#r$cCsR|j\}}|d|d}tj||ftj|jd}|j||||jd|d}|j S)r%rrBrAT)r"r{r|r}rr#) rIr~expandr.rzr0r?rrrU)rr"r#rrZseq_lenr|r{rTr r r!r&s  z,TorchExportableModuleWithHybridCache.forwardrW) rYrZr[r\rr]rr.r^r&rdr r rr!r~s$ ,rrr2r3r'r(cCststdddl}|f|dur|n |jdgg|jd}|dur%|n|jdg|jd}tdrH|jjt|||fi||durC|ndd}n#|durQt d |durZt d |jj j t||fd |id dd }|WdS1swwYdS)a Convert a `PreTrainedModel` into an exportable module and export it using `torch.export`, ensuring the exported model is compatible with `ExecuTorch`. Args: model (`PreTrainedModel`): The pretrained model to be exported. example_input_ids (`Optional[torch.Tensor]`): Example input token id used by `torch.export`. example_cache_position (`Optional[torch.Tensor]`): Example current cache position used by `torch.export`. dynamic_shapes(`Optional[dict]`): Dynamic shapes used by `torch.export`. strict(`Optional[bool]`): Flag to instruct `torch.export` to use `torchdynamo`. Returns: Exported program (`torch.export.ExportedProgram`): The exported program generated via `torch.export`. ztorch >= 2.3 is required.rNr r)z2.6.0Tr+zWDynamic shapes spec will be ignored by convert_and_export_with_cache for torch < 2.6.0.zSThe strict flag will be ignored by convert_and_export_with_cache for torch < 2.6.0.r#F)r,r-Z pre_dispatchr() r ImportErrorZtorch.export._traceno_gradr/r0r r1rrwarning_traceZ_export)rr2r3r'r(r.r8r r r!convert_and_export_with_caches>  $rc(eZdZdZfddZddZZS) Seq2SeqLMEncoderExportableModulez A wrapper module designed to make a Seq2Seq LM encoder exportable with `torch.export`. This module ensures that the exported encoder model is compatible with ExecuTorch. cst||_dS)N)rrencoder)rZ encoder_modelrr r!rs  z)Seq2SeqLMEncoderExportableModule.__init__cCs|j|djS)N)r")rZlast_hidden_state)rr"r r r!r& sz(Seq2SeqLMEncoderExportableModule.forwardrYrZr[r\rr&rdr r rr!rs rcr)/Seq2SeqLMDecoderExportableModuleWithStaticCachez A wrapper module designed to make a Seq2Seq LM decoder exportable with `torch.export`, specifically for use with static caching. This module ensures the exported decoder is compatible with ExecuTorch. cst||_|j|_|j|_t|j||dtjd|_ t t |j j D] }|j d||j j |dd|j d||j j|ddq&dS)Nr7rfrgFrhrj)rrZ get_decoderdecoderlm_headrrr.float32rsrHrtrurvrw)rrmax_static_cache_lengthrrrRrr r!r+s    z8Seq2SeqLMDecoderExportableModuleWithStaticCache.__init__cCs(|j|||jd|d}||d}|S)NT)r"encoder_hidden_statesr}rr#r)rrsr)rdecoder_input_idsrr#rTZ lm_logitsr r r!r&As z7Seq2SeqLMDecoderExportableModuleWithStaticCache.forwardrr r rr!r$s rcsBeZdZ dfdd ZddZd d Zdd d ZddZZS)Seq2SeqLMExportableModuler r recsPt||_||_|j|_||_td||||dd|_d|_ d|_ dS)NT)rrr)r max_lengthrrq) rr full_modelZ get_encoderrrmax_hidden_seq_lengthrroexported_encoderexported_decoder)rrrrrrZmax_cache_lengthrr r!rRs   z"Seq2SeqLMExportableModule.__init__cCsrt|jd}tjjd|jd}ttjj||fdd|iidd}Wd|S1s2wY|S)Nr7Zencoder_seq_lengthmaxr"r Tr'r() rrrFevalr.r1Dimrr)rencoder_input_idsZwrapped_encoderZ seq_len_dimrr r r!_export_encodergs  z)Seq2SeqLMExportableModule._export_encodercCst|j|jjj|jjjdd}tj j d|j d}t tj j ||||fdd|idddd}Wd|S1s?wY|S) N)rrrrr7Zencoder_hidden_seq_lengthrr )rrr#Tr) rrrorqrrrrFrr.r1rrr)rrrr#Zwrapped_decoderZencoder_seq_len_dimrr r r!_export_decoderus.    z)Seq2SeqLMExportableModule._export_decoderNc Cs|dur|ntjdtjd}|dur|n tjdggtjd}|dur$|ntjdgtjd}|dur3|ntj|jjjd|jj ftj d}| ||_ | ||||_|S)N)r r)rr)r.rzr0r/ZzerosrorqrrrZd_modelrrrrr) rrrrr#Zexample_encoder_input_idsZexample_decoder_input_idsr3Zexample_encoder_hidden_statesr r r!r1s  z Seq2SeqLMExportableModule.exportc Cst`|j|}tjdggtjd}dg}t|dD];}|j||tj|gtjd}tj|dddddfdd }| |tj|ggtjd}||j j krZnq|WdS1sgwYdS)Nrr)r rBrC) r.rrrEr/r0rHrrLrNrrrO) rrr:Zencoder_outputrrQrRrUZ next_tokenr r r!rVs $  $z"Seq2SeqLMExportableModule.generate)r r rerrX) rYrZr[rrrr1rVrdr r rr!rQs rrX)rtypingrr.Z+transformers.generation.configuration_utilsrZutils.import_utilsrZ transformersrrrZtransformers.pytorch_utilsr r nnModuler rrr^r_r`rrrrr r r r!s>   IZ @-