Method replace_moe_layer

deepspeed/module_inject/auto_ep.py:528–540

Replace a single MoE module with AutoEPMoELayer in-place on the model.

(
        self,
        spec: MoELayerSpec,
        ep_size: int,
        ep_rank: int,
    )

        )

    def replace_moe_layer(
        self,
        spec: MoELayerSpec,
        ep_size: int,
        ep_rank: int,
    ) -> None:
        """Replace a single MoE module with AutoEPMoELayer in-place on the model."""
        replacement = self._replace_moe_layer_without_retarget(spec, ep_size, ep_rank)
        self._retarget_transformers_output_recorders(spec, replacement)

        logger.info(f"AutoEP: replaced '{spec.moe_module_name}' with AutoEPMoELayer "
                    f"(ep_size={ep_size}, ep_rank={ep_rank}, "
                    f"local_experts={replacement.num_local_experts})")

    def replace_moe_layers(
        self,