Malay/move expandable flag

[malay-nagda]

What does this PR do ?

Add expandable_segments flag for all perf jobs except for those using FSDP and nccl_ub.

Changelog

    "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
    "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
        )

GitHub Actions CI

See the CI sectionin the Contributing doc for how to trigger the CI. A Nvidia developer will need to approve and trigger the CI for external contributors.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• 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?

If you haven't finished some of the above items you can still open "Draft" PR.

Additional Information

• Related to # (issue)

[claude]

Light Code Review

This PR promotes PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True and NCCL_GRAPH_REGISTER=0 from per-model conditionals into the global PERF_ENV_VARS default dict, then removes the now-redundant model-specific blocks. This is a good simplification.

Observations

1. FSDP pop is llama-only (perf_plugins.py:290): The new block that removes PYTORCH_CUDA_ALLOC_CONF and NCCL_GRAPH_REGISTER for FSDP is gated on model_family_name == llama. Currently only llama configs set use_megatron_fsdp=True, so this is correct today. If a non-llama family enables FSDP in the future, it will silently keep both env vars. Consider dropping the llama guard and using just workload_base_config.use_megatron_fsdp is True so the cleanup applies generically.

2. PR description is empty - the changelog and summary sections are still boilerplate. Please fill these in before merge so reviewers and future readers understand the rationale.

3. Missing blank line (perf_plugins.py:316-317): The old code had a blank line before if del_cudnn_ln: which was removed. Very minor, but the surrounding style uses blank lines between logical blocks.

Suggested test cases

No perf config files were changed, so no specific perf test names are impacted. However, since the env var change affects all perf runs globally, the following should be verified:

• Any llama FSDP config (e.g. llama3_8b / llama3_70b with use_megatron_fsdp=True) to confirm the env vars are correctly popped
• A non-llama model (e.g. deepseek_v3, qwen3_next_80b_a3b) to confirm the new defaults propagate and training still converges
• A non-FSDP llama config to confirm it keeps both env vars

Generated with Claude Code