Skip to content

SPEC.md

Latest commit

 

History

History
283 lines (242 loc) · 16.6 KB

SPEC.md

File metadata and controls

283 lines (242 loc) · 16.6 KB
Raw

u-moec-tiny GUI Overhaul — Specification Document

Project Name

u-moec-tiny GUI: Full Training & Evaluation Control Interface

Description

Overhaul the existing Dash-based GUI from a stub-level dashboard into the primary interface for configuring, launching, monitoring, and evaluating the entire u-moec-tiny neural compression pipeline. The GUI must allow users to edit all YAML config parameters, control training (start/stop/resume), monitor live training metrics, inspect MoE routing and MoD gating, visualize activations and hidden-state flow, run compression/decompression jobs, and execute benchmarks with roundtrip verification.

Technology Stack

  • Python 3.10+ (matching project requirement)
  • Dash 2.18+ + dash-bootstrap-components 1.6+ (existing framework)
  • Plotly 5.20+ (existing charting)
  • Pandas 2.0+ (existing data layer)
  • SQLite 3 (existing metrics DB, with new tables)
  • PyYAML 6.0+ (config read/write)

Workstreams

WS1: Full Config Editor (Sidebar Overhaul)

Target file: gui/panels/config_panel.py

Replace the minimal 4-field config sidebar with a comprehensive form covering every section of the YAML config. Use dbc.Accordion with the following sections:

  1. Model — d_model (number), d_state (number), d_conv (number), n_blocks (number), block_pattern (text display, read from YAML), context_len (number)
  2. MoE — n_routed (number), n_shared (number), top_k (number), d_ff (number), bias_lr (number)
  3. MoD — enabled (switch), capacity (slider 0-1), gate_kind (dropdown: scalar)
  4. Early Exit — enabled (switch), blocks (text input "4" or "4,7"), exit_threshold (number 0-1), exit_loss_weight (number)
  5. Cascade — enabled (switch), threshold (number), d_hidden (number)
  6. Conditioning — ada_ln_zero (switch), ada_ln_blocks (text "0,1"), cross_attn (switch)
  7. Patches — primary (number), scales (text "16" or "4,16,64"), multi_token_predict (switch), mtp_weight (number)
  8. Head — d_globals (number), n_early_exit_heads (number)
  9. Training — steps (number), batch_size (number), seq_len (number), grad_clip (number), warmup_steps (number), cosine_schedule (switch), precision (dropdown: bf16, bf16+fp8)
  10. Optimizer Muon — lr (number), momentum (number), ns_steps (number)
  11. Optimizer AdamW — lr (number), betas (text "0.9,0.95"), weight_decay (number)
  12. Loss — ce_weight (number), mtp_weight (number)
  13. Data — pack_documents (switch), seed (number)
  14. LoRA — enabled (switch), rank (number), alpha (number), target (dropdown: routed/all), trigger_bytes (number), adapt_steps (number), adapt_lr (number)
  15. Entropy — coder (dropdown: rans/tans), block_size (number), precision (number), prob_precision (number), probability_sharpening (switch)
  16. Logging — log_every (number), eval_every (number), save_every (number)
  17. Hardware — gpus (number), compile (switch)

Additional controls:

  • Tier dropdown (existing, keep at top)
  • Component checklist (existing, keep for quick toggles)
  • "Load Config" button: opens a text input for a YAML file path, reads it, and populates all form fields
  • "Save Config" button: writes current form state to a user-specified YAML file path
  • "Reset to Default" button: reloads the selected tier's default YAML
  • All form fields should have IDs like config-model-d-model, config-moe-n-routed, etc.
  • The existing config-start-btn, config-stop-btn, and config-status must remain

WS2: Training Panel Upgrade

Target files: gui/panels/training_panel.py, gui/db.py

  • Wire training_panel.kpis() into the refresh callback. Add a dcc.Graph or html.Div with id="training-kpi-cards" to the layout, and in the refresh callback compute the KPIs from the latest metrics row.
  • Add ETA display: html.Div(id="training-eta") showing estimated time remaining based on (target_steps - current_step) / throughput
  • Add checkpoint browser: html.Div(id="checkpoint-list") — scan the checkpoint directory for step*.pt files, display as a list with step number and file size
  • Add resume-from-checkpoint: dcc.Dropdown(id="resume-checkpoint") + button id="resume-btn"
  • Add live log tail: html.Pre(id="training-log-tail", style={"maxHeight": "300px", "overflowY": "auto", "fontSize": "11px"}) — read last 100 lines of training.log each refresh cycle
  • Add per-head loss table: html.Div(id="training-head-losses") — small table showing current loss for final, exit_4, etc.

WS3: Routing Panel — Fix Missing Callbacks

Target files: gui/panels/routing_panel.py, gui/viz_utils.py, gui/db.py

  • Populate routing-bias-history graph with time-series data from db.moe_routing_history(). Create viz_utils.moe_routing_history_fig() that plots expert load over time for all experts (line chart).
  • Add load-balance metrics: compute Gini coefficient = 1 - 2 * sum(cumsum(sorted(load)) / sum(load)) / n and max/min ratio for each logged step, plot over time. Create viz_utils.load_balance_fig().
  • Add click-to-drill-down on heatmap: use clickData from routing-heatmap to show a single expert's load over time. Add a new graph id="routing-expert-detail".

WS4: MoD Panel — History + Metrics

Target files: gui/panels/mod_panel.py, gui/viz_utils.py

  • Replace "TBD" mod-history with a proper time-series. Create viz_utils.mod_gate_history_fig() that reads from db.mod_gates_history() and plots fraction_processed per block over training steps.
  • Add capacity utilization line: overlay target capacity as a dashed horizontal line.
  • Add metric cards: average fraction processed, deviation from capacity.

WS5: Activations + Brain Panels — Real Data Flow

Target files: bin/train.py, bin/simulate_training.py, gui/db.py, gui/panels/activations_panel.py, gui/panels/brain_panel.py

New DB tables:

activation_stats(step INTEGER, block_idx INTEGER, mean REAL, std REAL,
                 p1 REAL, p25 REAL, p50 REAL, p75 REAL, p99 REAL,
                 PRIMARY KEY (step, block_idx));
hidden_states(step INTEGER, block_idx INTEGER, sample_idx INTEGER,
              pc1 REAL, pc2 REAL,
              PRIMARY KEY (step, block_idx, sample_idx));

Trainer changes (bin/train.py):

  • Add insert_activation_stat() and insert_hidden_state() to TrainingDB
  • Every eval_every steps, on a single batch with torch.no_grad():
    • Run a forward pass
    • For each block, detach the hidden state, compute per-block activation stats (mean, std, percentiles) on CPU, write to activation_stats
    • For a sample of 200 tokens per block (or all if fewer), PCA-project the hidden states to 2D (use SVD of centered data), write to hidden_states

Simulate training changes (bin/simulate_training.py):

  • Every log_every steps, generate synthetic activation_stats rows with realistic distributions (mean ~1.0, std ~0.3, with slight variation per block)
  • Every log_every * 5 steps, generate synthetic hidden_states rows (2D Gaussian clusters per block, with block 0 being a single cluster and later blocks separating into sub-clusters)

DB queries (gui/db.py):

  • activation_stats_at(conn, step=None) — returns DataFrame of activation stats for latest or given step
  • hidden_states_at(conn, step=None) — returns DataFrame of hidden states for latest or given step

Panel wiring:

  • Activations panel: add callback on refresh-interval that reads activation_stats_at(), feeds to activation_dist_fig() and activation_heatmap_fig()
  • Brain panel: add callback on refresh-interval that reads hidden_states_at(), feeds to pca_fig()

WS6: Compression + Evaluation Panels

Target files: gui/panels/compression_panel.py, new gui/panels/eval_panel.py, gui/db.py, gui/viz_utils.py

Compression Panel:

  • File path input: dcc.Input(id="compress-file-path", type="text", placeholder="/path/to/file.bin")
  • File upload: dcc.Upload(id="compress-file-upload") — stores uploaded file in /tmp/u-moec-uploads/
  • Profile selector: dcc.Dropdown(id="compress-profile", options=["turbo", "fast", "balanced"], value="turbo")
  • Checkpoint selector: dcc.Dropdown(id="compress-checkpoint") — populated from checkpoint dir scan
  • "Compress" button: html.Button("Compress", id="compress-btn") — runs python -m bin.compress {input} {output}.umc --config {config} --checkpoint {ckpt} --profile {profile} via subprocess
  • Results table: html.Table(id="compression-results-table") with columns: File, Raw Size, Compressed Size, Ratio, Time
  • Comparison row: run gzip -c and zstd -22 on the same file (if available), show their sizes
  • "Decompress & Verify" button: html.Button("Verify Roundtrip", id="compress-verify-btn") — runs decompress, then sha256sum compare
  • Roundtrip result: html.Div(id="compress-verify-result") showing "PASS" (green) or "FAIL" (red)

New compression_jobs DB table:

compression_jobs(job_id INTEGER PRIMARY KEY AUTOINCREMENT,
                 ts REAL, input_path TEXT, output_path TEXT,
                 profile TEXT, checkpoint TEXT, status TEXT,
                 raw_bytes INTEGER, compressed_bytes INTEGER, elapsed_s REAL);

Evaluation Panel (new tab):

  • "Run enwik8 benchmark" button: html.Button("Run enwik8 Eval", id="eval-enwik8-btn") — runs ./scripts/eval_enwik8.sh or equivalent via subprocess
  • Results table: html.Table(id="eval-results-table") — columns: Step, Dataset, BPB, Model Size, vs gzip (2.7), vs zstd-22 (1.95), vs CMIX (1.23), Delta from CMIX
  • Roundtrip verification status: html.Div(id="eval-roundtrip-status") — shows PASS/FAIL
  • Historical chart: dcc.Graph(id="eval-bpb-history") — bpb over evaluations, with reference lines

New eval_results DB table:

eval_results(step INTEGER, ts REAL, dataset TEXT, bpb REAL,
             compressed_bytes INTEGER, original_bytes INTEGER,
             roundtrip_ok INTEGER, PRIMARY KEY (step, dataset));

Simulate training changes:

  • Write synthetic eval_results: bpb decreasing toward target, roundtrip_ok=1

WS7: Callback Architecture + Robustness

Target files: gui/callbacks/__init__.py, new callback modules, gui/app.py

New callback modules (one per tab, each with a register_callbacks(app) function):

  • gui/callbacks/training_callbacks.py — refresh charts, KPI cards, log tail, ETA, checkpoint list
  • gui/callbacks/arch_callbacks.py — architecture diagram update on tier change
  • gui/callbacks/routing_callbacks.py — heatmap, bias history, expert drill-down, load balance
  • gui/callbacks/mod_callbacks.py — current MoD, history
  • gui/callbacks/activation_callbacks.py — activation dist + heatmap
  • gui/callbacks/brain_callbacks.py — PCA visualization
  • gui/callbacks/compression_callbacks.py — compress/decompress/verify
  • gui/callbacks/eval_callbacks.py — benchmark running, results display
  • gui/callbacks/growth_callbacks.py — growth sliders
  • gui/callbacks/config_callbacks.py — config CRUD, training start/stop/resume, export/import

gui/callbacks/__init__.py:

from .training_callbacks import register as reg_training
from .arch_callbacks import register as reg_arch
# ... etc

def register_all_callbacks(app):
    for reg in [reg_training, reg_arch, reg_routing, reg_mod,
                reg_activation, reg_brain, reg_compression, reg_eval,
                reg_growth, reg_config]:
        reg(app)

gui/app.py simplification:

  • Build layout (same structure, add new Evaluation tab)
  • Call register_all_callbacks(app)
  • Run server
  • No inline callback definitions

Error handling:

  • All DB reads: try/except with html.Div(f"Error: {e}", className="text-danger") returns
  • All subprocess calls: catch FileNotFoundError, subprocess.TimeoutExpired, non-zero return codes
  • Display errors in status divs

Process health monitoring:

  • In config_callbacks.py, check training_process.poll() each refresh cycle
  • If process has exited with non-zero code, update status to "Crashed (exit code N)"
  • If process has exited with code 0, update status to "Completed"

New DB Schema Additions

-- Activation statistics per block (periodic snapshots)
CREATE TABLE IF NOT EXISTS activation_stats (
    step INTEGER, block_idx INTEGER,
    mean REAL, std REAL,
    p1 REAL, p25 REAL, p50 REAL, p75 REAL, p99 REAL,
    PRIMARY KEY (step, block_idx)
);

-- PCA-projected hidden states for brain visualization
CREATE TABLE IF NOT EXISTS hidden_states (
    step INTEGER, block_idx INTEGER, sample_idx INTEGER,
    pc1 REAL, pc2 REAL,
    PRIMARY KEY (step, block_idx, sample_idx)
);

-- Evaluation results
CREATE TABLE IF NOT EXISTS eval_results (
    step INTEGER, ts REAL, dataset TEXT,
    bpb REAL, compressed_bytes INTEGER, original_bytes INTEGER,
    roundtrip_ok INTEGER,
    PRIMARY KEY (step, dataset)
);

-- Compression job history
CREATE TABLE IF NOT EXISTS compression_jobs (
    job_id INTEGER PRIMARY KEY AUTOINCREMENT,
    ts REAL, input_path TEXT, output_path TEXT,
    profile TEXT, checkpoint TEXT, status TEXT,
    raw_bytes INTEGER, compressed_bytes INTEGER, elapsed_s REAL
);

New Viz Utilities

  • moe_routing_history_fig(routing_history_df) — multi-line time series of expert loads
  • load_balance_fig(routing_history_df) — Gini coefficient + max/min ratio over time
  • mod_gate_history_fig(gates_history_df) — fraction_processed per block over time
  • compression_comparison_fig(jobs_df) — bar chart comparing compressed sizes
  • eval_bpb_history_fig(eval_df) — bpb over evaluations with reference lines

Layout Changes

Add new tab to the dbc.Tabs in app.py:

dbc.Tab(eval_panel.layout(), label="🏆 Evaluation", tab_id="tab-eval"),

File Structure After Changes

gui/
├── app.py                    (simplified: layout + callback registration only)
├── db.py                     (expanded: new table queries)
├── viz_utils.py              (expanded: 5 new figure functions)
├── requirements.txt          (unchanged)
├── callbacks/
│   ├── __init__.py           (rewrite: import and register all)
│   ├── training_callbacks.py (new: refresh, KPIs, ETA, log tail, checkpoints)
│   ├── arch_callbacks.py     (new: architecture diagram)
│   ├── routing_callbacks.py  (new: heatmap, history, drill-down, balance)
│   ├── mod_callbacks.py      (new: current + history)
│   ├── activation_callbacks.py (new: activation dist + heatmap)
│   ├── brain_callbacks.py    (new: PCA visualization)
│   ├── compression_callbacks.py (new: compress/decompress/verify)
│   ├── eval_callbacks.py     (new: benchmark running + results)
│   ├── growth_callbacks.py   (new: growth sliders)
│   └── config_callbacks.py   (new: config CRUD, start/stop/resume)
└── panels/
    ├── __init__.py           (update: add eval_panel)
    ├── config_panel.py       (major rewrite)
    ├── training_panel.py     (expand: KPIs, ETA, checkpoint, log, head losses)
    ├── architecture_panel.py (minor: cleanup)
    ├── routing_panel.py      (expand: add history, detail, balance)
    ├── mod_panel.py          (expand: history, metrics)
    ├── activations_panel.py  (minor: already has good layout)
    ├── brain_panel.py        (minor: already has good layout)
    ├── compression_panel.py  (major rewrite: full compression UX)
    ├── eval_panel.py         (NEW: evaluation tab)
    └── growth_panel.py       (minor: cleanup)

Acceptance Criteria

  1. Config sidebar: every field from tier1_micro.yaml is editable; Load/Save/Reset work
  2. Training control: Start launches training, Stop kills it, Resume restarts from checkpoint; sidebar values propagate to CLI args
  3. Training panel: KPI cards render, ETA shows, log tail updates, checkpoint list populates
  4. Routing panel: heatmap + history + drill-down + balance plots all render with data
  5. MoD panel: current + history both render with real time-series data
  6. Activations panel: dist + heatmap render from DB data
  7. Brain panel: PCA scatter renders from DB data
  8. Compression panel: file path input + upload; compress runs; results table shows; verify works
  9. Evaluation panel: enwik8 eval button works; results table renders; roundtrip status shows
  10. All callbacks in separate modules; app.py has no inline callbacks
  11. Error handling: DB errors, subprocess errors, missing files all show user-visible messages
  12. Process health: crashed/completed training shows correct status
  13. python -m bin.simulate_training --steps 5000 populates ALL new tables
  14. python gui/app.py starts without error and all tabs render