codingutf8 2020present the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license torch utilities for the trainer class this is used to suppress an undesired warning emitted by pytorch versions 1 4 21 7 0 concatenates tensor1 and tensor2 on first axis applying padding on the second if necessary tensor1 atleast1dtensor1 tensor2 atleast1dtensor2 if lentensor1 shape 1 or tensor1 shape1 tensor2 shape1 return torch cattensor1 tensor2 dim0 let s figure out the new shape newshape tensor1 shape0 tensor2 shape0 maxtensor1 shape1 tensor2 shape1 tensor1 shape2 now let s fill the result tensor result tensor1 newfullnewshape paddingindex result tensor1 shape0 tensor1 shape1 tensor1 resulttensor1 shape0 tensor2 shape1 tensor2 return result def numpypadandconcatenatearray1 array2 paddingindex100 let s figure out the new shape now let s fill the result tensor concat the newtensors to tensors on the first dim and pad them on the second if needed works for tensors or nested listtuplesdict of tensors find the first dimension of a tensor in a nested listtupledict of tensors as of numpy 1 21 4 numpy does not support bfloat16 see https github comnumpynumpybloba47ecdea856986cd60eabbd53265c2ca5916ad5ddocsourceuserbasics types rst until numpy adds bfloat16 we must convert float32 truncate the dummy elements added by sequentialdistributedsampler truncate the dummy elements added by sequentialdistributedsampler reissue warnings that are not the savestatewarning decorator to make all processes in distributed training wait for each localmaster to do something args localrank int the rank of the local process 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 batchsize samples args dataset torch utils data dataset dataset used for sampling batchsize int the batch size used with this sampler kwargs dictstr any optional all other keyword arguments passed to distributedsampler distributedsampler already added samples from the beginning to make the number of samples a round multiple of the world size so we skip those 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 add extra samples to make numsamples a multiple of batchsize if passed add extra samples to make it evenly divisible subsample create the same nested structure as arrays with a first dimension always at numsamples if isinstancearrays list tuple return typearraysnestednewlikex numsamples for x in arrays return np fulllikearrays paddingindex shapenumsamples arrays shape1 def expandlikearrays newseqlength paddingindex100 a class responsible for properly gathering tensors or nested listtuple 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 listtuple 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 worldsize int the number of processes used in the distributed training numsamples int the number of samples in our dataset makemultipleof 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 paddingindex int optional defaults to 100 the padding index to use if the arrays don t all have the same sequence length 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 expand the array on the fly if needed return the properly gathered arrays and truncate to the number of samples since the sampler added some extras to get each process a dataset of the same length adds labelsmoothing on a precomputed output from a transformers model args epsilon float optional defaults to 0 1 the label smoothing factor ignoreindex int optional defaults to 100 the index in the labels to ignore when computing the loss in case the ignoreindex is 100 the gather will fail so we replace labels by 0 the paddingmask will ignore them in any case works for fp16 input tensor too by internally upcasting it to fp32 take the mean over the label dimensions then divide by the number of active elements i e notpadded return a list of indices so that each slice of batchsize consecutive indices correspond to elements of similar lengths to do this the indices are randomly permuted grouped in megabatches of size megabatchmult batchsize sorted by length in each megabatch the result is the concatenation of all megabatches with the batch of batchsize containing the element of maximum length placed first so that an oom happens sooner rather than later default for megabatchmult 50 or the number to get 4 megabatches whichever is smaller just in case for tiny datasets we need to use torch for the random part as a distributed sampler will set the random seed for torch the rest is to get the biggest batch first since each megabatch is sorted by descending length the longest element is the first switch to put the longest element in first position def init self batchsize int dataset optionaldataset none lengths optionallistint none modelinputname optionalstr none generatornone if dataset is none and lengths is none raise valueerrorone of dataset and lengths must be provided self batchsize batchsize if lengths is none modelinputname modelinputname if modelinputname is not none else inputids if not isinstancedataset0 dict or isinstancedataset0 batchencoding or modelinputname not in dataset0 raise valueerror can only automatically infer lengths for datasets whose items are dictionaries with an f modelinputname key lengths lenfeaturemodelinputname for feature in dataset elif isinstancelengths torch tensor logger info if lengths is a torch tensor lengthgroupedsampler will be slow converting lengths to listint lengths lengths tolist self lengths lengths self generator generator def lenself return lenself lengths def iterself indices getlengthgroupedindicesself lengths self batchsize generatorself generator return iterindices class distributedlengthgroupedsamplerdistributedsampler r 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 copied and adapted from pytorch distributedsampler if the dataset length is evenly divisible by of replicas then there is no need to drop any data since the dataset will be split equally split to nearest available length that is evenly divisible this is to ensure each rank receives the same amount of data when using this sampler deterministically shuffle based on epoch and seed add extra samples to make it evenly divisible remove tail of data to make it evenly divisible subsample sampler that shards batches between several processes dispatches indices batch by batch on 2 processes with batch size 4 the first two batches are 0 1 2 3 4 5 6 7 and 8 9 10 11 12 13 14 15 which shard into 0 1 2 3 and 8 9 10 11 for gpu0 and 4 5 6 7 and 12 13 14 15 for gpu1 the sampler thus yields 0 1 2 3 8 9 10 11 on gpu0 and 4 5 6 7 12 13 14 15 on gpu1 add extra samples to make it evenly divisible while loop is there in the edge case we have a tiny dataset and it needs to be done several times each shard only sees a fraction of totalnumsamples wraps a pytorch iterabledataset to generate samples for one of the processes only instances of this class will always yield a number of samples that is a round multiple of the actual batch size which is batchsize x numprocesses depending on the value of the droplast attribute it will either stop the iteration at the first batch that would be too small or loop with indices from the beginning on two processes with an iterable dataset yielding of 0 1 2 3 4 5 6 7 8 9 10 11 with a batch size of 2 the shard on process 0 will yield 0 1 4 5 8 9 so will see batches 0 1 4 5 8 9 the shard on process 1 will yield 2 3 6 7 10 11 so will see batches 2 3 6 7 10 11 tip warningtrue 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 trainerptutils iterabledatasetshard setepoch 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 setepoch method in your iterable dataset to deal with this tip args dataset torch utils data iterabledataset the batch sampler to split in several shards batchsize int optional defaults to 1 the size of the batches per shard droplast 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 numprocesses int optional defaults to 1 the number of processes running concurrently processindex 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 trainerptutils iterabledatasetshard setepoch wait to have a full batch before yielding elements finished if droplast is true otherwise complete the last batch with elements from the beginning will raise an error if the underlying dataset is not sized in order to keep trainer py compact and easy to understand place any secondary pt trainer helper methods here with deepspeed s fp16 and dynamic loss scale enabled the optimizerscheduler steps may not run for the first few dozen steps while loss scale is too large and thus during that time getlastlr will fail if called during that warm up stage so work around it convert seconds to hh mm ss msec msecs rounded to 2 decimals reformat trainer metrics values to a humanreadable format args metrics dictstr float the metrics returned from trainevaluatepredict returns metrics dictstr float the reformatted metrics log metrics in a specially formatted way under distributed environment this is done only for a process with rank 0 args split str modesplit name one of train eval test metrics dictstr float the metrics returned from trainevaluatepredictmetrics 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 initmemcpuallocdelta 1301mb initmemcpupeakeddelta 154mb initmemgpuallocdelta 230mb initmemgpupeakeddelta 0mb trainmemcpuallocdelta 1345mb trainmemcpupeakeddelta 0mb trainmemgpuallocdelta 693mb trainmemgpupeakeddelta 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 allocdelta is the difference in the usedallocated memory counter between the end and the start of the stage it can be negative if a function released more memory than it allocated peakeddelta 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 allocdelta peakeddelta 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 memoryallocated and torch cuda maxmemoryallocated this metric reports only deltas for pytorchspecific 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 maxmemoryallocated is a single counter so if it gets reset by a nested eval call train s tracker will report incorrect info if this pytorch issuehttps github compytorchpytorchissues16266 gets resolved it will be possible to change this class to be reentrant 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 resetpeakmemorystats 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 resetpeakmemorystats themselves for best performance you may want to consider turning the memory profiling off for production runs save metrics into a json file for that split e g trainresults json under distributed environment this is done only for a process with rank 0 args split str modesplit name one of train eval test all metrics dictstr float the metrics returned from trainevaluatepredict combined bool optional defaults to true creates combined metrics by updating allresults json with metrics of this call to understand the metrics please read the docstring of trainer logmetrics the only difference is that raw unformatted numbers are saved in the current method saves the trainer state since trainer savemodel saves only the tokenizer with the model under distributed environment this is done only for a process with rank 0 calculate model s total param count if trainableonly is true then count only those requiring grads returns the names of the model parameters that are not inside a forbidden layer add model specific parameters defined with nn parameter since they are not in any child 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 it doesn t seem possible to check here if tensor is a stepoutput because stepoutput lives in smp step which is also the name of the decorator so python is confused coding utf 8 2020 present the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license torch utilities for the trainer class this is used to suppress an undesired warning emitted by pytorch versions 1 4 2 1 7 0 concatenates tensor1 and tensor2 on first axis applying padding on the second if necessary let s figure out the new shape now let s fill the result tensor concatenates array1 and array2 on first axis applying padding on the second if necessary let s figure out the new shape now let s fill the result tensor 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 find the first dimension of a tensor in a nested list tuple dict of tensors as of numpy 1 21 4 numpy does not support bfloat16 see https github com numpy numpy blob a47ecdea856986cd60eabbd53265c2ca5916ad5d doc source user basics types rst until numpy adds bfloat16 we must convert float32 truncate the dummy elements added by sequentialdistributedsampler truncate the dummy elements added by sequentialdistributedsampler reissue warnings that are not the save_state_warning 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 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 distributedsampler already added samples from the beginning to make the number of samples a round multiple of the world size so we skip those 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 add extra samples to make num_samples a multiple of batch_size if passed add extra samples to make it evenly divisible subsample create the same nested structure as arrays with a first dimension always at num_samples expand the arrays so that the second dimension grows to new_seq_length uses padding_index for padding 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 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 expand the array on the fly if needed return the properly gathered arrays and truncate to the number of samples since the sampler added some extras to get each process a dataset of the same length 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 in case the ignore_index is 100 the gather will fail so we replace labels by 0 the padding_mask will ignore them in any case works for fp16 input tensor too by internally upcasting it to fp32 take the mean over the label dimensions then divide by the number of active elements i e not padded 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 default for mega_batch_mult 50 or the number to get 4 megabatches whichever is smaller just in case for tiny datasets we need to use torch for the random part as a distributed sampler will set the random seed for torch the rest is to get the biggest batch first since each megabatch is sorted by descending length the longest element is the first switch to put the longest element in first position 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 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 copied and adapted from pytorch distributedsampler if the dataset length is evenly divisible by of replicas then there is no need to drop any data since the dataset will be split equally split to nearest available length that is evenly divisible this is to ensure each rank receives the same amount of data when using this sampler deterministically shuffle based on epoch and seed add extra samples to make it evenly divisible remove tail of data to make it evenly divisible subsample sampler that shards batches between several processes dispatches indices batch by batch on 2 processes with batch size 4 the first two batches are 0 1 2 3 4 5 6 7 and 8 9 10 11 12 13 14 15 which shard into 0 1 2 3 and 8 9 10 11 for gpu 0 and 4 5 6 7 and 12 13 14 15 for gpu 1 the sampler thus yields 0 1 2 3 8 9 10 11 on gpu 0 and 4 5 6 7 12 13 14 15 on gpu 1 add extra samples to make it evenly divisible while loop is there in the edge case we have a tiny dataset and it needs to be done several times each shard only sees a fraction of total_num_samples wraps a pytorch iterabledataset to generate samples for one of the processes only instances of this class will always yield a number of samples that is a round multiple of the actual batch size which is batch_size x num_processes depending on the value of the drop_last attribute it will either stop the iteration at the first batch that would be too small or loop with indices from the beginning on two processes with an iterable dataset yielding of 0 1 2 3 4 5 6 7 8 9 10 11 with a batch size of 2 the shard on process 0 will yield 0 1 4 5 8 9 so will see batches 0 1 4 5 8 9 the shard on process 1 will yield 2 3 6 7 10 11 so will see batches 2 3 6 7 10 11 tip warning true 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 tip 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 wait to have a full batch before yielding elements finished if drop_last is true otherwise complete the last batch with elements from the beginning will raise an error if the underlying dataset is not sized in order to keep trainer py compact and easy to understand place any secondary pt trainer helper methods here with deepspeed s fp16 and dynamic loss scale enabled the optimizer scheduler steps may not run for the first few dozen steps while loss scale is too large and thus during that time get_last_lr will fail if called during that warm up stage so work around it convert seconds to hh mm ss msec msecs rounded to 2 decimals 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 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 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 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 calculate model s total param count if trainable_only is true then count only those requiring grads returns the names of the model parameters that are not inside a forbidden layer add model specific parameters defined with nn parameter since they are not in any child 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 it doesn t seem possible to check here if tensor is a stepoutput because stepoutput lives in smp step which is also the name of the decorator so python is confused
import datetime
import json
import math
import os
import sys
import warnings
from collections.abc import Mapping
from contextlib import contextmanager
from dataclasses import dataclass
from logging import StreamHandler
from typing import Any, Dict, Iterator, List, Optional, Union

import numpy as np
import torch
import torch.distributed as dist
from torch import nn
from torch.utils.data import Dataset, IterableDataset, RandomSampler, Sampler
from torch.utils.data.distributed import DistributedSampler

from .integrations.deepspeed import is_deepspeed_zero3_enabled
from .tokenization_utils_base import BatchEncoding
from .utils import is_sagemaker_mp_enabled, is_torch_tpu_available, is_training_run_on_sagemaker, logging


if is_training_run_on_sagemaker():
    logging.add_handler(StreamHandler(sys.stdout))

if is_torch_tpu_available(check_device=False):
    import torch_xla.core.xla_model as xm


try:
    from torch.optim.lr_scheduler import SAVE_STATE_WARNING
except ImportError:
    SAVE_STATE_WARNING = ""

logger = logging.get_logger(__name__)


def get_dataloader_sampler(dataloader):
    if hasattr(dataloader, "batch_sampler") and dataloader.batch_sampler is not None:
        return get_dataloader_sampler(dataloader.batch_sampler)
    elif hasattr(dataloader, "sampler"):
        return dataloader.sampler


def atleast_1d(tensor_or_array: Union[torch.Tensor, np.ndarray]):
    if isinstance(tensor_or_array, torch.Tensor):
        if hasattr(torch, "atleast_1d"):
            tensor_or_array = torch.atleast_1d(tensor_or_array)
        elif tensor_or_array.ndim < 1:
            tensor_or_array = tensor_or_array[None]
    else:
        tensor_or_array = np.atleast_1d(tensor_or_array)
    return tensor_or_array


def torch_pad_and_concatenate(tensor1, tensor2, padding_index=-100):
    
    tensor1 = atleast_1d(tensor1)
    tensor2 = atleast_1d(tensor2)

    if len(tensor1.shape) == 1 or tensor1.shape[1] == tensor2.shape[1]:
        return torch.cat((tensor1, tensor2), dim=0)

    
    new_shape = (tensor1.shape[0] + tensor2.shape[0], max(tensor1.shape[1], tensor2.shape[1])) + tensor1.shape[2:]

    
    result = tensor1.new_full(new_shape, padding_index)
    result[: tensor1.shape[0], : tensor1.shape[1]] = tensor1
    result[tensor1.shape[0] :, : tensor2.shape[1]] = tensor2
    return result


def numpy_pad_and_concatenate(array1, array2, padding_index=-100):
    
    array1 = atleast_1d(array1)
    array2 = atleast_1d(array2)

    if len(array1.shape) == 1 or array1.shape[1] == array2.shape[1]:
        return np.concatenate((array1, array2), axis=0)

    
    new_shape = (array1.shape[0] + array2.shape[0], max(array1.shape[1], array2.shape[1])) + array1.shape[2:]

    
    result = np.full_like(array1, padding_index, shape=new_shape)
    result[: array1.shape[0], : array1.shape[1]] = array1
    result[array1.shape[0] :, : array2.shape[1]] = array2
    return result


def nested_concat(tensors, new_tensors, padding_index=-100):
    
    assert type(tensors) == type(
        new_tensors
    ), f"Expected `tensors` and `new_tensors` to have the same type but found {type(tensors)} and {type(new_tensors)}."
    if isinstance(tensors, (list, tuple)):
        return type(tensors)(nested_concat(t, n, padding_index=padding_index) for t, n in zip(tensors, new_tensors))
    elif isinstance(tensors, torch.Tensor):
        return torch_pad_and_concatenate(tensors, new_tensors, padding_index=padding_index)
    elif isinstance(tensors, Mapping):
        return type(tensors)(
            {k: nested_concat(t, new_tensors[k], padding_index=padding_index) for k, t in tensors.items()}
        )
    elif isinstance(tensors, np.ndarray):
        return numpy_pad_and_concatenate(tensors, new_tensors, padding_index=padding_index)
    else:
        raise TypeError(f"Unsupported type for concatenation: got {type(tensors)}")


def find_batch_size(tensors):
    
    if isinstance(tensors, (list, tuple)):
        for t in tensors:
            result = find_batch_size(t)
            if result is not None:
                return result
    elif isinstance(tensors, Mapping):
        for key, value in tensors.items():
            result = find_batch_size(value)
            if result is not None:
                return result
    elif isinstance(tensors, torch.Tensor):
        return tensors.shape[0] if len(tensors.shape) >= 1 else None
    elif isinstance(tensors, np.ndarray):
        return tensors.shape[0] if len(tensors.shape) >= 1 else None


def nested_numpify(tensors):
    "Numpify `tensors` (even if it's a nested list/tuple/dict of tensors)."
    if isinstance(tensors, (list, tuple)):
        return type(tensors)(nested_numpify(t) for t in tensors)
    if isinstance(tensors, Mapping):
        return type(tensors)({k: nested_numpify(t) for k, t in tensors.items()})

    t = tensors.cpu()
    if t.dtype == torch.bfloat16:
        
        
        
        t = t.to(torch.float32)
    return t.numpy()


def nested_detach(tensors):
    "Detach `tensors` (even if it's a nested list/tuple/dict of tensors)."
    if isinstance(tensors, (list, tuple)):
        return type(tensors)(nested_detach(t) for t in tensors)
    elif isinstance(tensors, Mapping):
        return type(tensors)({k: nested_detach(t) for k, t in tensors.items()})
    return tensors.detach()


def nested_xla_mesh_reduce(tensors, name):
    if is_torch_tpu_available():
        import torch_xla.core.xla_model as xm

        if isinstance(tensors, (list, tuple)):
            return type(tensors)(nested_xla_mesh_reduce(t, f"{name}_{i}") for i, t in enumerate(tensors))
        if isinstance(tensors, Mapping):
            return type(tensors)(
                {k: nested_xla_mesh_reduce(t, f"{name}_{i}") for i, (k, t) in enumerate(tensors.items())}
            )

        tensors = atleast_1d(tensors)
        return xm.mesh_reduce(name, tensors, torch.cat)
    else:
        raise ImportError("Torch xla must be installed to use `nested_xla_mesh_reduce`")


def distributed_concat(tensor: Any, num_total_examples: Optional[int] = None) -> Any:
    try:
        if isinstance(tensor, (tuple, list)):
            return type(tensor)(distributed_concat(t, num_total_examples) for t in tensor)
        if isinstance(tensor, Mapping):
            return type(tensor)({k: distributed_concat(t, num_total_examples) for k, t in tensor.items()})
        tensor = atleast_1d(tensor).contiguous()
        output_tensors = [tensor.clone() for _ in range(dist.get_world_size())]
        dist.all_gather(output_tensors, tensor)
        concat = torch.cat(output_tensors, dim=0)

        
        if num_total_examples is not None:
            concat = concat[:num_total_examples]
        return concat
    except AssertionError:
        raise AssertionError("Not currently using distributed training")


def distributed_broadcast_scalars(
    scalars: List[Union[int, float]],
    num_total_examples: Optional[int] = None,
    device: Optional[torch.device] = torch.device("cuda"),
) -> torch.Tensor:
    try:
        tensorized_scalar = torch.tensor(scalars).to(device)
        output_tensors = [tensorized_scalar.clone() for _ in range(dist.get_world_size())]
        dist.all_gather(output_tensors, tensorized_scalar)
        concat = torch.cat(output_tensors, dim=0)

        
        if num_total_examples is not None:
            concat = concat[:num_total_examples]
        return concat
    except AssertionError:
        raise AssertionError("Not currently using distributed training")


def reissue_pt_warnings(caught_warnings):
    
    if len(caught_warnings) > 1:
        for w in caught_warnings:
            if w.category != UserWarning or w.message != SAVE_STATE_WARNING:
                warnings.warn(w.message, w.category)


@contextmanager
def torch_distributed_zero_first(local_rank: int):
    
    if local_rank not in [-1, 0]:
        dist.barrier()
    yield
    if local_rank == 0:
        dist.barrier()


class DistributedSamplerWithLoop(DistributedSampler):
    

    def __init__(self, dataset, batch_size, **kwargs):
        super().__init__(dataset, **kwargs)
        self.batch_size = batch_size

    def __iter__(self):
        indices = list(super().__iter__())
        remainder = 0 if len(indices) % self.batch_size == 0 else self.batch_size - len(indices) % self.batch_size
        
        
        start_remainder = 1 if self.rank < len(self.dataset) % self.num_replicas else 0
        indices += indices[start_remainder : start_remainder + remainder]
        return iter(indices)


class SequentialDistributedSampler(Sampler):
    

    def __init__(self, dataset, num_replicas=None, rank=None, batch_size=None):
        warnings.warn(
            "SequentialDistributedSampler is deprecated and will be removed in v5 of Transformers.",
            FutureWarning,
        )
        if num_replicas is None:
            if not dist.is_available():
                raise RuntimeError("Requires distributed package to be available")
            num_replicas = dist.get_world_size()
        if rank is None:
            if not dist.is_available():
                raise RuntimeError("Requires distributed package to be available")
            rank = dist.get_rank()
        self.dataset = dataset
        self.num_replicas = num_replicas
        self.rank = rank
        num_samples = len(self.dataset)
        
        if batch_size is not None:
            self.num_samples = int(math.ceil(num_samples / (batch_size * num_replicas))) * batch_size
        else:
            self.num_samples = int(math.ceil(num_samples / num_replicas))
        self.total_size = self.num_samples * self.num_replicas
        self.batch_size = batch_size

    def __iter__(self):
        indices = list(range(len(self.dataset)))

        
        indices += indices[: (self.total_size - len(indices))]
        assert (
            len(indices) == self.total_size
        ), f"Indices length {len(indices)} and total size {self.total_size} mismatched"

        
        indices = indices[self.rank * self.num_samples : (self.rank + 1) * self.num_samples]
        assert (
            len(indices) == self.num_samples
        ), f"Indices length {len(indices)} and sample number {self.num_samples} mismatched"

        return iter(indices)

    def __len__(self):
        return self.num_samples


def get_tpu_sampler(dataset: torch.utils.data.Dataset, batch_size: int):
    if xm.xrt_world_size() <= 1:
        return RandomSampler(dataset)
    return DistributedSampler(dataset, num_replicas=xm.xrt_world_size(), rank=xm.get_ordinal())


def nested_new_like(arrays, num_samples, padding_index=-100):
    
    if isinstance(arrays, (list, tuple)):
        return type(arrays)(nested_new_like(x, num_samples) for x in arrays)
    return np.full_like(arrays, padding_index, shape=(num_samples, *arrays.shape[1:]))


def expand_like(arrays, new_seq_length, padding_index=-100):
    
    result = np.full_like(arrays, padding_index, shape=(arrays.shape[0], new_seq_length) + arrays.shape[2:])
    result[:, : arrays.shape[1]] = arrays
    return result


def nested_truncate(tensors, limit):
    "Truncate `tensors` at `limit` (even if it's a nested list/tuple/dict of tensors)."
    if isinstance(tensors, (list, tuple)):
        return type(tensors)(nested_truncate(t, limit) for t in tensors)
    if isinstance(tensors, Mapping):
        return type(tensors)({k: nested_truncate(t, limit) for k, t in tensors.items()})

    return tensors[:limit]


class DistributedTensorGatherer:
    

    def __init__(self, world_size, num_samples, make_multiple_of=None, padding_index=-100):
        warnings.warn(
            "DistributedTensorGatherer is deprecated and will be removed in v5 of Transformers.",
            FutureWarning,
        )
        self.world_size = world_size
        self.num_samples = num_samples
        total_size = world_size if make_multiple_of is None else world_size * make_multiple_of
        self.total_samples = int(np.ceil(num_samples / total_size)) * total_size
        self.process_length = self.total_samples // world_size
        self._storage = None
        self._offsets = None
        self.padding_index = padding_index

    def add_arrays(self, arrays):
        
        if arrays is None:
            return
        if self._storage is None:
            self._storage = nested_new_like(arrays, self.total_samples, padding_index=self.padding_index)
            self._offsets = list(range(0, self.total_samples, self.process_length))

        slice_len, self._storage = self._nested_set_tensors(self._storage, arrays)
        for i in range(self.world_size):
            self._offsets[i] += slice_len

    def _nested_set_tensors(self, storage, arrays):
        if isinstance(arrays, (list, tuple)):
            result = [self._nested_set_tensors(x, y) for x, y in zip(storage, arrays)]
            return result[0][0], type(arrays)(r[1] for r in result)
        assert (
            arrays.shape[0] % self.world_size == 0
        ), f"Arrays passed should all have a first dimension multiple of {self.world_size}, found {arrays.shape[0]}."

        slice_len = arrays.shape[0] // self.world_size
        for i in range(self.world_size):
            if len(arrays.shape) == 1:
                storage[self._offsets[i] : self._offsets[i] + slice_len] = arrays[i * slice_len : (i + 1) * slice_len]
            else:
                
                if len(storage.shape) > 1 and storage.shape[1] < arrays.shape[1]:
                    storage = expand_like(storage, arrays.shape[1], padding_index=self.padding_index)
                storage[self._offsets[i] : self._offsets[i] + slice_len, : arrays.shape[1]] = arrays[
                    i * slice_len : (i + 1) * slice_len
                ]
        return slice_len, storage

    def finalize(self):
        
        if self._storage is None:
            return
        if self._offsets[0] != self.process_length:
            logger.warning("Not all data has been set. Are you sure you passed all values?")
        return nested_truncate(self._storage, self.num_samples)


@dataclass
class LabelSmoother:
    

    epsilon: float = 0.1
    ignore_index: int = -100

    def __call__(self, model_output, labels, shift_labels=False):
        logits = model_output["logits"] if isinstance(model_output, dict) else model_output[0]
        if shift_labels:
            logits = logits[..., :-1, :].contiguous()
            labels = labels[..., 1:].contiguous()

        log_probs = -nn.functional.log_softmax(logits, dim=-1)
        if labels.dim() == log_probs.dim() - 1:
            labels = labels.unsqueeze(-1)

        padding_mask = labels.eq(self.ignore_index)
        
        
        labels = torch.clamp(labels, min=0)
        nll_loss = log_probs.gather(dim=-1, index=labels)
        
        smoothed_loss = log_probs.sum(dim=-1, keepdim=True, dtype=torch.float32)

        nll_loss.masked_fill_(padding_mask, 0.0)
        smoothed_loss.masked_fill_(padding_mask, 0.0)

        
        num_active_elements = padding_mask.numel() - padding_mask.long().sum()
        nll_loss = nll_loss.sum() / num_active_elements
        smoothed_loss = smoothed_loss.sum() / (num_active_elements * log_probs.shape[-1])
        return (1 - self.epsilon) * nll_loss + self.epsilon * smoothed_loss


def get_length_grouped_indices(lengths, batch_size, mega_batch_mult=None, generator=None):
    
    
    if mega_batch_mult is None:
        mega_batch_mult = min(len(lengths) // (batch_size * 4), 50)
        
        if mega_batch_mult == 0:
            mega_batch_mult = 1

    
    indices = torch.randperm(len(lengths), generator=generator)
    megabatch_size = mega_batch_mult * batch_size
    megabatches = [indices[i : i + megabatch_size].tolist() for i in range(0, len(lengths), megabatch_size)]
    megabatches = [sorted(megabatch, key=lambda i: lengths[i], reverse=True) for megabatch in megabatches]

    
    
    megabatch_maximums = [lengths[megabatch[0]] for megabatch in megabatches]
    max_idx = torch.argmax(torch.tensor(megabatch_maximums)).item()
    
    megabatches[0][0], megabatches[max_idx][0] = megabatches[max_idx][0], megabatches[0][0]

    return [i for megabatch in megabatches for i in megabatch]


class LengthGroupedSampler(Sampler):
    r

    def __init__(
        self,
        batch_size: int,
        dataset: Optional[Dataset] = None,
        lengths: Optional[List[int]] = None,
        model_input_name: Optional[str] = None,
        generator=None,
    ):
        if dataset is None and lengths is None:
            raise ValueError("One of dataset and lengths must be provided.")

        self.batch_size = batch_size
        if lengths is None:
            model_input_name = model_input_name if model_input_name is not None else "input_ids"
            if (
                not (isinstance(dataset[0], dict) or isinstance(dataset[0], BatchEncoding))
                or model_input_name not in dataset[0]
            ):
                raise ValueError(
                    "Can only automatically infer lengths for datasets whose items are dictionaries with an "
                    f"'{model_input_name}' key."
                )
            lengths = [len(feature[model_input_name]) for feature in dataset]
        elif isinstance(lengths, torch.Tensor):
            logger.info(
                "If lengths is a torch.Tensor, LengthGroupedSampler will be slow. Converting lengths to List[int]..."
            )
            lengths = lengths.tolist()

        self.lengths = lengths
        self.generator = generator

    def __len__(self):
        return len(self.lengths)

    def __iter__(self):
        indices = get_length_grouped_indices(self.lengths, self.batch_size, generator=self.generator)
        return iter(indices)


class DistributedLengthGroupedSampler(DistributedSampler):
    r

    
    def __init__(
        self,
        batch_size: int,
        dataset: Optional[Dataset] = None,
        num_replicas: Optional[int] = None,
        rank: Optional[int] = None,
        seed: int = 0,
        drop_last: bool = False,
        lengths: Optional[List[int]] = None,
        model_input_name: Optional[str] = None,
    ):
        if dataset is None and lengths is None:
            raise ValueError("One of dataset and lengths must be provided.")
        if num_replicas is None:
            if not dist.is_available():
                raise RuntimeError("Requires distributed package to be available")
            num_replicas = dist.get_world_size()
        if rank is None:
            if not dist.is_available():
                raise RuntimeError("Requires distributed package to be available")
            rank = dist.get_rank()

        self.batch_size = batch_size
        self.num_replicas = num_replicas
        self.rank = rank
        self.epoch = 0
        self.drop_last = drop_last

        if lengths is None:
            model_input_name = model_input_name if model_input_name is not None else "input_ids"
            if (
                not (isinstance(dataset[0], dict) or isinstance(dataset[0], BatchEncoding))
                or model_input_name not in dataset[0]
            ):
                raise ValueError(
                    "Can only automatically infer lengths for datasets whose items are dictionaries with an "
                    f"'{model_input_name}' key."
                )
            lengths = [len(feature[model_input_name]) for feature in dataset]
        elif isinstance(lengths, torch.Tensor):
            logger.info(
                "If lengths is a torch.Tensor, DistributedLengthGroupedSampler will be slow. Converting lengths to"
                " List[int]..."
            )
            lengths = lengths.tolist()

        self.lengths = lengths

        
        
        if self.drop_last and len(self.lengths) % self.num_replicas != 0:
            
            
            
            self.num_samples = math.ceil((len(self.lengths) - self.num_replicas) / self.num_replicas)
        else:
            self.num_samples = math.ceil(len(self.lengths) / self.num_replicas)
        self.total_size = self.num_samples * self.num_replicas
        self.seed = seed

    def __iter__(self) -> Iterator:
        
        g = torch.Generator()
        g.manual_seed(self.seed + self.epoch)
        indices = get_length_grouped_indices(self.lengths, self.batch_size, generator=g)

        if not self.drop_last:
            
            indices += indices[: (self.total_size - len(indices))]
        else:
            
            indices = indices[: self.total_size]
        assert len(indices) == self.total_size

        
        indices = indices[self.rank : self.total_size : self.num_replicas]
        assert len(indices) == self.num_samples

        return iter(indices)


class ShardSampler(Sampler):
    

    def __init__(
        self,
        dataset: Dataset,
        batch_size: int = 1,
        drop_last: bool = False,
        num_processes: int = 1,
        process_index: int = 0,
    ):
        self.dataset = dataset
        self.batch_size = batch_size
        self.drop_last = drop_last
        self.num_processes = num_processes
        self.process_index = process_index

        self.total_batch_size = total_batch_size = batch_size * num_processes

        num_batches = len(dataset) // total_batch_size if drop_last else math.ceil(len(dataset) / total_batch_size)
        self.total_num_samples = num_batches * total_batch_size

    def __iter__(self):
        indices = list(range(len(self.dataset)))

        
        
        while len(indices) < self.total_num_samples:
            indices += indices[: (self.total_num_samples - len(indices))]

        result = []
        for batch_start in range(self.batch_size * self.process_index, self.total_num_samples, self.total_batch_size):
            result += indices[batch_start : batch_start + self.batch_size]

        return iter(result)

    def __len__(self):
        
        return self.total_num_samples // self.num_processes


class IterableDatasetShard(IterableDataset):
    

    def __init__(
        self,
        dataset: IterableDataset,
        batch_size: int = 1,
        drop_last: bool = False,
        num_processes: int = 1,
        process_index: int = 0,
        seed: int = 0,
    ):
        self.dataset = dataset
        self.batch_size = batch_size
        self.drop_last = drop_last
        self.num_processes = num_processes
        self.process_index = process_index
        self.seed = seed
        self.epoch = 0
        self.num_examples = 0

    def set_epoch(self, epoch):
        self.epoch = epoch
        if hasattr(self.dataset, "set_epoch"):
            self.dataset.set_epoch(epoch)

    def __iter__(self):
        self.num_examples = 0
        if (
            not hasattr(self.dataset, "set_epoch")
            and hasattr(self.dataset, "generator")
            and isinstance(self.dataset.generator, torch.Generator)
        ):
            self.dataset.generator.manual_seed(self.seed + self.epoch)
        real_batch_size = self.batch_size * self.num_processes
        process_slice = range(self.process_index * self.batch_size, (self.process_index + 1) * self.batch_size)

        first_batch = None
        current_batch = []
        for element in self.dataset:
            self.num_examples += 1
            current_batch.append(element)
            
            if len(current_batch) == real_batch_size:
                for i in process_slice:
                    yield current_batch[i]
                if first_batch is None:
                    first_batch = current_batch.copy()
                current_batch = []

        
        if not self.drop_last and len(current_batch) > 0:
            if first_batch is None:
                first_batch = current_batch.copy()
            while len(current_batch) < real_batch_size:
                current_batch += first_batch
            for i in process_slice:
                yield current_batch[i]

    def __len__(self):
        
        if self.drop_last:
            return (len(self.dataset) // (self.batch_size * self.num_processes)) * self.batch_size
        else:
            return math.ceil(len(self.dataset) / (self.batch_size * self.num_processes)) * self.batch_size






def _get_learning_rate(self):
    if self.is_deepspeed_enabled:
        
        
        
        try:
            last_lr = self.lr_scheduler.get_last_lr()[0]
        except AssertionError as e:
            if "need to call step" in str(e):
                logger.warning("tried to get lr value before scheduler/optimizer started stepping, returning lr=0")
                last_lr = 0
            else:
                raise
    else:
        if isinstance(self.lr_scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
            last_lr = self.optimizer.param_groups[0]["lr"]
        else:
            last_lr = self.lr_scheduler.get_last_lr()[0]
        if torch.is_tensor(last_lr):
            last_lr = last_lr.item()
    return last_lr


def _secs2timedelta(secs):
    

    msec = int(abs(secs - int(secs)) * 100)
    return f"{datetime.timedelta(seconds=int(secs))}.{msec:02d}"


def metrics_format(self, metrics: Dict[str, float]) -> Dict[str, float]:
    

    metrics_copy = metrics.copy()
    for k, v in metrics_copy.items():
        if "_mem_" in k:
            metrics_copy[k] = f"{ v >> 20 }MB"
        elif "_runtime" in k:
            metrics_copy[k] = _secs2timedelta(v)
        elif k == "total_flos":
            metrics_copy[k] = f"{ int(v) >> 30 }GF"
        elif isinstance(metrics_copy[k], float):
            metrics_copy[k] = round(v, 4)

    return metrics_copy


def log_metrics(self, split, metrics):
    
    if not self.is_world_process_zero():
        return

    print(f"***** {split} metrics *****")
    metrics_formatted = self.metrics_format(metrics)
    k_width = max(len(str(x)) for x in metrics_formatted.keys())
    v_width = max(len(str(x)) for x in metrics_formatted.values())
    for key in sorted(metrics_formatted.keys()):
        print(f"  {key: <{k_width}} = {metrics_formatted[key]:>{v_width}}")


def save_metrics(self, split, metrics, combined=True):
    
    if not self.is_world_process_zero():
        return

    path = os.path.join(self.args.output_dir, f"{split}_results.json")
    with open(path, "w") as f:
        json.dump(metrics, f, indent=4, sort_keys=True)

    if combined:
        path = os.path.join(self.args.output_dir, "all_results.json")
        if os.path.exists(path):
            with open(path, "r") as f:
                all_metrics = json.load(f)
        else:
            all_metrics = {}

        all_metrics.update(metrics)
        with open(path, "w") as f:
            json.dump(all_metrics, f, indent=4, sort_keys=True)


def save_state(self):
    
    if not self.is_world_process_zero():
        return

    path = os.path.join(self.args.output_dir, "trainer_state.json")
    self.state.save_to_json(path)


def get_model_param_count(model, trainable_only=False):
    
    if is_deepspeed_zero3_enabled():

        def numel(p):
            return p.ds_numel if hasattr(p, "ds_numel") else p.numel()

    else:

        def numel(p):
            return p.numel()

    return sum(numel(p) for p in model.parameters() if not trainable_only or p.requires_grad)


def get_parameter_names(model, forbidden_layer_types):
    
    result = []
    for name, child in model.named_children():
        result += [
            f"{name}.{n}"
            for n in get_parameter_names(child, forbidden_layer_types)
            if not isinstance(child, tuple(forbidden_layer_types))
        ]
    
    result += list(model._parameters.keys())
    return result


def get_module_class_from_name(module, name):
    
    modules_children = list(module.children())
    if module.__class__.__name__ == name:
        return module.__class__
    elif len(modules_children) == 0:
        return
    else:
        for child_module in modules_children:
            module_class = get_module_class_from_name(child_module, name)
            if module_class is not None:
                return module_class


def remove_dummy_checkpoint(is_main_process, output_dir, filenames):
    if is_main_process:
        for filename in filenames:
            file = os.path.join(output_dir, filename)
            if os.path.isfile(file):
                os.remove(file)


if is_sagemaker_mp_enabled():
    import smdistributed.modelparallel.torch as smp

    @smp.step()
    def smp_forward_backward(model, inputs, gradient_accumulation_steps=1):
        outputs = model(**inputs)
        loss = outputs["loss"] if isinstance(outputs, dict) else outputs[0]
        loss /= gradient_accumulation_steps
        model.backward(loss)
        return loss

    @smp.step()
    def smp_forward_only(model, inputs):
        return model(**inputs)

    def smp_gather(tensor):
        if isinstance(tensor, (list, tuple)):
            return type(tensor)(smp_gather(t) for t in tensor)
        elif isinstance(tensor, dict):
            return type(tensor)({k: smp_gather(v) for k, v in tensor.items()})
        elif not isinstance(tensor, torch.Tensor):
            raise TypeError(
                f"Can't gather the values of type {type(tensor)}, only of nested list/tuple/dicts of tensors."
            )
        all_tensors = smp.allgather(tensor, smp.CommGroup.DP_GROUP)
        all_tensors = [atleast_1d(t) for t in all_tensors]
        return torch.cat([t.cpu() for t in all_tensors], dim=0)

    def smp_nested_concat(tensor):
        if isinstance(tensor, (list, tuple)):
            return type(tensor)(smp_nested_concat(t) for t in tensor)
        elif isinstance(tensor, dict):
            return type(tensor)({k: smp_nested_concat(v) for k, v in tensor.items()})
        
        
        return tensor.concat().detach().cpu()