diff --git a/src/megatron/bridge/perf_recipes/deepseek/__init__.py b/src/megatron/bridge/perf_recipes/deepseek/__init__.py index a784910aff..458fdf9532 100644 --- a/src/megatron/bridge/perf_recipes/deepseek/__init__.py +++ b/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, diff --git a/src/megatron/bridge/perf_recipes/deepseek/vrnvl72/deepseek_v3.py b/src/megatron/bridge/perf_recipes/deepseek/vrnvl72/deepseek_v3.py index d48159ac96..cc16ca9564 100644 --- a/src/megatron/bridge/perf_recipes/deepseek/vrnvl72/deepseek_v3.py +++ b/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 diff --git a/tests/unit_tests/recipes/test_deepseek_v3_proxy_recipes.py b/tests/unit_tests/recipes/test_deepseek_v3_proxy_recipes.py index 0e39115e16..4193f41dff 100644 --- a/tests/unit_tests/recipes/test_deepseek_v3_proxy_recipes.py +++ b/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 diff --git a/tests/unit_tests/recipes/test_gb300_ncclep_perf_recipes.py b/tests/unit_tests/recipes/test_gb300_ncclep_perf_recipes.py index 9fb2d0d503..7363944e9e 100644 --- a/tests/unit_tests/recipes/test_gb300_ncclep_perf_recipes.py +++ b/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"), diff --git a/tests/unit_tests/recipes/test_moe_full_iteration_recipes.py b/tests/unit_tests/recipes/test_moe_full_iteration_recipes.py index 3f34f5d418..0063fc6f9c 100644 --- a/tests/unit_tests/recipes/test_moe_full_iteration_recipes.py +++ b/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"), [ diff --git a/tests/unit_tests/recipes/test_perf_recipe_environment.py b/tests/unit_tests/recipes/test_perf_recipe_environment.py index a9f75dbd17..3dccbd9428 100644 --- a/tests/unit_tests/recipes/test_perf_recipe_environment.py +++ b/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",