| 3811 | |
| 3812 | @staticmethod |
| 3813 | def load_moe_state_dict(checkpoint_path, |
| 3814 | tag, |
| 3815 | state_dict, |
| 3816 | old_moe_load, |
| 3817 | model=None, |
| 3818 | mpu=None, |
| 3819 | num_experts=1, |
| 3820 | checkpoint_engine=TorchCheckpointEngine(), |
| 3821 | autoep_layers=None, |
| 3822 | folding_spec=None): |
| 3823 | try: |
| 3824 | from deepspeed.module_inject.auto_ep_layer import AutoEPMoELayer as _AutoEPMoELayer |
| 3825 | except ImportError: |
| 3826 | _AutoEPMoELayer = None |
| 3827 | |
| 3828 | has_autoep_layers = _AutoEPMoELayer is not None and model is not None and any( |
| 3829 | isinstance(m, _AutoEPMoELayer) for _, m in model.named_modules()) |
| 3830 | folded_autoep_tp = folding_spec is not None and folding_spec.tp_size > 1 |
| 3831 | |
| 3832 | if old_moe_load: |
| 3833 | if has_autoep_layers: |
| 3834 | raise RuntimeError("Legacy checkpoint format (old_moe_load) is incompatible with AutoEP layers. " |
| 3835 | "Use Universal Checkpointing to convert the checkpoint first.") |
| 3836 | expp_rank = groups._get_expert_data_parallel_rank(groups._get_max_expert_size_name()) |
| 3837 | |
| 3838 | num_local_experts = max(num_experts) // groups._get_expert_parallel_world_size( |
| 3839 | groups._get_max_expert_size_name()) |
| 3840 | for local_expert_id in range(num_local_experts): |
| 3841 | global_expert_id = expp_rank * num_local_experts + local_expert_id |
| 3842 | expert_state_dict = checkpoint_engine.load( |
| 3843 | DeepSpeedEngine._get_expert_ckpt_name( |
| 3844 | checkpoint_path, |
| 3845 | -1, # -1 means ignore layer_id |
| 3846 | global_expert_id, |
| 3847 | tag, |
| 3848 | mpu), |
| 3849 | map_location=torch.device('cpu')) |
| 3850 | |
| 3851 | # Updating global -> local expert ids |
| 3852 | moe_str_prefix = '.deepspeed_moe.experts.deepspeed_experts.' |
| 3853 | for key in list(expert_state_dict.keys()): |
| 3854 | local_key = key.replace(f'{moe_str_prefix}{global_expert_id}', |
| 3855 | f'{moe_str_prefix}{local_expert_id}') |
| 3856 | expert_state_dict[local_key] = expert_state_dict.pop(key) |
| 3857 | state_dict.update(expert_state_dict) |
| 3858 | |
| 3859 | else: |
| 3860 | # Validate AutoEP metadata if present |
| 3861 | if autoep_layers is not None: |
| 3862 | if not isinstance(autoep_layers, list): |
| 3863 | raise RuntimeError( |
| 3864 | f"ds_autoep_layers metadata is malformed: expected list, got {type(autoep_layers).__name__}") |
| 3865 | seen_ids = set() |
| 3866 | required_fields = { |
| 3867 | 'moe_layer_id', 'module_path', 'num_experts', 'num_local_experts', 'ep_size', 'expert_key_prefix' |
| 3868 | } |
| 3869 | for entry in autoep_layers: |
| 3870 | if not isinstance(entry, dict): |