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 os16from pathlib import Path17from typing import List, Optional18 19import nemo_run as run20import pandas as pd21from numpy import nan22 23from nemo.collections.llm.gpt.data.mock import MockDataModule24from nemo.collections.llm.recipes.precision.mixed_precision import (25    bf16_with_fp8_current_scaling_mixed,26    bf16_with_fp8_mixed,27    bf16_with_fp8_subchannel_scaling_mixed,28    bf16_with_mxfp8_mixed,29)30from nemo.lightning.pytorch.callbacks.flops_callback import FLOPsMeasurementCallback31from nemo.lightning.pytorch.callbacks.model_checkpoint import ModelCheckpoint32from nemo.utils import logging33 34from .utils import get_comm_overlap_callback_idx35 36 37def get_csv_configs(gpu: str, task: str, model_name: str, model_size: str, args) -> pd.DataFrame:38    """39    Get recommended configs tuned for performance from a csv file.40    User (command line) provided args override the recommended configs.41    """42    script_dir = str(Path(__file__).parent.absolute())43    recommended_configs_csv = os.path.join(script_dir, "recommended_model_configs", f"model_configs_{gpu}.csv")44    logging.info(f"Using {recommended_configs_csv} for loading default recommended model configs")45 46    config_df = pd.DataFrame()47    if os.path.isfile(recommended_configs_csv):48        df = pd.read_csv(recommended_configs_csv)49        config_df = df[50            (df["task"] == task)51            & (df["model"] == model_name)52            & (df["size"] == model_size)53            & (df["dtype"] == args.compute_dtype)54            & (args.num_gpus is None or df['num_gpus'] == args.num_gpus)55        ]56        config_df = config_df.replace({nan: None})57        if len(config_df) == 0:58            logging.warning(f"Missing performance configs for {task}-{model_name}-{model_size}-{args.compute_dtype}")59            logging.warning("Make sure you provide all necessary arguments in the command line")60 61    config = config_df.to_dict(orient='records')[0] if len(config_df) > 0 else {}62 63    return config64 65 66def get_user_configs(gpu: str, task: str, model_name: str, model_size: str, args) -> List[int]:67    """68    Choose recommended configs tuned for performance from a csv file if available.69    User (command line) provided args override the recommended configs.70 71    NOTE: pre-train and PEFT recommended configs available for H100 and B200.72 73    Args:74        gpu (str): target GPU machine for experiment. Options- ['h100', 'b200']75        task (str): experiment task. Options- ['pre_train', 'sft', 'lora']76        model_name (str): target model for experiment. E.g.: 'llama3', 'mixtral'77        model_size (str): size of target model. E.g.: '8b' (for llama3)78    """79    config = get_csv_configs(gpu.lower(), task, model_name, model_size, args)80 81    if gpu.lower() == "gb200" and args.gpus_per_node > 4:82        args.gpus_per_node = 483        logging.warning("GB200 has 4 GPUs per node. Setting gpus_per_node to 4.")84    num_gpus = config.get("num_gpus") if args.num_gpus is None else args.num_gpus85    num_nodes = -(num_gpus // -args.gpus_per_node)  # ceil division86    mbs = config.get("mbs") if args.micro_batch_size is None else args.micro_batch_size87    gbs = config.get("gbs") if args.global_batch_size is None else args.global_batch_size88    tp_size = config.get("tp_size") if args.tensor_parallel_size is None else args.tensor_parallel_size89    pp_size = config.get("pp_size") if args.pipeline_parallel_size is None else args.pipeline_parallel_size90    cp_size = config.get("cp_size") if args.context_parallel_size is None else args.context_parallel_size91    ep_size = config.get("ep_size") if args.expert_parallel_size is None else args.expert_parallel_size92    vp_size = args.virtual_pipeline_parallel_size93    vp_size = config.get("vp_size") if vp_size is None else vp_size94    etp_size = args.expert_tensor_parallel_size95    etp_size = config.get("etp_size") if etp_size is None else etp_size96 97    enable_cuda_graphs = config.get("cuda_graphs") if args.cuda_graphs is None else args.cuda_graphs98    enable_cuda_graphs = False if enable_cuda_graphs is None else bool(int(enable_cuda_graphs))99 100    use_mcore_fsdp = config.get("use_mcore_fsdp") if args.use_mcore_fsdp is None else args.use_mcore_fsdp101    use_mcore_fsdp = False if use_mcore_fsdp is None else bool(int(use_mcore_fsdp))102 103    recompute_layers = config.get("recompute_layers") if args.recompute_layers is None else args.recompute_layers104    recompute_layers = 0 if recompute_layers is None else int(recompute_layers)105    activation_offload_layers = (106        config.get("activation_offload_layers")107        if args.activation_offload_layers is None108        else args.activation_offload_layers109    )110    activation_offload_layers = 0 if activation_offload_layers is None else int(activation_offload_layers)111 112    if args.recompute_modules is not None:113        recompute_modules = args.recompute_modules114        assert isinstance(recompute_modules, list), "recompute_modules must be a list"115    elif config.get("recompute_modules") is not None:116        recompute_modules = config.get("recompute_modules").split('/')117    else:118        recompute_modules = None119 120    keep_fsdp_fp8_transpose_cache = (121        config.get("keep_fsdp_fp8_transpose_cache")122        if args.keep_fsdp_fp8_transpose_cache is None123        else args.keep_fsdp_fp8_transpose_cache124    )125    keep_fsdp_fp8_transpose_cache = (126        False if keep_fsdp_fp8_transpose_cache is None else bool(int(keep_fsdp_fp8_transpose_cache))127    )128 129    use_user_buffer_registration = (130        config.get("use_user_buffer_registration")131        if args.use_user_buffer_registration is None132        else args.use_user_buffer_registration133    )134    use_user_buffer_registration = (135        False if use_user_buffer_registration is None else bool(int(use_user_buffer_registration))136    )137 138    use_sharp = config.get("use_sharp") if args.use_sharp is None else args.use_sharp139    use_sharp = False if use_sharp is None else bool(int(use_sharp))140 141    kwargs = num_nodes, mbs, gbs, tp_size, pp_size, cp_size, vp_size, ep_size, etp_size142    kwargs = [int(arg) if arg is not None else arg for arg in kwargs]143    kwargs += [144        enable_cuda_graphs,145        use_mcore_fsdp,146        recompute_layers,147        activation_offload_layers,148        recompute_modules,149        keep_fsdp_fp8_transpose_cache,150        use_user_buffer_registration,151        use_sharp,152    ]153 154    # print the received arguments for users to debug155    logging.info("Received model parallel configs: ")156    logging.info(f"{num_nodes=}")157    logging.info(f"num_gpus_per_node={args.gpus_per_node}")158    logging.info(f"{mbs=}")159    logging.info(f"{gbs=}")160    logging.info(f"{tp_size=}")161    logging.info(f"{pp_size=}")162    logging.info(f"{cp_size=}")163    logging.info(f"{vp_size=}")164    logging.info(f"{ep_size=}")165    logging.info(f"{etp_size=}")166    logging.info(f"{enable_cuda_graphs=}")167    logging.info(f"{use_mcore_fsdp=}")168    logging.info(f"{recompute_layers=}")169    logging.info(f"{activation_offload_layers=}")170    logging.info(f"{recompute_modules=}")171    logging.info(f"{keep_fsdp_fp8_transpose_cache=}")172    logging.info(f"{use_user_buffer_registration=}")173    logging.info(f"{use_sharp=}")174 175    return kwargs176 177 178def set_mcore_fsdp_configs(recipe, comm_overlap_callback_idx: int | None, tp_size: int | None):179    """180    Set Mcore FSDP related configs.181    """182    recipe.model.config.init_model_with_meta_device = True183    recipe.trainer.strategy.fsdp = "megatron"184    recipe.trainer.strategy.ddp.data_parallel_sharding_strategy = "optim_grads_params"185    # At fp32 gradient, `recipe.trainer.strategy.ddp.gradient_reduce_div_fusion` is used for fusion186    if recipe.trainer.plugins.grad_reduce_in_fp32:187        recipe.trainer.strategy.ddp.average_in_collective = False188    recipe.trainer.strategy.ddp.keep_fp8_transpose_cache = False189 190    try:191        recipe.trainer.strategy.ddp.keep_fp8_transpose_cache = False192    except AttributeError:193        recipe.trainer.strategy.ddp.keep_fp8_transpose_cache_when_using_custom_fsdp = False194        logging.warning(195            "Deprecation Notice: `keep_fp8_transpose_cache_when_using_custom_fsdp` "196            "will be deprecated in M-Core 0.14. "197            "Please use `keep_fsdp_fp8_transpose_cache` instead."198        )199    recipe.model.config.gradient_accumulation_fusion = False200    if (201        comm_overlap_callback_idx is not None202        and recipe.trainer.callbacks[comm_overlap_callback_idx].defer_embedding_wgrad_compute203    ):204        logging.warning("Disabling deferring embedding wgrad compute because it cannot work with FSDP together.")205        recipe.trainer.callbacks[comm_overlap_callback_idx].defer_embedding_wgrad_compute = False206 207    return recipe208 209 210def set_precision_configs(recipe, compute_dtype: str, fp8_recipe: str | None = None):211    """212    Set precision related configs.213    """214    if compute_dtype is None:215        return recipe216 217    if compute_dtype.lower() == "bf16":218        recipe.optim.config.use_precision_aware_optimizer = True219 220    if compute_dtype is not None and compute_dtype.lower() == "fp8":221        if fp8_recipe is None:222            fp8_recipe = "ds"223        if fp8_recipe.lower() == "ds":224            recipe.trainer.plugins = bf16_with_fp8_mixed()225        elif fp8_recipe.lower() == "cs":226            recipe.trainer.plugins = bf16_with_fp8_current_scaling_mixed()227            # disable first/last layer bf16 for benchmarking228            recipe.trainer.plugins.first_last_layers_bf16 = False229        elif fp8_recipe.lower() == "mxfp8":230            recipe.trainer.plugins = bf16_with_mxfp8_mixed()231        elif fp8_recipe.lower() == "ss":232            recipe.trainer.plugins = bf16_with_fp8_subchannel_scaling_mixed()233 234    recipe.trainer.plugins.grad_reduce_in_fp32 = False235 236    # Enable reuse_grad_buf_for_mxfp8_param_ag for MXFP8 and disable AG overlap237    # because it is not supported with reuse_grad_buf_for_mxfp8_param_ag238    if compute_dtype.lower() == "fp8" and fp8_recipe.lower() == "mxfp8":239        comm_overlap_callback_idx = get_comm_overlap_callback_idx(recipe.trainer.callbacks)240        if comm_overlap_callback_idx is not None:241            recipe.trainer.callbacks[comm_overlap_callback_idx].overlap_param_gather = False242        logging.warning(243            "When using MXFP8, to reduce memory usage, we use reuse_grad_buf_for_mxfp8_param_ag. "244            "Disabling AG overlap because it is not supported with reuse_grad_buf_for_mxfp8_param_ag."245        )246 247    return recipe248 249 250def set_recompute_configs(251    recipe,252    recompute_layers: int,253    activation_offload_layers: int,254    recompute_modules: Optional[List[str]],255):256    """257    Set activation recomputing and offloading related configs.258    """259    if recompute_layers > 0:260        recipe.model.config.recompute_granularity = "full"261        recipe.model.config.recompute_method = "block"262        recipe.model.config.recompute_num_layers = recompute_layers263 264    # Activation cpu offloading265    if activation_offload_layers > 0:266        recipe.model.config.cpu_offloading = True267        recipe.model.config.cpu_offloading_weights = False268        recipe.model.config.cpu_offloading_num_layers = activation_offload_layers269 270    # Activation recompute configs271    if recompute_modules is not None:272        recipe.model.config.recompute_modules = recompute_modules273        assert (274            recipe.model.config.recompute_granularity == "selective"275        ), "recompute_granularity must be selective when recompute_modules is provided"276        assert (277            recipe.model.config.recompute_num_layers is None278        ), "recompute_num_layers must be None when recompute_modules is provided"279 280    return recipe281 282 283def set_cuda_graph_configs(recipe, enable_cuda_graphs: bool, task: str):284    """285    Set CUDA graph related configs.286    """287    recipe.model.config.enable_cuda_graph = enable_cuda_graphs288    recipe.trainer.strategy.use_te_rng_tracker = enable_cuda_graphs289    if (290        task in ["none", "lora"]291        and hasattr(recipe.data, "packed_sequence_specs")292        and recipe.data.packed_sequence_specs is not None293    ):294        recipe.data.packed_sequence_specs.pad_cu_seqlens = enable_cuda_graphs295 296    return recipe297 298 299def set_full_iteration_cuda_graph_configs(recipe, pp_size: int | None, vp_size: int | None):300    """301    Set optimizations required for full iteration CUDA graphs based on specific conditions.302    """303    if not (304        hasattr(recipe.model, 'config')305        and hasattr(recipe.model.config, 'cuda_graph_scope')306        and recipe.model.config.cuda_graph_scope == 'full_iteration'307    ):308        return recipe309 310    cuda_graph_configs = []311 312    if recipe.trainer.strategy.ddp.check_for_nan_in_grad != False:313        recipe.trainer.strategy.ddp.check_for_nan_in_grad = False314        cuda_graph_configs.append("check_for_nan_in_grad=False")315        logging.warning("For full iteration CUDA graphs, we need to disable check_for_nan_in_grad")316 317    if pp_size and pp_size > 1:318        if recipe.model.config.variable_seq_lengths != False:319            recipe.model.config.variable_seq_lengths = False320            cuda_graph_configs.append("variable_seq_lengths=False")321            logging.warning("For full iteration CUDA graphs, we need to disable variable_seq_lengths")322 323        if recipe.model.config.batch_p2p_sync != False:324            recipe.model.config.batch_p2p_sync = False325            cuda_graph_configs.append("batch_p2p_sync=False")326            logging.warning("For full iteration CUDA graphs, we need to disable batch_p2p_sync")327 328    comm_overlap_callback_idx = get_comm_overlap_callback_idx(recipe.trainer.callbacks)329    if comm_overlap_callback_idx is not None:330        callback = recipe.trainer.callbacks[comm_overlap_callback_idx]331 332        if pp_size and pp_size > 1:333            if callback.batch_p2p_comm != False:334                callback.batch_p2p_comm = False335                cuda_graph_configs.append("batch_p2p_comm=False")336                logging.warning("For full iteration CUDA graphs, disabling batch_p2p_comm would improve memory usage")337 338        if vp_size and vp_size > 1:339            if callback.overlap_param_gather_with_optimizer_step != False:340                callback.overlap_param_gather_with_optimizer_step = False341                cuda_graph_configs.append("overlap_param_gather_with_optimizer_step=False")342                logging.warning(343                    "For full iteration CUDA graphs, we need to disable overlap_param_gather_with_optimizer_step"344                )345    else:346        logging.warning("MegatronCommOverlapCallback not found in recipe.trainer.callbacks")347 348    # Log all applied configurations349    if cuda_graph_configs:350        logging.info(f"Applied full iteration CUDA graph optimizations: {', '.join(cuda_graph_configs)}")351 352    return recipe353 354 355def set_perf_optimization_configs(356    recipe,357    use_mcore_fsdp: bool,358    enable_cuda_graphs: bool,359    task: str,360    tp_size: int | None,361    pp_size: int | None,362    vp_size: int | None,363    compute_dtype: str,364    fp8_recipe: str | None,365    recompute_layers: int,366    activation_offload_layers: int,367    recompute_modules: Optional[List[str]],368    use_fsdp_double_buffer: Optional[bool] = None,369    use_user_buffer_registration: Optional[bool] = None,370    use_sharp: Optional[bool] = None,371    keep_fsdp_fp8_transpose_cache: Optional[bool] = None,372):373    """374    Set performance optimization related configs.375    """376    # enable cross entropy fusion with TE kernel377    recipe.model.config.cross_entropy_fusion_impl = "te"378 379    if use_fsdp_double_buffer:380        assert use_mcore_fsdp == True, "use_fsdp_double_buffer requires use_mcore_fsdp to be True"381 382    if use_mcore_fsdp and enable_cuda_graphs:383        logging.warning("Currently, cuda graphs are not supported with FSDP. Disabling cuda graphs.")384        enable_cuda_graphs = False385    recipe = set_cuda_graph_configs(recipe, enable_cuda_graphs, task)386 387    if enable_cuda_graphs:388        recipe = set_full_iteration_cuda_graph_configs(recipe, pp_size, vp_size)389 390    if use_mcore_fsdp:391        comm_overlap_callback_idx = get_comm_overlap_callback_idx(recipe.trainer.callbacks)392        recipe = set_mcore_fsdp_configs(recipe, comm_overlap_callback_idx, tp_size)393 394    recipe = set_precision_configs(recipe, compute_dtype, fp8_recipe)395 396    recipe = set_recompute_configs(recipe, recompute_layers, activation_offload_layers, recompute_modules)397 398    recipe.trainer.strategy.use_sharp = bool(use_sharp)399 400    is_ddp_obj = hasattr(recipe.trainer.strategy, "ddp") and not isinstance(recipe.trainer.strategy.ddp, str)401    if use_user_buffer_registration and not is_ddp_obj:402        logging.warning("DDP is not configured. Cannot use user buffer registration.")403    if is_ddp_obj:404        # Disable local gradient checker at non-debugging mode405        recipe.trainer.strategy.ddp.check_for_nan_in_grad = False406        recipe.trainer.strategy.ddp.check_for_large_grads = False407        recipe.trainer.strategy.ddp.nccl_ub = bool(use_user_buffer_registration)408        recipe.trainer.strategy.ddp.fsdp_double_buffer = bool(use_fsdp_double_buffer)409        try:410            recipe.trainer.strategy.ddp.keep_fp8_transpose_cache = bool(keep_fsdp_fp8_transpose_cache)411        except AttributeError:412            recipe.trainer.strategy.ddp.keep_fp8_transpose_cache_when_using_custom_fsdp = bool(413                keep_fsdp_fp8_transpose_cache414            )415            logging.warning(416                "Deprecation Notice: `keep_fp8_transpose_cache_when_using_custom_fsdp` "417                "will be deprecated in M-Core 0.14. "418                "Please use `keep_fsdp_fp8_transpose_cache` instead."419            )420 421    return recipe422 423 424def set_primary_perf_configs(425    recipe,426    task: str,427    num_nodes: int,428    num_gpus_per_node: int,429    mbs: int,430    gbs: int,431    max_steps: int,432    tp_size: int,433    pp_size: int,434    cp_size: int,435    vp_size: int,436    ep_size: int,437    etp_size: Optional[int] = None,438    enable_cuda_graphs: bool = False,439    use_mcore_fsdp: bool = False,440    use_fsdp_double_buffer: Optional[bool] = None,441    use_user_buffer_registration: Optional[bool] = None,442    use_sharp: Optional[bool] = None,443    recompute_layers: int = 0,444    activation_offload_layers: int = 0,445    compute_dtype: str = None,446    fp8_recipe: str = None,447    recompute_modules: Optional[List[str]] = None,448    nccl_communicator_config_path: str = None,449    keep_fsdp_fp8_transpose_cache: Optional[bool] = None,450    use_te_op_fuser: Optional[bool] = None,451    use_te_act_func: Optional[bool] = None,452    act_func_fp8_input_store: Optional[bool] = None,453):454    """Set experiment configs we usually tune for performance of all models."""455    # nemo.lightning.Trainer configs456    recipe.trainer.num_nodes = num_nodes457    recipe.trainer.devices = num_gpus_per_node458    recipe.trainer.max_steps = max_steps459 460    recipe.trainer.val_check_interval = max_steps461    recipe.trainer.limit_val_batches = 0462 463    # lightning.pytorch.LightningDataModule configs464    recipe.data.micro_batch_size = mbs465    recipe.data.global_batch_size = gbs466    if recipe.data.__fn_or_cls__ == MockDataModule:467        recipe.data.num_train_samples = max_steps * gbs  # ensure only 1 epoch for whole run468 469    # parallelism configs470    recipe.trainer.strategy.tensor_model_parallel_size = tp_size471    recipe.trainer.strategy.pipeline_model_parallel_size = pp_size472    recipe.trainer.strategy.context_parallel_size = cp_size473    recipe.trainer.strategy.virtual_pipeline_model_parallel_size = None if vp_size == 1 else vp_size474    recipe.trainer.strategy.expert_model_parallel_size = ep_size475    recipe.trainer.strategy.expert_tensor_parallel_size = etp_size476    recipe.trainer.strategy.sequence_parallel = bool(tp_size > 1)477    if nccl_communicator_config_path is not None:478        recipe.trainer.strategy.nccl_communicator_config_path = nccl_communicator_config_path479 480    # callback configs481    comm_overlap_callback_idx = get_comm_overlap_callback_idx(recipe.trainer.callbacks)482    dp_size = (num_nodes * num_gpus_per_node) / (tp_size * pp_size * cp_size)483    if comm_overlap_callback_idx is not None:484        # WARNING: If True, checkpointing (if enabled) might not work485        recipe.trainer.callbacks[comm_overlap_callback_idx].overlap_param_gather_with_optimizer_step = bool(486            dp_size > 1 and pp_size > 1 and vp_size and vp_size > 1487        )488 489    # te op fuser for MLP part490    if use_te_op_fuser:491        assert recipe.model.config.num_moe_experts is None, "use_te_op_fuser is not supported for MOE models"492        if hasattr(recipe.model.config, "use_transformer_engine_op_fuser"):493            recipe.model.config.use_transformer_engine_op_fuser = True494        else:495            logging.warning("use_transformer_engine_op_fuser is not supported for this version of MCORE.")496 497    # te activation function for MLP part498    recipe.model.config.use_te_activation_func = use_te_act_func or False499    assert (500        not act_func_fp8_input_store501    ) or use_te_act_func, "act_func_fp8_input_store requires use_te_act_func to be True"502    recipe.model.config.activation_func_fp8_input_store = act_func_fp8_input_store or False503 504    recipe = set_perf_optimization_configs(505        recipe=recipe,506        use_mcore_fsdp=use_mcore_fsdp,507        enable_cuda_graphs=enable_cuda_graphs,508        task=task,509        tp_size=tp_size,510        pp_size=pp_size,511        vp_size=vp_size,512        compute_dtype=compute_dtype,513        fp8_recipe=fp8_recipe,514        recompute_layers=recompute_layers,515        activation_offload_layers=activation_offload_layers,516        recompute_modules=recompute_modules,517        use_fsdp_double_buffer=use_fsdp_double_buffer,518        use_user_buffer_registration=use_user_buffer_registration,519        use_sharp=use_sharp,520        keep_fsdp_fp8_transpose_cache=keep_fsdp_fp8_transpose_cache,521    )522 523    return recipe524 525 526def set_exp_logging_configs(527    recipe,528    task: str,529    domain: str,530    model_name: str,531    enable_tb: bool,532    enable_wd: bool,533    wandb_prj_name: str,534    wandb_job_name: str,535):536    """Set experiment logging configs."""537    if task == "pre_train" and domain == "llm":538        recipe.trainer.callbacks.append(539            run.Config(540                FLOPsMeasurementCallback,541                model_config=recipe.model.config,542                data_config=recipe.data,543                model_name=model_name,544            )545        )546 547    if not enable_tb:  # tensorboard adds performance overhead.548        recipe.log.tensorboard = None549        recipe.trainer.logger = False550    else:551        # default path is NOT intuitive- `<log_dir>/code/nemo_experiments/tb_logs/default/<tfevents_file>`552        recipe.log.log_dir = "/nemo_run/lightning_logs"  # saves file at- `<log_dir>/lightning_logs/tb_logs553    if enable_wd:554        from nemo.collections.llm.recipes.log.default import wandb_logger555 556        recipe.log.wandb = wandb_logger(project=wandb_prj_name, name=wandb_job_name)557 558    # Misc. for overall faster experiment runtime559    recipe.log.ckpt = None560 561    # disable checkpointing if no ModelCheckpoint callback is found562    callbacks = recipe.trainer.callbacks563    checkpoint_callback_idx = None564    if callbacks:  # default is None in lightning565        for idx, callback in enumerate(callbacks):566            if callback.__fn_or_cls__ == ModelCheckpoint:567                checkpoint_callback_idx = idx568                break569    recipe.trainer.enable_checkpointing = checkpoint_callback_idx is not None570    recipe.trainer.log_every_n_steps = 1571 572    return recipe573 574 575def args_sanity_check(args: dict) -> None:576    """577    Check the sanity of argument settings578    """579    if args.wandb:580        assert args.wandb_key is not None, "wandb logger needs \"wandb_key\""581        assert args.wandb_prj_name is not None, "wandb logger needs \"wandb_prj_name\""582        assert args.wandb_job_name is not None, "wandb logger needs \"wandb_job_name\""583 584 585def build_perf_env_plugin(args, pp_size: int | None = None, user_buffer_registration: Optional[bool] = None):586    """587    Create a PerfEnvPlugin with consistent defaults across scripts.588 589    - enable_vboost only when gpu is h100590    - set nccl_pp_comm_chunksize when pipeline parallelism is used591    - set gpu_sm100_or_newer when gpu is in ['b200', 'gb200']592 593    Args:594        args: Parsed CLI args that include `gpu`.595        pp_size: Pipeline parallel size to decide comm chunk size.596        user_buffer_registration: Optional flag to enable user buffer registration.597    """598    from nemo.lightning.run.plugins import PerfEnvPlugin599 600    gpu_str = getattr(args, "gpu", "").lower()601    enable_vboost = args.enable_vboost602    gpu_sm100_or_newer = gpu_str in ["b200", "gb200"]603    nccl_pp_comm_chunksize = 2097152 if (pp_size is not None and pp_size > 1) else None604    user_buf = bool(user_buffer_registration) if user_buffer_registration is not None else False605 606    return PerfEnvPlugin(607        enable_vboost=enable_vboost,608        nccl_pp_comm_chunksize=nccl_pp_comm_chunksize,609        gpu_sm100_or_newer=gpu_sm100_or_newer,610        user_buffer_registration=user_buf,611    )612