(self, save_dir, tag, client_state={}, exclude_frozen_parameters=False)
| 4851 | ds_version=version) |
| 4852 | |
| 4853 | def _save_moe_checkpoint(self, save_dir, tag, client_state={}, exclude_frozen_parameters=False): |
| 4854 | save_path = self._get_ckpt_name(save_dir, tag) |
| 4855 | |
| 4856 | try: |
| 4857 | from deepspeed.module_inject.auto_ep_layer import AutoEPMoELayer as _AutoEPMoELayer |
| 4858 | except ImportError: |
| 4859 | _AutoEPMoELayer = None |
| 4860 | |
| 4861 | folding_spec = getattr(self, "_autoep_folding_spec", None) |
| 4862 | folded_autoep_tp = folding_spec is not None and folding_spec.tp_size > 1 |
| 4863 | |
| 4864 | def folding_metadata(*, |
| 4865 | family, |
| 4866 | ep_rank, |
| 4867 | zero_partition_group, |
| 4868 | zero_partition_rank, |
| 4869 | zero_partition_count, |
| 4870 | param_families=None): |
| 4871 | if not folded_autoep_tp: |
| 4872 | return None |
| 4873 | return DeepSpeedEngine._make_autoep_folding_metadata(folding_spec, |
| 4874 | family=family, |
| 4875 | ep_rank=ep_rank, |
| 4876 | zero_partition_group=zero_partition_group, |
| 4877 | zero_partition_rank=zero_partition_rank, |
| 4878 | zero_partition_count=zero_partition_count, |
| 4879 | param_families=param_families) |
| 4880 | |
| 4881 | def autoep_expert_writer() -> bool: |
| 4882 | if folded_autoep_tp: |
| 4883 | return groups._get_data_parallel_rank() < folding_spec.ep_size |
| 4884 | return self.checkpoint_engine.is_data_parallel_writer(exp_dp_rank) |
| 4885 | |
| 4886 | # A hack to save the checkpointing directory. Pipeline parallelism overrides |
| 4887 | # module_state_dict() and uses this path to save the model. module_state_dict() |
| 4888 | # then instead just returns None. |
| 4889 | |
| 4890 | # Using layer_#_export_# to save the model's expert state_dict |
| 4891 | autoep_layer_info = [] |
| 4892 | autoep_group_names = set() |
| 4893 | moe_layer_id = 0 |
| 4894 | found_native_moe = False |
| 4895 | found_autoep = False |
| 4896 | for n_module, module in self.module.named_modules(): |
| 4897 | if isinstance(module, MoE): # and deepspeed.comm.get_rank() == 0: |
| 4898 | found_native_moe = True |
| 4899 | if self.zero_optimization_partition_weights() and found_autoep: |
| 4900 | raise RuntimeError("AutoEP with ZeRO Stage 3 checkpointing does not support models that also " |
| 4901 | "contain native DeepSpeed MoE layers.") |
| 4902 | group_name = module.expert_group_name |
| 4903 | num_local_experts = module.num_local_experts |
| 4904 | expp_rank = groups._get_expert_parallel_rank(group_name) |
| 4905 | exp_dp_rank = groups._get_expert_data_parallel_rank(group_name) |
| 4906 | # print(expp_rank, exp_dp_rank) |
| 4907 | # if exp_dp_rank != 0: |
| 4908 | if not self.checkpoint_engine.is_data_parallel_writer(exp_dp_rank): |
| 4909 | moe_layer_id += 1 |
| 4910 | continue |
no test coverage detected