§
    ŠŠtjWI  ã            	       ó  — U d dl Z d dlZd dlmZmZ d dlmZ d dlmZm	Z	 d dl
Z
d dlmZ d dlmZ d dlmZ d dlmZ d dlmZ d d	lmZ d
dlmZmZmZmZmZmZmZ d
dl m!Z!m"Z"m#Z#m$Z$m%Z%m&Z& d
dl'm(Z(m)Z) dgZ*e+e,         e-d<   de#de#de.fd„Z/de+e#         de.fd„Z0de+e#         de+e#         de+e#         fd„Z1de
j2        defd„Z3dedefd„Z4dedede$fd„Z5de,dede%fd„Z6de,dedede%fd „Z7de,de
j2        de%fd!„Z8de,d"efd#„Z9d$„ Z:d%„ Z;de,d&ed'e+e         de+e"         fd(„Z<d)ede#fd*„Z=de,d+ede+e%         fd,„Z>dedefd-„Z?de
j2        de+e         fd.„Z@de,d/ed0ede+e"         fd1„ZAd)eBe,ef         defd2„ZCd3ed4ed5ed6efd7„ZDdS )8é    N)Úbisect_rightÚinsort)ÚCallable)ÚAnyÚcast)Ú_get_device_module)ÚShardMetadata)ÚShardedTensor)ÚDTensor)Ú%compute_local_shape_and_global_offseté   )ÚBytesStorageMetadataÚChunkStorageMetadataÚMetadataIndexÚSTATE_DICT_TYPEÚSTORAGE_TYPESÚTensorPropertiesÚTensorStorageMetadata)ÚLoadItemTypeÚReadItemÚSavePlanÚTensorWriteDataÚ	WriteItemÚWriteItemType)Ú"_check_shard_metadata_pair_overlapÚ+_shards_get_overlap_region_wrt_saved_tensorÚ create_read_items_for_chunk_listÚ__all__ÚplanÚ
other_planÚreturnc                 óJ  — | j         |j         k    rdS t          | j        ¦  «        t          |j        ¦  «        k    rdS t          | j        |j        ¦  «        D ]É\  }}|j        |j        k    r dS |j        }|j        }|j        |j        k    s |j        |j        k    s|j        |j        k    r dS |j        }|j        }|r|r|s|r dS |rU|rS|j	        |j	        k    r dS |j
        }|j
        }	|r|	r|s|	r dS |r%|	r#|j        |	j        k    s|j        |	j        k    r dS ŒÊdS )a  
    Compare the two Save plans and return True if they are equal.

    Args:
        plan (SavePlan): First SavePlan to compare.
        other_plan (SavePlan): Second SavePlan to compare.

    Returns:
       True if the two plans are equal, False otherwise.
    FT)ÚusableÚlenÚitemsÚzipÚtypeÚindexÚfqnÚoffsetÚtensor_dataÚsizeÚchunkÚoffsetsÚsizes)
r   r    Ú	plan_itemÚother_plan_itemÚplan_metadata_indexÚother_plan_metadata_indexr+   Úother_tensor_datar-   Úother_chunks
             új/var/www/html/CA-Chatbot/venv/lib/python3.11/site-packages/torch/distributed/checkpoint/planner_helpers.pyÚ_compare_save_plansr7   *   s£  € ð „{�jÔ'Ò'Ð'Øˆuõ ˆ4Œ:�„�#˜jÔ.Ñ/Ô/Ò/Ð/Øˆuõ '*¨$¬*°jÔ6FÑ&GÔ&Gð )!ð )!Ñ"ˆ	�?àŒ>˜_Ô1Ò1Ð1Ø�5�5à'œoÐØ$3Ô$9Ð!ð  Ô#Ð'@Ô'DÒDÐDØ"Ô)Ð-FÔ-MÒMÐMØ"Ô(Ð,EÔ,KÒKÐKà�5�5ð  Ô+ˆØ+Ô7ÐØð 	Ð 1ð 	Øð	Ø 1ð	ð �5�5àð 	!Ð,ð 	!àÔÐ#4Ô#9Ò9Ð9Ø�u�uð  Ô%ˆEØ+Ô1ˆKØð ˜kð °5ð ¸[ð Ø�u�uð ð !˜ð !à”M [Ô%8Ò8Ð8Ø”{ kÔ&7Ò7Ð7à ˜5˜5øàˆ4ó    Údelta_plansc                 ó4   — t          d„ | D ¦   «         ¦  «        S )zæ
    Check if any delta plan is usable, indicating the plan has changed.

    Args:
        delta_plans (List[SavePlan]): A list of delta plans to check.
    Returns:
        True if any delta plan is usable, False otherwise.
    c              3   ó(   K  — | ]}|o|j         V — Œd S ©N)r#   )Ú.0Ú
delta_plans     r6   ú	<genexpr>z(_contains_usable_plan.<locals>.<genexpr>t   s*   è è € ÐNÐN°JˆzÐ/˜jÔ/ÐNÐNÐNÐNÐNÐNr8   )Úany)r9   s    r6   Ú_contains_usable_planrA   k   s!   € õ ÐNÐNÀ+ÐNÑNÔNÑNÔNÐNr8   Úcached_plansc                 óž   — g }t          | |¦  «        D ]9\  }}|r|j        s|                     |¦  «         Œ$|                     |¦  «         Œ:|S )ac  
    Merge a list of delta plans into a single plan.

    Args:
        cached_plans (List[SavePlan]): A list of cached plans.
        delta_plans (List[SavePlan]): A list of delta plans to merge. It can contain empty plans

    Returns:
        A single merged plan. If a delta plan is not usable, use the cached plan. Otherwise, use the delta plan.
    )r&   r#   Úappend)rB   r9   Úmerged_plansÚcached_planr>   s        r6   Ú_merge_delta_local_plansrG   w   sm   € ð €Lå#& |°[Ñ#AÔ#Að ,ð ,Ñˆ�ZØð 	,˜jÔ/ð 	,Ø×Ò Ñ,Ô,Ð,Ð,à×Ò 
Ñ+Ô+Ð+Ð+àÐr8   Útensorc           	      ó²   — t          t          j        dgt          |                      ¦   «         ¦  «        z  ¦  «        |                      ¦   «         ¬¦  «        S )Nr   ©r.   r/   )r   ÚtorchÚSizer$   r,   )rH   s    r6   Ú_create_chunk_from_tensorrM   �   sG   € ÝÝ”
˜A˜3¥ V§[¢[¡]¤]Ñ!3Ô!3Ñ3Ñ4Ô4¸F¿KºK¹M¼Mðñ ô ð r8   Úshard_mdc                 ó€   — t          t          j        | j        ¦  «        t          j        | j        ¦  «        ¬¦  «        S ©NrJ   )r   rK   rL   Úshard_offsetsÚshard_sizes)rN   s    r6   Ú_chunk_for_shardrS   –   s9   € ÝÝ”
˜8Ô1Ñ2Ô2ÝŒj˜Ô-Ñ.Ô.ðñ ô ð r8   Úsharded_tensorc                 óú   — |                       ¦   «         j        }t          |j        |j        |j        |j        |j        ¬¦  «        }t          t          |¦  «        ||                       ¦   «         j
        ¬¦  «        S )N)ÚdtypeÚlayoutÚrequires_gradÚmemory_formatÚ
pin_memory©r-   Ú
propertiesr,   )ÚmetadataÚtensor_propertiesr   rV   rW   rX   rY   rZ   r   rS   r,   )rT   rN   Úshard_propertiesr\   s       r6   Ú_sharded_tensor_metadatar`   �   s„   € ð &×.Ò.Ñ0Ô0ÔBÐå!ØÔ$ØÔ&Ø&Ô4Ø&Ô4Ø#Ô.ðñ ô €Jõ Ý˜xÑ(Ô(ØØ×$Ò$Ñ&Ô&Ô+ðñ ô ð r8   r)   c                 ó–  — t          |j        |j        |j        ¦  «        \  }}t	          j        |¦  «        t	          j        |¦  «        }}t          t          | |¦  «        t          j	        t          t          ||¬¦  «        t          j        |                     ¦   «         ¦  «        |                     ¦   «         ¬¦  «        ¬¦  «        S )NrJ   r[   ©r(   r'   r+   )r   ÚshapeÚdevice_meshÚ
placementsrK   rL   r   r   r   ÚSHARDr   r   r   Úcreate_from_tensorÚto_localr,   )r)   rH   r/   r.   s       r6   Ú_create_write_items_for_dtensorri   ±   s¼   € Ý:ØŒ�fÔ(¨&Ô*;ñô �N€Eˆ7õ ”Z Ñ&Ô&­¬
°7Ñ(;Ô(;ˆ7€EåÝ˜C Ñ)Ô)ÝÔ Ý#Ý&ØØðñ ô õ (Ô:¸6¿?º?Ñ;LÔ;LÑMÔMØ—’‘”ð
ñ 
ô 
ðñ ô ð r8   c                 ó¤   — t          j        |j        ¦  «        }t          t	          | |¦  «        t
          j        t          ||¦  «        ¬¦  «        S )Nrb   )rK   rL   rQ   r   r   r   rf   r`   )r)   rT   rN   r.   s       r6   Ú_create_write_item_for_shardrk   Å   sM   € õ Œj˜Ô/Ñ0Ô0€GÝÝ˜C Ñ)Ô)ÝÔ Ý,¨^¸XÑFÔFðñ ô ð r8   c                 ón  — t          j        dgt          |                     ¦   «         ¦  «        z  ¦  «        }t	          t          | |¦  «        t          j        t          t          ||                     ¦   «         ¬¦  «        t          j        |¦  «        |                     ¦   «         ¬¦  «        ¬¦  «        S )Nr   rJ   r[   rb   )rK   rL   r$   r,   r   r   r   ÚTENSORr   r   r   rg   )r)   rH   r.   s      r6   Ú_create_write_item_for_tensorrn   Ð   s•   € ÝŒj˜!˜�s 6§;¢;¡=¤=Ñ1Ô1Ñ1Ñ2Ô2€GÝÝ˜C Ñ)Ô)ÝÔ!Ý#Ý&¨w¸f¿kºk¹m¼mÐLÑLÔLÝ'Ô:¸6ÑBÔBØ—’‘”ð
ñ 
ô 
ðñ ô ð r8   Úbytesc                 óR   — t          t          | ¦  «        t          j        ¬¦  «        S )N)r(   r'   )r   r   r   ÚBYTE_IO)r)   ro   s     r6   Ú_create_write_item_for_bytesiorr   Ý   s*   € ÝÝ˜CÑ Ô ÝÔ"ðñ ô ð r8   c           
      ó²   — t          t          j        | t          j        |f¦  «        |t          j        |f¦  «        t          j        |f¦  «        ¬¦  «        S ©N)r'   Ú
dest_indexÚdest_offsetsÚstorage_indexÚstorage_offsetsÚlengths)r   r   rq   rK   rL   ©ru   Údest_offsetrw   Ústorage_offsetÚlengths        r6   Ú_create_read_item_for_byteior~   ä   sU   € õ ÝÔ!ØÝ”Z  Ñ/Ô/Ø#Ýœ
 NÐ#4Ñ5Ô5Ý”
˜F˜9Ñ%Ô%ðñ ô ð r8   c           
      ó¬   — t          t          j        | t          j        |¦  «        |t          j        |¦  «        t          j        |¦  «        ¬¦  «        S rt   )r   r   rm   rK   rL   ©ru   rv   rw   rx   ry   s        r6   Ú_create_read_item_for_tensorr�   ñ   sN   € õ ÝÔ ØÝ”Z Ñ-Ô-Ø#Ýœ
 ?Ñ3Ô3Ý”
˜7Ñ#Ô#ðñ ô ð r8   Úcheckpoint_mdÚlocal_chunksc                 ó  ‡‡‡ ‡!— g }|j         }|r|s|S t          |d         j        ¦  «        }dŠ!|dk    rLd}t          |¦  «        D ]:Št	          ˆfd„t          j        ||¦  «        D ¦   «         ¦  «        }||k    r|}‰Š!Œ;|dk    r'dgt          |¦  «        z  Š dgt          |¦  «        z  Šnˆ!fd„|D ¦   «         Š ˆ!fd„|D ¦   «         Št          t          t          |¦  «        ¦  «        ˆ fd„¬¦  «        }t          t          t          |¦  «        ¦  «        ˆfd	„¬¦  «        }	g }
d}t          |¦  «        }|	D �]B}||         }‰|         \  }}t          |
|d
f¦  «        }|r|
d|…= ||k     r?||         }||         }‰ |         \  }}||k    rnt          |
||f¦  «         |dz  }||k     °?|
D ]Ë\  }}||         }t          ||¦  «        sŒg }g }g }t          ||¬¦  «        D ]F\  }}}}|                     |¦  «         |                     |¦  «         |                     |¦  «         ŒG|                     t          t          | |j        |¦  «        |t          | |j        |¦  «        ||¬¦  «        ¦  «         ŒÌ�ŒD|S )aW  
    Create a list of ``ReadItem`` based on the checkpoint and local chunks.

    This applies the resharding algorithm and computes the reads needed
    to satisfy ``local_chunks`` with a checkpoint described by ``checkpoint_md``.

    Args:
        fqn (str) : The state_dict FQN to pass to ``ReadItem``.
        checkpoint_md (TensorStorageMetadata): metadata for a given tensor
            from a checkpoint.
        local_chunks (List[ChunkStorageMetadata]): Local chunks that needs to be
            loaded.

    Returns:
        A list of ``ReadItem`` that will satisfy all input chunks.
    r   r   c              3   óN   •K  — | ]}|j         ‰         |j        ‰         z   V — Œ d S r<   rJ   )r=   r-   Údims     €r6   r?   z3create_read_items_for_chunk_list.<locals>.<genexpr>   sJ   øè è € ð ð àð ”˜cÔ" U¤[°Ô%5Ñ5ðð ð ð ð ð r8   )r   r   c                 ó`   •— g | ]*}|j         ‰         |j         ‰         |j        ‰         z   f‘Œ+S © rJ   ©r=   ÚcÚ	sweep_dims     €r6   ú
<listcomp>z4create_read_items_for_chunk_list.<locals>.<listcomp>.  óH   ø€ ð 
ð 
ð 
àð ŒY�yÔ! 1¤9¨YÔ#7¸!¼'À)Ô:LÑ#LÐMð
ð 
ð 
r8   c                 ó`   •— g | ]*}|j         ‰         |j         ‰         |j        ‰         z   f‘Œ+S rˆ   rJ   r‰   s     €r6   rŒ   z4create_read_items_for_chunk_list.<locals>.<listcomp>2  r�   r8   c                 ó    •— ‰|          d         S ©Nr   rˆ   )ÚidxÚsaved_boundss    €r6   ú<lambda>z2create_read_items_for_chunk_list.<locals>.<lambda>9  ó   ø€ ˜ SÔ)¨!Ô,€ r8   )Úkeyc                 ó    •— ‰|          d         S r�   rˆ   )r‘   Úlocal_boundss    €r6   r“   z2create_read_items_for_chunk_list.<locals>.<lambda>=  r”   r8   éÿÿÿÿN)Úsaved_shardÚcurrent_shardr€   )Úchunksr$   r.   ÚrangeÚmaxÚ	itertoolsÚchainÚsortedr   r   r   r   rD   r�   r   )"r)   r‚   rƒ   Ú
read_itemsÚsaved_chunksÚnum_dimsÚmax_sizeÚdim_sizeÚsaved_sorted_indicesÚlocal_sorted_indicesÚactive_savedÚ	saved_ptrÚ	num_savedÚ	local_idxÚlocal_chunkÚlocal_startÚ	local_endÚcutoffÚstorage_idxÚstorage_chunkÚsaved_startÚ	saved_endÚ_rx   rv   ry   Ú_dimÚoffset_for_saved_tensorÚoffset_for_current_tensorr}   r†   r—   r’   r‹   s"                                 @@@@r6   r   r   þ   s«  øøøø€ ð* "$€JØ Ô'€Làð ˜|ð ØÐå�< ”?Ô*Ñ+Ô+€Hð €IØ�!‚|€|ØˆÝ˜‘?”?ð 	 ð 	 ˆCÝð ð ð ð å&œ_¨\¸<ÑHÔHðñ ô ñ ô ˆHð ˜(Ò"Ð"Ø#�Ø�	øð �1‚}€}Ø�x¥# lÑ"3Ô"3Ñ3ˆØ�x¥# lÑ"3Ô"3Ñ3ˆˆð
ð 
ð 
ð 
à!ð
ñ 
ô 
ˆð
ð 
ð 
ð 
à!ð
ñ 
ô 
ˆõ
 "Ý�c�,ÑÔÑ Ô Ø,Ð,Ð,Ð,ðñ ô Ðõ "Ý�c�,ÑÔÑ Ô Ø,Ð,Ð,Ð,ðñ ô Ðð
 +-€LØ€IÝÐ(Ñ)Ô)€Ià)ð 1ñ 1ˆ	Ø" 9Ô-ˆØ!-¨iÔ!8Ñˆ�Yå˜l¨[¸"Ð,=Ñ>Ô>ˆØð 	&Ø˜W˜f˜WÐ%à˜)Ò#Ð#Ø.¨yÔ9ˆKØ(¨Ô5ˆMØ%1°+Ô%>Ñ"ˆK˜à˜iÒ'Ð'Øå�< )¨[Ð!9Ñ:Ô:Ð:Ø˜‰NˆIð ˜)Ò#Ð#ð +ð 	ð 	‰NˆAˆ{Ø(¨Ô5ˆMÝ5°kÀ=ÑQÔQð Øà ˆOØˆLØˆGõ =Ø)¸ðñ ô ð
'ð 
'ñ ØØ'Ø)Øð  ×&Ò&Ð'>Ñ?Ô?Ð?Ø×#Ò#Ð$=Ñ>Ô>Ð>Ø—’˜vÑ&Ô&Ð&Ð&à×ÒÝ,Ý,¨S°+Ô2EÀyÑQÔQØ!-Ý"/Ø˜]Ô2°Kñ#ô #ð %4Ø#ðñ ô ñ
ô 
ð 
ð 
ñ)	ð> Ðr8   Ú
state_dictc                 ó.  ‡‡— g }|                       ¦   «         D ]î\  ŠŠt          ‰t          ¦  «        r$|                     t	          ‰‰¦  «        ¦  «         Œ>t          ‰t
          ¦  «        r:|                     ˆˆfd„‰                     ¦   «         j        D ¦   «         ¦  «         Œ�t          ‰t          j
        ¦  «        r$|                     t          ‰‰¦  «        ¦  «         ŒË|                     t          ‰‰¦  «        ¦  «         Œït          |¦  «        S )Nc              3   ó:   •K  — | ]}t          ‰‰|¦  «        V — Œd S r<   )rk   )r=   rN   r)   Úobjs     €€r6   r?   z5_create_default_metadata_only_plan.<locals>.<genexpr>  sC   øè è € ð ð àõ -¨S°#°xÑ@Ô@ðð ð ð ð ð r8   )r%   Ú
isinstancer   rD   ri   r
   Úextendr]   Úshards_metadatarK   ÚTensorrn   rr   r   )r¸   Úrequestsr)   r»   s     @@r6   Ú"_create_default_metadata_only_planrÁ   y  s&  øø€ Ø€HØ×$Ò$Ñ&Ô&ð Fð F‰ˆˆSÝ�c�7Ñ#Ô#ð 
	FØ�OŠOÕ;¸CÀÑEÔEÑFÔFÐFÐFÝ˜�]Ñ+Ô+ð 	FØ�OŠOð ð ð ð ð à #§¢¡¤Ô >ðñ ô ñ ô ð ð õ ˜�Uœ\Ñ*Ô*ð 	FØ�OŠOÕ9¸#¸sÑCÔCÑDÔDÐDÐDà�OŠOÕ:¸3ÀÑDÔDÑEÔEÐEÐEÝ�HÑÔÐr8   Úobjectc                 ó6  ‡ ‡— t          ‰d¦  «        r‰                     ‰ ‰¦  «        S t          ‰t          ¦  «        r!ˆ ˆfd„‰                     ¦   «         D ¦   «         S t          ‰t
          j        ¦  «        rt          ‰ ‰¦  «        gS t          ‰ ‰¦  «        gS )NÚ__create_write_items__c                 ó<   •— g | ]}t          ‰‰|j        ¦  «        ‘ŒS rˆ   )rk   r]   )r=   Úshardr)   rÂ   s     €€r6   rŒ   z'_create_write_items.<locals>.<listcomp>�  s7   ø€ ð 
ð 
ð 
àõ )¨¨f°e´nÑEÔEð
ð 
ð 
r8   )	ÚhasattrrÄ   r¼   r
   Úlocal_shardsrK   r¿   rn   rr   )r)   rÂ   s   ``r6   Ú_create_write_itemsrÉ   Š  s·   øø€ ÝˆvÐ/Ñ0Ô0ð =à×,Ò,¨S°&Ñ9Ô9Ð9Ý	�F�MÑ	*Ô	*ð =ð
ð 
ð 
ð 
ð 
à×,Ò,Ñ.Ô.ð
ñ 
ô 
ð 	
õ 
�F�EœLÑ	)Ô	)ð =Ý-¨c°6Ñ:Ô:Ð;Ð;å.¨s°FÑ;Ô;Ð<Ð<r8   c                 óº   — t          | j        | j        | j        ¦  «        \  }}t	          j        |¦  «        t	          j        |¦  «        }}t          ||¬¦  «        S rP   )r   rc   rd   re   rK   rL   r   )rH   r/   r.   s      r6   Ú_create_chunk_from_dtensorrË   ™  s`   € Ý:ØŒ�fÔ(¨&Ô*;ñô �N€Eˆ7õ ”Z Ñ&Ô&­¬
°7Ñ(;Ô(;ˆ7€EÝØØðñ ô ð r8   c                 óL  — t          | d¦  «        r|                      ¦   «         }n~t          | t          ¦  «        rd„ |                      ¦   «         D ¦   «         }nJt          | t
          j        ¦  «        rt          | ¦  «        g}nt          dt          | ¦  «        › �¦  «        ‚|S )NÚ__create_chunk_list__c                 ó6   — g | ]}t          |j        ¦  «        ‘ŒS rˆ   )rS   r]   )r=   rÆ   s     r6   rŒ   z&_create_chunk_list.<locals>.<listcomp>©  s0   € ð 
ð 
ð 
Ø16Õ˜Uœ^Ñ,Ô,ð
ð 
ð 
r8   zMUnsupported Type, expecting one of [Tensor, DTensor, ShardedTensor] ,but got )
rÇ   rÍ   r¼   r
   rÈ   rK   r¿   rM   Ú
ValueErrorr'   )rH   rƒ   s     r6   Ú_create_chunk_listrÐ   ¤  sÀ   € ÝˆvÐ.Ñ/Ô/ð 
à×3Ò3Ñ5Ô5ˆˆÝ	�F�MÑ	*Ô	*ð 

ð
ð 
Ø:@×:MÒ:MÑ:OÔ:Oð
ñ 
ô 
ˆˆõ 
�F�EœLÑ	)Ô	)ð 
Ý1°&Ñ9Ô9Ð:ˆˆåð'Ý˜V™œð'ð 'ñ
ô 
ð 	
ð
 Ðr8   Úmdr»   c                 ó@  — t          |t          ¦  «        s[	 t          |¦  «        }n9# t          $ r,}t          d| › d�dt	          |¦  «        › �z   ¦  «        |‚d }~ww xY wt          | ||¦  «        S t          t          | ¦  «        dt          | ¦  «        dd¬¦  «        gS )Nz Invalid checkpoint metadata for z, z(expected BytesStorageMetadata but found r   rz   )r¼   r   rÐ   rÏ   r'   r   r~   r   )r)   rÑ   r»   rƒ   Úexs        r6   Ú_create_read_itemsrÔ   ·  sÕ   € Ý�bÕ.Ñ/Ô/ð 
ð	Ý-¨cÑ2Ô2ˆLˆLøÝð 	ð 	ð 	ÝØ:°3Ð:Ð:Ð:ØG½TÀ"¹X¼XÐGÐGñHñô ð ðøøøøð	øøøõ 0°°R¸ÑFÔFÐFõ )Ý(¨Ñ-Ô-ØÝ+¨CÑ0Ô0Ø Øðñ ô ð
ð 	
s   —' §
A±'AÁAc                 óv   — dt           fd„}dt          fd„}dt          j        fd„}t	          | |||¦  «         dS )zP
    Initializes meta tensor if the meta tensor is DTensor or torch.Tensor.
    Úvaluec                 óò  — t          | dd ¦  «        }|t          j        d¦  «        k    rÍt          j                             ¦   «         j        }t          t          j        t          |¦  «         	                    ¦   «         ¦  «        }t          j
        |                      ¦   «         |¬¦  «        }t          j        || j        | j        |                      ¦   «         |                      ¦   «         ¬¦  «        }|S | S )NÚdeviceÚmeta©rØ   )rd   re   rc   Ústride)ÚgetattrrK   rØ   ÚdistÚdistributed_c10dÚ_get_pg_default_devicer'   r   r   Úcurrent_deviceÚ
empty_likerh   r   Ú
from_localrd   re   r,   rÛ   )rÖ   rØ   Údevice_typeÚnew_local_tensorÚdtensors        r6   Údtensor_funcz&_init_state_dict.<locals>.dtensor_funcÓ  sÖ   € Ý˜ ¨$Ñ/Ô/ˆØ•U”\ &Ñ)Ô)Ò)Ð)ÝÔ/×FÒFÑHÔHÔMˆKÝÝ”Õ0°Ñ=Ô=×LÒLÑNÔNñô ˆFõ  %Ô/°·²Ñ0@Ô0@ÈÐPÑPÔPÐõ Ô(Ø Ø!Ô-Ø Ô+Ø—j’j‘l”lØ—|’|‘~”~ðñ ô ˆGð ˆNàˆLr8   c                 ó˜   — t          | dd ¦  «        }|t          j        d¦  «        k    r t          dt	          | ¦  «        › d�¦  «        ‚| S )NrØ   rÙ   zFound unsupported type z for meta device loading.)rÜ   rK   rØ   ÚRuntimeErrorr'   )rÖ   rØ   s     r6   Úsharded_tensor_funcz-_init_state_dict.<locals>.sharded_tensor_funcè  sU   € Ý˜ ¨$Ñ/Ô/ˆØ•U”\ &Ñ)Ô)Ò)Ð)ÝØP­$¨u©+¬+ÐPÐPÐPñô ð ð ˆLr8   c                 ó@  — t          | dd ¦  «        }|t          j        d¦  «        k    rtt          j                             ¦   «         j        }t          t          j        t          |¦  «         	                    ¦   «         ¦  «        }t          j
        | |¬¦  «        }|S | S )NrØ   rÙ   rÚ   )rÜ   rK   rØ   rÝ   rÞ   rß   r'   r   r   rà   rá   )rÖ   rØ   rã   rH   s       r6   Útensor_funcz%_init_state_dict.<locals>.tensor_funcñ  s‹   € Ý˜ ¨$Ñ/Ô/ˆØ•U”\ &Ñ)Ô)Ò)Ð)ÝÔ/×FÒFÑHÔHÔMˆKÝÝ”Õ0°Ñ=Ô=×LÒLÑNÔNñô ˆFõ Ô% e°FÐ;Ñ;Ô;ˆFØˆMàˆLr8   N)r   r   rK   r¿   Ú_iterate_state_dict)r¸   ræ   ré   rë   s       r6   Ú_init_state_dictrí   Î  s   € ð
�Gð ð ð ð ð*¥3ð ð ð ð ð
�5œ<ð 
ð 
ð 
ð 
õ ØØØØñ	ô ð ð ð r8   Úiter_objectræ   ré   rë   c                 ót  ‡‡‡— t          | t          ¦  «        r ‰| ¦  «        S t          | t          ¦  «        r ‰| ¦  «        S t          | t          j        ¦  «        r ‰| ¦  «        S t          | t
          t          t          t          t          j
        f¦  «        s| €| S t          | t          ¦  «        r1|                      ¦   «         D ]\  }}t          |‰‰‰¦  «        | |<   Œ| S t          | t          t          f¦  «        r6ˆˆˆfd„| D ¦   «         }t          | t          ¦  «        rt          |¦  «        }|S dS )a$  
    Iterate through the state dict, applying the given functions to each tensor type
    and update the state dict in place.

    Args:
        iter_object (Any): the target state_dict.
        sharded_tensor_func (Callable): the function to apply to ShardedTensor
        dtensor_func (Callable): the function to apply to DTensor
        tensor_func (Callable): the function to apply to Tensor

    # TODO: let state_dict_util._iterate_state_dict() to support in place option
    so we don't need to have two versions of _iterate_state_dict.
    Nc                 ó4   •— g | ]}t          |‰‰‰¦  «        ‘ŒS rˆ   )rì   )r=   Úvræ   ré   rë   s     €€€r6   rŒ   z'_iterate_state_dict.<locals>.<listcomp>+  s8   ø€ ð 
ð 
ð 
àõ    <Ð1DÀkÑRÔRð
ð 
ð 
r8   )r¼   r   r
   rK   r¿   ÚintÚfloatÚstrro   ÚioÚBytesIOÚdictr%   rì   ÚlistÚtuple)rî   ræ   ré   rë   r•   rÖ   Úrets    ```   r6   rì   rì     sk  øøø€ õ( �+�wÑ'Ô'ð Øˆ|˜KÑ(Ô(Ð(Ý	�K¥Ñ	/Ô	/ð Ø"Ð" ;Ñ/Ô/Ð/Ý	�K¥¤Ñ	.Ô	.ð Øˆ{˜;Ñ'Ô'Ð'å�;¥¥e­Sµ%½¼Ð DÑEÔEðàÐàÐÝ	�K¥Ñ	&Ô	&ð Ø%×+Ò+Ñ-Ô-ð 	ð 	‰JˆC�Ý2Ø�|Ð%8¸+ñ ô  ˆK˜ÑÐð ÐÝ	�K¥$­ Ñ	/Ô	/ð ð
ð 
ð 
ð 
ð 
ð 
à ð
ñ 
ô 
ˆõ �k¥5Ñ)Ô)ð 	Ý˜‘*”*ˆCØˆ
ðð r8   )Erõ   rž   Úbisectr   r   Úcollections.abcr   Útypingr   r   rK   Útorch.distributedÚdistributedrÝ   Útorch._utilsr   Ú!torch.distributed._shard.metadatar	   Ú'torch.distributed._shard.sharded_tensorr
   Útorch.distributed.tensorr   Útorch.distributed.tensor._utilsr   r]   r   r   r   r   r   r   r   Úplannerr   r   r   r   r   r   Ú
reshardingr   r   r   rø   rô   Ú__annotations__Úboolr7   rA   rG   r¿   rM   rS   r`   ri   rk   rn   rr   r~   r�   r   rÁ   rÉ   rË   rÐ   rÔ   r÷   rí   rì   rˆ   r8   r6   ú<module>r	     s  ðà 	€	€	€	€	Ø Ð Ð Ð Ø 'Ð 'Ð 'Ð 'Ð 'Ð 'Ð 'Ð 'Ø $Ð $Ð $Ð $Ð $Ð $Ø Ð Ð Ð Ð Ð Ð Ð à €€€Ø  Ð  Ð  Ð  Ð  Ð  Ø +Ð +Ð +Ð +Ð +Ð +Ø ;Ð ;Ð ;Ð ;Ð ;Ð ;Ø AÐ AÐ AÐ AÐ AÐ AØ ,Ð ,Ð ,Ð ,Ð ,Ð ,Ø QÐ QÐ QÐ QÐ QÐ Qðð ð ð ð ð ð ð ð ð ð ð ð ð ð ð ð ð ðð ð ð ð ð ð ð ð ð ð ð ð ð ð ð ðð ð ð ð ð ð ð ð 9Ð9€ˆˆcŒÐ 9Ð 9Ñ 9ð>˜hð >°Hð >Àð >ð >ð >ð >ðB	O t¨H¤~ð 	O¸$ð 	Oð 	Oð 	Oð 	OðØ�x”.ðà�h”ðð 
ˆ(„^ðð ð ð ð2 e¤lð Ð7Kð ð ð ð ð˜}ð Ð1Eð ð ð ð ðØ!ðØ-:ðàðð ð ð ð(¨ð °gð À)ð ð ð ð ð(Ø	ðØ+ðØ7Dðàðð ð ð ð
 sð 
°E´Lð 
ÀYð 
ð 
ð 
ð 
ð¨ð °Cð ð ð ð ð
ð 
ð 
ð
ð 
ð 
ðxØ	ðxà(ðxð Ð+Ô,ðxð 
ˆ(„^ð	xð xð xð xðv°?ð Àxð ð ð ð ð"=˜Sð =¨#ð =°$°y´/ð =ð =ð =ð =ð wð Ð3Gð ð ð ð ð˜uœ|ð °Ð5IÔ0Jð ð ð ð ð&
˜Cð 
 ]ð 
¸ð 
ÀÀhÄð 
ð 
ð 
ð 
ð.4  c¨3 h¤ð 4°Cð 4ð 4ð 4ð 4ðn,Øð,àð,ð "ð,ð ð	,ð ,ð ,ð ,ð ,ð ,r8   