Enhance run explorations

explorations/default.yaml

@@ -19,16 +19,17 @@ named_static_groups:
       use_rotary_embeddings: [true]
       use_abs_pos_embeddings: [false]
 
-  # Embedding Norm
-  - named_group: "rmsnorm_wte"
-    named_group_settings:
-      norm_variant_wte: ["rmsnorm"]
-
   # MLP Activation
   - named_group: "squared_relu"
     named_group_settings:
       activation_variant: ["squared_relu"]
 
+named_variation_groups:
+  # Embedding norm sweep. `default` omits the option entirely; YAML `null`
+  # preserves a Python None in the generated/logged config and omits the flag.
+  - named_group: "wte_norm"
+    norm_variant_wte: [default, "hyperspherenorm", null]
+
 common_group:
   dataset: ["minipile"]
   eval_interval: [1000]
@@ -41,5 +42,6 @@ parameter_groups:
       - "qk_norm"
       - "peri_ln"
       - "rotary"
-      - "rmsnorm_wte"
       - "squared_relu"
+    named_group_variations:
+      - "wte_norm"

optimization_and_search/run_experiments.py

@@ -130,6 +130,16 @@ def load_configurations(path: str, fmt: str) -> list[dict]:
     "distillation_source_run_name",
     "run_name_override",
 }
+DEFAULT_SENTINEL = "default"
+
+
+def _without_default_values(config: dict) -> dict:
+    """Drop parameters whose selected sweep value is the ``default`` sentinel."""
+    return {
+        key: value
+        for key, value in config.items()
+        if value != DEFAULT_SENTINEL
+    }
 
 
 def expand_range(val):
@@ -426,7 +436,8 @@ def _extract_common_group(cfg: dict) -> tuple[dict, set[str]]:
             raise ValueError(
                 "Values in 'common_group' cannot contain nested option structures"
             )
-        common[key] = normalized
+        if normalized != DEFAULT_SENTINEL:
+            common[key] = normalized
 
     return common, set(common)
 
@@ -520,7 +531,7 @@ def recurse(cfg: dict):
                 for key, value in metadata.items():
                     combo_dict[key] = deepcopy(value)
                 for final in _apply_conditionals(combo_dict, conditionals):
-                    yield final
+                    yield _without_default_values(final)
 
     for combo in recurse(cfg):
         merged = dict(common_values)
@@ -717,6 +728,11 @@ def build_command(combo: dict) -> list[str]:
     for k, v in combo.items():
         if k.startswith('_') or k in RESERVED_CONFIG_KEYS:
             continue
+        # A YAML null means that the Python setting should retain its None
+        # default. Omitting the CLI option is the only type-safe way to convey
+        # that through argparse (rather than passing the string "None").
+        if v is None:
+            continue
         if isinstance(v, bool):
             cmd.append(f"--{'' if v else 'no-'}{k}")
         elif isinstance(v, list):

requirements_cpu.txt

@@ -9,7 +9,7 @@ tiktoken==0.7.0
 torchinfo==1.8.0
 transformers==4.44.2
 pyspellchecker==0.8.3
-wandb==0.18.3
+wandb==0.19.10
 onnx==1.17.0
 beautifulsoup4==4.13.4
 dragonmapper==0.3.0

run_exploration_monitor.py

@@ -18,13 +18,15 @@
   E     - export CSV with prompt for custom name
   s     - save current layout
   p     - shows help menu
+  I     - view the associated exploration YAML file
   g     - graphs first two rows
   L     - graph & connect points sharing the 3rd column value
   1–9   - graph & connect points sharing merged columns 3..(2+N)
   q # # - multibarcharts - `q [1-9] [1-9]` - e.g. 'q 3 2' will create bar charts for columns 1 2 and 3, the next two columns (column 4 and column 5) as merged labels.
   z # # - Δ-bar chart (trim baseline) – e.g. ‘z 3 2’
   r–y   - barcharts with labels merged (r=1, y=3)
   c     - cycle colour-map for current column (high→low, low→high, off)
+  D     - remove colour-maps from all columns
   C # # - correlation + scatter for columns (1-based indexes, e.g. C 1 2)
   w     - toggle column width to fit largest visible cell
   u     - unsort / remove current column from the sort stack
@@ -46,8 +48,8 @@
 import yaml
 import math
 from textual.app import App, ComposeResult
-from textual.containers import Container
-from textual.widgets import DataTable, Footer, Header, Input, Label, Button
+from textual.containers import Container, VerticalScroll
+from textual.widgets import DataTable, Footer, Header, Input, Label, Button, Static
 from textual import events, on, work
 from textual.screen import Screen
 
@@ -84,12 +86,14 @@ def load_runs(log_file: Path) -> List[Dict]:
     "g: graph first two columns (matplotlib)\n"
     "g: graph first two columns (opens a Plotly window)\n"
     "p: shows help menu\n"
+    "I: view the associated exploration YAML file\n"
     "L: graph & connect points sharing the 3rd column value\n"
     "1–9: graph & connect points sharing merged columns 3..(2+N)\n"
     "q # #: multibarcharts - `q [1-9] [1-9]` - e.g. 'q 3 2' will create bar charts for columns 1 2 and 3, the next two columns (column 4 and column 5) as merged labels\n"
     "z # #: Δ-bar chart (trim baseline) – e.g. ‘z 3 2’\n"
     "r–y: barcharts with labels merged (r=1, y=3)\n"
     "c: cycle colour-map for current column (high→low, low→high, off)\n"
+    "D: remove colour-maps from all columns\n"
     "C # #: correlation + scatter (1-based indexes, e.g. C 1 2)\n"
     "w: toggle column width to fit largest visible cell\n"
     "u: unsort / remove current column from the sort stack\n"
@@ -126,10 +130,46 @@ def _cancel(self) -> None:                                  #  Cancel btn
         self.dismiss(None)
 
 
+class ExplorationConfigScreen(Screen):
+    """Read-only view of the YAML configuration associated with a run log."""
+
+    BINDINGS = [("escape", "dismiss", "Close")]
+
+    def __init__(self, config_file: Path) -> None:
+        super().__init__()
+        self.config_file = config_file
+
+    def compose(self) -> ComposeResult:
+        yield Label(f"Exploration YAML: {self.config_file}", id="config-title")
+        if self.config_file.exists():
+            contents = self.config_file.read_text()
+        else:
+            contents = "Associated exploration YAML file was not found."
+        with VerticalScroll(id="config-scroll"):
+            yield Static(contents, id="config-contents", markup=False)
+        yield Button("Close", id="close-config")
+
+    def action_dismiss(self) -> None:
+        self.dismiss()
+
+    @on(Button.Pressed, "#close-config")
+    def _close(self) -> None:
+        self.dismiss()
+
+
 class MonitorApp(App):
     CSS = """
     Screen { align: center middle; }
     Container { height: 1fr; }
+    #config-title { width: 90%; height: 1; text-style: bold; }
+    #config-scroll {
+        width: 90%;
+        height: 1fr;
+        border: round $accent;
+        padding: 1;
+    }
+    #config-contents { width: auto; }
+    #close-config { margin-top: 1; }
     DataTable#table {
         height: 1fr;
         width: 1fr;
@@ -141,6 +181,11 @@ class MonitorApp(App):
     def __init__(self, log_file: Path, interval: float, csv_dir: str) -> None:
         super().__init__()
         self.log_file = log_file
+        self.title = log_file.name
+        self.sub_title = str(log_file)
+        self.exploration_config_file = (
+            Path(__file__).resolve().parent / "explorations" / log_file.name
+        )
         self.interval = interval
         # Use JSON config file with same base name as YAML log file
         self.config_file = log_file.parent / f"{log_file.name}_monitor.json"
@@ -294,15 +339,41 @@ def _move_column_to_edge(
         edge_name = self.columns[0] if move_left else self.columns[-1]
         if col_name == edge_name:
             return
+        previous_columns = self.columns.copy()
         updated = [col for col in self.all_columns if col != col_name]
         edge_index = updated.index(edge_name)
         insert_index = edge_index if move_left else edge_index + 1
         updated.insert(insert_index, col_name)
         self.all_columns = updated
         self.columns = [col for col in self.all_columns if col not in self.hidden_cols]
+        self._remap_indexed_column_settings(previous_columns)
         new_cursor = self.columns.index(col_name) if move_cursor else col_index
         self.refresh_table(new_cursor=new_cursor)
 
+    def _remap_indexed_column_settings(self, previous_columns: List[str]) -> None:
+        """Keep colour and sort settings attached to columns after reordering."""
+        colour_by_name = {
+            previous_columns[index]: mode
+            for index, mode in self.colour_columns.items()
+            if 0 <= index < len(previous_columns)
+        }
+        sort_by_name = [
+            (previous_columns[index], ascending)
+            for index, ascending in self.sort_stack
+            if 0 <= index < len(previous_columns)
+        ]
+        new_indices = {name: index for index, name in enumerate(self.columns)}
+        self.colour_columns = {
+            new_indices[name]: mode
+            for name, mode in colour_by_name.items()
+            if name in new_indices
+        }
+        self.sort_stack = [
+            (new_indices[name], ascending)
+            for name, ascending in sort_by_name
+            if name in new_indices
+        ]
+
     def get_cell(self, entry: Dict, col_name: str):
         """Retrieve the value for a given column in an entry."""
         if col_name in entry:
@@ -779,6 +850,7 @@ async def on_key(self, event: events.Key) -> None:
             # Move column
             t = c - 1 if key == "h" else c + 1
             if 0 <= t < len(self.columns):
+                previous_columns = self.columns.copy()
                 n1, n2 = self.columns[c], self.columns[t]
                 i1, i2 = self.all_columns.index(n1), self.all_columns.index(n2)
                 self.all_columns[i1], self.all_columns[i2] = (
@@ -788,6 +860,7 @@ async def on_key(self, event: events.Key) -> None:
                 self.columns = [
                     col for col in self.all_columns if col not in self.hidden_cols
                 ]
+                self._remap_indexed_column_settings(previous_columns)
                 self.refresh_table(new_cursor=t)
         elif key in ("k", "j", "v", "m"):
             maxr = len(self.current_entries) - 1
@@ -856,6 +929,8 @@ async def on_key(self, event: events.Key) -> None:
             self.refresh_table(new_cursor=0)
         elif key == "p":
             self._msg(HOTKEYS_TEXT, timeout=10.0)
+        elif key == "I":
+            self.push_screen(ExplorationConfigScreen(self.exploration_config_file))
         elif key == "g":
             # ── Graph using first two visible columns: col[0] ⇒ Y, col[1] ⇒ X ──
             try:
@@ -891,6 +966,13 @@ async def on_key(self, event: events.Key) -> None:
                 self.colour_columns.pop(cur, None)
                 self._msg(f"Colour OFF for {self.columns[cur]}")
             self.refresh_table()
+        elif key == "D":
+            if self.colour_columns:
+                self.colour_columns.clear()
+                self.refresh_table(new_cursor=c)
+                self._msg("All column colours removed")
+            else:
+                self._msg("No active column colours")
         elif key == "w":
             col = self.columns[c]
             if col in self.auto_fit_columns:

tests/test_run_experiments.py

@@ -0,0 +1,53 @@
+import importlib.util
+from pathlib import Path
+import unittest
+
+import yaml
+
+
+MODULE_PATH = (
+    Path(__file__).parents[1] / "optimization_and_search" / "run_experiments.py"
+)
+SPEC = importlib.util.spec_from_file_location("run_experiments", MODULE_PATH)
+run_experiments = importlib.util.module_from_spec(SPEC)
+assert SPEC.loader is not None
+SPEC.loader.exec_module(run_experiments)
+
+
+class SweepDefaultAndNullTests(unittest.TestCase):
+    def test_unquoted_default_omits_parameter_and_null_preserves_none(self):
+        config = yaml.safe_load(
+            'norm_variant_wte: [default, "hyperspherenorm", null]'
+        )
+
+        combinations = [
+            combo for combo, _common_keys in run_experiments.generate_combinations(config)
+        ]
+
+        self.assertEqual(
+            combinations,
+            [
+                {},
+                {"norm_variant_wte": "hyperspherenorm"},
+                {"norm_variant_wte": None},
+            ],
+        )
+
+    def test_null_is_not_rendered_as_a_cli_string(self):
+        command = run_experiments.build_command(
+            {
+                "norm_variant_wte": None,
+                "activation_variant": "gelu",
+            }
+        )
+
+        self.assertEqual(command, ["python3", "train.py", "--activation_variant", "gelu"])
+
+    def test_default_in_common_group_is_omitted(self):
+        config = {"common_group": {"norm_variant_wte": ["default"]}}
+
+        self.assertEqual(list(run_experiments.generate_combinations(config)), [({}, set())])
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_run_exploration_monitor.py

@@ -0,0 +1,49 @@
+import unittest
+from pathlib import Path
+
+from run_exploration_monitor import ExplorationConfigScreen, MonitorApp
+
+
+class ColumnSettingRemapTests(unittest.TestCase):
+    def test_title_and_exploration_config_use_log_yaml_name(self):
+        app = MonitorApp(
+            log_file=Path("exploration_logs/default.yaml"),
+            interval=30.0,
+            csv_dir="rem_csv_exports",
+        )
+
+        self.assertEqual(app.title, "default.yaml")
+        self.assertEqual(app.sub_title, "exploration_logs/default.yaml")
+        self.assertEqual(app.exploration_config_file.name, "default.yaml")
+        self.assertEqual(app.exploration_config_file.parent.name, "explorations")
+
+    def test_colour_and_sort_settings_follow_reordered_columns(self):
+        app = object.__new__(MonitorApp)
+        app.columns = ["gamma", "alpha", "beta"]
+        app.colour_columns = {0: "high_low", 2: "low_high"}
+        app.sort_stack = [(1, True), (0, False)]
+
+        app._remap_indexed_column_settings(["alpha", "beta", "gamma"])
+
+        self.assertEqual(app.colour_columns, {1: "high_low", 0: "low_high"})
+        self.assertEqual(app.sort_stack, [(2, True), (1, False)])
+
+
+class ExplorationConfigScreenTests(unittest.IsolatedAsyncioTestCase):
+    async def test_hotkey_opens_associated_yaml_contents(self):
+        app = MonitorApp(
+            log_file=Path("exploration_logs/default.yaml"),
+            interval=3600.0,
+            csv_dir="rem_csv_exports",
+        )
+
+        async with app.run_test() as pilot:
+            await pilot.press("I")
+
+            self.assertIsInstance(app.screen, ExplorationConfigScreen)
+            contents = app.screen.query_one("#config-contents").content
+            self.assertIn("norm_variant_wte", str(contents))
+
+
+if __name__ == "__main__":
+    unittest.main()