fix(performance): preserve recipe env self-exec

Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>

Author: yaoyu-33

docs/fern/versions/nightly/pages/recipe-usage.mdx

@@ -317,7 +317,7 @@ cfg.env_vars.update(
 
 Values may be strings, integers, floats, or booleans and are converted to strings when exported. Existing shell, container, or launcher values take precedence over recipe defaults. This lets the same mechanism work for library recipes under `megatron.bridge.recipes` and flat performance recipes under `megatron.bridge.perf_recipes` without preventing cluster-specific overrides.
 
-When launching through `scripts/performance/setup_experiment.py`, a rank-local pre-exec wrapper resolves library or flat performance recipe defaults inside the worker container, then replaces itself with the selected training entry point. This keeps recipe imports off dependency-light launcher nodes. Explicit `--env`, `--custom_env_vars`, shell, and container values take precedence.
+When launching through `scripts/performance/setup_experiment.py`, the selected rank-local training script resolves its library or flat performance recipe defaults inside the worker container, exports them, and then replaces itself with a clean interpreter running that same script. This keeps recipe imports off dependency-light launcher nodes and ensures training-framework imports observe recipe-owned process settings. Explicit `--env`, `--custom_env_vars`, shell, and container values take precedence.
 
 Environment variables are serialized with the rest of the recipe and can be added or overridden through a Hydra-style override such as `'++env_vars={TORCHINDUCTOR_WORKER_START:fork,QUANTIZATION_TYPE_DEBUG:1}'`. Do not store credentials or other secrets in recipe configs.
 

docs/recipe-usage.md

@@ -318,7 +318,7 @@ cfg.env_vars.update(
 
 Values may be strings, integers, floats, or booleans and are converted to strings when exported. Existing shell, container, or launcher values take precedence over recipe defaults. This lets the same mechanism work for library recipes under `megatron.bridge.recipes` and flat performance recipes under `megatron.bridge.perf_recipes` without preventing cluster-specific overrides.
 
-When launching through `scripts/performance/setup_experiment.py`, a rank-local pre-exec wrapper resolves library or flat performance recipe defaults inside the worker container, then replaces itself with the selected training entry point. This keeps recipe imports off dependency-light launcher nodes. Explicit `--env`, `--custom_env_vars`, shell, and container values take precedence.
+When launching through `scripts/performance/setup_experiment.py`, the selected rank-local training script resolves its library or flat performance recipe defaults inside the worker container, exports them, and then replaces itself with a clean interpreter running that same script. This keeps recipe imports off dependency-light launcher nodes and ensures training-framework imports observe recipe-owned process settings. Explicit `--env`, `--custom_env_vars`, shell, and container values take precedence.
 
 Environment variables are serialized with the rest of the recipe and can be added or overridden through a Hydra-style override such as `'++env_vars={TORCHINDUCTOR_WORKER_START:fork,QUANTIZATION_TYPE_DEBUG:1}'`. Do not store credentials or other secrets in recipe configs.
 

scripts/performance/run_recipe.py

@@ -21,30 +21,13 @@
 import os
 import sys
 
-import torch
 from argument_parser import parse_cli_args
-from utils.datasets import (
-    create_mock_dataset_config,
-    create_rp2_dataset_config,
-    create_squad_dataset_config,
-)
-from utils.utils import (
-    apply_library_argparse_overrides,
-    apply_library_target_topology_environment,
-    explicit_environment_override_names,
-    finalize_library_config_overrides,
-    get_library_recipe,
-)
-
-from megatron.bridge.recipes.utils.determinism_utils import apply_determinism_overrides
-from megatron.bridge.training.config import apply_environment_variables
-from megatron.bridge.training.utils.omegaconf_utils import process_config_with_overrides
-from megatron.bridge.utils.common_utils import get_rank_safe
 
 
 # Diffusion model families manage their own dataset configs and require
 # a dedicated forward step function rather than the standard GPT step.
 DIFFUSION_FAMILIES = frozenset({"flux", "wan"})
+ENV_BOOTSTRAP_MARKER = "_MB_PERF_ENV_BOOTSTRAPPED"
 
 
 def _get_diffusion_step(model_family_name: str):
@@ -62,6 +45,9 @@ def _get_diffusion_step(model_family_name: str):
 
 def set_user_overrides(config, args):
     """Apply CLI arguments to ConfigContainer fields."""
+    from utils.datasets import create_mock_dataset_config, create_rp2_dataset_config, create_squad_dataset_config
+    from utils.utils import apply_library_argparse_overrides
+
     is_diffusion = args.model_family_name in DIFFUSION_FAMILIES
     config = apply_library_argparse_overrides(config, args)
 
@@ -194,6 +180,8 @@ def set_user_overrides(config, args):
         config.logger.wandb_save_dir = args.wandb_save_dir
 
     if args.deterministic:
+        from megatron.bridge.recipes.utils.determinism_utils import apply_determinism_overrides
+
         apply_determinism_overrides(config)
 
     # Handle convergence mode configuration
@@ -220,37 +208,87 @@ def set_user_overrides(config, args):
     return config
 
 
-def main():
-    """Main entry point for the training script."""
-
-    # Parse known args and capture unknown ones for Hydra-style config overrides
-    # (e.g. model.hidden_size=15360 model.num_moe_experts=8)
-    parser = parse_cli_args()
-    args, cli_overrides = parser.parse_known_args()
+def _get_library_recipe(args):
+    """Resolve the selected library recipe without importing it at module load time."""
+    from utils.utils import get_library_recipe
 
-    recipe = get_library_recipe(
+    return get_library_recipe(
         model_family_name=args.model_family_name,
         model_recipe_name=args.model_recipe_name,
         train_task=args.task,
         wandb_experiment_name=args.wandb_experiment_name,
     )
 
-    base_env_vars = dict(recipe.env_vars)
-    recipe = set_user_overrides(recipe, args)
 
+def _process_library_hydra_overrides(recipe, cli_overrides: list[str]):
+    """Apply Hydra overrides using the same ordering in both self-exec passes."""
+    from megatron.bridge.training.utils.omegaconf_utils import process_config_with_overrides
+
+    return process_config_with_overrides(recipe, cli_overrides=cli_overrides)
+
+
+def _apply_library_recipe_overrides(recipe, cli_overrides: list[str], args):
+    """Apply env-relevant argparse and Hydra overrides before self-exec."""
+    from utils.utils import apply_library_argparse_overrides, finalize_library_config_overrides
+
+    recipe = apply_library_argparse_overrides(recipe, args)
     if cli_overrides:
-        logging.info("Applying %d CLI config override(s)", len(cli_overrides))
-        recipe = process_config_with_overrides(recipe, cli_overrides=cli_overrides)
+        recipe = _process_library_hydra_overrides(recipe, cli_overrides)
+    return finalize_library_config_overrides(recipe)
+
+
+def _apply_target_environment(recipe, base_env_vars: dict, cli_overrides: list[str], args) -> None:
+    """Finalize target-dependent recipe environment defaults."""
+    from utils.utils import apply_library_target_topology_environment, explicit_environment_override_names
 
-    recipe = finalize_library_config_overrides(recipe)
     protected_env_names = explicit_environment_override_names(cli_overrides, base_env_vars, recipe.env_vars)
     apply_library_target_topology_environment(
         recipe,
         gpu=args.gpu,
         protected_env_names=protected_env_names,
     )
+
+
+def _apply_recipe_environment(recipe) -> None:
+    """Project recipe env defaults onto the current process."""
+    from megatron.bridge.training.config import apply_environment_variables
+
     apply_environment_variables(recipe)
 
+
+def _bootstrap_recipe_environment(args, cli_overrides: list[str]) -> None:
+    """Apply the effective recipe env and replace this process with a clean interpreter."""
+    recipe = _get_library_recipe(args)
+    base_env_vars = dict(recipe.env_vars)
+    recipe = _apply_library_recipe_overrides(recipe, cli_overrides, args)
+    _apply_target_environment(recipe, base_env_vars, cli_overrides, args)
+    _apply_recipe_environment(recipe)
+
+    environment = dict(os.environ)
+    environment[ENV_BOOTSTRAP_MARKER] = str(os.getpid())
+    os.execvpe(sys.executable, [sys.executable, __file__, *sys.argv[1:]], environment)
+
+
+def _run_training(args, cli_overrides: list[str]) -> None:
+    """Build the final recipe and run training after environment bootstrap."""
+    import torch
+
+    from megatron.bridge.utils.common_utils import get_rank_safe
+
+    recipe = _get_library_recipe(args)
+    base_env_vars = dict(recipe.env_vars)
+    recipe = set_user_overrides(recipe, args)
+
+    if cli_overrides:
+        logging.info("Applying %d CLI config override(s)", len(cli_overrides))
+        recipe = _process_library_hydra_overrides(recipe, cli_overrides)
+
+    from utils.utils import finalize_library_config_overrides
+
+    recipe = finalize_library_config_overrides(recipe)
+    _apply_target_environment(recipe, base_env_vars, cli_overrides, args)
+    _apply_recipe_environment(recipe)
+
     if args.dryrun:
         save_path = args.save_config_filepath or "ConfigContainer.yaml"
         os.makedirs(os.path.dirname(os.path.abspath(save_path)), exist_ok=True)
@@ -291,5 +329,16 @@ def main():
         torch.distributed.destroy_process_group()
 
 
+def main() -> None:
+    """Bootstrap recipe env on the first pass and train in the self-exec process."""
+    parser = parse_cli_args()
+    args, cli_overrides = parser.parse_known_args()
+
+    if os.environ.get(ENV_BOOTSTRAP_MARKER) != str(os.getpid()):
+        _bootstrap_recipe_environment(args, cli_overrides)
+    else:
+        _run_training(args, cli_overrides)
+
+
 if __name__ == "__main__":
     main()

scripts/performance/run_script.py

@@ -136,10 +136,16 @@ def __init__(self) -> None:
         self.env_vars = os.environ
 
 
-def _bootstrap_recipe_environment(args, cli_overrides: list[str]) -> None:
-    """Install recipe env and re-exec this script in a clean interpreter."""
+def _apply_perf_recipe_overrides(recipe, cli_overrides: list[str], args):
+    """Apply the same CLI and argparse overrides in both self-exec passes."""
     from utils.overrides import set_cli_overrides, set_user_overrides
 
+    recipe = set_cli_overrides(recipe, cli_overrides)
+    return set_user_overrides(recipe, args)
+
+
+def _bootstrap_recipe_environment(args, cli_overrides: list[str]) -> None:
+    """Install recipe env and re-exec this script in a clean interpreter."""
     recipe = get_perf_recipe_for_environment(
         model_recipe_name=args.model_recipe_name,
         task=args.task,
@@ -148,8 +154,7 @@ def _bootstrap_recipe_environment(args, cli_overrides: list[str]) -> None:
         precision=args.compute_dtype,
         config_variant=args.config_variant,
     )
-    recipe = set_cli_overrides(recipe, cli_overrides)
-    recipe = set_user_overrides(recipe, args)
+    recipe = _apply_perf_recipe_overrides(recipe, cli_overrides, args)
     workload_base_config = _workload_base_config_from_recipe(recipe, num_gpus=args.num_gpus)
     plugin = PerfEnvPlugin(
         deterministic=args.deterministic,
@@ -166,7 +171,6 @@ def _bootstrap_recipe_environment(args, cli_overrides: list[str]) -> None:
 def _run_training(args, cli_overrides: list[str]) -> None:
     """Import the training stack after env bootstrap and run the workload."""
     import torch
-    from utils.overrides import set_cli_overrides, set_user_overrides
 
     from megatron.bridge.diffusion.models.wan.wan_step import WanForwardStep
     from megatron.bridge.models.qwen_vl.qwen3_vl_step import forward_step as qwen3_vl_forward_step
@@ -184,8 +188,7 @@ def _run_training(args, cli_overrides: list[str]) -> None:
         config_variant=getattr(args, "config_variant", None),
     )
 
-    recipe = set_cli_overrides(recipe, cli_overrides)
-    recipe = set_user_overrides(recipe, args)
+    recipe = _apply_perf_recipe_overrides(recipe, cli_overrides, args)
     apply_environment_variables(recipe)
 
     if args.dump_env:

scripts/performance/setup_experiment.py

@@ -66,7 +66,7 @@
 def _filter_run_script_args(argv: List[str]) -> List[str]:
     """Drop launcher-only args before forwarding argv to the rank-local script.
 
-    The launcher (this script) and the rank-local pre-exec wrapper share one
+    The launcher (this script) and the rank-local training scripts share one
     parser, but some args are meaningful only to the launcher and must not
     reach the rank-local scripts:
 

scripts/performance/utils/executors.py

@@ -238,8 +238,8 @@ def kubeflow_executor(
     """
     # K8s/Kubeflow jobs deliberately do NOT inherit the Slurm orchestration
     # defaults in PERF_ENV_VARS. Recipe-owned process settings are applied by
-    # the worker wrapper; the cluster supplies fabric tuning explicitly via
-    # KUBEFLOW_ENV_LIST_JSON (-> custom_env_vars / env_list).
+    # each training script's self-exec bootstrap; the cluster supplies fabric
+    # tuning explicitly via KUBEFLOW_ENV_LIST_JSON (-> custom_env_vars / env_list).
     env_vars: Dict[str, str] = {}
     if wandb_key is not None:
         env_vars["WANDB_API_KEY"] = wandb_key

scripts/performance/utils/utils.py

@@ -216,8 +216,8 @@ def apply_library_target_topology_environment(
 def apply_library_argparse_overrides(config: Any, args: Any) -> Any:
     """Apply argparse values that affect library recipe configuration.
 
-    Both the worker pre-exec wrapper and the final library runner call this
-    helper before Hydra overrides so their effective training configs match.
+    Both passes of the library runner's self-exec bootstrap call this helper
+    before Hydra overrides so their effective training configs match.
     """
     if getattr(args, "nccl_ub", False):
         config.ddp.nccl_ub = True

tests/functional_tests/test_groups/recipes/test_perf_config_integration.py

@@ -194,7 +194,7 @@ def test_workload_base_config_derives_from_flat_recipe(self):
             (
                 "llama",
                 "megatron.bridge.perf_recipes.llama.h100.llama3:llama3_8b_pretrain_8gpu_h100_fp8cs_config",
-                {"NVTE_FWD_LAYERNORM_SM_MARGIN": 16, "NCCL_CTA_POLICY": 1},
+                {"NVTE_FWD_LAYERNORM_SM_MARGIN": 20, "NCCL_CTA_POLICY": 1},
             ),
             (
                 "nemotronh",
@@ -224,7 +224,7 @@ def test_workload_base_config_derives_from_flat_recipe(self):
             (
                 "wan",
                 "megatron.bridge.perf_recipes.wan.h100.wan:wan_14b_pretrain_32gpu_h100_bf16_config",
-                {"NVTE_FWD_LAYERNORM_SM_MARGIN": 16, "CUDA_DEVICE_MAX_CONNECTIONS": 1},
+                {"NVTE_FWD_LAYERNORM_SM_MARGIN": 20, "CUDA_DEVICE_MAX_CONNECTIONS": 1},
             ),
         ],
     )