(self, client_optimizer, model_parameters)
| 1980 | |
| 1981 | # Configure optimizer |
| 1982 | def _configure_optimizer(self, client_optimizer, model_parameters): |
| 1983 | if client_optimizer is None: |
| 1984 | if self.has_moe_layers: |
| 1985 | model_parameters = configure_moe_param_groups(model_parameters) |
| 1986 | basic_optimizer = self._configure_basic_optimizer(model_parameters) |
| 1987 | log_dist(f"Using DeepSpeed Optimizer param name {self.optimizer_name()} as basic optimizer", ranks=[0]) |
| 1988 | else: |
| 1989 | if isinstance(client_optimizer, tuple(self._supported_optims())): |
| 1990 | basic_optimizer = client_optimizer |
| 1991 | log_dist('Using client Optimizer as basic optimizer', ranks=[0]) |
| 1992 | else: |
| 1993 | basic_optimizer = client_optimizer(model_parameters) |
| 1994 | log_dist('Using client callable to create basic optimizer', ranks=[0]) |
| 1995 | |
| 1996 | if (self.zero_use_cpu_optimizer() and not isinstance(basic_optimizer, deepspeed.ops.adam.DeepSpeedCPUAdam) |
| 1997 | and not isinstance(basic_optimizer, deepspeed.ops.lion.DeepSpeedCPULion)): |
| 1998 | if self.zero_force_ds_cpu_optimizer(): |
| 1999 | msg = f'You are using ZeRO-Offload with a client provided optimizer ({type(basic_optimizer)}) which in most cases will yield poor performance. Please either use deepspeed.ops.adam.DeepSpeedCPUAdam or set an optimizer in your ds-config (https://www.deepspeed.ai/docs/config-json/#optimizer-parameters). If you really want to use a custom optimizer w. ZeRO-Offload and understand the performance impacts you can also set <"zero_force_ds_cpu_optimizer": false> in your configuration file.' |
| 2000 | raise ZeRORuntimeException(msg) |
| 2001 | |
| 2002 | basic_optimizer.param_groups[:] = [pg for pg in basic_optimizer.param_groups if len(pg["params"]) != 0] |
| 2003 | log_dist("Removing param_group that has no 'params' in the basic Optimizer", ranks=[0]) |
| 2004 | |
| 2005 | self._check_for_duplicates(basic_optimizer) |
| 2006 | |
| 2007 | self.basic_optimizer = basic_optimizer |
| 2008 | log_dist(f"DeepSpeed Basic Optimizer = {basic_optimizer.__class__.__name__}", ranks=[0]) |
| 2009 | |
| 2010 | optimizer_wrapper = self._do_optimizer_sanity_check(basic_optimizer) |
| 2011 | |
| 2012 | if optimizer_wrapper == ZERO_OPTIMIZATION: |
| 2013 | self.optimizer = self._configure_zero_optimizer(basic_optimizer) |
| 2014 | elif optimizer_wrapper == AMP: |
| 2015 | amp_params = self.amp_params() |
| 2016 | log_dist(f"Initializing AMP with these params: {amp_params}", ranks=[0]) |
| 2017 | model, self.optimizer = amp.initialize(self.module, basic_optimizer, **amp_params) |
| 2018 | self._set_client_model(model) |
| 2019 | self._broadcast_model() |
| 2020 | # TODO: maybe need to broadcast experts differently? |
| 2021 | elif optimizer_wrapper in [FP16, DDP_BFLOAT16]: |
| 2022 | lp_dtype = torch.float16 if optimizer_wrapper == FP16 else torch.bfloat16 |
| 2023 | self.optimizer = self._configure_fp16_optimizer(basic_optimizer, lp_dtype) |
| 2024 | elif optimizer_wrapper == BFLOAT16: |
| 2025 | self.optimizer = self._configure_bf16_optimizer(basic_optimizer) |
| 2026 | else: |
| 2027 | self.optimizer = basic_optimizer |
| 2028 | |
| 2029 | self._configure_autoep_folding_optimizer_gradient_reduction() |
| 2030 | log_dist("DeepSpeed Final Optimizer = {}".format(self.optimizer.__class__.__name__), ranks=[0]) |
| 2031 | |
| 2032 | self.compression_scheduler = self._configure_compression_scheduler() |
| 2033 | self.quantizer = self._configure_quantization() |
| 2034 | |
| 2035 | def _configure_autoep_folding_optimizer_gradient_reduction(self): |
| 2036 | configure = getattr(self.optimizer, "configure_autoep_folding_tp_gradient_reduction", None) |
no test coverage detected