U 9%e@sdZddlZddlZddlZddlZddlZddlZddlmZddl m Z ddl m Z ddl mZddlmZmZmZmZmZmZddlZddlZddlmZddlmZdd lmZmZm Z m!Z!dd l"m#Z#d d l$m%Z%d d l&m'Z'd dl(m)Z)m*Z*m+Z+m Z e+re ,eej-e*ddr6ddl.m/m0Z1zddl2m3Z3Wne4k r`dZ3YnXe 5e6Z7ddZ8eej9ej:fdddZ;dlddZdd Z?d!d"Z@d#d$ZAd%d&ZBdoeeeCed'd(d)ZDdeEd*feeeCeFfeeCeejEej9d+d,d-ZGd.d/ZHe eCd0d1d2ZIGd3d4d4e#ZJGd5d6d6e!ZKej(jLjeCd7d8d9ZMdpd:d;ZNdqdd?ZPGd@dAdAZQe GdBdCdCZRdrdDdEZSGdFdGdGe!ZTGdHdIdIe#ZUGdJdKdKe!ZVGdLdMdMeZWdNdOZXdPdQZYeeZeFfeeZeFfdRdSdTZ[dUdVZ\dsdXdYZ]dZd[Z^dtd\d]Z_d^d_Z`d`daZadbdcZbe)rddlcmdmZeeefdudddeZgeefdfdgZhdhdiZidjdkZjdS)vz( Torch utilities for the Trainer class. N)Mapping)contextmanager) dataclass) StreamHandler)AnyDictIteratorListOptionalUnion)nn)DatasetIterableDataset RandomSamplerSampler)DistributedSampler)is_deepspeed_zero3_enabled) BatchEncoding)is_sagemaker_mp_enabledis_torch_tpu_availableis_training_run_on_sagemakerloggingF)Z check_device)SAVE_STATE_WARNINGcCs2t|dr|jdk rt|jSt|dr.|jSdS)N batch_samplersampler)hasattrrget_dataloader_samplerr)Z dataloaderr\/var/www/html/Darija-Ai-API/env/lib/python3.8/site-packages/transformers/trainer_pt_utils.pyr:s  rZtensor_or_arraycCsDt|tjr6ttdr"t|}q@|jdkr@|d}n t|}|S)N atleast_1dr) isinstancetorchTensorrr"ndimnpr!rrr r"As      r"cCst|}t|}t|jdks2|jd|jdkrDtj||fddS|jd|jdt|jd|jdf|jdd}|||}||d|jdd|jdf<|||jddd|jdf<|S)z`Concatenates `tensor1` and `tensor2` on first axis, applying padding on the second if necessary.rrdimN)r"lenshaper$catmaxZnew_full)Ztensor1Ztensor2 padding_index new_shaperesultrrr torch_pad_and_concatenateLs"8   r3cCst|}t|}t|jdks2|jd|jdkrDtj||fddS|jd|jdt|jd|jdf|jdd}tj|||d}||d|jdd|jdf<|||jddd|jdf<|S)z^Concatenates `array1` and `array2` on first axis, applying padding on the second if necessary.rr)Zaxisr+Nr-)r"r,r-r'Z concatenater/ full_like)Zarray1Zarray2r0r1r2rrr numpy_pad_and_concatenate^s"8  r6cst|tks.tdt|dtdt|ttfr\t|fddt|DSt|tjrvt|dSt|t rt|fdd| DSt|t j rt |dStd t|d S) z Concat the `new_tensors` to `tensors` on the first dim and pad them on the second if needed. Works for tensors or nested list/tuples/dict of tensors. zEExpected `tensors` and `new_tensors` to have the same type but found z and .c3s |]\}}t||dVqdS)r0N nested_concat).0tnr8rr ysz nested_concat..r8cs$i|]\}}|t||dqS)r8r9r;kr< new_tensorsr0rr ~sz!nested_concat..z(Unsupported type for concatenation: got N)typeAssertionErrorr#listtuplezipr$r%r3ritemsr'ndarrayr6 TypeError)tensorsrBr0rrAr r:ps     r:cCst|ttfr2|D]}t|}|dk r|Sqnt|trh|D] \}}t|}|dk rD|SqDnPt|tjrt|j dkr|j dSdSt|t j rt|j dkr|j dSdSdS)zV Find the first dimension of a tensor in a nested list/tuple/dict of tensors. Nrr) r#rFrGfind_batch_sizerrIr$r%r,r-r'rJ)rLr<r2keyvaluerrr rMs     rMcCspt|ttfr$t|dd|DSt|trHt|dd|DS|}|jtj krh| tj }| S)zENumpify `tensors` (even if it's a nested list/tuple/dict of tensors).css|]}t|VqdSNnested_numpifyr;r<rrr r>sz!nested_numpify..cSsi|]\}}|t|qSrrQr?rrr rCsz"nested_numpify..) r#rFrGrDrrIcpudtyper$Zbfloat16tofloat32numpy)rLr<rrr rRs   rRcCsPt|ttfr$t|dd|DSt|trHt|dd|DS|S)zDDetach `tensors` (even if it's a nested list/tuple/dict of tensors).css|]}t|VqdSrP nested_detachrSrrr r>sz nested_detach..cSsi|]\}}|t|qSrrYr?rrr rCsz!nested_detach..)r#rFrGrDrrIdetach)rLrrr rZs  rZcstrddlmm}t|ttfrDt|fddt|DSt|t rpt|fddt| DSt |}| |t jStddS)Nrc3s&|]\}}t|d|VqdS)_Nnested_xla_mesh_reduce)r;ir<namerr r>sz)nested_xla_mesh_reduce..cs*i|]"\}\}}|t|d|qS)r\r])r;r_r@r<r`rr rCs z*nested_xla_mesh_reduce..z;Torch xla must be installed to use `nested_xla_mesh_reduce`)rtorch_xla.core.xla_modelcore xla_modelr#rFrGrD enumeraterrIr"Z mesh_reducer$r. ImportError)rLraxmrr`r r^s r^)tensornum_total_examplesreturncsztttfr,tfddDWSttrVtfddDWStfddtt D}t |t j |dd}dk r|d}|WStk rtd YnXdS) Nc3s|]}t|VqdSrPdistributed_concatrSrirr r>sz%distributed_concat..csi|]\}}|t|qSrrkr?rmrr rCsz&distributed_concat..csg|] }qSrcloner;r\rhrr sz&distributed_concat..rr)(Not currently using distributed training)r#rGrFrDrrIr" contiguousrangedistget_world_size all_gatherr$r.rE)rhrioutput_tensorsconcatr)rirhr rls     rlcuda)scalarsridevicerjcsz^t||fddttD}t|tj|dd}|dk rZ|d|}|WStk rztdYnXdS)Ncsg|] }qSrrnrpZtensorized_scalarrr rrsz1distributed_broadcast_scalars..rr)rs) r$rhrVrurvrwrxr.rE)r|rir}ryrzrr~r distributed_broadcast_scalarss  rcCs>t|dkr:|D](}|jtks(|jtkrt|j|jqdS)Nr)r,category UserWarningmessagerwarningswarn)Zcaught_warningswrrr reissue_pt_warningss rZ local_rankccs*|dkrtdV|dkr&tdS)z Decorator to make all processes in distributed training wait for each local_master to do something. Args: local_rank (`int`): The rank of the local process. )rNr)rvZbarrierrrrr torch_distributed_zero_firsts rcs,eZdZdZfddZfddZZS)DistributedSamplerWithLoopa Like a torch.utils.data.distributed.DistributedSampler` but loops at the end back to the beginning of the shuffled samples to make each process have a round multiple of batch_size samples. Args: dataset (`torch.utils.data.Dataset`): Dataset used for sampling. batch_size (`int`): The batch size used with this sampler kwargs (`Dict[str, Any]`, *optional*): All other keyword arguments passed to `DistributedSampler`. c stj|f|||_dSrP)super__init__ batch_size)selfdatasetrkwargs __class__rr rsz#DistributedSamplerWithLoop.__init__csrtt}t||jdkr$dn|jt||j}|jt|j|jkrRdnd}|||||7}t|S)Nrr) rFr__iter__r,rrankr num_replicasiter)rindices remainderZstart_remainderrrr rs *z#DistributedSamplerWithLoop.__iter__)__name__ __module__ __qualname____doc__rr __classcell__rrrr rs rc@s*eZdZdZd ddZddZddZdS) SequentialDistributedSamplera Distributed Sampler that subsamples indices sequentially, making it easier to collate all results at the end. Even though we only use this sampler for eval and predict (no training), which means that the model params won't have to be synced (i.e. will not hang for synchronization even if varied number of forward passes), we still add extra samples to the sampler to make it evenly divisible (like in `DistributedSampler`) to make it easy to `gather` or `reduce` resulting tensors at the end of the loop. NcCstdt|dkr,ts$tdt}|dkrLtsDtdt}||_||_ ||_ t |j}|dk rt t |||||_nt t |||_|j|j |_||_dS)NzUSequentialDistributedSampler is deprecated and will be removed in v5 of Transformers.,Requires distributed package to be available)rr FutureWarningrv is_available RuntimeErrorrwget_rankrrrr,intmathceil num_samples total_sizer)rrrrrrrrr r's* z%SequentialDistributedSampler.__init__cCsttt|j}||d|jt|7}t||jksVtdt|d|jd||j|j|jd|j}t||jkstdt|d|jdt|S)NzIndices length z and total size z mismatchedrz and sample number ) rFrur,rrrErrrrrrrr r@s   z%SequentialDistributedSampler.__iter__cCs|jSrPrrrrr __len__Qsz$SequentialDistributedSampler.__len__)NNN)rrrrrrrrrrr rs rrrcCs*tdkrt|St|ttdS)Nr)rr)rgZxrt_world_sizerrZ get_ordinalrrrr get_tpu_samplerUs rcsHt|ttfr(t|fdd|DStj||f|jdddS)z\Create the same nested structure as `arrays` with a first dimension always at `num_samples`.c3s|]}t|VqdSrP)nested_new_liker;xrrr r>^sz"nested_new_like..rNr4)r#rFrGrDr'r5r-)arraysrr0rrr r[srcCsFtj|||jd|f|jddd}||ddd|jdf<|S)zmExpand the `arrays` so that the second dimension grows to `new_seq_length`. Uses `padding_index` for padding.rr+Nr4r)r'r5r-)rZnew_seq_lengthr0r2rrr expand_likebs(rcs\t|ttfr(t|fdd|DSt|trPt|fdd|DS|dS)zQTruncate `tensors` at `limit` (even if it's a nested list/tuple/dict of tensors).c3s|]}t|VqdSrPnested_truncaterSlimitrr r>lsz"nested_truncate..csi|]\}}|t|qSrrr?rrr rCnsz#nested_truncate..N)r#rFrGrDrrI)rLrrrr ris  rc@s2eZdZdZd ddZddZdd Zd d ZdS) DistributedTensorGathereraR A class responsible for properly gathering tensors (or nested list/tuple of tensors) on the CPU by chunks. If our dataset has 16 samples with a batch size of 2 on 3 processes and we gather then transfer on CPU at every step, our sampler will generate the following indices: `[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 0, 1]` to get something of size a multiple of 3 (so that each process gets the same dataset length). Then process 0, 1 and 2 will be responsible of making predictions for the following samples: - P0: `[0, 1, 2, 3, 4, 5]` - P1: `[6, 7, 8, 9, 10, 11]` - P2: `[12, 13, 14, 15, 0, 1]` The first batch treated on each process will be - P0: `[0, 1]` - P1: `[6, 7]` - P2: `[12, 13]` So if we gather at the end of the first batch, we will get a tensor (nested list/tuple of tensor) corresponding to the following indices: `[0, 1, 6, 7, 12, 13]` If we directly concatenate our results without taking any precautions, the user will then get the predictions for the indices in this order at the end of the prediction loop: `[0, 1, 6, 7, 12, 13, 2, 3, 8, 9, 14, 15, 4, 5, 10, 11, 0, 1]` For some reason, that's not going to roll their boat. This class is there to solve that problem. Args: world_size (`int`): The number of processes used in the distributed training. num_samples (`int`): The number of samples in our dataset. make_multiple_of (`int`, *optional*): If passed, the class assumes the datasets passed to each process are made to be a multiple of this argument (by adding samples). padding_index (`int`, *optional*, defaults to -100): The padding index to use if the arrays don't all have the same sequence length. Nr(cCsftdt||_||_|dkr$|n||}tt||||_|j||_ d|_ d|_ ||_ dS)NzRDistributedTensorGatherer is deprecated and will be removed in v5 of Transformers.) rrr world_sizerrr'r total_samplesprocess_length_storage_offsetsr0)rrrZmake_multiple_ofr0rrrr rs z"DistributedTensorGatherer.__init__cCsz|dkr dS|jdkr@t||j|jd|_ttd|j|j|_||j|\}|_t|j D]}|j||7<q^dS)z Add `arrays` to the internal storage, Will initialize the storage to the full size at the first arrays passed so that if we're bound to get an OOM, it happens at the beginning. Nr8r) rrrr0rFrurr_nested_set_tensorsr)rr slice_lenr_rrr add_arrayss z$DistributedTensorGatherer.add_arrayscs\t|ttfrHfddt||D}|ddt|dd|DfS|jdjdksztdjd|jdd|jdj}tjD]}t |jd kr||||d ||j |j ||<qt |jd kr|jd |jd krt ||jd j d }||||d ||j |j ||d|jd f<q||fS) Ncsg|]\}}||qSr)r)r;ryrrr rrszADistributedTensorGatherer._nested_set_tensors..rcss|]}|dVqdS)rNr)r;rrrr r>sz@DistributedTensorGatherer._nested_set_tensors..zNot all data has been set. Are you sure you passed all values?)rrrloggerwarningrrrrrr finalizes   z"DistributedTensorGatherer.finalize)Nr()rrrrrrrrrrrr rss - rc@s4eZdZUdZdZeed<dZeed<d ddZ d S) LabelSmoothera@ Adds label-smoothing on a pre-computed output from a Transformers model. Args: epsilon (`float`, *optional*, defaults to 0.1): The label smoothing factor. ignore_index (`int`, *optional*, defaults to -100): The index in the labels to ignore when computing the loss. g?epsilonr( ignore_indexFc Cst|tr|dn|d}|rL|dddddf}|dddf}tjj|dd }||dkr||d}||j }t j |dd}|j d|d}|j dd t jd }||d ||d || } | | }| | |jd}d|j||j|S) Nlogitsr.rrr))min)r*indexT)r*ZkeepdimrUg)r#dictrtr Z functionalZ log_softmaxr*Z unsqueezeeqrr$clampgathersumrWZ masked_fill_numellongr-r) rZ model_outputlabelsZ shift_labelsrZ log_probsZ padding_maskZnll_lossZ smoothed_lossZnum_active_elementsrrr __call__s"     zLabelSmoother.__call__N)F) rrrrrfloat__annotations__rrrrrrr rs   rcs|dkr*tt|dd}|dkr*d}tjt|d||fddtdtD}fd d|D}fd d|D}tt|}||d|dd|dd<||d<d d|DS) a Return a list of indices so that each slice of `batch_size` consecutive indices correspond to elements of similar lengths. To do this, the indices are: - randomly permuted - grouped in mega-batches of size `mega_batch_mult * batch_size` - sorted by length in each mega-batch The result is the concatenation of all mega-batches, with the batch of `batch_size` containing the element of maximum length placed first, so that an OOM happens sooner rather than later. N2rr generatorcs g|]}||qSr)tolist)r;r_)rmegabatch_sizerr rr sz.get_length_grouped_indices..cs"g|]}t|fddddqS)cs|SrPr)r_lengthsrr !z7get_length_grouped_indices...T)rNreverse)sortedr; megabatchrrr rr!scsg|]}|dqS)rrrrrr rr%scSsg|]}|D]}|q qSrr)r;rr_rrr rr*s)rr,r$ZrandpermruZargmaxrhitem)rrZmega_batch_multrZ megabatchesZmegabatch_maximumsZmax_idxr)rrrr get_length_grouped_indices s  *rc@sFeZdZdZd eeeeeeeedddZ ddZ dd Z dS) LengthGroupedSamplerz Sampler that samples indices in a way that groups together features of the dataset of roughly the same length while keeping a bit of randomness. N)rrrmodel_input_namecs|dkr|dkrtd||_|dkrdk r2ndt|dtsRt|dtr^|dkrntddfdd|D}nt|tjrtd| }||_ ||_ dS) N,One of dataset and lengths must be provided. input_idsrXCan only automatically infer lengths for datasets whose items are dictionaries with an '' key.csg|]}t|qSrr,r;featurerrr rrIsz1LengthGroupedSampler.__init__..zcIf lengths is a torch.Tensor, LengthGroupedSampler will be slow. Converting lengths to List[int]...) ValueErrorrr#rrr$r%rinforrr)rrrrrrrrr r3s,     zLengthGroupedSampler.__init__cCs t|jSrP)r,rrrrr rSszLengthGroupedSampler.__len__cCst|j|j|jd}t|SNr)rrrrrrrrr rVszLengthGroupedSampler.__iter__)NNNN) rrrrrr r r strrrrrrrr r-s  rc @sTeZdZdZd eeeeeeeeeeeeee dddZ e dd d Z dS) DistributedLengthGroupedSamplerz Distributed Sampler that samples indices in a way that groups together features of the dataset of roughly the same length while keeping a bit of randomness. NrF)rrrrseed drop_lastrrc sl|dkr|dkrtd|dkr8ts0tdt}|dkrXtsPtdt}||_||_||_d|_ ||_ |dkrڈdk rndt |dt st |dt r|dkrtddfdd|D}nt |tjrtd |}||_|j r.znIf lengths is a torch.Tensor, DistributedLengthGroupedSampler will be slow. Converting lengths to List[int]...)rrvrrrwrrrrepochrr#rrr$r%rrrrr,rrrrr) rrrrrrrrrrrr rbsL       z(DistributedLengthGroupedSampler.__init__)rjcCst}||j|jt|j|j|d}|jsN||d|j t |7}n|d|j }t ||j ksnt ||j |j |j }t ||jkst t|Sr)r$ Generator manual_seedrrrrrrrr,rErrrr)rgrrrr rsz(DistributedLengthGroupedSampler.__iter__)NNNrFNN) rrrrrr r boolr rrrrrrrr r[s&   If your IterableDataset implements some randomization that needs to be applied the same way on all processes (for instance, a shuffling), you should use a `torch.Generator` in a `generator` attribute of the `dataset` to generate your random numbers and call the [`~trainer_pt_utils.IterableDatasetShard.set_epoch`] method of this object. It will set the seed of this `generator` to `seed + epoch` on all processes before starting the iteration. Alternatively, you can also implement a `set_epoch()` method in your iterable dataset to deal with this. Args: dataset (`torch.utils.data.IterableDataset`): The batch sampler to split in several shards. batch_size (`int`, *optional*, defaults to 1): The size of the batches per shard. drop_last (`bool`, *optional*, defaults to `False`): Whether or not to drop the last incomplete batch or complete the last batches by using the samples from the beginning. num_processes (`int`, *optional*, defaults to 1): The number of processes running concurrently. process_index (`int`, *optional*, defaults to 0): The index of the current process. seed (`int`, *optional*, defaults to 0): A random seed that will be used for the random number generation in [`~trainer_pt_utils.IterableDatasetShard.set_epoch`]. rFr)rrrrrrcCs4||_||_||_||_||_||_d|_d|_dS)Nr)rrrrrrr num_examples)rrrrrrrrrr r s zIterableDatasetShard.__init__cCs"||_t|jdr|j|dS)N set_epoch)rrrr)rrrrr rs zIterableDatasetShard.set_epochccs&d|_t|jdsDt|jdrDt|jjtjrD|jj|j|j |j |j }t |j |j |j d|j }d}g}|jD]P}|jd7_||t||kr||D]}||Vq|dkr|}g}q||js"t|dkr"|dkr|}t||kr ||7}q|D]}||VqdS)Nrrrr)rrrr#rr$rrrrrrrurappendr,copyr)rZreal_batch_sizeZ process_sliceZ first_batchZ current_batchelementr_rrr r"s8        zIterableDatasetShard.__iter__cCsH|jr"t|j|j|j|jStt|j|j|j|jSdSrP)rr,rrrrrrrrr rCszIterableDatasetShard.__len__N)rFrrr) rrrrrrrrrrrrrrr rs"+ !rc Cs|jr\z|jd}Wqtk rX}z"dt|krFtdd}nW5d}~XYqXnDt|jtj jj r|j j dd}n|jd}t |r|}|S)Nrzneed to call stepzQtried to get lr value before scheduler/optimizer started stepping, returning lr=0lr)Zis_deepspeed_enabledZ lr_schedulerZ get_last_lrrErrrr#r$ZoptimZReduceLROnPlateauZ optimizerZ param_groupsZ is_tensorr)rZlast_lrerrr _get_learning_rateOs   r cCs4tt|t|d}tjt|dd|dS)zG convert seconds to hh:mm:ss.msec, msecs rounded to 2 decimals d)secondsr702d)rabsdatetime timedelta)ZsecsZmsecrrr _secs2timedeltafsr)metricsrjcCs|}|D]x\}}d|kr4|d?d||<qd|krJt|||<q|dkrjt|d?d||<qt||tkrt|d||<q|S) z Reformat Trainer metrics values to a human-readable format Args: metrics (`Dict[str, float]`): The metrics returned from train/evaluate/predict Returns: metrics (`Dict[str, float]`): The reformatted metrics Z_mem_MBZ_runtimeZ total_flosZGFr)rrIrrrDrround)rrZ metrics_copyr@vrrr metrics_formatos rcCs|s dStd|d||}tdd|D}tdd|D}t|D],}td|d|d ||d |q^dS) a@ Log metrics in a specially formatted way Under distributed environment this is done only for a process with rank 0. Args: split (`str`): Mode/split name: one of `train`, `eval`, `test` metrics (`Dict[str, float]`): The metrics returned from train/evaluate/predictmetrics: metrics dict Notes on memory reports: In order to get memory usage report you need to install `psutil`. You can do that with `pip install psutil`. Now when this method is run, you will see a report that will include: : ``` init_mem_cpu_alloc_delta = 1301MB init_mem_cpu_peaked_delta = 154MB init_mem_gpu_alloc_delta = 230MB init_mem_gpu_peaked_delta = 0MB train_mem_cpu_alloc_delta = 1345MB train_mem_cpu_peaked_delta = 0MB train_mem_gpu_alloc_delta = 693MB train_mem_gpu_peaked_delta = 7MB ``` **Understanding the reports:** - the first segment, e.g., `train__`, tells you which stage the metrics are for. Reports starting with `init_` will be added to the first stage that gets run. So that if only evaluation is run, the memory usage for the `__init__` will be reported along with the `eval_` metrics. - the third segment, is either `cpu` or `gpu`, tells you whether it's the general RAM or the gpu0 memory metric. - `*_alloc_delta` - is the difference in the used/allocated memory counter between the end and the start of the stage - it can be negative if a function released more memory than it allocated. - `*_peaked_delta` - is any extra memory that was consumed and then freed - relative to the current allocated memory counter - it is never negative. When you look at the metrics of any stage you add up `alloc_delta` + `peaked_delta` and you know how much memory was needed to complete that stage. The reporting happens only for process of rank 0 and gpu 0 (if there is a gpu). Typically this is enough since the main process does the bulk of work, but it could be not quite so if model parallel is used and then other GPUs may use a different amount of gpu memory. This is also not the same under DataParallel where gpu0 may require much more memory than the rest since it stores the gradient and optimizer states for all participating GPUS. Perhaps in the future these reports will evolve to measure those too. The CPU RAM metric measures RSS (Resident Set Size) includes both the memory which is unique to the process and the memory shared with other processes. It is important to note that it does not include swapped out memory, so the reports could be imprecise. The CPU peak memory is measured using a sampling thread. Due to python's GIL it may miss some of the peak memory if that thread didn't get a chance to run when the highest memory was used. Therefore this report can be less than reality. Using `tracemalloc` would have reported the exact peak memory, but it doesn't report memory allocations outside of python. So if some C++ CUDA extension allocated its own memory it won't be reported. And therefore it was dropped in favor of the memory sampling approach, which reads the current process memory usage. The GPU allocated and peak memory reporting is done with `torch.cuda.memory_allocated()` and `torch.cuda.max_memory_allocated()`. This metric reports only "deltas" for pytorch-specific allocations, as `torch.cuda` memory management system doesn't track any memory allocated outside of pytorch. For example, the very first cuda call typically loads CUDA kernels, which may take from 0.5 to 2GB of GPU memory. Note that this tracker doesn't account for memory allocations outside of [`Trainer`]'s `__init__`, `train`, `evaluate` and `predict` calls. Because `evaluation` calls may happen during `train`, we can't handle nested invocations because `torch.cuda.max_memory_allocated` is a single counter, so if it gets reset by a nested eval call, `train`'s tracker will report incorrect info. If this [pytorch issue](https://github.com/pytorch/pytorch/issues/16266) gets resolved it will be possible to change this class to be re-entrant. Until then we will only track the outer level of `train`, `evaluate` and `predict` methods. Which means that if `eval` is called during `train`, it's the latter that will account for its memory usage and that of the former. This also means that if any other tool that is used along the [`Trainer`] calls `torch.cuda.reset_peak_memory_stats`, the gpu peak memory stats could be invalid. And the [`Trainer`] will disrupt the normal behavior of any such tools that rely on calling `torch.cuda.reset_peak_memory_stats` themselves. For best performance you may want to consider turning the memory profiling off for production runs. Nz***** z metrics *****css|]}tt|VqdSrPr,rrrrr r>szlog_metrics..css|]}tt|VqdSrPrrrrr r>sz z )is_world_process_zeroprintrr/keysvaluesr)rsplitrZmetrics_formattedZk_widthZv_widthrNrrr log_metricssO r Tc Cs|s dStj|jj|d}t|d}tj||dddW5QRX|rtj|jjd}tj |rt|d}t |}W5QRXni}| |t|d}tj||dddW5QRXdS) a Save metrics into a json file for that split, e.g. `train_results.json`. Under distributed environment this is done only for a process with rank 0. Args: split (`str`): Mode/split name: one of `train`, `eval`, `test`, `all` metrics (`Dict[str, float]`): The metrics returned from train/evaluate/predict combined (`bool`, *optional*, defaults to `True`): Creates combined metrics by updating `all_results.json` with metrics of this call To understand the metrics please read the docstring of [`~Trainer.log_metrics`]. The only difference is that raw unformatted numbers are saved in the current method. Nz _results.jsonrrT)indent sort_keyszall_results.jsonr) rospathjoinargs output_diropenjsondumpexistsloadupdate)rrrcombinedr$fZ all_metricsrrr save_metricss     r0cCs.|s dStj|jjd}|j|dS)z Saves the Trainer state, since Trainer.save_model saves only the tokenizer with the model Under distributed environment this is done only for a process with rank 0. Nztrainer_state.json)rr#r$r%r&r'stateZ save_to_json)rr$rrr save_state sr2cs4trddnddtfdd|DS)zn Calculate model's total param count. If trainable_only is True then count only those requiring grads cSst|dr|jS|S)Nds_numel)rr3rprrr rsz$get_model_param_count..numelcSs|SrP)rr4rrr r!sc3s |]}r|jr|VqdSrP)Z requires_grad)r;r5rtrainable_onlyrr r>$sz(get_model_param_count..)rr parameters)modelr7rr6r get_model_param_counts r:csLg}|D](\|fddtD7}q |t|j7}|S)zZ Returns the names of the model parameters that are not inside a forbidden layer. cs(g|] }ttsd|qS)r7)r#rG)r;r=childforbidden_layer_typesrarr rr-sz'get_parameter_names..)Znamed_childrenget_parameter_namesrF _parametersr)r9r=r2rr;r r>'s r>cCsVt|}|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)rFchildrenrrr,get_module_class_from_name)moduleraZmodules_childrenZ child_moduleZ module_classrrr rA7s    rAcCs6|r2|D](}tj||}tj|rt|qdSrP)r#r$r%isfileremove)Zis_main_processr' filenamesfilenamefilerrr remove_dummy_checkpointKs  rHcCs:|f|}t|tr|dn|d}||}|||S)Nlossr)r#rZbackward)r9inputsZgradient_accumulation_stepsoutputsrIrrr smp_forward_backwardVs   rLcCs |f|SrPr)r9rJrrr smp_forward_only^srMcCst|ttfr$t|dd|DSt|trHt|dd|DSt|tjshtdt|dt |t j j }dd|D}tj d d|Dd d S) Ncss|]}t|VqdSrP smp_gatherrSrrr r>dszsmp_gather..cSsi|]\}}|t|qSrrNr;r@rrrr rCfszsmp_gather..z Can't gather the values of type z-, only of nested list/tuple/dicts of tensors.cSsg|] }t|qSr)r"rSrrr rrlszsmp_gather..cSsg|] }|qSr)rTrSrrr rrmsrr))r#rFrGrDrrIr$r%rKsmpZ allgatherZ CommGroupZDP_GROUPr.)rhZ all_tensorsrrr rObs  rOcCsXt|ttfr$t|dd|DSt|trHt|dd|DS|S)Ncss|]}t|VqdSrPsmp_nested_concatrSrrr r>qsz$smp_nested_concat..cSsi|]\}}|t|qSrrRrPrrr rCssz%smp_nested_concat..) r#rFrGrDrrIrzr[rTrqrrr rSos  rS)r()r()r()N)r()r()NN)T)F)r)krrr)rr#sysrcollections.abcr contextlibr dataclassesrrrtypingrrrr r r rXr'r$Ztorch.distributed distributedrvr Ztorch.utils.datar rrrZtorch.utils.data.distributedrZintegrations.deepspeedrZtokenization_utils_baserutilsrrr add_handlerstdoutrbrcrdrgZtorch.optim.lr_schedulerrrfZ get_loggerrrrr%rJr"r3r6r:rMrRrZr^rrlr}rrrrrrdatarrrrrrrrrrrr rrrr r0r2r:r>rArHZ!smdistributed.modelparallel.torchZ modelparallelrQsteprLrMrOrSrrrr s                  8   l* #.X/m Z &