text: # Copyright 2021 The HuggingFace Team. 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.
  
  import importlib
  import os
  import random
  from collections.abc import Mapping
  from dataclasses import dataclass, field
  from enum import Enum
  from typing import List, Optional, Union
  
  import numpy as np
  import torch
  
  from .state import AcceleratorState, DistributedType, is_deepspeed_available, is_tpu_available
  
  
  if is_tpu_available():
      import torch_xla.core.xla_model as xm
  
  
  def is_boto3_available():
      return importlib.util.find_spec("boto3") is not None
  
  
  def is_sagemaker_available():
      return importlib.util.find_spec("sagemaker") is not None
  
  
  if is_deepspeed_available():
      from deepspeed import DeepSpeedEngine
  
  
  class RNGType(Enum):
      TORCH = "torch"
      CUDA = "cuda"
      XLA = "xla"
      GENERATOR = "generator"
  
  
  @dataclass
  class TensorInformation:
      shape: torch.Size
      dtype: torch.dtype
  
  
  def set_seed(seed: int):
      """
      Helper function for reproducible behavior to set the seed in ``random``, ``numpy``, ``torch``.
  
      Args:
          seed (:obj:`int`): The seed to set.
      """
      random.seed(seed)
      np.random.seed(seed)
      torch.manual_seed(seed)
      torch.cuda.manual_seed_all(seed)
      # ^^ safe to call this function even if cuda is not available
      if is_tpu_available():
          xm.set_rng_state(seed)
  
  
  def synchronize_rng_state(rng_type: Optional[RNGType] = None, generator: Optional[torch.Generator] = None):
      # Get the proper rng state
      if rng_type == RNGType.TORCH:
          rng_state = torch.get_rng_state()
      elif rng_type == RNGType.CUDA:
          rng_state = torch.cuda.get_rng_state()
      elif rng_type == RNGType.XLA:
          assert is_tpu_available(), "Can't synchronize XLA seeds on an environment without TPUs."
          rng_state = torch.tensor(xm.get_rng_state())
      elif rng_type == RNGType.GENERATOR:
          assert generator is not None, "Need a generator to synchronize its seed."
          rng_state = generator.get_state()
  
      # Broadcast the rng state from device 0 to other devices
      state = AcceleratorState()
      if state.distributed_type == DistributedType.TPU:
          rng_state = xm.mesh_reduce("random_seed", rng_state, lambda x: x[0])
      elif state.distributed_type == DistributedType.MULTI_GPU:
          rng_state = rng_state.to(state.device)
          torch.distributed.broadcast(rng_state, 0)
          rng_state = rng_state.cpu()
      elif state.distributed_type == DistributedType.MULTI_CPU:
          torch.distributed.broadcast(rng_state, 0)
  
      # Set the broadcast rng state
      if rng_type == RNGType.TORCH:
          torch.set_rng_state(rng_state)
      elif rng_type == RNGType.CUDA:
          torch.cuda.set_rng_state(rng_state)
      elif rng_type == RNGType.XLA:
          xm.set_rng_state(rng_state.item())
      elif rng_type == RNGType.GENERATOR:
          generator.set_state(rng_state)
  
  
  def synchronize_rng_states(rng_types: List[Union[str, RNGType]], generator: Optional[torch.Generator] = None):
      for rng_type in rng_types:
          synchronize_rng_state(RNGType(rng_type), generator=generator)
  
  
  def honor_type(obj, generator):
      """
      Cast a generator to the same type as obj (list, tuple or namedtuple)
      """
      # There is no direct check whether an object if of type namedtuple sadly, this is a workaround.
      if isinstance(obj, tuple) and hasattr(obj, "_fields"):
          # Can instantiate a namedtuple from a generator directly, contrary to a tuple/list.
          return type(obj)(*list(generator))
      return type(obj)(generator)
  
  
  def is_torch_tensor(tensor):
      return isinstance(tensor, torch.Tensor)
  
  
  def is_tensor_information(tensor_info):
      return isinstance(tensor_info, TensorInformation)
  
  
  def recursively_apply(func, data, *args, test_type=is_torch_tensor, error_on_other_type=False, **kwargs):
      """
      Recursively apply a function on a data structure that is a nested list/tuple/dictionary of a given base type.
  
      Args:
          func (:obj:`callable`):
              The function to recursively apply.
          data (nested list/tuple/dictionary of :obj:`main_type`):
              The data on which to apply :obj:`func`
          *args:
              Positional arguments that will be passed to :obj:`func` when applied on the unpacked data.
          main_type (:obj:`type`, `optional`, defaults to :obj:`torch.Tensor`):
              The base type of the objects to which apply :obj:`func`.
          error_on_other_type (:obj:`bool`, `optional`, defaults to :obj:`False`):
              Whether to return an error or not if after unpacking :obj:`data`, we get on an object that is not of type
              :obj:`main_type`. If :obj:`False`, the function will leave objects of types different than :obj:`main_type`
              unchanged.
          **kwargs:
              Keyword arguments that will be passed to :obj:`func` when applied on the unpacked data.
  
      Returns:
          The same data structure as :obj:`data` with :obj:`func` applied to every object of type :obj:`main_type`.
      """
      if isinstance(data, (tuple, list)):
          return honor_type(
              data,
              (
                  recursively_apply(
                      func, o, *args, test_type=test_type, error_on_other_type=error_on_other_type, **kwargs
                  )
                  for o in data
              ),
          )
      elif isinstance(data, Mapping):
          return type(data)(
              {
                  k: recursively_apply(
                      func, v, *args, test_type=test_type, error_on_other_type=error_on_other_type, **kwargs
                  )
                  for k, v in data.items()
              }
          )
      elif test_type(data):
          return func(data, *args, **kwargs)
      elif error_on_other_type:
          raise TypeError(
              f"Can't apply {func.__name__} on object of type {type(data)}, only of nested list/tuple/dicts of objects "
              f"that satisfy {test_type.__name__}."
          )
      return data
  
  
  def send_to_device(tensor, device):
      """
      Recursively sends the elements in a nested list/tuple/dictionary of tensors to a given device.
  
      Args:
          tensor (nested list/tuple/dictionary of :obj:`torch.Tensor`):
              The data to send to a given device.
          device (:obj:`torch.device`):
              The device to send the data to
  
      Returns:
          The same data structure as :obj:`tensor` with all tensors sent to the proper device.
      """
  
      def _send_to_device(t, device):
          return t.to(device)
  
      def _has_to_method(t):
          return hasattr(t, "to")
  
      return recursively_apply(_send_to_device, tensor, device, test_type=_has_to_method)
  
  
  def get_data_structure(data):
      """
      Recursively gathers the information needed to rebuild a nested list/tuple/dictionary of tensors.
  
      Args:
          data (nested list/tuple/dictionary of :obj:`torch.Tensor`):
              The data to send to analyze.
  
      Returns:
          The same data structure as :obj:`data` with :class:`~accelerate.utils.TensorInformation` instead of tensors.
      """
  
      def _get_data_structure(tensor):
          return TensorInformation(shape=tensor.shape, dtype=tensor.dtype)
  
      return recursively_apply(_get_data_structure, data)
  
  
  def initialize_tensors(data_structure):
      """
      Recursively initializes tensors from a nested list/tuple/dictionary of
      :class:`~accelerate.utils.TensorInformation`.
  
      Returns:
          The same data structure as :obj:`data` with tensors instead of :class:`~accelerate.utils.TensorInformation`.
      """
  
      def _initialize_tensor(tensor_info):
          return torch.empty(*tensor_info.shape, dtype=tensor_info.dtype)
  
      return recursively_apply(_initialize_tensor, data_structure, test_type=is_tensor_information)
  
  
  def convert_to_fp32(tensor):
      """
      Recursively converts the elements nested list/tuple/dictionary of tensors in FP16 precision to FP32.
  
      Args:
          tensor (nested list/tuple/dictionary of :obj:`torch.Tensor`):
              The data to convert from FP16 to FP32.
  
      Returns:
          The same data structure as :obj:`tensor` with all tensors that were in FP16 precision converted to FP32.
      """
  
      def _convert_to_fp32(tensor):
          return tensor.float()
  
      def _is_fp16_tensor(tensor):
          return hasattr(tensor, "dtype") and tensor.dtype == torch.float16
  
      return recursively_apply(_convert_to_fp32, tensor, test_type=_is_fp16_tensor)
  
  
  def convert_outputs_to_fp32(model_forward):
      """
      Decorator to apply to a function outputing tensors (like a model forward pass) that ensures the outputs in FP16
      precision will be convert back to FP32.
  
      Args:
          model_forward (:obj:`Callable`):
              The function which outputs we want to treat.
  
      Returns:
          The same function as :obj:`model_forward` but with converted outputs.
      """
  
      def convert_outputs(*args, **kwargs):
          outputs = model_forward(*args, **kwargs)
          return convert_to_fp32(outputs)
  
      return convert_outputs
  
  
  def extract_model_from_parallel(model):
      """
      Extract a model from its distributed containers.
  
      Args:
          model (:obj:`torch.nn.Module`): The model to extract.
  
      Returns:
          :obj:`torch.nn.Module`: The extracted model.
      """
      options = (torch.nn.parallel.DistributedDataParallel, torch.nn.DataParallel)
      if is_deepspeed_available():
          options += (DeepSpeedEngine,)
  
      while isinstance(model, options):
          model = model.module
      return model
  
  
  def _tpu_gather(tensor, name="gather tensor"):
      if isinstance(tensor, (list, tuple)):
          return honor_type(tensor, (_tpu_gather(t, name=f"{name}_{i}") for i, t in enumerate(tensor)))
      elif isinstance(tensor, Mapping):
          return type(tensor)({k: _tpu_gather(v, name=f"{name}_{k}") for k, v in tensor.items()})
      elif not isinstance(tensor, torch.Tensor):
          raise TypeError(f"Can't gather the values of type {type(tensor)}, only of nested list/tuple/dicts of tensors.")
      if tensor.ndim == 0:
          tensor = tensor.clone()[None]
      return xm.mesh_reduce(name, tensor, torch.cat)
  
  
  def _gpu_gather(tensor):
      def _gpu_gather_one(tensor):
          if tensor.ndim == 0:
              tensor = tensor.clone()[None]
          output_tensors = [tensor.clone() for _ in range(torch.distributed.get_world_size())]
          torch.distributed.all_gather(output_tensors, tensor)
          return torch.cat(output_tensors, dim=0)
  
      return recursively_apply(_gpu_gather_one, tensor, error_on_other_type=True)
  
  
  _cpu_gather = _gpu_gather
  
  
  def gather(tensor):
      """
      Recursively gather tensor in a nested list/tuple/dictionary of tensors from all devices.
  
      Args:
          tensor (nested list/tuple/dictionary of :obj:`torch.Tensor`):
              The data to gather.
  
      Returns:
          The same data structure as :obj:`tensor` with all tensors sent to the proper device.
      """
      if AcceleratorState().distributed_type == DistributedType.TPU:
          return _tpu_gather(tensor, name="accelerate.utils.gather")
      elif AcceleratorState().distributed_type == DistributedType.MULTI_GPU:
          return _gpu_gather(tensor)
      elif AcceleratorState().distributed_type == DistributedType.MULTI_CPU:
          return _cpu_gather(tensor)
      else:
          return tensor
  
  
  def _gpu_broadcast(data, src=0):
      def _gpu_broadcast_one(tensor, src=0):
          torch.distributed.broadcast(tensor, src=src)
          return tensor
  
      return recursively_apply(_gpu_broadcast_one, data, error_on_other_type=True, src=src)
  
  
  def _tpu_broadcast(tensor, src=0, name="broadcast tensor"):
      if isinstance(tensor, (list, tuple)):
          return honor_type(tensor, (_tpu_broadcast(t, name=f"{name}_{i}") for i, t in enumerate(tensor)))
      elif isinstance(tensor, Mapping):
          return type(tensor)({k: _tpu_broadcast(v, name=f"{name}_{k}") for k, v in tensor.items()})
      return xm.mesh_reduce(name, tensor, lambda x: x[src])
  
  
  def broadcast(tensor, from_process: int = 0):
      """
      Recursively broadcast tensor in a nested list/tuple/dictionary of tensors to all devices.
  
      Args:
          tensor (nested list/tuple/dictionary of :obj:`torch.Tensor`):
              The data to gather.
          from_process (:obj:`int`, `optional`, defaults to 0):
              The process from which to send the data
  
      Returns:
          The same data structure as :obj:`tensor` with all tensors broadcasted to the proper device.
      """
      if AcceleratorState().distributed_type == DistributedType.TPU:
          return _tpu_broadcast(tensor, src=from_process, name="accelerate.utils.broadcast")
      elif AcceleratorState().distributed_type == DistributedType.MULTI_GPU:
          return _gpu_broadcast(tensor, src=from_process)
      elif AcceleratorState().distributed_type == DistributedType.MULTI_CPU:
          return _gpu_broadcast(tensor, src=from_process)
      else:
          return tensor
  
  
  def broadcast_object_list(object_list, from_process: int = 0):
      """
      Broadcast a list of picklable objects form one process to the others.
  
      Args:
          object_list (list of picklable objects):
              The list of objects to broadcast. This list will be modified inplace.
          from_process (:obj:`int`, `optional`, defaults to 0):
              The process from which to send the data.
  
      Returns:
          The same list containing the objects from process 0.
      """
      if AcceleratorState().distributed_type == DistributedType.TPU:
          for i, obj in enumerate(object_list):
              object_list[i] = xm.mesh_reduce("accelerate.utils.broadcast_object_list", obj, lambda x: x[from_process])
      elif AcceleratorState().distributed_type == DistributedType.MULTI_GPU:
          torch.distributed.broadcast_object_list(object_list, src=from_process)
      elif AcceleratorState().distributed_type == DistributedType.MULTI_CPU:
          torch.distributed.broadcast_object_list(object_list, src=from_process)
      return object_list
  
  
  def slice_tensors(data, tensor_slice):
      """
      Recursively takes a slice in a nested list/tuple/dictionary of tensors.
  
      Args:
          data (nested list/tuple/dictionary of :obj:`torch.Tensor`):
              The data to slice.
          tensor_slice (:obj:`slice`):
              The slice to take.
  
      Returns:
          The same data structure as :obj:`data` with all the tensors slices.
      """
  
      def _slice_tensor(tensor, tensor_slice):
          return tensor[tensor_slice]
  
      return recursively_apply(_slice_tensor, data, tensor_slice)
  
  
  def find_batch_size(data):
      """
      Recursively finds the batch size in a nested list/tuple/dictionary of lists of tensors.
  
      Args:
          data (nested list/tuple/dictionary of :obj:`torch.Tensor`): The data from which to find the batch size.
  
      Returns:
          :obj:`int`: The batch size.
      """
      if isinstance(data, (tuple, list)):
          return find_batch_size(data[0])
      elif isinstance(data, Mapping):
          for k in data.keys():
              return find_batch_size(data[k])
      elif not isinstance(data, torch.Tensor):
          raise TypeError(f"Can only find the batch size of tensors but got {type(data)}.")
      return data.shape[0]
  
  
  def concatenate(data, dim=0):
      """
      Recursively concatenate the tensors in a nested list/tuple/dictionary of lists of tensors with the same shape.
  
      Args:
          data (nested list/tuple/dictionary of lists of tensors :obj:`torch.Tensor`):
              The data to concatenate.
          dim (:obj:`int`, `optional`, defaults to 0):
              The dimension on which to concatenate.
  
      Returns:
          The same data structure as :obj:`data` with all the tensors concatenated.
      """
      if isinstance(data[0], (tuple, list)):
          return honor_type(data[0], (concatenate([d[i] for d in data], dim=dim) for i in range(len(data[0]))))
      elif isinstance(data[0], Mapping):
          return type(data[0])({k: concatenate([d[k] for d in data], dim=dim) for k in data[0].keys()})
      elif not isinstance(data[0], torch.Tensor):
          raise TypeError(f"Can only concatenate tensors but got {type(data[0])}")
      return torch.cat(data, dim=dim)
  
  
  def pad_across_processes(tensor, dim=0, pad_index=0, pad_first=False):
      """
      Recursively pad the tensors in a nested list/tuple/dictionary of tensors from all devices to the same size so they
      can safely be gathered.
  
      Args:
          tensor (nested list/tuple/dictionary of :obj:`torch.Tensor`):
              The data to gather.
          dim (:obj:`int`, `optional`, defaults to 0):
              The dimension on which to pad.
          pad_index (:obj:`int`, `optional`, defaults to 0):
              The value with which to pad.
          pad_first (:obj:`bool`, `optional`, defaults to :obj:`False`):
              Whether to pad at the beginning or the end.
      """
  
      def _pad_across_processes(tensor, dim=0, pad_index=0, pad_first=False):
          if dim >= len(tensor.shape):
              return tensor
  
          # Gather all sizes
          size = torch.tensor(tensor.shape, device=tensor.device)[None]
          sizes = gather(size).cpu()
          # Then pad to the maximum size
          max_size = max(s[dim] for s in sizes)
          if max_size == tensor.shape[dim]:
              return tensor
  
          old_size = tensor.shape
          new_size = list(old_size)
          new_size[dim] = max_size
          new_tensor = tensor.new_zeros(tuple(new_size)) + pad_index
          if pad_first:
              indices = tuple(
                  slice(max_size - old_size[dim], max_size) if i == dim else slice(None) for i in range(len(new_size))
              )
          else:
              indices = tuple(slice(0, old_size[dim]) if i == dim else slice(None) for i in range(len(new_size)))
          new_tensor[indices] = tensor
          return new_tensor
  
      return recursively_apply(
          _pad_across_processes, tensor, error_on_other_type=True, dim=dim, pad_index=pad_index, pad_first=pad_first
      )
  
  
  def wait_for_everyone():
      """
      Introduces a blocking point in the script, making sure all processes have reached this point before continuing.
  
      Warning::
  
          Make sure all processes will reach this instruction otherwise one of your processes will hang forever.
      """
      if (
          AcceleratorState().distributed_type == DistributedType.MULTI_GPU
          or AcceleratorState().distributed_type == DistributedType.MULTI_CPU
          or AcceleratorState().distributed_type == DistributedType.DEEPSPEED
      ):
          torch.distributed.barrier()
      elif AcceleratorState().distributed_type == DistributedType.TPU:
          xm.rendezvous("accelerate.utils.wait_for_everyone")
  
  
  def save(obj, f):
      """
      Save the data to disk. Use in place of :obj:`torch.save()`.
  
      Args:
          obj: The data to save
          f: The file (or file-like object) to use to save the data
      """
      if AcceleratorState().distributed_type == DistributedType.TPU:
          xm.save(obj, f)
      elif AcceleratorState().local_process_index == 0:
          torch.save(obj, f)
  
  
  class PrepareForLaunch:
      """
      Prepare a function that will launched in a distributed setup.
  
      Args:
          launcher (:obj:`Callable`):
              The function to launch.
          distributed_type (:class:`~accelerate.state.DistributedType`):
              The distributed type to prepare for.
      """
  
      def __init__(self, launcher, distributed_type="NO"):
          self.launcher = launcher
          self.distributed_type = DistributedType(distributed_type)
  
      def __call__(self, index, *args):
          if self.distributed_type == DistributedType.MULTI_GPU or self.distributed_type == DistributedType.MULTI_CPU:
              # Prepare the environment for torch.distributed
              os.environ["LOCAL_RANK"] = str(index)
              os.environ["RANK"] = str(index)
  
          self.launcher(*args)
  
  
  @dataclass
  class DeepSpeedPlugin:
  
      gradient_accumulation_steps: int = field(
          default=None, metadata={"help": "Number of steps to accumulate gradients before updating optimizer states"}
      )
      zero_stage: int = field(
          default=None,
          metadata={"help": "Possible options are 0,1,2,3; Default will be taken from environment variable"},
      )
      is_train_batch_min: str = field(
          default=True,
          metadata={"help": "If both train & eval dataloaders are specified, this will decide the train_batch_size"},
      )
  
      auto_opt_mapping: bool = field(
          default=True,
          metadata={"help": "whether to map torch.adam to deepspeed optimizer version of adam based on config"},
      )
  
      offload_optimizer_device: bool = field(default=None, metadata={"help": "Possible options are none|cpu|nvme"})
  
      def __post_init__(self):
  
          if self.gradient_accumulation_steps is None:
              self.gradient_accumulation_steps = int(os.environ.get("GRADIENT_ACCUMULATION_STEPS", 1))
  
          if self.zero_stage is None:
              self.zero_stage = int(os.environ.get("DEEPSPEED_ZERO_STAGE", 2))
  
          if self.offload_optimizer_device is None:
              self.offload_optimizer_device = os.environ.get("DEEPSPEED_OFFLOAD_OPTIMIZER_DEVICE", "none")
  
          self.deepspeed_config = {
              "train_batch_size": None,
              "gradient_accumulation_steps": self.gradient_accumulation_steps,
              "zero_optimization": {
                  "stage": self.zero_stage,
                  "offload_optimizer": {
                      "device": self.offload_optimizer_device,
                  },
              },
              "steps_per_print": float("inf"),  # this will stop deepspeed from logging @ stdout
              "zero_allow_untested_optimizer": True,
          }
  