| 3936 | |
| 3937 | @staticmethod |
| 3938 | def load_moe_state_dict(checkpoint_path, |
| 3939 | tag, |
| 3940 | state_dict, |
| 3941 | old_moe_load, |
| 3942 | model=None, |
| 3943 | mpu=None, |
| 3944 | num_experts=1, |
| 3945 | checkpoint_engine=TorchCheckpointEngine(), |
| 3946 | autoep_layers=None, |
| 3947 | folding_spec=None): |
| 3948 | try: |
| 3949 | from deepspeed.module_inject.auto_ep_layer import AutoEPMoELayer as _AutoEPMoELayer |
| 3950 | except ImportError: |
| 3951 | _AutoEPMoELayer = None |
| 3952 | |
| 3953 | has_autoep_layers = _AutoEPMoELayer is not None and model is not None and any( |
| 3954 | isinstance(m, _AutoEPMoELayer) for _, m in model.named_modules()) |
| 3955 | folded_autoep_tp = folding_spec is not None and folding_spec.tp_size > 1 |
| 3956 | |
| 3957 | if old_moe_load: |
| 3958 | if has_autoep_layers: |
| 3959 | raise RuntimeError("Legacy checkpoint format (old_moe_load) is incompatible with AutoEP layers. " |
| 3960 | "Use Universal Checkpointing to convert the checkpoint first.") |
| 3961 | expp_rank = groups._get_expert_data_parallel_rank(groups._get_max_expert_size_name()) |
| 3962 | |
| 3963 | num_local_experts = max(num_experts) // groups._get_expert_parallel_world_size( |
| 3964 | groups._get_max_expert_size_name()) |
| 3965 | for local_expert_id in range(num_local_experts): |
| 3966 | global_expert_id = expp_rank * num_local_experts + local_expert_id |
| 3967 | expert_state_dict = checkpoint_engine.load( |
| 3968 | DeepSpeedEngine._get_expert_ckpt_name( |
| 3969 | checkpoint_path, |
| 3970 | -1, # -1 means ignore layer_id |
| 3971 | global_expert_id, |
| 3972 | tag, |
| 3973 | mpu), |
| 3974 | map_location=torch.device('cpu')) |
| 3975 | |
| 3976 | # Updating global -> local expert ids |
| 3977 | moe_str_prefix = '.deepspeed_moe.experts.deepspeed_experts.' |
| 3978 | for key in list(expert_state_dict.keys()): |
| 3979 | local_key = key.replace(f'{moe_str_prefix}{global_expert_id}', |
| 3980 | f'{moe_str_prefix}{local_expert_id}') |
| 3981 | expert_state_dict[local_key] = expert_state_dict.pop(key) |
| 3982 | state_dict.update(expert_state_dict) |
| 3983 | |
| 3984 | else: |
| 3985 | # Validate AutoEP metadata if present |
| 3986 | if autoep_layers is not None: |
| 3987 | if not isinstance(autoep_layers, list): |
| 3988 | raise RuntimeError( |
| 3989 | f"ds_autoep_layers metadata is malformed: expected list, got {type(autoep_layers).__name__}") |
| 3990 | seen_ids = set() |
| 3991 | required_fields = { |
| 3992 | 'moe_layer_id', 'module_path', 'num_experts', 'num_local_experts', 'ep_size', 'expert_key_prefix' |
| 3993 | } |
| 3994 | for entry in autoep_layers: |
| 3995 | if not isinstance(entry, dict): |