o I h^ @sPddlZddlZddlmZddlZddlZddlmZddlm Z m Z ddl m Z GdddeZ Gd d d eZejd efd d Zdejd e fddZdejd efddZdejd efddZGdddeZGdddeZGdddeZdgdggZdgdggdgdggdgd gggZgd!gd"gd"gZdejd efd#d$ZdS)%N)IntEnum)ir)get_dtype_size sympy_product)Vc@eZdZdZdZdZdS) NCCL_COLLrrN)__name__ __module__ __qualname__ ALL_REDUCE ALL_GATHERREDUCE_SCATTERrrQ/var/www/vscode/kcb/lib/python3.10/site-packages/torch/_inductor/comm_analysis.pyr r c@r)NVIDIA_GPU_TYPErrr N)r r r VOLTAAMPEREHOPPERrrrrrrrreturncCsHtjjtjjjp d}d|vrtjSd|vrtjSd|vr!tjStjS)NV100A100H100) torchutils collect_env get_gpu_inforunrrrr)gpu_inforrr get_gpu_typesr#nodecCsdt|tjs td||j}|dusJd|vrtjSd|vr$tjSd|vr+tjStd|)Nz!node is not a collective kernel: all_reduce all_gatherreduce_scatterzUnsupported collective kernel: ) isinstancer_CollectiveKernel ValueErrorpython_kernel_namer rrr)r$ kernel_namerrrget_collective_type(s  r-cCs\d}|jD]&}t|jj}t|tjrt|}n tj j j |dd}||t |jj 7}q|S)Nr)fallback)inputsrlayoutsizer(sympyIntegerintrgraphsizevars size_hintrdtype)r$sz_bytesinpnumelrrrget_collective_input_size_bytes8s    r<cCs6t|tjkrddlm}||jdStd|)Nr)_get_group_size_by_namezUnsupported collective type: )typerr)"torch.distributed.distributed_c10dr= constant_args TypeError)r$r=rrrget_collective_group_sizeEs rCc@r)NCCL_HWrrr N)r r r NVLINKPCINETrrrrrDSrrDc@seZdZdZdZdS) NCCL_ALGOrrN)r r r TREERINGrrrrrHYsrHc@seZdZdZdS) NCCL_PROTOrN)r r r LLrrrrrK^srKg333333@gffffff@g333333??g@g@)C@rNgffffff4@)gU@g6@g3@cCst|}|ddd}d}t|}t||}|}|dkr!dStj}tj}t|} t j j j } t j j j } t} |dkr@|dnd} |dkrH| nd}t|| }|dkrV| n| }d}||}t|||dksj| tjkrldnd}| tjkr|d|d}n | tjtjfvr|d}d||}||}|d }tj}| tjkr|dkrd|}nd}n | tjtjfvr|d}t||}t|||}ttj||}d }|dkrd}t||}||||||7}|d }||}||S) a9 Returns estimated NCCL collective runtime in nanoseconds (ns). The following heuristics are copied from https://github.com/NVIDIA/nccl/blob/master/src/graph/tuning.cc. We aim to estimate the runtime as accurately as possible. Assumptions: - only ring algorithm (NCCL_ALGO_RING) is used - only Low-Latency protocol (NCCL_PROTO_LL) is used, i.e. Simple or LL128 is not used - 8 gpus per node # TODO: Need to find a way to get accurate "gpus per node" and "# nodes" info. - collective is one of: allreduce, reducescatter, allgather irrr g?gUUUUUU?rMgeAgg@@)r<rCmathceilrHrJrKrLr-r _inductorconfig intra_node_bw inter_node_bwr#llMaxBwsminr rrrrDrEbaseLathwLatrGmax)r$tensor_storage_size_bytestensor_storage_size_GBnum_gpus_per_node group_sizenNodesnRanks nccl_algo nccl_protocollbwIntrabwInter compCapIndexindex2index1llMaxBwbw nChannelsbusBwnstepsratio bandwidthbandwidth_GB_per_nsintraHw nInterStepslatencyintraLatinterLat netOverhead latency_ns transport_nsrrr estimate_nccl_collective_runtimesf          ry) functoolsrPenumrr2rrrrrr virtualizedrr r lru_cacher#IRNoder-r4r<rCrDrHrKrXrYrVfloatryrrrrsL