fix(deepseek): use HybridEP for text inference

Signed-off-by: Chen Cui <chcui@nvidia.com>

Author: cuichenx

examples/conversion/hf_to_megatron_generate_text.py

@@ -161,6 +161,29 @@ def _decode_completion(tokenizer, generated_ids: torch.Tensor, prompt_length: in
     )
 
 
+def _moe_dispatcher_overrides(args: argparse.Namespace) -> dict[str, str | int | bool]:
+    """Build explicit flex-dispatcher overrides for providers and loaded checkpoints."""
+    if args.moe_flex_dispatcher_backend is None:
+        if args.moe_flex_dispatcher_num_sms is not None or args.hybridep_permute_fusion is not None:
+            raise ValueError(
+                "--moe-flex-dispatcher-num-sms and --[no-]hybridep-permute-fusion require "
+                "--moe-flex-dispatcher-backend"
+            )
+        return {}
+
+    overrides: dict[str, str | int | bool] = {
+        "moe_token_dispatcher_type": "flex",
+        "moe_flex_dispatcher_backend": args.moe_flex_dispatcher_backend,
+    }
+    if args.moe_flex_dispatcher_num_sms is not None:
+        overrides["moe_flex_dispatcher_num_sms"] = args.moe_flex_dispatcher_num_sms
+    if args.hybridep_permute_fusion is not None:
+        if args.moe_flex_dispatcher_backend != "hybridep":
+            raise ValueError("--[no-]hybridep-permute-fusion requires --moe-flex-dispatcher-backend=hybridep")
+        overrides["moe_permute_fusion_into_hybridep"] = args.hybridep_permute_fusion
+    return overrides
+
+
 def main(args) -> None:
     """Main function for text generation from HuggingFace or Megatron models.
 
@@ -177,6 +200,7 @@ def main(args) -> None:
     pp = args.pp
     ep = args.ep
     etp = args.etp
+    dispatcher_overrides = _moe_dispatcher_overrides(args)
 
     # Choose loading method based on arguments
     if args.megatron_model_path:
@@ -201,6 +225,8 @@ def main(args) -> None:
         model_provider.expert_model_parallel_size = ep
         model_provider.expert_tensor_parallel_size = etp
         model_provider.pipeline_dtype = torch.bfloat16
+        for name, value in dispatcher_overrides.items():
+            setattr(model_provider, name, value)
 
         # Read pipeline layout from checkpoint for PP > 1
         if args.pipeline_model_parallel_layout is not None:
@@ -232,6 +258,7 @@ def main(args) -> None:
             "expert_model_parallel_size": ep,
             "expert_tensor_parallel_size": etp,
             "pipeline_dtype": torch.bfloat16,
+            **dispatcher_overrides,
         }
         if args.pipeline_model_parallel_layout is not None:
             mp_overrides["pipeline_model_parallel_layout"] = args.pipeline_model_parallel_layout
@@ -259,6 +286,8 @@ def main(args) -> None:
         model_provider.expert_model_parallel_size = ep
         model_provider.expert_tensor_parallel_size = etp
         model_provider.pipeline_dtype = torch.bfloat16
+        for name, value in dispatcher_overrides.items():
+            setattr(model_provider, name, value)
         if args.pipeline_model_parallel_layout is not None:
             model_provider.pipeline_model_parallel_layout = args.pipeline_model_parallel_layout
 
@@ -425,6 +454,22 @@ def build_parser() -> argparse.ArgumentParser:
     )
     parser.add_argument("--ep", type=int, default=1, help="Expert parallelism size")
     parser.add_argument("--etp", type=int, default=1, help="Expert tensor parallelism size")
+    parser.add_argument(
+        "--moe-flex-dispatcher-backend",
+        choices=("deepep", "hybridep"),
+        help="Override the MoE token dispatcher with the selected flex backend.",
+    )
+    parser.add_argument(
+        "--moe-flex-dispatcher-num-sms",
+        type=int,
+        help="Number of SMs assigned to the selected flex dispatcher backend.",
+    )
+    parser.add_argument(
+        "--hybridep-permute-fusion",
+        action=argparse.BooleanOptionalAction,
+        default=None,
+        help="Enable or disable HybridEP permutation fusion when HybridEP is selected.",
+    )
     parser.add_argument("--megatron_model_path", type=str, default=None, help="Path to the Megatron model checkpoint")
     parser.add_argument("--trust-remote-code", action="store_true", help="if trust_remote_code")
     return parser

examples/models/deepseek_v4/README.md

@@ -61,6 +61,18 @@ performance features, so it is not a supported runtime for that recipe.
 - `slurm_pretrain.sh` runs the DeepSeek-V4-Flash pretraining recipes.
 - `slurm_sft.sh` runs DeepSeek-V4-Flash full SFT end to end (import, then fine-tune) on Hopper or Blackwell, with MTP on or off.
 
+The inference script selects the `flex`/HybridEP MoE dispatcher with 16
+dispatcher SMs and unfused HybridEP permutation for both direct-HF and
+Megatron-checkpoint loading. The explicit checkpoint override prevents an older
+serialized `alltoall` setting from restoring the standard dispatcher. Custom
+launchers need a runtime containing the HybridEP package. HybridEP auto-detects
+accessible ranks and fabric support when its topology environment variables are
+unset. To override detection, set
+`NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN`, `NVLINK_DOMAIN_SIZE`, and
+`USE_MNNVL` to values that describe the actual deployment. This selection is
+inference-specific; training recipes retain their independently qualified
+dispatcher defaults.
+
 Run `bash conversion.sh` after setting `WORKSPACE` and `MODEL_VARIANT`. See each script's header comments for the expected environment variables and `#SBATCH` directives to edit before submitting.
 
 ## Pretraining Recipes

examples/models/deepseek_v4/inference.sh

@@ -22,8 +22,13 @@
 #   EP: expert-parallel size (default: 4 for Flash, 8 for Pro)
 #   PP: pipeline-parallel size (default: 1 for Flash, 4 for Pro)
 #   PROMPT: prompt string (default: "Explain hyper-connections in transformer models.")
+#   NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN: optional HybridEP topology override
+#   NVLINK_DOMAIN_SIZE: optional physical NVLink-domain size override
+#   USE_MNNVL: optional HybridEP fabric override (0 or 1)
 #
 # Defaults below are for GB200 (192 GB). For H100 (80 GB) configs, see README.md.
+# HybridEP auto-detects accessible ranks and fabric support when the topology
+# overrides are unset. If set, the values must describe the actual deployment.
 
 set -xeuo pipefail
 
@@ -54,6 +59,9 @@ uv run python -m torch.distributed.run --nproc_per_node=$((PP * EP)) \
     --prompt "${PROMPT}" \
     --max_new_tokens 100 \
     --tp ${TP} --pp ${PP} --ep ${EP} \
+    --moe-flex-dispatcher-backend hybridep \
+    --moe-flex-dispatcher-num-sms 16 \
+    --no-hybridep-permute-fusion \
     --legacy-full-prefix \
     --trust-remote-code
 
@@ -67,6 +75,9 @@ if [[ -d "${MEGATRON_DIR}/iter_0000000" ]]; then
         --prompt "${PROMPT}" \
         --max_new_tokens 100 \
         --tp ${TP} --pp ${PP} --ep ${EP} \
+        --moe-flex-dispatcher-backend hybridep \
+        --moe-flex-dispatcher-num-sms 16 \
+        --no-hybridep-permute-fusion \
         --legacy-full-prefix \
         --trust-remote-code
 fi

src/megatron/bridge/models/model_provider.py

@@ -557,6 +557,10 @@ class ModelParallelKwargs(TypedDict, total=False):
     hierarchical_context_parallel_sizes: list[int] | None
     pipeline_model_parallel_layout: list[list[str]] | None
     pipeline_dtype: torch.dtype
+    moe_token_dispatcher_type: str
+    moe_flex_dispatcher_backend: str | None
+    moe_flex_dispatcher_num_sms: int | None
+    moe_permute_fusion_into_hybridep: bool
 
 
 def get_model(

tests/unit_tests/examples/test_hf_to_megatron_generate_text.py

@@ -27,6 +27,7 @@
 _decode_completion = _SCRIPT_GLOBALS["_decode_completion"]
 _hf_revision_kwargs = _SCRIPT_GLOBALS["_hf_revision_kwargs"]
 _maybe_gather_tensor_parallel_logits = _SCRIPT_GLOBALS["_maybe_gather_tensor_parallel_logits"]
+_moe_dispatcher_overrides = _SCRIPT_GLOBALS["_moe_dispatcher_overrides"]
 _run_megatron_forward = _SCRIPT_GLOBALS["_run_megatron_forward"]
 _text_forward_step = _SCRIPT_GLOBALS["text_forward_step"]
 _tokenize_prompt = _SCRIPT_GLOBALS["_tokenize_prompt"]
@@ -107,6 +108,63 @@ def test_legacy_full_prefix_disables_inference_context() -> None:
     assert _build_inference_context(input_ids, legacy_full_prefix=True) is None
 
 
+@pytest.mark.unit
+def test_hybridep_cli_builds_checkpoint_safe_dispatcher_overrides() -> None:
+    args = _build_parser().parse_args(
+        [
+            "--hf_model_path",
+            "org/model",
+            "--moe-flex-dispatcher-backend",
+            "hybridep",
+            "--moe-flex-dispatcher-num-sms",
+            "16",
+            "--no-hybridep-permute-fusion",
+        ]
+    )
+
+    assert _moe_dispatcher_overrides(args) == {
+        "moe_token_dispatcher_type": "flex",
+        "moe_flex_dispatcher_backend": "hybridep",
+        "moe_flex_dispatcher_num_sms": 16,
+        "moe_permute_fusion_into_hybridep": False,
+    }
+
+
+@pytest.mark.unit
+def test_flex_dispatcher_cli_is_opt_in() -> None:
+    args = _build_parser().parse_args(["--hf_model_path", "org/model"])
+
+    assert _moe_dispatcher_overrides(args) == {}
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize("option", ["--moe-flex-dispatcher-num-sms", "--no-hybridep-permute-fusion"])
+def test_flex_dispatcher_dependent_options_require_backend(option: str) -> None:
+    command = ["--hf_model_path", "org/model", option]
+    if option == "--moe-flex-dispatcher-num-sms":
+        command.append("16")
+    args = _build_parser().parse_args(command)
+
+    with pytest.raises(ValueError, match="require --moe-flex-dispatcher-backend"):
+        _moe_dispatcher_overrides(args)
+
+
+@pytest.mark.unit
+def test_hybridep_permute_flag_rejects_other_backends() -> None:
+    args = _build_parser().parse_args(
+        [
+            "--hf_model_path",
+            "org/model",
+            "--moe-flex-dispatcher-backend",
+            "deepep",
+            "--no-hybridep-permute-fusion",
+        ]
+    )
+
+    with pytest.raises(ValueError, match="requires --moe-flex-dispatcher-backend=hybridep"):
+        _moe_dispatcher_overrides(args)
+
+
 @pytest.mark.unit
 def test_kv_cache_builds_static_inference_context() -> None:
     input_ids = torch.ones((2, 3), dtype=torch.long)