U 0-e@sdZddlZddlZddlZddlZddlZddlZddlZddlm Z ddl m Z m Z ddl mZddlmZmZmZmZmZmZmZddlZddlmZmZmZdd lmZdd lmZGd d d Ze Gd ddeZ e GdddeZ!e GdddeZ"e GdddeZ#e GdddeZ$Gddde%ej&Z'Gddde%ej&Z(Gddde%ej&Z)Gddde%ej&Z*Gdd d ej+Z,Gd!d"d"ej&e,d#Z-Gd$d%d%e-Z.Gd&d'd'e-Z/Gd(d)d)e-Z0Gd*d+d+ej&Z1e Gd,d-d-Z2e Gd.d/d/Z3e Gd0d1d1eZ4e Gd2d3d3eZ5e Gd4d5d5Z6e Gd6d7d7Z7e Gd8d9d9Z8e Gd:d;d;Z9dS)8s z+KwargsHandler.to_kwargs..)otherr __class__ritems)rrZ this_dictrr"r to_kwargs.s  zKwargsHandler.to_kwargsN)__name__ __module__ __qualname____doc__rr'rrrrr&src@s*eZdZUdZdZeed<dZeed<dS)AutocastKwargsa Use this object in your [`Accelerator`] to customize how `torch.autocast` behaves. Please refer to the documentation of this [context manager](https://pytorch.org/docs/stable/amp.html#torch.autocast) for more information on each argument. Example: ```python from accelerate import Accelerator from accelerate.utils import AutocastKwargs kwargs = AutocastKwargs(cache_enabled=True) accelerator = Accelerator(kwargs_handlers=[kwargs]) ``` TenabledN cache_enabled)r(r)r*r+r-bool__annotations__r.rrrrr,;s  r,c@sfeZdZUdZdZeed<dZeed<dZ eed<dZ eed <dZ eed <dZ eed <dZ eed <d S)DistributedDataParallelKwargsa Use this object in your [`Accelerator`] to customize how your model is wrapped in a `torch.nn.parallel.DistributedDataParallel`. Please refer to the documentation of this [wrapper](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html) for more information on each argument. `gradient_as_bucket_view` is only available in PyTorch 1.7.0 and later versions. `static_graph` is only available in PyTorch 1.11.0 and later versions. Example: ```python from accelerate import Accelerator from accelerate.utils import DistributedDataParallelKwargs kwargs = DistributedDataParallelKwargs(find_unused_parameters=True) accelerator = Accelerator(kwargs_handlers=[kwargs]) ``` rdimTbroadcast_buffers bucket_cap_mbFfind_unused_parameterscheck_reductiongradient_as_bucket_view static_graphN)r(r)r*r+r2intr0r3r/r5r6r7r8r9rrrrr1Qs       r1c@sNeZdZUdZdZeed<dZeed<dZeed<dZ e ed <d Z e ed <d S) GradScalerKwargsa Use this object in your [`Accelerator`] to customize the behavior of mixed precision, specifically how the `torch.cuda.amp.GradScaler` used is created. Please refer to the documentation of this [scaler](https://pytorch.org/docs/stable/amp.html?highlight=gradscaler) for more information on each argument. `GradScaler` is only available in PyTorch 1.5.0 and later versions. Example: ```python from accelerate import Accelerator from accelerate.utils import GradScalerKwargs kwargs = GradScalerKwargs(backoff_filter=0.25) accelerator = Accelerator(kwargs_handlers=[kwargs]) ``` g@ init_scaleg@ growth_factorg?backoff_factorigrowth_intervalTr-N) r(r)r*r+r<floatr0r=r>r?r:r-r/rrrrr;us     r;c@sDeZdZUdZdZeeed<dZeeed<e ddZ e ed<dS) InitProcessGroupKwargsaB Use this object in your [`Accelerator`] to customize the initialization of the distributed processes. Please refer to the documentation of this [method](https://pytorch.org/docs/stable/distributed.html#torch.distributed.init_process_group) for more information on each argument. ```python from datetime import timedelta from accelerate import Accelerator from accelerate.utils import InitProcessGroupKwargs kwargs = InitProcessGroupKwargs(timeout=timedelta(seconds=800)) accelerator = Accelerator(kwargs_handlers=[kwargs]) ``` ZncclbackendN init_methodi)secondstimeout) r(r)r*r+rBr strr0rCrrErrrrrAs rAc@sleZdZUdZdZeed<dZeed<dZe ed<dZ eed<d Z e ed <d Z e eeefed <d dZdS)FP8RecipeKwargsa[ Use this object in your [`Accelerator`] to customize the initialization of the recipe for FP8 mixed precision training. Please refer to the documentation of this [class](https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/api/common.html#transformer_engine.common.recipe.DelayedScaling) for more information on each argument. ```python from accelerate import Accelerator from accelerate.utils import FP8RecipeKwargs kwargs = FP8RecipeKwargs(fp8_format="HYBRID") accelerator = Accelerator(mixed_precision="fp8", kwargs_handlers=[kwargs]) ``` rmarginr intervalE4M3 fp8_formatamax_history_len most_recentamax_compute_algo)FFFoverride_linear_precisioncCs4|j|_|jdkrtd|jdkr0tddS)N)rJZHYBRIDz(`fp8_format` must be 'E4M3' or 'HYBRID'.)maxrMz2`amax_compute_algo` must be 'max' or 'most_recent')rKupper ValueErrorrNrrrr __post_init__s    zFP8RecipeKwargs.__post_init__N)r(r)r*r+rHr:r0rIrKrFrLrNrOr r/rSrrrrrGs      rGc@s4eZdZdZdZdZdZdZdZdZ dZ d Z d Z d S) DistributedTypea Represents a type of distributed environment. Values: - **NO** -- Not a distributed environment, just a single process. - **MULTI_CPU** -- Distributed on multiple CPU nodes. - **MULTI_GPU** -- Distributed on multiple GPUs. - **MULTI_NPU** -- Distributed on multiple NPUs. - **MULTI_XPU** -- Distributed on multiple XPUs. - **DEEPSPEED** -- Using DeepSpeed. - **TPU** -- Distributed on TPUs. NO MULTI_CPU MULTI_GPU MULTI_NPU MULTI_XPU DEEPSPEEDFSDPTPU MEGATRON_LMN) r(r)r*r+rUrVrWrXrYrZr[r\r]rrrrrTsrTc@seZdZdZdZdZdZdS)SageMakerDistributedTypea+ Represents a type of distributed environment. Values: - **NO** -- Not a distributed environment, just a single process. - **DATA_PARALLEL** -- using sagemaker distributed data parallelism. - **MODEL_PARALLEL** -- using sagemaker distributed model parallelism. rU DATA_PARALLELMODEL_PARALLELN)r(r)r*r+rUr_r`rrrrr^s r^c@seZdZdZdZdZdS)ComputeEnvironmentz Represents a type of the compute environment. Values: - **LOCAL_MACHINE** -- private/custom cluster hardware. - **AMAZON_SAGEMAKER** -- Amazon SageMaker as compute environment. LOCAL_MACHINEAMAZON_SAGEMAKERN)r(r)r*r+rbrcrrrrras rac@s<eZdZdZdZdZdZdZdZdZ dZ d Z d Z d Z d Zd S) DynamoBackenday Represents a dynamo backend (see https://github.com/pytorch/torchdynamo). Values: - **NO** -- Do not use torch dynamo. - **EAGER** -- Uses PyTorch to run the extracted GraphModule. This is quite useful in debugging TorchDynamo issues. - **AOT_EAGER** -- Uses AotAutograd with no compiler, i.e, just using PyTorch eager for the AotAutograd's extracted forward and backward graphs. This is useful for debugging, and unlikely to give speedups. - **INDUCTOR** -- Uses TorchInductor backend with AotAutograd and cudagraphs by leveraging codegened Triton kernels. [Read more](https://dev-discuss.pytorch.org/t/torchinductor-a-pytorch-native-compiler-with-define-by-run-ir-and-symbolic-shapes/747) - **NVFUSER** -- nvFuser with TorchScript. [Read more](https://dev-discuss.pytorch.org/t/tracing-with-primitives-update-1-nvfuser-and-its-primitives/593) - **AOT_NVFUSER** -- nvFuser with AotAutograd. [Read more](https://dev-discuss.pytorch.org/t/tracing-with-primitives-update-1-nvfuser-and-its-primitives/593) - **AOT_CUDAGRAPHS** -- cudagraphs with AotAutograd. [Read more](https://github.com/pytorch/torchdynamo/pull/757) - **OFI** -- Uses Torchscript optimize_for_inference. Inference only. [Read more](https://pytorch.org/docs/stable/generated/torch.jit.optimize_for_inference.html) - **FX2TRT** -- Uses Nvidia TensorRT for inference optimizations. Inference only. [Read more](https://github.com/pytorch/TensorRT/blob/master/docsrc/tutorials/getting_started_with_fx_path.rst) - **ONNXRT** -- Uses ONNXRT for inference on CPU/GPU. Inference only. [Read more](https://onnxruntime.ai/) - **IPEX** -- Uses IPEX for inference on CPU. Inference only. [Read more](https://github.com/intel/intel-extension-for-pytorch). rUEAGER AOT_EAGERINDUCTORNVFUSER AOT_NVFUSERAOT_CUDAGRAPHSOFIFX2TRTONNXRTIPEXN)r(r)r*r+rUrerfrgrhrirjrkrlrmrnrrrrrdsrdc@seZdZdZddZdS)EnumWithContainsz\A metaclass that adds the ability to check if `self` contains an item with the `in` operatorcCs(z ||Wntk r"YdSXdS)NFT)rR)clsitemrrr __contains__5s  zEnumWithContains.__contains__N)r(r)r*r+rrrrrrro2sroc@s$eZdZdZddZeddZdS)BaseEnumzDAn enum class that can get the value of an item with `str(Enum.key)`cCs|jSr)valuerrrr__str__@szBaseEnum.__str__cCsttt|S)z.Method to list all the possible items in `cls`)listmaprF)rprrrrvCsz BaseEnum.listN)r(r)r*r+ru classmethodrvrrrrrs=srs) metaclassc@s(eZdZdZdZdZdZdZdZdZ dS) LoggerTypeaIRepresents a type of supported experiment tracker Values: - **ALL** -- all available trackers in the environment that are supported - **TENSORBOARD** -- TensorBoard as an experiment tracker - **WANDB** -- wandb as an experiment tracker - **COMETML** -- comet_ml as an experiment tracker allZaimZ tensorboardZwandbZcomet_mlZmlflowN) r(r)r*r+ALLZAIMZ TENSORBOARDZWANDBZCOMETMLZMLFLOWrrrrrzIs rzc@s eZdZdZdZdZdZdZdS) PrecisionTypezRepresents a type of precision used on floating point values Values: - **NO** -- using full precision (FP32) - **FP16** -- using half precision - **BF16** -- using brain floating point precision nofp8fp16bf16N)r(r)r*r+rUFP8ZFP16ZBF16rrrrr}\s  r}c@s$eZdZdZdZdZdZdZdZdS)RNGTypetorchcudaZnpuZxlaZxpu generatorN) r(r)r*ZTORCHCUDAZNPUZXLAZXPU GENERATORrrrrrls rc@seZdZdZdZdZdS) CustomDtypezd An enum that contains multiple custom dtypes that can be used for `infer_auto_device_map`. rZint4N)r(r)r*r+rINT4rrrrrusrc@s"eZdZUejed<ejed<dS)TensorInformationshapedtypeN)r(r)r*rSizer0rrrrrrs  rc@seZdZUdZedddidZeed<edddidZeed<ed dd idZ e ed <eddd idZ e ed <edddidZ e ed<dedddZddZdS)ProjectConfigurationzP Configuration for the Accelerator object based on inner-project needs. Nhelpz'A path to a directory for storing data.defaultmetadata project_dirziA path to a directory for storing logs of locally-compatible loggers. If None, defaults to `project_dir`. logging_dirFz?Whether saved states should be automatically iteratively named.automatic_checkpoint_namingz1The maximum number of total saved states to keep. total_limitrzThe current save iteration. iteration)rcCs||_|jdkr||_dS)zISets `self.project_dir` and `self.logging_dir` to the appropriate values.N)rr)rrrrrset_directoriess z$ProjectConfiguration.set_directoriescCs||jdSr)rrrrrrrSsz"ProjectConfiguration.__post_init__)N)r(r)r*r+rrrFr0rrr/rr:rrrSrrrrrs, rc@sZeZdZUdZedddidZeed<edddidZe ed <eddd idZ e ed <dS) GradientAccumulationPluginz? A plugin to configure gradient accumulation behavior. Nrz0The number of steps to accumulate gradients for.r num_stepsTzWhether to adjust the scheduler steps to account for the number of steps being accumulated. Should be `True` if the used scheduler was not adjusted for gradient accumulation.adjust_schedulerzWhether to synchronize setting the gradients when at the end of the dataloader. Should only be set to `False` if you know what you're doing.sync_with_dataloader) r(r)r*r+rrr:r0rr/rrrrrrs rc@seZdZUdZedddddeDidZeed<eddd idZe ed <eddd idZ e ed <eddd idZ e ed<edddidZ eed<edddidZe ed<ddZddZdS)TorchDynamoPluginzA This plugin is used to compile a model with PyTorch 2.0 NrzPossible options are cCsg|]}|jqSr)rtlower)rbrrr szTorchDynamoPlugin.rrBzCPossible options are 'default', 'reduce-overhead' or 'max-autotune'modez6Whether it is ok to break model into several subgraphs fullgraphz(Whether to use dynamic shape for tracingdynamicz/A dictionary of options to pass to the backend.optionsFz-Turn torch.compile() into a no-op for testingdisablecCsd}|jdkr"tj|dd|_t|j|_|jdkrPtj|dd|_|jdkrvttj|dddk|_|j dkrttj|d ddk|_ dS) NZACCELERATE_DYNAMO_ZBACKENDr~ZMODErZ USE_FULLGRAPHFalser Z USE_DYNAMIC) rBosenvirongetrdrQrrrrrprefixrrrrSs    zTorchDynamoPlugin.__post_init__cCs"t|j}|dj|d<|S)NrB)rrrrtr)rZ dynamo_configrrrrs zTorchDynamoPlugin.to_dict)r(r)r*r+rrdrBr0rrFrr/rrrrrSrrrrrrs  rc@sdeZdZUdZedddidZeed<edddidZe ed<eddd idZ e ed <eddd idZ e ed <ed ddidZ eed<edddidZeed<edddidZeed<edddidZeed<edddidZeed<edddidZeed<edddidZeed<ddZd-ddZd.d d!Zd"d#Zd$d%Zd&d'Zed/d)d*Zd+d,ZdS)0DeepSpeedPluginz5 This plugin is used to integrate DeepSpeed. Nrzkpath to DeepSpeed config file or dict or an object of class `accelerate.utils.deepspeed.HfDeepSpeedConfig`.r hf_ds_configzNumber of steps to accumulate gradients before updating optimizer states. If not set, will use the value from the `Accelerator` directly.gradient_accumulation_stepsz#Enable gradient clipping with valuegradient_clippingzMPossible options are 0,1,2,3; Default will be taken from environment variable zero_stageTzUIf both train & eval dataloaders are specified, this will decide the train_batch_sizeis_train_batch_minzMPossible options are none|cpu|nvme. Only applicable with ZeRO Stages 2 and 3.offload_optimizer_devicezFPossible options are none|cpu|nvme. Only applicable with ZeRO Stage 3.offload_param_devicezJPossible options are /nvme|/local_nvme. Only applicable with ZeRO Stage 3.offload_optimizer_nvme_pathoffload_param_nvme_pathz{Flag to indicate whether to enable `deepspeed.zero.Init` for constructing massive models.Only applicable with ZeRO Stage-3.zero3_init_flagzQFlag to indicate whether to save 16-bit model. Only applicable with ZeRO Stage-3.zero3_save_16bit_modelc s(ddlm}jdkr:tjdd}|r4t|n|_jdkrdtjdd}|dkrdt |_j dkrttjdd_ j dkrtjd d_ j dkrtjd d_ j dkrtjd d_ jdkrtjd d_jdkr tjd ddk_jdkr&tjdd_tjts\tjtrNjdks\tj|rLtj|sv|j_djjkrdjjd<djjkrtdddddddddd}fdd|D}|D]}j|f|ddiqj|D]6\}}j|}|dk r|dkrt||qnnddjj j j d krnj ndd!j j d krjndd!jd"d#} jrj| d<|| _jj_t d$jd%<jdkrttjd&tj dk_jr$j s$t!"d'd_dS)(Nr HfDeepSpeedConfig&ACCELERATE_GRADIENT_ACCUMULATION_STEPSautoACCELERATE_GRADIENT_CLIPPINGnoneACCELERATE_DEEPSPEED_ZERO_STAGE-ACCELERATE_DEEPSPEED_OFFLOAD_OPTIMIZER_DEVICE)ACCELERATE_DEEPSPEED_OFFLOAD_PARAM_DEVICE0ACCELERATE_DEEPSPEED_OFFLOAD_OPTIMIZER_NVME_PATH,ACCELERATE_DEEPSPEED_OFFLOAD_PARAM_NVME_PATH+ACCELERATE_DEEPSPEED_ZERO3_SAVE_16BIT_MODELfalsetrueZ ACCELERATE_DEEPSPEED_CONFIG_FILErzero_optimizationzDPlease specify the ZeRO optimization config in the DeepSpeed config.rzzero_optimization.stagez*zero_optimization.offload_optimizer.devicez&zero_optimization.offload_param.devicez)zero_optimization.offload_param.nvme_pathz-zero_optimization.offload_optimizer.nvme_pathz;zero_optimization.stage3_gather_16bit_weights_on_model_save)rrrrrrrrcs*i|]"\}}t|dk r|t|qSr)getattrrrrrr#Wsz1DeepSpeedPlugin.__post_init__.. must_matchFZnvme)deviceZ nvme_path)ZstageZoffload_optimizerZ offload_paramZ)stage3_gather_16bit_weights_on_model_save)train_batch_sizetrain_micro_batch_size_per_gpurrinfZsteps_per_printZACCELERATE_DEEPSPEED_ZERO3_INITzSDeepSpeed Zero3 Init flag is only applicable for ZeRO Stage 3. Setting it to False.)#Z deepspeedrrrrrisdigitr:rr@rrrrrrr isinstancedictrFconfigrR_deepspeed_config_checksr&keys fill_matchZset_stage_and_offload get_valuesetattrdeepspeed_configrrZis_zero3warningswarn) rrgasrZplugin_to_config_mappingkwargskeyrtZ config_valuerrrrrSs                          zDeepSpeedPlugin.__post_init__c Ks|dkr gn|}|j|\}}|dkr,dS||dkrh||krR||||<dStd|d|d|spdS||}|dk r||kr|||kr|d|d|d|d||dS)Nr`z'` not found in kwargs. Please specify `zY` without `auto`(set to correct value) in the DeepSpeed config file or pass it in kwargs.z- ds =z vs arg )rZfind_config_noderrRappend)rZ ds_key_long mismatchesrrrZds_keyZds_valrrrrs"   zDeepSpeedPlugin.fill_matchc Ks|dkr gn|}|dkr|j}|D]R\}}t|tr\|jf||d|||d|q&|j|||fd|i|q&t|dkr|dkrd|}td|d dS) z=Process the DeepSpeed config with the values from the kwargs.N.)rrrrrrr zSPlease correct the following DeepSpeed config values that mismatch kwargs values: zF The easiest method is to set these DeepSpeed config values to 'auto'.) rr&rrdeepspeed_config_processrlenjoinrR) rrrrrrrrtZmismatches_msgrrrrs&    z(DeepSpeedPlugin.deepspeed_config_processcCs|j}|dk|dkd}|dkr8d|krTddd|d<n|dkrTd|krTddi|d<|dkr|dkrhdnd}t||iddd krtd |d |d d D]}||krddi||<q|jdddi||jdddi|dS)Nrr) fp16.enabled bf16.enabledT)r-Z auto_castr-r~rrz*`--mixed_precision` arg cannot be set to `z` when `z&` is set in the DeepSpeed config file.)rrFrrr)r)r)rrFrrrRr)rmixed_precision ds_configrZ diff_dtyperrrrset_mixed_precisions*  z#DeepSpeedPlugin.set_mixed_precisioncCsddlm}|jr|s tdt|j}d|ks@|ddkrHd|d<d|ks\|ddkrdd|d<|ddkrv|d=tdd d rd d lm }n d d l m }|||_ dS) Nr )is_transformers_availablezoWhen `zero3_init_flag` is set, it requires Transformers to be installed. Please run `pip install transformers`.rrrrZ transformers.ZACCELERATE_CONFIG_DS_FIELDSr,c3s|]}|kVqdSrrrZenv_variable_names_to_ignorerr sz;DeepSpeedPlugin._deepspeed_config_checks..z_When using `deepspeed_config_file`, the following accelerate config variables will be ignored: am. Please specify them appropriately in the DeepSpeed config file. If you are using an accelerate config file, remove others config variables mentioned in the above specified list. The easiest method is to create a new config following the questionnaire via `accelerate config`. It will only ask for the necessary config variables when using `deepspeed_config_file`.)rrrsplitanyrR)rZ'deepspeed_fields_from_accelerate_configrrrrs$  z(DeepSpeedPlugin._deepspeed_config_checks)NT)rNNT)F)r(r)r*r+rrrr0rr:rr@rrrFrr/rrrrrrSrrrrrrrrrrrrrsr h   rc@seZdZUdZedddidZded<edddidZded <eddd idZded <eddd idZ e e ed <edddidZ ded<edddidZ e eejjed<edddidZded<edddidZded<edddidZded<edddidZeed<edddidZeed<edddidZe e ejjgdfed<eddd idZeed!<eddd"idZeed#<eddd$idZeed%<d&d'Zed(d)Zd*d+Zd,d-Z d.d/Z!dS)0FullyShardedDataParallelPluginzG This plugin is used to enable fully sharded data parallelism. NrzdFSDP Sharding Strategy of type `torch.distributed.fsdp.fully_sharded_data_parallel.ShardingStrategy`rz typing.Anysharding_strategyzdFSDP Backward Prefetch of type `torch.distributed.fsdp.fully_sharded_data_parallel.BackwardPrefetch`backward_prefetchaA config to enable mixed precision training with FullyShardedDataParallel. The 3 flags that are set are `param_dtype`, `reduce_dtype`, `buffer_dtype`. Each flag expects `torch.dtype` as the value. It is of type `torch.distributed.fsdp.fully_sharded_data_parallel.MixedPrecision`.mixed_precision_policyzCA callable specifying a policy to recursively wrap layers with FSDPauto_wrap_policyzDecides Whether to offload parameters and gradients to CPU. It is of type `torch.distributed.fsdp.fully_sharded_data_parallel.CPUOffload`. cpu_offloadz%A list of modules to ignore for FSDP.ignored_modulesz_FSDP State Dict Type of type `torch.distributed.fsdp.fully_sharded_data_parallel.StateDictType`state_dict_typezcFSDP State Dict Config of type `torch.distributed.fsdp.fully_sharded_data_parallel.StateDictConfig`state_dict_configzrFSDP Optimizer State Dict Config of type `torch.distributed.fsdp.fully_sharded_data_parallel.OptimStateDictConfig`optim_state_dict_configFaRIf False, then FSDP allows the CPU thread to schedule all-gathers without any extra synchronization. If True, then FSDP explicitly synchronizes the CPU thread to prevent too many in-flight all-gathers. This bool only affects the sharded strategies that schedule all-gathers. Enabling this can help lower the number of CUDA malloc retries.limit_all_gathersaBIf True, allows non-uniform `requires_grad` during init, which means support for interspersed frozen and trainable paramteres. Useful in cases such as parameter-efficient fine-tuning. Please refer this [blog](https://dev-discuss.pytorch.org/t/rethinking-pytorch-fully-sharded-data-parallel-fsdp-from-first-principles/1019)use_orig_paramszA Callable[torch.nn.Module] -> None that specifies how modules that are currently on the meta device should be initialized onto an actual device. param_init_fnTzIf True, each individually wrapped FSDP unit will broadcast module parameters from rank 0 to ensure they are the same across all ranks after initializationsync_module_stateszIf True, then FSDP explicitly prefetches the next upcoming all-gather while executing in the forward pass. only use with Static graphs.forward_prefetchzIf True, activation checkpointing is a technique to reduce memory usage by clearing activations of certain layers and recomputing them during a backward pass. Effectively, this trades extra computation time for reduced memory usage.activation_checkpointingcCsdddlm}m}m}d}|jdkr>|ttj|dd|_|j dkr|t tj|dddkrp|dd |_ n |d d |_ |j dkrtj|d d }|t d kr|t |d|_ |jdkrtj|dd}||t tj|dddk|_t tj|dddk|_t tj|dddk|_t tj|dddk|_|jr`dd|_dS)Nr)BackwardPrefetch CPUOffloadShardingStrategyZFSDP_ZSHARDING_STRATEGYr ZOFFLOAD_PARAMSrT)Zoffload_paramsFZBACKWARD_PREFETCHZ NO_PREFETCHZSTATE_DICT_TYPEFULL_STATE_DICTZUSE_ORIG_PARAMSZSYNC_MODULE_STATESTrueZFORWARD_PREFETCHZACTIVATION_CHECKPOINTINGcSs|jtjddS)NF)rrecurse)Zto_emptyrrZcurrent_device)xrrrz>FullyShardedDataParallelPlugin.__post_init__..)2torch.distributed.fsdp.fully_sharded_data_parallelrrrrr:rrrrrrrindexrset_state_dict_typer r r r r )rrrrrZprefetch_policystate_dict_type_policyrrrrS|s*       z,FullyShardedDataParallelPlugin.__post_init__cCsXt|}|jj|kr|jSt|dkr.dS|D] }t||}|dk r2|Sq2dS)z Gets a class from a module by its name. Args: module (`torch.nn.Module`): The module to get the class from. name (`str`): The name of the class. rN)rvchildrenr%r(rrget_module_class_from_name)modulerZmodules_childrenZ child_moduleZ module_classrrrrs    z9FullyShardedDataParallelPlugin.get_module_class_from_namec Csddlm}m}t|dddk r,d|jnd}|jdkrtj dd}|t dkrtj d| d}t }|D],}t ||} | dkrtd qr|| qrtj||d |_n6|t d krttj d d} | dkrtj|| d |_dS)Nr)size_based_auto_wrap_policytransformer_auto_wrap_policy_no_split_modulesrrrZNO_WRAPZFSDP_TRANSFORMER_CLS_TO_WRAPz@Could not find the transformer layer class to wrap in the model.)Ztransformer_layer_clsr ZFSDP_MIN_NUM_PARAMS)min_num_params)Ztorch.distributed.fsdp.wraprr rrr!rrrrrrsetrrradd functoolspartialr:) rmodelrr Z%default_transformer_cls_names_to_wraprZtransformer_cls_names_to_wrapZtransformer_cls_to_wrapZ layer_classZtransformer_clsr"rrrset_auto_wrap_policys:       z3FullyShardedDataParallelPlugin.set_auto_wrap_policycCsX|dkrtj}n|dkr tj}ntd|ddlm}|jdkrT||||d|_dS)NrrzUnknown mixed precision value: r)MixedPrecision)Z param_dtypeZ reduce_dtypeZ buffer_dtype)rfloat16bfloat16rRrr)r)rrrr)rrrrs  z2FullyShardedDataParallelPlugin.set_mixed_precisioncCshddlm}m}m}|t|d|_|j|jkrd|jdkrL|ddd|_|j dkrd|ddd|_ dS)Nr)FullOptimStateDictConfigFullStateDictConfig StateDictTyper T)Zoffload_to_cpuZ rank0_only) rr,r-r.rrrrrr)rrr,r-r.rrrrs   z2FullyShardedDataParallelPlugin.set_state_dict_type)"r(r)r*r+rrr0rrrr rrrr rnnModulerrrrr/r r r r r rS staticmethodrr(rrrrrrrs        rc@s<eZdZUdZedddidZeed<edddidZeed<eddd idZ eed <eddd idZ e ed <eddd idZ e ed<edddidZe ed<edddidZe ed<edddidZeed<edddidZeed<edddidZeed<edddidZeed<edddidZeed<edddidZeed <eddd!idZe ed"<eddd#idZe ed$<ed%dd&idZeed'<eddd(idZeed)<eddd*idZeed+<eddd,idZeed-<eddd.idZeed/<eddd0idZe ed1<ed2dd3idZe ed4<eddd5idZ e!eed6<eddd7idZ"e#e$ed8<eddd9idZ%e#e$ed:<ed;dddd?idZ'e ed@<edddAidZ(eedB<edddCidZ)eedD<edddEidZ*eedF<edddGidZ+eedH<ed>ddIidZ,e edJ<edKddLidZ-eedM<edNddOidZ.eedP<ed>ddQidZ/e edR<edddSidZ0e#e1edT<edddUidZ2e#e3ee1fedV<edddWidZ4e#e$edX<edddYidZ5e#e$edZ<eddd[idZ6e#e3ee1fed\<d]d^Z7dkd_d`Z8dadbZ9dcddZ:dedfZ;dgdhZgradient clipping value based on global L2 Norm (0 to disable)rzenable sequence parallelismsequence_parallelismz)enable selective activation recomputationrecompute_activationzenable distributed optimizeruse_distributed_optimizerz/Rank where encoder and decoder should be split."pipeline_model_parallel_split_rankz,Number of layers per virtual pipeline stage.%num_layers_per_virtual_pipeline_stageTzUIf both train & eval dataloaders are specified, this will decide the micro_batch_sizerzTotal number of iterations to train over all training runs. Note that either train-iters or train-samples should be provided when using `MegatronLMDummyScheduler` train_iterszTotal number of samples to train over all training runs. Note that either train-iters or train-samples should be provided when using `MegatronLMDummyScheduler` train_samplesZconstantzKWeight decay increment function. choices=["constant", "linear", "cosine"]. weight_decay_incr_stylez7Initial weight decay coefficient for L2 regularization.start_weight_decayz:End of run weight decay coefficient for L2 regularization.end_weight_decayZlinearzGLearning rate decay function. choices=['constant', 'linear', 'cosine'].lr_decay_stylezPNumber of iterations for learning rate decay. If None defaults to `train_iters`.lr_decay_iterszONumber of samples for learning rate decay. If None defaults to `train_samples`.lr_decay_samplesz;number of iterations to linearly warmup learning rate over.lr_warmup_itersz8number of samples to linearly warmup learning rate over.lr_warmup_sampleszLfraction of lr-warmup-(iters/samples) to linearly warmup learning rate over.lr_warmup_fractionrzPMinumum value for learning rate. The scheduler clip values below this threshold.min_lrz^Number of samples consumed in the same order as the dataloaders to `accelerator.prepare` call.consumed_samplesz"Condition to disable weight decay.no_wd_decay_condz!Condition to scale learning rate. scale_lr_cond?zLearning rate multiplier.lr_multFzUWhether the format of dataset follows Megatron-LM Indexed/Cached/MemoryMapped format.megatron_dataset_flagz#Maximum sequence length to process. seq_lengthz3Maximum sequence length to process for the encoder.encoder_seq_lengthz3Maximum sequence length to process for the decoder.decoder_seq_lengthzPath to save tensorboard logs.tensorboard_dirz#Whether to set all logging options.set_all_logging_optionsdz?Number of iterations to run for evaluation validation/test for. eval_itersiz6Interval between running evaluation on validation set. eval_intervalz(Whether to return logits from the model. return_logitszCustom train step class.custom_train_step_classzCustom train step kwargs.custom_train_step_kwargszCustom model provider function.custom_model_provider_functionzCustom prepare model function.custom_prepare_model_functionz5Other Megatron-LM arguments. Please refer Megatron-LMother_megatron_argscCs d}|jdkr&ttj|dd|_|jdkrHttj|dd|_|jdkrjttj|dd|_|jdkrttj|dd|_|j dkrt tj|dd dk|_ |j dkrt tj|d d dk|_ |j dkrt tj|d d dk|_ |jdks|j rd |_ nd |_ |jdk rjt|jdkrN|jddgnt|jdkrj|jd|j|j|j|j|j |j |j |j|j|j|j|j|j|j|j|jd|_|j rd|jd<|jdk r|j|jd<|jr||jdk r|j|jdS)NZ MEGATRON_LM_Z TP_DEGREEr Z PP_DEGREEZNUM_MICRO_BATCHESZGRADIENT_CLIPPINGrJZRECOMPUTE_ACTIVATIONrZUSE_DISTRIBUTED_OPTIMIZERZSEQUENCE_PARALLELISMlocalrrr)Ztensor_model_parallel_sizeZpipeline_model_parallel_sizer9r:DDP_implr8Zsequence_parallelZ clip_gradr5rGrHrIrKrLrSrTZ selectiveZrecompute_granularityrP) r3r:rrrr4r5rr@r7rr8r6r\rGrextendrr9r:rHrIrKrLrSrTmegatron_lm_default_argsrPrQset_tensorboard_logging_optionsrZupdaterrrrrSsb             zMegatronLMPlugin.__post_init__c Csd|jjkrd}|jj}|jj}|jj}|jj}|jj}|jj} d|j j krXd} |j dk r|j dk rvt d|j |_ n4|j dk r|j |_ n |dk r|djd|_ n||_ |j |jd<nd |jjkrd }|jj}|jj}|jj}|jj}|jj} d} |j dk r2|jdk r(t d |j|_ n8|jdk rH|j|_ n"|dk rd|djd|_ n||_ |j |jd<|j|jd <d |jd<nd|jjkrXd}|jj}|jj}|jj}t|jdr|jjnd}|jj} d} |j dkr|dk r |djd|_ n||_ |jdkr>|dk r8|djd|_n||_|j |jd<|j|jd<ntd||jd<||jd<||jd<||jd<||jd<| |jd<| |jd<|jj|jd<|dkr||jd<dS)Nz megatron-bertZbertZmaskedlmTzOBoth `seq_length` and `encoder_seq_length` are set. Using `encoder_seq_length`.Z input_idsr rMZgpt2ZgptzOBoth `seq_length` and `decoder_seq_length` are set. Using `decoder_seq_length`.rUZGPT2BPETokenizerZtokenizer_typeZt5 n_positionsilabelsrNrOu🤗 Accelerate Megatron-LM integration supports only BERT, GPT and T5 model. Please check the model you are using is one of those.model_type_name num_layers hidden_sizenum_attention_headsmax_position_embeddingspretraining_flagorig_vocab_sizeZmodel_return_dict num_labels)rZ model_typerZnum_hidden_layersrerfrgrjZ vocab_sizer%r(rMrNrrrr^Zn_layerZn_embdZn_headrarOrUrdZd_modelZ num_headshasattrrRZ return_dict) rr'Z batch_datarcrdrerfrgrjrirhrrrset_network_size_argss                             z&MegatronLMPlugin.set_network_size_argscCs<|dkrd|jd<n$|dkr8d|jd<d|_|j|jd<dS)NrTrr[r\)r^r\)rrrrrrs   z$MegatronLMPlugin.set_mixed_precisioncCsD||_||_|||j|_|j|jd<|j|jd<|j|jd<dS)Ndata_parallel_sizemicro_batch_sizeglobal_batch_size)rmrnr5ror^)rrnZ dp_degreerrrset_training_argss   z"MegatronLMPlugin.set_training_argscCs|jj}d|krXd|jd<|jdd|jd<|jdd|jd<|jd|jd <n4d |kr|d |jd<|jd |jd <ntd |d|jd|jd<|jd|jd<dS)NZadam optimizerZbetasrZ adam_beta1r Z adam_beta2ZepsZadam_epsZsgdZmomentumZ sgd_momentumz Optimizer z is not supported by Megatron-LMlrZ weight_decay)r%r(rr^defaultsrR)rrqZoptimizer_namerrrset_optimizer_type%s   z#MegatronLMPlugin.set_optimizer_typecCs|jdkr6|j|jd|_|jdk r6d|_td|jdkrl|j|jd|_|jdk rftdd|_|j|jd<|j|jd<|j|jd<|j|jd<|j |jd <|j |jd <|j |jd <|j |jd <|j |jd <|j|jd<|j|jd<|j|jd<dS)NrmzXIgnoring `train_samples` as `train_iters` based on scheduler is being used for training.z`Ignoring `lr_warmup_samples` as `lr_warmup_iters` based on scheduler is being used for training.rr;rCr<rDrArBrEr@r=r>r?rF)r;Ztotal_num_stepsr^r<rrrCZwarmup_num_stepsrDrArBrEr@r=r>r?rF)rZ schedulerrrrset_scheduler_args5s4               z#MegatronLMPlugin.set_scheduler_argscCs|ddlm}t}||}|}t|d|_|jD]:\}}|drZd|j |<q<|drr(r)r*r+rr3r:r0r4r5rr@r6r/r7r8r9r:rrFr;r<r=r>r?r@rArBrCrDrErFrGr rHr rrIrKrLrMrNrOrPrQrSrTrUrVrrWrrXrYrZrSrlrrprtrur_rrrrr2s8 9 Sr2c@seZdZUdZedddidZeed<edddidZe ed <eddd idZ eed <ed dd idZ e ed<edddidZ eed<edddidZeed<edddidZejed<edddidZee ed<edddidZee ed<ddZdS)BnbQuantizationConfigzD A plugin to enable BitsAndBytes 4bit and 8bit quantization Frzenable 8bit quantization.r load_in_8bit@zEvalue of the outliner threshold. only relevant when load_in_8bit=Truellm_int8_thresholdzenable 4bit quantization. load_in_4bitfp4z\set the quantization data type in the `bnb.nn.Linear4Bit` layers. Options are {'fp4','np4'}.bnb_4bit_quant_typezlenable nested quantization where the quantization constants from the first quantization are quantized again.bnb_4bit_use_double_quantrzThis sets the computational type which might be different than the input time. For example, inputs might be fp32, but computation can be set to bf16 for speedups. Options are {'fp32','fp16','bf16'}.bnb_4bit_compute_dtypeNzthis sets the dtype of the remaining non quantized layers. `bitsandbytes` library suggests to set the valueto `torch.float16` for 8 bit model and use the same dtype as the compute dtype for 4 bit model torch_dtypezian explicit list of the modules that we don't quantize. The dtype of these modules will be `torch_dtype`. skip_moduleszXan explicit list of the modules that we don't quantize. We keep them in `torch.float32`.keep_in_fp32_modulescCsVt|jtstdt|jts(td|jr<|jrsr  $  #,    %#-Xx