From a139a3ed7619362ba0ac8af5042315b39403e314 Mon Sep 17 00:00:00 2001 From: "PD User (shared)" Date: Tue, 14 Jul 2026 04:24:01 +0000 Subject: [PATCH] =?UTF-8?q?feat(campaign-view):=20pd-campaign-view=20?= =?UTF-8?q?=E2=80=94=20the=20standard=20campaign=20workspace=20view,=20as?= =?UTF-8?q?=20code?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A campaign (one wandb group + a saved workspace view) has hand-built its view per campaign since the torch-era create_workspace_view died in the JAX migration (the wandb-workspaces dep survived, unused). pd-campaign-view saves the standard view — referee / recon quality / train losses / health / schedules sections over the canonical metric keys, runset-filtered on the group — so the layout is versioned and evolves by PR instead of UI fiddling. Verified end-to-end against the smooth-l0-matrix-100k campaign: https://wandb.ai/goodfire/param-decomp?nw=kb666b2nyzk Crew-Address: task/tidy-wandb-logging-structure Co-Authored-By: Claude Fable 5 --- CLAUDE.md | 1 + param_decomp_lab/experiments/CLAUDE.md | 5 + param_decomp_lab/experiments/campaign_view.py | 152 ++++++++++++++++++ param_decomp_lab/pyproject.toml | 2 + 4 files changed, 160 insertions(+) create mode 100644 param_decomp_lab/experiments/campaign_view.py diff --git a/CLAUDE.md b/CLAUDE.md index fae0420fb..d5e8b7b84 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -253,6 +253,7 @@ via `pd-lm`. Slow/plot eval is in-loop only (no CLI). | `pd-lm` | `experiments/lm/launch.py` | Launch a decomposition trainer run; config-driven via `runtime.dp` (`dp=N` → snapshot + pinned launch config + sbatch across `N//8` nodes, per-node job-side venv; `dp=null` → inline) | | `pd-pretrain` | `experiments/lm/pretrain/launch.py` | Launch a pretrainer run; config-driven via `dp` (`dp=N` → sbatch; `dp=null` → inline) | | `pd-tms` / `pd-resid-mlp` | `experiments/{tms,resid_mlp}/run.py` | The CPU toy decomposition CLIs | +| `pd-campaign-view` | `experiments/campaign_view.py` | Save the standard wandb workspace view for a campaign group | | `pd-harvest` | `harvest/scripts/run_slurm_cli.py` | Submit harvest SLURM job | | `pd-autointerp` | `autointerp/scripts/run_slurm_cli.py` | Submit autointerp SLURM job | | `pd-clustering` / `pd-cluster-merge` / `pd-cluster-distances` | `clustering/scripts/` | Clustering ensemble / merge / consensus distances | diff --git a/param_decomp_lab/experiments/CLAUDE.md b/param_decomp_lab/experiments/CLAUDE.md index 246e41758..f6b27cb89 100644 --- a/param_decomp_lab/experiments/CLAUDE.md +++ b/param_decomp_lab/experiments/CLAUDE.md @@ -130,3 +130,8 @@ Every `pd-*` run command accepts `--group ` and `--tags a,b,c` (no-ops when - **`--group`** sets wandb's first-class `group` field — used by the UI's native collapsing and matched by workspace filters via `ws.Metric("Group")`. - **`--tags`** adds wandb tags — orthogonal to `group`, many per run, user-defined. + +A campaign (one group + its saved workspace view) gets its view from +`pd-campaign-view [--project ...]` (`experiments/campaign_view.py`): the +standard section layout (referee / recon quality / train losses / health / schedules), +coded in the repo so it evolves by PR instead of per-campaign UI hand-building. diff --git a/param_decomp_lab/experiments/campaign_view.py b/param_decomp_lab/experiments/campaign_view.py new file mode 100644 index 000000000..8894a7bdf --- /dev/null +++ b/param_decomp_lab/experiments/campaign_view.py @@ -0,0 +1,152 @@ +"""`pd-campaign-view ` — save the standard wandb workspace view for a campaign. + +A campaign is a coordinated multi-run matrix sharing one wandb `group` (set at launch +via `--group`). Each campaign used to hand-build its saved workspace view; this CLI +codes the standard instead, so the view layout evolves by PR. It creates a saved view +named `campaign ` in the target project, runset-filtered to the group, with +sections mirroring the canonical metric namespaces (`param_decomp.run` documents the +tree). + +Panels list the config-stable keys; run-shaped families (per-site `eval/l0/*`, per-site +grad-norm leaves, campaign-specific recon terms beyond the common set) stay in the +UI's hidden-panels pool — pull the ones a campaign cares about into its sections by +hand after seeding. Re-running always creates a NEW view (the wandb API exposes no +list-views surface to update by name); delete stale ones in the UI. +""" + +import fire +import wandb_workspaces.reports.v2 as wr +import wandb_workspaces.workspaces as ws + +from param_decomp_lab.infra.wandb import DEFAULT_WANDB_PROJECT, get_wandb_entity + +# Masked-forward variants of the fast CE/KL eval (`param_decomp.eval`); `zero_masked` +# has no `ce_difference_*` key. +_CE_KL_VARIANTS = ( + "ci_masked", + "stoch_masked", + "unmasked", + "random_masked", + "rounded_masked", + "zero_masked", +) + +_REFEREE_CONVENTION = ( + "**Reading the referee**: `eval/loss/PGDReconLoss` is noisy per-eval — compare arms " + "on last-5/10-eval **medians**, and prefer mid-window reads (e.g. step 35k of 40k) " + "over endpoints." +) + + +def _plot(title: str, keys: tuple[str, ...], log_y: bool = False) -> wr.LinePlot: + return wr.LinePlot(title=title, x="Step", y=list(keys), log_y=log_y or None) + + +def _campaign_sections() -> list[ws.Section]: + referee = ws.Section( + name="Referee", + panels=[ + _plot("PGD recon (the referee)", ("eval/loss/PGDReconLoss",), log_y=True), + _plot("KL (CI-masked)", ("eval/ce_kl/kl_ci_masked",), log_y=True), + _plot("CE difference (CI-masked)", ("eval/ce_kl/ce_difference_ci_masked",)), + wr.MarkdownPanel(markdown=_REFEREE_CONVENTION), + ], + is_open=True, + ) + recon_quality = ws.Section( + name="Recon quality (ce_kl)", + panels=[ + _plot( + "KL by masking variant", + tuple(f"eval/ce_kl/kl_{v}" for v in _CE_KL_VARIANTS), + log_y=True, + ), + _plot( + "CE difference by masking variant", + tuple( + f"eval/ce_kl/ce_difference_{v}" for v in _CE_KL_VARIANTS if v != "zero_masked" + ), + ), + ], + is_open=True, + ) + train_losses = ws.Section( + name="Train losses", + panels=[ + _plot("total", ("train/loss/total",), log_y=True), + _plot( + "faith + imp-min", + ( + "train/loss/FaithfulnessLoss", + "train/loss/ImportanceMinimalityLoss", + "train/loss/SmoothL0ImportanceMinimalityLoss", + "train/loss/FrequencyMinimalityLoss", + ), + log_y=True, + ), + _plot( + "recon terms (common set)", + ( + "train/loss/ChunkwiseSubsetReconLoss", + "train/loss/PersistentPGDReconLoss", + "train/loss/StochasticReconSubsetLoss", + "train/loss/PGDReconLoss", + ), + log_y=True, + ), + ], + ) + health = ws.Section( + name="Health", + panels=[ + _plot( + "grad-norm summaries (median)", + ( + "train/grad_norms/summary/components/median", + "train/grad_norms/summary/ci_fns/median", + "train/grad_norms/summary/total/median", + ), + log_y=True, + ), + _plot("peak memory (GB/rank)", ("train/mem/peak_gb_per_rank",)), + _plot("step time (s)", ("train/perf/step_time_s",)), + ], + ) + schedules = ws.Section( + name="Schedules", + panels=[ + _plot( + "learning rates", + ( + "train/schedules/lr/components", + "train/schedules/lr/ci_fn", + "train/schedules/lr/src", + ), + log_y=True, + ), + _plot("imp-min anneal", ("train/schedules/p_imp", "train/schedules/gamma_imp")), + ], + ) + return [referee, recon_quality, train_losses, health, schedules] + + +def make_campaign_view(group: str, project: str, entity: str) -> ws.Workspace: + """Pure builder: the standard campaign workspace, runset-filtered to `group`.""" + return ws.Workspace( + entity=entity, + project=project, + name=f"campaign {group}", + sections=_campaign_sections(), + settings=ws.WorkspaceSettings(max_runs=50), + runset_settings=ws.RunsetSettings(filters=[ws.Metric("Group") == group]), + ) + + +def main(group: str, project: str = DEFAULT_WANDB_PROJECT, entity: str | None = None) -> None: + workspace = make_campaign_view(group, project, entity or get_wandb_entity()) + workspace.save_as_new_view() + print(f"saved campaign view for group {group!r}: {workspace.url}", flush=True) + + +def cli() -> None: + fire.Fire(main) diff --git a/param_decomp_lab/pyproject.toml b/param_decomp_lab/pyproject.toml index cd16cd26d..182b5a81e 100644 --- a/param_decomp_lab/pyproject.toml +++ b/param_decomp_lab/pyproject.toml @@ -45,6 +45,8 @@ pd-resid-mlp = "param_decomp_lab.experiments.resid_mlp.run:cli" # `python -m pretrain.train` (or `--local`). The core pretrain trainer lives at the root # sibling `pretrain/`. pd-pretrain = "param_decomp_lab.experiments.lm.pretrain.launch:cli" +# Save the standard wandb workspace view for a campaign group (see experiments/campaign_view.py) +pd-campaign-view = "param_decomp_lab.experiments.campaign_view:cli" # Post-processing pipeline CLIs (SLURM-driven; lab subtrees). pd-clustering = "param_decomp_lab.clustering.scripts.run_pipeline:cli" pd-cluster-merge = "param_decomp_lab.clustering.scripts.run_merge:cli"