(self, save_dir, tag, client_state={}, exclude_frozen_parameters=False)
| 4726 | ds_version=version) |
| 4727 | |
| 4728 | def _save_moe_checkpoint(self, save_dir, tag, client_state={}, exclude_frozen_parameters=False): |
| 4729 | save_path = self._get_ckpt_name(save_dir, tag) |
| 4730 | |
| 4731 | try: |
| 4732 | from deepspeed.module_inject.auto_ep_layer import AutoEPMoELayer as _AutoEPMoELayer |
| 4733 | except ImportError: |
| 4734 | _AutoEPMoELayer = None |
| 4735 | |
| 4736 | folding_spec = getattr(self, "_autoep_folding_spec", None) |
| 4737 | folded_autoep_tp = folding_spec is not None and folding_spec.tp_size > 1 |
| 4738 | |
| 4739 | def folding_metadata(*, |
| 4740 | family, |
| 4741 | ep_rank, |
| 4742 | zero_partition_group, |
| 4743 | zero_partition_rank, |
| 4744 | zero_partition_count, |
| 4745 | param_families=None): |
| 4746 | if not folded_autoep_tp: |
| 4747 | return None |
| 4748 | return DeepSpeedEngine._make_autoep_folding_metadata(folding_spec, |
| 4749 | family=family, |
| 4750 | ep_rank=ep_rank, |
| 4751 | zero_partition_group=zero_partition_group, |
| 4752 | zero_partition_rank=zero_partition_rank, |
| 4753 | zero_partition_count=zero_partition_count, |
| 4754 | param_families=param_families) |
| 4755 | |
| 4756 | def autoep_expert_writer() -> bool: |
| 4757 | if folded_autoep_tp: |
| 4758 | return groups._get_data_parallel_rank() < folding_spec.ep_size |
| 4759 | return self.checkpoint_engine.is_data_parallel_writer(exp_dp_rank) |
| 4760 | |
| 4761 | # A hack to save the checkpointing directory. Pipeline parallelism overrides |
| 4762 | # module_state_dict() and uses this path to save the model. module_state_dict() |
| 4763 | # then instead just returns None. |
| 4764 | |
| 4765 | # Using layer_#_export_# to save the model's expert state_dict |
| 4766 | autoep_layer_info = [] |
| 4767 | autoep_group_names = set() |
| 4768 | moe_layer_id = 0 |
| 4769 | found_native_moe = False |
| 4770 | found_autoep = False |
| 4771 | for n_module, module in self.module.named_modules(): |
| 4772 | if isinstance(module, MoE): # and deepspeed.comm.get_rank() == 0: |
| 4773 | found_native_moe = True |
| 4774 | if self.zero_optimization_partition_weights() and found_autoep: |
| 4775 | raise RuntimeError("AutoEP with ZeRO Stage 3 checkpointing does not support models that also " |
| 4776 | "contain native DeepSpeed MoE layers.") |
| 4777 | group_name = module.expert_group_name |
| 4778 | num_local_experts = module.num_local_experts |
| 4779 | expp_rank = groups._get_expert_parallel_rank(group_name) |
| 4780 | exp_dp_rank = groups._get_expert_data_parallel_rank(group_name) |
| 4781 | # print(expp_rank, exp_dp_rank) |
| 4782 | # if exp_dp_rank != 0: |
| 4783 | if not self.checkpoint_engine.is_data_parallel_writer(exp_dp_rank): |
| 4784 | moe_layer_id += 1 |
| 4785 | continue |
no test coverage detected