o ٷil@s~ddlZddlZddlZddlZddlZddlZddlZddlmZddl Z ddl Z ddl Z ddl m Z ddlmZddlmZddlmZddlmZddlmZdd lmZdd lmZmZmZmZee Z!gd Z"e j#d e j$d e j%diZ&GdddeZ'GdddeZ(GdddeZ)GdddeZ*GdddeZ+e*ddfe+ddfe)ddfdZ,GdddZ-Gd d!d!Z.dS)"N)Path) Precision)float_to_float16_max_diff) FusionOptions)IOBindingHelper) OnnxModel)optimize_model)torch_onnx_export) GPT2ConfigGPT2LMHeadModel GPT2Model TFGPT2Model) distilgpt2gpt2z gpt2-mediumz gpt2-largezgpt2-xlMb@?g?g@c,eZdZdZfddZfddZZS)GPT2ModelNoPastState2Here we wrap a class to disable past state output.ct|dSNsuper__init__selfconfig __class__d/home/ubuntu/.local/lib/python3.10/site-packages/onnxruntime/transformers/models/gpt2/gpt2_helper.pyr*zGPT2ModelNoPastState.__init__cstj|dddS)NF) use_cache return_dict)rforwardr input_idsrrrr#-szGPT2ModelNoPastState.forward__name__ __module__ __qualname____doc__rr# __classcell__rrrrr' rcr)TFGPT2ModelNoPastStatercsd|_t|dS)NF)r!rrrrrrr4szTFGPT2ModelNoPastState.__init__cstj|ddS)NF)r!)rcallr$rrrr#8r zTFGPT2ModelNoPastState.forwardr&rrrrr-1s r-cs8eZdZdZfddZeddZfddZZS) MyGPT2ModelzMHere we wrap a class for Onnx model conversion for GPT2Model with past state.crrrrrrrr?r zMyGPT2Model.__init__cCst|ddttfrNt|d|krt|dddksJg}t|D] }|tj|d|dd|d|ddfddq%|dt|fS|S)Nr)dim) isinstancetuplelistlenrangeappendtorchcat unsqueeze)result num_layerpresentirrr post_processBs( *zMyGPT2Model.post_processc&tj||||dd}t||jjSNF) position_idsattention_maskpast_key_valuesr"rr#r/r@rn_layerrr%rCrDpastr<rrrr#UszMyGPT2Model.forward) r'r(r)r*r staticmethodr@r#r+rrrrr/<s   r/cr)MyGPT2LMHeadModelzSHere we wrap a class for Onnx model conversion for GPT2LMHeadModel with past state.crrrrrrrrcr zMyGPT2LMHeadModel.__init__crArBrFrHrrrr#fszMyGPT2LMHeadModel.forwardr&rrrrrK`r,rKcr)MyGPT2LMHeadModel_NoPaddingaHere we wrap a class for Onnx model conversion for GPT2LMHeadModel with past state and no padding. When you always use batch_size=1 in inference, there is no padding in inputs. In such case, position_ids and attention_mask need no be in inputs. crrrrrrrrxr z$MyGPT2LMHeadModel_NoPadding.__init__cs"tj||dd}t||jjS)NF)rEr"rF)rr%rIr<rrrr#{sz#MyGPT2LMHeadModel_NoPadding.forwardr&rrrrrLrs rLlogitsTF last_state)r GPT2LMHeadModel_NoPaddingr c@s8eZdZddZdefddZdefddZdd Zd S) Gpt2InputscCs||_||_||_||_dSr)r%rCrDrI)rr%rCrDrIrrrrs zGpt2Inputs.__init__returncCs0dd|j|j|jfD}|jr||j|S)NcSsg|]}|dur|qSrr.0vrrr z&Gpt2Inputs.to_list..)r%rCrDrIextend)r input_listrrrto_lists zGpt2Inputs.to_listcCs"tdd|j|j|j|jfDS)Ncss|] }|dur|VqdSrrrRrrr sz&Gpt2Inputs.to_tuple..)r4r%rCrDrI)rrrrto_tuples"zGpt2Inputs.to_tuplecCsTd}|jdur|jjtjkr|jjtjdn|j}dd|jD}t|j|j ||S)NdtypecSsg|] }|jtjdqS)r\)tor9float32rSprrrrUsz&Gpt2Inputs.to_fp32..) rDr]r9float16r^r_rIrPr%rC)rrDrIrrrto_fp32s zGpt2Inputs.to_fp32N) r'r(r)rr5rYr4r[rcrrrrrPs  rPc"@seZdZdZedddejejejdfdededededed ed ed ejd e d e de dej dej dej de de f ddZ e dXdededede dedeeeeff ddZeddZedYddZedYdd ZedZd"d#Zed[d%d&Zeddddejejejfd'ed(e d)e d e de dej dej dej fd*d+Ze   ,d\d-d.Zegd/fd0ed1eefd2d3Zed]d4e d5efd6d7Zed]d4e d5efd8d9Zed:d;Zed^deeejfd?eeeefd5ed@e dAe f dBdCZ edDdEZ!edFdGZ"eddHdHdIdJddddejejejd,ddfdKdLZ#eddMddddejejejdNdJdOf dPdQZ$ed`dRdSZ%edddgdTfdefdUdVZ&dWS)a Gpt2HelperzEA helper class for Gpt2 model conversion, inference and verification.FT batch_sizepast_sequence_lengthsequence_lengthnum_attention_heads hidden_sizer= vocab_sizedevicerbhas_position_idshas_attention_maskinput_ids_dtypeposition_ids_dtypeattention_mask_dtypeleft_side_paddingrQcs"|rtjntjd|||t||gfddt|D}tjd|d||f| d}d}| rh||}tj||g| d}|dkrht|D]}td|d}|r]d||d|f<qHd||||df<qHd}| r| d d}| |dkd|dd|df | }t ||||S) zCreate random inputs for GPT2 model. Returns torch tensors of input_ids, position_ids, attention_mask and a list of past state tensors. r1cs$g|]}tjdddqS)r]rkg@?)r9rand)rS_rk float_type past_shaperrrUs$z/Gpt2Helper.get_dummy_inputs..rr0)lowhighsizer]rkNrr) r9rbr_intr7randintonesrandomlongcumsum masked_fill_r^rP)rerfrgrhrir=rjrkrbrlrmrnrorprqrIr%rDtotal_sequence_lengthr?padding_lengthrCrrvrget_dummy_inputssF  zGpt2Helper.get_dummy_inputsr r model_classcCs~|j}|j}|j}|j}t|d} ||| dkr|n|g} d||||t||g} | | i} t|D] } | | dt| <q2| S)zAReturns a dictionary with output name as key, and shape as value.r0rMr1present_)rhrinum_hidden_layersrj MODEL_CLASSESr}r7str)rerfrgrrrhrir=rj output_namelast_state_shapepresent_state_shape output_shapesr?rrrget_output_shapess&    zGpt2Helper.get_output_shapescCsZ|D](}||vs J||}t|||kr*tjt|||j|jd||<qdS)Nrr)numpyprodnelementr9emptyr]rk)output_buffersrkeybufferrrrauto_increase_buffer_sizes   z$Gpt2Helper.auto_increase_buffer_sizecCsD|rtjntj}i}|D]\}}tjt|||d||<q|S)zpReturns a dictionary of output name as key, and 1D tensor as value. The tensor has enough space for given shape.rr)r9rbr_itemsrrr)rrk is_float16 data_typernameshaperrrget_output_bufferss zGpt2Helper.get_output_bufferscCsH|d}t||d}|rt|t|dSt|S)zGReturns the maximum difference between PyTorch and OnnxRuntime outputs.rư>)cpurabsamax) torch_outputs ort_outputsrelativeexpected_outputsdiffrrr diff_outputs%s  zGpt2Helper.diff_outputsMbP?c Kstj|d|d||d}td||}t|d}t|D])}tj|d||d|||d}td|d|d||oM|}q%|s`t||} t d| d |S) zReturns True if torch and ORT outputs are close for given thresholds, and False otherwise. Note: need kwargs since Gpt2BeamSearchHelper.compare_outputs has an extra parameter model_class r)rtolatolz9PyTorch and OnnxRuntime output 0 (last_state) are close: r0zPyTorch and OnnxRuntime layer z state (present_z ) are close:z@PyTorch and OnnxRuntime results are not all close: max_abs_diff=.5f) rallcloserloggerdebugr6r7rdrinfo) rrrrkwargsis_close is_all_close num_layerslayer max_abs_diffrrrcompare_outputs/s""     zGpt2Helper.compare_outputsrc Csd}d}g}g}tt|D]}||}|dkr|dn|d|d} tj|| |dd} |tt| ||oA| }t|  rRt d|dt |  rbt d|dt| rrt d |dt | rt d |dt || } t| | j} |d | | d d | d || d dt| | d |dkrttj|dd|j} ttj| dd| j}t| |}q|t|}|t||||fS)aCompare outputs from PyTorch and OnnxRuntime Args: torch_outputs (Tuple[Torch.Tensor]): PyTorch model output ort_outputs (List[numpy.ndarray]): OnnxRuntime output atol (float, optional): Absolute tollerance. Defaults to 1e-06. Returns: is_all_close(bool): whether all elements are close. max_abs_diff(float): maximum absolute difference. messages(str): a list of debug message for each output TFrr0)rrzPyTorch output z has nanz has infz ORT output zdiff=z.9fz index=z ort=z torch=N)axis)r7r6rrrr8rrisnananyrrisinffabs unravel_indexargmaxrfloat array_equalindexmax)rrrris_top1_matched max_diffsmessagesr? ort_output torch_outputrridx ort_max_indextorch_max_indexmax_diff_output_indexrrrcompare_outputs_v2JsF(0 zGpt2Helper.compare_outputs_v2onnx_model_pathverboseuse_external_data_formatc Cs|j} | j} tjddd| j| j| | j|d||||| d} | } t || }Wdn1s3wYddt | D}ddt | D}|dj d | jks`|dj d | jks`J|dj d | jkrld nd g|}d d dd|dd ddi}|D] }d dd||<q|D] }d dd||<qd g}|rd dd|d<| d|rd dd|d<| d| |t|d krt|d| ksJtd| jj d| jdj d|dj d|ddj t|jjddd|rAt;}tj|d}t|jjdddt|t| |d|||ddd|d tj|dd}tj ||dddWddS1s:wYdSt|t| |d|||ddd|d dS) z1Export GPT-2 model with past state to ONNX model.r0F)rerfrgrhrir=rjrkrbrlrmrnrorpNcSg|]}d|qS)past_rrSr?rrrrUz*Gpt2Helper.export_onnx..cSr)rrrrrrrUrrr1rMrNr%reseq_len)rr0 past_seq_len)r0 total_seq_lenrCrDzShapes: input_ids=z past=z output=z present=T)parentsexist_okz gpt2.onnx ) argsf export_params input_names output_names dynamic_axes opset_versiondo_constant_foldingrr)load_external_data)save_as_external_dataall_tensors_to_one_file)!rrGrdrrhrirjrYr9no_gradr7rr8rWr6rrr%rIrparentmkdirtempfileTemporaryDirectoryospathjoinr r4onnx load_modelrsave)modelrkrrrrlrmrnrorprr= dummy_inputsrXoutputs past_names present_namesrrrr tmp_dir_nametemp_onnx_model_pathrrr export_onnxs  ,"     6 $ zGpt2Helper.export_onnxrc Ksftd} t|d||d| dd} |r+|rt| nd|vr!d|d<| jd ddi|| ||| S) zHOptimize ONNX model with an option to convert it to use mixed precision.rrF) model_type num_headsri opt_leveloptimization_optionsuse_gpu keep_io_typesuse_symbolic_shape_inferTNr)rrrdauto_mixed_precisionconvert_float_to_float16save_model_to_file) roptimized_model_pathrrhrirrstagerrmrrr optimize_onnxs$    zGpt2Helper.optimize_onnx)AddLayerNormalizationSkipLayerNormalizationFastGeluEmbedLayerNormalization onnx_model op_block_listcCsPdd|D}t|}||}td|d||jdj}d}|}||vs1J||}d} |j dkro|} td |jd} |j D] } | | } | durYnqLt | } t d |jd | | d k}n td |j d|jg} g}|s| dur|g} | jg}| |||d}td||jdddi||S)a?Convert GPT-2 model to mixed precision. It detects whether original model has fp16 weights, and set parameters for float16 conversion automatically. Args: onnx_model (OnnxModel): optimized ONNX model op_block_list (List[str], optional): operators to compute in fp32. Defaults to ["Add", "LayerNormalization", "SkipLayerNormalization", "FastGelu", "EmbedLayerNormalization"] Returns: parameters(dict): a dictionary of parameters used in float16 conversion cSsh|]}|jqSr)op_type)rSnoderrr 0sz2Gpt2Helper.auto_mixed_precision..z fp32 op: z fp16 op: rFNMatMulz#Found last MatMul node for logits: z3max diff of converting weights in last MatMul node : rz-Failed to find MatMul node for logits. Found z of node )rr node_block_listforce_fp16_initializersz!auto_mixed_precision parameters: rTr)nodesset differencerrgraphoutputroutput_name_to_noderinputget_initializerrrwarningr)r r  op_full_set fp32_op_set fp16_op_setlogits_output_nameis_weight_fp16_precisionrrlast_matmul_node initializerrmax_diffrr parametersrrrrsH       zGpt2Helper.auto_mixed_precisioninputs total_runsc Cstd|}t ||}Wdn1swY|dkr)|Sg}t t|D]}t}||}|t|q4Wdn1sRwYt |dt |}td t |d||fS)zfRun inference of PyTorch model, and returns average latency in ms when total_runs > 0 besides outputs.zstart pytorch_inferenceNrzPyTorch inference time = {} ms.2f) rrrcrYr9rr7timer8sumr6format) rr'r(rXrlatencyrustartaverage_latencyrrrpytorch_inferencebs$      zGpt2Helper.pytorch_inferencec Cs"tddt|ji}|jdur.t|jD]\}}t||d|<q|jdur?t|j|d<|j durPt|j |d<| d|}|dkr\|Sg}t |D]}t } | d|}| t | qbt|dt|} td t| d || fS) zcRun inference of ONNX model, and returns average latency in ms when total_runs > 0 besides outputs.zstart onnxruntime_inferencer%NrrDrCrr)z"OnnxRuntime Inference time = {} msr*)rrrascontiguousarrayr%rrI enumeraterDrCrunr7r+r8r,r6r-) ort_sessionr'r( ort_inputsr?past_irr.rur/r0rrronnxruntime_inference|s(       z Gpt2Helper.onnxruntime_inferencec Cst|||||||S)z)Returnas IO binding object for a session.)rprepare_io_binding)r5r%rCrDrIrrrrrr9s zGpt2Helper.prepare_io_bindingcCst||||S)z3Copy results to cpu. Returns a list of numpy array.)r"get_outputs_from_io_binding_buffer)r5rr return_numpyrrrr:sz-Gpt2Helper.get_outputs_from_io_binding_bufferrrr;include_copy_output_latencyc Cstdt||j|j|j|j||}||t ||||}|dkr'|Sg} t |D]} t } |||rBt ||||} | t | q-t | dt| } td| || fS)zUInference with IO binding. Returns outputs, and optional latency when total_runs > 0.z*start onnxruntime_inference_with_binded_iorr)z4OnnxRuntime with IO binding inference time = %.2f ms)rrrdr9r%rCrDrIrun_with_iobindingr:r7r+r8r,r6) r5r'rrr(r;r< io_bindingrr.rur/r0rrr$onnxruntime_inference_with_binded_ios8     z/Gpt2Helper.onnxruntime_inference_with_binded_iocCstd|dd}t||Wdn1swYtd|dtd|dd}t||Wdn1sBwYtd|ddS)N ort_outputs_.picklewbz$ORT output are saved to ort_outputs_torch_outputs_z(Torch output are saved to torch_outputs_openpickledumprr)r?rrrrrr save_outputsszGpt2Helper.save_outputscCsTtd|dd}t||Wdn1swYtd|ddS)N dummy_inputs_rArBz!inputs are saved to dummy_inputs_rD)r?rrrrrrr save_inputsszGpt2Helper.save_inputsri'r0c,s(|j}td|dd|d|d| d|d d}d }d }d }|r5t||||| }t|||}d }d }gd g|}|}t|D]}t|}t d |}|d kr\d nt d |}t d |} t d| d|dtj | |||j |j |j|j||| | | | |dd}!t||!}"|rt||!}#nt| |||| }$t||!||$}#tj|"|#|d\}%}&}'}(})t|&sÈ|&|%r|d 7}|)r|d 7}||d 7<|r |%s td|d| d|d|d|& t|(D]\}}*td|d||jd|*q|r*t|&s|&d|kr*t||!t||#|"qHr8fdddD}+ndddD}+|d ||+d!<fd"d#|D|+d$<|d ||+d%<|td ||+d&<td'|d(|d)|td*||d+|krtd,t|d|d-d.|+S)/zKGenerate random inputs and compare the results of PyTorch and Onnx Runtime.zRunning parity test (atol=z , test_cases=z, runs=z, use_io_binding=z, model_class=z , is_float16=z) ...r1Nrr0z#Running parity test for batch_size=z past_sequence_length=z...T)rnrorprq)rz test_case=z batch_size=z sequence_length=z MaxDiff= z: Name=z, dcs$i|]}d|t|dqS)max_diff_percentile_r)r percentiler`)max_abs_diff_listrr osz*Gpt2Helper.test_parity..)2Z_ccSsi|]}d|dqS)rOnanrr`rrrrRsrVrstop1_match_ratecsg|]}|dqS)rsr)rSx)test_cases_per_runrrrUvrVz*Gpt2Helper.test_parity..top1_match_rate_per_rundiff_pass_ratenan_ratezParity Test Cases=z ; Passed=z; Nan=z; Top1_Matched=gffffff?zParity is good: passed rate=z.0f%)rrrrdrrr7r}rr~rrrhrirGrjr1r8r?rrrr8r3 get_outputsrrJrHr6),r5rrkrrrrZr(use_io_bindingrrlrmrnrorprrenable_pickle_outputrmax_batch_sizemax_past_seq_len max_seq_lenrmax_output_shapespassed_test_casestop1_matched_casestop1_matched_cases_per_runtotal_test_casesr?run_idrgrfrerrrrrrrrrmessager<r)rQrZr test_paritys(            (   " zGpt2Helper.test_parityrNrK cCs|j}d}|rt| || ||}t|||}tj| || |j|j|j|j||||| | | d}|r;t |||\}}|St |||||\}}|S)zCGenerate random inputs and measure average latency of Onnx Runtime.N)rnrorp) rrdrrrrhrirGrjr8r?)r5rrkrr(r`rrlrmrnrorprergrfrrrrrur.rrrtest_performances<  zGpt2Helper.test_performancec Cs:tjddd|j|j|j|j|d||d }tj ||S)zJIT trace for TorchScript.r0F) rerfrgrhrir=rjrkrbrlrm) rdrrhrirGrjrYr9jittrace)rrrkrlrmrXrrr torchscripts   zGpt2Helper.torchscriptrawfp32fp16int8c Cs|}tj|rt|jd}n|dd|dkr!|d|7}|r'|d7}|rdddd d }d D]P}tj||||} tj| r||vrwzt | t d | Wq2t yv} zt d | d | j WYd} ~ q2d} ~ wwt d|d| q2tjtj|||dtjtj||d|dtjtj||d|dtjtj||d |dd Stj||dtj||dtj||dtj||dd S)z=Build a path name for given model based on given attributes.r|/r ru_past_fp32_fp16_int8rrzRemoved the existed directory: zFailed to remove the directory rNzDirectory for z existed: z.onnxz _fp32.onnxz _fp16.onnxz _int8.onnx)rrisdirrpartssplitrexistsshutilrmtreerrOSErrorstrerror) output_dirmodel_name_or_pathrhas_past new_folderremove_existing model_namesuffixrnew_direrrrget_onnx_pathssT    $zGpt2Helper.get_onnx_pathsN)r )F)rr)r)FFr)r)T)rTF)TT)'r'r(r)r*rJr9int32r}rkboolr]rPrr rdictr5rrrrrrrrrrr1r8r9r:Tensorr?rHrJrlrnrqrrrrrrds`       @ "     5  w # E    2    6 rd)/loggingrrFrrrr+pathlibrrrr9benchmark_helperrrbrfusion_optionsrio_binding_helperrr r optimizerrtorch_onnx_export_helperr transformersr r r r getLoggerr'rPRETRAINED_GPT2_MODELSFLOAT32FLOAT16INT8DEFAULT_TOLERANCErr-r/rKrLrrPrdrrrrsH           $