diff --git a/.gitignore b/.gitignore index b2b19a4..d2460a8 100644 --- a/.gitignore +++ b/.gitignore @@ -211,6 +211,9 @@ jobs outputs assets runs +figs +analysis +wandb # Generated plots *.png diff --git a/README.md b/README.md index bc24074..97d67e0 100644 --- a/README.md +++ b/README.md @@ -77,15 +77,6 @@ maxp/ ├── dag.py # DAG builder — traces PM-to-PM data flow ├── trace.py # Operation tracer — records matmul-like ops └── diagnose.py # Coord-check diagnostics — sweep widths, plot scaling - -experiments/lm/ # LLaMA-3 pre-training experiments (torchtitan) -├── train.py # Main training entry point (MaxPTrainer) -├── maxp_llama3.py # LLaMA-3 scale configs (debug, s1–s5) + compute_steps -├── maxp_converter.py # Post-optimizer-build hook; wires up Parametrization -├── launch_sweep.py # Generate and submit SLURM jobs for full LR sweep -├── run.sh # Submit a sweep for one scale -├── run_debug.sh # Quick single-GPU debug run -└── coord_check.py # Coord check for the parametrized LLaMA-3 model ``` **Phase 1** (static, at init): discover PMs → reinit weights → solve LP → build param groups. @@ -143,6 +134,11 @@ param = Parametrization( solve_interval=1, # re-solve every N steps sample_size=32, # max batch size for alignment measurement c_ema=0.0, # EMA smoothing for c values (0 = instant) + alignment_ema=0.0, # EMA smoothing for measured alignment values + resample_w0=False, # re-sample w0 snapshot each solve + use_training_activations=False, # use activations from training forward pass + solver=None, # custom PuLP solver (default: CBC) + warm_start=False, # warm-start LP from previous c solution ) # Param groups for optimizer @@ -270,8 +266,8 @@ def make_input(width): return torch.randint(0, vocab_size, (4, seq_len)) all_ops, affected, act_stats = diagnose_axis(make_model, make_input, widths) -print_axis(all_ops, affected, act_stats, widths) -plot_axis(all_ops, affected, act_stats, widths, path="coord_check.png") +print_axis("width", all_ops, affected, act_stats, widths) +plot_axis("width", all_ops, affected, act_stats, widths, filename="coord_check.png") ``` `print_axis` shows how the RMS of each op's output scales with width. For a correctly parametrized model, activations should be roughly **constant** (slope ≈ 0 in log-log) at init and remain stable after a few training steps. @@ -290,6 +286,21 @@ The `experiments/lm/` directory contains a full LLaMA-3 pre-training pipeline bu Training steps are computed automatically as `20 × non-embed params / tokens-per-step`. +### Files + +``` +experiments/lm/ +├── train.py # Main training entry point (MaxPTrainer) +├── maxp_llama3.py # LLaMA-3 scale configs (debug, s1–s5) + compute_steps +├── maxp_converter.py # Post-optimizer-build hook; wires up Parametrization +├── launch_sweep.py # Generate and submit SLURM jobs for full LR sweep +├── fineweb.py # FineWeb-Edu dataset registration for torchtitan +├── download_hf_assets.py # HuggingFace asset downloader (from torchtitan) +├── run.sh # Submit a sweep for one scale +├── run_debug.sh # Quick single-GPU debug run +└── coord_check.py # Coord check for the parametrized LLaMA-3 model +``` + ### Debug run (single GPU) ```bash @@ -305,6 +316,39 @@ bash experiments/lm/run_debug.sh \ bash experiments/lm/run.sh s3 # submits 7 LRs × 3 methods × 2 seeds = 42 jobs ``` +## Vision Experiments + +The `experiments/vision/` directory contains ViT and MLP pre-training experiments using [timm](https://github.com/huggingface/pytorch-image-models) with streaming HuggingFace datasets. It trains four ViT scales and four MLP scales with the same three methods as the LLM experiments. + +### Files + +``` +experiments/vision/ +├── train.py # Main training entry point (single-GPU, streaming HF data) +├── maxp_timm.py # timm model registry + automatic ParametrizedModule wrappers +├── hf_vision_data.py # HF streaming train/val pipeline + transforms +├── utils.py # Shared helpers (LR logging, arg parsing, checkpointing) +├── launch_sweep.py # Generate and submit SLURM jobs for full LR sweep +├── run.sh # Submit a sweep for one scale +├── run_debug.sh # Quick tiny-model debug run (CPU/GPU) +├── coord_check_vit.py # Coord-style diagnostic for ViT models +└── coord_check_mlp.py # Coord-style diagnostic for MLP models +``` + +Available scales: `debug`, `vit-s`, `vit-b`, `vit-l`, `mlp-s`, `mlp-m`, `mlp-b`, `mlp-l`. + +### Debug run (CPU/GPU) + +```bash +bash experiments/vision/run_debug.sh --steps 20 +``` + +### SLURM sweep + +```bash +bash experiments/vision/run.sh vit-s # submits 9 LRs × 3 methods × 3 seeds = 81 jobs +``` + ## Running Tests ```bash diff --git a/experiments/vision/coord_check_mlp.py b/experiments/vision/coord_check_mlp.py new file mode 100644 index 0000000..43e6404 --- /dev/null +++ b/experiments/vision/coord_check_mlp.py @@ -0,0 +1,103 @@ +#!/usr/bin/env python3 +"""Coord-style diagnostic for ScalableMLP. + +Width axis is `hidden` — the dimension that scales with model size. +""" + +from __future__ import annotations + +import argparse + +import torch +import torch.nn.functional as F + +from maxp import Parametrization, diagnose_axis, plot_axis, print_axis + +from maxp_timm import ScalableMLP, _install_mlp_wrappers + + +WIDTHS = [64, 128, 256, 512, 1024] +IMAGE_SIZE = 32 +NUM_CLASSES = 1000 +DEPTH = 6 + + +def _make_model(hidden: int, parametrized: bool): + model = ScalableMLP( + hidden=hidden, + depth=DEPTH, + num_classes=NUM_CLASSES, + dropout=0.0, + image_size=IMAGE_SIZE, + patch_size=4, + ) + if not parametrized: + return model, None + _install_mlp_wrappers(model) + sample_input = torch.randn(1, 3, IMAGE_SIZE, IMAGE_SIZE) + param = Parametrization( + model, + lr_prefactor=1e-3, + optimizer_type="adam", + alignment="full", + sample_input=sample_input, + ) + return model, param.param_groups + + +def _make_input(hidden: int) -> torch.Tensor: + return torch.randn(8, 3, IMAGE_SIZE, IMAGE_SIZE) + + +def _make_train_step(model, param_groups): + opt = ( + torch.optim.AdamW(param_groups) + if param_groups is not None + else torch.optim.AdamW(model.parameters(), lr=1e-3) + ) + x_fixed = torch.randn(8, 3, IMAGE_SIZE, IMAGE_SIZE) + y_fixed = torch.randint(0, NUM_CLASSES, (8,)) + + def step(model, step_idx): + del step_idx + logits = model(x_fixed) + loss = F.cross_entropy(logits, y_fixed) + loss.backward() + opt.step() + opt.zero_grad() + + return step + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Coord check for ScalableMLP (continuous hidden width axis)", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument("--parametrized", action="store_true") + parser.add_argument("--steps", type=int, default=10) + parser.add_argument("--seeds", type=int, default=5) + parser.add_argument("--plot", action="store_true") + args = parser.parse_args() + + variant = "parametrized" if args.parametrized else "plain" + print(f"ScalableMLP coord check ({variant}) — widths={WIDTHS}") + + ops, affected, act_stats = diagnose_axis( + make_model_fn=lambda w: _make_model(w, args.parametrized), + make_input_fn=_make_input, + widths=WIDTHS, + n_steps=args.steps, + n_seeds=args.seeds, + train_step_fn=_make_train_step, + ) + print_axis("hidden", ops, affected, act_stats, WIDTHS) + + if args.plot: + out_name = f"coord_check_mlp_{variant}.png" + plot_axis("hidden", ops, affected, act_stats, WIDTHS, out_name, plot_every=1) + print(f"Saved {out_name}") + + +if __name__ == "__main__": + main() diff --git a/experiments/vision/coord_check_vit.py b/experiments/vision/coord_check_vit.py new file mode 100644 index 0000000..178bdf2 --- /dev/null +++ b/experiments/vision/coord_check_vit.py @@ -0,0 +1,111 @@ +#!/usr/bin/env python3 +"""Coord-style diagnostic for wrapped timm ViT models.""" + +from __future__ import annotations + +import argparse + +import torch +import torch.nn.functional as F + +from maxp import Parametrization, diagnose_axis, plot_axis, print_axis + +from maxp_timm import SCALE_CONFIGS, create_model, install_pm_wrappers + + +WIDTH_TO_SCALE = { + 192: "debug", # vit_tiny + 384: "vit-s", # vit_small + 768: "vit-b", # vit_base +} + + +def _make_model(width: int, parametrized: bool): + scale = WIDTH_TO_SCALE[width] + cfg = SCALE_CONFIGS[scale] + model = create_model( + scale=scale, + num_classes=1000, + image_size=cfg.image_size, + ) + if not parametrized: + return model, None + install_pm_wrappers(model) + sample_input = torch.randn(1, 3, cfg.image_size, cfg.image_size) + param = Parametrization( + model, + lr_prefactor=1e-3, + optimizer_type="adam", + alignment="full", + sample_input=sample_input, + ) + return model, param.param_groups + + +def _make_input(width: int) -> torch.Tensor: + scale = WIDTH_TO_SCALE[width] + cfg = SCALE_CONFIGS[scale] + return torch.randn(32, 3, cfg.image_size, cfg.image_size) + + +def _make_train_step(model, param_groups): + if param_groups is not None: + opt = torch.optim.AdamW(param_groups) + else: + opt = torch.optim.AdamW(model.parameters(), lr=1e-2) + + x_fixed = torch.randn(32, 3, 224, 224) + y_fixed = torch.randint(0, 1000, (32,)) + + def step(model, step_idx): + del step_idx + logits = model(x_fixed) + loss = F.cross_entropy(logits, y_fixed) + loss.backward() + opt.step() + opt.zero_grad() + + return step + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Coord-style diagnostic for timm ViT wrappers", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument("--parametrized", action="store_true") + parser.add_argument("--steps", type=int, default=5) + parser.add_argument("--seeds", type=int, default=2) + parser.add_argument("--plot", action="store_true") + args = parser.parse_args() + + widths = sorted(WIDTH_TO_SCALE) + variant = "parametrized" if args.parametrized else "plain" + print(f"ViT coord check ({variant}) — widths={widths}") + + ops, affected, act_stats = diagnose_axis( + make_model_fn=lambda w: _make_model(w, args.parametrized), + make_input_fn=_make_input, + widths=widths, + n_steps=args.steps, + n_seeds=args.seeds, + train_step_fn=_make_train_step, + ) + print_axis("embed_dim", ops, affected, act_stats, widths) + + if args.plot: + out_name = f"coord_check_vit_{variant}.png" + plot_axis( + "embed_dim", + ops, + affected, + act_stats, + widths, + out_name, + plot_every=1, + ) + print(f"Saved {out_name}") + + +if __name__ == "__main__": + main() diff --git a/experiments/vision/hf_vision_data.py b/experiments/vision/hf_vision_data.py new file mode 100644 index 0000000..718f679 --- /dev/null +++ b/experiments/vision/hf_vision_data.py @@ -0,0 +1,127 @@ +"""HF non-streaming data pipeline for vision experiments.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Callable + +import torch +from datasets import load_dataset +from torch.utils.data import DataLoader, Dataset + + +@dataclass(frozen=True, slots=True) +class DatasetPreset: + dataset_id: str + train_split: str + val_split: str | None + num_classes: int + train_samples: int + val_samples: int | None + image_key: str + label_key: str + + +DATASET_CONFIGS: dict[str, DatasetPreset] = { + "imagenet12k": DatasetPreset( + dataset_id="timm/imagenet-12k-wds", + train_split="train", + val_split="validation", + num_classes=11821, + train_samples=12_129_687, + val_samples=472_840, + image_key="jpg", + label_key="cls", + ), + "beans": DatasetPreset( + dataset_id="AI-Lab-Makerere/beans", + train_split="train", + val_split="validation", + num_classes=3, + train_samples=1034, + val_samples=133, + image_key="image", + label_key="labels", + ), +} + + +class HFVisionDataset(Dataset): + """Map-style wrapper around a downloaded HF Dataset.""" + + def __init__(self, hf_dataset, *, preset: DatasetPreset, transform: Callable) -> None: + super().__init__() + self._hf = hf_dataset + self._preset = preset + self._transform = transform + + def __len__(self) -> int: + return len(self._hf) + + def __getitem__(self, idx: int): + sample = self._hf[idx] + image = sample[self._preset.image_key].convert("RGB") + label = sample[self._preset.label_key] + x = self._transform(image) + y = torch.tensor(label, dtype=torch.long) + return x, y + + +def make_loader( + ds: Dataset, + *, + batch_size: int, + num_workers: int, + is_train: bool, + prefetch_factor: int, +) -> DataLoader: + kwargs: dict = { + "dataset": ds, + "batch_size": batch_size, + "num_workers": num_workers, + "pin_memory": torch.cuda.is_available(), + "drop_last": is_train, + "persistent_workers": (num_workers > 0), + } + if num_workers > 0: + kwargs["prefetch_factor"] = prefetch_factor + return DataLoader(**kwargs) + + +def build_dataloaders( + *, + dataset_name: str, + batch_size: int, + num_workers: int, + prefetch_factor: int = 2, + train_transform=None, + eval_transform=None, +) -> tuple[DataLoader, DataLoader | None, DatasetPreset]: + if train_transform is None or eval_transform is None: + raise ValueError( + "build_dataloaders requires train/eval transforms (pass timm create_transform outputs)." + ) + + preset = DATASET_CONFIGS.get(dataset_name) + + train_hf = load_dataset(preset.dataset_id, split=preset.train_split, streaming=False) + train_ds = HFVisionDataset(train_hf, preset=preset, transform=train_transform) + train_loader = make_loader( + train_ds, + batch_size=batch_size, + num_workers=num_workers, + is_train=True, + prefetch_factor=prefetch_factor, + ) + + val_hf = load_dataset(preset.dataset_id, split=preset.val_split, streaming=False) + val_ds = HFVisionDataset(val_hf, preset=preset, transform=eval_transform) + val_loader = make_loader( + val_ds, + batch_size=batch_size, + num_workers=num_workers, + is_train=False, + prefetch_factor=prefetch_factor, + ) + + return train_loader, val_loader, preset diff --git a/experiments/vision/launch_sweep.py b/experiments/vision/launch_sweep.py new file mode 100755 index 0000000..a0dcfe8 --- /dev/null +++ b/experiments/vision/launch_sweep.py @@ -0,0 +1,230 @@ +#!/usr/bin/env python3 +"""Generate and submit SLURM jobs for vision sweeps.""" + +from __future__ import annotations + +import argparse +import subprocess +from datetime import date +from pathlib import Path +from string import Template + + +ALL_LRS = [3e-4, 1e-3, 3e-3, 1e-2, 3e-2, 1e-1, 3e-1] +TRANSFER_LR = [1e-2] + +SCALE_CONFIGS = { + "debug": { + "wall": "01:00:00", + "gpus": 1, + "methods": ["maxP"], + "lrs": [3e-3], + "seeds": [1], + "image_size": 224, + "drop_path_rate": 0.0, + }, + "vit-s": { + "wall": "24:00:00", + "gpus": 1, + "methods": ["maxP", "mup-full", "mup-no"], + "lrs": ALL_LRS, + "seeds": [1, 2, 3], + "image_size": 224, + "drop_path_rate": 0.1, + }, + "vit-b": { + "wall": "36:00:00", + "gpus": 1, + "methods": ["maxP", "mup-full", "mup-no"], + "lrs": TRANSFER_LR, + "seeds": [1], + "image_size": 224, + "drop_path_rate": 0.2, + }, + "vit-l": { + "wall": "48:00:00", + "gpus": 1, + "methods": ["maxP", "mup-full", "mup-no"], + "lrs": TRANSFER_LR, + "seeds": [1], + "image_size": 224, + "drop_path_rate": 0.4, + }, + "mlp-s": { + "wall": "6:00:00", + "gpus": 1, + "methods": ["maxP", "mup-full", "mup-no"], + "lrs": ALL_LRS, + "seeds": [1, 2, 3], + "image_size": 224, + "drop_path_rate": 0.1, + }, + "mlp-m": { + "wall": "12:00:00", + "gpus": 1, + "methods": ["maxP", "mup-full", "mup-no"], + "lrs": ALL_LRS, + "seeds": [1], + "image_size": 224, + "drop_path_rate": 0.2, + }, + "mlp-b": { + "wall": "12:00:00", + "gpus": 1, + "methods": ["maxP", "mup-full", "mup-no"], + "lrs": TRANSFER_LR, + "seeds": [1], + "image_size": 224, + "drop_path_rate": 0.3, + }, + "mlp-l": { + "wall": "24:00:00", + "gpus": 1, + "methods": ["maxP", "mup-full", "mup-no"], + "lrs": TRANSFER_LR, + "seeds": [1], + "image_size": 224, + "drop_path_rate": 0.4, + }, +} + + +SLURM_TEMPLATE = Template( + """\ +#!/bin/bash -l +#SBATCH --job-name maxp_${scale}_${method_tag}_lr${lr_tag}_s${seed} +#SBATCH --nodes 1 +#SBATCH --cpus-per-gpu 72 +#SBATCH --mem-per-gpu 118GB +#SBATCH --time ${wall_time} +#SBATCH --account plgadlers-gpu-gh200 +#SBATCH --partition plgrid-gpu-gh200 +#SBATCH --gres gpu:1 +#SBATCH --output ${output_dir}/maxp_${scale}_${method_tag}_lr${lr_tag}_s${seed}.out +#SBATCH --error ${output_dir}/maxp_${scale}_${method_tag}_lr${lr_tag}_s${seed}.err + +module add ML-bundle/25.10 +source "${venv_path}/bin/activate" +cd "${repo_path}" + +export OMP_NUM_THREADS=16 +export HF_HOME="$${HF_HOME:-/net/scratch/hscra/plgrid/plgmwojnar/hf}" +export WANDB_PROJECT="$${WANDB_PROJECT:-maxP-vision}" +export WANDB_RUN_NAME="maxp_${scale}_${method_tag}_lr${lr_tag}_s${seed}" + +python experiments/vision/train.py \\ + --scale ${scale} \\ + --method ${method} \\ + --lr ${lr} \\ + --seed ${seed} \\ + --dataset ${dataset} \\ + --epochs ${epochs} \\ + --batch-size ${batch_size} \\ + --num-workers ${num_workers} \\ + --val-steps ${val_steps} \\ + --output-dir ${output_dir} \\ + ${extra_train_args} +""" +) + + +def _lr_tag(lr: float) -> str: + return f"{lr:.0e}".replace("+", "") + + +def _method_tag(method: str) -> str: + return method.replace("-", "") + + +def _has_checkpoint(out_dir: Path) -> bool: + ckpt_dir = out_dir / "checkpoint" + return ckpt_dir.is_dir() and any(ckpt_dir.iterdir()) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Launch maxP timm vision sweeps via SLURM", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument("--scale", choices=list(SCALE_CONFIGS), required=True) + parser.add_argument("--methods", nargs="+", default=None) + parser.add_argument("--lrs", type=float, nargs="+", default=None) + parser.add_argument("--seeds", type=int, nargs="+", default=None) + + parser.add_argument("--runs-dir", required=True) + parser.add_argument("--venv-path", required=True) + parser.add_argument("--repo-path", required=True) + + parser.add_argument("--dataset", default="imagenet12k") + parser.add_argument("--epochs", type=int, default=1) + parser.add_argument("--batch-size", type=int, default=128) + parser.add_argument("--num-workers", type=int, default=8) + parser.add_argument("--val-steps", type=int, default=500) + + parser.add_argument("--extra-train-args", default="") + parser.add_argument("--dry-run", action="store_true") + parser.add_argument("--resume", action="store_true") + return parser.parse_args() + + +def main() -> None: + args = parse_args() + sc = SCALE_CONFIGS[args.scale] + methods = args.methods or sc["methods"] + lrs = args.lrs or sc["lrs"] + seeds = args.seeds or sc["seeds"] + + today = date.today().strftime("%Y-%m-%d") + submitted = skipped = 0 + + for method in methods: + for lr in lrs: + for seed in seeds: + run_name = f"{today}_{args.scale}_{_method_tag(method)}_lr{_lr_tag(lr)}_s{seed}" + out_dir = Path(args.runs_dir) / run_name + + if not args.resume and _has_checkpoint(out_dir): + print(f"[skip] {run_name}") + skipped += 1 + continue + + out_dir.mkdir(parents=True, exist_ok=True) + + script = SLURM_TEMPLATE.substitute( + scale=args.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, + dataset=args.dataset, + epochs=args.epochs, + batch_size=args.batch_size, + num_workers=args.num_workers, + val_steps=args.val_steps, + extra_train_args=args.extra_train_args, + ) + script_path = out_dir / "job.sh" + script_path.write_text(script) + + if args.dry_run: + print(f"[dry] sbatch {script_path}") + else: + result = subprocess.run( + ["sbatch", str(script_path)], + capture_output=True, + text=True, + ) + print(f"[submit] {run_name} → {result.stdout.strip()}") + 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/vision/maxp_timm.py b/experiments/vision/maxp_timm.py new file mode 100644 index 0000000..c460963 --- /dev/null +++ b/experiments/vision/maxp_timm.py @@ -0,0 +1,303 @@ +"""timm model registry and ParametrizedModule wrappers for vision experiments.""" + +from __future__ import annotations +from dataclasses import dataclass +from typing import Optional + +import timm +import torch +import torch.nn as nn +import torch.nn.functional as F +from timm.layers import Attention, PatchEmbed, maybe_add_mask, resolve_self_attn_mask +from timm.layers.format import Format, nchw_to + +from maxp import ParametrizedModule + + +@dataclass(frozen=True, slots=True) +class ScaleConfig: + model_name: str + family: str + image_size: int + drop_path_rate: float = 0.0 + hidden: int = 0 + depth: int = 0 + dropout: float = 0.0 + + +SCALE_CONFIGS: dict[str, ScaleConfig] = { + "debug": ScaleConfig( + model_name="vit_tiny_patch16_224", family="vit", image_size=224, drop_path_rate=0.0 + ), + "vit-s": ScaleConfig( + model_name="vit_small_patch16_224", family="vit", image_size=224, drop_path_rate=0.1 + ), + "vit-b": ScaleConfig( + model_name="vit_base_patch16_224", family="vit", image_size=224, drop_path_rate=0.2 + ), + "vit-l": ScaleConfig( + model_name="vit_large_patch16_224", family="vit", image_size=224, drop_path_rate=0.4 + ), + "mlp-s": ScaleConfig( + model_name="mlp", family="mlp", image_size=224, hidden=256, depth=4, dropout=0.0 + ), + "mlp-m": ScaleConfig( + model_name="mlp", family="mlp", image_size=224, hidden=512, depth=6, dropout=0.1 + ), + "mlp-b": ScaleConfig( + model_name="mlp", family="mlp", image_size=224, hidden=1024, depth=8, dropout=0.2 + ), + "mlp-l": ScaleConfig( + model_name="mlp", family="mlp", image_size=224, hidden=2048, depth=8, dropout=0.3 + ), +} + + +class MLPBlock(nn.Module): + def __init__(self, hidden: int, dropout: float) -> None: + super().__init__() + self.norm = nn.LayerNorm(hidden) + self.linear = nn.Linear(hidden, hidden) + self.act = nn.GELU() + self.drop = nn.Dropout(dropout) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.drop(self.act(self.linear(self.norm(x)))) + + +class Patchify(nn.Module): + def __init__(self, patch_size: tuple[int, int]): + super().__init__() + self.patch_size = patch_size + + def forward(self, x: torch.Tensor) -> torch.Tensor: + B, C, H, W = x.shape + p_h, p_w = self.patch_size + assert H % p_h == 0 and W % p_w == 0, "Image dimensions must be divisible by patch size." + x = x.reshape(B, C, H // p_h, p_h, W // p_w, p_w) + x = x.permute(0, 2, 4, 3, 5, 1).reshape(B, (H // p_h) * (W // p_w), C * p_h * p_w) + return x + + +class ScalableMLP(nn.Module): + def __init__( + self, + hidden: int, + depth: int, + num_classes: int, + dropout: float = 0.0, + image_size: int = 224, + patch_size: int = 16, + ) -> None: + super().__init__() + self.patch_size = (patch_size, patch_size) + in_features = 3 * patch_size * patch_size + + self.patchify = Patchify(self.patch_size) + self.embed = nn.Linear(in_features, hidden) + self.blocks = nn.ModuleList([MLPBlock(hidden, dropout) for _ in range(depth)]) + self.norm = nn.LayerNorm(hidden) + self.head = nn.Linear(hidden, num_classes) + + def set_grad_checkpointing(self, enable: bool = False) -> None: + pass + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.patchify(x) + x = self.embed(x) + x = x.mean(dim=1) + + for block in self.blocks: + x = block(x) + + x = self.norm(x) + return self.head(x) + + +class _ScaledAttention(Attention): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + d_model = kwargs.get("dim", args[0] if len(args) > 0 else None) + num_heads = kwargs.get("num_heads", args[1] if len(args) > 1 else None) + head_dim = kwargs.get("attn_head_dim", args[2] if len(args) > 2 else d_model // num_heads) + + self.scale = 1.0 / head_dim + self.qk_score = ParametrizedModule( + lambda q, k: (q, k), width_dim=head_dim, + layer_type="readout", scale_output=False + ) + + def forward( + self, + x: torch.Tensor, + attn_mask: Optional[torch.Tensor] = None, + is_causal: bool = False, + ) -> torch.Tensor: + B, N, C = x.shape + qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4) + q, k, v = qkv.unbind(0) + q, k = self.qk_score(q, k) + q, k = self.q_norm(q), self.k_norm(k) + + if self.fused_attn: + x = F.scaled_dot_product_attention( + q, k, v, + attn_mask=attn_mask, + dropout_p=self.attn_drop.p if self.training else 0., + is_causal=is_causal, + scale=self.scale, + ) + else: + q = q * self.scale + attn = q @ k.transpose(-2, -1) + attn_bias = resolve_self_attn_mask(N, attn, attn_mask, is_causal) + attn = maybe_add_mask(attn, attn_bias) + attn = attn.softmax(dim=-1) + attn = self.attn_drop(attn) + x = attn @ v + + x = x.transpose(1, 2).reshape(B, N, self.attn_dim) + x = self.norm(x) + x = self.proj(x) + x = self.proj_drop(x) + return x + + +class _LinearPatchEmbed(PatchEmbed): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + patch_dim = self.proj.in_channels * self.patch_size[0] * self.patch_size[1] + self.proj = nn.Linear(patch_dim, self.proj.out_channels) + self.patchify = Patchify(self.patch_size) + + def forward(self, x): + _, _, H, W = x.shape + + if self.dynamic_img_pad: + pad_h = (self.patch_size[0] - H % self.patch_size[0]) % self.patch_size[0] + pad_w = (self.patch_size[1] - W % self.patch_size[1]) % self.patch_size[1] + x = F.pad(x, (0, pad_w, 0, pad_h)) + + x = self.patchify(x) + x = self.proj(x) + x = x.transpose(1, 2).reshape(x.size(0), -1, H // self.patch_size[0], W // self.patch_size[1]) + + if self.flatten: + x = x.flatten(2).transpose(1, 2) # NCHW -> NLC + elif self.output_fmt != Format.NCHW: + x = nchw_to(x, self.output_fmt) + + x = self.norm(x) + return x + + +def create_model( + *, + scale: str, + num_classes: int, + image_size: int | None = None, +) -> nn.Module: + cfg = SCALE_CONFIGS.get(scale) + img_size = image_size if image_size is not None else cfg.image_size + + if cfg.family == "mlp": + return ScalableMLP( + hidden=cfg.hidden, + depth=cfg.depth, + num_classes=num_classes, + dropout=cfg.dropout, + image_size=img_size, + ) + + kwargs: dict = { + "pretrained": False, + "num_classes": num_classes, + "drop_path_rate": cfg.drop_path_rate, + "img_size": img_size, + "attn_layer": _ScaledAttention, + "embed_layer": _LinearPatchEmbed, + } + return timm.create_model(cfg.model_name, **kwargs) + + +def _get_child(module: nn.Module, key: str) -> nn.Module: + if key.isdigit(): + return module[int(key)] # type: ignore[index] + return getattr(module, key) + + +def _set_child(module: nn.Module, key: str, child: nn.Module) -> None: + if key.isdigit(): + module[int(key)] = child # type: ignore[index] + return + setattr(module, key, child) + + +def _resolve_parent_and_key(root: nn.Module, module_path: str) -> tuple[nn.Module, str]: + parts = module_path.split(".") + if not parts: + raise ValueError("Empty module path.") + parent = root + for key in parts[:-1]: + parent = _get_child(parent, key) + return parent, parts[-1] + + +def _get_module(root: nn.Module, module_path: str) -> nn.Module: + parent, key = _resolve_parent_and_key(root, module_path) + return _get_child(parent, key) + + +def _wrap_module(model: nn.Module, module_path: str, width_dim: int, layer_type: str, **pm_kw) -> bool: + parent, key = _resolve_parent_and_key(model, module_path) + module = _get_child(parent, key) + if isinstance(module, ParametrizedModule): + return False + wrapped = ParametrizedModule(module, width_dim=width_dim, layer_type=layer_type, **pm_kw) + _set_child(parent, key, wrapped) + return True + + +def _install_vit_wrappers(model: nn.Module): + embed_dim = int(getattr(model, "embed_dim")) + + for block in model.blocks: + attn = block.attn + ffn = block.mlp + + _wrap_module(attn, "qkv", embed_dim, "hidden") + _wrap_module(attn, "proj", embed_dim, "hidden") + + _wrap_module(ffn, "fc1", block.mlp.fc1.in_features, "hidden") + _wrap_module(ffn, "fc2", block.mlp.fc2.in_features, "hidden") + + _wrap_module(model, "patch_embed.proj", embed_dim, "embedding") + # _wrap_module(model, "pos_embed", embed_dim, "embedding") + # _wrap_module(model, "cls_token", embed_dim, "embedding") + _wrap_module(model, "head", embed_dim, "readout", a=0.0) + + +def _install_mlp_wrappers(model: ScalableMLP) -> None: + hidden = model.embed.out_features + + for block in model.blocks: + _wrap_module(block, "linear", hidden, "hidden") + + _wrap_module(model, "embed", hidden, "embedding") + _wrap_module(model, "head", hidden, "readout", a=0.0) + + +def install_pm_wrappers(model: nn.Module): + if isinstance(model, ScalableMLP): + _install_mlp_wrappers(model) + elif "visiontransformer" in model.__class__.__name__.lower(): + _install_vit_wrappers(model) + else: + raise TypeError( + f"Unsupported model type: {model.__class__.__name__}. Expected ViT or ScalableMLP." + ) + + +def count_trainable_params(model: nn.Module) -> int: + return sum(p.numel() for p in model.parameters() if p.requires_grad) diff --git a/experiments/vision/run.sh b/experiments/vision/run.sh new file mode 100755 index 0000000..0d13d7c --- /dev/null +++ b/experiments/vision/run.sh @@ -0,0 +1,39 @@ +#!/usr/bin/env bash +# Submit a vision sweep for one scale. +# +# Usage: +# bash experiments/vision/run.sh vit-s [extra launch_sweep.py flags ...] +set -euo pipefail + +SCALE="${1:?Usage: $0 [extra launch_sweep.py flags]}" +shift + +case "$SCALE" in + debug) METHODS="maxP"; LRS="3e-3"; SEEDS="1" ;; + vit-s) METHODS="maxP mup-full mup-no"; LRS="3e-4 1e-3 3e-3 1e-2 3e-2 1e-1 3e-1 1.0 3.0"; SEEDS="1 2 3" ;; + vit-b) METHODS="maxP mup-full mup-no"; LRS="3e-4 1e-3 3e-3 1e-2 3e-2 1e-1 3e-1 1.0 3.0"; SEEDS="1" ;; + vit-l) METHODS="maxP mup-full mup-no"; LRS="3e-4 1e-3 3e-3 1e-2 3e-2 1e-1 3e-1 1.0 3.0"; SEEDS="1" ;; + mlp-s) METHODS="maxP mup-full mup-no"; LRS="3e-4 1e-3 3e-3 1e-2 3e-2 1e-1 3e-1 1.0 3.0"; SEEDS="1 2 3" ;; + mlp-m) METHODS="maxP mup-full mup-no"; LRS="3e-4 1e-3 3e-3 1e-2 3e-2 1e-1 3e-1 1.0 3.0"; SEEDS="1 2 3" ;; + mlp-b) METHODS="maxP mup-full mup-no"; LRS="3e-4 1e-3 3e-3 1e-2 3e-2 1e-1 3e-1 1.0 3.0"; SEEDS="1 2" ;; + mlp-l) METHODS="maxP mup-full mup-no"; LRS="3e-4 1e-3 3e-3 1e-2 3e-2 1e-1 3e-1 1.0 3.0"; SEEDS="1" ;; + *) echo "Unknown scale '$SCALE'. Choose from debug vit-s vit-b vit-l mlp-s mlp-m mlp-b mlp-l."; exit 1 ;; +esac + +RUNS_DIR="${RUNS_DIR:-/net/storage/pr3/plgrid/plggadlers/maxP/runs}" +VENV_PATH="${VENV_PATH:-/net/storage/pr3/plgrid/plggadlers/maxP/.venv}" +REPO_PATH="${REPO_PATH:-/net/storage/pr3/plgrid/plggadlers/maxP}" + +cd "$REPO_PATH" + +python experiments/vision/launch_sweep.py \ + --scale "$SCALE" \ + --methods $METHODS \ + --lrs $LRS \ + --seeds $SEEDS \ + --runs-dir "$RUNS_DIR" \ + --venv-path "$VENV_PATH" \ + --repo-path "$REPO_PATH" \ + --epochs 10 \ + --batch-size 3072 \ + "${@}" diff --git a/experiments/vision/run_debug.sh b/experiments/vision/run_debug.sh new file mode 100755 index 0000000..abb1223 --- /dev/null +++ b/experiments/vision/run_debug.sh @@ -0,0 +1,50 @@ +#!/usr/bin/env bash +# Quick local debug run (CPU/GPU) with a tiny model and tiny dataset slice. +# +# Usage: +# bash experiments/vision/run_debug.sh [--steps N] [--method METHOD] +set -euo pipefail + +REPO="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +VENV="${REPO}/.venv" + +SCALE="${SCALE:-debug}" +METHOD="${METHOD:-maxP}" +DATASET="${DATASET:-beans}" +STEPS="${STEPS:-200}" +BATCH_SIZE="${BATCH_SIZE:-8}" +OUTPUT_DIR="${OUTPUT_DIR:-/tmp/maxp_vision_debug}" + +while [[ $# -gt 0 ]]; do + case "$1" in + --scale) SCALE="$2"; shift 2 ;; + --method) METHOD="$2"; shift 2 ;; + --dataset) DATASET="$2"; shift 2 ;; + --steps) STEPS="$2"; shift 2 ;; + --batch-size) BATCH_SIZE="$2"; shift 2 ;; + --output-dir) OUTPUT_DIR="$2"; shift 2 ;; + *) echo "Unknown flag: $1"; exit 1 ;; + esac +done + +source "${VENV}/bin/activate" +cd "${REPO}" + +mkdir -p "${OUTPUT_DIR}" + +python experiments/vision/train.py \ + --scale "${SCALE}" \ + --method "${METHOD}" \ + --dataset "${DATASET}" \ + --lr 0.03 \ + --batch-size "${BATCH_SIZE}" \ + --num-workers 0 \ + --max-steps "${STEPS}" \ + --val-steps 4 \ + --no-compile \ + --debug \ + --alignment-warmup 5 \ + --solve-interval 5 \ + --log-interval 5 \ + --val-interval 50 \ + --output-dir "${OUTPUT_DIR}" diff --git a/experiments/vision/train.py b/experiments/vision/train.py new file mode 100644 index 0000000..868c707 --- /dev/null +++ b/experiments/vision/train.py @@ -0,0 +1,379 @@ +#!/usr/bin/env python3 +"""maxP vision pre-training entrypoint using timm models and HF streaming data.""" + +from __future__ import annotations + +import argparse +import json +import math +import os +import time +from pathlib import Path +from typing import Any + +import torch +import torch.nn as nn +import torch.nn.functional as F +import wandb +from timm.data import create_transform, resolve_model_data_config +from timm.loss import LabelSmoothingCrossEntropy +from timm.scheduler.cosine_lr import CosineLRScheduler +from torch.utils.tensorboard import SummaryWriter + +from maxp import Parametrization + +from hf_vision_data import DATASET_CONFIGS, build_dataloaders +from maxp_timm import SCALE_CONFIGS, count_trainable_params, create_model, install_pm_wrappers +from utils import append_json, collect_alignments, collect_layer_lrs, save_checkpoint, write_json + + +def evaluate( + *, + model: nn.Module, + loader, + device: torch.device, + amp_enabled: bool, + max_steps: int | None = None, + num_classes: int | None = None, +) -> dict[str, float]: + + model.eval() + total_loss = 0.0 + total_correct = 0 + total_top5 = 0 + total_count = 0 + steps = 0 + + with torch.no_grad(): + for xb, yb in loader: + xb = xb.to(device, non_blocking=True) + yb = yb.to(device, non_blocking=True) + + with torch.autocast(device.type, torch.bfloat16, amp_enabled): + logits = model(xb) + loss = F.cross_entropy(logits, yb) + + preds = logits.argmax(dim=1) + total_loss += float(loss.item()) * xb.size(0) + total_correct += int((preds == yb).sum().item()) + if num_classes is not None and num_classes >= 5: + total_top5 += int( + (logits.topk(5, dim=1).indices == yb.unsqueeze(1)).any(dim=1).sum().item() + ) + total_count += int(xb.size(0)) + steps += 1 + + if max_steps is not None and steps >= max_steps: + break + + model.train() + + if total_count == 0: + return {} + + return { + "loss/val_loss": total_loss / total_count, + "loss/val_top1": total_correct / total_count, + "loss/val_top5": total_top5 / total_count, + "loss/val_samples": float(total_count), + } + + +def current_lr_prefactor(optimizer: torch.optim.Optimizer) -> float: + for group in optimizer.param_groups: + if group.get("layer_name") == "_other": + return float(group["lr"]) + return float(optimizer.param_groups[0]["lr"]) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--scale", type=str, default="vit-s") + parser.add_argument("--method", choices=["maxP", "mup-full", "mup-no"], default="maxP") + parser.add_argument("--lr", type=float, default=1e-2) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--dataset", type=str, default="imagenet12k") + parser.add_argument("--batch-size", type=int, default=128) + parser.add_argument("--epochs", type=int, default=5) + parser.add_argument("--num-workers", type=int, default=8) + parser.add_argument("--prefetch-factor", type=int, default=4) + parser.add_argument("--max-steps", type=int, default=None) + parser.add_argument("--val-interval", type=int, default=500) + parser.add_argument("--val-steps", type=int, default=50) + parser.add_argument("--alignment-warmup", type=int, default=100) + parser.add_argument("--solve-interval", type=int, default=200) + parser.add_argument("--sample-size", type=int, default=32) + parser.add_argument("--lr-warmup", type=int, default=None) + parser.add_argument("--label-smoothing", type=float, default=0.1) + parser.add_argument("--no-compile", action="store_true", default=False) + parser.add_argument("--debug", action="store_true", default=False) + parser.add_argument("--log-interval", type=int, default=20) + parser.add_argument("--checkpoint-interval", type=int, default=None) + parser.add_argument("--grad-clip", type=float, default=1.0) + parser.add_argument("--resume", default=None) + parser.add_argument("--output-dir", type=str, default="runs") + return parser.parse_args() + + +def main() -> None: + args = parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + if torch.cuda.is_available(): + torch.backends.cudnn.benchmark = True + amp_enabled = True + else: + amp_enabled = False + + scale_cfg = SCALE_CONFIGS.get(args.scale) + dataset_cfg = DATASET_CONFIGS.get(args.dataset) + + if args.max_steps is not None: + total_steps = args.max_steps + else: + steps_per_epoch = max(1, dataset_cfg.train_samples // args.batch_size) + total_steps = steps_per_epoch * args.epochs + + num_classes = dataset_cfg.num_classes + model = create_model( + scale=args.scale, + num_classes=num_classes, + image_size=scale_cfg.image_size, + ) + model.set_grad_checkpointing(enable=True) + + install_pm_wrappers(model) + model = model.to(device) + model.train() + + if scale_cfg.family == "mlp": + from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD + model_data_cfg = { + "input_size": (3, scale_cfg.image_size, scale_cfg.image_size), + "mean": IMAGENET_DEFAULT_MEAN, + "std": IMAGENET_DEFAULT_STD, + "interpolation": "bicubic", + "crop_pct": 0.875, + } + else: + model_data_cfg = resolve_model_data_config(model) + train_transform = create_transform(**model_data_cfg, is_training=True) + eval_transform = create_transform(**model_data_cfg, is_training=False) + + train_loader, val_loader, _ = build_dataloaders( + dataset_name=args.dataset, + batch_size=args.batch_size, + num_workers=args.num_workers, + prefetch_factor=args.prefetch_factor, + train_transform=train_transform, + eval_transform=eval_transform, + ) + + sample_input = torch.randn(1, 3, scale_cfg.image_size, scale_cfg.image_size, device=device) + alignment_mode = "full" if "full" in args.method else "no" + dynamic = args.method == "maxP" + + param = Parametrization( + model, + optimizer_type="adam", + alignment=alignment_mode, + lr_prefactor=args.lr, + sample_input=sample_input, + warmup_steps=args.alignment_warmup, + solve_interval=args.solve_interval, + sample_size=args.sample_size, + ) + + optimizer = torch.optim.AdamW( + param.param_groups, + lr=args.lr, + betas=(0.9, 0.999), + weight_decay=0.05, + ) + lr_warmup = args.lr_warmup or int(0.05 * total_steps) + scheduler = CosineLRScheduler( + optimizer, + t_initial=total_steps - lr_warmup, + warmup_t=lr_warmup, + warmup_prefix=True, + ) + + smooth_xe_loss = LabelSmoothingCrossEntropy(smoothing=args.label_smoothing) + + if not args.no_compile: + model = torch.compile(model) + + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + metrics_path = output_dir / "metrics.json" + + run_config = vars(args).copy() + write_json(output_dir / "config.json", run_config) + + if not args.debug: + wandb_project = os.getenv("WANDB_PROJECT", "maxP-vision") + wandb_run_name = os.getenv("WANDB_RUN_NAME", output_dir.name) + wandb.init(project=wandb_project, name=wandb_run_name, config=run_config) + + log_dir = output_dir / "tb" + log_dir.mkdir(parents=True, exist_ok=True) + tb_writer = SummaryWriter(log_dir=str(log_dir)) + + print("\n=== maxP vision run ===") + print(f" scale: {args.scale} ({scale_cfg.model_name})") + print(f" params: {count_trainable_params(model)}") + print(f" method: {args.method}") + print(f" dataset: {args.dataset}") + print(f" train_samples: {dataset_cfg.train_samples}") + print(f" val_samples: {dataset_cfg.val_samples}") + print(f" image_size: {scale_cfg.image_size}") + print(f" num_classes: {num_classes}") + print(f" batch_size: {args.batch_size}") + print(f" num_workers: {args.num_workers}") + print(f" prefetch_factor: {args.prefetch_factor}") + print(f" device: {device}") + print(f" output_dir: {output_dir}") + print("") + + + global_step = start_epoch = total_seen = 0 + align_sample: torch.Tensor | None = None + + t0 = time.time() + last_log_time = t0 + last_log_seen = total_seen + stop_training = False + + for epoch in range(start_epoch, args.epochs): + for xb, yb in train_loader: + xb = xb.to(device, non_blocking=True) + yb = yb.to(device, non_blocking=True) + + if dynamic and align_sample is None: + align_sample = xb[:args.sample_size].detach().clone().to(device) + param.capture_initial(align_sample) + + optimizer.zero_grad() + + with torch.autocast(device.type, torch.bfloat16, amp_enabled): + logits = model(xb) + loss = smooth_xe_loss(logits, yb) + + loss.backward() + torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) + optimizer.step() + scheduler.step(global_step) + + if dynamic: + param.lr_prefactor = current_lr_prefactor(optimizer) + param.step(align_sample, optimizer) + + global_step += 1 + total_seen += int(xb.size(0)) + + if global_step % args.log_interval == 0: + now = time.time() + elapsed = max(now - t0, 1e-6) + interval_time = max(now - last_log_time, 1e-6) + interval_samples = max(total_seen - last_log_seen, 1) + interval_sps = interval_samples / interval_time + last_log_time = now + last_log_seen = total_seen + + if global_step % args.val_interval == 0: + eval_metrics = evaluate( + model=model, + loader=val_loader, + device=device, + amp_enabled=amp_enabled, + max_steps=args.val_steps, + num_classes=num_classes, + ) + else: + eval_metrics = {} + + row: dict[str, Any] = { + "step": global_step, + "epoch": epoch, + "loss/train_loss": float(loss.item()), + "perf/samples_seen": total_seen, + "perf/samples_per_sec": total_seen / elapsed, + "perf/samples_per_sec_interval": interval_sps, + "lr/lr_prefactor": current_lr_prefactor(optimizer), + **collect_layer_lrs(param), + **eval_metrics, + } + if dynamic: + row.update(collect_alignments(param)) + append_json(metrics_path, row) + + if not args.debug: + wandb.log(row, step=global_step) + + for key, value in row.items(): + tb_writer.add_scalar(key, value, global_step) + + val_loss_print = row.get("loss/val_loss") + val_txt = f" val_loss={val_loss_print:.4f}" if isinstance(val_loss_print, float) else "" + ts = time.strftime("%Y-%m-%d %H:%M:%S") + print( + f"[{ts}] [step {global_step:7d}] " + f"loss={row['loss/train_loss']:.4f}" + f"{val_txt} " + f"sps={interval_sps:8.1f}" + ) + + if args.checkpoint_interval and global_step % args.checkpoint_interval == 0: + save_checkpoint( + output_dir / "checkpoint" / f"step_{global_step}.pt", + model=model, + optimizer=optimizer, + step=global_step, + epoch=epoch, + samples_seen=total_seen, + args=args, + ) + + if args.max_steps is not None and global_step >= args.max_steps: + stop_training = True + break + + if stop_training: + break + + eval_metrics = evaluate( + model=model, + loader=val_loader, + device=device, + amp_enabled=amp_enabled, + num_classes=num_classes, + ) + elapsed = time.time() - t0 + + write_json(output_dir / "final_metrics.json", eval_metrics) + + if not args.debug: + wandb.log(eval_metrics, step=global_step) + wandb.finish() + + for key, value in eval_metrics.items(): + tb_writer.add_scalar(key, value, global_step) + tb_writer.close() + + save_checkpoint( + output_dir / "checkpoint" / "final.pt", + model=model, + optimizer=optimizer, + step=global_step, + epoch=max(0, args.epochs - 1), + samples_seen=total_seen, + args=args, + ) + + print(f"\nDone. Total training time: {elapsed:.1f} seconds.") + print(json.dumps(eval_metrics, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/experiments/vision/utils.py b/experiments/vision/utils.py new file mode 100644 index 0000000..a393868 --- /dev/null +++ b/experiments/vision/utils.py @@ -0,0 +1,68 @@ +"""Utility helpers shared across vision experiment scripts.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import torch +import torch.nn as nn + +from maxp import Parametrization + + +def collect_layer_lrs(param: Parametrization) -> dict[str, float]: + out: dict[str, float] = {} + for group in param.param_groups: + layer_name = group.get("layer_name") + if layer_name: + out[f"lrs/{layer_name}"] = float(group["lr"]) + return out + + +def collect_alignments(param: Parametrization) -> dict[str, dict[str, float]]: + out: dict[str, float] = {} + for name, pm in param._pms: + if pm.weight is None: + continue + if pm.align_z0_dW is None: + continue + out[f"align/z0_dW/{name}"] = float(pm.align_z0_dW) + out[f"align/dZ_w0/{name}"] = float(pm.align_dZ_w0) + out[f"align/dZ_dW/{name}"] = float(pm.align_dZ_dW) + return out + + +def save_checkpoint( + ckpt_path: Path, + *, + model: nn.Module, + optimizer: torch.optim.Optimizer, + step: int, + epoch: int, + samples_seen: int, + args: argparse.Namespace, +) -> None: + ckpt_path.parent.mkdir(parents=True, exist_ok=True) + model_for_io = model._orig_mod if hasattr(model, "_orig_mod") else model + torch.save( + { + "model": model_for_io.state_dict(), + "optimizer": optimizer.state_dict(), + "step": step, + "epoch": epoch, + "samples_seen": samples_seen, + "args": vars(args), + }, + ckpt_path, + ) + + +def write_json(path: Path, obj: dict) -> None: + path.write_text(json.dumps(obj, indent=2, sort_keys=True)) + + +def append_json(path: Path, row: dict) -> None: + with path.open("a", encoding="utf-8") as f: + f.write(json.dumps(row, sort_keys=True) + "\n") diff --git a/maxp/trace.py b/maxp/trace.py index 6dbc219..dc2de58 100644 --- a/maxp/trace.py +++ b/maxp/trace.py @@ -125,13 +125,21 @@ def _propagate_pm_tags(cls, func, args, wrapped_out): """ from maxp.dag import MergeType - # Collect distinct tag sets from tensor inputs + # Collect distinct tag sets from tensor inputs (recurse into list/tuple + # args so that ops like torch.cat, which pass tensors in a list, are handled). input_tag_sets: list[frozenset[str]] = [] - for arg in args: - if isinstance(arg, _TracingTensor): - tags = getattr(arg, '_pm_tags', frozenset()) + + def _collect(obj): + if isinstance(obj, _TracingTensor): + tags = getattr(obj, '_pm_tags', frozenset()) if tags and tags not in input_tag_sets: input_tag_sets.append(tags) + elif isinstance(obj, (tuple, list)): + for item in obj: + _collect(item) + + for arg in args: + _collect(arg) if not input_tag_sets: return diff --git a/paramR b/paramR deleted file mode 160000 index f1934d9..0000000 --- a/paramR +++ /dev/null @@ -1 +0,0 @@ -Subproject commit f1934d99f5df75dba063f36f9a740064ee6482f4 diff --git a/pyproject.toml b/pyproject.toml index 461dd27..20d553a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,6 +44,7 @@ dev = [ "pytest-cov>=7.0.0", "pyyaml>=6.0.3", "tiktoken>=0.7.0", + "timm~=1.0.26", "torchtitan", "torchvision>=0.24.1", "tqdm>=4.60.0",