Proxy 1 node setup for DSv3

src/megatron/bridge/perf_recipes/deepseek/__init__.py

@@ -54,6 +54,8 @@
     deepseek_v3_pretrain_1024gpu_h100_fp8sc_large_scale_config,
 )
 from megatron.bridge.perf_recipes.deepseek.vrnvl72.deepseek_v3 import (
+    deepseek_v3_pretrain_4gpu_vrnvl72_fp8mx_proxy_config,
+    deepseek_v3_pretrain_4gpu_vrnvl72_nvfp4_proxy_config,
     deepseek_v3_pretrain_64gpu_vrnvl72_fp8mx_proxy_config,
     deepseek_v3_pretrain_64gpu_vrnvl72_nvfp4_proxy_config,
     deepseek_v3_pretrain_128gpu_vrnvl72_bf16_config,

src/megatron/bridge/perf_recipes/deepseek/vrnvl72/deepseek_v3.py

@@ -446,3 +446,123 @@ def deepseek_v3_pretrain_64gpu_vrnvl72_nvfp4_proxy_config() -> ConfigContainer:
         "NVTE_DPA_FP8_FORMAT": "E4M3",
     }
     return cfg
+
+
+def _apply_deepseek_v3_4gpu_vrnvl72_proxy(cfg: ConfigContainer) -> None:
+    """Shrink a 256-GPU VRNVL72 DeepSeek V3 recipe onto one 4-GPU VRNVL72 node.
+
+    Both 4-GPU proxies go through this helper so that every model, parallelism and batch
+    setting matches and only the quantization recipe differs between their profiles.
+    """
+    # Three leading dense layers and ten MoE layers plus one MTP layer, as in the 64-GPU proxies.
+    cfg.model.num_layers = 13
+    cfg.model.moe_layer_freq = [0] * 3 + [1] * (cfg.model.num_layers - 3)
+    # EP=4 over 32 experts keeps the parent's per-rank expert work (256 experts over EP=32):
+    # 8 local experts, each receiving 4096 tokens per microbatch under forced load balancing.
+    cfg.model.num_moe_experts = 32
+    cfg.model.expert_model_parallel_size = 4
+    cfg.model.pipeline_model_parallel_size = 1
+    cfg.model.virtual_pipeline_model_parallel_size = None
+    # Clear the inherited interleaved layout; embedding, decoder and MTP are colocated.
+    set_deepseek_v3_pipeline_model_parallel_layout(cfg.model)
+    # Dense DP=4 and expert DP=1; 32 microbatches per rank, as in the 256-GPU parent.
+    cfg.train.global_batch_size = 128
+
+    # The parents differ in these two knobs; pin both so the proxies share one graph structure.
+    # MXFP8 is the only recipe with a quantized LM head, so keep the LM head in BF16 for both.
+    cfg.model.fp8_output_proj = False
+    # Fused input LayerNorm and q/kv down-projection works with either recipe (NVFP4 parent setting).
+    cfg.model.mla_down_proj_fusion = True
+
+
+def deepseek_v3_pretrain_4gpu_vrnvl72_fp8mx_proxy_config() -> ConfigContainer:
+    """DeepSeek V3 MXFP8 profiling proxy: 13 decoder layers and 32 experts on one 4-GPU VRNVL72 node.
+
+    Matches ``deepseek_v3_pretrain_4gpu_vrnvl72_nvfp4_proxy_config`` in everything except the
+    quantization recipe, so the two profiles can be compared directly. Per-layer shapes and
+    per-rank expert work follow the 256-GPU parent: TP=1, PP=1, EP=4, MBS=1, GBS=128, one MTP
+    layer, full-iteration CUDA graphs, HybridEP and combined-1F1B MoE overlap. Attention and all
+    GEMMs run in MXFP8; the LM head stays in BF16. This proxy does not reproduce at-scale
+    communication, convergence or memory use.
+    """
+    cfg = deepseek_v3_pretrain_256gpu_vrnvl72_fp8mx_config()
+    _apply_deepseek_v3_4gpu_vrnvl72_proxy(cfg)
+
+    # Keep process settings next to the recipe so users can see the exact benchmark environment.
+    # Identical to the NVFP4 proxy apart from the NVFP4-only entries at the end of that recipe.
+    cfg.env_vars = {
+        **COMMON_PERF_ENV_VARS,
+        # CUDA stream scheduling for this model and parallel layout.
+        "CUDA_DEVICE_MAX_CONNECTIONS": 32,
+        # CUDA graph and allocator behavior for this recipe.
+        "NCCL_GRAPH_REGISTER": 0,
+        "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True,graph_capture_record_stream_reuse:True",
+        "TORCH_NCCL_AVOID_RECORD_STREAMS": 0,
+        # NCCL user-buffer and launch settings.
+        "NCCL_NVLS_ENABLE": 0,
+        # HybridEP topology: the whole EP=4 group sits in one node.
+        "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 4,
+        "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128,
+        "NVLINK_DOMAIN_SIZE": 72,
+        "USE_MNNVL": 1,
+        # Transformer Engine overlap settings for this model.
+        "CUDNNFE_CLUSTER_OVERLAP_MARGIN": 8,
+        "NVTE_BWD_LAYERNORM_SM_MARGIN": 20,
+        "NVTE_CUTEDSL_FUSED_GROUPED_MLP": 1,
+        "NVTE_FWD_LAYERNORM_SM_MARGIN": 20,
+        # Use cuDNN LayerNorm in both 4-GPU proxies.
+        "NVTE_NORM_BWD_USE_CUDNN": 1,
+        "NVTE_NORM_FWD_USE_CUDNN": 1,
+        # Keep DeepSeek kernel selection aligned with the 256-GPU parent recipes.
+        "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0,
+    }
+    return cfg
+
+
+def deepseek_v3_pretrain_4gpu_vrnvl72_nvfp4_proxy_config() -> ConfigContainer:
+    """DeepSeek V3 NVFP4 profiling proxy: 13 decoder layers and 32 experts on one 4-GPU VRNVL72 node.
+
+    Matches ``deepseek_v3_pretrain_4gpu_vrnvl72_fp8mx_proxy_config`` in everything except the
+    quantization recipe, so the two profiles can be compared directly. Per-layer shapes and
+    per-rank expert work follow the 256-GPU parent: TP=1, PP=1, EP=4, MBS=1, GBS=128, one MTP
+    layer, full-iteration CUDA graphs, HybridEP and combined-1F1B MoE overlap. GEMMs run in
+    NVFP4, attention in MXFP8 E4M3 as in the MXFP8 proxy, and the LM head in BF16. This proxy
+    does not reproduce at-scale communication, convergence or memory use.
+    """
+    cfg = deepseek_v3_pretrain_256gpu_vrnvl72_nvfp4_config()
+    _apply_deepseek_v3_4gpu_vrnvl72_proxy(cfg)
+
+    # Keep process settings next to the recipe so users can see the exact benchmark environment.
+    # Identical to the MXFP8 proxy apart from the NVFP4-only entries at the end.
+    cfg.env_vars = {
+        **COMMON_PERF_ENV_VARS,
+        # CUDA stream scheduling for this model and parallel layout.
+        "CUDA_DEVICE_MAX_CONNECTIONS": 32,
+        # CUDA graph and allocator behavior for this recipe.
+        "NCCL_GRAPH_REGISTER": 0,
+        "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True,graph_capture_record_stream_reuse:True",
+        "TORCH_NCCL_AVOID_RECORD_STREAMS": 0,
+        # NCCL user-buffer and launch settings.
+        "NCCL_NVLS_ENABLE": 0,
+        # HybridEP topology: the whole EP=4 group sits in one node.
+        "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 4,
+        "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128,
+        "NVLINK_DOMAIN_SIZE": 72,
+        "USE_MNNVL": 1,
+        # Transformer Engine overlap settings for this model.
+        "CUDNNFE_CLUSTER_OVERLAP_MARGIN": 8,
+        "NVTE_BWD_LAYERNORM_SM_MARGIN": 20,
+        "NVTE_CUTEDSL_FUSED_GROUPED_MLP": 1,
+        "NVTE_FWD_LAYERNORM_SM_MARGIN": 20,
+        # Use cuDNN LayerNorm in both 4-GPU proxies.
+        "NVTE_NORM_BWD_USE_CUDNN": 1,
+        "NVTE_NORM_FWD_USE_CUDNN": 1,
+        # Keep DeepSeek kernel selection aligned with the 256-GPU parent recipes.
+        "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0,
+        # NVFP4 fast-math path.
+        "NVTE_USE_FAST_MATH": 1,
+        # Run attention in MXFP8 E4M3, as the MXFP8 proxy does natively.
+        "NVTE_DPA_FP8_RECIPE": "MXFP8BlockScaling",
+        "NVTE_DPA_FP8_FORMAT": "E4M3",
+    }
+    return cfg

tests/unit_tests/recipes/test_deepseek_v3_proxy_recipes.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Guard the 64-GPU debugging proxies against drift from their 256-GPU parents."""
+"""Guard the 64-GPU and 4-GPU proxies against drift from their 256-GPU parents and each other."""
 
 from collections.abc import Callable
 from dataclasses import fields
@@ -29,6 +29,31 @@
 
 pytestmark = pytest.mark.unit
 
+# Model fields the 4-GPU proxies change relative to their 256-GPU parents.
+_4GPU_PROXY_MODEL_FIELDS = {
+    "num_layers",
+    "moe_layer_freq",
+    "num_moe_experts",
+    "expert_model_parallel_size",
+    "pipeline_model_parallel_size",
+    "virtual_pipeline_model_parallel_size",
+    "pipeline_model_parallel_layout",
+    "fp8_output_proj",
+    "mla_down_proj_fusion",
+}
+# Mixed-precision fields that differ between the MXFP8 and NVFP4 presets; ``*_param`` mirrors ``*_param_gather``.
+_QUANTIZATION_FIELDS = {
+    "fp8",
+    "fp8_recipe",
+    "fp8_param",
+    "fp8_param_gather",
+    "reuse_grad_buf_for_mxfp8_param_ag",
+    "fp4",
+    "fp4_param",
+    "fp4_param_gather",
+}
+_NVFP4_ONLY_ENV = {"NVTE_USE_FAST_MATH": 1, "NVTE_DPA_FP8_RECIPE": "MXFP8BlockScaling", "NVTE_DPA_FP8_FORMAT": "E4M3"}
+
 
 @pytest.fixture(autouse=True)
 def _offline_recipes(monkeypatch: pytest.MonkeyPatch) -> None:
@@ -110,14 +135,110 @@ def test_proxy_preserves_parent_except_depth_and_layout(precision: str) -> None:
 
 
 @pytest.mark.parametrize("precision", ["fp8mx", "nvfp4"])
-def test_proxy_overlap_setup_without_pipeline_parallelism(monkeypatch: pytest.MonkeyPatch, precision: str) -> None:
-    cfg = getattr(deepseek, f"deepseek_v3_pretrain_64gpu_vrnvl72_{precision}_proxy_config")()
+def test_4gpu_proxy_only_resizes_its_parent(precision: str) -> None:
+    parent = getattr(deepseek, f"deepseek_v3_pretrain_256gpu_vrnvl72_{precision}_config")()
+    proxy = getattr(deepseek, f"deepseek_v3_pretrain_4gpu_vrnvl72_{precision}_proxy_config")()
+
+    assert {key: value for key, value in vars(proxy.model).items() if key not in _4GPU_PROXY_MODEL_FIELDS} == {
+        key: value for key, value in vars(parent.model).items() if key not in _4GPU_PROXY_MODEL_FIELDS
+    }
+    assert {**vars(proxy.train), "global_batch_size": None} == {**vars(parent.train), "global_batch_size": None}
+    for field in fields(parent):
+        if field.name not in {"model", "train", "env_vars"}:
+            assert getattr(proxy, field.name) == getattr(parent, field.name), field.name
+
+    model = proxy.model
+    assert model.num_layers == 13
+    assert model.moe_layer_freq == [0] * 3 + [1] * 10
+    assert model.mtp_num_layers == 1
+    assert (model.tensor_model_parallel_size, model.context_parallel_size) == (1, 1)
+    assert (model.pipeline_model_parallel_size, model.virtual_pipeline_model_parallel_size) == (1, None)
+    assert model.pipeline_model_parallel_layout is None
+    assert (model.expert_model_parallel_size, model.expert_tensor_parallel_size) == (4, 1)
+    assert (parent.model.pipeline_model_parallel_size, parent.model.virtual_pipeline_model_parallel_size) == (2, 8)
+    # With the same MBS, sequence length and top-k, equal local-expert counts give equal tokens per expert.
+    deepseek_v3_num_experts = 256
+    assert model.num_moe_experts == 32
+    assert (
+        model.num_moe_experts // model.expert_model_parallel_size
+        == deepseek_v3_num_experts // parent.model.expert_model_parallel_size
+        == 8
+    )
+    for cfg, world, expected_dp, expected_edp in ((proxy, 4, 4, 1), (parent, 256, 128, 4)):
+        mesh_model = cfg.model
+        dense_mesh = (
+            mesh_model.tensor_model_parallel_size
+            * mesh_model.context_parallel_size
+            * mesh_model.pipeline_model_parallel_size
+        )
+        expert_mesh = (
+            mesh_model.expert_tensor_parallel_size
+            * mesh_model.expert_model_parallel_size
+            * mesh_model.pipeline_model_parallel_size
+        )
+        assert world % dense_mesh == world % expert_mesh == 0
+        assert world // dense_mesh == expected_dp
+        assert world // expert_mesh == expected_edp
+        assert cfg.train.micro_batch_size == 1
+        assert cfg.train.global_batch_size % expected_dp == 0
+        assert cfg.train.global_batch_size // expected_dp == 32
+
+    assert (
+        get_forward_backward_func(
+            pp_size=model.pipeline_model_parallel_size, vp_size=model.virtual_pipeline_model_parallel_size
+        )
+        is forward_backward_no_pipelining
+    )
+
+
+def test_4gpu_proxies_differ_only_in_quantization() -> None:
+    mxfp8 = deepseek.deepseek_v3_pretrain_4gpu_vrnvl72_fp8mx_proxy_config()
+    nvfp4 = deepseek.deepseek_v3_pretrain_4gpu_vrnvl72_nvfp4_proxy_config()
+
+    # Same model, parallelism, batch, CUDA graph, dispatcher, overlap, optimizer and runtime settings.
+    assert vars(mxfp8.model) == vars(nvfp4.model)
+    assert mxfp8.model.cuda_graph_impl == "full_iteration"
+    assert mxfp8.model.fp8_output_proj is False
+    assert mxfp8.model.mla_down_proj_fusion is True
+    for field in fields(mxfp8):
+        if field.name not in {"model", "mixed_precision", "env_vars"}:
+            assert getattr(mxfp8, field.name) == getattr(nvfp4, field.name), field.name
+
+    differing = {
+        field.name
+        for field in fields(mxfp8.mixed_precision)
+        if getattr(mxfp8.mixed_precision, field.name) != getattr(nvfp4.mixed_precision, field.name)
+    }
+    assert differing == _QUANTIZATION_FIELDS
+    assert (mxfp8.mixed_precision.fp8, mxfp8.mixed_precision.fp8_recipe, mxfp8.mixed_precision.fp4) == (
+        "e4m3",
+        "mxfp8",
+        None,
+    )
+    assert (nvfp4.mixed_precision.fp8, nvfp4.mixed_precision.fp4, nvfp4.mixed_precision.fp4_recipe) == (
+        None,
+        "e2m1",
+        "nvfp4",
+    )
+    # Both run MXFP8 E4M3 attention: natively under MXFP8, through the DPA recipe override under NVFP4.
+    assert mxfp8.mixed_precision.fp8_dot_product_attention is nvfp4.mixed_precision.fp8_dot_product_attention is True
+
+    assert mxfp8.env_vars.keys().isdisjoint(_NVFP4_ONLY_ENV)
+    assert nvfp4.env_vars == {**mxfp8.env_vars, **_NVFP4_ONLY_ENV}
+
+
+@pytest.mark.parametrize("num_gpus", [64, 4])
+@pytest.mark.parametrize("precision", ["fp8mx", "nvfp4"])
+def test_proxy_overlap_setup_without_pipeline_parallelism(
+    monkeypatch: pytest.MonkeyPatch, precision: str, num_gpus: int
+) -> None:
+    cfg = getattr(deepseek, f"deepseek_v3_pretrain_{num_gpus}gpu_vrnvl72_{precision}_proxy_config")()
     # Validate the topology without requiring the GPU libraries in this unit test.
     monkeypatch.setattr("megatron.bridge.training.comm_overlap.is_te_min_version", lambda *_: True)
     monkeypatch.setattr("megatron.bridge.training.comm_overlap.is_torch_min_version", lambda *_: True)
     cfg.mixed_precision.setup(cfg.model, cfg.optimizer, cfg.ddp)
     assert cfg.comm_overlap is not None
-    cfg.comm_overlap.data_parallel_size = cfg.get_data_parallel_size(64)
+    cfg.comm_overlap.data_parallel_size = cfg.get_data_parallel_size(num_gpus)
     cfg.comm_overlap.finalize()
     cfg.comm_overlap.setup(cfg.model, cfg.optimizer, cfg.ddp)
 
@@ -132,13 +253,17 @@ def test_proxy_overlap_setup_without_pipeline_parallelism(monkeypatch: pytest.Mo
     assert cfg.ddp.align_param_gather is False
 
 
+@pytest.mark.parametrize(("num_gpus", "global_batch_size", "expert_parallel_size"), [(64, 4096, 32), (4, 128, 4)])
 @pytest.mark.parametrize("precision", ["fp8_mx", "nvfp4"])
-def test_proxy_is_available_through_performance_selector(monkeypatch: pytest.MonkeyPatch, precision: str) -> None:
+def test_proxy_is_available_through_performance_selector(
+    monkeypatch: pytest.MonkeyPatch, precision: str, num_gpus: int, global_batch_size: int, expert_parallel_size: int
+) -> None:
     monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[3] / "scripts" / "performance"))
     from utils.utils import get_perf_recipe_by_name, list_available_config_variants
 
-    cfg = get_perf_recipe_by_name("deepseek_v3", "pretrain", 64, "vrnvl72", precision, config_variant="proxy")
-    assert cfg.train.global_batch_size == 4096
+    cfg = get_perf_recipe_by_name("deepseek_v3", "pretrain", num_gpus, "vrnvl72", precision, config_variant="proxy")
+    assert cfg.train.global_batch_size == global_batch_size
+    assert cfg.model.expert_model_parallel_size == expert_parallel_size
     assert cfg.model.num_layers == 13
     assert cfg.model.pipeline_model_parallel_size == 1
     assert cfg.model.virtual_pipeline_model_parallel_size is None

tests/unit_tests/recipes/test_gb300_ncclep_perf_recipes.py

@@ -73,6 +73,8 @@
     ("nemotronh.vrnvl72.nemotronh", "nemotron_3_5_lightning_pretrain_8gpu_vrnvl72_fp8mx_config"),
     ("deepseek.vrnvl72.deepseek_v3", "deepseek_v3_pretrain_64gpu_vrnvl72_fp8mx_proxy_config"),
     ("deepseek.vrnvl72.deepseek_v3", "deepseek_v3_pretrain_64gpu_vrnvl72_nvfp4_proxy_config"),
+    ("deepseek.vrnvl72.deepseek_v3", "deepseek_v3_pretrain_4gpu_vrnvl72_fp8mx_proxy_config"),
+    ("deepseek.vrnvl72.deepseek_v3", "deepseek_v3_pretrain_4gpu_vrnvl72_nvfp4_proxy_config"),
     ("qwen.vrnvl72.qwen3_moe", "qwen3_235b_a22b_pretrain_256gpu_vrnvl72_bf16_config"),
     ("qwen.vrnvl72.qwen3_moe", "qwen3_235b_a22b_pretrain_256gpu_vrnvl72_fp8mx_config"),
     ("qwen.vrnvl72.qwen3_moe", "qwen3_235b_a22b_pretrain_256gpu_vrnvl72_nvfp4_config"),

tests/unit_tests/recipes/test_moe_full_iteration_recipes.py

@@ -18,7 +18,11 @@
 
 import pytest
 
-from megatron.bridge.perf_recipes.deepseek import deepseek_v3_pretrain_256gpu_vrnvl72_nvfp4_config
+from megatron.bridge.perf_recipes.deepseek import (
+    deepseek_v3_pretrain_4gpu_vrnvl72_fp8mx_proxy_config,
+    deepseek_v3_pretrain_4gpu_vrnvl72_nvfp4_proxy_config,
+    deepseek_v3_pretrain_256gpu_vrnvl72_nvfp4_config,
+)
 from megatron.bridge.perf_recipes.nemotronh import (
     nemotron_3_super_pretrain_64gpu_gb200_fp8mx_config,
     nemotron_3_super_pretrain_64gpu_gb200_nvfp4_config,
@@ -119,6 +123,28 @@ def test_deepseek_v3_256gpu_vrnvl72_nvfp4_uses_full_iteration_stack() -> None:
     ) == (1, 2, 8, 32)
 
 
+@pytest.mark.parametrize(
+    "recipe",
+    [deepseek_v3_pretrain_4gpu_vrnvl72_fp8mx_proxy_config, deepseek_v3_pretrain_4gpu_vrnvl72_nvfp4_proxy_config],
+)
+def test_deepseek_v3_4gpu_vrnvl72_proxies_use_full_iteration_stack(recipe: Callable[[], ConfigContainer]) -> None:
+    cfg = recipe()
+
+    _assert_full_iteration_hybridep(cfg)
+    assert cfg.model.high_priority_a2a_comm_stream is True
+    assert cfg.comm_overlap is not None
+    assert cfg.comm_overlap.overlap_moe_expert_parallel_comm is True
+    assert cfg.comm_overlap.delay_wgrad_compute is True
+    assert cfg.env_vars["TORCH_NCCL_AVOID_RECORD_STREAMS"] == 0
+    assert cfg.env_vars["NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN"] == cfg.model.expert_model_parallel_size == 4
+    assert (
+        cfg.model.tensor_model_parallel_size,
+        cfg.model.pipeline_model_parallel_size,
+        cfg.model.virtual_pipeline_model_parallel_size,
+        cfg.model.expert_model_parallel_size,
+    ) == (1, 1, None, 4)
+
+
 @pytest.mark.parametrize(
     ("recipe", "expected_tp", "expected_ep"),
     [

tests/unit_tests/recipes/test_perf_recipe_environment.py

@@ -47,6 +47,8 @@
     "TORCHINDUCTOR_WORKER_START",
 }
 _VRNVL72_CUDNN_LAYERNORM_RECIPES = {
+    "deepseek_v3_pretrain_4gpu_vrnvl72_fp8mx_proxy_config",
+    "deepseek_v3_pretrain_4gpu_vrnvl72_nvfp4_proxy_config",
     "deepseek_v3_pretrain_64gpu_vrnvl72_fp8mx_proxy_config",
     "deepseek_v3_pretrain_128gpu_vrnvl72_fp8mx_config",
     "deepseek_v3_pretrain_256gpu_vrnvl72_fp8mx_config",