fix(run): exit after profile.async_test — donation has consumed state

param_decomp/configs.py

@@ -625,8 +625,9 @@ class ProfileConfig(BaseConfig):
 
     `mem_profile` (static + runtime memory analysis, then exits), `time_steps` (per-step wall
     breakdown), `trace` (perfetto trace over `trace_start`..`trace_start+trace_steps`),
-    `profile_max_events` (raise the perfetto GPU-activity event cap), `async_test`,
-    `leaf_bench`, `no_checkpoint` (skip ALL saves — throwaway profiling only).
+    `profile_max_events` (raise the perfetto GPU-activity event cap), `async_test`
+    (dispatch-asynchrony probe, then exits), `leaf_bench`, `no_checkpoint` (skip ALL
+    saves — throwaway profiling only).
     """
 
     mem_profile: bool = False
@@ -645,6 +646,14 @@ class ProfileConfig(BaseConfig):
     leaf_bench: bool = False
     no_checkpoint: bool = False
 
+    @model_validator(mode="after")
+    def validate_exclusive_exiting_hooks(self) -> Self:
+        assert not (self.async_test and self.mem_profile), (
+            "async_test and mem_profile each consume `state` via step_fn donation and exit "
+            "before the loop — run them one at a time"
+        )
+        return self
+
 
 class LaunchEnv(BaseConfig):
     """The process-environment surface a SLURM-launched rank runs with — the XLA *client*

param_decomp/run.py

@@ -599,6 +599,7 @@ def _bench_dispatch(tree: object, n: int = 15) -> float:
                 f"sync/host-bound => each call big)",
                 flush=True,
             )
+        os._exit(0)  # profiling-only path; donation has consumed `state`, so don't enter the loop
 
     if profile.mem_profile:
         _gib = 1024**3

param_decomp/tests/test_config.py

@@ -16,6 +16,7 @@
     ImportanceMinimalityLossConfig,
     PDConfig,
     PersistentPGDReconLossConfig,
+    ProfileConfig,
 )
 from param_decomp.recon import (
     build_loss_terms,
@@ -89,6 +90,15 @@ def test_legacy_top_level_n_mask_samples_pushes_onto_stochastic_terms():
     assert by_type["StochasticReconSubsetLoss"].n_mask_samples == 7  # pyright: ignore[reportAttributeAccessIssue]
 
 
+def test_profile_exiting_hooks_are_mutually_exclusive():
+    # Both hooks consume `state` via step_fn donation and os._exit(0) before the loop,
+    # so a config combining them would silently skip whichever runs second.
+    ProfileConfig(async_test=True)
+    ProfileConfig(mem_profile=True)
+    with pytest.raises(ValidationError, match="one at a time"):
+        ProfileConfig(async_test=True, mem_profile=True)
+
+
 def test_b128_config_converts():
     converted, raw = load_config(CONFIGS / "llama8b_l18_b128_cmp32.yaml", RUN_ID)
     assert raw["pd"]["batch_size"] == 128