o iqO@sddlZddlZddlmZddlmZddlmZmZddl m Z m Z m Z m Z mZmZmZddlZddlmZmZmZddlmZddlmZd Zzdd lmZmZmZd ZWn eycd ZYnwzdd l m!Z!dd l"m#Z#m$Z$Wneyej%j&Z!GdddZ'e'Z#e'Z$Ynwz ddl"m(Z)d a*Wneyddl"m)Z)d a*Ynwej+j,j-j.ej+j/j0j1ej+j,j2j.hZ3e)e3BZ4eGdddZ5d=ddZ6GdddeZ7ddZ8Gddde#Z9GdddeZ:d=ddZ;d dd d!e fd"d#Zd(e?d!e fd)d*Z@d+ejAd,ejAd-e?d.eeBd/eeeBd0fd1eeBd2eCd!ejAfd3d4ZDGd5d6d6ZEGd7d8d8e!ZF  d>d9ej%j&d(ee?d:ee fd;d<ZGdS)?N) defaultdict)deepcopy)astuple dataclass)AnyCallableContextManagerDictListOptionalTuple) is_inplaceis_inplace_view_fn is_view_fn)TorchDispatchMode)tree_mapF)BoundsLinearConstraintmilpT)ActivationWrapper)_CachedTorchDispatchMode_CachingTorchDispatchModec@seZdZddZdS) _NotAvailablecOstd)NzNeed PyTorch >= 2.2) RuntimeErrorselfargskwargsrG/home/ubuntu/.local/lib/python3.10/site-packages/xformers/checkpoint.py__init__-sz_NotAvailable.__init__N)__name__ __module__ __qualname__r rrrrr,s r)SAC_IGNORED_OPS) _ignored_opsc@sVeZdZUeed<eed<eed<eed<eed<eeefed<e ed<e ed<d S) ProfileMetadataname time_taken memory_usedcurr_idx output_ids inplace_info is_view_like is_rand_opN) r!r"r#str__annotations__floatintrr boolrrrrr&Fs  r&cs$gd}dur |fdd}|S)N)z4xformers.efficient_attention_forward_cutlass.defaultz xformers_flash.flash_fwd.defaultzaten.addmm.defaultzaten.mm.defaultcs t|vSN)r/)ctxfuncrr allow_listrr_default_policy\s z,_get_default_policy.._default_policyr)r8_default_allow_listr9rr7r_get_default_policyRs  r;c@seZdZddZdddZdS)VerboseTorchDispatchModecCs g|_dSr4) operators)rrrrr cs z!VerboseTorchDispatchMode.__init__rNcCs&|duri}|j|||i|Sr4)r=appendrr6typesrrrrr__torch_dispatch__fs z+VerboseTorchDispatchMode.__torch_dispatch__rN)r!r"r#r rArrrrr<bsr<cOsBt}|||i|Wd|jS1swY|jS)zZ Returns the list of operators used inside `function` with *args and **kwargs N)r<r=)functionrr verbose_moderrrlist_operatorsms rEcs$eZdZfddZddZZS)CachedTorchDispatchModecs*tr t|||dSt||dSr4)_PT_HAS_NEW_IMPLsuperr )r policy_fnstorageallow_cache_entry_mutation __class__rrr ysz CachedTorchDispatchMode.__init__cCs(|j|r |j|dS||i|SNr)rJpop)rr6rrrrrpop_from_storages z(CachedTorchDispatchMode.pop_from_storage)r!r"r#r rP __classcell__rrrLrrFxs rFc@seZdZdddZdS)NullTorchDispatchModerNcCs|duri}||i|Sr4rr?rrrrAsz(NullTorchDispatchMode.__torch_dispatch__rB)r!r"r#rArrrrrRsrRcCsr|durt}nt|trt|}nt|sJdtt}tr*tt||}nt }t t||d}||fS)aAn activation checkpoint context_fn for selectively deciding what to store and what to recompute. Accepts a custom policy. Args: policy_fn(Union[List[Op], callable]): policy for deciding what to store (instead of recompute). If it's a function, it should be of form (func, *args, **kwargs) -> bool which indicates if func outputs with *args and **kwargs should be stored or not. Additionally, a list[Op] is also supported for easier cases. The op should be in the format `torch.ops.***`, where the `***` names of operators can be obtained with `list_operators`. Nz,policy_fn should be None, list or a callableT) r; isinstancelistcallablertorchis_grad_enabledrrrRrF)rI temp_storage caching_mode cached_moderrrselective_checkpoint_context_fns   r[)preserve_rng_staterIreturncOs,tjjj|g|Rd|tt|d|S)a8Wrapper around torch.utils.checkpoint that accepts a custom policy function for selectively deciding what to store and what to recompute Args: function: describes what to run in the forward pass of the model or part of the model. It should also know how to handle the inputs passed as the tuple. For example, in LSTM, if user passes ``(activation, hidden)``, :attr:`function` should correctly use the first input as ``activation`` and the second input as ``hidden`` preserve_rng_state(bool, optional): Omit stashing and restoring the RNG state during each checkpoint. Default: ``True`` policy_fn(Union[List[Op], callable]): policy for deciding what to store (instead of recompute). If it's a function, it should be of form (func, *args, **kwargs) -> bool which indicates if func outputs with *args and **kwargs should be stored or not. Additionally, a list[Op] is also supported for easier cases. The op should be in the format `torch.ops.***`, where the `***` names of operators can be obtained with `list_operators`. *args: Arguments to pass in to the given ``function``. **kwargs: Keyword arguments to pass into the given ``function``. F) use_reentrantr\ context_fn)rVutils checkpoint functoolspartialr[)rCr\rIrrrrrras rac@sJeZdZd deddfddZdeeeeedfffdd Zdd d ZdS)!ProfileOperatorsTorchDispatchMode num_runsr]NcCsg|_||_dSr4)datarf)rrfrrrr s z*ProfileOperatorsTorchDispatchMode.__init__.c Cst|j}dd}t||}t|s||dfS|}d}t|jD]\}} | j} t| ttt fs2| gn| } || vr<|}nq |dkrC|}||f} ||| fS)NcSst|tjr |SdSr4)rSrVTensoruntyped_storagedata_ptr)errr get_tensor_idszNProfileOperatorsTorchDispatchMode._get_inplace_metadata..get_tensor_idrr) lenrgrr enumerater+rSrTtupledict) rr6outr*rlr+op_id op_parent_ididpast_output_idsr,rrr_get_inplace_metadatas*    z7ProfileOperatorsTorchDispatchMode._get_inplace_metadatarc Cs |duri}||i|}|||\}}}t|pt|} tjj|jv} |jjdkr3| dddk} tj t } t |jD] } ||i|qAtj t | |j} tj tj d}||i|tj d}|jt|| |||||| | |S)N#_scaled_dot_product_flash_attention dropout_pri)rxrrrVTagnondeterministic_seededtagsoverloadpacketr!getcuda synchronizetimerangerfreset_peak_memory_statsmax_memory_allocatedrgr>r&)rr6r@rrrrr*r+r,r-r.trur(mem1mem2rrrrAs>     z4ProfileOperatorsTorchDispatchMode.__torch_dispatch__)rerB)r!r"r#r2r r rxrArrrrrds rdcGs<t}| ||Wdn1swY|j}|S)a Use ProfileOperatorsTorchDispatchMode to get runtime and memory info. Args: function: The function to optimize which will be selectively checkpointed. Usually the forward pass of the model. *args: Arguments to pass in to the given ``function``. Returns: A list of tuples, where each tuples contains the name of the operator, the runtime of the operator, and the memory usage of the operator. N)rdrg)rCr profile_opsrgrrr_analyze_operatorss  r memory_budgetc s\tstd|dks|dkrtd|dt|g|R}dd|D}tdd|D\}}}}}} } tj|tjd } tj|tjd } d dt| D} d dt| D}fd d|D}t |d}t | t d d|DB}t t |}t |D] }||kr|d8}q~d| |<|| }tdd|D}t| | || |||d}t|dS)a Given a function, its arguments, and the maximum amount of memory available, find the subset of operators that can be optimized to reduce runtime while still fitting within the memory budget. Args: function: The function to optimize which will be selectively checkpointed. Usually the forward pass of the model. *args: Arguments to pass in to the given ``function``. memory_budget (float): A float between 0 and 1 which describes what percentage of the total memory to use. Returns: A callable policy which can be passed to xformers.checkpoint() Raises: RuntimeError: If `scipy` is not available. ValueError: If `memory_budget` is not a float between 0 and 1. zlPlease install scipy 1.9.0+ to use `get_optimal_checkpoint_policy`. You can do so using `pip install scipy`.rz5`memory_budget` must be a float between 0 and 1. Got .cSsg|] }|jtvr|qSr)r'OPS_TO_ALWAYS_SKIP.0xrrr Uz1get_optimal_checkpoint_policy..cSsg|]}t|qSr)rrrrrrX)dtypecSg|]\}}|r|qSrrrrurrrrr\cSrrrrrrrr]rcs g|] }|rttj|qSr)rpmapindexrnew_idsrrr`s cSsg|]}|dqS)rrrrrrrircSsg|] }t|tj qSr)rSrVrhrrrrrtr)memoryruntimes max_memory view_like_ops inplace_ops random_opsforce_store_random) optim_output)_scipy_is_availabler ValueErrorrziprVtensorfloat64rornsetsortedrTreversedsumitemall#_optimize_runtime_with_given_memory_OptimalPolicy)rCrrrgops runtimes_memory__ inplace_ops_view_like_ops_ rand_ops_rrrrand_opsrlast_op skip_ops_skip_opsoprrrrrrget_optimal_checkpoint_policy5sL      rrrrrr.rrcCs | }t||d}|g} |D]} t|} d| | <| t| dddq|D].\} } t|} | | krFd| | <d| | <| t| dddq&d| | <| t| dddq&|D]} t|} d| | <t|}| t| ||dqWt|}t|| |tddd}|jst dt |j }|S)aD Given a list of operator names, their corresponding runtimes, and the maximum amount of memory available, find the subset of operators that can be optimized to reduce runtime while still fitting within the memory budget. Uses https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.milp.html Args: memory (torch.Tensor): Tensor containing the memory usage of each operator. runtimes (torch.Tensor): Tensor containing the runtime of each operator. max_memory (float): Maximum amount of memory to use. view_like_ops ([List[int]): Indices of the view-like ops. inplace_ops (List[Tuple[int, int]]): Tuple with the pair of inplace op -> parent of inplace op. This will be used to add the constraint that in-place ops need to either be stored in memory with the previous op, or recomputed with the previous op. random_ops ([List[int]): Indices of the random ops, which will always be recomputed. force_store_random (bool): force random ops to always be stored (instead of recomputed) )Aubrr)rlbrrm)c constraints integralityboundszThe problem is infeasible, and probably due to a change in xformers that makes random ops always be stored. Try passing a larger memory_budget. This will be fixed once https://github.com/pytorch/pytorch/issues/121212 is solved) rrV zeros_liker>r2 ones_likerrsuccessr from_numpyr)rrrrrrrrmemory_constraintrrurr op_parentvalrresrrrrrs<       rc@s*eZdZdejfddZdefddZdS)rrcCsd|_||_dSrN)countertolistr)rrrrrr sz_OptimalPolicy.__init__r]cOs.|tvrdS|j}|jd7_|j|dkS)NFr)rrr)rr5r6rrcountrrr__call__s z_OptimalPolicy.__call__N)r!r"r#rVrhr r3rrrrrrsrcs>eZdZd fdd ZejjddZddZdd Z Z S) SelectiveCheckpointWrapperNcsVt||du|duAstd||_||_zdtjj_WdSt y*YdSw)Nz1Need to specify either policy_fn or memory_budgetT) rHr rrrIrV_dynamoconfig:_experimental_support_context_fn_in_torch_utils_checkpointAttributeError)rmodrrIrLrrr s  z#SelectiveCheckpointWrapper.__init__cOstsgStjt|jg|Ri|d|ji}Wdn1s'wYtjrLtj rLtj dkrL|g}tjj |dd|d}|S)Nrrr)src) rVrWrandomfork_rngr_checkpoint_wrapped_moduler distributed is_availableis_initializedget_world_sizebroadcast_object_list)rrrrIobjectsrrr_get_policy_fns,  z)SelectiveCheckpointWrapper._get_policy_fncOs"|jdur|j|i||_|jSr4)rIrrrrr get_policy_fns z(SelectiveCheckpointWrapper.get_policy_fncOs0|j|i|}t|jg|Ri|d|iS)NrI)rrar)rrrrIrrrforwardsz"SelectiveCheckpointWrapper.forwardNN) r!r"r#r rVcompilerdisablerrrrQrrrLrrs  rmodulerIcCs t|||S)a Wrap a module with selective activation checkpointing. It behaves similarly to PyTorch's checkpoint_wrapper, but gives the possibility to the user to either specify a handcrafted policy_fn, or to let an optimization algorithm to select the policy given a user-specified memory_budget. The user should either specify the memory_budget argument or the policy_fn. The memory_budget is a float value between 0 (recompute everything in the backward) or 1 (store everything for backward). Using a value of 0 should be similar to PyTorch's activation checkpoint, while 1 should be similar to the behavior of not using any activation checkpointing. )r)rrrIrrrselective_checkpoint_wrappers rr4r)Hrbr collectionsrcopyr dataclassesrrtypingrrrr r r r rV,torch.testing._internal.composite_compliancer rrtorch.utils._python_dispatchrtorch.utils._pytreerrscipy.optimizerrr ImportError;torch.distributed.algorithms._checkpoint.checkpoint_wrapperrtorch.utils.checkpointrrnnModulerr$r%rGraten lift_freshdefaultprofiler_record_function_exit_RecordFunctionclone_additional_ignored_opsrr&r;r<rErFrRr[rardrr1rrhr2r3rrrrrrrrs  $                 "NM K6