argument_parser.py

#
# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import logging
import os
import re
from pathlib import Path

from nemo_run.config import get_nemorun_home


logger: logging.Logger = logging.getLogger(__name__)

DEFAULT_NEMO_HOME = os.getenv("NEMO_HOME", Path.home() / ".cache" / "nemo")
VALID_CUDA_GRAPH_IMPLS = ["none", "local", "transformer_engine"]
VALID_CUDA_GRAPH_SCOPES = ["full_iteration", "attn", "mlp", "moe", "moe_router", "moe_preprocess", "mamba"]

NUM_GPUS_PER_NODE_MAP = {
    "h100": 8,
    "b200": 8,
    "b300": 8,
    "gb200": 4,
    "gb300": 4,
    "vrnvl72": 4,
    "r100": 1,
}

# Deprecated --gpu names and their replacements. When the deprecation window ends, drop an entry together with
# its counterpart in LEGACY_HARDWARE_ALIASES (scripts/training/recipe_metadata.py).
LEGACY_GPU_ALIASES = {"vr200": "vrnvl72"}


def list_of_strings(arg):
    """Split a comma-separated string into a list of substrings."""
    return arg.split(",")


def list_of_ints(arg):
    """Split a comma-separated string into a list of integers."""
    if arg is None:
        raise argparse.ArgumentTypeError("empty argument list")
    try:
        result = [int(p, 10) for p in list_of_strings(arg)]
    except Exception:
        raise argparse.ArgumentTypeError(f"invalid comma-separated integer list: {arg!r}") from None

    return result


def to_dict(arg):
    """Split a comma-separated string into a dictionary of key-value pairs."""
    return dict(item.split("=") for item in arg.split(","))


def parse_kv(s: str):
    """Parse a key-value pair from a string."""
    KEY_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")  # Useful check for errors like hyphen in var names
    if "=" not in s:
        raise argparse.ArgumentTypeError(f"Expected KEY=VALUE, got {s!r}")

    key, value = s.split("=", 1)

    if not KEY_RE.match(key):
        raise argparse.ArgumentTypeError(f"Invalid env var name: {key!r}")

    return key, value


def lower_str(arg):
    """Lowercase a CLI string argument with a runtime type check."""
    assert isinstance(arg, str), f"Argument {arg} is not a string"
    return arg.lower()


def canonical_gpu(arg: str) -> str:
    """Map a deprecated GPU name to its replacement before argparse checks the choices."""
    if arg in LEGACY_GPU_ALIASES:
        # Every rank re-parses the forwarded argv, so only the launcher and global rank 0 warn.
        if os.environ.get("RANK", os.environ.get("SLURM_PROCID", "0")) == "0":
            logger.warning("--gpu %s is deprecated; use --gpu %s.", arg, LEGACY_GPU_ALIASES[arg])
        return LEGACY_GPU_ALIASES[arg]
    return arg


def bool_arg(arg):
    """Convert a string CLI value to a boolean."""
    if arg.lower() in ["true", "1", "t", "yes", "y"]:
        return True
    elif arg.lower() in ["false", "0", "f", "no", "n"]:
        return False
    else:
        raise ValueError(f"Invalid value for boolean argument: {arg}")


def is_cuda_graph_impl_valid(arg):
    """Validate and normalize the CUDA graph implementation argument."""
    if arg in VALID_CUDA_GRAPH_IMPLS:
        return arg
    else:
        raise ValueError(f"Invalid value for cuda_graph_impl: {arg}. Valid options are: {VALID_CUDA_GRAPH_IMPLS}")


def is_cuda_graph_scope_valid(arg):
    """Validate the CUDA graph scope argument."""
    args = arg.split(",")
    if all(a in VALID_CUDA_GRAPH_SCOPES for a in args):
        return args
    else:
        raise ValueError(
            f"Invalid value for cuda_graph_scope: {arg}. Valid options are: {VALID_CUDA_GRAPH_SCOPES}. "
            "Comma separated list of scopes is allowed."
        )


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


def parse_cli_args():
    """
    Command line arguments for running pre-training and fine-tuning experiments.
    """
    parser = argparse.ArgumentParser(
        description="Megatron-Bridge Pretraining and Fine-Tuning",
        argument_default=None,
    )
    parser.add_argument(
        "--domain",
        type=lower_str,
        choices=["llm", "vlm", "qwen3vl", "diffusion"],
        help="Domain to use for experiment.",
        default="llm",
    )
    parser.add_argument(
        "-m",
        "--model_family_name",
        type=lower_str,
        help="Model family name to use for experiment. E.g. `--model_family_name llama` (not llama3)",
        required=True,
    )
    parser.add_argument(
        "-mr",
        "--model_recipe_name",
        type=lower_str,
        help="Model recipe name to use for experiment. E.g. `--model_recipe_name llama31_405b`",
        required=True,
    )
    parser.add_argument(
        "--use_recipes",
        action="store_true",
        help="Use model recipes instead of flat performance recipes. Disabled by default.",
        default=False,
    )
    parser.add_argument(
        "-nh",
        "--nemo_home",
        type=str,
        help=" ".join(
            [
                "Sets env var `NEMO_HOME` (on compute node using sbatch script)- directory where NeMo searches",
                "for models and checkpoints. This saves a lot of time (especially for bigger models) if checkpoints already",
                f"exist here. Missing files will be downloaded here from HuggingFace. Defaults to {DEFAULT_NEMO_HOME}",
            ]
        ),
        default=DEFAULT_NEMO_HOME,
    )
    parser.add_argument(
        "--detach",
        help="Detach the experiment from the terminal. Enabled by default",
        type=bool_arg,
        default=True,
    )
    parser.add_argument(
        "--max_retries",
        type=int,
        help="Maximum number of retries. Defaults to 2",
        default=2,
    )
    parser.add_argument(
        "-ng",
        "--num_gpus",
        type=int,
        help="Number of gpus.",
        required=True,
    )
    parser.add_argument(
        "--hidden_size",
        type=int,
        help="Hidden size to use for the experiment. Defaults to None.",
        required=False,
        default=None,
    )
    parser.add_argument(
        "--num_layers",
        type=int,
        help="Number of layers to use for the experiment. Defaults to None.",
        required=False,
        default=None,
    )
    parser.add_argument(
        "--num_moe_experts",
        type=int,
        help="Number of experts to use for the experiment. Defaults to None.",
        required=False,
        default=None,
    )
    parser.add_argument(
        "--pipeline_model_parallel_layout",
        type=str,
        help="Pipeline model parallel layout to use for the experiment. Defaults to None.",
        required=False,
        default=None,
    )
    parser.add_argument(
        "--first_k_dense_replace",
        type=int,
        help="Number of MoE layers to be converted to dense layers. Defaults to None.",
        required=False,
        default=None,
    )
    parser.add_argument(
        "-d",
        "--dryrun",
        help="Dry-run mode. In setup_experiment.py: prints the sbatch script without launching. "
        "In bootstrap.py: selects a training entrypoint, builds the full ConfigContainer with all overrides, "
        "saves it to --save_config_filepath (default: ConfigContainer.yaml), and exits without training.",
        required=False,
        action="store_true",
    )

    # Training
    training_args = parser.add_argument_group("Training arguments")
    training_args.add_argument(
        "--task",
        choices=["pretrain", "sft", "peft"],
        help="Task to run. Defaults to 'pretrain'",
        default="pretrain",
    )
    training_args.add_argument(
        "-ms",
        "--max_steps",
        type=int,
        help="Maximum number of steps to run the experiment for. Defaults to 50.",
    )
    training_args.add_argument(
        "-gb",
        "--global_batch_size",
        type=int,
    )
    training_args.add_argument(
        "-mb",
        "--micro_batch_size",
        type=int,
    )
    training_args.add_argument(
        "-sl",
        "--seq_length",
        type=int,
    )
    training_args.add_argument(
        "--distributed_timeout_minutes",
        type=int,
        help=(
            "Process-group init timeout in minutes (dist.distributed_timeout_minutes). "
            "Widen above the 10-minute default when a one-time dataset index build on rank 0 "
            "can outlast the NCCL collective watchdog while other ranks wait at the setup barrier."
        ),
    )

    # Optimizer
    optimizer_args = parser.add_argument_group("Optimizer arguments")
    optimizer_args.add_argument("--lr", type=float, help="Learning rate")
    optimizer_args.add_argument("--min_lr", type=float, help="Minimum learning rate")
    optimizer_args.add_argument("--warmup_iters", type=int, help="Warmup iterations", default=10)

    # Checkpointing
    checkpointing_args = parser.add_argument_group("Checkpointing arguments")
    checkpointing_args.add_argument("--pretrained_checkpoint", type=str, help="Path to pretrained checkpoint")
    checkpointing_args.add_argument("--save_dir", type=str, help="Directory to save checkpoints")
    checkpointing_args.add_argument("--load_dir", type=str, help="Directory to load checkpoints")
    checkpointing_args.add_argument("--save_interval", type=int, help="Number of iterations between checkpoint saves")
    checkpointing_args.add_argument("--most_recent_k", type=int, help="Number of latest checkpoints to keep")

    # Data
    data_args = parser.add_argument_group("Data arguments")
    data_args.add_argument(
        "--data",
        type=str,
        default="mock",
        choices=["mock", "rp2", "squad", "squad_packed", "c4"],
        help="Dataset type to use",
    )
    data_args.add_argument(
        "--c4_root",
        type=str,
        help="C4 preproc root dir (containing c4-train.en_<shard>_text_document.{bin,idx} and "
        "c4-validation-91205-samples.en_text_document.{bin,idx}). Used when --data c4. "
        "Mirrors the layout used by NVIDIA's MLPerf DSV3 671B reference.",
    )
    data_args.add_argument(
        "--c4_train_shards",
        nargs="+",
        type=int,
        default=[6],
        help="C4 train shard indices to blend (e.g. 6 for the 8b dataset; 6 7 for the full dataset). Default: [6].",
    )
    data_args.add_argument("--dataset_paths", nargs="*", help="Dataset paths (for rp2 dataset)")
    data_args.add_argument("--dataset_root", type=str, help="Dataset root directory (for squad datasets)")
    data_args.add_argument("--index_mapping_dir", type=str, help="Index mapping directory (for rp2 dataset)")
    data_args.add_argument("--dataset_name", type=str, help="Dataset name (deprecated)")
    data_args.add_argument("--packed_sequence", action="store_true", help="Use packed sequences")
    data_args.add_argument("--head_only", action="store_true", help="Use only head data (for rp2 dataset)")
    data_args.add_argument(
        "--diffusion_dataset_path",
        type=str,
        help="WebDataset root path for diffusion recipes (FLUX, WAN). When unset, recipes fall back to mock data.",
    )

    # Tokenizer configuration
    tokenizer_args = parser.add_argument_group("Tokenizer arguments")
    tokenizer_args.add_argument(
        "--tokenizer_type",
        type=str,
        choices=["NullTokenizer", "HuggingFaceTokenizer", "SentencePieceTokenizer"],
    )
    tokenizer_args.add_argument(
        "--tokenizer_model", type=str, help="Path to tokenizer model (automatically provided by launcher)"
    )
    tokenizer_args.add_argument("--vocab_size", type=int, default=32000, help="Vocabulary size for NullTokenizer")
    hf_mode = tokenizer_args.add_mutually_exclusive_group()
    hf_mode.add_argument(
        "-hf",
        "--hf_token",
        type=str,
        help="HuggingFace token. Defaults to None. Required for accessing tokenizers and checkpoints.",
    )
    hf_mode.add_argument(
        "--offline",
        action="store_true",
        help="Enable offline HuggingFace Hub mode by setting HF_HUB_OFFLINE=1.",
    )

    # Parallelism
    parallelism_args = parser.add_argument_group("Parallelism arguments")
    parallelism_args.add_argument(
        "-tp",
        "--tensor_model_parallel_size",
        type=int,
        help="Intra-layer model parallelism. Splits tensors across GPU ranks.",
    )
    parallelism_args.add_argument(
        "-pp",
        "--pipeline_model_parallel_size",
        type=int,
        help="Inter-layer model parallelism. Splits transformer layers across GPU ranks.",
    )
    parallelism_args.add_argument(
        "-cp",
        "--context_parallel_size",
        type=int,
        help="Splits network input along sequence dimension across GPU ranks.",
    )
    parallelism_args.add_argument(
        "-vp",
        "--virtual_pipeline_model_parallel_size",
        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.",
        default=-1,
    )
    parallelism_args.add_argument(
        "-ep",
        "--expert_model_parallel_size",
        type=int,
        help="Distributes Moe Experts across sub data parallel dimension.",
    )
    parallelism_args.add_argument(
        "-et",
        "--expert_tensor_parallel_size",
        type=lambda x: int(x) if x is not None else None,
        nargs="?",
        const=None,
        help="Intra-layer tensor model parallelsm for expert layer. Splits tensors across GPU ranks.\
            Use -et/--expert_tensor_parallel_size <space> for None or -et/--expert_tensor_parallel_size <int>",
    )

    # Slurm
    slurm_args = parser.add_argument_group("Slurm arguments")
    slurm_args.add_argument(
        "-a",
        "--account",
        type=str,
        help="Slurm account to use for experiment",
    )
    slurm_args.add_argument(
        "-p",
        "--partition",
        type=str,
        help="Slurm partition to use for experiment",
    )
    slurm_args.add_argument(
        "-t",
        "--time_limit",
        type=str,
        help="Maximum time limit to run experiment for. Defaults to 30 minutes (format- 'HH:MM:SS')",
        default="00:30:00",
    )
    slurm_args.add_argument(
        "--preempt_time",
        type=int,
        help=" ".join(
            [
                "For long-convergence runs only: seconds before the Slurm time limit at which to send",
                "SIGTERM so the training loop can save a checkpoint and exit gracefully before the hard",
                "walltime kill, letting the next slice resume from it. Must exceed the checkpoint flush",
                "time. Defaults to 60.",
            ]
        ),
        default=60,
    )
    slurm_args.add_argument(
        "-gn",
        "--gpus_per_node",
        type=int,
        help="Number of gpus per node. Defaults to None. If not provided, will be inferred from the GPU type.",
        default=None,
    )
    slurm_args.add_argument(
        "-i",
        "--container_image",
        type=str,
        help=" ".join(
            [
                "NeMo container to use for experiment. Defaults to latest dev container- 'nvcr.io/nvidia/nemo:dev'",
                "Make sure your NGC credentials are accessible in your environment.",
            ]
        ),
        default="nvcr.io/nvidia/nemo:dev",
    )
    slurm_args.add_argument(
        "-cm",
        "--custom_mounts",
        type=list_of_strings,
        help="Comma separated string of mounts",
        default=[],
    )
    slurm_args.add_argument(
        "-ce",
        "--custom_env_vars",
        type=to_dict,
        help="Comma separated string of environment variables",
        default={},
    )
    slurm_args.add_argument(
        "-E",
        "--env",
        action="append",
        type=parse_kv,
        metavar="KEY=VALUE",
        help="Set environment variable (repeatable arg). This is an alternative to --custom_env_vars \
        (--custom_env_vars is preferred for most cases). Example: -E var1=value1,value2 -E var2=value3",
    )
    slurm_args.add_argument(
        "-cs",
        "--custom_srun_args",
        type=list_of_strings,
        help="Comma separated string of srun arguments",
        default=[],
    )
    slurm_args.add_argument(
        "-cb",
        "--custom_bash_cmds",
        nargs="*",
        action="append",
        help="List of bash commands to execute before the main command",
        default=None,
    )
    slurm_args.add_argument(
        "--gres",
        type=str,
        help="Slurm generic resources to request (e.g., 'gpu:4').",
        required=False,
        default=None,
    )
    slurm_args.add_argument(
        "--additional_slurm_params",
        type=parse_additional_slurm_params,
        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,
    )
    slurm_args.add_argument(
        "--packager",
        type=str,
        choices=["git", "none"],
        default="git",
        help="How code is packaged for the job. 'git' snapshots the repo at submission time (default). "
        "'none' skips snapshotting — use when code is pre-installed in the container image or available via a shared filesystem.",
        required=False,
    )
    slurm_args.add_argument(
        "--enable_pct_binding",
        type=bool_arg,
        help="Enable PCT binding. Enabled by default.",
        required=False,
        default=True,
    )

    # Kubeflow
    kubeflow_args = parser.add_argument_group("Kubeflow arguments")
    kubeflow_args.add_argument(
        "--kubeflow_namespace",
        type=str,
        help="Kubernetes namespace for Kubeflow TrainJob. When set, uses the Kubeflow executor instead of Slurm.",
        required=False,
    )
    kubeflow_args.add_argument(
        "--csp",
        type=str,
        choices=["aws", "gcp"],
        default=None,
        help="Cloud provider of the Kubeflow cluster. Selects the CSP fabric plugin "
        "(aws -> EKSEnvPlugin/EFA, gcp -> GKEEnvPlugin/gIB). Omit on Slurm or when no "
        "CSP-specific fabric config is needed.",
        required=False,
    )
    kubeflow_args.add_argument(
        "--kubeflow_workdir_pvc",
        type=str,
        help="PVC name for syncing job workdir (launch scripts, packaged code) to the cluster before launch.",
        required=False,
    )
    kubeflow_args.add_argument(
        "--kubeflow_workdir_pvc_path",
        type=str,
        help="Mount path for the workdir PVC inside the training pod. Defaults to '/nemo_run'.",
        default="/nemo_run",
        required=False,
    )
    kubeflow_args.add_argument(
        "--kubeflow_workdir_local_path",
        type=str,
        help="Local directory whose contents nemo-run's KubeflowExecutor.package() "
        "rsyncs into the workdir PVC via a temporary alpine pod before launch. "
        "Used to overlay a --mbridge-ref checkout onto /opt/Megatron-Bridge in the "
        "trainer container without rebuilding the image.",
        required=False,
        default=None,
    )
    kubeflow_args.add_argument(
        "--kubeflow_image_pull_secrets",
        type=list_of_strings,
        help="Comma-separated list of Kubernetes image pull secret names.",
        required=False,
        default=[],
    )
    kubeflow_args.add_argument(
        "--kubeflow_volumes_json",
        type=str,
        help="JSON-encoded list of Kubernetes Volume dicts attached to the training pod "
        "(e.g. PVCs, emptyDir, hostPath).",
        required=False,
        default=None,
    )
    kubeflow_args.add_argument(
        "--kubeflow_volume_mounts_json",
        type=str,
        help="JSON-encoded list of Kubernetes VolumeMount dicts applied to the training container.",
        required=False,
        default=None,
    )
    kubeflow_args.add_argument(
        "--kubeflow_tolerations_json",
        type=str,
        help="JSON-encoded list of Kubernetes Toleration dicts applied to the training pods. "
        "Required for landing on tainted GPU nodes.",
        required=False,
        default=None,
    )
    kubeflow_args.add_argument(
        "--kubeflow_affinity_json",
        type=str,
        help="JSON-encoded Kubernetes Affinity dict applied to the training pods. "
        "Used to pin pods to specific nodes via nodeAffinity.",
        required=False,
        default=None,
    )
    kubeflow_args.add_argument(
        "--kubeflow_env_list_json",
        type=str,
        help="JSON-encoded list of Kubernetes EnvVar dicts (e.g. valueFrom.secretKeyRef entries) "
        "appended to the training container in addition to flat env_vars.",
        required=False,
        default=None,
    )
    kubeflow_args.add_argument(
        "--kubeflow_extra_resource_requests_json",
        type=str,
        help="JSON-encoded dict of extra container resource requests (e.g. {'vpc.amazonaws.com/efa': '32'}).",
        required=False,
        default=None,
    )
    kubeflow_args.add_argument(
        "--kubeflow_extra_resource_limits_json",
        type=str,
        help="JSON-encoded dict of extra container resource limits.",
        required=False,
        default=None,
    )
    kubeflow_args.add_argument(
        "--kubeflow_pod_spec_overrides_json",
        type=str,
        help="JSON-encoded dict merged into the pod spec (escape hatch for nodeSelector, hostNetwork, etc.).",
        required=False,
        default=None,
    )
    kubeflow_args.add_argument(
        "--kubeflow_pod_annotations_json",
        type=str,
        help="JSON-encoded dict of annotations applied to the trainer pod template metadata "
        "(e.g. networking.gke.io/interfaces to attach GKE RDMA NICs for gIB).",
        required=False,
        default=None,
    )
    kubeflow_args.add_argument(
        "--kubeflow_container_kwargs_json",
        type=str,
        help="JSON-encoded dict of extra fields set on the training container "
        "(e.g. {'securityContext': {'privileged': true}} for EFA/RDMA).",
        required=False,
        default=None,
    )
    kubeflow_args.add_argument(
        "--kubeflow_labels_json",
        type=str,
        help="JSON-encoded dict of labels applied to the TrainJob's pods.",
        required=False,
        default=None,
    )

    # Nvcre
    nvcre_args = parser.add_argument_group("Nvcre arguments")
    nvcre_args.add_argument(
        "--nvcre_namespace",
        type=str,
        help="Kubernetes namespace for Nvcre WorkloadRun. When set, uses the Nvcre executor instead of Slurm.",
        required=False,
    )
    nvcre_args.add_argument(
        "--nvcre_image_pull_secret",
        type=str,
        help="Kubernetes image pull secret name for pulling the container image.",
        required=False,
        default=None,
    )
    nvcre_args.add_argument(
        "--nvcre_workdir_pvc",
        type=str,
        help="PVC name for syncing job workdir to the cluster before launch.",
        required=False,
        default=None,
    )
    nvcre_args.add_argument(
        "--nvcre_workdir_pvc_path",
        type=str,
        help="Mount path for the workdir PVC inside the training pod.",
        default="/nemo_run",
        required=False,
    )
    nvcre_args.add_argument(
        "--nvcre_workdir_local_path",
        type=str,
        help="Local directory rsynced into the workdir PVC before launch.",
        required=False,
        default=None,
    )
    nvcre_args.add_argument(
        "--nvcre_node_selector_json",
        type=str,
        help="JSON-encoded dict of node selector labels for targeting specific nodes.",
        required=False,
        default=None,
    )
    nvcre_args.add_argument(
        "--nvcre_volumes_json",
        type=str,
        help="JSON-encoded list of Kubernetes Volume dicts.",
        required=False,
        default=None,
    )
    nvcre_args.add_argument(
        "--nvcre_volume_mounts_json",
        type=str,
        help="JSON-encoded list of Kubernetes VolumeMount dicts.",
        required=False,
        default=None,
    )
    nvcre_args.add_argument(
        "--nvcre_timeout_per_job",
        type=str,
        help="Per-job timeout passed to Nvcre orchestration (e.g. '24h').",
        required=False,
        default="24h",
    )

    nvcre_args.add_argument(
        "--nvcre_kubeconfig",
        type=str,
        help="Path to kubeconfig file for nvcrectl/kubectl.",
        required=False,
        default=None,
    )
    nvcre_args.add_argument(
        "--nvcre_kube_context",
        type=str,
        help="Kubernetes context to use with nvcrectl/kubectl.",
        required=False,
        default=None,
    )
    nvcre_args.add_argument(
        "--nvcre_gang_scheduler_name",
        type=str,
        help="Gang scheduler name for the WorkloadRun (e.g. 'kai-scheduler').",
        required=False,
        default=None,
    )
    nvcre_args.add_argument(
        "--nvcre_hf_token_secret_name",
        type=str,
        help="Name of the Kubernetes Secret holding the HuggingFace token.",
        required=False,
        default=None,
    )
    nvcre_args.add_argument(
        "--nvcre_hf_token_secret_key",
        type=str,
        help="Key inside the Kubernetes Secret that holds the HuggingFace token value.",
        required=False,
        default="token",
    )

    # For performance
    performance_args = parser.add_argument_group("Performance arguments")
    performance_args.add_argument(
        "-g",
        "--gpu",
        type=canonical_gpu,
        choices=NUM_GPUS_PER_NODE_MAP.keys(),
        help="Target gpu type. vr200 is a deprecated alias for vrnvl72.",
        required=True,
    )
    performance_args.add_argument(
        "-c",
        "--compute_dtype",
        type=str,
        choices=["bf16", "fp8_cs", "fp8_mx", "fp8_sc", "nvfp4"],
        help="Compute precision. Options- bf16 or fp8. Defaults to bf16",
        required=False,
        default="bf16",
    )
    performance_args.add_argument(
        "--optimizer_type",
        type=str,
        choices=["adam", "muon"],
        help="Optimizer type for recipes that support it (e.g. Kimi-K2). Defaults to muon.",
        required=False,
        default="muon",
    )
    performance_args.add_argument(
        "-vb",
        "--enable_vboost",
        help="Enable VBoost which steers more power towards tensor cores. Disabled by default",
        type=bool_arg,
        required=False,
    )
    performance_args.add_argument(
        "-lgc",
        "--lock_gpu_freq",
        help="Lock GPU graphics clock to the specified frequency in MHz via "
        "`sudo nvidia-smi -lgc <freq>`. Runs once per node before training. "
        "Use `nvidia-smi -rgc` to reset after the job.",
        type=int,
        required=False,
        default=None,
    )
    performance_args.add_argument(
        "-lmc",
        "--peak_mem_clk",
        help="Lock GPU memory clock to the specified peak frequency in MHz via "
        "`sudo nvidia-smi -lmc <freq>,<freq>`. Runs once per node before training. "
        "Defaults to 4752 MHz for VRNVL72 and is disabled by default for other GPUs. "
        "Pass `-lmc -1` or `--peak_mem_clk -1` to disable the VRNVL72 default.",
        type=int,
        required=False,
        default=None,
    )
    performance_args.add_argument(
        "-en",
        "--enable_nsys",
        help="Enable Nsys profiling. Disabled by default",
        action="store_true",
    )
    performance_args.add_argument(
        "--export_nsys_sqlite",
        help="Export a SQLite report after Nsys profiling finishes. Requires --enable_nsys.",
        action="store_true",
    )
    performance_args.add_argument(
        "-pyp",
        "--pytorch_profiler",
        type=bool_arg,
        help="Enable PyTorch profiler. Disabled by default",
        required=False,
        default=False,
    )
    performance_args.add_argument(
        "--profiling_start_step",
        type=int,
        help="Defines start step for profiling",
        required=False,
        default=10,
    )
    performance_args.add_argument(
        "--profiling_stop_step",
        type=int,
        help="Defines stop step for profiling",
        required=False,
        default=11,
    )
    performance_args.add_argument(
        "-mh",
        "--record_memory_history",
        type=bool_arg,
        help="Enable PyTorch profiler memory history recording. Enabled by default (if pytorch_profiler is enabled)",
        required=False,
        default=True,
    )
    performance_args.add_argument(
        "--profiling_gpu_metrics",
        help="Enable nsys gpu metrics. Disabled by default.",
        action="store_true",
    )
    performance_args.add_argument(
        "--profiling_ranks",
        type=list_of_ints,
        metavar="N[,N...]",
        help="List of ranks to target for profiling (defaults to just first rank)",
        required=False,
        default=None,
    )
    performance_args.add_argument(
        "--nsys_trace",
        type=list_of_strings,
        metavar="TRACE[,TRACE...]",
        help="Comma-separated list of events to trace during nsys profiling (e.g., 'cuda,nvtx'). Defaults to nemo_run defaults.",
        required=False,
        default=None,
    )
    performance_args.add_argument(
        "--nsys_extra_args",
        type=list_of_strings,
        metavar="ARG[,ARG...]",
        help="Comma-separated list of additional nsys arguments. Will be combined with default args.",
        required=False,
        default=None,
    )
    performance_args.add_argument(
        "--use_tokendrop",
        help="Use token drop. Disabled by default. Currently only supported for DeepSeek v3",
        type=bool_arg,
        required=False,
    )
    performance_args.add_argument(
        "--use_megatron_fsdp",
        help="Use Megatron FSDP. Disabled by default.",
        type=bool_arg,
        required=False,
    )
    performance_args.add_argument(
        "--nccl_ub",
        help="Enable NCCL user buffer for FSDP communication. Disabled by default.",
        type=bool_arg,
        required=False,
    )
    performance_args.add_argument(
        "--cuda_graph_impl",
        help=f"Cuda graph implementation. Options- {', '.join(VALID_CUDA_GRAPH_IMPLS)}.",
        type=is_cuda_graph_impl_valid,
        required=False,
    )
    performance_args.add_argument(
        "--cuda_graph_scope",
        help=(
            f"Cuda graph scope. Options- {VALID_CUDA_GRAPH_SCOPES}. Comma separated list of "
            "scopes is allowed. Scoped captures such as attn or mlp require "
            "cuda_graph_impl=transformer_engine; local is reserved for full_iteration."
        ),
        type=is_cuda_graph_scope_valid,
        required=False,
    )
    performance_args.add_argument(
        "--moe_a2a_overlap",
        type=bool_arg,
        required=False,
    )
    performance_args.add_argument(
        "-rl",
        "--recompute_num_layers",
        type=int,
        help="Number of Transformer layers to recompute, where all the intermediate "
        "activations of a Transformer layer are computed. Defaults to None",
        required=False,
    )
    performance_args.add_argument(
        "-ol",
        "--activation_offload_layers",
        type=int,
        help="Number of Transformer layers to offload to the CPU memory. Defaults to None",
        required=False,
    )
    performance_args.add_argument(
        "--recompute_modules",
        type=list_of_strings,
        help="Comma separated list of modules to recompute. Defaults to None",
        required=False,
    )
    performance_args.add_argument(
        "--moe_flex_dispatcher_backend",
        type=lambda x: None if x == "None" else x,
        help="MoE flex dispatcher backend. Options- deepep, hybridep, ncclep, None. If None, will use alltoall dispatcher.",
        choices=["deepep", "hybridep", "ncclep", None],
        required=False,
        # -1 means the option was omitted and the recipe backend must be kept;
        # None means the user explicitly requested the alltoall dispatcher.
        default=-1,
    )
    performance_args.add_argument(
        "--deterministic",
        help="Enable bit-exact deterministic training. Sets NCCL/cuBLAS/TE env vars "
        "and disables fused cross-entropy loss and TP comm overlap.",
        action="store_true",
    )

    # Logging
    logging_args = parser.add_argument_group("Logging arguments")
    logging_args.add_argument(
        "-wdk",
        "--wandb_key",
        type=str,
        help="wandb key. Needed for wandb logger projection to server",
        required=False,
    )
    logging_args.add_argument(
        "-wdp",
        "--wandb_project_name",
        type=str,
        help="wandb project name",
        required=False,
    )
    logging_args.add_argument(
        "-wde",
        "--wandb_entity_name",
        type=str,
        help="wandb project name",
        required=False,
    )
    logging_args.add_argument(
        "-wdj",
        "--wandb_experiment_name",
        type=str,
        help="wandb job name",
        required=False,
    )
    logging_args.add_argument(
        "-wds",
        "--wandb_save_dir",
        type=str,
        help="wandb save directory",
        required=False,
    )
    logging_args.add_argument(
        "-l",
        "--log_dir",
        type=str,
        help=f"Directory for logging experiment results. Defaults to {get_nemorun_home()} or NEMORUN_HOME envvar",
        required=False,
        default=None,
    )
    logging_args.add_argument("--save_config_filepath", type=str, help="Path to save the task configuration file")
    logging_args.add_argument(
        "--dump_env",
        action="store_true",
        help="Write environment variables to /nemo_run/env_<SLURM_JOB_ID>.log on rank 0. "
        "Useful for post-run debugging of NCCL, CUDA, and SLURM settings.",
    )

    # Config variant selection
    config_variant_args = parser.add_argument_group("Config variant arguments")
    config_variant_args.add_argument(
        "-cv",
        "--config_variant",
        type=str,
        help=(
            "Flat perf recipe variant for setup_experiment.py. Omit to use the suffix-less canonical recipe. "
            "Named variants such as large_scale are supported when a matching recipe exists."
        ),
    )
    config_variant_args.add_argument(
        "--list_config_variants",
        action="store_true",
        help="List available config variants for the specified model/task/gpu/dtype and interactively select one (with 15s timeout).",
    )

    _testing_args(parser)

    return parser


def _testing_args(parser):
    """Add testing-related arguments to the parser."""
    testing_args = parser.add_argument_group("Testing arguments")
    testing_args.add_argument(
        "--is_long_convergence_run",
        action="store_true",
        help="If true, runs a long convergence run.",
        required=False,
        default=False,
    )
    testing_args.add_argument(
        "--retry_on_testing_failure",
        action="store_true",
        help=(
            "If true, when the post-training convergence/perf check fails, the "
            "experiment is retrained from scratch and re-tested in the next "
            "outer-loop iteration (bounded by --max_retries). Off by default: "
            "a single failed test surfaces immediately as AssertionError, "
            "which is the right behavior for CI."
        ),
        required=False,
        default=False,
    )
    testing_args.add_argument(
        "--golden_values_path",
        type=str,
        help="Path to golden values file",
        required=False,
    )
    testing_args.add_argument(
        "--timing_threshold", type=float, default=0.05, help="Step timing validation threshold (default: 0.05 = 5%%)"
    )
    testing_args.add_argument(
        "--skip_first_percent_time",
        type=float,
        default=0.70,
        help="Percentage of iterations to skip for timing comparison (default: 0.75 = 75%%)",
    )
    testing_args.add_argument(
        "--correlation_threshold", type=float, default=0.95, help="Correlation threshold for loss curve validation"
    )
    testing_args.add_argument(
        "--high_loss_tolerance", type=float, default=0.10, help="Tolerance for high loss values (>2.0)"
    )
    testing_args.add_argument(
        "--medium_loss_tolerance", type=float, default=0.05, help="Tolerance for medium loss values (0.5-2.0)"
    )
    testing_args.add_argument(
        "--low_loss_tolerance", type=float, default=0.02, help="Tolerance for low loss values (<0.5)"
    )
    testing_args.add_argument(
        "--final_loss_tolerance", type=float, default=0.05, help="Tolerance for final loss value"
    )
    testing_args.add_argument("--max_outlier_ratio", type=float, default=0.1, help="Maximum ratio of outliers allowed")
    testing_args.add_argument(
        "--outlier_threshold", type=float, default=3.0, help="Outlier detection threshold (sigma)"
    )
    testing_args.add_argument(
        "--skip_first_percent_loss",
        type=float,
        default=0.20,
        help="Percentage of loss points to skip from beginning for convergence analysis",
    )
    testing_args.add_argument(
        "--memory_threshold", type=float, default=0.05, help="Memory validation threshold (default: 0.05 = 5%%)"
    )
    testing_args.add_argument(
        "--eval_time_start_step",
        type=int,
        default=None,
        help="Start step (0-indexed, inclusive) for timing average window. Overrides skip_first_percent_time when set.",
    )
    testing_args.add_argument(
        "--eval_time_end_step",
        type=int,
        default=None,
        help="End step (0-indexed, exclusive) for timing average window. If None, averages to end.",
    )