diff --git a/scripts/performance/perf_plugins.py b/scripts/performance/perf_plugins.py
index 89f2133b88..3537d75ca2 100644
--- a/scripts/performance/perf_plugins.py
+++ b/scripts/performance/perf_plugins.py
@@ -37,9 +37,9 @@
 
 
 try:
-    from utils.utils import get_workload_base_config
+    from utils.utils import WorkloadBaseConfig, get_workload_base_config
 except (ImportError, ModuleNotFoundError):
-    from .utils.utils import get_workload_base_config
+    from .utils.utils import WorkloadBaseConfig, get_workload_base_config
 
 logger: logging.Logger = logging.getLogger(__name__)
 NSYS_SQLITE_EXPORT_ARG = "--export=sqlite"
@@ -276,6 +276,7 @@ def _set_model_specific_environment_variables(
         self,
         task: Union["run.Partial", "run.Script"],
         executor: "run.Executor",
+        workload_base_config: WorkloadBaseConfig,
         model_family_name: str,
         model_recipe_name: str,
         gpu: str,
@@ -283,61 +284,16 @@ def _set_model_specific_environment_variables(
         train_task: str,
     ):
         """Set model-specific environment variables"""
-        if (
-            model_family_name in ["llama"]
-            and model_recipe_name in ["llama31_405b"]
-            and train_task == "pretrain"
-            and gpu in ["gb200", "gb300"]
-        ):
-            if compute_dtype in ["fp8_cs", "fp8_mx", "nvfp4"]:
-                executor.env_vars["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
-                executor.env_vars["NCCL_GRAPH_REGISTER"] = "0"
-        elif (
-            model_family_name in ["deepseek"]
-            and model_recipe_name in ["deepseek_v3"]
-            and train_task == "pretrain"
-            and gpu in ["h100", "gb200", "b200"]
-        ):
-            executor.env_vars["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
-            if gpu in ["gb200", "b200"]:
-                executor.env_vars["NCCL_GRAPH_REGISTER"] = "0"
-        elif (
-            model_family_name in ["qwen"]
-            and model_recipe_name in ["qwen3_next_80b_a3b"]
-            and train_task == "pretrain"
-            and gpu in ["h100"]
-            and compute_dtype == "fp8_cs"
-        ):
-            # NCCL 2.29.7 increases memory pressure on H100, causing allocator
-            # fragmentation OOM. expandable_segments lets the allocator reclaim
-            # fragmented physical memory and avoids the OOM without disabling
-            # any NCCL algorithms.
-            executor.env_vars["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
-        elif (
-            model_family_name in ["nemotronh"]
-            and model_recipe_name in ["nemotron_3_nano"]
-            and train_task == "pretrain"
-            and gpu in ["h100"]
-            and compute_dtype == "fp8_cs"
-        ):
-            # pytorch:26.04-py3 base + NCCL 2.30.4 bump (vs 26.02 / 2.29.3) tightens
-            # memory headroom on H100 fp8_cs and the optimizer's
-            # _copy_main_params_to_model_params spike at iter 3 (post CUDA-Graph capture)
-            # OOMs from allocator fragmentation. expandable_segments lets the allocator
-            # reclaim fragmented physical memory; NCCL_GRAPH_REGISTER=0 is its required
-            # partner under CUDA Graphs (MCore guards against the unsafe combo and asserts
-            # at cuda_graphs.py:1750 if expandable_segments:True is set without it).
-            executor.env_vars["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
-            executor.env_vars["NCCL_GRAPH_REGISTER"] = "0"
-
         if model_family_name in ["deepseek"]:
             executor.env_vars["NVTE_ALLOW_NONDETERMINISTIC_ALGO"] = "0"
-        if model_recipe_name in ["llama3_70b"]:
-            if compute_dtype in ["fp8_cs", "fp8_mx"]:
-                if train_task in ["sft", "lora"]:
-                    if gpu in ["gb300", "gb200", "h100"]:
-                        executor.env_vars["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
-                        executor.env_vars["NCCL_GRAPH_REGISTER"] = "0"
+
+        remove_allocator_env_vars = workload_base_config.nccl_ub is True or (
+            model_family_name == "llama" and workload_base_config.use_megatron_fsdp is True
+        )
+        if remove_allocator_env_vars:
+            for env_var in ("PYTORCH_CUDA_ALLOC_CONF", "NCCL_GRAPH_REGISTER"):
+                executor.env_vars.pop(env_var, None)
+
         del_cudnn_ln = True
         if gpu in ["h100"]:
             if model_family_name == "llama" and model_recipe_name == "llama3_8b" and train_task == "pretrain":
@@ -359,14 +315,9 @@ def _set_model_specific_environment_variables(
                 if compute_dtype == "fp8_mx":
                     del_cudnn_ln = False
         if model_family_name in ["llama"] and train_task in ["sft"]:
-            # TODO: Verify for H100 and 8b
             del_cudnn_ln = False
-            if gpu in ["h100"] and model_recipe_name in ["llama3_70b"] and compute_dtype == "fp8_cs":
-                executor.env_vars["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
-                executor.env_vars["NCCL_GRAPH_REGISTER"] = "0"
         if model_recipe_name in ["nemotron_3_nano"]:
             del_cudnn_ln = False
-
         if del_cudnn_ln:
             if "NVTE_NORM_FWD_USE_CUDNN" in executor.env_vars:
                 executor.env_vars.pop("NVTE_NORM_FWD_USE_CUDNN")
@@ -584,6 +535,7 @@ def setup(self, task: Union["run.Partial", "run.Script"], executor: "run.Executo
         self._set_model_specific_environment_variables(
             task,
             executor,
+            workload_base_config,
             self.model_family_name,
             self.model_recipe_name,
             self.gpu,
diff --git a/scripts/performance/utils/executors.py b/scripts/performance/utils/executors.py
index 091fd8ee3d..82bda32000 100644
--- a/scripts/performance/utils/executors.py
+++ b/scripts/performance/utils/executors.py
@@ -45,6 +45,8 @@
     "NVTE_NORM_BWD_USE_CUDNN": "1",
     "TORCH_NCCL_HIGH_PRIORITY": "1",
     "HF_HUB_OFFLINE": "0",  # Keep HF Hub online by default; --offline flips this to 1.
+    "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
+    "NCCL_GRAPH_REGISTER": "0",
 }