feat(multipool): forward profiling/trace env vars to the SLURM job

The launch exports only a curated env dict (DDP_ENV + THREE_POOL_SLURM_ENV), so
PD_TORCH_PROFILE_* / PD_MEMORY_PROFILE_* / PD_TRACE* set at submit were dropped and
the documented profiler env-var interface was a no-op over SLURM. profiling_env_passthrough()
collects those from the submit environment; both 2-pool submit sites merge it in.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: ocg-goodfire

param_decomp_lab/experiments/lm/three_pool_run.py

@@ -107,6 +107,16 @@
 # cross-pool collectives).
 THREE_POOL_SLURM_ENV: dict[str, str] = {"PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True"}
 
+# Profiling / tracing env vars (read on the compute node by the profiler / trace hooks)
+# only reach the job if forwarded explicitly — the SLURM script exports a curated env dict,
+# not the submitter's whole environment. `profiling_env_passthrough` collects those set at
+# submit so `PD_TORCH_PROFILE_* / PD_MEMORY_PROFILE_* / PD_TRACE*` flow through to the run.
+_PROFILE_ENV_PREFIXES = ("PD_TORCH_PROFILE_", "PD_MEMORY_PROFILE_", "PD_TRACE")
+
+
+def profiling_env_passthrough() -> dict[str, str]:
+    return {k: v for k, v in os.environ.items() if k.startswith(_PROFILE_ENV_PREFIXES)}
+
 
 class ThreePoolRuntimeConfig(PooledRuntimeConfig):
     """`PooledRuntimeConfig` (core substrate scalars + the `compile_*` / `checkpoint_ci_fn`

param_decomp_lab/experiments/lm/two_pool_run.py

@@ -57,6 +57,7 @@
     _install_first_fail_marker,
     _maybe_build_torch_profiler,
     _maybe_enable_memory_profile,
+    profiling_env_passthrough,
 )
 from param_decomp_lab.experiments.utils import (
     EXPERIMENT_CONFIG_FILENAME,
@@ -525,7 +526,9 @@ def _submit_slurm(
         requeue=True,
     )
     script = generate_script(
-        slurm_config, launch.command, env={**launch.env, **THREE_POOL_SLURM_ENV}
+        slurm_config,
+        launch.command,
+        env={**launch.env, **THREE_POOL_SLURM_ENV, **profiling_env_passthrough()},
     )
     result = submit_slurm_job(script, "lm")
 
@@ -631,7 +634,9 @@ def submit_slurm_async_consolidate_and_eval(
         comment=f"async-consol-eval:{train_run_id}@{step}",
     )
     script = generate_script(
-        slurm_config, launch.command, env={**launch.env, **THREE_POOL_SLURM_ENV}
+        slurm_config,
+        launch.command,
+        env={**launch.env, **THREE_POOL_SLURM_ENV, **profiling_env_passthrough()},
     )
     result = submit_slurm_job(script, "lm")
     logger.info(