Aluode/PerceptionLabPortable

1# Copyright 2020-present the HuggingFace Inc. team.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14"""15PyTorch-independent utilities for the Trainer class.16"""17 18import copy19import functools20import gc21import inspect22import os23import random24import re25import threading26import time27from typing import Any, Callable, NamedTuple, Optional, Union28 29import numpy as np30 31from .utils import (32    ExplicitEnum,33    is_psutil_available,34    is_tf_available,35    is_torch_available,36    is_torch_cuda_available,37    is_torch_hpu_available,38    is_torch_mlu_available,39    is_torch_mps_available,40    is_torch_musa_available,41    is_torch_npu_available,42    is_torch_xla_available,43    is_torch_xpu_available,44    requires_backends,45)46 47 48if is_torch_available():49    import torch50 51 52def seed_worker(worker_id: int, num_workers: int, rank: int):53    """54    Helper function to set worker seed during Dataloader initialization.55    """56    init_seed = torch.initial_seed() % 2**3257    worker_seed = num_workers * rank + init_seed58    set_seed(worker_seed)59 60 61def enable_full_determinism(seed: int, warn_only: bool = False):62    """63    Helper function for reproducible behavior during distributed training. See64    - https://pytorch.org/docs/stable/notes/randomness.html for pytorch65    - https://www.tensorflow.org/api_docs/python/tf/config/experimental/enable_op_determinism for tensorflow66    """67    # set seed first68    set_seed(seed)69 70    if is_torch_available():71        # Enable PyTorch deterministic mode. This potentially requires either the environment72        # variable 'CUDA_LAUNCH_BLOCKING' or 'CUBLAS_WORKSPACE_CONFIG' to be set,73        # depending on the CUDA version, so we set them both here74        os.environ["CUDA_LAUNCH_BLOCKING"] = "1"75        os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8"76        # The environment variable required to enable deterministic mode on Ascend NPUs.77        os.environ["ASCEND_LAUNCH_BLOCKING"] = "1"78        os.environ["HCCL_DETERMINISTIC"] = "1"79 80        os.environ["FLASH_ATTENTION_DETERMINISTIC"] = "1"81        torch.use_deterministic_algorithms(True, warn_only=warn_only)82 83        # Enable CUDNN deterministic mode84        torch.backends.cudnn.deterministic = True85        torch.backends.cudnn.benchmark = False86 87    if is_tf_available():88        import tensorflow as tf89 90        tf.config.experimental.enable_op_determinism()91 92 93def set_seed(seed: int, deterministic: bool = False):94    """95    Helper function for reproducible behavior to set the seed in `random`, `numpy`, `torch` and/or `tf` (if installed).96 97    Args:98        seed (`int`):99            The seed to set.100        deterministic (`bool`, *optional*, defaults to `False`):101            Whether to use deterministic algorithms where available. Can slow down training.102    """103    random.seed(seed)104    np.random.seed(seed)105    if is_torch_available():106        torch.manual_seed(seed)107        torch.cuda.manual_seed_all(seed)108        # ^^ safe to call this function even if cuda is not available109        if deterministic:110            torch.use_deterministic_algorithms(True)111    if is_torch_mlu_available():112        torch.mlu.manual_seed_all(seed)113    if is_torch_musa_available():114        torch.musa.manual_seed_all(seed)115    if is_torch_npu_available():116        torch.npu.manual_seed_all(seed)117    if is_torch_hpu_available():118        torch.hpu.manual_seed_all(seed)119    if is_torch_xpu_available():120        torch.xpu.manual_seed_all(seed)121    if is_tf_available():122        import tensorflow as tf123 124        tf.random.set_seed(seed)125        if deterministic:126            tf.config.experimental.enable_op_determinism()127 128 129def neftune_post_forward_hook(module, input, output):130    """131    Implements the NEFTune forward pass for the model using forward hooks. Note this works only for torch.nn.Embedding132    layers. This method is slightly adapted from the original source code that can be found here:133    https://github.com/neelsjain/NEFTune Simply add it to your model as follows:134    ```python135    model = ...136    model.embed_tokens.neftune_noise_alpha = 0.1137    model.embed_tokens.register_forward_hook(neftune_post_forward_hook)138    ```139    Args:140        module (`torch.nn.Module`):141            The embedding module where the hook is attached. Note that you need to set `module.neftune_noise_alpha` to142            the desired noise alpha value.143        input (`torch.Tensor`):144            The input tensor to the model.145        output (`torch.Tensor`):146            The output tensor of the model (i.e. the embeddings).147    """148    if module.training:149        dims = torch.tensor(output.size(1) * output.size(2))150        mag_norm = module.neftune_noise_alpha / torch.sqrt(dims)151        output = output + torch.zeros_like(output).uniform_(-mag_norm, mag_norm)152    return output153 154 155class EvalPrediction:156    """157    Evaluation output (always contains labels), to be used to compute metrics.158 159    Parameters:160        predictions (`np.ndarray`): Predictions of the model.161        label_ids (`np.ndarray`): Targets to be matched.162        inputs (`np.ndarray`, *optional*): Input data passed to the model.163        losses (`np.ndarray`, *optional*): Loss values computed during evaluation.164    """165 166    def __init__(167        self,168        predictions: Union[np.ndarray, tuple[np.ndarray]],169        label_ids: Union[np.ndarray, tuple[np.ndarray]],170        inputs: Optional[Union[np.ndarray, tuple[np.ndarray]]] = None,171        losses: Optional[Union[np.ndarray, tuple[np.ndarray]]] = None,172    ):173        self.predictions = predictions174        self.label_ids = label_ids175        self.inputs = inputs176        self.losses = losses177        self.elements = (self.predictions, self.label_ids)178        if self.inputs is not None:179            self.elements += (self.inputs,)180        if self.losses is not None:181            self.elements += (self.losses,)182 183    def __iter__(self):184        return iter(self.elements)185 186    def __getitem__(self, idx):187        if idx < 0 or idx >= len(self.elements):188            raise IndexError("tuple index out of range")189        return self.elements[idx]190 191 192class EvalLoopOutput(NamedTuple):193    predictions: Union[np.ndarray, tuple[np.ndarray]]194    label_ids: Optional[Union[np.ndarray, tuple[np.ndarray]]]195    metrics: Optional[dict[str, float]]196    num_samples: Optional[int]197 198 199class PredictionOutput(NamedTuple):200    predictions: Union[np.ndarray, tuple[np.ndarray]]201    label_ids: Optional[Union[np.ndarray, tuple[np.ndarray]]]202    metrics: Optional[dict[str, float]]203 204 205class TrainOutput(NamedTuple):206    global_step: int207    training_loss: float208    metrics: dict[str, float]209 210 211PREFIX_CHECKPOINT_DIR = "checkpoint"212_re_checkpoint = re.compile(r"^" + PREFIX_CHECKPOINT_DIR + r"\-(\d+)$")213 214 215def get_last_checkpoint(folder):216    content = os.listdir(folder)217    checkpoints = [218        path219        for path in content220        if _re_checkpoint.search(path) is not None and os.path.isdir(os.path.join(folder, path))221    ]222    if len(checkpoints) == 0:223        return224    return os.path.join(folder, max(checkpoints, key=lambda x: int(_re_checkpoint.search(x).groups()[0])))225 226 227class IntervalStrategy(ExplicitEnum):228    NO = "no"229    STEPS = "steps"230    EPOCH = "epoch"231 232 233class SaveStrategy(ExplicitEnum):234    NO = "no"235    STEPS = "steps"236    EPOCH = "epoch"237    BEST = "best"238 239 240class EvaluationStrategy(ExplicitEnum):241    NO = "no"242    STEPS = "steps"243    EPOCH = "epoch"244 245 246class HubStrategy(ExplicitEnum):247    END = "end"248    EVERY_SAVE = "every_save"249    CHECKPOINT = "checkpoint"250    ALL_CHECKPOINTS = "all_checkpoints"251 252 253class BestRun(NamedTuple):254    """255    The best run found by a hyperparameter search (see [`~Trainer.hyperparameter_search`]).256 257    Parameters:258        run_id (`str`):259            The id of the best run (if models were saved, the corresponding checkpoint will be in the folder ending260            with run-{run_id}).261        objective (`float`):262            The objective that was obtained for this run.263        hyperparameters (`dict[str, Any]`):264            The hyperparameters picked to get this run.265        run_summary (`Optional[Any]`):266            A summary of tuning experiments. `ray.tune.ExperimentAnalysis` object for Ray backend.267    """268 269    run_id: str270    objective: Union[float, list[float]]271    hyperparameters: dict[str, Any]272    run_summary: Optional[Any] = None273 274 275def default_compute_objective(metrics: dict[str, float]) -> float:276    """277    The default objective to maximize/minimize when doing an hyperparameter search. It is the evaluation loss if no278    metrics are provided to the [`Trainer`], the sum of all metrics otherwise.279 280    Args:281        metrics (`dict[str, float]`): The metrics returned by the evaluate method.282 283    Return:284        `float`: The objective to minimize or maximize285    """286    metrics = copy.deepcopy(metrics)287    loss = metrics.pop("eval_loss", None)288    _ = metrics.pop("epoch", None)289    # Remove speed metrics290    speed_metrics = [291        m for m in metrics if m.endswith("_runtime") or m.endswith("_per_second") or m.endswith("_compilation_time")292    ]293    for sm in speed_metrics:294        _ = metrics.pop(sm, None)295    return loss if len(metrics) == 0 else sum(metrics.values())296 297 298def default_hp_space_optuna(trial) -> dict[str, float]:299    from .integrations import is_optuna_available300 301    assert is_optuna_available(), "This function needs Optuna installed: `pip install optuna`"302    return {303        "learning_rate": trial.suggest_float("learning_rate", 1e-6, 1e-4, log=True),304        "num_train_epochs": trial.suggest_int("num_train_epochs", 1, 5),305        "seed": trial.suggest_int("seed", 1, 40),306        "per_device_train_batch_size": trial.suggest_categorical("per_device_train_batch_size", [4, 8, 16, 32, 64]),307    }308 309 310def default_hp_space_ray(trial) -> dict[str, Any]:311    from .integrations import is_ray_tune_available312 313    assert is_ray_tune_available(), "This function needs ray installed: `pip install ray[tune]`"314    from ray import tune315 316    return {317        "learning_rate": tune.loguniform(1e-6, 1e-4),318        "num_train_epochs": tune.choice(list(range(1, 6))),319        "seed": tune.uniform(1, 40),320        "per_device_train_batch_size": tune.choice([4, 8, 16, 32, 64]),321    }322 323 324def default_hp_space_sigopt(trial):325    return [326        {"bounds": {"min": 1e-6, "max": 1e-4}, "name": "learning_rate", "type": "double", "transformation": "log"},327        {"bounds": {"min": 1, "max": 6}, "name": "num_train_epochs", "type": "int"},328        {"bounds": {"min": 1, "max": 40}, "name": "seed", "type": "int"},329        {330            "categorical_values": ["4", "8", "16", "32", "64"],331            "name": "per_device_train_batch_size",332            "type": "categorical",333        },334    ]335 336 337def default_hp_space_wandb(trial) -> dict[str, Any]:338    from .integrations import is_wandb_available339 340    if not is_wandb_available():341        raise ImportError("This function needs wandb installed: `pip install wandb`")342 343    return {344        "method": "random",345        "metric": {"name": "objective", "goal": "minimize"},346        "parameters": {347            "learning_rate": {"distribution": "uniform", "min": 1e-6, "max": 1e-4},348            "num_train_epochs": {"distribution": "int_uniform", "min": 1, "max": 6},349            "seed": {"distribution": "int_uniform", "min": 1, "max": 40},350            "per_device_train_batch_size": {"values": [4, 8, 16, 32, 64]},351        },352    }353 354 355class HPSearchBackend(ExplicitEnum):356    OPTUNA = "optuna"357    RAY = "ray"358    SIGOPT = "sigopt"359    WANDB = "wandb"360 361 362def is_main_process(local_rank):363    """364    Whether or not the current process is the local process, based on `xr.global_ordinal()` (for TPUs) first, then on365    `local_rank`.366    """367    if is_torch_xla_available():368        import torch_xla.runtime as xr369 370        return xr.global_ordinal() == 0371    return local_rank in [-1, 0]372 373 374def total_processes_number(local_rank):375    """376    Return the number of processes launched in parallel. Works with `torch.distributed` and TPUs.377    """378    if is_torch_xla_available():379        import torch_xla.runtime as xr380 381        return xr.world_size()382    elif local_rank != -1 and is_torch_available():383        import torch384 385        return torch.distributed.get_world_size()386    return 1387 388 389def speed_metrics(split, start_time, num_samples=None, num_steps=None, num_tokens=None):390    """391    Measure and return speed performance metrics.392 393    This function requires a time snapshot `start_time` before the operation to be measured starts and this function394    should be run immediately after the operation to be measured has completed.395 396    Args:397    - split: name to prefix metric (like train, eval, test...)398    - start_time: operation start time399    - num_samples: number of samples processed400    - num_steps: number of steps processed401    - num_tokens: number of tokens processed402    """403    runtime = time.time() - start_time404    result = {f"{split}_runtime": round(runtime, 4)}405    if runtime == 0:406        return result407    if num_samples is not None:408        samples_per_second = num_samples / runtime409        result[f"{split}_samples_per_second"] = round(samples_per_second, 3)410    if num_steps is not None:411        steps_per_second = num_steps / runtime412        result[f"{split}_steps_per_second"] = round(steps_per_second, 3)413    if num_tokens is not None:414        tokens_per_second = num_tokens / runtime415        result[f"{split}_tokens_per_second"] = round(tokens_per_second, 3)416    return result417 418 419class SchedulerType(ExplicitEnum):420    """421    Scheduler names for the parameter `lr_scheduler_type` in [`TrainingArguments`].422    By default, it uses "linear". Internally, this retrieves `get_linear_schedule_with_warmup` scheduler from [`Trainer`].423    Scheduler types:424       - "linear" = [`get_linear_schedule_with_warmup`]425       - "cosine" = [`get_cosine_schedule_with_warmup`]426       - "cosine_with_restarts" = [`get_cosine_with_hard_restarts_schedule_with_warmup`]427       - "polynomial" = [`get_polynomial_decay_schedule_with_warmup`]428       - "constant" =  [`get_constant_schedule`]429       - "constant_with_warmup" = [`get_constant_schedule_with_warmup`]430       - "inverse_sqrt" = [`get_inverse_sqrt_schedule`]431       - "reduce_lr_on_plateau" = [`get_reduce_on_plateau_schedule`]432       - "cosine_with_min_lr" = [`get_cosine_with_min_lr_schedule_with_warmup`]433       - "cosine_warmup_with_min_lr" = [`get_cosine_with_min_lr_schedule_with_warmup_lr_rate`]434       - "warmup_stable_decay" = [`get_wsd_schedule`]435    """436 437    LINEAR = "linear"438    COSINE = "cosine"439    COSINE_WITH_RESTARTS = "cosine_with_restarts"440    POLYNOMIAL = "polynomial"441    CONSTANT = "constant"442    CONSTANT_WITH_WARMUP = "constant_with_warmup"443    INVERSE_SQRT = "inverse_sqrt"444    REDUCE_ON_PLATEAU = "reduce_lr_on_plateau"445    COSINE_WITH_MIN_LR = "cosine_with_min_lr"446    COSINE_WARMUP_WITH_MIN_LR = "cosine_warmup_with_min_lr"447    WARMUP_STABLE_DECAY = "warmup_stable_decay"448 449 450class TrainerMemoryTracker:451    """452    A helper class that tracks cpu and gpu memory.453 454    This class will silently skip unless `psutil` is available. Install with `pip install psutil`.455 456    When a stage completes, it can pass metrics dict to update with the memory metrics gathered during this stage.457 458    Example :459 460    ```python461    self._memory_tracker = TrainerMemoryTracker(self.args.skip_memory_metrics)462    self._memory_tracker.start()463    # code ...464    metrics = {"train_runtime": 10.5}465    self._memory_tracker.stop_and_update_metrics(metrics)466    ```467 468    At the moment GPU tracking is only for `pytorch`, but can be extended to support `tensorflow`.469 470    To understand this class' intricacies please read the documentation of [`~Trainer.log_metrics`].471    """472 473    # map trainer methods to metrics prefix474    stages = {475        "__init__": "init",476        "train": "train",477        "_inner_training_loop": "train",478        "evaluate": "eval",479        "predict": "test",480    }481 482    def __init__(self, skip_memory_metrics=False):483        self.skip_memory_metrics = skip_memory_metrics484 485        if not is_psutil_available():486            # soft dependency on psutil487            self.skip_memory_metrics = True488 489        if self.skip_memory_metrics:490            return491 492        import psutil493 494        if is_torch_cuda_available() or is_torch_mlu_available() or is_torch_musa_available():495            import torch496 497            self.torch = torch498            self.gpu = {}499        elif is_torch_mps_available():500            import torch501 502            self.torch = torch503            self.gpu = {}504        elif is_torch_xpu_available():505            import torch506 507            self.torch = torch508            self.gpu = {}509        elif is_torch_npu_available():510            import torch511 512            self.torch = torch513            self.gpu = {}514        elif is_torch_hpu_available():515            import torch516 517            self.torch = torch518            self.gpu = {}519        else:520            self.torch = None521 522        self.process = psutil.Process()523 524        self.cur_stage = None525        self.cpu = {}526        self.init_reported = False527 528    def derive_stage(self):529        """derives the stage/caller name automatically"""530        caller = inspect.currentframe().f_back.f_back.f_code.co_name531        if caller in self.stages:532            return self.stages[caller]533        else:534            raise ValueError(535                f"was called from {caller}, but only expect to be called from one of {self.stages.keys()}"536            )537 538    def cpu_mem_used(self):539        """get resident set size memory for the current process"""540        return self.process.memory_info().rss541 542    def peak_monitor_func(self):543        self.cpu_mem_used_peak = -1544 545        while True:546            self.cpu_mem_used_peak = max(self.cpu_mem_used(), self.cpu_mem_used_peak)547 548            # can't sleep or will not catch the peak right (this comment is here on purpose)549            # time.sleep(0.001) # 1msec550 551            if not self.peak_monitoring:552                break553 554    def start(self):555        """start tracking for the caller's stage"""556        if self.skip_memory_metrics:557            return558 559        stage = self.derive_stage()560        # deal with nested calls of eval during train - simply ignore those561        if self.cur_stage is not None and self.cur_stage != stage:562            return563 564        self.cur_stage = stage565 566        gc.collect()567 568        if self.torch is not None:569            if torch.cuda.is_available():570                self.torch.cuda.reset_peak_memory_stats()571                self.torch.cuda.empty_cache()572            elif is_torch_mlu_available():573                self.torch.mlu.reset_peak_memory_stats()574                self.torch.mlu.empty_cache()575            elif is_torch_musa_available():576                self.torch.musa.reset_peak_memory_stats()577                self.torch.musa.empty_cache()578            elif is_torch_xpu_available():579                self.torch.xpu.reset_peak_memory_stats()580                self.torch.xpu.empty_cache()581            elif is_torch_npu_available():582                self.torch.npu.reset_peak_memory_stats()583                self.torch.npu.empty_cache()584            elif is_torch_hpu_available():585                self.torch.hpu.reset_peak_memory_stats()586                # not available on hpu as it reserves all device memory for the current process587                # self.torch.hpu.empty_cache()588            elif is_torch_mps_available():589                self.torch.mps.empty_cache()590 591        # gpu592        if self.torch is not None:593            if torch.cuda.is_available():594                self.gpu_mem_used_at_start = self.torch.cuda.memory_allocated()595            elif is_torch_mlu_available():596                self.gpu_mem_used_at_start = self.torch.mlu.memory_allocated()597            elif is_torch_musa_available():598                self.gpu_mem_used_at_start = self.torch.musa.memory_allocated()599            elif is_torch_xpu_available():600                self.gpu_mem_used_at_start = self.torch.xpu.memory_allocated()601            elif is_torch_npu_available():602                self.gpu_mem_used_at_start = self.torch.npu.memory_allocated()603            elif is_torch_hpu_available():604                self.gpu_mem_used_at_start = self.torch.hpu.memory_allocated()605            elif is_torch_mps_available():606                self.gpu_mem_used_at_start = self.torch.mps.current_allocated_memory()607 608        # cpu609        self.cpu_mem_used_at_start = self.cpu_mem_used()610 611        self.peak_monitoring = True612        peak_monitor_thread = threading.Thread(target=self.peak_monitor_func)613        peak_monitor_thread.daemon = True614        peak_monitor_thread.start()615 616    def stop(self, stage):617        """stop tracking for the passed stage"""618 619        # deal with nested calls of eval during train - simply ignore those620        if self.cur_stage is not None and self.cur_stage != stage:621            return622 623        # this sends a signal to peak_monitor_func to complete its loop624        self.peak_monitoring = False625 626        # first ensure all objects get collected and their memory is freed627        gc.collect()628 629        if self.torch is not None:630            if torch.cuda.is_available():631                self.torch.cuda.empty_cache()632            elif is_torch_mlu_available():633                self.torch.mlu.empty_cache()634            elif is_torch_musa_available():635                self.torch.musa.empty_cache()636            elif is_torch_xpu_available():637                self.torch.xpu.empty_cache()638            elif is_torch_npu_available():639                self.torch.npu.empty_cache()640            elif is_torch_hpu_available():641                # not available on hpu as it reserves all device memory for the current process642                # self.torch.npu.empty_cache()643                pass644            elif is_torch_mps_available():645                self.torch.mps.empty_cache()646 647        # concepts:648        # - alloc_delta:  the difference of allocated memory between the end and the start649        # - peaked_delta: the difference between the peak memory and the current memory650        # in order to know how much memory the measured code consumed one needs to sum these two651 652        # gpu653        if self.torch is not None:654            if torch.cuda.is_available():655                self.gpu_mem_used_now = self.torch.cuda.memory_allocated()656                self.gpu_mem_used_peak = self.torch.cuda.max_memory_allocated()657            elif is_torch_mlu_available():658                self.gpu_mem_used_now = self.torch.mlu.memory_allocated()659                self.gpu_mem_used_peak = self.torch.mlu.max_memory_allocated()660            elif is_torch_musa_available():661                self.gpu_mem_used_now = self.torch.musa.memory_allocated()662                self.gpu_mem_used_peak = self.torch.musa.max_memory_allocated()663            elif is_torch_xpu_available():664                self.gpu_mem_used_now = self.torch.xpu.memory_allocated()665                self.gpu_mem_used_peak = self.torch.xpu.max_memory_allocated()666            elif is_torch_npu_available():667                self.gpu_mem_used_now = self.torch.npu.memory_allocated()668                self.gpu_mem_used_peak = self.torch.npu.max_memory_allocated()669            elif is_torch_hpu_available():670                self.gpu_mem_used_now = self.torch.hpu.memory_allocated()671                self.gpu_mem_used_peak = self.torch.hpu.max_memory_allocated()672            elif is_torch_mps_available():673                self.gpu_mem_used_now = self.torch.mps.current_allocated_memory()674                # self.torch.mps.max_memory_allocated() does not exist yet675                self.gpu_mem_used_peak = None676 677            else:678                raise ValueError("No available GPU device found!")679 680            self.gpu[self.cur_stage] = {681                "begin": self.gpu_mem_used_at_start,682                "end": self.gpu_mem_used_now,683                "alloc": (self.gpu_mem_used_now - self.gpu_mem_used_at_start),684            }685            if self.gpu_mem_used_peak is not None:686                self.gpu[self.cur_stage]["peaked"] = max(0, self.gpu_mem_used_peak - self.gpu_mem_used_now)687            else:688                self.gpu[self.cur_stage]["peaked"] = "Not available"689 690        # cpu691        self.cpu_mem_used_now = self.cpu_mem_used()692        self.cpu[self.cur_stage] = {693            "begin": self.cpu_mem_used_at_start,694            "end": self.cpu_mem_used_now,695            "alloc": (self.cpu_mem_used_now - self.cpu_mem_used_at_start),696            "peaked": max(0, self.cpu_mem_used_peak - self.cpu_mem_used_now),697        }698 699        # reset - cycle finished700        self.cur_stage = None701 702    def update_metrics(self, stage, metrics):703        """updates the metrics"""704        if self.skip_memory_metrics:705            return706 707        # deal with nested calls of eval during train - simply ignore those708        if self.cur_stage is not None and self.cur_stage != stage:709            return710 711        # since we don't have a way to return init metrics, we push them into the first of train/val/predict712        stages = [stage]713        if not self.init_reported:714            stages.insert(0, "init")715            self.init_reported = True716 717        for stage in stages:718            for t in ["alloc", "peaked"]:719                if stage in self.cpu and t in self.cpu[stage]:720                    metrics[f"{stage}_mem_cpu_{t}_delta"] = self.cpu[stage][t]721                if self.torch is not None and stage in self.gpu and t in self.gpu[stage]:722                    metrics[f"{stage}_mem_gpu_{t}_delta"] = self.gpu[stage][t]723            # if we need additional debug info, enable the following724            # for t in ["begin", "end"]:725            #     if stage in self.cpu and t in self.cpu[stage]:726            #         metrics[f"{stage}_mem_cpu_{t}"] = self.cpu[stage][t]727            #     if self.torch is not None and stage in self.gpu and t in self.gpu[stage]:728            #         metrics[f"{stage}_mem_gpu_{t}"] = self.gpu[stage][t]729 730        # since memory can be allocated before init, and it might be difficult to track overall731        # memory usage, in particular for GPU, let's report memory usage at the point init was called732        if stages[0] == "init":733            metrics["before_init_mem_cpu"] = self.cpu["init"]["begin"]734            if self.torch is not None:735                metrics["before_init_mem_gpu"] = self.gpu["init"]["begin"]736            # if we also wanted to report any additional memory allocations in between init and737            # whatever the next stage was we could also report this:738            # if self.cpu["init"]["end"] != self.cpu[stage]["begin"]:739            #     metrics[f"after_init_mem_cpu_delta"] = self.cpu[stage]["begin"] - self.cpu["init"]["end"]740            # if self.torch is not None and self.gpu["init"]["end"] != self.gpu[stage]["begin"]:741            #     metrics[f"after_init_mem_gpu_delta"] = self.gpu[stage]["begin"] - self.gpu["init"]["end"]742 743    def stop_and_update_metrics(self, metrics=None):744        """combine stop and metrics update in one call for simpler code"""745        if self.skip_memory_metrics:746            return747 748        stage = self.derive_stage()749        self.stop(stage)750 751        # init doesn't have metrics to update so we just save that data for later stages to retrieve752        if metrics is not None:753            self.update_metrics(stage, metrics)754 755 756def has_length(dataset):757    """758    Checks if the dataset implements __len__() and it doesn't raise an error759    """760    try:761        return len(dataset) is not None762    except TypeError:763        # TypeError: len() of unsized object764        return False765    except AttributeError:766        # Ray DataSets raises an AttributeError: https://github.com/ray-project/ray/blob/master/python/ray/data/dataset.py#L5616767        return False768 769 770def denumpify_detensorize(metrics):771    """772    Recursively calls `.item()` on the element of the dictionary passed773    """774    if isinstance(metrics, (list, tuple)):775        return type(metrics)(denumpify_detensorize(m) for m in metrics)776    elif isinstance(metrics, dict):777        return type(metrics)({k: denumpify_detensorize(v) for k, v in metrics.items()})778    elif isinstance(metrics, np.generic):779        return metrics.item()780    elif is_torch_available() and isinstance(metrics, torch.Tensor) and metrics.numel() == 1:781        return metrics.item()782    return metrics783 784 785def number_of_arguments(func):786    """787    Return the number of arguments of the passed function, even if it's a partial function.788    """789    if isinstance(func, functools.partial):790        total_args = len(inspect.signature(func.func).parameters)791        return total_args - len(func.args) - len(func.keywords)792    return len(inspect.signature(func).parameters)793 794 795def find_executable_batch_size(796    function: Optional[Callable] = None, starting_batch_size: int = 128, auto_find_batch_size: bool = False797):798    """799    Args:800    A basic decorator that will try to execute `function`. If it fails from exceptions related to out-of-memory or801    CUDNN, the batch size is multiplied by 0.9 and passed to `function`. `function` must take in a `batch_size` parameter as802    its first argument.803        function (`Callable`, *optional*)804            A function to wrap805        starting_batch_size (`int`, *optional*)806            The batch size to try and fit into memory807        auto_find_batch_size (`bool`, *optional*)808            If False, will just execute `function`809    """810    if function is None:811        return functools.partial(812            find_executable_batch_size,813            starting_batch_size=starting_batch_size,814            auto_find_batch_size=auto_find_batch_size,815        )816 817    if auto_find_batch_size:818        requires_backends(find_executable_batch_size, "accelerate")819        from accelerate.utils import find_executable_batch_size as accelerate_find_executable_batch_size820 821        return accelerate_find_executable_batch_size(function=function, starting_batch_size=starting_batch_size)822 823    return functools.partial(function, batch_size=starting_batch_size)824 825 826class FSDPOption(ExplicitEnum):827    FULL_SHARD = "full_shard"828    SHARD_GRAD_OP = "shard_grad_op"829    NO_SHARD = "no_shard"830    HYBRID_SHARD = "hybrid_shard"831    HYBRID_SHARD_ZERO2 = "hybrid_shard_zero2"832    OFFLOAD = "offload"833    AUTO_WRAP = "auto_wrap"834 835 836class RemoveColumnsCollator:837    """Wrap the data collator to remove unused columns before they are passed to the collator."""838 839    def __init__(840        self,841        data_collator,842        signature_columns,843        logger=None,844        model_name: Optional[str] = None,845        description: Optional[str] = None,846    ):847        self.data_collator = data_collator848        self.signature_columns = signature_columns849        self.logger = logger850        self.description = description851        self.model_name = model_name852        self.message_logged = False853 854    def _remove_columns(self, feature: dict) -> dict:855        if not isinstance(feature, dict):856            return feature857        if not self.message_logged and self.logger and self.model_name:858            ignored_columns = list(set(feature.keys()) - set(self.signature_columns))859            if len(ignored_columns) > 0:860                dset_description = "" if self.description is None else f"in the {self.description} set"861                self.logger.info(862                    f"The following columns {dset_description} don't have a corresponding argument in "863                    f"`{self.model_name}.forward` and have been ignored: {', '.join(ignored_columns)}."864                    f" If {', '.join(ignored_columns)} are not expected by `{self.model_name}.forward`, "865                    " you can safely ignore this message."866                )867                self.message_logged = True868        return {k: v for k, v in feature.items() if k in self.signature_columns}869 870    def __call__(self, features: list[dict]):871        features = [self._remove_columns(feature) for feature in features]872        return self.data_collator(features)873 874 875def check_target_module_exists(optim_target_modules, key: str, return_is_regex: bool = False):876    """A helper method to check if the passed module's key name matches any of the target modules in the optim_target_modules.877 878    Args:879        optim_target_modules (`Union[str, list[str]]`):880            A list of strings to try to match. Can be also a full string.881        key (`str`):882            A key to search any matches in optim_target_modules883        return_is_regex (`bool`):884            If set to `True`, the method will return whether the passed `optim_target_modules`885            is a regex or not.886 887    Returns:888        `bool` : True of match object if key matches any target modules from config, False or889        None if no match found890        `bool` : If the matched target module is a regex to silence out the warnings in Trainer891        for extra modules being found (only if `target_module_found=True` for an array of regex).892    """893    target_module_found = False894    is_regex = False895 896    if isinstance(optim_target_modules, str):897        target_module_found = bool(re.fullmatch(optim_target_modules, key))898        is_regex = optim_target_modules != key899    elif key in optim_target_modules:  # from here, target_module_found must be a list of str900        # this module is specified directly in target_modules901        target_module_found = True902    elif any(target_key in key for target_key in optim_target_modules):903        target_module_found = True904    elif any(bool(re.fullmatch(optim_target_module, key)) for optim_target_module in optim_target_modules):905        target_module_found = True906        is_regex = True907 908    if return_is_regex:909        return target_module_found, is_regex910 911    return target_module_found912