Repository: NVIDIA/NeMo
Pull request #12947: Add dsv3 pretrain script, support flops calculation
Author: login: guyueh1
Labels: name: r2.3.0
description: Pick this label for auto-cherrypicking into v2.3.0
name: RC4
Created: 2025-04-09T17:49:49Z | Merged: 2025-04-22T17:43:54Z

> [!IMPORTANT]  
> The `Update branch` button must only be pressed in very rare occassions.
> An outdated branch is never blocking the merge of a PR.
> Please reach out to the automation team before pressing that button.

# What does this PR do ?

Add deepseek v3 pretrain performance testing script.

**Collection**: [Note which collection this PR will affect]

# Changelog

- Add specific line by line info of high level changes in this PR.

# Usage

- You can potentially add a usage example below

```python
# Add a code snippet demonstrating how to use this 
```

# GitHub Actions CI

The Jenkins CI system has been replaced by GitHub Actions self-hosted runners.

The GitHub Actions CI will run automatically when the "Run CICD" label is added to the PR.
To re-run CI remove and add the label again.
To run CI on an untrusted fork, a NeMo user with write access must first click "Approve and run".

# Before your PR is "Ready for review"

**Pre checks**:

- [ ] Make sure you read and followed [Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md)
- [ ] Did you write any new necessary tests?
- [ ] Did you add or update any necessary documentation?
- [ ] Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  - [ ] Reviewer: Does the PR have correct import guards for all optional libraries?
  
**PR Type**:

- [ ] New Feature
- [ ] Bugfix
- [ ] Documentation

If you haven't finished some of the above items you can still open "Draft" PR.

## Who can review?

Anyone in the community is free to review the PR once the checks have passed.
[Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md) contains specific people who can review PRs to various areas.

# Additional Information

- Related to # (issue)



Diff:
diff --git a/nemo/lightning/pytorch/callbacks/flops_callback.py b/nemo/lightning/pytorch/callbacks/flops_callback.py
index 32e4ee25c174..c85791558c74 100644
--- a/nemo/lightning/pytorch/callbacks/flops_callback.py
+++ b/nemo/lightning/pytorch/callbacks/flops_callback.py
@@ -35,6 +35,7 @@
     "mixtral": flops_formulas.mixtral,
     "bert": flops_formulas.bert,
     "hyena": hyena,
+    "deepseekv3": flops_formulas.deepseekv3,
     "transformer": flops_formulas.transformer,
 }
 
@@ -81,18 +82,33 @@ def __init__(
         if query_groups is None:
             query_groups = attention_heads
 
-        self.flops_config = flops_formulas.FLOPSConfig(
-            gbs=gbs,
-            enc_seq_len=enc_seq_len,
-            hs=hs,
-            layers=layers,
-            ffn_hs=ffn_hs,
-            attention_heads=attention_heads,
-            moe_router_topk=moe_router_topk,
-            query_groups=query_groups,
-            vocab_size=vocab_size,
-            model_pattern=model_pattern,
-        )
+        config_kwargs = {
+            "gbs": gbs,
+            "enc_seq_len": enc_seq_len,
+            "hs": hs,
+            "layers": layers,
+            "ffn_hs": ffn_hs,
+            "attention_heads": attention_heads,
+            "moe_router_topk": moe_router_topk,
+            "query_groups": query_groups,
+            "vocab_size": vocab_size,
+            "model_pattern": model_pattern,
+        }
+
+        from megatron.core.transformer.transformer_config import MLATransformerConfig
+
+        if isinstance(self.model_cfg, MLATransformerConfig):
+            config_kwargs["qk_head_dim"] = self.model_cfg.qk_head_dim
+            config_kwargs["qk_pos_emb_head_dim"] = self.model_cfg.qk_pos_emb_head_dim
+            config_kwargs["v_head_dim"] = self.model_cfg.v_head_dim
+            config_kwargs["q_lora_rank"] = self.model_cfg.q_lora_rank
+            config_kwargs["kv_lora_rank"] = self.model_cfg.kv_lora_rank
+        config_kwargs["moe_layer_freq"] = self.model_cfg.moe_layer_freq
+        config_kwargs["moe_shared_expert_intermediate_size"] = self.model_cfg.moe_shared_expert_intermediate_size
+        config_kwargs["moe_ffn_hidden_size"] = self.model_cfg.moe_ffn_hidden_size
+        config_kwargs["mtp_num_layers"] = self.model_cfg.mtp_num_layers
+
+        self.flops_config = flops_formulas.FLOPSConfig(**config_kwargs)
 
         self.model = self.model.lower() if self.model is not None else self.model
 
diff --git a/nemo/utils/flops_formulas.py b/nemo/utils/flops_formulas.py
index 95035f45a8d6..f180d2ea1ee1 100644
--- a/nemo/utils/flops_formulas.py
+++ b/nemo/utils/flops_formulas.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from dataclasses import dataclass
-from typing import Optional
+from typing import List, Optional, Union
 
 from nemo.collections.common.parts.perf_metrics_utils import LLM_VOCAB_SIZE_MAP
 
@@ -42,6 +42,15 @@ class FLOPSConfig:
     vocab_size: Optional[int] = None
     model_channels: Optional[int] = None
     vec_in_dim: Optional[int] = None
+    q_lora_rank: Optional[int] = None
+    kv_lora_rank: Optional[int] = None
+    qk_head_dim: Optional[int] = None
+    qk_pos_emb_head_dim: Optional[int] = None
+    v_head_dim: Optional[int] = None
+    moe_layer_freq: Union[int, List[int]] = None
+    moe_shared_expert_intermediate_size: Optional[int] = None
+    moe_ffn_hidden_size: Optional[int] = None
+    mtp_num_layers: Optional[int] = None
     causal_self_attn: Optional[bool] = None
 
 
@@ -166,17 +175,17 @@ def transformer(config: FLOPSConfig):
 
     if vocab_size is None:
         raise ValueError("vocab_size is required for transformer FLOPs calculation")
-    
+
     # Handle optional parameters with reasonable defaults
     query_groups = config.query_groups if config.query_groups is not None else num_attention_heads
     causal_self_attn = config.causal_self_attn if config.causal_self_attn is not None else False
     moe_router_topk = config.moe_router_topk if config.moe_router_topk is not None else 0
     kv_channels = hidden_size // num_attention_heads  # Standard dimension per head
-    
+
     # Calculate query projection size and ratio
     query_projection_size = kv_channels * num_attention_heads
     query_projection_to_hidden_size_ratio = query_projection_size / hidden_size
-    
+
     # MoE parameters - simplified for NeMo config
     # In this implementation, we assume all layers are dense if num_experts is None
     if moe_router_topk == 0:
@@ -189,11 +198,11 @@ def transformer(config: FLOPSConfig):
         num_moe_layers = num_layers // 2  # Simplified assumption
         num_dense_layers = num_layers - num_moe_layers
         num_experts_routed_to = moe_router_topk
-    
+
     # Handle SwiGLU vs standard GELU/ReLU
     # Default to standard activation (no SwiGLU)
     gated_linear_multiplier = 1
-    
+
     # Define the expansion factor as described in the paper
     # 3x: Each GEMM needs forward pass, backward wgrad, and backward dgrad
     # 2x: GEMMs are stacked twice in standard Transformer architectures
@@ -202,22 +211,18 @@ def transformer(config: FLOPSConfig):
     # Attention
     if not causal_self_attn:
         attention_component = (
-                    (
-                        1
-                        + (query_groups / num_attention_heads)
-                        # Only half of the attention matrix is non-zero and needs to be multiplied with V
-                        + (seq_length / hidden_size) # If causal self attn -> divide by 2.
-                    ) * query_projection_to_hidden_size_ratio
-                ) 
+            1
+            + (query_groups / num_attention_heads)
+            # Only half of the attention matrix is non-zero and needs to be multiplied with V
+            + (seq_length / hidden_size)  # If causal self attn -> divide by 2.
+        ) * query_projection_to_hidden_size_ratio
     else:
         attention_component = (
-                (
-                    1
-                    + (query_groups / num_attention_heads)
-                    # Only half of the attention matrix is non-zero and needs to be multiplied with V
-                    + (seq_length / hidden_size / 2) # If causal self attn -> divide by 2.
-                ) * query_projection_to_hidden_size_ratio
-            ) 
+            1
+            + (query_groups / num_attention_heads)
+            # Only half of the attention matrix is non-zero and needs to be multiplied with V
+            + (seq_length / hidden_size / 2)  # If causal self attn -> divide by 2.
+        ) * query_projection_to_hidden_size_ratio
 
     # Calculate total FLOPs
     total_flops = (
@@ -228,29 +233,32 @@ def transformer(config: FLOPSConfig):
         * hidden_size
         * hidden_size
         * (
-            
             attention_component
             # MLP component
             + (
                 (
                     # Dense layers
-                    (ffn_hidden_size * num_dense_layers) +
+                    (ffn_hidden_size * num_dense_layers)
+                    +
                     # MoE layers
                     (
                         (
                             # Routed experts
-                            ffn_hidden_size * num_experts_routed_to
+                            ffn_hidden_size
+                            * num_experts_routed_to
                             # Note: Shared experts are not implemented in this version
                         )
                         * num_moe_layers
                     )
-                ) * gated_linear_multiplier / (num_layers * hidden_size)
+                )
+                * gated_linear_multiplier
+                / (num_layers * hidden_size)
             )
             # Logit component
             + (vocab_size / (2 * num_layers * hidden_size))
         )
     )
-    
+
     return total_flops
 
 
@@ -325,3 +333,58 @@ def flux(config: FLOPSConfig):
     )
 
     return joint_layer_flops + single_layer_flops + other_flops
+
+
+def deepseekv3(config: FLOPSConfig):
+    """Model FLOPs for DeepSeek V3"""
+
+    # self-attention flops
+    bmm1_flops = (
+        0.5 * (config.qk_head_dim + config.qk_pos_emb_head_dim) * config.attention_heads * (config.enc_seq_len**2)
+    )
+    bmm2_flops = 0.5 * config.v_head_dim * config.attention_heads * (config.enc_seq_len**2)
+    per_input_attention_flops = 6 * (bmm1_flops + bmm2_flops) * config.layers
+    if config.mtp_num_layers is not None:
+        per_input_attention_flops += 6 * (bmm1_flops + bmm2_flops) * config.mtp_num_layers
+
+    # linear layer flops
+    per_layer_mla_params = config.hs * config.q_lora_rank + config.q_lora_rank * (
+        (config.qk_head_dim + config.qk_pos_emb_head_dim) * config.attention_heads
+    )  # Q
+    per_layer_mla_params += config.hs * config.qk_pos_emb_head_dim  # K^R
+    per_layer_mla_params += config.hs * config.kv_lora_rank + config.kv_lora_rank * (
+        (config.qk_head_dim + config.v_head_dim) * config.attention_heads
+    )  # K^C and V^C
+    per_layer_mla_params += config.v_head_dim * config.attention_heads * config.hs  # Proj
+    mla_params = per_layer_mla_params * config.layers
+    if config.mtp_num_layers is not None:
+        mla_params += per_layer_mla_params * config.mtp_num_layers
+
+    dense_layer_ffn_params = config.hs * config.ffn_hs * 3  # gated linear unit
+    per_shared_expert_params = config.hs * config.moe_shared_expert_intermediate_size * 3
+    per_selected_expert_params = config.hs * config.moe_ffn_hidden_size * 3
+    ffn_params = 0
+
+    if isinstance(config.moe_layer_freq, int):
+        moe_layer_pattern = [1 if (i % config.moe_layer_freq == 0) else 0 for i in range(config.layers)]
+    else:
+        moe_layer_pattern = config.moe_layer_freq
+    for i in moe_layer_pattern:
+        if i == 0:
+            ffn_params += dense_layer_ffn_params
+        else:
+            ffn_params += per_shared_expert_params + (per_selected_expert_params * config.moe_router_topk)
+    if config.mtp_num_layers is not None:
+        for i in range(config.mtp_num_layers):
+            ffn_params += per_shared_expert_params + (per_selected_expert_params * config.moe_router_topk)
+    per_input_params = mla_params + ffn_params
+    per_input_linear_flops = 6 * per_input_params * config.enc_seq_len
+
+    # vocab flops
+    per_input_vocab_flops = 6 * config.vocab_size * config.hs * config.enc_seq_len
+    if config.mtp_num_layers is not None:
+        for i in range(config.mtp_num_layers):
+            per_input_vocab_flops += 6 * config.vocab_size * config.hs * config.enc_seq_len
+            per_input_vocab_flops += 6 * config.hs * 2 * config.hs * config.enc_seq_len
+
+    return (per_input_attention_flops + per_input_linear_flops + per_input_vocab_flops) * config.gbs
diff --git a/scripts/performance/argument_parser.py b/scripts/performance/argument_parser.py
index 0ac34b39b8be..f07e1e93afe2 100644
--- a/scripts/performance/argument_parser.py
+++ b/scripts/performance/argument_parser.py
@@ -298,6 +298,15 @@ def parse_cli_args():
     def list_of_strings(arg):
         return arg.split(',')
 
+    parser.add_argument(
+        "-rm",
+        "--recompute_modules",
+        type=list_of_strings,
+        help="Comma-separated string of modules to recompute. Defaults to None",
+        required=False,
+        default=None,
+    )
+
     parser.add_argument(
         "-cm",
         "--custom_mounts",
diff --git a/scripts/performance/diffusion/pretrain_flux_12b.py b/scripts/performance/diffusion/pretrain_flux_12b.py
index 0d5ac4497d69..9fe2931c62c9 100644
--- a/scripts/performance/diffusion/pretrain_flux_12b.py
+++ b/scripts/performance/diffusion/pretrain_flux_12b.py
@@ -84,7 +84,7 @@ def override_recipe_configs(
     args_sanity_check(args)
 
     kwargs = get_user_configs(args.gpu.lower(), "pre_train", "flux", "12b", args)
-    num_nodes, mbs, gbs, tp_size, pp_size, cp_size, vp_size, ep_size, _, enable_cuda_graphs = kwargs
+    num_nodes, mbs, gbs, tp_size, pp_size, cp_size, vp_size, ep_size, _, enable_cuda_graphs = kwargs[:10]
 
     recipe = override_recipe_configs(
         args, num_nodes, mbs, gbs, tp_size, pp_size, cp_size, vp_size, ep_size, enable_cuda_graphs
diff --git a/scripts/performance/llm/finetune_deepseek_v3.py b/scripts/performance/llm/finetune_deepseek_v3.py
index fe6bdcfe0f9a..a9df11fbd8c1 100644
--- a/scripts/performance/llm/finetune_deepseek_v3.py
+++ b/scripts/performance/llm/finetune_deepseek_v3.py
@@ -134,7 +134,7 @@ def override_recipe_configs(
     args_sanity_check(args)
 
     kwargs = get_user_configs(args.gpu.lower(), args.finetuning, "deepseek", "v3", args)
-    num_nodes, mbs, gbs, tp_size, pp_size, cp_size, vp_size, ep_size, _, enable_cuda_graphs = kwargs
+    num_nodes, mbs, gbs, tp_size, pp_size, cp_size, vp_size, ep_size, _, enable_cuda_graphs = kwargs[:10]
 
     recipe = override_recipe_configs(
         args, num_nodes, mbs, gbs, tp_size, pp_size, cp_size, vp_size, ep_size, enable_cuda_graphs
@@ -158,7 +158,13 @@ def override_recipe_configs(
         wandb_key=args.wandb_key,
     )
 
-    plugins = [PerfEnvPlugin(enable_vboost=True, nccl_pp_comm_chunksize=2097152 if pp_size > 1 else None)]
+    plugins = [
+        PerfEnvPlugin(
+            enable_vboost=True,
+            nccl_pp_comm_chunksize=2097152 if pp_size > 1 else None,
+            gpu_sm100_or_newer=(args.gpu.lower() in ['b200', 'gb200']),
+        )
+    ]
     if args.enable_nsys:
         plugins.append(NsysPlugin(start_step=10, end_step=12, gen_shape=True))
 
diff --git a/scripts/performance/llm/finetune_llama31_405b.py b/scripts/performance/llm/finetune_llama31_405b.py
index 82249133bd5a..aff7e9c81ae4 100644
--- a/scripts/performance/llm/finetune_llama31_405b.py
+++ b/scripts/performance/llm/finetune_llama31_405b.py
@@ -153,7 +153,7 @@ def override_recipe_configs(
         use_mcore_fsdp,
         recompute_layers,
         activation_offload_layers,
-    ) = kwargs
+    ) = kwargs[:13]
 
     recipe = override_recipe_configs(
         args,
diff --git a/scripts/performance/llm/finetune_llama3_70b.py b/scripts/performance/llm/finetune_llama3_70b.py
index 05f980d876f7..68f9b1ee3b80 100644
--- a/scripts/performance/llm/finetune_llama3_70b.py
+++ b/scripts/performance/llm/finetune_llama3_70b.py
@@ -160,7 +160,7 @@ def override_recipe_configs(
         use_mcore_fsdp,
         recompute_layers,
         activation_offload_layers,
-    ) = kwargs
+    ) = kwargs[:13]
 
     recipe = override_recipe_configs(
         args,
diff --git a/scripts/performance/llm/finetune_llama3_8b.py b/scripts/performance/llm/finetune_llama3_8b.py
index b5fa71770b9e..28a83851ef92 100644
--- a/scripts/performance/llm/finetune_llama3_8b.py
+++ b/scripts/performance/llm/finetune_llama3_8b.py
@@ -111,7 +111,7 @@ def override_recipe_configs(
     args_sanity_check(args)
 
     kwargs = get_user_configs(args.gpu.lower(), args.finetuning, "llama3", "8b", args)
-    num_nodes, mbs, gbs, tp_size, pp_size, cp_size, vp_size, ep_size, _, enable_cuda_graphs, _, _, _ = kwargs
+    num_nodes, mbs, gbs, tp_size, pp_size, cp_size, vp_size, ep_size, _, enable_cuda_graphs = kwargs[:10]
 
     recipe = override_recipe_configs(
         args, num_nodes, mbs, gbs, tp_size, pp_size, cp_size, vp_size, ep_size, enable_cuda_graphs
diff --git a/scripts/performance/llm/pretrain_automodel_llama3_8b.py b/scripts/performance/llm/pretrain_automodel_llama3_8b.py
index 43dbff3b4479..f985bf625b8b 100644
--- a/scripts/performance/llm/pretrain_automodel_llama3_8b.py
+++ b/scripts/performance/llm/pretrain_automodel_llama3_8b.py
@@ -74,7 +74,7 @@ def override_recipe_configs(
     args_sanity_check(args)
 
     kwargs = get_user_configs(args.gpu.lower(), "pre_train", "llama3", "8b", args)
-    num_nodes, mbs, gbs, tp_size, pp_size, cp_size, vp_size, ep_size, _, _, _, _, _ = kwargs
+    num_nodes, mbs, gbs, tp_size, pp_size, cp_size, vp_size, ep_size = kwargs[:8]
 
     recipe = override_recipe_configs(
         args,
diff --git a/scripts/performance/llm/pretrain_deepseek_v3.py b/scripts/performance/llm/pretrain_deepseek_v3.py
new file mode 100644
index 000000000000..5385c5d5f884
--- /dev/null
+++ b/scripts/performance/llm/pretrain_deepseek_v3.py
@@ -0,0 +1,215 @@
+# 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.
+
+from os.path import basename, splitext
+from typing import List, Optional
+
+import nemo_run as run
+
+from nemo.collections.llm.recipes.deepseek_v3 import pretrain_recipe
+from nemo.lightning.pytorch.callbacks.garbage_collection import GarbageCollectionCallback
+from nemo.lightning.pytorch.callbacks.megatron_comm_overlap import MegatronCommOverlapCallback
+from nemo.lightning.pytorch.callbacks.moe_token_drop import MegatronTokenDropCallback
+from nemo.lightning.run.plugins import NsysPlugin, PerfEnvPlugin
+
+from ..argument_parser import parse_cli_args
+from ..utils import (
+    args_sanity_check,
+    get_user_configs,
+    hf_tokenizer,
+    set_exp_logging_configs,
+    set_primary_perf_configs,
+    slurm_executor,
+)
+
+HF_MODEL_URI = "deepseek-ai/DeepSeek-V3-Base"
+
+
+# Use token drop callback
+USE_TOKEN_DROP = True
+
+
+def override_recipe_configs(
+    args: str,
+    num_nodes: int,
+    mbs: int,
+    gbs: int,
+    tp_size: int,
+    pp_size: int,
+    cp_size: int,
+    vp_size: int,
+    ep_size: int,
+    etp_size: int,
+    enable_cuda_graphs: bool,
+    use_mcore_fsdp: bool,
+    recompute_layers: int,
+    activation_offload_layers: int,
+    recompute_modules: Optional[List[str]] = None,
+):
+    """
+    DeepSeek V3 pre-train recipe aimed at achieving best possible performance.
+    """
+    recipe = pretrain_recipe()
+    recipe.model.config.moe_permute_fusion = True
+    recipe.model.config.recompute_granularity = "selective"
+    recipe.model.config.recompute_num_layers = None
+    recipe.model.config.recompute_method = None
+    recipe.model.config.cross_entropy_fusion_impl = "te"
+    recipe.model.config.moe_router_dtype = 'fp32'
+
+    assert pp_size is None or pp_size == 1 or (64 % pp_size == 0), "pp_size must be 1 or a factor of 64"
+    num_layers_in_middle_pipeline_stages = 64 // pp_size
+
+    recipe.trainer.strategy.num_layers_in_first_pipeline_stage = num_layers_in_middle_pipeline_stages - 1
+    recipe.trainer.strategy.num_layers_in_last_pipeline_stage = num_layers_in_middle_pipeline_stages - 2
+
+    callbacks = []
+    if USE_TOKEN_DROP:
+        callbacks.append(run.Config(MegatronTokenDropCallback))
+    garbage_collection_callback = run.Config(
+        GarbageCollectionCallback,
+        gc_interval_train=60,
+        gc_interval_val=60,
+    )
+    comm_overlap_callback = run.Config(
+        MegatronCommOverlapCallback,
+        tp_comm_overlap=False,
+    )
+    callbacks.extend([garbage_collection_callback, comm_overlap_callback])
+    recipe.trainer.callbacks.extend(callbacks)
+
+    recipe = set_primary_perf_configs(
+        recipe,
+        "pre_train",
+        num_nodes,
+        args.gpus_per_node,
+        mbs,
+        gbs,
+        args.max_steps,
+        tp_size,
+        pp_size,
+        cp_size,
+        vp_size,
+        ep_size,
+        etp_size,
+        enable_cuda_graphs=enable_cuda_graphs,
+        use_mcore_fsdp=use_mcore_fsdp,
+        recompute_layers=recompute_layers,
+        activation_offload_layers=activation_offload_layers,
+        compute_dtype=args.compute_dtype,
+        fp8_recipe=args.fp8_recipe,
+        recompute_modules=recompute_modules,
+    )
+    recipe = set_exp_logging_configs(
+        recipe,
+        "pre_train",
+        "llm",
+        "deepseekv3",
+        args.tensorboard,
+        args.wandb,
+        args.wandb_prj_name,
+        args.wandb_job_name,
+    )
+
+    # data module configs
+    recipe.data.tokenizer = hf_tokenizer(HF_MODEL_URI)
+    recipe.model.tokenizer = recipe.data.tokenizer
+
+    # compute dtype configs
+    if args.compute_dtype.lower() == "fp8":
+        raise ValueError("Deepseek FP8 recipe requires subchannel scaling which will be supported soon.")
+
+    return recipe
+
+
+if __name__ == "__main__":
+    args = parse_cli_args().parse_args()
+    args_sanity_check(args)
+
+    kwargs = get_user_configs(args.gpu.lower(), "pre_train", "deepseek", "v3", args)
+    (
+        num_nodes,
+        mbs,
+        gbs,
+        tp_size,
+        pp_size,
+        cp_size,
+        vp_size,
+        ep_size,
+        etp_size,
+        enable_cuda_graphs,
+        use_mcore_fsdp,
+        recompute_layers,
+        activation_offload_layers,
+        recompute_modules,
+    ) = kwargs
+
+    recipe = override_recipe_configs(
+        args,
+        num_nodes,
+        mbs,
+        gbs,
+        tp_size,
+        pp_size,
+        cp_size,
+        vp_size,
+        ep_size,
+        etp_size,
+        enable_cuda_graphs,
+        use_mcore_fsdp,
+        recompute_layers,
+        activation_offload_layers,
+        recompute_modules,
+    )
+
+    exp_config = f"{num_nodes}nodes_tp{tp_size}_pp{pp_size}_cp{cp_size}_vp{vp_size}_ep{ep_size}_{mbs}mbs_{gbs}gbs"
+    exp_name = f"{splitext(basename(__file__))[0]}_{args.compute_dtype}_{exp_config}"
+
+    executor = slurm_executor(
+        args.account,
+        args.partition,
+        args.log_dir,
+        num_nodes,
+        args.gpus_per_node,
+        args.time_limit,
+        args.container_image,
+        custom_mounts=args.custom_mounts,
+        custom_env_vars={},
+        hf_token=args.hf_token,
+        nemo_home=args.nemo_home,
+        wandb_key=args.wandb_key,
+    )
+
+    plugins = [
+        PerfEnvPlugin(
+            enable_vboost=True,
+            nccl_pp_comm_chunksize=2097152 if pp_size > 1 else None,
+            gpu_sm100_or_newer=(args.gpu.lower() in ['b200', 'gb200']),
+        )
+    ]
+    if args.enable_nsys:
+        plugins.append(NsysPlugin(start_step=5, end_step=6))
+
+    with run.Experiment(exp_name) as exp:
+        exp.add(
+            recipe,
+            executor=executor,
+            name=exp_name,
+            plugins=plugins,
+        )
+
+        if not args.dryrun:
+            exp.run(sequential=True, detach=True)
+        else:
+            exp.dryrun()
diff --git a/scripts/performance/llm/pretrain_gpt3_175b.py b/scripts/performance/llm/pretrain_gpt3_175b.py
index 7aaa9ea735f4..f0faf3470cc2 100644
--- a/scripts/performance/llm/pretrain_gpt3_175b.py
+++ b/scripts/performance/llm/pretrain_gpt3_175b.py
@@ -135,7 +135,7 @@ def override_recipe_configs(
         use_mcore_fsdp,
         recompute_layers,
         activation_offload_layers,
-    ) = kwargs
+    ) = kwargs[:13]
 
     recipe = override_recipe_configs(
         args,
diff --git a/scripts/performance/llm/pretrain_llama31_405b.py b/scripts/performance/llm/pretrain_llama31_405b.py
index 6f6f280518f0..8e458163a33f 100644
--- a/scripts/performance/llm/pretrain_llama31_405b.py
+++ b/scripts/performance/llm/pretrain_llama31_405b.py
@@ -134,7 +134,7 @@ def override_recipe_configs(
         use_mcore_fsdp,
         recompute_layers,
         activation_offload_layers,
-    ) = kwargs
+    ) = kwargs[:13]
 
     recipe = override_recipe_configs(
         args,
diff --git a/scripts/performance/llm/pretrain_llama3_70b.py b/scripts/performance/llm/pretrain_llama3_70b.py
index aa8116cf0e61..72f4b9be0d45 100644
--- a/scripts/performance/llm/pretrain_llama3_70b.py
+++ b/scripts/performance/llm/pretrain_llama3_70b.py
@@ -134,7 +134,7 @@ def override_recipe_configs(
         use_mcore_fsdp,
         recompute_layers,
         activation_offload_layers,
-    ) = kwargs
+    ) = kwargs[:13]
 
     recipe = override_recipe_configs(
         args,
diff --git a/scripts/performance/llm/pretrain_llama3_8b.py b/scripts/performance/llm/pretrain_llama3_8b.py
index c85f81744525..781df30da53b 100644
--- a/scripts/performance/llm/pretrain_llama3_8b.py
+++ b/scripts/performance/llm/pretrain_llama3_8b.py
@@ -81,7 +81,7 @@ def override_recipe_configs(
     args_sanity_check(args)
 
     kwargs = get_user_configs(args.gpu.lower(), "pre_train", "llama3", "8b", args)
-    num_nodes, mbs, gbs, tp_size, pp_size, cp_size, vp_size, ep_size, _, enable_cuda_graphs, _, _, _ = kwargs
+    num_nodes, mbs, gbs, tp_size, pp_size, cp_size, vp_size, ep_size, _, enable_cuda_graphs = kwargs[:10]
 
     recipe = override_recipe_configs(
         args, num_nodes, mbs, gbs, tp_size, pp_size, cp_size, vp_size, ep_size, enable_cuda_graphs
diff --git a/scripts/performance/llm/pretrain_mixtral_8x22b.py b/scripts/performance/llm/pretrain_mixtral_8x22b.py
index bdc7773120cb..54335a8b5773 100644
--- a/scripts/performance/llm/pretrain_mixtral_8x22b.py
+++ b/scripts/performance/llm/pretrain_mixtral_8x22b.py
@@ -108,7 +108,7 @@ def override_recipe_configs(
         use_mcore_fsdp,
         recompute_layers,
         activation_offload_layers,
-    ) = kwargs
+    ) = kwargs[:13]
 
     recipe = override_recipe_configs(
         args,
diff --git a/scripts/performance/llm/pretrain_mixtral_8x7b.py b/scripts/performance/llm/pretrain_mixtral_8x7b.py
index f4cb8c89bee4..4a3ffe6c81d9 100644
--- a/scripts/performance/llm/pretrain_mixtral_8x7b.py
+++ b/scripts/performance/llm/pretrain_mixtral_8x7b.py
@@ -83,7 +83,7 @@ def override_recipe_configs(
     args_sanity_check(args)
 
     kwargs = get_user_configs(args.gpu.lower(), "pre_train", "mixtral", "8x7b", args)
-    num_nodes, mbs, gbs, tp_size, pp_size, cp_size, vp_size, ep_size, etp_size, enable_cuda_graphs, _, _, _ = kwargs
+    num_nodes, mbs, gbs, tp_size, pp_size, cp_size, vp_size, ep_size, etp_size, enable_cuda_graphs = kwargs[:10]
 
     recipe = override_recipe_configs(
         args, num_nodes, mbs, gbs, tp_size, pp_size, cp_size, vp_size, ep_size, etp_size, enable_cuda_graphs
diff --git a/scripts/performance/llm/pretrain_nemotron3_22b.py b/scripts/performance/llm/pretrain_nemotron3_22b.py
index b8d2abc36512..32f0669edc04 100644
--- a/scripts/performance/llm/pretrain_nemotron3_22b.py
+++ b/scripts/performance/llm/pretrain_nemotron3_22b.py
@@ -103,7 +103,7 @@ def override_recipe_configs(
         use_mcore_fsdp,
         recompute_layers,
         activation_offload_layers,
-    ) = kwargs
+    ) = kwargs[:13]
 
     recipe = override_recipe_configs(
         args,
diff --git a/scripts/performance/llm/pretrain_nemotron3_8b.py b/scripts/performance/llm/pretrain_nemotron3_8b.py
index 6e65e94c0319..4ba1fdd7828f 100644
--- a/scripts/performance/llm/pretrain_nemotron3_8b.py
+++ b/scripts/performance/llm/pretrain_nemotron3_8b.py
@@ -82,7 +82,7 @@ def override_recipe_configs(
     args_sanity_check(args)
 
     kwargs = get_user_configs(args.gpu.lower(), "pre_train", "nemotron3", "8b", args)
-    num_nodes, mbs, gbs, tp_size, pp_size, cp_size, vp_size, ep_size, _, enable_cuda_graphs, _, _, _ = kwargs
+    num_nodes, mbs, gbs, tp_size, pp_size, cp_size, vp_size, ep_size, _, enable_cuda_graphs = kwargs[:10]
 
     recipe = override_recipe_configs(
         args, num_nodes, mbs, gbs, tp_size, pp_size, cp_size, vp_size, ep_size, enable_cuda_graphs
diff --git a/scripts/performance/llm/pretrain_nemotron4_15b.py b/scripts/performance/llm/pretrain_nemotron4_15b.py
index a359205ba884..d27926d71035 100644
--- a/scripts/performance/llm/pretrain_nemotron4_15b.py
+++ b/scripts/performance/llm/pretrain_nemotron4_15b.py
@@ -97,7 +97,7 @@ def override_recipe_configs(
     args_sanity_check(args)
 
     kwargs = get_user_configs(args.gpu.lower(), "pre_train", "nemotron4", "15b", args)
-    num_nodes, mbs, gbs, tp_size, pp_size, cp_size, vp_size, ep_size, _, enable_cuda_graphs, _, _, _ = kwargs
+    num_nodes, mbs, gbs, tp_size, pp_size, cp_size, vp_size, ep_size, _, enable_cuda_graphs = kwargs[:10]
 
     recipe = override_recipe_configs(
         args, num_nodes, mbs, gbs, tp_size, pp_size, cp_size, vp_size, ep_size, enable_cuda_graphs
diff --git a/scripts/performance/llm/pretrain_nemotron4_340b.py b/scripts/performance/llm/pretrain_nemotron4_340b.py
index f3cdb689286f..0f422c5a1ac8 100644
--- a/scripts/performance/llm/pretrain_nemotron4_340b.py
+++ b/scripts/performance/llm/pretrain_nemotron4_340b.py
@@ -125,7 +125,7 @@ def override_recipe_configs(
         use_mcore_fsdp,
         recompute_layers,
         activation_offload_layers,
-    ) = kwargs
+    ) = kwargs[:13]
 
     recipe = override_recipe_configs(
         args,
diff --git a/scripts/performance/recommended_model_configs/model_configs_b200.csv b/scripts/performance/recommended_model_configs/model_configs_b200.csv
index 854800b94747..e507687394f9 100644
--- a/scripts/performance/recommended_model_configs/model_configs_b200.csv
+++ b/scripts/performance/recommended_model_configs/model_configs_b200.csv
@@ -1,27 +1,30 @@
-task,model,size,system,dtype,num_gpus,seq_len,tp_size,pp_size,cp_size,ep_size,vp_size,mbs,gbs,etp_size,cuda_graphs,use_mcore_fsdp,recompute_layers,activation_offload_layers
-lora,llama3,70b,b200,fp8,8,4096,1,4,1,1,20,1,32,,0,0,0,0
-lora,llama3,70b,b200,bf16,8,4096,1,4,1,1,20,1,32,,0,0,0,0
-lora,llama3,8b,b200,bf16,8,16384,1,1,1,1,1,1,8,,1,0,0,0
-lora,llama3,8b,b200,fp8,8,16384,1,1,1,1,1,1,8,,0,0,0,0
-lora,llama31,405b,b200,fp8,32,2048,4,4,1,1,4,1,32,,0,0,0,0
-lora,llama31,405b,b200,bf16,32,4096,4,4,1,1,4,1,32,,0,0,0,0
-pre_train,gpt3,175b,b200,bf16,128,2048,4,4,1,1,12,2,256,,0,0,0,0
-pre_train,gpt3,175b,b200,fp8,128,2048,4,4,1,1,12,2,256,,0,0,0,0
-pre_train,gpt3,175b,b200,fp8,512,2048,4,4,1,1,6,2,2048,,0,0,0,0
-pre_train,llama3,8b,b200,bf16,8,8192,1,1,1,1,1,2,128,,1,0,0,0
-pre_train,llama3,8b,b200,fp8,8,8192,1,1,1,1,1,2,128,,1,0,0,0
-pre_train,llama3,70b,b200,fp8,64,8192,2,4,2,1,5,1,128,,1,0,0,0
-pre_train,llama3,70b,b200,bf16,64,8192,2,4,2,1,5,1,128,,1,0,0,0
-pre_train,llama31,405b,b200,bf16,128,8192,4,8,2,1,8,1,64,,0,0,0,0
-pre_train,llama31,405b,b200,fp8,128,8192,4,8,2,1,8,1,64,,0,0,0,0
-pre_train,mixtral,8x7b,b200,bf16,64,4096,1,1,1,8,1,2,256,,1,0,0,0
-pre_train,mixtral,8x7b,b200,fp8,64,4096,1,1,1,8,1,2,256,,1,0,0,0
-pre_train,mixtral,8x22b,b200,bf16,256,65536,2,4,8,8,14,1,64,1,0,0,0,0
-pre_train,mixtral,8x22b,b200,fp8,256,65536,2,4,8,8,14,1,64,1,0,0,0,0
-pre_train,nemotron4,15b,b200,bf16,64,4096,1,1,1,1,1,2,256,,1,0,0,0
-pre_train,nemotron4,340b,b200,bf16,128,4096,8,4,1,1,12,1,32,,1,0,0,0
-pre_train,nemotron4,340b,b200,fp8,128,4096,8,4,1,1,12,1,32,,1,0,0,0
-pre_train,nemotron4,15b,b200,fp8,64,4096,1,1,1,1,1,2,256,,1,0,0,0
+task,model,size,system,dtype,num_gpus,seq_len,tp_size,pp_size,cp_size,ep_size,vp_size,mbs,gbs,etp_size,cuda_graphs,use_mcore_fsdp,recompute_layers,activation_offload_layers,recompute_modules
+lora,llama3,70b,b200,fp8,8,4096,1,4,1,1,20,1,32,,0,0,0,0,
+lora,llama3,70b,b200,bf16,8,4096,1,4,1,1,20,1,32,,0,0,0,0,
+lora,llama3,8b,b200,bf16,8,16384,1,1,1,1,1,1,8,,1,0,0,0,
+lora,llama3,8b,b200,fp8,8,16384,1,1,1,1,1,1,8,,0,0,0,0,
+lora,llama31,405b,b200,fp8,32,2048,4,4,1,1,4,1,32,,0,0,0,0,
+lora,llama31,405b,b200,bf16,32,4096,4,4,1,1,4,1,32,,0,0,0,0,
+pre_train,gpt3,175b,b200,bf16,128,2048,4,4,1,1,12,2,256,,0,0,0,0,
+pre_train,gpt3,175b,b200,fp8,128,2048,4,4,1,1,12,2,256,,0,0,0,0,
+pre_train,gpt3,175b,b200,fp8,512,2048,4,4,1,1,6,2,2048,,0,0,0,0,
+pre_train,llama3,8b,b200,bf16,8,8192,1,1,1,1,1,2,128,,1,0,0,0,
+pre_train,llama3,8b,b200,fp8,8,8192,1,1,1,1,1,2,128,,1,0,0,0,
+pre_train,llama3,70b,b200,fp8,64,8192,2,4,2,1,5,1,128,,1,0,0,0,
+pre_train,llama3,70b,b200,bf16,64,8192,2,4,2,1,5,1,128,,1,0,0,0,
+pre_train,llama31,405b,b200,bf16,128,8192,4,8,2,1,8,1,64,,0,0,0,0,
+pre_train,llama31,405b,b200,fp8,128,8192,4,8,2,1,8,1,64,,0,0,0,0,
+pre_train,mixtral,8x7b,b200,bf16,64,4096,1,1,1,8,1,2,256,,1,0,0,0,
+pre_train,mixtral,8x7b,b200,fp8,64,4096,1,1,1,8,1,2,256,,1,0,0,0,
+pre_train,mixtral,8x22b,b200,bf16,256,65536,2,4,8,8,14,1,64,1,0,0,0,0,
+pre_train,mixtral,8x22b,b200,fp8,256,65536,2,4,8,8,14,1,64,1,0,0,0,0,
+pre_train,nemotron4,15b,b200,bf16,64,4096,1,1,1,1,1,2,256,,1,0,0,0,
+pre_train,nemotron4,340b,b200,bf16,128,4096,8,4,1,1,12,1,32,,1,0,0,0,
+pre_train,nemotron4,340b,b200,fp8,128,4096,8,4,1,1,12,1,32,,1,0,0,0,
+pre_train,nemotron4,15b,b200,fp8,64,4096,1,1,1,1,1,2,256,,1,0,0,0,
+pre_train,deepseek,v3,b200,bf16,1024,4096,1,16,1,8,1,1,8192,,0,0,0,0,core_attn/mla_up_proj
+pre_train,deepseek,v3,b200,bf16,256,4096,1,16,1,8,1,1,2048,,0,0,0,0,core_attn/mla_up_proj
+pre_train,deepseek,v3,b200,bf16,128,4096,2,16,1,8,1,1,1024,1,0,0,0,0,core_attn
 sft,llama3,8b,b200,fp8,8,16384,1,1,1,1,1,1,8,,0,0,0,0
 sft,llama3,70b,b200,bf16,32,4096,2,4,1,1,5,1,32,,1,0,0,0
 sft,llama3,70b,b200,fp8,32,4096,2,4,1,1,5,1,32,,1,0,0,0
diff --git a/scripts/performance/recommended_model_configs/model_configs_gb200.csv b/scripts/performance/recommended_model_configs/model_configs_gb200.csv
index fdc1fee42329..c36d1eb2ecd1 100644
--- a/scripts/performance/recommended_model_configs/model_configs_gb200.csv
+++ b/scripts/performance/recommended_model_configs/model_configs_gb200.csv
@@ -1,25 +1,28 @@
-task,model,size,system,dtype,num_gpus,seq_len,tp_size,pp_size,cp_size,ep_size,vp_size,mbs,gbs,etp_size,cuda_graphs,use_mcore_fsdp,recompute_layers,activation_offload_layers
-lora,llama3,70b,gb200,fp8,8.0,2048.0,1.0,4.0,1.0,1.0,20.0,1.0,64.0,,1,0,0,0
-lora,llama3,70b,gb200,bf16,8.0,2048.0,1.0,4.0,1.0,1.0,20.0,1.0,64.0,,1,0,0,0
-lora,llama3,8b,gb200,bf16,8.0,16384.0,1.0,1.0,1.0,1.0,1.0,1.0,8.0,,1,0,0,0
-lora,llama3,8b,gb200,fp8,8.0,16384.0,1.0,1.0,1.0,1.0,1.0,1.0,8.0,,1,0,0,0
-lora,llama31,405b,gb200,fp8,32.0,2048.0,4.0,4.0,1.0,1.0,16.0,1.0,32.0,,1,0,0,0
-lora,llama31,405b,gb200,bf16,32.0,2048.0,4.0,4.0,1.0,1.0,16.0,1.0,32.0,,1,0,0,0
-pre_train,gpt3,175b,gb200,bf16,128.0,2048.0,4.0,4.0,1.0,1.0,12.0,2.0,256.0,,1,0,0,0
-pre_train,gpt3,175b,gb200,fp8,128.0,2048.0,4.0,4.0,1.0,1.0,12.0,2.0,256.0,,1,0,0,0
-pre_train,llama3,8b,gb200,bf16,8.0,8192.0,1.0,1.0,1.0,1.0,1.0,2.0,128.0,,1,0,0,0
-pre_train,llama3,8b,gb200,fp8,8.0,8192.0,1.0,1.0,1.0,1.0,1.0,4.0,128.0,,1,0,0,0
-pre_train,llama3,70b,gb200,fp8,64.0,8192.0,1.0,1.0,1.0,1.0,1.0,1.0,128.0,,0,1,0,0
-pre_train,llama3,70b,gb200,bf16,64.0,8192.0,1.0,1.0,1.0,1.0,1.0,1.0,128.0,,0,1,10,0
-pre_train,llama31,405b,gb200,bf16,128.0,8192.0,4.0,8.0,2.0,1.0,8.0,1.0,64.0,,0,0,0,0
-pre_train,llama31,405b,gb200,fp8,128.0,8192.0,4.0,1.0,1.0,1.0,1.0,1.0,64.0,,0,1,0,0
-pre_train,mixtral,8x7b,gb200,bf16,64.0,4096.0,1.0,1.0,1.0,8.0,1.0,2.0,256.0,,1,0,0,0
-pre_train,mixtral,8x7b,gb200,fp8,64.0,4096.0,1.0,1.0,1.0,8.0,1.0,2.0,256.0,,1,0,0,0
-pre_train,nemotron4,15b,gb200,bf16,64.0,4096.0,1.0,1.0,1.0,1.0,1.0,2.0,256.0,,1,0,0,0
-pre_train,nemotron4,15b,gb200,fp8,64.0,4096.0,1.0,1.0,1.0,1.0,1.0,2.0,256.0,,1,0,0,0
-pre_train,nemotron4,340b,gb200,bf16,128.0,4096.0,4.0,1.0,1.0,1.0,1.0,1.0,32.0,,0,1,0,0
-pre_train,nemotron4,340b,gb200,fp8,128.0,4096.0,8.0,4.0,1.0,1.0,12.0,1.0,32.0,,1,0,0,0
-sft,llama3,8b,gb200,fp8,8.0,16384.0,1.0,1.0,1.0,1.0,1.0,1.0,8.0,,1,0,0,0
-sft,llama3,8b,gb200,bf16,8.0,16384.0,1.0,1.0,1.0,1.0,1.0,1.0,8.0,,1,0,0,0
-sft,llama3,70b,gb200,bf16,32.0,4096.0,2.0,4.0,1.0,1.0,5.0,1.0,32.0,,1,0,0,0
-sft,llama3,70b,gb200,fp8,32.0,4096.0,2.0,4.0,1.0,1.0,5.0,1.0,32.0,,1,0,0,0
+task,model,size,system,dtype,num_gpus,seq_len,tp_size,pp_size,cp_size,ep_size,vp_size,mbs,gbs,etp_size,cuda_graphs,use_mcore_fsdp,recompute_layers,activation_offload_layers,recompute_modules
+lora,llama3,70b,gb200,fp8,8.0,2048.0,1.0,4.0,1.0,1.0,20.0,1.0,64.0,,1,0,0,0,
+lora,llama3,70b,gb200,bf16,8.0,2048.0,1.0,4.0,1.0,1.0,20.0,1.0,64.0,,1,0,0,0,
+lora,llama3,8b,gb200,bf16,8.0,16384.0,1.0,1.0,1.0,1.0,1.0,1.0,8.0,,1,0,0,0,
+lora,llama3,8b,gb200,fp8,8.0,16384.0,1.0,1.0,1.0,1.0,1.0,1.0,8.0,,1,0,0,0,
+lora,llama31,405b,gb200,fp8,32.0,2048.0,4.0,4.0,1.0,1.0,16.0,1.0,32.0,,1,0,0,0,
+lora,llama31,405b,gb200,bf16,32.0,2048.0,4.0,4.0,1.0,1.0,16.0,1.0,32.0,,1,0,0,0,
+pre_train,gpt3,175b,gb200,bf16,128.0,2048.0,4.0,4.0,1.0,1.0,12.0,2.0,256.0,,1,0,0,0,
+pre_train,gpt3,175b,gb200,fp8,128.0,2048.0,4.0,4.0,1.0,1.0,12.0,2.0,256.0,,1,0,0,0,
+pre_train,llama3,8b,gb200,bf16,8.0,8192.0,1.0,1.0,1.0,1.0,1.0,2.0,128.0,,1,0,0,0,
+pre_train,llama3,8b,gb200,fp8,8.0,8192.0,1.0,1.0,1.0,1.0,1.0,4.0,128.0,,1,0,0,0,
+pre_train,llama3,70b,gb200,fp8,64.0,8192.0,1.0,1.0,1.0,1.0,1.0,1.0,128.0,,0,1,0,0,
+pre_train,llama3,70b,gb200,bf16,64.0,8192.0,1.0,1.0,1.0,1.0,1.0,1.0,128.0,,0,1,10,0,
+pre_train,llama31,405b,gb200,bf16,128.0,8192.0,4.0,8.0,2.0,1.0,8.0,1.0,64.0,,0,0,0,0,
+pre_train,llama31,405b,gb200,fp8,128.0,8192.0,4.0,1.0,1.0,1.0,1.0,1.0,64.0,,0,1,0,0,
+pre_train,mixtral,8x7b,gb200,bf16,64.0,4096.0,1.0,1.0,1.0,8.0,1.0,2.0,256.0,,1,0,0,0,
+pre_train,mixtral,8x7b,gb200,fp8,64.0,4096.0,1.0,1.0,1.0,8.0,1.0,2.0,256.0,,1,0,0,0,
+pre_train,nemotron4,15b,gb200,bf16,64.0,4096.0,1.0,1.0,1.0,1.0,1.0,2.0,256.0,,1,0,0,0,
+pre_train,nemotron4,15b,gb200,fp8,64.0,4096.0,1.0,1.0,1.0,1.0,1.0,2.0,256.0,,1,0,0,0,
+pre_train,nemotron4,340b,gb200,bf16,128.0,4096.0,4.0,1.0,1.0,1.0,1.0,1.0,32.0,,0,1,0,0,
+pre_train,nemotron4,340b,gb200,fp8,128.0,4096.0,8.0,4.0,1.0,1.0,12.0,1.0,32.0,,1,0,0,0,
+pre_train,deepseek,v3,gb200,bf16,1024.0,4096.0,1.0,16.0,1.0,64.0,1.0,1.0,8192.0,,0,0,0,0,core_attn/mla_up_proj
+pre_train,deepseek,v3,gb200,bf16,256.0,4096.0,2.0,4.0,1.0,64.0,1.0,1.0,2048.0,1,0,0,0,0,core_attn
+pre_train,deepseek,v3,gb200,bf16,128.0,4096.0,2.0,4.0,1.0,32.0,1.0,1.0,1024.0,1,0,0,0,0,core_attn
+sft,llama3,8b,gb200,fp8,8.0,16384.0,1.0,1.0,1.0,1.0,1.0,1.0,8.0,,1,0,0,0,
+sft,llama3,8b,gb200,bf16,8.0,16384.0,1.0,1.0,1.0,1.0,1.0,1.0,8.0,,1,0,0,0,
+sft,llama3,70b,gb200,bf16,32.0,4096.0,2.0,4.0,1.0,1.0,5.0,1.0,32.0,,1,0,0,0,
+sft,llama3,70b,gb200,fp8,32.0,4096.0,2.0,4.0,1.0,1.0,5.0,1.0,32.0,,1,0,0,0,
diff --git a/scripts/performance/recommended_model_configs/model_configs_h100.csv b/scripts/performance/recommended_model_configs/model_configs_h100.csv
index 07af44c92ac9..48b568dcf339 100644
--- a/scripts/performance/recommended_model_configs/model_configs_h100.csv
+++ b/scripts/performance/recommended_model_configs/model_configs_h100.csv
@@ -1,35 +1,36 @@
-task,model,size,system,dtype,num_gpus,seq_len,tp_size,pp_size,cp_size,ep_size,vp_size,mbs,gbs,etp_size,cuda_graphs,use_mcore_fsdp,recompute_layers,activation_offload_layers
-lora,llama3,70b,h100,fp8,8.0,4096.0,2.0,4.0,1.0,1.0,20.0,1.0,32.0,,0,0,0,0
-lora,llama3,8b,h100,bf16,8.0,4096.0,1.0,1.0,1.0,1.0,1.0,1.0,32.0,,0,0,0,0
-lora,llama3,8b,h100,fp8,8.0,4096.0,1.0,1.0,1.0,1.0,1.0,1.0,32.0,,0,0,0,0
-lora,llama3,70b,h100,bf16,8.0,4096.0,2.0,4.0,1.0,1.0,20.0,1.0,32.0,,0,0,0,0
-lora,llama31,405b,h100,bf16,32.0,2048.0,8.0,6.0,1.0,1.0,8.0,1.0,32.0,,0,0,0,0
-lora,llama31,405b,h100,fp8,32.0,2048.0,8.0,6.0,1.0,1.0,8.0,1.0,32.0,,0,0,0,0
-pre_train,gpt3,175b,h100,fp8,512.0,2048.0,4.0,8.0,1.0,1.0,6.0,2.0,2048.0,,0,0,0,0
-pre_train,gpt3,175b,h100,fp8,128.0,2048.0,4.0,8.0,1.0,1.0,6.0,1.0,256.0,,0,0,0,0
-pre_train,gpt3,175b,h100,bf16,128.0,2048.0,4.0,8.0,1.0,1.0,6.0,1.0,256.0,,0,0,0,0
-pre_train,llama31,405b,h100,fp8,1024.0,8192.0,8.0,8.0,2.0,1.0,8.0,1.0,512.0,,0,0,0,0
-pre_train,llama31,405b,h100,bf16,1024.0,8192.0,8.0,8.0,2.0,1.0,8.0,1.0,512.0,,0,0,0,0
-pre_train,llama3,70b,h100,fp8,64.0,8192.0,4.0,8.0,1.0,1.0,5.0,1.0,128.0,,0,0,0,0
-pre_train,llama3,70b,h100,bf16,64.0,8192.0,4.0,4.0,2.0,1.0,5.0,1.0,128.0,,0,0,0,0
-pre_train,llama3,8b,h100,fp8,8.0,8192.0,1.0,1.0,2.0,1.0,1.0,1.0,128.0,,0,0,0,0
-pre_train,llama3,8b,h100,bf16,8.0,8192.0,1.0,1.0,2.0,1.0,1.0,1.0,128.0,,0,0,0,0
-pre_train,mixtral,8x7b,h100,fp8,64.0,4096.0,1.0,4.0,1.0,8.0,8.0,1.0,256.0,,0,0,0,0
-pre_train,mixtral,8x7b,h100,bf16,64.0,4096.0,1.0,4.0,1.0,8.0,8.0,1.0,256.0,,0,0,0,0
-pre_train,mixtral,8x22b,h100,bf16,256.0,65536.0,4.0,4.0,8.0,8.0,14.0,1.0,256.0,1.0,0,0,0,0
-pre_train,mixtral,8x22b,h100,fp8,256.0,65536.0,4.0,4.0,8.0,8.0,14.0,1.0,256.0,1.0,0,0,0,0
-pre_train,nemotron3,22b,h100,bf16,64.0,4096.0,2.0,4.0,1.0,1.0,10.0,1.0,256.0,,0,0,0,0
-pre_train,nemotron3,22b,h100,fp8,64.0,4096.0,1.0,4.0,1.0,1.0,10.0,1.0,256.0,,0,0,0,0
-pre_train,nemotron3,8b,h100,bf16,8.0,4096.0,2.0,1.0,1.0,1.0,1.0,2.0,32.0,,0,0,0,0
-pre_train,nemotron3,8b,h100,fp8,8.0,4096.0,2.0,1.0,1.0,1.0,1.0,2.0,32.0,,0,0,0,0
-pre_train,nemotron4,15b,h100,bf16,64.0,4096.0,2.0,1.0,1.0,1.0,1.0,2.0,256.0,,0,0,0,0
-pre_train,nemotron4,15b,h100,fp8,64.0,4096.0,2.0,1.0,1.0,1.0,1.0,2.0,256.0,,0,0,0,0
-pre_train,nemotron4,340b,h100,bf16,256.0,4096.0,8.0,8.0,1.0,1.0,12.0,1.0,64.0,,0,0,0,0
-pre_train,nemotron4,340b,h100,fp8,256.0,4096.0,8.0,8.0,1.0,1.0,12.0,1.0,64.0,,0,0,0,0
-sft,llama3,8b,h100,bf16,8.0,4096.0,1.0,1.0,1.0,1.0,1.0,1.0,32.0,,0,0,0,0
-sft,llama3,8b,h100,fp8,8.0,4096.0,1.0,1.0,1.0,1.0,1.0,1.0,32.0,,0,0,0,0
-sft,llama3,70b,h100,bf16,32.0,4096.0,4.0,4.0,1.0,1.0,5.0,1.0,32.0,,0,0,0,0
-sft,llama3,70b,h100,fp8,32.0,4096.0,4.0,4.0,1.0,1.0,5.0,1.0,32.0,,0,0,0,0
-pre_train,neva_llama3,8b,h100,bf16,8.0,8192.0,1.0,1.0,2.0,1.0,1.0,1.0,128.0,,0,0,0,0
-pre_train,flux,12b,h100,bf16,8.0,4096.0,1.0,1.0,1.0,1.0,1.0,1.0,64.0,,0,0,0,0
-sft,deepseek,v3,h100,bf16,992.0,4096.0,1.0,31.0,1.0,8.0,2.0,1.0,3968.0,,0,0,0,0
+task,model,size,system,dtype,num_gpus,seq_len,tp_size,pp_size,cp_size,ep_size,vp_size,mbs,gbs,etp_size,cuda_graphs,use_mcore_fsdp,recompute_layers,activation_offload_layers,recompute_modules
+lora,llama3,70b,h100,fp8,8.0,4096.0,2.0,4.0,1.0,1.0,20.0,1.0,32.0,,0,0,0,0,
+lora,llama3,8b,h100,bf16,8.0,4096.0,1.0,1.0,1.0,1.0,1.0,1.0,32.0,,0,0,0,0,
+lora,llama3,8b,h100,fp8,8.0,4096.0,1.0,1.0,1.0,1.0,1.0,1.0,32.0,,0,0,0,0,
+lora,llama3,70b,h100,bf16,8.0,4096.0,2.0,4.0,1.0,1.0,20.0,1.0,32.0,,0,0,0,0,
+lora,llama31,405b,h100,bf16,32.0,2048.0,8.0,6.0,1.0,1.0,8.0,1.0,32.0,,0,0,0,0,
+lora,llama31,405b,h100,fp8,32.0,2048.0,8.0,6.0,1.0,1.0,8.0,1.0,32.0,,0,0,0,0,
+pre_train,gpt3,175b,h100,fp8,512.0,2048.0,4.0,8.0,1.0,1.0,6.0,2.0,2048.0,,0,0,0,0,
+pre_train,gpt3,175b,h100,fp8,128.0,2048.0,4.0,8.0,1.0,1.0,6.0,1.0,256.0,,0,0,0,0,
+pre_train,gpt3,175b,h100,bf16,128.0,2048.0,4.0,8.0,1.0,1.0,6.0,1.0,256.0,,0,0,0,0,
+pre_train,llama31,405b,h100,fp8,1024.0,8192.0,8.0,8.0,2.0,1.0,8.0,1.0,512.0,,0,0,0,0,
+pre_train,llama31,405b,h100,bf16,1024.0,8192.0,8.0,8.0,2.0,1.0,8.0,1.0,512.0,,0,0,0,0,
+pre_train,llama3,70b,h100,fp8,64.0,8192.0,4.0,8.0,1.0,1.0,5.0,1.0,128.0,,0,0,0,0,
+pre_train,llama3,70b,h100,bf16,64.0,8192.0,4.0,4.0,2.0,1.0,5.0,1.0,128.0,,0,0,0,0,
+pre_train,llama3,8b,h100,fp8,8.0,8192.0,1.0,1.0,2.0,1.0,1.0,1.0,128.0,,0,0,0,0,
+pre_train,llama3,8b,h100,bf16,8.0,8192.0,1.0,1.0,2.0,1.0,1.0,1.0,128.0,,0,0,0,0,
+pre_train,mixtral,8x7b,h100,fp8,64.0,4096.0,1.0,4.0,1.0,8.0,8.0,1.0,256.0,,0,0,0,0,
+pre_train,mixtral,8x7b,h100,bf16,64.0,4096.0,1.0,4.0,1.0,8.0,8.0,1.0,256.0,,0,0,0,0,
+pre_train,mixtral,8x22b,h100,bf16,256.0,65536.0,4.0,4.0,8.0,8.0,14.0,1.0,256.0,1.0,0,0,0,0,
+pre_train,mixtral,8x22b,h100,fp8,256.0,65536.0,4.0,4.0,8.0,8.0,14.0,1.0,256.0,1.0,0,0,0,0,
+pre_train,nemotron3,22b,h100,bf16,64.0,4096.0,2.0,4.0,1.0,1.0,10.0,1.0,256.0,,0,0,0,0,
+pre_train,nemotron3,22b,h100,fp8,64.0,4096.0,1.0,4.0,1.0,1.0,10.0,1.0,256.0,,0,0,0,0,
+pre_train,nemotron3,8b,h100,bf16,8.0,4096.0,2.0,1.0,1.0,1.0,1.0,2.0,32.0,,0,0,0,0,
+pre_train,nemotron3,8b,h100,fp8,8.0,4096.0,2.0,1.0,1.0,1.0,1.0,2.0,32.0,,0,0,0,0,
+pre_train,nemotron4,15b,h100,bf16,64.0,4096.0,2.0,1.0,1.0,1.0,1.0,2.0,256.0,,0,0,0,0,
+pre_train,nemotron4,15b,h100,fp8,64.0,4096.0,2.0,1.0,1.0,1.0,1.0,2.0,256.0,,0,0,0,0,
+pre_train,nemotron4,340b,h100,bf16,256.0,4096.0,8.0,8.0,1.0,1.0,12.0,1.0,64.0,,0,0,0,0,
+pre_train,nemotron4,340b,h100,fp8,256.0,4096.0,8.0,8.0,1.0,1.0,12.0,1.0,64.0,,0,0,0,0,
+pre_train,deepseek,v3,h100,bf16,1024.0,4096.0,2.0,16.0,1.0,64.0,1.0,1.0,8192.0,1.0,0,0,0,0,mla_up_proj
+sft,llama3,8b,h100,bf16,8.0,4096.0,1.0,1.0,1.0,1.0,1.0,1.0,32.0,,0,0,0,0,
+sft,llama3,8b,h100,fp8,8.0,4096.0,1.0,1.0,1.0,1.0,1.0,1.0,32.0,,0,0,0,0,
+sft,llama3,70b,h100,bf16,32.0,4096.0,4.0,4.0,1.0,1.0,5.0,1.0,32.0,,0,0,0,0,
+sft,llama3,70b,h100,fp8,32.0,4096.0,4.0,4.0,1.0,1.0,5.0,1.0,32.0,,0,0,0,0,
+pre_train,neva_llama3,8b,h100,bf16,8.0,8192.0,1.0,1.0,2.0,1.0,1.0,1.0,128.0,,0,0,0,0,
+pre_train,flux,12b,h100,bf16,8.0,4096.0,1.0,1.0,1.0,1.0,1.0,1.0,64.0,,0,0,0,0,
+sft,deepseek,v3,h100,bf16,992.0,4096.0,1.0,31.0,1.0,8.0,2.0,1.0,3968.0,,0,0,0,0,
diff --git a/scripts/performance/utils.py b/scripts/performance/utils.py
index 8201a93f6087..3a2ef2f3416b 100644
--- a/scripts/performance/utils.py
+++ b/scripts/performance/utils.py
@@ -211,12 +211,21 @@ def get_user_configs(gpu: str, task: str, model_name: str, model_size: str, args
     )
     activation_offload_layers = 0 if activation_offload_layers is None else int(activation_offload_layers)
 
+    if args.recompute_modules is not None:
+        recompute_modules = args.recompute_modules
+        assert isinstance(recompute_modules, list), "recompute_modules must be a list"
+    elif config.get("recompute_modules") is not None:
+        recompute_modules = config.get("recompute_modules").split('/')
+    else:
+        recompute_modules = None
+
     kwargs = num_nodes, mbs, gbs, tp_size, pp_size, cp_size, vp_size, ep_size, etp_size
     kwargs = [int(arg) if arg is not None else arg for arg in kwargs] + [
         enable_cuda_graphs,
         use_mcore_fsdp,
         recompute_layers,
         activation_offload_layers,
+        recompute_modules,
     ]
 
     return kwargs
@@ -242,6 +251,7 @@ def set_primary_perf_configs(
     activation_offload_layers: int = 0,
     compute_dtype: str = None,
     fp8_recipe: str = None,
+    recompute_modules: Optional[List[str]] = None,
 ):
     """Set experiment configs we usually tune for performance of all models."""
     # nemo.lightning.Trainer configs
@@ -343,6 +353,16 @@ def set_primary_perf_configs(
             logging.warning("Currently FSDP does not support FP8 param gather. Disabling fp8 param gather.")
             recipe.trainer.plugins.fp8_param_gather = False
 
+    # Activation recompute configs
+    if recompute_modules is not None:
+        recipe.model.config.recompute_modules = recompute_modules
+        assert (
+            recipe.model.config.recompute_granularity == "selective"
+        ), "recompute_granularity must be selective when recompute_modules is provided"
+        assert (
+            recipe.model.config.recompute_num_layers is None
+        ), "recompute_num_layers must be None when recompute_modules is provided"
+
     return recipe
 
 
diff --git a/scripts/performance/vlm/finetune_neva_8b.py b/scripts/performance/vlm/finetune_neva_8b.py
index f2cde29a780c..2cf248a905e1 100644
--- a/scripts/performance/vlm/finetune_neva_8b.py
+++ b/scripts/performance/vlm/finetune_neva_8b.py
@@ -89,7 +89,7 @@ def override_recipe_configs(
     args_sanity_check(args)
 
     kwargs = get_user_configs(args.gpu.lower(), "pre_train", "neva_llama3", "8b", args)
-    num_nodes, mbs, gbs, tp_size, pp_size, cp_size, vp_size, ep_size, _, enable_cuda_graphs, _, _, _ = kwargs
+    num_nodes, mbs, gbs, tp_size, pp_size, cp_size, vp_size, ep_size, _, enable_cuda_graphs = kwargs[:10]
 
     recipe = override_recipe_configs(
         args, num_nodes, mbs, gbs, tp_size, pp_size, cp_size, vp_size, ep_size, enable_cuda_graphs


Reviews:
user: login: github-advanced-security[bot]
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-09T17:57:06Z
user: login: gdengk
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-14T18:41:36Z
user: login: cuichenx
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-14T20:50:24Z
user: login: erhoo82
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-22T04:10:09Z
user: login: erhoo82
body: LGTM
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-22T04:29:15Z

Review comments:
diff_hunk: @@ -221,3 +229,46 @@
     )
 
     return joint_layer_flops + single_layer_flops + other_flops
+
+
+def deepseekv3(config: FLOPSConfig):
+    """Model FLOPs for DeepSeek V3"""
+
+    # self-attention flops
+    bmm1_flops = (config.qk_head_dim + config.qk_pos_emb_head_dim) * (config.enc_seq_len**2)
+    bmm2_flops = config.v_head_dim * (config.enc_seq_len**2)
+    per_input_attention_flops = 6 * (bmm1_flops + bmm2_flops) * config.num_layers
+
+    # linear layer flops
+    mla_params = (
+        config.hs * config.q_lora_rank
+        + config.q_lora_rank * (config.qk_head_dim + config.qk_pos_emb_head_dim) * config.attention_heads
+    )  # Q
+    mla_params += (
+        config.hs * config.kv_lora_rank
+        + config.kv_lora_rank * (config.qk_head_dim + config.qk_pos_emb_head_dim) * config.attention_heads
+    )  # K
+    mla_params += (
+        config.hs * config.kv_lora_rank + config.kv_lora_rank * config.v_head_dim * config.attention_heads
+    )  # V
+    mla_params += config.v_head_dim * config.attention_heads * config.hs  # Proj
+
+    dense_layer_ffn_params = config.hs * config.ffn_hidden_size * 3  # gated linear unit
+    per_shared_expert_params = config.hs * config.moe_shared_expert_intermediate_size * 3
+    per_selected_expert_params = config.hs * config.moe_ffn_hidden_size * 3
+    ffn_params = 0
+
+    if isinstance(config.moe_layer_freq, int):
+        moe_layer_pattern = [1 if (i % config.moe_layer_freq == 0) else 0 for i in range(config.layers)]
+    else:
+        moe_layer_pattern = config.moe_layer_freq
+    for i in moe_layer_pattern:
+        if i == 0:
+            ffn_params += dense_layer_ffn_params
+        else:
+            ffn_params += per_shared_expert_params + (per_selected_expert_params * config.moe_router_topk)
+    per_input_params = mla_params + ffn_params
+    per_input_linear_flops = 6 * per_input_params * config.enc_seq_len
path: nemo/utils/flops_formulas.py
user: login: github-advanced-security[bot]
body: ## Unreachable code

This statement is unreachable.

[Show more details](https://github.com/NVIDIA/NeMo/security/code-scanning/13928)
created_at: 2025-04-09T17:57:05Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,177 @@
+# 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.
+
+from os.path import basename, splitext
+
+import fiddle as fdl
path: scripts/performance/llm/pretrain_deepseek_v3.py
user: login: github-advanced-security[bot]
body: ## Unused import

Import of 'fdl' is not used.

[Show more details](https://github.com/NVIDIA/NeMo/security/code-scanning/13926)
created_at: 2025-04-09T17:57:05Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,177 @@
+# 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.
+
+from os.path import basename, splitext
+
+import fiddle as fdl
+import fiddle._src.experimental.dataclasses as fdl_dc
path: scripts/performance/llm/pretrain_deepseek_v3.py
user: login: github-advanced-security[bot]
body: ## Unused import

Import of 'fdl_dc' is not used.

[Show more details](https://github.com/NVIDIA/NeMo/security/code-scanning/13927)
created_at: 2025-04-09T17:57:05Z
author_association: CONTRIBUTOR
diff_hunk: @@ -221,3 +229,50 @@ def flux(config: FLOPSConfig):
     )
 
     return joint_layer_flops + single_layer_flops + other_flops
+
+
+def deepseekv3(config: FLOPSConfig):
+    """Model FLOPs for DeepSeek V3"""
+
+    # self-attention flops
+    bmm1_flops = (config.qk_head_dim + config.qk_pos_emb_head_dim) * (config.enc_seq_len**2)
+    bmm2_flops = config.v_head_dim * (config.enc_seq_len**2)
+    per_input_attention_flops = 6 * (bmm1_flops + bmm2_flops) * config.layers
+
+    # linear layer flops
+    per_layer_mla_params = (
+        config.hs * config.q_lora_rank
+        + config.q_lora_rank * (config.qk_head_dim + config.qk_pos_emb_head_dim) * config.attention_heads
+    )  # Q
+    per_layer_mla_params += (
+        config.hs * config.kv_lora_rank
+        + config.kv_lora_rank * (config.qk_head_dim + config.qk_pos_emb_head_dim) * config.attention_heads
path: nemo/utils/flops_formulas.py
user: login: gdengk
body: For k pos embed, the project gemm dim is `config.qk_pos_emb_head_dim` without `config.attention_heads` since the output will be broadcast concat to the K tensor.
created_at: 2025-04-14T18:39:52Z
author_association: COLLABORATOR
diff_hunk: @@ -221,3 +229,50 @@ def flux(config: FLOPSConfig):
     )
 
     return joint_layer_flops + single_layer_flops + other_flops
+
+
+def deepseekv3(config: FLOPSConfig):
+    """Model FLOPs for DeepSeek V3"""
+
+    # self-attention flops
+    bmm1_flops = (config.qk_head_dim + config.qk_pos_emb_head_dim) * (config.enc_seq_len**2)
path: nemo/utils/flops_formulas.py
user: login: gdengk
body: To align with the nums from Mcore team, we actually put the casualmasking factor of 0.5 into consideration. 

created_at: 2025-04-14T18:41:22Z
author_association: COLLABORATOR
diff_hunk: @@ -80,18 +81,32 @@ def __init__(
         if query_groups is None:
             query_groups = attention_heads
 
-        self.flops_config = flops_formulas.FLOPSConfig(
-            gbs=gbs,
-            enc_seq_len=enc_seq_len,
-            hs=hs,
-            layers=layers,
-            ffn_hs=ffn_hs,
-            attention_heads=attention_heads,
-            moe_router_topk=moe_router_topk,
-            query_groups=query_groups,
-            vocab_size=vocab_size,
-            model_pattern=model_pattern,
-        )
+        config_kwargs = {
+            "gbs": gbs,
+            "enc_seq_len": enc_seq_len,
+            "hs": hs,
+            "layers": layers,
+            "ffn_hs": ffn_hs,
+            "attention_heads": attention_heads,
+            "moe_router_topk": moe_router_topk,
+            "query_groups": query_groups,
+            "vocab_size": vocab_size,
+            "model_pattern": model_pattern,
+        }
+
+        from megatron.core.transformer.transformer_config import MLATransformerConfig
+
+        if isinstance(self.model_cfg, MLATransformerConfig):
+            config_kwargs["qk_head_dim"] = self.model_cfg.qk_head_dim
+            config_kwargs["qk_pos_emb_head_dim"] = self.model_cfg.qk_pos_emb_head_dim
+            config_kwargs["v_head_dim"] = self.model_cfg.v_head_dim
+            config_kwargs["q_lora_rank"] = self.model_cfg.q_lora_rank
+            config_kwargs["kv_lora_rank"] = self.model_cfg.kv_lora_rank
+            config_kwargs["moe_layer_freq"] = self.model_cfg.moe_layer_freq
+            config_kwargs["moe_shared_expert_intermediate_size"] = self.model_cfg.moe_shared_expert_intermediate_size
+            config_kwargs["moe_ffn_hidden_size"] = self.model_cfg.moe_ffn_hidden_size
path: nemo/lightning/pytorch/callbacks/flops_callback.py
user: login: cuichenx
body: The last three are not specific to MLA. Should we pull them out of the if block?
created_at: 2025-04-14T20:45:52Z
author_association: COLLABORATOR
diff_hunk: @@ -298,6 +298,15 @@ def parse_cli_args():
     def list_of_strings(arg):
         return arg.split(',')
 
+    parser.add_argument(
+        "-rm",
+        "--recompute_modules",
+        type=list_of_strings,
+        help="Comma-separated string of modules to recompute. Defaults to None",
path: scripts/performance/argument_parser.py
user: login: erhoo82
body: `Comma-separated` should be `Slash-separated`
created_at: 2025-04-22T04:08:29Z
author_association: COLLABORATOR

PR comments:
user: login: guyueh1
created_at: 2025-04-10T01:27:09Z
body: @cuichenx @gdengk This is still WIP, but can you review, especially the flops calculator?
author_association: COLLABORATOR
user: login: guyueh1
created_at: 2025-04-21T15:03:50Z
body: This PR is ready, please review @cuichenx @gdengk @erhoo82 
author_association: COLLABORATOR