x-jD$ddlZddlZddlmZddlmZddlZddlmZddl m Z e e ej ddZe e ej d dZejeejd Zd Z d'dZd(dZdZdZdZdZdZdZGddZGddeZGddZGdd Z d!Z!d"Z"Gd#d$ej#j$Z%Gd%d&ej&j'Z(dS))N)deque)Enum) log_helper) CUDAGraph+PADDLE_DEBUG_ENABLE_CUDAGRAPH_LAYER_LOGGING01PADDLE_DEBUG_CUDAGRAPHEDLAYER_FALLBACK_TO_DEFAULTz[%(levelname)s] %(message)s)fmtcLtsdSt|dSN)enable_debug_printloggerinfoxs e/var/www/html/banglarbhumi/venv/lib/python3.11/site-packages/paddle/device/cuda/cuda_graphed_layer.py debug_printr&s"  KKNNNNNUnnamedTFcg}|r|||d}|t|ddSt|tjr|r:|d|j|d|j|r7|dt||r!|d||td |dSdS)Ncpt|ddzS)Nfloat32i)floatastypesum)ts rzprint_tensor..8s+ 3 3d :??AABBrz is Nonezshape = zplace = zptr = zhash = z | ) appendr isinstancepaddleTensorshapeplacehexdata_ptrjoin)rname print_meta print_ptr print_hashhashoutputs r print_tensorr.,s0F  d |BByt%%%&&&&& Av} % %(  0 MM.QW.. / / / MM.QW.. / / /  8 MM63qzz||#4#466 7 7 7  / MM-DDGG-- . . .EJJv&&'''''((rprinterctsdSt|ddtt|dS)Nd-)rrcenterrecursive_applyr.)rbanners rr/r/GsA  c3''(((L!$$$$$rc*t|trfd|DSt|trtfd|DSt|tr fd|DS|S)Nc0g|]}t|Sr4.0itemfunctions r z#recursive_apply..Rs#FFFD$//FFFrc38K|]}t|VdSr r9r:s r z"recursive_apply..Ts-KK_Xt44KKKKKKrc8i|]\}}|t|Sr8r9)r;keyvaluer=s r z#recursive_apply..Vs9   U 511   r)r listtupledictitems)r= input_vars` rr4r4Ps)T"" #FFFFIFFFF Iu % %#KKKKKKKKKK It $ $#    'oo//    x """rc~t|tjr"|}|j|_|S|Sr )r r!r"detach stop_gradient)tensordetached_tensors r detach_tensorrO^s9&&-(( --//(.(<% Mrc6gfd}t||S)Ncjt|tjr|dSdSr )r r!r"r)argrets rrz!recursive_flatten..appendls6 c6= ) )  JJsOOOOO  rr9)targetrrSs @rrecursive_flattenrUis8 CFF### Jrcgt|tt|Sr )rUrFvalues)argskwargss rrecursive_flatten_args_kwargsrZ{s<  4   511 2 2 rc,tt|Sr )r4rOrs rrrs?=!44rcZt|tjr|jrdS|jSdS)zWReturns the gradient of a Paddle Tensor if it's a tensor; otherwise, returns the input.N)r r!r"rLgradrs rget_grad_tensorr^s1!V]## ? 46M 4rc2eZdZdZdZdZdZdZdZdS)CUDAGraphWithStaticInputOutputc||_t|_d|_d|_d|_d|_d|_d|_dS)NF) num_warmup_stepsrgraph has_recordedhas_preserved_inputs args_static kwargs_static inputs_staticoutputs_staticselfrbs r__init__z'CUDAGraphWithStaticInputOutput.__init__sM 0[[ !$)!!""rc|js6d|_||_||_t|j|j|_dSt||}t |j|D]\}}||ddS)a For the CUDA Graph, it is crucial that the buffer remains address-stable, meaning that the buffer addresses for any inputs to the CUDA Graph should not change. One solution to achieve this is to preserve all input tensors. This function attempts to recursively flatten the input arguments and keyword arguments to identify all tensors passed to the layer (though it may still miss some due to other implicit ways inputs can be passed to a layer). It then preserves references to these input tensors as `self.inputs_static` so that the buffer pointers can be reused later. When this method is called subsequently, it copies the values back to the preserved input tensors to ensure the buffers are reused. TN)rerfrgrZrhzipcopy_)rkrXrYinputsx_staticrs rpreserve_or_copyz/CUDAGraphWithStaticInputOutput.preserve_or_copys( ((,D %#D !'D !> $"4""D   34@@F"4#5v>> ( ( !q$'''' ( (rc,||||j||ji|j|_|jtd|jd|_ |jS)NzF[CUDAGraph] Record-Replay Start (Graph is replayed for the first time)T) rrrc capture_beginrfrgri capture_endrreplayrd)rkfrXrYs rrecordz%CUDAGraphWithStaticInputOutput.records dF+++   """a!1HT5GHH     T     ""rc||_dSr )ri)rkris rset_output_staticz0CUDAGraphWithStaticInputOutput.set_output_statics,rc|jstd|||td|j|jS)NzGraph should be recorded firstz[CUDAGraph] Replay Start)rd RuntimeErrorrrrrcrvri)rkrXrYs rrvz%CUDAGraphWithStaticInputOutput.replays_  A?@@ @ dF+++./// ""rchtjd||j|dS)Nzsave graph to )loggingrrcprint_to_dot_files)rkr(s rsavez#CUDAGraphWithStaticInputOutput.saves7 ,d,,--- %%d+++++rN) __name__ __module__ __qualname__rlrrrxrzrvrr8rrr`r`sn # # #(((4 # # #---###,,,,,rr`ceZdZdZdZdZdZdS)CUDAGraphLayerStatusz3Enum to represent the status of a CUDA Graph Layer.rN)rrr__doc__WARMUPRECORD CUDAGRAPHr8rrrrs#== F FIIIrrc&eZdZdZdZdZdZdS)CUDAGraphForwardBackwardcxt||_t||_tj|_dSr )r` forward_graphbackward_graphrrstatusrjs rrlz!CUDAGraphForwardBackward.__init__s2; t A % %+D,ABB B#++-- -rc:|j|dSr )rr)rkgs r reuse_graphzCUDAGraphContext.reuse_graphs """""rc:|j|dSr )rr)rkrXs r push_datazCUDAGraphContext.push_datas t$$$$$rc4|jSr )rrrs rpop_datazCUDAGraphContext.pop_datas&&(((rcl|xjdz c_|j|jkrtj|_dSdS)Nr)rrbrrrrs r warmup_stepzCUDAGraphContext.warmup_step s7 a :. . ..8DKKK / .rc,|jtjkSr )rrrrs ris_warmup_stepzCUDAGraphContext.is_warmup_step%rrc,|jtjkSr rrs rrz#CUDAGraphContext.is_cuda_graph_step(rrN) rrrrrlrrrrrrrr8rrrrs ###*... ###%%%)))999 :::=====rrc|d\}}t|tjr ||d}}nYt|ttfr=t ||D],\}}t|tjr |js||}}n-t|tjrt|tjsJ||fS)N)NNr)r r!r"rErFrnrL)ysdysydyvdvs rselect_y_with_gradr,s EAr"fm$$CF2 Bu & &S\\  EAr!V]++ Q_ 22 a ' 'IJr6=,I,III I b5Lrc0|\}}g}t||D]p\}}|jsO|j-|t j|j@||j[|dqt|Sr )rnrLr]rr!zerosr#rF)rp grad_inputsdetached_grad_inputs args_gradr detached_xs r get_args_gradr<s(.%K%I[*>?? # # : #&  j.>!?!?@@@@  1111   T " " " "   rc>eZdZdZedZedZdS)_CUDAGraphedLayerz A custom layer that integrates CUDA Graph recording and execution into PaddlePaddle's autograd system. It handles forward and backward operations differently based on the CUDA graph layer status. c|\}}t|}t|}t||}||f}t|dstrkt dt j5j|i|}dddn #1swxYwY tj d||fnډ } | r^t dt| fd} | jj| g|Ri|} tj| ||fnTt dt| | jj|i|} tj| d|ft d|t|dt|S) z Handles the forward pass of the layer. It operates differently based on the context's status: warmup, recording, or CUDA graph step. z Forward inputz"[CUDAGraph] Forward Step (Default)Nz%[CUDAGraph] Forward Step (Record) id cxtj5j|i|cdddS#1swxYwYdSr )r! enable_gradr)rXrYcontexts rforwardz*_CUDAGraphedLayer.forward..forwardss+-->>,w}d=f==>>>>>>>>>>>>>>>>>>s /33z$[CUDAGraph] Forward Step (Graph) id z[CUDAGraph] Forward Step EndzForward output)rKrZr/r*debug_cudagraphedlayer_fallback_to_defaultrr!rrrrrrridrrxrrvrsave_for_backward) ctxr arg_tuplerrXrYrrprrcrs ` rrz_CUDAGraphedLayer.forwardTsm ! fd||>>>>/E'.wHHHHHH!!)0%CN2e99NNOOO.E'.???!!)3UD!D 2333 g&&&#$$$ayysB  B$'B$c|\}|\}}}}t||\}}t||fd|tjkrIt d||t|} | n|tj krt dt|d} |j | ||t|} |j | |||no|tjkrPt dt||j ||} ||nt'dt dt| d| S) z Handles the backward pass of the layer. Similar to forward, it handles backward based on the context's status: warmup, record, or CUDAGraph. zBackward inputz#[CUDAGraph] Backward Step (Default)z&[CUDAGraph] Backward Step (Record) id c0||dSr )backward)rrs rrz,_CUDAGraphedLayer.backward..backwards 2rz%[CUDAGraph] Backward Step (Graph) id zUnknown cuda graph statusz[CUDAGraph] Backward Step EndzBackward output) saved_tensorrrr/rrrrrrrrrrxrzrurrrvr|) rrrrrcrprrrrrs rrz_CUDAGraphedLayer.backwards %%'' &-&6&6&8&8#"2s++2B)*** )0 0 0 = > > > JJrNNN%f--I    ! ! ! ! +2 2 2 LELL M M M     ' '!R 8 8 8&f--I  2 29 = = =          & & & & +5 5 5 K5 KK L L L,33Ar::I    & & & &:;; ;3444 ,---rN)rrrr staticmethodrrr8rrrrNsU 33\3j00\000rrcPeZdZdZddejjffd ZdZdZ dZ xZ S) CUDAGraphedLayera CUDAGraphedLayer: A PaddlePaddle Layer to convert an eager mode model to utilize CUDA Graphs. CUDA Graphs provide a way to capture kernel-level operations of a model and play them back efficiently, allowing for potential speedups in repetitive computations, such as those during training iterations. This layer is a wrapper that enables the usage of CUDA Graphs with PaddlePaddle models. Overview: - The layer encapsulates another layer (the model to be converted). - During the first few (num_warmup_steps) iterations, the layer operates in eager mode without any CUDA Graphs. - After the warmup steps, the layer captures the forward and backward computations and replays them using CUDA Graphs in subsequent iterations. Usage: model = Model() graphed_model = CUDAGraphedLayer(model) Parameters: - layer (paddle.nn.Layer): The PaddlePaddle model/layer to be converted. - num_warmup_steps (int): The number of iterations before the CUDA Graph capture begins. Default is 3. Notes: - Restrictions: * CPU-GPU Synchronization: Operations that synchronize the CPU with the GPU, like device to host transfers, are not allowed. * CPU Work: Any operations on the CPU within the captured graph are not recorded. * Memory Address (Pointer) Consistency: Replays consistently read from and write to identical virtual memory addresses. * Dynamic Operations: - Control Flow: Dynamic control flows, especially those based on CPU data like if/else statements, are prohibited. - Tensor Shapes: Dynamic tensor shapes are not supported. - Allowed Operations: * CUDA RNG Operations: CUDA-based Random Number Generation operations are allowed. rrctt|||_|dt |j|dS)NzGraphed )superrlrr add_sublayertyper)rkrrb __class__s rrlzCUDAGraphedLayer.__init__sY '/?@@  ;T%[[%9;;UCCCCCrcVt||}tj|j||fg|RSr )rZrapplyr)rkrXrYrs rrzCUDAGraphedLayer.forwards>3D&AA  & L4. +6    rc4|jSr )rrrs rrzCUDAGraphedLayer.is_warmup_steps|**,,,rc4|jSr )rrrs rrz#CUDAGraphedLayer.is_cuda_graph_steps|..000r)r) rrrrr!nnLayerrlrrr __classcell__)rs@rrrs##JDDfioDDDDDD    ---1111111rr)rTFTN)r/))r~os collectionsrenumrr! paddle.basergraphsrboolintgetenvrr get_loggerrINFOrrr.r/r4rOrUrZrKr^r`rrrrrautogradPyLayerrrrrr8rrrs  """"""TC  ?EEFF.2TC  EsKKLL..*   gl =      ((((6%%%% # # #$ 5 4G,G,G,G,G,G,G,G,V4 = = = = = = = = 6=6=6=6=6=6=6=6=r    $mmmmm/mmm`6161616161vy6161616161r