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 the trainer class to easily train a transformers from scratch or finetune it on a new task integrations must be imported before ml frameworks isort off isort on name of the files used for checkpointing trainer is a simple but featurecomplete training and eval loop for pytorch optimized for transformers args model pretrainedmodel or torch nn module optional the model to train evaluate or use for predictions if not provided a modelinit must be passed tip trainer is optimized to work with the pretrainedmodel provided by the library you can still use your own models defined as torch nn module as long as they work the same way as the transformers models tip args trainingarguments optional the arguments to tweak for training will default to a basic instance of trainingarguments with the outputdir set to a directory named tmptrainer in the current directory if not provided datacollator datacollator optional the function to use to form a batch from a list of elements of traindataset or evaldataset will default to defaultdatacollator if no tokenizer is provided an instance of datacollatorwithpadding otherwise traindataset torch utils data dataset or torch utils data iterabledataset optional the dataset to use for training if it is a datasets dataset columns not accepted by the model forward method are automatically removed note that if it s a torch utils data iterabledataset with some randomization and you are training in a distributed fashion your iterable dataset should either use a internal attribute generator that is a torch generator for the randomization that must be identical on all processes and the trainer will manually set the seed of this generator at each epoch or have a setepoch method that internally sets the seed of the rngs used evaldataset uniontorch utils data dataset dictstr torch utils data dataset optional the dataset to use for evaluation if it is a datasets dataset columns not accepted by the model forward method are automatically removed if it is a dictionary it will evaluate on each dataset prepending the dictionary key to the metric name tokenizer pretrainedtokenizerbase optional the tokenizer used to preprocess the data if provided will be used to automatically pad the inputs to the maximum length when batching inputs and it will be saved along the model to make it easier to rerun an interrupted training or reuse the finetuned model modelinit callable pretrainedmodel optional a function that instantiates the model to be used if provided each call to trainer train will start from a new instance of the model as given by this function the function may have zero argument or a single one containing the optunaray tunesigopt trial object to be able to choose different architectures according to hyper parameters such as layer count sizes of inner layers dropout probabilities etc computemetrics callableevalprediction dict optional the function that will be used to compute metrics at evaluation must take a evalprediction and return a dictionary string to metric values callbacks list of trainercallback optional a list of callbacks to customize the training loop will add those to the list of default callbacks detailed in herecallback if you want to remove one of the default callbacks used use the trainer removecallback method optimizers tupletorch optim optimizer torch optim lrscheduler lambdalr optional defaults to none none a tuple containing the optimizer and the scheduler to use will default to an instance of adamw on your model and a scheduler given by getlinearschedulewithwarmup controlled by args preprocesslogitsformetrics callabletorch tensor torch tensor torch tensor optional a function that preprocess the logits right before caching them at each evaluation step must take two tensors the logits and the labels and return the logits once processed as desired the modifications made by this function will be reflected in the predictions received by computemetrics note that the labels second parameter will be none if the dataset does not have them important attributes model always points to the core model if using a transformers model it will be a pretrainedmodel subclass modelwrapped always points to the most external model in case one or more other modules wrap the original model this is the model that should be used for the forward pass for example under deepspeed the inner model is wrapped in deepspeed and then again in torch nn distributeddataparallel if the inner model hasn t been wrapped then self modelwrapped is the same as self model ismodelparallel whether or not a model has been switched to a model parallel mode different from data parallelism this means some of the model layers are split on different gpus placemodelondevice whether or not to automatically place the model on the device it will be set to false if model parallel or deepspeed is used or if the default trainingarguments placemodelondevice is overridden to return false isintrain whether or not a model is currently running train e g when evaluate is called while in train those are used as methods of the trainer in examples seed must be set before instantiating the model when using model memory metrics must set up as early as possible set the correct log level depending on the node force device and distributed setup init explicitly warn users at this stage the model is already loaded one place to sort out whether to place the model on device or not postpone switching model to cuda when 1 mp since we are trying to fit a much bigger than 1 gpu model 2 fp16enabled deepspeed loads the model in half the size and it doesn t need to anyway and we only use deepspeed for training at the moment 3 full bf16 or fp16 eval since the model needs to be cast to the right dtype first 4 fsdp same as mp bnb quantized models doesn t support to operation force ngpu to 1 to avoid dataparallel as mp will manage the gpus later use self model is self modelwrapped to check if it s wrapped or not will be set to true by self setuploggers on first call to self log create distant repo and output directory if needed mixed precision setup mixed precision setup for sagemaker model parallel bf16 model parallelism in sagemaker currently not supported raise an error when there s mismatch between smp config and trainer argument use smp config as truth smp 1 10 does not support fp16 in trainer deepspeed and sagemaker model parallel manage their own half precision label smoothing internal variable to count flos in each process will be accumulated in self state totalflos then returned to 0 every time flos need to be logged internal variables to help with automatic batch size reduction very last torch compile unwrappedmodel unwrapmodelmodel if ispeftavailable and isinstanceunwrappedmodel peftmodel embeddings unwrappedmodel basemodel model getinputembeddings else embeddings unwrappedmodel getinputembeddings del unwrappedmodel embeddings neftunenoisealpha self neftunenoisealpha hookhandle embeddings registerforwardhookneftunepostforwardhook self neftunehookhandle hookhandle return model def deactivateneftuneself model if not hasattrself neftunehookhandle raise valueerrorneftune is not activated make sure to call trainer activateneftune first unwrappedmodel unwrapmodelmodel if ispeftavailable and isinstanceunwrappedmodel peftmodel embeddings unwrappedmodel basemodel model getinputembeddings else embeddings unwrappedmodel getinputembeddings self neftunehookhandle remove del embeddings neftunenoisealpha unwrappedmodel def addcallbackself callback self callbackhandler addcallbackcallback def popcallbackself callback return self callbackhandler popcallbackcallback def removecallbackself callback self callbackhandler removecallbackcallback def movemodeltodeviceself model device model model todevice moving a model to an xla device disconnects the tied weights so we have to retie them if self args parallelmode parallelmode tpu and hasattrmodel tieweights model tieweights def setsignaturecolumnsifneededself if self signaturecolumns is none inspect model forward signature to keep only the arguments it accepts signature inspect signatureself model forward self signaturecolumns listsignature parameters keys labels may be named label or labelids the default data collator handles that self signaturecolumns listsetlabel labelids self labelnames def removeunusedcolumnsself dataset datasets dataset description optionalstr none if not self args removeunusedcolumns return dataset self setsignaturecolumnsifneeded signaturecolumns self signaturecolumns ignoredcolumns listsetdataset columnnames setsignaturecolumns if lenignoredcolumns 0 dsetdescription if description is none else fin the description set logger info fthe following columns dsetdescription don t have a corresponding argument in fself model class name forward and have been ignored joinignoredcolumns f if joinignoredcolumns are not expected by self model class name forward you can safely ignore this message columns k for k in signaturecolumns if k in dataset columnnames if version parsedatasets version version parse1 4 0 dataset setformat typedataset formattype columnscolumns formatkwargsdataset formatformatkwargs return dataset else return dataset removecolumnsignoredcolumns def getcollatorwithremovedcolumns self datacollator callable description optionalstr none callable build the sampler returns the training torch utils data dataloader will use no sampler if traindataset does not implement len a random sampler adapted to distributed training if necessary otherwise subclass and override this method if you want to inject some custom behavior deprecated code returns the evaluation torch utils data dataloader subclass and override this method if you want to inject some custom behavior args evaldataset torch utils data dataset optional if provided will override self evaldataset if it is a datasets dataset columns not accepted by the model forward method are automatically removed it must implement len returns the test torch utils data dataloader subclass and override this method if you want to inject some custom behavior args testdataset torch utils data dataset optional the test dataset to use if it is a datasets dataset columns not accepted by the model forward method are automatically removed it must implement len we use the same batchsize as for eval setup the optimizer and the learning rate scheduler we provide a reasonable default that works well if you want to use something else you can pass a tuple in the trainer s init through optimizers or subclass and override this method or createoptimizer andor createscheduler in a subclass if smp 1 10 and fp16 is enabled we unwrap the optimizer get all parameter names that weight decay will be applied to note that some models implement their own layernorm instead of calling nn layernorm weight decay could still apply to those modules since this function only filter out instance of nn layernorm setup the optimizer we provide a reasonable default that works well if you want to use something else you can pass a tuple in the trainer s init through optimizers or subclass and override this method in a subclass returns the optimizer class and optimizer parameters based on the training arguments args args transformers trainingargs trainingarguments the training arguments for the training session parse args optimargs todo change dtypes back to mfp32 var bf16 kahan false once they can be cast together in torchdistx setup the scheduler the optimizer of the trainer must have been set up either before this method is called or passed as an argument args numtrainingsteps int the number of training steps to do helper to get number of samples in a torch utils data dataloader by accessing its dataset when dataloader dataset does not exist or has no length estimates as best it can special case for iterabledatasetshard we need to dig deeper helper to get number of tokens in a torch utils data dataloader by enumerating dataloader hp search setup code self trial trial if self hpsearchbackend is none or trial is none return if self hpsearchbackend hpsearchbackend optuna params self hpspacetrial elif self hpsearchbackend hpsearchbackend ray params trial params popwandb none elif self hpsearchbackend hpsearchbackend sigopt params k intv if isinstancev str else v for k v in trial assignments items elif self hpsearchbackend hpsearchbackend wandb params trial for key value in params items if not hasattrself args key logger warning ftrying to set key in the hyperparameter search but there is no corresponding field in trainingarguments continue oldattr getattrself args key none casting value to the proper type if oldattr is not none value typeoldattrvalue setattrself args key value if self hpsearchbackend hpsearchbackend optuna logger infoftrial trial params if self hpsearchbackend hpsearchbackend sigopt logger infofsigopt assignments trial assignments if self hpsearchbackend hpsearchbackend wandb logger infofwb sweep parameters trial if self isdeepspeedenabled if self args deepspeed is none raise valueerrorfor sweeps with deepspeed args deepspeed must be set rebuild the deepspeed config to reflect the updated training parameters from accelerate utils import deepspeedplugin from transformers integrations deepspeed import hftrainerdeepspeedconfig self args hfdeepspeedconfig hftrainerdeepspeedconfigself args deepspeed self args hfdeepspeedconfig trainerconfigprocessself args self args deepspeedplugin deepspeedpluginhfdsconfigself args hfdeepspeedconfig self createacceleratorandpostprocess def reporttohpsearchself trial unionoptuna trial dictstr any step int metrics dictstr float if self hpsearchbackend is none or trial is none return self objective self computeobjectivemetrics copy if self hpsearchbackend hpsearchbackend optuna import optuna if not trial study ismultiobjective trial reportself objective step if trial shouldprune self callbackhandler ontrainendself args self state self control raise optuna trialpruned elif self hpsearchbackend hpsearchbackend ray from ray import tune if self control shouldsave self tunesavecheckpoint tune reportobjectiveself objective metrics def tunesavecheckpointself from ray import tune if not self usetunecheckpoints return with tune checkpointdirstepself state globalstep as checkpointdir outputdir os path joincheckpointdir fprefixcheckpointdirself state globalstep self savemodeloutputdir internalcalltrue if self args shouldsave self state savetojsonos path joinoutputdir trainerstatename torch saveself optimizer statedict os path joinoutputdir optimizername torch saveself lrscheduler statedict os path joinoutputdir schedulername def callmodelinitself trialnone modelinitargcount numberofargumentsself modelinit if modelinitargcount 0 model self modelinit elif modelinitargcount 1 model self modelinittrial else raise runtimeerrormodelinit should have 0 or 1 argument if model is none raise runtimeerrormodelinit should not return none return model def torchjitmodelevalself model dataloader trainingfalse if not training if dataloader is none logger warningfailed to use pytorch jit mode due to current dataloader is none return model examplebatch nextiterdataloader examplebatch self prepareinputsexamplebatch try jitmodel copy copymodel jitmodel eval originalforward jitmodel dict poporiginalforward none remove mixed precision hooks from the model if originalforward jitmodel forward originalforward with self accelerator autocastcacheenabledfalse torch nograd if version parseversion parsetorch version baseversion version parse2 0 0 if isinstanceexamplebatch dict jitmodel torch jit tracejitmodel examplekwarginputsexamplebatch strictfalse else jitmodel torch jit trace jitmodel examplekwarginputskey examplebatchkey for key in examplebatch strictfalse else jitinputs for key in examplebatch exampletensor torch oneslikeexamplebatchkey jitinputs appendexampletensor jitinputs tuplejitinputs jitmodel torch jit tracejitmodel jitinputs strictfalse jitmodel torch jit freezejitmodel with torch nograd jitmodelexamplebatch jitmodelexamplebatch model jitmodel self usecpuamp false except runtimeerror typeerror valueerror nameerror indexerror as e logger warningffailed to use pytorch jit mode due to e return model def ipexoptimizemodelself model trainingfalse dtypetorch float32 if not isipexavailable raise importerror using ipex but ipex is not installed or ipex s version does not match current pytorch please refer to https github comintelintelextensionforpytorch import intelextensionforpytorch as ipex if not training model eval dtype torch bfloat16 if not self isintrain and self args bf16fulleval else dtype convbnfolding is disabled as it fails in symbolic tracing resulting in ipex warnings model ipex optimizemodel dtypedtype levelo1 convbnfoldingfalse inplacenot self isintrain else if not model training model train model self optimizer ipex optimize model dtypedtype optimizerself optimizer inplacetrue levelo1 return model def wrapmodelself model trainingtrue dataloadernone if self args useipex dtype torch bfloat16 if self usecpuamp else torch float32 model self ipexoptimizemodelmodel training dtypedtype if issagemakermpenabled wrapping the base model twice in a distributedmodel will raise an error if isinstanceself modelwrapped smp model distributedmodel return self modelwrapped return smp distributedmodelmodel backwardpassesperstepself args gradientaccumulationsteps traineval could be run multipletimes if already wrapped don t rewrap it again if unwrapmodelmodel is not model return model mixed precision training with apex torch 1 6 if self useapex and training model self optimizer amp initializemodel self optimizer optlevelself args fp16optlevel multigpu training should be after apex fp16 initialization 8bit models does not support ddp if self args ngpu 1 and not getattrmodel isloadedin8bit false model nn dataparallelmodel if self args jitmodeeval starttime time time model self torchjitmodelevalmodel dataloader training self jitcompilationtime roundtime time starttime 4 note in torch distributed mode there s no point in wrapping the model inside a distributeddataparallel as we ll be under nograd anyways if not training return model distributed training should be after apex fp16 initialization distributed training using pytorch fsdp if self isfsdpxlaenabled try from torchxla distributed fsdp import xlafullyshardeddataparallel as fsdp from torchxla distributed fsdp import checkpointmodule from torchxla distributed fsdp wrap import sizebasedautowrappolicy transformerautowrappolicy except importerror raise importerrormissing xla fsdp related module please make sure to use torchxla 2 0 autowrappolicy none autowrappercallable none defaulttransformerclsnamestowrap getattrmodel nosplitmodules none fsdptransformerlayerclstowrap self args fsdpconfig get transformerlayerclstowrap defaulttransformerclsnamestowrap if self args fsdpconfigminnumparams 0 autowrappolicy functools partial sizebasedautowrappolicy minnumparamsself args fsdpconfigminnumparams elif fsdptransformerlayerclstowrap is not none transformerclstowrap set for layerclass in fsdptransformerlayerclstowrap transformercls getmoduleclassfromnamemodel layerclass if transformercls is none raise exceptioncould not find the transformer layer class to wrap in the model else transformerclstowrap addtransformercls autowrappolicy functools partial transformerautowrappolicy transformer layer class to wrap transformerlayerclstransformerclstowrap fsdpkwargs self args xlafsdpconfig if self args fsdpconfigxlafsdpgradckpt apply gradient checkpointing to autowrapped submodules if specified def autowrappercallablem args kwargs return fsdpcheckpointmodulem args kwargs wrap the base model with an outer fsdp wrapper self model model fsdp model autowrappolicyautowrappolicy autowrappercallableautowrappercallable fsdpkwargs patch xm optimizerstep should not reduce gradients in this case as fsdp does not need gradient reduction over sharded parameters def patchedoptimizerstepoptimizer barrierfalse optimizerargs loss optimizer stepoptimizerargs if barrier xm markstep return loss xm optimizerstep patchedoptimizerstep elif issagemakerdpenabled model nn parallel distributeddataparallel model deviceidsintos getenvsmdataparallellocalrank elif self args parallelmode parallelmode distributed if istorchneuroncoreavailable return model kwargs if self args ddpfindunusedparameters is not none kwargsfindunusedparameters self args ddpfindunusedparameters elif isinstancemodel pretrainedmodel findunusedparameters breaks checkpointing as per https github comhuggingfacetransformerspull4659issuecomment643356021 kwargsfindunusedparameters not model isgradientcheckpointing else kwargsfindunusedparameters true if self args ddpbucketcapmb is not none kwargsbucketcapmb self args ddpbucketcapmb if self args ddpbroadcastbuffers is not none kwargsbroadcastbuffers self args ddpbroadcastbuffers self accelerator ddphandler distributeddataparallelkwargskwargs return model def train self resumefromcheckpoint optionalunionstr bool none trial unionoptuna trial dictstr any none ignorekeysforeval optionalliststr none kwargs if resumefromcheckpoint is false resumefromcheckpoint none memory metrics must set up as early as possible self memorytracker start args self args self isintrain true attach neftune hooks if necessary if self neftunenoisealpha is not none self model self activateneftuneself model dotrain is not a reliable argument as it might not be set and train still called so the following is a workaround if args fp16fulleval or args bf16fulleval and not args dotrain self movemodeltodeviceself model args device if modelpath in kwargs resumefromcheckpoint kwargs popmodelpath warnings warn modelpath is deprecated and will be removed in a future version use resumefromcheckpoint instead futurewarning if lenkwargs 0 raise typeerrorftrain received got unexpected keyword arguments joinlistkwargs keys this might change the seed so needs to run first self hpsearchsetuptrial self trainbatchsize self args trainbatchsize model reinit modelreloaded false if self modelinit is not none seed must be set before instantiating the model when using modelinit enablefulldeterminismself args seed if self args fulldeterminism else setseedself args seed self model self callmodelinittrial modelreloaded true reinitializes optimizer and scheduler self optimizer self lrscheduler none none load potential model checkpoint if isinstanceresumefromcheckpoint bool and resumefromcheckpoint resumefromcheckpoint getlastcheckpointargs outputdir if resumefromcheckpoint is none raise valueerrorfno valid checkpoint found in output directory args outputdir if resumefromcheckpoint is not none and not issagemakermpenabled and not self isdeepspeedenabled and not self isfsdpenabled self loadfromcheckpointresumefromcheckpoint if model was reinitialized put it on the right device and update self modelwrapped if modelreloaded if self placemodelondevice self movemodeltodeviceself model args device self modelwrapped self model innertrainingloop findexecutablebatchsize self innertrainingloop self trainbatchsize args autofindbatchsize if args pushtohub try disable progress bars when uploading models during checkpoints to avoid polluting stdout hfhubutils disableprogressbars return innertrainingloop argsargs resumefromcheckpointresumefromcheckpoint trialtrial ignorekeysforevalignorekeysforeval finally hfhubutils enableprogressbars else return innertrainingloop argsargs resumefromcheckpointresumefromcheckpoint trialtrial ignorekeysforevalignorekeysforeval def innertrainingloop self batchsizenone argsnone resumefromcheckpointnone trialnone ignorekeysforevalnone self accelerator freememory self trainbatchsize batchsize logger debugfcurrently training with a batch size of self trainbatchsize data loader and number of training steps traindataloader self gettraindataloader setting up training control variables number of training epochs numtrainepochs number of training steps per epoch numupdatestepsperepoch total number of training steps to execute maxsteps totaltrainbatchsize self trainbatchsize args gradientaccumulationsteps args worldsize lendataloader none numtraintokens none if haslengthtraindataloader lendataloader lentraindataloader numupdatestepsperepoch lendataloader args gradientaccumulationsteps numupdatestepsperepoch maxnumupdatestepsperepoch 1 numexamples self numexamplestraindataloader if args maxsteps 0 maxsteps args maxsteps numtrainepochs args maxsteps numupdatestepsperepoch int args maxsteps numupdatestepsperepoch 0 may be slightly incorrect if the last batch in the training dataloader has a smaller size but it s the best we can do numtrainsamples args maxsteps totaltrainbatchsize if args includetokenspersecond numtraintokens self numtokenstraindataloader args maxsteps args gradientaccumulationsteps else maxsteps math ceilargs numtrainepochs numupdatestepsperepoch numtrainepochs math ceilargs numtrainepochs numtrainsamples self numexamplestraindataloader args numtrainepochs if args includetokenspersecond numtraintokens self numtokenstraindataloader args numtrainepochs elif args maxsteps 0 rely on maxsteps when dataloader does not have a working size maxsteps args maxsteps setting a very large number of epochs so we go as many times as necessary over the iterator numtrainepochs sys maxsize numupdatestepsperepoch maxsteps numexamples totaltrainbatchsize args maxsteps numtrainsamples args maxsteps totaltrainbatchsize if args includetokenspersecond numtraintokens self numtokenstraindataloader args maxsteps args gradientaccumulationsteps else raise valueerror args maxsteps must be set to a positive value if dataloader does not have a length was f args maxsteps if debugoption underflowoverflow in self args debug if self args ngpu 1 nn dataparallelmodel replicates the model creating new variables and module references registered here no longer work on other gpus breaking the module raise valueerror currently debug underflowoverflow is not supported under dp please use ddp torchrun or torch distributed launch deprecated else debugoverflow debugunderflowoverflowself model noqa delayoptimizercreation issagemakermpenabled or self isfsdpxlaenabled or self isfsdpenabled we need to reset the scheduler as its parameters may be different on subsequent calls if self createdlrscheduler self lrscheduler none self createdlrscheduler false if self isdeepspeedenabled self optimizer self lrscheduler deepspeedinitself numtrainingstepsmaxsteps if not delayoptimizercreation self createoptimizerandschedulernumtrainingstepsmaxsteps self state trainerstate self state ishyperparamsearch trial is not none compute absolute values for logging eval and save if given as ratio if args loggingsteps is not none if args loggingsteps 1 self state loggingsteps math ceilmaxsteps args loggingsteps else self state loggingsteps args loggingsteps if args evalsteps is not none if args evalsteps 1 self state evalsteps math ceilmaxsteps args evalsteps else self state evalsteps args evalsteps if args savesteps is not none if args savesteps 1 self state savesteps math ceilmaxsteps args savesteps else self state savesteps args savesteps activate gradient checkpointing if needed if args gradientcheckpointing if args gradientcheckpointingkwargs is none gradientcheckpointingkwargs else gradientcheckpointingkwargs args gradientcheckpointingkwargs self model gradientcheckpointingenablegradientcheckpointingkwargsgradientcheckpointingkwargs model self wrapmodelself modelwrapped as the model is wrapped don t use accelerator prepare this is for unhandled cases such as fsdpxla sagemaker mpdp dataparallel ipex useacceleratorprepare true if model is self model else false if delayoptimizercreation self createoptimizerandschedulernumtrainingstepsmaxsteps prepare using accelerator prepare if useacceleratorprepare self model train if hasattrself lrscheduler step if self useapex model self accelerator prepareself model else model self optimizer self accelerator prepareself model self optimizer else to handle cases wherein we pass dummyscheduler such as when it is specified in deepspeed config model self optimizer self lrscheduler self accelerator prepare self model self optimizer self lrscheduler if self isfsdpenabled self model self modelwrapped model for the rest of this function model is the outside model whether it was wrapped or not if model is not self model self modelwrapped model backward compatibility if self isdeepspeedenabled self deepspeed self modelwrapped ckpt loading if resumefromcheckpoint is not none if self isdeepspeedenabled deepspeedloadcheckpointself modelwrapped resumefromcheckpoint elif issagemakermpenabled or self isfsdpenabled self loadfromcheckpointresumefromcheckpoint self modelwrapped check if saved optimizer or scheduler states exist self loadoptimizerandschedulerresumefromcheckpoint important at this point self model is the transformers model self modelwrapped is ddptransformers model deepspeedtransformers model fsdptransformers model dynamo optimized moduletransformers model etc train logger info running training logger infof num examples numexamples logger infof num epochs numtrainepochs logger infof instantaneous batch size per device self args perdevicetrainbatchsize if self args perdevicetrainbatchsize self trainbatchsize logger infof training with dataparallel so batch size has been adjusted to self trainbatchsize logger infof total train batch size w parallel distributed accumulation totaltrainbatchsize logger infof gradient accumulation steps args gradientaccumulationsteps logger infof total optimization steps maxsteps logger infof number of trainable parameters getmodelparamcountmodel trainableonlytrue self state epoch 0 starttime time time epochstrained 0 stepstrainedincurrentepoch 0 stepstrainedprogressbar none check if continuing training from a checkpoint if resumefromcheckpoint is not none and os path isfile os path joinresumefromcheckpoint trainerstatename self state trainerstate loadfromjsonos path joinresumefromcheckpoint trainerstatename epochstrained self state globalstep numupdatestepsperepoch if not args ignoredataskip stepstrainedincurrentepoch self state globalstep numupdatestepsperepoch stepstrainedincurrentepoch args gradientaccumulationsteps else stepstrainedincurrentepoch 0 logger info continuing training from checkpoint will skip to saved globalstep logger infof continuing training from epoch epochstrained logger infof continuing training from global step self state globalstep if not args ignoredataskip logger info f will skip the first epochstrained epochs then the first f stepstrainedincurrentepoch batches in the first epoch update the references self callbackhandler model self model self callbackhandler optimizer self optimizer self callbackhandler lrscheduler self lrscheduler self callbackhandler traindataloader traindataloader if self hpname is not none and self trial is not none use self trial because the sigoptoptuna hpo only call hpsearchsetuptrial instead of passing trial parameter to train when using ddp self state trialname self hpnameself trial if trial is not none assignments trial assignments if self hpsearchbackend hpsearchbackend sigopt else trial self state trialparams hpparamsassignments else self state trialparams none this should be the same if the state has been saved but in case the training arguments changed it s safer to set this after the load self state maxsteps maxsteps self state numtrainepochs numtrainepochs self state islocalprocesszero self islocalprocesszero self state isworldprocesszero self isworldprocesszero trloss is a tensor to avoid synchronization of tpus through item trloss torch tensor0 0 toargs device totallossscalar is updated everytime item has to be called on trloss and stores the sum of all losses self totallossscalar 0 0 self globalsteplastlogged self state globalstep model zerograd self control self callbackhandler ontrainbeginargs self state self control skip the first epochstrained epochs to get the random state of the dataloader at the right point if not args ignoredataskip for epoch in rangeepochstrained sampler getdataloadersamplertraindataloader samplerkinds randomsampler if version parseaccelerateversion version parse0 23 0 samplerkinds appendseedablerandomsampler israndomsampler isinstancesampler tuplesamplerkinds if istorchlessthan111 or not israndomsampler we just need to begin an iteration to create the randomization of the sampler for in traindataloader break else otherwise we need to call the whooooole sampler cause there is some random operation added at the very end sampler sampler if sampler is not none else listsampler totalbatchedsamples 0 for epoch in rangeepochstrained numtrainepochs epochiterator traindataloader if hasattrepochiterator setepoch epochiterator setepochepoch reset the past mems state at the beginning of each epoch if necessary if args pastindex 0 self past none stepsinepoch lenepochiterator if lendataloader is not none else args maxsteps args gradientaccumulationsteps self control self callbackhandler onepochbeginargs self state self control if epoch epochstrained and resumefromcheckpoint is not none and stepstrainedincurrentepoch 0 self loadrngstateresumefromcheckpoint rngtosync false stepsskipped 0 if stepstrainedincurrentepoch 0 epochiterator skipfirstbatchesepochiterator stepstrainedincurrentepoch stepsskipped stepstrainedincurrentepoch stepstrainedincurrentepoch 0 rngtosync true step 1 for step inputs in enumerateepochiterator totalbatchedsamples 1 if self args includenuminputtokensseen maininputname getattrself model maininputname inputids if maininputname not in inputs logger warning tried to track the number of tokens seen however the current model is not configured properly to know what item is the input to fix this add a maininputname attribute to the model class you are using else self state numinputtokensseen self accelerator gatherinputsmaininputname numel if rngtosync self loadrngstateresumefromcheckpoint rngtosync false skip past any already trained steps if resuming training if stepstrainedincurrentepoch 0 stepstrainedincurrentepoch 1 if stepstrainedprogressbar is not none stepstrainedprogressbar update1 if stepstrainedincurrentepoch 0 self loadrngstateresumefromcheckpoint continue elif stepstrainedprogressbar is not none stepstrainedprogressbar close stepstrainedprogressbar none if step args gradientaccumulationsteps 0 self control self callbackhandler onstepbeginargs self state self control with self accelerator accumulatemodel trlossstep self trainingstepmodel inputs if args loggingnaninffilter and not istorchtpuavailable and torch isnantrlossstep or torch isinftrlossstep if loss is nan or inf simply add the average of previous logged losses trloss trloss 1 self state globalstep self globalsteplastlogged else trloss trlossstep self currentflos floatself floatingpointopsinputs islaststepandstepslessthangradacc stepsinepoch args gradientaccumulationsteps and step 1 stepsinepoch if totalbatchedsamples args gradientaccumulationsteps 0 or last step in epoch but step is always smaller than gradientaccumulationsteps islaststepandstepslessthangradacc the or condition of islaststepandstepslessthangradacc is not covered in accelerate so explicitly enable sync gradients to true in that case if islaststepandstepslessthangradacc self accelerator gradientstate setsyncgradientstrue gradient clipping if args maxgradnorm is not none and args maxgradnorm 0 deepspeed does its own clipping if issagemakermpenabled and args fp16 self optimizer clipmastergradsargs maxgradnorm elif self useapex revert to normal clipping otherwise handling apex or full precision nn utils clipgradnorm amp masterparamsself optimizer args maxgradnorm else self accelerator clipgradnorm model parameters args maxgradnorm optimizer step self optimizer step optimizerwasrun not self accelerator optimizerstepwasskipped if optimizerwasrun delay optimizer scheduling until metrics are generated if not isinstanceself lrscheduler torch optim lrscheduler reducelronplateau self lrscheduler step model zerograd self state globalstep 1 self state epoch epoch step 1 stepsskipped stepsinepoch self control self callbackhandler onstependargs self state self control self maybelogsaveevaluatetrloss model trial epoch ignorekeysforeval else self control self callbackhandler onsubstependargs self state self control if self control shouldepochstop or self control shouldtrainingstop break if step 0 logger warning there seems to be not a single sample in your epochiterator stopping training at step f self state globalstep this is expected if you re using an iterabledataset and set f numsteps maxsteps higher than the number of available samples self control shouldtrainingstop true self control self callbackhandler onepochendargs self state self control self maybelogsaveevaluatetrloss model trial epoch ignorekeysforeval if debugoption tpumetricsdebug in self args debug if istorchtpuavailable tpucomment logging debug metrics for pytorchxla compile execute times ops etc xm masterprintmet metricsreport else logger warning you enabled pytorchxla debug metrics but you don t have a tpu configured check your training configuration if this is unexpected if self control shouldtrainingstop break if args pastindex and hasattrself past clean the state at the end of training delattrself past logger infonntraining completed do not forget to share your model on huggingface comodels nn if args loadbestmodelatend and self state bestmodelcheckpoint is not none wait for everyone to get here so we are sure the model has been saved by process 0 if istorchtpuavailable xm rendezvousloadbestmodelatend elif args parallelmode parallelmode distributed dist barrier elif issagemakermpenabled smp barrier self loadbestmodel add remaining trloss self totallossscalar trloss item trainloss self totallossscalar self state globalstep metrics speedmetrics train starttime numsamplesnumtrainsamples numstepsself state maxsteps numtokensnumtraintokens self storeflos metricstotalflos self state totalflos metricstrainloss trainloss self isintrain false self memorytracker stopandupdatemetricsmetrics self logmetrics rundir self getoutputdirtrial checkpointssorted self sortedcheckpointsusemtimefalse outputdirrundir delete the last checkpoint when savetotallimit1 if it s different from the best checkpoint and process allowed to save if self args shouldsave and self state bestmodelcheckpoint is not none and self args savetotallimit 1 for checkpoint in checkpointssorted if not os path samefilecheckpoint self state bestmodelcheckpoint logger infofdeleting older checkpoint checkpoint due to args savetotallimit shutil rmtreecheckpoint self control self callbackhandler ontrainendargs self state self control wait for the checkpoint to be uploaded self finishcurrentpush after training we make sure to retrieve back the original forward pass method for the embedding layer by removing the forward post hook if self neftunenoisealpha is not none self deactivateneftuneself model return trainoutputself state globalstep trainloss metrics def getoutputdirself trial if self hpsearchbackend is not none and trial is not none if self hpsearchbackend hpsearchbackend optuna runid trial number elif self hpsearchbackend hpsearchbackend ray from ray import tune runid tune gettrialid elif self hpsearchbackend hpsearchbackend sigopt runid trial id elif self hpsearchbackend hpsearchbackend wandb import wandb runid wandb run id runname self hpnametrial if self hpname is not none else frunrunid rundir os path joinself args outputdir runname else rundir self args outputdir return rundir def loadfromcheckpointself resumefromcheckpoint modelnone if model is none model self model configfile os path joinresumefromcheckpoint configname adapterweightsfile os path joinresumefromcheckpoint adapterweightsname adaptersafeweightsfile os path joinresumefromcheckpoint adaptersafeweightsname weightsfile os path joinresumefromcheckpoint weightsname weightsindexfile os path joinresumefromcheckpoint weightsindexname safeweightsfile os path joinresumefromcheckpoint safeweightsname safeweightsindexfile os path joinresumefromcheckpoint safeweightsindexname isfsdpckpt os path isdirresumefromcheckpoint and any fsdpmodelname in foldername for foldername in os listdirresumefromcheckpoint if os path isdiros path joinresumefromcheckpoint foldername if isfsdpckpt and not self isfsdpenabled raise valueerrorfcheckpoint found at resumefromcheckpoint is only supported when using pytorch fsdp if not any os path isfilef for f in weightsfile safeweightsfile weightsindexfile safeweightsindexfile adapterweightsfile adaptersafeweightsfile or isfsdpckpt raise valueerrorfcan t find a valid checkpoint at resumefromcheckpoint logger infofloading model from resumefromcheckpoint if os path isfileconfigfile config pretrainedconfig fromjsonfileconfigfile checkpointversion config transformersversion if checkpointversion is not none and checkpointversion version logger warning fyou are resuming training from a checkpoint trained with checkpointversion of ftransformers but your current version is version this is not recommended and could yield to errors or unwanted behaviors if os path isfileweightsfile or os path isfilesafeweightsfile or isfsdpckpt if the model is on the gpu it still works if issagemakermpenabled if os path isfileos path joinresumefromcheckpoint usercontent pt if the usercontent pt file exists load with the new smp api checkpoint must have been saved with the new smp api smp resumefromcheckpoint pathresumefromcheckpoint tagweightsname partialfalse loadoptimizerfalse else if the usercontent pt file does not exist load with the old smp api checkpoint must have been saved with the old smp api if hasattrself args fp16 and self args fp16 is true logger warning enabling fp16 and loading from smp 1 10 checkpoint together is not suppported statedict torch loadweightsfile maplocationcpu required for smp to not autotranslate statedict from hf to smp is already smp statedictsmpispartial false loadresult model loadstatedictstatedict stricttrue release memory del statedict elif self isfsdpenabled loadfsdpmodelself accelerator state fsdpplugin self accelerator model resumefromcheckpoint else we load the model state dict on the cpu to avoid an oom error if self args savesafetensors and os path isfilesafeweightsfile statedict safetensors torch loadfilesafeweightsfile devicecpu else statedict torch loadweightsfile maplocationcpu workaround for fsdp bug https github compytorchpytorchissues82963 which takes args instead of kwargs loadresult model loadstatedictstatedict false release memory del statedict self issuewarningsafterloadloadresult load adapters following pr 24096 elif ispeftavailable and isinstancemodel peftmodel if train a model using peft lora assume that adapter have been saved properly if hasattrmodel activeadapter and hasattrmodel loadadapter if os path existsresumefromcheckpoint model loadadapterresumefromcheckpoint model activeadapter istrainabletrue else logger warning the intermediate checkpoints of peft may not be saved correctly fconsider using a custom callback to save adapterweightsname in corresponding saving folders check some examples here https github comhuggingfacepeftissues96 else logger warningcould not load adapter model make sure to have peft0 3 0 installed else we load the sharded checkpoint loadresult loadshardedcheckpoint model resumefromcheckpoint strictissagemakermpenabled prefersafeself args savesafetensors if not issagemakermpenabled self issuewarningsafterloadloadresult def loadbestmodelself logger infofloading best model from self state bestmodelcheckpoint score self state bestmetric bestmodelpath os path joinself state bestmodelcheckpoint weightsname bestsafemodelpath os path joinself state bestmodelcheckpoint safeweightsname bestadaptermodelpath os path joinself state bestmodelcheckpoint adapterweightsname bestsafeadaptermodelpath os path joinself state bestmodelcheckpoint adaptersafeweightsname model self modelwrapped if issagemakermpenabled else self model if self isdeepspeedenabled deepspeedloadcheckpointself modelwrapped self state bestmodelcheckpoint elif self isfsdpenabled loadresult loadfsdpmodel self accelerator state fsdpplugin self accelerator model self state bestmodelcheckpoint elif os path existsbestmodelpath or os path existsbestsafemodelpath or os path existsbestadaptermodelpath or os path existsbestsafeadaptermodelpath hasbeenloaded true if issagemakermpenabled if os path isfileos path joinself state bestmodelcheckpoint usercontent pt if the usercontent pt file exists load with the new smp api checkpoint must have been saved with the new smp api smp resumefromcheckpoint pathself state bestmodelcheckpoint tagweightsname partialfalse loadoptimizerfalse else if the usercontent pt file does not exist load with the old smp api checkpoint must have been saved with the old smp api if self args savesafetensors and os path isfilebestsafemodelpath statedict safetensors torch loadfilebestsafemodelpath devicecpu else statedict torch loadbestmodelpath maplocationcpu statedictsmpispartial false loadresult model loadstatedictstatedict stricttrue else if ispeftavailable and isinstancemodel peftmodel if train a model using peft lora assume that adapter have been saved properly if hasattrmodel activeadapter and hasattrmodel loadadapter if os path existsbestadaptermodelpath or os path existsbestsafeadaptermodelpath model loadadapterself state bestmodelcheckpoint model activeadapter loadadapter has no return value present modify it when appropriate from torch nn modules module import incompatiblekeys loadresult incompatiblekeys else logger warning the intermediate checkpoints of peft may not be saved correctly fconsider using a custom callback to save adapterweightsname in corresponding saving folders check some examples here https github comhuggingfacepeftissues96 hasbeenloaded false else logger warningcould not load adapter model make sure to have peft0 3 0 installed hasbeenloaded false else we load the model state dict on the cpu to avoid an oom error if self args savesafetensors and os path isfilebestsafemodelpath statedict safetensors torch loadfilebestsafemodelpath devicecpu else statedict torch loadbestmodelpath maplocationcpu if the model is on the gpu it still works workaround for fsdp bug https github compytorchpytorchissues82963 which takes args instead of kwargs loadresult model loadstatedictstatedict false if not issagemakermpenabled and hasbeenloaded self issuewarningsafterloadloadresult elif os path existsos path joinself state bestmodelcheckpoint weightsindexname loadresult loadshardedcheckpoint model self state bestmodelcheckpoint strictissagemakermpenabled if not issagemakermpenabled self issuewarningsafterloadloadresult else logger warning fcould not locate the best model at bestmodelpath if you are running a distributed training on multiple nodes you should activate saveoneachnode def issuewarningsafterloadself loadresult if lenloadresult missingkeys 0 if self model keystoignoreonsave is not none and setloadresult missingkeys set self model keystoignoreonsave self model tieweights else logger warningfthere were missing keys in the checkpoint model loaded loadresult missingkeys if lenloadresult unexpectedkeys 0 logger warning fthere were unexpected keys in the checkpoint model loaded loadresult unexpectedkeys def maybelogsaveevaluateself trloss model trial epoch ignorekeysforeval if self control shouldlog if istorchtpuavailable xm markstep logs dictstr float allgather mean to get average loss over all processes trlossscalar self nestedgathertrloss mean item reset trloss to zero trloss trloss logsloss roundtrlossscalar self state globalstep self globalsteplastlogged 4 logslearningrate self getlearningrate self totallossscalar trlossscalar self globalsteplastlogged self state globalstep self storeflos self loglogs metrics none if self control shouldevaluate if isinstanceself evaldataset dict metrics for evaldatasetname evaldataset in self evaldataset items datasetmetrics self evaluate evaldatasetevaldataset ignorekeysignorekeysforeval metrickeyprefixfevalevaldatasetname metrics updatedatasetmetrics else metrics self evaluateignorekeysignorekeysforeval self reporttohpsearchtrial self state globalstep metrics run delayed lr scheduler now that metrics are populated if isinstanceself lrscheduler torch optim lrscheduler reducelronplateau metrictocheck self args metricforbestmodel if not metrictocheck startswitheval metrictocheck fevalmetrictocheck self lrscheduler stepmetricsmetrictocheck if self control shouldsave self savecheckpointmodel trial metricsmetrics self control self callbackhandler onsaveself args self state self control def loadrngstateself checkpoint load rng states from checkpoint if checkpoint is none return if self args worldsize 1 processindex self args processindex rngfile os path joincheckpoint frngstateprocessindex pth if not os path isfilerngfile logger info fdidn t find an rng file for process processindex if you are resuming a training that wasn t launched in a distributed fashion reproducibility is not guaranteed return else rngfile os path joincheckpoint rngstate pth if not os path isfilerngfile logger info didn t find an rng file if you are resuming a training that was launched in a distributed fashion reproducibility is not guaranteed return checkpointrngstate torch loadrngfile random setstatecheckpointrngstatepython np random setstatecheckpointrngstatenumpy torch random setrngstatecheckpointrngstatecpu if torch cuda isavailable if self args parallelmode parallelmode distributed torch cuda random setrngstateallcheckpointrngstatecuda else try torch cuda random setrngstatecheckpointrngstatecuda except exception as e logger info fdidn t manage to set back the rng states of the gpu because of the following error n e nthis won t yield the same results as if the training had not been interrupted if istorchtpuavailable xm setrngstatecheckpointrngstatexla if istorchnpuavailable if self args parallelmode parallelmode distributed torch npu random setrngstateallcheckpointrngstatenpu else try torch npu random setrngstatecheckpointrngstatenpu except exception as e logger info fdidn t manage to set back the rng states of the npu because of the following error n e nthis won t yield the same results as if the training had not been interrupted def savecheckpointself model trial metricsnone in all cases including ddpdpdeepspeed self model is always a reference to the model we want to save except fullyshardedddp assert unwrapmodelmodel is self model internal model should be a reference to self model save model checkpoint checkpointfolder fprefixcheckpointdirself state globalstep if self hpsearchbackend is none and trial is none self storeflos rundir self getoutputdirtrialtrial outputdir os path joinrundir checkpointfolder self savemodeloutputdir internalcalltrue if not self args saveonlymodel save optimizer and scheduler self saveoptimizerandscheduleroutputdir save rng state self saverngstateoutputdir determine the new best metric best model checkpoint if metrics is not none and self args metricforbestmodel is not none metrictocheck self args metricforbestmodel if not metrictocheck startswitheval metrictocheck fevalmetrictocheck metricvalue metricsmetrictocheck operator np greater if self args greaterisbetter else np less if self state bestmetric is none or self state bestmodelcheckpoint is none or operatormetricvalue self state bestmetric self state bestmetric metricvalue self state bestmodelcheckpoint outputdir save the trainer state if self args shouldsave self state savetojsonos path joinoutputdir trainerstatename if self args pushtohub self pushfromcheckpointoutputdir maybe delete some older checkpoints if self args shouldsave self rotatecheckpointsusemtimetrue outputdirrundir def saverngstateself outputdir save rng state in nondistributed training rngstates python random getstate numpy np random getstate cpu torch random getrngstate if torch cuda isavailable if self args parallelmode parallelmode distributed in non distributed we save the global cuda rng state will take care of dataparallel rngstatescuda torch cuda random getrngstateall else rngstatescuda torch cuda random getrngstate if istorchtpuavailable rngstatesxla xm getrngstate if istorchnpuavailable if self args parallelmode parallelmode distributed rngstatesnpu torch npu random getrngstateall else rngstatesnpu torch npu random getrngstate a process can arrive here before the process 0 has a chance to save the model in which case outputdir may not yet exist os makedirsoutputdir existoktrue if self args worldsize 1 torch saverngstates os path joinoutputdir rngstate pth else torch saverngstates os path joinoutputdir frngstateself args processindex pth def saveoptimizerandschedulerself outputdir if istorchtpuavailable xm rendezvoussavingoptimizerstates xm saveself optimizer statedict os path joinoutputdir optimizername with warnings catchwarningsrecordtrue as caughtwarnings xm saveself lrscheduler statedict os path joinoutputdir schedulername reissueptwarningscaughtwarnings elif issagemakermpenabled optstatedict self optimizer localstatedictgatherifshardfalse smp barrier if smp rdprank 0 or smp state cfg shardoptimizerstate smp save optstatedict os path joinoutputdir optimizername partialtrue v3smp state cfg shardoptimizerstate elif self isdeepspeedenabled under zero3 model file itself doesn t get saved since it s bogus unless deepspeed config stage3gather16bitweightsonmodelsave is true self modelwrapped savecheckpointoutputdir elif self isfsdpenabled save fsdp specific ckpt for resuming from ckpt savefsdpmodelself accelerator state fsdpplugin self accelerator self model outputdir savefsdpoptimizer self accelerator state fsdpplugin self accelerator self optimizer self model outputdir elif self args shouldsave deepspeed savecheckpoint above saves modeloptimsched torch saveself optimizer statedict os path joinoutputdir optimizername save scheduler scaler isdeepspeedcustomscheduler self isdeepspeedenabled and not isinstance self lrscheduler deepspeedschedulerwrapper if self args shouldsave and not self isdeepspeedenabled or isdeepspeedcustomscheduler and not istorchtpuavailable with warnings catchwarningsrecordtrue as caughtwarnings torch saveself lrscheduler statedict os path joinoutputdir schedulername reissueptwarningscaughtwarnings def loadoptimizerandschedulerself checkpoint deepspeed loads optimizerlrscheduler together with the model in deepspeedinit load in optimizer and scheduler states on tpu we have to take some extra precautions to properly load the states on the right device optimizer checkpoint was saved with smp 1 10 optimizer checkpoint was saved with smp 1 10 we use the cpu when training on one gpu to avoid oom for gpu ram when training big models in distributed training however we load directly on each gpu and risk the gpu oom as it s more likely to get oom on cpu since we load numgpu times the optimizer state launch an hyperparameter search using optuna or ray tune or sigopt the optimized quantity is determined by computeobjective which defaults to a function returning the evaluation loss when no metric is provided the sum of all metrics otherwise tip warningtrue to use this method you need to have provided a modelinit when initializing your trainer we need to reinitialize the model at each new run this is incompatible with the optimizers argument so you need to subclass trainer and override the method trainer createoptimizerandscheduler for custom optimizerscheduler tip args hpspace callableoptuna trial dictstr float optional a function that defines the hyperparameter search space will default to trainerutils defaulthpspaceoptuna or trainerutils defaulthpspaceray or trainerutils defaulthpspacesigopt depending on your backend computeobjective callabledictstr float float optional a function computing the objective to minimize or maximize from the metrics returned by the evaluate method will default to trainerutils defaultcomputeobjective ntrials int optional defaults to 100 the number of trial runs to test direction str or liststr optional defaults to minimize if it s single objective optimization direction is str can be minimize or maximize you should pick minimize when optimizing the validation loss maximize when optimizing one or several metrics if it s multi objectives optimization direction is liststr can be list of minimize and maximize you should pick minimize when optimizing the validation loss maximize when optimizing one or several metrics backend str or trainingutils hpsearchbackend optional the backend to use for hyperparameter search will default to optuna or ray tune or sigopt depending on which one is installed if all are installed will default to optuna hpname callableoptuna trial str optional a function that defines the trialrun name will default to none kwargs dictstr any optional additional keyword arguments passed along to optuna createstudy or ray tune run for more information see the documentation of optuna createstudyhttps optuna readthedocs ioenstablereferencegeneratedoptuna study createstudy html the documentation of tune runhttps docs ray ioenlatesttuneapidocsexecution htmltunerun the documentation of sigopthttps app sigopt comdocsendpointsexperimentscreate returns trainerutils bestrun or listtrainerutils bestrun all the information about the best run or best runs for multiobjective optimization experiment summary can be found in runsummary attribute for ray backend log logs on the various objects watching training subclass and override this method to inject custom behavior args logs dictstr float the values to log prepares one data before feeding it to the model be it a tensor or a nested listdictionary of tensors nlp models inputs are intuint and those get adjusted to the right dtype of the embedding other models such as wav2vec2 s inputs are already float and thus may need special handling to match the dtypes of the model prepare inputs before feeding them to the model converting them to tensors if they are not already and handling potential state a helper wrapper to group together context managers a helper wrapper that creates an appropriate context manager for autocast while feeding it the desired arguments depending on the situation perform a training step on a batch of inputs subclass and override to inject custom behavior args model nn module the model to train inputs dictstr uniontorch tensor any the inputs and targets of the model the dictionary will be unpacked before being fed to the model most models expect the targets under the argument labels check your model s documentation for all accepted arguments return torch tensor the tensor with training loss on this batch how the loss is computed by trainer by default all models return the loss in the first element subclass and override for custom behavior save past state if it exists todo this needs to be fixed and made cleaner later we don t use loss here since the model may return tuples instead of modeloutput whether or not this process is the local e g on one machine if training in a distributed fashion on several machines main process whether or not this process is the global main process when training in a distributed fashion on several machines this is only going to be true for one process special case for sagemaker modelparallel since there processindex is dpprocessindex not the global process index will save the model so you can reload it using frompretrained will only save from the main process calling the statedict needs to be done on the wrapped model and on all processes usercontent pt indicates model statedict saved with smp 1 10 remove the dummy statedict push to the hub when savemodel is called by the user save a trained model and configuration using savepretrained they can then be reloaded using frompretrained if we are executing this function we are the process zero so we don t check for that save a trained model and configuration using savepretrained they can then be reloaded using frompretrained good practice save your training arguments together with the trained model storing the number of floatingpoint operations that went into the model make sure we don t delete the best model check if we should delete older checkpoints if savetotallimit1 with loadbestmodelatendtrue we could end up deleting the last checkpoint which we don t do to allow resuming run evaluation and returns metrics the calling script will be responsible for providing a method to compute metrics as they are taskdependent pass it to the init computemetrics argument you can also subclass and override this method to inject custom behavior args evaldataset dataset optional pass a dataset if you wish to override self evaldataset if it is a datasets dataset columns not accepted by the model forward method are automatically removed it must implement the len method ignorekeys liststr optional a list of keys in the output of your model if it is a dictionary that should be ignored when gathering predictions metrickeyprefix str optional defaults to eval an optional prefix to be used as the metrics key prefix for example the metrics bleu will be named evalbleu if the prefix is eval default returns a dictionary containing the evaluation loss and the potential metrics computed from the predictions the dictionary also contains the epoch number which comes from the training state memory metrics must set up as early as possible no point gathering the predictions if there are no metrics otherwise we defer to self args predictionlossonly tpucomment logging debug metrics for pytorchxla compile execute times ops etc run prediction and returns predictions and potential metrics depending on the dataset and your use case your test dataset may contain labels in that case this method will also return metrics like in evaluate args testdataset dataset dataset to run the predictions on if it is an datasets dataset columns not accepted by the model forward method are automatically removed has to implement the method len ignorekeys liststr optional a list of keys in the output of your model if it is a dictionary that should be ignored when gathering predictions metrickeyprefix str optional defaults to test an optional prefix to be used as the metrics key prefix for example the metrics bleu will be named testbleu if the prefix is test default tip if your predictions or labels have different sequence length for instance because you re doing dynamic padding in a token classification task the predictions will be padded on the right to allow for concatenation into one array the padding index is 100 tip returns namedtuple a namedtuple with the following keys predictions np ndarray the predictions on testdataset labelids np ndarray optional the labels if the dataset contained some metrics dictstr float optional the potential dictionary of metrics if the dataset contained labels memory metrics must set up as early as possible predictionevaluation loop shared by trainer evaluate and trainer predict works both with or without labels if eval is called wo train handle model prep here for the rest of this function model is the outside model whether it was wrapped or not backward compatibility if full fp16 or bf16 eval is wanted and this evaluation or predict isn t called while train is running cast it to the right dtype first and then put on device do this before wrapping initialize containers lossespredslabels on gputpu accumulated for evalaccumulationsteps lossespredslabels on cpu final containers will be useful when we have an iterable dataset so don t know its length main evaluation loop update the observed num examples for batch samplers batchsize is not known by the dataloader in advance prediction step update containers on host gather all tensors and put them back on the cpu if we have done enough accumulation steps set back to none to begin a new accumulation after all calls to gatherfunction reset to gatherformetrics clean the state at the end of the evaluation loop gather all remaining tensors and put them back on the cpu number of samples the instance check is weird and does not actually check for the type but whether the dataset has the right methods therefore we need to make sure it also has the attribute metrics to be jsonserializable we need to remove numpy types or zerod tensors prefix all keys with metrickeyprefix gather value of tensors tensor or listtuple of nested tensors and convert them to numpy before concatenating them to gathered perform an evaluation step on model using inputs subclass and override to inject custom behavior args model nn module the model to evaluate inputs dictstr uniontorch tensor any the inputs and targets of the model the dictionary will be unpacked before being fed to the model most models expect the targets under the argument labels check your model s documentation for all accepted arguments predictionlossonly bool whether or not to return the loss only ignorekeys liststr optional a list of keys in the output of your model if it is a dictionary that should be ignored when gathering predictions return tupleoptionaltorch tensor optionaltorch tensor optionaltorch tensor a tuple with the loss logits and labels each being optional for cliplike models capable of returning loss values if returnloss is not specified or being none in inputs we check if the default value of returnloss is true in model forward labels may be popped when computing the loss label smoothing for instance so we grab them first todo this needs to be fixed and made cleaner later for models that inherit from pretrainedmodel uses that method to compute the number of floating point operations for every backward forward pass if using another model either implement such a method in the model or subclass and override this method args inputs dictstr uniontorch tensor any the inputs and targets of the model returns int the number of floatingpoint operations initializes a git repo in self args hubmodelid only on process zero creates a draft of a model card using the information available to the trainer args language str optional the language of the model if applicable license str optional the license of the model will default to the license of the pretrained model used if the original model given to the trainer comes from a repo on the hub tags str or liststr optional some tags to be included in the metadata of the model card modelname str optional the name of the model finetunedfrom str optional the name of the model used to finetune this one if applicable will default to the name of the repo of the original model given to the trainer if it comes from the hub tasks str or liststr optional one or several task identifiers to be included in the metadata of the model card datasettags str or liststr optional one or several dataset tags to be included in the metadata of the model card dataset str or liststr optional one or several dataset identifiers to be included in the metadata of the model card datasetargs str or liststr optional one or several dataset arguments to be included in the metadata of the model card only push from one node if we haven t finished the last push we don t do this one unless args hubalwayspushtrue to avoid a new synchronization of all model weights we just copy the file from the checkpoint folder saving the tokenizer is fast and we don t know how many files it may have spawned so we resave it to be sure same for the training arguments upload self model and self tokenizer to the model hub on the repo self args hubmodelid parameters commitmessage str optional defaults to end of training message to commit while pushing blocking bool optional defaults to true whether the function should return only when the git push has finished kwargs dictstr any optional additional keyword arguments passed along to trainer createmodelcard returns the url of the repository where the model was pushed if blockingfalse or a future object tracking the progress of the commit if blockingtrue in case the user calls this method with args pushtohub false needs to be executed on all processes for tpu training but will only save on the processed determined by self args shouldsave only push from one node wait for the current upload to be finished deprecated code predictionevaluation loop shared by trainer evaluate and trainer predict works both with or without labels if eval is called wo train handle model prep here for the rest of this function model is the outside model whether it was wrapped or not backward compatibility if full fp16 or bf16 eval is wanted and this evaluation or predict isn t called while train is running cast it to the right dtype first and then put on device the actual number of evalsample can be greater than numexamples in distributed settings when we pass a batch size to the sampler gather all tensors and put them back on the cpu if we have done enough accumulation steps set back to none to begin a new accumulation clean the state at the end of the evaluation loop gather all remaining tensors and put them back on the cpu to be jsonserializable we need to remove numpy types or zerod tensors prefix all keys with metrickeyprefix gather value of tensors tensor or listtuple of nested tensors and convert them to numpy before concatenating them to gathered add sagemaker checkpointing patterns to gitignore file make sure we only do this on the main process if not self isworldprocesszero return patterns sagemakeruploading sagemakeruploaded get current gitignore content if os path existsos path joinself repo localdir gitignore with openos path joinself repo localdir gitignore r as f currentcontent f read else currentcontent add the patterns to gitignore content currentcontent for pattern in patterns if pattern not in content if content endswithn content pattern else content fnpattern write the gitignore file if it has changed if content currentcontent with openos path joinself repo localdir gitignore w as f logger debugfwriting gitignore file content content f writecontent self repo gitadd gitignore avoid race condition with git status time sleep0 5 if not self repo isrepoclean self repo gitcommitadd sagemaker patterns to gitignore self repo gitpush def createacceleratorandpostprocessself gradacckwargs numsteps self args gradientaccumulationsteps gradacckwargssyncwithdataloader false gradientaccumulationplugin gradientaccumulationplugingradacckwargs create accelerator object self accelerator accelerator dispatchbatchesself args dispatchbatches splitbatchesself args splitbatches deepspeedpluginself args deepspeedplugin gradientaccumulationplugingradientaccumulationplugin some trainer classes need to use gather instead of gatherformetrics thus we store a flag self gatherfunction self accelerator gatherformetrics deepspeed and accelerate flags covering both trainer args and accelerate launcher self isdeepspeedenabled getattrself accelerator state deepspeedplugin none is not none self isfsdpenabled getattrself accelerator state fsdpplugin none is not none post accelerator creation setup if self isfsdpenabled fsdpplugin self accelerator state fsdpplugin fsdpplugin limitallgathers self args fsdpconfig get limitallgathers fsdpplugin limitallgathers if isaccelerateavailable0 23 0 fsdpplugin activationcheckpointing self args fsdpconfig get activationcheckpointing fsdpplugin activationcheckpointing if fsdpplugin activationcheckpointing and self args gradientcheckpointing raise valueerror the activationcheckpointing in fsdp config and the gradientcheckpointing in training arg can t be set to true simultaneously please use fsdp s activationcheckpointing logic when using fsdp if self isdeepspeedenabled if getattrself args hfdeepspeedconfig none is none from transformers integrations deepspeed import hftrainerdeepspeedconfig dsplugin self accelerator state deepspeedplugin dsplugin hfdsconfig hftrainerdeepspeedconfigdsplugin hfdsconfig config dsplugin deepspeedconfig dsplugin hfdsconfig config dsplugin hfdsconfig trainerconfigprocessself args 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 the trainer class to easily train a transformers from scratch or finetune it on a new task integrations must be imported before ml frameworks isort off isort on name of the files used for checkpointing trainer is a simple but feature complete training and eval loop for pytorch optimized for transformers args model pretrainedmodel or torch nn module optional the model to train evaluate or use for predictions if not provided a model_init must be passed tip trainer is optimized to work with the pretrainedmodel provided by the library you can still use your own models defined as torch nn module as long as they work the same way as the transformers models tip args trainingarguments optional the arguments to tweak for training will default to a basic instance of trainingarguments with the output_dir set to a directory named tmp_trainer in the current directory if not provided data_collator datacollator optional the function to use to form a batch from a list of elements of train_dataset or eval_dataset will default to default_data_collator if no tokenizer is provided an instance of datacollatorwithpadding otherwise train_dataset torch utils data dataset or torch utils data iterabledataset optional the dataset to use for training if it is a datasets dataset columns not accepted by the model forward method are automatically removed note that if it s a torch utils data iterabledataset with some randomization and you are training in a distributed fashion your iterable dataset should either use a internal attribute generator that is a torch generator for the randomization that must be identical on all processes and the trainer will manually set the seed of this generator at each epoch or have a set_epoch method that internally sets the seed of the rngs used eval_dataset union torch utils data dataset dict str torch utils data dataset optional the dataset to use for evaluation if it is a datasets dataset columns not accepted by the model forward method are automatically removed if it is a dictionary it will evaluate on each dataset prepending the dictionary key to the metric name tokenizer pretrainedtokenizerbase optional the tokenizer used to preprocess the data if provided will be used to automatically pad the inputs to the maximum length when batching inputs and it will be saved along the model to make it easier to rerun an interrupted training or reuse the fine tuned model model_init callable pretrainedmodel optional a function that instantiates the model to be used if provided each call to trainer train will start from a new instance of the model as given by this function the function may have zero argument or a single one containing the optuna ray tune sigopt trial object to be able to choose different architectures according to hyper parameters such as layer count sizes of inner layers dropout probabilities etc compute_metrics callable evalprediction dict optional the function that will be used to compute metrics at evaluation must take a evalprediction and return a dictionary string to metric values callbacks list of trainercallback optional a list of callbacks to customize the training loop will add those to the list of default callbacks detailed in here callback if you want to remove one of the default callbacks used use the trainer remove_callback method optimizers tuple torch optim optimizer torch optim lr_scheduler lambdalr optional defaults to none none a tuple containing the optimizer and the scheduler to use will default to an instance of adamw on your model and a scheduler given by get_linear_schedule_with_warmup controlled by args preprocess_logits_for_metrics callable torch tensor torch tensor torch tensor optional a function that preprocess the logits right before caching them at each evaluation step must take two tensors the logits and the labels and return the logits once processed as desired the modifications made by this function will be reflected in the predictions received by compute_metrics note that the labels second parameter will be none if the dataset does not have them important attributes model always points to the core model if using a transformers model it will be a pretrainedmodel subclass model_wrapped always points to the most external model in case one or more other modules wrap the original model this is the model that should be used for the forward pass for example under deepspeed the inner model is wrapped in deepspeed and then again in torch nn distributeddataparallel if the inner model hasn t been wrapped then self model_wrapped is the same as self model is_model_parallel whether or not a model has been switched to a model parallel mode different from data parallelism this means some of the model layers are split on different gpus place_model_on_device whether or not to automatically place the model on the device it will be set to false if model parallel or deepspeed is used or if the default trainingarguments place_model_on_device is overridden to return false is_in_train whether or not a model is currently running train e g when evaluate is called while in train those are used as methods of the trainer in examples seed must be set before instantiating the model when using model memory metrics must set up as early as possible set the correct log level depending on the node force device and distributed setup init explicitly warn users at this stage the model is already loaded one place to sort out whether to place the model on device or not postpone switching model to cuda when 1 mp since we are trying to fit a much bigger than 1 gpu model 2 fp16 enabled deepspeed loads the model in half the size and it doesn t need to anyway and we only use deepspeed for training at the moment 3 full bf16 or fp16 eval since the model needs to be cast to the right dtype first 4 fsdp same as mp bnb quantized models doesn t support to operation force n_gpu to 1 to avoid dataparallel as mp will manage the gpus later use self model is self model_wrapped to check if it s wrapped or not will be set to true by self _setup_loggers on first call to self log create distant repo and output directory if needed mixed precision setup mixed precision setup for sagemaker model parallel bf16 model parallelism in sagemaker currently not supported raise an error when there s mismatch between smp config and trainer argument use smp config as truth smp 1 10 does not support fp16 in trainer deepspeed and sagemaker model parallel manage their own half precision label smoothing internal variable to count flos in each process will be accumulated in self state total_flos then returned to 0 every time flos need to be logged internal variables to help with automatic batch size reduction very last torch compile activates the neftune as presented in this code https github com neelsjain neftune and paper https arxiv org abs 2310 05914 deactivates the neftune method make sure to call _activate_neftune first add a callback to the current list of transformers trainercallback args callback type or transformers trainercallback a transformers trainercallback class or an instance of a transformers trainercallback in the first case will instantiate a member of that class remove a callback from the current list of transformers trainercallback and returns it if the callback is not found returns none and no error is raised args callback type or transformers trainercallback a transformers trainercallback class or an instance of a transformers trainercallback in the first case will pop the first member of that class found in the list of callbacks returns transformers trainercallback the callback removed if found remove a callback from the current list of transformers trainercallback args callback type or transformers trainercallback a transformers trainercallback class or an instance of a transformers trainercallback in the first case will remove the first member of that class found in the list of callbacks moving a model to an xla device disconnects the tied weights so we have to retie them inspect model forward signature to keep only the arguments it accepts labels may be named label or label_ids the default data collator handles that wrap the data collator in a callable removing unused columns build the sampler returns the training torch utils data dataloader will use no sampler if train_dataset does not implement __len__ a random sampler adapted to distributed training if necessary otherwise subclass and override this method if you want to inject some custom behavior deprecated code returns the evaluation torch utils data dataloader subclass and override this method if you want to inject some custom behavior args eval_dataset torch utils data dataset optional if provided will override self eval_dataset if it is a datasets dataset columns not accepted by the model forward method are automatically removed it must implement __len__ returns the test torch utils data dataloader subclass and override this method if you want to inject some custom behavior args test_dataset torch utils data dataset optional the test dataset to use if it is a datasets dataset columns not accepted by the model forward method are automatically removed it must implement __len__ we use the same batch_size as for eval setup the optimizer and the learning rate scheduler we provide a reasonable default that works well if you want to use something else you can pass a tuple in the trainer s init through optimizers or subclass and override this method or create_optimizer and or create_scheduler in a subclass if smp 1 10 and fp16 is enabled we unwrap the optimizer get all parameter names that weight decay will be applied to note that some models implement their own layernorm instead of calling nn layernorm weight decay could still apply to those modules since this function only filter out instance of nn layernorm setup the optimizer we provide a reasonable default that works well if you want to use something else you can pass a tuple in the trainer s init through optimizers or subclass and override this method in a subclass returns the optimizer class and optimizer parameters based on the training arguments args args transformers training_args trainingarguments the training arguments for the training session parse args optim_args todo change dtypes back to m fp32 var bf16 kahan false once they can be cast together in torchdistx setup the scheduler the optimizer of the trainer must have been set up either before this method is called or passed as an argument args num_training_steps int the number of training steps to do helper to get number of samples in a torch utils data dataloader by accessing its dataset when dataloader dataset does not exist or has no length estimates as best it can special case for iterabledatasetshard we need to dig deeper no dataset or length estimate by length of dataloader helper to get number of tokens in a torch utils data dataloader by enumerating dataloader hp search setup code casting value to the proper type rebuild the deepspeed config to reflect the updated training parameters remove mixed precision hooks from the model conv_bn_folding is disabled as it fails in symbolic tracing resulting in ipex warnings wrapping the base model twice in a distributedmodel will raise an error train eval could be run multiple times if already wrapped don t re wrap it again mixed precision training with apex torch 1 6 multi gpu training should be after apex fp16 initialization 8bit models does not support ddp note in torch distributed mode there s no point in wrapping the model inside a distributeddataparallel as we ll be under no_grad anyways distributed training should be after apex fp16 initialization distributed training using pytorch fsdp transformer layer class to wrap apply gradient checkpointing to auto wrapped sub modules if specified wrap the base model with an outer fsdp wrapper patch xm optimizer_step should not reduce gradients in this case as fsdp does not need gradient reduction over sharded parameters find_unused_parameters breaks checkpointing as per https github com huggingface transformers pull 4659 issuecomment 643356021 main training entry point args resume_from_checkpoint str or bool optional if a str local path to a saved checkpoint as saved by a previous instance of trainer if a bool and equals true load the last checkpoint in args output_dir as saved by a previous instance of trainer if present training will resume from the model optimizer scheduler states loaded here trial optuna trial or dict str any optional the trial run or the hyperparameter dictionary for hyperparameter search ignore_keys_for_eval list str optional a list of keys in the output of your model if it is a dictionary that should be ignored when gathering predictions for evaluation during the training kwargs dict str any optional additional keyword arguments used to hide deprecated arguments memory metrics must set up as early as possible attach neftune hooks if necessary do_train is not a reliable argument as it might not be set and train still called so the following is a workaround this might change the seed so needs to run first model re init seed must be set before instantiating the model when using model_init reinitializes optimizer and scheduler load potential model checkpoint if model was re initialized put it on the right device and update self model_wrapped disable progress bars when uploading models during checkpoints to avoid polluting stdout data loader and number of training steps setting up training control variables number of training epochs num_train_epochs number of training steps per epoch num_update_steps_per_epoch total number of training steps to execute max_steps may be slightly incorrect if the last batch in the training dataloader has a smaller size but it s the best we can do rely on max_steps when dataloader does not have a working size setting a very large number of epochs so we go as many times as necessary over the iterator nn dataparallel model replicates the model creating new variables and module references registered here no longer work on other gpus breaking the module noqa we need to reset the scheduler as its parameters may be different on subsequent calls compute absolute values for logging eval and save if given as ratio activate gradient checkpointing if needed as the model is wrapped don t use accelerator prepare this is for unhandled cases such as fsdp xla sagemaker mp dp dataparallel ipex prepare using accelerator prepare to handle cases wherein we pass dummyscheduler such as when it is specified in deepspeed config for the rest of this function model is the outside model whether it was wrapped or not backward compatibility ckpt loading check if saved optimizer or scheduler states exist important at this point self model is the transformers model self model_wrapped is ddp transformers model deepspeed transformers model fsdp transformers model dynamo optimized module transformers model etc train check if continuing training from a checkpoint update the references use self _trial because the sigopt optuna hpo only call _hp_search_setup trial instead of passing trial parameter to train when using ddp this should be the same if the state has been saved but in case the training arguments changed it s safer to set this after the load tr_loss is a tensor to avoid synchronization of tpus through item _total_loss_scalar is updated everytime item has to be called on tr_loss and stores the sum of all losses skip the first epochs_trained epochs to get the random state of the dataloader at the right point we just need to begin an iteration to create the randomization of the sampler otherwise we need to call the whooooole sampler cause there is some random operation added at the very end reset the past mems state at the beginning of each epoch if necessary skip past any already trained steps if resuming training if loss is nan or inf simply add the average of previous logged losses last step in epoch but step is always smaller than gradient_accumulation_steps the or condition of is_last_step_and_steps_less_than_grad_acc is not covered in accelerate so explicitly enable sync gradients to true in that case gradient clipping deepspeed does its own clipping revert to normal clipping otherwise handling apex or full precision optimizer step delay optimizer scheduling until metrics are generated tpu comment logging debug metrics for pytorch xla compile execute times ops etc clean the state at the end of training wait for everyone to get here so we are sure the model has been saved by process 0 add remaining tr_loss delete the last checkpoint when save_total_limit 1 if it s different from the best checkpoint and process allowed to save wait for the checkpoint to be uploaded after training we make sure to retrieve back the original forward pass method for the embedding layer by removing the forward post hook if the model is on the gpu it still works if the user_content pt file exists load with the new smp api checkpoint must have been saved with the new smp api if the user_content pt file does not exist load with the old smp api checkpoint must have been saved with the old smp api required for smp to not auto translate state_dict from hf to smp is already smp release memory we load the model state dict on the cpu to avoid an oom error workaround for fsdp bug https github com pytorch pytorch issues 82963 which takes args instead of kwargs release memory load adapters following pr 24096 if train a model using peft lora assume that adapter have been saved properly we load the sharded checkpoint if the user_content pt file exists load with the new smp api checkpoint must have been saved with the new smp api if the user_content pt file does not exist load with the old smp api checkpoint must have been saved with the old smp api if train a model using peft lora assume that adapter have been saved properly load_adapter has no return value present modify it when appropriate we load the model state dict on the cpu to avoid an oom error if the model is on the gpu it still works workaround for fsdp bug https github com pytorch pytorch issues 82963 which takes args instead of kwargs all_gather mean to get average loss over all processes reset tr_loss to zero run delayed lr scheduler now that metrics are populated load rng states from checkpoint in all cases including ddp dp deepspeed self model is always a reference to the model we want to save except fullyshardedddp assert unwrap_model model is self model internal model should be a reference to self model save model checkpoint save optimizer and scheduler save rng state determine the new best metric best model checkpoint save the trainer state maybe delete some older checkpoints save rng state in non distributed training in non distributed we save the global cuda rng state will take care of dataparallel a process can arrive here before the process 0 has a chance to save the model in which case output_dir may not yet exist under zero3 model file itself doesn t get saved since it s bogus unless deepspeed config stage3_gather_16bit_weights_on_model_save is true save fsdp specific ckpt for resuming from ckpt deepspeed save_checkpoint above saves model optim sched save scheduler scaler if optimizer and scheduler states exist load them deepspeed loads optimizer lr_scheduler together with the model in deepspeed_init load in optimizer and scheduler states on tpu we have to take some extra precautions to properly load the states on the right device optimizer checkpoint was saved with smp 1 10 optimizer checkpoint was saved with smp 1 10 we use the cpu when training on one gpu to avoid oom for gpu ram when training big models in distributed training however we load directly on each gpu and risk the gpu oom as it s more likely to get oom on cpu since we load num_gpu times the optimizer state launch an hyperparameter search using optuna or ray tune or sigopt the optimized quantity is determined by compute_objective which defaults to a function returning the evaluation loss when no metric is provided the sum of all metrics otherwise tip warning true to use this method you need to have provided a model_init when initializing your trainer we need to reinitialize the model at each new run this is incompatible with the optimizers argument so you need to subclass trainer and override the method trainer create_optimizer_and_scheduler for custom optimizer scheduler tip args hp_space callable optuna trial dict str float optional a function that defines the hyperparameter search space will default to trainer_utils default_hp_space_optuna or trainer_utils default_hp_space_ray or trainer_utils default_hp_space_sigopt depending on your backend compute_objective callable dict str float float optional a function computing the objective to minimize or maximize from the metrics returned by the evaluate method will default to trainer_utils default_compute_objective n_trials int optional defaults to 100 the number of trial runs to test direction str or list str optional defaults to minimize if it s single objective optimization direction is str can be minimize or maximize you should pick minimize when optimizing the validation loss maximize when optimizing one or several metrics if it s multi objectives optimization direction is list str can be list of minimize and maximize you should pick minimize when optimizing the validation loss maximize when optimizing one or several metrics backend str or training_utils hpsearchbackend optional the backend to use for hyperparameter search will default to optuna or ray tune or sigopt depending on which one is installed if all are installed will default to optuna hp_name callable optuna trial str optional a function that defines the trial run name will default to none kwargs dict str any optional additional keyword arguments passed along to optuna create_study or ray tune run for more information see the documentation of optuna create_study https optuna readthedocs io en stable reference generated optuna study create_study html the documentation of tune run https docs ray io en latest tune api_docs execution html tune run the documentation of sigopt https app sigopt com docs endpoints experiments create returns trainer_utils bestrun or list trainer_utils bestrun all the information about the best run or best runs for multi objective optimization experiment summary can be found in run_summary attribute for ray backend log logs on the various objects watching training subclass and override this method to inject custom behavior args logs dict str float the values to log prepares one data before feeding it to the model be it a tensor or a nested list dictionary of tensors nlp models inputs are int uint and those get adjusted to the right dtype of the embedding other models such as wav2vec2 s inputs are already float and thus may need special handling to match the dtypes of the model prepare inputs before feeding them to the model converting them to tensors if they are not already and handling potential state a helper wrapper to group together context managers a helper wrapper that creates an appropriate context manager for autocast while feeding it the desired arguments depending on the situation perform a training step on a batch of inputs subclass and override to inject custom behavior args model nn module the model to train inputs dict str union torch tensor any the inputs and targets of the model the dictionary will be unpacked before being fed to the model most models expect the targets under the argument labels check your model s documentation for all accepted arguments return torch tensor the tensor with training loss on this batch mean to average on multi gpu parallel training how the loss is computed by trainer by default all models return the loss in the first element subclass and override for custom behavior save past state if it exists todo this needs to be fixed and made cleaner later we don t use loss here since the model may return tuples instead of modeloutput whether or not this process is the local e g on one machine if training in a distributed fashion on several machines main process whether or not this process is the global main process when training in a distributed fashion on several machines this is only going to be true for one process special case for sagemaker modelparallel since there process_index is dp_process_index not the global process index will save the model so you can reload it using from_pretrained will only save from the main process calling the state_dict needs to be done on the wrapped model and on all processes user_content pt indicates model state_dict saved with smp 1 10 remove the dummy state_dict push to the hub when save_model is called by the user save a trained model and configuration using save_pretrained they can then be reloaded using from_pretrained if we are executing this function we are the process zero so we don t check for that save a trained model and configuration using save_pretrained they can then be reloaded using from_pretrained good practice save your training arguments together with the trained model storing the number of floating point operations that went into the model make sure we don t delete the best model check if we should delete older checkpoint s if save_total_limit 1 with load_best_model_at_end true we could end up deleting the last checkpoint which we don t do to allow resuming run evaluation and returns metrics the calling script will be responsible for providing a method to compute metrics as they are task dependent pass it to the init compute_metrics argument you can also subclass and override this method to inject custom behavior args eval_dataset dataset optional pass a dataset if you wish to override self eval_dataset if it is a datasets dataset columns not accepted by the model forward method are automatically removed it must implement the __len__ method ignore_keys list str optional a list of keys in the output of your model if it is a dictionary that should be ignored when gathering predictions metric_key_prefix str optional defaults to eval an optional prefix to be used as the metrics key prefix for example the metrics bleu will be named eval_bleu if the prefix is eval default returns a dictionary containing the evaluation loss and the potential metrics computed from the predictions the dictionary also contains the epoch number which comes from the training state memory metrics must set up as early as possible no point gathering the predictions if there are no metrics otherwise we defer to self args prediction_loss_only tpu comment logging debug metrics for pytorch xla compile execute times ops etc run prediction and returns predictions and potential metrics depending on the dataset and your use case your test dataset may contain labels in that case this method will also return metrics like in evaluate args test_dataset dataset dataset to run the predictions on if it is an datasets dataset columns not accepted by the model forward method are automatically removed has to implement the method __len__ ignore_keys list str optional a list of keys in the output of your model if it is a dictionary that should be ignored when gathering predictions metric_key_prefix str optional defaults to test an optional prefix to be used as the metrics key prefix for example the metrics bleu will be named test_bleu if the prefix is test default tip if your predictions or labels have different sequence length for instance because you re doing dynamic padding in a token classification task the predictions will be padded on the right to allow for concatenation into one array the padding index is 100 tip returns namedtuple a namedtuple with the following keys predictions np ndarray the predictions on test_dataset label_ids np ndarray optional the labels if the dataset contained some metrics dict str float optional the potential dictionary of metrics if the dataset contained labels memory metrics must set up as early as possible prediction evaluation loop shared by trainer evaluate and trainer predict works both with or without labels if eval is called w o train handle model prep here for the rest of this function model is the outside model whether it was wrapped or not backward compatibility if full fp16 or bf16 eval is wanted and this evaluation or predict isn t called while train is running cast it to the right dtype first and then put on device do this before wrapping initialize containers losses preds labels on gpu tpu accumulated for eval_accumulation_steps losses preds labels on cpu final containers will be useful when we have an iterable dataset so don t know its length main evaluation loop update the observed num examples for batch samplers batch_size is not known by the dataloader in advance prediction step update containers on host gather all tensors and put them back on the cpu if we have done enough accumulation steps set back to none to begin a new accumulation after all calls to gather_function reset to gather_for_metrics clean the state at the end of the evaluation loop gather all remaining tensors and put them back on the cpu number of samples the instance check is weird and does not actually check for the type but whether the dataset has the right methods therefore we need to make sure it also has the attribute both len dataloader dataset and len dataloader fail metrics to be json serializable we need to remove numpy types or zero d tensors prefix all keys with metric_key_prefix _ gather value of tensors tensor or list tuple of nested tensors and convert them to numpy before concatenating them to gathered perform an evaluation step on model using inputs subclass and override to inject custom behavior args model nn module the model to evaluate inputs dict str union torch tensor any the inputs and targets of the model the dictionary will be unpacked before being fed to the model most models expect the targets under the argument labels check your model s documentation for all accepted arguments prediction_loss_only bool whether or not to return the loss only ignore_keys list str optional a list of keys in the output of your model if it is a dictionary that should be ignored when gathering predictions return tuple optional torch tensor optional torch tensor optional torch tensor a tuple with the loss logits and labels each being optional for clip like models capable of returning loss values if return_loss is not specified or being none in inputs we check if the default value of return_loss is true in model forward labels may be popped when computing the loss label smoothing for instance so we grab them first todo this needs to be fixed and made cleaner later for models that inherit from pretrainedmodel uses that method to compute the number of floating point operations for every backward forward pass if using another model either implement such a method in the model or subclass and override this method args inputs dict str union torch tensor any the inputs and targets of the model returns int the number of floating point operations initializes a git repo in self args hub_model_id only on process zero creates a draft of a model card using the information available to the trainer args language str optional the language of the model if applicable license str optional the license of the model will default to the license of the pretrained model used if the original model given to the trainer comes from a repo on the hub tags str or list str optional some tags to be included in the metadata of the model card model_name str optional the name of the model finetuned_from str optional the name of the model used to fine tune this one if applicable will default to the name of the repo of the original model given to the trainer if it comes from the hub tasks str or list str optional one or several task identifiers to be included in the metadata of the model card dataset_tags str or list str optional one or several dataset tags to be included in the metadata of the model card dataset str or list str optional one or several dataset identifiers to be included in the metadata of the model card dataset_args str or list str optional one or several dataset arguments to be included in the metadata of the model card only push from one node if we haven t finished the last push we don t do this one unless args hub_always_push true to avoid a new synchronization of all model weights we just copy the file from the checkpoint folder saving the tokenizer is fast and we don t know how many files it may have spawned so we resave it to be sure same for the training arguments upload self model and self tokenizer to the model hub on the repo self args hub_model_id parameters commit_message str optional defaults to end of training message to commit while pushing blocking bool optional defaults to true whether the function should return only when the git push has finished kwargs dict str any optional additional keyword arguments passed along to trainer create_model_card returns the url of the repository where the model was pushed if blocking false or a future object tracking the progress of the commit if blocking true in case the user calls this method with args push_to_hub false needs to be executed on all processes for tpu training but will only save on the processed determined by self args should_save only push from one node wait for the current upload to be finished deprecated code prediction evaluation loop shared by trainer evaluate and trainer predict works both with or without labels if eval is called w o train handle model prep here for the rest of this function model is the outside model whether it was wrapped or not backward compatibility if full fp16 or bf16 eval is wanted and this evaluation or predict isn t called while train is running cast it to the right dtype first and then put on device the actual number of eval_sample can be greater than num_examples in distributed settings when we pass a batch size to the sampler gather all tensors and put them back on the cpu if we have done enough accumulation steps set back to none to begin a new accumulation clean the state at the end of the evaluation loop gather all remaining tensors and put them back on the cpu to be json serializable we need to remove numpy types or zero d tensors prefix all keys with metric_key_prefix _ gather value of tensors tensor or list tuple of nested tensors and convert them to numpy before concatenating them to gathered add sagemaker checkpointing patterns to gitignore file make sure we only do this on the main process get current gitignore content add the patterns to gitignore write the gitignore file if it has changed avoid race condition with git status create accelerator object some trainer classes need to use gather instead of gather_for_metrics thus we store a flag deepspeed and accelerate flags covering both trainer args and accelerate launcher post accelerator creation setup
import contextlib
import copy
import functools
import glob
import importlib.metadata
import inspect
import math
import os
import random
import re
import shutil
import sys
import time
import warnings
from collections.abc import Mapping
from pathlib import Path
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union




from .integrations import (
    get_reporting_integration_callbacks,
    hp_params,
)



import huggingface_hub.utils as hf_hub_utils
import numpy as np
import torch
import torch.distributed as dist
from huggingface_hub import create_repo, upload_folder
from packaging import version
from torch import nn
from torch.utils.data import DataLoader, Dataset, RandomSampler, SequentialSampler

from . import __version__
from .configuration_utils import PretrainedConfig
from .data.data_collator import DataCollator, DataCollatorWithPadding, default_data_collator
from .debug_utils import DebugOption, DebugUnderflowOverflow
from .hyperparameter_search import ALL_HYPERPARAMETER_SEARCH_BACKENDS, default_hp_search_backend
from .integrations.deepspeed import deepspeed_init, deepspeed_load_checkpoint, is_deepspeed_available
from .modelcard import TrainingSummary
from .modeling_utils import PreTrainedModel, load_sharded_checkpoint, unwrap_model
from .models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES, MODEL_MAPPING_NAMES
from .optimization import Adafactor, get_scheduler
from .pytorch_utils import ALL_LAYERNORM_LAYERS, is_torch_less_than_1_11
from .tokenization_utils_base import PreTrainedTokenizerBase
from .trainer_callback import (
    CallbackHandler,
    DefaultFlowCallback,
    PrinterCallback,
    ProgressCallback,
    TrainerCallback,
    TrainerControl,
    TrainerState,
)
from .trainer_pt_utils import (
    DistributedTensorGatherer,
    IterableDatasetShard,
    LabelSmoother,
    LengthGroupedSampler,
    SequentialDistributedSampler,
    distributed_broadcast_scalars,
    distributed_concat,
    find_batch_size,
    get_dataloader_sampler,
    get_model_param_count,
    get_module_class_from_name,
    get_parameter_names,
    nested_concat,
    nested_detach,
    nested_numpify,
    nested_xla_mesh_reduce,
    reissue_pt_warnings,
    remove_dummy_checkpoint,
)
from .trainer_utils import (
    PREFIX_CHECKPOINT_DIR,
    BestRun,
    EvalLoopOutput,
    EvalPrediction,
    HPSearchBackend,
    HubStrategy,
    IntervalStrategy,
    PredictionOutput,
    RemoveColumnsCollator,
    TrainerMemoryTracker,
    TrainOutput,
    default_compute_objective,
    denumpify_detensorize,
    enable_full_determinism,
    find_executable_batch_size,
    get_last_checkpoint,
    has_length,
    neftune_post_forward_hook,
    number_of_arguments,
    seed_worker,
    set_seed,
    speed_metrics,
)
from .training_args import OptimizerNames, ParallelMode, TrainingArguments
from .utils import (
    ADAPTER_CONFIG_NAME,
    ADAPTER_SAFE_WEIGHTS_NAME,
    ADAPTER_WEIGHTS_NAME,
    CONFIG_NAME,
    SAFE_WEIGHTS_INDEX_NAME,
    SAFE_WEIGHTS_NAME,
    WEIGHTS_INDEX_NAME,
    WEIGHTS_NAME,
    PushInProgress,
    can_return_loss,
    find_labels,
    is_accelerate_available,
    is_apex_available,
    is_bitsandbytes_available,
    is_datasets_available,
    is_in_notebook,
    is_ipex_available,
    is_peft_available,
    is_safetensors_available,
    is_sagemaker_dp_enabled,
    is_sagemaker_mp_enabled,
    is_torch_compile_available,
    is_torch_neuroncore_available,
    is_torch_npu_available,
    is_torch_tpu_available,
    logging,
    strtobool,
)
from .utils.quantization_config import QuantizationMethod


DEFAULT_CALLBACKS = [DefaultFlowCallback]
DEFAULT_PROGRESS_CALLBACK = ProgressCallback

if is_in_notebook():
    from .utils.notebook import NotebookProgressCallback

    DEFAULT_PROGRESS_CALLBACK = NotebookProgressCallback

if is_apex_available():
    from apex import amp

if is_datasets_available():
    import datasets

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


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

    IS_SAGEMAKER_MP_POST_1_10 = version.parse(SMP_VERSION) >= version.parse("1.10")

    from .trainer_pt_utils import smp_forward_backward, smp_forward_only, smp_gather, smp_nested_concat
else:
    IS_SAGEMAKER_MP_POST_1_10 = False


if is_safetensors_available():
    import safetensors.torch


if is_peft_available():
    from peft import PeftModel


if is_accelerate_available():
    from accelerate import Accelerator, skip_first_batches
    from accelerate import __version__ as accelerate_version
    from accelerate.utils import (
        DistributedDataParallelKwargs,
        GradientAccumulationPlugin,
        load_fsdp_model,
        load_fsdp_optimizer,
        save_fsdp_model,
        save_fsdp_optimizer,
    )

    DATA_SAMPLERS = [RandomSampler]
    if version.parse(accelerate_version) > version.parse("0.23.0"):
        from accelerate.data_loader import SeedableRandomSampler

        DATA_SAMPLERS += [SeedableRandomSampler]

    if is_deepspeed_available():
        from accelerate.utils import DeepSpeedSchedulerWrapper


if TYPE_CHECKING:
    import optuna


logger = logging.get_logger(__name__)



TRAINING_ARGS_NAME = "training_args.bin"
TRAINER_STATE_NAME = "trainer_state.json"
OPTIMIZER_NAME = "optimizer.pt"
OPTIMIZER_NAME_BIN = "optimizer.bin"
SCHEDULER_NAME = "scheduler.pt"
SCALER_NAME = "scaler.pt"
FSDP_MODEL_NAME = "pytorch_model_fsdp"


class Trainer:
    

    
    from .trainer_pt_utils import _get_learning_rate, log_metrics, metrics_format, save_metrics, save_state

    def __init__(
        self,
        model: Union[PreTrainedModel, nn.Module] = None,
        args: TrainingArguments = None,
        data_collator: Optional[DataCollator] = None,
        train_dataset: Optional[Dataset] = None,
        eval_dataset: Optional[Union[Dataset, Dict[str, Dataset]]] = None,
        tokenizer: Optional[PreTrainedTokenizerBase] = None,
        model_init: Optional[Callable[[], PreTrainedModel]] = None,
        compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,
        callbacks: Optional[List[TrainerCallback]] = None,
        optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),
        preprocess_logits_for_metrics: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None,
    ):
        if args is None:
            output_dir = "tmp_trainer"
            logger.info(f"No `TrainingArguments` passed, using `output_dir={output_dir}`.")
            args = TrainingArguments(output_dir=output_dir)
        self.args = args
        
        enable_full_determinism(self.args.seed) if self.args.full_determinism else set_seed(self.args.seed)
        self.hp_name = None
        self.deepspeed = None
        self.is_in_train = False

        self.create_accelerator_and_postprocess()

        
        self._memory_tracker = TrainerMemoryTracker(self.args.skip_memory_metrics)
        self._memory_tracker.start()

        
        log_level = args.get_process_log_level()
        logging.set_verbosity(log_level)

        
        args._setup_devices

        if model is None:
            if model_init is not None:
                self.model_init = model_init
                model = self.call_model_init()
            else:
                raise RuntimeError("`Trainer` requires either a `model` or `model_init` argument")
        else:
            if model_init is not None:
                warnings.warn(
                    "`Trainer` requires either a `model` or `model_init` argument, but not both. `model_init` will"
                    " overwrite your model when calling the `train` method. This will become a fatal error in the next"
                    " release.",
                    FutureWarning,
                )
            self.model_init = model_init

        if model.__class__.__name__ in MODEL_MAPPING_NAMES:
            raise ValueError(
                f"The model you have picked ({model.__class__.__name__}) cannot be used as is for training: it only "
                "computes hidden states and does not accept any labels. You should choose a model with a head "
                "suitable for your task like any of the `AutoModelForXxx` listed at "
                "https://huggingface.co/docs/transformers/model_doc/auto"
            )

        if hasattr(model, "is_parallelizable") and model.is_parallelizable and model.model_parallel:
            self.is_model_parallel = True
        else:
            self.is_model_parallel = False

        if getattr(model, "hf_device_map", None) is not None:
            devices = [device for device in set(model.hf_device_map.values()) if device not in ["cpu", "disk"]]
            if len(devices) > 1:
                self.is_model_parallel = True
            elif len(devices) == 1:
                self.is_model_parallel = self.args.device != torch.device(devices[0])
            else:
                self.is_model_parallel = False

            
            if self.is_model_parallel:
                logger.info(
                    "You have loaded a model on multiple GPUs. `is_model_parallel` attribute will be force-set"
                    " to `True` to avoid any unexpected behavior such as device placement mismatching."
                )

        _is_peft_model = is_peft_available() and isinstance(model, PeftModel)
        _is_quantized_and_base_model = getattr(model, "is_quantized", False) and not getattr(
            model, "_hf_peft_config_loaded", False
        )

        
        if _is_quantized_and_base_model and not _is_peft_model:
            raise ValueError(
                "You cannot perform fine-tuning on purely quantized models. Please attach trainable adapters on top of"
                " the quantized model to correctly perform fine-tuning. Please see: https://huggingface.co/docs/transformers/peft"
                " for more details"
            )
        elif _is_quantized_and_base_model and not getattr(model, "_is_quantized_training_enabled", False):
            raise ValueError(
                "The model you want to train is loaded in 8-bit precision.  if you want to fine-tune an 8-bit"
                " model, please make sure that you have installed `bitsandbytes>=0.37.0`. "
            )

        self.is_fsdp_xla_enabled = args.fsdp_config["xla"]
        if len(args.fsdp) > 0:
            if self.is_deepspeed_enabled:
                raise ValueError(
                    "Using --fsdp xxx together with --deepspeed is not possible, deactivate one of those flags."
                )
            if not args.fsdp_config["xla"] and args.parallel_mode != ParallelMode.DISTRIBUTED:
                raise ValueError("Using fsdp only works in distributed training.")

        
        
        
        
        
        
        
        self.place_model_on_device = args.place_model_on_device
        if (
            self.is_model_parallel
            or self.is_deepspeed_enabled
            or ((args.fp16_full_eval or args.bf16_full_eval) and not args.do_train)
            or self.is_fsdp_xla_enabled
            or self.is_fsdp_enabled
        ):
            self.place_model_on_device = False

        default_collator = default_data_collator if tokenizer is None else DataCollatorWithPadding(tokenizer)
        self.data_collator = data_collator if data_collator is not None else default_collator
        self.train_dataset = train_dataset
        self.eval_dataset = eval_dataset
        self.tokenizer = tokenizer

        
        if (
            self.place_model_on_device
            and not getattr(model, "quantization_method", None) == QuantizationMethod.BITS_AND_BYTES
        ):
            self._move_model_to_device(model, args.device)

        
        if self.is_model_parallel:
            self.args._n_gpu = 1

        
        self.model_wrapped = model
        self.model = model

        self.neftune_noise_alpha = args.neftune_noise_alpha

        self.compute_metrics = compute_metrics
        self.preprocess_logits_for_metrics = preprocess_logits_for_metrics
        self.optimizer, self.lr_scheduler = optimizers
        if model_init is not None and (self.optimizer is not None or self.lr_scheduler is not None):
            raise RuntimeError(
                "Passing a `model_init` is incompatible with providing the `optimizers` argument. "
                "You should subclass `Trainer` and override the `create_optimizer_and_scheduler` method."
            )
        if is_torch_tpu_available() and self.optimizer is not None:
            for param in self.model.parameters():
                model_device = param.device
                break
            for param_group in self.optimizer.param_groups:
                if len(param_group["params"]) > 0:
                    optimizer_device = param_group["params"][0].device
                    break
            if model_device != optimizer_device:
                raise ValueError(
                    "The model and the optimizer parameters are not on the same device, which probably means you"
                    " created an optimizer around your model **before** putting on the device and passing it to the"
                    " `Trainer`. Make sure the lines `import torch_xla.core.xla_model as xm` and"
                    " `model.to(xm.xla_device())` is performed before the optimizer creation in your script."
                )
        if (self.is_deepspeed_enabled or self.is_fsdp_xla_enabled or self.is_fsdp_enabled) and (
            self.optimizer is not None or self.lr_scheduler is not None
        ):
            raise RuntimeError(
                "Passing `optimizers` is not allowed if Deepspeed or PyTorch FSDP is enabled. "
                "You should subclass `Trainer` and override the `create_optimizer_and_scheduler` method."
            )
        default_callbacks = DEFAULT_CALLBACKS + get_reporting_integration_callbacks(self.args.report_to)
        callbacks = default_callbacks if callbacks is None else default_callbacks + callbacks
        self.callback_handler = CallbackHandler(
            callbacks, self.model, self.tokenizer, self.optimizer, self.lr_scheduler
        )
        self.add_callback(PrinterCallback if self.args.disable_tqdm else DEFAULT_PROGRESS_CALLBACK)

        
        self._loggers_initialized = False

        
        self.hub_model_id = None
        if self.args.push_to_hub:
            self.init_hf_repo()
        if self.args.should_save:
            os.makedirs(self.args.output_dir, exist_ok=True)

        if not callable(self.data_collator) and callable(getattr(self.data_collator, "collate_batch", None)):
            raise ValueError("The `data_collator` should be a simple callable (function, class with `__call__`).")

        if args.max_steps > 0:
            logger.info("max_steps is given, it will override any value given in num_train_epochs")

        if train_dataset is not None and not has_length(train_dataset) and args.max_steps <= 0:
            raise ValueError(
                "The train_dataset does not implement __len__, max_steps has to be specified. "
                "The number of steps needs to be known in advance for the learning rate scheduler."
            )

        if (
            train_dataset is not None
            and isinstance(train_dataset, torch.utils.data.IterableDataset)
            and args.group_by_length
        ):
            raise ValueError("the `--group_by_length` option is only available for `Dataset`, not `IterableDataset")

        self._signature_columns = None

        
        self.use_apex = False
        self.use_cpu_amp = False

        
        if is_sagemaker_mp_enabled():
            
            if args.bf16:
                raise ValueError("SageMaker Model Parallelism does not support BF16 yet. Please use FP16 instead ")

            if IS_SAGEMAKER_MP_POST_1_10:
                
                if args.fp16 != smp.state.cfg.fp16:
                    logger.warning(
                        f"FP16 provided in SM_HP_MP_PARAMETERS is {smp.state.cfg.fp16}, "
                        f"but FP16 provided in trainer argument is {args.fp16}, "
                        f"setting to {smp.state.cfg.fp16}"
                    )
                    args.fp16 = smp.state.cfg.fp16
            else:
                
                if hasattr(smp.state.cfg, "fp16"):
                    logger.warning(
                        f"FP16 provided in SM_HP_MP_PARAMETERS is {smp.state.cfg.fp16}, "
                        "but SageMaker Model Parallelism < 1.10 does not support FP16 in trainer."
                    )
        if (args.fp16 or args.bf16) and args.half_precision_backend == "auto":
            if args.device == torch.device("cpu"):
                if args.fp16:
                    raise ValueError("Tried to use `fp16` but it is not supported on cpu")
                else:
                    args.half_precision_backend = "cpu_amp"
            logger.info(f"Using {args.half_precision_backend} half precision backend")

        if (args.fp16 or args.bf16) and not (self.is_deepspeed_enabled or is_sagemaker_mp_enabled()):
            
            if args.half_precision_backend == "cpu_amp":
                self.use_cpu_amp = True
                self.amp_dtype = torch.bfloat16
            elif args.half_precision_backend == "apex":
                if not is_apex_available():
                    raise ImportError(
                        "Using FP16 with APEX but APEX is not installed, please refer to"
                        " https://www.github.com/nvidia/apex."
                    )
                self.use_apex = True

        
        if self.args.label_smoothing_factor != 0:
            self.label_smoother = LabelSmoother(epsilon=self.args.label_smoothing_factor)
        else:
            self.label_smoother = None

        self.state = TrainerState(
            is_local_process_zero=self.is_local_process_zero(),
            is_world_process_zero=self.is_world_process_zero(),
        )

        self.control = TrainerControl()
        
        
        self.current_flos = 0
        self.hp_search_backend = None
        self.use_tune_checkpoints = False
        default_label_names = find_labels(self.model.__class__)
        self.label_names = default_label_names if self.args.label_names is None else self.args.label_names
        self.can_return_loss = can_return_loss(self.model.__class__)
        self.control = self.callback_handler.on_init_end(self.args, self.state, self.control)

        
        self._train_batch_size = args.train_batch_size
        self._created_lr_scheduler = False

        
        self._memory_tracker.stop_and_update_metrics()

        
        if args.torch_compile and not is_torch_compile_available():
            raise RuntimeError("Using torch.compile requires PyTorch 2.0 or higher.")

    def _activate_neftune(self, model):
        r
        unwrapped_model = unwrap_model(model)

        if is_peft_available() and isinstance(unwrapped_model, PeftModel):
            embeddings = unwrapped_model.base_model.model.get_input_embeddings()
        else:
            embeddings = unwrapped_model.get_input_embeddings()

        del unwrapped_model

        embeddings.neftune_noise_alpha = self.neftune_noise_alpha
        hook_handle = embeddings.register_forward_hook(neftune_post_forward_hook)
        self.neftune_hook_handle = hook_handle
        return model

    def _deactivate_neftune(self, model):
        
        if not hasattr(self, "neftune_hook_handle"):
            raise ValueError("Neftune is not activated make sure to call `trainer._activate_neftune()` first")

        unwrapped_model = unwrap_model(model)

        if is_peft_available() and isinstance(unwrapped_model, PeftModel):
            embeddings = unwrapped_model.base_model.model.get_input_embeddings()
        else:
            embeddings = unwrapped_model.get_input_embeddings()

        self.neftune_hook_handle.remove()
        del embeddings.neftune_noise_alpha, unwrapped_model

    def add_callback(self, callback):
        
        self.callback_handler.add_callback(callback)

    def pop_callback(self, callback):
        
        return self.callback_handler.pop_callback(callback)

    def remove_callback(self, callback):
        
        self.callback_handler.remove_callback(callback)

    def _move_model_to_device(self, model, device):
        model = model.to(device)
        
        if self.args.parallel_mode == ParallelMode.TPU and hasattr(model, "tie_weights"):
            model.tie_weights()

    def _set_signature_columns_if_needed(self):
        if self._signature_columns is None:
            
            signature = inspect.signature(self.model.forward)
            self._signature_columns = list(signature.parameters.keys())
            
            self._signature_columns += list(set(["label", "label_ids"] + self.label_names))

    def _remove_unused_columns(self, dataset: "datasets.Dataset", description: Optional[str] = None):
        if not self.args.remove_unused_columns:
            return dataset
        self._set_signature_columns_if_needed()
        signature_columns = self._signature_columns

        ignored_columns = list(set(dataset.column_names) - set(signature_columns))
        if len(ignored_columns) > 0:
            dset_description = "" if description is None else f"in the {description} set"
            logger.info(
                f"The following columns {dset_description} don't have a corresponding argument in "
                f"`{self.model.__class__.__name__}.forward` and have been ignored: {', '.join(ignored_columns)}."
                f" If {', '.join(ignored_columns)} are not expected by `{self.model.__class__.__name__}.forward`, "
                " you can safely ignore this message."
            )

        columns = [k for k in signature_columns if k in dataset.column_names]

        if version.parse(datasets.__version__) < version.parse("1.4.0"):
            dataset.set_format(
                type=dataset.format["type"], columns=columns, format_kwargs=dataset.format["format_kwargs"]
            )
            return dataset
        else:
            return dataset.remove_columns(ignored_columns)

    def _get_collator_with_removed_columns(
        self, data_collator: Callable, description: Optional[str] = None
    ) -> Callable:
        
        if not self.args.remove_unused_columns:
            return data_collator
        self._set_signature_columns_if_needed()
        signature_columns = self._signature_columns

        remove_columns_collator = RemoveColumnsCollator(
            data_collator=data_collator,
            signature_columns=signature_columns,
            logger=logger,
            description=description,
            model_name=self.model.__class__.__name__,
        )
        return remove_columns_collator

    def _get_train_sampler(self) -> Optional[torch.utils.data.Sampler]:
        if self.train_dataset is None or not has_length(self.train_dataset):
            return None

        
        if self.args.group_by_length:
            if is_datasets_available() and isinstance(self.train_dataset, datasets.Dataset):
                lengths = (
                    self.train_dataset[self.args.length_column_name]
                    if self.args.length_column_name in self.train_dataset.column_names
                    else None
                )
            else:
                lengths = None
            model_input_name = self.tokenizer.model_input_names[0] if self.tokenizer is not None else None
            return LengthGroupedSampler(
                self.args.train_batch_size * self.args.gradient_accumulation_steps,
                dataset=self.train_dataset,
                lengths=lengths,
                model_input_name=model_input_name,
            )

        else:
            return RandomSampler(self.train_dataset)

    def get_train_dataloader(self) -> DataLoader:
        
        if self.train_dataset is None:
            raise ValueError("Trainer: training requires a train_dataset.")

        train_dataset = self.train_dataset
        data_collator = self.data_collator
        if is_datasets_available() and isinstance(train_dataset, datasets.Dataset):
            train_dataset = self._remove_unused_columns(train_dataset, description="training")
        else:
            data_collator = self._get_collator_with_removed_columns(data_collator, description="training")

        dataloader_params = {
            "batch_size": self._train_batch_size,
            "collate_fn": data_collator,
            "num_workers": self.args.dataloader_num_workers,
            "pin_memory": self.args.dataloader_pin_memory,
            "persistent_workers": self.args.dataloader_persistent_workers,
        }

        if not isinstance(train_dataset, torch.utils.data.IterableDataset):
            dataloader_params["sampler"] = self._get_train_sampler()
            dataloader_params["drop_last"] = self.args.dataloader_drop_last
            dataloader_params["worker_init_fn"] = seed_worker

        return self.accelerator.prepare(DataLoader(train_dataset, **dataloader_params))

    def _get_eval_sampler(self, eval_dataset: Dataset) -> Optional[torch.utils.data.Sampler]:
        
        if self.args.use_legacy_prediction_loop:
            if is_torch_tpu_available():
                return SequentialDistributedSampler(
                    eval_dataset, num_replicas=xm.xrt_world_size(), rank=xm.get_ordinal()
                )
            elif is_sagemaker_mp_enabled():
                return SequentialDistributedSampler(
                    eval_dataset,
                    num_replicas=smp.dp_size(),
                    rank=smp.dp_rank(),
                    batch_size=self.args.per_device_eval_batch_size,
                )
            else:
                return SequentialSampler(eval_dataset)

        if self.args.world_size <= 1:
            return SequentialSampler(eval_dataset)
        else:
            return None

    def get_eval_dataloader(self, eval_dataset: Optional[Dataset] = None) -> DataLoader:
        
        if eval_dataset is None and self.eval_dataset is None:
            raise ValueError("Trainer: evaluation requires an eval_dataset.")
        eval_dataset = eval_dataset if eval_dataset is not None else self.eval_dataset
        data_collator = self.data_collator

        if is_datasets_available() and isinstance(eval_dataset, datasets.Dataset):
            eval_dataset = self._remove_unused_columns(eval_dataset, description="evaluation")
        else:
            data_collator = self._get_collator_with_removed_columns(data_collator, description="evaluation")

        dataloader_params = {
            "batch_size": self.args.eval_batch_size,
            "collate_fn": data_collator,
            "num_workers": self.args.dataloader_num_workers,
            "pin_memory": self.args.dataloader_pin_memory,
            "persistent_workers": self.args.dataloader_persistent_workers,
        }

        if not isinstance(eval_dataset, torch.utils.data.IterableDataset):
            dataloader_params["sampler"] = self._get_eval_sampler(eval_dataset)
            dataloader_params["drop_last"] = self.args.dataloader_drop_last

        return self.accelerator.prepare(DataLoader(eval_dataset, **dataloader_params))

    def get_test_dataloader(self, test_dataset: Dataset) -> DataLoader:
        
        data_collator = self.data_collator

        if is_datasets_available() and isinstance(test_dataset, datasets.Dataset):
            test_dataset = self._remove_unused_columns(test_dataset, description="test")
        else:
            data_collator = self._get_collator_with_removed_columns(data_collator, description="test")

        dataloader_params = {
            "batch_size": self.args.eval_batch_size,
            "collate_fn": data_collator,
            "num_workers": self.args.dataloader_num_workers,
            "pin_memory": self.args.dataloader_pin_memory,
            "persistent_workers": self.args.dataloader_persistent_workers,
        }

        if not isinstance(test_dataset, torch.utils.data.IterableDataset):
            dataloader_params["sampler"] = self._get_eval_sampler(test_dataset)
            dataloader_params["drop_last"] = self.args.dataloader_drop_last

        
        return self.accelerator.prepare(DataLoader(test_dataset, **dataloader_params))

    def create_optimizer_and_scheduler(self, num_training_steps: int):
        
        self.create_optimizer()
        if IS_SAGEMAKER_MP_POST_1_10 and smp.state.cfg.fp16:
            
            optimizer = self.optimizer.optimizer
        else:
            optimizer = self.optimizer
        self.create_scheduler(num_training_steps=num_training_steps, optimizer=optimizer)

    def get_decay_parameter_names(self, model) -> List[str]:
        
        decay_parameters = get_parameter_names(model, ALL_LAYERNORM_LAYERS)
        decay_parameters = [name for name in decay_parameters if "bias" not in name]
        return decay_parameters

    def create_optimizer(self):
        
        opt_model = self.model_wrapped if is_sagemaker_mp_enabled() else self.model

        if self.optimizer is None:
            decay_parameters = self.get_decay_parameter_names(opt_model)
            optimizer_grouped_parameters = [
                {
                    "params": [
                        p for n, p in opt_model.named_parameters() if (n in decay_parameters and p.requires_grad)
                    ],
                    "weight_decay": self.args.weight_decay,
                },
                {
                    "params": [
                        p for n, p in opt_model.named_parameters() if (n not in decay_parameters and p.requires_grad)
                    ],
                    "weight_decay": 0.0,
                },
            ]

            optimizer_cls, optimizer_kwargs = Trainer.get_optimizer_cls_and_kwargs(self.args)

            self.optimizer = optimizer_cls(optimizer_grouped_parameters, **optimizer_kwargs)
            if optimizer_cls.__name__ == "Adam8bit":
                import bitsandbytes

                manager = bitsandbytes.optim.GlobalOptimManager.get_instance()

                skipped = 0
                for module in opt_model.modules():
                    if isinstance(module, nn.Embedding):
                        skipped += sum({p.data_ptr(): p.numel() for p in module.parameters()}.values())
                        logger.info(f"skipped {module}: {skipped/2**20}M params")
                        manager.register_module_override(module, "weight", {"optim_bits": 32})
                        logger.debug(f"bitsandbytes: will optimize {module} in fp32")
                logger.info(f"skipped: {skipped/2**20}M params")

        if is_sagemaker_mp_enabled():
            self.optimizer = smp.DistributedOptimizer(self.optimizer)

        return self.optimizer

    @staticmethod
    def get_optimizer_cls_and_kwargs(args: TrainingArguments) -> Tuple[Any, Any]:
        

        
        optim_args = {}
        if args.optim_args:
            for mapping in args.optim_args.replace(" ", "").split(","):
                key, value = mapping.split("=")
                optim_args[key] = value

        optimizer_kwargs = {"lr": args.learning_rate}

        adam_kwargs = {
            "betas": (args.adam_beta1, args.adam_beta2),
            "eps": args.adam_epsilon,
        }
        if args.optim == OptimizerNames.ADAFACTOR:
            optimizer_cls = Adafactor
            optimizer_kwargs.update({"scale_parameter": False, "relative_step": False})
        elif args.optim == OptimizerNames.ADAMW_HF:
            from .optimization import AdamW

            optimizer_cls = AdamW
            optimizer_kwargs.update(adam_kwargs)
        elif args.optim in [OptimizerNames.ADAMW_TORCH, OptimizerNames.ADAMW_TORCH_FUSED]:
            from torch.optim import AdamW

            optimizer_cls = AdamW
            optimizer_kwargs.update(adam_kwargs)
            if args.optim == OptimizerNames.ADAMW_TORCH_FUSED:
                optimizer_kwargs.update({"fused": True})
        elif args.optim == OptimizerNames.ADAMW_TORCH_XLA:
            try:
                from torch_xla.amp.syncfree import AdamW

                optimizer_cls = AdamW
                optimizer_kwargs.update(adam_kwargs)
            except ImportError:
                raise ValueError("Trainer failed to import syncfree AdamW from torch_xla.")
        elif args.optim == OptimizerNames.ADAMW_TORCH_NPU_FUSED:
            try:
                from torch_npu.optim import NpuFusedAdamW

                optimizer_cls = NpuFusedAdamW
                optimizer_kwargs.update(adam_kwargs)
            except ImportError:
                raise ValueError("Trainer failed to import FusedAdamW from torch_npu.")
        elif args.optim == OptimizerNames.ADAMW_APEX_FUSED:
            try:
                from apex.optimizers import FusedAdam

                optimizer_cls = FusedAdam
                optimizer_kwargs.update(adam_kwargs)
            except ImportError:
                raise ValueError("Trainer tried to instantiate apex FusedAdam but apex is not installed!")
        elif args.optim in [
            OptimizerNames.ADAMW_BNB,
            OptimizerNames.ADAMW_8BIT,
            OptimizerNames.PAGED_ADAMW,
            OptimizerNames.PAGED_ADAMW_8BIT,
            OptimizerNames.LION,
            OptimizerNames.LION_8BIT,
            OptimizerNames.PAGED_LION,
            OptimizerNames.PAGED_LION_8BIT,
        ]:
            try:
                from bitsandbytes.optim import AdamW, Lion

                is_paged = False
                optim_bits = 32
                optimizer_cls = None
                additional_optim_kwargs = adam_kwargs
                if "paged" in args.optim:
                    is_paged = True
                if "8bit" in args.optim:
                    optim_bits = 8
                if "adam" in args.optim:
                    optimizer_cls = AdamW
                elif "lion" in args.optim:
                    optimizer_cls = Lion
                    additional_optim_kwargs = {"betas": (args.adam_beta1, args.adam_beta2)}

                bnb_kwargs = {"is_paged": is_paged, "optim_bits": optim_bits}
                optimizer_kwargs.update(additional_optim_kwargs)
                optimizer_kwargs.update(bnb_kwargs)
            except ImportError:
                raise ValueError("Trainer tried to instantiate bnb optimizer but bnb is not installed!")
            if is_bitsandbytes_available() and version.parse(
                importlib.metadata.version("bitsandbytes")
            ) < version.parse("0.41.1"):
                logger.warning(
                    "You are using 8-bit optimizers with a version of `bitsandbytes` < 0.41.1. "
                    "It is recommended to update your version as a major bug has been fixed in 8-bit optimizers."
                )
        elif args.optim == OptimizerNames.ADAMW_ANYPRECISION:
            try:
                from torchdistx.optimizers import AnyPrecisionAdamW

                optimizer_cls = AnyPrecisionAdamW
                optimizer_kwargs.update(adam_kwargs)

                
                optimizer_kwargs.update(
                    {
                        "use_kahan_summation": strtobool(optim_args.get("use_kahan_summation", "False")),
                        "momentum_dtype": getattr(torch, optim_args.get("momentum_dtype", "float32")),
                        "variance_dtype": getattr(torch, optim_args.get("variance_dtype", "float32")),
                        "compensation_buffer_dtype": getattr(
                            torch, optim_args.get("compensation_buffer_dtype", "bfloat16")
                        ),
                    }
                )
            except ImportError:
                raise ValueError("Please install https://github.com/pytorch/torchdistx")
        elif args.optim == OptimizerNames.SGD:
            optimizer_cls = torch.optim.SGD
        elif args.optim == OptimizerNames.ADAGRAD:
            optimizer_cls = torch.optim.Adagrad
        elif args.optim == OptimizerNames.RMSPROP:
            optimizer_cls = torch.optim.RMSprop
        else:
            raise ValueError(f"Trainer cannot instantiate unsupported optimizer: {args.optim}")
        return optimizer_cls, optimizer_kwargs

    def create_scheduler(self, num_training_steps: int, optimizer: torch.optim.Optimizer = None):
        
        if self.lr_scheduler is None:
            self.lr_scheduler = get_scheduler(
                self.args.lr_scheduler_type,
                optimizer=self.optimizer if optimizer is None else optimizer,
                num_warmup_steps=self.args.get_warmup_steps(num_training_steps),
                num_training_steps=num_training_steps,
                scheduler_specific_kwargs=self.args.lr_scheduler_kwargs,
            )
            self._created_lr_scheduler = True
        return self.lr_scheduler

    def num_examples(self, dataloader: DataLoader) -> int:
        
        try:
            dataset = dataloader.dataset
            
            if isinstance(dataset, IterableDatasetShard):
                return len(dataloader.dataset.dataset)
            return len(dataloader.dataset)
        except (NameError, AttributeError, TypeError):  
            return len(dataloader) * self.args.per_device_train_batch_size

    def num_tokens(self, train_dl: DataLoader, max_steps: Optional[int] = None) -> int:
        
        train_tokens = 0
        try:
            for step, batch in enumerate(train_dl):
                tokens = batch["input_ids"].numel()
                if max_steps is not None:
                    return tokens * max_steps
                train_tokens += tokens
            return train_tokens
        except KeyError:
            logger.warning("Cannot get num_tokens from dataloader")
            return train_tokens

    def _hp_search_setup(self, trial: Union["optuna.Trial", Dict[str, Any]]):
        
        self._trial = trial

        if self.hp_search_backend is None or trial is None:
            return
        if self.hp_search_backend == HPSearchBackend.OPTUNA:
            params = self.hp_space(trial)
        elif self.hp_search_backend == HPSearchBackend.RAY:
            params = trial
            params.pop("wandb", None)
        elif self.hp_search_backend == HPSearchBackend.SIGOPT:
            params = {k: int(v) if isinstance(v, str) else v for k, v in trial.assignments.items()}
        elif self.hp_search_backend == HPSearchBackend.WANDB:
            params = trial

        for key, value in params.items():
            if not hasattr(self.args, key):
                logger.warning(
                    f"Trying to set {key} in the hyperparameter search but there is no corresponding field in"
                    " `TrainingArguments`."
                )
                continue
            old_attr = getattr(self.args, key, None)
            
            if old_attr is not None:
                value = type(old_attr)(value)

            setattr(self.args, key, value)
        if self.hp_search_backend == HPSearchBackend.OPTUNA:
            logger.info(f"Trial: {trial.params}")
        if self.hp_search_backend == HPSearchBackend.SIGOPT:
            logger.info(f"SigOpt Assignments: {trial.assignments}")
        if self.hp_search_backend == HPSearchBackend.WANDB:
            logger.info(f"W&B Sweep parameters: {trial}")
        if self.is_deepspeed_enabled:
            if self.args.deepspeed is None:
                raise ValueError("For sweeps with deepspeed, `args.deepspeed` must be set")
            
            from accelerate.utils import DeepSpeedPlugin

            from transformers.integrations.deepspeed import HfTrainerDeepSpeedConfig

            self.args.hf_deepspeed_config = HfTrainerDeepSpeedConfig(self.args.deepspeed)
            self.args.hf_deepspeed_config.trainer_config_process(self.args)
            self.args.deepspeed_plugin = DeepSpeedPlugin(hf_ds_config=self.args.hf_deepspeed_config)

        self.create_accelerator_and_postprocess()

    def _report_to_hp_search(self, trial: Union["optuna.Trial", Dict[str, Any]], step: int, metrics: Dict[str, float]):
        if self.hp_search_backend is None or trial is None:
            return
        self.objective = self.compute_objective(metrics.copy())
        if self.hp_search_backend == HPSearchBackend.OPTUNA:
            import optuna

            if not trial.study._is_multi_objective():
                trial.report(self.objective, step)
                if trial.should_prune():
                    self.callback_handler.on_train_end(self.args, self.state, self.control)
                    raise optuna.TrialPruned()
        elif self.hp_search_backend == HPSearchBackend.RAY:
            from ray import tune

            if self.control.should_save:
                self._tune_save_checkpoint()
            tune.report(objective=self.objective, **metrics)

    def _tune_save_checkpoint(self):
        from ray import tune

        if not self.use_tune_checkpoints:
            return
        with tune.checkpoint_dir(step=self.state.global_step) as checkpoint_dir:
            output_dir = os.path.join(checkpoint_dir, f"{PREFIX_CHECKPOINT_DIR}-{self.state.global_step}")
            self.save_model(output_dir, _internal_call=True)
            if self.args.should_save:
                self.state.save_to_json(os.path.join(output_dir, TRAINER_STATE_NAME))
                torch.save(self.optimizer.state_dict(), os.path.join(output_dir, OPTIMIZER_NAME))
                torch.save(self.lr_scheduler.state_dict(), os.path.join(output_dir, SCHEDULER_NAME))

    def call_model_init(self, trial=None):
        model_init_argcount = number_of_arguments(self.model_init)
        if model_init_argcount == 0:
            model = self.model_init()
        elif model_init_argcount == 1:
            model = self.model_init(trial)
        else:
            raise RuntimeError("model_init should have 0 or 1 argument.")

        if model is None:
            raise RuntimeError("model_init should not return None.")

        return model

    def torch_jit_model_eval(self, model, dataloader, training=False):
        if not training:
            if dataloader is None:
                logger.warning("failed to use PyTorch jit mode due to current dataloader is none.")
                return model
            example_batch = next(iter(dataloader))
            example_batch = self._prepare_inputs(example_batch)
            try:
                jit_model = copy.copy(model)
                jit_model.eval()
                original_forward = jit_model.__dict__.pop("_original_forward", None)
                
                if original_forward:
                    jit_model.forward = original_forward
                with self.accelerator.autocast(cache_enabled=False), torch.no_grad():
                    if version.parse(version.parse(torch.__version__).base_version) >= version.parse("2.0.0"):
                        if isinstance(example_batch, dict):
                            jit_model = torch.jit.trace(jit_model, example_kwarg_inputs=example_batch, strict=False)
                        else:
                            jit_model = torch.jit.trace(
                                jit_model,
                                example_kwarg_inputs={key: example_batch[key] for key in example_batch},
                                strict=False,
                            )
                    else:
                        jit_inputs = []
                        for key in example_batch:
                            example_tensor = torch.ones_like(example_batch[key])
                            jit_inputs.append(example_tensor)
                        jit_inputs = tuple(jit_inputs)
                        jit_model = torch.jit.trace(jit_model, jit_inputs, strict=False)
                jit_model = torch.jit.freeze(jit_model)
                with torch.no_grad():
                    jit_model(**example_batch)
                    jit_model(**example_batch)
                model = jit_model
                self.use_cpu_amp = False
            except (RuntimeError, TypeError, ValueError, NameError, IndexError) as e:
                logger.warning(f"failed to use PyTorch jit mode due to: {e}.")

        return model

    def ipex_optimize_model(self, model, training=False, dtype=torch.float32):
        if not is_ipex_available():
            raise ImportError(
                "Using IPEX but IPEX is not installed or IPEX's version does not match current PyTorch, please refer"
                " to https://github.com/intel/intel-extension-for-pytorch."
            )

        import intel_extension_for_pytorch as ipex

        if not training:
            model.eval()
            dtype = torch.bfloat16 if not self.is_in_train and self.args.bf16_full_eval else dtype
            
            model = ipex.optimize(model, dtype=dtype, level="O1", conv_bn_folding=False, inplace=not self.is_in_train)
        else:
            if not model.training:
                model.train()
            model, self.optimizer = ipex.optimize(
                model, dtype=dtype, optimizer=self.optimizer, inplace=True, level="O1"
            )

        return model

    def _wrap_model(self, model, training=True, dataloader=None):
        if self.args.use_ipex:
            dtype = torch.bfloat16 if self.use_cpu_amp else torch.float32
            model = self.ipex_optimize_model(model, training, dtype=dtype)

        if is_sagemaker_mp_enabled():
            
            if isinstance(self.model_wrapped, smp.model.DistributedModel):
                return self.model_wrapped
            return smp.DistributedModel(model, backward_passes_per_step=self.args.gradient_accumulation_steps)

        
        if unwrap_model(model) is not model:
            return model

        
        if self.use_apex and training:
            model, self.optimizer = amp.initialize(model, self.optimizer, opt_level=self.args.fp16_opt_level)

        
        if self.args.n_gpu > 1 and not getattr(model, "is_loaded_in_8bit", False):
            model = nn.DataParallel(model)

        if self.args.jit_mode_eval:
            start_time = time.time()
            model = self.torch_jit_model_eval(model, dataloader, training)
            self.jit_compilation_time = round(time.time() - start_time, 4)

        
        
        if not training:
            return model

        
        
        if self.is_fsdp_xla_enabled:
            try:
                from torch_xla.distributed.fsdp import XlaFullyShardedDataParallel as FSDP
                from torch_xla.distributed.fsdp import checkpoint_module
                from torch_xla.distributed.fsdp.wrap import (
                    size_based_auto_wrap_policy,
                    transformer_auto_wrap_policy,
                )
            except ImportError:
                raise ImportError("Missing XLA FSDP related module; please make sure to use torch-xla >= 2.0.")
            auto_wrap_policy = None
            auto_wrapper_callable = None
            default_transformer_cls_names_to_wrap = getattr(model, "_no_split_modules", None)
            fsdp_transformer_layer_cls_to_wrap = self.args.fsdp_config.get(
                "transformer_layer_cls_to_wrap", default_transformer_cls_names_to_wrap
            )

            if self.args.fsdp_config["min_num_params"] > 0:
                auto_wrap_policy = functools.partial(
                    size_based_auto_wrap_policy, min_num_params=self.args.fsdp_config["min_num_params"]
                )
            elif fsdp_transformer_layer_cls_to_wrap is not None:
                transformer_cls_to_wrap = set()
                for layer_class in fsdp_transformer_layer_cls_to_wrap:
                    transformer_cls = get_module_class_from_name(model, layer_class)
                    if transformer_cls is None:
                        raise Exception("Could not find the transformer layer class to wrap in the model.")
                    else:
                        transformer_cls_to_wrap.add(transformer_cls)

                auto_wrap_policy = functools.partial(
                    transformer_auto_wrap_policy,
                    
                    transformer_layer_cls=transformer_cls_to_wrap,
                )
            fsdp_kwargs = self.args.xla_fsdp_config
            if self.args.fsdp_config["xla_fsdp_grad_ckpt"]:
                
                def auto_wrapper_callable(m, *args, **kwargs):
                    return FSDP(checkpoint_module(m), *args, **kwargs)

            
            self.model = model = FSDP(
                model,
                auto_wrap_policy=auto_wrap_policy,
                auto_wrapper_callable=auto_wrapper_callable,
                **fsdp_kwargs,
            )

            
            
            def patched_optimizer_step(optimizer, barrier=False, optimizer_args={}):
                loss = optimizer.step(**optimizer_args)
                if barrier:
                    xm.mark_step()
                return loss

            xm.optimizer_step = patched_optimizer_step
        elif is_sagemaker_dp_enabled():
            model = nn.parallel.DistributedDataParallel(
                model, device_ids=[int(os.getenv("SMDATAPARALLEL_LOCAL_RANK"))]
            )
        elif self.args.parallel_mode == ParallelMode.DISTRIBUTED:
            if is_torch_neuroncore_available():
                return model
            kwargs = {}
            if self.args.ddp_find_unused_parameters is not None:
                kwargs["find_unused_parameters"] = self.args.ddp_find_unused_parameters
            elif isinstance(model, PreTrainedModel):
                
                
                kwargs["find_unused_parameters"] = not model.is_gradient_checkpointing
            else:
                kwargs["find_unused_parameters"] = True

            if self.args.ddp_bucket_cap_mb is not None:
                kwargs["bucket_cap_mb"] = self.args.ddp_bucket_cap_mb

            if self.args.ddp_broadcast_buffers is not None:
                kwargs["broadcast_buffers"] = self.args.ddp_broadcast_buffers

            self.accelerator.ddp_handler = DistributedDataParallelKwargs(**kwargs)

        return model

    def train(
        self,
        resume_from_checkpoint: Optional[Union[str, bool]] = None,
        trial: Union["optuna.Trial", Dict[str, Any]] = None,
        ignore_keys_for_eval: Optional[List[str]] = None,
        **kwargs,
    ):
        
        if resume_from_checkpoint is False:
            resume_from_checkpoint = None

        
        self._memory_tracker.start()

        args = self.args

        self.is_in_train = True

        
        if self.neftune_noise_alpha is not None:
            self.model = self._activate_neftune(self.model)

        
        
        if (args.fp16_full_eval or args.bf16_full_eval) and not args.do_train:
            self._move_model_to_device(self.model, args.device)

        if "model_path" in kwargs:
            resume_from_checkpoint = kwargs.pop("model_path")
            warnings.warn(
                "`model_path` is deprecated and will be removed in a future version. Use `resume_from_checkpoint` "
                "instead.",
                FutureWarning,
            )
        if len(kwargs) > 0:
            raise TypeError(f"train() received got unexpected keyword arguments: {', '.join(list(kwargs.keys()))}.")
        
        self._hp_search_setup(trial)
        self._train_batch_size = self.args.train_batch_size

        
        model_reloaded = False
        if self.model_init is not None:
            
            enable_full_determinism(self.args.seed) if self.args.full_determinism else set_seed(self.args.seed)
            self.model = self.call_model_init(trial)
            model_reloaded = True
            
            self.optimizer, self.lr_scheduler = None, None

        
        if isinstance(resume_from_checkpoint, bool) and resume_from_checkpoint:
            resume_from_checkpoint = get_last_checkpoint(args.output_dir)
            if resume_from_checkpoint is None:
                raise ValueError(f"No valid checkpoint found in output directory ({args.output_dir})")

        if (
            resume_from_checkpoint is not None
            and not is_sagemaker_mp_enabled()
            and not self.is_deepspeed_enabled
            and not self.is_fsdp_enabled
        ):
            self._load_from_checkpoint(resume_from_checkpoint)

        
        if model_reloaded:
            if self.place_model_on_device:
                self._move_model_to_device(self.model, args.device)
            self.model_wrapped = self.model

        inner_training_loop = find_executable_batch_size(
            self._inner_training_loop, self._train_batch_size, args.auto_find_batch_size
        )
        if args.push_to_hub:
            try:
                
                hf_hub_utils.disable_progress_bars()
                return inner_training_loop(
                    args=args,
                    resume_from_checkpoint=resume_from_checkpoint,
                    trial=trial,
                    ignore_keys_for_eval=ignore_keys_for_eval,
                )
            finally:
                hf_hub_utils.enable_progress_bars()
        else:
            return inner_training_loop(
                args=args,
                resume_from_checkpoint=resume_from_checkpoint,
                trial=trial,
                ignore_keys_for_eval=ignore_keys_for_eval,
            )

    def _inner_training_loop(
        self, batch_size=None, args=None, resume_from_checkpoint=None, trial=None, ignore_keys_for_eval=None
    ):
        self.accelerator.free_memory()
        self._train_batch_size = batch_size
        logger.debug(f"Currently training with a batch size of: {self._train_batch_size}")
        
        train_dataloader = self.get_train_dataloader()

        
        
        
        
        total_train_batch_size = self._train_batch_size * args.gradient_accumulation_steps * args.world_size

        len_dataloader = None
        num_train_tokens = None
        if has_length(train_dataloader):
            len_dataloader = len(train_dataloader)
            num_update_steps_per_epoch = len_dataloader // args.gradient_accumulation_steps
            num_update_steps_per_epoch = max(num_update_steps_per_epoch, 1)
            num_examples = self.num_examples(train_dataloader)
            if args.max_steps > 0:
                max_steps = args.max_steps
                num_train_epochs = args.max_steps // num_update_steps_per_epoch + int(
                    args.max_steps % num_update_steps_per_epoch > 0
                )
                
                
                num_train_samples = args.max_steps * total_train_batch_size
                if args.include_tokens_per_second:
                    num_train_tokens = (
                        self.num_tokens(train_dataloader, args.max_steps) * args.gradient_accumulation_steps
                    )
            else:
                max_steps = math.ceil(args.num_train_epochs * num_update_steps_per_epoch)
                num_train_epochs = math.ceil(args.num_train_epochs)
                num_train_samples = self.num_examples(train_dataloader) * args.num_train_epochs
                if args.include_tokens_per_second:
                    num_train_tokens = self.num_tokens(train_dataloader) * args.num_train_epochs
        elif args.max_steps > 0:  
            max_steps = args.max_steps
            
            num_train_epochs = sys.maxsize
            num_update_steps_per_epoch = max_steps
            num_examples = total_train_batch_size * args.max_steps
            num_train_samples = args.max_steps * total_train_batch_size
            if args.include_tokens_per_second:
                num_train_tokens = self.num_tokens(train_dataloader, args.max_steps) * args.gradient_accumulation_steps
        else:
            raise ValueError(
                "args.max_steps must be set to a positive value if dataloader does not have a length, was"
                f" {args.max_steps}"
            )

        if DebugOption.UNDERFLOW_OVERFLOW in self.args.debug:
            if self.args.n_gpu > 1:
                
                
                raise ValueError(
                    "Currently --debug underflow_overflow is not supported under DP. Please use DDP"
                    " (torchrun or torch.distributed.launch (deprecated))."
                )
            else:
                debug_overflow = DebugUnderflowOverflow(self.model)  

        delay_optimizer_creation = is_sagemaker_mp_enabled() or self.is_fsdp_xla_enabled or self.is_fsdp_enabled

        
        if self._created_lr_scheduler:
            self.lr_scheduler = None
            self._created_lr_scheduler = False

        if self.is_deepspeed_enabled:
            self.optimizer, self.lr_scheduler = deepspeed_init(self, num_training_steps=max_steps)

        if not delay_optimizer_creation:
            self.create_optimizer_and_scheduler(num_training_steps=max_steps)

        self.state = TrainerState()
        self.state.is_hyper_param_search = trial is not None

        
        if args.logging_steps is not None:
            if args.logging_steps < 1:
                self.state.logging_steps = math.ceil(max_steps * args.logging_steps)
            else:
                self.state.logging_steps = args.logging_steps
        if args.eval_steps is not None:
            if args.eval_steps < 1:
                self.state.eval_steps = math.ceil(max_steps * args.eval_steps)
            else:
                self.state.eval_steps = args.eval_steps
        if args.save_steps is not None:
            if args.save_steps < 1:
                self.state.save_steps = math.ceil(max_steps * args.save_steps)
            else:
                self.state.save_steps = args.save_steps

        
        if args.gradient_checkpointing:
            if args.gradient_checkpointing_kwargs is None:
                gradient_checkpointing_kwargs = {}
            else:
                gradient_checkpointing_kwargs = args.gradient_checkpointing_kwargs

            self.model.gradient_checkpointing_enable(gradient_checkpointing_kwargs=gradient_checkpointing_kwargs)

        model = self._wrap_model(self.model_wrapped)

        
        
        
        use_accelerator_prepare = True if model is self.model else False

        if delay_optimizer_creation:
            self.create_optimizer_and_scheduler(num_training_steps=max_steps)

        
        if use_accelerator_prepare:
            self.model.train()
            if hasattr(self.lr_scheduler, "step"):
                if self.use_apex:
                    model = self.accelerator.prepare(self.model)
                else:
                    model, self.optimizer = self.accelerator.prepare(self.model, self.optimizer)
            else:
                
                model, self.optimizer, self.lr_scheduler = self.accelerator.prepare(
                    self.model, self.optimizer, self.lr_scheduler
                )

        if self.is_fsdp_enabled:
            self.model = self.model_wrapped = model

        
        if model is not self.model:
            self.model_wrapped = model

        
        if self.is_deepspeed_enabled:
            self.deepspeed = self.model_wrapped

        
        if resume_from_checkpoint is not None:
            if self.is_deepspeed_enabled:
                deepspeed_load_checkpoint(self.model_wrapped, resume_from_checkpoint)
            elif is_sagemaker_mp_enabled() or self.is_fsdp_enabled:
                self._load_from_checkpoint(resume_from_checkpoint, self.model_wrapped)

        
        self._load_optimizer_and_scheduler(resume_from_checkpoint)

        
        
        
        

        
        logger.info("***** Running training *****")
        logger.info(f"  Num examples = {num_examples:,}")
        logger.info(f"  Num Epochs = {num_train_epochs:,}")
        logger.info(f"  Instantaneous batch size per device = {self.args.per_device_train_batch_size:,}")
        if self.args.per_device_train_batch_size != self._train_batch_size:
            logger.info(f"  Training with DataParallel so batch size has been adjusted to: {self._train_batch_size:,}")
        logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_train_batch_size:,}")
        logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
        logger.info(f"  Total optimization steps = {max_steps:,}")
        logger.info(f"  Number of trainable parameters = {get_model_param_count(model, trainable_only=True):,}")

        self.state.epoch = 0
        start_time = time.time()
        epochs_trained = 0
        steps_trained_in_current_epoch = 0
        steps_trained_progress_bar = None

        
        if resume_from_checkpoint is not None and os.path.isfile(
            os.path.join(resume_from_checkpoint, TRAINER_STATE_NAME)
        ):
            self.state = TrainerState.load_from_json(os.path.join(resume_from_checkpoint, TRAINER_STATE_NAME))
            epochs_trained = self.state.global_step // num_update_steps_per_epoch
            if not args.ignore_data_skip:
                steps_trained_in_current_epoch = self.state.global_step % (num_update_steps_per_epoch)
                steps_trained_in_current_epoch *= args.gradient_accumulation_steps
            else:
                steps_trained_in_current_epoch = 0

            logger.info("  Continuing training from checkpoint, will skip to saved global_step")
            logger.info(f"  Continuing training from epoch {epochs_trained}")
            logger.info(f"  Continuing training from global step {self.state.global_step}")
            if not args.ignore_data_skip:
                logger.info(
                    f"  Will skip the first {epochs_trained} epochs then the first"
                    f" {steps_trained_in_current_epoch} batches in the first epoch."
                )

        
        self.callback_handler.model = self.model
        self.callback_handler.optimizer = self.optimizer
        self.callback_handler.lr_scheduler = self.lr_scheduler
        self.callback_handler.train_dataloader = train_dataloader
        if self.hp_name is not None and self._trial is not None:
            
            
            self.state.trial_name = self.hp_name(self._trial)
        if trial is not None:
            assignments = trial.assignments if self.hp_search_backend == HPSearchBackend.SIGOPT else trial
            self.state.trial_params = hp_params(assignments)
        else:
            self.state.trial_params = None
        
        
        self.state.max_steps = max_steps
        self.state.num_train_epochs = num_train_epochs
        self.state.is_local_process_zero = self.is_local_process_zero()
        self.state.is_world_process_zero = self.is_world_process_zero()

        
        tr_loss = torch.tensor(0.0).to(args.device)
        
        self._total_loss_scalar = 0.0
        self._globalstep_last_logged = self.state.global_step
        model.zero_grad()

        self.control = self.callback_handler.on_train_begin(args, self.state, self.control)

        
        if not args.ignore_data_skip:
            for epoch in range(epochs_trained):
                sampler = get_dataloader_sampler(train_dataloader)
                sampler_kinds = [RandomSampler]
                if version.parse(accelerate_version) > version.parse("0.23.0"):
                    sampler_kinds.append(SeedableRandomSampler)
                is_random_sampler = isinstance(sampler, tuple(sampler_kinds))
                if is_torch_less_than_1_11 or not is_random_sampler:
                    
                    for _ in train_dataloader:
                        break
                else:
                    
                    
                    sampler = sampler if sampler is not None else []
                    _ = list(sampler)

        total_batched_samples = 0
        for epoch in range(epochs_trained, num_train_epochs):
            epoch_iterator = train_dataloader
            if hasattr(epoch_iterator, "set_epoch"):
                epoch_iterator.set_epoch(epoch)

            
            if args.past_index >= 0:
                self._past = None

            steps_in_epoch = (
                len(epoch_iterator)
                if len_dataloader is not None
                else args.max_steps * args.gradient_accumulation_steps
            )
            self.control = self.callback_handler.on_epoch_begin(args, self.state, self.control)

            if epoch == epochs_trained and resume_from_checkpoint is not None and steps_trained_in_current_epoch == 0:
                self._load_rng_state(resume_from_checkpoint)

            rng_to_sync = False
            steps_skipped = 0
            if steps_trained_in_current_epoch > 0:
                epoch_iterator = skip_first_batches(epoch_iterator, steps_trained_in_current_epoch)
                steps_skipped = steps_trained_in_current_epoch
                steps_trained_in_current_epoch = 0
                rng_to_sync = True

            step = -1
            for step, inputs in enumerate(epoch_iterator):
                total_batched_samples += 1

                if self.args.include_num_input_tokens_seen:
                    main_input_name = getattr(self.model, "main_input_name", "input_ids")
                    if main_input_name not in inputs:
                        logger.warning(
                            "Tried to track the number of tokens seen, however the current model is "
                            "not configured properly to know what item is the input. To fix this, add "
                            "a `main_input_name` attribute to the model class you are using."
                        )
                    else:
                        self.state.num_input_tokens_seen += self.accelerator.gather(inputs[main_input_name]).numel()
                if rng_to_sync:
                    self._load_rng_state(resume_from_checkpoint)
                    rng_to_sync = False

                
                if steps_trained_in_current_epoch > 0:
                    steps_trained_in_current_epoch -= 1
                    if steps_trained_progress_bar is not None:
                        steps_trained_progress_bar.update(1)
                    if steps_trained_in_current_epoch == 0:
                        self._load_rng_state(resume_from_checkpoint)
                    continue
                elif steps_trained_progress_bar is not None:
                    steps_trained_progress_bar.close()
                    steps_trained_progress_bar = None

                if step % args.gradient_accumulation_steps == 0:
                    self.control = self.callback_handler.on_step_begin(args, self.state, self.control)

                with self.accelerator.accumulate(model):
                    tr_loss_step = self.training_step(model, inputs)

                if (
                    args.logging_nan_inf_filter
                    and not is_torch_tpu_available()
                    and (torch.isnan(tr_loss_step) or torch.isinf(tr_loss_step))
                ):
                    
                    tr_loss += tr_loss / (1 + self.state.global_step - self._globalstep_last_logged)
                else:
                    tr_loss += tr_loss_step

                self.current_flos += float(self.floating_point_ops(inputs))

                is_last_step_and_steps_less_than_grad_acc = (
                    steps_in_epoch <= args.gradient_accumulation_steps and (step + 1) == steps_in_epoch
                )

                if (
                    total_batched_samples % args.gradient_accumulation_steps == 0
                    or
                    
                    is_last_step_and_steps_less_than_grad_acc
                ):
                    
                    
                    if is_last_step_and_steps_less_than_grad_acc:
                        self.accelerator.gradient_state._set_sync_gradients(True)

                    
                    if args.max_grad_norm is not None and args.max_grad_norm > 0:
                        

                        if is_sagemaker_mp_enabled() and args.fp16:
                            self.optimizer.clip_master_grads(args.max_grad_norm)
                        elif self.use_apex:
                            
                            nn.utils.clip_grad_norm_(
                                amp.master_params(self.optimizer),
                                args.max_grad_norm,
                            )
                        else:
                            self.accelerator.clip_grad_norm_(
                                model.parameters(),
                                args.max_grad_norm,
                            )

                    
                    self.optimizer.step()
                    optimizer_was_run = not self.accelerator.optimizer_step_was_skipped
                    if optimizer_was_run:
                        
                        if not isinstance(self.lr_scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
                            self.lr_scheduler.step()

                    model.zero_grad()
                    self.state.global_step += 1
                    self.state.epoch = epoch + (step + 1 + steps_skipped) / steps_in_epoch
                    self.control = self.callback_handler.on_step_end(args, self.state, self.control)

                    self._maybe_log_save_evaluate(tr_loss, model, trial, epoch, ignore_keys_for_eval)
                else:
                    self.control = self.callback_handler.on_substep_end(args, self.state, self.control)

                if self.control.should_epoch_stop or self.control.should_training_stop:
                    break
            if step < 0:
                logger.warning(
                    "There seems to be not a single sample in your epoch_iterator, stopping training at step"
                    f" {self.state.global_step}! This is expected if you're using an IterableDataset and set"
                    f" num_steps ({max_steps}) higher than the number of available samples."
                )
                self.control.should_training_stop = True

            self.control = self.callback_handler.on_epoch_end(args, self.state, self.control)
            self._maybe_log_save_evaluate(tr_loss, model, trial, epoch, ignore_keys_for_eval)

            if DebugOption.TPU_METRICS_DEBUG in self.args.debug:
                if is_torch_tpu_available():
                    
                    xm.master_print(met.metrics_report())
                else:
                    logger.warning(
                        "You enabled PyTorch/XLA debug metrics but you don't have a TPU "
                        "configured. Check your training configuration if this is unexpected."
                    )
            if self.control.should_training_stop:
                break

        if args.past_index and hasattr(self, "_past"):
            
            delattr(self, "_past")

        logger.info("\n\nTraining completed. Do not forget to share your model on huggingface.co/models =)\n\n")
        if args.load_best_model_at_end and self.state.best_model_checkpoint is not None:
            
            if is_torch_tpu_available():
                xm.rendezvous("load_best_model_at_end")
            elif args.parallel_mode == ParallelMode.DISTRIBUTED:
                dist.barrier()
            elif is_sagemaker_mp_enabled():
                smp.barrier()

            self._load_best_model()

        
        self._total_loss_scalar += tr_loss.item()
        train_loss = self._total_loss_scalar / self.state.global_step

        metrics = speed_metrics(
            "train",
            start_time,
            num_samples=num_train_samples,
            num_steps=self.state.max_steps,
            num_tokens=num_train_tokens,
        )
        self.store_flos()
        metrics["total_flos"] = self.state.total_flos
        metrics["train_loss"] = train_loss

        self.is_in_train = False

        self._memory_tracker.stop_and_update_metrics(metrics)

        self.log(metrics)

        run_dir = self._get_output_dir(trial)
        checkpoints_sorted = self._sorted_checkpoints(use_mtime=False, output_dir=run_dir)

        
        if self.args.should_save and self.state.best_model_checkpoint is not None and self.args.save_total_limit == 1:
            for checkpoint in checkpoints_sorted:
                if not os.path.samefile(checkpoint, self.state.best_model_checkpoint):
                    logger.info(f"Deleting older checkpoint [{checkpoint}] due to args.save_total_limit")
                    shutil.rmtree(checkpoint)

        self.control = self.callback_handler.on_train_end(args, self.state, self.control)

        
        self._finish_current_push()

        
        
        if self.neftune_noise_alpha is not None:
            self._deactivate_neftune(self.model)

        return TrainOutput(self.state.global_step, train_loss, metrics)

    def _get_output_dir(self, trial):
        if self.hp_search_backend is not None and trial is not None:
            if self.hp_search_backend == HPSearchBackend.OPTUNA:
                run_id = trial.number
            elif self.hp_search_backend == HPSearchBackend.RAY:
                from ray import tune

                run_id = tune.get_trial_id()
            elif self.hp_search_backend == HPSearchBackend.SIGOPT:
                run_id = trial.id
            elif self.hp_search_backend == HPSearchBackend.WANDB:
                import wandb

                run_id = wandb.run.id
            run_name = self.hp_name(trial) if self.hp_name is not None else f"run-{run_id}"
            run_dir = os.path.join(self.args.output_dir, run_name)
        else:
            run_dir = self.args.output_dir
        return run_dir

    def _load_from_checkpoint(self, resume_from_checkpoint, model=None):
        if model is None:
            model = self.model

        config_file = os.path.join(resume_from_checkpoint, CONFIG_NAME)
        adapter_weights_file = os.path.join(resume_from_checkpoint, ADAPTER_WEIGHTS_NAME)
        adapter_safe_weights_file = os.path.join(resume_from_checkpoint, ADAPTER_SAFE_WEIGHTS_NAME)
        weights_file = os.path.join(resume_from_checkpoint, WEIGHTS_NAME)
        weights_index_file = os.path.join(resume_from_checkpoint, WEIGHTS_INDEX_NAME)
        safe_weights_file = os.path.join(resume_from_checkpoint, SAFE_WEIGHTS_NAME)
        safe_weights_index_file = os.path.join(resume_from_checkpoint, SAFE_WEIGHTS_INDEX_NAME)
        is_fsdp_ckpt = os.path.isdir(resume_from_checkpoint) and any(
            FSDP_MODEL_NAME in folder_name
            for folder_name in os.listdir(resume_from_checkpoint)
            if os.path.isdir(os.path.join(resume_from_checkpoint, folder_name))
        )

        if is_fsdp_ckpt and not self.is_fsdp_enabled:
            raise ValueError(f"Checkpoint found at {resume_from_checkpoint} is only supported when using PyTorch FSDP")

        if not (
            any(
                os.path.isfile(f)
                for f in [
                    weights_file,
                    safe_weights_file,
                    weights_index_file,
                    safe_weights_index_file,
                    adapter_weights_file,
                    adapter_safe_weights_file,
                ]
            )
            or is_fsdp_ckpt
        ):
            raise ValueError(f"Can't find a valid checkpoint at {resume_from_checkpoint}")

        logger.info(f"Loading model from {resume_from_checkpoint}.")

        if os.path.isfile(config_file):
            config = PretrainedConfig.from_json_file(config_file)
            checkpoint_version = config.transformers_version
            if checkpoint_version is not None and checkpoint_version != __version__:
                logger.warning(
                    f"You are resuming training from a checkpoint trained with {checkpoint_version} of "
                    f"Transformers but your current version is {__version__}. This is not recommended and could "
                    "yield to errors or unwanted behaviors."
                )

        if os.path.isfile(weights_file) or os.path.isfile(safe_weights_file) or is_fsdp_ckpt:
            
            if is_sagemaker_mp_enabled():
                if os.path.isfile(os.path.join(resume_from_checkpoint, "user_content.pt")):
                    
                    
                    smp.resume_from_checkpoint(
                        path=resume_from_checkpoint, tag=WEIGHTS_NAME, partial=False, load_optimizer=False
                    )
                else:
                    
                    
                    if hasattr(self.args, "fp16") and self.args.fp16 is True:
                        logger.warning(
                            "Enabling FP16 and loading from smp < 1.10 checkpoint together is not suppported."
                        )
                    state_dict = torch.load(weights_file, map_location="cpu")
                    
                    state_dict["_smp_is_partial"] = False
                    load_result = model.load_state_dict(state_dict, strict=True)
                    
                    del state_dict
            elif self.is_fsdp_enabled:
                load_fsdp_model(self.accelerator.state.fsdp_plugin, self.accelerator, model, resume_from_checkpoint)
            else:
                
                if self.args.save_safetensors and os.path.isfile(safe_weights_file):
                    state_dict = safetensors.torch.load_file(safe_weights_file, device="cpu")
                else:
                    state_dict = torch.load(weights_file, map_location="cpu")

                
                
                load_result = model.load_state_dict(state_dict, False)
                
                del state_dict
                self._issue_warnings_after_load(load_result)

        
        elif is_peft_available() and isinstance(model, PeftModel):
            
            if hasattr(model, "active_adapter") and hasattr(model, "load_adapter"):
                if os.path.exists(resume_from_checkpoint):
                    model.load_adapter(resume_from_checkpoint, model.active_adapter, is_trainable=True)
                else:
                    logger.warning(
                        "The intermediate checkpoints of PEFT may not be saved correctly, "
                        f"consider using a custom callback to save {ADAPTER_WEIGHTS_NAME} in corresponding saving folders. "
                        "Check some examples here: https://github.com/huggingface/peft/issues/96"
                    )
            else:
                logger.warning("Could not load adapter model, make sure to have `peft>=0.3.0` installed")
        else:
            
            load_result = load_sharded_checkpoint(
                model, resume_from_checkpoint, strict=is_sagemaker_mp_enabled(), prefer_safe=self.args.save_safetensors
            )
            if not is_sagemaker_mp_enabled():
                self._issue_warnings_after_load(load_result)

    def _load_best_model(self):
        logger.info(f"Loading best model from {self.state.best_model_checkpoint} (score: {self.state.best_metric}).")
        best_model_path = os.path.join(self.state.best_model_checkpoint, WEIGHTS_NAME)
        best_safe_model_path = os.path.join(self.state.best_model_checkpoint, SAFE_WEIGHTS_NAME)
        best_adapter_model_path = os.path.join(self.state.best_model_checkpoint, ADAPTER_WEIGHTS_NAME)
        best_safe_adapter_model_path = os.path.join(self.state.best_model_checkpoint, ADAPTER_SAFE_WEIGHTS_NAME)

        model = self.model_wrapped if is_sagemaker_mp_enabled() else self.model
        if self.is_deepspeed_enabled:
            deepspeed_load_checkpoint(self.model_wrapped, self.state.best_model_checkpoint)
        elif self.is_fsdp_enabled:
            load_result = load_fsdp_model(
                self.accelerator.state.fsdp_plugin, self.accelerator, model, self.state.best_model_checkpoint
            )
        elif (
            os.path.exists(best_model_path)
            or os.path.exists(best_safe_model_path)
            or os.path.exists(best_adapter_model_path)
            or os.path.exists(best_safe_adapter_model_path)
        ):
            has_been_loaded = True
            if is_sagemaker_mp_enabled():
                if os.path.isfile(os.path.join(self.state.best_model_checkpoint, "user_content.pt")):
                    
                    
                    smp.resume_from_checkpoint(
                        path=self.state.best_model_checkpoint,
                        tag=WEIGHTS_NAME,
                        partial=False,
                        load_optimizer=False,
                    )
                else:
                    
                    
                    if self.args.save_safetensors and os.path.isfile(best_safe_model_path):
                        state_dict = safetensors.torch.load_file(best_safe_model_path, device="cpu")
                    else:
                        state_dict = torch.load(best_model_path, map_location="cpu")

                    state_dict["_smp_is_partial"] = False
                    load_result = model.load_state_dict(state_dict, strict=True)
            else:
                if is_peft_available() and isinstance(model, PeftModel):
                    
                    if hasattr(model, "active_adapter") and hasattr(model, "load_adapter"):
                        if os.path.exists(best_adapter_model_path) or os.path.exists(best_safe_adapter_model_path):
                            model.load_adapter(self.state.best_model_checkpoint, model.active_adapter)
                            
                            from torch.nn.modules.module import _IncompatibleKeys

                            load_result = _IncompatibleKeys([], [])
                        else:
                            logger.warning(
                                "The intermediate checkpoints of PEFT may not be saved correctly, "
                                f"consider using a custom callback to save {ADAPTER_WEIGHTS_NAME} in corresponding saving folders. "
                                "Check some examples here: https://github.com/huggingface/peft/issues/96"
                            )
                            has_been_loaded = False
                    else:
                        logger.warning("Could not load adapter model, make sure to have `peft>=0.3.0` installed")
                        has_been_loaded = False
                else:
                    
                    if self.args.save_safetensors and os.path.isfile(best_safe_model_path):
                        state_dict = safetensors.torch.load_file(best_safe_model_path, device="cpu")
                    else:
                        state_dict = torch.load(best_model_path, map_location="cpu")

                    
                    
                    
                    load_result = model.load_state_dict(state_dict, False)
                if not is_sagemaker_mp_enabled() and has_been_loaded:
                    self._issue_warnings_after_load(load_result)
        elif os.path.exists(os.path.join(self.state.best_model_checkpoint, WEIGHTS_INDEX_NAME)):
            load_result = load_sharded_checkpoint(
                model, self.state.best_model_checkpoint, strict=is_sagemaker_mp_enabled()
            )
            if not is_sagemaker_mp_enabled():
                self._issue_warnings_after_load(load_result)
        else:
            logger.warning(
                f"Could not locate the best model at {best_model_path}, if you are running a distributed training "
                "on multiple nodes, you should activate `--save_on_each_node`."
            )

    def _issue_warnings_after_load(self, load_result):
        if len(load_result.missing_keys) != 0:
            if self.model._keys_to_ignore_on_save is not None and set(load_result.missing_keys) == set(
                self.model._keys_to_ignore_on_save
            ):
                self.model.tie_weights()
            else:
                logger.warning(f"There were missing keys in the checkpoint model loaded: {load_result.missing_keys}.")
        if len(load_result.unexpected_keys) != 0:
            logger.warning(
                f"There were unexpected keys in the checkpoint model loaded: {load_result.unexpected_keys}."
            )

    def _maybe_log_save_evaluate(self, tr_loss, model, trial, epoch, ignore_keys_for_eval):
        if self.control.should_log:
            if is_torch_tpu_available():
                xm.mark_step()

            logs: Dict[str, float] = {}

            
            tr_loss_scalar = self._nested_gather(tr_loss).mean().item()

            
            tr_loss -= tr_loss

            logs["loss"] = round(tr_loss_scalar / (self.state.global_step - self._globalstep_last_logged), 4)
            logs["learning_rate"] = self._get_learning_rate()

            self._total_loss_scalar += tr_loss_scalar
            self._globalstep_last_logged = self.state.global_step
            self.store_flos()

            self.log(logs)

        metrics = None
        if self.control.should_evaluate:
            if isinstance(self.eval_dataset, dict):
                metrics = {}
                for eval_dataset_name, eval_dataset in self.eval_dataset.items():
                    dataset_metrics = self.evaluate(
                        eval_dataset=eval_dataset,
                        ignore_keys=ignore_keys_for_eval,
                        metric_key_prefix=f"eval_{eval_dataset_name}",
                    )
                    metrics.update(dataset_metrics)
            else:
                metrics = self.evaluate(ignore_keys=ignore_keys_for_eval)
            self._report_to_hp_search(trial, self.state.global_step, metrics)

            
            if isinstance(self.lr_scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
                metric_to_check = self.args.metric_for_best_model
                if not metric_to_check.startswith("eval_"):
                    metric_to_check = f"eval_{metric_to_check}"
                self.lr_scheduler.step(metrics[metric_to_check])

        if self.control.should_save:
            self._save_checkpoint(model, trial, metrics=metrics)
            self.control = self.callback_handler.on_save(self.args, self.state, self.control)

    def _load_rng_state(self, checkpoint):
        
        if checkpoint is None:
            return

        if self.args.world_size > 1:
            process_index = self.args.process_index
            rng_file = os.path.join(checkpoint, f"rng_state_{process_index}.pth")
            if not os.path.isfile(rng_file):
                logger.info(
                    f"Didn't find an RNG file for process {process_index}, if you are resuming a training that "
                    "wasn't launched in a distributed fashion, reproducibility is not guaranteed."
                )
                return
        else:
            rng_file = os.path.join(checkpoint, "rng_state.pth")
            if not os.path.isfile(rng_file):
                logger.info(
                    "Didn't find an RNG file, if you are resuming a training that was launched in a distributed "
                    "fashion, reproducibility is not guaranteed."
                )
                return

        checkpoint_rng_state = torch.load(rng_file)
        random.setstate(checkpoint_rng_state["python"])
        np.random.set_state(checkpoint_rng_state["numpy"])
        torch.random.set_rng_state(checkpoint_rng_state["cpu"])
        if torch.cuda.is_available():
            if self.args.parallel_mode == ParallelMode.DISTRIBUTED:
                torch.cuda.random.set_rng_state_all(checkpoint_rng_state["cuda"])
            else:
                try:
                    torch.cuda.random.set_rng_state(checkpoint_rng_state["cuda"])
                except Exception as e:
                    logger.info(
                        f"Didn't manage to set back the RNG states of the GPU because of the following error:\n {e}"
                        "\nThis won't yield the same results as if the training had not been interrupted."
                    )
        if is_torch_tpu_available():
            xm.set_rng_state(checkpoint_rng_state["xla"])
        if is_torch_npu_available():
            if self.args.parallel_mode == ParallelMode.DISTRIBUTED:
                torch.npu.random.set_rng_state_all(checkpoint_rng_state["npu"])
            else:
                try:
                    torch.npu.random.set_rng_state(checkpoint_rng_state["npu"])
                except Exception as e:
                    logger.info(
                        f"Didn't manage to set back the RNG states of the NPU because of the following error:\n {e}"
                        "\nThis won't yield the same results as if the training had not been interrupted."
                    )

    def _save_checkpoint(self, model, trial, metrics=None):
        
        
        

        
        checkpoint_folder = f"{PREFIX_CHECKPOINT_DIR}-{self.state.global_step}"

        if self.hp_search_backend is None and trial is None:
            self.store_flos()

        run_dir = self._get_output_dir(trial=trial)
        output_dir = os.path.join(run_dir, checkpoint_folder)
        self.save_model(output_dir, _internal_call=True)

        if not self.args.save_only_model:
            
            self._save_optimizer_and_scheduler(output_dir)
            
            self._save_rng_state(output_dir)

        
        if metrics is not None and self.args.metric_for_best_model is not None:
            metric_to_check = self.args.metric_for_best_model
            if not metric_to_check.startswith("eval_"):
                metric_to_check = f"eval_{metric_to_check}"
            metric_value = metrics[metric_to_check]

            operator = np.greater if self.args.greater_is_better else np.less
            if (
                self.state.best_metric is None
                or self.state.best_model_checkpoint is None
                or operator(metric_value, self.state.best_metric)
            ):
                self.state.best_metric = metric_value
                self.state.best_model_checkpoint = output_dir

        
        if self.args.should_save:
            self.state.save_to_json(os.path.join(output_dir, TRAINER_STATE_NAME))

        if self.args.push_to_hub:
            self._push_from_checkpoint(output_dir)

        
        if self.args.should_save:
            self._rotate_checkpoints(use_mtime=True, output_dir=run_dir)

    def _save_rng_state(self, output_dir):
        
        rng_states = {
            "python": random.getstate(),
            "numpy": np.random.get_state(),
            "cpu": torch.random.get_rng_state(),
        }
        if torch.cuda.is_available():
            if self.args.parallel_mode == ParallelMode.DISTRIBUTED:
                
                rng_states["cuda"] = torch.cuda.random.get_rng_state_all()
            else:
                rng_states["cuda"] = torch.cuda.random.get_rng_state()

        if is_torch_tpu_available():
            rng_states["xla"] = xm.get_rng_state()

        if is_torch_npu_available():
            if self.args.parallel_mode == ParallelMode.DISTRIBUTED:
                rng_states["npu"] = torch.npu.random.get_rng_state_all()
            else:
                rng_states["npu"] = torch.npu.random.get_rng_state()

        
        
        os.makedirs(output_dir, exist_ok=True)

        if self.args.world_size <= 1:
            torch.save(rng_states, os.path.join(output_dir, "rng_state.pth"))
        else:
            torch.save(rng_states, os.path.join(output_dir, f"rng_state_{self.args.process_index}.pth"))

    def _save_optimizer_and_scheduler(self, output_dir):
        if is_torch_tpu_available():
            xm.rendezvous("saving_optimizer_states")
            xm.save(self.optimizer.state_dict(), os.path.join(output_dir, OPTIMIZER_NAME))
            with warnings.catch_warnings(record=True) as caught_warnings:
                xm.save(self.lr_scheduler.state_dict(), os.path.join(output_dir, SCHEDULER_NAME))
                reissue_pt_warnings(caught_warnings)
        elif is_sagemaker_mp_enabled():
            opt_state_dict = self.optimizer.local_state_dict(gather_if_shard=False)
            smp.barrier()
            if smp.rdp_rank() == 0 or smp.state.cfg.shard_optimizer_state:
                smp.save(
                    opt_state_dict,
                    os.path.join(output_dir, OPTIMIZER_NAME),
                    partial=True,
                    v3=smp.state.cfg.shard_optimizer_state,
                )
        elif self.is_deepspeed_enabled:
            
            
            self.model_wrapped.save_checkpoint(output_dir)
        elif self.is_fsdp_enabled:
            
            save_fsdp_model(self.accelerator.state.fsdp_plugin, self.accelerator, self.model, output_dir)
            save_fsdp_optimizer(
                self.accelerator.state.fsdp_plugin, self.accelerator, self.optimizer, self.model, output_dir
            )
        elif self.args.should_save:
            
            torch.save(self.optimizer.state_dict(), os.path.join(output_dir, OPTIMIZER_NAME))

        
        is_deepspeed_custom_scheduler = self.is_deepspeed_enabled and not isinstance(
            self.lr_scheduler, DeepSpeedSchedulerWrapper
        )
        if (
            self.args.should_save
            and (not self.is_deepspeed_enabled or is_deepspeed_custom_scheduler)
            and not is_torch_tpu_available()
        ):
            with warnings.catch_warnings(record=True) as caught_warnings:
                torch.save(self.lr_scheduler.state_dict(), os.path.join(output_dir, SCHEDULER_NAME))
            reissue_pt_warnings(caught_warnings)

    def _load_optimizer_and_scheduler(self, checkpoint):
        
        if checkpoint is None:
            return

        if self.is_deepspeed_enabled:
            
            if not isinstance(self.lr_scheduler, DeepSpeedSchedulerWrapper):
                with warnings.catch_warnings(record=True) as caught_warnings:
                    self.lr_scheduler.load_state_dict(torch.load(os.path.join(checkpoint, SCHEDULER_NAME)))
                reissue_pt_warnings(caught_warnings)
            return

        checkpoint_file_exists = (
            glob.glob(os.path.join(checkpoint, OPTIMIZER_NAME) + "_*")
            if is_sagemaker_mp_enabled()
            else (
                os.path.isfile(os.path.join(checkpoint, OPTIMIZER_NAME))
                or os.path.isfile(os.path.join(checkpoint, OPTIMIZER_NAME_BIN))
                or (
                    os.path.isdir(checkpoint)
                    and any(
                        OPTIMIZER_NAME_BIN.split(".")[0] in folder_name
                        for folder_name in os.listdir(checkpoint)
                        if os.path.isdir(os.path.join(checkpoint, folder_name))
                    )
                )
            )
        )
        if checkpoint_file_exists and os.path.isfile(os.path.join(checkpoint, SCHEDULER_NAME)):
            
            if is_torch_tpu_available():
                
                optimizer_state = torch.load(os.path.join(checkpoint, OPTIMIZER_NAME), map_location="cpu")
                with warnings.catch_warnings(record=True) as caught_warnings:
                    lr_scheduler_state = torch.load(os.path.join(checkpoint, SCHEDULER_NAME), map_location="cpu")
                reissue_pt_warnings(caught_warnings)

                xm.send_cpu_data_to_device(optimizer_state, self.args.device)
                xm.send_cpu_data_to_device(lr_scheduler_state, self.args.device)

                self.optimizer.load_state_dict(optimizer_state)
                self.lr_scheduler.load_state_dict(lr_scheduler_state)
            else:
                if is_sagemaker_mp_enabled():
                    if os.path.isfile(os.path.join(checkpoint, "user_content.pt")):
                        
                        def opt_load_hook(mod, opt):
                            opt.load_state_dict(smp.load(os.path.join(checkpoint, OPTIMIZER_NAME), partial=True))

                    else:
                        
                        def opt_load_hook(mod, opt):
                            if IS_SAGEMAKER_MP_POST_1_10:
                                opt.load_state_dict(
                                    smp.load(os.path.join(checkpoint, OPTIMIZER_NAME), partial=True, back_compat=True)
                                )
                            else:
                                opt.load_state_dict(smp.load(os.path.join(checkpoint, OPTIMIZER_NAME), partial=True))

                    self.model_wrapped.register_post_step_hook(opt_load_hook)
                else:
                    
                    
                    
                    map_location = self.args.device if self.args.world_size > 1 else "cpu"
                    if self.is_fsdp_enabled:
                        load_fsdp_optimizer(
                            self.accelerator.state.fsdp_plugin,
                            self.accelerator,
                            self.optimizer,
                            self.model,
                            checkpoint,
                        )
                    else:
                        self.optimizer.load_state_dict(
                            torch.load(os.path.join(checkpoint, OPTIMIZER_NAME), map_location=map_location)
                        )
                with warnings.catch_warnings(record=True) as caught_warnings:
                    self.lr_scheduler.load_state_dict(torch.load(os.path.join(checkpoint, SCHEDULER_NAME)))
                reissue_pt_warnings(caught_warnings)

    def hyperparameter_search(
        self,
        hp_space: Optional[Callable[["optuna.Trial"], Dict[str, float]]] = None,
        compute_objective: Optional[Callable[[Dict[str, float]], float]] = None,
        n_trials: int = 20,
        direction: Union[str, List[str]] = "minimize",
        backend: Optional[Union["str", HPSearchBackend]] = None,
        hp_name: Optional[Callable[["optuna.Trial"], str]] = None,
        **kwargs,
    ) -> Union[BestRun, List[BestRun]]:
        
        if backend is None:
            backend = default_hp_search_backend()
        backend = HPSearchBackend(backend)
        backend_obj = ALL_HYPERPARAMETER_SEARCH_BACKENDS[backend]()
        backend_obj.ensure_available()
        self.hp_search_backend = backend
        if self.model_init is None:
            raise RuntimeError(
                "To use hyperparameter search, you need to pass your model through a model_init function."
            )

        self.hp_space = backend_obj.default_hp_space if hp_space is None else hp_space
        self.hp_name = hp_name
        self.compute_objective = default_compute_objective if compute_objective is None else compute_objective

        best_run = backend_obj.run(self, n_trials, direction, **kwargs)

        self.hp_search_backend = None
        return best_run

    def log(self, logs: Dict[str, float]) -> None:
        
        if self.state.epoch is not None:
            logs["epoch"] = round(self.state.epoch, 2)
        if self.args.include_num_input_tokens_seen:
            logs["num_input_tokens_seen"] = self.state.num_input_tokens_seen

        output = {**logs, **{"step": self.state.global_step}}
        self.state.log_history.append(output)
        self.control = self.callback_handler.on_log(self.args, self.state, self.control, logs)

    def _prepare_input(self, data: Union[torch.Tensor, Any]) -> Union[torch.Tensor, Any]:
        
        if isinstance(data, Mapping):
            return type(data)({k: self._prepare_input(v) for k, v in data.items()})
        elif isinstance(data, (tuple, list)):
            return type(data)(self._prepare_input(v) for v in data)
        elif isinstance(data, torch.Tensor):
            kwargs = {"device": self.args.device}
            if self.is_deepspeed_enabled and (torch.is_floating_point(data) or torch.is_complex(data)):
                
                
                
                kwargs.update({"dtype": self.accelerator.state.deepspeed_plugin.hf_ds_config.dtype()})
            return data.to(**kwargs)
        return data

    def _prepare_inputs(self, inputs: Dict[str, Union[torch.Tensor, Any]]) -> Dict[str, Union[torch.Tensor, Any]]:
        
        inputs = self._prepare_input(inputs)
        if len(inputs) == 0:
            raise ValueError(
                "The batch received was empty, your model won't be able to train on it. Double-check that your "
                f"training dataset contains keys expected by the model: {','.join(self._signature_columns)}."
            )
        if self.args.past_index >= 0 and self._past is not None:
            inputs["mems"] = self._past

        return inputs

    def compute_loss_context_manager(self):
        
        return self.autocast_smart_context_manager()

    def autocast_smart_context_manager(self, cache_enabled: Optional[bool] = True):
        
        if self.use_cpu_amp:
            ctx_manager = torch.cpu.amp.autocast(cache_enabled=cache_enabled, dtype=self.amp_dtype)
        else:
            ctx_manager = contextlib.nullcontext()

        return ctx_manager

    def training_step(self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]]) -> torch.Tensor:
        
        model.train()
        inputs = self._prepare_inputs(inputs)

        if is_sagemaker_mp_enabled():
            loss_mb = smp_forward_backward(model, inputs, self.args.gradient_accumulation_steps)
            return loss_mb.reduce_mean().detach().to(self.args.device)

        with self.compute_loss_context_manager():
            loss = self.compute_loss(model, inputs)

        if self.args.n_gpu > 1:
            loss = loss.mean()  

        if self.use_apex:
            with amp.scale_loss(loss, self.optimizer) as scaled_loss:
                scaled_loss.backward()
        else:
            self.accelerator.backward(loss)

        return loss.detach() / self.args.gradient_accumulation_steps

    def compute_loss(self, model, inputs, return_outputs=False):
        
        if self.label_smoother is not None and "labels" in inputs:
            labels = inputs.pop("labels")
        else:
            labels = None
        outputs = model(**inputs)
        
        
        if self.args.past_index >= 0:
            self._past = outputs[self.args.past_index]

        if labels is not None:
            unwrapped_model = unwrap_model(model)
            if is_peft_available() and isinstance(unwrapped_model, PeftModel):
                model_name = unwrapped_model.base_model.model._get_name()
            else:
                model_name = unwrapped_model._get_name()
            if model_name in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES.values():
                loss = self.label_smoother(outputs, labels, shift_labels=True)
            else:
                loss = self.label_smoother(outputs, labels)
        else:
            if isinstance(outputs, dict) and "loss" not in outputs:
                raise ValueError(
                    "The model did not return a loss from the inputs, only the following keys: "
                    f"{','.join(outputs.keys())}. For reference, the inputs it received are {','.join(inputs.keys())}."
                )
            
            loss = outputs["loss"] if isinstance(outputs, dict) else outputs[0]

        return (loss, outputs) if return_outputs else loss

    def is_local_process_zero(self) -> bool:
        
        return self.args.local_process_index == 0

    def is_world_process_zero(self) -> bool:
        
        
        
        if is_sagemaker_mp_enabled():
            return smp.rank() == 0
        else:
            return self.args.process_index == 0

    def save_model(self, output_dir: Optional[str] = None, _internal_call: bool = False):
        

        if output_dir is None:
            output_dir = self.args.output_dir

        if is_torch_tpu_available():
            self._save_tpu(output_dir)
        elif is_sagemaker_mp_enabled():
            
            os.makedirs(output_dir, exist_ok=True)
            state_dict = self.model_wrapped.state_dict()
            if self.args.should_save:
                self._save(output_dir, state_dict=state_dict)
            if IS_SAGEMAKER_MP_POST_1_10:
                
                Path(os.path.join(output_dir, "user_content.pt")).touch()
        elif self.is_fsdp_enabled:
            if ("FULL_STATE_DICT" in str(self.accelerator.state.fsdp_plugin.state_dict_type)) and (
                version.parse(accelerate_version) > version.parse("0.24.1")
            ):
                state_dict = self.accelerator.get_state_dict(self.model)
                if self.args.should_save:
                    self._save(output_dir, state_dict=state_dict)
        elif self.is_deepspeed_enabled:
            try:
                state_dict = self.accelerator.get_state_dict(self.deepspeed)
                if self.args.should_save:
                    self._save(output_dir, state_dict=state_dict)
            except ValueError:
                logger.warning(
                    " stage3_gather_16bit_weights_on_model_save=false. Saving the full checkpoint instead, use"
                    " zero_to_fp32.py to recover weights"
                )
                if self.args.should_save:
                    self._save(output_dir, state_dict={})
                
                remove_dummy_checkpoint(self.args.should_save, output_dir, [WEIGHTS_NAME, SAFE_WEIGHTS_NAME])
                self.model_wrapped.save_checkpoint(output_dir)

        elif self.args.should_save:
            self._save(output_dir)

        
        if self.args.push_to_hub and not _internal_call:
            self.push_to_hub(commit_message="Model save")

    def _save_tpu(self, output_dir: Optional[str] = None):
        output_dir = output_dir if output_dir is not None else self.args.output_dir
        logger.info(f"Saving model checkpoint to {output_dir}")

        if xm.is_master_ordinal():
            os.makedirs(output_dir, exist_ok=True)
            torch.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME))

        
        
        xm.rendezvous("saving_checkpoint")
        if not isinstance(self.model, PreTrainedModel):
            if isinstance(unwrap_model(self.model), PreTrainedModel):
                unwrap_model(self.model).save_pretrained(
                    output_dir,
                    is_main_process=self.args.should_save,
                    state_dict=self.model.state_dict(),
                    save_function=xm.save,
                )
            else:
                logger.info("Trainer.model is not a `PreTrainedModel`, only saving its state dict.")
                state_dict = self.model.state_dict()
                xm.save(state_dict, os.path.join(output_dir, WEIGHTS_NAME))
        else:
            self.model.save_pretrained(output_dir, is_main_process=self.args.should_save, save_function=xm.save)
        if self.tokenizer is not None and self.args.should_save:
            self.tokenizer.save_pretrained(output_dir)

    def _save(self, output_dir: Optional[str] = None, state_dict=None):
        
        output_dir = output_dir if output_dir is not None else self.args.output_dir
        os.makedirs(output_dir, exist_ok=True)
        logger.info(f"Saving model checkpoint to {output_dir}")

        supported_classes = (PreTrainedModel,) if not is_peft_available() else (PreTrainedModel, PeftModel)
        
        
        if not isinstance(self.model, supported_classes):
            if state_dict is None:
                state_dict = self.model.state_dict()

            if isinstance(unwrap_model(self.model), supported_classes):
                unwrap_model(self.model).save_pretrained(
                    output_dir, state_dict=state_dict, safe_serialization=self.args.save_safetensors
                )
            else:
                logger.info("Trainer.model is not a `PreTrainedModel`, only saving its state dict.")
                if self.args.save_safetensors:
                    safetensors.torch.save_file(state_dict, os.path.join(output_dir, SAFE_WEIGHTS_NAME))
                else:
                    torch.save(state_dict, os.path.join(output_dir, WEIGHTS_NAME))
        else:
            self.model.save_pretrained(
                output_dir, state_dict=state_dict, safe_serialization=self.args.save_safetensors
            )

        if self.tokenizer is not None:
            self.tokenizer.save_pretrained(output_dir)

        
        torch.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME))

    def store_flos(self):
        
        if self.args.parallel_mode == ParallelMode.DISTRIBUTED:
            self.state.total_flos += (
                distributed_broadcast_scalars([self.current_flos], device=self.args.device).sum().item()
            )
            self.current_flos = 0
        else:
            self.state.total_flos += self.current_flos
            self.current_flos = 0

    def _sorted_checkpoints(
        self, output_dir=None, checkpoint_prefix=PREFIX_CHECKPOINT_DIR, use_mtime=False
    ) -> List[str]:
        ordering_and_checkpoint_path = []

        glob_checkpoints = [str(x) for x in Path(output_dir).glob(f"{checkpoint_prefix}-*") if os.path.isdir(x)]

        for path in glob_checkpoints:
            if use_mtime:
                ordering_and_checkpoint_path.append((os.path.getmtime(path), path))
            else:
                regex_match = re.match(f".*{checkpoint_prefix}-([0-9]+)", path)
                if regex_match is not None and regex_match.groups() is not None:
                    ordering_and_checkpoint_path.append((int(regex_match.groups()[0]), path))

        checkpoints_sorted = sorted(ordering_and_checkpoint_path)
        checkpoints_sorted = [checkpoint[1] for checkpoint in checkpoints_sorted]
        
        if (
            self.state.best_model_checkpoint is not None
            and str(Path(self.state.best_model_checkpoint)) in checkpoints_sorted
        ):
            best_model_index = checkpoints_sorted.index(str(Path(self.state.best_model_checkpoint)))
            for i in range(best_model_index, len(checkpoints_sorted) - 2):
                checkpoints_sorted[i], checkpoints_sorted[i + 1] = checkpoints_sorted[i + 1], checkpoints_sorted[i]
        return checkpoints_sorted

    def _rotate_checkpoints(self, use_mtime=False, output_dir=None) -> None:
        if self.args.save_total_limit is None or self.args.save_total_limit <= 0:
            return

        
        checkpoints_sorted = self._sorted_checkpoints(use_mtime=use_mtime, output_dir=output_dir)
        if len(checkpoints_sorted) <= self.args.save_total_limit:
            return

        
        
        save_total_limit = self.args.save_total_limit
        if (
            self.state.best_model_checkpoint is not None
            and self.args.save_total_limit == 1
            and checkpoints_sorted[-1] != self.state.best_model_checkpoint
        ):
            save_total_limit = 2

        number_of_checkpoints_to_delete = max(0, len(checkpoints_sorted) - save_total_limit)
        checkpoints_to_be_deleted = checkpoints_sorted[:number_of_checkpoints_to_delete]
        for checkpoint in checkpoints_to_be_deleted:
            logger.info(f"Deleting older checkpoint [{checkpoint}] due to args.save_total_limit")
            shutil.rmtree(checkpoint, ignore_errors=True)

    def evaluate(
        self,
        eval_dataset: Optional[Dataset] = None,
        ignore_keys: Optional[List[str]] = None,
        metric_key_prefix: str = "eval",
    ) -> Dict[str, float]:
        
        
        self._memory_tracker.start()

        eval_dataloader = self.get_eval_dataloader(eval_dataset)
        start_time = time.time()

        eval_loop = self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
        output = eval_loop(
            eval_dataloader,
            description="Evaluation",
            
            
            prediction_loss_only=True if self.compute_metrics is None else None,
            ignore_keys=ignore_keys,
            metric_key_prefix=metric_key_prefix,
        )

        total_batch_size = self.args.eval_batch_size * self.args.world_size
        if f"{metric_key_prefix}_jit_compilation_time" in output.metrics:
            start_time += output.metrics[f"{metric_key_prefix}_jit_compilation_time"]
        output.metrics.update(
            speed_metrics(
                metric_key_prefix,
                start_time,
                num_samples=output.num_samples,
                num_steps=math.ceil(output.num_samples / total_batch_size),
            )
        )

        self.log(output.metrics)

        if DebugOption.TPU_METRICS_DEBUG in self.args.debug:
            
            xm.master_print(met.metrics_report())

        self.control = self.callback_handler.on_evaluate(self.args, self.state, self.control, output.metrics)

        self._memory_tracker.stop_and_update_metrics(output.metrics)

        return output.metrics

    def predict(
        self, test_dataset: Dataset, ignore_keys: Optional[List[str]] = None, metric_key_prefix: str = "test"
    ) -> PredictionOutput:
        
        
        self._memory_tracker.start()

        test_dataloader = self.get_test_dataloader(test_dataset)
        start_time = time.time()

        eval_loop = self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
        output = eval_loop(
            test_dataloader, description="Prediction", ignore_keys=ignore_keys, metric_key_prefix=metric_key_prefix
        )
        total_batch_size = self.args.eval_batch_size * self.args.world_size
        if f"{metric_key_prefix}_jit_compilation_time" in output.metrics:
            start_time += output.metrics[f"{metric_key_prefix}_jit_compilation_time"]
        output.metrics.update(
            speed_metrics(
                metric_key_prefix,
                start_time,
                num_samples=output.num_samples,
                num_steps=math.ceil(output.num_samples / total_batch_size),
            )
        )

        self.control = self.callback_handler.on_predict(self.args, self.state, self.control, output.metrics)
        self._memory_tracker.stop_and_update_metrics(output.metrics)

        return PredictionOutput(predictions=output.predictions, label_ids=output.label_ids, metrics=output.metrics)

    def evaluation_loop(
        self,
        dataloader: DataLoader,
        description: str,
        prediction_loss_only: Optional[bool] = None,
        ignore_keys: Optional[List[str]] = None,
        metric_key_prefix: str = "eval",
    ) -> EvalLoopOutput:
        
        args = self.args

        prediction_loss_only = prediction_loss_only if prediction_loss_only is not None else args.prediction_loss_only

        
        if self.is_deepspeed_enabled and self.deepspeed is None:
            _, _ = deepspeed_init(self, num_training_steps=0, inference=True)

        model = self._wrap_model(self.model, training=False, dataloader=dataloader)

        if len(self.accelerator._models) == 0 and model is self.model:
            model = (
                self.accelerator.prepare(model)
                if self.is_deepspeed_enabled
                else self.accelerator.prepare_model(model, evaluation_mode=True)
            )

            if self.is_fsdp_enabled:
                self.model = model

            
            if model is not self.model:
                self.model_wrapped = model

            
            if self.is_deepspeed_enabled:
                self.deepspeed = self.model_wrapped

        
        
        if not self.is_in_train:
            if args.fp16_full_eval:
                model = model.to(dtype=torch.float16, device=args.device)
            elif args.bf16_full_eval:
                model = model.to(dtype=torch.bfloat16, device=args.device)

        batch_size = self.args.eval_batch_size

        logger.info(f"***** Running {description} *****")
        if has_length(dataloader):
            logger.info(f"  Num examples = {self.num_examples(dataloader)}")
        else:
            logger.info("  Num examples: Unknown")
        logger.info(f"  Batch size = {batch_size}")

        model.eval()

        self.callback_handler.eval_dataloader = dataloader
        
        eval_dataset = getattr(dataloader, "dataset", None)

        if args.past_index >= 0:
            self._past = None

        
        
        losses_host = None
        preds_host = None
        labels_host = None
        inputs_host = None

        
        all_losses = None
        all_preds = None
        all_labels = None
        all_inputs = None
        

        observed_num_examples = 0
        
        for step, inputs in enumerate(dataloader):
            
            observed_batch_size = find_batch_size(inputs)
            if observed_batch_size is not None:
                observed_num_examples += observed_batch_size
                
                if batch_size is None:
                    batch_size = observed_batch_size

            
            loss, logits, labels = self.prediction_step(model, inputs, prediction_loss_only, ignore_keys=ignore_keys)
            main_input_name = getattr(self.model, "main_input_name", "input_ids")
            inputs_decode = self._prepare_input(inputs[main_input_name]) if args.include_inputs_for_metrics else None

            if is_torch_tpu_available():
                xm.mark_step()

            
            if loss is not None:
                losses = self.gather_function((loss.repeat(batch_size)))
                losses_host = losses if losses_host is None else nested_concat(losses_host, losses, padding_index=-100)
            if labels is not None:
                labels = self.accelerator.pad_across_processes(labels, dim=1, pad_index=-100)
            if inputs_decode is not None:
                inputs_decode = self.accelerator.pad_across_processes(inputs_decode, dim=1, pad_index=-100)
                inputs_decode = self.gather_function((inputs_decode))
                inputs_host = (
                    inputs_decode
                    if inputs_host is None
                    else nested_concat(inputs_host, inputs_decode, padding_index=-100)
                )
            if logits is not None:
                logits = self.accelerator.pad_across_processes(logits, dim=1, pad_index=-100)
                if self.preprocess_logits_for_metrics is not None:
                    logits = self.preprocess_logits_for_metrics(logits, labels)
                logits = self.gather_function((logits))
                preds_host = logits if preds_host is None else nested_concat(preds_host, logits, padding_index=-100)

            if labels is not None:
                labels = self.gather_function((labels))
                labels_host = labels if labels_host is None else nested_concat(labels_host, labels, padding_index=-100)

            self.control = self.callback_handler.on_prediction_step(args, self.state, self.control)

            
            if args.eval_accumulation_steps is not None and (step + 1) % args.eval_accumulation_steps == 0:
                if losses_host is not None:
                    losses = nested_numpify(losses_host)
                    all_losses = losses if all_losses is None else np.concatenate((all_losses, losses), axis=0)
                if preds_host is not None:
                    logits = nested_numpify(preds_host)
                    all_preds = logits if all_preds is None else nested_concat(all_preds, logits, padding_index=-100)
                if inputs_host is not None:
                    inputs_decode = nested_numpify(inputs_host)
                    all_inputs = (
                        inputs_decode
                        if all_inputs is None
                        else nested_concat(all_inputs, inputs_decode, padding_index=-100)
                    )
                if labels_host is not None:
                    labels = nested_numpify(labels_host)
                    all_labels = (
                        labels if all_labels is None else nested_concat(all_labels, labels, padding_index=-100)
                    )

                
                losses_host, preds_host, inputs_host, labels_host = None, None, None, None

        
        self.gather_function = self.accelerator.gather_for_metrics
        if args.past_index and hasattr(self, "_past"):
            
            delattr(self, "_past")

        
        if losses_host is not None:
            losses = nested_numpify(losses_host)
            all_losses = losses if all_losses is None else np.concatenate((all_losses, losses), axis=0)
        if preds_host is not None:
            logits = nested_numpify(preds_host)
            all_preds = logits if all_preds is None else nested_concat(all_preds, logits, padding_index=-100)
        if inputs_host is not None:
            inputs_decode = nested_numpify(inputs_host)
            all_inputs = (
                inputs_decode if all_inputs is None else nested_concat(all_inputs, inputs_decode, padding_index=-100)
            )
        if labels_host is not None:
            labels = nested_numpify(labels_host)
            all_labels = labels if all_labels is None else nested_concat(all_labels, labels, padding_index=-100)

        
        if has_length(eval_dataset):
            num_samples = len(eval_dataset)
        
        
        elif isinstance(eval_dataset, IterableDatasetShard) and getattr(eval_dataset, "num_examples", 0) > 0:
            num_samples = eval_dataset.num_examples
        else:
            if has_length(dataloader):
                num_samples = self.num_examples(dataloader)
            else:  
                num_samples = observed_num_examples
        if num_samples == 0 and observed_num_examples > 0:
            num_samples = observed_num_examples

        
        if self.compute_metrics is not None and all_preds is not None and all_labels is not None:
            if args.include_inputs_for_metrics:
                metrics = self.compute_metrics(
                    EvalPrediction(predictions=all_preds, label_ids=all_labels, inputs=all_inputs)
                )
            else:
                metrics = self.compute_metrics(EvalPrediction(predictions=all_preds, label_ids=all_labels))
        else:
            metrics = {}

        
        metrics = denumpify_detensorize(metrics)

        if all_losses is not None:
            metrics[f"{metric_key_prefix}_loss"] = all_losses.mean().item()
        if hasattr(self, "jit_compilation_time"):
            metrics[f"{metric_key_prefix}_jit_compilation_time"] = self.jit_compilation_time

        
        for key in list(metrics.keys()):
            if not key.startswith(f"{metric_key_prefix}_"):
                metrics[f"{metric_key_prefix}_{key}"] = metrics.pop(key)

        return EvalLoopOutput(predictions=all_preds, label_ids=all_labels, metrics=metrics, num_samples=num_samples)

    def _nested_gather(self, tensors, name=None):
        
        if tensors is None:
            return
        if is_torch_tpu_available():
            if name is None:
                name = "nested_gather"
            tensors = nested_xla_mesh_reduce(tensors, name)
        elif is_sagemaker_mp_enabled():
            tensors = smp_gather(tensors)
        elif (self.args.distributed_state is not None and self.args.distributed_state.distributed_type != "NO") or (
            self.args.distributed_state is None and self.args.local_rank != -1
        ):
            tensors = distributed_concat(tensors)
        return tensors

    def prediction_step(
        self,
        model: nn.Module,
        inputs: Dict[str, Union[torch.Tensor, Any]],
        prediction_loss_only: bool,
        ignore_keys: Optional[List[str]] = None,
    ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]:
        
        has_labels = False if len(self.label_names) == 0 else all(inputs.get(k) is not None for k in self.label_names)
        
        
        
        return_loss = inputs.get("return_loss", None)
        if return_loss is None:
            return_loss = self.can_return_loss
        loss_without_labels = True if len(self.label_names) == 0 and return_loss else False

        inputs = self._prepare_inputs(inputs)
        if ignore_keys is None:
            if hasattr(self.model, "config"):
                ignore_keys = getattr(self.model.config, "keys_to_ignore_at_inference", [])
            else:
                ignore_keys = []

        
        if has_labels or loss_without_labels:
            labels = nested_detach(tuple(inputs.get(name) for name in self.label_names))
            if len(labels) == 1:
                labels = labels[0]
        else:
            labels = None

        with torch.no_grad():
            if is_sagemaker_mp_enabled():
                raw_outputs = smp_forward_only(model, inputs)
                if has_labels or loss_without_labels:
                    if isinstance(raw_outputs, dict):
                        loss_mb = raw_outputs["loss"]
                        logits_mb = tuple(v for k, v in raw_outputs.items() if k not in ignore_keys + ["loss"])
                    else:
                        loss_mb = raw_outputs[0]
                        logits_mb = raw_outputs[1:]

                    loss = loss_mb.reduce_mean().detach().cpu()
                    logits = smp_nested_concat(logits_mb)
                else:
                    loss = None
                    if isinstance(raw_outputs, dict):
                        logits_mb = tuple(v for k, v in raw_outputs.items() if k not in ignore_keys)
                    else:
                        logits_mb = raw_outputs
                    logits = smp_nested_concat(logits_mb)
            else:
                if has_labels or loss_without_labels:
                    with self.compute_loss_context_manager():
                        loss, outputs = self.compute_loss(model, inputs, return_outputs=True)
                    loss = loss.mean().detach()

                    if isinstance(outputs, dict):
                        logits = tuple(v for k, v in outputs.items() if k not in ignore_keys + ["loss"])
                    else:
                        logits = outputs[1:]
                else:
                    loss = None
                    with self.compute_loss_context_manager():
                        outputs = model(**inputs)
                    if isinstance(outputs, dict):
                        logits = tuple(v for k, v in outputs.items() if k not in ignore_keys)
                    else:
                        logits = outputs
                    
                    if self.args.past_index >= 0:
                        self._past = outputs[self.args.past_index - 1]

        if prediction_loss_only:
            return (loss, None, None)

        logits = nested_detach(logits)
        if len(logits) == 1:
            logits = logits[0]

        return (loss, logits, labels)

    def floating_point_ops(self, inputs: Dict[str, Union[torch.Tensor, Any]]):
        
        if hasattr(self.model, "floating_point_ops"):
            return self.model.floating_point_ops(inputs)
        else:
            return 0

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

        if self.args.hub_model_id is None:
            repo_name = Path(self.args.output_dir).absolute().name
        else:
            repo_name = self.args.hub_model_id

        repo_url = create_repo(repo_name, token=self.args.hub_token, private=self.args.hub_private_repo, exist_ok=True)
        self.hub_model_id = repo_url.repo_id
        self.push_in_progress = None

    def create_model_card(
        self,
        language: Optional[str] = None,
        license: Optional[str] = None,
        tags: Union[str, List[str], None] = None,
        model_name: Optional[str] = None,
        finetuned_from: Optional[str] = None,
        tasks: Union[str, List[str], None] = None,
        dataset_tags: Union[str, List[str], None] = None,
        dataset: Union[str, List[str], None] = None,
        dataset_args: Union[str, List[str], None] = None,
    ):
        
        if not self.is_world_process_zero():
            return

        training_summary = TrainingSummary.from_trainer(
            self,
            language=language,
            license=license,
            tags=tags,
            model_name=model_name,
            finetuned_from=finetuned_from,
            tasks=tasks,
            dataset_tags=dataset_tags,
            dataset=dataset,
            dataset_args=dataset_args,
        )
        model_card = training_summary.to_model_card()
        with open(os.path.join(self.args.output_dir, "README.md"), "w") as f:
            f.write(model_card)

    def _push_from_checkpoint(self, checkpoint_folder):
        
        if not self.is_world_process_zero() or self.args.hub_strategy == HubStrategy.END:
            return
        
        if not self.args.hub_always_push and self.push_in_progress is not None and not self.push_in_progress.is_done():
            return

        output_dir = self.args.output_dir
        
        modeling_files = [CONFIG_NAME, WEIGHTS_NAME, SAFE_WEIGHTS_NAME]
        if is_peft_available():
            modeling_files.extend([ADAPTER_CONFIG_NAME, ADAPTER_WEIGHTS_NAME, ADAPTER_SAFE_WEIGHTS_NAME])
        for modeling_file in modeling_files:
            if os.path.isfile(os.path.join(checkpoint_folder, modeling_file)):
                shutil.copy(os.path.join(checkpoint_folder, modeling_file), os.path.join(output_dir, modeling_file))
        
        if self.tokenizer is not None:
            self.tokenizer.save_pretrained(output_dir)
        
        torch.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME))

        if self.args.save_strategy == IntervalStrategy.STEPS:
            commit_message = f"Training in progress, step {self.state.global_step}"
        else:
            commit_message = f"Training in progress, epoch {int(self.state.epoch)}"

        model_push_job = upload_folder(
            repo_id=self.hub_model_id,
            folder_path=output_dir,
            commit_message=commit_message,
            token=self.args.hub_token,
            run_as_future=True,
            ignore_patterns=["_*", f"{PREFIX_CHECKPOINT_DIR}-*"],
        )

        push_jobs = [model_push_job]

        if self.args.hub_strategy in [HubStrategy.CHECKPOINT, HubStrategy.ALL_CHECKPOINTS]:
            path_in_repo = (
                "last-checkpoint" if self.args.hub_strategy == HubStrategy.CHECKPOINT else Path(checkpoint_folder).name
            )
            checkpoint_push = upload_folder(
                repo_id=self.hub_model_id,
                folder_path=checkpoint_folder,
                path_in_repo=path_in_repo,
                commit_message=commit_message + ", checkpoint",
                token=self.args.hub_token,
                run_as_future=True,
            )
            push_jobs.append(checkpoint_push)

        if self.push_in_progress is None or self.push_in_progress.is_done():
            self.push_in_progress = PushInProgress(push_jobs)
        else:
            self.push_in_progress.jobs.extend(push_jobs)

    def _finish_current_push(self):
        if not hasattr(self, "push_in_progress"):
            return
        if self.push_in_progress is not None and not self.push_in_progress.is_done():
            logger.info("Waiting for the current checkpoint push to be finished, this might take a couple of minutes.")
            self.push_in_progress.wait_until_done()

    def push_to_hub(self, commit_message: Optional[str] = "End of training", blocking: bool = True, **kwargs) -> str:
        
        model_name = kwargs.pop("model_name", None)
        if model_name is None and self.args.should_save:
            if self.args.hub_model_id is None:
                model_name = Path(self.args.output_dir).name
            else:
                model_name = self.args.hub_model_id.split("/")[-1]

        
        if self.hub_model_id is None:
            self.init_hf_repo()

        
        
        self.save_model(_internal_call=True)

        
        if not self.is_world_process_zero():
            return

        self.create_model_card(model_name=model_name, **kwargs)

        
        self._finish_current_push()
        return upload_folder(
            repo_id=self.hub_model_id,
            folder_path=self.args.output_dir,
            commit_message=commit_message,
            token=self.args.hub_token,
            run_as_future=not blocking,
            ignore_patterns=["_*", f"{PREFIX_CHECKPOINT_DIR}-*"],
        )

    
    
    

    def prediction_loop(
        self,
        dataloader: DataLoader,
        description: str,
        prediction_loss_only: Optional[bool] = None,
        ignore_keys: Optional[List[str]] = None,
        metric_key_prefix: str = "eval",
    ) -> EvalLoopOutput:
        
        args = self.args

        if not has_length(dataloader):
            raise ValueError("dataloader must implement a working __len__")

        prediction_loss_only = prediction_loss_only if prediction_loss_only is not None else args.prediction_loss_only

        
        if self.is_deepspeed_enabled and self.deepspeed is None:
            _, _ = deepspeed_init(self, num_training_steps=0, inference=True)

        model = self._wrap_model(self.model, training=False, dataloader=dataloader)

        if len(self.accelerator._models) == 0 and model is self.model:
            model = (
                self.accelerator.prepare(model)
                if self.is_deepspeed_enabled
                else self.accelerator.prepare_model(model, evaluation_mode=True)
            )

            if self.is_fsdp_enabled:
                self.model = model

            
            if model is not self.model:
                self.model_wrapped = model

            
            if self.is_deepspeed_enabled:
                self.deepspeed = self.model_wrapped

        
        
        if not self.is_in_train:
            if args.fp16_full_eval:
                model = model.to(dtype=torch.float16, device=args.device)
            elif args.bf16_full_eval:
                model = model.to(dtype=torch.bfloat16, device=args.device)

        batch_size = dataloader.batch_size
        num_examples = self.num_examples(dataloader)
        logger.info(f"***** Running {description} *****")
        logger.info(f"  Num examples = {num_examples}")
        logger.info(f"  Batch size = {batch_size}")
        losses_host: torch.Tensor = None
        preds_host: Union[torch.Tensor, List[torch.Tensor]] = None
        labels_host: Union[torch.Tensor, List[torch.Tensor]] = None
        inputs_host: Union[torch.Tensor, List[torch.Tensor]] = None

        world_size = max(1, args.world_size)

        eval_losses_gatherer = DistributedTensorGatherer(world_size, num_examples, make_multiple_of=batch_size)
        if not prediction_loss_only:
            
            
            make_multiple_of = None
            if hasattr(dataloader, "sampler") and isinstance(dataloader.sampler, SequentialDistributedSampler):
                make_multiple_of = dataloader.sampler.batch_size
            preds_gatherer = DistributedTensorGatherer(world_size, num_examples, make_multiple_of=make_multiple_of)
            labels_gatherer = DistributedTensorGatherer(world_size, num_examples, make_multiple_of=make_multiple_of)
            inputs_gatherer = DistributedTensorGatherer(world_size, num_examples, make_multiple_of=make_multiple_of)

        model.eval()

        if args.past_index >= 0:
            self._past = None

        self.callback_handler.eval_dataloader = dataloader

        for step, inputs in enumerate(dataloader):
            loss, logits, labels = self.prediction_step(model, inputs, prediction_loss_only, ignore_keys=ignore_keys)
            main_input_name = getattr(self.model, "main_input_name", "input_ids")
            inputs_decode = self._prepare_input(inputs[main_input_name]) if args.include_inputs_for_metrics else None

            if loss is not None:
                losses = loss.repeat(batch_size)
                losses_host = losses if losses_host is None else torch.cat((losses_host, losses), dim=0)
            if logits is not None:
                preds_host = logits if preds_host is None else nested_concat(preds_host, logits, padding_index=-100)
            if labels is not None:
                labels_host = labels if labels_host is None else nested_concat(labels_host, labels, padding_index=-100)
            if inputs_decode is not None:
                inputs_host = (
                    inputs_decode
                    if inputs_host is None
                    else nested_concat(inputs_host, inputs_decode, padding_index=-100)
                )
            self.control = self.callback_handler.on_prediction_step(args, self.state, self.control)

            
            if args.eval_accumulation_steps is not None and (step + 1) % args.eval_accumulation_steps == 0:
                eval_losses_gatherer.add_arrays(self._gather_and_numpify(losses_host, "eval_losses"))
                if not prediction_loss_only:
                    preds_gatherer.add_arrays(self._gather_and_numpify(preds_host, "eval_preds"))
                    labels_gatherer.add_arrays(self._gather_and_numpify(labels_host, "eval_label_ids"))
                    inputs_gatherer.add_arrays(self._gather_and_numpify(inputs_host, "eval_inputs_ids"))

                
                losses_host, preds_host, labels_host, inputs_host = None, None, None, None

        if args.past_index and hasattr(self, "_past"):
            
            delattr(self, "_past")

        
        eval_losses_gatherer.add_arrays(self._gather_and_numpify(losses_host, "eval_losses"))
        if not prediction_loss_only:
            preds_gatherer.add_arrays(self._gather_and_numpify(preds_host, "eval_preds"))
            labels_gatherer.add_arrays(self._gather_and_numpify(labels_host, "eval_label_ids"))
            inputs_gatherer.add_arrays(self._gather_and_numpify(inputs_host, "eval_inputs_ids"))

        eval_loss = eval_losses_gatherer.finalize()
        preds = preds_gatherer.finalize() if not prediction_loss_only else None
        label_ids = labels_gatherer.finalize() if not prediction_loss_only else None
        inputs_ids = inputs_gatherer.finalize() if not prediction_loss_only else None

        if self.compute_metrics is not None and preds is not None and label_ids is not None:
            if args.include_inputs_for_metrics:
                metrics = self.compute_metrics(
                    EvalPrediction(predictions=preds, label_ids=label_ids, inputs=inputs_ids)
                )
            else:
                metrics = self.compute_metrics(EvalPrediction(predictions=preds, label_ids=label_ids))
        else:
            metrics = {}

        
        metrics = denumpify_detensorize(metrics)

        if eval_loss is not None:
            metrics[f"{metric_key_prefix}_loss"] = eval_loss.mean().item()

        
        for key in list(metrics.keys()):
            if not key.startswith(f"{metric_key_prefix}_"):
                metrics[f"{metric_key_prefix}_{key}"] = metrics.pop(key)

        return EvalLoopOutput(predictions=preds, label_ids=label_ids, metrics=metrics, num_samples=num_examples)

    def _gather_and_numpify(self, tensors, name):
        
        if tensors is None:
            return
        if is_torch_tpu_available():
            tensors = nested_xla_mesh_reduce(tensors, name)
        elif is_sagemaker_mp_enabled():
            tensors = smp_gather(tensors)
        elif self.args.parallel_mode == ParallelMode.DISTRIBUTED:
            tensors = distributed_concat(tensors)

        return nested_numpify(tensors)

    def _add_sm_patterns_to_gitignore(self) -> None:
        
        
        if not self.is_world_process_zero():
            return

        patterns = ["*.sagemaker-uploading", "*.sagemaker-uploaded"]

        
        if os.path.exists(os.path.join(self.repo.local_dir, ".gitignore")):
            with open(os.path.join(self.repo.local_dir, ".gitignore"), "r") as f:
                current_content = f.read()
        else:
            current_content = ""

        
        content = current_content
        for pattern in patterns:
            if pattern not in content:
                if content.endswith("\n"):
                    content += pattern
                else:
                    content += f"\n{pattern}"

        
        if content != current_content:
            with open(os.path.join(self.repo.local_dir, ".gitignore"), "w") as f:
                logger.debug(f"Writing .gitignore file. Content: {content}")
                f.write(content)

        self.repo.git_add(".gitignore")

        
        time.sleep(0.5)

        if not self.repo.is_repo_clean():
            self.repo.git_commit("Add *.sagemaker patterns to .gitignore.")
            self.repo.git_push()

    def create_accelerator_and_postprocess(self):
        grad_acc_kwargs = {"num_steps": self.args.gradient_accumulation_steps}
        grad_acc_kwargs["sync_with_dataloader"] = False
        gradient_accumulation_plugin = GradientAccumulationPlugin(**grad_acc_kwargs)

        
        self.accelerator = Accelerator(
            dispatch_batches=self.args.dispatch_batches,
            split_batches=self.args.split_batches,
            deepspeed_plugin=self.args.deepspeed_plugin,
            gradient_accumulation_plugin=gradient_accumulation_plugin,
        )
        
        self.gather_function = self.accelerator.gather_for_metrics

        
        self.is_deepspeed_enabled = getattr(self.accelerator.state, "deepspeed_plugin", None) is not None
        self.is_fsdp_enabled = getattr(self.accelerator.state, "fsdp_plugin", None) is not None

        
        if self.is_fsdp_enabled:
            fsdp_plugin = self.accelerator.state.fsdp_plugin
            fsdp_plugin.limit_all_gathers = self.args.fsdp_config.get(
                "limit_all_gathers", fsdp_plugin.limit_all_gathers
            )
            if is_accelerate_available("0.23.0"):
                fsdp_plugin.activation_checkpointing = self.args.fsdp_config.get(
                    "activation_checkpointing", fsdp_plugin.activation_checkpointing
                )
                if fsdp_plugin.activation_checkpointing and self.args.gradient_checkpointing:
                    raise ValueError(
                        "The activation_checkpointing in FSDP config and the gradient_checkpointing in training arg "
                        "can't be set to True simultaneously. Please use FSDP's activation_checkpointing logic "
                        "when using FSDP."
                    )

        if self.is_deepspeed_enabled:
            if getattr(self.args, "hf_deepspeed_config", None) is None:
                from transformers.integrations.deepspeed import HfTrainerDeepSpeedConfig

                ds_plugin = self.accelerator.state.deepspeed_plugin

                ds_plugin.hf_ds_config = HfTrainerDeepSpeedConfig(ds_plugin.hf_ds_config.config)
                ds_plugin.deepspeed_config = ds_plugin.hf_ds_config.config
                ds_plugin.hf_ds_config.trainer_config_process(self.args)