# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.## 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.importglobimportosimportsignalimportsubprocessimportsysimporttimeimportwarningsfromcollectionsimportdefaultdictfromdataclassesimportdataclass,fieldfromdatetimeimporttimedeltafrompathlibimportPathfromshutilimportcopy,movefromtypingimportAny,Collection,Dict,List,Optional,Tuple,Unionimportlightning.pytorchimporttorchfromhydra.core.hydra_configimportHydraConfigfromhydra.utilsimportget_original_cwdfromlightning.pytorch.callbacksimportCallback,ModelCheckpointfromlightning.pytorch.callbacks.early_stoppingimportEarlyStoppingfromlightning.pytorch.callbacks.timerimportInterval,Timerfromlightning.pytorch.loggersimportMLFlowLogger,NeptuneLogger,TensorBoardLogger,WandbLoggerfromlightning.pytorch.loopsimport_TrainingEpochLoopfromlightning.pytorch.strategies.ddpimportDDPStrategyfromlightning.pytorch.trainer.connectors.checkpoint_connectorimport_CheckpointConnectorfromomegaconfimportDictConfig,OmegaConf,open_dictfromnemo.collections.common.callbacksimportEMAfromnemo.collections.common.callbacks.ipl_epoch_stopperimportIPLEpochStopperfromnemo.constantsimportNEMO_ENV_VARNAME_TESTING,NEMO_ENV_VARNAME_VERSIONfromnemo.utilsimportlogging,timersfromnemo.utils.app_stateimportAppStatefromnemo.utils.callbacksimportNeMoModelCheckpoint,PreemptionCallbackfromnemo.utils.env_var_parsingimportget_envboolfromnemo.utils.exceptionsimportNeMoBaseExceptionfromnemo.utils.get_rankimportis_global_rank_zerofromnemo.utils.import_utilsimportsafe_import_fromfromnemo.utils.lightning_logger_patchimportadd_filehandlers_to_pl_loggerfromnemo.utils.loggersimportClearMLLogger,ClearMLParams,DLLogger,DLLoggerParams,MLFlowParamsfromnemo.utils.mcore_loggerimportadd_handlers_to_mcore_loggerfromnemo.utils.model_utilsimportuninject_model_parallel_rankfromnemo.utils.msc_utilsimportimport_multistorageclient,is_multistorageclient_urlget_current_global_batch_size,HAVE_MCORE_MBATCH_CALCULATOR=safe_import_from("megatron.core.num_microbatches_calculator","get_current_global_batch_size")try:# `ptl_resiliency` is included in `gwe_resiliency_pkg` packagefromptl_resiliencyimportStragglerDetectionCallbackHAVE_STRAGGLER_DET=Trueexcept(ImportError,ModuleNotFoundError):HAVE_STRAGGLER_DET=Falsetry:fromptl_resiliencyimportFaultToleranceCallbackHAVE_FT=Trueexcept(ImportError,ModuleNotFoundError):HAVE_FT=FalseclassNotFoundError(NeMoBaseException):"""Raised when a file or folder is not found"""classLoggerMisconfigurationError(NeMoBaseException):"""Raised when a mismatch between trainer.logger and exp_manager occurs"""def__init__(self,message):message=(message+" You can disable lighning's trainer from creating a logger by passing logger=False to its constructor.")super().__init__(message)classCheckpointMisconfigurationError(NeMoBaseException):"""Raised when a mismatch between trainer.callbacks and exp_manager occurs"""@dataclassclassEarlyStoppingParams:"""EarlyStoppingParams POD"""# The metric that early stopping should consider.monitor:str="val_loss"# inform early stopping whether to look for increase or decrease in monitor.mode:str="min"min_delta:float=0.001# smallest change to consider as improvement.# how many (continuous) validation cycles to wait with no improvement and stopping training.patience:int=10verbose:bool=Truestrict:bool=Truecheck_finite:bool=Truestopping_threshold:Optional[float]=Nonedivergence_threshold:Optional[float]=Nonecheck_on_train_epoch_end:Optional[bool]=Nonelog_rank_zero_only:bool=False@dataclassclassIPLEpochStopperParams:""" Parameters for the IPLEpochStopper callback used in iterative pseudo-label training. This is part of the TopIPL pipeline, a semi-supervised training method for ASR that uses iterative pseudo-labeling (IPL) — periodically stopping training to generate pseudo-labels for unlabeled data and fine-tuning the model on them. For more details, see: 🔗 Top-IPL: Top-N Pseudo-Label Averaging for Iterative ASR Training https://arxiv.org/abs/2506.07659 Attributes: enable_stop (bool): If True, enables the stopping behavior in the callback. stop_every_n_epochs (int): Specifies how many epochs to train before stopping. """# Flag that allows stoppingenable_stop:bool=Truestop_every_n_epochs:int=1@dataclassclassCallbackParams:"""CallbackParams POD"""filepath:Optional[str]=None# Deprecated# If None, exp_manager will attempt to handle the filepathdirpath:Optional[str]=None# If None, exp_manager will attempt to handle the filepathfilename:Optional[str]=Nonemonitor:Optional[str]="val_loss"verbose:Optional[bool]=Truesave_last:Optional[bool]=Truesave_top_k:Optional[int]=3save_weights_only:Optional[bool]=Falsemode:Optional[str]="min"auto_insert_metric_name:bool=Trueevery_n_epochs:Optional[int]=1every_n_train_steps:Optional[int]=Nonetrain_time_interval:Optional[Any]=None# If None, exp_manager will attempt to handle the filepathprefix:Optional[str]=Nonepostfix:str=".nemo"save_best_model:bool=Falsealways_save_nemo:bool=False# Whether to automatically save .nemo file durin on_train_end hooksave_nemo_on_train_end:Optional[bool]=True# tensor parallel size * pipeline parallel sizemodel_parallel_size:Optional[int]=None# Save after training, not after validationsave_on_train_epoch_end:Optional[bool]=Falseasync_save:Optional[bool]=False# save the checkpoint asynchronously# a number of last checkpoints to be saved with optimizer statessave_last_n_optim_states:Optional[int]=-1@dataclassclassStepTimingParams:"""StepTimingParams POD"""reduction:Optional[str]="mean"# if True torch.cuda.synchronize() is called on start/stopsync_cuda:Optional[bool]=False# if positive, defines the size of a sliding window for computing meanbuffer_size:Optional[int]=1@dataclassclassEMAParams:"""EMAParams POD"""enable:Optional[bool]=Falsedecay:Optional[float]=0.999cpu_offload:Optional[bool]=Falsevalidate_original_weights:Optional[bool]=Falseevery_n_steps:int=1@dataclassclassStragglerDetectionParams:"""StragglerDetectionParams POD"""report_time_interval:float=300calc_relative_gpu_perf:bool=Truecalc_individual_gpu_perf:bool=Truenum_gpu_perf_scores_to_log:int=5gpu_relative_perf_threshold:float=0.7gpu_individual_perf_threshold:float=0.7stop_if_detected:bool=False@dataclassclassFaultToleranceParams:"""FaultToleranceParams POD"""# NOTE: This config section is also read by the launcher.# NOTE: Default values should match fault_tolerance.FaultToleranceConfig.workload_check_interval:float=5.0initial_rank_heartbeat_timeout:Optional[float]=60.0*60.0rank_heartbeat_timeout:Optional[float]=45.0*60.0calculate_timeouts:bool=Truesafety_factor:float=5.0rank_termination_signal:signal.Signals=signal.SIGKILLifos.name!='nt'elsesignal.SIGTERMlog_level:str='INFO'max_rank_restarts:int=0max_subsequent_job_failures:int=0additional_ft_launcher_args:str=''simulated_fault:Optional[Any]=None
@dataclassclassExpManagerConfig:"""Experiment Manager config for validation of passed arguments."""# Log dir creation parametersexplicit_log_dir:Optional[str]=Noneexp_dir:Optional[str]=Nonename:Optional[str]=Noneversion:Optional[str]=Noneuse_datetime_version:Optional[bool]=Trueresume_if_exists:Optional[bool]=Falseresume_past_end:Optional[bool]=Falseresume_ignore_no_checkpoint:Optional[bool]=Falseresume_from_checkpoint:Optional[str]=None# Logging parameterscreate_tensorboard_logger:Optional[bool]=Truesummary_writer_kwargs:Optional[Dict[Any,Any]]=Nonecreate_wandb_logger:Optional[bool]=Falsewandb_logger_kwargs:Optional[Dict[Any,Any]]=Nonecreate_mlflow_logger:Optional[bool]=Falsemlflow_logger_kwargs:Optional[MLFlowParams]=field(default_factory=lambda:MLFlowParams())create_dllogger_logger:Optional[bool]=Falsedllogger_logger_kwargs:Optional[DLLoggerParams]=field(default_factory=lambda:DLLoggerParams())create_clearml_logger:Optional[bool]=Falseclearml_logger_kwargs:Optional[ClearMLParams]=field(default_factory=lambda:ClearMLParams())create_neptune_logger:Optional[bool]=Falseneptune_logger_kwargs:Optional[Dict[Any,Any]]=None# Checkpointing parameterscreate_checkpoint_callback:Optional[bool]=Truecheckpoint_callback_params:Optional[CallbackParams]=field(default_factory=lambda:CallbackParams())create_early_stopping_callback:Optional[bool]=Falsecreate_ipl_epoch_stopper_callback:Optional[bool]=Falseearly_stopping_callback_params:Optional[EarlyStoppingParams]=field(default_factory=lambda:EarlyStoppingParams())ipl_epoch_stopper_callback_params:Optional[IPLEpochStopperParams]=field(default_factory=lambda:IPLEpochStopperParams())create_preemption_callback:Optional[bool]=True# Additional exp_manager argumentsfiles_to_copy:Optional[List[str]]=None# logs timing of train/val/test stepslog_step_timing:Optional[bool]=True# log step time with nemo logger instead of lightning logger to avoid lightning logger overheadlog_delta_step_timing:Optional[bool]=Falsestep_timing_kwargs:Optional[StepTimingParams]=field(default_factory=lambda:StepTimingParams())# disable initial validation when resuming from a checkpoint saved during validationdisable_validation_on_resume:Optional[bool]=Trueema:Optional[EMAParams]=field(default_factory=lambda:EMAParams())# Wall clock time limitmax_time_per_run:Optional[str]=None# time to sleep non 0 ranks during initializationseconds_to_sleep:float=5# Straggler detectioncreate_straggler_detection_callback:Optional[bool]=Falsestraggler_detection_params:Optional[StragglerDetectionParams]=field(default_factory=StragglerDetectionParams)# Fault tolrancecreate_fault_tolerance_callback:Optional[bool]=Falsefault_tolerance:Optional[FaultToleranceParams]=field(default_factory=FaultToleranceParams)# logs TFLOPs per sec per gpulog_tflops_per_sec_per_gpu:Optional[bool]=True
classTimingCallback(Callback):""" Logs execution time of train/val/test steps """def__init__(self,log_tokens_per_sec:bool=False,timer_kwargs={}):"""init for TimitCallback Args: log_tokens_per_sec (bool, optional): _description_. Defaults to False. timer_kwargs (dict, optional): _description_. Defaults to {}. """self.log_tokens_per_sec=log_tokens_per_secself.timer=timers.NamedTimer(**timer_kwargs)def_on_batch_start(self,name):"""Setup the timer Args: name (_type_): name of timer """# reset only if we do not return mean of a sliding windowifself.timer.buffer_size<=0:self.timer.reset(name)ifself.timer.is_active(name):logging.warning(f"Timer `{name}` was not correctly stopped, suggesting a ""possible issue. The timer will be reset for now.")self.timer.reset(name)self.timer.start(name)def_on_batch_end(self,name,pl_module):"""end of the callback log Args: name (_type_): _description_ pl_module (_type_): _description_ """try:self.timer.stop(name)exceptRuntimeError:logging.warning(f"Missing timer '{name}' in exp_manager's _on_batch_end callback - not logging.")return# Set the `batch_size=1` as WAR for `dataloader_iter`, which is not used for any metricpl_module.log(name+' in s',torch.as_tensor(self.timer[name]),on_step=True,on_epoch=False,batch_size=1,prog_bar=(name=="train_step_timing"),)defon_train_batch_start(self,trainer,pl_module,batch,batch_idx):"""wrapper Args: trainer (_type_): _description_ pl_module (_type_): _description_ batch (_type_): _description_ batch_idx (_type_): _description_ """self._on_batch_start("train_step_timing")defon_train_batch_end(self,trainer,pl_module,outputs,batch,batch_idx):"""wrapper Args: trainer (_type_): _description_ pl_module (_type_): _description_ outputs (_type_): _description_ batch (_type_): _description_ batch_idx (_type_): _description_ """self._on_batch_end("train_step_timing",pl_module)ifself.log_tokens_per_sec:if"text"inbatch:batch['tokens']=batch['text']tokens_per_gpu=((get_current_global_batch_size()//trainer.accumulate_grad_batches)*batch["tokens"].shape[1]/torch.distributed.get_world_size())pl_module.log("tokens_per_sec_per_gpu",tokens_per_gpu/(torch.as_tensor(self.timer["train_step_timing"])),on_step=True,on_epoch=False,batch_size=1,prog_bar=True,)defon_validation_batch_start(self,trainer,pl_module,batch,batch_idx,dataloader_idx=0):"""on_validation_batch_start"""self._on_batch_start("validation_step_timing")defon_validation_batch_end(self,trainer,pl_module,outputs,batch,batch_idx,dataloader_idx=0):"""on_validation_batch_end"""self._on_batch_end("validation_step_timing",pl_module)defon_test_batch_start(self,trainer,pl_module,batch,batch_idx,dataloader_idx=0):"""on_test_batch_start"""self._on_batch_start("test_step_timing")defon_test_batch_end(self,trainer,pl_module,outputs,batch,batch_idx,dataloader_idx=0):"""on_test_batch_end"""self._on_batch_end("test_step_timing",pl_module)defon_before_backward(self,trainer,pl_module,loss):"""on_before_backward"""self._on_batch_start("train_backward_timing")defon_after_backward(self,trainer,pl_module):"""on_after_backward"""self._on_batch_end("train_backward_timing",pl_module)classDeltaTimingCallback(Callback):""" Logs execution time of train/val/test steps using nemo logger. Calculates time from previous batch end to current batch end. This ensures accuracy. Note: step time will only be printed in stdout. If you have initialized loggers like TensorBoard, WandB, etc, step time will not be recorded there. Use this callback instead of 'TimingCallback' to avoid logging overhead with lightning logger used in the latter. """def__init__(self,timer_kwargs={}):"""init Args: timer_kwargs (dict, optional): _description_. Defaults to {}. """self._sync_cuda=timer_kwargs.get("sync_cuda",False)self.timers=defaultdict(defaultdict)def_on_epoch_start(self,name,trainer,pl_module):"""_on_epoch_start"""# synchronize pytorch cuda execution if supportedifself._sync_cudaandtorch.cuda.is_initialized():torch.cuda.synchronize()self.timers[name]["step"]=0self.timers[name]["start"]=time.time()def_on_batch_end(self,name,trainer,pl_module):"""_on_epoch_start"""# synchronize pytorch cuda execution if supportedifself._sync_cudaandtorch.cuda.is_initialized():torch.cuda.synchronize()end=time.time()dt=end-self.timers[name]["start"]logging.info(f'Step {self.timers[name]["step"]}: {name} in s={dt}')self.timers[name]["step"]+=1self.timers[name]["start"]=enddefon_train_epoch_start(self,trainer,pl_module):"""on_train_epoch_start"""self._on_epoch_start("train_step_timing in s",trainer,pl_module)defon_validation_epoch_start(self,trainer,pl_module):"""on_validation_epoch_start"""self._on_epoch_start("validation_step_timing in s",trainer,pl_module)defon_train_batch_end(self,trainer,pl_module,outputs,batch,batch_idx):"""on_train_batch_end"""self._on_batch_end("train_step_timing in s",trainer,pl_module)defon_validation_batch_end(self,trainer,pl_module,outputs,batch,batch_idx):"""on_validation_batch_end"""self._on_batch_end("validation_step_timing in s",trainer,pl_module)
defexp_manager(trainer:'lightning.pytorch.Trainer',cfg:Optional[Union[DictConfig,Dict]]=None)->Optional[Path]:""" exp_manager is a helper function used to manage folders for experiments. It follows the pytorch lightning paradigm of exp_dir/model_or_experiment_name/version. If the lightning trainer has a logger, exp_manager will get exp_dir, name, and version from the logger. Otherwise it will use the exp_dir and name arguments to create the logging directory. exp_manager also allows for explicit folder creation via explicit_log_dir. The version can be a datetime string or an integer. Datestime version can be disabled if use_datetime_version is set to False. It optionally creates TensorBoardLogger, WandBLogger, DLLogger, MLFlowLogger, ClearMLLogger, ModelCheckpoint objects from pytorch lightning. It copies sys.argv, and git information if available to the logging directory. It creates a log file for each process to log their output into. exp_manager additionally has a resume feature (resume_if_exists) which can be used to continuing training from the constructed log_dir. When you need to continue the training repeatedly (like on a cluster which you need multiple consecutive jobs), you need to avoid creating the version folders. Therefore from v1.0.0, when resume_if_exists is set to True, creating the version folders is ignored. Args: trainer (lightning.pytorch.Trainer): The lightning trainer. cfg (DictConfig, dict): Can have the following keys: - explicit_log_dir (str, Path): Can be used to override exp_dir/name/version folder creation. Defaults to None, which will use exp_dir, name, and version to construct the logging directory. - exp_dir (str, Path): The base directory to create the logging directory. Defaults to None, which logs to ./nemo_experiments. - name (str): The name of the experiment. Defaults to None which turns into "default" via name = name or "default". - version (str): The version of the experiment. Defaults to None which uses either a datetime string or lightning's TensorboardLogger system of using version_{int}. - use_datetime_version (bool): Whether to use a datetime string for version. Defaults to True. - resume_if_exists (bool): Whether this experiment is resuming from a previous run. If True, it sets trainer._checkpoint_connector._ckpt_path so that the trainer should auto-resume. exp_manager will move files under log_dir to log_dir/run_{int}. Defaults to False. From v1.0.0, when resume_if_exists is True, we would not create version folders to make it easier to find the log folder for next runs. - resume_past_end (bool): exp_manager errors out if resume_if_exists is True and a checkpoint matching ``*end.ckpt`` indicating a previous training run fully completed. This behaviour can be disabled, in which case the ``*end.ckpt`` will be loaded by setting resume_past_end to True. Defaults to False. - resume_ignore_no_checkpoint (bool): exp_manager errors out if resume_if_exists is True and no checkpoint could be found. This behaviour can be disabled, in which case exp_manager will print a message and continue without restoring, by setting resume_ignore_no_checkpoint to True. Defaults to False. - resume_from_checkpoint (str): Can be used to specify a path to a specific checkpoint file to load from. This will override any checkpoint found when resume_if_exists is True. Defaults to None. - create_tensorboard_logger (bool): Whether to create a tensorboard logger and attach it to the pytorch lightning trainer. Defaults to True. - summary_writer_kwargs (dict): A dictionary of kwargs that can be passed to lightning's TensorboardLogger class. Note that log_dir is passed by exp_manager and cannot exist in this dict. Defaults to None. - create_wandb_logger (bool): Whether to create a Weights and Baises logger and attach it to the pytorch lightning trainer. Defaults to False. - wandb_logger_kwargs (dict): A dictionary of kwargs that can be passed to lightning's WandBLogger class. Note that name and project are required parameters if create_wandb_logger is True. Defaults to None. - create_mlflow_logger (bool): Whether to create an MLFlow logger and attach it to the pytorch lightning training. Defaults to False - mlflow_logger_kwargs (dict): optional parameters for the MLFlow logger - create_dllogger_logger (bool): Whether to create an DLLogger logger and attach it to the pytorch lightning training. Defaults to False - dllogger_logger_kwargs (dict): optional parameters for the DLLogger logger - create_clearml_logger (bool): Whether to create an ClearML logger and attach it to the pytorch lightning training. Defaults to False - clearml_logger_kwargs (dict): optional parameters for the ClearML logger - create_checkpoint_callback (bool): Whether to create a ModelCheckpoint callback and attach it to the pytorch lightning trainer. The ModelCheckpoint saves the top 3 models with the best "val_loss", the most recent checkpoint under ``*last.ckpt``, and the final checkpoint after training completes under ``*end.ckpt``. Defaults to True. - create_early_stopping_callback (bool): Flag to decide if early stopping should be used to stop training. Default is False. See EarlyStoppingParams dataclass above. - create_preemption_callback (bool): Flag to decide whether to enable preemption callback to save checkpoints and exit training immediately upon preemption. Default is True. - create_straggler_detection_callback (bool): Use straggler detection callback. Default is False. - create_fault_tolerance_callback (bool): Use fault tolerance callback. Default is False. - files_to_copy (list): A list of files to copy to the experiment logging directory. Defaults to None which copies no files. - max_time (str): The maximum wall clock time *per run*. This is intended to be used on clusters where you want a checkpoint to be saved after this specified time and be able to resume from that checkpoint. Defaults to None. - seconds_to_sleep (float): seconds to sleep non rank 0 processes for. Used to give enough time for rank 0 to initialize - train_time_interval (timedelta): pass an object of timedelta to save the model every timedelta. Defaults to None. (use _target_ with hydra to achieve this) returns: log_dir (Path): The final logging directory where logging files are saved. Usually the concatenation of exp_dir, name, and version. """# Add rank information to logger# Note: trainer.global_rank and trainer.is_global_zero are not set until trainer.fit, so have to hack around itlocal_rank=int(os.environ.get("LOCAL_RANK",0))global_rank=trainer.node_rank*trainer.num_devices+local_ranklogging.rank=global_rankifcfgisNone:logging.error("exp_manager did not receive a cfg argument. It will be disabled.")returniftrainer.fast_dev_run:logging.info("Trainer was called with fast_dev_run. exp_manager will return without any functionality.")return# Ensure passed cfg is compliant with ExpManagerConfigschema=OmegaConf.structured(ExpManagerConfig)# TODO: remove this checkifis_global_rank_zero():logging.info('ExpManager schema')logging.info(schema)ifisinstance(cfg,dict):cfg=OmegaConf.create(cfg)elifnotisinstance(cfg,DictConfig):raiseValueError(f"cfg was type: {type(cfg)}. Expected either a dict or a DictConfig")cfg=OmegaConf.create(OmegaConf.to_container(cfg,resolve=True))cfg=OmegaConf.merge(schema,cfg)# type: ExpManagerConfig# Ensures that trainer options are compliant with NeMo and exp_manager argumentserror_checks(trainer,cfg)log_dir,exp_dir,name,version=get_log_dir(trainer=trainer,exp_dir=cfg.exp_dir,name=cfg.name,version=cfg.version,explicit_log_dir=cfg.explicit_log_dir,use_datetime_version=cfg.use_datetime_version,resume_if_exists=cfg.resume_if_exists,)check_resume(trainer,log_dir,cfg.resume_if_exists,cfg.resume_past_end,cfg.resume_ignore_no_checkpoint,cfg.checkpoint_callback_params.dirpath,cfg.resume_from_checkpoint,)checkpoint_name=name# If name returned from get_log_dir is "", use cfg.name for checkpointingifcheckpoint_nameisNoneorcheckpoint_name=='':checkpoint_name=cfg.nameor"default"# Set mlflow name if it's not set, before the main name is erasedifcfg.create_mlflow_loggerand(notcfg.mlflow_logger_kwargs.get("experiment_name",None)):cfg.mlflow_logger_kwargs.experiment_name=cfg.namelogging.warning('mlflow logger specified but no experiment name set. Using the same as Tensorboard: %s',cfg.mlflow_logger_kwargs.experiment_name,)cfg.name=name# Used for configure_loggers so that the log_dir is properly set even if name is ""cfg.version=version# update app_state with log_dir, exp_dir, etcapp_state=AppState()app_state.log_dir=log_dirapp_state.exp_dir=exp_dirapp_state.name=nameapp_state.version=versionapp_state.checkpoint_name=checkpoint_nameapp_state.create_checkpoint_callback=cfg.create_checkpoint_callbackapp_state.checkpoint_callback_params=cfg.checkpoint_callback_params# Create the logging directory if it does not exist# Cannot limit creation to global zero as all ranks write to own log fileos.makedirs(log_dir,exist_ok=True)logging.info(f'Experiments will be logged at {log_dir}')trainer._default_root_dir=log_dir# Only log on all ranks when NEMO_TESTING is Trueifget_envbool(NEMO_ENV_VARNAME_TESTING,False):log_file=log_dir/f'nemo_log_globalrank-{global_rank}_localrank-{local_rank}.txt'logging.add_file_handler(log_file)# For some reason, LearningRateLogger requires trainer to have a logger. Safer to create logger on all ranks# not just global rank 0.if(cfg.create_tensorboard_loggerorcfg.create_wandb_loggerorcfg.create_mlflow_loggerorcfg.create_dllogger_loggerorcfg.create_clearml_loggerorcfg.create_neptune_logger):configure_loggers(trainer,exp_dir,log_dir,cfg.name,cfg.version,cfg.checkpoint_callback_params,cfg.create_tensorboard_logger,cfg.summary_writer_kwargs,cfg.create_wandb_logger,cfg.wandb_logger_kwargs,cfg.create_mlflow_logger,cfg.mlflow_logger_kwargs,cfg.create_dllogger_logger,cfg.dllogger_logger_kwargs,cfg.create_clearml_logger,cfg.clearml_logger_kwargs,cfg.create_neptune_logger,cfg.neptune_logger_kwargs,)# add loggers timing callbacksifcfg.log_delta_step_timing:timing_callback=DeltaTimingCallback(timer_kwargs=cfg.step_timing_kwargsor{})trainer.callbacks.insert(0,timing_callback)elifcfg.log_step_timing:timing_callback=TimingCallback(timer_kwargs=cfg.step_timing_kwargsor{})trainer.callbacks.insert(0,timing_callback)ifcfg.ema.enable:ema_callback=EMA(decay=cfg.ema.decay,validate_original_weights=cfg.ema.validate_original_weights,cpu_offload=cfg.ema.cpu_offload,every_n_steps=cfg.ema.every_n_steps,)trainer.callbacks.append(ema_callback)ifcfg.create_early_stopping_callback:early_stop_callback=EarlyStopping(**cfg.early_stopping_callback_params)trainer.callbacks.append(early_stop_callback)ifcfg.create_ipl_epoch_stopper_callback:ipl_epoch_stopper_callback=IPLEpochStopper(**cfg.ipl_epoch_stopper_callback_params)trainer.callbacks.append(ipl_epoch_stopper_callback)ifcfg.create_checkpoint_callback:configure_checkpointing(trainer,log_dir,checkpoint_name,cfg.resume_if_exists,cfg.checkpoint_callback_params,cfg.create_preemption_callback,)ifcfg.disable_validation_on_resume:# extend training loop to skip initial validation when resuming from checkpointconfigure_no_restart_validation_training_loop(trainer)# Setup a stateless timer for use on clusters.ifcfg.max_time_per_runisnotNone:found_ptl_timer=Falseforidx,callbackinenumerate(trainer.callbacks):ifisinstance(callback,Timer):# NOTE: PTL does not expose a `trainer.max_time`. By the time we are in this function,# PTL has already setup a timer if the user specifies `trainer.max_time` so best we# can do is replace that.# Working: If only `trainer.max_time` is set - it behaves as a normal PTL timer.# If only `exp_manager.max_time_per_run` is set - it behaves as a StateLessTimer.# If both are set, it also behaves as a StateLessTimer.logging.warning('Found a PTL Timer callback, replacing with a StatelessTimer callback. ''This will happen if you set trainer.max_time as well as exp_manager.max_time_per_run.')trainer.callbacks[idx]=StatelessTimer(cfg.max_time_per_run)found_ptl_timer=Truebreakifnotfound_ptl_timer:trainer.max_time=cfg.max_time_per_runtrainer.callbacks.append(StatelessTimer(cfg.max_time_per_run))ifcfg.create_straggler_detection_callback:ifHAVE_STRAGGLER_DET:logging.info("Enabling straggler detection...")straggler_det_args_dict=dict(cfg.straggler_detection_params)straggler_det_callback=StragglerDetectionCallback(**straggler_det_args_dict)trainer.callbacks.append(straggler_det_callback)else:raiseValueError("`create_straggler_detection_callback` is True, but there is no Straggler Det. ""package installed.")ifcfg.create_fault_tolerance_callback:ifHAVE_FT:logging.info("Enabling fault tolerance...")ft_params=cfg.fault_tolerance# job failures are handled by the ft_launcher,# here we only need to know if the autoresume is enabled.ft_use_autoresume=ft_params.max_subsequent_job_failures>0fault_tol_callback=FaultToleranceCallback(# log_dir is "<run name>/results/"exp_dir=Path(log_dir).parent,autoresume=ft_use_autoresume,calculate_timeouts=ft_params.calculate_timeouts,simulated_fault_params=ft_params.simulated_fault,)trainer.callbacks.append(fault_tol_callback)else:raiseValueError('FaultToleranceCallback was enabled with create_fault_tolerance_callback, ''but fault_tolerance package is not installed.')ifcfg.log_tflops_per_sec_per_gpu:logging.info("TFLOPs per sec per GPU will be calculated, conditioned on supported models. ""Defaults to -1 upon failure.")ifis_global_rank_zero():# Move files_to_copy to folder and add git information if presentifcfg.files_to_copy:for_fileincfg.files_to_copy:copy(Path(_file),log_dir)# Create files for cmd args and git infowithopen(log_dir/'cmd-args.log','w',encoding='utf-8')as_file:_file.write(" ".join(sys.argv))# Try to get git hashgit_repo,git_hash=get_git_hash()ifgit_repo:withopen(log_dir/'git-info.log','a',encoding='utf-8')as_file:_file.write(f'commit hash: {git_hash}')_file.write(get_git_diff())# Add err_file logging to global_rank zerologging.add_err_file_handler(log_dir/'nemo_error_log.txt')# Add lightning file logging to global_rank zeroadd_filehandlers_to_pl_logger(log_dir/'lightning_logs.txt',log_dir/'nemo_error_log.txt')eliftrainer.num_nodes*trainer.num_devices>1:# sleep other ranks so rank 0 can finish# doing the initialization such as moving filestime.sleep(cfg.seconds_to_sleep)add_handlers_to_mcore_logger()returnlog_dir
deferror_checks(trainer:'lightning.pytorch.Trainer',cfg:Optional[Union[DictConfig,Dict]]=None):""" Checks that the passed trainer is compliant with NeMo and exp_manager's passed configuration. Checks that: - Throws error when hydra has changed the working directory. This causes issues with lightning's DDP - Throws error when trainer has loggers defined but create_tensorboard_logger or create_wandB_logger or create_mlflow_logger or create_dllogger_logger is True - Prints error messages when 1) run on multi-node and not Slurm, and 2) run on multi-gpu without DDP """ifHydraConfig.initialized()andget_original_cwd()!=os.getcwd():raiseValueError("Hydra changed the working directory. This interferes with ExpManger's functionality."" Please pass hydra.run.dir=. to your python script.")iftrainer.loggerisnotNoneand(cfg.create_tensorboard_loggerorcfg.create_wandb_loggerorcfg.create_mlflow_logger):raiseLoggerMisconfigurationError("The pytorch lightning trainer that was passed to exp_manager contained a logger, ""and either "f"create_tensorboard_logger: {cfg.create_tensorboard_logger} or create_wandb_logger: "f"{cfg.create_wandb_logger} or create_mlflow_logger: {cfg.create_mlflow_logger}"f"or create_dllogger_logger: {cfg.create_mlflow_logger} was set to True. ""These can only be used if trainer does not already have a logger.")iftrainer.num_nodes>1andnotcheck_slurm(trainer):logging.error("You are running multi-node training without SLURM handling the processes."" Please note that this is not tested in NeMo and could result in errors.")iftrainer.num_devices>1andnotisinstance(trainer.strategy,DDPStrategy):logging.error("You are running multi-gpu without ddp.Please note that this is not tested in NeMo and ""could result in errors.")def_filter_out_unfinished_checkpoints(checkpoint_paths:Collection[Union[Path,str]])->Collection[Union[Path,str]]:"""_filter_out_unfinished_checkpoints"""res=[]forchkpt_pathincheckpoint_paths:ifNeMoModelCheckpoint.is_checkpoint_unfinished(chkpt_path):logging.warning(f'Checkpoint {chkpt_path} has the unfinished marker set - skipped while looking ''for the last one.')else:res.append(chkpt_path)returnresdefcheck_resume(trainer:'lightning.pytorch.Trainer',log_dir:str,resume_if_exists:bool=False,resume_past_end:bool=False,resume_ignore_no_checkpoint:bool=False,dirpath:str=None,resume_from_checkpoint:str=None,):"""Checks that resume=True was used correctly with the arguments pass to exp_manager. Sets trainer._checkpoint_connector._ckpt_path as necessary. Returns: log_dir (Path): The log_dir exp_dir (str): The base exp_dir without name nor version name (str): The name of the experiment version (str): The version of the experiment Raises: NotFoundError: If resume is True, resume_ignore_no_checkpoint is False, and checkpoints could not be found. ValueError: If resume is True, and there were more than 1 checkpoint could found. """ifnotlog_dir:raiseValueError(f"Resuming requires the log_dir {log_dir} to be passed to exp_manager")# is_s3_url from here has no dependency requirementsfromnemo.utils.s3_dirpath_utilsimportis_s3_urltry:# when using an s3 dirpath, we rely on optional dependencies in the S3Utils class.ifdirpathisnotNoneandis_s3_url(dirpath):fromnemo.utils.s3_utilsimportS3UtilsexceptImportErroraserr:returnFalse,"Detected S3 dirpath while missing required dependencies.\n{}\n".format(err.output.decode("utf-8"))checkpoint=Noneifresume_from_checkpoint:checkpoint=resume_from_checkpointifresume_if_exists:''' attach valid checkpoint path to trainer if current rank is rank zero of any data parallel groups this limit to only global rank 0 process calling s3, instead of all processes calling s3 '''# If we are using S3 checkpointing, we want check_resume to only execute on a single rank# to avoid throttling S3.ifis_global_rank_zero()ornot(is_s3_url(dirpath)andis_multistorageclient_url(dirpath)):checkpoint_dir_exists=Falseifis_s3_url(dirpath):checkpoint_dir=dirpathcheckpoint_dir_exists=S3Utils.s3_path_exists(checkpoint_dir,match_directory=True)ifcheckpoint_dir_exists:# max number of last.ckpt files: save_last_k_checkpoints * tp * pp = 5*8*40.# If optim states is saved distributedly, multiply by dp_sizeall_keys=S3Utils.find_files_with_suffix(checkpoint_dir,suffix=None,return_key_only=False)end_checkpoints=[kforkinall_keysifk.endswith('end.ckpt')]last_checkpoints=[kforkinall_keysifk.endswith('last.ckpt')]else:end_checkpoints=[]last_checkpoints=[]elifis_multistorageclient_url(dirpath):msc=import_multistorageclient()checkpoint_dir=dirpathall_keys=msc.glob(f"{dirpath}**/*.ckpt")checkpoint_dir_exists=Trueifall_keyselseFalseifall_keys:end_checkpoints=sorted([kforkinall_keysifk.endswith('end.ckpt')],reverse=True)last_checkpoints=sorted([kforkinall_keysifk.endswith('last.ckpt')],reverse=True)else:end_checkpoints=[]last_checkpoints=[]else:# default non-s3 implementation# Use <log_dir>/checkpoints/ unless `dirpath` is setcheckpoint_dir=Path(dirpath)ifdirpathelsePath(Path(log_dir)/"checkpoints")checkpoint_dir_exists=checkpoint_dir.exists()# when using distributed checkpointing, checkpoint_dir is a directory of directories# we check for this heredist_checkpoints=[dfordinlist(checkpoint_dir.glob("*"))ifd.is_dir()]end_dist_checkpoints=[dfordindist_checkpointsifd.match("*end")]last_dist_checkpoints=[dfordindist_checkpointsifd.match("*last")]end_checkpoints=(end_dist_checkpointsifend_dist_checkpointselselist(checkpoint_dir.rglob("*end.ckpt")))end_chkpt_cnt=len(end_checkpoints)end_checkpoints=_filter_out_unfinished_checkpoints(end_checkpoints)finished_end_chkpt_cnt=len(end_checkpoints)ifend_chkpt_cnt>0andfinished_end_chkpt_cnt==0:raiseValueError("End checkpoint is unfinished and cannot be used to resume the training."" Please remove the checkpoint manually to avoid unexpected cosequences, such as"" restarting from scratch.")last_checkpoints=(last_dist_checkpointsiflast_dist_checkpointselselist(checkpoint_dir.rglob("*last.ckpt")))last_chkpt_cnt=len(last_checkpoints)last_checkpoints=_filter_out_unfinished_checkpoints(last_checkpoints)finished_last_chkpt_cnt=len(last_checkpoints)iflast_chkpt_cnt>0andfinished_last_chkpt_cnt==0:raiseValueError("Last checkpoint is unfinished and cannot be used to resume the training."" Please remove the checkpoint manually to avoid unexpected cosequences, "" such as restarting from scratch. Hint: Iteration number can be added "" to the checkpoint name pattern"" to maximize chance that there is at least one finished last checkpoint to"" resume from.")ifnotcheckpoint_dir_existsor(notlen(end_checkpoints)>0andnotlen(last_checkpoints)>0):ifresume_ignore_no_checkpoint:warn=(f"There were no checkpoints found in checkpoint_dir or no checkpoint "f"folder at checkpoint_dir :{checkpoint_dir}. ")ifcheckpointisNone:warn+="Training from scratch."elifcheckpoint==resume_from_checkpoint:warn+=f"Training from {resume_from_checkpoint}."logging.warning(warn)else:raiseNotFoundError(f"There were no checkpoints found in checkpoint_dir or no checkpoint "f"folder at checkpoint_dir :{checkpoint_dir}. Cannot resume.")eliflen(end_checkpoints)>0:ifresume_past_end:iflen(end_checkpoints)>1:if'mp_rank'instr(end_checkpoints[0]):checkpoint=end_checkpoints[0]else:raiseValueError(f"Multiple checkpoints {end_checkpoints} that matches *end.ckpt.")else:raiseValueError(f"Found {end_checkpoints[0]} indicating that the last training run has already completed.")eliflen(last_checkpoints)>1:ifany([sforsin['mp_rank','tp_rank','fsdp_shard']ifsinstr(last_checkpoints[0])]):checkpoint=last_checkpoints[0]checkpoint=uninject_model_parallel_rank(checkpoint)else:raiseValueError(f"Multiple checkpoints {last_checkpoints} that matches *last.ckpt.")else:checkpoint=last_checkpoints[0]# PTL 2.0 supports ckpt_path instead of resume_from_checkpoint as the trainer flagifcheckpointisnotNone:trainer.ckpt_path=str(checkpoint)logging.info(f'Resuming training from checkpoint: {trainer.ckpt_path}')ifis_global_rank_zero():# Check to see if any files exist that need to be movedfiles_to_move=[]ifPath(log_dir).exists():forchildinPath(log_dir).iterdir():ifchild.is_file()andnotchild.name.startswith("events.out.tfevents"):files_to_move.append(child)iflen(files_to_move)>0:# Move old files to a new folderother_run_dirs=Path(log_dir).glob("run_*")run_count=0forfoldinother_run_dirs:iffold.is_dir():run_count+=1new_run_dir=Path(Path(log_dir)/f"run_{run_count}")new_run_dir.mkdir()for_fileinfiles_to_move:move(str(_file),str(new_run_dir))defcheck_explicit_log_dir(trainer:'lightning.pytorch.Trainer',explicit_log_dir:Union[Path,str],exp_dir:str,name:str,version:str)->Tuple[Path,str,str,str]:"""Checks that the passed arguments are compatible with explicit_log_dir. Returns: log_dir (Path): the log_dir exp_dir (str): the base exp_dir without name nor version name (str): The name of the experiment version (str): The version of the experiment Raise: LoggerMisconfigurationError """iftrainer.loggerisnotNone:raiseLoggerMisconfigurationError("The pytorch lightning trainer that was passed to exp_manager contained a ""logger and explicit_log_dir: "f"{explicit_log_dir} was pass to exp_manager. ""Please remove the logger from the lightning trainer.")# Checking only (explicit_log_dir) vs (exp_dir and version).# The `name` will be used as the actual name of checkpoint/archive.ifexp_dirorversion:logging.error(f"exp_manager received explicit_log_dir: {explicit_log_dir} and at least "f"one of exp_dir: {exp_dir}, "f"or version: {version}. Please note that exp_dir, name, and version will be ignored.")ifis_global_rank_zero()andPath(explicit_log_dir).exists():logging.warning(f"Exp_manager is logging to {explicit_log_dir}, but it already exists.")returnPath(explicit_log_dir),str(explicit_log_dir),"",""defget_log_dir(trainer:'lightning.pytorch.Trainer',exp_dir:str=None,name:str=None,version:str=None,explicit_log_dir:str=None,use_datetime_version:bool=True,resume_if_exists:bool=False,)->Tuple[Path,str,str,str]:""" Obtains the log_dir used for exp_manager. Returns: log_dir (Path): the log_dir exp_dir (str): the base exp_dir without name nor version name (str): The name of the experiment version (str): The version of the experiment explicit_log_dir (str): The explicit path to the log folder. Defaults to False. use_datetime_version (bool): Uses date and time as the version of the log folder. Defaults to True. resume_if_exists (bool): if resume_if_exists of the exp_manager's config is enabled or not. When enabled, the version folders would not get created. Raise: LoggerMisconfigurationError: If trainer is incompatible with arguments NotFoundError: If resume is True, resume_ignore_no_checkpoint is False, and checkpoints could not be found. ValueError: If resume is True, and there were more than 1 checkpoint could found. """ifexplicit_log_dir:# If explicit log_dir was passed, short circuitreturncheck_explicit_log_dir(trainer,explicit_log_dir,exp_dir,name,version)# Default exp_dir to ./nemo_experiments if None was passed_exp_dir=exp_dirifexp_dirisNone:_exp_dir=str(Path.cwd()/'nemo_experiments')# If the user has already defined a logger for the trainer,# use the logger defaults for logging directoryiftrainer.loggerisnotNone:iftrainer.logger.save_dir:ifexp_dir:raiseLoggerMisconfigurationError("The pytorch lightning trainer that was passed to exp_manager contained a ""logger, the logger's "f"save_dir was not None, and exp_dir ({exp_dir}) was not None. ""If trainer.logger.save_dir ""exists, exp_manager will use trainer.logger.save_dir as the ""logging directory and exp_dir ""must be None.")_exp_dir=trainer.logger.save_dirifname:raiseLoggerMisconfigurationError("The pytorch lightning trainer that was passed to exp_manager ""contained a logger, and name: "f"{name} was also passed to exp_manager. If the trainer contains a ""logger, exp_manager will use trainer.logger.name, and name passed ""to exp_manager must be None.")name=trainer.logger.nameversion=f"version_{trainer.logger.version}"# Use user-defined exp_dir, project_name, exp_name, and versioning optionselse:name=nameor"default"version=versionoros.environ.get(NEMO_ENV_VARNAME_VERSION,None)ifnotversion:ifresume_if_exists:logging.warning("No version folders would be created under the log folder as ""'resume_if_exists' is enabled.")version=Noneelifis_global_rank_zero():ifuse_datetime_version:version=time.strftime('%Y-%m-%d_%H-%M-%S')else:tensorboard_logger=TensorBoardLogger(save_dir=Path(_exp_dir),name=name,version=version)version=f"version_{tensorboard_logger.version}"os.environ[NEMO_ENV_VARNAME_VERSION]=""ifversionisNoneelseversionlog_dir=Path(_exp_dir)/Path(str(name))/Path(""ifversionisNoneelsestr(version))returnlog_dir,str(_exp_dir),name,versiondefget_git_hash():""" Helper function that tries to get the commit hash if running inside a git folder returns: Bool: Whether the git subprocess ran without error str: git subprocess output or error message """try:return(True,subprocess.check_output(['git','rev-parse','HEAD'],stderr=subprocess.STDOUT).decode(),)except(subprocess.CalledProcessError,FileNotFoundError)aserr:returnFalse,"{}\n".format(err)defget_git_diff():""" Helper function that tries to get the git diff if running inside a git folder returns: Bool: Whether the git subprocess ran without error str: git subprocess output or error message """try:returnsubprocess.check_output(['git','diff'],stderr=subprocess.STDOUT).decode()exceptsubprocess.CalledProcessErroraserr:return"{}\n".format(err.output.decode("utf-8"))defconfigure_loggers(trainer:'lightning.pytorch.Trainer',exp_dir:[Path,str],log_dir:[Path,str],name:str,version:str,checkpoint_callback_params:dict,create_tensorboard_logger:bool,summary_writer_kwargs:dict,create_wandb_logger:bool,wandb_kwargs:dict,create_mlflow_logger:bool,mlflow_kwargs:dict,create_dllogger_logger:bool,dllogger_kwargs:dict,create_clearml_logger:bool,clearml_kwargs:dict,create_neptune_logger:bool,neptune_kwargs:dict,):""" Creates TensorboardLogger and/or WandBLogger / MLFlowLogger / DLlogger / ClearMLLogger and attach them to trainer. Raises ValueError if summary_writer_kwargs or wandb_kwargs are misconfigured. """# Potentially create tensorboard logger and/or WandBLogger / MLFlowLogger / DLLoggerlogger_list=[]ifcreate_tensorboard_logger:ifsummary_writer_kwargsisNone:summary_writer_kwargs={}elif"log_dir"insummary_writer_kwargs:raiseValueError("You cannot pass `log_dir` as part of `summary_writer_kwargs`. `log_dir` ""is handled by lightning's ""TensorBoardLogger logger.")tensorboard_logger=TensorBoardLogger(save_dir=exp_dir,name=name,version=version,**summary_writer_kwargs)logger_list.append(tensorboard_logger)logging.info("TensorboardLogger has been set up")ifcreate_wandb_logger:ifwandb_kwargsisNone:wandb_kwargs={}if"name"notinwandb_kwargsand"project"notinwandb_kwargs:raiseValueError("name and project are required for wandb_logger")# Update the wandb save_dirifwandb_kwargs.get('save_dir',None)isNone:wandb_kwargs['save_dir']=exp_diros.makedirs(wandb_kwargs['save_dir'],exist_ok=True)wandb_logger=WandbLogger(version=version,**wandb_kwargs)logger_list.append(wandb_logger)logging.info("WandBLogger has been set up")ifcreate_mlflow_logger:mlflow_logger=MLFlowLogger(run_name=version,**mlflow_kwargs)logger_list.append(mlflow_logger)logging.info("MLFlowLogger has been set up")ifcreate_dllogger_logger:dllogger_logger=DLLogger(**dllogger_kwargs)logger_list.append(dllogger_logger)logging.info("DLLogger has been set up")ifcreate_clearml_logger:clearml_logger=ClearMLLogger(clearml_cfg=clearml_kwargs,log_dir=log_dir,prefix=name,save_best_model=checkpoint_callback_params.save_best_model,)logger_list.append(clearml_logger)logging.info("ClearMLLogger has been set up")ifcreate_neptune_logger:ifneptune_kwargsisNone:neptune_kwargs={}if"name"notinneptune_kwargsand"project"notinneptune_kwargs:raiseValueError("name and project are required for neptune_logger")if"api_key"notinneptune_kwargsandnotos.getenv("NEPTUNE_API_TOKEN",None):raiseValueError("either api_key should be set in neptune_kwargs or NEPTUNE_API_TOKEN should ""be set in environment variable for neptune_logger")neptune_logger=NeptuneLogger(**neptune_kwargs)logger_list.append(neptune_logger)logging.info("NeptuneLogger has been set up")trainer._logger_connector.configure_logger(logger_list)classNeMoCheckpointConnector(_CheckpointConnector):""" Wrapper around Lightning's _CheckpointConnector to use broadcasted checkpoint path in distributed training settings to pre-load checkpoint. """defresume_start(self,checkpoint_path=None)->None:"""resume_start"""checkpoint_path=self.trainer.ckpt_pathifcheckpoint_pathisnotNone:logging.info(f'Resuming from checkpoint {checkpoint_path}, rank {torch.distributed.get_rank()}')start_time=time.perf_counter()super().resume_start(checkpoint_path)ifcheckpoint_pathisnotNone:logging.info('Time elapsed loading checkpoint/optimizer states: 'f'{(time.perf_counter()-start_time):.2f} seconds, 'f'rank {torch.distributed.get_rank()}')defconfigure_checkpointing(trainer:'lightning.pytorch.Trainer',log_dir:Path,name:str,resume:bool,params:'DictConfig',create_preemption_callback:bool,):"""Adds ModelCheckpoint to trainer. Raises CheckpointMisconfigurationError if trainer already has a ModelCheckpoint callback """forcallbackintrainer.callbacks:ifisinstance(callback,ModelCheckpoint):raiseCheckpointMisconfigurationError("The pytorch lightning trainer that was passed to exp_manager ""contained a ModelCheckpoint ""and create_checkpoint_callback was set to True. ""Please either set create_checkpoint_callback ""to False, or remove ModelCheckpoint from the lightning trainer")# Create the callback and attach it to trainerif"filepath"inparams:ifparams.filepathisnotNone:logging.warning("filepath is deprecated. Please switch to dirpath and filename instead")ifparams.dirpathisNone:params.dirpath=Path(params.filepath).parentifparams.filenameisNone:params.filename=Path(params.filepath).namewithopen_dict(params):delparams["filepath"]ifparams.dirpathisNone:params.dirpath=Path(log_dir/'checkpoints')ifparams.filenameisNone:params.filename=f'{name}--{{{params.monitor}:.4f}}-{{epoch}}'ifparams.prefixisNone:params.prefix=nameifparams.always_save_nemo:app_state=AppState()if((app_state.tensor_model_parallel_sizeisnotNoneandapp_state.tensor_model_parallel_size>1)or(app_state.pipeline_model_parallel_sizeisnotNoneandapp_state.pipeline_model_parallel_size>1)or(app_state.context_parallel_sizeisnotNoneandapp_state.context_parallel_size>1)):raiseLoggerMisconfigurationError("always_save_nemo is set to True, please ensure that model parallel is not used."f"tensor_model_parallel_size: {app_state.tensor_model_parallel_size},"f"pipeline_model_parallel_size: {app_state.pipeline_model_parallel_size},"f"context_parallel_size: {app_state.context_parallel_size},")NeMoModelCheckpoint.CHECKPOINT_NAME_LAST=params.filename+'-last'logging.debug(params.dirpath)logging.debug(params.filename)logging.debug(params.prefix)if"val"inparams.monitor:if(trainer.max_epochsisnotNoneandtrainer.max_epochs!=-1andtrainer.max_epochs<trainer.check_val_every_n_epoch):logging.error("The checkpoint callback was told to monitor a validation value but ""trainer.max_epochs("f"{trainer.max_epochs}) was less than "f"trainer.check_val_every_n_epoch({trainer.check_val_every_n_epoch}"f"). It is very likely this run will fail with "f"ModelCheckpoint(monitor='{params.monitor}') not found ""in the returned metrics. Please ensure that validation is run within trainer.max_epochs.")eliftrainer.max_stepsisnotNoneandtrainer.max_steps!=-1:logging.warning("The checkpoint callback was told to monitor a validation value and trainer's"" max_steps was set to "f"{trainer.max_steps}. Please ensure that max_steps will run for at least "f"{trainer.check_val_every_n_epoch} epochs to ensure that checkpointing"" will not error out.")checkpoint_callback=NeMoModelCheckpoint(n_resume=resume,**params)checkpoint_callback.last_model_path=trainer.ckpt_pathor""if'mp_rank'incheckpoint_callback.last_model_pathor'tp_rank'incheckpoint_callback.last_model_path:checkpoint_callback.last_model_path=uninject_model_parallel_rank(checkpoint_callback.last_model_path)trainer.callbacks.append(checkpoint_callback)ifcreate_preemption_callback:# Check if cuda is avialable as preemption is supported only on GPUsiftorch.cuda.is_available():# By default PreemptionCallback handles SIGTERM. To handle other signals pass the# signal in the call as below:# PreemptionCallback(checkpoint_callback, signal.SIGCHLD)preemption_callback=PreemptionCallback(checkpoint_callback)trainer.callbacks.append(preemption_callback)else:logging.info("Preemption is supported only on GPUs, disabling preemption")defcheck_slurm(trainer):"""check_slurm"""try:returntrainer.accelerator_connector.is_slurm_managing_tasksexceptAttributeError:returnFalseclassStatelessTimer(Timer):"""Extension of PTL timers to be per run."""def__init__(self,duration:timedelta=None,interval:str=Interval.step,verbose:bool=True,)->None:"""stateless timer Args: duration (timedelta, optional): _description_. Defaults to None. interval (str, optional): _description_. Defaults to Interval.step. verbose (bool, optional): _description_. Defaults to True. """super().__init__(duration,interval,verbose)# Override PTL Timer's state dict to not store elapsed time information so that we can# restore and continue training.defstate_dict(self)->Dict[str,Any]:"""state_dict"""return{}defload_state_dict(self,state_dict:Dict[str,Any])->None:"""load_state_dict"""returndef_check_time_remaining(self,trainer:lightning.pytorch.Trainer)->None:"""_check_time_remaining"""super()._check_time_remaining(trainer)iftrainer.should_stop:# PTL's TrainingEpochLoop.advance() calls the on_train_batch_end hooks (which is where# Timer._check_time_remaining fires) BEFORE batch_progress.increment_completed(). The# current batch's optim step has already advanced global_step, so saving here would# capture batch_progress.current.completed lagging one behind optim_progress. On# resume, reset_on_restart rewinds batch_progress to .completed, PTL replays the# in-flight batch, and its optim step runs a second time — double-counting one# global_step per wall-time resume. Flush the in-flight batch first to keep the# saved state self-consistent._flush_in_flight_batch_progress(trainer)checkpoint_callback:Optional[NeMoModelCheckpoint]=trainer.checkpoint_callbackifcheckpoint_callback:monitor_candidates=checkpoint_callback._monitor_candidates(trainer)checkpoint_callback._save_last_checkpoint(trainer,monitor_candidates)# Throw this exception to signal to Lightning to terminate gracefully.fromlightning.pytorch.utilities.exceptionsimport_TunerExitExceptionraise_TunerExitException()def_flush_in_flight_batch_progress(trainer:lightning.pytorch.Trainer)->None:"""Bring batch_progress.current.completed up to .ready if a batch is in flight. Meant to be called from an ``on_train_batch_end`` hook before a checkpoint save, where PTL has not yet incremented ``batch_progress.current.completed`` but the batch's optim step has already advanced ``global_step``. See :meth:`StatelessTimer._check_time_remaining` for the off-by-one it avoids. """try:batch_progress=trainer.fit_loop.epoch_loop.batch_progressexceptAttributeError:returnifbatch_progress.current.ready>batch_progress.current.completed:batch_progress.increment_completed()defconfigure_no_restart_validation_training_loop(trainer:lightning.pytorch.Trainer)->None:"""configure_no_restart_validation_training_loop"""iftype(trainer.fit_loop.epoch_loop)!=_TrainingEpochLoop:warnings.warn("Detected custom epoch loop. Skipping no validation on restart support.",UserWarning)return# Pass trainer object to avoid trainer getting overwritten as Noneloop=SkipResumeTrainingValidationLoop(trainer,trainer.min_steps,trainer.max_steps)trainer.fit_loop.epoch_loop=loopclassSkipResumeTrainingValidationLoop(_TrainingEpochLoop):""" Extend the PTL Epoch loop to skip validating when resuming. This happens when resuming a checkpoint that has already run validation, but loading restores the training state before validation has run. """def_should_check_val_fx(self,data_fetcher)->bool:"""_should_check_val_fx"""ifself.restarting:returnFalsereturnsuper()._should_check_val_fx(data_fetcher)defclean_exp_ckpt(exp_log_dir:Union[str,Path],remove_ckpt:bool=True,remove_nemo:bool=False):""" Helper method that removes Pytorch Lightning .ckpt files or NeMo .nemo files from the checkpoint directory Args: exp_log_dir: str path to the root directory of the current experiment. remove_ckpt: bool, whether to remove all *.ckpt files in the checkpoints directory. remove_nemo: bool, whether to remove all *.nemo files in the checkpoints directory. """exp_log_dir=str(exp_log_dir)ifremove_ckpt:logging.info("Deleting *.ckpt files ...")ckpt_files=glob.glob(os.path.join(exp_log_dir,"checkpoints","*.ckpt"))forfilepathinckpt_files:os.remove(filepath)logging.info(f"Deleted file : {filepath}")ifremove_nemo:logging.info("Deleting *.nemo files ...")nemo_files=glob.glob(os.path.join(exp_log_dir,"checkpoints","*.nemo"))forfilepathinnemo_files:os.remove(filepath)logging.info(f"Deleted file : {filepath}")