a h@sddlZddlZddlZddlZddlZddlmZmZddlm Z m Z ddl m Z ddl mZmZmZmZmZddlZddlmZddlmZddlmZmZddlmZdd lmZd d l m!Z!d d l"m#Z#m$Z$m%Z%d d l&m'Z'gdZ(e)e*Z+Gddde Z,e,j-Z-e,j.Z.e,j/Z/e,j0Z0e,j1Z1e,j2Z2e,j3Z3e,j4Z4e,j5Z5e,j6Z6e-Z7e.Z8e/Z9e6Z:e;dZe?e@ee=fee?eAdddZBGdddeZCdNe@ejDeeAe@ejEdddZFdOe@ejDeeAe>e?e@ejEfdddZGe@ejEddd ZHGd!d"d"eCZIGd#d$d$eIZJGd%d&d&eIZKGd'd(d(eIZLdPe@ee=e?e@e=d*d+d,ZMe@ee=e@e=d-d.d/ZNe>e?e@e=fee?ge?fe?e>e?e@e=fd0d1d2ZOe>e?e@ee=fe?e?e?e>e?e?fd3d4d5ZPGd6d7d7eCZQGd8d9d9eQZRGd:d;d;eQZSdQd=d>ZTGd?d@d@eQZUGdAdBdBeQZVGdCdDdDeQZWeAdEdFdGZXee?ge?fe?dHdIdJZYdKdLZZdS)RN)ABCabstractmethod)Counter defaultdict)Enum)AnyCallable NamedTupleOptionalUnion)OptimizedModule) FSDPModule UnshardHandle)_Loss)record_function)generate_stage_to_rank_mapping) merge_chunkssplit_args_kwargs_into_chunksTensorChunkSpec)_PipelineStageBase) get_schedule_classPipelineScheduleSinglePipelineScheduleMulti Schedule1F1B ScheduleGPipeScheduleInterleaved1F1BScheduleLoopedBFSScheduleInterleavedZeroBubbleScheduleZBVZeroBubblec@sHeZdZdZdZdZdZdZdZdZ dZ d Z d Z d d Z ed dZdS)_ComputationTyper cCsHtjdtjdtjdtjdtjdtjdtjdtjdtj d tj d i }||S) NFIWUNSHARDRESHARDSEND_FRECV_FSEND_BRECV_BB) r FORWARDBACKWARD_INPUTBACKWARD_WEIGHTr-r.r/r0r1r2 FULL_BACKWARD)selfZstr_mapr9T/var/www/auris/lib/python3.9/site-packages/torch/distributed/pipelining/schedules.py__str__6s z_ComputationType.__str__cCs|dkrtjS|dkrtjS|dkr*tjS|dkr8tjS|dkrFtjS|dkrTtjS|dkrbtjS|dkrptjS|d kr~tj S|d krtj St d |dS) Nr*r+r,r-r.r/r0r1r2r3zInvalid computation type ) r r4r5r6r-r.r/r0r1r2r7 RuntimeErroractionr9r9r:from_strEs*z_ComputationType.from_strN)__name__ __module__ __qualname__r4r5r6r-r.r/r0r1r2r7r; staticmethodr?r9r9r9r:r )sr z?(\d+)(F|I|B|W|UNSHARD|RESHARD|SEND_F|RECV_F|SEND_B|RECV_B)(\d*)c@sHeZdZUeed<eed<dZeeed<ddZe e ddd Z dS) _Action stage_indexcomputation_typeNmicrobatch_indexcCs4t|j}|t|j7}|jdur0|t|j7}|SN)strrErFrG)r8reprr9r9r:__repr__{s   z_Action.__repr__) action_stringcCsj|}t|}rJ|\}}}tt|t|t|rDt|ndS|dkrVdSt d|ddS)z Reverse of __repr__ String should be formatted as [stage][action type][(microbatch)] e.g. `2F0`, `1UNSHARD`, `3SEND_F1` NzInvalid action string: zD, should be formatted as [stage][action type][(microbatch)] e.g. 2F0) strip _action_regexmatchgroupsrDintr r?lenr<)rLrPrErFrGr9r9r:r?s z_Action.from_str) r@rArBrR__annotations__r rGr rKrCrIr?r9r9r9r:rDvs rD)pipeline_ordererror_step_numberreturnc s6tD]6}tt|D] }||dur"d||<q"qtddDfddtD}fddtD}ttj |ddi}t}d dt|D}d dt |g|RDd t|d d d fddt |D} fddt ||D} | dd | d} | S)z Formats the pipeline order in a timestep (row) x rank (column) grid of actions and returns the formatted string. If `error_step_number` is passed in, an additional label will be added to signify which step that it is erroring on. NrMcss|]}t|VqdSrH)rS).0actionsr9r9r: z)_format_pipeline_order..c s*g|]"}dt|ttdqS)zStep r)rIzfillrSrXi) num_stepsr9r: sz*_format_pipeline_order..csg|]}|dgqS)rM)get)rXkey)r_rUr9r:r`s fillvaluecSsg|]}dt|qS)zRank rIr]r9r9r:r`r[cSsg|]}tdd|DqS)css&|]}|durtt|ndVqdS)Nr)rSrI)rXitemr9r9r:rZr[4_format_pipeline_order...)max)rXcolr9r9r:r`s rr!c3s$|]\}}|d|VqdS.N Expecting z losses but got ) isinstanceryanyrSrrr<clearextend)r8stageslossesZcontains_last_stager9r9r:_update_losses s   z _PipelineSchedule._update_lossesarg_mbs kwarg_mbsrrcCstdS)z Run one iteration of the pipeline schedule with list of microbatches. Will go through all the microbatches according to the schedule implementation. Args: microbatches: list of microbatch args. NNotImplementedError)r8rrrrr9r9r:_step_microbatches"sz$_PipelineSchedule._step_microbatchestargetrrcOstdSa Run one iteration of the pipeline schedule with *whole-batch* input. Will chunk the input into microbatches automatically, and go through the microbatches according to the schedule implementation. args: positional arguments to the model (as in non-pipeline case). kwargs: keyword arguments to the model (as in non-pipeline case). target: target for the loss function. losses: a list to store the losses for each microbatch. Nr)r8rrargskwargsr9r9r:step4s z_PipelineSchedule.stepcstdfdd }|dur&||dn dgj}|durF||dn igj}|durd||d|durt|tstd t|||fS) z* Pre-process/check inputs )namecsRt|ts t|dt|t|jkrNtdjd|dt|dS)Nz must be a list but got a rriz but got )rry TypeErrortyperSr ValueError)Zmbsrr8r9r:check_type_and_lenMs  z;_PipelineSchedule._check_inputs..check_type_and_lenNrr9rrz losses must be a list but got a )rIrrryrr)r8rrrrrr9rr: _check_inputsBs       z_PipelineSchedule._check_inputscCs |||SrH)r)r8rrr9r9r:rhsz_PipelineSchedule._compute_loss)rrcCsF|s|r*t|||j|j|j\}}||fSdg|jig|jfSdS)zj Splits a full-batch input into chunks (i.e. microbatches) and returns the chunks r9N)rrrr)r8rr args_split kwargs_splitr9r9r: _split_inputsks z_PipelineSchedule._split_inputs) output_chunksrWcCs t||jS)z Merge output chunks back to a batch state. If output_merge_spec is None, the utility will merge output chunks by dimension 0 (batch dim). )rr)r8rr9r9r:_merge_outputssz _PipelineSchedule._merge_outputs)NNNNT)NNNN)NNNN)N)r@rArBrRr rtorchZTensortuplerdictrIr rboolrrrrrryrrrrrrr9r9r9r:rs\ "  &  r)p2p_opsdescrWcCs:t|dkrgS|r|dnd}td||t|S)zt Simple wrapper over batch_isend_irecv from torch.distributed, which just adds a descriptive logger on top. rz, rMzbatch_p2p %s%s)rSrdebugdistZbatch_isend_irecv)rrZdesc_strr9r9r: _batch_p2ps  rcCs`tt}i}t|dkr|S|D]}||j|q t|D]\}}t||d||<qB|S)z Sorts the list of P2P ops by the peer rank, and then calls batch_isend_irecv. Return a dictionary of works by peer rank. This function helps us avoid hangs in case of skip connections. rr)rryrSpeerrrxitemsr)rrZ ops_by_peerZ work_by_peeropropsr9r9r:_sorted_batch_p2ps  r)workcCs|D] }|qdS)zX Waits for a list of dist.Work (typically from _batch_p2p / _sorted_batch_p2p). N)wait)rwr9r9r:_wait_batch_p2psrc seZdZdZdeeeeeee dfee e e fee e e e fee fedfdd Zdd Zddd eed d d Zee eeeefdddZZS)ra Base class for single-stage schedules. Implements the `step` method. Derived classes should implement `_step_microbatches`. Gradients are scaled by num_microbatches depending on the `scale_grads` argument, defaulting to True. This setting should match the configuration of your loss_fn, which may either average losses (scale_grads=True) or sum losses (scale_grads=False). NT.)rrrrrrrcsftj||||||d||_|j|_|j|j_d|_||jkrXtd|d|jd| |_ dS)NrFzNumber of microbatches (z9) must be greater than or equal to the number of stages (z).) superr_stage num_stages _num_stagesr has_backward_stage_initializedr_get_pipeline_orderrU)r8rrrrrrrrr9r:rs(    zPipelineScheduleSingle.__init__cCs0|j|j|||jr&|j|jd|_dSNT)r_prepare_forward_infrarr_prepare_backward_infrar)r8rrr9r9r:_initialize_stagesz(PipelineScheduleSingle._initialize_stagerrcOsh|j|||\}}|dur6tt||j}nd}||||||jjr`| |jj SdSdSr) rclear_runtime_statesrryr tensor_splitrrrrr)r8rrrrrr targets_splitr9r9r:rs zPipelineScheduleSingle.steprWcCsdS)a Returns the pipeline execution order as a schedule IR. The returned IR is a dictionary mapping rank IDs to lists of actions. Each action is either an _Action object representing computation to perform, or None representing a deliberate idle step. The None values are used to represent pipeline bubbles where a rank must wait for dependencies from other ranks before proceeding. However during execution, with the _PipelineScheduleRuntime, these Nones are skipped since the relevant communication (send/recv) will be scheduled and waited on. Returns: A dictionary mapping rank -> list of actions Nr9rr9r9r:rsz*PipelineScheduleSingle._get_pipeline_order)NNNNT)r@rArB__doc__rrRr rrrrrIr rrrrryrrDr __classcell__r9r9rr:rs$$!rc@s6eZdZdZdeeeeeeeedddZdS)_ScheduleForwardOnlyzo The forward-only schedule. Will go through all the microbatches and perform only the forward pass Nrc Cs|dus|durtd|||||\}}|jsF||d|dg}t|jD]}td|||j|}t |dd}| D] } t | q|j ||||||j |}t |dd}|| Wdn1s0Ytd|jj|qT|D]} t | q dS) z< Run one iteration of the pipeline schedule Nz7Forward-only schedule does not support loss computationrForward fwd_recvrfwd_send[%s] Forwarded microbatch %s)r<rrrrvrrrget_fwd_recv_opsrrwrforward_one_chunkget_fwd_send_opsrrrrE) r8rrrrfwd_sends_to_waitr^rworksrr9r9r:r,s*       ,z'_ScheduleForwardOnly._step_microbatches)NNNN)r@rArBrr ryrr9r9r9r:r&src@sXeZdZdZd eeeeeeeedddZeeeeee fdddZ dS) rz^ The GPipe schedule. Will go through all the microbatches in a fill-drain manner. Nrc Cs:|||||\}}|js.||d|dg}t|jD]}td|||j|}t|dd}| D] } t | qp|j |||||} |j |}t|dd}| | Wdn1s0Ytd|jj|||j| ||q<|D]} t | q|js dSg} t|jD]}td||j|}t|d d}| D]} t | qb||j|} |jj|| ||jd kd |j|}t|d d}| | Wdn1s0Ytd |jj|q.|jj|jr |jnd d||j|| D]} t | q&dS)z Run one iteration of the pipeline schedule with list of microbatches. Will go through all the microbatches according to the GPipe schedule. Args: microbatches: list of microbatch args. rrrrrNrz Backward bwd_recvrr last_backwardbwd_sendz[%s] Backwarded microbatch %sgrad_scale_factor)rrrrvrrrrrrwrrrrrrrErrget_bwd_recv_opsrbackward_one_chunkget_bwd_send_opsrr) r8rrrrrr^rrrrZbwd_sends_to_waitrr9r9r:r_sV      ,        .z ScheduleGPipe._step_microbatchesrcCsi}|j}t|D]}g}|}|dg|t|jD]}|t|tj|q8d|d|}|dg|t|jD]}|t|tj|q||||<q|S)z Returns the pipeline order for GPipe schedule. See base method in PipelineScheduleSingle for details on the schedule IR format. Nr"r) rrvrrrrDr r4r7)r8rU pp_group_sizer}rYZ warmup_delaymb_idxZbackward_delayr9r9r:rs  z!ScheduleGPipe._get_pipeline_order)NNNN r@rArBrr ryrrrRrDrr9r9r9r:rYs Trc@sXeZdZdZd eeeeeeeedddZeeeeee fdddZ dS) rzo The 1F1B schedule. Will perform one forward and one backward on the microbatches in steady state. NrcCsj|||||\}}|js.||d|dt|j|j|jj}d}d}g}g} t|D]~} |j |} t t | dd|j |||||} t ||j |} ||dkrt | dd}||j| |||d7}q\|j|} t t | | dd||j|}|jj||||jdkd|j|}|d7}||jkrJq|j |} t t || dd|j |||||} ||j| |||j |} |d7}qt |d d}||jkr4|j|} t t | d d||j|}|jj||||jdkdt ||j|}t |d d}|d7}q|jj|jrH|jndd t |||j|d S) z Run one iteration of the pipeline schedule with list of microbatches. Will go through all the microbatches according to the 1F1B schedule. Args: microbatches: list of microbatch args. rrrrrZfwd_send_bwd_recvrZbwd_send_fwd_recvrrrN)rrrminrrrrErvrrrrrrrrrrrr)r8rrrrZ warmup_chunks fwd_mb_index bwd_mb_indexZ send_workZ fwd_sends_Z fwd_recvsrZ bwd_recvsrZ bwd_sendsr9r9r:rs|                     zSchedule1F1B._step_microbatchesrc Cspi}|j}t|D]V}g}|dg||d|}d}t|D]}|t|tj||}qDtdd|d|}|dg|d} |j|} | dkr|d7}|t|tj|| d8} |t|tj | | d7} q|j| } | dkrb||dkr<|d| dkr`|t|tj | | d7} | d8} q|t|tj | | d7} | d8} q|||<q|S)z Returns the pipeline order for 1F1B schedule. See base method in PipelineScheduleSingle for details on the schedule IR format. Nrrr!) rrvrrrDr r4rgrr7) r8rUrr}rYZ num_forwardZ forward_mbr^Z wait_for_1f1bZ backward_mbZremaining_forwardZremaining_backwardr9r9r:rfsN              z Schedule1F1B._get_pipeline_order)NNNNrr9r9r9r:rs  rr")compute_actionsmax_active_stagesrWc sttttttddd}tgtdfdd }tdfdd }t|D]~\}}|d urjqX||||d ttfd d }ttfd d }|D] } || q|D] } || q|qXS) aQGiven a basic schedule involving only compute actions (F,B,W), add UNSHARD/RESHARD actions for FSDP. UNSHARD refers to fetching the full contents of an FSDP-sharded layer, requiring an all-gather operation. RESHARD does the opposite, releasing memory (but doing no communication) We abandon the "timestep lock" during lowering max_active_stages controls how many prefetches we allow. It should be measured in mb and tuneable but in practice 3 stages is probably the thing we want? (to account for having one f and one b active, and something else prefetching?) )count next_actionsrWcSsRt}g}|D]>}|dur|j|vr||j||jt||krqNq|S)zdRemove duplicates (same stage, different microbatch), find next 'count' stages that will do compute.N)setrEaddrrS)rrseenretar9r9r:next_stage_indicess   z0_add_unshard_reshard..next_stage_indicesrEcs |t|tddSrH)rrrDr-r active_stagesfsdp_aware_actionsr9r:_unshards z&_add_unshard_reshard.._unshardcs |t|tddSrH)removerrDr.rr r9r:_reshards z&_add_unshard_reshard.._reshardNcs|vSrHr9s)r r9r:r[z&_add_unshard_reshard..cs|vSrHr9r)next_nr9r:rr[)rRryr rDrrpfilterr) rrrr rr^r>fetchZevictrr9)r r rr:_add_unshard_reshards&      r)rrWcCsg}|r|d}|durqt|r@|d}dur@|dq|jtkr|dur|jtkr|j|jkr|j|jkr|t|jt |j|dq||q|S)a9Given a basic schedule involving only compute actions (F,I,W), merge adjacent I and W ops into B ops. (note: I = BACKWARD_INPUT, W = BACKWARD_WEIGHT, B = FULL_BACKWARD) B refers to running the whole backward (not separating grad_input and grad_weight), which can be more efficient in some cases. rN) poprSrFr5r6rErGrrDr7)rZmerged_actionsr>Z next_actionr9r9r: _merge_bws.       r)r stage_to_rankrrWc sdd|D}dd|D}ttdfdd ttttfdfdd }tttttd fd d }|rd }t|D]}t||d ksJd|dt||||d } || ||sq~| durJ||| ||| | rJ|| \} } ||| ||| || j | || j | || d t||d krp||=d}q~|slJdql|S)NcSsi|] }|gqSr9r9rXr}r9r9r: r[z"_add_send_recv..cSsi|] }|tqSr9rrr9r9r:rr[r>rWcsd|jtkr0|jdko.|jd|jkS|jttfvr`|jdko^|jd|jkSdS)NrrF)rFr*rEr5r7r=)rrr9r: _has_commss  z"_add_send_recv.._has_commscsx|sJ|d|j}|j}|j}t||tkr8tnt|}|tkrP|dn|d}t||tkrhtnt|}||fS)Nz is not a valid comm actionr) rErFrGrDr*r/r1r0r2)r> stage_idxctypersendZrecv_stage_idxrecv)rr9r: _get_comms'sz"_add_send_recv.._get_comms)r> prev_actionsrWcs|dur dS|jtkrX|jdksXt|jt|j|vr8dSt|jdt|j|vrTdSdS|jttfvr|jdkst|jt|j|vrdSt|jdt|j|vrdSt|jdt|j|vrdSdSdSdS)aWe don't put our own recv ops in the schedule, we let a sender on another rank put our recv ops in place. This helps ensure a sane (non-hanging) ordering of sends and recvs. But it also means we might not be able to schedule our next compute action yet. NTrrF) rFr*rErDr0rGr5r7r2)r>r#)rr9r:_ready_to_schedule1sF  z*_add_send_recv.._ready_to_scheduleFrrank=z, len(compute_actions[rank])=Tz6Malformed compute schedule, can't schedule sends/recvs) rDrrr rrxrSrrrEr) rrrZ comm_actionsr#r"r$progressr}r>r r!r9)rrrr:_add_send_recvs>  ,     r')rYrrnum_microbatchesrWc Cst||ks$Jd|dt|t|D]}||vs,Jd|q,ddt|D}i}|D]r}||D]b}|durqpt|tsJd|d|j}|j} |j} | tkr||t| n| t kr | ||tvsJd|d | d ||t | n| t krR| ||tvs>Jd |d | d ||t | nD| t kr| ||t vsJd |d | d ||t | ||vr|||<qp||} || kspJd|d|d| qpqb|D]}t||t} t||t } t||t }t||t }| |ksJJd| dtd|d|| ||d|ksJd|d|d| d|d| q|S)Nz2Schedule has incorrect number of ranks - expected z , actual z%Schedule is missing actions for rank c Ss*i|]"}|ttttttttiqSr9)r*rr3r+r,)rXZstage_idr9r9r:rsz&_validate_schedule..zGot an invalid action: z, expected instance of _Actionz Running Full Backward for stage z , microbatch z without first running Forwardz!Running Backward Input for stage z"Running Backward Weight for stage z% without first running Backward InputzStage z is assigned to both rank z and rank zGot riz microbatches for stage z , expected r!z(Invalid backward microbatches for stage z : expected z( total backwards, but got B=z, I=z, W=) rSrvrrDrErFrGr*rr3r+r,)rYrrr(r}Z stage_actionsZstage_index_to_rank_mappingr>Zs_idrZmb_idZ existing_rankZf_mbZb_mbZi_mbZw_mbr9r9r:_validate_schedule}sx            r)c seZdZdZdeeeeeee e dfee e e fee e e efe efeeedfdd Ze edfdd d Ze eeeefdd d d ZddZdddZdddeedddZdeeeeeeeedddZZS)raX Base class for multi-stage schedules. Implements the `step` method. Gradients are scaled by num_microbatches depending on the `scale_grads` argument, defaulting to True. This setting should match the configuration of your loss_fn, which may either average losses (scale_grads=True) or sum losses (scale_grads=False). NT.)rrrrrruse_full_backwardrc stj||||||d||_|dj|_|dj|_|dj|_t |j|j|_ |jD] } |j | _ qX|jD] } |j | _ qld|_ |jdufdd|_i|_|durtddS)NrrFcs |joSrHr)rZhas_lossr9r:rr[z0PipelineScheduleMulti.__init__..zDeprecation warning: 'use_full_backward' is no longer supported. Simply stop passing it, and everything should still work fine.)rr_stagesrr group_sizer group_rankr}rstage_index_to_group_rankrr_stages_initializedrZ_should_compute_lossrUrwarning) r8rrrrrrr*rrrr+r:rs6          zPipelineScheduleMulti.__init__)rcCsVt}|jD]>}|jr(||j||}n||j||}|jr ||jq d|_dSr)rr,is_firstrrrrr0)r8rrZnext_stage_argsrr9r9r:_initialize_stagess z(PipelineScheduleMulti._initialize_stages)rYrWcCs.t||j|j|j|_|jD] }|j|_qdS)z] Allocates the stage index to rank mapping which is needed for communication N)r)rrrr/r,)r8rYrr9r9r:_validate_and_set_stage_mappings z5PipelineScheduleMulti._validate_and_set_stage_mappingcCsXt|ddd6}t|}|jD]}||j|q Wdn1sJ0YdS)zQDump a CSV representation of the schedule into a file with the provided filename.rrMnewlineN)opencsvwriterrUwriterowr8filenamecsvfiler9r}r9r9r: _dump_csv&s  zPipelineScheduleMulti._dump_csv compute_onlycCsx|dks Jt|dd@}t|}t|D]\}}dd|D|j|<q,Wdn1s^0Y||jdS)zLoad a CSV representation of the schedule from a file with the provided filename. This API will most likely get renamed/refactored so is marked as internal for now. format must be "compute_only" for PipelineScheduleMulti. r?rMr5cSsg|]}t|qSr9rDr?rXrr9r9r:r`7r[z3PipelineScheduleMulti._load_csv..N)r7r8readerrprUr4)r8r<formatr=rBr}rrr9r9r: _load_csv-s   4zPipelineScheduleMulti._load_csvrrc Osz|jD] }|q|||\}}|dur@tt||j}nd}||||||jD]}|jrZ| |j SqZdSr) r,rrryrrrrrrr) r8rrrrrrrrr9r9r:r=s   zPipelineScheduleMulti.steprc CsZ|||||\}}|js.||d|ddd|jD}t}t}|D]B}|dkrr||j|d||jdkrR||j|dqRt } t |j |j D]\} } z4g} | durr| j } | j}| j}|dusJd| tjkr8||}||||||}|||||| ||n:| tjkr||}|||}| |d7<| ||jk}|jr|jnd}|j||d|d|r||| ||n| tjkr||}|||}|j||d d d| ||nr| tjkrd||}| |d7<| ||jk}|jr@|jnd}|j||d |rr||ntd | |D]}|j |}d}| t |kr|| }|durv|j } |j}|j}|dusJd| tjkr|d|vr$||d}| |!|n | tttfvrntd | qv|D]}|j |}d}| t |krT|| }|dur,|j } |j}|j}|dusJd| ttfvrnH| ttfvr|d|vr||d}| |"|ntd | q,t#t$| Wqt%yD}z>t&'d |j |j(j)| | t&'d t*|j | d|WYd}~qd}~00q|+|j|dS) Operate on the microbatches for looped schedules (multiple stages on each rank). TODO: Does not use sorted_batch_isend_irecv(). As a result, this schedule does not support models with skip connections. rcSsi|] }|j|qSr9rrr9r9r:rssz.rNzCAll currently supported action types require valid microbatch_indexTrZ full_backwardrFrzUnknown computation type zy[Rank %s] pipeline schedule %s caught the following exception at time_step %s when running action %sz%srV),rr0r3r,rkeysrr/rrrprUr}rFrGrEr r4rrrrr7rrrrrr5r6backward_weight_one_chunkrrSrrrr Exceptionrerrorrr@r~r)r8rrrrstage_index_to_stageZall_prev_ranksZall_next_ranksrEbackward_counter time_stepr>rrFrrrrrrZ prev_rankZ prev_rank_opsZprev_rank_actionZ next_rankZ next_rank_opsZnext_rank_actioner9r9r:r_s                       z(PipelineScheduleMulti._step_microbatches)NNNNNT)r?)NNNN)r@rArBrryrrRr rrrrrIr rrrr3rDr4r>rDrrrr9r9rr:rsD 2  $rcseZdZdZdeeeeefe dfdd Z de e dfdd Z e d d d Z d d Z deeeeeeeedddZZS)_PipelineScheduleRuntimea% Provides a simple runtime that requires a 'schedule IR' including specified communication operations. Can be instantiated directly by creating _PipelineScheduleRuntime and calling load_csv, or can be subclassed and the subclass can be responsible for creating a schedule IR. r?)rYrCcst|i_|dkrZ|D]8}gj|<||D] }|dusDJj||q4qnR|dkr|D]}t||j|<qftjfddjd_ntd|ddS) z Given an in-memory representation for a simple compute-only schedule, lower it to a complex schedule including communication actions. Stores the schedule in self, and must be called before running step_mo() compute_commsNr?cs j|SrHr/rrr9r:r@r[z8_PipelineScheduleRuntime._load_actions..rrformat= is not implemented)rr4pipeline_order_with_commsrrr'rr)r8rYrCr}r>rrr: _load_actions#s&       z&_PipelineScheduleRuntime._load_actions)r<rCcs|dkr"t|||jn|dkri}t|ddL}t|}t|D]\}}dd|D||<qN|j||dWdq1s0Yntd |d dS) a Loads a csv in simple format and then lowers it to include communication actions format must be either "compute_only" or "compute_comms". If compute_only, the lowering passes will automatically be run to generate a compute_comms schedule. r?rRrMr5cSsg|]}t|qSr9r@rAr9r9r:r`Vr[z6_PipelineScheduleRuntime._load_csv..)rCNrUrV) rrDrXrUr7r8rBrpr)r8r<rCrYr=rBr}rrrr9r:rDFs  .z"_PipelineScheduleRuntime._load_csv)r<cCsj|jdusJdt|ddd6}t|}|jD]}||j|q2Wdn1s\0YdS)zaDump a CSV representation of the compute + comms schedule into a file with the provided filename.Nz6Must initialize compute_comms schedule before dump_csvrrMr5)rWr7r8r9r:r;r9r9r:r>[s   z"_PipelineScheduleRuntime._dump_csvcstjfddjS)Ncs j|SrHrSrrr9r:rjr[z4_PipelineScheduleRuntime._simulate..)_simulate_comms_computerWrrr9rr: _simulategs  z"_PipelineScheduleRuntime._simulateNrc sZ|||||\}}|js.||d|ddd|jD}|jdusPJdi}i}g}ittdfdd } t} t|j|j D]\} } z.| j } | j dur| j nd }|dks| t t fvsJd | d | j}||}t|jt}|d |v}|d |v}td | | | tkr>|t||n| tkr`|t||nh| tkr||f|vsJdt|||||f<n.| tkr||f|vsJdt|||||f<n| t kr|r|vr|vsJd|d|jjdd|<n| t krp|r|vsHJd|d|vsbJd|d|jnX| t kr |r| ||j!s|s||f|vsJd| dt"|#||f|$|||||}|%|||||r||d &||n| t'kr|r"| ||j(s`|s`||f|vsNJd| dt"|#||f|)||}| |d 7<| ||j*k}|j+r|j*nd }|j,||d|d|r|+||r||d -|.||n| t/krv|r| ||j(s6|s6||f|vs$Jd| dt"|#||f|)||}|j,||ddd|r||d -|.||nR| t0kr|r| || |d 7<|j1|| ||j*kdnt2d | dWqt3y}z0t4d| | t5t6|j| d|WYd}~qd}~00qt7|r2t"|#qt7dksHJd |8|j|dS)!rErcSsi|] }|j|qSr9rrr9r9r:rsz?_PipelineScheduleRuntime._step_microbatches..Nz=Must call _load_actions() before calling _step_microbatches()rcs>|vr$||=||vs:Jd|dS)zQIf an unshard is active for `stage_idx`, wait() it and mark `stage_idx` unshared.z*Attempted to compute on sharded stage_idx=N)rrr[Z unshard_opsZunsharded_stagesr9r:_assert_unshardeds   zF_PipelineScheduleRuntime._step_microbatches.._assert_unshardedzaction=z missing mb_indexrz8_PipelineScheduleRuntime running time_step %d, action %szARecv twice for {stage_idx=} {mb_index=} without executing forwardzBRecv twice for {stage_idx=} {mb_index=} without executing backwardzUnsharding the same stage_idx=z twiceT)Zasync_opzResharding stage_idx=z without unshardingz before finishing unshardzComputing action=z before receiving inputz Attempted to run compute action=rFFrGz is unknown or unsupportedz\_PipelineScheduleRuntime caught exception at step %s when running action %s. Full Schedule:rHzUnused unshard operations)9rr0r3r,rWrrRrrpr}rFrGr-r.rErsubmodr rrr/rrrr1rr0rr2rZunshardZreshardr4r2rrrrZset_local_fwd_inputr7rrrrrZset_local_bwd_inputZget_local_bwd_outputr5r6rJrrKrLprintr~rSr)r8rrrrrMZ bwd_recv_opsZ fwd_recv_opsZsend_opsr]rNrOr>Z comp_typerrrZstage_uses_fsdpZis_next_stage_on_this_rankZis_prev_stage_on_this_rankrrrrrPr9r\r:rnsx                                   z+_PipelineScheduleRuntime._step_microbatches)r?)r?)NNNN)r@rArBrrrRryr rDrIrXrDr>rZrrr9r9rr:rQs& #  rQc s`eZdZdZd eeeeee e feee e e fee fedfdd ZddZZS) rai Breadth-First Pipeline Parallelism. See https://arxiv.org/abs/2211.05953 for details. Similar to Interleaved 1F1B, Looped BFS supports multiple stages per rank. What is different is that when microbatches are ready for multiple local stages, Loops BFS will prioritizes the earlier stage, running all available microbatches at once. NTrrrrrcsDtj|||||di|_t|jD]}||}||j|<q&dS)Nra)rrrUrvr!_calculate_single_rank_operations)r8rrrrrr}rank_opsrr9r:rks  zScheduleLoopedBFS.__init__cst|j}t||j||j}ddt|D}|D]"|fddt|jDq4d|jd|}|dg|t|D]&|fddtt|jDq|S)NcSsg|]}dqSrHr9rXrr9r9r:r`r[zGScheduleLoopedBFS._calculate_single_rank_operations..c3s|]}ttj|VqdSrH)rDr r4rXrrr9r:rZszFScheduleLoopedBFS._calculate_single_rank_operations..r!rc3s|]}ttj|VqdSrH)rDr r7rerr9r:rZs)rSr,rvrrrreversed)r8r}n_local_stagesZ stage_indicesrcpost_warmup_opsr9rr:rbs     z3ScheduleLoopedBFS._calculate_single_rank_operations)NNT)r@rArBrryrrRr r rrrrIrrrrrbrr9r9rr:ras rFc Cstt} tt} tt} ddt|D} ||d|d|||}| rZ||d}|||}g}d}| rvtnt}t|D]}||kr||}| |}d| |<| t|tj|||dkr| dg|q||kr||krnn||}| |}d| |<| t|tj|||}| |}d| |<| t|||||| rJ|||krJ|||}| |}d| |<| t|tj ||d7}q| s| d||}| |}d| |<| t|||||| r|||kr|||}| |}d| |<| t|tj ||d7}q| r|t |kr|||}| |}d| |<| t|tj ||d7}qL| S)NcSsg|]}dqSrHr9rdr9r9r:r`r[z&_get_1f1b_rank_ops..r!rr) rrRrvr5r7rrDr r4rr6rS)rgr warmup_ops fwd_bwd_ops cooldown_opsr}forward_stage_indexbackward_stage_indexnum_1f1b_microbatchesenable_zero_bubbleZfwd_stage_mb_indexZbwd_stage_mb_indexZweight_stage_mb_indexrcrh total_opsZbackward_op_idsZweight_op_countZFULL_BACKWARD_OR_BACKWARD_INPUTrZfwd_stage_indexrrZbwd_stage_indexrZweight_stage_indexZweight_mb_indexr9r9r:_get_1f1b_rank_opss                        rqc seZdZdZd eeeeeee e dfee e e fee e e efe efedfdd Zeeedd d ZZS) ra The Interleaved 1F1B schedule. See https://arxiv.org/pdf/2104.04473 for details. Will perform one forward and one backward on the microbatches in steady state and supports multiple stages per rank. When microbatches are ready for multiple local stages, Interleaved 1F1B prioritizes the earlier microbatch (also called "depth first"). This schedule is mostly similar to the original paper. It differs by being relaxing the requirement of num_microbatch % pp_size == 0. Using the flex_pp schedule, we will have num_rounds = max(1, n_microbatches // pp_group_size) and it works as long as n_microbatches % num_rounds is 0. As a few examples, support 1. pp_group_size = 4, n_microbatches = 10. We will have num_rounds = 2 and n_microbatches % 2 is 0. 2. pp_group_size = 4, n_microbatches = 3. We will have num_rounds = 1 and n_microbatches % 1 is 0. NT.rrrrrrrc s|dj|_tj|||||||dt||_|dj|_td||j|_ ||j |_ ||j dkrt d|j d|di|_ t |jD]}||} | |j |<qdS)Nrrrrz_Interleaved 1F1B requires the number of microbatches to be a multiple of the number of rounds ( ), but got .)r-rrrrSrgr.r}rgnumber_of_roundsmicrobatches_per_roundrrUrvrb) r8rrrrrrrr}rcrr9r:r<s4     z ScheduleInterleaved1F1B.__init__rc sfdd}|jj}|}||}||}td|||fdd}fdd}tjj||||S)Ncs<jdj}d}||jd|}t|jjS)Nrr!rgrvrrrr}Zwarmups_ops_last_stageZmultiply_factorrirr9r:get_rank_warmup_opscs zVScheduleInterleaved1F1B._calculate_single_rank_operations..get_rank_warmup_ops=rank %s, warmup_ops %s, 1f1b %s, cooldown_ops %s total_ops %scs|jj}|jSrHrvrgrrZ local_indexr}r8r9r:rlszVScheduleInterleaved1F1B._calculate_single_rank_operations..forward_stage_indexcs,jd|jj}|jSNrrgrvrr|r}r8rir9r:rms zWScheduleInterleaved1F1B._calculate_single_rank_operations..backward_stage_indexrgrrrrqr) r8r}rymicrobatch_opsrjrkrprlrmr9rr:rbbs4     z9ScheduleInterleaved1F1B._calculate_single_rank_operations)NNNNTr@rArBrryrrRr rrrrrIr rrrrDrbrr9r9rr:r*s &rc seZdZdZd eeeeeee e dfee e e fee e e efe efedfdd Zeeedd d Zd d ZZS)raw The Interleaved Zero Bubble schedule. See https://arxiv.org/pdf/2401.10241 for details. Will perform one forward and one backward on inputs for the microbatches in steady state and supports multiple stages per rank. Uses the backward for weights to fill in the pipeline bubble. In particular this is implementing the ZB1P schedule in the paper. NT.rrc s|D]}t|jtrtdq|dj|_tj|||||||dt||_ |dj |_ t d||j|_ ||j |_||j dkrtd|j d|di|_t|jD]} || } | |j| <q||j |j|_dS)NzYThe Zero Bubble schedule is not supported with stage modules that have used torch.compilerrrrzZZero bubble requires the number of microbatches to be a multiple of the number of rounds (rsrt)rr_r r<r-rrrrSrgr.r}rgrurvrrUrvrb_add_bubbles_to_actions r8rrrrrrrrr}rcrr9r:rsD         z&ScheduleInterleavedZeroBubble.__init__rc sfdd}|jj}|}||}||}td|||fdd}fdd}} tjj||||| dd S) Ncs<jdj}d}||jd|}t|jjSr~rwrxrr9r:rys z\ScheduleInterleavedZeroBubble._calculate_single_rank_operations..get_rank_warmup_opsrzcs|jj}|jSrHr{r|r}r9r:rl sz\ScheduleInterleavedZeroBubble._calculate_single_rank_operations..forward_stage_indexcs,jd|jj}|jSr~rr|rr9r:rm s z]ScheduleInterleavedZeroBubble._calculate_single_rank_operations..backward_stage_indexT)ror) r8r}ryrrjrkrprlrmrnr9rr:rbs:     z?ScheduleInterleavedZeroBubble._calculate_single_rank_operationscCsx|j}dd}t}i}i}i}d}t|jD]} g|| <d|| <d|| <q.d} t} t|jD]} || } | t|| kr~q`d} || | dur(|| | } | dusJ| \}}}||||||s|| || | |dur| |||f|| d7<n|| d|| d7<q`|| d7<|| dq`|| | rLq\qL|dkrtt d|||S)NcSsf|tjkr*|dkrb|d||f|vrbdSn8|tjkrb||dkrP|tj|f|vS|d||f|vSdS)NrrTF)r r4r7)rr microbatchnum_stages_globalseen_opsr9r9r: need_bubble" s   zJScheduleInterleavedZeroBubble._add_bubbles_to_actions..need_bubblerTFrz?Non zero bubbles added: total_bubbles_added=%s bubbles_added=%s) rUrrvrrSrrupdaterr1)r8rrYrrresultZ next_pointerZ bubbles_addedZtotal_bubbles_addedr}Z should_stopZ temp_seen_ops timestampZ temp_actionrErrr9r9r:r sV        z5ScheduleInterleavedZeroBubble._add_bubbles_to_actions)NNNNT)r@rArBrryrrRr rrrrrIr rrrrDrbrrr9r9rr:rs"6@rc seZdZdZd eeeeeee e dfee e e fee e e efe efedfdd Zeeedd d ZZS) ra The Zero Bubble schedule (ZBV variant). See https://arxiv.org/pdf/2401.10241 Section 6 for details. This schedules requires exactly two stages per rank. This schedule will perform one forward and one backward on inputs for the microbatches in steady state and supports multiple stages per rank. Uses backward with respect to weights to fill in the pipeline bubble. This ZB-V schedule would have the "zero bubble" property only if time forward == time backward input == time backward weights. In practice, this is not likely true for real models so alternatively a greedy scheduler could be implemented for unequal/unbalanced time. NT.rrc s|dj|_tj|||||||dt|j|jdd|_|jD] }|j|_q@t||_ |j dkrtt d|j d|dj |_ |dj |_ i|_t|jD]} || } | |j| <qdS)Nrrrv)styler!z0ZBV requires exactly 2 stages per rank, but got rt)r-rrrrrr/r,rSrgrr.r}rrUrvrbrrr9r:rr s8          zScheduleZBVZeroBubble.__init__rcs,tdjdj}ddt|D}d\}}}}dj|d}|} jd|} t|D] } |t| t|d|d7}q`|} t| D]<} |t| t|d|d7}|t| t|d|d7}qj|} t| D]P} |t| t|d|d7}|t| t|d|t| t |d|d7}q||ksD||kr||krj|t| t|d|d7}|t| t|d|t| t |d|d7}|t| t|d|d7}|t| t|d|t| t |d|d7}q0||}}|}t|D]>} |t| t|d|d7}|t| t|d|d7}qj|}t|D]>} |t| t|d|d7}|t| t |d|d7}qR||kr|t| t |d|d7}q||kr|t| t |d|d7}q||kr||ksJ||kr||ksJfdd|D}|S)Nr!rcSsg|]}dqSrHr9rdr9r9r:r` r[zKScheduleZBVZeroBubble._calculate_single_rank_operations..)rrrr)rFrGcs2g|]*}|dur*|jdur*|jjkr*|ndqSrH)rGr)rXr>rr9r:r` s ) rgrrrvrrrDr*r+r,)r8r}Zn_microrcZf0_cntZf1_cntZb0_cntZb1_cntZ warmup_n1Zstage_id_chunk0Zstage_id_chunk1rZ warmup_n2Z warmup_n3Zw0_cntZw1_cntZ cooldown_n1Z cooldown_n2r9rr:rb s                                       z7ScheduleZBVZeroBubble._calculate_single_rank_operations)NNNNTrr9r9rr:rb s ,r) schedule_namec Cs`ttttttttd}dd|D}| }||vrTt d|dt ||||S)z Maps a schedule name (case insensitive) to its corresponding class object. Args: schedule_name (str): The name of the schedule. )Z1F1BZInterleaved1F1BZGPipeZ LoopedBFSZInterleavedZeroBubblerrZ ZBVZeroBubblecSsi|]}||qSr9)lower)rXkr9r9r:r& r[z&get_schedule_class..zUnknown schedule name 'z'. The valid options are ) rrrrrrrrrIrrry)rZ schedule_mapZlowercase_keysZlowercase_schedule_namer9r9r:r s  rrTc sfddtDddtDddDtttdfdd }tttdfd d }rd }tD]Z}t|d krq~|d }||r|d ur||||d d}q~||d q~tddD]}t|d kr|=qtD]z}t|d kr$q |dd ur:q |d }||r |d urt||d<|||d q tddD]}t|d kr|=q|sltdt D]"}td|d|d qt dqlS)aThis function dry-run simulates the actions in the schedule from the perspective of all ranks, and flags any deadlocks caused by missing or misordered communications. It also simulates any bubbles in time where a rank can not execute any action due to waiting for unmet dependencies. The total number of simulator steps can be used as a metric for unit tests involving IR optimization passes as reordering and merging of IR can reduce the number of simulated steps. The simulation is not high-fidelity and does not model overlapping of compute and communication, or cuda streams. Future work may be to enhance this and model the compute time, comms overlap, and even memory. cs i|]}|dd|DqS)cSsg|]}|dur|qSrHr9)rXrr9r9r:r`< r[z6_simulate_comms_compute...r9r)rUr9r:r; sz+_simulate_comms_compute..cSsi|] }|gqSr9r9rr9r9r:r? scSsi|] }|tqSr9rrr9r9r:rC r[r}r>cs(|||dur$||dSrH)rrr)_prev_ops_rank _scheduler9r:add_to_scheduleE sz0_simulate_comms_compute..add_to_schedulercs|dur dS|j}|}|jtkrn|jdkr6dSt|jt|j|vrNdSt|jdt|j|vrjdSdS|jttfvr|jdkrdSt|jt|j|vrdSt|jdt|j|vrdSt|jdt|j|vrdSdS|jt krdS|jt krt|jt|j}||vS|jtkrF|d}t|t |j}||vS|jt krt|jt|j}t|jt|j}||vp||vS|jtkr|d}t|t |j}||vSt d|dS)NTrrFzUnsupported action type ) rErFr*rDr0rGr5r7r2r6r/r1r)r>rZprev_opsZ expected_fZpeer_stage_idxZ expected_sendZ expected_bZ expected_bw)rrrr9r:r$J sh          z3_simulate_comms_compute.._ready_to_scheduleFrNT)reverser^zWIP comms schedule: r%z next action= zSchedule is not progressing) rxrRr rDrrSrrr`r~r) rUrrrr$r&r}r>r^r9)rrrrUrr:rY/ sX :            rYc Csg}t|D]V}t||D]D\}}|dur.q|t||jtttfvrLdndd|||ddqq ddl}t |d }| d |i|Wdn1s0YdS) a This function dumps a schedule IR into a chrometrace format so it can be visualized. It is currently very basic and only serves as a graphical alternative to dumping the schedule IR as text. As future work we may extend this to include more accurate heuristics for durations, or let users input durations, add 'flow events' to let the UI show the connection between sends and recvs, and model cuda streams for comm/compute as separate streams on the chrometrace view. NZ computationZ communicationXr)rcatphpidtidtsZdurrrZ traceEvents) rxrprrIrFr*r3r,jsonr7dump)Zscheduler<eventsr}Ztimestepr>rfr9r9r:_dump_chrometrace s(   r)N)N)N)r")rF)[rtr8rzloggingreabcrr collectionsrrenumrtypingrrr r r rZtorch.distributedZ distributedrZ torch._dynamor Ztorch.distributed.fsdpr rZtorch.nn.modules.lossrZtorch.profilerr_utilsrrrrrrr__all__ getLoggerr@rr r4r5r6r-r.r/r0r1r2r7r*r+r,r3compilerOrDrrRryrIr~rZP2POpZWorkrrrrrrrrrr'r)rrQrrqrrrrrYrr9r9r9r:s        6% >8  i3{Z  G  $  j LTHI  tE4