Vision experiments

.gitignore

@@ -211,6 +211,9 @@ jobs
 outputs
 assets
 runs
+figs
+analysis
+wandb
 
 # Generated plots
 *.png

README.md

@@ -77,15 +77,6 @@ maxp/
 ├── dag.py             # DAG builder — traces PM-to-PM data flow
 ├── trace.py           # Operation tracer — records matmul-like ops
 └── diagnose.py        # Coord-check diagnostics — sweep widths, plot scaling
-
-experiments/lm/        # LLaMA-3 pre-training experiments (torchtitan)
-├── train.py           # Main training entry point (MaxPTrainer)
-├── maxp_llama3.py     # LLaMA-3 scale configs (debug, s1–s5) + compute_steps
-├── maxp_converter.py  # Post-optimizer-build hook; wires up Parametrization
-├── launch_sweep.py    # Generate and submit SLURM jobs for full LR sweep
-├── run.sh             # Submit a sweep for one scale
-├── run_debug.sh       # Quick single-GPU debug run
-└── coord_check.py     # Coord check for the parametrized LLaMA-3 model
 ```
 
 **Phase 1** (static, at init): discover PMs → reinit weights → solve LP → build param groups.
@@ -143,6 +134,11 @@ param = Parametrization(
     solve_interval=1,            # re-solve every N steps
     sample_size=32,              # max batch size for alignment measurement
     c_ema=0.0,                   # EMA smoothing for c values (0 = instant)
+    alignment_ema=0.0,           # EMA smoothing for measured alignment values
+    resample_w0=False,           # re-sample w0 snapshot each solve
+    use_training_activations=False,  # use activations from training forward pass
+    solver=None,                 # custom PuLP solver (default: CBC)
+    warm_start=False,            # warm-start LP from previous c solution
 )
 
 # Param groups for optimizer
@@ -270,8 +266,8 @@ def make_input(width):
     return torch.randint(0, vocab_size, (4, seq_len))
 
 all_ops, affected, act_stats = diagnose_axis(make_model, make_input, widths)
-print_axis(all_ops, affected, act_stats, widths)
-plot_axis(all_ops, affected, act_stats, widths, path="coord_check.png")
+print_axis("width", all_ops, affected, act_stats, widths)
+plot_axis("width", all_ops, affected, act_stats, widths, filename="coord_check.png")
 ```
 
 `print_axis` shows how the RMS of each op's output scales with width. For a correctly parametrized model, activations should be roughly **constant** (slope ≈ 0 in log-log) at init and remain stable after a few training steps.
@@ -290,6 +286,21 @@ The `experiments/lm/` directory contains a full LLaMA-3 pre-training pipeline bu
 
 Training steps are computed automatically as `20 × non-embed params / tokens-per-step`.
 
+### Files
+
+```
+experiments/lm/
+├── train.py               # Main training entry point (MaxPTrainer)
+├── maxp_llama3.py         # LLaMA-3 scale configs (debug, s1–s5) + compute_steps
+├── maxp_converter.py      # Post-optimizer-build hook; wires up Parametrization
+├── launch_sweep.py        # Generate and submit SLURM jobs for full LR sweep
+├── fineweb.py             # FineWeb-Edu dataset registration for torchtitan
+├── download_hf_assets.py  # HuggingFace asset downloader (from torchtitan)
+├── run.sh                 # Submit a sweep for one scale
+├── run_debug.sh           # Quick single-GPU debug run
+└── coord_check.py         # Coord check for the parametrized LLaMA-3 model
+```
+
 ### Debug run (single GPU)
 
 ```bash
@@ -305,6 +316,39 @@ bash experiments/lm/run_debug.sh \
 bash experiments/lm/run.sh s3   # submits 7 LRs × 3 methods × 2 seeds = 42 jobs
 ```
 
+## Vision Experiments
+
+The `experiments/vision/` directory contains ViT and MLP pre-training experiments using [timm](https://github.com/huggingface/pytorch-image-models) with streaming HuggingFace datasets. It trains four ViT scales and four MLP scales with the same three methods as the LLM experiments.
+
+### Files
+
+```
+experiments/vision/
+├── train.py           # Main training entry point (single-GPU, streaming HF data)
+├── maxp_timm.py       # timm model registry + automatic ParametrizedModule wrappers
+├── hf_vision_data.py  # HF streaming train/val pipeline + transforms
+├── utils.py           # Shared helpers (LR logging, arg parsing, checkpointing)
+├── launch_sweep.py    # Generate and submit SLURM jobs for full LR sweep
+├── run.sh             # Submit a sweep for one scale
+├── run_debug.sh       # Quick tiny-model debug run (CPU/GPU)
+├── coord_check_vit.py # Coord-style diagnostic for ViT models
+└── coord_check_mlp.py # Coord-style diagnostic for MLP models
+```
+
+Available scales: `debug`, `vit-s`, `vit-b`, `vit-l`, `mlp-s`, `mlp-m`, `mlp-b`, `mlp-l`.
+
+### Debug run (CPU/GPU)
+
+```bash
+bash experiments/vision/run_debug.sh --steps 20
+```
+
+### SLURM sweep
+
+```bash
+bash experiments/vision/run.sh vit-s   # submits 9 LRs × 3 methods × 3 seeds = 81 jobs
+```
+
 ## Running Tests
 
 ```bash

experiments/vision/coord_check_mlp.py

@@ -0,0 +1,103 @@
+#!/usr/bin/env python3
+"""Coord-style diagnostic for ScalableMLP.
+
+Width axis is `hidden` — the dimension that scales with model size.
+"""
+
+from __future__ import annotations
+
+import argparse
+
+import torch
+import torch.nn.functional as F
+
+from maxp import Parametrization, diagnose_axis, plot_axis, print_axis
+
+from maxp_timm import ScalableMLP, _install_mlp_wrappers
+
+
+WIDTHS = [64, 128, 256, 512, 1024]
+IMAGE_SIZE = 32
+NUM_CLASSES = 1000
+DEPTH = 6
+
+
+def _make_model(hidden: int, parametrized: bool):
+    model = ScalableMLP(
+        hidden=hidden,
+        depth=DEPTH,
+        num_classes=NUM_CLASSES,
+        dropout=0.0,
+        image_size=IMAGE_SIZE,
+        patch_size=4,
+    )
+    if not parametrized:
+        return model, None
+    _install_mlp_wrappers(model)
+    sample_input = torch.randn(1, 3, IMAGE_SIZE, IMAGE_SIZE)
+    param = Parametrization(
+        model,
+        lr_prefactor=1e-3,
+        optimizer_type="adam",
+        alignment="full",
+        sample_input=sample_input,
+    )
+    return model, param.param_groups
+
+
+def _make_input(hidden: int) -> torch.Tensor:
+    return torch.randn(8, 3, IMAGE_SIZE, IMAGE_SIZE)
+
+
+def _make_train_step(model, param_groups):
+    opt = (
+        torch.optim.AdamW(param_groups)
+        if param_groups is not None
+        else torch.optim.AdamW(model.parameters(), lr=1e-3)
+    )
+    x_fixed = torch.randn(8, 3, IMAGE_SIZE, IMAGE_SIZE)
+    y_fixed = torch.randint(0, NUM_CLASSES, (8,))
+
+    def step(model, step_idx):
+        del step_idx
+        logits = model(x_fixed)
+        loss = F.cross_entropy(logits, y_fixed)
+        loss.backward()
+        opt.step()
+        opt.zero_grad()
+
+    return step
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="Coord check for ScalableMLP (continuous hidden width axis)",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument("--parametrized", action="store_true")
+    parser.add_argument("--steps", type=int, default=10)
+    parser.add_argument("--seeds", type=int, default=5)
+    parser.add_argument("--plot", action="store_true")
+    args = parser.parse_args()
+
+    variant = "parametrized" if args.parametrized else "plain"
+    print(f"ScalableMLP coord check ({variant}) — widths={WIDTHS}")
+
+    ops, affected, act_stats = diagnose_axis(
+        make_model_fn=lambda w: _make_model(w, args.parametrized),
+        make_input_fn=_make_input,
+        widths=WIDTHS,
+        n_steps=args.steps,
+        n_seeds=args.seeds,
+        train_step_fn=_make_train_step,
+    )
+    print_axis("hidden", ops, affected, act_stats, WIDTHS)
+
+    if args.plot:
+        out_name = f"coord_check_mlp_{variant}.png"
+        plot_axis("hidden", ops, affected, act_stats, WIDTHS, out_name, plot_every=1)
+        print(f"Saved {out_name}")
+
+
+if __name__ == "__main__":
+    main()

experiments/vision/coord_check_vit.py

@@ -0,0 +1,111 @@
+#!/usr/bin/env python3
+"""Coord-style diagnostic for wrapped timm ViT models."""
+
+from __future__ import annotations
+
+import argparse
+
+import torch
+import torch.nn.functional as F
+
+from maxp import Parametrization, diagnose_axis, plot_axis, print_axis
+
+from maxp_timm import SCALE_CONFIGS, create_model, install_pm_wrappers
+
+
+WIDTH_TO_SCALE = {
+    192: "debug",   # vit_tiny
+    384: "vit-s",   # vit_small
+    768: "vit-b",   # vit_base
+}
+
+
+def _make_model(width: int, parametrized: bool):
+    scale = WIDTH_TO_SCALE[width]
+    cfg = SCALE_CONFIGS[scale]
+    model = create_model(
+        scale=scale,
+        num_classes=1000,
+        image_size=cfg.image_size,
+    )
+    if not parametrized:
+        return model, None
+    install_pm_wrappers(model)
+    sample_input = torch.randn(1, 3, cfg.image_size, cfg.image_size)
+    param = Parametrization(
+        model,
+        lr_prefactor=1e-3,
+        optimizer_type="adam",
+        alignment="full",
+        sample_input=sample_input,
+    )
+    return model, param.param_groups
+
+
+def _make_input(width: int) -> torch.Tensor:
+    scale = WIDTH_TO_SCALE[width]
+    cfg = SCALE_CONFIGS[scale]
+    return torch.randn(32, 3, cfg.image_size, cfg.image_size)
+
+
+def _make_train_step(model, param_groups):
+    if param_groups is not None:
+        opt = torch.optim.AdamW(param_groups)
+    else:
+        opt = torch.optim.AdamW(model.parameters(), lr=1e-2)
+
+    x_fixed = torch.randn(32, 3, 224, 224)
+    y_fixed = torch.randint(0, 1000, (32,))
+
+    def step(model, step_idx):
+        del step_idx
+        logits = model(x_fixed)
+        loss = F.cross_entropy(logits, y_fixed)
+        loss.backward()
+        opt.step()
+        opt.zero_grad()
+
+    return step
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="Coord-style diagnostic for timm ViT wrappers",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument("--parametrized", action="store_true")
+    parser.add_argument("--steps", type=int, default=5)
+    parser.add_argument("--seeds", type=int, default=2)
+    parser.add_argument("--plot", action="store_true")
+    args = parser.parse_args()
+
+    widths = sorted(WIDTH_TO_SCALE)
+    variant = "parametrized" if args.parametrized else "plain"
+    print(f"ViT coord check ({variant}) — widths={widths}")
+
+    ops, affected, act_stats = diagnose_axis(
+        make_model_fn=lambda w: _make_model(w, args.parametrized),
+        make_input_fn=_make_input,
+        widths=widths,
+        n_steps=args.steps,
+        n_seeds=args.seeds,
+        train_step_fn=_make_train_step,
+    )
+    print_axis("embed_dim", ops, affected, act_stats, widths)
+
+    if args.plot:
+        out_name = f"coord_check_vit_{variant}.png"
+        plot_axis(
+            "embed_dim",
+            ops,
+            affected,
+            act_stats,
+            widths,
+            out_name,
+            plot_every=1,
+        )
+        print(f"Saved {out_name}")
+
+
+if __name__ == "__main__":
+    main()