diff --git a/param_decomp/configs.py b/param_decomp/configs.py index 455375223..826bbe9b9 100644 --- a/param_decomp/configs.py +++ b/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* diff --git a/param_decomp/run.py b/param_decomp/run.py index 17ebca0eb..43c0f78dd 100644 --- a/param_decomp/run.py +++ b/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 diff --git a/param_decomp/tests/test_config.py b/param_decomp/tests/test_config.py index 38a0519dd..3d396d0a2 100644 --- a/param_decomp/tests/test_config.py +++ b/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