(self, client_optimizer, model_parameters)
| 1899 | |
| 1900 | # Configure optimizer |
| 1901 | def _configure_optimizer(self, client_optimizer, model_parameters): |
| 1902 | if client_optimizer is None: |
| 1903 | if self.has_moe_layers: |
| 1904 | model_parameters = configure_moe_param_groups(model_parameters) |
| 1905 | basic_optimizer = self._configure_basic_optimizer(model_parameters) |
| 1906 | log_dist(f"Using DeepSpeed Optimizer param name {self.optimizer_name()} as basic optimizer", ranks=[0]) |
| 1907 | else: |
| 1908 | if isinstance(client_optimizer, tuple(self._supported_optims())): |
| 1909 | basic_optimizer = client_optimizer |
| 1910 | log_dist('Using client Optimizer as basic optimizer', ranks=[0]) |
| 1911 | else: |
| 1912 | basic_optimizer = client_optimizer(model_parameters) |
| 1913 | log_dist('Using client callable to create basic optimizer', ranks=[0]) |
| 1914 | |
| 1915 | if (self.zero_use_cpu_optimizer() and not isinstance(basic_optimizer, deepspeed.ops.adam.DeepSpeedCPUAdam) |
| 1916 | and not isinstance(basic_optimizer, deepspeed.ops.lion.DeepSpeedCPULion)): |
| 1917 | if self.zero_force_ds_cpu_optimizer(): |
| 1918 | 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.' |
| 1919 | raise ZeRORuntimeException(msg) |
| 1920 | |
| 1921 | basic_optimizer.param_groups[:] = [pg for pg in basic_optimizer.param_groups if len(pg["params"]) != 0] |
| 1922 | log_dist("Removing param_group that has no 'params' in the basic Optimizer", ranks=[0]) |
| 1923 | |
| 1924 | self._check_for_duplicates(basic_optimizer) |
| 1925 | |
| 1926 | self.basic_optimizer = basic_optimizer |
| 1927 | log_dist(f"DeepSpeed Basic Optimizer = {basic_optimizer.__class__.__name__}", ranks=[0]) |
| 1928 | |
| 1929 | optimizer_wrapper = self._do_optimizer_sanity_check(basic_optimizer) |
| 1930 | |
| 1931 | if optimizer_wrapper == ZERO_OPTIMIZATION: |
| 1932 | self.optimizer = self._configure_zero_optimizer(basic_optimizer) |
| 1933 | elif optimizer_wrapper == AMP: |
| 1934 | amp_params = self.amp_params() |
| 1935 | log_dist(f"Initializing AMP with these params: {amp_params}", ranks=[0]) |
| 1936 | model, self.optimizer = amp.initialize(self.module, basic_optimizer, **amp_params) |
| 1937 | self._set_client_model(model) |
| 1938 | self._broadcast_model() |
| 1939 | # TODO: maybe need to broadcast experts differently? |
| 1940 | elif optimizer_wrapper in [FP16, DDP_BFLOAT16]: |
| 1941 | lp_dtype = torch.float16 if optimizer_wrapper == FP16 else torch.bfloat16 |
| 1942 | self.optimizer = self._configure_fp16_optimizer(basic_optimizer, lp_dtype) |
| 1943 | elif optimizer_wrapper == BFLOAT16: |
| 1944 | self.optimizer = self._configure_bf16_optimizer(basic_optimizer) |
| 1945 | else: |
| 1946 | self.optimizer = basic_optimizer |
| 1947 | |
| 1948 | self._configure_autoep_folding_optimizer_gradient_reduction() |
| 1949 | log_dist("DeepSpeed Final Optimizer = {}".format(self.optimizer.__class__.__name__), ranks=[0]) |
| 1950 | |
| 1951 | self.compression_scheduler = self._configure_compression_scheduler() |
| 1952 | self.quantizer = self._configure_quantization() |
| 1953 | |
| 1954 | def _configure_autoep_folding_optimizer_gradient_reduction(self): |
| 1955 | configure = getattr(self.optimizer, "configure_autoep_folding_tp_gradient_reduction", None) |
no test coverage detected