Update performance scripts to align with llmb requirements

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

Author: scsudhakaran

scripts/performance/argument_parser.py

@@ -261,6 +261,14 @@ def parse_cli_args():
         required=False,
         default=None,
     )
+    parser.add_argument(
+        "--moe_flex_dispatcher_backend",
+        type=str,
+        choices=["deepep", "hybridep"],
+        help="MoE flex dispatcher backend to use. Defaults to None",
+        required=False,
+        default=None,
+    )
     parser.add_argument(
         "--use_megatron_fsdp",
         help="Use Megatron FSDP. Disabled by default.",
@@ -398,6 +406,74 @@ def parse_cli_args():
         action="store_false",
         dest="detach",
     )
-
+    parser.add_argument(
+        "--profiling_start_step", type=int, help="Defines start step for profiling", required=False, default=10
+    )
+    parser.add_argument(
+        "--profiling_stop_step", type=int, help="Defines stop step for profiling", required=False, default=11
+    )
+    parser.add_argument(
+        "-pgm",
+        "--profiling_gpu_metrics",
+        help="Enable nsys gpu metrics. Disabled by default.",
+        action="store_true",
+    )
+    parser.add_argument(
+        "--additional_slurm_params",
+        type=str,
+        help="Additional SLURM parameters as key=value pairs. "
+        "Use semicolons (;) to separate parameters when values contain commas. "
+        "Examples: 'nodelist=node001,node002;constraint=gpu' or 'reservation=my_res;exclusive'",
+        required=False,
+        default=None,
+    )
     args, cli_dotlist_overrides = parser.parse_known_args()
     return args, cli_dotlist_overrides
+
+def parse_additional_slurm_params(params_str):
+    """
+    Parse additional SLURM parameters from a string of key=value pairs.
+    This function handles different separator formats:
+    1. Semicolon-separated: "key1=value1;key2=value2" (recommended for multiple parameters)
+    2. Space-separated: "key1=value1 key2=value2"
+    3. Single parameter: "key1=value1,value2" (no separators = single parameter)
+    Args:
+        params_str (str): String with parameters
+    Returns:
+        dict: Dictionary of parameters, or None if params_str is None/empty
+    Example:
+        parse_additional_slurm_params("nodelist=node001,node002")
+        returns {"nodelist": "node001,node002"}
+        parse_additional_slurm_params("nodelist=node001,node002;constraint=gpu")
+        returns {"nodelist": "node001,node002", "constraint": "gpu"}
+        parse_additional_slurm_params("reservation=my_res;constraint=gpu")
+        returns {"reservation": "my_res", "constraint": "gpu"}
+    """
+    if not params_str:
+        return None
+
+    params = {}
+
+    # Try semicolon separation first (most reliable for complex values)
+    if ';' in params_str:
+        parts = params_str.split(';')
+    # Try space separation next
+    elif ' ' in params_str:
+        parts = params_str.split()
+    # No separators found - treat as single parameter
+    else:
+        parts = [params_str]
+
+    for part in parts:
+        part = part.strip()
+        if not part:
+            continue
+
+        if '=' in part:
+            key, value = part.split('=', 1)
+            params[key.strip()] = value.strip()
+        else:
+            # Boolean flag (no value)
+            params[part] = True
+
+    return params if params else None

scripts/performance/perf_plugins.py

@@ -107,6 +107,7 @@ class NsysPlugin(Plugin):
     profile_step_end: int
     profile_ranks: Optional[list[int]] = None
     nsys_trace: Optional[list[str]] = None
+    nsys_extra_args: Optional[list[str]] = None
     record_shapes: bool = False
     nsys_gpu_metrics: bool = False
     script_args_converter_fn: Optional[Callable[[NsysPluginScriptArgs], List[str]]] = None
@@ -116,6 +117,7 @@ def setup(self, task: Union["run.Partial", "run.Script"], executor: "run.Executo
         launcher = executor.get_launcher()
         launcher.nsys_profile = True
         launcher.nsys_trace = self.nsys_trace or ["nvtx", "cuda"]
+        launcher.nsys_extra_args = self.nsys_extra_args or launcher.nsys_extra_args
 
         if isinstance(executor, SlurmExecutor):
             # NOTE: DO NOT change to f-string, `%q{}` is Slurm placeholder
@@ -195,6 +197,7 @@ class PerfEnvPlugin(Plugin):
     pp_size: int = 1
     script_args_converter_fn: Optional[Callable[[PerfEnvPluginScriptArgs], List[str]]] = None
     moe_a2a_overlap: bool = False
+    moe_flex_dispatcher_backend: str
     model_name: str
     model_size: str
     gpu: str

scripts/performance/setup_experiment.py

@@ -18,10 +18,10 @@
 
 
 try:
-    from argument_parser import parse_cli_args
+    from argument_parser import parse_additional_slurm_params, parse_cli_args
     from utils.executors import slurm_executor
 except (ImportError, ModuleNotFoundError):
-    from .argument_parser import parse_cli_args
+    from .argument_parser import parse_additional_slurm_params, parse_cli_args
     from .utils.executors import slurm_executor
 
 import nemo_run as run
@@ -50,6 +50,7 @@ def main(
     task: str,
     compute_dtype: str,
     gpu: str,
+    num_gpus: int,
     hf_token: str,
     custom_mounts: List[str],
     detach: bool,
@@ -58,12 +59,20 @@ def main(
     enable_nsys: bool,
     use_tokendrop: bool,
     moe_a2a_overlap: bool,
+    moe_flex_dispatcher_backend: str,
     tp_size: Optional[int],
     pp_size: Optional[int],
     cp_size: Optional[int],
+    vp_size: Optional[int],
+    ep_size: Optional[int],
+    mbs: Optional[int],
+    gbs: Optional[int],
     wandb_key: str,
     wandb_prj_name: str,
     wandb_exp_name: str,
+    profiling_start_step: int,
+    profiling_stop_step: int,
+    profiling_gpu_metrics: bool,
     executor: run.Executor,
 ):
     """Sets up the experiment and runs it."""
@@ -87,6 +96,7 @@ def main(
         PerfEnvPlugin(
             enable_vboost=enable_vboost,
             moe_a2a_overlap=moe_a2a_overlap,
+            moe_flex_dispatcher_backend=moe_flex_dispatcher_backend,
             tp_size=tp_size,
             pp_size=pp_size,
             cp_size=cp_size,
@@ -98,7 +108,20 @@ def main(
         )
     )
     if enable_nsys:
-        plugins.append(NsysPlugin(profile_step_start=10, profile_step_end=11))
+        plugins.append(NsysPlugin(
+            profile_step_start=profiling_start_step,
+            profile_step_end=profiling_stop_step,
+            profile_ranks=list(range(num_gpus)),
+            nsys_gpu_metrics=profiling_gpu_metrics,
+            nsys_trace=['cuda'],
+            nsys_extra_args=[
+                "--force-overwrite=true",
+                "--capture-range=cudaProfilerApi",
+                "--capture-range-end=stop",
+                "--cuda-graph-trace=node",
+                "--cuda-event-trace=false",
+                "--nvtx-domain-include=NCCL",
+            ]))
 
     executor.container_mounts.extend(
         custom_mounts
@@ -109,9 +132,12 @@ def main(
     )
     logger.info(f"Custom mounts: {executor.container_mounts}")
 
-    exp_name = f"{model_name}_{model_size}_{domain}_{task}" + (
-        "_bf16" if compute_dtype == "bf16" else f"_{compute_dtype}"
+    exp_name = (
+        f"{task}_{model_name}_{model_size}_{compute_dtype}"
+        f"_gpus{num_gpus}_tp{tp_size}_pp{pp_size}_cp{cp_size}"
+        f"_vp{vp_size}_ep{ep_size}_mbs{mbs}_gbs{gbs}"
     )
+
     logger.debug(
         run.Script(
             path=str(RUN_SCRIPT_PATH),
@@ -146,6 +172,11 @@ def main(
 if __name__ == "__main__":
     args, _ = parse_cli_args()
 
+    # Parse additional SLURM parameters if provided
+    additional_slurm_params = None
+    if hasattr(args, 'additional_slurm_params') and args.additional_slurm_params:
+        additional_slurm_params = parse_additional_slurm_params(args.additional_slurm_params)
+
     main(
         script_name=SCRIPT_NAME,
         model_name=args.model_name,
@@ -154,6 +185,7 @@ def main(
         task=args.task,
         compute_dtype=args.compute_dtype,
         gpu=args.gpu,
+        num_gpus=args.num_gpus,
         hf_token=args.hf_token,
         custom_mounts=args.custom_mounts,
         detach=args.detach,
@@ -162,12 +194,20 @@ def main(
         enable_nsys=args.enable_nsys,
         use_tokendrop=args.use_tokendrop,
         moe_a2a_overlap=args.moe_a2a_overlap,
+        moe_flex_dispatcher_backend=args.moe_flex_dispatcher_backend,
         tp_size=args.tensor_model_parallel_size,
         pp_size=args.pipeline_model_parallel_size,
         cp_size=args.context_parallel_size,
+        vp_size=args.virtual_pipeline_model_parallel_size,
+        ep_size=args.expert_model_parallel_size,
+        mbs=args.micro_batch_size,
+        gbs=args.global_batch_size,
         wandb_key=args.wandb_key,
         wandb_prj_name=args.wandb_prj_name,
         wandb_exp_name=args.wandb_exp_name,
+        profiling_start_step=args.profiling_start_step,
+        profiling_stop_step=args.profiling_stop_step,
+        profiling_gpu_metrics=args.profiling_gpu_metrics,
         executor=slurm_executor(
             args.gpu,
             args.account,
@@ -181,5 +221,6 @@ def main(
             hf_token=args.hf_token,
             nemo_home=args.nemo_home,
             wandb_key=args.wandb_key,
+            additional_slurm_params=additional_slurm_params,
         ),
     )

scripts/performance/utils/executors.py

@@ -15,7 +15,7 @@
 import os
 import sys
 from pathlib import Path
-from typing import Dict, List
+from typing import Any, Dict, List
 
 import nemo_run as run
 from nemo_run.config import get_nemorun_home
@@ -43,6 +43,7 @@
     "NVTE_NORM_FWD_USE_CUDNN": "1",
     "NVTE_NORM_BWD_USE_CUDNN": "1",
     "TORCH_NCCL_HIGH_PRIORITY": "1",
+    "HF_HUB_OFFLINE": "0",
 }
 
 
@@ -63,10 +64,18 @@ def slurm_executor(
     wandb_key: str = None,
     network: str = None,
     custom_bash_cmds: List[str] = None,
+    additional_slurm_params: Dict[str, Any] = None,
 ) -> run.SlurmExecutor:
     """
     Slurm cluster definition with appropriate cluster params and NeMo container params needed for pre-training
     and fine-tuning experiments
+
+    Args:
+        additional_slurm_params: Dict[str, Any], optional
+            Additional SLURM parameters to pass to sbatch. These will be converted to #SBATCH directives.
+            Example: {"nodelist": "node001,node002", "constraint": "gpu"} will generate:
+                #SBATCH --nodelist=node001,node002
+                #SBATCH --constraint=gpu
     """
     custom_bash_cmds = [] if custom_bash_cmds is None else custom_bash_cmds
     err_msgs = []
@@ -135,6 +144,7 @@ def slurm_executor(
         segment=segment,
         network=network,
         launcher=launcher,
+        additional_parameters=additional_slurm_params,
     )
 
     return executor