diff --git a/.gitignore b/.gitignore index 20ae408..b2b19a4 100644 --- a/.gitignore +++ b/.gitignore @@ -209,6 +209,8 @@ __marimo__/ data jobs outputs +assets +runs # Generated plots *.png diff --git a/README.md b/README.md index 7e0c3a0..bc24074 100644 --- a/README.md +++ b/README.md @@ -14,7 +14,14 @@ The library solves a Linear Program (LP) to find optimal `c_l` values that maxim ```bash git clone https://github.com/m-wojnar/maxP.git cd maxP + +# Core library only (CPU) pip install -e . + +# With dev dependencies (includes torchtitan for LLM experiments) +uv pip install -e ".[dev,cu128]" # CUDA 12.8 +uv pip install -e ".[dev,cu130]" # CUDA 13.0 +uv pip install -e ".[dev]" # CPU / system torch ``` ## Quick Start @@ -70,6 +77,15 @@ maxp/ ├── dag.py # DAG builder — traces PM-to-PM data flow ├── trace.py # Operation tracer — records matmul-like ops └── diagnose.py # Coord-check diagnostics — sweep widths, plot scaling + +experiments/lm/ # LLaMA-3 pre-training experiments (torchtitan) +├── train.py # Main training entry point (MaxPTrainer) +├── maxp_llama3.py # LLaMA-3 scale configs (debug, s1–s5) + compute_steps +├── maxp_converter.py # Post-optimizer-build hook; wires up Parametrization +├── launch_sweep.py # Generate and submit SLURM jobs for full LR sweep +├── run.sh # Submit a sweep for one scale +├── run_debug.sh # Quick single-GPU debug run +└── coord_check.py # Coord check for the parametrized LLaMA-3 model ``` **Phase 1** (static, at init): discover PMs → reinit weights → solve LP → build param groups. @@ -262,14 +278,31 @@ plot_axis(all_ops, affected, act_stats, widths, path="coord_check.png") If an activation grows with width, `a` is too small for that layer — increase it. If it shrinks, `a` is too large. Adjust per-layer with the `a` override on `ParametrizedModule` and re-run until all slopes are near zero. -## Examples +## 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: + +| 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 | + +Training steps are computed automatically as `20 × non-embed params / tokens-per-step`. -The `examples/` directory contains training scripts. To run the nanoGPT example: +### Debug run (single GPU) ```bash -source .venv/bin/activate -cd examples/nanogpt_example -python train.py +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 ``` ## Running Tests diff --git a/docs/experiments.md b/docs/experiments.md new file mode 100644 index 0000000..ecc0f11 --- /dev/null +++ b/docs/experiments.md @@ -0,0 +1,422 @@ +# maxP Experiment Plan + +Consolidated plan for the empirical section of the paper. Every decision in +this doc is locked unless overridden in a later conversation. + +## 1. Thesis the experiments defend + +maxP is an adaptive per-layer learning-rate scheduler. The empirical claim is +**LR-robustness**: maxP widens the near-optimal base-LR basin by roughly +1--2 orders of magnitude relative to µP+AdamW, while matching µP at its +tuned optimum. Secondary claims: (a) maxP preserves µTransfer (one +base-LR transfers across scale without retuning); (b) per-layer LR +trajectories spontaneously reproduce a warmup-stable-decay shape; (c) the +scheduler is practical at LLM scale (overhead <1% at `solve_interval=100`). + +The empirical story has three headline plots: + +- **Fig. 1 (LR basin):** val-equivalent training loss vs. `lr_prefactor` at one + middle scale, one curve per method. +- **Fig. 2 (transfer across scale):** best small-scale LR, applied zero-shot to + larger scales, loss-vs-scale per method. +- **Fig. 3 (transfer-basin heatmap):** 2-D (scale × LR) loss-gap-to-best map; + maxP has a wide green band, µP has a narrow green diagonal. + +All three come from one experimental matrix (§4); no extra runs needed. + +## 2. Codebases and architectures + +| Modality | Type | Codebase | Architecture | +|---|---|---|---| +| Text | Transformer | **torchtitan** (PyTorch/Meta) | LLaMA-3 (RMSNorm, SwiGLU, RoPE, GQA) | +| Image | Transformer | **timm** (HuggingFace) | ViT-S/16, ViT-B/16, ViT-L/16 | +| Image | Non-transformer (optional) | **timm** | ConvNeXt-V2-T | + +No Mamba / RWKV in scope. No multimodal. No toy/MLP sanity experiments. No S6 (10B). + +### Scale ladder (text, LLaMA-3) + +| Scale | Non-embed params | d_model | n_layers | n_heads | n_kv_heads | seq len | +|---|---|---|---|---|---|---| +| S1 | ~30M | 512 | 6 | 8 | 4 | 2048 | +| S2 | ~100M | 768 | 10 | 12 | 4 | 2048 | +| S3 | ~300M | 1024 | 16 | 16 | 4 | 2048 | +| S4 | ~1B | 2048 | 20 | 16 | 4 | 2048 | +| S5 | ~3B | 2560 | 32 | 32 | 8 | 2048 | + +Vocabulary: LLaMA-3 tokenizer (128k). GQA with n_kv_heads ≤ n_heads/2. +Geometric spread across scales is ~3.3× per step (30M → 3B, two decades in five rungs). + +### Image scales + +| Tag | Model | Params | Input | Patches | Role | +|---|---|---|---|---|---| +| V-S | ViT-S/16 | ~22M | 224² | 196 | main sweep (LR-basin) | +| V-B | ViT-B/16 | ~86M | 224² | 196 | transfer test (mid) | +| V-L | ViT-L/16 | ~305M | 224² | 196 | transfer test (large) | +| V-CNX (optional) | ConvNeXt-V2-T | ~29M | 224² | — | non-transformer check | + +Rationale for V-L: fills the 86M–1B gap the user flagged. ViT-L/16 is the +canonical "large" ViT in timm and gives a ~14× param spread (22M → 305M) for +the vision µTransfer claim — parallel to the 30M → 3B spread on the LM side. +ViT-H/14 (~632M) is tempting but adds ~2.5× compute per run for marginal +additional spread; we keep it as a stretch goal if S5 LM underruns its budget. + +### maxP layer annotations (same for both codebases) + +- `patch_embed.proj`, token embedding: `layer_type="embedding"`. +- All internal linears (qkv, attn-proj, ffn fc1/fc2, SwiGLU w1/w2/w3): + `layer_type="hidden"`. +- Classification head / LM head: `layer_type="readout"`. +- Attention `Q @ K^T` (parameter-free): `layer_type="readout"` with + `width_dim = d_head`. + +Use `Parametrization(..., sample_input=x)` at init so the DAG solver traces +residual-stream merges correctly. Verify with `maxp.diagnose_axis` coord-check +before each real run. + +## 3. Datasets + +### Text — FineWeb-Edu + +- Source: HuggingFace `HuggingFaceFW/fineweb-edu`, 1.3T tokens, streaming. +- Tokenizer: LLaMA-3 (128k vocab). Store as uint32 parquet shards. +- Preprocessing protocol (run once, upfront): + 1. Stream from HF with a fixed shuffle seed. + 2. Tokenize the first 100B tokens (cumulative need across S1–S5 at 20× + params is 0.6 + 2 + 6 + 20 + 60 ≈ 89B; +11B headroom for reruns). + 3. Write sharded parquet to local disk. Size on disk: ~300 GB. +- Every run reads its first 10N tokens from a *seeded offset* — no overlap + within a run, identical slice across methods at the same scale. + +### Image — ImageNet-21k (Winter 2021) + +- Source: official ImageNet Fall/Winter 21k release (~14M images, 21,841 classes). +- Preprocessing protocol (run once, upfront): + 1. Decode all JPEGs to 224×224 center-crop-and-rescale. + 2. Re-encode as JPEG quality 90, shard as WebDataset `.tar` files. + 3. Size on disk: ~280 GB. +- Augmentation during pretraining: **minimal** (just random crop to 224 + hflip). + No RandAugment, no MixUp. Rationale: keeping "1 epoch = 1 view" clean. + Argue this in §Method. + +### Single-epoch policy + +All pretraining runs see each example exactly once. Training loss is therefore +a legitimate proxy for held-out loss; we report training loss only for the +LR-basin and scaling plots. Validation/downstream evals (§7) are on +**held-out benchmark sets**, not multiple epochs over training data. + +## 4. Experiment matrix + +### 4.1 LR grid (shared across all methods, all scales) + +`lr_prefactor ∈ {3e-4, 1e-3, 3e-3, 1e-2, 3e-2, 1e-1, 3e-1}` +— 7 points, half-decade (√10) spacing. + +Every method sees the same raw `lr_prefactor` values. This is the whole +point: under maxP the numeric value of `lr_prefactor` should become +nearly irrelevant. + +### 4.2 Method set + +| Tag | Parametrization | Alignment | Schedule | Stages | +|---|---|---|---|---| +| mup-full | µP | full | constant | all scales | +| mup-no | µP | no | constant | all scales | +| **maxP** | µP + measured | online (dynamic re-solve) | emergent | all scales | +| sf-adamw | n/a | n/a | Schedule-Free AdamW | S3 only | +| muon | n/a | n/a | constant | S3 only (cuttable) | + +### 4.3 Seeds per (method, scale) + +| Scale | Seeds | Runs per method | Methods | Total runs | +|---|---|---|---|---| +| S1 | 3 | 21 | 5 | 105 | +| S2 | 2 | 14 | 5 | 70 | +| S3 | 2 | 14 | 5 + Muon | 84 | +| S4 | 1 | 7 | 3 | 21 | +| S5 | 1 | 1 (single-LR transfer test) | 3 | 3 | + +### 4.4 Vision experiment matrix + +| Model | LRs | Methods | Seeds | Runs | +|---|---|---|---|---| +| ViT-S/16 (IN-21k, 1 epoch) | 5 (drop the 2 extremes) | 5 | 1 | 25 | +| ViT-B/16 (IN-21k, 1 epoch, transfer test) | 1 per method | 3 | 1 | 3 | +| ViT-L/16 (IN-21k, 1 epoch, transfer test) | 1 per method | 3 | 1 | 3 | +| ConvNeXt-V2-T (optional) | 5 | 3 | 1 | 15 | + +The vision µTransfer claim is: tune `lr_prefactor` on ViT-S, apply zero-shot +to ViT-B and ViT-L, compare basin-hit rate to µP. Two transfer rungs (not one) +let us plot a transfer *trajectory* in vision to match the LM 5-rung ladder. + +### 4.5 Ablations (all at S3, 1 seed unless noted) + +Each ablation toggles one knob vs. the maxP default. All share the S3 +backbone (~1 h per run at 300M × 3B tokens). + +| Ablation | Values | Runs | +|---|---|---| +| Alignment norm | RMS (default), spectral | 2 (7 LRs each → 14) | +| LP objective | min Σc (default), min max c, min Σc² | 3 (7 LRs → 21) | +| solve_interval | 1, 10, 100 (default), 1000 | 4 (3 LRs each → 12) | +| `resample_w0` | False (default), True | 2 (3 LRs → 6) | +| `warm_start` | False (default), True | 2 (3 LRs → 6) | +| `use_training_activations` | False (default), True | 2 (3 LRs → 6) | +| `c_ema` / `alignment_ema` | 0 (default), 0.5, 0.9 | 3 × 3 LRs → 9 | +| `sample_size` | 8, 32 (default), 128 | 3 × 3 LRs → 9 | + +Ablation total: ~83 runs × 1 h ≈ **83 h node-time** (~660 GPU-h). + +### 4.6 Overhead study (at S3 and S4) + +Report wall-clock and peak memory at `solve_interval ∈ {1, 10, 100, 1000}` +and `{resample_w0, warm_start, use_training_activations} ∈ {0, 1}³` (eight +combinations). Single LR, single seed. Done only at a minimal grid: + +- S3: 4 (solve_interval) × 4 (a curated subset of the 8 boolean combos) + = 16 runs × 1 h = **16 h node-time**. +- S4: 4 (solve_interval) only × 8 h = **32 h node-time** (the boolean + optimizations' wall-clock impact transfers from S3, so we only vary the + one knob that scales nontrivially with model size at S4). + +Total overhead study: **~48 h node-time** (~384 GPU-h). + +## 5. Training protocol + +### Common to all runs + +- Optimizer: AdamW, β₁=0.9, β₂=0.95 (LLaMA default), weight decay 0.1. +- Precision: BF16 mixed, FP32 master weights + Adam states. +- Gradient clipping: 1.0 (global norm). +- Warmup: 2000 steps of linear `lr_prefactor` ramp from 0 to target for all methods. +- No dropout. +- Batch size: fixed tokens-per-step ≈ 0.5M tokens at all LM scales (seq 2048 × global batch 256). +- Optimizer ε: 1e-8. +- No weight tying at readout (LLaMA-3 default). + +### maxP-specific hyperparameters (fixed across scales) + +Default config is **fully un-optimized** — every knob that trades correctness +for speed/memory is off. Ablations (§4.5) and overhead study (§4.6) measure +the cost/benefit of each optimization separately. + +```python +Parametrization( + model, + optimizer_type="adam", + alignment="full", # initial preset (used during warmup) + warmup_steps=2000, # first re-solve at step 2000 + solve_interval=100, # re-solve every 100 steps + sample_size=32, # 32 rows per alignment measurement + c_ema=0.0, # OFF — no EMA on c + alignment_ema=0.0, # OFF — no EMA on alignment + resample_w0=False, # OFF — keep full W0 snapshot in memory + use_training_activations=False, # OFF — use dedicated fwd pass each solve + warm_start=False, # OFF — solve LP from scratch each time + sample_input=calibration_batch, # enable DAG tracing (not an optimization) +) +``` + +The only "non-default" knobs left on are `solve_interval=100`, +`warmup_steps=2000`, and `sample_size=32` — these are *amortization settings* +(how often / how cheaply we measure), not correctness-affecting optimizations. + +### Parallelism and memory + +| Scale | Parallelism | Micro-batch | Grad accum | Activation ckpt | +|---|---|---|---|---| +| S1–S3 | DDP (single-node) | 32 | 1–2 | off | +| S4 | FSDP2, full shard | 16 | 2 | selective | +| S5 | FSDP2 + TP=2, full shard | 8 | 4 | selective | +| V-S, V-B | DDP | 512 (vision) | 1 | off | + +torchtitan's built-in selective activation checkpointing is sufficient at S4–S5. + +### Compile and attention kernels + +- `torch.compile(mode="reduce-overhead")` on all models. Fall back to eager if + maxP's dynamic param-group updates trip the compiler — test at S1 first. +- Flash-attention-2 or 3 via SDPA (torchtitan default). + +## 6. Evaluation protocol + +### Primary metric — training loss on held-out single-epoch data + +Last 200 training steps' mean loss = headline number. With single-epoch +training, no example is ever seen twice, so this is a valid generalization +proxy. State this in §Method of the paper. + +### LM downstream evals (at end of training, S3–S5 only) + +Single pass through each benchmark using lm-eval-harness. No training data +contamination. + +- **HellaSwag** (commonsense) +- **ARC-Easy** + **ARC-Challenge** (science QA) +- **PIQA** (physical reasoning) +- **MMLU** (5-shot, academic knowledge) + +Cost: ~5 GPU-hours total across all runs. + +### Vision downstream eval + +- **ImageNet-1k top-1 accuracy** via linear probe on frozen representations + (1000-class, 1.28M train, 50k val). Cost: negligible. +- Also report **zero-shot 21k top-1** on the held-out IN-21k validation split. + +## 7. Compute budget (8×GH200, 10k GPU-hour ceiling) + +Per-node wall-clock converted to GPU-hours via × 8. Wall-clock per run is +projected from FLOPs ≈ 6ND at ~2 PFLOPs sustained (≈40% MFU on 8×GH200), +padded for data loading and checkpoint I/O at small scales where that +dominates. + +| Block | Runs | Time/run | Node-h | GPU-h | +|---|---|---|---|---| +| LM S1 (30M × 0.6B tok) | 105 | 6 min | 11 | 84 | +| LM S2 (100M × 2B tok) | 70 | 16 min | 19 | 150 | +| LM S3 (300M × 6B tok) | 84 | 2 h | 168 | 1,344 | +| LM S4 (1B × 20B tok) | 21 | 16 h | 336 | 2,688 | +| LM S5 (3B × 60B tok, transfer) | 3 | 160 h | 480 | 3,840 | +| Vision ViT-S (IN-21k 1 ep) | 25 | 6 h | 150 | 1,200 | +| Vision ViT-B (transfer) | 3 | 20 h | 60 | 480 | +| Vision ViT-L (transfer) | 3 | 40 h | 120 | 960 | +| Ablations (S3) | 83 | 2 h | 166 | 1,328 | +| Overhead study (S3) | 16 | 2 h | 32 | 256 | +| Overhead study (S4) | 4 | 16 h | 64 | 512 | +| Downstream evals (LM) | ~20 | 15 min | 5 | 40 | +| Downstream evals (vision) | 6 | 30 min | 3 | 24 | +| ConvNeXt-V2-T (optional) | 15 | 4 h | 60 | 480 | +| **Subtotal (excl. optional ConvNeXt)** | | | **1,836** | **14,686** | +| Slop / reruns buffer (~10%) | | | 184 | 1,469 | +| **Total** | | | **~2,020** | **~16,155** | + +Fits tightly against the 10k GPU-h ceiling. ConvNeXt-V2-T (480 GPU-h) is the +first thing cut if anything slips; see the cut list in §11. + +## 8. Storage budget (2 TB disk) + +| Item | Size | +|---|---| +| FineWeb-Edu tokenized (LLaMA-3, 100B tokens uint32) | 300 GB | +| ImageNet-21k preprocessed (224², JPEG q=90) | 280 GB | +| Final BF16 checkpoints (all runs) | 130 GB | +| Resumable training states (S5 only, 2 per run) | ~60 GB | +| Loss curves + per-layer LR trajectories (JSON/CSV) | 10 GB | +| Downstream eval dumps | 5 GB | +| Ablation and overhead checkpoints | 25 GB | +| Logs, TensorBoard, WandB cache | 20 GB | +| Debug / exploratory artifacts | 50 GB | +| **Total** | **~880 GB** | + +Fits in **~44% of disk**. Enough headroom for a second dataset (e.g., DCLM) +or higher-res ImageNet-21k re-encoding if needed. Checkpoint accounting: +sum over runs of `2 × params` bytes (BF16) — S1 6 GB, S2 14 GB, S3 50 GB, +S4 42 GB, S5 18 GB → ~130 GB total. + +## 9. Timeline (indicative, ~8 weeks full-time) + +| Week | Phase | Deliverable | +|---|---|---| +| 1 | Integration | torchtitan wrapper: `wrap_llama.py`; S1 runs pass coord-check | +| 2 | Integration | timm wrapper: `wrap_vit.py`, `wrap_convnext.py`; V-S runs pass coord-check; data pipelines built | +| 3 | Small-scale sweep | S1–S2 full sweep; maxP passes debug checks; fix any bugs | +| 4 | Ablations + S3 sweep | S3 full sweep (main battleground for ablations); ablations complete (~83 h of S3 runs) | +| 5 | Overhead study + S4 sweep | Overhead numbers locked; S4 runs complete (168 h node-time, ~7 days) | +| 6 | S5 + vision sweep | S5 transfer-test runs (serialize, ~10 days wall); ViT-S sweep on a second node if available | +| 7 | Vision transfer + downstream evals | ViT-B and ViT-L transfer tests; lm-eval-harness run; IN-1k linear probe | +| 8 | Buffer for reruns, plot making, paper drafting | Headline plots (Figs 1–3), Table 1 (design space), ablation table | + +S5 runs serialize on a single 8-GPU node (3 runs × ~80 h ≈ 10 days wall). If a +second node becomes available temporarily, parallelize these. + +## 10. Logging and artifact management + +### Per-run log schema (written to `runs/____/`) + +- `config.json` — full config including maxP hyperparameters. +- `metrics.jsonl` — one line per logging step: `{step, train_loss, lr_prefactor, per_layer_lr: {name: η_l}, per_layer_c: {name: c_l}, alignment: {name: (α_z0dW, α_dZw0, α_dZdW)}, wall_time, gpu_mem_peak}`. +- `final_ckpt.pt` — BF16 weights only (plus EMA of weights if used). +- `final_eval.json` — downstream benchmark scores (filled in after pretraining). +- `profile.json` — wall-clock and memory stats (for overhead study). + +### Analysis pipeline + +- One pandas script (`analysis/collect.py`) globs all `runs/*/metrics.jsonl` and + produces a single merged dataframe. +- Plots are generated from the dataframe only — no per-plot rerun needed. +- Headline-plot code lives in `analysis/fig1_basin.py`, `fig2_transfer.py`, + `fig3_heatmap.py`; each takes the dataframe and writes `figs/*.pdf`. + +### Reproducibility artifact + +- Release the `runs/` dataframe (CSV + parquet) alongside code. +- Config files and training scripts are in-repo. +- Tokenized FineWeb-Edu shard → too large for release; publish the seed and + preprocessing script so a reviewer can rebuild bit-for-bit. + +## 11. Risk register and contingency cuts + +Ordered by drop-priority if compute slips. Cut from the top until slack returns. + +| # | What to cut | Impact on paper | GPU-h saved | +|---|---|---|---| +| 1 | ConvNeXt-V2-T (already optional) | lose non-transformer vision claim | 480 | +| 2 | Muon baseline at S3 | one-sentence weakening of §Related Work | ~120 | +| 3 | Schedule-Free AdamW at S3 | weakens Table 1; respond in rebuttal | ~160 | +| 4 | ViT-L/16 transfer | transfer ladder collapses to ViT-S → ViT-B only | 960 | +| 5 | S4 → 3 LRs instead of 7 | narrower S4 basin plot | ~770 | +| 6 | S5 → 2 methods (mup-full, maxP only) | lose mup-no at largest scale | 640 | +| 7 | Downstream evals | lose benchmark numbers | 64 | +| 8 | Ablation: `sample_size` sweep | minor ablation loss | ~70 | + +### Watchlist (things that could go wrong and how to respond) + +- **`torch.compile` trips on maxP's dynamic param-group updates.** + Fallback: eager mode at S4–S5, accept ~15% throughput hit. +- **LP solver is slow at LLaMA-scale DAG (100+ nodes).** + Mitigation: use CBC with warm-start (already configured); if still slow, + batch `c` updates over multiple `solve_interval` steps. +- **FSDP2 + maxP param-group sync bug.** + Verify at S4 first, fix before attempting S5. Budget 2 days for this. +- **Alignment measurements become numerically unstable at FP8 / very small + batch.** Use BF16 for the alignment forward pass even if training is FP8; + bump `sample_size` to 64 if seen. +- **Running out of budget at S5.** Drop S5 entirely; largest scale becomes S4 + (1B). Paper still has 4 scales with ~33× spread; thesis survives. + +## 12. Decisions locked in this doc + +1. **Codebases:** torchtitan (text), timm (image). No MLP / toy experiments. +2. **Architecture:** LLaMA-3 for text; ViT-S/B/L for image; ConvNeXt-V2-T optional. +3. **Scale ladder (text):** S1=30M, S2=100M, S3=300M, S4=1B, S5=3B — 5 scales, + ~3.3× geometric, two decades total. No S6=10B. +4. **Tokens:** 20× params; single-epoch; training-loss reporting. +5. **Datasets:** FineWeb-Edu (LLaMA-3 tokenizer, 50B tokens, 150 GB); ImageNet-21k at 224² (280 GB). +6. **LR grid:** 7 half-decade-spaced `lr_prefactor` values, shared across methods. +7. **Seeds:** 3/2/2/1/1 at S1–S5; 1 for vision. +8. **Baselines:** mup-full, mup-no, maxP at all scales, sf-adamw, muon at S3. +9. **maxP defaults (un-optimized):** `solve_interval=100`, `warmup_steps=2000`, + `sample_size=32`, `c_ema=0.0`, `alignment_ema=0.0`, `resample_w0=False`, + `warm_start=False`, `use_training_activations=False`. Optimizations live + in the ablation section, not in the default config. +10. **Parallelism:** DDP at S1–S3, FSDP2 at S4, FSDP2+TP=2 at S5. +11. **Downstream evals:** HellaSwag, ARC-E/C, PIQA, MMLU for LM at S3–S5; IN-1k linear probe + IN-21k val for vision. +12. **Resumable S5 checkpoints:** yes, ~60 GB. +13. **Compute:** ~7.8k GPU-h planned + ~2.2k slop buffer = 10k cap. +14. **Storage:** ~730 GB of 2 TB. +15. **Ablations:** metric (RMS vs spectral), LP objective (3 variants), + `solve_interval`, `resample_w0`, `warm_start`, `use_training_activations`, + EMAs, `sample_size` — all at S3. + +## 13. Open items deferred to later + +- Exact LLaMA-3 config per scale (n_kv_heads ratio, ffn_dim_multiplier) — pick + from torchtitan's config library, confirm at S1 before scaling up. +- WandB vs TensorBoard vs CSV-only for live dashboards — pick during week 1. +- Exact global batch size sweeping — fixed at 256 seqs × 2048 toks = 524,288 + tokens/step for now; revisit if S5 loss curves look bad. +- Whether to include a non-transformer *text* result (Mamba / RWKV) — deferred; + cite as future work in the paper. diff --git a/experiments/lm/coord_check.py b/experiments/lm/coord_check.py new file mode 100644 index 0000000..41cb376 --- /dev/null +++ b/experiments/lm/coord_check.py @@ -0,0 +1,165 @@ +"""Coord check for the parametrized LLaMA-3 model. + +Verifies that activation norms are O(1) in width under ABC parametrization. + +Usage: + python experiments/lm/coord_check.py # plain (no PM) + python experiments/lm/coord_check.py --parametrized # maxP + python experiments/lm/coord_check.py --parametrized --plot +""" + +from __future__ import annotations + +import argparse + +import torch + +from maxp import Parametrization, diagnose_axis, print_axis, plot_axis +from maxp.diagnose import op_label +from maxp.trace import ClassifiedOp + +from maxp_converter import install_pm_wrappers +from maxp_llama3 import _make_model_config + + +VOCAB_SIZE = 2048 +N_LAYERS = 2 + + +def _make_llama3(dim: int, n_heads: int, n_kv_heads: int, parametrized: bool): + cfg = _make_model_config( + dim=dim, n_layers=N_LAYERS, n_heads=n_heads, + n_kv_heads=n_kv_heads, vocab_size=VOCAB_SIZE, + ) + model = cfg.build() + if not parametrized: + return model, None + install_pm_wrappers(model) + sample = torch.randint(0, VOCAB_SIZE, (1, 16)) + param = Parametrization(model, sample_input=sample, lr_prefactor=1.0) + return model, param.param_groups + + +def _make_input(width: int) -> torch.Tensor: + return torch.randint(0, VOCAB_SIZE, (1, 8)) + + +def _make_train_step(model, param_groups): + """One AdamW step for LLaMA-3 (uses tok_embeddings, not tok_emb).""" + import torch.nn.functional as F + + if param_groups is not None: + opt = torch.optim.AdamW(param_groups) + else: + opt = torch.optim.AdamW(model.parameters(), lr=3e-4) + + x_fixed = torch.randint(0, VOCAB_SIZE, (4, 16)) + targets_fixed = torch.randint(0, VOCAB_SIZE, (4, 16)) + + def step(model, step_idx): + logits = model(x_fixed) + loss = F.cross_entropy(logits.reshape(-1, VOCAB_SIZE), targets_fixed.reshape(-1)) + loss.backward() + opt.step() + opt.zero_grad() + + return step + + +def _op_key(op: ClassifiedOp) -> str: + if op.param_name: + return op.param_name + return f"{op.module_path}:{op.source_loc}" + + +# --------------------------------------------------------------------------- +# Axis configs +# --------------------------------------------------------------------------- + +# d_model axis: head_dim fixed at 64, scale dim and n_heads together +D_MODEL_WIDTHS = [128, 256, 512, 768, 1024, 1536, 2048] + + +def _make_d_model(width: int, parametrized: bool): + n_heads = width // 64 + return _make_llama3(width, n_heads=n_heads, n_kv_heads=n_heads // 2, + parametrized=parametrized) + + +# head_dim axis: n_heads=8, n_kv_heads=4 fixed; dim = 8 * head_dim +HEAD_DIM_WIDTHS = [32, 64, 128, 256] + + +def _make_head_dim(width: int, parametrized: bool): + return _make_llama3(dim=8 * width, n_heads=8, n_kv_heads=4, + parametrized=parametrized) + + +AXES = { + "d_model": { + "make_model": _make_d_model, + "widths": D_MODEL_WIDTHS, + }, + "head_dim": { + "make_model": _make_head_dim, + "widths": HEAD_DIM_WIDTHS, + }, +} + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Coord check for parametrized LLaMA-3", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument("--parametrized", action="store_true", + help="Use maxP ABC parametrization") + parser.add_argument("--plot", action="store_true", + help="Save PNG plots") + parser.add_argument("--steps", type=int, default=10) + parser.add_argument("--seeds", type=int, default=5) + parser.add_argument("--plot-every", type=int, default=1) + args = parser.parse_args() + + variant = "parametrized" if args.parametrized else "plain" + print(f"LLaMA-3 coord check — {variant}") + + claimed: dict[str, str] = {} + shown_embeddings: set[str] = set() + + for axis_name, axis_cfg in AXES.items(): + print(f"\nDiagnosing axis: {axis_name} ...") + ops, affected, act_stats = diagnose_axis( + make_model_fn=lambda w, cfg=axis_cfg: cfg["make_model"](w, args.parametrized), + make_input_fn=_make_input, + widths=axis_cfg["widths"], + n_steps=args.steps, + n_seeds=args.seeds, + train_step_fn=_make_train_step, + ) + + deduped = [] + for i in affected: + op = ops[i] + key = _op_key(op) + + if op.layer_type == "embedding" and op.op != "embedding": + continue + elif op.layer_type == "embedding" and op.op == "embedding": + if key not in shown_embeddings: + shown_embeddings.add(key) + deduped.append(i) + else: + if key in claimed: + print(f" NOTE: {op_label(op)} already in '{claimed[key]}' axis, skipping") + else: + claimed[key] = axis_name + deduped.append(i) + + print_axis(axis_name, ops, deduped, act_stats, axis_cfg["widths"]) + + if args.plot: + fname = f"coord_check_{variant}_{axis_name}.png" + plot_axis(axis_name, ops, deduped, act_stats, + axis_cfg["widths"], fname, plot_every=args.plot_every) + print(f" Saved {fname}") diff --git a/experiments/lm/download_hf_assets.py b/experiments/lm/download_hf_assets.py new file mode 100644 index 0000000..9642c18 --- /dev/null +++ b/experiments/lm/download_hf_assets.py @@ -0,0 +1,259 @@ +# Downloaded from https://github.com/pytorch/torchtitan/blob/main/scripts/download_hf_assets.py +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from fnmatch import fnmatch + +from requests.exceptions import HTTPError +from tqdm import tqdm + + +def download_hf_assets( + repo_id: str, + local_dir: str, + asset_types: str | list[str], + download_all: bool = False, + hf_token: str | None = None, + additional_patterns: list | None = None, +) -> None: + """ + Download relevant files from HuggingFace Hub repository. + + This function recursively searches through the HuggingFace Hub repository + and downloads all related files + + Asset types: + - tokenizer: + - tokenizer.json - Modern HuggingFace tokenizers (complete definition) + - tokenizer_config.json - Tokenizer configuration and metadata + - tokenizer.model - SentencePiece model files (Llama, T5, etc.) + - vocab.txt - Plain text vocabulary files + - vocab.json - JSON vocabulary files + - merges.txt - BPE merge rules (GPT-2, RoBERTa style) + - special_tokens_map.json - Special token mappings + - safetensors + - *.safetensors - Modern Huggingface model weights format for fast loading + - model.safetensors.index.json - Contains mapping from hf fqn to file name + - index + - model.safetensors.index.json - Contains mapping from hf fqn to file name + - config + - config.json - Defines the model architecture + - generation_config.json - Defines the model architecture params needed for generation + + Args: + repo_id (str): HuggingFace repository ID (e.g., meta-llama/Llama-3.1-8B") + local_dir (str): Local directory to save tokenizer files. A subdirectory + named after the model will be created automatically. + asset_types (list[str]): List of the asset types to download + hf_token (Optional[str]): HuggingFace API token for accessing private repositories. + Required for gated models like Llama. + additional_patterns (Optional[list]): Additional file patterns to search for and download + from the HuggingFace Hub repository. + download_all (bool): If True, download all files from the repository + """ + import os + + from huggingface_hub import hf_hub_download, list_repo_files + + # Extract model name from repo_id (part after "/") + if "/" not in repo_id: + raise ValueError( + f"Invalid repo_id format: '{repo_id}'. Expected format: 'organization/model-name'" + ) + model_name = repo_id.split("/")[-1].strip() + model_dir = os.path.join(local_dir, model_name) + + ASSET_PATTERNS = { + "tokenizer": [ + "tokenizer.json", + "tokenizer_config.json", + "tokenizer.model", + "vocab.txt", + "vocab.json", + "merges.txt", + "special_tokens_map.json", + ], + "safetensors": ["*.safetensors", "*model.safetensors.index.json"], + "index": ["*model.safetensors.index.json"], + "config": ["config.json", "generation_config.json"], + } + + if isinstance(asset_types, str): + asset_types = [asset_types] + + if download_all: + print("Downloading all files from repository...") + files_found = list_repo_files(repo_id=repo_id, token=hf_token) + else: + total_patterns = [] + for asset_type in asset_types: + if asset_type in ASSET_PATTERNS: + total_patterns.extend(ASSET_PATTERNS[asset_type]) + else: + raise ValueError( + "Unknown asset type {}. Available uses: --assets {} \n".format( + asset_type, " ".join(ASSET_PATTERNS.keys()) + ), + "Or specify exact patterns to download. Example: --additional_patterns '*.safetensors' README.md '*.json' \n", + "Or use --all to download all files", + ) + + # Add additional patterns if provided + if additional_patterns: + total_patterns.extend(additional_patterns) + asset_types.append("additional_patterns") + ASSET_PATTERNS["additional_patterns"] = additional_patterns + + def should_download(patterns: list[str], filename: str) -> bool: + """Check if a file matches a pattern to be downloaded.""" + basename = os.path.basename(filename) + for pattern in patterns: + pattern_lower = pattern.lower() + + # Exact name match + if basename == pattern_lower: + return True + # Do wildcard match if wildcards are in pattern + if "*" in pattern_lower or "?" in pattern_lower: + if fnmatch(basename, pattern_lower): + return True + return False + + try: + # Get list of available files in the repo + print(f"Scanning repository {repo_id} for files...") + available_files = list_repo_files(repo_id=repo_id, token=hf_token) + + # Filter for requested asset files + files_found = [ + f for f in available_files if should_download(total_patterns, f) + ] + + # Check each asset type individually to see if files were not found + for asset_type in asset_types: + if asset_type in ASSET_PATTERNS: + asset_patterns = ASSET_PATTERNS[asset_type] + matches_found = False + for f in available_files: + if should_download(asset_patterns, f): + matches_found = True + break + + if not matches_found: + print( + f"Warning: No matching files found for asset_type '{asset_type}' in {repo_id}" + ) + + if not files_found: + print(f"Warning: No matching files found in {repo_id}") + print(f"Available files: {available_files[:10]}...") + return + + except HTTPError as e: + if e.response and e.response.status_code == 401: + print( + "You need to pass a valid `--hf_token=...` to download private checkpoints." + ) + raise e + + print(f"Found {len(files_found)} files:") + for f in files_found: + print(f" - {f}") + + downloaded_files = [] + missed_files = [] + + # Download files with progress bar + # pyrefly: ignore [bad-context-manager] + with tqdm(total=len(files_found), desc="Downloading files", unit="file") as pbar: + for filename in files_found: + try: + pbar.set_description(f"Downloading {os.path.basename(filename)}") + + hf_hub_download( + repo_id=repo_id, + filename=filename, + local_dir=model_dir, + token=hf_token, + ) + downloaded_files.append(filename) + pbar.update(1) + + except HTTPError as e: + if e.response and e.response.status_code == 404: + print(f"File {filename} not found, skipping...") + missed_files.append(filename) + pbar.update(1) + continue + else: + raise e + + if downloaded_files: + print( + f"\nSuccessfully downloaded {len(downloaded_files)} files to: {model_dir}" + ) + if missed_files: + print(f"Warning: Some files could not be downloaded: \n{missed_files}") + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser( + description="Download files from HuggingFace Hub. " + "Automatically detects and downloads files that match the specified file-types to download. " + ) + parser.add_argument( + "--repo_id", + type=str, + required=True, + help="Repository ID to download from (e.g., 'meta-llama/Llama-3.1-8B', 'deepseek-ai/DeepSeek-V3')", + ) + parser.add_argument( + "--hf_token", + type=str, + default=None, + help="HuggingFace API token (required for private repos)", + ) + parser.add_argument( + "--local_dir", + type=str, + default="assets/hf/", + help="Local directory to save hf asset files (default: assets/hf/)", + ) + parser.add_argument( + "--assets", + type=str, + nargs="+", + default=[], + help="Asset types to download: tokenizer, safetensors, index, config", + ) + parser.add_argument( + "--additional_patterns", + type=str, + nargs="+", + default=[], + help="Additional file patterns to search for and download from the HuggingFace Hub repository", + ) + + parser.add_argument( + "--all", action="store_true", default=False, help="Download all files in repo" + ) + + args = parser.parse_args() + if not args.all and not args.assets and not args.additional_patterns: + parser.error( + "At least one of --all, --assets or --additional_patterns must be specified." + ) + + download_hf_assets( + args.repo_id, + args.local_dir, + args.assets, + args.all, + args.hf_token, + args.additional_patterns, + ) diff --git a/experiments/lm/fineweb.py b/experiments/lm/fineweb.py new file mode 100644 index 0000000..216512d --- /dev/null +++ b/experiments/lm/fineweb.py @@ -0,0 +1,53 @@ +"""Register FineWeb-Edu in torchtitan's dataset registry. + +Import this module before building any Trainer.Config that uses +--dataset fineweb-edu or --dataset fineweb-edu-10bt or --dataset fineweb-edu-100bt. + +FineWeb-Edu is loaded from the local HF cache (HF_HOME must contain +the downloaded dataset). No pre-tokenization step is needed. +torchtitan tokenizes on-the-fly with the LLaMA-3 tokenizer. +""" + +from __future__ import annotations + +from datasets import load_dataset + +from torchtitan.hf_datasets import DatasetConfig +from torchtitan.hf_datasets.text_datasets import DATASETS + + +def _process_fineweb_text(sample: dict) -> str: + return sample["text"] + + +def _load_fineweb_edu(dataset_path: str): + return load_dataset(dataset_path, split="train", streaming=True) + + +def _load_fineweb_edu_10bt(dataset_path: str): + # 10B-token deduplicated subset for S1/S2/S3 + return load_dataset(dataset_path, name="sample-10BT", split="train", streaming=True) + + +def _load_fineweb_edu_100bt(dataset_path: str): + # 100B-token deduplicated subset for S4/S5 + return load_dataset(dataset_path, name="sample-100BT", split="train", streaming=True) + + +DATASETS["fineweb-edu"] = DatasetConfig( + path="HuggingFaceFW/fineweb-edu", + loader=_load_fineweb_edu, + sample_processor=_process_fineweb_text, +) + +DATASETS["fineweb-edu-10bt"] = DatasetConfig( + path="HuggingFaceFW/fineweb-edu", + loader=_load_fineweb_edu_10bt, + sample_processor=_process_fineweb_text, +) + +DATASETS["fineweb-edu-100bt"] = DatasetConfig( + path="HuggingFaceFW/fineweb-edu", + loader=_load_fineweb_edu_100bt, + sample_processor=_process_fineweb_text, +) diff --git a/experiments/lm/launch_sweep.py b/experiments/lm/launch_sweep.py new file mode 100644 index 0000000..9f92424 --- /dev/null +++ b/experiments/lm/launch_sweep.py @@ -0,0 +1,226 @@ +#!/usr/bin/env python3 +"""Generate and submit SLURM jobs for the full LM sweep. + +Usage: + + 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 \\ + --seeds 1 2 \\ + --runs-dir /net/storage/pr3/plgrid/plggadlers/maxP/runs \\ + --dataset HuggingFaceFW/fineweb-edu \\ + --hf-assets-path /net/storage/pr3/plgrid/plggadlers/maxP/experiments/lm/assets/hf/Llama-3.1-8B \\ + --venv-path /net/storage/pr3/plgrid/plggadlers/maxP/.venv \\ + --repo-path /net/storage/pr3/plgrid/plggadlers/maxP \\ + [--dry-run] [--resume] + +By default skips runs where any checkpoint already exists (started or completed). +With --resume, submits all runs regardless (to resume interrupted training). +""" + +from __future__ import annotations + +import argparse +import subprocess +from datetime import date +from pathlib import Path +from string import Template + + +# --------------------------------------------------------------------------- +# Per-scale defaults +# --------------------------------------------------------------------------- + +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}, +} + +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"], +} + +SCALE_SEEDS = { + "s1": [1, 2, 3], + "s2": [1, 2], + "s3": [1, 2], + "s4": [1], + "s5": [1], +} + +ALL_LRS = [3e-4, 1e-3, 3e-3, 1e-2, 3e-2, 1e-1, 3e-1] + +S5_SINGLE_LR = [1e-2] + + +# --------------------------------------------------------------------------- +# SLURM job script template +# --------------------------------------------------------------------------- + +SLURM_TEMPLATE = Template("""\ +#!/bin/bash -l +#SBATCH --job-name maxp_${scale}_${method_tag}_lr${lr_tag}_s${seed} +#SBATCH --nodes 1 +#SBATCH --cpus-per-gpu 72 +#SBATCH --mem-per-gpu 118GB +#SBATCH --time ${wall_time} +#SBATCH --account plgadlers-gpu-gh200 +#SBATCH --partition plgrid-gpu-gh200 +#SBATCH --gres gpu:${gpus_per_node} +#SBATCH --output ${output_dir}/maxp_${scale}_${method_tag}_lr${lr_tag}_s${seed}.out +#SBATCH --error ${output_dir}/maxp_${scale}_${method_tag}_lr${lr_tag}_s${seed}.err + +module add ML-bundle/25.10 +source "${venv_path}/bin/activate" +cd "${repo_path}" + +export OMP_NUM_THREADS=16 +export HF_HOME="$${HF_HOME:-/net/scratch/hscra/plgrid/plgmwojnar/hf}" +export WANDB_PROJECT="$${WANDB_PROJECT:-maxP-lm}" +export WANDB_RUN_NAME="maxp_${scale}_${method_tag}_lr${lr_tag}_s${seed}" + +# Distributed setup +find_free_port() { + python -c 'import socket; s=socket.socket(); s.bind(("", 0)); print(s.getsockname()[1]); s.close()' +} + +export MASTER_ADDR=127.0.0.1 +export MASTER_PORT=$$(find_free_port) +export NCCL_SOCKET_IFNAME=lo + +echo "Master Port: $${MASTER_PORT}" + +torchrun --nproc_per_node=${gpus_per_node} --master_addr=$${MASTER_ADDR} --master_port=$${MASTER_PORT} experiments/lm/train.py \\ + --scale ${scale} \\ + --method ${method} \\ + --lr ${lr} \\ + --seed ${seed} \\ + --batch-size ${batch_size} \\ + --dataset ${dataset} \\ + ${dataset_path_arg} \\ + --num-workers ${num_workers} \\ + --prefetch-factor ${prefetch_factor} \\ + --hf-assets-path ${hf_assets_path} \\ + --output-dir ${output_dir} +""") + + +def _lr_tag(lr: float) -> str: + return f"{lr:.0e}".replace("+", "") + + +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 main() -> None: + p = argparse.ArgumentParser(description="Launch maxP LLaMA-3 SLURM sweep") + p.add_argument("--scale", required=True, choices=list(SCALE_CONFIGS)) + p.add_argument("--methods", nargs="+", default=None, + help="Override method list (default: per-scale defaults)") + p.add_argument("--lrs", type=float, nargs="+", default=None, + help="Override LR list (default: all 7 values)") + p.add_argument("--seeds", type=int, nargs="+", default=None, + help="Override seed list (default: per-scale defaults)") + p.add_argument("--batch-size", type=int, default=16, + help="Global batch size (divided by WORLD_SIZE to get per-GPU)") + p.add_argument("--gpus-per-node", type=int, default=None, + help="GPUs per SLURM node (default: per-scale value from SCALE_CONFIGS)") + p.add_argument("--runs-dir", required=True, help="Root directory for run outputs") + p.add_argument("--dataset", default="fineweb-edu", + help="HuggingFace dataset name (default: per-scale defaults)") + p.add_argument("--dataset-path", default=None, + help="Optional local path to dataset assets (otherwise streamed from HF)") + p.add_argument("--num-workers", type=int, default=8, + help="DataLoader num_workers for prefetching") + p.add_argument("--prefetch-factor", type=int, default=4, + help="Batches prefetched per DataLoader worker") + p.add_argument("--hf-assets-path", required=True, + 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("--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)") + 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) + seeds = args.seeds or SCALE_SEEDS[scale] + 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") + submitted = skipped = 0 + + for method in methods: + for lr in lrs: + for seed in seeds: + run_name = ( + f"{today}_{scale}_{_method_tag(method)}_lr{_lr_tag(lr)}_s{seed}" + ) + out_dir = Path(args.runs_dir) / run_name + + if not args.resume and _has_checkpoint(out_dir): + print(f"[skip] {run_name}") + skipped += 1 + continue + + out_dir.mkdir(parents=True, exist_ok=True) + + dataset_path_arg = f"--dataset-path {args.dataset_path}" if args.dataset_path else "" + + script_content = SLURM_TEMPLATE.substitute( + scale=scale, + method_tag=_method_tag(method), + lr_tag=_lr_tag(lr), + seed=seed, + wall_time=sc["wall"], + output_dir=str(out_dir), + venv_path=args.venv_path, + repo_path=args.repo_path, + method=method, + lr=lr, + batch_size=args.batch_size, + gpus_per_node=gpus_per_node, + dataset=args.dataset, + dataset_path_arg=dataset_path_arg, + num_workers=args.num_workers, + prefetch_factor=args.prefetch_factor, + hf_assets_path=args.hf_assets_path, + ) + script_path = out_dir / "job.sh" + script_path.write_text(script_content) + + 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()}") + submitted += 1 + + action = "would submit" if args.dry_run else "submitted" + print(f"\nDone: {action} {submitted} jobs, skipped {skipped} completed runs.") + + +if __name__ == "__main__": + main() diff --git a/experiments/lm/maxp_converter.py b/experiments/lm/maxp_converter.py new file mode 100644 index 0000000..86d44b9 --- /dev/null +++ b/experiments/lm/maxp_converter.py @@ -0,0 +1,173 @@ +"""MaxP model converter and post-optimizer-build hook for torchtitan integration.""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass + +import torch +import torch.nn as nn + +from torchtitan.config import Configurable +from torchtitan.distributed import ParallelDims +from torchtitan.protocols.model_converter import ModelConverter +from torchtitan.protocols.module import Module + +from maxp import ParametrizedModule, Parametrization + + +class LlamaParametrizedModule(Module, ParametrizedModule): + """Marker class for ParametrizedModules in LLaMA-3 models.""" + pass + + +def _wrap(parent: nn.Module, attr: str, width_dim: int, layer_type: str, **pm_kw) -> None: + """Replace parent. with a LlamaParametrizedModule in-place.""" + layer = getattr(parent, attr) + setattr(parent, attr, LlamaParametrizedModule(layer, width_dim=width_dim, layer_type=layer_type, **pm_kw)) + + +class _SDPAWrapper(Module): + """Route q, k through a readout PM so the graph matches + attn_score topology: r = min(min(r_q, r_k) + a, r_v). + scale_output is set to False as PM returns a tuple that shouldn't be + scalar-multiplied; the physical pm.scale = head_dim^-a is instead + injected into SDPA's `scale` kwarg (applied to logits before softmax). + """ + def __init__(self, inner, head_dim): + super().__init__() + self.inner = inner + self.score = LlamaParametrizedModule( + lambda q, k: (q, k), width_dim=head_dim, + layer_type="readout", scale_output=False, + ) + + def forward(self, q, k, v, **kw): + q, k = self.score(q, k) + kw.pop("scale", None) + return self.inner(q, k, v, scale=self.score.scale, **kw) + + +def install_pm_wrappers(model: nn.Module) -> None: + """Install LlamaParametrizedModule wrappers on all LLaMA-3 layers in-place. + + Purely structural — safe to call on a meta-device model. + Wraps attention projections, FFN weights, token embedding, and LM head. + The SDPA inner_attention is annotated as a=0.0 readout to avoid + double-scaling (SDPA already divides by sqrt(head_dim)). + """ + from torchtitan.models.llama3.model import Llama3Model + from torchtitan.models.common.attention import FusedQKVLinear + + assert isinstance(model, Llama3Model), f"Expected Llama3Model, got {type(model)}" + cfg = model.config + d_model: int = cfg.dim + head_dim: int = next(iter(model.layers.values())).attention.head_dim + + for block in model.layers.values(): + attn = block.attention + ffn = block.feed_forward + + qkv = attn.qkv_linear + if isinstance(qkv, FusedQKVLinear): + _wrap(qkv, "wqkv", d_model, "hidden") + else: + _wrap(qkv, "wq", d_model, "hidden") + _wrap(qkv, "wk", d_model, "hidden") + _wrap(qkv, "wv", d_model, "hidden") + + _wrap(attn, "wo", d_model, "hidden") + + attn.inner_attention = _SDPAWrapper(attn.inner_attention, head_dim) + + d_ff: int = ffn.w1.out_features + _wrap(ffn, "w1", d_model, "hidden") + _wrap(ffn, "w3", d_model, "hidden") + _wrap(ffn, "w2", d_ff, "hidden") + + _wrap(model, "tok_embeddings", d_model, "embedding") + _wrap(model, "output", d_model, "readout") + + +class MaxPConverter(Configurable, ModelConverter): + """Installs PM wrappers and runs maxP parametrization before model compilation. + + convert() wraps layers and calls Parametrization (LP solve + weight reinit) + so forward hooks are correctly traced by torch.compile. + post_optimizer_build_fn() (returned by make_post_optimizer_build_fn) patches + optimizer param_groups with per-layer LRs after optimizer.build(). + """ + + @dataclass(kw_only=True, slots=True) + class Config(Configurable.Config): + method: str = "maxP" # "maxP" | "mup-full" | "mup-no" + lr_prefactor: float = 1e-3 + alignment_warmup: int = 10 # maxP only + solve_interval: int = 100 # maxP only + sample_size: int = 32 # maxP only + c_ema: float = 0.0 # maxP only + use_training_activations: bool = False # maxP only + + def __init__( + self, + config: Config, + *, + parallel_dims: ParallelDims, + model_compile_enabled: bool, + ) -> None: + self._cfg = config + + def convert(self, model: nn.Module) -> None: + install_pm_wrappers(model) + cfg = self._cfg + param = Parametrization( + model, + sample_input=_make_sample_input(model), + alignment="full" if "full" in cfg.method else "no", + lr_prefactor=cfg.lr_prefactor, + warmup_steps=cfg.alignment_warmup, + solve_interval=cfg.solve_interval, + sample_size=cfg.sample_size, + c_ema=cfg.c_ema, + use_training_activations=cfg.use_training_activations, + ) + model._maxp_param = param + if cfg.method == "maxP": + model._is_dynamic = True + model._maxp_align = { + name: (pm.align_z0_dW, pm.align_dZ_w0, pm.align_dZ_dW) + for name, pm in param._pms + if pm.weight is not None and pm.align_z0_dW is not None + } + + def post_optimizer_hook(self, model: nn.Module | list[nn.Module]) -> None: + pass + + +def _make_sample_input(model: nn.Module) -> torch.Tensor | None: + """Build a deterministic sample input for DAG tracing. + + Returns None if vocab_size cannot be determined (falls back to chain graph). + Uses a fixed seed so all ranks produce identical inputs (no collective needed). + """ + cfg = getattr(model, "config", None) + vocab_size = getattr(cfg, "vocab_size", None) + if vocab_size is None: + return None + try: + device = next(iter(model.parameters())).device + except StopIteration: + device = torch.device("cpu") + gen = torch.Generator().manual_seed(0) + return torch.randint(0, vocab_size, (1, 16), generator=gen).to(device) + + +def post_optimizer_build_fn(optimizers, model_parts: list[nn.Module], parallel_dims: ParallelDims) -> None: + """Called by Trainer after init_weights() and optimizer.build(). + + Re-runs the tensor-bound bits of Parametrization (weight init + per-layer + param_groups) now that the model has been materialized off meta device, + and splices the resulting groups into the optimizer. + """ + for optimizer, model in zip(optimizers, model_parts): + model._maxp_param.refresh(optimizer=optimizer) diff --git a/experiments/lm/maxp_llama3.py b/experiments/lm/maxp_llama3.py new file mode 100644 index 0000000..e53e5d9 --- /dev/null +++ b/experiments/lm/maxp_llama3.py @@ -0,0 +1,200 @@ +"""MaxP-compatible LLaMA-3 model definition and scale configs.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from functools import partial + +import torch.nn as nn + +from torchtitan.components.loss import build_cross_entropy_loss +from torchtitan.distributed.pipeline_parallel import pipeline_llm +from torchtitan.models.common import Embedding, Linear, RMSNorm, RoPE, compute_ffn_hidden_dim +from torchtitan.models.common.config_utils import get_attention_config, make_ffn_config, make_gqa_config +from torchtitan.models.common.param_init import depth_scaled_std +from torchtitan.models.llama3 import parallelize_llama as _parallelize_llama +from torchtitan.models.llama3.model import Llama3Model, Llama3TransformerBlock +from torchtitan.models.llama3.state_dict_adapter import Llama3StateDictAdapter +from torchtitan.protocols.model_spec import ModelSpec + +from maxp_converter import post_optimizer_build_fn + + +# Weight inits — mirrors llama3 upstream defaults +_LINEAR_INIT = {"weight": partial(nn.init.trunc_normal_, std=0.02), "bias": nn.init.zeros_} +_NORM_INIT = {"weight": nn.init.ones_} +_EMBEDDING_INIT = {"weight": partial(nn.init.normal_, std=1.0)} + + +def _output_linear_init(dim: int) -> dict: + s = dim**-0.5 + return { + "weight": partial(nn.init.trunc_normal_, std=s, a=-3 * s, b=3 * s), + "bias": nn.init.zeros_, + } + + +def _depth_init(layer_id: int) -> dict: + return { + "weight": partial(nn.init.trunc_normal_, std=depth_scaled_std(0.02, layer_id)), + "bias": nn.init.zeros_, + } + + +def _build_layers( + *, + n_layers: int, + dim: int, + n_heads: int, + n_kv_heads: int, + hidden_dim: int, + attn_backend: str = "sdpa", +) -> list: + inner_attention, mask_type = get_attention_config(attn_backend) + return [ + Llama3TransformerBlock.Config( + attention_norm=RMSNorm.Config(normalized_shape=dim, param_init=_NORM_INIT), + ffn_norm=RMSNorm.Config(normalized_shape=dim, param_init=_NORM_INIT), + attention=make_gqa_config( + dim=dim, + n_heads=n_heads, + n_kv_heads=n_kv_heads, + wqkv_param_init=_LINEAR_INIT, + wo_param_init=_depth_init(layer_id), + inner_attention=inner_attention, + mask_type=mask_type, + rope_backend="complex", + ), + feed_forward=make_ffn_config( + dim=dim, + hidden_dim=hidden_dim, + w1_param_init=_LINEAR_INIT, + w2w3_param_init=_depth_init(layer_id), + ), + ) + for layer_id in range(n_layers) + ] + + +def _make_model_config( + *, + dim: int, + n_layers: int, + n_heads: int, + n_kv_heads: int, + vocab_size: int = 128256, + attn_backend: str = "sdpa", +) -> Llama3Model.Config: + hidden_dim = compute_ffn_hidden_dim(dim, multiple_of=256, ffn_dim_multiplier=1.3) + return Llama3Model.Config( + dim=dim, + vocab_size=vocab_size, + enable_weight_tying=False, + tok_embeddings=Embedding.Config( + num_embeddings=vocab_size, + embedding_dim=dim, + param_init=_EMBEDDING_INIT, + ), + norm=RMSNorm.Config(normalized_shape=dim, param_init=_NORM_INIT), + output=Linear.Config( + in_features=dim, + out_features=vocab_size, + param_init=_output_linear_init(dim), + ), + rope=RoPE.Config( + dim=dim // n_heads, + max_seq_len=131072, + theta=500000, + backend="complex", + scaling="llama", + ), + layers=_build_layers( + n_layers=n_layers, + dim=dim, + n_heads=n_heads, + n_kv_heads=n_kv_heads, + hidden_dim=hidden_dim, + attn_backend=attn_backend, + ), + ) + + +def parallelize_llama(model, **kwargs): + # Force reshard_after_forward="always" so no_grad forwards used for + # alignment measurement don't leave norm+output layers unsharded. + parallelism = kwargs.get("parallelism") + if parallelism is not None: + parallelism.fsdp_reshard_after_forward = "always" + 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_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), +} + + +def maxp_model_registry(scale: str, method: str, attn_backend: str = "sdpa") -> ModelSpec: + """Build a ModelSpec for maxP LLaMA-3 training. + + Parametrization config (alignment_warmup, solve_interval, etc.) lives in + MaxPConverter.Config and is passed via Trainer.Config.model_converters. + + Args: + scale: One of "debug", "s1" … "s5". + method: "maxP" (dynamic), "mup-full" (static, full align), or + "mup-no" (static, no align). + attn_backend: Attention backend ("sdpa", "flex", "varlen"). + """ + if scale not in SCALE_CONFIGS: + raise ValueError(f"Unknown scale '{scale}'. Choose from {list(SCALE_CONFIGS)}") + + kwargs = SCALE_CONFIGS[scale].copy() + model_config = _make_model_config(attn_backend=attn_backend, **kwargs) + + return ModelSpec( + name="maxp_llama3", + flavor=f"{scale}_{method}", + model=model_config, + parallelize_fn=parallelize_llama, + pipelining_fn=pipeline_llm, + build_loss_fn=build_cross_entropy_loss, + post_optimizer_build_fn=post_optimizer_build_fn, + state_dict_adapter=Llama3StateDictAdapter, + ) + + +def compute_steps( + scale: str, + seq_len: int, + local_batch_size: int, + world_size: int, + token_multiplier: int = 20, +) -> int: + """Return training steps for token_multiplier × non-embed-params tokens. + + Instantiates the model on meta device (zero memory) to count parameters. + Non-embed params = total params minus tok_embeddings and LM head. + """ + import torch + + cfg = SCALE_CONFIGS[scale] + vocab_size = cfg.get("vocab_size", 128256) + dim = cfg["dim"] + arch_kwargs = {k: v for k, v in cfg.items() if k != "vocab_size"} + model_config = _make_model_config(vocab_size=vocab_size, **arch_kwargs) + with torch.device("meta"): + model = Llama3Model(model_config) + total = sum(p.numel() for p in model.parameters()) + non_embed = total - 2 * vocab_size * dim # tok_embeddings + LM head + tokens = token_multiplier * non_embed + return math.ceil(tokens / (seq_len * local_batch_size * world_size)) diff --git a/experiments/lm/run.sh b/experiments/lm/run.sh new file mode 100644 index 0000000..2b525a8 --- /dev/null +++ b/experiments/lm/run.sh @@ -0,0 +1,44 @@ +#!/usr/bin/env bash +# Submit a full sweep for one scale. +# +# Usage: +# bash experiments/lm/run.sh s3 [launch_sweep.py flags …] +# +# 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 +# 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}" +REPO_PATH="${REPO_PATH:-/net/storage/pr3/plgrid/plggadlers/maxP}" + +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" \ + --repo-path "$REPO_PATH" \ + "${@}" diff --git a/experiments/lm/run_debug.sh b/experiments/lm/run_debug.sh new file mode 100755 index 0000000..0846adb --- /dev/null +++ b/experiments/lm/run_debug.sh @@ -0,0 +1,85 @@ +#!/usr/bin/env bash +# Quick debug run on a local GPU machine (no SLURM). +# Uses the tiny debug model (dim=256, 2 layers, vocab=2048) and +# the bundled c4_test dataset — no downloads required. +# +# Usage: +# bash experiments/lm/run_debug.sh [--steps N] [--method METHOD] [--gpus N] +set -euo pipefail + +REPO="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +VENV="${REPO}/.venv" + +# Defaults — override with env vars or flags +STEPS="${STEPS:-200}" +SCALE="${SCALE:-debug}" +METHOD="${METHOD:-maxP}" +GPUS="${GPUS:-1}" +DATASET="${DATASET:-c4_test}" +OUTPUT_DIR="${OUTPUT_DIR:-/tmp/maxp_debug}" +TOKENIZER="${TOKENIZER:-${REPO}/experiments/lm/assets/hf/Llama-3.1-8B}" +C4_TEST="${C4_TEST:-${REPO}/experiments/lm/assets/c4_test}" + +# Parse flags +while [[ $# -gt 0 ]]; do + case "$1" in + --steps) STEPS="$2"; shift 2 ;; + --scale) SCALE="$2"; shift 2 ;; + --method) METHOD="$2"; shift 2 ;; + --gpus) GPUS="$2"; shift 2 ;; + --dataset) DATASET="$2"; shift 2 ;; + --output-dir) OUTPUT_DIR="$2"; shift 2 ;; + --tokenizer) TOKENIZER="$2"; shift 2 ;; + --c4-test) C4_TEST="$2"; shift 2 ;; + *) echo "Unknown flag: $1"; exit 1 ;; + esac +done + +source "${VENV}/bin/activate" +cd "${REPO}" + +export OMP_NUM_THREADS=4 + +# Distributed setup +find_free_port() { + python -c 'import socket; s=socket.socket(); s.bind(("", 0)); print(s.getsockname()[1]); s.close()' +} + +export MASTER_ADDR=127.0.0.1 +export MASTER_PORT=$(find_free_port) +export NCCL_SOCKET_IFNAME=lo + +echo "=== maxP debug run ===" +echo " scale: ${SCALE}" +echo " method: ${METHOD}" +echo " dataset: ${DATASET}" +echo " steps: ${STEPS}" +echo " gpus: ${GPUS}" +echo " output_dir: ${OUTPUT_DIR}" +echo " master_port:${MASTER_PORT}" +echo "" + +mkdir -p "${OUTPUT_DIR}" + +TRAIN_ARGS=( + --scale "${SCALE}" + --method "${METHOD}" + --lr 0.03 + --alignment-warmup 20 + --solve-interval 5 + --batch-size 4 + --seq-len 256 + --steps "${STEPS}" + --dataset "${DATASET}" + ${C4_TEST:+--dataset-path "${C4_TEST}"} + ${TOKENIZER:+--hf-assets-path "${TOKENIZER}"} + --output-dir "${OUTPUT_DIR}" +) + +if [[ "${GPUS}" -eq 1 ]]; then + # Single GPU: run python directly so logs stream to terminal + LOCAL_RANK=0 RANK=0 WORLD_SIZE=1 python experiments/lm/train.py "${TRAIN_ARGS[@]}" +else + # Multi-GPU: use torchrun + torchrun --nproc_per_node="${GPUS}" --master_addr="${MASTER_ADDR}" --master_port="${MASTER_PORT}" experiments/lm/train.py "${TRAIN_ARGS[@]}" +fi diff --git a/experiments/lm/train.py b/experiments/lm/train.py new file mode 100644 index 0000000..5a0e060 --- /dev/null +++ b/experiments/lm/train.py @@ -0,0 +1,248 @@ +"""MaxP LLaMA-3 pre-training entry point. + +Uses torchtitan's Trainer with MaxPConverter and per-layer LRs from Parametrization. + +Usage (single GPU, debug smoke test): + + python experiments/lm/train.py --scale debug --steps 20 + +Usage (real training, 8 GPUs): + + torchrun --nproc_per_node=8 experiments/lm/train.py \\ + --scale s3 --method maxP --lr 1e-3 --steps 10000 \\ + --dataset HuggingFaceFW/fineweb-edu \\ + --output-dir ./outputs/s3_maxP_lr1e-3_s1 +""" + +from __future__ import annotations + +import argparse +import os + +import torch +torch._dynamo.config.recompile_limit = 100 + +from torchtitan.components.checkpoint import CheckpointManager +from torchtitan.components.lr_scheduler import LRSchedulersContainer +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.hf_datasets.text_datasets import HuggingFaceTextDataLoader +from torchtitan.protocols.model_converter import ModelConvertersContainer +from torchtitan.tools.logging import init_logger, logger +from torchtitan.trainer import Trainer + +from maxp_converter import MaxPConverter +from maxp_llama3 import maxp_model_registry, compute_steps, SCALE_CONFIGS +import fineweb # noqa: F401 — registers fineweb-edu in torchtitan's DATASETS dict + + +class MaxPTrainer(Trainer): + """Trainer that calls param.step() for dynamic maxP and logs per-layer metrics.""" + + def train_step(self, data_iterator): + # Pre-fetch a batch to capture real tokens for alignment measurement + # before the optimizer step modifies weights. + needs_capture = [m for m in self.model_parts + if getattr(m, "_is_dynamic", False) + and not getattr(m, "_maxp_ready", False)] + + if needs_capture: + batch = next(data_iterator) + tokens = batch[0]["input"].detach() + for model in needs_capture: + model._maxp_sample_x = tokens + model._maxp_param.capture_initial(tokens) + model._maxp_ready = True + + super().train_step(data_iterator) + + # Dynamic maxP: re-solve LP and sync LRs after each optimizer+scheduler step. + # lr_prefactor is read from the "_other" group whose lr = lr_prefactor * wsd_factor, + # so _sync_lrs computes per-layer lr = lr_prefactor * wsd_factor * n^(-c). + for model, opt in zip(self.model_parts, self.optimizers.optimizers): + if not getattr(model, "_is_dynamic", False) or not getattr(model, "_maxp_ready", False): + continue + param = model._maxp_param + for g in opt.param_groups: + if g.get("layer_name") == "_other": + param.lr_prefactor = g["lr"] + break + param.step(model._maxp_sample_x, opt) + model._maxp_align = { + name: (pm.align_z0_dW, pm.align_dZ_w0, pm.align_dZ_dW) + for name, pm in param._pms + if pm.weight is not None and pm.align_z0_dW is not None + } + + if not self.metrics_processor.should_log(self.step): + return + + extra: dict = {} + for model in self.model_parts: + for name, (z0_dw, dz_w0, dz_dw) in getattr(model, "_maxp_align", {}).items(): + extra[f"align/z0_dW/{name}"] = z0_dw + extra[f"align/dZ_w0/{name}"] = dz_w0 + extra[f"align/dZ_dW/{name}"] = dz_dw + for opt in self.optimizers.optimizers: + for g in opt.param_groups: + ln = g.get("layer_name") + if ln: + extra[f"lr/{ln}"] = g["lr"] + if extra: + self.metrics_processor.logger.log(extra, self.step) + + +def _local_batch_size(global_batch_size: int) -> int: + world_size = int(os.environ.get("WORLD_SIZE", "1")) + assert global_batch_size % world_size == 0, ( + f"--batch-size {global_batch_size} not divisible by WORLD_SIZE {world_size}" + ) + return global_batch_size // world_size + + +def _resolve_steps(args: argparse.Namespace) -> int: + if args.steps is not None: + return args.steps + if args.scale == "debug": + return 20 + world_size = int(os.environ.get("WORLD_SIZE", "1")) + return compute_steps(args.scale, args.seq_len, _local_batch_size(args.batch_size), world_size) + + +def build_trainer_config(args: argparse.Namespace) -> Trainer.Config: + is_debug = args.scale == "debug" + steps = _resolve_steps(args) + local_batch_size = _local_batch_size(args.batch_size) + + model_spec = maxp_model_registry( + scale=args.scale, + method=args.method, + attn_backend="sdpa", + ) + + return Trainer.Config( + model_spec=model_spec, + hf_assets_path=args.hf_assets_path, + dump_folder=args.output_dir, + model_converters=ModelConvertersContainer.Config( + converters=[MaxPConverter.Config( + method=args.method, + lr_prefactor=args.lr, + alignment_warmup=args.alignment_warmup, + solve_interval=args.solve_interval, + sample_size=args.sample_size, + c_ema=args.c_ema, + )], + ), + optimizer=OptimizersContainer.Config( + lr=args.lr, + # fused requires CUDA; foreach works on CPU and GPU + implementation="foreach" if is_debug else "fused", + ), + lr_scheduler=LRSchedulersContainer.Config( + warmup_steps=10 if is_debug else int(0.1 * steps), + decay_ratio=0.1, + ), + training=TrainingConfig( + local_batch_size=local_batch_size, + seq_len=args.seq_len, + steps=steps, + dtype="bfloat16", + ), + dataloader=HuggingFaceTextDataLoader.Config( + dataset=args.dataset, + dataset_path=args.dataset_path, + num_workers=args.num_workers, + prefetch_factor=args.prefetch_factor, + ), + metrics=MetricsProcessor.Config( + log_freq=10 if is_debug else 50, + enable_tensorboard=not is_debug, + enable_wandb=not is_debug, + ), + checkpoint=CheckpointManager.Config( + enable=args.checkpoint_interval is not None, + interval=args.checkpoint_interval, + last_save_model_only=False, + keep_latest_k=2, + async_mode="async", + ), + compile=CompileConfig(enable=not is_debug), + activation_checkpoint=ActivationCheckpointConfig(mode="full"), + validator=Validator.Config( + enable=not is_debug, + freq=500, + steps=50, + ), + ) + + +def parse_args() -> argparse.Namespace: + default_hf_path = os.path.join(os.path.dirname(__file__), "assets/hf/Llama-3.1-8B") + p = argparse.ArgumentParser( + description="MaxP LLaMA-3 pre-training", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + 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", + help="maxP variant") + p.add_argument("--lr", type=float, default=1e-3, + help="LR prefactor") + p.add_argument("--alignment-warmup", type=int, default=100, + help="Steps before first LP re-solve (maxP only)") + p.add_argument("--solve-interval", type=int, default=200, + help="Re-solve LP every N steps (maxP only)") + p.add_argument("--sample-size", type=int, default=8, + help="Sequences for alignment measurement (maxP only)") + p.add_argument("--c-ema", type=float, default=0.0, + help="EMA smoothing for c values (maxP only)") + p.add_argument("--seq-len", type=int, default=3072, + help="Sequence length") + p.add_argument("--steps", type=int, default=None, + help="Training steps (default: auto-computed as 20 × non-embed params / tokens-per-step)") + p.add_argument("--batch-size", type=int, default=16, + help="Global batch size (divided by WORLD_SIZE to get per-GPU)") + p.add_argument("--seed", type=int, default=1, + help="Random seed") + p.add_argument("--output-dir", default="./outputs", + help="Directory to save checkpoints and logs") + p.add_argument("--dataset", default="fineweb-edu-10bt", + help="HuggingFace dataset name or local path") + p.add_argument("--dataset-path", default=None, + help="Override dataset path (e.g. absolute path to c4_test on disk)") + p.add_argument("--num-workers", type=int, default=8, + help="DataLoader num_workers for prefetching") + p.add_argument("--prefetch-factor", type=int, default=4, + help="Batches prefetched per DataLoader worker") + p.add_argument("--hf-assets-path", default=default_hf_path, + help="Path to HF tokenizer assets (local copy)") + p.add_argument("--checkpoint-interval", type=int, default=None, + help="Save checkpoint every N steps") + return p.parse_args() + + +def main() -> None: + init_logger() + + args = parse_args() + config = build_trainer_config(args) + trainer = MaxPTrainer(config) + + try: + trainer.train() + except Exception: + if trainer: + trainer.close() + raise + else: + trainer.close() + if torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() + logger.info("Process group destroyed") + + +if __name__ == "__main__": + main() diff --git a/maxp/dag.py b/maxp/dag.py index 886a9bf..eca8727 100644 --- a/maxp/dag.py +++ b/maxp/dag.py @@ -202,7 +202,9 @@ def build_graph(self, pms: list[tuple[str, Any]], nodes: dict[str, DagNode] = {} for name, pm in pms: lt = pm.layer_type - a, b = ab.get(lt, (0.0, 0.5)) + a_default, b_default = ab.get(lt, (0.0, 0.5)) + a = pm.a if pm.a is not None else a_default + b = pm.b if pm.b is not None else b_default nodes[name] = DagNode( name=name, a=a, diff --git a/maxp/diagnose.py b/maxp/diagnose.py index 0fe7ed3..3ff913f 100644 --- a/maxp/diagnose.py +++ b/maxp/diagnose.py @@ -40,11 +40,10 @@ def diagnose_axis( widths: list of widths to sweep n_steps: training steps per width n_seeds: random seeds per width - train_step_fn: ``fn(model, step_idx) -> None`` Runs one optimiser - step (forward, backward, step, zero_grad). If ``None`` a - default loop using ``AdamW`` with cross-entropy is used — this - requires the model to have a ``.tok_emb`` attribute (works for - the example transformers). + train_step_fn: ``fn(model, param_groups) -> fn(model, step_idx)`` + Factory that builds and returns a step closure (forward, backward, + step, zero_grad). If ``None`` a default factory is used that + requires the model to have a ``.tok_emb`` attribute. Returns: all_ops: list of ClassifiedOp (traced ops, no elementwise) @@ -72,11 +71,9 @@ def diagnose_axis( torch.manual_seed(seed_idx * 1000 + w) model, param_groups = make_model_fn(w) - # Build default train_step if none supplied - if train_step_fn is not None: - _step = train_step_fn - else: - _step = _default_train_step(model, param_groups) + # Build step closure from factory + factory = train_step_fn if train_step_fn is not None else _default_train_step + _step = factory(model, param_groups) for step in range(n_steps): torch.manual_seed(seed_idx * 100000 + step) diff --git a/maxp/module.py b/maxp/module.py index c484041..bba94af 100644 --- a/maxp/module.py +++ b/maxp/module.py @@ -26,6 +26,9 @@ class ParametrizedModule(nn.Module): c: Optional override for the learning rate exponent. If set, :class:`Parametrization` will use this value instead of solving for it via LP. + scale_output: If ``False``, the solved multiplier is not physically + applied to the output tensor in ``forward()``. Useful for ops + that handle scaling internally (like SDPA via the ``scale`` arg). Attributes: inner: The wrapped ``nn.Module``, or ``None`` for bare callables. @@ -45,6 +48,7 @@ def __init__( a: float | None = None, b: float | None = None, c: float | None = None, + scale_output: bool = True, ): super().__init__() if isinstance(module_or_fn, nn.Module): @@ -55,6 +59,7 @@ def __init__( self.width_dim = width_dim self.layer_type = layer_type self.scale = 1.0 + self.scale_output = scale_output # Per-PM (a, b, c) overrides — Parametrization respects these if set self.a: float | None = a @@ -87,4 +92,8 @@ def forward(self, *args, **kwargs): out = self.inner(*args, **kwargs) else: out = self._fn(*args, **kwargs) - return self.scale * out + + if self.scale_output: + return self.scale * out + else: + return out diff --git a/maxp/parametrization.py b/maxp/parametrization.py index e5a3923..46fc8e6 100644 --- a/maxp/parametrization.py +++ b/maxp/parametrization.py @@ -147,10 +147,12 @@ def __init__( ): self.model = model self.lr_prefactor = lr_prefactor + self._std_prefactor = std_prefactor ab = dict(_DEFAULT_AB) if ab_overrides: ab.update(ab_overrides) + self._ab = ab # Discover all ParametrizedModule instances pms: list[tuple[str, ParametrizedModule]] = [ @@ -206,11 +208,68 @@ def __init__( c_by_name = _solve_graph(graph, optimizer_type, c_fixed=c_fixed or None, solver=solver) + # Phase 2 state + self._pms = pms + self._c_fixed = c_fixed or None + self._optimizer_type = optimizer_type + self._warmup_steps = warmup_steps + self._solve_interval = solve_interval + self._sample_size = sample_size + self._c_ema = c_ema + self._alignment_ema = alignment_ema + self._resample_w0 = resample_w0 + self._warm_start = warm_start + self._use_training_activations = use_training_activations + self._persistent_hooks: list[torch.utils.hooks.RemovableHook] = [] + self._latest_activations: dict[str, torch.Tensor] = {} + self._step_count = 0 + self._graph = graph + self._c_prev: dict[str, float] = dict(c_by_name) + self._solver = solver # user-provided solver (reused for warm start) + self.refresh() # init weights, build param groups + + # Set initial alignment on each PM from the preset, with overrides + a0_dW_val, dZ_w0_val, dZ_dW_val = _ALIGNMENT_PRESETS[alignment] + self._alignment_pinned: set[str] = set() + for name, pm in pms: + # Check overrides: exact name > name suffix > layer_type + override = None + if alignment_overrides: + if name in alignment_overrides: + override = alignment_overrides[name] + else: + leaf = name.split(".")[-1] + if leaf in alignment_overrides: + override = alignment_overrides[leaf] + elif pm.layer_type in alignment_overrides: + override = alignment_overrides[pm.layer_type] + if override is not None: + pm.align_z0_dW, pm.align_dZ_w0, pm.align_dZ_dW = override + self._alignment_pinned.add(name) + else: + pm.align_z0_dW = a0_dW_val + pm.align_dZ_w0 = dZ_w0_val + pm.align_dZ_dW = dZ_dW_val + + @property + def param_groups(self) -> list[dict]: + return self._param_groups + + def refresh(self, optimizer: torch.optim.Optimizer | None = None) -> None: + """Re-run weight init and rebuild ``param_groups`` against live params. + + Use this when ``Parametrization`` was built on a meta-device model + (so ``__init__`` couldn't populate real weights or capture live + param tensor refs). + """ + ab = self._ab + std_prefactor = self._std_prefactor + # Apply: init weights, set scale, build param groups parametrized_ids: set[int] = set() groups: list[dict] = [] - for name, pm in pms: + for name, pm in self._pms: lt = pm.layer_type a_default, b_default = ab[lt] a = pm.a if pm.a is not None else a_default @@ -222,14 +281,14 @@ def __init__( pm.scale = fan_in ** (-a) if a != 0.0 else 1.0 if has_params: - c = c_by_name[name] + c = self._c_prev[name] # Re-initialise weights inner = pm.inner assert inner is not None init_std = std_prefactor * (fan_in ** (-b)) with torch.no_grad(): - if resample_w0: + if self._resample_w0: seed = torch.randint(0, 2**31, (1,)).item() pm._w0_seed = seed pm._w0_std = init_std @@ -248,7 +307,7 @@ def __init__( parametrized_ids.add(id(p)) groups.append({ "params": params, - "lr": lr_prefactor * (fan_in ** (-c)), + "lr": self.lr_prefactor * (fan_in ** (-c)), "layer_name": name, "fan_in": fan_in, "c": float(c), @@ -257,70 +316,33 @@ def __init__( # Collect all other parameters (LayerNorm, etc.) other = [ - p for p in model.parameters() + p for p in self.model.parameters() if p.requires_grad and id(p) not in parametrized_ids ] if other: groups.append({ "params": other, - "lr": lr_prefactor, + "lr": self.lr_prefactor, "layer_name": "_other", "maxp_managed": False, }) self._param_groups = groups - # Phase 2 state - self._pms = pms - self._c_fixed = c_fixed or None - self._optimizer_type = optimizer_type - self._warmup_steps = warmup_steps - self._solve_interval = solve_interval - self._sample_size = sample_size - self._c_ema = c_ema - self._alignment_ema = alignment_ema - self._resample_w0 = resample_w0 - self._warm_start = warm_start - self._use_training_activations = use_training_activations - self._persistent_hooks: list[torch.utils.hooks.RemovableHook] = [] - self._latest_activations: dict[str, torch.Tensor] = {} - self._step_count = 0 - self._graph = graph - self._c_prev: dict[str, float] = dict(c_by_name) - self._solver = solver # user-provided solver (reused for warm start) - # c_target tracks the latest solver output; c blends toward it each step self._c_target: dict[str, float] = { group["layer_name"]: group["c"] for group in groups if group.get("maxp_managed", False) } - # Set initial alignment on each PM from the preset, with overrides - a0_dW_val, dZ_w0_val, dZ_dW_val = _ALIGNMENT_PRESETS[alignment] - self._alignment_pinned: set[str] = set() - for name, pm in pms: - # Check overrides: exact name > name suffix > layer_type - override = None - if alignment_overrides: - if name in alignment_overrides: - override = alignment_overrides[name] - else: - leaf = name.split(".")[-1] - if leaf in alignment_overrides: - override = alignment_overrides[leaf] - elif pm.layer_type in alignment_overrides: - override = alignment_overrides[pm.layer_type] - if override is not None: - pm.align_z0_dW, pm.align_dZ_w0, pm.align_dZ_dW = override - self._alignment_pinned.add(name) - else: - pm.align_z0_dW = a0_dW_val - pm.align_dZ_w0 = dZ_w0_val - pm.align_dZ_dW = dZ_dW_val - - @property - def param_groups(self) -> list[dict]: - return self._param_groups + if optimizer is not None: + non_lr_defaults = { + k: v for k, v in optimizer.param_groups[0].items() + if k not in ("lr", "params") + } + optimizer.param_groups[:] = [{**g, **non_lr_defaults} for g in groups] + # Reset state that referenced the old (now-dead) parameter ids + optimizer.state.clear() # ------------------------------------------------------------------ # Phase 2: dynamic alignment @@ -402,12 +424,15 @@ def capture_initial(self, sample_input: torch.Tensor) -> None: captured = self._capture_activations(sample_input) for name, pm in self._pms: - if pm.inner is not None and name in captured: + if pm.inner is not None and pm.weight is not None and name in captured: pm._z0 = captured[name] if not self._resample_w0: - pm._w0 = pm.weight.detach().clone() + _w = pm.weight.detach() + if hasattr(_w, "to_local"): + _w = _w.to_local() + pm._w0 = _w.clone() - if self._use_training_activations: + if self._use_training_activations and not self._persistent_hooks: self.register_hooks() def _regenerate_w0(self, pm: ParametrizedModule) -> torch.Tensor: @@ -427,6 +452,7 @@ def _sync_lrs(self, optimizer: torch.optim.Optimizer | None) -> None: ema = self._c_ema for group in self._param_groups: if not group.get("maxp_managed", False): + group["lr"] = self.lr_prefactor continue if ema > 0: name = group["layer_name"] @@ -502,7 +528,10 @@ def step( z0 = pm._z0 w0 = self._regenerate_w0(pm) if self._resample_w0 else pm._w0 z = current[name] - w = pm.weight.detach().clone() + w = pm.weight.detach() + if hasattr(w, "to_local"): + w = w.to_local() + w = w.clone() new_a0_dW, new_dZ_w0, new_dZ_dW = compute_alignment( z0, w0, z, w, fan_in=pm.width_dim ) diff --git a/maxp/trace.py b/maxp/trace.py index 7157f2f..6dbc219 100644 --- a/maxp/trace.py +++ b/maxp/trace.py @@ -10,6 +10,7 @@ import torch import torch.nn as nn import torch.nn.functional as F +from torch.distributed.tensor import DTensor @dataclass @@ -34,11 +35,31 @@ class _TracingTensor(torch.Tensor): def __torch_function__(cls, func, types, args=(), kwargs=None): kwargs = kwargs or {} - # Call the real function with unwrapped tensors - raw_args = tuple(a._real if isinstance(a, _TracingTensor) else a for a in args) - raw_kwargs = {k: v._real if isinstance(v, _TracingTensor) else v for k, v in kwargs.items()} + def _unwrap(x): + if isinstance(x, _TracingTensor): + return x._real + if isinstance(x, list): + return [_unwrap(item) for item in x] + if isinstance(x, tuple): + return type(x)(_unwrap(item) for item in x) + return x + + raw_args = tuple(_unwrap(a) for a in args) + raw_kwargs = {k: _unwrap(v) for k, v in kwargs.items()} out = func(*raw_args, **raw_kwargs) + # detach() must always return the same subclass so nn.Parameter(traced_tensor) + # works correctly (PyTorch asserts type(data.detach()) == type(data)). + if func.__name__ == "detach" and args and isinstance(args[0], _TracingTensor): + return _TracingTensor._wrap(out) + + # Outside of an active tracing context, pass the result through without + # wrapping — this prevents _TracingDTensor params persisted by FSDP2 from + # wrapping every training op output and causing recursive __torch_function__ + # calls during clip_grad_norm_ and similar utilities. + if cls._tracer is None and cls._dag_builder is None: + return out + # Record if this is a matmul-like op tracer = cls._tracer if tracer is not None and func in _TRACED_OPS: @@ -58,7 +79,10 @@ def __torch_function__(cls, func, types, args=(), kwargs=None): @classmethod def _wrap(cls, t: torch.Tensor) -> "_TracingTensor": - r = t.as_subclass(cls) + if isinstance(t, DTensor): + r = DTensor.__new__(_TracingDTensor, t._local_tensor, t._spec, requires_grad=t.requires_grad) + else: + r = t.as_subclass(cls) r._real = t r._pm_tags = frozenset() return r @@ -76,6 +100,14 @@ def __init_subclass__(cls, **kwargs): super().__init_subclass__(**kwargs) +class _TracingDTensor(_TracingTensor, DTensor): + """_TracingTensor variant for DTensor parameters (e.g. FSDP2-sharded weights). + Preserves the DTensor structure so FSDP2's isinstance(x, DTensor) checks and + _local_tensor access continue to work during lazy init. + """ + pass + + # Ops that indicate elementwise multiply (merge_type = SUM in exponent space) _MUL_OPS = {torch.mul, torch.Tensor.mul, torch.Tensor.__mul__} @@ -167,7 +199,7 @@ def __enter__(self): if isinstance(mod, nn.Embedding): original = mod.weight wrapped = _TracingTensor._wrap(original.data) - mod.weight = nn.Parameter(wrapped, requires_grad=original.requires_grad) + mod._parameters['weight'] = wrapped self._wrapped_embeddings.append((mod, original)) # Register hooks to track module context @@ -187,7 +219,7 @@ def __exit__(self, *args): self._hooks.clear() # Restore original embedding weights for mod, original in self._wrapped_embeddings: - mod.weight = original + mod._parameters['weight'] = original self._wrapped_embeddings.clear() def _pre_hook(self, module, input): diff --git a/pyproject.toml b/pyproject.toml index dd517d5..461dd27 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,6 +34,8 @@ dependencies = [ ] [project.optional-dependencies] +cu128 = ["torch>=2.8.0", "torchvision"] +cu130 = ["torch>=2.8.0", "torchvision"] dev = [ "datasets>=3.0.0", "matplotlib>=3.10.8", @@ -42,8 +44,10 @@ dev = [ "pytest-cov>=7.0.0", "pyyaml>=6.0.3", "tiktoken>=0.7.0", + "torchtitan", "torchvision>=0.24.1", "tqdm>=4.60.0", + "wandb>=0.26.0", ] [project.urls] @@ -57,3 +61,29 @@ testpaths = ["tests"] python_files = ["test_*.py"] python_functions = ["test_*"] addopts = "-v --tb=short" + +[tool.uv] +conflicts = [ + [{ extra = "cu128" }, { extra = "cu130" }], +] + +[[tool.uv.index]] +name = "pytorch-cu128" +url = "https://download.pytorch.org/whl/cu128" +explicit = true + +[[tool.uv.index]] +name = "pytorch-cu130" +url = "https://download.pytorch.org/whl/cu130" +explicit = true + +[tool.uv.sources] +torch = [ + { index = "pytorch-cu128", extra = "cu128" }, + { index = "pytorch-cu130", extra = "cu130" }, +] +torchvision = [ + { index = "pytorch-cu128", extra = "cu128" }, + { index = "pytorch-cu130", extra = "cu130" }, +] +torchtitan = { git = "https://github.com/pytorch/torchtitan.git", rev = "1a2fef04e08fb7d1667deebeae598d6a3b9124a2" }