feat(2pool): wire torch profiler into the 2-pool run path

two_pool_run never passed a profiler to trainer.run (only the 3-pool did), so
PD_TORCH_PROFILE_* was a no-op for 2-pool runs. Broaden _maybe_build_torch_profiler
to accept either trainer (it only reads ctx.kind / ctx.role.rank) and pass it on
both the fresh and resume run calls. Drop the now-unused ThreePoolRuntimeConfig
import (TwoPoolRuntimeConfig subclasses PooledRuntimeConfig directly).

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

Author: ocg-goodfire

param_decomp_lab/experiments/lm/three_pool_run.py

@@ -96,6 +96,7 @@
 from param_decomp_lab.three_pool.config import PooledRuntimeConfig
 from param_decomp_lab.three_pool.consolidate import SNAPSHOT_SCRATCH_DIRNAME, ppgd_shard_dirname
 from param_decomp_lab.three_pool.pd_config import ThreePoolConstrainedPDConfig
+from param_decomp_lab.three_pool.two_pool_optimize import TwoPoolTrainer
 
 # Cross-pool NCCL p2p deadlock guard. An asymmetric topology (fanout>1 on a cross-pool
 # edge) at long sequence lengths wedges when a rendezvous-blocked NCCL send/recv kernel
@@ -291,14 +292,17 @@ def _excepthook(exctype: type[BaseException], value: BaseException, tb: Any) ->
     sys.excepthook = _excepthook
 
 
-def _maybe_build_torch_profiler(trainer: ThreePoolTrainer) -> "torch.profiler.profile | None":
+def _maybe_build_torch_profiler(
+    trainer: ThreePoolTrainer | TwoPoolTrainer,
+) -> "torch.profiler.profile | None":
     """Opt-in `torch.profiler.profile` for the listed ranks (typically one per pool).
 
     Env: `PD_TORCH_PROFILE_RANKS=0,96,100` + `PD_TORCH_PROFILE_OUT=/abs/dir`, plus
     schedule / instrumentation knobs (`PD_TORCH_PROFILE_SKIP_FIRST` default 20,
     `_ACTIVE` default 3, `_MEMORY` default on, `_STACK` / `_MODULES` / `_SHAPES` off).
     Returns the profile context (caller passes it to `trainer.run(profiler=...)`) or
     `None` if this rank isn't profiled. Verified at 104 ranks (gpt2-xl) 2026-05-28.
+    Works for both the 3-pool and 2-pool trainers (only reads `ctx.kind` / `ctx.role.rank`).
     """
     prof_ranks_env = os.environ.get("PD_TORCH_PROFILE_RANKS", "").strip()
     if not prof_ranks_env:

param_decomp_lab/experiments/lm/two_pool_run.py

@@ -53,9 +53,9 @@
 )
 from param_decomp_lab.experiments.lm.three_pool_run import (
     THREE_POOL_SLURM_ENV,
-    ThreePoolRuntimeConfig,
     _agree_on_run_id,
     _install_first_fail_marker,
+    _maybe_build_torch_profiler,
     _maybe_enable_memory_profile,
 )
 from param_decomp_lab.experiments.utils import (
@@ -84,17 +84,15 @@
 )
 from param_decomp_lab.run_sink import ThreePoolSink
 from param_decomp_lab.seed import set_seed
+from param_decomp_lab.three_pool.config import PooledRuntimeConfig
 from param_decomp_lab.three_pool.consolidate import SNAPSHOT_SCRATCH_DIRNAME, ppgd_shard_dirname
 from param_decomp_lab.three_pool.pd_config import ThreePoolConstrainedPDConfig
 from param_decomp_lab.three_pool.two_pool_config import TwoPoolTopology
 from param_decomp_lab.three_pool.two_pool_optimize import TwoPoolTrainer
 
 
-class TwoPoolRuntimeConfig(ThreePoolRuntimeConfig):
-    """Core's substrate scalars + a 2-pool ``topology``. Subclasses
-    ``ThreePoolRuntimeConfig`` only to narrow the topology field's type."""
-
-    topology: TwoPoolTopology  # pyright: ignore[reportIncompatibleVariableOverride]
+class TwoPoolRuntimeConfig(PooledRuntimeConfig):
+    topology: TwoPoolTopology
 
 
 class TwoPoolLMExperimentConfig(BaseConfig):
@@ -221,12 +219,13 @@ def main(
         )
         return
 
-    if resume is not None:
-        assert config_path is None, "pass either config_path or --resume, not both"
-        _resume_main(Path(resume), group=group, tags=tags, run_id=run_id)
-    else:
-        assert config_path is not None, "must provide either config_path or --resume"
-        _fresh_or_requeue_main(Path(config_path), group=group, tags=tags, run_id=run_id)
+    match (resume, config_path):
+        case (None, str(config_path)):
+            _fresh_or_requeue_main(Path(config_path), group=group, tags=tags, run_id=run_id)
+        case (str(resume), None):
+            _resume_main(Path(resume), group=group, tags=tags, run_id=run_id)
+        case _:
+            raise ValueError("must provide either config_path or --resume")
 
 
 def _fresh_or_requeue_main(
@@ -318,7 +317,14 @@ def _fresh_main(
             runtime_config=cfg.runtime,
             two_pool_config=cfg.runtime.topology,
         )
-        trainer.run(train_loader, sink, cfg.cadence, scratch_dir=scratch_dir, eval_loop=eval_loop)
+        trainer.run(
+            train_loader,
+            sink,
+            cfg.cadence,
+            scratch_dir=scratch_dir,
+            eval_loop=eval_loop,
+            profiler=_maybe_build_torch_profiler(trainer),
+        )
     finally:
         sink.finish()
 
@@ -449,7 +455,12 @@ def _run_resume(
             ppgd_shard_dir=from_run / ppgd_shard_dirname(resolved_step),
         )
         trainer.run(
-            train_loader, sink, effective_cfg.cadence, scratch_dir=scratch_dir, eval_loop=eval_loop
+            train_loader,
+            sink,
+            effective_cfg.cadence,
+            scratch_dir=scratch_dir,
+            eval_loop=eval_loop,
+            profiler=_maybe_build_torch_profiler(trainer),
         )
     finally:
         sink.finish()