From 6dbcc5cab39b7e7b83d64d63af222821e6227d6e Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Tue, 21 Apr 2026 20:36:16 +0200 Subject: [PATCH 01/29] Add LLaMA-3 LM experiments and update core library MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - experiments/lm: full torchtitan-based training pipeline (MaxPTrainer, maxp_converter, maxp_llama3 scale configs, SLURM launch sweep, debug runner) - Steps auto-computed as 20× non-embed params via meta-device param count - Replace torchtitan submodule with torchtitan>=0.2.2 pip dependency; remove all sys.path hacks from scripts - Fix DTensor/Tensor mix in alignment computation (FSDP2 compat) - Method rename: mp-full→mup-full, mp-noalign→mup-no - docs/experiments.md: consolidated experiment plan - pyproject.toml: add torchtitan, trim unused dev deps Co-Authored-By: Claude Sonnet 4.6 --- docs/experiments.md | 422 +++++++++++++++++++++++++++++++ experiments/lm/coord_check.py | 165 ++++++++++++ experiments/lm/fineweb.py | 42 +++ experiments/lm/launch_sweep.py | 206 +++++++++++++++ experiments/lm/maxp_converter.py | 192 ++++++++++++++ experiments/lm/maxp_llama3.py | 247 ++++++++++++++++++ experiments/lm/run.sh | 48 ++++ experiments/lm/run_debug.sh | 69 +++++ experiments/lm/train.py | 243 ++++++++++++++++++ maxp/diagnose.py | 17 +- maxp/parametrization.py | 12 +- maxp/trace.py | 44 +++- pyproject.toml | 2 + 13 files changed, 1690 insertions(+), 19 deletions(-) create mode 100644 docs/experiments.md create mode 100644 experiments/lm/coord_check.py create mode 100644 experiments/lm/fineweb.py create mode 100644 experiments/lm/launch_sweep.py create mode 100644 experiments/lm/maxp_converter.py create mode 100644 experiments/lm/maxp_llama3.py create mode 100644 experiments/lm/run.sh create mode 100755 experiments/lm/run_debug.sh create mode 100644 experiments/lm/train.py 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/fineweb.py b/experiments/lm/fineweb.py new file mode 100644 index 0000000..02ea36d --- /dev/null +++ b/experiments/lm/fineweb.py @@ -0,0 +1,42 @@ +"""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. + +FineWeb-Edu is streamed directly from HuggingFace (or from a local cache +if HF_DATASETS_CACHE is set). 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 — faster to iterate for S1/S2 + return load_dataset(dataset_path, name="sample-10BT", 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, +) diff --git a/experiments/lm/launch_sweep.py b/experiments/lm/launch_sweep.py new file mode 100644 index 0000000..7e92dd9 --- /dev/null +++ b/experiments/lm/launch_sweep.py @@ -0,0 +1,206 @@ +#!/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/plgwifillm/maxP/runs \\ + --dataset HuggingFaceFW/fineweb-edu \\ + --hf-assets-path /net/storage/plgwifillm/maxP/tokenizer \\ + --venv-path /net/storage/plgwifillm/maxP/.venv \\ + --repo-path /net/storage/plgwifillm/maxP \\ + [--dry-run] + +Skips runs where outputs//checkpoint/step-/ already exists. +""" + +from __future__ import annotations + +import argparse +import os +import subprocess +from datetime import date +from pathlib import Path +from string import Template + +from maxp_llama3 import compute_steps + + +# --------------------------------------------------------------------------- +# Per-scale defaults +# --------------------------------------------------------------------------- + +SCALE_CONFIGS = { + "s1": {"wall": "02:00:00", "nodes": 1}, + "s2": {"wall": "02:00:00", "nodes": 1}, + "s3": {"wall": "06:00:00", "nodes": 1}, + "s4": {"wall": "24:00:00", "nodes": 1}, + "s5": {"wall": "48:00:00", "nodes": 1}, +} + +# Supported methods (currently implemented) +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], +} + +# Full LR grid from experiments.md §4.1 +ALL_LRS = [3e-4, 1e-3, 3e-3, 1e-2, 3e-2, 1e-1, 3e-1] + +S5_SINGLE_LR = [1e-2] # transfer test: single best LR from S4 + + +# --------------------------------------------------------------------------- +# 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 plgwifillm-gpu-gh200 +#SBATCH --partition plgrid-gpu-gh200 +#SBATCH --gres gpu:8 +#SBATCH --output ${output_dir}/slurm.out +#SBATCH --error ${output_dir}/slurm.err + +module add ML-bundle/25.10 +source "${venv_path}/bin/activate" +cd "${repo_path}" + +export OMP_NUM_THREADS=8 +export HF_DATASETS_CACHE="$${HF_DATASETS_CACHE:-/net/storage/plgwifillm/maxP/hf_cache}" +export WANDB_PROJECT="$${WANDB_PROJECT:-maxP-lm}" +export WANDB_RUN_NAME="maxp_${scale}_${method_tag}_lr${lr_tag}_s${seed}" + +torchrun --nproc_per_node=${gpus_per_node} experiments/lm/train.py \\ + --scale ${scale} \\ + --method ${method} \\ + --lr ${lr} \\ + --seed ${seed} \\ + --batch-size ${batch_size} \\ + --dataset ${dataset} \\ + --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 _run_complete(out_dir: Path, steps: int) -> bool: + """True if a checkpoint at the final step already exists.""" + ckpt_dir = out_dir / "checkpoint" / f"step-{steps}" + return ckpt_dir.exists() + + +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=8, + help="Local batch size per GPU") + p.add_argument("--gpus-per-node", type=int, default=8, + help="GPUs per SLURM node (also sets --nproc_per_node)") + p.add_argument("--seq-len", type=int, default=2048) + p.add_argument("--runs-dir", required=True, help="Root directory for run outputs") + p.add_argument("--dataset", required=True, + help="HuggingFace dataset name (e.g. HuggingFaceFW/fineweb-edu)") + 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") + 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] + world_size = sc["nodes"] * args.gpus_per_node + steps = compute_steps(scale, args.seq_len, args.batch_size, world_size) + + 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 _run_complete(out_dir, steps): + print(f"[skip] {run_name}") + skipped += 1 + continue + + out_dir.mkdir(parents=True, exist_ok=True) + + 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=args.gpus_per_node, + dataset=args.dataset, + 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..79f33d4 --- /dev/null +++ b/experiments/lm/maxp_converter.py @@ -0,0 +1,192 @@ +"""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 maxp import ParametrizedModule, Parametrization + + +def _wrap(parent: nn.Module, attr: str, width_dim: int, layer_type: str, **pm_kw) -> None: + """Replace parent. with a ParametrizedModule in-place.""" + layer = getattr(parent, attr) + setattr( + parent, + attr, + ParametrizedModule(layer, width_dim=width_dim, layer_type=layer_type, **pm_kw), + ) + + +def install_pm_wrappers(model: nn.Module) -> None: + """Install ParametrizedModule 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") + + # SDPA already scales by 1/sqrt(head_dim), so set a=0 to avoid double-scaling + attn.inner_attention = ParametrizedModule( + attn.inner_attention, + width_dim=head_dim, + layer_type="readout", + a=0.0, + ) + + 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): + """ModelConverter that installs ParametrizedModule wrappers on a LLaMA-3 model. + + install_pm_wrappers() is called during convert() while the model is still + on meta device — purely structural, no tensor operations. + Parametrization() (LP solve + weight reinit) runs later in post_optimizer_build_fn, + after init_weights() has materialized real tensors. + """ + + @dataclass(kw_only=True, slots=True) + class Config(Configurable.Config): + pass + + def __init__( + self, + config: Config, + *, + parallel_dims: ParallelDims, + model_compile_enabled: bool, + ) -> None: + pass # stateless + + def convert(self, model: nn.Module) -> None: + install_pm_wrappers(model) + + def post_optimizer_hook(self, model: nn.Module | list[nn.Module]) -> None: + pass # static maxP: no dynamic re-solve after optimizer steps + + +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 make_post_optimizer_build_fn( + method: str, + lr_prefactor: float, + alignment_warmup: int = 10, + solve_interval: int = 100, + sample_size: int = 32, + c_ema: float = 0.0, +) -> Callable: + """Return a post_optimizer_build_fn for maxP initialization. + + The returned function is called by Trainer after init_weights() and + optimizer.build(), but before lr_scheduler.build(). It runs + Parametrization (LP solve + weight reinit) and replaces each inner + AdamW's param_groups with maxP per-layer groups, so LRSchedulersContainer + records our per-layer initial_lr values correctly. + + For method="maxP" (dynamic), the Parametrization object is stored on the + model as ``_maxp_param`` so MaxPTrainer.train_step can call param.step() + each iteration. For "mup-full"/"mup-no" (static), LP is solved once + and never updated. + + Args: + method: "maxP" (dynamic), "mup-full" (static, full align), or + "mup-no" (static, no align). + lr_prefactor: Base LR multiplier (e.g. 1e-3). + alignment_warmup: Steps before first dynamic re-solve (maxP only). + solve_interval: Re-solve LP every N steps (maxP only). + sample_size: Number of sequences for alignment measurement (maxP only). + c_ema: EMA smoothing for c values toward LP targets (maxP only). + """ + is_dynamic = method == "maxP" + alignment = "no" if method == "mup-no" else "full" + + def fn( + optimizers, + model_parts: list[nn.Module], + parallel_dims: ParallelDims, + ) -> None: + for i, model in enumerate(model_parts): + sample_input = _make_sample_input(model) + param = Parametrization( + model, + sample_input=sample_input, + alignment=alignment, + lr_prefactor=lr_prefactor, + warmup_steps=alignment_warmup if is_dynamic else 0, + solve_interval=solve_interval, + sample_size=sample_size, + c_ema=c_ema, + ) + optimizer = optimizers.optimizers[i] + non_lr_defaults = {k: v for k, v in optimizer.defaults.items() if k != "lr"} + merged = [ + {"params": g["params"], "lr": g["lr"], + "layer_name": g.get("layer_name", "_other"), **non_lr_defaults} + for g in param.param_groups + ] + # Replace the inner AdamW's param_groups in-place. + # LRSchedulersContainer (built next) records these as initial_lr. + optimizer.param_groups[:] = merged + + if is_dynamic: + model._maxp_param = param + + 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 + } + + return fn diff --git a/experiments/lm/maxp_llama3.py b/experiments/lm/maxp_llama3.py new file mode 100644 index 0000000..08ef61d --- /dev/null +++ b/experiments/lm/maxp_llama3.py @@ -0,0 +1,247 @@ +"""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 +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 torchtitan.protocols.module import Module + +from maxp import ParametrizedModule +from maxp_converter import make_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) + ] + + +class MaxPLlama3Model(Llama3Model): + """Llama3Model that permits ParametrizedModule wrappers in verify_module_protocol. + + ParametrizedModule is an nn.Module (not a torchtitan Module), so the base + verify_module_protocol() would reject it. We skip ParametrizedModule instances + since they wrap a real torchtitan Module internally. + """ + + @dataclass(kw_only=True, slots=True) + class Config(Llama3Model.Config): + pass + + def verify_module_protocol(self) -> None: + failures: list[tuple[str, str]] = [] + for fqn, mod in self.named_modules(): + if isinstance(mod, ParametrizedModule): + continue # wraps a torchtitan Module internally + if not isinstance(mod, Module): + failures.append((fqn, type(mod).__name__)) + if failures: + details = ", ".join(f"'{fqn}' ({cls})" for fqn, cls in failures) + raise RuntimeError( + f"The following modules do not satisfy the Module protocol: {details}" + ) + + +def _make_model_config( + *, + dim: int, + n_layers: int, + n_heads: int, + n_kv_heads: int, + vocab_size: int = 128256, + attn_backend: str = "sdpa", +) -> MaxPLlama3Model.Config: + hidden_dim = compute_ffn_hidden_dim(dim, multiple_of=256, ffn_dim_multiplier=1.3) + return MaxPLlama3Model.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, + ), + ) + + +# 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, vocab_size=2048), + "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, + lr_prefactor: float, + attn_backend: str = "sdpa", + alignment_warmup: int = 10, + solve_interval: int = 100, + sample_size: int = 32, + c_ema: float = 0.0, +) -> ModelSpec: + """Build a ModelSpec for maxP LLaMA-3 training. + + Args: + scale: One of "debug", "s1" … "s5". + method: "maxP" (dynamic), "mup-full" (static, full align), or + "mup-no" (static, no align). + lr_prefactor: Base LR multiplier; per-layer LRs scale from this. + attn_backend: Attention backend ("sdpa", "flex", "varlen"). + alignment_warmup: Steps before first LP re-solve (maxP only). + solve_interval: Re-solve LP every N steps (maxP only). + sample_size: Sequences for alignment measurement (maxP only). + c_ema: EMA smoothing for c values toward LP targets (maxP only). + """ + 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=make_post_optimizer_build_fn( + method=method, + lr_prefactor=lr_prefactor, + alignment_warmup=alignment_warmup, + solve_interval=solve_interval, + sample_size=sample_size, + c_ema=c_ema, + ), + 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 = MaxPLlama3Model(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..88502d1 --- /dev/null +++ b/experiments/lm/run.sh @@ -0,0 +1,48 @@ +#!/usr/bin/env bash +# Submit a full sweep for one scale. +# +# Usage: +# bash experiments/lm/slurm/run.sh s3 [launch_sweep.py flags …] +# +# Per-scale defaults (seeds, LRs, methods) match experiments.md §4.1. +# Override any of them by passing extra flags after the scale argument. +# +# Examples: +# bash experiments/lm/slurm/run.sh s3 +# bash experiments/lm/slurm/run.sh s5 --lrs 3e-3 # transfer LR +# bash experiments/lm/slurm/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/plgwifillm/maxP/runs}" +DATASET="${DATASET:-fineweb-edu}" +HF_ASSETS_PATH="${HF_ASSETS_PATH:-/net/storage/plgwifillm/maxP/tokenizer}" +VENV_PATH="${VENV_PATH:-/net/storage/plgwifillm/maxP/.venv}" +REPO_PATH="${REPO_PATH:-/net/storage/plgwifillm/maxP}" + +cd "$REPO_PATH" +source "$VENV_PATH/bin/activate" + +python experiments/lm/launch_sweep.py \ + --scale "$SCALE" \ + --methods maxP mup-full mup-no \ + --lrs $LRS \ + --seeds $SEEDS \ + --runs-dir "$RUNS_DIR" \ + --dataset "$DATASET" \ + --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..176e0fe --- /dev/null +++ b/experiments/lm/run_debug.sh @@ -0,0 +1,69 @@ +#!/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:-20}" +SCALE="${SCALE:-debug}" +METHOD="${METHOD:-mup-no}" +GPUS="${GPUS:-1}" +OUTPUT_DIR="${OUTPUT_DIR:-/tmp/maxp_debug}" +TOKENIZER="${TOKENIZER:-}" +C4_TEST="${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 ;; + --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 + +echo "=== maxP debug run ===" +echo " scale: ${SCALE}" +echo " method: ${METHOD}" +echo " steps: ${STEPS}" +echo " gpus: ${GPUS}" +echo " output_dir: ${OUTPUT_DIR}" +echo "" + +mkdir -p "${OUTPUT_DIR}" + +TRAIN_ARGS=( + --scale "${SCALE}" + --method "${METHOD}" + --lr 1e-3 + --steps "${STEPS}" + --dataset c4_test + ${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 MASTER_ADDR=localhost MASTER_PORT=29500 \ + python experiments/lm/train.py "${TRAIN_ARGS[@]}" +else + # Multi-GPU: use torchrun + torchrun --nproc_per_node="${GPUS}" 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..7ee8a8e --- /dev/null +++ b/experiments/lm/train.py @@ -0,0 +1,243 @@ +"""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 + +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, + DebugConfig, + TrainingConfig, +) +from torchtitan.hf_datasets.text_datasets import HuggingFaceTextDataLoader +from torchtitan.protocols.model_converter import ModelConvertersContainer +from torchtitan.trainer import Trainer + +from torchtitan.tools.logging import init_logger, logger + +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, "_maxp_param", None) is not None + 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): + param = getattr(model, "_maxp_param", None) + if param is None or not getattr(model, "_maxp_ready", False): + continue + 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 _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, args.batch_size, world_size) + + +def build_trainer_config(args: argparse.Namespace) -> Trainer.Config: + is_debug = args.scale == "debug" + steps = _resolve_steps(args) + + model_spec = maxp_model_registry( + scale=args.scale, + method=args.method, + lr_prefactor=args.lr, + attn_backend="sdpa", + alignment_warmup=args.alignment_warmup, + solve_interval=args.solve_interval, + sample_size=args.sample_size, + c_ema=args.c_ema, + ) + + 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()], + ), + 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 2000, + decay_ratio=0.1, + ), + training=TrainingConfig( + local_batch_size=2 if is_debug else args.batch_size, + seq_len=args.seq_len, + steps=steps, + ), + dataloader=HuggingFaceTextDataLoader.Config( + dataset=args.dataset, + dataset_path=args.dataset_path, + ), + metrics=MetricsProcessor.Config( + log_freq=10 if is_debug else 50, + enable_tensorboard=not is_debug, + enable_wandb=not is_debug, + ), + checkpoint=CheckpointManager.Config( + interval=500, + last_save_model_only=False, + ), + compile=CompileConfig(enable=not is_debug), + activation_checkpoint=ActivationCheckpointConfig(mode="selective"), + debug=DebugConfig(seed=args.seed), + validator=Validator.Config( + enable=not is_debug, + freq=500, + steps=50, + ), + ) + + +def parse_args() -> argparse.Namespace: + 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=10, + help="Steps before first LP re-solve (maxP only)") + p.add_argument("--solve-interval", type=int, default=100, + help="Re-solve LP every N steps (maxP only)") + p.add_argument("--sample-size", type=int, default=32, + 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=2048) + 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=8, help="Local batch size per GPU") + p.add_argument("--seed", type=int, default=1) + p.add_argument("--output-dir", default="./outputs") + p.add_argument( + "--dataset", + default="c4_test", + 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( + "--hf-assets-path", + default=None, + help="Path to HF tokenizer assets (local copy)", + ) + 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/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/parametrization.py b/maxp/parametrization.py index e5a3923..8207e71 100644 --- a/maxp/parametrization.py +++ b/maxp/parametrization.py @@ -402,10 +402,13 @@ 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: self.register_hooks() @@ -502,7 +505,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..02996a9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -42,8 +42,10 @@ dev = [ "pytest-cov>=7.0.0", "pyyaml>=6.0.3", "tiktoken>=0.7.0", + "torchtitan>=0.2.2", "torchvision>=0.24.1", "tqdm>=4.60.0", + "wandb>=0.26.0", ] [project.urls] From f4b7b7ad9c3d8c6d7e39b2ae642a5c1bd48bb13e Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Tue, 21 Apr 2026 20:50:53 +0200 Subject: [PATCH 02/29] Update README and pyproject: CUDA extras, torchtitan nightly - README: add installation with cu128/cu130/cpu extras, LLM experiments section - pyproject: cu128/cu130 extras for CUDA index selection, remove torchaudio - torchtitan pinned to nightly 1a2fef04 (2026-04-21) Co-Authored-By: Claude Sonnet 4.6 --- README.md | 43 ++++++++++++++++++++++++++++++++++++++----- pyproject.toml | 25 ++++++++++++++++++++++++- 2 files changed, 62 insertions(+), 6 deletions(-) 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/pyproject.toml b/pyproject.toml index 02996a9..76a5d15 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,7 +44,7 @@ dev = [ "pytest-cov>=7.0.0", "pyyaml>=6.0.3", "tiktoken>=0.7.0", - "torchtitan>=0.2.2", + "torchtitan", "torchvision>=0.24.1", "tqdm>=4.60.0", "wandb>=0.26.0", @@ -59,3 +61,24 @@ testpaths = ["tests"] python_files = ["test_*.py"] python_functions = ["test_*"] addopts = "-v --tb=short" + +[[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" } From 6119ab3802afebe6037eb1208c6b34f0f134f702 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Tue, 21 Apr 2026 20:53:32 +0200 Subject: [PATCH 03/29] Fix uv CUDA extras: add conflicts to prevent cu128/cu130 co-resolution Co-Authored-By: Claude Sonnet 4.6 --- pyproject.toml | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 76a5d15..461dd27 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -62,6 +62,11 @@ 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" From 409dbfe3aac226685cdbf17a5f3da56e0b207841 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Wed, 22 Apr 2026 10:47:23 +0200 Subject: [PATCH 04/29] Add download_hf_assets script and fix SDPA wrapper --- .gitignore | 1 + experiments/lm/download_hf_assets.py | 259 +++++++++++++++++++++++++++ experiments/lm/launch_sweep.py | 2 +- experiments/lm/maxp_converter.py | 28 ++- experiments/lm/run.sh | 2 +- experiments/lm/run_debug.sh | 9 +- experiments/lm/train.py | 63 +++---- maxp/dag.py | 4 +- maxp/module.py | 11 +- 9 files changed, 327 insertions(+), 52 deletions(-) create mode 100644 experiments/lm/download_hf_assets.py diff --git a/.gitignore b/.gitignore index 20ae408..59dab27 100644 --- a/.gitignore +++ b/.gitignore @@ -209,6 +209,7 @@ __marimo__/ data jobs outputs +assets # Generated plots *.png 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/launch_sweep.py b/experiments/lm/launch_sweep.py index 7e92dd9..ae70081 100644 --- a/experiments/lm/launch_sweep.py +++ b/experiments/lm/launch_sweep.py @@ -1,7 +1,7 @@ #!/usr/bin/env python3 """Generate and submit SLURM jobs for the full LM sweep. -Usage:: +Usage: python experiments/lm/launch_sweep.py \\ --scale s3 \\ diff --git a/experiments/lm/maxp_converter.py b/experiments/lm/maxp_converter.py index 79f33d4..32666f2 100644 --- a/experiments/lm/maxp_converter.py +++ b/experiments/lm/maxp_converter.py @@ -24,6 +24,26 @@ def _wrap(parent: nn.Module, attr: str, width_dim: int, layer_type: str, **pm_kw ParametrizedModule(layer, width_dim=width_dim, layer_type=layer_type, **pm_kw), ) +class _SDPAWrapper(nn.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 = ParametrizedModule( + 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 ParametrizedModule wrappers on all LLaMA-3 layers in-place. @@ -55,13 +75,7 @@ def install_pm_wrappers(model: nn.Module) -> None: _wrap(attn, "wo", d_model, "hidden") - # SDPA already scales by 1/sqrt(head_dim), so set a=0 to avoid double-scaling - attn.inner_attention = ParametrizedModule( - attn.inner_attention, - width_dim=head_dim, - layer_type="readout", - a=0.0, - ) + attn.inner_attention = _SDPAWrapper(attn.inner_attention, head_dim) d_ff: int = ffn.w1.out_features _wrap(ffn, "w1", d_model, "hidden") diff --git a/experiments/lm/run.sh b/experiments/lm/run.sh index 88502d1..cda3eb5 100644 --- a/experiments/lm/run.sh +++ b/experiments/lm/run.sh @@ -28,7 +28,7 @@ esac RUNS_DIR="${RUNS_DIR:-/net/storage/plgwifillm/maxP/runs}" DATASET="${DATASET:-fineweb-edu}" -HF_ASSETS_PATH="${HF_ASSETS_PATH:-/net/storage/plgwifillm/maxP/tokenizer}" +HF_ASSETS_PATH="${HF_ASSETS_PATH:-$(pwd)/assets/hf/Llama-3.1-8B}" VENV_PATH="${VENV_PATH:-/net/storage/plgwifillm/maxP/.venv}" REPO_PATH="${REPO_PATH:-/net/storage/plgwifillm/maxP}" diff --git a/experiments/lm/run_debug.sh b/experiments/lm/run_debug.sh index 176e0fe..37d1ea5 100755 --- a/experiments/lm/run_debug.sh +++ b/experiments/lm/run_debug.sh @@ -15,9 +15,10 @@ STEPS="${STEPS:-20}" SCALE="${SCALE:-debug}" METHOD="${METHOD:-mup-no}" GPUS="${GPUS:-1}" +DATASET="${DATASET:-c4_test}" OUTPUT_DIR="${OUTPUT_DIR:-/tmp/maxp_debug}" -TOKENIZER="${TOKENIZER:-}" -C4_TEST="${C4_TEST:-}" +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 @@ -26,6 +27,7 @@ while [[ $# -gt 0 ]]; do --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 ;; @@ -41,6 +43,7 @@ export OMP_NUM_THREADS=4 echo "=== maxP debug run ===" echo " scale: ${SCALE}" echo " method: ${METHOD}" +echo " dataset: ${DATASET}" echo " steps: ${STEPS}" echo " gpus: ${GPUS}" echo " output_dir: ${OUTPUT_DIR}" @@ -53,7 +56,7 @@ TRAIN_ARGS=( --method "${METHOD}" --lr 1e-3 --steps "${STEPS}" - --dataset c4_test + --dataset "${DATASET}" ${C4_TEST:+--dataset-path "${C4_TEST}"} ${TOKENIZER:+--hf-assets-path "${TOKENIZER}"} --output-dir "${OUTPUT_DIR}" diff --git a/experiments/lm/train.py b/experiments/lm/train.py index 7ee8a8e..eec6dbe 100644 --- a/experiments/lm/train.py +++ b/experiments/lm/train.py @@ -2,11 +2,11 @@ Uses torchtitan's Trainer with MaxPConverter and per-layer LRs from Parametrization. -Usage (single GPU, debug smoke test):: +Usage (single GPU, debug smoke test): python experiments/lm/train.py --scale debug --steps 20 -Usage (real training, 8 GPUs):: +Usage (real training, 8 GPUs): torchrun --nproc_per_node=8 experiments/lm/train.py \\ --scale s3 --method maxP --lr 1e-3 --steps 10000 \\ @@ -34,9 +34,8 @@ ) from torchtitan.hf_datasets.text_datasets import HuggingFaceTextDataLoader from torchtitan.protocols.model_converter import ModelConvertersContainer -from torchtitan.trainer import Trainer - 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 @@ -98,8 +97,6 @@ def train_step(self, data_iterator): self.metrics_processor.logger.log(extra, self.step) - - def _resolve_steps(args: argparse.Namespace) -> int: if args.steps is not None: return args.steps @@ -170,23 +167,18 @@ def build_trainer_config(args: argparse.Namespace) -> Trainer.Config: def parse_args() -> argparse.Namespace: + default_hf_path = os.path.join(os.path.dirname(__file__), "assets/hf/Llama-3.1-8B") + default_dataset_path = os.path.join(os.path.dirname(__file__), "assets/c4_test") 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("--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=10, help="Steps before first LP re-solve (maxP only)") p.add_argument("--solve-interval", type=int, default=100, @@ -195,27 +187,22 @@ def parse_args() -> argparse.Namespace: 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=2048) + p.add_argument("--seq-len", type=int, default=2048, + 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=8, help="Local batch size per GPU") - p.add_argument("--seed", type=int, default=1) - p.add_argument("--output-dir", default="./outputs") - p.add_argument( - "--dataset", - default="c4_test", - 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( - "--hf-assets-path", - default=None, - help="Path to HF tokenizer assets (local copy)", - ) + p.add_argument("--batch-size", type=int, default=8, + help="Local batch size 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="c4_test", + help="HuggingFace dataset name or local path") + p.add_argument("--dataset-path", default=default_dataset_path, + help="Override dataset path (e.g. absolute path to c4_test on disk)") + p.add_argument("--hf-assets-path", default=default_hf_path, + help="Path to HF tokenizer assets (local copy)") return p.parse_args() 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/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 From f82c57ac164439948973923dc85df81474eb5c90 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Wed, 22 Apr 2026 11:13:21 +0200 Subject: [PATCH 05/29] Introduce LlamaParametrizedModule, fix compilation --- experiments/lm/maxp_converter.py | 18 ++++++++---- experiments/lm/maxp_llama3.py | 48 ++++---------------------------- experiments/lm/train.py | 2 +- 3 files changed, 18 insertions(+), 50 deletions(-) diff --git a/experiments/lm/maxp_converter.py b/experiments/lm/maxp_converter.py index 32666f2..426d493 100644 --- a/experiments/lm/maxp_converter.py +++ b/experiments/lm/maxp_converter.py @@ -11,20 +11,26 @@ 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 ParametrizedModule in-place.""" + """Replace parent. with a LlamaParametrizedModule in-place.""" layer = getattr(parent, attr) setattr( parent, attr, - ParametrizedModule(layer, width_dim=width_dim, layer_type=layer_type, **pm_kw), + LlamaParametrizedModule(layer, width_dim=width_dim, layer_type=layer_type, **pm_kw), ) -class _SDPAWrapper(nn.Module): +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 @@ -34,7 +40,7 @@ class _SDPAWrapper(nn.Module): def __init__(self, inner, head_dim): super().__init__() self.inner = inner - self.score = ParametrizedModule( + self.score = LlamaParametrizedModule( lambda q, k: (q, k), width_dim=head_dim, layer_type="readout", scale_output=False, ) @@ -46,7 +52,7 @@ def forward(self, q, k, v, **kw): def install_pm_wrappers(model: nn.Module) -> None: - """Install ParametrizedModule wrappers on all LLaMA-3 layers in-place. + """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. @@ -87,7 +93,7 @@ def install_pm_wrappers(model: nn.Module) -> None: class MaxPConverter(Configurable, ModelConverter): - """ModelConverter that installs ParametrizedModule wrappers on a LLaMA-3 model. + """ModelConverter that installs LlamaParametrizedModule wrappers on a LLaMA-3 model. install_pm_wrappers() is called during convert() while the model is still on meta device — purely structural, no tensor operations. diff --git a/experiments/lm/maxp_llama3.py b/experiments/lm/maxp_llama3.py index 08ef61d..2f040d8 100644 --- a/experiments/lm/maxp_llama3.py +++ b/experiments/lm/maxp_llama3.py @@ -10,26 +10,14 @@ 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 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 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 torchtitan.protocols.module import Module -from maxp import ParametrizedModule from maxp_converter import make_post_optimizer_build_fn @@ -89,32 +77,6 @@ def _build_layers( ] -class MaxPLlama3Model(Llama3Model): - """Llama3Model that permits ParametrizedModule wrappers in verify_module_protocol. - - ParametrizedModule is an nn.Module (not a torchtitan Module), so the base - verify_module_protocol() would reject it. We skip ParametrizedModule instances - since they wrap a real torchtitan Module internally. - """ - - @dataclass(kw_only=True, slots=True) - class Config(Llama3Model.Config): - pass - - def verify_module_protocol(self) -> None: - failures: list[tuple[str, str]] = [] - for fqn, mod in self.named_modules(): - if isinstance(mod, ParametrizedModule): - continue # wraps a torchtitan Module internally - if not isinstance(mod, Module): - failures.append((fqn, type(mod).__name__)) - if failures: - details = ", ".join(f"'{fqn}' ({cls})" for fqn, cls in failures) - raise RuntimeError( - f"The following modules do not satisfy the Module protocol: {details}" - ) - - def _make_model_config( *, dim: int, @@ -123,9 +85,9 @@ def _make_model_config( n_kv_heads: int, vocab_size: int = 128256, attn_backend: str = "sdpa", -) -> MaxPLlama3Model.Config: +) -> Llama3Model.Config: hidden_dim = compute_ffn_hidden_dim(dim, multiple_of=256, ffn_dim_multiplier=1.3) - return MaxPLlama3Model.Config( + return Llama3Model.Config( dim=dim, vocab_size=vocab_size, enable_weight_tying=False, @@ -240,7 +202,7 @@ def compute_steps( 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 = MaxPLlama3Model(model_config) + 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 diff --git a/experiments/lm/train.py b/experiments/lm/train.py index eec6dbe..34d34d0 100644 --- a/experiments/lm/train.py +++ b/experiments/lm/train.py @@ -155,7 +155,7 @@ def build_trainer_config(args: argparse.Namespace) -> Trainer.Config: interval=500, last_save_model_only=False, ), - compile=CompileConfig(enable=not is_debug), + compile=CompileConfig(enable=True), activation_checkpoint=ActivationCheckpointConfig(mode="selective"), debug=DebugConfig(seed=args.seed), validator=Validator.Config( From 8bdb1672189a94fdf9ffa8504246d1fb3eca875b Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Wed, 22 Apr 2026 12:44:26 +0200 Subject: [PATCH 06/29] Fix vocab size in debug run --- experiments/lm/maxp_llama3.py | 2 +- experiments/lm/train.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/experiments/lm/maxp_llama3.py b/experiments/lm/maxp_llama3.py index 2f040d8..8870820 100644 --- a/experiments/lm/maxp_llama3.py +++ b/experiments/lm/maxp_llama3.py @@ -125,7 +125,7 @@ def _make_model_config( # 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, vocab_size=2048), + "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), diff --git a/experiments/lm/train.py b/experiments/lm/train.py index 34d34d0..eec6dbe 100644 --- a/experiments/lm/train.py +++ b/experiments/lm/train.py @@ -155,7 +155,7 @@ def build_trainer_config(args: argparse.Namespace) -> Trainer.Config: interval=500, last_save_model_only=False, ), - compile=CompileConfig(enable=True), + compile=CompileConfig(enable=not is_debug), activation_checkpoint=ActivationCheckpointConfig(mode="selective"), debug=DebugConfig(seed=args.seed), validator=Validator.Config( From 8c0a1c6de8da14b5017995f2d7264ec5575b675c Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Wed, 22 Apr 2026 18:30:07 +0200 Subject: [PATCH 07/29] Fix PL-Grid paths --- experiments/lm/launch_sweep.py | 12 ++++++------ experiments/lm/run.sh | 6 +++--- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/experiments/lm/launch_sweep.py b/experiments/lm/launch_sweep.py index ae70081..1358011 100644 --- a/experiments/lm/launch_sweep.py +++ b/experiments/lm/launch_sweep.py @@ -8,11 +8,11 @@ --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/plgwifillm/maxP/runs \\ + --runs-dir /net/storage/pr3/plgrid/plggadlers/maxP/runs \\ --dataset HuggingFaceFW/fineweb-edu \\ - --hf-assets-path /net/storage/plgwifillm/maxP/tokenizer \\ - --venv-path /net/storage/plgwifillm/maxP/.venv \\ - --repo-path /net/storage/plgwifillm/maxP \\ + --hf-assets-path /net/storage/pr3/plgrid/plggadlers/maxP/tokenizer \\ + --venv-path /net/storage/pr3/plgrid/plggadlers/maxP/.venv \\ + --repo-path /net/storage/pr3/plgrid/plggadlers/maxP \\ [--dry-run] Skips runs where outputs//checkpoint/step-/ already exists. @@ -76,7 +76,7 @@ #SBATCH --cpus-per-gpu 72 #SBATCH --mem-per-gpu 118GB #SBATCH --time ${wall_time} -#SBATCH --account plgwifillm-gpu-gh200 +#SBATCH --account plgadlers-gpu-gh200 #SBATCH --partition plgrid-gpu-gh200 #SBATCH --gres gpu:8 #SBATCH --output ${output_dir}/slurm.out @@ -87,7 +87,7 @@ cd "${repo_path}" export OMP_NUM_THREADS=8 -export HF_DATASETS_CACHE="$${HF_DATASETS_CACHE:-/net/storage/plgwifillm/maxP/hf_cache}" +export HF_DATASETS_CACHE="$${HF_DATASETS_CACHE:-/net/storage/pr3/plgrid/plggadlers/maxP/hf_cache}" export WANDB_PROJECT="$${WANDB_PROJECT:-maxP-lm}" export WANDB_RUN_NAME="maxp_${scale}_${method_tag}_lr${lr_tag}_s${seed}" diff --git a/experiments/lm/run.sh b/experiments/lm/run.sh index cda3eb5..fddca34 100644 --- a/experiments/lm/run.sh +++ b/experiments/lm/run.sh @@ -26,11 +26,11 @@ case "$SCALE" in *) echo "Unknown scale '$SCALE'. Choose from s1 s2 s3 s4 s5."; exit 1 ;; esac -RUNS_DIR="${RUNS_DIR:-/net/storage/plgwifillm/maxP/runs}" +RUNS_DIR="${RUNS_DIR:-/net/storage/pr3/plgrid/plggadlers/maxP/runs}" DATASET="${DATASET:-fineweb-edu}" HF_ASSETS_PATH="${HF_ASSETS_PATH:-$(pwd)/assets/hf/Llama-3.1-8B}" -VENV_PATH="${VENV_PATH:-/net/storage/plgwifillm/maxP/.venv}" -REPO_PATH="${REPO_PATH:-/net/storage/plgwifillm/maxP}" +VENV_PATH="${VENV_PATH:-/net/storage/pr3/plgrid/plggadlers/maxP/.venv}" +REPO_PATH="${REPO_PATH:-/net/storage/pr3/plgrid/plggadlers/maxP}" cd "$REPO_PATH" source "$VENV_PATH/bin/activate" From 3bc59f37f2ddeb5daac3273ab577edd10bc57b30 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Wed, 22 Apr 2026 18:46:29 +0200 Subject: [PATCH 08/29] Fix launch script --- experiments/lm/launch_sweep.py | 22 +++++++++------------- experiments/lm/run.sh | 8 ++++---- 2 files changed, 13 insertions(+), 17 deletions(-) diff --git a/experiments/lm/launch_sweep.py b/experiments/lm/launch_sweep.py index 1358011..bef9645 100644 --- a/experiments/lm/launch_sweep.py +++ b/experiments/lm/launch_sweep.py @@ -13,22 +13,20 @@ --hf-assets-path /net/storage/pr3/plgrid/plggadlers/maxP/tokenizer \\ --venv-path /net/storage/pr3/plgrid/plggadlers/maxP/.venv \\ --repo-path /net/storage/pr3/plgrid/plggadlers/maxP \\ - [--dry-run] + [--dry-run] [--resume] -Skips runs where outputs//checkpoint/step-/ already exists. +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 os import subprocess from datetime import date from pathlib import Path from string import Template -from maxp_llama3 import compute_steps - # --------------------------------------------------------------------------- # Per-scale defaults @@ -111,10 +109,9 @@ def _method_tag(method: str) -> str: return method.replace("-", "") -def _run_complete(out_dir: Path, steps: int) -> bool: - """True if a checkpoint at the final step already exists.""" - ckpt_dir = out_dir / "checkpoint" / f"step-{steps}" - return ckpt_dir.exists() +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: @@ -130,7 +127,6 @@ def main() -> None: help="Local batch size per GPU") p.add_argument("--gpus-per-node", type=int, default=8, help="GPUs per SLURM node (also sets --nproc_per_node)") - p.add_argument("--seq-len", type=int, default=2048) p.add_argument("--runs-dir", required=True, help="Root directory for run outputs") p.add_argument("--dataset", required=True, help="HuggingFace dataset name (e.g. HuggingFaceFW/fineweb-edu)") @@ -140,6 +136,8 @@ def main() -> None: 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 @@ -147,8 +145,6 @@ def main() -> None: 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] - world_size = sc["nodes"] * args.gpus_per_node - steps = compute_steps(scale, args.seq_len, args.batch_size, world_size) today = date.today().strftime("%Y-%m-%d") submitted = skipped = 0 @@ -161,7 +157,7 @@ def main() -> None: ) out_dir = Path(args.runs_dir) / run_name - if _run_complete(out_dir, steps): + if not args.resume and _has_checkpoint(out_dir): print(f"[skip] {run_name}") skipped += 1 continue diff --git a/experiments/lm/run.sh b/experiments/lm/run.sh index fddca34..eb71214 100644 --- a/experiments/lm/run.sh +++ b/experiments/lm/run.sh @@ -2,15 +2,15 @@ # Submit a full sweep for one scale. # # Usage: -# bash experiments/lm/slurm/run.sh s3 [launch_sweep.py flags …] +# bash experiments/lm/run.sh s3 [launch_sweep.py flags …] # # Per-scale defaults (seeds, LRs, methods) match experiments.md §4.1. # Override any of them by passing extra flags after the scale argument. # # Examples: -# bash experiments/lm/slurm/run.sh s3 -# bash experiments/lm/slurm/run.sh s5 --lrs 3e-3 # transfer LR -# bash experiments/lm/slurm/run.sh s2 --dry-run +# 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]}" From 41f817df147de693caaef6dabec4fbc1ea18fc78 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Wed, 22 Apr 2026 18:57:13 +0200 Subject: [PATCH 09/29] Fix launch script --- experiments/lm/launch_sweep.py | 10 +++++----- experiments/lm/run.sh | 2 -- 2 files changed, 5 insertions(+), 7 deletions(-) diff --git a/experiments/lm/launch_sweep.py b/experiments/lm/launch_sweep.py index bef9645..511bbe7 100644 --- a/experiments/lm/launch_sweep.py +++ b/experiments/lm/launch_sweep.py @@ -76,16 +76,16 @@ #SBATCH --time ${wall_time} #SBATCH --account plgadlers-gpu-gh200 #SBATCH --partition plgrid-gpu-gh200 -#SBATCH --gres gpu:8 -#SBATCH --output ${output_dir}/slurm.out -#SBATCH --error ${output_dir}/slurm.err +#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=8 -export HF_DATASETS_CACHE="$${HF_DATASETS_CACHE:-/net/storage/pr3/plgrid/plggadlers/maxP/hf_cache}" +export HF_DATASETS_CACHE="$${HF_DATASETS_CACHE:-/net/storage/pr3/plgrid/plggadlers/hf_cache}" export WANDB_PROJECT="$${WANDB_PROJECT:-maxP-lm}" export WANDB_RUN_NAME="maxp_${scale}_${method_tag}_lr${lr_tag}_s${seed}" @@ -125,7 +125,7 @@ def main() -> None: help="Override seed list (default: per-scale defaults)") p.add_argument("--batch-size", type=int, default=8, help="Local batch size per GPU") - p.add_argument("--gpus-per-node", type=int, default=8, + p.add_argument("--gpus-per-node", type=int, default=1, help="GPUs per SLURM node (also sets --nproc_per_node)") p.add_argument("--runs-dir", required=True, help="Root directory for run outputs") p.add_argument("--dataset", required=True, diff --git a/experiments/lm/run.sh b/experiments/lm/run.sh index eb71214..dd88de9 100644 --- a/experiments/lm/run.sh +++ b/experiments/lm/run.sh @@ -29,11 +29,9 @@ esac RUNS_DIR="${RUNS_DIR:-/net/storage/pr3/plgrid/plggadlers/maxP/runs}" DATASET="${DATASET:-fineweb-edu}" HF_ASSETS_PATH="${HF_ASSETS_PATH:-$(pwd)/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" -source "$VENV_PATH/bin/activate" python experiments/lm/launch_sweep.py \ --scale "$SCALE" \ From 62eb8a12adcde804107faddf7f10e5973a0ae624 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Wed, 22 Apr 2026 18:57:50 +0200 Subject: [PATCH 10/29] Fix launch script --- experiments/lm/run.sh | 1 + 1 file changed, 1 insertion(+) diff --git a/experiments/lm/run.sh b/experiments/lm/run.sh index dd88de9..c43f577 100644 --- a/experiments/lm/run.sh +++ b/experiments/lm/run.sh @@ -29,6 +29,7 @@ esac RUNS_DIR="${RUNS_DIR:-/net/storage/pr3/plgrid/plggadlers/maxP/runs}" DATASET="${DATASET:-fineweb-edu}" HF_ASSETS_PATH="${HF_ASSETS_PATH:-$(pwd)/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" From 97c6e80642cf0a9ebefdccbd6ec18d2699833f47 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Wed, 22 Apr 2026 19:01:19 +0200 Subject: [PATCH 11/29] Fix launch script --- experiments/lm/launch_sweep.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/experiments/lm/launch_sweep.py b/experiments/lm/launch_sweep.py index 511bbe7..9a1eca8 100644 --- a/experiments/lm/launch_sweep.py +++ b/experiments/lm/launch_sweep.py @@ -76,7 +76,7 @@ #SBATCH --time ${wall_time} #SBATCH --account plgadlers-gpu-gh200 #SBATCH --partition plgrid-gpu-gh200 -#SBATCH --gres gpu:{gpus_per_node} +#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 From 3a3d9821f5348948be5ec8f2dd1ae9bb0b791ef9 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Wed, 22 Apr 2026 19:36:19 +0200 Subject: [PATCH 12/29] Fix launch script --- experiments/lm/launch_sweep.py | 4 ++-- experiments/lm/run.sh | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/experiments/lm/launch_sweep.py b/experiments/lm/launch_sweep.py index 9a1eca8..80830c0 100644 --- a/experiments/lm/launch_sweep.py +++ b/experiments/lm/launch_sweep.py @@ -10,7 +10,7 @@ --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/tokenizer \\ + --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] @@ -85,7 +85,7 @@ cd "${repo_path}" export OMP_NUM_THREADS=8 -export HF_DATASETS_CACHE="$${HF_DATASETS_CACHE:-/net/storage/pr3/plgrid/plggadlers/hf_cache}" +export HF_HOME="$${HF_HOME:-/net/storage/pr3/plgrid/plggadlers/hf_cache}" export WANDB_PROJECT="$${WANDB_PROJECT:-maxP-lm}" export WANDB_RUN_NAME="maxp_${scale}_${method_tag}_lr${lr_tag}_s${seed}" diff --git a/experiments/lm/run.sh b/experiments/lm/run.sh index c43f577..b2d912a 100644 --- a/experiments/lm/run.sh +++ b/experiments/lm/run.sh @@ -28,7 +28,7 @@ esac RUNS_DIR="${RUNS_DIR:-/net/storage/pr3/plgrid/plggadlers/maxP/runs}" DATASET="${DATASET:-fineweb-edu}" -HF_ASSETS_PATH="${HF_ASSETS_PATH:-$(pwd)/assets/hf/Llama-3.1-8B}" +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}" From b9deee523d56538eab79199b297fcd0f1d419b0b Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Wed, 22 Apr 2026 20:47:52 +0200 Subject: [PATCH 13/29] Init parametrization before compile (to enable hooks) --- experiments/lm/fineweb.py | 2 +- experiments/lm/maxp_converter.py | 127 +++++++++++-------------------- experiments/lm/maxp_llama3.py | 30 ++------ experiments/lm/train.py | 14 ++-- 4 files changed, 59 insertions(+), 114 deletions(-) diff --git a/experiments/lm/fineweb.py b/experiments/lm/fineweb.py index 02ea36d..40bef25 100644 --- a/experiments/lm/fineweb.py +++ b/experiments/lm/fineweb.py @@ -4,7 +4,7 @@ --dataset fineweb-edu or --dataset fineweb-edu-10bt. FineWeb-Edu is streamed directly from HuggingFace (or from a local cache -if HF_DATASETS_CACHE is set). No pre-tokenization step is needed. +if HF_HOME is set). No pre-tokenization step is needed. torchtitan tokenizes on-the-fly with the LLaMA-3 tokenizer. """ diff --git a/experiments/lm/maxp_converter.py b/experiments/lm/maxp_converter.py index 426d493..88f6a83 100644 --- a/experiments/lm/maxp_converter.py +++ b/experiments/lm/maxp_converter.py @@ -24,11 +24,8 @@ class LlamaParametrizedModule(Module, ParametrizedModule): 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), - ) + 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 @@ -93,17 +90,22 @@ def install_pm_wrappers(model: nn.Module) -> None: class MaxPConverter(Configurable, ModelConverter): - """ModelConverter that installs LlamaParametrizedModule wrappers on a LLaMA-3 model. + """Installs PM wrappers and runs maxP parametrization before model compilation. - install_pm_wrappers() is called during convert() while the model is still - on meta device — purely structural, no tensor operations. - Parametrization() (LP solve + weight reinit) runs later in post_optimizer_build_fn, - after init_weights() has materialized real tensors. + 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): - pass + 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 def __init__( self, @@ -112,13 +114,32 @@ def __init__( parallel_dims: ParallelDims, model_compile_enabled: bool, ) -> None: - pass # stateless + 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, + ) + model._maxp_param_groups = param.param_groups + if cfg.method == "maxP": + model._maxp_param = param + 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 # static maxP: no dynamic re-solve after optimizer steps + pass def _make_sample_input(model: nn.Module) -> torch.Tensor | None: @@ -139,74 +160,14 @@ def _make_sample_input(model: nn.Module) -> torch.Tensor | None: return torch.randint(0, vocab_size, (1, 16), generator=gen).to(device) -def make_post_optimizer_build_fn( - method: str, - lr_prefactor: float, - alignment_warmup: int = 10, - solve_interval: int = 100, - sample_size: int = 32, - c_ema: float = 0.0, -) -> Callable: - """Return a post_optimizer_build_fn for maxP initialization. - - The returned function is called by Trainer after init_weights() and - optimizer.build(), but before lr_scheduler.build(). It runs - Parametrization (LP solve + weight reinit) and replaces each inner - AdamW's param_groups with maxP per-layer groups, so LRSchedulersContainer - records our per-layer initial_lr values correctly. - - For method="maxP" (dynamic), the Parametrization object is stored on the - model as ``_maxp_param`` so MaxPTrainer.train_step can call param.step() - each iteration. For "mup-full"/"mup-no" (static), LP is solved once - and never updated. - - Args: - method: "maxP" (dynamic), "mup-full" (static, full align), or - "mup-no" (static, no align). - lr_prefactor: Base LR multiplier (e.g. 1e-3). - alignment_warmup: Steps before first dynamic re-solve (maxP only). - solve_interval: Re-solve LP every N steps (maxP only). - sample_size: Number of sequences for alignment measurement (maxP only). - c_ema: EMA smoothing for c values toward LP targets (maxP only). +def post_optimizer_build_fn(optimizers, model_parts: list[nn.Module], parallel_dims: ParallelDims) -> None: + """Called by Trainer after init_weights() and optimizer.build(). + + Reads param_groups set by MaxPConverter.convert(). """ - is_dynamic = method == "maxP" - alignment = "no" if method == "mup-no" else "full" - - def fn( - optimizers, - model_parts: list[nn.Module], - parallel_dims: ParallelDims, - ) -> None: - for i, model in enumerate(model_parts): - sample_input = _make_sample_input(model) - param = Parametrization( - model, - sample_input=sample_input, - alignment=alignment, - lr_prefactor=lr_prefactor, - warmup_steps=alignment_warmup if is_dynamic else 0, - solve_interval=solve_interval, - sample_size=sample_size, - c_ema=c_ema, - ) - optimizer = optimizers.optimizers[i] - non_lr_defaults = {k: v for k, v in optimizer.defaults.items() if k != "lr"} - merged = [ - {"params": g["params"], "lr": g["lr"], - "layer_name": g.get("layer_name", "_other"), **non_lr_defaults} - for g in param.param_groups - ] - # Replace the inner AdamW's param_groups in-place. - # LRSchedulersContainer (built next) records these as initial_lr. - optimizer.param_groups[:] = merged - - if is_dynamic: - model._maxp_param = param - - 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 - } - - return fn + for optimizer, model in zip(optimizers, model_parts): + param_groups = getattr(model, "_maxp_param_groups", None) + if param_groups is None: + continue + non_lr_defaults = {k: v for k, v in optimizer.defaults.items() if k != "lr"} + optimizer.param_groups[:] = [{**g, **non_lr_defaults} for g in param_groups] diff --git a/experiments/lm/maxp_llama3.py b/experiments/lm/maxp_llama3.py index 8870820..4664f66 100644 --- a/experiments/lm/maxp_llama3.py +++ b/experiments/lm/maxp_llama3.py @@ -18,7 +18,7 @@ from torchtitan.models.llama3.state_dict_adapter import Llama3StateDictAdapter from torchtitan.protocols.model_spec import ModelSpec -from maxp_converter import make_post_optimizer_build_fn +from maxp_converter import post_optimizer_build_fn # Weight inits — mirrors llama3 upstream defaults @@ -134,28 +134,17 @@ def _make_model_config( } -def maxp_model_registry( - scale: str, - method: str, - lr_prefactor: float, - attn_backend: str = "sdpa", - alignment_warmup: int = 10, - solve_interval: int = 100, - sample_size: int = 32, - c_ema: float = 0.0, -) -> ModelSpec: +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). - lr_prefactor: Base LR multiplier; per-layer LRs scale from this. attn_backend: Attention backend ("sdpa", "flex", "varlen"). - alignment_warmup: Steps before first LP re-solve (maxP only). - solve_interval: Re-solve LP every N steps (maxP only). - sample_size: Sequences for alignment measurement (maxP only). - c_ema: EMA smoothing for c values toward LP targets (maxP only). """ if scale not in SCALE_CONFIGS: raise ValueError(f"Unknown scale '{scale}'. Choose from {list(SCALE_CONFIGS)}") @@ -170,14 +159,7 @@ def maxp_model_registry( parallelize_fn=parallelize_llama, pipelining_fn=pipeline_llm, build_loss_fn=build_cross_entropy_loss, - post_optimizer_build_fn=make_post_optimizer_build_fn( - method=method, - lr_prefactor=lr_prefactor, - alignment_warmup=alignment_warmup, - solve_interval=solve_interval, - sample_size=sample_size, - c_ema=c_ema, - ), + post_optimizer_build_fn=post_optimizer_build_fn, state_dict_adapter=Llama3StateDictAdapter, ) diff --git a/experiments/lm/train.py b/experiments/lm/train.py index eec6dbe..51ace71 100644 --- a/experiments/lm/train.py +++ b/experiments/lm/train.py @@ -113,12 +113,7 @@ def build_trainer_config(args: argparse.Namespace) -> Trainer.Config: model_spec = maxp_model_registry( scale=args.scale, method=args.method, - lr_prefactor=args.lr, attn_backend="sdpa", - alignment_warmup=args.alignment_warmup, - solve_interval=args.solve_interval, - sample_size=args.sample_size, - c_ema=args.c_ema, ) return Trainer.Config( @@ -126,7 +121,14 @@ def build_trainer_config(args: argparse.Namespace) -> Trainer.Config: hf_assets_path=args.hf_assets_path, dump_folder=args.output_dir, model_converters=ModelConvertersContainer.Config( - converters=[MaxPConverter.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, From 5707e90d73439bb0bd9ed76dfa38b225ccdf062c Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Wed, 22 Apr 2026 21:23:56 +0200 Subject: [PATCH 14/29] bfloat16 training and fix default port --- experiments/lm/launch_sweep.py | 2 +- experiments/lm/run_debug.sh | 2 +- experiments/lm/train.py | 1 + 3 files changed, 3 insertions(+), 2 deletions(-) diff --git a/experiments/lm/launch_sweep.py b/experiments/lm/launch_sweep.py index 80830c0..c2615f9 100644 --- a/experiments/lm/launch_sweep.py +++ b/experiments/lm/launch_sweep.py @@ -89,7 +89,7 @@ export WANDB_PROJECT="$${WANDB_PROJECT:-maxP-lm}" export WANDB_RUN_NAME="maxp_${scale}_${method_tag}_lr${lr_tag}_s${seed}" -torchrun --nproc_per_node=${gpus_per_node} experiments/lm/train.py \\ +torchrun --nproc_per_node=${gpus_per_node} --master_port=0 experiments/lm/train.py \\ --scale ${scale} \\ --method ${method} \\ --lr ${lr} \\ diff --git a/experiments/lm/run_debug.sh b/experiments/lm/run_debug.sh index 37d1ea5..32d2e03 100755 --- a/experiments/lm/run_debug.sh +++ b/experiments/lm/run_debug.sh @@ -64,7 +64,7 @@ TRAIN_ARGS=( if [[ "${GPUS}" -eq 1 ]]; then # Single GPU: run python directly so logs stream to terminal - LOCAL_RANK=0 RANK=0 WORLD_SIZE=1 MASTER_ADDR=localhost MASTER_PORT=29500 \ + LOCAL_RANK=0 RANK=0 WORLD_SIZE=1 MASTER_ADDR=localhost MASTER_PORT=0 \ python experiments/lm/train.py "${TRAIN_ARGS[@]}" else # Multi-GPU: use torchrun diff --git a/experiments/lm/train.py b/experiments/lm/train.py index 51ace71..9a3eb0f 100644 --- a/experiments/lm/train.py +++ b/experiments/lm/train.py @@ -143,6 +143,7 @@ def build_trainer_config(args: argparse.Namespace) -> Trainer.Config: local_batch_size=2 if is_debug else args.batch_size, seq_len=args.seq_len, steps=steps, + dtype="bfloat16", ), dataloader=HuggingFaceTextDataLoader.Config( dataset=args.dataset, From 7c1e23baf0db991d317cf48cbba8a3044405c45b Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Thu, 23 Apr 2026 00:13:59 +0200 Subject: [PATCH 15/29] Modify optimizer defaults handling and number of warmup steps --- .gitignore | 1 + experiments/lm/maxp_converter.py | 2 +- experiments/lm/train.py | 2 +- 3 files changed, 3 insertions(+), 2 deletions(-) diff --git a/.gitignore b/.gitignore index 59dab27..b2b19a4 100644 --- a/.gitignore +++ b/.gitignore @@ -210,6 +210,7 @@ data jobs outputs assets +runs # Generated plots *.png diff --git a/experiments/lm/maxp_converter.py b/experiments/lm/maxp_converter.py index 88f6a83..fe6d092 100644 --- a/experiments/lm/maxp_converter.py +++ b/experiments/lm/maxp_converter.py @@ -169,5 +169,5 @@ def post_optimizer_build_fn(optimizers, model_parts: list[nn.Module], parallel_d param_groups = getattr(model, "_maxp_param_groups", None) if param_groups is None: continue - non_lr_defaults = {k: v for k, v in optimizer.defaults.items() if k != "lr"} + non_lr_defaults = {k: v for k, v in optimizer.param_groups[0].items() if k != "lr" and k != "params"} optimizer.param_groups[:] = [{**g, **non_lr_defaults} for g in param_groups] diff --git a/experiments/lm/train.py b/experiments/lm/train.py index 9a3eb0f..33e7a4d 100644 --- a/experiments/lm/train.py +++ b/experiments/lm/train.py @@ -136,7 +136,7 @@ def build_trainer_config(args: argparse.Namespace) -> Trainer.Config: implementation="foreach" if is_debug else "fused", ), lr_scheduler=LRSchedulersContainer.Config( - warmup_steps=10 if is_debug else 2000, + warmup_steps=10 if is_debug else int(0.1 * steps), decay_ratio=0.1, ), training=TrainingConfig( From 989dbf5ecaf3f3ac45aa2655193498aae3e0b12d Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Thu, 23 Apr 2026 11:01:26 +0200 Subject: [PATCH 16/29] Fix maxP optimizer param groups --- experiments/lm/maxp_converter.py | 18 ++--- experiments/lm/run_debug.sh | 4 +- experiments/lm/train.py | 6 +- maxp/parametrization.py | 127 ++++++++++++++++++------------- 4 files changed, 89 insertions(+), 66 deletions(-) diff --git a/experiments/lm/maxp_converter.py b/experiments/lm/maxp_converter.py index fe6d092..f2ae5b8 100644 --- a/experiments/lm/maxp_converter.py +++ b/experiments/lm/maxp_converter.py @@ -129,9 +129,9 @@ def convert(self, model: nn.Module) -> None: sample_size=cfg.sample_size, c_ema=cfg.c_ema, ) - model._maxp_param_groups = param.param_groups + model._maxp_param = param if cfg.method == "maxP": - model._maxp_param = param + 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 @@ -161,13 +161,11 @@ def _make_sample_input(model: nn.Module) -> torch.Tensor | None: def post_optimizer_build_fn(optimizers, model_parts: list[nn.Module], parallel_dims: ParallelDims) -> None: - """Called by Trainer after init_weights() and optimizer.build(). - - Reads param_groups set by MaxPConverter.convert(). + """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): - param_groups = getattr(model, "_maxp_param_groups", None) - if param_groups is None: - continue - non_lr_defaults = {k: v for k, v in optimizer.param_groups[0].items() if k != "lr" and k != "params"} - optimizer.param_groups[:] = [{**g, **non_lr_defaults} for g in param_groups] + model._maxp_param.refresh(optimizer=optimizer) diff --git a/experiments/lm/run_debug.sh b/experiments/lm/run_debug.sh index 32d2e03..936fba1 100755 --- a/experiments/lm/run_debug.sh +++ b/experiments/lm/run_debug.sh @@ -54,7 +54,9 @@ mkdir -p "${OUTPUT_DIR}" TRAIN_ARGS=( --scale "${SCALE}" --method "${METHOD}" - --lr 1e-3 + --lr 0.03 + --alignment-warmup 20 + --solve-interval 5 --steps "${STEPS}" --dataset "${DATASET}" ${C4_TEST:+--dataset-path "${C4_TEST}"} diff --git a/experiments/lm/train.py b/experiments/lm/train.py index 33e7a4d..cc011d6 100644 --- a/experiments/lm/train.py +++ b/experiments/lm/train.py @@ -49,7 +49,7 @@ 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, "_maxp_param", None) is not None + if getattr(m, "_is_dynamic", False) and not getattr(m, "_maxp_ready", False)] if needs_capture: batch = next(data_iterator) @@ -65,9 +65,9 @@ def train_step(self, data_iterator): # 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): - param = getattr(model, "_maxp_param", None) - if param is None or not getattr(model, "_maxp_ready", False): + 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"] diff --git a/maxp/parametrization.py b/maxp/parametrization.py index 8207e71..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 @@ -410,7 +432,7 @@ def capture_initial(self, sample_input: torch.Tensor) -> None: _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: @@ -430,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"] From dc6343a46145c292f6a110d0cc280f195ad1b1de Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Thu, 23 Apr 2026 14:36:41 +0200 Subject: [PATCH 17/29] Try to fix torchrun issues --- experiments/lm/launch_sweep.py | 13 ++++++++++++- experiments/lm/run_debug.sh | 15 ++++++++++++--- 2 files changed, 24 insertions(+), 4 deletions(-) diff --git a/experiments/lm/launch_sweep.py b/experiments/lm/launch_sweep.py index c2615f9..5ca3159 100644 --- a/experiments/lm/launch_sweep.py +++ b/experiments/lm/launch_sweep.py @@ -89,7 +89,18 @@ export WANDB_PROJECT="$${WANDB_PROJECT:-maxP-lm}" export WANDB_RUN_NAME="maxp_${scale}_${method_tag}_lr${lr_tag}_s${seed}" -torchrun --nproc_per_node=${gpus_per_node} --master_port=0 experiments/lm/train.py \\ +# 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} \\ diff --git a/experiments/lm/run_debug.sh b/experiments/lm/run_debug.sh index 936fba1..0c3f569 100755 --- a/experiments/lm/run_debug.sh +++ b/experiments/lm/run_debug.sh @@ -40,6 +40,15 @@ 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}" @@ -47,6 +56,7 @@ echo " dataset: ${DATASET}" echo " steps: ${STEPS}" echo " gpus: ${GPUS}" echo " output_dir: ${OUTPUT_DIR}" +echo " master_port:${MASTER_PORT}" echo "" mkdir -p "${OUTPUT_DIR}" @@ -66,9 +76,8 @@ TRAIN_ARGS=( if [[ "${GPUS}" -eq 1 ]]; then # Single GPU: run python directly so logs stream to terminal - LOCAL_RANK=0 RANK=0 WORLD_SIZE=1 MASTER_ADDR=localhost MASTER_PORT=0 \ - python experiments/lm/train.py "${TRAIN_ARGS[@]}" + 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}" experiments/lm/train.py "${TRAIN_ARGS[@]}" + torchrun --nproc_per_node="${GPUS}" --master_addr="${MASTER_ADDR}" --master_port="${MASTER_PORT}" experiments/lm/train.py "${TRAIN_ARGS[@]}" fi From 968c5c251d29f9d2883eb46b49c51818cb55e397 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Thu, 23 Apr 2026 15:48:24 +0200 Subject: [PATCH 18/29] Fix dataset and full activation checkpointing --- experiments/lm/launch_sweep.py | 8 +++++++- experiments/lm/train.py | 8 +++----- 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/experiments/lm/launch_sweep.py b/experiments/lm/launch_sweep.py index 5ca3159..98b086a 100644 --- a/experiments/lm/launch_sweep.py +++ b/experiments/lm/launch_sweep.py @@ -85,7 +85,7 @@ cd "${repo_path}" export OMP_NUM_THREADS=8 -export HF_HOME="$${HF_HOME:-/net/storage/pr3/plgrid/plggadlers/hf_cache}" +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}" @@ -107,6 +107,7 @@ --seed ${seed} \\ --batch-size ${batch_size} \\ --dataset ${dataset} \\ + ${dataset_path_arg} \\ --hf-assets-path ${hf_assets_path} \\ --output-dir ${output_dir} """) @@ -141,6 +142,8 @@ def main() -> None: p.add_argument("--runs-dir", required=True, help="Root directory for run outputs") p.add_argument("--dataset", required=True, help="HuggingFace dataset name (e.g. HuggingFaceFW/fineweb-edu)") + p.add_argument("--dataset-path", default=None, + help="Optional local path to dataset assets (otherwise streamed from HF)") 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") @@ -175,6 +178,8 @@ def main() -> None: 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), @@ -189,6 +194,7 @@ def main() -> None: batch_size=args.batch_size, gpus_per_node=args.gpus_per_node, dataset=args.dataset, + dataset_path_arg=dataset_path_arg, hf_assets_path=args.hf_assets_path, ) script_path = out_dir / "job.sh" diff --git a/experiments/lm/train.py b/experiments/lm/train.py index cc011d6..c815dfa 100644 --- a/experiments/lm/train.py +++ b/experiments/lm/train.py @@ -159,8 +159,7 @@ def build_trainer_config(args: argparse.Namespace) -> Trainer.Config: last_save_model_only=False, ), compile=CompileConfig(enable=not is_debug), - activation_checkpoint=ActivationCheckpointConfig(mode="selective"), - debug=DebugConfig(seed=args.seed), + activation_checkpoint=ActivationCheckpointConfig(mode="full"), validator=Validator.Config( enable=not is_debug, freq=500, @@ -171,7 +170,6 @@ def build_trainer_config(args: argparse.Namespace) -> Trainer.Config: def parse_args() -> argparse.Namespace: default_hf_path = os.path.join(os.path.dirname(__file__), "assets/hf/Llama-3.1-8B") - default_dataset_path = os.path.join(os.path.dirname(__file__), "assets/c4_test") p = argparse.ArgumentParser( description="MaxP LLaMA-3 pre-training", formatter_class=argparse.ArgumentDefaultsHelpFormatter, @@ -200,9 +198,9 @@ def parse_args() -> argparse.Namespace: help="Random seed") p.add_argument("--output-dir", default="./outputs", help="Directory to save checkpoints and logs") - p.add_argument("--dataset", default="c4_test", + p.add_argument("--dataset", default="fineweb-edu", help="HuggingFace dataset name or local path") - p.add_argument("--dataset-path", default=default_dataset_path, + p.add_argument("--dataset-path", default=None, help="Override dataset path (e.g. absolute path to c4_test on disk)") p.add_argument("--hf-assets-path", default=default_hf_path, help="Path to HF tokenizer assets (local copy)") From cd4af1362f796f0bf97d2e4e08b0dee7f2ea3139 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Thu, 23 Apr 2026 20:00:41 +0200 Subject: [PATCH 19/29] Fix torch recompile limit --- experiments/lm/run_debug.sh | 6 ++++-- experiments/lm/train.py | 1 + 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/experiments/lm/run_debug.sh b/experiments/lm/run_debug.sh index 0c3f569..0846adb 100755 --- a/experiments/lm/run_debug.sh +++ b/experiments/lm/run_debug.sh @@ -11,9 +11,9 @@ REPO="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" VENV="${REPO}/.venv" # Defaults — override with env vars or flags -STEPS="${STEPS:-20}" +STEPS="${STEPS:-200}" SCALE="${SCALE:-debug}" -METHOD="${METHOD:-mup-no}" +METHOD="${METHOD:-maxP}" GPUS="${GPUS:-1}" DATASET="${DATASET:-c4_test}" OUTPUT_DIR="${OUTPUT_DIR:-/tmp/maxp_debug}" @@ -67,6 +67,8 @@ TRAIN_ARGS=( --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}"} diff --git a/experiments/lm/train.py b/experiments/lm/train.py index c815dfa..f2dfad4 100644 --- a/experiments/lm/train.py +++ b/experiments/lm/train.py @@ -20,6 +20,7 @@ import os import torch +torch._dynamo.config.recompile_limit = 100 from torchtitan.components.checkpoint import CheckpointManager from torchtitan.components.lr_scheduler import LRSchedulersContainer From c1c252a20f1c59830c665bd7d3b32b37d95c2f13 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Thu, 23 Apr 2026 22:16:40 +0200 Subject: [PATCH 20/29] Update training config --- experiments/lm/fineweb.py | 8 ++++---- experiments/lm/launch_sweep.py | 17 +++++++++-------- experiments/lm/train.py | 10 ++++++---- 3 files changed, 19 insertions(+), 16 deletions(-) diff --git a/experiments/lm/fineweb.py b/experiments/lm/fineweb.py index 40bef25..0ecedf2 100644 --- a/experiments/lm/fineweb.py +++ b/experiments/lm/fineweb.py @@ -3,8 +3,8 @@ Import this module before building any Trainer.Config that uses --dataset fineweb-edu or --dataset fineweb-edu-10bt. -FineWeb-Edu is streamed directly from HuggingFace (or from a local cache -if HF_HOME is set). No pre-tokenization step is needed. +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. """ @@ -21,12 +21,12 @@ def _process_fineweb_text(sample: dict) -> str: def _load_fineweb_edu(dataset_path: str): - return load_dataset(dataset_path, split="train", streaming=True) + return load_dataset(dataset_path, split="train", streaming=False) def _load_fineweb_edu_10bt(dataset_path: str): # 10B-token deduplicated subset — faster to iterate for S1/S2 - return load_dataset(dataset_path, name="sample-10BT", split="train", streaming=True) + return load_dataset(dataset_path, name="sample-10BT", split="train", streaming=False) DATASETS["fineweb-edu"] = DatasetConfig( diff --git a/experiments/lm/launch_sweep.py b/experiments/lm/launch_sweep.py index 98b086a..9aaeafa 100644 --- a/experiments/lm/launch_sweep.py +++ b/experiments/lm/launch_sweep.py @@ -33,11 +33,11 @@ # --------------------------------------------------------------------------- SCALE_CONFIGS = { - "s1": {"wall": "02:00:00", "nodes": 1}, - "s2": {"wall": "02:00:00", "nodes": 1}, - "s3": {"wall": "06:00:00", "nodes": 1}, - "s4": {"wall": "24:00:00", "nodes": 1}, - "s5": {"wall": "48:00:00", "nodes": 1}, + "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}, } # Supported methods (currently implemented) @@ -137,8 +137,8 @@ def main() -> None: help="Override seed list (default: per-scale defaults)") p.add_argument("--batch-size", type=int, default=8, help="Local batch size per GPU") - p.add_argument("--gpus-per-node", type=int, default=1, - help="GPUs per SLURM node (also sets --nproc_per_node)") + 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", required=True, help="HuggingFace dataset name (e.g. HuggingFaceFW/fineweb-edu)") @@ -159,6 +159,7 @@ def main() -> None: 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 @@ -192,7 +193,7 @@ def main() -> None: method=method, lr=lr, batch_size=args.batch_size, - gpus_per_node=args.gpus_per_node, + gpus_per_node=gpus_per_node, dataset=args.dataset, dataset_path_arg=dataset_path_arg, hf_assets_path=args.hf_assets_path, diff --git a/experiments/lm/train.py b/experiments/lm/train.py index f2dfad4..a992e5e 100644 --- a/experiments/lm/train.py +++ b/experiments/lm/train.py @@ -156,8 +156,10 @@ def build_trainer_config(args: argparse.Namespace) -> Trainer.Config: enable_wandb=not is_debug, ), checkpoint=CheckpointManager.Config( - interval=500, + enable=True, + interval=5000, last_save_model_only=False, + keep_latest_k=1, ), compile=CompileConfig(enable=not is_debug), activation_checkpoint=ActivationCheckpointConfig(mode="full"), @@ -181,11 +183,11 @@ def parse_args() -> argparse.Namespace: help="maxP variant") p.add_argument("--lr", type=float, default=1e-3, help="LR prefactor") - p.add_argument("--alignment-warmup", type=int, default=10, + 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=100, + 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=32, + p.add_argument("--sample-size", type=int, default=128, 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)") From 41c9be98ed8e3316f4ba317c0435a75910ab50d3 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Thu, 23 Apr 2026 22:31:25 +0200 Subject: [PATCH 21/29] Async checkpointing --- experiments/lm/train.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/experiments/lm/train.py b/experiments/lm/train.py index a992e5e..d116f83 100644 --- a/experiments/lm/train.py +++ b/experiments/lm/train.py @@ -160,6 +160,7 @@ def build_trainer_config(args: argparse.Namespace) -> Trainer.Config: interval=5000, last_save_model_only=False, keep_latest_k=1, + async_mode="async", ), compile=CompileConfig(enable=not is_debug), activation_checkpoint=ActivationCheckpointConfig(mode="full"), @@ -187,7 +188,7 @@ def parse_args() -> argparse.Namespace: 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=128, + 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)") From 0b86bec283f299ed14d2f53515737e4dc261abf3 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Fri, 24 Apr 2026 09:41:25 +0200 Subject: [PATCH 22/29] Add fineweb 100BT dataset --- experiments/lm/fineweb.py | 15 +++++++++++++-- experiments/lm/launch_sweep.py | 17 +++++++++++++---- experiments/lm/run.sh | 2 -- experiments/lm/train.py | 2 +- 4 files changed, 27 insertions(+), 9 deletions(-) diff --git a/experiments/lm/fineweb.py b/experiments/lm/fineweb.py index 0ecedf2..4a6a04a 100644 --- a/experiments/lm/fineweb.py +++ b/experiments/lm/fineweb.py @@ -1,7 +1,7 @@ """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. +--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. @@ -25,10 +25,15 @@ def _load_fineweb_edu(dataset_path: str): def _load_fineweb_edu_10bt(dataset_path: str): - # 10B-token deduplicated subset — faster to iterate for S1/S2 + # 10B-token deduplicated subset for S1/S2/S3 return load_dataset(dataset_path, name="sample-10BT", split="train", streaming=False) +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=False) + + DATASETS["fineweb-edu"] = DatasetConfig( path="HuggingFaceFW/fineweb-edu", loader=_load_fineweb_edu, @@ -40,3 +45,9 @@ def _load_fineweb_edu_10bt(dataset_path: str): 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 index 9aaeafa..0f66232 100644 --- a/experiments/lm/launch_sweep.py +++ b/experiments/lm/launch_sweep.py @@ -40,7 +40,6 @@ "s5": {"wall": "48:00:00", "nodes": 1, "gpus": 4}, } -# Supported methods (currently implemented) SCALE_METHODS = { "s1": ["maxP", "mup-full", "mup-no"], "s2": ["maxP", "mup-full", "mup-no"], @@ -57,6 +56,15 @@ "s5": [1], } +SCALE_DATASETS = { + "s1": "fineweb-edu-10bt", + "s2": "fineweb-edu-10bt", + "s3": "fineweb-edu-10bt", + "s4": "fineweb-edu-100bt", + "s5": "fineweb-edu-100bt", +} + + # Full LR grid from experiments.md §4.1 ALL_LRS = [3e-4, 1e-3, 3e-3, 1e-2, 3e-2, 1e-1, 3e-1] @@ -140,8 +148,8 @@ def main() -> None: 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", required=True, - help="HuggingFace dataset name (e.g. HuggingFaceFW/fineweb-edu)") + p.add_argument("--dataset", default=None, + 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("--hf-assets-path", required=True, @@ -159,6 +167,7 @@ def main() -> None: 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] + dataset = args.dataset or SCALE_DATASETS[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") @@ -194,7 +203,7 @@ def main() -> None: lr=lr, batch_size=args.batch_size, gpus_per_node=gpus_per_node, - dataset=args.dataset, + dataset=dataset, dataset_path_arg=dataset_path_arg, hf_assets_path=args.hf_assets_path, ) diff --git a/experiments/lm/run.sh b/experiments/lm/run.sh index b2d912a..0c196f1 100644 --- a/experiments/lm/run.sh +++ b/experiments/lm/run.sh @@ -27,7 +27,6 @@ case "$SCALE" in esac RUNS_DIR="${RUNS_DIR:-/net/storage/pr3/plgrid/plggadlers/maxP/runs}" -DATASET="${DATASET:-fineweb-edu}" 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}" @@ -40,7 +39,6 @@ python experiments/lm/launch_sweep.py \ --lrs $LRS \ --seeds $SEEDS \ --runs-dir "$RUNS_DIR" \ - --dataset "$DATASET" \ --hf-assets-path "$HF_ASSETS_PATH" \ --venv-path "$VENV_PATH" \ --repo-path "$REPO_PATH" \ diff --git a/experiments/lm/train.py b/experiments/lm/train.py index d116f83..ddb3b50 100644 --- a/experiments/lm/train.py +++ b/experiments/lm/train.py @@ -202,7 +202,7 @@ def parse_args() -> argparse.Namespace: help="Random seed") p.add_argument("--output-dir", default="./outputs", help="Directory to save checkpoints and logs") - p.add_argument("--dataset", default="fineweb-edu", + 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)") From 18cefa464f288de97a7b6ef967f8d85ff3dc8ba8 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Fri, 24 Apr 2026 09:57:12 +0200 Subject: [PATCH 23/29] Fix checkpointing --- experiments/lm/train.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/experiments/lm/train.py b/experiments/lm/train.py index ddb3b50..e7f4784 100644 --- a/experiments/lm/train.py +++ b/experiments/lm/train.py @@ -159,7 +159,7 @@ def build_trainer_config(args: argparse.Namespace) -> Trainer.Config: enable=True, interval=5000, last_save_model_only=False, - keep_latest_k=1, + keep_latest_k=2, async_mode="async", ), compile=CompileConfig(enable=not is_debug), From 6508da8664892c5c4cfaf863107afe7721bd2038 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Fri, 24 Apr 2026 10:43:20 +0200 Subject: [PATCH 24/29] Prefetch dataset --- experiments/lm/launch_sweep.py | 10 +++++++++- experiments/lm/train.py | 6 ++++++ 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/experiments/lm/launch_sweep.py b/experiments/lm/launch_sweep.py index 0f66232..3d215dd 100644 --- a/experiments/lm/launch_sweep.py +++ b/experiments/lm/launch_sweep.py @@ -92,7 +92,7 @@ source "${venv_path}/bin/activate" cd "${repo_path}" -export OMP_NUM_THREADS=8 +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}" @@ -116,6 +116,8 @@ --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} """) @@ -152,6 +154,10 @@ def main() -> None: 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") @@ -205,6 +211,8 @@ def main() -> None: gpus_per_node=gpus_per_node, dataset=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" diff --git a/experiments/lm/train.py b/experiments/lm/train.py index e7f4784..1e8151e 100644 --- a/experiments/lm/train.py +++ b/experiments/lm/train.py @@ -149,6 +149,8 @@ def build_trainer_config(args: argparse.Namespace) -> Trainer.Config: 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, @@ -206,6 +208,10 @@ def parse_args() -> argparse.Namespace: 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=1, + help="DataLoader num_workers for prefetching") + p.add_argument("--prefetch-factor", type=int, default=1, + help="Batches prefetched per DataLoader worker") p.add_argument("--hf-assets-path", default=default_hf_path, help="Path to HF tokenizer assets (local copy)") return p.parse_args() From 9d80521a9966ec3225d577e3a77eb2d211e6a150 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Fri, 24 Apr 2026 12:25:49 +0200 Subject: [PATCH 25/29] Improve data loading and training config - Switch fineweb loaders to streaming=True for multi-worker support - Use fineweb-edu base dataset across all scales for fair comparison - Make --batch-size global (divided by WORLD_SIZE for local batch size) - Add --num-workers (default 8) and --prefetch-factor (default 4) args - Set OMP_NUM_THREADS=16, HF_DATASETS_OFFLINE=1 in SLURM template - Remove SCALE_DATASETS dict, default dataset to fineweb-edu Co-Authored-By: Claude Sonnet 4.6 --- experiments/lm/fineweb.py | 6 +++--- experiments/lm/launch_sweep.py | 22 ++++++---------------- experiments/lm/train.py | 23 ++++++++++++++++------- 3 files changed, 25 insertions(+), 26 deletions(-) diff --git a/experiments/lm/fineweb.py b/experiments/lm/fineweb.py index 4a6a04a..216512d 100644 --- a/experiments/lm/fineweb.py +++ b/experiments/lm/fineweb.py @@ -21,17 +21,17 @@ def _process_fineweb_text(sample: dict) -> str: def _load_fineweb_edu(dataset_path: str): - return load_dataset(dataset_path, split="train", streaming=False) + 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=False) + 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=False) + return load_dataset(dataset_path, name="sample-100BT", split="train", streaming=True) DATASETS["fineweb-edu"] = DatasetConfig( diff --git a/experiments/lm/launch_sweep.py b/experiments/lm/launch_sweep.py index 3d215dd..5f5f30e 100644 --- a/experiments/lm/launch_sweep.py +++ b/experiments/lm/launch_sweep.py @@ -35,8 +35,8 @@ 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}, + "s3": {"wall": "24:00:00", "nodes": 1, "gpus": 2}, + "s4": {"wall": "48:00:00", "nodes": 1, "gpus": 2}, "s5": {"wall": "48:00:00", "nodes": 1, "gpus": 4}, } @@ -56,15 +56,6 @@ "s5": [1], } -SCALE_DATASETS = { - "s1": "fineweb-edu-10bt", - "s2": "fineweb-edu-10bt", - "s3": "fineweb-edu-10bt", - "s4": "fineweb-edu-100bt", - "s5": "fineweb-edu-100bt", -} - - # Full LR grid from experiments.md §4.1 ALL_LRS = [3e-4, 1e-3, 3e-3, 1e-2, 3e-2, 1e-1, 3e-1] @@ -145,12 +136,12 @@ def main() -> 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=8, - help="Local batch size per GPU") + 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=None, + 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)") @@ -173,7 +164,6 @@ def main() -> None: 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] - dataset = args.dataset or SCALE_DATASETS[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") @@ -209,7 +199,7 @@ def main() -> None: lr=lr, batch_size=args.batch_size, gpus_per_node=gpus_per_node, - dataset=dataset, + dataset=args.dataset, dataset_path_arg=dataset_path_arg, num_workers=args.num_workers, prefetch_factor=args.prefetch_factor, diff --git a/experiments/lm/train.py b/experiments/lm/train.py index 1e8151e..9d5d0f4 100644 --- a/experiments/lm/train.py +++ b/experiments/lm/train.py @@ -98,18 +98,27 @@ def train_step(self, data_iterator): 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, args.batch_size, world_size) + 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, @@ -141,7 +150,7 @@ def build_trainer_config(args: argparse.Namespace) -> Trainer.Config: decay_ratio=0.1, ), training=TrainingConfig( - local_batch_size=2 if is_debug else args.batch_size, + local_batch_size=local_batch_size, seq_len=args.seq_len, steps=steps, dtype="bfloat16", @@ -194,12 +203,12 @@ def parse_args() -> argparse.Namespace: 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=2048, + 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=8, - help="Local batch size per GPU") + 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", @@ -208,9 +217,9 @@ def parse_args() -> argparse.Namespace: 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=1, + p.add_argument("--num-workers", type=int, default=8, help="DataLoader num_workers for prefetching") - p.add_argument("--prefetch-factor", type=int, default=1, + 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)") From bcc54d2b4bed3bcaf282d44c4fdc1fde9b32f2db Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Fri, 24 Apr 2026 14:23:51 +0200 Subject: [PATCH 26/29] Remove checkpointing --- experiments/lm/launch_sweep.py | 25 +++++-------------------- experiments/lm/run.sh | 1 - experiments/lm/train.py | 15 +-------------- 3 files changed, 6 insertions(+), 35 deletions(-) diff --git a/experiments/lm/launch_sweep.py b/experiments/lm/launch_sweep.py index 5f5f30e..ec39dee 100644 --- a/experiments/lm/launch_sweep.py +++ b/experiments/lm/launch_sweep.py @@ -13,10 +13,8 @@ --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] + [--dry-run] -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 @@ -36,7 +34,7 @@ "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": 2}, - "s4": {"wall": "48:00:00", "nodes": 1, "gpus": 2}, + "s4": {"wall": "48:00:00", "nodes": 1, "gpus": 4}, "s5": {"wall": "48:00:00", "nodes": 1, "gpus": 4}, } @@ -56,10 +54,9 @@ "s5": [1], } -# Full LR grid from experiments.md §4.1 ALL_LRS = [3e-4, 1e-3, 3e-3, 1e-2, 3e-2, 1e-1, 3e-1] -S5_SINGLE_LR = [1e-2] # transfer test: single best LR from S4 +S5_SINGLE_LR = [1e-2] # --------------------------------------------------------------------------- @@ -122,11 +119,6 @@ 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)) @@ -155,8 +147,6 @@ def main() -> None: 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 @@ -167,7 +157,7 @@ def main() -> None: 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 + submitted = 0 for method in methods: for lr in lrs: @@ -177,11 +167,6 @@ def main() -> None: ) 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 "" @@ -220,7 +205,7 @@ def main() -> None: submitted += 1 action = "would submit" if args.dry_run else "submitted" - print(f"\nDone: {action} {submitted} jobs, skipped {skipped} completed runs.") + print(f"\nDone: {action} {submitted} jobs.") if __name__ == "__main__": diff --git a/experiments/lm/run.sh b/experiments/lm/run.sh index 0c196f1..2b525a8 100644 --- a/experiments/lm/run.sh +++ b/experiments/lm/run.sh @@ -4,7 +4,6 @@ # Usage: # bash experiments/lm/run.sh s3 [launch_sweep.py flags …] # -# Per-scale defaults (seeds, LRs, methods) match experiments.md §4.1. # Override any of them by passing extra flags after the scale argument. # # Examples: diff --git a/experiments/lm/train.py b/experiments/lm/train.py index 9d5d0f4..d867935 100644 --- a/experiments/lm/train.py +++ b/experiments/lm/train.py @@ -22,17 +22,11 @@ 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, - DebugConfig, - TrainingConfig, -) +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 @@ -166,13 +160,6 @@ def build_trainer_config(args: argparse.Namespace) -> Trainer.Config: enable_tensorboard=not is_debug, enable_wandb=not is_debug, ), - checkpoint=CheckpointManager.Config( - enable=True, - interval=5000, - 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( From c3ae59d57f6c9aca4fd48d75c35c997c7e8c5d37 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Mon, 27 Apr 2026 10:42:42 +0200 Subject: [PATCH 27/29] Restore checkpointing, FSDP2-safe capture Re-enable async checkpointing via --checkpoint-interval. Resync optimizer param refs by FQN after no_grad capture forwards so FSDP2 reshard-induced Parameter swaps don't break get_optimizer_state_dict. Add use_training_activations flag to MaxPConverter config. --- experiments/lm/maxp_converter.py | 2 ++ experiments/lm/train.py | 18 ++++++++-- maxp/parametrization.py | 57 ++++++++++++++++++++++++++++++-- 3 files changed, 71 insertions(+), 6 deletions(-) diff --git a/experiments/lm/maxp_converter.py b/experiments/lm/maxp_converter.py index f2ae5b8..86d44b9 100644 --- a/experiments/lm/maxp_converter.py +++ b/experiments/lm/maxp_converter.py @@ -106,6 +106,7 @@ class Config(Configurable.Config): 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, @@ -128,6 +129,7 @@ def convert(self, model: nn.Module) -> None: 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": diff --git a/experiments/lm/train.py b/experiments/lm/train.py index d867935..6837518 100644 --- a/experiments/lm/train.py +++ b/experiments/lm/train.py @@ -22,6 +22,7 @@ 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 @@ -46,12 +47,14 @@ def train_step(self, data_iterator): 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() + model_to_opt = dict(zip(self.model_parts, self.optimizers.optimizers)) for model in needs_capture: model._maxp_sample_x = tokens - model._maxp_param.capture_initial(tokens) + model._maxp_param.capture_initial(tokens, optimizer=model_to_opt[model]) model._maxp_ready = True super().train_step(data_iterator) @@ -160,6 +163,13 @@ def build_trainer_config(args: argparse.Namespace) -> Trainer.Config: 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( @@ -178,9 +188,9 @@ 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-full", "mup-no"], default="maxP", help="maxP variant") - p.add_argument("--lr", type=float, default=1e-3, + 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)") @@ -210,6 +220,8 @@ def parse_args() -> argparse.Namespace: 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() diff --git a/maxp/parametrization.py b/maxp/parametrization.py index 46fc8e6..638e9ae 100644 --- a/maxp/parametrization.py +++ b/maxp/parametrization.py @@ -371,6 +371,51 @@ def _hook(mod, inp, out, _name=name): return captured + def _capture_and_resync( + self, + sample_input: torch.Tensor, + optimizer: torch.optim.Optimizer | None, + ) -> dict[str, torch.Tensor]: + """Capture activations, then resync optimizer param refs by FQN. + + A ``torch.no_grad()`` forward can trigger FSDP2 to reshard a unit + and replace ``nn.Parameter`` Python objects. The optimizer then + holds stale refs that don't match ``model.named_parameters()`` by + identity — which breaks ``get_optimizer_state_dict`` at checkpoint + time (``KeyError`` on integer param index). + + We snapshot ``{id(p): fqn}`` before the forward; after, any group + param whose old id maps to a known fqn but whose current tensor at + that fqn is a different Python object is rebound, and its entry + in ``optimizer.state`` migrated to the new key. + """ + if optimizer is None: + return self._capture_activations(sample_input) + + old_id_to_name = {id(p): name for name, p in self.model.named_parameters()} + captured = self._capture_activations(sample_input) + current_by_name = {name: p for name, p in self.model.named_parameters()} + + swaps: list[tuple[torch.nn.Parameter, torch.nn.Parameter]] = [] + for group in optimizer.param_groups: + new_params = [] + for p in group["params"]: + name = old_id_to_name.get(id(p)) + if name is None: + new_params.append(p) + continue + new_p = current_by_name.get(name, p) + if new_p is not p: + swaps.append((p, new_p)) + new_params.append(new_p) + group["params"] = new_params + + for old_p, new_p in swaps: + if old_p in optimizer.state: + optimizer.state[new_p] = optimizer.state.pop(old_p) + + return captured + def register_hooks(self) -> None: """Install persistent forward hooks to capture activations during training. @@ -408,7 +453,11 @@ def remove_hooks(self) -> None: self._persistent_hooks.clear() self._latest_activations.clear() - def capture_initial(self, sample_input: torch.Tensor) -> None: + def capture_initial( + self, + sample_input: torch.Tensor, + optimizer: torch.optim.Optimizer | None = None, + ) -> None: """Capture initial (z_0, w_0) for alignment measurement. Must be called before training starts (after ``__init__``). @@ -420,8 +469,10 @@ def capture_initial(self, sample_input: torch.Tensor) -> None: Args: sample_input: A batch of inputs to run through the model. Only the first ``sample_size`` samples are kept. + optimizer: If provided, optimizer param_groups + state are + resynced after the no_grad forward (FSDP2-safe). """ - captured = self._capture_activations(sample_input) + captured = self._capture_and_resync(sample_input, optimizer) for name, pm in self._pms: if pm.inner is not None and pm.weight is not None and name in captured: @@ -515,7 +566,7 @@ def step( raise ValueError( "sample_input is required when use_training_activations=False." ) - current = self._capture_activations(sample_input) + current = self._capture_and_resync(sample_input, optimizer) # 2. Compute alignment per PM, write back to PM from maxp.alignment import compute_alignment From a926a918b49c38fae5df98d8cc0cc637a2f7801d Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Mon, 27 Apr 2026 10:45:39 +0200 Subject: [PATCH 28/29] Restore --resume flag in launch_sweep Skip runs with existing checkpoints by default; --resume bypasses the skip to resubmit interrupted training. --- experiments/lm/launch_sweep.py | 20 +++++++++++++++++--- 1 file changed, 17 insertions(+), 3 deletions(-) diff --git a/experiments/lm/launch_sweep.py b/experiments/lm/launch_sweep.py index ec39dee..0690ee5 100644 --- a/experiments/lm/launch_sweep.py +++ b/experiments/lm/launch_sweep.py @@ -13,8 +13,10 @@ --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] + [--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 @@ -119,6 +121,11 @@ 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)) @@ -147,6 +154,8 @@ def main() -> None: 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 @@ -157,7 +166,7 @@ def main() -> None: 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 = 0 + submitted = skipped = 0 for method in methods: for lr in lrs: @@ -167,6 +176,11 @@ def main() -> None: ) 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 "" @@ -205,7 +219,7 @@ def main() -> None: submitted += 1 action = "would submit" if args.dry_run else "submitted" - print(f"\nDone: {action} {submitted} jobs.") + print(f"\nDone: {action} {submitted} jobs, skipped {skipped} completed runs.") if __name__ == "__main__": From c87506048cbbae994d512d1db795dbcb220a743b Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Tue, 28 Apr 2026 10:30:19 +0200 Subject: [PATCH 29/29] Fix FSDP2 stale params via parallelize_llama wrapper Force reshard_after_forward=always to prevent no_grad forwards from leaving layers unsharded; remove now-unnecessary _capture_and_resync. Co-Authored-By: Claude Sonnet 4.6 --- experiments/lm/launch_sweep.py | 2 +- experiments/lm/maxp_llama3.py | 11 ++++++- experiments/lm/train.py | 3 +- maxp/parametrization.py | 57 ++-------------------------------- 4 files changed, 15 insertions(+), 58 deletions(-) diff --git a/experiments/lm/launch_sweep.py b/experiments/lm/launch_sweep.py index 0690ee5..9f92424 100644 --- a/experiments/lm/launch_sweep.py +++ b/experiments/lm/launch_sweep.py @@ -35,7 +35,7 @@ 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": 2}, + "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}, } diff --git a/experiments/lm/maxp_llama3.py b/experiments/lm/maxp_llama3.py index 4664f66..e53e5d9 100644 --- a/experiments/lm/maxp_llama3.py +++ b/experiments/lm/maxp_llama3.py @@ -13,7 +13,7 @@ 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 +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 @@ -120,6 +120,15 @@ def _make_model_config( ) +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) diff --git a/experiments/lm/train.py b/experiments/lm/train.py index 6837518..5a0e060 100644 --- a/experiments/lm/train.py +++ b/experiments/lm/train.py @@ -51,10 +51,9 @@ def train_step(self, data_iterator): if needs_capture: batch = next(data_iterator) tokens = batch[0]["input"].detach() - model_to_opt = dict(zip(self.model_parts, self.optimizers.optimizers)) for model in needs_capture: model._maxp_sample_x = tokens - model._maxp_param.capture_initial(tokens, optimizer=model_to_opt[model]) + model._maxp_param.capture_initial(tokens) model._maxp_ready = True super().train_step(data_iterator) diff --git a/maxp/parametrization.py b/maxp/parametrization.py index 638e9ae..46fc8e6 100644 --- a/maxp/parametrization.py +++ b/maxp/parametrization.py @@ -371,51 +371,6 @@ def _hook(mod, inp, out, _name=name): return captured - def _capture_and_resync( - self, - sample_input: torch.Tensor, - optimizer: torch.optim.Optimizer | None, - ) -> dict[str, torch.Tensor]: - """Capture activations, then resync optimizer param refs by FQN. - - A ``torch.no_grad()`` forward can trigger FSDP2 to reshard a unit - and replace ``nn.Parameter`` Python objects. The optimizer then - holds stale refs that don't match ``model.named_parameters()`` by - identity — which breaks ``get_optimizer_state_dict`` at checkpoint - time (``KeyError`` on integer param index). - - We snapshot ``{id(p): fqn}`` before the forward; after, any group - param whose old id maps to a known fqn but whose current tensor at - that fqn is a different Python object is rebound, and its entry - in ``optimizer.state`` migrated to the new key. - """ - if optimizer is None: - return self._capture_activations(sample_input) - - old_id_to_name = {id(p): name for name, p in self.model.named_parameters()} - captured = self._capture_activations(sample_input) - current_by_name = {name: p for name, p in self.model.named_parameters()} - - swaps: list[tuple[torch.nn.Parameter, torch.nn.Parameter]] = [] - for group in optimizer.param_groups: - new_params = [] - for p in group["params"]: - name = old_id_to_name.get(id(p)) - if name is None: - new_params.append(p) - continue - new_p = current_by_name.get(name, p) - if new_p is not p: - swaps.append((p, new_p)) - new_params.append(new_p) - group["params"] = new_params - - for old_p, new_p in swaps: - if old_p in optimizer.state: - optimizer.state[new_p] = optimizer.state.pop(old_p) - - return captured - def register_hooks(self) -> None: """Install persistent forward hooks to capture activations during training. @@ -453,11 +408,7 @@ def remove_hooks(self) -> None: self._persistent_hooks.clear() self._latest_activations.clear() - def capture_initial( - self, - sample_input: torch.Tensor, - optimizer: torch.optim.Optimizer | None = None, - ) -> None: + def capture_initial(self, sample_input: torch.Tensor) -> None: """Capture initial (z_0, w_0) for alignment measurement. Must be called before training starts (after ``__init__``). @@ -469,10 +420,8 @@ def capture_initial( Args: sample_input: A batch of inputs to run through the model. Only the first ``sample_size`` samples are kept. - optimizer: If provided, optimizer param_groups + state are - resynced after the no_grad forward (FSDP2-safe). """ - captured = self._capture_and_resync(sample_input, optimizer) + captured = self._capture_activations(sample_input) for name, pm in self._pms: if pm.inner is not None and pm.weight is not None and name in captured: @@ -566,7 +515,7 @@ def step( raise ValueError( "sample_input is required when use_training_activations=False." ) - current = self._capture_and_resync(sample_input, optimizer) + current = self._capture_activations(sample_input) # 2. Compute alignment per PM, write back to PM from maxp.alignment import compute_alignment