Skip to content

Malay/move expandable flag #3781

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
dingqingy-nv merged 3 commits into from
May 12, 2026
Merged

Malay/move expandable flag #3781

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
af5e7f6
move expandable flag
malay-nagda May 12, 2026
b45986e
cleanup
malay-nagda May 12, 2026
84c6142
remove vars for nccl_ub
malay-nagda May 12, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
72 changes: 12 additions & 60 deletions scripts/performance/perf_plugins.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -276,68 +276,24 @@ 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,
compute_dtype: str,
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":
Expand All @@ -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")
Expand Down Expand Up @@ -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,
Expand Down
2 changes: 2 additions & 0 deletions scripts/performance/utils/executors.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
}


Expand Down
Loading