diff --git a/CLAUDE.md b/CLAUDE.md index c5f906a..b81d03f 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -1,50 +1,35 @@ -# CLAUDE.md — maxP project instructions +# CLAUDE.md — maxP -## What This Project Is +maxP is a PyTorch library for **abc-parametrization** (Everett et al. 2024, arXiv:2407.05872), extended with **measured dynamic alignment** (https://iejmac.github.io/2025/03/26/alignments.html). Each layer `l` has exponents `(a_l, b_l, c_l)`: output `×n^{-a}`, init variance `n^{-2b}`, learning rate `lr_prefactor·n^{-c}`. Static parametrization and dynamic alignment/LP solving are both implemented; current focus is width-ladder muTransfer experiments (LM + vision). -maxP is a PyTorch library for neural network parametrization. It implements the -abc-parametrization framework from "Scaling Exponents Across Parameterizations and -Optimizers" (Everett et al., 2024, arXiv:2407.05872) and extends it with dynamic -alignment measurement from https://iejmac.github.io/2025/03/26/alignments.html. +## Layout -Core idea: each layer l has exponents (a_l, b_l, c_l) controlling output multiplier -(n^{-a}), init variance (n^{-2b}), and learning rate (lr_prefactor * n^{-c}). - -## Status - -Active rewrite in progress. Working in two phases: -- **Phase 1** (current): Static parametrizations with robust tests -- **Phase 2**: Dynamic alignment measurement and LP solving - -See `docs/DESIGN.md` for full design. All working docs live in `docs/`. +``` +maxp/ # the library + module.py # ParametrizedModule — wraps a layer (output scale + per-layer init/LR) + parametrization.py # Parametrization — entry point: wrap → solve c → optimizer param_groups + solver.py # LP solver for the c-exponents + alignment.py dag.py trace.py diagnose.py # measure / classify-DAG / coord-checks +tests/ # pytest, CPU-only +examples/ # small standalone demos (MLP / ViT / nanoGPT) +experiments/ + lm/ # torchtitan LLaMA-3 width-ladder (SLURM, server) + vision/ # timm ViT width-ladder (SLURM, server) +docs/ + walkthrough.md # per-module code map (start here) + parametrization_policy.md # layer-classification / abc theory + research_context.md # paper + blog theory summary + verify.md / verify.py # component-by-component correctness evidence + width_ladder_results.md # LM results + vision_experiment_plan.md # vision plan +``` -## Environment +## Dev ```bash -source .venv/bin/activate # Python 3.13 via Homebrew +source .venv/bin/activate pip install -e .[dev] -python -m pytest tests/ # run tests (CPU only, no GPU) -``` - -## Key Dependencies - -- torch >= 2.8, numpy >= 2.0, pulp >= 3.0 -- Dev: pytest, pytest-cov, matplotlib, pandas, pyyaml, torchvision - -## Project Layout - -``` -maxp/ # Main package -tests/ # pytest suite -examples/ # Training scripts (MLP, ViT on CIFAR-10) -docs/ # Design docs and notes - DESIGN.md # Target design - research_context.md # Paper/blog theory summary - current_state_of_main.md # Pre-refactor snapshot +python -m pytest tests/ -v --tb=short # CPU-only, no GPU; CI on py3.10–3.14 ``` -## Conventions - -- Tests must run on CPU (no GPU assumed during development) -- CI runs on Python 3.10-3.14 -- pytest with `-v --tb=short` +Experiments run on GPU via SLURM (server), never locally. diff --git a/README.md b/README.md index 97d67e0..ed1182e 100644 --- a/README.md +++ b/README.md @@ -1,13 +1,15 @@ # maxP -**maxP** is a PyTorch library for neural network parametrization implementing the abc-parametrization framework from [Everett et al., 2024](https://arxiv.org/abs/2407.05872) with dynamic alignment measurement from [this blog post](https://iejmac.github.io/2025/03/26/alignments.html). +**maxP** is a PyTorch library for width-robust neural network parametrization. It implements the abc-parametrization framework from [Everett et al., 2024](https://arxiv.org/abs/2407.05872) and extends it with measured alignment from [this blog post](https://iejmac.github.io/2025/03/26/alignments.html). + +> **Tune the learning rate once at a small width — reuse it at billions of parameters.** Wrap your width-sensitive layers, and maxP sets each layer's init, output scale, and LR so the optimal LR *transfers across width*. With one passive alignment measurement, it then beats µP by ~0.1 nat at every scale from 77M to 3.7B params — see [`docs/width_ladder_results.md`](docs/width_ladder_results.md). Each layer `l` has three exponents controlling its width-scaling behavior: - **`a_l`**: Output multiplier — layer output is scaled by `n^{-a_l}` - **`b_l`**: Init variance — weights initialized as `N(0, n^{-2b_l})` - **`c_l`**: Learning rate — `lr_l = lr_prefactor * n^{-c_l}` -The library solves a Linear Program (LP) to find optimal `c_l` values that maximize per-layer learning rates while maintaining numerical stability. Optionally, it measures actual alignment between initial and current weights/activations during training and re-solves the LP dynamically. +The library solves a Linear Program (LP) to find optimal `c_l` values that maximize per-layer learning rates while maintaining numerical stability. Alignment can be measured rather than assumed worst-case. The validated recipe takes a single passive alignment measurement at a small scale, freezes it into a per-layer `c` table, and applies that table unchanged across the width ladder. Online per-step re-solving is also supported but is exploratory. ## Installation @@ -47,25 +49,21 @@ class MLP(nn.Module): model = MLP(width=256) -# 2. Apply parametrization (re-inits weights, solves LP, builds param groups) +# 2. Solve per-layer LRs + re-init weights, then build optimizer param groups. +# alignment="full" is µP. The lr_prefactor you tune here transfers across width. param = Parametrization(model, lr_prefactor=1e-3, alignment="full") - -# 3. Create optimizer from param_groups (each layer gets its own LR) optimizer = torch.optim.AdamW(param.param_groups) -# 4. Capture initial state BEFORE training (needed for dynamic alignment) -X_init = next(iter(train_loader))[0] -param.capture_initial(X_init) - -# 5. Training loop +# 3. Train normally — nothing maxP-specific in the loop. for X, y in train_loader: optimizer.zero_grad() loss = criterion(model(X), y) loss.backward() optimizer.step() - param.step(X, optimizer) # measure alignment, re-solve LP, update LRs ``` +That's µP, done in two lines. The `lr_prefactor` tuned at width 256 stays optimal at width 8192 — no retuning. To beat µP, measure alignment once and freeze it into a per-layer `c` table via `alignment_overrides` — see [Dynamic Alignment](#dynamic-alignment-phase-2) for measurement and [`docs/width_ladder_results.md`](docs/width_ladder_results.md) for results. + ## Design Overview ``` @@ -276,77 +274,77 @@ If an activation grows with width, `a` is too small for that layer — increase ## LLM Experiments -The `experiments/lm/` directory contains a full LLaMA-3 pre-training pipeline built on [torchtitan](https://github.com/pytorch/torchtitan). It trains five scales (30M–3B parameters) with three methods: +The `experiments/lm/` directory contains a full LLaMA-3 pre-training pipeline built on [torchtitan](https://github.com/pytorch/torchtitan). It runs a width-only ladder (depth 12, FineWeb-Edu train / c4 val), comparing two arms: + +| Method | Parametrization | Alignment | +|---|---|---| +| `mup-no` | µP | none (permissive preset) — baseline | +| `maxP-meas` | µP + measured | static `c`-table from one passive measurement at s2, applied unchanged across scales | -| Method | Parametrization | Alignment | Schedule | -|---|---|---|---| -| `maxP` | µP + measured | online (dynamic re-solve) | emergent | -| `mup-full` | µP | full (worst-case preset) | constant | -| `mup-no` | µP | none (permissive preset) | constant | +| scale | s1 | s2 | s3 | s4 | s5 | +|---|---|---|---|---|---| +| **embed_dim** | 256 | 512 | 1024 | 2048 | 4096 | +| **total params** | 77M | 172M | 426M | 1.18B | 3.67B | -Training steps are computed automatically as `20 × non-embed params / tokens-per-step`. +Training steps are computed automatically as `20 × non-embed params / tokens-per-step`. Results and analysis: `docs/width_ladder_results.md`. ### Files ``` experiments/lm/ -├── train.py # Main training entry point (MaxPTrainer) +├── train.py # 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 +├── launch_sweep.py # Generate/submit SLURM jobs for one LR sweep +├── pipeline.sh # Full submit-only DAG: measure → export → transfer → verify +├── export_alignment.py # Build the static c-table from a measured run +├── pipeline_analyze.py # Pick best LR, gate checks, transfer/E1 verdicts ├── 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 +├── run.sh / run_debug.sh # Single-scale sweep / quick debug run └── coord_check.py # Coord check for the parametrized LLaMA-3 model ``` -### Debug run (single GPU) +### Run ```bash -bash experiments/lm/run_debug.sh \ - --tokenizer /path/to/tokenizer \ - --c4-test /path/to/c4_test \ - --steps 200 --method maxP -``` - -### SLURM sweep - -```bash -bash experiments/lm/run.sh s3 # submits 7 LRs × 3 methods × 2 seeds = 42 jobs +bash experiments/lm/run_debug.sh --tokenizer /path/to/tokenizer \ + --c4-test /path/to/c4_test --steps 200 --method mup-no # single-GPU smoke +bash experiments/lm/pipeline.sh # full ladder DAG (server) ``` ## 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. +The `experiments/vision/` directory mirrors the LM protocol for ViT on [timm](https://github.com/huggingface/pytorch-image-models) with ImageNet-12k. It runs a width-only ladder (depth 12, head_dim 64, embed 256→4096 held otherwise fixed) with the same two arms as LM (`mup-no` baseline, `maxP-meas`). + +| scale | s1 | s2 | s3 | s4 | s5 | +|---|---|---|---|---|---| +| **embed_dim** | 256 | 512 | 1024 | 2048 | 4096 | +| **total params** | 12.8M | 44.4M | 164M | 630M | 2.47B | ### 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 +├── train.py # Training entry point (single-GPU, map-style HF data) +├── maxp_timm.py # timm ViT registry (s1–s5) + automatic ParametrizedModule wrappers +├── hf_vision_data.py # Non-streaming HF train/val pipeline + transforms +├── utils.py # Shared helpers (LR logging, checkpointing, RNG) +├── launch_sweep.py # Generate/submit SLURM jobs for one LR sweep +├── pipeline.sh # Full submit-only DAG: measure → export → transfer → verify +├── export_alignment.py # Build the static c-table from a measured run +├── pipeline_analyze.py # Pick best LR, gate checks, transfer/E1 verdicts +├── run.sh / run_debug.sh # Single-scale sweep / quick tiny-model debug run +└── coord_check_vit.py / coord_check_mlp.py # Coord-check diagnostics ``` -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 -``` +Scales: `debug`, `s1`–`s5` (ViT only). -### SLURM sweep +### Run ```bash -bash experiments/vision/run.sh vit-s # submits 9 LRs × 3 methods × 3 seeds = 81 jobs +bash experiments/vision/run_debug.sh --steps 20 # tiny-model smoke (CPU/GPU) +bash experiments/vision/pipeline.sh # full ladder DAG (server) ``` ## Running Tests diff --git a/docs/DESIGN.md b/docs/DESIGN.md deleted file mode 100644 index c526085..0000000 --- a/docs/DESIGN.md +++ /dev/null @@ -1,265 +0,0 @@ -# DESIGN.md — maxP v2 - -**Status**: DRAFT — for discussion - ---- - -## Goals - -1. **Generalize beyond MLPs**: Support arbitrary architectures by letting the user specify - which layers to parametrize and what role each has. -2. **Clean API**: Parametrization setup mutates the model and returns param groups for the optimizer. -3. **Composable with PyTorch ecosystem**: Work alongside standard LR schedulers, not replace them. -4. **Robust tests**: Verify correctness against known parametrization results, not just API shape. -5. **Phased implementation**: Start with static parametrizations (no dynamic alignment/solver), - build robust tests, then add dynamic `.step()`. - ---- - -## Target API - -### Phase 1: Static Parametrizations - -```python -import maxp - -model = MyModel(...) - -param = maxp.Parametrization(model, config) -# This: -# 1. Finds specified layers in the model -# 2. Re-initializes their weights according to b_l -# 3. Wraps them to apply layer multipliers according to a_l -# 4. Computes per-layer learning rates according to c_l -# 5. Exposes .param_groups for the optimizer - -optimizer = AdamW(param.param_groups) -lr_sched = CosineAnnealingLR(optimizer, T_max=1000) - -for batch in dataloader: - loss = loss_fn(model(batch.x), batch.y) - loss.backward() - optimizer.step() - lr_sched.step() - optimizer.zero_grad() -``` - -### Phase 2: Dynamic Alignment (future) - -```python -param = maxp.Parametrization(model, config) -optimizer = AdamW(param.param_groups) -lr_sched = CosineAnnealingLR(optimizer, T_max=1000) - -for batch in dataloader: - loss = loss_fn(model(batch.x), batch.y) - loss.backward() - lr_sched.step() - param.step() # measure alignment, solve LP, adjust per-layer LRs - optimizer.step() - optimizer.zero_grad() -``` - ---- - -## Layer Specification - -The user specifies which layers to parametrize and what role each has. No magic -heuristics, no hardcoded skip logic. Any module type can be parametrized — Linear, -Conv2d, LayerNorm, whatever has parameters and produces output that should be scaled. - -### Layer Roles - -Each parametrized layer gets a **role** that determines its constraint type in the LP: - -| Role | Meaning | LP Constraint | -|------|---------|---------------| -| `embedding` | Input doesn't change during training | `r_l = a_l + c_l` | -| `hidden` | Both activations and weights change | `r_l = min(a_l+c_l-alpha, a_l+c_l+r_{l-1}-U, 0.5+r_{l-1}-omega)` | -| `readout` | Final output / classification layer | Stable-logits variant | - -### Config specifies layers via patterns - -```python -config = maxp.Config( - parametrization="mup", - lr_prefactor=1e-3, - layers={ - # glob pattern on module name → role - "embed": "embedding", - "blocks.*.ffn.up": "hidden", - "blocks.*.ffn.down": "hidden", - "blocks.*.attn.qkv": "hidden", - "blocks.*.attn.proj": "hidden", - "head": "readout", - }, -) -``` - -Layers not listed are **not parametrized** — they go into the optimizer with default -lr_prefactor, no ABC treatment. This gives the user full control over exactly which -modules get ABC scaling and which don't. - -### MLP convenience - -For simple MLPs where all Linear layers should be parametrized in order -(first=embedding, middle=hidden, last=readout): - -```python -config = maxp.Config( - parametrization="mup", - lr_prefactor=1e-3, - # layers=None → auto-detect: find all nn.Linear in definition order, - # assign first=embedding, last=readout, rest=hidden -) -``` - -Auto-detection is opt-in (the default when `layers` is omitted), not a hidden heuristic. - ---- - -## What Parametrization.__init__ Does - -1. **Resolves layers**: Walks model modules, matches against `layers` patterns (or - auto-detects). Produces an ordered list of (name, module, role) tuples. - Any module type is valid — not limited to nn.Linear. - -2. **Resolves ABC values**: From named parametrization or custom al/bl/cl. - Each matched layer gets its own (a_l, b_l, c_l). For named parametrizations, - values are derived from the role. For custom, the user provides one triplet per - matched layer. - -3. **Initializes weights**: For each parametrized layer, re-init with - `std = std_prefactor * n^{-b_l}` where n = fan_in. - -4. **Applies multipliers**: Wraps each parametrized layer so its output is scaled by - `n^{-a_l}`. (How: replace the module in the model with a wrapper that scales output.) - Works for any module type since it just wraps the forward pass. - -5. **Builds param_groups**: One group per parametrized layer with - `lr = lr_prefactor * n^{-c_l}`, plus one group for all other parameters with - `lr = lr_prefactor`. - ---- - -## Config - -```python -@dataclass -class Config: - # --- Parametrization --- - # Named (mutually exclusive with custom al/bl/cl) - parametrization: str | None = None # "mup", "sp", "ntk", "mfp" - optimizer_type: str | None = None # "adam" or "sgd" — needed for ABC table lookup - alignment: str | None = None # "full" or "no" — alignment assumption for ABC tables - - # Custom (mutually exclusive with named parametrization) - al: list[float] | None = None - bl: list[float] | None = None - cl: list[float] | None = None - - # --- Learning rate --- - lr_prefactor: float = 1e-3 - - # --- Layer specification --- - # Dict of glob pattern → role ("embedding", "hidden", "readout") - # None = auto-detect (all Linear layers, first=embedding, last=readout, rest=hidden) - layers: dict[str, str] | None = None - - # --- Weight init --- - std_prefactor: float = 1.0 - apply_multipliers: bool = True -``` - -Note: solver/alignment/tracer config omitted — that's Phase 2. - ---- - -## Module Structure - -``` -maxp/ - __init__.py # Public API: Parametrization, Config - parametrization.py # Parametrization class — main entry point - tables.py # Named parametrization ABC tables (mup, sp, ntk, mfp) - solver.py # LP solver (Phase 2) - alignment.py # Alignment computation (Phase 2) - tracer.py # Hook-based layer state capture (Phase 2) -``` - -Phase 1 only needs: `__init__.py`, `parametrization.py`, `tables.py`. - ---- - -## Named Parametrization Tables - -The ABC values for named parametrizations depend on: -- Which parametrization (mup, sp, ntk, mfp) -- Which optimizer (adam, sgd) -- Which alignment assumption (full, no) -- Layer role (embedding, hidden, readout) -- Number of layers (for some parametrizations) - -These are pure data — lookup tables derived from the paper. Stored in `tables.py`. - ---- - -## Tests (Phase 1) - -### 1. ABC Tables -- Verify values match the paper (Table 1 from Everett et al.) for all combinations -- Verify consistency: n_layers=2 edge case, n_layers=100 -- Verify stability-at-initialization conditions hold: a_0+b_0=0, a_l+b_l=0.5 (hidden), etc. - -### 2. Weight Initialization -- Variance matches n^{-2b_l} for each layer -- Works with different model architectures (not just Sequential) -- Biases reset to zero -- Only specified layers are re-initialized - -### 3. Layer Multipliers -- Output is scaled by n^{-a_l} -- Forward pass produces correct values -- Gradient flows through correctly - -### 4. Param Groups -- One group per parametrized layer with correct lr = lr_prefactor * n^{-c_l} -- Non-parametrized params get default lr -- Groups contain correct parameters - -### 5. Layer Matching -- Glob patterns match expected modules -- Auto-detection assigns correct roles -- Unmatched layers excluded from parametrization -- Error on empty match (no layers found) - -### 6. Integration / Correctness -- For named parametrizations under full alignment, the static c_l values should equal - the known muP/SP/NTK/MFP learning rate exponents -- Width transfer: train at width 64, find good lr_prefactor, verify it works at width 512 -- Compare against standalone muP implementation (if available) - -### 7. Solver (Phase 2) -- With full alignment inputs, recover named parametrization c_l values -- Constraint satisfaction: r_l >= 0 -- Relaxed alignment → larger LRs (smaller c_l) -- Edge cases, infeasibility - ---- - -## Open Questions - -1. **Multiplier wrapper generality**: Current code has ScaledLinear (Linear-specific). - Since we now support any module type, we need a generic ScaledModule wrapper that - scales the output of any module. This is straightforward — just wrap forward() and - multiply output by scale. Anything tricky here? - -2. **fan_in for non-Linear layers**: For nn.Linear, fan_in = in_features and the scaling - formulas are well-defined. For Conv2d, fan_in = in_channels * kernel_h * kernel_w. - For LayerNorm, what is fan_in? We may need the user to specify width `n` explicitly - when parametrizing non-standard layer types. - -3. **Constraint chain ordering (Phase 2)**: When layers are specified via patterns, the - LP chains r_l values sequentially. We'll use module definition order. The user controls - which layers are included and their roles, so the ordering follows naturally from the - model definition. Not relevant for Phase 1 (static c_l, no LP). diff --git a/docs/lr_scheduler_compatibility.md b/docs/lr_scheduler_compatibility.md deleted file mode 100644 index 32271da..0000000 --- a/docs/lr_scheduler_compatibility.md +++ /dev/null @@ -1,76 +0,0 @@ -# LR Scheduler Compatibility Analysis - -## Problem - -`Parametrization.step()` overwrites LRs with absolute values: - -```python -group["lr"] = self.lr_prefactor * (fan_in ** (-c)) -``` - -Any external PyTorch LR scheduler adjustment gets blown away on the next re-solve. - -## Conceptual Decomposition - -MaxP controls **per-layer LR ratios** (`n^{-c}`). A scheduler controls the **global LR envelope**. These compose as a product: - -``` -effective_lr(t) = schedule_multiplier(t) * lr_prefactor * fan_in^(-c_l) -``` - -The two concerns are orthogonal — maxP picks the ratios, the scheduler picks the overall scale over time. - -## Old Implementation (`maxp/scheduler.py`) - -The old package solved this two ways: - -### 1. Built-in WSD (Warmup-Stable-Decay) - -`MaxPScheduler` has `_get_wsd_multiplier()` returning a scalar envelope: - -```python -lr = lr_prefactor * (fan_in ** (-c)) * wsd_mult -``` - -- Warmup: linear ramp from `wsd_min_factor` to 1.0 -- Stable: multiplier = 1.0 (LP solver active) -- Decay: LP solver stops, LRs frozen, only decay multiplier applied - -### 2. ChainedMaxPScheduler (for external PyTorch schedulers) - -Wraps `MaxPScheduler` + arbitrary `torch.optim.lr_scheduler` instances: - -```python -# Each step: -lr_before = optimizer.param_groups[managed_idx]["lr"] -for sched in external_schedulers: - sched.step() -lr_after = optimizer.param_groups[managed_idx]["lr"] - -# Apply relative change to prefactor -maxp_scheduler.lr_prefactor *= (lr_after / lr_before) - -# Then maxP computes per-layer LRs with updated prefactor -maxp_scheduler.step(X) -``` - -This lets CosineAnnealingLR, LinearLR, etc. control the envelope while maxP handles per-layer ratios. - -## Plan for `maxp` - -Simplest approach: add a `schedule_multiplier` to `Parametrization.step()` or adopt the `ChainedMaxPScheduler` pattern. The key change in `step()`: - -```python -# Current: -group["lr"] = self.lr_prefactor * (fan_in ** (-c)) - -# With scheduler support: -group["lr"] = self.lr_prefactor * (fan_in ** (-c)) * multiplier -``` - -Where `multiplier` comes from either: -- A built-in WSD schedule (like the old code) -- An external scheduler via the chaining pattern -- A user-supplied callback - -The old `maxp/scheduler.py` is the reference implementation — it already handles WSD, chaining, state_dict/load_state_dict, and decay-phase LR freezing. diff --git a/docs/research_context.md b/docs/research_context.md index 6c3e5d5..af63e6e 100644 --- a/docs/research_context.md +++ b/docs/research_context.md @@ -101,25 +101,11 @@ r_{L-1} >= 0 --- -## paramR vs maxP Comparison +## Implementation status -paramR is the research prototype; maxP is the library extraction. +The theory above is current; the refactor it motivated is done. maxP now wraps arbitrary +layers via `ParametrizedModule`, classifies/solves over a traced op-DAG (`dag.py`/`trace.py`), +and applies the result through a single `Parametrization` entry point — no separate scheduler +class. For the per-module code map, see `walkthrough.md`; for layer-classification rules, see +`parametrization_policy.md`. -Key differences relevant to the refactor: -- paramR: tightly coupled to MLP model class, closure-based scheduler -- maxP: model-agnostic via nn.Linear scanning, proper scheduler class -- Both use identical LP formulation -- maxP added: WSD support, ChainedMaxPScheduler, checkpointing, two alignment norm modes -- paramR has a metric registry system that maxP dropped -- paramR's alignment uses ratio-style; maxP defaults to log-scale (configurable) - ---- - -## What Needs to Change - -1. **Solver generalization**: Currently hardcoded to embed/hidden/readout structure. - Need to support arbitrary layer topologies. -2. **API simplification**: Three-step setup (init weights, create param groups, create scheduler) - is error-prone. Should be a single entry point. -3. **Architecture support**: Need to map LP semantics onto transformers, convnets, etc. -4. **Tests**: Need tests that verify against known parametrization results, not just API smoke tests. diff --git a/docs/verify.py b/docs/verify.py index a64df85..1649b80 100644 --- a/docs/verify.py +++ b/docs/verify.py @@ -571,7 +571,7 @@ def forward(self, x): m5b.h0.align_z0_dW, m5b.h0.align_dZ_w0, m5b.h0.align_dZ_dW = 0.8, 0.3, 0.6 m5b.h1.align_z0_dW, m5b.h1.align_dZ_w0, m5b.h1.align_dZ_dW = 1.0, 0.5, 1.0 # full m5b.h2.align_z0_dW, m5b.h2.align_dZ_w0, m5b.h2.align_dZ_dW = 0.5, 0.2, 0.3 -c_by_name = p5b._resolve_chain() +c_by_name = p5b._resolve() print(f" Per-PM c values: {c_by_name}") assert_check( "hidden layers have different c", diff --git a/docs/vision_experiment_plan.md b/docs/vision_experiment_plan.md new file mode 100644 index 0000000..1d4aaf0 --- /dev/null +++ b/docs/vision_experiment_plan.md @@ -0,0 +1,255 @@ +# Vision muTransfer experiment plan (ViT) — maxP-meas vs muP + +Goal: confirm the maxP-meas result on a **second modality**. Replicate the Llama3 +LM pipeline on a Vision Transformer width ladder: a single passive alignment +measurement → static per-layer c-table → does it (a) beat standard muP and (b) +keep the LR optimum width-stable, transferring the table both down and up the +ladder. + +This plan is the acceptance gate. **No training runs launch until accepted.** + +## Scope (confirmed) + +- **ViT only.** ViT is the transformer analog of the LM experiment. MLP deferred. +- **Two arms only:** `mup-no` (standard muP baseline) and `maxP-meas` (static + c-table from a measured alignment table). Dynamic `maxP` and `mup-full` dropped. +- **Dataset: imagenet-12k** (`timm/imagenet-12k-wds`, 11821 classes). The maxP-meas + win is readout-dominated, so a large head (= num_classes) is required or the + effect is muted — 11821 classes satisfies this. +- **Fixed epoch budget** across all widths (same #epochs = same #samples at every + scale; textbook muTransfer setup): a few real epochs of imagenet-12k. Tests + width-stability of the optimal LR. Batch size held constant across all widths. + +## Width ladder (width-only, mirrors LM s1–s5) + +Width axis = `embed_dim`. **`head_dim = 64` and `depth = 12` held constant**; +`num_heads = embed_dim / 64`. `drop_path = dropout = 0` at every scale (so the +only thing varying is width — regularization that scales with width would +confound the transfer test). Built via timm geometry overrides on +`vit_base_patch16_224`, 224px / patch16 (196 tokens). + +**Biases kept (original ViT).** Linear biases stay at timm ViT defaults +(qkv/proj/fc1/fc2/head all biased); LayerNorm affine kept. (The LM arm is bias-free +because torchtitan `Linear` defaults to `bias=False`; here we keep the stock ViT.) +Each bias lives *inside* its Linear's `ParametrizedModule` (`…qkv.inner.bias`), so +it shares that layer's per-layer LR (`lr_prefactor · width_dim^(−c)`) — identical +to how the original (pre-refactor) ViT wrapper already handled them. Both arms wrap +biases identically, so the only between-arm difference is still the per-layer `c`; +the transfer comparison is unaffected. Empirically verified: the parametrized +coord check stays width-stable across embed 128→1024 *with* biases (same numbers as +bias-free), so the muP scaling is intact. `pos_embed` / `cls_token` are additive +O(1) terms in the `_other` group (fixed-std init + constant LR — conventional muP +for additive embeddings). + +| scale | embed_dim | heads | params (11821 cls) | optim state (AdamW, 16 B/param) | seeds | LR grid | +|-------|-----------|-------|--------------------|--------------------------------|-------|---------| +| s1 | 256 | 4 | 12.8M | 0.2 GB | 3 | full (7) | +| s2 | 512 | 8 | 44.4M | 0.7 GB | 3 | full (7) — **measure base** | +| s3 | 1024 | 16 | 164.3M | 2.6 GB | 2 | narrow (3) | +| s4 | 2048 | 32 | 630.5M | 10.1 GB | 2 | narrow (3) | +| s5 | 4096 | 64 | 2.47B | 39.5 GB | 1 | transfer LR only (1) | + +Seed counts and LR-grid narrowing match the LM ladder. Full grid = +`[3e-4, 1e-3, 3e-3, 1e-2, 3e-2, 1e-1, 3e-1]`. **No scale hardcodes its LR set** — +every scale defaults to the full grid; the coordinators narrow s3/s4 to a 3-point +bracket (argmin ± 1 grid step) and s5 to the single transferred optimum *at +runtime*, by releasing per-arm sentinel files. So the "narrow (3)" / "transfer (1)" +column above is the effective set the gates release, not a value baked into the +config — the optimum is always data-driven, never assumed. + +## Single-GPU feasibility (no FSDP) + +The vision trainer is **single-GPU** (`model.to(device)` + `torch.compile`, no +DDP/FSDP). Each model must fit one GH200 (96 GB HBM). + +- **Memory: fits at every scale.** Worst case s5: 39.5 GB optimizer state + + activations (grad-checkpointing enabled) leaves comfortable headroom under 96 GB. +- **The s5 risk is wall-time, made survivable by checkpoint chaining.** A 2.47B + ViT on one GPU with no sharding and an epoch budget (~95k steps) won't finish in + one 48 h wall. Handled by **resume chains** (s4 = 2 links, s5 = 4 links): each + link checkpoints every 1000 steps, the next link resumes from the latest + checkpoint via `--dependency=afterany`, and once `final_metrics.json` exists the + remaining links no-op. So a wall-time TIMEOUT is recoverable; at most ~1000 steps + are lost per interruption. Still worth measuring sec/step on the server to size + the chain length, but it is no longer a single-shot gate. (s1–s3 run in one job, + chain = 1.) + +## Automated pipeline (DAG, mirrors LM `pipeline.sh`) + +`bash experiments/vision/pipeline.sh` (login node, `DRY=1` to preview) submits the +whole DAG upfront — submit-only, exits in seconds. Decisions unknowable at submit +time are enforced by runtime sentinel files (`launch_sweep --require`), so the +later stages are pre-submitted and self-gate when their coordinator releases them. + +- **Stage A** — s1 + s2 `mup-no --measure-only` full LR sweeps: the alignment + sources. Logs per-layer `(z0·dW, dZ·w0, dZ·dW)` to `metrics.json` (LRs never + changed; measurement is passive). +- **Coord B** (afterany A) — pick best mup-no LR per scale (`pipeline_analyze + best-lr`, lowest val loss), **export** the alignment table from that LR's runs, + write gate report (G1 seed stability, G2 source-LR stability, G3 s1-vs-s2). +- **Stage C** (afterany B, gated on `s2_align.json`) — s1 + s2 `maxP-meas` sweeps: + prefactor re-tune under the measured table. +- **Coord C** (afterany C) — best maxP-meas LR per scale, s1-vs-s2 consistency, + releases the 3-point s3 LR bracket (argmin ± 1 grid step) per arm via sentinels. +- **Stage D** (afterany C, per-LR gated) — s3 full-grid candidates both arms; only + the released bracket runs, the rest exit in ~1 s. +- **Coord D** (afterany D) — transfer proof (s3 argmin == transferred LR?) + **E1 + verdict** (transfer ≤ baseline + tol → continue, else stop); releases s4/s5 + sentinels. E1 fail ⇒ s4/s5 self-skip. +- **Stage E** (afterany D, per-LR gated) — s4 (chain 2) + s5 (chain 4) candidates, + chained checkpoint-resume so they complete across wall-times. + +**Export keying — deliberate divergence from LM:** table keyed by **exact PM name** +(depth fixed ⇒ block names identical across the ladder ⇒ transfers verbatim; +`patch_embed` excluded). LM pooled per op-type suffix; vision keys per exact block +(forced by the `proj` leaf collision, slightly more faithful). Does not affect the +within-vision comparison. + +**Manual fallback:** `run.sh measure|export|sweep ` runs the same phases by +hand without the DAG/gating, for debugging. + +## Training settings + +- Optimizer AdamW (β=0.9/0.999, weight_decay 0.05), bf16 autocast, grad clip 1.0. +- Cosine LR schedule, warmup 5% of steps. +- Label smoothing 0.1, batch size 512 (constant across all widths), **4 epochs** + of imagenet-12k (fixed sample budget; ≈95k steps/run). grad-checkpointing on, + compile on. +- **Checkpointing (all scales):** every 1000 steps, keep latest 1 (+ a `final.pt`), + atomic write, model + optimizer + RNG saved. `final_metrics.json` is the + completion sentinel. **Resume chains** for s4 (2 links) and s5 (4 links) via + `--dependency=afterany`; `--resume` loads the latest checkpoint and the run + continues to the same optimizer-step budget. **Resume-boundary honesty:** the + train loader now shuffles (deterministic per-pass reseed), but within-pass + position is not restored — on resume the interrupted pass restarts, so up to one + pass (epoch) of *reshuffled* data is re-seen and, because the step budget is + fixed, an equal tail of the final pass is dropped. Bounded, roughly symmetric + across arms, optimizer-step count exact. Moot for s1–s3 (chain=1); a small + bounded perturbation per interruption for s4/s5. +- Metric: validation top-1 / top-5 / loss on the imagenet-12k validation split + (in-distribution, unlike the LM's OOD c4 val). Validation cadence + `--val-interval 500 --val-steps 50` (~190 val points over a 4-epoch run, ~10% + compute overhead; the original `--val-steps 500` default would have burned ~⅓ of + every run on eval — fixed). `--val-interval` must be a multiple of + `--log-interval` (20); 500 satisfies this. +- wandb runs **online** (matches the LM arm, which also sets no `WANDB_MODE`); + `pipeline_analyze` reads `metrics.json`, not wandb, so the decision path never + depends on wandb connectivity. +- LR per layer = `lr_prefactor · width_dim^(−c)`; the two arms differ only in the + per-layer `c` (mup-no: uniform 0.5 alignment; maxP-meas: measured table). + +## Verification already done (local, CPU) + +- **coord_check_vit PASSES** — every activation width-stable across embed + 128→1024 (qkv 0.0125, attn.proj 0.0009, fc1 0.798, fc2 ~0.52, head ~0.79, all + flat). muP parametrization is correct. +- **Width-only geometry confirmed** — head_dim 64 and depth 12 constant at all + five scales (verified by build+forward). +- **Bug found & fixed: train loader never shuffled** — `make_loader` set + `drop_last` from `is_train` but never `shuffle`, so imagenet-12k was walked in + identical order every epoch (no SGD shuffling). Now `shuffle=is_train` with a + per-pass-reseeded generator. (An obsolete-assumption defect, exactly what to hunt.) +- **measure → export → apply smoke** — measure-only logs alignment; export builds + a 9-layer table (block PMs + head; `patch_embed` excluded); maxP-meas loads it + and solves a distinct per-layer c-table (head a≈0.85, readout-dominant as in LM). +- **launch_sweep dry-run** — correct job counts (s4 = 2 methods × 3 LRs × 2 seeds + = 12), `--method maxP-meas --alignment-table … --epochs … --batch-size 512` + wired into job.sh, and the "maxP-meas requires --alignment-table" guard fires. + Idempotent resubmit via a `final_metrics.json` completion guard. + +- **checkpoint round-trip** — save → load restores weights bit-exact (max diff + 0.0), optimizer moments + step + RNG + progress (step/epoch/samples) restored; + `latest_checkpoint` selects the highest step; chained dry-run emits the correct + `afterany` links (s4 = 2, s5 = 4; s1–s3 carry no `--resume`). +- **full pipeline `DRY=1`** — all 10 stage-sweeps submit, coordinators B/C/D + generate, dependencies chain (afterany), s3/s4/s5 job.sh carry the runtime + sentinel guard, s4/s5 emit chained links. `pipeline_analyze` smoke: `best-lr` + picks the min-val-loss LR, `e1` returns PASS/exit 0, `gates` reports G1/G2. + +**Not yet exercised (gates before the full sweep):** the smoke validated the +*library* path but rebuilt the loop in a scratch script — `train.py`'s `main()` +has **never run end-to-end** (arg parsing, the in-main table load, +`build_dataloaders`, `evaluate`, wandb/tb logging, and a real +checkpoint→interrupt→resume cycle through the training loop), and `--dry-run` +only writes `job.sh`. So the integrated path is unproven. + +## Code changes made (experiments/vision) + +- `maxp_timm.py` — replaced the obsolete named-model ladder (vit-s/b/l, depth + drifted to 24 at large, dropout/drop-path scaling with width) with a clean + width-only s1–s5 ladder; `create_model` passes embed_dim/depth/num_heads. + Biases kept at stock timm ViT defaults (land in `_other`, constant LR). +- `train.py` — added `maxP-meas`, `--measure-only`, `--alignment-table`, + `--c-ema`; dropped dynamic `maxP`/`mup-full`; `dynamic = measure_only`; + alignment overrides + measure_only wired into `Parametrization`. Implemented + `--resume` (was parsed but ignored): loads latest checkpoint, restores + weights/optimizer/RNG/progress, completes to a step budget. Checkpointing on by + default (every 1000 steps, keep latest 1, atomic save). +- `export_alignment.py` (new) — exact-name table from `metrics.json` (avoids the + tensorboard-on-aarch64 problem from the LM run). +- `coord_check_vit.py` — width-only widths + `alignment="no"` (was the wrong + `"full"`); now the correctness gate. +- `launch_sweep.py` — two-arm ladder, fixed `--epochs` budget, constant batch + size, completion guard, method-specific args (`--alignment-table` / + `--measure-only`), and **resume chaining** (`--chain`; s4 = 2, s5 = 4 links via + `--dependency=afterany`, with sbatch-failure abort and job-id capture). No + `--run-date`. `utils.py` — atomic checkpoint save/prune + `load_checkpoint` / + `latest_checkpoint`. +- `hf_vision_data.py` — **fixed missing train shuffle** (`shuffle=is_train` + + optional generator for reproducible per-pass order). +- `pipeline.sh` (new) — full submit-only DAG (stages A–E + coordinators B/C/D, + sentinel gating), mirroring `experiments/lm/pipeline.sh`. +- `pipeline_analyze.py` (new) — coordinator decision helpers (`best-lr`, `e1`, + `gates`) reading `metrics.json` (no tensorboard dependency). +- `launch_sweep.py` — added the DAG primitives: `--tag`, `--job-ids-file`, + `--dependency` (first-link DAG edge), `--require` (runtime sentinel guards with + `{lr}` templating), plus the measure-only+chain guard. +- `run.sh` / `run_debug.sh` — updated to the measure/export/sweep phases. + +## Open items / decisions for you + +1. **s5 wall-time** — the hard gate (see Single-GPU feasibility); validate + sec/step on the server before committing s5; fallbacks = fewer epochs / smaller + batch / enable checkpointing / cap at s4. +2. **Budget magnitude** — 4 epochs × batch 512 are proposed defaults. More epochs + = better absolute accuracy but more compute. Adjust if you want longer/shorter. +3. **Measure base = s2** (as in LM). Could measure at s1 instead (cheaper); s2 + chosen to match the LM protocol. + +## Not done until accepted — launch gates + +No SLURM jobs submitted. On acceptance: +1. scp changed files to the server. +2. **GATE #1 — dataset loads map-style + is pre-staged.** `timm/imagenet-12k-wds` + is WebDataset; confirm `load_dataset(..., streaming=False)` returns a map-style + `Dataset` (`hasattr(ds, "__len__")`, `ds[0]["jpg"]` decodes) — if it returns an + `IterableDataset`, `len`/indexing/`shuffle=True` all break and the data layer + needs a rewrite. Also **pre-stage the dataset into `HF_HOME`** (the non-streaming + prepare is a ~1 TB download + Arrow conversion); do it once before launch so the + 84 stage-A/C jobs don't each trigger or race the prepare (and so the train.py + smoke doesn't hang for hours on a cold cache). +3. **Disk budget.** Checkpoints are fp32 model+Adam: s5 ≈ 30 GB, s4 ≈ 7.5 GB each. + `keep_latest_k` is 1 at every scale (set by launch_sweep; atomic save means one + complete checkpoint always survives an interrupt), but concurrent s4/s5 + runs + the ~1 TB dataset can still total hundreds of GB — check the PLG quota + before launch (a disk-full stalls the run; the atomic save protects the prior + checkpoint). +4. **Real end-to-end run** (not just `launch_sweep --dry-run`): execute `train.py` + once per path on a tiny config (debug/s1, ~50 steps) and confirm by eye — + measure-only populates `align/*` rows in `metrics.json`; maxP-meas loads + `s2_align.json`, applies per-layer `c`, and steps. +5. **Resume verification** (mirrors the LM `verify_resume.sh`): run a short job, + interrupt it mid-training, relaunch with `--resume`, and confirm it loads the + latest checkpoint and trains through to `final_metrics.json`. Validates the + full checkpoint→interrupt→resume cycle that s4/s5 chains depend on. +6. **s5 throughput projection** — ~100 s5 steps, sec/step → project the full run, + size the chain length (s4 = 2, s5 = 4). Measure with the production + `--val-interval 500 --val-steps 50` (validation is ~10% overhead at those + settings; do not project off a config with heavier eval). +7. Then launch phase 1 (measure → export → sweep). + +**Post-run completeness guard:** before treating any results as final, check that +**every** run dir has `final_metrics.json`. If a chain was too short (last link +hit the wall mid-train, no successor), the run is silently incomplete — nothing +alerts. Re-running `launch_sweep` for that scale (with `--chain`) appends links +and the guard no-ops the already-complete runs. diff --git a/docs/walkthrough.md b/docs/walkthrough.md index cdcc5b6..b113f52 100644 --- a/docs/walkthrough.md +++ b/docs/walkthrough.md @@ -660,7 +660,7 @@ model definitions and training scripts. | `mlp.py` | Vanilla MLP (no parametrization, baseline) | | `parametrized_mlp.py` | MLP with `ParametrizedModule` wrappers | | `train.py` | LR transfer demo: SP vs muP across widths | -| `mup_vs_maxp.py` | Conservative vs muP static alignment comparison | +| `mup_vs_conservative.py` | Conservative vs muP static alignment comparison | ### `vit_example/` — ViT baselines @@ -718,22 +718,11 @@ python -m pytest tests/ -v --tb=short | File | Tests | What it covers | |------|-------|---------------| -| `test_dag.py` | 23 | DAG tracing, DAG solver, Parametrization+DAG integration | -| `test_dag_solver_properties.py` | 18 | Analytical correctness, optimality, per-op differentiation | -| `test_alignment_new.py` | 11 | `compute_alignment()` edge cases and properties | -| `test_step.py` | 11 | `capture_initial()`, `step()`, warmup, interval, optimizer sync | - -### Legacy tests (old maxp package) - -| File | Tests | What it covers | -|------|-------|---------------| -| `test_parametrization.py` | 14+ | `get_abc_parametrization`, `create_param_groups` | -| `test_scheduler.py` | 20+ | `MaxPScheduler`, WSD warmup/decay | -| `test_init_weights.py` | — | `ScaledLinear`, `initialize_abc_weights` | -| `test_smoke.py` | — | End-to-end smoke tests | -| `test_solver.py` | 2 | Old `find_c_adam`, `find_c_sgd` | -| `test_tracer.py` | 3 | Old `Tracer` class | -| `test_alignment.py` | 3 | Old alignment computation | +| `test_dag.py` | 22 | DAG tracing, DAG solver, Parametrization+DAG integration | +| `test_dag_solver_properties.py` | 19 | Analytical correctness, optimality, per-op differentiation | +| `test_step.py` | 47 | `capture_initial()`, `step()`, warmup, interval, optimizer sync | +| `test_alignment.py` | 5 | `compute_alignment()` edge cases and properties | +| `test_optimizations_integration.py` | 10 | End-to-end optimization integration | --- diff --git a/docs/width_ladder_results.md b/docs/width_ladder_results.md new file mode 100644 index 0000000..56c672d --- /dev/null +++ b/docs/width_ladder_results.md @@ -0,0 +1,106 @@ +# Width-ladder muTransfer results: maxP-meas vs muP (mup-no) + +Run root (helios): `/net/storage/pr3/plgrid/plggadlers/maxP/runs/2026-06-12_*` +Status: **s1–s5 complete** (2026-06-25). + +## Setup + +- Architecture: LLaMA-3, width-only ladder, `n_layers=12`, `seq_len=3072`, global batch 16. +- Compute budget: `steps = 20 × non_embed_params / tokens_per_step` (Chinchilla-ish, fixed tokens-per-param). +- Train corpus: fineweb-edu. Validation corpus: **c4-validation** (OOD vs train — fair across arms, but not in-distribution). +- Two arms, identical everything except parametrization: + - **mup-no** = standard muP (alignment = 0.5 uniform). + - **maxP-meas** = static per-layer c-table from a single passive alignment measurement at **s2**, applied unchanged at every scale (`transfer-s2`). No online re-solving. +- Metric: validation loss (last `validate` point) and mean train loss over last 1k steps. LR grid swept per scale; s5 ran only the transferred optimum (lr 3e-2). + +| scale | non-embed params | total steps | +|-------|------------------|-------------| +| s1 | 11.4M | — | +| s2 | 40.9M | — | +| s3 | 163.6M | — | +| s4 | 654.4M | — | +| s5 | 2.62B | 1,065,002 | + +Seeds: s1/s2 ×3, s3/s4 ×2, s5 ×1. + +## Results — validation loss (best LR per arm, mean over seeds) + +| scale | params | mup-no | maxP-meas | **Δ (nat)** | +|-------|--------|--------|-----------|-------------| +| s1 | 11M | 4.847 (3e-2) | 4.756 (3e-2) | **0.091** | +| s2 | 41M | 4.319 (3e-2) | 4.223 (3e-2) | **0.096** | +| s3 | 164M | 3.945 (3e-2) | 3.822 (3e-2) | **0.123** | +| s4 | 654M | 3.583 (1e-2)\* | 3.487 (3e-2) | **0.096** | +| s5 | 2.62B | 3.322 (3e-2 only)† | 3.240 (3e-2) | **0.083** | + +\*s4 mup-no best (1e-2) sits at the **grid edge** (lowest LR run) — unbracketed, so the s4 margin is an upper bound; a 3e-3 point could narrow it. At matched 3e-2, s4 Δ = 0.102. + +†s5 ran **only** lr3e-2 for both arms (not a swept best). By s4, muP's own optimum had slid below 3e-2 (own-best 1e-2), so the s5 muP number is muP at a likely-too-high LR while maxP-meas is at its genuine optimum. Δ_s5 = 0.083 is therefore an **upper bound on the loss-floor gap** (muP at its own optimum would land ~3.31, narrowing it to ~0.07). Conversely, that same fact is itself a result — see conclusion 3. + +## Results — train loss (mean last 1k steps, lr 3e-2; s4 at each arm's own best) + +| scale | mup-no | maxP-meas | **Δ (nat)** | +|-------|--------|-----------|-------------| +| s1 | 4.442 | 4.301 | **0.141** | +| s2 | 3.809 | 3.678 | **0.131** | +| s3 | 3.308 | 3.178 | **0.130** | +| s4 | 2.921\* | 2.826 | **0.096** | +| s5 | 2.685 | 2.595 | **0.090** | + +\*s4 mup-no own-best = lr1e-2 (2.921); lr3e-2 = 2.928 (gap 0.007). maxP-meas own-best clearly lr3e-2 at s4/s5. + +## Conclusions + +1. **maxP-meas beats muP at every scale, ~0.1 nat, roughly flat across 230× params (11M → 2.6B).** + Val Δ: 0.091 / 0.096 / 0.123 / 0.096 / 0.083. Train Δ: 0.141 / 0.131 / 0.130 / 0.096 / 0.090. + Val is flat-ish/humped (s3 bump = one 2-seed point, treat as noise). Train Δ is **monotone declining** — 0.141 → 0.131 → 0.130 → 0.096 → 0.090, ~36% relative drop s1→s5 — but stays clearly positive, settling to a ~0.08–0.09 nat floor at 2.6B. Honest phrasing: **positive at every scale, gently decaying in train, ~0.08–0.09 nat at 2.6B — not decaying to zero**. + +2. **Strongest finding — one alignment measurement transfers both directions.** The single s2-measured alignment table, applied unchanged, wins at s1 (down), s3/s4/s5 (up). One passive measurement → width-correct c-table that holds across the full ladder. + +3. **LR optimum is width-stable for maxP-meas — partly a muTransfer-robustness result.** maxP-meas optimum stays 3e-2, cleanly bracketed at s1–s4 (s4: 1e-2 > 3e-2 < 1e-1) and still optimal-and-winning at s5. muP optimum holds 3e-2 through s3, then slides to the grid edge by s4 (best at lowest LR run) — i.e. **muP's transferred LR stops being optimal at scale while maxP-meas's does not**. So part of the s5 head-to-head is not "maxP-meas reaches a lower floor" but "muP failed to keep its muTransfer-predicted 3e-2 optimal" — a robustness claim, stronger and distinct from the loss-floor gap. + +4. **Mechanism still open — not adjudicated by this data.** No run varies readout LR independently of global LR. The maxP-meas c-table is readout-dominated (output Δc ≈ +0.228 vs muP; all other layers within ±0.04). Width-consistency is *consistent with* width-correct per-layer scaling, but a flat-enough constant readout-LR cut could also hold ~0.1 nat. The clean falsification arm (**mup-no + tuned readout LR**, 2D sweep) was not run. Claim the ~0.1 nat as real and width-stable; leave the *mechanism* (width-correct per-layer scaling vs a constant readout cut) unproven. + +## Is the preserved LR optimum (mup-no ↔ maxP-meas) a guarantee or luck? + +Both arms share the **same** prefactor optimum at the tuning scales — `mup-no s1=s2=3e-2`, `maxP-meas s1=s2=3e-2` (pipeline `report.txt`, `[consistent]`), and it `[TRANSFERS]` to s3. So switching mup-no → maxP-meas re-scales per-layer LRs but does **not** move the global-LR argmin. Is that guaranteed? **Neither a free theorem nor luck — a conditional structural result, and the condition is empirically met here.** + +**Reduction.** At fixed tuning width `n`, the c-table only changes the width-*exponent*, so it multiplies each layer's LR by a **constant** `k_l = n^(−Δc_l)`. The question is purely: rescale per-layer LRs `{η·m_l} → {η·m_l·k_l}`, does the optimal global `η*` stay put? + +**In general — no guarantee.** Rescaling *relative* per-layer LRs deforms the loss-vs-η landscape; `η*` can move. No theorem makes it invariant under arbitrary per-layer rescaling. + +**Why it held here — two semi-principled facts:** + +1. **The correction is readout-dominated** — `Δc_readout ≈ +0.228`, every other layer within `±0.04` → `k_l ≈ 1` except the readout. Predicted, not accidental: muP's `α=0.5` alignment assumption is approximately right for the **bulk hidden layers** (muP's basis) and wrong mainly at the **boundary** (readout), so measured corrections concentrate there. +2. **Readout LR is weakly coupled to `η*`** — the usable-LR window (hence the argmin) is set by the hidden stack's feature-learning dynamics; the readout is one top linear layer, and its c is solver-**pinned** by the output-stability constraint. Rescaling it lowers the loss **floor** but barely moves the **location** of the optimum. + +**Deeper near-principle.** Both arms are maximal-update-class parametrizations. The optimal global LR is an intrinsic property of the optimization problem (arch + optimizer + data) that muP factors the width-scaling out of; two *correct* MUP-family parametrizations target the same joint-"all-layers-maximal" `η`, differing only in **how** correct they are → they differ in the **loss floor**, not the **argmin**. The residual gap between them *is* mup-no's alignment error — small and concentrated → small `η*` shift. + +**Evidence the optimum is decisive (not a coarse-grid near-tie).** Per-LR loss at s2 (`.out`, seed s1): + +| LR | mup-no | maxP-meas | +|----|--------|-----------| +| 1e-2 | 3.988 | 4.040 | +| **3e-2** | **3.786** | **3.675** | +| 1e-1 | 3.945 | 3.740 | + +3e-2 wins by ~0.1–0.2 nat over both neighbours in **both** arms — well above seed noise. + +**Where it breaks (paper-honest boundary):** + +- If a regime makes muP's alignment assumption badly wrong across **many bulk layers** (large Δc spread through depth, not just readout), `η*` **will** move → re-tuning is mandatory. +- "Preserved" is only **within one grid step** (×3 spacing); a sub-3× drift is real-but-invisible. +- Assumes both stay in the stable feature-learning class. + +**Implication.** The small-scale maxP-meas re-tune (Stage C) is **not ceremony** — it is the *test* that the readout-concentrated + decoupled precondition holds. It is **not** a priori safe to skip for a new arch/dataset (e.g. **vision** — its c-table may not be readout-dominated; verify before trusting a single tune). Correct claim for the paper: *"the LR optimum is preserved because the measured correction is readout-concentrated and the readout LR is decoupled from the global optimum"* — **not** *"maxP-meas preserves the LR optimum"* unconditionally. + +## Caveats + +- **OOD validation**: val = c4, train = fineweb-edu. Fair across arms (identical), but val is not in-distribution. Train margin runs ~0.03–0.05 larger than val at small scale (maxP-meas fits train harder; part doesn't transfer to c4) — expected. +- **Coarse LR grid at large scale**: s3/s4 swept ~3 LRs, s5 only 1 → "stable optimum" claims rest on a coarse grid. muP s4 optimum unbracketed (see note above). +- **Thin seeds at scale**: s5 single seed → no error bar on the headline 2.6B number. +- One s4 cell (mupno lr3e-2 s2) lost its `.out` train history to a resume-job guard-skip truncation; val/checkpoint intact, seed-s1 twin used for train. + +## Operational note + +18 leftover SLURM jobs (4 × s4, 14 × s5) remain PD — chain links whose target dirs already have `COMPLETED`; they hit the launch_sweep guard and exit on start. Safe to `scancel` to free scheduling slots. diff --git a/examples/mlp_example/mup_vs_maxp.py b/examples/mlp_example/mup_vs_conservative.py similarity index 95% rename from examples/mlp_example/mup_vs_maxp.py rename to examples/mlp_example/mup_vs_conservative.py index 96bed7d..aca906c 100644 --- a/examples/mlp_example/mup_vs_maxp.py +++ b/examples/mlp_example/mup_vs_conservative.py @@ -7,8 +7,8 @@ 2. muP — static "full" alignment (alpha=1, omega=0.5, u=1). Usage: - python examples/mlp_example/mup_vs_maxp.py - python examples/mlp_example/mup_vs_maxp.py --width 256 --steps 1000 + python examples/mlp_example/mup_vs_conservative.py + python examples/mlp_example/mup_vs_conservative.py --width 256 --steps 1000 """ import argparse @@ -91,7 +91,7 @@ def smooth(values, window=20): return np.convolve(values, kernel, mode="valid").tolist() -def plot_results(cons_losses, mup_losses, filename="mup_vs_maxp.png", window=20): +def plot_results(cons_losses, mup_losses, filename="mup_vs_conservative.png", window=20): import matplotlib.pyplot as plt fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(13, 5)) @@ -135,7 +135,7 @@ def main(): parser.add_argument("--lr-prefactor", type=float, default=0.01) parser.add_argument("--seed", type=int, default=42) parser.add_argument("--no-plot", action="store_true") - parser.add_argument("--output", type=str, default="mup_vs_maxp.png") + parser.add_argument("--output", type=str, default="mup_vs_conservative.png") args = parser.parse_args() device = get_device() diff --git a/experiments/lm/export_alignment.py b/experiments/lm/export_alignment.py new file mode 100644 index 0000000..fd041c8 --- /dev/null +++ b/experiments/lm/export_alignment.py @@ -0,0 +1,94 @@ +#!/usr/bin/env python3 +"""Export a measured alignment table from a run's tensorboard logs. + +Reads align/{z0_dW,dZ_w0,dZ_dW}/ scalars from a source run (typically +``--method mup-no --measure-only``), averages each layer over the last +``--last-frac`` of training, aggregates per op-type suffix (median over +layers), and writes a JSON table consumable by ``--alignment-table``: + + {"wq": [a, w, u], "wk": [...], ..., "output": [...]} + +Keys are leaf suffixes matched by Parametrization.alignment_overrides. +tok_embeddings is excluded: its alignment is never measured (fixed fan-in), +so it keeps the "no" preset. + +Usage: + + python experiments/lm/export_alignment.py runs/ -o s1_align.json + python experiments/lm/export_alignment.py runs/ runs/ -o avg.json +""" + +from __future__ import annotations + +import argparse +import glob +import json +import re + +import numpy as np + +METRICS = ("z0_dW", "dZ_w0", "dZ_dW") +SUFFIXES = ("wq", "wk", "wv", "wqkv", "wo", "w1", "w2", "w3", "output") + + +def _suffix(layer: str) -> str | None: + if layer == "output": + return "output" + m = re.search(r"(wq|wk|wv|wqkv|wo|w1|w2|w3)$", layer) + return m.group(1) if m else None + + +def collect(run_dir: str, last_frac: float) -> dict[tuple[str, str], list[float]]: + """Return (metric, suffix) -> per-layer steady-state means for one run.""" + from tensorboard.backend.event_processing.event_accumulator import EventAccumulator + + events = glob.glob(f"{run_dir}/tb/*/events.*") + if not events: + raise FileNotFoundError(f"No tensorboard events under {run_dir}/tb/") + + ea = EventAccumulator(events[0], size_guidance={"scalars": 0}) + ea.Reload() + out: dict[tuple[str, str], list[float]] = {} + for tag in ea.Tags()["scalars"]: + if not tag.startswith("align/"): + continue + _, metric, layer = tag.split("/", 2) + suffix = _suffix(layer) + if suffix is None: + continue + vals = np.array([e.value for e in ea.Scalars(tag)]) + steady = vals[int(len(vals) * (1 - last_frac)):] + out.setdefault((metric, suffix), []).append(float(steady.mean())) + if not out: + raise ValueError(f"No align/* scalars in {run_dir} — was it run with --measure-only or --method maxP?") + return out + + +def main() -> None: + p = argparse.ArgumentParser(description="Export alignment table from run logs") + p.add_argument("run_dirs", nargs="+", help="Source run director(ies); multiple are pooled") + p.add_argument("-o", "--output", required=True, help="Output JSON path") + p.add_argument("--last-frac", type=float, default=0.25, + help="Fraction of training (from the end) to average over") + args = p.parse_args() + + pooled: dict[tuple[str, str], list[float]] = {} + for run_dir in args.run_dirs: + for k, v in collect(run_dir, args.last_frac).items(): + pooled.setdefault(k, []).extend(v) + + table: dict[str, list[float]] = {} + for suffix in SUFFIXES: + if (METRICS[0], suffix) not in pooled: + continue + table[suffix] = [round(float(np.median(pooled[(m, suffix)])), 4) for m in METRICS] + + with open(args.output, "w") as f: + json.dump(table, f, indent=2) + print(f"Wrote {args.output}:") + for k, v in table.items(): + print(f" {k:8s} a={v[0]:.4f} w={v[1]:.4f} u={v[2]:.4f}") + + +if __name__ == "__main__": + main() diff --git a/experiments/lm/launch_sweep.py b/experiments/lm/launch_sweep.py index 9f92424..5ac9b76 100644 --- a/experiments/lm/launch_sweep.py +++ b/experiments/lm/launch_sweep.py @@ -5,8 +5,8 @@ python experiments/lm/launch_sweep.py \\ --scale s3 \\ - --methods maxP mup-full mup-no \\ - --lrs 3e-4 1e-3 3e-3 1e-2 3e-2 1e-1 3e-1 \\ + --methods mup-no \\ + --lrs 1e-3 3e-3 1e-2 3e-2 1e-1 3e-1 \\ --seeds 1 2 \\ --runs-dir /net/storage/pr3/plgrid/plggadlers/maxP/runs \\ --dataset HuggingFaceFW/fineweb-edu \\ @@ -32,33 +32,41 @@ # Per-scale defaults # --------------------------------------------------------------------------- +# Wall times: total train FLOPs at 30% MFU on GH200 (494 TFLOPS BF16 dense +# -> ~148 TF/GPU effective), x1.5 safety margin, rounded up: +# s1 0.2h | s2 1.7h | s3 18.4h | s4 (4 GPU) 55h | s5 (4 GPU) ~730h +# The cluster caps wall time at 48h, so s4/s5 checkpoint periodically and are +# submitted as a chain of dependent jobs ("chain" below): each job resumes +# from the latest checkpoint and exits when the wall limit hits; the next in +# the chain continues. torchtitan auto-resumes when a checkpoint exists. SCALE_CONFIGS = { - "s1": {"wall": "02:00:00", "nodes": 1, "gpus": 1}, - "s2": {"wall": "04:00:00", "nodes": 1, "gpus": 1}, - "s3": {"wall": "24:00:00", "nodes": 1, "gpus": 1}, - "s4": {"wall": "48:00:00", "nodes": 1, "gpus": 4}, - "s5": {"wall": "48:00:00", "nodes": 1, "gpus": 4}, + "s1": {"wall": "01:00:00", "nodes": 1, "gpus": 1, "ckpt": None, "chain": 1}, + "s2": {"wall": "03:00:00", "nodes": 1, "gpus": 1, "ckpt": None, "chain": 1}, + "s3": {"wall": "24:00:00", "nodes": 1, "gpus": 1, "ckpt": None, "chain": 1}, + "s4": {"wall": "48:00:00", "nodes": 1, "gpus": 4, "ckpt": 5000, "chain": 2}, + "s5": {"wall": "48:00:00", "nodes": 1, "gpus": 4, "ckpt": 2500, "chain": 12}, } SCALE_METHODS = { - "s1": ["maxP", "mup-full", "mup-no"], - "s2": ["maxP", "mup-full", "mup-no"], - "s3": ["maxP", "mup-full", "mup-no"], - "s4": ["maxP", "mup-full", "mup-no"], - "s5": ["maxP", "mup-full", "mup-no"], + "s1": ["mup-no"], + "s2": ["mup-no"], + "s3": ["mup-no"], + "s4": ["mup-no"], + "s5": ["mup-no"], } SCALE_SEEDS = { "s1": [1, 2, 3], - "s2": [1, 2], + "s2": [1, 2, 3], "s3": [1, 2], - "s4": [1], + "s4": [1, 2], "s5": [1], } -ALL_LRS = [3e-4, 1e-3, 3e-3, 1e-2, 3e-2, 1e-1, 3e-1] - -S5_SINGLE_LR = [1e-2] +# 7-point grid at half-decade spacing; architecture-agnostic default for +# every scale. Large-scale sweeps should pass an explicit --lrs subset +# (the pipeline gates per-LR via --require sentinels instead). +ALL_LRS = [1e-3, 3e-3, 1e-2, 3e-2, 1e-1, 3e-1, 1e0] # --------------------------------------------------------------------------- @@ -84,6 +92,8 @@ export OMP_NUM_THREADS=16 export HF_HOME="$${HF_HOME:-/net/scratch/hscra/plgrid/plgmwojnar/hf}" +export HF_HUB_DOWNLOAD_TIMEOUT="$${HF_HUB_DOWNLOAD_TIMEOUT:-120}" +export HF_HUB_ETAG_TIMEOUT="$${HF_HUB_ETAG_TIMEOUT:-120}" export WANDB_PROJECT="$${WANDB_PROJECT:-maxP-lm}" export WANDB_RUN_NAME="maxp_${scale}_${method_tag}_lr${lr_tag}_s${seed}" @@ -98,6 +108,12 @@ echo "Master Port: $${MASTER_PORT}" +# Chained jobs: skip if an earlier link already finished this run. +if [ -f "${output_dir}/COMPLETED" ]; then + echo "Run already completed — skipping." + exit 0 +fi +${runtime_guards} torchrun --nproc_per_node=${gpus_per_node} --master_addr=$${MASTER_ADDR} --master_port=$${MASTER_PORT} experiments/lm/train.py \\ --scale ${scale} \\ --method ${method} \\ @@ -109,6 +125,7 @@ --num-workers ${num_workers} \\ --prefetch-factor ${prefetch_factor} \\ --hf-assets-path ${hf_assets_path} \\ + ${extra_args} \\ --output-dir ${output_dir} """) @@ -152,20 +169,50 @@ def main() -> None: help="Path to local HF tokenizer assets directory") p.add_argument("--venv-path", required=True, help="Path to Python venv") p.add_argument("--repo-path", required=True, help="Path to maxP repo root") + p.add_argument("--measure-only", action="store_true", + help="Pass --measure-only to train.py (alignment source runs)") + p.add_argument("--alignment-table", default=None, + help="Pass --alignment-table to train.py (maxP-meas runs)") + p.add_argument("--tag", default=None, + help="Suffix appended to run names (e.g. 'meas', 'transfer-s1')") + p.add_argument("--chain", type=int, default=None, + help="Submit N dependent jobs per run (default: per-scale value); " + "later jobs resume from the latest checkpoint") + p.add_argument("--job-ids-file", default=None, + help="Append submitted SLURM job IDs (one per line) for " + "downstream dependency wiring") + p.add_argument("--dependency", default=None, + help="sbatch --dependency spec applied to each run's first " + "job (e.g. 'afterany:123:456')") + p.add_argument("--require", action="append", default=[], + help="File that must exist at job runtime or the job exits " + "cleanly; '{lr}' is replaced by the run's lr tag. " + "Repeatable.") p.add_argument("--dry-run", action="store_true", help="Print sbatch commands without submitting") p.add_argument("--resume", action="store_true", help="Submit all runs even if a checkpoint exists (resume interrupted training)") + p.add_argument("--run-date", default=None, + help="Override the YYYY-MM-DD prefix in run names (default: " + "today). Use to resubmit into an existing run's dir so " + "training resumes from its checkpoint.") args = p.parse_args() scale = args.scale sc = SCALE_CONFIGS[scale] methods = args.methods or SCALE_METHODS[scale] - lrs = args.lrs or (S5_SINGLE_LR if scale == "s5" else ALL_LRS) + if "maxP-meas" in methods and not args.alignment_table: + p.error("--methods maxP-meas requires --alignment-table") + lrs = args.lrs or ALL_LRS seeds = args.seeds or SCALE_SEEDS[scale] + chain = args.chain if args.chain is not None else sc["chain"] + if chain > 1 and (args.measure_only or "maxP" in methods): + p.error("--measure-only and dynamic maxP require runs that finish in " + "one job (chain=1): resuming re-snapshots z0/w0 from the " + "checkpoint and corrupts alignment measurement") gpus_per_node = args.gpus_per_node if args.gpus_per_node is not None else sc["gpus"] - today = date.today().strftime("%Y-%m-%d") + today = args.run_date or date.today().strftime("%Y-%m-%d") submitted = skipped = 0 for method in methods: @@ -174,6 +221,8 @@ def main() -> None: run_name = ( f"{today}_{scale}_{_method_tag(method)}_lr{_lr_tag(lr)}_s{seed}" ) + if args.tag: + run_name += f"_{args.tag}" out_dir = Path(args.runs_dir) / run_name if not args.resume and _has_checkpoint(out_dir): @@ -185,6 +234,25 @@ def main() -> None: dataset_path_arg = f"--dataset-path {args.dataset_path}" if args.dataset_path else "" + extra = [] + if args.measure_only: + extra.append("--measure-only") + if args.alignment_table: + extra.append(f"--alignment-table {args.alignment_table}") + if sc["ckpt"]: + extra.append(f"--checkpoint-interval {sc['ckpt']}") + extra_args = " ".join(extra) + + # Runtime gating: lets jobs be pre-submitted before the + # decision that enables them exists (login node submits only). + guards = [] + for path in args.require: + path = path.replace("{lr}", _lr_tag(lr)) + guards.append( + f'if [ ! -f "{path}" ]; then ' + f'echo "gate: {path} missing — skipping."; exit 0; fi') + runtime_guards = ("\n" + "\n".join(guards) + "\n") if guards else "" + script_content = SLURM_TEMPLATE.substitute( scale=scale, method_tag=_method_tag(method), @@ -203,19 +271,40 @@ def main() -> None: num_workers=args.num_workers, prefetch_factor=args.prefetch_factor, hf_assets_path=args.hf_assets_path, + extra_args=extra_args, + runtime_guards=runtime_guards, ) script_path = out_dir / "job.sh" script_path.write_text(script_content) - if args.dry_run: - print(f"[dry] sbatch {script_path}") - else: + prev_job_id = None + for link in range(chain): + if prev_job_id: + dep = ["--dependency", f"afterany:{prev_job_id}"] + elif args.dependency: + dep = ["--dependency", args.dependency] + else: + dep = [] + if args.dry_run: + dep_str = f" {' '.join(dep)}" if dep else "" + print(f"[dry] sbatch{dep_str} {script_path}") + continue result = subprocess.run( - ["sbatch", str(script_path)], + ["sbatch", *dep, str(script_path)], capture_output=True, text=True, ) - print(f"[submit] {run_name} → {result.stdout.strip()}") + if result.returncode != 0: + raise SystemExit( + f"sbatch failed for {run_name} " + f"(link {link + 1}/{chain}): {result.stderr.strip()}") + out = result.stdout.strip() + prev_job_id = out.split()[-1] if out else None + if args.job_ids_file and prev_job_id: + with open(args.job_ids_file, "a") as f: + f.write(f"{prev_job_id}\n") + chain_tag = f" [chain {link + 1}/{chain}]" if chain > 1 else "" + print(f"[submit] {run_name}{chain_tag} → {out}") submitted += 1 action = "would submit" if args.dry_run else "submitted" diff --git a/experiments/lm/maxp_converter.py b/experiments/lm/maxp_converter.py index 86d44b9..c7c03f0 100644 --- a/experiments/lm/maxp_converter.py +++ b/experiments/lm/maxp_converter.py @@ -100,13 +100,15 @@ class MaxPConverter(Configurable, ModelConverter): @dataclass(kw_only=True, slots=True) class Config(Configurable.Config): - method: str = "maxP" # "maxP" | "mup-full" | "mup-no" + method: str = "maxP" # "maxP" | "mup-no" | "maxP-meas" lr_prefactor: float = 1e-3 - alignment_warmup: int = 10 # maxP only - solve_interval: int = 100 # maxP only - sample_size: int = 32 # maxP only + alignment_warmup: int = 10 # maxP / measure_only + solve_interval: int = 100 # maxP / measure_only + sample_size: int = 32 # maxP / measure_only c_ema: float = 0.0 # maxP only use_training_activations: bool = False # maxP only + measure_only: bool = False # measure + log alignment, never touch LRs + alignment_table: str | None = None # JSON path; required for maxP-meas def __init__( self, @@ -120,19 +122,33 @@ def __init__( def convert(self, model: nn.Module) -> None: install_pm_wrappers(model) cfg = self._cfg + if cfg.method not in ("maxP", "mup-no", "maxP-meas"): + raise ValueError(f"Unknown method '{cfg.method}'") + + if cfg.measure_only and cfg.method == "maxP": + raise ValueError("measure_only contradicts dynamic method 'maxP'; use 'mup-no'") + + alignment_overrides = None + if cfg.method == "maxP-meas": + if cfg.alignment_table is None: + raise ValueError("method 'maxP-meas' requires alignment_table") + alignment_overrides = _load_alignment_table(cfg.alignment_table) + param = Parametrization( model, sample_input=_make_sample_input(model), - alignment="full" if "full" in cfg.method else "no", + alignment="no", lr_prefactor=cfg.lr_prefactor, + alignment_overrides=alignment_overrides, warmup_steps=cfg.alignment_warmup, solve_interval=cfg.solve_interval, sample_size=cfg.sample_size, c_ema=cfg.c_ema, + measure_only=cfg.measure_only, use_training_activations=cfg.use_training_activations, ) model._maxp_param = param - if cfg.method == "maxP": + if cfg.method == "maxP" or cfg.measure_only: model._is_dynamic = True model._maxp_align = { name: (pm.align_z0_dW, pm.align_dZ_w0, pm.align_dZ_dW) @@ -144,6 +160,19 @@ def post_optimizer_hook(self, model: nn.Module | list[nn.Module]) -> None: pass +def _load_alignment_table(path: str) -> dict[str, tuple[float, float, float]]: + """Load measured alignment from JSON written by export_alignment.py. + + Keys are exact PM names, leaf suffixes (e.g. "wo") or layer types; + values are [align_z0_dW, align_dZ_w0, align_dZ_dW]. + """ + import json + + with open(path) as f: + table = json.load(f) + return {k: tuple(v) for k, v in table.items()} + + def _make_sample_input(model: nn.Module) -> torch.Tensor | None: """Build a deterministic sample input for DAG tracing. diff --git a/experiments/lm/maxp_llama3.py b/experiments/lm/maxp_llama3.py index e53e5d9..94040c1 100644 --- a/experiments/lm/maxp_llama3.py +++ b/experiments/lm/maxp_llama3.py @@ -129,17 +129,17 @@ def parallelize_llama(model, **kwargs): return _parallelize_llama(model, **kwargs) -# Scale definitions: (dim, n_layers, n_heads, n_kv_heads[, vocab_size]) -# Approximate parameter counts (no weight tying, vocab=128256): -# debug: tiny 4-layer (CPU smoke tests, vocab=2048) -# s1: ~21M s2: ~81M s3: ~218M s4: ~1.09B s5: ~2.71B +# Scale definitions: width-only ladder for muP-style width scaling. +# Non-embed parameter counts (vocab=128256): +# debug: tiny 4-layer (CPU smoke tests) +# s1: 11.4M s2: 40.9M s3: 163.6M s4: 654.4M s5: 2.62B SCALE_CONFIGS: dict[str, dict] = { "debug": dict(dim=256, n_layers=4, n_heads=4, n_kv_heads=2), - "s1": dict(dim=512, n_layers=6, n_heads=8, n_kv_heads=4), - "s2": dict(dim=768, n_layers=10, n_heads=12, n_kv_heads=4), - "s3": dict(dim=1024, n_layers=16, n_heads=16, n_kv_heads=4), - "s4": dict(dim=2048, n_layers=20, n_heads=16, n_kv_heads=4), - "s5": dict(dim=2560, n_layers=32, n_heads=20, n_kv_heads=4), + "s1": dict(dim=256, n_layers=12, n_heads=4, n_kv_heads=1), + "s2": dict(dim=512, n_layers=12, n_heads=8, n_kv_heads=2), + "s3": dict(dim=1024, n_layers=12, n_heads=16, n_kv_heads=4), + "s4": dict(dim=2048, n_layers=12, n_heads=32, n_kv_heads=8), + "s5": dict(dim=4096, n_layers=12, n_heads=64, n_kv_heads=16), } @@ -151,8 +151,8 @@ def maxp_model_registry(scale: str, method: str, attn_backend: str = "sdpa") -> Args: scale: One of "debug", "s1" … "s5". - method: "maxP" (dynamic), "mup-full" (static, full align), or - "mup-no" (static, no align). + method: "maxP" (dynamic), "mup-no" (static, no align), or + "maxP-meas" (static, c solved from a measured alignment table). attn_backend: Attention backend ("sdpa", "flex", "varlen"). """ if scale not in SCALE_CONFIGS: diff --git a/experiments/lm/pipeline.sh b/experiments/lm/pipeline.sh new file mode 100644 index 0000000..cde5c46 --- /dev/null +++ b/experiments/lm/pipeline.sh @@ -0,0 +1,247 @@ +#!/usr/bin/env bash +# Unattended pipeline — "real-life" muTransfer protocol: +# tune everything at small scale (s1+s2), only VERIFY at s3+. +# +# LOGIN-NODE SUBMIT-ONLY DESIGN: this script only calls sbatch (no compute, +# no polling) and exits in seconds. The whole DAG is submitted upfront: +# +# stage A s1 + s2 mup-no sweeps (--measure-only) ........ 42 GPU jobs +# coord B (afterany: A) best mup-no LR per scale, export +# alignment tables, gate report G1-G3 .......... 1 small job +# stage C s1 + s2 maxP-meas sweeps — the prefactor +# RE-TUNE under the measured table ............. 42 GPU jobs +# coord C (afterany: C) best maxP-meas LR per scale, +# s1-vs-s2 consistency check, writes s3 bracket +# sentinels (3-pt bracket per arm) ............. 1 small job +# stage D s3 full-grid candidates; coord C releases only +# the 3-pt bracket around each arm's s2 optimum 28 jobs (12 train) +# coord D (afterany: D) transfer proof (s3 argmin == +# transferred LR?) + E1 verdict; releases s4/s5 1 small job +# stage E s4 full-grid candidates (chain=2, bracket +# released by coord D) + s5 full-grid candidates +# (chain=11, single LR released) ................ 210 jobs (26 train) +# +# Decisions that can't be known at submit time are enforced by runtime +# sentinel files (see launch_sweep --require). +# +# PREREQUISITE: run experiments/lm/verify_resume.sh manually and confirm +# PASS before submitting — s4/s5 rely on chained checkpoint resume. +# +# Usage (login node): bash experiments/lm/pipeline.sh +# Dry run (no sbatch): DRY=1 bash experiments/lm/pipeline.sh +# +# Idempotent-ish: launch_sweep skips runs that already have checkpoints, +# but coordinator jobs are resubmitted on every invocation — do not run +# this script twice while the pipeline is in flight. +set -euo pipefail + +# --- config (env-overridable) ---------------------------------------------- +REPO="${REPO:-/net/storage/pr3/plgrid/plggadlers/maxP}" +RUNS_DIR="${RUNS_DIR:-$REPO/runs}" +VENV="${VENV:-$REPO/.venv}" +HF_ASSETS="${HF_ASSETS:-$REPO/experiments/lm/assets/hf/Llama-3.1-8B}" +DATASET="${DATASET:-fineweb-edu}" +STATE="${STATE:-$RUNS_DIR/.pipeline}" +GATE_TOL="${GATE_TOL:-0.02}" +ACCOUNT="${ACCOUNT:-plgadlers-gpu-gh200}" +PARTITION="${PARTITION:-plgrid-gpu-gh200}" +DRY="${DRY:-0}" +RUN_DATE="${RUN_DATE:-}" + +# Candidate LR grid for s3/s4/s5 (must match launch_sweep ALL_LRS; tags in +# launch_sweep's printf %.0e format so sentinel and run names line up). +# No hard-coded optima: jobs are pre-submitted at EVERY grid LR and the +# coordinators release only the 3-point bracket (argmin +/- 1 grid step) +# computed from this experiment's own small-scale results. Non-released +# candidates exit in ~1 s at runtime. +GRID_LRS="${GRID_LRS:-1e-03 3e-03 1e-02 3e-02 1e-01 3e-01 1e00}" + +cd "$REPO" +mkdir -p "$STATE" +log() { echo "[pipeline] $*"; } + +# Login node has no python by default; load a stdlib interpreter for launch_sweep.py. +command -v python >/dev/null 2>&1 || module add GCCcore/13.2.0 Python/3.11.5 + +LAUNCH=(python experiments/lm/launch_sweep.py + --runs-dir "$RUNS_DIR" --hf-assets-path "$HF_ASSETS" + --venv-path "$VENV" --repo-path "$REPO" --dataset "$DATASET") +[ "$DRY" = "1" ] && LAUNCH+=(--dry-run) +# Pin run-date prefix so a resubmit reuses existing run dirs (skip done / resume chains). +[ -n "$RUN_DATE" ] && LAUNCH+=(--run-date "$RUN_DATE") + +submit() { # submit a standalone job script, echo job id + if [ "$DRY" = "1" ]; then echo "DRY"; log "[dry] sbatch $1" >&2; return; fi + sbatch "$@" | awk '{print $NF}' +} + +dep_from() { # ids file -> afterany:id:id:... + [ "$DRY" = "1" ] && { echo "afterany:DRY"; return; } + [ -s "$1" ] || { echo "ERROR: no job ids in $1 — refusing to wire dependencies (re-run on a dirty runs dir?)" >&2; exit 1; } + echo "afterany:$(paste -sd: "$1")" +} + +coord_header() { # $1 job name, $2 dependency + cat < "$STATE/A.ids" +log "stage A: s1 + s2 mup-no sweeps (sources)" +"${LAUNCH[@]}" --scale s1 --measure-only --tag meas --job-ids-file "$STATE/A.ids" +"${LAUNCH[@]}" --scale s2 --measure-only --tag meas --job-ids-file "$STATE/A.ids" + +# ============ COORD B: pick source LRs, export tables, gates ================ +{ + coord_header maxp_coordB "$(dep_from "$STATE/A.ids")" + cat < "$STATE/best_lr_mupno_s2.txt" +echo "\$BEST_S1" > "$STATE/best_lr_mupno_s1.txt" +python experiments/lm/export_alignment.py $RUNS_DIR/*_s2_mupno_lr\${BEST_S2}_s*_meas -o "$STATE/s2_align.json" +python experiments/lm/export_alignment.py $RUNS_DIR/*_s1_mupno_lr\${BEST_S1}_s*_meas -o "$STATE/s1_align.json" +# adjacent LR tag for gate G2 (grid must match launch_sweep ALL_LRS) +ADJ_S2=\$(python -c "g='$GRID_LRS'.split(); i=g.index('\$BEST_S2'); print(g[i+1] if i+1 < len(g) else g[i-1])") +{ + echo "=== Gate report (\$(date)) ===" + echo "mup-no optima: s1=\$BEST_S1 s2=\$BEST_S2 \$([ "\$BEST_S1" = "\$BEST_S2" ] && echo '[consistent]' || echo '[WARN: differ across scale]')" + python experiments/lm/pipeline_analyze.py gates \\ + --runs-glob "$RUNS_DIR/*_s2_mupno_lr\${BEST_S2}_s*_meas" \\ + --adjacent-glob "$RUNS_DIR/*_s2_mupno_lr\${ADJ_S2}_s*_meas" + python - "$STATE/s1_align.json" "$STATE/s2_align.json" <<'PYEOF' +import json, sys +a, b = (json.load(open(p)) for p in sys.argv[1:3]) +worst = max(abs(a[k][i] - b[k][i]) for k in set(a) & set(b) for i in range(3)) +ok = "OK - tiny-proxy story holds" if worst < 0.05 else "WARN - keep s2 as canonical source" +print(f"G3 s1-vs-s2: max |alignment delta| = {worst:.4f} [{ok}]") +PYEOF +} >> "$STATE/report.txt" 2>&1 +EOF +} > "$STATE/coordB.sh" +BJOB=$(submit "$STATE/coordB.sh") +log "coord B: $BJOB" + +# ============ STAGE C: s1+s2 maxP-meas sweeps (prefactor re-tune) =========== +: > "$STATE/C.ids" +log "stage C: s1 + s2 maxP-meas sweeps (re-tune under measured table)" +for SC in s1 s2; do + "${LAUNCH[@]}" --scale "$SC" --methods maxP-meas \ + --alignment-table "$STATE/s2_align.json" --measure-only --tag transfer-s2 \ + --dependency "afterany:$BJOB" --require "$STATE/s2_align.json" \ + --job-ids-file "$STATE/C.ids" +done + +# ============ COORD C: pick transfer LR at small scale, release s3 ========== +{ + coord_header maxp_coordC "$(dep_from "$STATE/C.ids")" + cat < "$STATE/best_lr_meas_s2.txt" +echo "\$T_S1" > "$STATE/best_lr_meas_s1.txt" +{ + echo "maxP-meas optima: s1=\$T_S1 s2=\$T_S2 \$([ "\$T_S1" = "\$T_S2" ] && echo '[consistent]' || echo '[WARN: differ across scale]')" +} >> "$STATE/report.txt" +# Release the 3-point s3 bracket (argmin +/- 1 grid step) around each arm's +# s2 optimum (the tuned scale closest to target). Edge-of-grid centers get a +# truncated 2-point bracket + warning (optimum may lie outside the grid). +release() { # \$1 stage prefix (s3), \$2 arm name, \$3 center + BRACKET=\$(python -c "g='$GRID_LRS'.split(); i=g.index('\$3'); print(' '.join(g[max(0, i-1):i+2]))") + for lr in \$BRACKET; do touch "$STATE/\${1}_go_\${2}_\${lr}"; done + EDGE="" + python -c "g='$GRID_LRS'.split(); import sys; sys.exit(0 if g.index('\$3') in (0, len(g)-1) else 1)" \ + && EDGE=" [WARN: center at grid edge - optimum may lie outside grid]" + echo "\$1 \$2 bracket released: \$BRACKET (center \$3)\$EDGE" >> "$STATE/report.txt" +} +release s3 base "\$B_S2" +release s3 transfer "\$T_S2" +EOF +} > "$STATE/coordC.sh" +CJOB=$(submit "$STATE/coordC.sh") +log "coord C: $CJOB" + +# ============ STAGE D: s3 verification brackets ============================= +: > "$STATE/D.ids" +log "stage D: s3 brackets (both arms, gated per LR)" +"${LAUNCH[@]}" --scale s3 --lrs $GRID_LRS --measure-only --tag meas \ + --dependency "afterany:$CJOB" --require "$STATE/s3_go_base_{lr}" \ + --job-ids-file "$STATE/D.ids" +"${LAUNCH[@]}" --scale s3 --methods maxP-meas --alignment-table "$STATE/s2_align.json" \ + --lrs $GRID_LRS --measure-only --tag transfer-s2 \ + --dependency "afterany:$CJOB" --require "$STATE/s3_go_transfer_{lr}" \ + --job-ids-file "$STATE/D.ids" + +# ============ COORD D: transfer proof + E1 verdict, release s4/s5 =========== +{ + coord_header maxp_coordD "$(dep_from "$STATE/D.ids")" + cat <> "$STATE/report.txt" +if ! python experiments/lm/pipeline_analyze.py e1 \\ + --baseline-glob "$RUNS_DIR/*_s3_mupno_lr*_meas" \\ + --transfer-glob "$RUNS_DIR/*_s3_maxPmeas_lr*_transfer-s2" \\ + --tol $GATE_TOL >> "$STATE/report.txt" 2>&1; then + echo "E1 FAILED - s4/s5 will self-skip" >> "$STATE/report.txt" + exit 0 +fi +# s4 brackets re-centered on the s3 argmin (best knowledge at release time) +release() { # \$1 stage prefix, \$2 arm name, \$3 center + BRACKET=\$(python -c "g='$GRID_LRS'.split(); i=g.index('\$3'); print(' '.join(g[max(0, i-1):i+2]))") + for lr in \$BRACKET; do touch "$STATE/\${1}_go_\${2}_\${lr}"; done + EDGE="" + python -c "g='$GRID_LRS'.split(); import sys; sys.exit(0 if g.index('\$3') in (0, len(g)-1) else 1)" \ + && EDGE=" [WARN: center at grid edge - optimum may lie outside grid]" + echo "\$1 \$2 bracket released: \$BRACKET (center \$3)\$EDGE" >> "$STATE/report.txt" +} +release s4 base "\$BEST3_BASE" +release s4 transfer "\$BEST3_TRAN" +touch "$STATE/s5_go_base_\${BEST3_BASE}" +touch "$STATE/s5_go_transfer_\${BEST3_TRAN}" +echo "s5 released: base=\${BEST3_BASE} transfer=\${BEST3_TRAN}" >> "$STATE/report.txt" +EOF +} > "$STATE/coordD.sh" +DJOB=$(submit "$STATE/coordD.sh") +log "coord D: $DJOB" + +# ============ STAGE E: s4 brackets + s5 single-LR candidates ================ +log "stage E: s4 + s5 full-grid candidates (runtime-gated per LR by coord D)" +"${LAUNCH[@]}" --scale s4 --lrs $GRID_LRS --tag base \ + --dependency "afterany:$DJOB" --require "$STATE/s4_go_base_{lr}" +"${LAUNCH[@]}" --scale s4 --methods maxP-meas --alignment-table "$STATE/s2_align.json" \ + --lrs $GRID_LRS --tag transfer-s2 \ + --dependency "afterany:$DJOB" --require "$STATE/s4_go_transfer_{lr}" +# s5: pre-submit every candidate LR; only the one coord D blesses will run. +"${LAUNCH[@]}" --scale s5 --lrs $GRID_LRS --tag base \ + --dependency "afterany:$DJOB" --require "$STATE/s5_go_base_{lr}" +"${LAUNCH[@]}" --scale s5 --methods maxP-meas --alignment-table "$STATE/s2_align.json" \ + --lrs $GRID_LRS --tag transfer-s2 \ + --dependency "afterany:$DJOB" --require "$STATE/s5_go_transfer_{lr}" + +log "DAG fully submitted. Nothing else to do on the login node." +log "Progress: squeue --me | grep maxp" +log "Verdicts: cat $STATE/report.txt (after each coordinator runs)" diff --git a/experiments/lm/pipeline_analyze.py b/experiments/lm/pipeline_analyze.py new file mode 100644 index 0000000..0fc94ab --- /dev/null +++ b/experiments/lm/pipeline_analyze.py @@ -0,0 +1,150 @@ +#!/usr/bin/env python3 +"""Decision helpers for the unattended pipeline (run on the login node). + +Subcommands: + best-lr --glob 'runs/*_s2_mupno_lr*_meas' + Picks the LR with lowest final loss (median over seeds). + Prints ONLY the lr tag (e.g. "1e-02") on stdout; details on stderr. + + e1 --baseline-glob G1 --transfer-glob G2 [--tol 0.02] + Compares best (over LRs, median over seeds) final loss. + Exit 0 if transfer <= baseline + tol (continue to s4/s5), else 1. + + gates --runs-glob 'runs/*_s2_mupno_lr{BEST}_s*_meas' --adjacent-glob '...' + Seed-stability and source-LR-stability of exported alignment tables. + Report on stdout, always exit 0 (warnings only). +""" + +from __future__ import annotations + +import argparse +import glob as globlib +import re +import sys + +import numpy as np + + +def _final_loss(run_dir: str, last_k: int = 5) -> float | None: + from tensorboard.backend.event_processing.event_accumulator import EventAccumulator + + events = globlib.glob(f"{run_dir}/tb/*/events.*") + if not events: + return None + ea = EventAccumulator(events[0], size_guidance={"scalars": 0}) + ea.Reload() + tags = [t for t in ea.Tags()["scalars"] if "loss" in t and "validation" not in t] + if not tags: + return None + tag = next((t for t in tags if "global_avg" in t), tags[0]) + vals = [e.value for e in ea.Scalars(tag)] + return float(np.mean(vals[-last_k:])) if vals else None + + +def _by_lr(pattern: str) -> dict[str, list[float]]: + """lr tag -> final losses (one per seed/run).""" + out: dict[str, list[float]] = {} + for run_dir in sorted(globlib.glob(pattern)): + m = re.search(r"_lr([^_]+)_", run_dir) + if not m: + continue + loss = _final_loss(run_dir) + if loss is None or not np.isfinite(loss): + print(f" [warn] no usable loss in {run_dir}", file=sys.stderr) + continue + out.setdefault(m.group(1), []).append(loss) + if not out: + raise SystemExit(f"No usable runs match {pattern}") + return out + + +def cmd_best_lr(args) -> None: + by_lr = _by_lr(args.glob) + med = {lr: float(np.median(v)) for lr, v in by_lr.items()} + best = min(med, key=med.get) + for lr in sorted(med, key=lambda t: float(t)): + mark = " <-- best" if lr == best else "" + print(f" lr {lr}: median {med[lr]:.4f} over {len(by_lr[lr])} run(s){mark}", + file=sys.stderr) + tags = sorted(med, key=lambda t: float(t)) + if best in (tags[0], tags[-1]): + print(f" [warn] best LR {best} is at the grid edge — optimum may lie outside", + file=sys.stderr) + print(best) + + +def cmd_e1(args) -> None: + base = {lr: float(np.median(v)) for lr, v in _by_lr(args.baseline_glob).items()} + tran = {lr: float(np.median(v)) for lr, v in _by_lr(args.transfer_glob).items()} + b_lr, b = min(base.items(), key=lambda kv: kv[1]) + t_lr, t = min(tran.items(), key=lambda kv: kv[1]) + print(f"E1 verdict: baseline best {b:.4f} @ lr {b_lr} | transfer best {t:.4f} @ lr {t_lr}") + if t <= b + args.tol: + print(f"PASS (transfer within tol {args.tol} of baseline or better) — continue to s4/s5") + else: + print(f"FAIL (transfer worse by {t - b:.4f} > tol {args.tol}) — stopping before s4/s5") + raise SystemExit(1) + + +def _table_from_runs(pattern: str) -> dict[str, list[float]]: + from export_alignment import collect, METRICS, SUFFIXES + + pooled: dict[tuple[str, str], list[float]] = {} + for run_dir in sorted(globlib.glob(pattern)): + for k, v in collect(run_dir, last_frac=0.25).items(): + pooled.setdefault(k, []).extend(v) + return {s: [float(np.median(pooled[(m, s)])) for m in METRICS] + for s in SUFFIXES if (METRICS[0], s) in pooled} + + +def cmd_gates(args) -> None: + runs = sorted(globlib.glob(args.runs_glob)) + print(f"Gate report — source runs: {[r.split('/')[-1] for r in runs]}") + pooled = _table_from_runs(args.runs_glob) + + # G1: per-seed tables vs pooled + worst = 0.0 + for r in runs: + t = _table_from_runs(r) + for k in t: + for i in range(3): + worst = max(worst, abs(t[k][i] - pooled[k][i])) + status = "OK" if worst < args.warn else "WARN — noisy estimator, check per-seed tables" + print(f"G1 seed stability: max |alignment delta| seed-vs-pooled = {worst:.4f} [{status}]") + + # G2: adjacent source LR vs best + if args.adjacent_glob and globlib.glob(args.adjacent_glob): + adj = _table_from_runs(args.adjacent_glob) + worst = max(abs(adj[k][i] - pooled[k][i]) + for k in set(adj) & set(pooled) for i in range(3)) + status = "OK" if worst < args.warn else "WARN — source-LR sensitive" + print(f"G2 source-LR stability: max |alignment delta| = {worst:.4f} [{status}]") + else: + print("G2 source-LR stability: skipped (no adjacent-LR runs)") + + +def main() -> None: + p = argparse.ArgumentParser(description=__doc__) + sub = p.add_subparsers(dest="cmd", required=True) + + b = sub.add_parser("best-lr"); b.add_argument("--glob", required=True) + b.set_defaults(fn=cmd_best_lr) + + e = sub.add_parser("e1") + e.add_argument("--baseline-glob", required=True) + e.add_argument("--transfer-glob", required=True) + e.add_argument("--tol", type=float, default=0.02) + e.set_defaults(fn=cmd_e1) + + g = sub.add_parser("gates") + g.add_argument("--runs-glob", required=True) + g.add_argument("--adjacent-glob", default=None) + g.add_argument("--warn", type=float, default=0.05) + g.set_defaults(fn=cmd_gates) + + args = p.parse_args() + args.fn(args) + + +if __name__ == "__main__": + main() diff --git a/experiments/lm/run.sh b/experiments/lm/run.sh index 2b525a8..58a3101 100644 --- a/experiments/lm/run.sh +++ b/experiments/lm/run.sh @@ -7,24 +7,23 @@ # Override any of them by passing extra flags after the scale argument. # # Examples: -# bash experiments/lm/run.sh s3 -# bash experiments/lm/run.sh s5 --lrs 3e-3 # transfer LR +# # baseline mup-no sweep + passive alignment logging (source runs): +# bash experiments/lm/run.sh s1 --measure-only --tag meas +# +# # transfer arm: c solved from an exported alignment table +# # (--measure-only additionally logs realized alignment for free): +# bash experiments/lm/run.sh s2 --methods maxP-meas \ +# --alignment-table s1_align.json --measure-only --tag transfer-s1 +# +# # dynamic maxP (diagnostics only): +# bash experiments/lm/run.sh s2 --methods maxP +# # bash experiments/lm/run.sh s2 --dry-run set -euo pipefail SCALE="${1:?Usage: $0 [extra launch_sweep.py flags]}" shift -# Per-scale defaults -case "$SCALE" in - s1) SEEDS="1 2 3"; LRS="3e-4 1e-3 3e-3 1e-2 3e-2 1e-1 3e-1" ;; - s2) SEEDS="1 2"; LRS="3e-4 1e-3 3e-3 1e-2 3e-2 1e-1 3e-1" ;; - s3) SEEDS="1 2"; LRS="3e-4 1e-3 3e-3 1e-2 3e-2 1e-1 3e-1" ;; - s4) SEEDS="1"; LRS="3e-4 1e-3 3e-3 1e-2 3e-2 1e-1 3e-1" ;; - s5) SEEDS="1"; LRS="${TRANSFER_LR:-1e-2}" ;; - *) echo "Unknown scale '$SCALE'. Choose from s1 s2 s3 s4 s5."; exit 1 ;; -esac - RUNS_DIR="${RUNS_DIR:-/net/storage/pr3/plgrid/plggadlers/maxP/runs}" HF_ASSETS_PATH="${HF_ASSETS_PATH:-/net/storage/pr3/plgrid/plggadlers/maxP/experiments/lm/assets/hf/Llama-3.1-8B}" VENV_PATH="${VENV_PATH:-/net/storage/pr3/plgrid/plggadlers/maxP/.venv}" @@ -34,9 +33,6 @@ cd "$REPO_PATH" python experiments/lm/launch_sweep.py \ --scale "$SCALE" \ - --methods maxP mup-full mup-no \ - --lrs $LRS \ - --seeds $SEEDS \ --runs-dir "$RUNS_DIR" \ --hf-assets-path "$HF_ASSETS_PATH" \ --venv-path "$VENV_PATH" \ diff --git a/experiments/lm/train.py b/experiments/lm/train.py index 5a0e060..efd11b9 100644 --- a/experiments/lm/train.py +++ b/experiments/lm/train.py @@ -27,7 +27,7 @@ from torchtitan.components.metrics import MetricsProcessor from torchtitan.components.optimizer import OptimizersContainer from torchtitan.components.validate import Validator -from torchtitan.config.configs import ActivationCheckpointConfig, CompileConfig, TrainingConfig +from torchtitan.config.configs import ActivationCheckpointConfig, CommConfig, CompileConfig, DebugConfig, TrainingConfig from torchtitan.hf_datasets.text_datasets import HuggingFaceTextDataLoader from torchtitan.protocols.model_converter import ModelConvertersContainer from torchtitan.tools.logging import init_logger, logger @@ -134,6 +134,8 @@ def build_trainer_config(args: argparse.Namespace) -> Trainer.Config: solve_interval=args.solve_interval, sample_size=args.sample_size, c_ema=args.c_ema, + measure_only=args.measure_only, + alignment_table=args.alignment_table, )], ), optimizer=OptimizersContainer.Config( @@ -171,6 +173,8 @@ def build_trainer_config(args: argparse.Namespace) -> Trainer.Config: ), compile=CompileConfig(enable=not is_debug), activation_checkpoint=ActivationCheckpointConfig(mode="full"), + comm=CommConfig(init_timeout_seconds=600, train_timeout_seconds=1800), + debug=DebugConfig(seed=args.seed), validator=Validator.Config( enable=not is_debug, freq=500, @@ -187,8 +191,12 @@ def parse_args() -> argparse.Namespace: ) p.add_argument("--scale", choices=list(SCALE_CONFIGS), default="s3", help="Model scale") - p.add_argument("--method", choices=["maxP", "mup-full", "mup-no"], default="maxP", + p.add_argument("--method", choices=["maxP", "mup-no", "maxP-meas"], default="maxP", help="maxP variant") + p.add_argument("--measure-only", action="store_true", + help="Measure + log alignment without changing LRs (mup-no source runs)") + p.add_argument("--alignment-table", default=None, + help="JSON alignment table from export_alignment.py (maxP-meas only)") p.add_argument("--lr", type=float, default=1e-3, help="LR prefactor") p.add_argument("--alignment-warmup", type=int, default=100, @@ -239,6 +247,10 @@ def main() -> None: raise else: trainer.close() + if int(os.environ.get("RANK", "0")) == 0: + # Sentinel for chained SLURM jobs: lets follow-up links skip + # without spinning up torchrun (see launch_sweep.py). + open(os.path.join(args.output_dir, "COMPLETED"), "w").close() if torch.distributed.is_initialized(): torch.distributed.destroy_process_group() logger.info("Process group destroyed") diff --git a/experiments/vision/coord_check_vit.py b/experiments/vision/coord_check_vit.py index 178bdf2..9617501 100644 --- a/experiments/vision/coord_check_vit.py +++ b/experiments/vision/coord_check_vit.py @@ -5,47 +5,53 @@ import argparse +import timm 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 +from maxp_timm import _LinearPatchEmbed, _ScaledAttention, install_pm_wrappers -WIDTH_TO_SCALE = { - 192: "debug", # vit_tiny - 384: "vit-s", # vit_small - 768: "vit-b", # vit_base -} +IMAGE_SIZE = 224 +DEPTH = 2 +NUM_CLASSES = 1000 -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, +def _build(width: int): + return timm.create_model( + "vit_base_patch16_224", + pretrained=False, + num_classes=NUM_CLASSES, + img_size=IMAGE_SIZE, + embed_dim=width, + depth=DEPTH, + num_heads=width // 64, + drop_path_rate=0.0, + attn_layer=_ScaledAttention, + embed_layer=_LinearPatchEmbed, ) + + +def _make_model(width: int, parametrized: bool): + model = _build(width) if not parametrized: return model, None install_pm_wrappers(model) - sample_input = torch.randn(1, 3, cfg.image_size, cfg.image_size) + sample_input = torch.randn(1, 3, IMAGE_SIZE, IMAGE_SIZE) param = Parametrization( model, lr_prefactor=1e-3, optimizer_type="adam", - alignment="full", + alignment="no", 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) + return torch.randn(32, 3, IMAGE_SIZE, IMAGE_SIZE) def _make_train_step(model, param_groups): @@ -79,7 +85,7 @@ def main() -> None: parser.add_argument("--plot", action="store_true") args = parser.parse_args() - widths = sorted(WIDTH_TO_SCALE) + widths = [128, 256, 512, 1024] variant = "parametrized" if args.parametrized else "plain" print(f"ViT coord check ({variant}) — widths={widths}") diff --git a/experiments/vision/export_alignment.py b/experiments/vision/export_alignment.py new file mode 100644 index 0000000..4a6157e --- /dev/null +++ b/experiments/vision/export_alignment.py @@ -0,0 +1,99 @@ +#!/usr/bin/env python3 +"""Export a measured alignment table from a vision run's metrics.json. + +Reads ``align/{z0_dW,dZ_w0,dZ_dW}/`` scalars logged per row by +``train.py --measure-only`` (a ``mup-no`` source run), averages each layer over +the last ``--last-frac`` of training, and writes a JSON table consumable by +``train.py --method maxP-meas --alignment-table``: + + {"blocks.0.attn.qkv": [a, w, u], ..., "head": [...]} + +Keys are EXACT PM names. The width ladder holds depth = 12 constant, so the +block names are identical at every scale and the table transfers verbatim +(down to s1, up to s3/s4). Exact-name keys also avoid the ``proj`` leaf +collision (``patch_embed.proj`` vs ``blocks.N.attn.proj``). + +``patch_embed.*`` (input embedding) is excluded: its alignment is never +measured (fixed fan-in), so it keeps the "no" preset. + +Usage: + python experiments/vision/export_alignment.py runs/ -o s2_align.json + python experiments/vision/export_alignment.py runs/ runs/ -o avg.json +""" + +from __future__ import annotations + +import argparse +import json +from collections import defaultdict +from pathlib import Path + +import numpy as np + +METRICS = ("z0_dW", "dZ_w0", "dZ_dW") + + +def collect(run_dir: str, last_frac: float) -> dict[tuple[str, str], float]: + """Return (metric, layer) -> steady-state mean over the last `last_frac`.""" + path = Path(run_dir) / "metrics.json" + if not path.exists(): + raise FileNotFoundError(f"No metrics.json under {run_dir}") + + series: dict[tuple[str, str], list[float]] = defaultdict(list) + with path.open() as f: + for line in f: + line = line.strip() + if not line: + continue + row = json.loads(line) + for key, val in row.items(): + if not key.startswith("align/"): + continue + _, metric, layer = key.split("/", 2) + if layer.startswith("patch_embed"): + continue + series[(metric, layer)].append(float(val)) + + if not series: + raise ValueError( + f"No align/* scalars in {run_dir}/metrics.json — " + "was it run with --measure-only?" + ) + + out: dict[tuple[str, str], float] = {} + for key, vals in series.items(): + arr = np.asarray(vals) + steady = arr[int(len(arr) * (1 - last_frac)):] + out[key] = float(steady.mean()) + return out + + +def main() -> None: + p = argparse.ArgumentParser(description="Export alignment table from vision run logs") + p.add_argument("run_dirs", nargs="+", help="Source run director(ies); multiple are pooled (e.g. seeds)") + p.add_argument("-o", "--output", required=True, help="Output JSON path") + p.add_argument("--last-frac", type=float, default=0.25, + help="Fraction of training (from the end) to average over") + args = p.parse_args() + + pooled: dict[tuple[str, str], list[float]] = defaultdict(list) + for run_dir in args.run_dirs: + for key, val in collect(run_dir, args.last_frac).items(): + pooled[key].append(val) + + layers = sorted({layer for (_, layer) in pooled}) + table: dict[str, list[float]] = {} + for layer in layers: + if any((m, layer) not in pooled for m in METRICS): + continue + table[layer] = [round(float(np.median(pooled[(m, layer)])), 4) for m in METRICS] + + with open(args.output, "w") as f: + json.dump(table, f, indent=2) + print(f"Wrote {args.output} ({len(table)} layers):") + for k, v in table.items(): + print(f" {k:28s} a={v[0]:.4f} w={v[1]:.4f} u={v[2]:.4f}") + + +if __name__ == "__main__": + main() diff --git a/experiments/vision/hf_vision_data.py b/experiments/vision/hf_vision_data.py index 718f679..101b2eb 100644 --- a/experiments/vision/hf_vision_data.py +++ b/experiments/vision/hf_vision_data.py @@ -62,7 +62,7 @@ def __getitem__(self, idx: int): sample = self._hf[idx] image = sample[self._preset.image_key].convert("RGB") label = sample[self._preset.label_key] - x = self._transform(image) + x = self._transform(image).to(torch.bfloat16) y = torch.tensor(label, dtype=torch.long) return x, y @@ -74,15 +74,19 @@ def make_loader( num_workers: int, is_train: bool, prefetch_factor: int, + generator: torch.Generator | None = None, ) -> DataLoader: kwargs: dict = { "dataset": ds, "batch_size": batch_size, "num_workers": num_workers, "pin_memory": torch.cuda.is_available(), + "shuffle": is_train, "drop_last": is_train, "persistent_workers": (num_workers > 0), } + if is_train and generator is not None: + kwargs["generator"] = generator if num_workers > 0: kwargs["prefetch_factor"] = prefetch_factor return DataLoader(**kwargs) @@ -96,6 +100,7 @@ def build_dataloaders( prefetch_factor: int = 2, train_transform=None, eval_transform=None, + generator: torch.Generator | None = None, ) -> tuple[DataLoader, DataLoader | None, DatasetPreset]: if train_transform is None or eval_transform is None: raise ValueError( @@ -112,6 +117,7 @@ def build_dataloaders( num_workers=num_workers, is_train=True, prefetch_factor=prefetch_factor, + generator=generator, ) val_hf = load_dataset(preset.dataset_id, split=preset.val_split, streaming=False) diff --git a/experiments/vision/launch_sweep.py b/experiments/vision/launch_sweep.py index a0dcfe8..0002b26 100755 --- a/experiments/vision/launch_sweep.py +++ b/experiments/vision/launch_sweep.py @@ -10,81 +10,50 @@ from string import Template +# Every scale defaults to the full LR grid. The pipeline narrows s3/s4/s5 +# DYNAMICALLY at runtime (coordinators release a per-arm bracket via --require +# sentinels — see pipeline.sh); no scale hardcodes a single LR, so the optimum +# is always data-driven, never assumed. ALL_LRS = [3e-4, 1e-3, 3e-3, 1e-2, 3e-2, 1e-1, 3e-1] -TRANSFER_LR = [1e-2] SCALE_CONFIGS = { "debug": { "wall": "01:00:00", - "gpus": 1, - "methods": ["maxP"], + "methods": ["mup-no"], "lrs": [3e-3], "seeds": [1], - "image_size": 224, - "drop_path_rate": 0.0, }, - "vit-s": { + "s1": { "wall": "24:00:00", - "gpus": 1, - "methods": ["maxP", "mup-full", "mup-no"], + "methods": ["mup-no", "maxP-meas"], "lrs": ALL_LRS, "seeds": [1, 2, 3], - "image_size": 224, - "drop_path_rate": 0.1, }, - "vit-b": { - "wall": "36:00:00", - "gpus": 1, - "methods": ["maxP", "mup-full", "mup-no"], - "lrs": TRANSFER_LR, - "seeds": [1], - "image_size": 224, - "drop_path_rate": 0.2, - }, - "vit-l": { - "wall": "48:00:00", - "gpus": 1, - "methods": ["maxP", "mup-full", "mup-no"], - "lrs": TRANSFER_LR, - "seeds": [1], - "image_size": 224, - "drop_path_rate": 0.4, - }, - "mlp-s": { - "wall": "6:00:00", - "gpus": 1, - "methods": ["maxP", "mup-full", "mup-no"], + "s2": { + "wall": "24:00:00", + "methods": ["mup-no", "maxP-meas"], "lrs": ALL_LRS, "seeds": [1, 2, 3], - "image_size": 224, - "drop_path_rate": 0.1, }, - "mlp-m": { - "wall": "12:00:00", - "gpus": 1, - "methods": ["maxP", "mup-full", "mup-no"], + "s3": { + "wall": "48:00:00", + "methods": ["mup-no", "maxP-meas"], "lrs": ALL_LRS, - "seeds": [1], - "image_size": 224, - "drop_path_rate": 0.2, + "seeds": [1, 2], }, - "mlp-b": { - "wall": "12:00:00", - "gpus": 1, - "methods": ["maxP", "mup-full", "mup-no"], - "lrs": TRANSFER_LR, - "seeds": [1], - "image_size": 224, - "drop_path_rate": 0.3, + "s4": { + "wall": "48:00:00", + "methods": ["mup-no", "maxP-meas"], + "lrs": ALL_LRS, + "seeds": [1, 2], + "chain": 2, }, - "mlp-l": { - "wall": "24:00:00", - "gpus": 1, - "methods": ["maxP", "mup-full", "mup-no"], - "lrs": TRANSFER_LR, + "s5": { + "wall": "48:00:00", + "methods": ["mup-no", "maxP-meas"], + "lrs": ALL_LRS, "seeds": [1], - "image_size": 224, - "drop_path_rate": 0.4, + "chain": 4, }, } @@ -107,8 +76,15 @@ source "${venv_path}/bin/activate" cd "${repo_path}" -export OMP_NUM_THREADS=16 +if [ -f "${output_dir}/final_metrics.json" ]; then + echo "[skip] ${output_dir} already complete (final_metrics.json present)" + exit 0 +fi +${runtime_guards} +export OMP_NUM_THREADS=1 export HF_HOME="$${HF_HOME:-/net/scratch/hscra/plgrid/plgmwojnar/hf}" +export HF_HUB_DOWNLOAD_TIMEOUT="$${HF_HUB_DOWNLOAD_TIMEOUT:-120}" +export HF_HUB_ETAG_TIMEOUT="$${HF_HUB_ETAG_TIMEOUT:-120}" export WANDB_PROJECT="$${WANDB_PROJECT:-maxP-vision}" export WANDB_RUN_NAME="maxp_${scale}_${method_tag}_lr${lr_tag}_s${seed}" @@ -121,9 +97,11 @@ --epochs ${epochs} \\ --batch-size ${batch_size} \\ --num-workers ${num_workers} \\ + --val-interval ${val_interval} \\ --val-steps ${val_steps} \\ + --keep-latest-k ${keep_latest_k} \\ --output-dir ${output_dir} \\ - ${extra_train_args} + ${resume_arg} ${method_args} ${extra_train_args} """ ) @@ -136,12 +114,12 @@ def _method_tag(method: str) -> str: return method.replace("-", "") -def _has_checkpoint(out_dir: Path) -> bool: - ckpt_dir = out_dir / "checkpoint" - return ckpt_dir.is_dir() and any(ckpt_dir.iterdir()) +def _is_complete(out_dir: Path) -> bool: + """A run is complete once train.py has written final_metrics.json.""" + return (out_dir / "final_metrics.json").is_file() -def parse_args() -> argparse.Namespace: +def main(): parser = argparse.ArgumentParser( description="Launch maxP timm vision sweeps via SLURM", formatter_class=argparse.ArgumentDefaultsHelpFormatter, @@ -156,40 +134,89 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--repo-path", required=True) parser.add_argument("--dataset", default="imagenet12k") - parser.add_argument("--epochs", type=int, default=1) - parser.add_argument("--batch-size", type=int, default=128) - parser.add_argument("--num-workers", type=int, default=8) - parser.add_argument("--val-steps", type=int, default=500) - + parser.add_argument("--epochs", type=int, default=4, + help="Epochs — identical at every width (fixed sample budget, muTransfer setup)") + parser.add_argument("--batch-size", type=int, default=512) + parser.add_argument("--num-workers", type=int, default=24) + parser.add_argument("--val-interval", type=int, default=500, + help="Run validation every N steps (must be a multiple of --log-interval=20)") + parser.add_argument("--val-steps", type=int, default=50, + help="Batches per validation pass (keep small — eval compute ≈ train compute)") + + parser.add_argument("--alignment-table", default=None, + help="JSON table for maxP-meas runs (from export_alignment.py)") + parser.add_argument("--measure-only", action="store_true", + help="Append --measure-only (mup-no alignment source run)") + parser.add_argument("--tag", default=None, + help="Suffix appended to run names (e.g. 'meas', 'transfer-s2')") + parser.add_argument("--job-ids-file", default=None, + help="Append submitted SLURM job IDs (one per line) for downstream deps") + parser.add_argument("--dependency", default=None, + help="sbatch --dependency spec applied to each run's first link") + parser.add_argument("--require", action="append", default=[], + help="File that must exist at job runtime or the job exits cleanly; " + "'{lr}' is replaced by the run's lr tag. Repeatable.") + parser.add_argument("--chain", type=int, default=None, + help="Number of resume-chain links per run (default: per-scale config). " + "Links run sequentially via --dependency=afterany; each resumes from " + "the latest checkpoint, and a completed run no-ops via the final_metrics guard.") parser.add_argument("--extra-train-args", default="") parser.add_argument("--dry-run", action="store_true") - parser.add_argument("--resume", action="store_true") - return parser.parse_args() + parser.add_argument("--resume", action="store_true", + help="Force resubmit even if final_metrics.json already exists") + parser.add_argument("--run-date", default=None, + help="Override the YYYY-MM-DD prefix in run names (default: " + "today). Use to resubmit into an existing run's dir so " + "training resumes from its checkpoint.") + args = parser.parse_args() - -def main() -> None: - args = parse_args() sc = SCALE_CONFIGS[args.scale] methods = args.methods or sc["methods"] lrs = args.lrs or sc["lrs"] seeds = args.seeds or sc["seeds"] - today = date.today().strftime("%Y-%m-%d") + if "maxP-meas" in methods and not args.measure_only and not args.alignment_table: + raise SystemExit("maxP-meas runs require --alignment-table (from export_alignment.py)") + + chain = args.chain if args.chain is not None else sc.get("chain", 1) + if chain > 1 and args.measure_only: + raise SystemExit("--measure-only requires chain==1: resume re-snapshots z0/w0 and " + "corrupts alignment measurement (use full-budget single jobs to measure)") + + today = args.run_date or date.today().strftime("%Y-%m-%d") submitted = skipped = 0 for method in methods: + method_args = "" + if method == "maxP-meas": + method_args = f"--alignment-table {args.alignment_table}" + elif args.measure_only: + method_args = "--measure-only" + for lr in lrs: for seed in seeds: run_name = f"{today}_{args.scale}_{_method_tag(method)}_lr{_lr_tag(lr)}_s{seed}" + if args.tag: + run_name += f"_{args.tag}" out_dir = Path(args.runs_dir) / run_name - if not args.resume and _has_checkpoint(out_dir): + if not args.resume and _is_complete(out_dir): print(f"[skip] {run_name}") skipped += 1 continue out_dir.mkdir(parents=True, exist_ok=True) + # Runtime gating: jobs can be pre-submitted before the decision + # that enables them exists; '{lr}' resolves to this run's lr tag. + guards = [] + for path in args.require: + path = path.replace("{lr}", _lr_tag(lr)) + guards.append( + f'if [ ! -f "{path}" ]; then ' + f'echo "gate: {path} missing — skipping."; exit 0; fi') + runtime_guards = ("\n" + "\n".join(guards)) if guards else "" + script = SLURM_TEMPLATE.substitute( scale=args.scale, method_tag=_method_tag(method), @@ -205,21 +232,45 @@ def main() -> None: epochs=args.epochs, batch_size=args.batch_size, num_workers=args.num_workers, + val_interval=args.val_interval, val_steps=args.val_steps, + keep_latest_k=1, + resume_arg="--resume" if chain > 1 else "", + method_args=method_args, extra_train_args=args.extra_train_args, + runtime_guards=runtime_guards, ) script_path = out_dir / "job.sh" script_path.write_text(script) - if args.dry_run: - print(f"[dry] sbatch {script_path}") - else: - result = subprocess.run( - ["sbatch", str(script_path)], - capture_output=True, - text=True, - ) - print(f"[submit] {run_name} → {result.stdout.strip()}") + prev = None + for link in range(chain): + if prev is not None: + dep_arg = f"--dependency=afterany:{prev}" + elif args.dependency: + dep_arg = f"--dependency={args.dependency}" + else: + dep_arg = None + if args.dry_run: + suffix = f" ({dep_arg})" if dep_arg else "" + print(f"[dry] sbatch link {link + 1}/{chain} {script_path}{suffix}") + prev = f"" + continue + cmd = ["sbatch"] + if dep_arg: + cmd.append(dep_arg) + cmd.append(str(script_path)) + result = subprocess.run(cmd, capture_output=True, text=True) + if result.returncode != 0: + raise SystemExit( + f"sbatch failed for {run_name} (link {link + 1}/{chain}): " + f"{result.stderr.strip()}" + ) + prev = result.stdout.strip().split()[-1] + if args.job_ids_file: + with open(args.job_ids_file, "a") as f: + f.write(f"{prev}\n") + print(f"[submit] {run_name} link {link + 1}/{chain} → job {prev}") submitted += 1 action = "would submit" if args.dry_run else "submitted" diff --git a/experiments/vision/maxp_timm.py b/experiments/vision/maxp_timm.py index c460963..dbe158b 100644 --- a/experiments/vision/maxp_timm.py +++ b/experiments/vision/maxp_timm.py @@ -23,32 +23,38 @@ class ScaleConfig: hidden: int = 0 depth: int = 0 dropout: float = 0.0 + embed_dim: int = 0 + num_heads: int = 0 +# ViT width-only ladder (mirrors the Llama3 s1..s5 ladder): +# width axis = embed_dim, with head_dim = 64 and depth = 12 held CONSTANT. +# drop_path / dropout = 0 at every scale so the only thing varying is width +# (regularization that scales with width would confound the muP transfer test). SCALE_CONFIGS: dict[str, ScaleConfig] = { "debug": ScaleConfig( - model_name="vit_tiny_patch16_224", family="vit", image_size=224, drop_path_rate=0.0 + model_name="vit_base_patch16_224", family="vit", image_size=224, + embed_dim=128, num_heads=2, depth=2, ), - "vit-s": ScaleConfig( - model_name="vit_small_patch16_224", family="vit", image_size=224, drop_path_rate=0.1 + "s1": ScaleConfig( + model_name="vit_base_patch16_224", family="vit", image_size=224, + embed_dim=256, num_heads=4, depth=12, ), - "vit-b": ScaleConfig( - model_name="vit_base_patch16_224", family="vit", image_size=224, drop_path_rate=0.2 + "s2": ScaleConfig( + model_name="vit_base_patch16_224", family="vit", image_size=224, + embed_dim=512, num_heads=8, depth=12, ), - "vit-l": ScaleConfig( - model_name="vit_large_patch16_224", family="vit", image_size=224, drop_path_rate=0.4 + "s3": ScaleConfig( + model_name="vit_base_patch16_224", family="vit", image_size=224, + embed_dim=1024, num_heads=16, depth=12, ), - "mlp-s": ScaleConfig( - model_name="mlp", family="mlp", image_size=224, hidden=256, depth=4, dropout=0.0 + "s4": ScaleConfig( + model_name="vit_base_patch16_224", family="vit", image_size=224, + embed_dim=2048, num_heads=32, depth=12, ), - "mlp-m": ScaleConfig( - model_name="mlp", family="mlp", image_size=224, hidden=512, depth=6, dropout=0.1 - ), - "mlp-b": ScaleConfig( - model_name="mlp", family="mlp", image_size=224, hidden=1024, depth=8, dropout=0.2 - ), - "mlp-l": ScaleConfig( - model_name="mlp", family="mlp", image_size=224, hidden=2048, depth=8, dropout=0.3 + "s5": ScaleConfig( + model_name="vit_base_patch16_224", family="vit", image_size=224, + embed_dim=4096, num_heads=64, depth=12, ), } @@ -215,6 +221,9 @@ def create_model( "num_classes": num_classes, "drop_path_rate": cfg.drop_path_rate, "img_size": img_size, + "embed_dim": cfg.embed_dim, + "depth": cfg.depth, + "num_heads": cfg.num_heads, "attn_layer": _ScaledAttention, "embed_layer": _LinearPatchEmbed, } diff --git a/experiments/vision/pipeline.sh b/experiments/vision/pipeline.sh new file mode 100644 index 0000000..5f643e4 --- /dev/null +++ b/experiments/vision/pipeline.sh @@ -0,0 +1,228 @@ +#!/usr/bin/env bash +# Unattended VISION pipeline (ViT width ladder) — same muTransfer protocol as the +# LM pipeline: tune everything at small scale (s1+s2), only VERIFY at s3+. +# +# LOGIN-NODE SUBMIT-ONLY DESIGN: this script only calls sbatch (no compute, no +# polling) and exits in seconds. The whole DAG is submitted upfront: +# +# stage A s1 + s2 mup-no sweeps (--measure-only) ........ alignment sources +# coord B (afterany: A) best mup-no LR per scale, export +# alignment tables, gate report G1-G3 .......... 1 small job +# stage C s1 + s2 maxP-meas sweeps — prefactor re-tune +# under the measured table +# coord C (afterany: C) best maxP-meas LR per scale, +# s1-vs-s2 consistency, writes s3 bracket sentinels +# stage D s3 full-grid candidates; coord C releases only +# the 3-pt bracket around each arm's s2 optimum +# coord D (afterany: D) transfer proof + E1 verdict; +# releases s4/s5 sentinels +# stage E s4 (chained) + s5 (chained) candidates, runtime-gated per LR +# +# Decisions unknowable at submit time are enforced by runtime sentinel files +# (launch_sweep --require). s4/s5 rely on chained checkpoint resume. +# +# PREREQUISITE: confirm a real train.py end-to-end run + a checkpoint→resume +# cycle on the server before submitting (see docs/vision_experiment_plan.md gates). +# +# Usage (login node): bash experiments/vision/pipeline.sh +# Dry run (no sbatch): DRY=1 bash experiments/vision/pipeline.sh +set -euo pipefail + +# --- config (env-overridable) ---------------------------------------------- +REPO="${REPO:-/net/storage/pr3/plgrid/plggadlers/maxP}" +RUNS_DIR="${RUNS_DIR:-$REPO/runs}" +VENV="${VENV:-$REPO/.venv}" +DATASET="${DATASET:-imagenet12k}" +EPOCHS="${EPOCHS:-4}" +BATCH="${BATCH:-512}" +STATE="${STATE:-$RUNS_DIR/.vpipeline}" +GATE_TOL="${GATE_TOL:-0.02}" +ACCOUNT="${ACCOUNT:-plgadlers-gpu-gh200}" +PARTITION="${PARTITION:-plgrid-gpu-gh200}" +DRY="${DRY:-0}" +RUN_DATE="${RUN_DATE:-}" + +# Candidate LR grid (tags must match launch_sweep ALL_LRS in %.0e form). +GRID_LRS="${GRID_LRS:-3e-04 1e-03 3e-03 1e-02 3e-02 1e-01 3e-01}" + +cd "$REPO" +mkdir -p "$STATE" +log() { echo "[vpipeline] $*"; } + +# Login node has no python by default; load a stdlib interpreter for launch_sweep.py. +command -v python >/dev/null 2>&1 || module add GCCcore/13.2.0 Python/3.11.5 + +LAUNCH=(python experiments/vision/launch_sweep.py + --runs-dir "$RUNS_DIR" --venv-path "$VENV" --repo-path "$REPO" + --dataset "$DATASET" --epochs "$EPOCHS" --batch-size "$BATCH") +[ "$DRY" = "1" ] && LAUNCH+=(--dry-run) +# Pin run-date prefix so a resubmit reuses existing run dirs (skip done / resume chains). +[ -n "$RUN_DATE" ] && LAUNCH+=(--run-date "$RUN_DATE") + +submit() { # submit a standalone coordinator script, echo job id + if [ "$DRY" = "1" ]; then echo "DRY"; log "[dry] sbatch $1" >&2; return; fi + sbatch "$@" | awk '{print $NF}' +} + +dep_from() { # ids file -> afterany:id:id:... + [ "$DRY" = "1" ] && { echo "afterany:DRY"; return; } + [ -s "$1" ] || { echo "ERROR: no job ids in $1 — refusing to wire dependencies" >&2; exit 1; } + echo "afterany:$(paste -sd: "$1")" +} + +coord_header() { # $1 job name, $2 dependency + cat < "$STATE/A.ids" +log "stage A: s1 + s2 mup-no sweeps (--measure-only sources)" +"${LAUNCH[@]}" --scale s1 --methods mup-no --measure-only --tag meas --job-ids-file "$STATE/A.ids" +"${LAUNCH[@]}" --scale s2 --methods mup-no --measure-only --tag meas --job-ids-file "$STATE/A.ids" + +# ============ COORD B: pick source LRs, export tables, gates ================ +{ + coord_header vmaxp_coordB "$(dep_from "$STATE/A.ids")" + cat < "$STATE/best_lr_mupno_s2.txt" +echo "\$BEST_S1" > "$STATE/best_lr_mupno_s1.txt" +$EXPORT $RUNS_DIR/*_s2_mupno_lr\${BEST_S2}_s*_meas -o "$STATE/s2_align.json" +$EXPORT $RUNS_DIR/*_s1_mupno_lr\${BEST_S1}_s*_meas -o "$STATE/s1_align.json" +ADJ_S2=\$(python -c "g='$GRID_LRS'.split(); i=g.index('\$BEST_S2'); print(g[i+1] if i+1 < len(g) else g[i-1])") +{ + echo "=== Gate report (\$(date)) ===" + echo "mup-no optima: s1=\$BEST_S1 s2=\$BEST_S2 \$([ "\$BEST_S1" = "\$BEST_S2" ] && echo '[consistent]' || echo '[WARN: differ across scale]')" + $ANALYZE gates \\ + --runs-glob "$RUNS_DIR/*_s2_mupno_lr\${BEST_S2}_s*_meas" \\ + --adjacent-glob "$RUNS_DIR/*_s2_mupno_lr\${ADJ_S2}_s*_meas" + python - "$STATE/s1_align.json" "$STATE/s2_align.json" <<'PYEOF' +import json, sys +a, b = (json.load(open(p)) for p in sys.argv[1:3]) +worst = max(abs(a[k][i] - b[k][i]) for k in set(a) & set(b) for i in range(3)) +ok = "OK - tiny-proxy story holds" if worst < 0.05 else "WARN - keep s2 as canonical source" +print(f"G3 s1-vs-s2: max |alignment delta| = {worst:.4f} [{ok}]") +PYEOF +} >> "$STATE/report.txt" 2>&1 +EOF +} > "$STATE/coordB.sh" +BJOB=$(submit "$STATE/coordB.sh") +log "coord B: $BJOB" + +# ============ STAGE C: s1+s2 maxP-meas sweeps (prefactor re-tune) =========== +: > "$STATE/C.ids" +log "stage C: s1 + s2 maxP-meas sweeps (re-tune under measured table)" +for SC in s1 s2; do + "${LAUNCH[@]}" --scale "$SC" --methods maxP-meas \ + --alignment-table "$STATE/s2_align.json" --tag transfer-s2 \ + --dependency "afterany:$BJOB" --require "$STATE/s2_align.json" \ + --job-ids-file "$STATE/C.ids" +done + +# ============ COORD C: pick transfer LR at small scale, release s3 ========== +{ + coord_header vmaxp_coordC "$(dep_from "$STATE/C.ids")" + cat < "$STATE/best_lr_meas_s2.txt" +echo "\$T_S1" > "$STATE/best_lr_meas_s1.txt" +echo "maxP-meas optima: s1=\$T_S1 s2=\$T_S2 \$([ "\$T_S1" = "\$T_S2" ] && echo '[consistent]' || echo '[WARN: differ across scale]')" >> "$STATE/report.txt" +release() { # \$1 stage prefix, \$2 arm name, \$3 center + BRACKET=\$(python -c "g='$GRID_LRS'.split(); i=g.index('\$3'); print(' '.join(g[max(0, i-1):i+2]))") + for lr in \$BRACKET; do touch "$STATE/\${1}_go_\${2}_\${lr}"; done + EDGE="" + python -c "g='$GRID_LRS'.split(); import sys; sys.exit(0 if g.index('\$3') in (0, len(g)-1) else 1)" \ + && EDGE=" [WARN: center at grid edge]" + echo "\$1 \$2 bracket released: \$BRACKET (center \$3)\$EDGE" >> "$STATE/report.txt" +} +release s3 base "\$B_S2" +release s3 transfer "\$T_S2" +EOF +} > "$STATE/coordC.sh" +CJOB=$(submit "$STATE/coordC.sh") +log "coord C: $CJOB" + +# ============ STAGE D: s3 verification brackets ============================ +: > "$STATE/D.ids" +log "stage D: s3 brackets (both arms, gated per LR)" +"${LAUNCH[@]}" --scale s3 --methods mup-no --lrs $GRID_LRS --tag base \ + --dependency "afterany:$CJOB" --require "$STATE/s3_go_base_{lr}" \ + --job-ids-file "$STATE/D.ids" +"${LAUNCH[@]}" --scale s3 --methods maxP-meas --alignment-table "$STATE/s2_align.json" \ + --lrs $GRID_LRS --tag transfer-s2 \ + --dependency "afterany:$CJOB" --require "$STATE/s3_go_transfer_{lr}" \ + --job-ids-file "$STATE/D.ids" + +# ============ COORD D: transfer proof + E1 verdict, release s4/s5 =========== +{ + coord_header vmaxp_coordD "$(dep_from "$STATE/D.ids")" + cat <> "$STATE/report.txt" +if ! $ANALYZE e1 \\ + --baseline-glob "$RUNS_DIR/*_s3_mupno_lr*_s*_base" \\ + --transfer-glob "$RUNS_DIR/*_s3_maxPmeas_lr*_s*_transfer-s2" \\ + --tol $GATE_TOL >> "$STATE/report.txt" 2>&1; then + echo "E1 FAILED - s4/s5 will self-skip" >> "$STATE/report.txt" + exit 0 +fi +release() { # \$1 stage prefix, \$2 arm name, \$3 center + BRACKET=\$(python -c "g='$GRID_LRS'.split(); i=g.index('\$3'); print(' '.join(g[max(0, i-1):i+2]))") + for lr in \$BRACKET; do touch "$STATE/\${1}_go_\${2}_\${lr}"; done + echo "\$1 \$2 bracket released: \$BRACKET (center \$3)" >> "$STATE/report.txt" +} +release s4 base "\$BEST3_BASE" +release s4 transfer "\$BEST3_TRAN" +touch "$STATE/s5_go_base_\${BEST3_BASE}" +touch "$STATE/s5_go_transfer_\${BEST3_TRAN}" +echo "s5 released: base=\${BEST3_BASE} transfer=\${BEST3_TRAN}" >> "$STATE/report.txt" +EOF +} > "$STATE/coordD.sh" +DJOB=$(submit "$STATE/coordD.sh") +log "coord D: $DJOB" + +# ============ STAGE E: s4 + s5 candidates (chained, runtime-gated) ========== +log "stage E: s4 + s5 candidates (per-LR gated by coord D, chained resume)" +"${LAUNCH[@]}" --scale s4 --methods mup-no --lrs $GRID_LRS --tag base \ + --dependency "afterany:$DJOB" --require "$STATE/s4_go_base_{lr}" +"${LAUNCH[@]}" --scale s4 --methods maxP-meas --alignment-table "$STATE/s2_align.json" \ + --lrs $GRID_LRS --tag transfer-s2 \ + --dependency "afterany:$DJOB" --require "$STATE/s4_go_transfer_{lr}" +"${LAUNCH[@]}" --scale s5 --methods mup-no --lrs $GRID_LRS --tag base \ + --dependency "afterany:$DJOB" --require "$STATE/s5_go_base_{lr}" +"${LAUNCH[@]}" --scale s5 --methods maxP-meas --alignment-table "$STATE/s2_align.json" \ + --lrs $GRID_LRS --tag transfer-s2 \ + --dependency "afterany:$DJOB" --require "$STATE/s5_go_transfer_{lr}" + +log "DAG fully submitted. Nothing else to do on the login node." +log "Progress: squeue --me | grep maxp" +log "Verdicts: cat $STATE/report.txt (after each coordinator runs)" diff --git a/experiments/vision/pipeline_analyze.py b/experiments/vision/pipeline_analyze.py new file mode 100644 index 0000000..61970ed --- /dev/null +++ b/experiments/vision/pipeline_analyze.py @@ -0,0 +1,161 @@ +#!/usr/bin/env python3 +"""Decision helpers for the unattended vision pipeline (run inside coordinator jobs). + +Reads validation loss / alignment from each run's metrics.json (jsonl) — no +tensorboard dependency. Lower val loss is better. + +Subcommands: + best-lr --glob 'runs/*_s2_mupno_lr*_meas' + Picks the LR with lowest final val loss (median over seeds). + Prints ONLY the lr tag (e.g. "1e-02") on stdout; details on stderr. + + e1 --baseline-glob G1 --transfer-glob G2 [--tol 0.02] + Compares best (over LRs, median over seeds) final val loss. + Exit 0 if transfer <= baseline + tol (continue to s4/s5), else 1. + + gates --runs-glob '...' [--adjacent-glob '...'] [--warn 0.05] + Seed-stability and source-LR-stability of exported alignment tables. + Report on stdout, always exit 0 (warnings only). +""" + +from __future__ import annotations + +import argparse +import glob as globlib +import json +import re +import sys + +import numpy as np + + +def _final_val_loss(run_dir: str, last_k: int = 5) -> float | None: + """Mean of the last `last_k` validation losses from metrics.json.""" + path = f"{run_dir}/metrics.json" + vals: list[float] = [] + try: + with open(path) as f: + for line in f: + line = line.strip() + if not line: + continue + row = json.loads(line) + v = row.get("loss/val_loss") + if isinstance(v, (int, float)): + vals.append(float(v)) + except FileNotFoundError: + return None + return float(np.mean(vals[-last_k:])) if vals else None + + +def _by_lr(pattern: str) -> dict[str, list[float]]: + """lr tag -> final val losses (one per seed/run).""" + out: dict[str, list[float]] = {} + for run_dir in sorted(globlib.glob(pattern)): + m = re.search(r"_lr([^_]+)_", run_dir) + if not m: + continue + loss = _final_val_loss(run_dir) + if loss is None or not np.isfinite(loss): + print(f" [warn] no usable val loss in {run_dir}", file=sys.stderr) + continue + out.setdefault(m.group(1), []).append(loss) + if not out: + raise SystemExit(f"No usable runs match {pattern}") + return out + + +def cmd_best_lr(args) -> None: + by_lr = _by_lr(args.glob) + med = {lr: float(np.median(v)) for lr, v in by_lr.items()} + best = min(med, key=med.get) + for lr in sorted(med, key=lambda t: float(t)): + mark = " <-- best" if lr == best else "" + print(f" lr {lr}: median {med[lr]:.4f} over {len(by_lr[lr])} run(s){mark}", + file=sys.stderr) + tags = sorted(med, key=lambda t: float(t)) + if best in (tags[0], tags[-1]): + print(f" [warn] best LR {best} is at the grid edge — optimum may lie outside", + file=sys.stderr) + print(best) + + +def cmd_e1(args) -> None: + base = {lr: float(np.median(v)) for lr, v in _by_lr(args.baseline_glob).items()} + tran = {lr: float(np.median(v)) for lr, v in _by_lr(args.transfer_glob).items()} + b_lr, b = min(base.items(), key=lambda kv: kv[1]) + t_lr, t = min(tran.items(), key=lambda kv: kv[1]) + print(f"E1 verdict: baseline best {b:.4f} @ lr {b_lr} | transfer best {t:.4f} @ lr {t_lr}") + if t <= b + args.tol: + print(f"PASS (transfer within tol {args.tol} of baseline or better) — continue to s4/s5") + else: + print(f"FAIL (transfer worse by {t - b:.4f} > tol {args.tol}) — stopping before s4/s5") + raise SystemExit(1) + + +def _table_from_runs(pattern: str) -> dict[str, list[float]]: + """Pool exact-name alignment tables across runs: 'layer' -> [a, w, u] medians.""" + from export_alignment import collect, METRICS + + pooled: dict[tuple[str, str], list[float]] = {} + for run_dir in sorted(globlib.glob(pattern)): + for key, val in collect(run_dir, last_frac=0.25).items(): + pooled.setdefault(key, []).append(val) + layers = sorted({layer for (_, layer) in pooled}) + return {layer: [float(np.median(pooled[(m, layer)])) for m in METRICS] + for layer in layers if all((m, layer) in pooled for m in METRICS)} + + +def cmd_gates(args) -> None: + runs = sorted(globlib.glob(args.runs_glob)) + print(f"Gate report — source runs: {[r.split('/')[-1] for r in runs]}") + pooled = _table_from_runs(args.runs_glob) + + # G1: per-seed tables vs pooled + worst = 0.0 + for r in runs: + t = _table_from_runs(r) + for k in t: + if k not in pooled: + continue + for i in range(3): + worst = max(worst, abs(t[k][i] - pooled[k][i])) + status = "OK" if worst < args.warn else "WARN — noisy estimator, check per-seed tables" + print(f"G1 seed stability: max |alignment delta| seed-vs-pooled = {worst:.4f} [{status}]") + + # G2: adjacent source LR vs best + if args.adjacent_glob and globlib.glob(args.adjacent_glob): + adj = _table_from_runs(args.adjacent_glob) + worst = max(abs(adj[k][i] - pooled[k][i]) + for k in set(adj) & set(pooled) for i in range(3)) + status = "OK" if worst < args.warn else "WARN — source-LR sensitive" + print(f"G2 source-LR stability: max |alignment delta| = {worst:.4f} [{status}]") + else: + print("G2 source-LR stability: skipped (no adjacent-LR runs)") + + +def main() -> None: + p = argparse.ArgumentParser(description=__doc__) + sub = p.add_subparsers(dest="cmd", required=True) + + b = sub.add_parser("best-lr"); b.add_argument("--glob", required=True) + b.set_defaults(fn=cmd_best_lr) + + e = sub.add_parser("e1") + e.add_argument("--baseline-glob", required=True) + e.add_argument("--transfer-glob", required=True) + e.add_argument("--tol", type=float, default=0.02) + e.set_defaults(fn=cmd_e1) + + g = sub.add_parser("gates") + g.add_argument("--runs-glob", required=True) + g.add_argument("--adjacent-glob", default=None) + g.add_argument("--warn", type=float, default=0.05) + g.set_defaults(fn=cmd_gates) + + args = p.parse_args() + args.fn(args) + + +if __name__ == "__main__": + main() diff --git a/experiments/vision/run.sh b/experiments/vision/run.sh index 0d13d7c..38e004f 100755 --- a/experiments/vision/run.sh +++ b/experiments/vision/run.sh @@ -1,39 +1,54 @@ #!/usr/bin/env bash -# Submit a vision sweep for one scale. +# Vision muTransfer pipeline (ViT width ladder), mirroring the Llama3 LM flow. +# +# Three phases: +# 1) measure : run s2 mup-no --measure-only (1 seed) to log alignment +# 2) export : build the alignment table from that run's metrics.json +# 3) sweep : mup-no + maxP-meas across the ladder (maxP-meas uses the table) # # Usage: -# bash experiments/vision/run.sh vit-s [extra launch_sweep.py flags ...] +# bash experiments/vision/run.sh measure +# bash experiments/vision/run.sh export +# bash experiments/vision/run.sh sweep [extra launch_sweep.py flags] set -euo pipefail -SCALE="${1:?Usage: $0 [extra launch_sweep.py flags]}" -shift - -case "$SCALE" in - debug) METHODS="maxP"; LRS="3e-3"; SEEDS="1" ;; - vit-s) METHODS="maxP mup-full mup-no"; LRS="3e-4 1e-3 3e-3 1e-2 3e-2 1e-1 3e-1 1.0 3.0"; SEEDS="1 2 3" ;; - vit-b) METHODS="maxP mup-full mup-no"; LRS="3e-4 1e-3 3e-3 1e-2 3e-2 1e-1 3e-1 1.0 3.0"; SEEDS="1" ;; - vit-l) METHODS="maxP mup-full mup-no"; LRS="3e-4 1e-3 3e-3 1e-2 3e-2 1e-1 3e-1 1.0 3.0"; SEEDS="1" ;; - mlp-s) METHODS="maxP mup-full mup-no"; LRS="3e-4 1e-3 3e-3 1e-2 3e-2 1e-1 3e-1 1.0 3.0"; SEEDS="1 2 3" ;; - mlp-m) METHODS="maxP mup-full mup-no"; LRS="3e-4 1e-3 3e-3 1e-2 3e-2 1e-1 3e-1 1.0 3.0"; SEEDS="1 2 3" ;; - mlp-b) METHODS="maxP mup-full mup-no"; LRS="3e-4 1e-3 3e-3 1e-2 3e-2 1e-1 3e-1 1.0 3.0"; SEEDS="1 2" ;; - mlp-l) METHODS="maxP mup-full mup-no"; LRS="3e-4 1e-3 3e-3 1e-2 3e-2 1e-1 3e-1 1.0 3.0"; SEEDS="1" ;; - *) echo "Unknown scale '$SCALE'. Choose from debug vit-s vit-b vit-l mlp-s mlp-m mlp-b mlp-l."; exit 1 ;; -esac +PHASE="${1:?Usage: $0 ...}" +shift || true RUNS_DIR="${RUNS_DIR:-/net/storage/pr3/plgrid/plggadlers/maxP/runs}" VENV_PATH="${VENV_PATH:-/net/storage/pr3/plgrid/plggadlers/maxP/.venv}" REPO_PATH="${REPO_PATH:-/net/storage/pr3/plgrid/plggadlers/maxP}" +ALIGN_TABLE="${ALIGN_TABLE:-${REPO_PATH}/experiments/vision/s2_align.json}" + +# Fixed epoch budget — identical #epochs (= #samples) at every width. +EPOCHS="${EPOCHS:-4}" +BATCH_SIZE="${BATCH_SIZE:-512}" +DATASET="${DATASET:-imagenet12k}" cd "$REPO_PATH" -python experiments/vision/launch_sweep.py \ - --scale "$SCALE" \ - --methods $METHODS \ - --lrs $LRS \ - --seeds $SEEDS \ - --runs-dir "$RUNS_DIR" \ - --venv-path "$VENV_PATH" \ - --repo-path "$REPO_PATH" \ - --epochs 10 \ - --batch-size 3072 \ - "${@}" +case "$PHASE" in + measure) + python experiments/vision/launch_sweep.py \ + --scale s2 --methods mup-no --lrs 3e-2 --seeds 1 \ + --measure-only \ + --runs-dir "$RUNS_DIR" --venv-path "$VENV_PATH" --repo-path "$REPO_PATH" \ + --dataset "$DATASET" --epochs "$EPOCHS" --batch-size "$BATCH_SIZE" "$@" + ;; + export) + SRC="${1:?Usage: $0 export }" + source "${VENV_PATH}/bin/activate" + python experiments/vision/export_alignment.py "$SRC" -o "$ALIGN_TABLE" + ;; + sweep) + SCALE="${1:?Usage: $0 sweep }"; shift + python experiments/vision/launch_sweep.py \ + --scale "$SCALE" \ + --alignment-table "$ALIGN_TABLE" \ + --runs-dir "$RUNS_DIR" --venv-path "$VENV_PATH" --repo-path "$REPO_PATH" \ + --dataset "$DATASET" --epochs "$EPOCHS" --batch-size "$BATCH_SIZE" "$@" + ;; + *) + echo "Unknown phase '$PHASE'. Choose measure | export | sweep."; exit 1 + ;; +esac diff --git a/experiments/vision/run_debug.sh b/experiments/vision/run_debug.sh index abb1223..5974c36 100755 --- a/experiments/vision/run_debug.sh +++ b/experiments/vision/run_debug.sh @@ -9,7 +9,7 @@ REPO="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" VENV="${REPO}/.venv" SCALE="${SCALE:-debug}" -METHOD="${METHOD:-maxP}" +METHOD="${METHOD:-mup-no}" DATASET="${DATASET:-beans}" STEPS="${STEPS:-200}" BATCH_SIZE="${BATCH_SIZE:-8}" @@ -35,6 +35,7 @@ mkdir -p "${OUTPUT_DIR}" python experiments/vision/train.py \ --scale "${SCALE}" \ --method "${METHOD}" \ + --measure-only \ --dataset "${DATASET}" \ --lr 0.03 \ --batch-size "${BATCH_SIZE}" \ diff --git a/experiments/vision/train.py b/experiments/vision/train.py index 868c707..1c83028 100644 --- a/experiments/vision/train.py +++ b/experiments/vision/train.py @@ -5,7 +5,6 @@ import argparse import json -import math import os import time from pathlib import Path @@ -24,7 +23,19 @@ from hf_vision_data import DATASET_CONFIGS, build_dataloaders from maxp_timm import SCALE_CONFIGS, count_trainable_params, create_model, install_pm_wrappers -from utils import append_json, collect_alignments, collect_layer_lrs, save_checkpoint, write_json +from utils import ( + append_json, + collect_alignments, + collect_layer_lrs, + latest_checkpoint, + load_checkpoint, + save_checkpoint, + write_json, +) + +# Device peak dense throughput for the MFU estimate (GH200 bf16 ≈ 989 TFLOP/s). +# Override via env for a different GPU/dtype, or MFU is off by the mis-spec ratio. +DEVICE_PEAK_FLOPS = float(os.getenv("DEVICE_PEAK_FLOPS", 989e12)) def evaluate( @@ -88,15 +99,20 @@ def current_lr_prefactor(optimizer: torch.optim.Optimizer) -> float: def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() - parser.add_argument("--scale", type=str, default="vit-s") - parser.add_argument("--method", choices=["maxP", "mup-full", "mup-no"], default="maxP") + parser.add_argument("--scale", type=str, default="s2") + parser.add_argument("--method", choices=["mup-no", "maxP-meas"], default="mup-no") + parser.add_argument("--measure-only", action="store_true", default=False, + help="Measure + log alignment without changing LRs (mup-no source runs)") + parser.add_argument("--alignment-table", default=None, + help="JSON alignment table from export_alignment.py (maxP-meas only)") + parser.add_argument("--c-ema", type=float, default=0.0, help="EMA smoothing for c (maxP only)") parser.add_argument("--lr", type=float, default=1e-2) parser.add_argument("--seed", type=int, default=42) parser.add_argument("--dataset", type=str, default="imagenet12k") - parser.add_argument("--batch-size", type=int, default=128) - parser.add_argument("--epochs", type=int, default=5) - parser.add_argument("--num-workers", type=int, default=8) - parser.add_argument("--prefetch-factor", type=int, default=4) + parser.add_argument("--batch-size", type=int, default=512) + parser.add_argument("--epochs", type=int, default=4) + parser.add_argument("--num-workers", type=int, default=16) + parser.add_argument("--prefetch-factor", type=int, default=2) parser.add_argument("--max-steps", type=int, default=None) parser.add_argument("--val-interval", type=int, default=500) parser.add_argument("--val-steps", type=int, default=50) @@ -108,9 +124,13 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--no-compile", action="store_true", default=False) parser.add_argument("--debug", action="store_true", default=False) parser.add_argument("--log-interval", type=int, default=20) - parser.add_argument("--checkpoint-interval", type=int, default=None) + parser.add_argument("--checkpoint-interval", type=int, default=1000, + help="Save a checkpoint every N steps (all scales)") + parser.add_argument("--keep-latest-k", type=int, default=2, + help="Number of recent step checkpoints to retain") parser.add_argument("--grad-clip", type=float, default=1.0) - parser.add_argument("--resume", default=None) + parser.add_argument("--resume", action="store_true", default=False, + help="Resume from the latest checkpoint in /checkpoint if present") parser.add_argument("--output-dir", type=str, default="runs") return parser.parse_args() @@ -129,10 +149,10 @@ def main() -> None: scale_cfg = SCALE_CONFIGS.get(args.scale) dataset_cfg = DATASET_CONFIGS.get(args.dataset) + steps_per_epoch = max(1, dataset_cfg.train_samples // args.batch_size) if args.max_steps is not None: total_steps = args.max_steps else: - steps_per_epoch = max(1, dataset_cfg.train_samples // args.batch_size) total_steps = steps_per_epoch * args.epochs num_classes = dataset_cfg.num_classes @@ -141,7 +161,7 @@ def main() -> None: num_classes=num_classes, image_size=scale_cfg.image_size, ) - model.set_grad_checkpointing(enable=True) + model.set_grad_checkpointing(enable=args.scale in ("s4", "s5")) install_pm_wrappers(model) model = model.to(device) @@ -161,6 +181,7 @@ def main() -> None: train_transform = create_transform(**model_data_cfg, is_training=True) eval_transform = create_transform(**model_data_cfg, is_training=False) + data_gen = torch.Generator() # reseeded per pass so the shuffle order is reproducible train_loader, val_loader, _ = build_dataloaders( dataset_name=args.dataset, batch_size=args.batch_size, @@ -168,11 +189,19 @@ def main() -> None: prefetch_factor=args.prefetch_factor, train_transform=train_transform, eval_transform=eval_transform, + generator=data_gen, ) sample_input = torch.randn(1, 3, scale_cfg.image_size, scale_cfg.image_size, device=device) - alignment_mode = "full" if "full" in args.method else "no" - dynamic = args.method == "maxP" + alignment_mode = "no" + dynamic = args.measure_only + + alignment_overrides = None + if args.method == "maxP-meas": + if args.alignment_table is None: + raise ValueError("method 'maxP-meas' requires --alignment-table") + with open(args.alignment_table) as f: + alignment_overrides = {k: tuple(v) for k, v in json.load(f).items()} param = Parametrization( model, @@ -180,9 +209,12 @@ def main() -> None: alignment=alignment_mode, lr_prefactor=args.lr, sample_input=sample_input, + alignment_overrides=alignment_overrides, warmup_steps=args.alignment_warmup, solve_interval=args.solve_interval, sample_size=args.sample_size, + c_ema=args.c_ema, + measure_only=args.measure_only, ) optimizer = torch.optim.AdamW( @@ -220,9 +252,14 @@ def main() -> None: log_dir.mkdir(parents=True, exist_ok=True) tb_writer = SummaryWriter(log_dir=str(log_dir)) + # For the 6ND MFU estimate: N = trainable params, D = tokens/sample. + # ViT tokens = (img/patch)^2 patches + 1 CLS (patch16 across all configs). + n_params = count_trainable_params(model) + tokens_per_sample = (scale_cfg.image_size // 16) ** 2 + 1 + print("\n=== maxP vision run ===") print(f" scale: {args.scale} ({scale_cfg.model_name})") - print(f" params: {count_trainable_params(model)}") + print(f" params: {n_params}") print(f" method: {args.method}") print(f" dataset: {args.dataset}") print(f" train_samples: {dataset_cfg.train_samples}") @@ -240,18 +277,35 @@ def main() -> None: global_step = start_epoch = total_seen = 0 align_sample: torch.Tensor | None = None + if args.resume: + ckpt_path = latest_checkpoint(output_dir / "checkpoint") + if ckpt_path is not None: + progress = load_checkpoint(ckpt_path, model=model, optimizer=optimizer, device=device) + global_step = progress["step"] + start_epoch = progress["epoch"] + total_seen = progress["samples_seen"] + print(f"[resume] loaded {ckpt_path} — step={global_step} epoch={start_epoch}") + else: + print(f"[resume] no checkpoint under {output_dir/'checkpoint'} — starting fresh") + + step_budget = total_steps + t0 = time.time() last_log_time = t0 last_log_seen = total_seen stop_training = False - for epoch in range(start_epoch, args.epochs): + pass_idx = start_epoch + while not stop_training: + # Reseed the shuffle each pass → reproducible, deterministic order. + data_gen.manual_seed(args.seed * 1_000_000 + pass_idx) + epoch = pass_idx for xb, yb in train_loader: xb = xb.to(device, non_blocking=True) yb = yb.to(device, non_blocking=True) if dynamic and align_sample is None: - align_sample = xb[:args.sample_size].detach().clone().to(device) + align_sample = xb[:args.sample_size].detach().clone().float().to(device) param.capture_initial(align_sample) optimizer.zero_grad() @@ -300,6 +354,7 @@ def main() -> None: "perf/samples_seen": total_seen, "perf/samples_per_sec": total_seen / elapsed, "perf/samples_per_sec_interval": interval_sps, + "perf/mfu": 6 * n_params * tokens_per_sample * interval_sps / DEVICE_PEAK_FLOPS, "lr/lr_prefactor": current_lr_prefactor(optimizer), **collect_layer_lrs(param), **eval_metrics, @@ -333,14 +388,14 @@ def main() -> None: epoch=epoch, samples_seen=total_seen, args=args, + keep_latest_k=args.keep_latest_k, ) - if args.max_steps is not None and global_step >= args.max_steps: + if global_step >= step_budget: stop_training = True break - if stop_training: - break + pass_idx += 1 eval_metrics = evaluate( model=model, @@ -372,6 +427,9 @@ def main() -> None: ) print(f"\nDone. Total training time: {elapsed:.1f} seconds.") + if torch.cuda.is_available(): + print(f"Peak CUDA mem: {torch.cuda.max_memory_allocated() / 1e9:.1f} GB / " + f"{torch.cuda.get_device_properties(device).total_memory / 1e9:.0f} GB") print(json.dumps(eval_metrics, indent=2, sort_keys=True)) diff --git a/experiments/vision/utils.py b/experiments/vision/utils.py index a393868..9648572 100644 --- a/experiments/vision/utils.py +++ b/experiments/vision/utils.py @@ -21,7 +21,7 @@ def collect_layer_lrs(param: Parametrization) -> dict[str, float]: return out -def collect_alignments(param: Parametrization) -> dict[str, dict[str, float]]: +def collect_alignments(param: Parametrization) -> dict[str, float]: out: dict[str, float] = {} for name, pm in param._pms: if pm.weight is None: @@ -43,9 +43,11 @@ def save_checkpoint( epoch: int, samples_seen: int, args: argparse.Namespace, + keep_latest_k: int = 2, ) -> None: ckpt_path.parent.mkdir(parents=True, exist_ok=True) model_for_io = model._orig_mod if hasattr(model, "_orig_mod") else model + tmp_path = ckpt_path.with_suffix(ckpt_path.suffix + ".tmp") torch.save( { "model": model_for_io.state_dict(), @@ -53,10 +55,65 @@ def save_checkpoint( "step": step, "epoch": epoch, "samples_seen": samples_seen, + "rng": torch.get_rng_state(), + "cuda_rng": torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None, "args": vars(args), }, - ckpt_path, + tmp_path, ) + # Atomic rename so an interrupted save never leaves a half-written ckpt. + tmp_path.replace(ckpt_path) + _prune_checkpoints(ckpt_path.parent, keep_latest_k) + + +def _prune_checkpoints(ckpt_dir: Path, keep_latest_k: int) -> None: + """Keep only the newest `keep_latest_k` step checkpoints (final.pt is never pruned).""" + steps = sorted( + ckpt_dir.glob("step_*.pt"), + key=lambda p: int(p.stem.split("_")[1]), + ) + for stale in steps[:-keep_latest_k]: + stale.unlink(missing_ok=True) + + +def latest_checkpoint(ckpt_dir: Path) -> Path | None: + """Return the highest-step checkpoint in `ckpt_dir`, or final.pt, else None.""" + if not ckpt_dir.is_dir(): + return None + steps = sorted( + ckpt_dir.glob("step_*.pt"), + key=lambda p: int(p.stem.split("_")[1]), + ) + if steps: + return steps[-1] + final = ckpt_dir / "final.pt" + return final if final.is_file() else None + + +def load_checkpoint( + ckpt_path: Path, + *, + model: nn.Module, + optimizer: torch.optim.Optimizer, + device: torch.device, +) -> dict: + """Restore model + optimizer + RNG in-place; return progress (step/epoch/samples_seen).""" + ckpt = torch.load(ckpt_path, map_location=device, weights_only=False) + model_for_io = model._orig_mod if hasattr(model, "_orig_mod") else model + model_for_io.load_state_dict(ckpt["model"]) + optimizer.load_state_dict(ckpt["optimizer"]) + if ckpt.get("rng") is not None: + torch.set_rng_state(ckpt["rng"].cpu() if hasattr(ckpt["rng"], "cpu") else ckpt["rng"]) + if ckpt.get("cuda_rng") is not None and torch.cuda.is_available(): + # map_location may have moved these onto the GPU; set_rng_state_all needs + # CPU ByteTensors. + cuda_rng = [s.cpu() if hasattr(s, "cpu") else s for s in ckpt["cuda_rng"]] + torch.cuda.set_rng_state_all(cuda_rng) + return { + "step": int(ckpt["step"]), + "epoch": int(ckpt["epoch"]), + "samples_seen": int(ckpt["samples_seen"]), + } def write_json(path: Path, obj: dict) -> None: diff --git a/maxp/parametrization.py b/maxp/parametrization.py index 46fc8e6..c3aaba4 100644 --- a/maxp/parametrization.py +++ b/maxp/parametrization.py @@ -103,6 +103,11 @@ class Parametrization: Each step: ``align = ema * align_old + (1 - ema) * align_new``. 0.0 means no smoothing (default). Useful when alignment measurements are noisy (e.g. with ``use_training_activations``). + measure_only: If True, :meth:`step` measures alignment (including + for layers pinned via *alignment_overrides*) and stores it on + each PM for logging, but never re-solves the LP or changes + learning rates. Use to passively record alignment during a + baseline run. resample_w0: If True, store a random seed per layer instead of cloning ``w_0``. Regenerates ``w_0`` on-the-fly during alignment measurement, saving memory proportional to the total weight size. @@ -140,6 +145,7 @@ def __init__( sample_size: int = 32, c_ema: float = 0.0, alignment_ema: float = 0.0, + measure_only: bool = False, resample_w0: bool = False, use_training_activations: bool = False, solver: "plp.LpSolver | None" = None, @@ -217,6 +223,7 @@ def __init__( self._sample_size = sample_size self._c_ema = c_ema self._alignment_ema = alignment_ema + self._measure_only = measure_only self._resample_w0 = resample_w0 self._warm_start = warm_start self._use_training_activations = use_training_activations @@ -521,7 +528,7 @@ def step( from maxp.alignment import compute_alignment for name, pm in self._pms: - if name in self._alignment_pinned: + if name in self._alignment_pinned and not self._measure_only: continue if pm._z0 is None or name not in current: continue @@ -543,6 +550,11 @@ def step( else: pm.align_z0_dW, pm.align_dZ_w0, pm.align_dZ_dW = new_a0_dW, new_dZ_w0, new_dZ_dW + # Measure-only: alignment recorded on PMs for logging, LRs untouched + if self._measure_only: + self._sync_lrs(optimizer) + return + # 3. Re-solve LP (skip c update if infeasible with current alignment) try: c_by_name = self._resolve() diff --git a/paper b/paper new file mode 160000 index 0000000..edaa611 --- /dev/null +++ b/paper @@ -0,0 +1 @@ +Subproject commit edaa611f28b8df07ea48349b9fc75a94d574f67c diff --git a/tests/test_step.py b/tests/test_step.py index 9fce023..e841610 100644 --- a/tests/test_step.py +++ b/tests/test_step.py @@ -989,3 +989,52 @@ def test_all_opts_training_loop(self): assert pm._w0 is None param.remove_hooks() + + +class TestMeasureOnly: + """measure_only=True: alignment measured and stored, c/LRs never change.""" + + def test_lrs_and_c_frozen(self): + model, param, optimizer, X = _setup(measure_only=True) + param.capture_initial(X) + initial = [(g["c"], g["lr"]) for g in param.param_groups if g.get("maxp_managed")] + + for _ in range(5): + optimizer.zero_grad() + model(X).sum().backward() + optimizer.step() + param.step(X, optimizer) + + after = [(g["c"], g["lr"]) for g in param.param_groups if g.get("maxp_managed")] + assert after == initial + + def test_alignment_still_measured(self): + model, param, optimizer, X = _setup(measure_only=True) + param.capture_initial(X) + + for _ in range(5): + optimizer.zero_grad() + model(X).sum().backward() + optimizer.step() + param.step(X, optimizer) + + # Training happened, so at least one PM should have alignment != preset + measured = [pm.align_z0_dW for _, pm in param._pms if pm.weight is not None] + assert any(not math.isclose(v, 1.0) for v in measured) + + def test_pinned_layers_also_measured(self): + model, param, optimizer, X = _setup( + measure_only=True, + alignment_overrides={"hidden": (0.9, 0.5, 0.9)}, + ) + param.capture_initial(X) + + for _ in range(5): + optimizer.zero_grad() + model(X).sum().backward() + optimizer.step() + param.step(X, optimizer) + + # Pinned value should have been overwritten by the realized measurement + pm = dict(param._pms)["hidden"] + assert not math.isclose(pm.align_z0_dW, 0.9, abs_tol=1e-6)