U
    0-eQ                     @   s@   d dl mZ ddlmZ e r*d dlmZ d	ddZdd ZdS )
    N   )is_fp8_availableTc           	      C   s  t  std|  D ]\}}t|tjr|r|rtdd |jjD rP dS |j	dk	}t
j|j|j||jjd}|jj |j_|r|j	j |j	_t| || qt|tjr
|r
|r
t
j|jd |j|jjd}|jj |j_|j	j |j	_t| || qt|t
jr||s||r||j	dk	}tj|j|j||jjd}|jj |j_|rn|j	j |j	_t| || qt|t
jr|s|rtj|jd |j|jjd}|jj |j_|j	j |j	_t| || qt||||d qdS )	zu
    Recursively converts the linear and layernorm layers of a model to their `transformers_engine` counterpart.
    zBUsing `convert_model` requires transformer_engine to be installed.c                 s   s   | ]}|d  dkV  qdS )   r   N ).0pr   r   d/var/www/html/Darija-Ai-Train/env/lib/python3.8/site-packages/accelerate/utils/transformer_engine.py	<genexpr>!   s     z convert_model.<locals>.<genexpr>N)biasparams_dtyper   )epsr   )to_transformer_engine_convert_linear_convert_ln)r   ImportErrorZnamed_children
isinstancennLinearanyweightshaper
   teZin_featuresZout_featuresZdtypedataclonesetattr	LayerNormZnormalized_shaper   convert_model)	modelr   r   r   namemoduleZhas_biasZ	te_moduleZ
new_moduler   r   r   r      sZ    
   
   r   c                 C   s8   t  std|  D ]}t|tjtjfr dS qdS )zS
    Returns whether a given model has some `transformer_engine` layer or not.
    zRUsing `has_transformer_engine_layers` requires transformer_engine to be installed.TF)r   r   modulesr   r   r   r   )r   mr   r   r   has_transformer_engine_layersK   s    r"   )TTT)	Ztorch.nnr   Zimportsr   Ztransformer_engine.pytorchZpytorchr   r   r"   r   r   r   r   <module>   s
   
3