






'''simple docstring'''


import importlib.util
import weakref
from functools import partialmethod

from .dependency_versions_check import dep_version_check
from .utils import is_accelerate_available, is_torch_available, logging


if is_torch_available():
  import torch

SCREAMING_SNAKE_CASE_				      =							logging.get_logger(__name__)






def   lowerCAmelCase__			(					):
 return importlib.util.find_spec('deepspeed'		) is not None


if is_accelerate_available() and is_deepspeed_available():
  from accelerate.utils.deepspeed import HfDeepSpeedConfig as DeepSpeedConfig
else:
  # Inherits from a dummy `object` if accelerate is not available, so that python succeeds to import this file.
  # Deepspeed glue code will never inherit this dummy object as it checks if accelerate is available.
  from builtins import object as DeepSpeedConfig




class 							lowerCAmelCase					(							lowerCamelCase__   ):







   """simple docstring"""







   def __init__( self   ,   _A			)    ->   List[Any]:
    # set global weakref object
    set_hf_deepspeed_config(self			)
    dep_version_check('accelerate'			)
    dep_version_check('deepspeed'			)
    super().__init__(_A			)




class 							lowerCAmelCase					(							lowerCamelCase__   ):







   """simple docstring"""







   def __init__( self   ,   _A			)    ->   Dict:
    super().__init__(_A			)
    __a    :	int									=						None
    __a    :	Any									=						[]



   def    __magic_name__    ( self			)    ->   Any:
    if self._dtype is None:
     raise ValueError('trainer_config_process() wasn\'t called yet to tell dtype'			)
    return self._dtype



   def    __magic_name__    ( self   ,   _A			)    ->   int:
    __a    :	Optional[int]									=						self.get_value(_A			)
    if val is None:
     return False
    else:
     return val == "auto"



   def    __magic_name__    ( self   ,   _A   ,   _A   ,   _A=None   ,   _A=True			)    ->   int:
    __a    :	Union[str, Any]									=						self.find_config_node(_A			)
    if config is None:
     return

    if config.get(_A			) == "auto":
     __a    :	Optional[Any]									=						hf_val
     return

    if not must_match:
     return

    __a    :	Optional[Any]									=						config.get(_A			)
    if ds_val is not None and ds_val != hf_val:
     self.mismatches.append(f'''- ds {ds_key_long}={ds_val} vs hf {hf_key}={hf_val}'''			)

   _A						=   partialmethod(lowerCamelCase__							,      must_match=lowerCamelCase__   )



   def    __magic_name__    ( self   ,   _A			)    ->   Optional[Any]:
    __a    :	str									=						args.world_size * args.per_device_train_batch_size * args.gradient_accumulation_steps
    self.fill_match(
        'train_micro_batch_size_per_gpu'   ,   args.per_device_train_batch_size   ,   'per_device_train_batch_size'			)
    self.fill_match('gradient_accumulation_steps'   ,   args.gradient_accumulation_steps   ,   'gradient_accumulation_steps'			)
    self.fill_match('train_batch_size'   ,   _A   ,   'train_batch_size (calculated)'			)
    self.fill_match('gradient_clipping'   ,   args.max_grad_norm   ,   'max_grad_norm'			)

    self.fill_match('optimizer.params.lr'   ,   args.learning_rate   ,   'learning_rate'			)
    self.fill_match('optimizer.params.betas'   ,   [args.adam_betaa, args.adam_betaa]   ,   'adam_beta1+adam_beta2'			)
    self.fill_match('optimizer.params.eps'   ,   args.adam_epsilon   ,   'adam_epsilon'			)
    self.fill_match('optimizer.params.weight_decay'   ,   args.weight_decay   ,   'weight_decay'			)

    self.fill_only('scheduler.params.warmup_min_lr'   ,   0			)  # not a trainer arg
    self.fill_match('scheduler.params.warmup_max_lr'   ,   args.learning_rate   ,   'learning_rate'			)
    # total_num_steps - will get set in trainer_config_finalize

    # fp16
    if args.fpaa or args.fpaa_full_eval:
     __a    :	Dict									=						'apex' if args.fpaa_backend == 'apex' else 'amp'
    else:
     __a    :	Optional[int]									=						None

    if args.save_on_each_node:
     # deepspeed uses shared storage by default. Let's override this setting if save_on_each_node == True
     __a    :	Dict									=						self.config.get('checkpoint'   ,   {}			)
     __a    :	str									=						args.save_on_each_node

    # amp: similar to the pytorch native amp - it has a bunch of optional params but we won't set
    # any here unless the user did the work
    self.fill_match(
        'fp16.enabled'   ,   ((args.fpaa or args.fpaa_full_eval) and fpaa_backend == 'amp')   ,   'fp16|fp16_full_eval+fp16_backend(amp)'   ,   )

    # apex: delegates amp work to apex (which needs to be available), but it cannot be used with any
    # ZeRO features
    self.fill_match('amp.enabled'   ,   fpaa_backend == 'apex'   ,   'fp16+fp16_backend(apex)'			)
    self.fill_match('amp.opt_level'   ,   args.fpaa_opt_level   ,   'fp16_opt_level'			)

    self.fill_match('bf16.enabled'   ,   (args.bfaa or args.bfaa_full_eval)   ,   'bf16|bf16_full_eval'			)

    # deepspeed's default mode is fp16 unless there is a config that says differently
    if self.is_true('bf16.enabled'			):
     __a    :	int									=						torch.bfloataa
    elif self.is_false('fp16.enabled'			):
     __a    :	Optional[int]									=						torch.floataa
    else:
     __a    :	Tuple									=						torch.floataa



   def    __magic_name__    ( self   ,   _A   ,   _A   ,   _A			)    ->   int:
    __a    :	List[str]									=						[
        'zero_optimization.reduce_bucket_size',
        'zero_optimization.stage3_prefetch_bucket_size',
        'zero_optimization.stage3_param_persistence_threshold',
    ]
    __a    :	Optional[int]									=						[x for x in hidden_size_based_keys if self.is_auto(_A			)]

    if len(_A			) > 0:
     if hasattr(model.config   ,   'hidden_size'			):
      __a    :	Dict									=						model.config.hidden_size
     elif hasattr(model.config   ,   'hidden_sizes'			):
      # if there are many hidden sizes pick the largest one
      __a    :	Union[str, Any]									=						max(model.config.hidden_sizes			)
     else:
      raise ValueError(
          'The model\'s config file has neither `hidden_size` nor `hidden_sizes` entry, '
          'therefore it\'s not possible to automatically fill out the following `auto` entries '
          f'''in the DeepSpeed config file: {hidden_size_auto_keys}. You can fix that by replacing '''
          '`auto` values for these keys with an integer value of your choice.'			)

     self.fill_only('zero_optimization.reduce_bucket_size'   ,   hidden_size * hidden_size			)
     if self.is_zeroa():
      # automatically assign the optimal config values based on model config
      self.fill_only('zero_optimization.stage3_prefetch_bucket_size'   ,   0.9 * hidden_size * hidden_size			)
      self.fill_only('zero_optimization.stage3_param_persistence_threshold'   ,   10 * hidden_size			)

        # scheduler
    self.fill_match('scheduler.params.total_num_steps'   ,   _A   ,   'num_training_steps (calculated)'			)
    self.fill_match('scheduler.params.warmup_num_steps'   ,   args.get_warmup_steps(_A			)   ,   'warmup_steps'			)

    if len(self.mismatches			) > 0:
     __a    :	Union[str, Any]									=						'\n'.join(self.mismatches			)
     raise ValueError(
         'Please correct the following DeepSpeed config values that mismatch TrainingArguments'
         f''' values:\n{mismatches}\nThe easiest method is to set these DeepSpeed config values to \'auto\'.'''			)
SCREAMING_SNAKE_CASE_				      =							None






def   lowerCAmelCase__			(					SCREAMING_SNAKE_CASE__		):
 # this is a special weakref global object to allow us to get to Deepspeed config from APIs
 # that don't have an easy way to get to the Deepspeed config outside of the Trainer domain.
 global _hf_deepspeed_config_weak_ref
 # will go away automatically when HfDeepSpeedConfig is destroyed (when TrainingArguments is destroyed)
 __a    :	List[Any]									=						weakref.ref(SCREAMING_SNAKE_CASE__		)






def   lowerCAmelCase__			(					):
 # useful for unit tests to ensure the global state doesn't leak - call from `tearDown` method
 global _hf_deepspeed_config_weak_ref
 __a    :	Union[str, Any]									=						None






def   lowerCAmelCase__			(					):
 if _hf_deepspeed_config_weak_ref is not None and _hf_deepspeed_config_weak_ref() is not None:
  return _hf_deepspeed_config_weak_ref().is_zeroa()
 else:
  return False






def   lowerCAmelCase__			(					):
 if _hf_deepspeed_config_weak_ref is not None and _hf_deepspeed_config_weak_ref() is not None:
  return _hf_deepspeed_config_weak_ref().config
 else:
  return None






def   lowerCAmelCase__			(					SCREAMING_SNAKE_CASE__       ,		SCREAMING_SNAKE_CASE__       ,		SCREAMING_SNAKE_CASE__       ,		SCREAMING_SNAKE_CASE__       ,		SCREAMING_SNAKE_CASE__		):
 from accelerate.utils import DummyOptim, DummyScheduler

 __a    :	Any									=						hf_deepspeed_config.config

 # Optimizer + Scheduler
 # Currently supported combos:
 # 1. DS scheduler + DS optimizer: Yes
 # 2. HF scheduler + HF optimizer: Yes
 # 3. DS scheduler + HF optimizer: Yes
 # 4. HF scheduler + DS optimizer: No
 #
 # Unless Offload is enabled in which case it's:
 # 1. DS scheduler + DS optimizer: Yes
 # 2. HF scheduler + HF optimizer: Mostly*
 # 3. DS scheduler + HF optimizer: Mostly*
 # 4. HF scheduler + DS optimizer: No
 #
 # Mostly*: All non-native DeepSpeed optimizers that have both CPU and GPU implementation should work (except LAMB)

 __a    :	int									=						None
 if "optimizer" in config:
  if args.adafactor:
   raise ValueError(
       '--adafactor was passed, but also found `optimizer` configured in the DeepSpeed config. '
       'Only one optimizer can be configured.'		)
  __a    :	List[Any]									=						DummyOptim(params=SCREAMING_SNAKE_CASE__		)
 else:
  if hf_deepspeed_config.is_offload():
   logger.info(
       'Detected ZeRO Offload and non-DeepSpeed optimizers: This combination should work as long as the'
       ' custom optimizer has both CPU and GPU implementation (except LAMB)'		)

  # ds supports Adam, OneBitAdam, and Lamb optimizers and can import other optimizers from torch.
  # But trainer uses AdamW by default.
  __a    :	List[Any]									=						trainer.create_optimizer()
  # To use other optimizers requires voiding warranty with: `zero_allow_untested_optimizer`
  __a    :	Optional[int]									=						True

 __a    :	str									=						None
 if "scheduler" in config:
  __a    :	str									=						DummyScheduler(SCREAMING_SNAKE_CASE__		)
 else:
  if isinstance(SCREAMING_SNAKE_CASE__       ,		SCREAMING_SNAKE_CASE__		):
   raise ValueError(
       'Found `optimizer` configured in the DeepSpeed config, but no `scheduler`. '
       'Please configure a scheduler in the DeepSpeed config.'		)
  __a    :	Optional[int]									=						trainer.create_scheduler(num_training_steps=SCREAMING_SNAKE_CASE__       ,		optimizer=SCREAMING_SNAKE_CASE__		)

 return optimizer, lr_scheduler






def   lowerCAmelCase__			(					SCREAMING_SNAKE_CASE__       ,		SCREAMING_SNAKE_CASE__       ,		SCREAMING_SNAKE_CASE__=False		):
 from deepspeed.utils import logger as ds_logger

 __a    :	List[Any]									=						trainer.model
 __a    :	int									=						trainer.args

 __a    :	Union[str, Any]									=						trainer.accelerator.state.deepspeed_plugin.hf_ds_config

 # resume config update - some bits like `model` and `num_training_steps` only become available during train
 hf_deepspeed_config.trainer_config_finalize(SCREAMING_SNAKE_CASE__       ,		SCREAMING_SNAKE_CASE__       ,		SCREAMING_SNAKE_CASE__		)

 # set the Deepspeed log level consistent with the Trainer
 ds_logger.setLevel(args.get_process_log_level()		)

 if inference:
  # only Z3 makes sense for the inference
  if not hf_deepspeed_config.is_zeroa():
   raise ValueError('ZeRO inference only makes sense with ZeRO Stage 3 - please adjust your config'		)

  # in case the training config is re-used for inference
  hf_deepspeed_config.del_config_sub_tree('optimizer'		)
  hf_deepspeed_config.del_config_sub_tree('lr_scheduler'		)
  __a    :	Optional[int]									=						None, None
  __a    :	Optional[int]									=						None
 else:
  __a    :	List[str]									=						None  # important for when deepspeed_init is used as re-init
  __a    :	Optional[int]									=						list(filter(lambda SCREAMING_SNAKE_CASE__		: p.requires_grad       ,		model.parameters()		)		)
  __a    :	Tuple									=						deepspeed_optim_sched(
      SCREAMING_SNAKE_CASE__       ,		SCREAMING_SNAKE_CASE__       ,		SCREAMING_SNAKE_CASE__       ,		SCREAMING_SNAKE_CASE__       ,		SCREAMING_SNAKE_CASE__		)

 # keep for quick debug:
 # from pprint import pprint; pprint(config)

 return optimizer, lr_scheduler






def   lowerCAmelCase__			(					SCREAMING_SNAKE_CASE__       ,		SCREAMING_SNAKE_CASE__		):
 # it's possible that the user is trying to resume from model_path, which doesn't necessarily
 # contain a deepspeed checkpoint. e.g. examples just check if the dir exists and assume it's
 # a resume from a checkpoint and not just a local pretrained weight. So we check here if the
 # path contains what looks like a deepspeed checkpoint
 import glob

 __a    :	Tuple									=						sorted(glob.glob(f'''{checkpoint_path}/global_step*'''		)		)

 if len(SCREAMING_SNAKE_CASE__		) > 0:
  logger.info(f'''Attempting to resume from {checkpoint_path}'''		)
  # this magically updates self.optimizer and self.lr_scheduler
  __a    :	List[str]									=						deepspeed_engine.load_checkpoint(
      SCREAMING_SNAKE_CASE__       ,		load_optimizer_states=SCREAMING_SNAKE_CASE__       ,		load_lr_scheduler_states=SCREAMING_SNAKE_CASE__		)
  if load_path is None:
   raise ValueError(f'''[deepspeed] failed to resume from checkpoint {checkpoint_path}'''		)
 else:
  raise ValueError(f'''Can\'t find a valid checkpoint at {checkpoint_path}'''		)












'''simple docstring'''


import unittest

import numpy as np
import torch

from diffusers import DDPMPipeline, DDPMScheduler, UNetaDModel
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, slow, torch_device


enable_full_determinism()




class 							lowerCAmelCase					(							unittest.TestCase   ):







   """simple docstring"""







   @property
   def    __magic_name__    ( self			)    ->   Any:
    torch.manual_seed(0			)
    __a    :	Optional[Any]									=						UNetaDModel(
        block_out_channels=(32, 64)   ,   layers_per_block=2   ,   sample_size=32   ,   in_channels=3   ,   out_channels=3   ,   down_block_types=('DownBlock2D', 'AttnDownBlock2D')   ,   up_block_types=('AttnUpBlock2D', 'UpBlock2D')   ,   )
    return model



   def    __magic_name__    ( self			)    ->   Tuple:
    __a    :	Union[str, Any]									=						'cpu'
    __a    :	Tuple									=						self.dummy_uncond_unet
    __a    :	str									=						DDPMScheduler()

    __a    :	int									=						DDPMPipeline(unet=_A   ,   scheduler=_A			)
    ddpm.to(_A			)
    ddpm.set_progress_bar_config(disable=_A			)

    __a    :	str									=						torch.Generator(device=_A			).manual_seed(0			)
    __a    :	int									=						ddpm(generator=_A   ,   num_inference_steps=2   ,   output_type='numpy'			).images

    __a    :	Union[str, Any]									=						torch.Generator(device=_A			).manual_seed(0			)
    __a    :	Optional[int]									=						ddpm(generator=_A   ,   num_inference_steps=2   ,   output_type='numpy'   ,   return_dict=_A			)[0]

    __a    :	List[str]									=						image[0, -3:, -3:, -1]
    __a    :	List[str]									=						image_from_tuple[0, -3:, -3:, -1]

    assert image.shape == (1, 32, 32, 3)
    __a    :	List[Any]									=						np.array(
        [9.9_5_6E-0_1, 5.7_8_5E-0_1, 4.6_7_5E-0_1, 9.9_3_0E-0_1, 0.0, 1.000, 1.1_9_9E-0_3, 2.6_4_8E-0_4, 5.1_0_1E-0_4]			)

    assert np.abs(image_slice.flatten() - expected_slice			).max() < 1E-2
    assert np.abs(image_from_tuple_slice.flatten() - expected_slice			).max() < 1E-2



   def    __magic_name__    ( self			)    ->   Tuple:
    __a    :	Optional[Any]									=						self.dummy_uncond_unet
    __a    :	List[str]									=						DDPMScheduler(prediction_type='sample'			)

    __a    :	List[str]									=						DDPMPipeline(unet=_A   ,   scheduler=_A			)
    ddpm.to(_A			)
    ddpm.set_progress_bar_config(disable=_A			)

    __a    :	Optional[Any]									=						torch.manual_seed(0			)
    __a    :	int									=						ddpm(generator=_A   ,   num_inference_steps=2   ,   output_type='numpy'			).images

    __a    :	Union[str, Any]									=						torch.manual_seed(0			)
    __a    :	int									=						ddpm(generator=_A   ,   num_inference_steps=2   ,   output_type='numpy'			)[0]

    __a    :	Dict									=						image[0, -3:, -3:, -1]
    __a    :	int									=						image_eps[0, -3:, -3:, -1]

    assert image.shape == (1, 32, 32, 3)
    __a    :	int									=						1E-2 if torch_device != 'mps' else 3E-2
    assert np.abs(image_slice.flatten() - image_eps_slice.flatten()			).max() < tolerance




@slow
@require_torch_gpu
class 							lowerCAmelCase					(							unittest.TestCase   ):







   """simple docstring"""







   def    __magic_name__    ( self			)    ->   List[str]:
    __a    :	Union[str, Any]									=						'google/ddpm-cifar10-32'

    __a    :	str									=						UNetaDModel.from_pretrained(_A			)
    __a    :	List[str]									=						DDPMScheduler.from_pretrained(_A			)

    __a    :	List[str]									=						DDPMPipeline(unet=_A   ,   scheduler=_A			)
    ddpm.to(_A			)
    ddpm.set_progress_bar_config(disable=_A			)

    __a    :	int									=						torch.manual_seed(0			)
    __a    :	Optional[Any]									=						ddpm(generator=_A   ,   output_type='numpy'			).images

    __a    :	int									=						image[0, -3:, -3:, -1]

    assert image.shape == (1, 32, 32, 3)
    __a    :	Optional[Any]									=						np.array([0.4_200, 0.3_588, 0.1_939, 0.3_847, 0.3_382, 0.2_647, 0.4_155, 0.3_582, 0.3_385]			)
    assert np.abs(image_slice.flatten() - expected_slice			).max() < 1E-2




