o 0i">@s~ddlZddlZddlmZddlmZmZmZmZddl m Z m Z m Z m Z ddlmZer4ddlmZeGdddZdS) N) dataclass) TYPE_CHECKINGLiteralOptionalUnion)DeepSpeedSequenceParallelConfigDistributedTypeTorchContextParallelConfigTorchTensorParallelConfig)is_torch_version) Acceleratorc@seZdZUdZdZeeed<dZeeed<dZ eeed<dZ eeed<dZ e ded<dZ eeed <dZe d ed <dZedefed <dZedefed <dZedefed<dZddZddZeddZeddZeddZeddZeddZeddZedd Z ed!d"Z!ed#d$Z"ed%d&Z#ed'd(Z$ed)d*Z%ed+d,Z&ed-d.Z'd/e(fd0d1Z)dBd/ee(fd2d3Z*d4e+e+ed5fe+e(d5fffd6d7Z,d8d9Z-d:e(d;efd 1` and `tp_size > 1`, `dp_shard_size` must also be greater than 1, as composing DDP + TP is currently not supported. tp_size (`int`, defaults to `1`): The size of the tensor parallel group. If `tp_size` is set to `1`, the tensor parallel group will not be used. tp_handler (`~utils.TorchTensorParallelConfig`, defaults to `None`): The handler for the tensor parallel group. cp_size (`int`, defaults to `1`): The size of the context parallel group. Currently not supported, but reserved for future use and enabled for downstream libraries. cp_backend (`str`, defaults to `torch`): Which CP backend to use: `torch` (FSDP2) sp_size (`int`, defaults to `1`): The size of the sequence parallel group. sp_backend (`str`, defaults to `deepspeed`): Which SP backend to use:`deepspeed` (ALST/Ulysses) You may obtain different distributed data parallel paradigms by configuring `dp_replicate_size` and `dp_shard_size` together: - `dp_replicate_size == 1` and `dp_shard_size > 1`, we obtain Fully Sharded Data Parallel (FSDP). - `dp_replicate_size > 1` and `dp_shard_size > 1`, we obtain Hybrid Sharded Data Parallel (HSDP). - `dp_replicate_size > 1` and `dp_shard_size == 1` is an invalid configuration, to use pure DP, use `DistributedDataParallelKwargs` instead. Ndp_replicate_size dp_shard_sizetp_sizecp_sizetorch cp_backendsp_size deepspeed sp_backend tp_handler cp_handler sp_handlercCsVd|jd|jd|jd|jd|jd|jd|jd|jd |jd |j d S) Nz'ParallelismConfig( dp_replicate_size=z, dp_shard_size=z , tp_size=z , cp_size=z, cp_backend=z , sp_size=z, sp_backend=z, total_size=z tp_handler=z, cp_handler=z) ) rrrrrrr total_sizerrselfr_/sda-disk/www/egybert/egybert_env/lib/python3.10/site-packages/accelerate/parallelism_config.py__repr__Us,  zParallelismConfig.__repr__cs2ddldgfdd|jDdS)Nr device_meshcs4i|]\}}|vr|t|dr|jn|qS)__dict__)hasattrdeepcopyr!).0kv_non_serializable_fieldscopyrr js z-ParallelismConfig.to_json..)r)r#r!itemsrrr'rto_jsonds zParallelismConfig.to_jsoncC(g}|jr |dg7}|jr|dg7}|S)zENames of enabled dimensions across which data parallelism is applied. dp_replicatedp_shard)dp_replicate_enableddp_shard_enabledrdimsrrr dp_dim_namesq   zParallelismConfig.dp_dim_namescC8g}|jr |dg7}|jr|dg7}|jr|dg7}|S)z]Names of enabled dimensions which will receive the same batch (non-data parallel dimensions).tpcpsp) tp_enabled cp_enabled sp_enabledr2rrrnon_dp_dim_names{   z"ParallelismConfig.non_dp_dim_namescCr-)zlNames of enabled dimensions which will be flattened into a joint mesh across which is model sharded in FSDP.r/r8)r1r;r2rrrdp_shard_cp_dim_namesr5z'ParallelismConfig.dp_shard_cp_dim_namescCr6)z@Names of enabled dimensions across which loss should be averagedr.r/r8)r0r1r;r2rrrdp_cp_dim_namesr>z!ParallelismConfig.dp_cp_dim_namescCs"g}|jr |dg7}|dg7}|S)z^Names of enabled dimensions across which FSDP is applied, including data parallel replication.r. dp_shard_cp)r0r2rrrfsdp_dim_namess   z ParallelismConfig.fsdp_dim_namescCs|j|j|j|j|jS)zSThe total size of the parallelism configuration, which is the product of all sizes.)rrrrrrrrrrszParallelismConfig.total_sizecCs|j|j|jS)zhThe size of the non-data parallel dimensions, which is the product of tensor and context parallel sizes.)rrrrrrrnon_data_parallel_sizesz(ParallelismConfig.non_data_parallel_sizecCs |j|jS)z_The size of the data parallel dimensions, which is the product of data parallel replication and)rrrrrrdata_parallel_size z$ParallelismConfig.data_parallel_sizecC |jdkS)zKTrue if data parallel replication is enabled, i.e. `dp_replicate_size > 1`.)rrrrrr0 z&ParallelismConfig.dp_replicate_enabledcCrF)zDTrue if data parallel sharding is enabled, i.e. `dp_shard_size > 1`.rG)rrrrrr1rHz"ParallelismConfig.dp_shard_enabledcCrF)z:True if tensor parallelism is enabled, i.e. `tp_size > 1`.rG)rrrrrr:rHzParallelismConfig.tp_enabledcCrF)z;True if context parallelism is enabled, i.e. `cp_size > 1`.rG)rrrrrr;rHzParallelismConfig.cp_enabledcCrF)z;True if context parallelism is enabled, i.e. `sp_size > 1`.rG)rrrrrr<rHzParallelismConfig.sp_enabledcCs |j|jS)z$Names of all active mesh dimensions.)r4r=rrrractive_mesh_dimsrEz"ParallelismConfig.active_mesh_dims device_typecCstddr ddlm}ntd|}t|dkrdS|\}}||||d}|jr2||jd|jr=||jd |j rH||j d |S) a!Builds a device mesh for the given device type based on the parallelism configuration. This method will also create required joint meshes (e.g. `dp_shard_cp`, `dp_cp`, `dp`). Args: device_type (`str`): The type of device for which to build the mesh, e z>=z2.2.0r)init_device_meshz4Building a device_mesh requires to have torch>=2.2.0N)mesh_dim_namesdprAdp_cp) r torch.distributed.device_meshrK RuntimeError _get_meshlenr4_flattenr?r@)rrJrKmeshrL mesh_shaper rrrbuild_device_meshs&  z#ParallelismConfig.build_device_meshcCs\|jdur|dur|||_|jSd|dur+|jj|kr+td|jjd|d|jS)Nz@You need to pass a device_type e.g cuda to build the device meshz4The device_mesh is already created with device type z@. However, you are trying to get a device mesh with device_type z<. Please check if you correctly initialized your device_mesh)r rVrJ ValueError)rrJrrrget_device_meshs    z!ParallelismConfig.get_device_meshreturn.cs@fddjD}gdt|fddd}tt|S)zQGenerate mesh shape and dimension names for torch.distributed.init_device_mesh().csi|]}|j|qSr)_sizes)r$ parallelismrrrr*sz/ParallelismConfig._get_mesh..)r.r/r8r9r7cs|dS)Nr)index)x) mesh_orderrr sz-ParallelismConfig._get_mesh..)key)rIsortedr+tuplezip)r mesh_dims sorted_itemsr)r^rrrQs  zParallelismConfig._get_meshcCs|jdurttjdd|_|jdurttjdd|_|jdur-ttjdd|_|jdur 1) cannot be used with pure data parallelism (dp_replicate_size > 1 and dp_shard_size == 1). Please set dp_shard_size > 1 and dp_replicate_size == 1 to compose FSDP + TP/CP for 2D parallel, or set dp_replicate_size == 1 and dp_shard_size > 1 to compose HSDP + TP/CP for 3D parallel.)r.r/r7r8r9)rintosenvirongetrrrrrrrr rr dict isinstancerWtyperrrZ)rcp_backends_config_mapvalid_cp_backendsvalid_sp_backendsrrr __post_init__ sp            "         0 zParallelismConfig.__post_init__r[sizecCsB||jvsJd|j||j|<t||d|dS)NzParallelism must be one of _size)rZkeyssetattr)rr[ryrrr _set_sizeQs" zParallelismConfig._set_size acceleratorr cCst}|js |jdkr dS|jdkr|d|j|j|jkr,td|jd|jd|jdkrF|jsF|jsF|jtj ksFtd|jd|j D]\}}|dkrjt ||dddurj| d |d |d qK|r}|jrtd d |tdSdSdS)NrGr.zParallelismConfig total_size (z ) does not match num_processes (zJ). Please adjust dp_replicate_size/ dp_shard_size/tp_size/cp_size/sp_size.zParallelismConfig is only compatible DistributedType.FSDP (version 2) or DistributedType.Multi{Device} or DistributedType.DEEPSPEED, but got ._handlerzParallelismConfig.z_handler is set, but z0_size is set to 1. This handler will be ignored.z.ParallelismConfig has the following warnings:  )set multi_devicerr} num_processesrWis_fsdp2distributed_typer DEEPSPEEDrZr+getattraddis_main_processwarningswarnjoin UserWarning)rr~ _warningsr[ryrrr_validate_acceleratorVs@        z'ParallelismConfig._validate_accelerator)N)r~r )0__name__ __module__ __qualname____doc__rrrn__annotations__rrrrrrrrrr rr rrr rr,propertyr4r=r?r@rBrrCrDr0r1r:r;r<rIstrrVrXrbrQrxr}rrrrrr !sb #              &Dr )ror dataclassesrtypingrrrraccelerate.utils.dataclassesrrr r accelerate.utils.versionsr accelerater r rrrrs