Multi-concept Gram/cosine dynamics across OLMo post-training - #3
Open
xzAscC wants to merge 6 commits into
Open
Conversation
Download HumanEval-X, MiniF2F, BeyondX, MATH-500, Belebele, WinoGender, SST-2, and LLM-LAT one source at a time and pin shared item IDs. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
Cover 20 HumanEval-X pairs, 20 Belebele pairs, MiniF2F, BeyondX, MATH-500 CoT-to-direct, WinoGender she-to-he, SST-2, and LLM-LAT. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
Use the slide layer schedule, optional chat-template wrapping, batched forward passes, and quick-mode checkpoint caps. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
Default to six OLMo trajectories, the full 46-concept catalogue, and results/concept_dynamics_multi output paths. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
Plot per-checkpoint entanglement matrices and within-concept checkpoint stability, plus a summary CLI entrypoint. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
Describe dataset download, DiM extraction, cosine/Gram analysis, and plotting commands for the weekly multi-concept setup. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
![cubic-dev-ai[bot]](https://avatars.githubusercontent.com/in/1082092?s=60&v=4){kind=small}
There was a problem hiding this comment.
10 issues found across 18 files
Prompt for AI agents (unresolved issues)
Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.
<file name="src/dataset_store.py">
<violation number="1" location="src/dataset_store.py:109">
P2: Concurrent dataset downloads writing the same JSON can fail or lose a completed write because every call shares `<path>.tmp`. Use a unique temporary file per invocation before `os.replace` so the atomic-write helper is safe for parallel callers.</violation>
</file>
<file name="experiments/run_concept_dynamics.py">
<violation number="1" location="experiments/run_concept_dynamics.py:345">
P1: The reported HumanEval-X preflight can pass while the Python/C++ programs actually extracted were never validated, because validation uses the first task IDs but extraction uses pinned sampled IDs. Validate the exact pinned subset and its hashes (or make extraction select the same subset) before calling this preflight strict.
(Based on your team's feedback about subset-only preflight validation.) [FEEDBACK_USED]</violation>
</file>
<file name="experiments/plot_concept_dynamics.py">
<violation number="1" location="experiments/plot_concept_dynamics.py:71">
P3: These three unused private helpers are dead code and suggest validation that the CLI does not perform. Remove them until a caller needs them.</violation>
<violation number="2" location="experiments/plot_concept_dynamics.py:201">
P2: An unknown `--concepts` value silently renders every Gram concept instead of the requested subset. Validate requested names or make `_filter_concepts` reject an empty match so mistyped filters do not produce misleading plots.</violation>
<violation number="3" location="experiments/plot_concept_dynamics.py:231">
P3: Supplying both `--concept` and `--checkpoint` silently produces a stability plot and ignores the checkpoint selector. Reject mutually incompatible selector combinations so an accidental extra flag cannot generate the wrong visualization.</violation>
</file>
<file name="src/contrastive_datasets.py">
<violation number="1" location="src/contrastive_datasets.py:103">
P3: A misspelled duplicate seed constant was introduced (`SAMPlE_SEED` next to `SAMPLE_SEED`). Keeping two near-identical constants increases drift risk if one is updated later, so collapsing to one correctly named constant would keep sampling behavior unambiguous.</violation>
<violation number="2" location="src/contrastive_datasets.py:950">
P3: An unused private helper (`_get_loader`) was added but never called. Keeping dead dispatch helpers makes the loader path harder to audit and easier to desynchronize from the actual `_LOADERS` usage.</violation>
</file>
<file name="src/visualization.py">
<violation number="1" location="src/visualization.py:1229">
P2: Summary plots silently omit requested concepts after the first three. Limit to three only for the default selection so `--concepts` renders every supplied concept.</violation>
</file>
<file name="src/concept_dynamics.py">
<violation number="1" location="src/concept_dynamics.py:186">
P2: Batched extraction can select the wrong token position for left-padded tokenizers, which skews concept vectors for those models. This comes from using `attention_mask.sum(dim=1)-1`, which assumes right padding; deriving the last non-pad index from the right edge avoids padding-side dependence.</violation>
</file>
<file name="experiments/download_datasets.py">
<violation number="1" location="experiments/download_datasets.py:668">
P2: Dataset download failures are reported but still produce a successful process status, so shell/CI pipelines proceed to extraction with missing or stale datasets. Return nonzero after all downloaders have been attempted when `failures` is nonempty.</violation>
</file>
Tip: cubic can generate docs of your entire codebase and keep them up to date. Try it here.
Re-trigger cubic
| f"{args.humaneval_report_path} (n_required={n_samples})" | ||
| strict_code = _preflight_concepts_needing_strict_check(concepts) | ||
| if strict_code: | ||
| _run_strict_humaneval_preflight(args.humaneval_report_path, n_samples) |
There was a problem hiding this comment.
P1: The reported HumanEval-X preflight can pass while the Python/C++ programs actually extracted were never validated, because validation uses the first task IDs but extraction uses pinned sampled IDs. Validate the exact pinned subset and its hashes (or make extraction select the same subset) before calling this preflight strict.
(Based on your team's feedback about subset-only preflight validation.)
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At experiments/run_concept_dynamics.py, line 345:
<comment>The reported HumanEval-X preflight can pass while the Python/C++ programs actually extracted were never validated, because validation uses the first task IDs but extraction uses pinned sampled IDs. Validate the exact pinned subset and its hashes (or make extraction select the same subset) before calling this preflight strict.
(Based on your team's feedback about subset-only preflight validation.) </comment>
<file context>
@@ -250,30 +340,23 @@ def main(argv: list[str] | None = None):
- f"{args.humaneval_report_path} (n_required={n_samples})"
+ strict_code = _preflight_concepts_needing_strict_check(concepts)
+ if strict_code:
+ _run_strict_humaneval_preflight(args.humaneval_report_path, n_samples)
+
+ relaxed_code = _preflight_concepts_needing_warning(concepts)
</file context>
| """ | ||
| p = Path(path) | ||
| p.parent.mkdir(parents=True, exist_ok=True) | ||
| tmp = p.with_suffix(p.suffix + ".tmp") |
There was a problem hiding this comment.
P2: Concurrent dataset downloads writing the same JSON can fail or lose a completed write because every call shares <path>.tmp. Use a unique temporary file per invocation before os.replace so the atomic-write helper is safe for parallel callers.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At src/dataset_store.py, line 109:
<comment>Concurrent dataset downloads writing the same JSON can fail or lose a completed write because every call shares `<path>.tmp`. Use a unique temporary file per invocation before `os.replace` so the atomic-write helper is safe for parallel callers.</comment>
<file context>
@@ -0,0 +1,234 @@
+ """
+ p = Path(path)
+ p.parent.mkdir(parents=True, exist_ok=True)
+ tmp = p.with_suffix(p.suffix + ".tmp")
+ with tmp.open("w", encoding="utf-8") as handle:
+ json.dump(obj, handle, ensure_ascii=False, indent=2, sort_keys=False)
</file context>
| plot_concept_dynamics_summary( | ||
| args.input, | ||
| models=models, | ||
| concepts=concepts, |
There was a problem hiding this comment.
P2: An unknown --concepts value silently renders every Gram concept instead of the requested subset. Validate requested names or make _filter_concepts reject an empty match so mistyped filters do not produce misleading plots.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At experiments/plot_concept_dynamics.py, line 201:
<comment>An unknown `--concepts` value silently renders every Gram concept instead of the requested subset. Validate requested names or make `_filter_concepts` reject an empty match so mistyped filters do not produce misleading plots.</comment>
<file context>
@@ -0,0 +1,239 @@
+ plot_concept_dynamics_summary(
+ args.input,
+ models=models,
+ concepts=concepts,
+ )
+ except FileNotFoundError as exc:
</file context>
| summary_concepts.extend(c for c in model_block if c not in summary_concepts) | ||
| if concepts: | ||
| summary_concepts = [c for c in summary_concepts if c in set(concepts)] | ||
| summary_concepts = summary_concepts[:3] |
There was a problem hiding this comment.
P2: Summary plots silently omit requested concepts after the first three. Limit to three only for the default selection so --concepts renders every supplied concept.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At src/visualization.py, line 1229:
<comment>Summary plots silently omit requested concepts after the first three. Limit to three only for the default selection so `--concepts` renders every supplied concept.</comment>
<file context>
@@ -813,6 +939,319 @@ def generate_all_plots(results_dir: Optional[str] = None):
+ summary_concepts.extend(c for c in model_block if c not in summary_concepts)
+ if concepts:
+ summary_concepts = [c for c in summary_concepts if c in set(concepts)]
+ summary_concepts = summary_concepts[:3]
+
+ for model in target_models:
</file context>
Suggested change
| summary_concepts = summary_concepts[:3] | |
| if concepts is None: | |
| summary_concepts = summary_concepts[:3] |
| dtype=torch.long, | ||
| ) | ||
| else: | ||
| last_indices = attention_mask.to(dtype=torch.long).sum(dim=1) - 1 |
There was a problem hiding this comment.
P2: Batched extraction can select the wrong token position for left-padded tokenizers, which skews concept vectors for those models. This comes from using attention_mask.sum(dim=1)-1, which assumes right padding; deriving the last non-pad index from the right edge avoids padding-side dependence.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At src/concept_dynamics.py, line 186:
<comment>Batched extraction can select the wrong token position for left-padded tokenizers, which skews concept vectors for those models. This comes from using `attention_mask.sum(dim=1)-1`, which assumes right padding; deriving the last non-pad index from the right edge avoids padding-side dependence.</comment>
<file context>
@@ -132,45 +144,85 @@ def extract_layer_activations(
+ dtype=torch.long,
+ )
+ else:
+ last_indices = attention_mask.to(dtype=torch.long).sum(dim=1) - 1
+ batch_idx = torch.arange(
</file context>
Suggested change
| last_indices = attention_mask.to(dtype=torch.long).sum(dim=1) - 1 | |
| mask = attention_mask.to(dtype=torch.long) | |
| last_indices = mask.shape[1] - 1 - mask.flip(dims=[1]).argmax(dim=1) |
| else: | ||
| _log("All requested downloads completed.") | ||
| _log("=" * 60) | ||
| return 0 |
There was a problem hiding this comment.
P2: Dataset download failures are reported but still produce a successful process status, so shell/CI pipelines proceed to extraction with missing or stale datasets. Return nonzero after all downloaders have been attempted when failures is nonempty.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At experiments/download_datasets.py, line 668:
<comment>Dataset download failures are reported but still produce a successful process status, so shell/CI pipelines proceed to extraction with missing or stale datasets. Return nonzero after all downloaders have been attempted when `failures` is nonempty.</comment>
<file context>
@@ -0,0 +1,672 @@
+ else:
+ _log("All requested downloads completed.")
+ _log("=" * 60)
+ return 0
+
+
</file context>
| return sorted(seen) | ||
|
|
||
|
|
||
| def _list_checkpoints(gram: dict, model: str) -> list[str]: |
There was a problem hiding this comment.
P3: These three unused private helpers are dead code and suggest validation that the CLI does not perform. Remove them until a caller needs them.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At experiments/plot_concept_dynamics.py, line 71:
<comment>These three unused private helpers are dead code and suggest validation that the CLI does not perform. Remove them until a caller needs them.</comment>
<file context>
@@ -0,0 +1,239 @@
+ return sorted(seen)
+
+
+def _list_checkpoints(gram: dict, model: str) -> list[str]:
+ block = gram.get(model, {}) if isinstance(gram, dict) else {}
+ return sorted(k for k, v in block.items() if isinstance(v, dict) and v)
</file context>
| ) | ||
| return 1 | ||
|
|
||
| if args.concept: |
There was a problem hiding this comment.
P3: Supplying both --concept and --checkpoint silently produces a stability plot and ignores the checkpoint selector. Reject mutually incompatible selector combinations so an accidental extra flag cannot generate the wrong visualization.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At experiments/plot_concept_dynamics.py, line 231:
<comment>Supplying both `--concept` and `--checkpoint` silently produces a stability plot and ignores the checkpoint selector. Reject mutually incompatible selector combinations so an accidental extra flag cannot generate the wrong visualization.</comment>
<file context>
@@ -0,0 +1,239 @@
+ )
+ return 1
+
+ if args.concept:
+ return _plot_single_stability(args, stab_path)
+ if args.checkpoint:
</file context>
| _LOADERS[_concept_key] = _if_loader_for(_concept_key) | ||
|
|
||
|
|
||
| def _get_loader(concept: str) -> Callable[[int], list[tuple[str, str]]]: |
There was a problem hiding this comment.
P3: An unused private helper (_get_loader) was added but never called. Keeping dead dispatch helpers makes the loader path harder to audit and easier to desynchronize from the actual _LOADERS usage.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At src/contrastive_datasets.py, line 950:
<comment>An unused private helper (`_get_loader`) was added but never called. Keeping dead dispatch helpers makes the loader path harder to audit and easier to desynchronize from the actual `_LOADERS` usage.</comment>
<file context>
@@ -471,42 +841,170 @@ def load_winogender_pairs(n_samples: int = 50) -> list[tuple[str, str]]:
+ _LOADERS[_concept_key] = _if_loader_for(_concept_key)
+
+
+def _get_loader(concept: str) -> Callable[[int], list[tuple[str, str]]]:
+ resolved = _resolve_concept(concept)
+ return _LOADERS[resolved]
</file context>
Comment on lines
+103
to
+104
| SAMPlE_SEED: int = 42 # local convenience alias (kept private) | ||
| SAMPLE_SEED: int = 42 |
There was a problem hiding this comment.
P3: A misspelled duplicate seed constant was introduced (SAMPlE_SEED next to SAMPLE_SEED). Keeping two near-identical constants increases drift risk if one is updated later, so collapsing to one correctly named constant would keep sampling behavior unambiguous.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At src/contrastive_datasets.py, line 103:
<comment>A misspelled duplicate seed constant was introduced (`SAMPlE_SEED` next to `SAMPLE_SEED`). Keeping two near-identical constants increases drift risk if one is updated later, so collapsing to one correctly named constant would keep sampling behavior unambiguous.</comment>
<file context>
@@ -52,82 +55,234 @@
+
+DEFAULT_N_SAMPLES: int = 50
+N_SHARED_ITEMS: int = 50
+SAMPlE_SEED: int = 42 # local convenience alias (kept private)
+SAMPLE_SEED: int = 42
+
</file context>
Suggested change
| SAMPlE_SEED: int = 42 # local convenience alias (kept private) | |
| SAMPLE_SEED: int = 42 | |
| SAMPLE_SEED: int = 42 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
datasets/*.json(HumanEval-X 5-lang, MiniF2F, BeyondX, MATH-500 CoT→direct, Belebele 5-lang, WinoGender, SST-2, LLM-LAT) with pinned shared item IDs.[3,6,9,11,14,17,20,22,25,28]), default 6 trajectories (Think-SFT + 5 RL-Zero), chat-template wrapping, and batched activation extraction.Pipeline
uv run python experiments/download_datasets.py experiments/run_concept_dynamics.sh full # → results/concept_dynamics_multi uv run python experiments/plot_concept_dynamics.py --input results/concept_dynamics_multiTest plan
uv run pytest tests/test_contrastive_datasets.py tests/test_concept_dynamics.py tests/test_experiment_config.pyresults/concept_dynamics_multiSummary by cubic
Adds a multi-concept concept-dynamics pipeline across OLMo post-training with 46 paired concepts, 6 trajectories, slide-layer sampling, and Gram/stability heatmaps. Datasets are streamed and materialized to JSON for deterministic, offline runs.
New Features
datasets/*.jsonwith pinned shared IDs; offline loaders read JSONs.experiments/download_datasets.py,experiments/run_concept_dynamics.py/.sh,experiments/plot_concept_dynamics.py; docs updated.Migration
uv run python experiments/download_datasets.py(once) to createdatasets/*.json.experiments/run_concept_dynamics.sh full(orquick); outputs underresults/concept_dynamics_multi.code_python_vs_cpp,if_eng_vs_fra,gender_she_vs_he); legacy aliases remain supported.Written for commit 7571332. Summary will update on new commits.