RASMUS/Finnish-ASR-Canary-v2

1# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14 15import argparse16 17from nemo_run.config import get_nemorun_home18 19from .utils import DEFAULT_NEMO_HOME20 21 22def parse_cli_args():23    """24    Command line arguments correspong to Slurm cluster and NeMo2.0 for running pre-training and25    fine-tuning experiments.26    """27    parser = argparse.ArgumentParser(description="NeMo2.0 Performance Pretraining and Fine-Tuning")28 29    parser.add_argument(30        "-a",31        "--account",32        type=str,33        help="Slurm account to use for experiment",34        required=True,35    )36    parser.add_argument(37        "-p",38        "--partition",39        type=str,40        help="Slurm partition to use for experiment",41        required=True,42    )43    parser.add_argument(44        "-g",45        "--gpu",46        type=str,47        choices=["h100", "b200", "gb200"],48        help="Target gpu type.",49        required=True,50    )51    parser.add_argument(52        "-l",53        "--log_dir",54        type=str,55        help=f"Directory for logging experiment results. Defaults to {get_nemorun_home()}",56        required=False,57        default=get_nemorun_home(),58    )59    parser.add_argument(60        "-t",61        "--time_limit",62        type=str,63        help="Maximum time limit to run experiment for. Defaults to 30 minutes (format- 'HH:MM:SS')",64        required=False,65        default="00:30:00",66    )67 68    parser.add_argument(69        "--additional_slurm_params",70        type=str,71        help="Additional SLURM parameters as key=value pairs. "72        "Use semicolons (;) to separate parameters when values contain commas. "73        "Examples: 'nodelist=node001,node002;constraint=gpu' or 'reservation=my_res;exclusive'",74        required=False,75        default=None,76    )77    container_img_msg = [78        "NeMo container to use for experiment. Defaults to latest dev container- 'nvcr.io/nvidia/nemo:dev'",79        "Make sure your NGC credentials are accessible in your environment.",80    ]81    parser.add_argument(82        "-i",83        "--container_image",84        type=str,85        help=" ".join(container_img_msg),86        required=False,87        default="nvcr.io/nvidia/nemo:dev",88    )89    parser.add_argument(90        "-c",91        "--compute_dtype",92        type=str,93        choices=["bf16", "fp8"],94        help="Compute precision. Options- bf16 or fp8. Defaults to bf16",95        required=False,96        default="bf16",97    )98    fp8_recipe_msg = (99        "FP8 recipe. Options- ds (per-tensor delayed scaling), cs (per-tensor current scaling), "100        "mxfp8, ss (subchannel scaling). Defaults to ds"101    )102    parser.add_argument(103        "-fr",104        "--fp8_recipe",105        type=str,106        choices=["ds", "cs", "mxfp8", "ss"],107        help=fp8_recipe_msg,108        required=False,109        default="ds",110    )111    parser.add_argument(112        "-en",113        "--enable_nsys",114        help="Enable Nsys profiling. Diabled by default",115        action="store_true",116    )117    parser.add_argument(118        "-em",119        "--enable_memory_profile",120        help="Enable memory usage profiling. Diabled by default",121        action="store_true",122    )123    parser.add_argument(124        "-mp",125        "--memory_profile_out_path",126        type=str,127        help="Path to the output file of memory profiling",128        required=False,129        default=None,130    )131    parser.add_argument(132        "-tb",133        "--tensorboard",134        help="Enable tensorboard logging. Disabled by default",135        action="store_true",136    )137    parser.add_argument(138        "-wd",139        "--wandb",140        help="Enable wandb logging. Disabled by default",141        action="store_true",142    )143    parser.add_argument(144        "-wdk",145        "--wandb_key",146        type=str,147        help="wandb key. Needed for wandb logger projetion to server",148        required=False,149        default=None,150    )151    parser.add_argument(152        "-wdp",153        "--wandb_prj_name",154        type=str,155        help="wandb project name",156        required=False,157        default=None,158    )159    parser.add_argument(160        "-wdj",161        "--wandb_job_name",162        type=str,163        help="wandb job name",164        required=False,165        default=None,166    )167    parser.add_argument(168        "-f",169        "--finetuning",170        choices=["sft", "lora"],171        help="Finetuning scheme to use. Defaults to 'lora'",172        default='lora',173    )174    parser.add_argument(175        "-hf",176        "--hf_token",177        type=str,178        help="HuggingFace token. Defaults to None. Required for accessing tokenizers and checkpoints.",179        default=None,180    )181    nemo_home_msg = [182        "Sets env var `NEMO_HOME` (on compute node using sbatch script)- directory where NeMo searches",183        "for models and datasets. This saves a lot of time (especially for bigger models) if checkpoints already",184        f"exist here. Missing files will be downloaded here from HuggingFace. Defaults to {DEFAULT_NEMO_HOME}",185    ]186    parser.add_argument(187        "-nh",188        "--nemo_home",189        type=str,190        help=" ".join(nemo_home_msg),191        default=DEFAULT_NEMO_HOME,192    )193    parser.add_argument(194        "-d",195        "--dryrun",196        help="If true, prints sbatch script to terminal without launching experiment.",197        required=False,198        action="store_true",199    )200    parser.add_argument(201        "-tp",202        "--tensor_parallel_size",203        type=int,204        help="Intra-layer model parallelism. Splits tensors across GPU ranks.",205        required=False,206        default=None,207    )208    parser.add_argument(209        "-pp",210        "--pipeline_parallel_size",211        type=int,212        help="Inter-layer model parallelism. Splits transformer layers across GPU ranks.",213        required=False,214        default=None,215    )216    parser.add_argument(217        "-cp",218        "--context_parallel_size",219        type=int,220        help="Splits network input along sequence dimension across GPU ranks.",221        required=False,222        default=None,223    )224    parser.add_argument(225        "-vp",226        "--virtual_pipeline_parallel_size",227        type=int,228        help="Number of virtual blocks per pipeline model parallel rank is the virtual model parallel size.",229        required=False,230        default=None,231    )232    parser.add_argument(233        "-ep",234        "--expert_parallel_size",235        type=int,236        help="Distributes Moe Experts across sub data parallel dimension.",237        required=False,238        default=None,239    )240    parser.add_argument(241        "-et",242        "--expert_tensor_parallel_size",243        type=lambda x: int(x) if x is not None else None,244        nargs="?",245        const=None,246        help="Intra-layer tensor model parallelsm for expert layer. Splits tensors across GPU ranks.\247            Use -et/--expert_tensor_parallel_size <space> for None or -et/--expert_tensor_parallel_size <int>",248        required=False,249        default=None,250    )251    parser.add_argument(252        "-mb",253        "--micro_batch_size",254        type=int,255        required=False,256        default=None,257    )258    parser.add_argument(259        "-gb",260        "--global_batch_size",261        type=int,262        required=False,263        default=None,264    )265    parser.add_argument(266        "-ng",267        "--num_gpus",268        type=int,269        help="Number of gpus.",270        required=False,271        default=None,272    )273    parser.add_argument(274        "-gn",275        "--gpus_per_node",276        type=int,277        help="Number of gpus per node. Defaults to 8",278        required=False,279        default=8,280    )281    parser.add_argument(282        "-ms",283        "--max_steps",284        type=int,285        help="Number of train steps. Defaults to 100",286        required=False,287        default=100,288    )289 290    def bool_arg(arg):291        if arg.lower() in ['true', '1', 't', 'yes', 'y']:292            return True293        elif arg.lower() in ['false', '0', 'f', 'no', 'n']:294            return False295        else:296            raise ValueError(f"Invalid value for boolean argument: {arg}")297 298    parser.add_argument(299        "-cg",300        "--cuda_graphs",301        help="Enable CUDA graphs. Disabled by default",302        type=bool_arg,303        required=False,304        default=None,  # NOTE: DO NOT SET DEFAULT TO FALSE, IT WILL BE OVERRIDDEN BY THE RECOMMENDED MODEL CONFIGS305    )306    parser.add_argument(307        "-fsdp",308        "--use_mcore_fsdp",309        help="Enable Megatron Core (Mcore) FSDP. Disabled by default",310        type=bool_arg,311        required=False,312        default=None,313    )314    parser.add_argument(315        "-fsdp_db",316        "--use_fsdp_double_buffer",317        help="Enable FSDP double buffer. Disabled by default",318        type=bool_arg,319        required=False,320        default=None,321    )322    parser.add_argument(323        "-ubr",324        "--use_user_buffer_registration",325        help="Enable user buffer registration. Disabled by default",326        type=bool_arg,327        required=False,328        default=None,329    )330    parser.add_argument(331        "-sharp",332        "--use_sharp",333        help="Enable sharp. Disabled by default",334        type=bool_arg,335        required=False,336        default=None,337    )338    parser.add_argument(339        "-rl",340        "--recompute_layers",341        type=int,342        help="Number of Transformer layers to recompute, where all the intermediate "343        "activations of a Transformer layer are computed. Defaults to None",344        required=False,345        default=None,346    )347    parser.add_argument(348        "-ol",349        "--activation_offload_layers",350        type=int,351        help="Number of Transformer layers to offload to the CPU memory. Defaults to None",352        required=False,353        default=None,354    )355    parser.add_argument(356        "--nccl_communicator_config_path",357        type=str,358        help="Path to NCCL communicator config yaml file",359        required=False,360        default=None,361    )362 363    def list_of_strings(arg):364        return arg.split(',')365 366    parser.add_argument(367        "-rm",368        "--recompute_modules",369        nargs="*",370        const=None,371        type=str,372        help="List of modules to perform selective activation recompute. "373        "Users can provide 0 or any number of arguments. Defaults to None",374        required=False,375        default=None,376    )377    parser.add_argument(378        "-cm",379        "--custom_mounts",380        type=list_of_strings,381        help="Comma separated string of mounts",382        required=False,383        default=[],384    )385    parser.add_argument(386        "--use_hf_tokenizer",387        help="Use HuggingFace tokenizer. Disabled by default. Null tokenizer will be used if not provided.",388        action="store_true",389        required=False,390    )391    parser.add_argument(392        "-dcdfr",393        "--dump_config_diff_from_base_recipe",394        help="Dump the config diff from the base recipe. Defaults to False",395        action="store_true",396        required=False,397        default=False,398    )399    parser.add_argument(400        "--keep_fsdp_fp8_transpose_cache",401        help="Keep FSDP FP8 transpose cache. Disabled by default",402        type=bool_arg,403        required=False,404        default=None,405    )406    parser.add_argument(407        "-vb",408        "--enable_vboost",409        help="Enable VBoost which steers more power towards tensor cores. Disabled by default",410        type=bool_arg,411        required=False,412        default=None,413    )414    parser.add_argument(415        "--use_te_op_fuser",416        help="Enable Transformer Engine's operation fuser. This feature is experimental and disabled by default",417        type=bool_arg,418        required=False,419        default=None,420    )421    parser.add_argument(422        "--use_te_act_func",423        help="Use TE activation function for the MLP part.",424        type=bool_arg,425        required=False,426        default=None,427    )428    parser.add_argument(429        "--act_func_fp8_input_store",430        help="Store input of activation function in FP8 (tensorwise recipe). Disabled by default",431        type=bool_arg,432        required=False,433        default=False,434    )435    parser.add_argument(436        "--detach",437        help="Detach from experiment. Default is True. Unset to keep the process running.",438        type=bool_arg,439        required=False,440        default=True,441    )442 443    return parser444 445 446def parse_additional_slurm_params(params_str):447    """448    Parse additional SLURM parameters from a string of key=value pairs.449 450    This function handles different separator formats:451    1. Semicolon-separated: "key1=value1;key2=value2" (recommended for multiple parameters)452    2. Space-separated: "key1=value1 key2=value2"453    3. Single parameter: "key1=value1,value2" (no separators = single parameter)454 455    Args:456        params_str (str): String with parameters457 458    Returns:459        dict: Dictionary of parameters, or None if params_str is None/empty460 461    Example:462        parse_additional_slurm_params("nodelist=node001,node002")463        returns {"nodelist": "node001,node002"}464 465        parse_additional_slurm_params("nodelist=node001,node002;constraint=gpu")466        returns {"nodelist": "node001,node002", "constraint": "gpu"}467 468        parse_additional_slurm_params("reservation=my_res;constraint=gpu")469        returns {"reservation": "my_res", "constraint": "gpu"}470    """471    if not params_str:472        return None473 474    params = {}475 476    # Try semicolon separation first (most reliable for complex values)477    if ';' in params_str:478        parts = params_str.split(';')479    # Try space separation next480    elif ' ' in params_str:481        parts = params_str.split()482    # No separators found - treat as single parameter483    else:484        parts = [params_str]485 486    for part in parts:487        part = part.strip()488        if not part:489            continue490 491        if '=' in part:492            key, value = part.split('=', 1)493            params[key.strip()] = value.strip()494        else:495            # Boolean flag (no value)496            params[part] = True497 498    return params if params else None499