Support scaling runs for DeepSeek-V3 (#1501)

Signed-off-by: Sanju C Sudhakaran <[email protected]>

Author: scsudhakaran

scripts/performance/argument_parser.py

@@ -317,10 +317,10 @@ def parse_cli_args():
     parser.add_argument(
         "-vp",
         "--virtual_pipeline_model_parallel_size",
-        type=int,
+        type=lambda x: None if x == "None" else int(x),
         help="Number of virtual blocks per pipeline model parallel rank is the virtual model parallel size.",
         required=False,
-        default=None,
+        default=-1,
     )
     parser.add_argument(
         "-ep",

scripts/performance/setup_experiment.py

@@ -137,6 +137,7 @@ def main(
         executor.container_mounts.extend([f"{megatron_ckpt_dir}:/mnt/megatron_ckpt"])
     logger.info(f"Custom mounts: {executor.container_mounts}")
 
+    vp_size = vp_size if vp_size != -1 else None
     exp_name = (
         f"{task}_{model_name}_{model_size}_{compute_dtype}"
         f"_gpus{num_gpus}_tp{tp_size}_pp{pp_size}_cp{cp_size}"

scripts/performance/utils/helpers.py

@@ -240,7 +240,7 @@ def set_user_overrides(recipe: ConfigContainer, kwargs: Dict[str, Any]) -> None:
         recipe.model.pipeline_model_parallel_size = kwargs.get("pipeline_model_parallel_size")
     if kwargs.get("context_parallel_size") is not None:
         recipe.model.context_parallel_size = kwargs.get("context_parallel_size")
-    if kwargs.get("virtual_pipeline_model_parallel_size") is not None:
+    if kwargs.get("virtual_pipeline_model_parallel_size") != -1:
         recipe.model.virtual_pipeline_model_parallel_size = kwargs.get("virtual_pipeline_model_parallel_size")
     if kwargs.get("expert_model_parallel_size") is not None:
         recipe.model.expert_model_parallel_size = kwargs.get("expert_model_parallel_size")
@@ -269,6 +269,23 @@ def set_user_overrides(recipe: ConfigContainer, kwargs: Dict[str, Any]) -> None:
         if hasattr(recipe, "comm_overlap") and isinstance(recipe.comm_overlap, CommOverlapConfig):
             recipe.comm_overlap.overlap_param_gather_with_optimizer_step = True
 
+def set_deepseek_v3_layout(recipe: ConfigContainer) -> None:
+    """Set the DeepSeek V3 layout."""
+    pp = recipe.model.pipeline_model_parallel_size
+    vp = recipe.model.virtual_pipeline_model_parallel_size or 1
+    mtp_layers = getattr(recipe.model, "mtp_num_layers", 1) or 0
+    last_layer = ["mtp"] * mtp_layers + ["loss"]
+
+    layout_map = {
+        (1, 1): None,
+        (4, 1): [["embedding"] + ["decoder"] * 16, ["decoder"] * 16, ["decoder"] * 16, ["decoder"] * 13 + last_layer],
+        (8, 1): [["embedding"] + ["decoder"] * 8] + [["decoder"] * 8] * 6 + [["decoder"] * 5 + last_layer],
+        (4, 2): [["embedding"] + ["decoder"] * 8] + [["decoder"] * 8] * 6 + [["decoder"] * 5 + last_layer],
+        (16, 1): [["embedding"] + ["decoder"] * 4] + [["decoder"] * 4] * 14 + [["decoder"] + last_layer],
+        (8, 2): [["embedding"] + ["decoder"] * 4] + [["decoder"] * 4] * 14 + [["decoder"] + last_layer],
+        (4, 4): [["embedding"] + ["decoder"] * 4] + [["decoder"] * 4] * 14 + [["decoder"] + last_layer],
+    }
+    recipe.model.pipeline_model_parallel_layout = layout_map[(pp, vp)]
 
 def get_model_recipe_with_user_overrides(**kwargs) -> ConfigContainer:
     """Get the model recipe with user overrides."""
@@ -284,6 +301,8 @@ def get_model_recipe_with_user_overrides(**kwargs) -> ConfigContainer:
     recipe = get_model_recipe(model_name, model_size, gpu, compute_dtype, domain, task)
     set_common_perf_overrides(recipe)
     set_user_overrides(recipe, kwargs)
+    if model_name == "deepseek" and model_size == "v3":
+        set_deepseek_v3_layout(recipe)
 
     # Scale global batch size based on the number of GPUs IF GBS is not specified by the use 0 r
     workload_base_config = get_workload_base_config(model_name, model_size, gpu, compute_dtype, domain, task)