[jG PddlZddlZddlmZmZeGddZdS)N)asdict dataclassceZdZUdZdZedzed<dZee e fdzed<dZ e ed<dZ e ed<dZ edzed<dZe ed <dZe ed <d Zed ed dfdZd efdZd e fdZde ejzfdZdZdS)DistributedConfigaK Configuration for native distributed training (FSDP2 + TP). Args: tp_size (`int`, *optional*): Number of devices for tensor parallelism. If `None` and `fsdp_size` is set, defaults to 1. tp_plan (`dict`, *optional*): Tensor parallel sharding plan. Leave as `None` to use the model's `base_model_tp_plan`. Set explicitly to override. enable_sequence_parallel (`bool`, *optional*, defaults to `False`): Reserved for sequence parallelism. Not wired up yet. enable_expert_parallel (`bool`, *optional*, defaults to `False`): Route MoE models through the expert-parallel path (``base_model_ep_plan``). fsdp_size (`int`, *optional*): Number of devices for FSDP (data parallelism). If `None` and `tp_size` is set, defaults to 1. fsdp_cpu_offload (`bool`, *optional*, defaults to `False`): Whether to enable CPU offloading for FSDP2. fsdp_mixed_precision (`bool`, *optional*, defaults to `False`): Whether to enable mixed precision for FSDP2. Ntp_sizetp_planFenable_sequence_parallelenable_expert_parallel fsdp_sizefsdp_cpu_offloadfsdp_mixed_precisionc|j |jdS|jd|_|jd|_|jdkr|jdkrtddSdS)NzFSDP+TP is not supported yet. Use DistributedConfig(fsdp_size=N) or DistributedConfig(tp_size=N), not both. 2D support will come soon.)rr ValueErrorselfs ZC:\Projects\DSCons\.venv\Lib\site-packages\transformers\distributed\configuration_utils.py __post_init__zDistributedConfig.__post_init__3sr < DN$: F < DL > !DN z.DistributedConfig.from_dict..FsHHHafHHHrc$i|] \}}|v || Srr)rkv valid_keyss r z/DistributedConfig.from_dict..Gs$IIItq!jarr)__dataclass_fields__valuesitems)clsrkwargsmergedr"s @r from_dictzDistributedConfig.from_dictCsg*K*6*HHc&>&E&E&G&GHHH sJJIIIIv||~~IIIJJJrc t|S)N)rrs rto_dictzDistributedConfig.to_dictIsd||rcXtj|ddzS)N)indent )jsondumpsr,rs rto_json_stringz DistributedConfig.to_json_stringLs$z$,,..333d::rjson_file_pathct|dd5}||ddddS#1swxYwYdS)Nwzutf-8)encoding)openwriter3)rr4rs r to_json_filezDistributedConfig.to_json_fileOs .# 8 8 8 +A GGD'')) * * * + + + + + + + + + + + + + + + + + +s(AA A cH|jjd|S)N ) __class____name__r3rs r__repr__zDistributedConfig.__repr__Ss'.)CCD,?,?,A,ACCCr)r> __module__ __qualname____doc__rint__annotations__rdictstrr boolr r r r r classmethodr*r,r3osPathLiker:r?rrrrrs[*GS4Z%)GT#s(^d ")))%*d***#(D((( IsTz   "d"""!&$&&& KDK7JKKK[K ;;;;;+3+<++++DDDDDrr)r1rI dataclassesrrrrrrrLsy )))))))) ?D?D?D?D?D?D?D ?D?D?Dr