From 3eb4fca215374a699a35f6317d20b2335fe23cd6 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Thu, 30 Apr 2026 10:34:33 +0200 Subject: [PATCH 01/11] ViT experiments --- .gitignore | 2 + README.md | 30 +- experiments/vision/coord_check_convnext.py | 234 +++++++++++++ experiments/vision/coord_check_vit.py | 111 +++++++ experiments/vision/hf_vision_data.py | 129 ++++++++ experiments/vision/launch_sweep.py | 248 ++++++++++++++ experiments/vision/maxp_timm.py | 236 +++++++++++++ experiments/vision/run.sh | 37 +++ experiments/vision/run_debug.sh | 46 +++ experiments/vision/train.py | 365 +++++++++++++++++++++ experiments/vision/utils.py | 68 ++++ maxp/trace.py | 17 +- pyproject.toml | 1 + 13 files changed, 1518 insertions(+), 6 deletions(-) create mode 100644 experiments/vision/coord_check_convnext.py create mode 100644 experiments/vision/coord_check_vit.py create mode 100644 experiments/vision/hf_vision_data.py create mode 100755 experiments/vision/launch_sweep.py create mode 100644 experiments/vision/maxp_timm.py create mode 100755 experiments/vision/run.sh create mode 100755 experiments/vision/run_debug.sh create mode 100644 experiments/vision/train.py create mode 100644 experiments/vision/utils.py diff --git a/.gitignore b/.gitignore index b2b19a4..993a2ad 100644 --- a/.gitignore +++ b/.gitignore @@ -211,6 +211,8 @@ jobs outputs assets runs +figs +analysis # Generated plots *.png diff --git a/README.md b/README.md index bc24074..4b81e43 100644 --- a/README.md +++ b/README.md @@ -83,15 +83,36 @@ experiments/lm/ # LLaMA-3 pre-training experiments (torchtitan) ├── 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 + +experiments/vision/ # ViT/ConvNeXt-V2 experiments (timm) +├── 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 +├── launch_sweep.py # Generate and submit SLURM jobs for full LR sweep +├── run.sh # Submit a sweep for one vision scale +├── run_debug.sh # Quick tiny-model debug run (CPU/GPU) +└── coord_check.py # Lightweight coord-style diagnostic for wrapped models ``` **Phase 1** (static, at init): discover PMs → reinit weights → solve LP → build param groups. **Phase 2** (dynamic, per step): capture activations → measure alignment → re-solve LP → update per-layer LRs. +### Vision experiment quick start + +```bash +# CPU/GPU debug smoke test (small model + tiny dataset slice) +bash experiments/vision/run_debug.sh --steps 20 + +# Launch full SLURM sweep for one scale +bash experiments/vision/run.sh vit-s +``` + ## ParametrizedModule @@ -143,6 +164,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 +296,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. diff --git a/experiments/vision/coord_check_convnext.py b/experiments/vision/coord_check_convnext.py new file mode 100644 index 0000000..c148926 --- /dev/null +++ b/experiments/vision/coord_check_convnext.py @@ -0,0 +1,234 @@ +#!/usr/bin/env python3 +"""Coord-style diagnostic for a custom width-scalable ConvNeXt model. + +Uses a hand-rolled ConvNeXt-V1 style model with a single base width `w` +(channel progression [w, 2w, 4w, 8w]) so the scaling axis is continuous +and interpretable — unlike discrete timm variants (T/S/B/L). + +Standard PyTorch init; maxP can reinitialize weights freely. +""" + +from __future__ import annotations + +import argparse + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from maxp import Parametrization, ParametrizedModule, diagnose_axis, plot_axis, print_axis + + +# --------------------------------------------------------------------------- +# Model +# --------------------------------------------------------------------------- + +class _LayerNorm2d(nn.Module): + """NCHW wrapper around nn.LayerNorm (no learnable params in the conv path).""" + + def __init__(self, dim: int) -> None: + super().__init__() + self.norm = nn.LayerNorm(dim) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.norm(x.permute(0, 2, 3, 1)).permute(0, 3, 1, 2) + + +class _ConvNeXtBlock(nn.Module): + def __init__(self, dim: int) -> None: + super().__init__() + self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim, bias=True) + self.norm = nn.LayerNorm(dim) + self.pwconv1 = nn.Linear(dim, 4 * dim, bias=True) + self.act = nn.GELU() + self.pwconv2 = nn.Linear(4 * dim, dim, bias=True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + residual = x + x = self.dwconv(x) + x = x.permute(0, 2, 3, 1) + x = self.norm(x) + x = self.pwconv1(x) + x = self.act(x) + x = self.pwconv2(x) + x = x.permute(0, 3, 1, 2) + return residual + x + + +class ScalableConvNeXt(nn.Module): + """ConvNeXt-V1 with configurable base width w. + + Channel progression: [w, 2w, 4w, 8w]. + No layer scale, no GRN — keeps init simple for maxP. + """ + + def __init__( + self, + w: int, + num_classes: int = 1000, + depths: tuple[int, ...] = (2, 2, 6, 2), + ) -> None: + super().__init__() + dims = [w, 2 * w, 4 * w, 8 * w] + + self.stem = nn.Sequential( + nn.Conv2d(3, dims[0], kernel_size=4, stride=4, bias=True), + _LayerNorm2d(dims[0]), + ) + + self.stages = nn.ModuleList([ + nn.Sequential(*[_ConvNeXtBlock(d) for _ in range(depth)]) + for d, depth in zip(dims, depths) + ]) + + # Between-stage downsampling (3 transitions for 4 stages) + self.downsamples = nn.ModuleList([ + nn.Sequential( + _LayerNorm2d(dims[i]), + nn.Conv2d(dims[i], dims[i + 1], kernel_size=2, stride=2, bias=True), + ) + for i in range(3) + ]) + + self.head_norm = nn.LayerNorm(dims[-1]) + self.head = nn.Linear(dims[-1], num_classes, bias=True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.stem(x) + for i, stage in enumerate(self.stages): + x = stage(x) + if i < len(self.downsamples): + x = self.downsamples[i](x) + x = x.mean([-2, -1]) + x = self.head_norm(x) + return self.head(x) + + +# --------------------------------------------------------------------------- +# ParametrizedModule installation +# --------------------------------------------------------------------------- + +def _conv_fan_in(m: nn.Conv2d) -> int: + # Depthwise: scaling axis is channel count, not spatial fan-in. + if m.groups == m.in_channels == m.out_channels: + return m.in_channels + return (m.in_channels // m.groups) * m.kernel_size[0] * m.kernel_size[1] + + +def install_wrappers(model: ScalableConvNeXt) -> list[str]: + wrapped: list[str] = [] + + # Stem Conv2d — embedding + stem_conv = model.stem[0] + model.stem[0] = ParametrizedModule(stem_conv, width_dim=_conv_fan_in(stem_conv), layer_type="embedding") + wrapped.append("stem.0") + + # Stage blocks: dwconv, pwconv1, pwconv2 + for si, stage in enumerate(model.stages): + for bi, block in enumerate(stage): + base = f"stages.{si}.{bi}" + + dw = block.dwconv + block.dwconv = ParametrizedModule(dw, width_dim=_conv_fan_in(dw), layer_type="hidden") + wrapped.append(f"{base}.dwconv") + + pw1 = block.pwconv1 + block.pwconv1 = ParametrizedModule(pw1, width_dim=pw1.in_features, layer_type="hidden") + wrapped.append(f"{base}.pwconv1") + + pw2 = block.pwconv2 + block.pwconv2 = ParametrizedModule(pw2, width_dim=pw2.in_features, layer_type="hidden") + wrapped.append(f"{base}.pwconv2") + + # Downsample Conv2d — hidden + for di, ds in enumerate(model.downsamples): + conv = ds[1] + model.downsamples[di][1] = ParametrizedModule(conv, width_dim=_conv_fan_in(conv), layer_type="hidden") + wrapped.append(f"downsamples.{di}.1") + + # Head Linear — readout + head = model.head + model.head = ParametrizedModule(head, width_dim=head.in_features, layer_type="readout") + wrapped.append("head") + + return wrapped + + +# --------------------------------------------------------------------------- +# Coord check harness +# --------------------------------------------------------------------------- + +WIDTHS = [32, 48, 64, 96, 128] # base channel widths → dims [w, 2w, 4w, 8w] + + +def _make_model(w: int, parametrized: bool): + model = ScalableConvNeXt(w=w, num_classes=1000) + if not parametrized: + return model, None + install_wrappers(model) + param = Parametrization( + model, + lr_prefactor=1e-3, + optimizer_type="adam", + alignment="full", + sample_input=None, + ) + return model, param.param_groups + + +def _make_input(w: int) -> torch.Tensor: + return torch.randn(2, 3, 224, 224) + + +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(2, 3, 224, 224) + y_fixed = torch.randint(0, 1000, (2,)) + + 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 scalable ConvNeXt (continuous width axis)", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument("--parametrized", action="store_true") + parser.add_argument("--steps", type=int, default=3) + parser.add_argument("--seeds", type=int, default=2) + parser.add_argument("--plot", action="store_true") + args = parser.parse_args() + + variant = "parametrized" if args.parametrized else "plain" + print(f"Scalable ConvNeXt 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("base_width", ops, affected, act_stats, WIDTHS) + + if args.plot: + out_name = f"coord_check_convnext_scalable_{variant}.png" + plot_axis("base_width", 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..96e124d --- /dev/null +++ b/experiments/vision/hf_vision_data.py @@ -0,0 +1,129 @@ +"""HF 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, IterableDataset + + +@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", + ), +} + + +class HFVisionIterableDataset(IterableDataset): + """Bridge HF streaming iterable datasets into torch IterableDataset.""" + + def __init__( + self, + hf_stream, + *, + preset: DatasetPreset, + transform: Callable, + ) -> None: + super().__init__() + self._hf_stream = hf_stream + self._preset = preset + self._transform = transform + + def __iter__(self): + for sample in self._hf_stream: + image = sample.get(self._preset.image_key) + image = image.convert("RGB") + label = sample.get(self._preset.label_key) + + x = self._transform(image) + y = torch.tensor(label, dtype=torch.long) + + yield x, y + + +def make_loader( + ds: IterableDataset, + *, + batch_size: int, + num_workers: int, + is_train: bool, + prefetch_factor: int, +) -> DataLoader: + kwargs = { + "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, + seed: int = 1, + train_shuffle_buffer: int = 10_000, + 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_stream = load_dataset(preset.dataset_id, split=preset.train_split, streaming=True) + train_stream = train_stream.shuffle(buffer_size=train_shuffle_buffer, seed=seed) + train_ds = HFVisionIterableDataset(train_stream, 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_stream = load_dataset(preset.dataset_id, split=preset.val_split, streaming=True) + val_ds = HFVisionIterableDataset(val_stream, 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..98460ef --- /dev/null +++ b/experiments/vision/launch_sweep.py @@ -0,0 +1,248 @@ +#!/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 + +from hf_vision_data import DATASET_CONFIGS + + +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, + }, + "cnx-t": { + "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, + }, + "cnx-s": { + "wall": "36:00:00", + "gpus": 1, + "methods": ["maxP", "mup-full", "mup-no"], + "lrs": ALL_LRS, + "seeds": [1], + "image_size": 224, + "drop_path_rate": 0.2, + }, + "cnx-b": { + "wall": "48:00:00", + "gpus": 1, + "methods": ["maxP", "mup-full", "mup-no"], + "lrs": TRANSFER_LR, + "seeds": [1], + "image_size": 224, + "drop_path_rate": 0.3, + }, + "cnx-l": { + "wall": "72: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 ${cpus_per_gpu} +#SBATCH --mem-per-gpu ${mem_per_gpu} +#SBATCH --time ${wall_time} +#SBATCH --account ${account} +#SBATCH --partition ${partition} +#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 + +source "${venv_path}/bin/activate" +cd "${repo_path}" + +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 ${train_dataset} \\ + --lr-warmup ${lr_warmup} \\ + --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("--train-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("--lr-warmup", type=int, default=None) + parser.add_argument("--warmup-ratio", type=float, default=0.05) + + parser.add_argument("--gpus-per-node", type=int, default=None) + parser.add_argument("--cpus-per-gpu", type=int, default=16) + parser.add_argument("--mem-per-gpu", default="64GB") + parser.add_argument("--account", default="plgadlers-gpu-gh200") + parser.add_argument("--partition", default="plgrid-gpu-gh200") + + 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"] + gpus_per_node = args.gpus_per_node if args.gpus_per_node is not None else sc["gpus"] + train_samples = DATASET_CONFIGS[args.train_dataset].train_samples + total_steps = max(1, train_samples // args.batch_size) * args.epochs + lr_warmup = args.lr_warmup or int(round(args.warmup_ratio * total_steps)) + + 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, + train_dataset=args.train_dataset, + lr_warmup=lr_warmup, + epochs=args.epochs, + batch_size=args.batch_size, + num_workers=args.num_workers, + val_steps=args.val_steps, + gpus_per_node=gpus_per_node, + cpus_per_gpu=args.cpus_per_gpu, + mem_per_gpu=args.mem_per_gpu, + account=args.account, + partition=args.partition, + 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..2fd0689 --- /dev/null +++ b/experiments/vision/maxp_timm.py @@ -0,0 +1,236 @@ +"""timm model registry and ParametrizedModule wrappers for vision experiments.""" + +from __future__ import annotations +from dataclasses import dataclass +from typing import Optional, Type + +import timm +import torch +import torch.nn as nn +import torch.nn.functional as F +from timm.layers import Attention, maybe_add_mask, resolve_self_attn_mask + +from maxp import ParametrizedModule + + +@dataclass(frozen=True, slots=True) +class ScaleConfig: + model_name: str + family: str + image_size: int + drop_path_rate: float + + +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 + ), + "cnx-t": ScaleConfig( + model_name="convnextv2_tiny", family="convnext", image_size=224, drop_path_rate=0.1 + ), + "cnx-s": ScaleConfig( + model_name="convnextv2_small", family="convnext", image_size=224, drop_path_rate=0.2 + ), + "cnx-b": ScaleConfig( + model_name="convnextv2_base", family="convnext", image_size=224, drop_path_rate=0.3 + ), + "cnx-l": ScaleConfig( + model_name="convnextv2_large", family="convnext", image_size=224, drop_path_rate=0.4 + ), +} + + +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 + + +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 + resolved_drop_path = cfg.drop_path_rate + + kwargs: dict = { + "pretrained": False, + "num_classes": num_classes, + "drop_path_rate": resolved_drop_path, + } + + if cfg.family == "vit": + kwargs["img_size"] = img_size + kwargs["attn_layer"] = _ScaledAttention + + model = timm.create_model(cfg.model_name, **kwargs) + return model + + +def _conv_width_dim(module: nn.Conv2d) -> int: + # For depthwise convs (groups == in_channels), fan-in stays constant (k*k), + # so use channel width to preserve scaling across model sizes. + if module.groups == module.in_channels and module.in_channels == module.out_channels: + return module.in_channels + k_h, k_w = module.kernel_size + return (module.in_channels // module.groups) * k_h * k_w + + +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 _convnext_layer_type(name: str) -> str: + if name.startswith("stem") or name.startswith("downsample_layers.0"): + return "embedding" + if name.startswith("head"): + return "readout" + return "hidden" + + +def _install_convnext_wrappers(model: nn.Module) -> list[str]: + wrapped: list[str] = [] + targets: list[tuple[str, int, str]] = [] + + for name, module in model.named_modules(): + if isinstance(module, ParametrizedModule): + continue + if isinstance(module, nn.Linear): + targets.append((name, _linear_width_dim(module), _convnext_layer_type(name))) + elif isinstance(module, nn.Conv2d): + targets.append((name, _conv_width_dim(module), _convnext_layer_type(name))) + + for module_path, width_dim, layer_type in targets: + if _wrap_module( + model, + module_path, + width_dim=width_dim, + layer_type=layer_type, + ): + wrapped.append(module_path) + + return wrapped + + +def install_pm_wrappers(model: nn.Module): + if "visiontransformer" in model.__class__.__name__.lower(): + _install_vit_wrappers(model) + elif "convnext" in model.__class__.__name__.lower(): + _install_convnext_wrappers(model) + else: + raise TypeError( + "Unsupported timm model type. Expected ViT or ConvNeXt-like model." + ) + + +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..615eb61 --- /dev/null +++ b/experiments/vision/run.sh @@ -0,0 +1,37 @@ +#!/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"; SEEDS="1 2 3" ;; + vit-b) METHODS="maxP mup-full mup-no"; LRS="${TRANSFER_LR:-1e-2}"; SEEDS="1" ;; + vit-l) METHODS="maxP mup-full mup-no"; LRS="${TRANSFER_LR:-1e-2}"; SEEDS="1" ;; + cnx-t) METHODS="maxP mup-full mup-no"; LRS="3e-4 1e-3 3e-3 1e-2 3e-2 1e-1 3e-1"; SEEDS="1 2 3" ;; + cnx-s) METHODS="maxP mup-full mup-no"; LRS="3e-4 1e-3 3e-3 1e-2 3e-2 1e-1 3e-1"; SEEDS="1" ;; + cnx-b) METHODS="maxP mup-full mup-no"; LRS="${TRANSFER_LR:-1e-2}"; SEEDS="1" ;; + cnx-l) METHODS="maxP mup-full mup-no"; LRS="${TRANSFER_LR:-1e-2}"; SEEDS="1" ;; + *) echo "Unknown scale '$SCALE'. Choose from debug vit-s vit-b vit-l cnx-t cnx-s cnx-b cnx-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" \ + "${@}" diff --git a/experiments/vision/run_debug.sh b/experiments/vision/run_debug.sh new file mode 100755 index 0000000..7a93cbc --- /dev/null +++ b/experiments/vision/run_debug.sh @@ -0,0 +1,46 @@ +#!/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}" +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 ;; + --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}" \ + --lr 1e-3 \ + --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 \ + --output-dir "${OUTPUT_DIR}" diff --git a/experiments/vision/train.py b/experiments/vision/train.py new file mode 100644 index 0000000..bf65cb0 --- /dev/null +++ b/experiments/vision/train.py @@ -0,0 +1,365 @@ +#!/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, +) -> 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()) + 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=4) + parser.add_argument("--prefetch-factor", type=int, default=2) + parser.add_argument("--max-steps", type=int, default=None) + parser.add_argument("--val-steps", type=int, default=None) + 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() + + 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, + seed=args.seed, + 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(set_to_none=True) + + with torch.autocast(device.type, torch.bfloat16, amp_enabled): + logits = model(xb) + loss = smooth_xe_loss(logits, yb) + + if not math.isfinite(float(loss.item())): + raise RuntimeError(f"Non-finite loss at step {global_step + 1}: {loss.item()}") + + loss.backward() + torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) + optimizer.step() + + if scheduler is not None: + scheduler.step() + + if dynamic: + if scheduler is not None: + 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 + + val_interval = evaluate( + model=model, + loader=val_loader, + device=device, + amp_enabled=amp_enabled, + max_steps=args.val_steps, + ) + + 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), + **val_interval, + } + 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("val_loss") + val_txt = f" val_loss={val_loss_print:.4f}" if isinstance(val_loss_print, float) else "" + print( + f"[step {global_step}] loss={row['train_loss']:.4f}{val_txt}" + f" sps={interval_sps:.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 + ) + 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..277e975 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 @@ -169,6 +177,7 @@ def _propagate_pm_tags(cls, func, args, wrapped_out): torch.addmm: "addmm", F.linear: "linear", F.embedding: "embedding", + F.conv2d: "conv2d", } 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", From b0264b34f9552c10b1cb7aa1917b01a409782536 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Thu, 30 Apr 2026 13:15:38 +0200 Subject: [PATCH 02/11] Refactor, fix patch embedding --- experiments/vision/hf_vision_data.py | 10 ++++++ experiments/vision/maxp_timm.py | 48 ++++++++++++++++++++++++++-- experiments/vision/run_debug.sh | 6 +++- experiments/vision/train.py | 46 +++++++++++++------------- 4 files changed, 85 insertions(+), 25 deletions(-) diff --git a/experiments/vision/hf_vision_data.py b/experiments/vision/hf_vision_data.py index 96e124d..c5f0df8 100644 --- a/experiments/vision/hf_vision_data.py +++ b/experiments/vision/hf_vision_data.py @@ -33,6 +33,16 @@ class DatasetPreset: 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", + ), } diff --git a/experiments/vision/maxp_timm.py b/experiments/vision/maxp_timm.py index 2fd0689..beeda4f 100644 --- a/experiments/vision/maxp_timm.py +++ b/experiments/vision/maxp_timm.py @@ -2,13 +2,14 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Optional, Type +from typing import Optional import timm import torch import torch.nn as nn import torch.nn.functional as F -from timm.layers import Attention, maybe_add_mask, resolve_self_attn_mask +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 @@ -97,6 +98,48 @@ def forward( x = self.proj(x) x = self.proj_drop(x) return 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 _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( @@ -118,6 +161,7 @@ def create_model( if cfg.family == "vit": kwargs["img_size"] = img_size kwargs["attn_layer"] = _ScaledAttention + kwargs["embed_layer"] = _LinearPatchEmbed model = timm.create_model(cfg.model_name, **kwargs) return model diff --git a/experiments/vision/run_debug.sh b/experiments/vision/run_debug.sh index 7a93cbc..abb1223 100755 --- a/experiments/vision/run_debug.sh +++ b/experiments/vision/run_debug.sh @@ -10,6 +10,7 @@ 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}" @@ -18,6 +19,7 @@ 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 ;; @@ -33,7 +35,8 @@ mkdir -p "${OUTPUT_DIR}" python experiments/vision/train.py \ --scale "${SCALE}" \ --method "${METHOD}" \ - --lr 1e-3 \ + --dataset "${DATASET}" \ + --lr 0.03 \ --batch-size "${BATCH_SIZE}" \ --num-workers 0 \ --max-steps "${STEPS}" \ @@ -43,4 +46,5 @@ python experiments/vision/train.py \ --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 index bf65cb0..16e71c6 100644 --- a/experiments/vision/train.py +++ b/experiments/vision/train.py @@ -34,6 +34,7 @@ def evaluate( device: torch.device, amp_enabled: bool, max_steps: int | None = None, + num_classes: int | None = None, ) -> dict[str, float]: model.eval() @@ -55,9 +56,10 @@ def evaluate( preds = logits.argmax(dim=1) total_loss += float(loss.item()) * xb.size(0) total_correct += int((preds == yb).sum().item()) - total_top5 += int( - (logits.topk(5, dim=1).indices == yb.unsqueeze(1)).any(dim=1).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 @@ -96,6 +98,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--num-workers", type=int, default=4) parser.add_argument("--prefetch-factor", type=int, default=2) parser.add_argument("--max-steps", type=int, default=None) + parser.add_argument("--val-interval", type=int, default=500) parser.add_argument("--val-steps", type=int, default=None) parser.add_argument("--alignment-warmup", type=int, default=100) parser.add_argument("--solve-interval", type=int, default=200) @@ -242,25 +245,19 @@ def main() -> None: align_sample = xb[:args.sample_size].detach().clone().to(device) param.capture_initial(align_sample) - optimizer.zero_grad(set_to_none=True) + optimizer.zero_grad() with torch.autocast(device.type, torch.bfloat16, amp_enabled): logits = model(xb) loss = smooth_xe_loss(logits, yb) - if not math.isfinite(float(loss.item())): - raise RuntimeError(f"Non-finite loss at step {global_step + 1}: {loss.item()}") - loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) optimizer.step() - - if scheduler is not None: - scheduler.step() + scheduler.step(global_step) if dynamic: - if scheduler is not None: - param.lr_prefactor = current_lr_prefactor(optimizer) + param.lr_prefactor = current_lr_prefactor(optimizer) param.step(align_sample, optimizer) global_step += 1 @@ -275,13 +272,17 @@ def main() -> None: last_log_time = now last_log_seen = total_seen - val_interval = evaluate( - model=model, - loader=val_loader, - device=device, - amp_enabled=amp_enabled, - max_steps=args.val_steps, - ) + if global_step % args.val_interval == 0: + val_interval = evaluate( + model=model, + loader=val_loader, + device=device, + amp_enabled=amp_enabled, + max_steps=args.val_steps, + num_classes=num_classes, + ) + else: + val_interval = {} row: dict[str, Any] = { "step": global_step, @@ -304,10 +305,10 @@ def main() -> None: for key, value in row.items(): tb_writer.add_scalar(key, value, global_step) - val_loss_print = row.get("val_loss") + val_loss_print = row.get("loss/val_loss") val_txt = f" val_loss={val_loss_print:.4f}" if isinstance(val_loss_print, float) else "" print( - f"[step {global_step}] loss={row['train_loss']:.4f}{val_txt}" + f"[step {global_step}] loss={row['loss/train_loss']:.4f}{val_txt}" f" sps={interval_sps:.1f}" ) @@ -333,7 +334,8 @@ def main() -> None: model=model, loader=val_loader, device=device, - amp_enabled=amp_enabled + amp_enabled=amp_enabled, + num_classes=num_classes, ) elapsed = time.time() - t0 From 2dab5c7f91a3f26ae3218a758280ae495ccf244d Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Thu, 30 Apr 2026 13:21:47 +0200 Subject: [PATCH 03/11] Fix slurm script --- experiments/vision/launch_sweep.py | 26 +++++++------------------- 1 file changed, 7 insertions(+), 19 deletions(-) diff --git a/experiments/vision/launch_sweep.py b/experiments/vision/launch_sweep.py index 98460ef..a1538bf 100755 --- a/experiments/vision/launch_sweep.py +++ b/experiments/vision/launch_sweep.py @@ -9,8 +9,6 @@ from pathlib import Path from string import Template -from hf_vision_data import DATASET_CONFIGS - ALL_LRS = [3e-4, 1e-3, 3e-3, 1e-2, 3e-2, 1e-1, 3e-1] TRANSFER_LR = [1e-2] @@ -96,18 +94,21 @@ #!/bin/bash -l #SBATCH --job-name maxp_${scale}_${method_tag}_lr${lr_tag}_s${seed} #SBATCH --nodes 1 -#SBATCH --cpus-per-gpu ${cpus_per_gpu} -#SBATCH --mem-per-gpu ${mem_per_gpu} +#SBATCH --cpus-per-gpu 72 +#SBATCH --mem-per-gpu 118GB #SBATCH --time ${wall_time} -#SBATCH --account ${account} -#SBATCH --partition ${partition} +#SBATCH --account plgadlers-gpu-gh200 +#SBATCH --partition plgrid-gpu-gh200 #SBATCH --gres gpu:${gpus_per_node} #SBATCH --output ${output_dir}/maxp_${scale}_${method_tag}_lr${lr_tag}_s${seed}.out #SBATCH --error ${output_dir}/maxp_${scale}_${method_tag}_lr${lr_tag}_s${seed}.err +module add ML-bundle/25.10 source "${venv_path}/bin/activate" cd "${repo_path}" +export OMP_NUM_THREADS=16 +export HF_HOME="$${HF_HOME:-/net/scratch/hscra/plgrid/plgmwojnar/hf}" export WANDB_PROJECT="$${WANDB_PROJECT:-maxP-vision}" export WANDB_RUN_NAME="maxp_${scale}_${method_tag}_lr${lr_tag}_s${seed}" @@ -117,7 +118,6 @@ --lr ${lr} \\ --seed ${seed} \\ --dataset ${train_dataset} \\ - --lr-warmup ${lr_warmup} \\ --epochs ${epochs} \\ --batch-size ${batch_size} \\ --num-workers ${num_workers} \\ @@ -160,14 +160,6 @@ def parse_args() -> argparse.Namespace: 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("--lr-warmup", type=int, default=None) - parser.add_argument("--warmup-ratio", type=float, default=0.05) - - parser.add_argument("--gpus-per-node", type=int, default=None) - parser.add_argument("--cpus-per-gpu", type=int, default=16) - parser.add_argument("--mem-per-gpu", default="64GB") - parser.add_argument("--account", default="plgadlers-gpu-gh200") - parser.add_argument("--partition", default="plgrid-gpu-gh200") parser.add_argument("--extra-train-args", default="") parser.add_argument("--dry-run", action="store_true") @@ -182,9 +174,6 @@ def main() -> None: lrs = args.lrs or sc["lrs"] seeds = args.seeds or sc["seeds"] gpus_per_node = args.gpus_per_node if args.gpus_per_node is not None else sc["gpus"] - train_samples = DATASET_CONFIGS[args.train_dataset].train_samples - total_steps = max(1, train_samples // args.batch_size) * args.epochs - lr_warmup = args.lr_warmup or int(round(args.warmup_ratio * total_steps)) today = date.today().strftime("%Y-%m-%d") submitted = skipped = 0 @@ -214,7 +203,6 @@ def main() -> None: method=method, lr=lr, train_dataset=args.train_dataset, - lr_warmup=lr_warmup, epochs=args.epochs, batch_size=args.batch_size, num_workers=args.num_workers, From a2899bb33638b416bccb6c772add0b86670256ff Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Thu, 30 Apr 2026 13:23:30 +0200 Subject: [PATCH 04/11] Fix slurm script --- experiments/vision/launch_sweep.py | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/experiments/vision/launch_sweep.py b/experiments/vision/launch_sweep.py index a1538bf..6b9a5be 100755 --- a/experiments/vision/launch_sweep.py +++ b/experiments/vision/launch_sweep.py @@ -117,7 +117,7 @@ --method ${method} \\ --lr ${lr} \\ --seed ${seed} \\ - --dataset ${train_dataset} \\ + --dataset ${dataset} \\ --epochs ${epochs} \\ --batch-size ${batch_size} \\ --num-workers ${num_workers} \\ @@ -155,7 +155,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--venv-path", required=True) parser.add_argument("--repo-path", required=True) - parser.add_argument("--train-dataset", default="imagenet12k") + 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) @@ -173,7 +173,6 @@ def main() -> None: methods = args.methods or sc["methods"] lrs = args.lrs or sc["lrs"] seeds = args.seeds or sc["seeds"] - 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 @@ -202,16 +201,11 @@ def main() -> None: repo_path=args.repo_path, method=method, lr=lr, - train_dataset=args.train_dataset, + dataset=args.dataset, epochs=args.epochs, batch_size=args.batch_size, num_workers=args.num_workers, val_steps=args.val_steps, - gpus_per_node=gpus_per_node, - cpus_per_gpu=args.cpus_per_gpu, - mem_per_gpu=args.mem_per_gpu, - account=args.account, - partition=args.partition, extra_train_args=args.extra_train_args, ) script_path = out_dir / "job.sh" From 5236de495a150c94fb04958b634b0cd33bebd7f4 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Thu, 30 Apr 2026 13:24:33 +0200 Subject: [PATCH 05/11] Fix slurm script --- experiments/vision/launch_sweep.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/experiments/vision/launch_sweep.py b/experiments/vision/launch_sweep.py index 6b9a5be..297212e 100755 --- a/experiments/vision/launch_sweep.py +++ b/experiments/vision/launch_sweep.py @@ -99,7 +99,7 @@ #SBATCH --time ${wall_time} #SBATCH --account plgadlers-gpu-gh200 #SBATCH --partition plgrid-gpu-gh200 -#SBATCH --gres gpu:${gpus_per_node} +#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 From 45731e71aa15a0bd0d03d0b23273d93be035bc56 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Thu, 30 Apr 2026 14:58:47 +0200 Subject: [PATCH 06/11] Adjust training parameters --- .gitignore | 1 + experiments/vision/hf_vision_data.py | 4 ++-- experiments/vision/run.sh | 10 ++++++---- experiments/vision/train.py | 4 ++-- 4 files changed, 11 insertions(+), 8 deletions(-) diff --git a/.gitignore b/.gitignore index 993a2ad..d2460a8 100644 --- a/.gitignore +++ b/.gitignore @@ -213,6 +213,7 @@ assets runs figs analysis +wandb # Generated plots *.png diff --git a/experiments/vision/hf_vision_data.py b/experiments/vision/hf_vision_data.py index c5f0df8..580fc49 100644 --- a/experiments/vision/hf_vision_data.py +++ b/experiments/vision/hf_vision_data.py @@ -115,7 +115,7 @@ def build_dataloaders( preset = DATASET_CONFIGS.get(dataset_name) - train_stream = load_dataset(preset.dataset_id, split=preset.train_split, streaming=True) + train_stream = load_dataset(preset.dataset_id, split=preset.train_split, streaming=False) train_stream = train_stream.shuffle(buffer_size=train_shuffle_buffer, seed=seed) train_ds = HFVisionIterableDataset(train_stream, preset=preset, transform=train_transform) train_loader = make_loader( @@ -126,7 +126,7 @@ def build_dataloaders( prefetch_factor=prefetch_factor, ) - val_stream = load_dataset(preset.dataset_id, split=preset.val_split, streaming=True) + val_stream = load_dataset(preset.dataset_id, split=preset.val_split, streaming=False) val_ds = HFVisionIterableDataset(val_stream, preset=preset, transform=eval_transform) val_loader = make_loader( val_ds, diff --git a/experiments/vision/run.sh b/experiments/vision/run.sh index 615eb61..e080268 100755 --- a/experiments/vision/run.sh +++ b/experiments/vision/run.sh @@ -10,11 +10,11 @@ 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"; SEEDS="1 2 3" ;; - vit-b) METHODS="maxP mup-full mup-no"; LRS="${TRANSFER_LR:-1e-2}"; 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="${TRANSFER_LR:-1e-2}"; SEEDS="1" ;; - cnx-t) METHODS="maxP mup-full mup-no"; LRS="3e-4 1e-3 3e-3 1e-2 3e-2 1e-1 3e-1"; SEEDS="1 2 3" ;; - cnx-s) METHODS="maxP mup-full mup-no"; LRS="3e-4 1e-3 3e-3 1e-2 3e-2 1e-1 3e-1"; SEEDS="1" ;; + cnx-t) 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" ;; + cnx-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" ;; cnx-b) METHODS="maxP mup-full mup-no"; LRS="${TRANSFER_LR:-1e-2}"; SEEDS="1" ;; cnx-l) METHODS="maxP mup-full mup-no"; LRS="${TRANSFER_LR:-1e-2}"; SEEDS="1" ;; *) echo "Unknown scale '$SCALE'. Choose from debug vit-s vit-b vit-l cnx-t cnx-s cnx-b cnx-l."; exit 1 ;; @@ -34,4 +34,6 @@ python experiments/vision/launch_sweep.py \ --runs-dir "$RUNS_DIR" \ --venv-path "$VENV_PATH" \ --repo-path "$REPO_PATH" \ + --epochs 10 \ + --batch-size 3072 \ "${@}" diff --git a/experiments/vision/train.py b/experiments/vision/train.py index 16e71c6..ca45c06 100644 --- a/experiments/vision/train.py +++ b/experiments/vision/train.py @@ -95,8 +95,8 @@ def parse_args() -> argparse.Namespace: 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=4) - parser.add_argument("--prefetch-factor", type=int, default=2) + 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=None) From 1d773e0b5af5fea90e3a8db0d56332a520743c97 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Thu, 30 Apr 2026 16:14:45 +0200 Subject: [PATCH 07/11] Add MLP --- experiments/vision/coord_check_convnext.py | 234 --------------------- experiments/vision/coord_check_mlp.py | 103 +++++++++ experiments/vision/hf_vision_data.py | 59 +++--- experiments/vision/launch_sweep.py | 16 +- experiments/vision/maxp_timm.py | 175 ++++++++------- experiments/vision/run.sh | 12 +- experiments/vision/train.py | 13 +- 7 files changed, 251 insertions(+), 361 deletions(-) delete mode 100644 experiments/vision/coord_check_convnext.py create mode 100644 experiments/vision/coord_check_mlp.py diff --git a/experiments/vision/coord_check_convnext.py b/experiments/vision/coord_check_convnext.py deleted file mode 100644 index c148926..0000000 --- a/experiments/vision/coord_check_convnext.py +++ /dev/null @@ -1,234 +0,0 @@ -#!/usr/bin/env python3 -"""Coord-style diagnostic for a custom width-scalable ConvNeXt model. - -Uses a hand-rolled ConvNeXt-V1 style model with a single base width `w` -(channel progression [w, 2w, 4w, 8w]) so the scaling axis is continuous -and interpretable — unlike discrete timm variants (T/S/B/L). - -Standard PyTorch init; maxP can reinitialize weights freely. -""" - -from __future__ import annotations - -import argparse - -import torch -import torch.nn as nn -import torch.nn.functional as F - -from maxp import Parametrization, ParametrizedModule, diagnose_axis, plot_axis, print_axis - - -# --------------------------------------------------------------------------- -# Model -# --------------------------------------------------------------------------- - -class _LayerNorm2d(nn.Module): - """NCHW wrapper around nn.LayerNorm (no learnable params in the conv path).""" - - def __init__(self, dim: int) -> None: - super().__init__() - self.norm = nn.LayerNorm(dim) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.norm(x.permute(0, 2, 3, 1)).permute(0, 3, 1, 2) - - -class _ConvNeXtBlock(nn.Module): - def __init__(self, dim: int) -> None: - super().__init__() - self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim, bias=True) - self.norm = nn.LayerNorm(dim) - self.pwconv1 = nn.Linear(dim, 4 * dim, bias=True) - self.act = nn.GELU() - self.pwconv2 = nn.Linear(4 * dim, dim, bias=True) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - residual = x - x = self.dwconv(x) - x = x.permute(0, 2, 3, 1) - x = self.norm(x) - x = self.pwconv1(x) - x = self.act(x) - x = self.pwconv2(x) - x = x.permute(0, 3, 1, 2) - return residual + x - - -class ScalableConvNeXt(nn.Module): - """ConvNeXt-V1 with configurable base width w. - - Channel progression: [w, 2w, 4w, 8w]. - No layer scale, no GRN — keeps init simple for maxP. - """ - - def __init__( - self, - w: int, - num_classes: int = 1000, - depths: tuple[int, ...] = (2, 2, 6, 2), - ) -> None: - super().__init__() - dims = [w, 2 * w, 4 * w, 8 * w] - - self.stem = nn.Sequential( - nn.Conv2d(3, dims[0], kernel_size=4, stride=4, bias=True), - _LayerNorm2d(dims[0]), - ) - - self.stages = nn.ModuleList([ - nn.Sequential(*[_ConvNeXtBlock(d) for _ in range(depth)]) - for d, depth in zip(dims, depths) - ]) - - # Between-stage downsampling (3 transitions for 4 stages) - self.downsamples = nn.ModuleList([ - nn.Sequential( - _LayerNorm2d(dims[i]), - nn.Conv2d(dims[i], dims[i + 1], kernel_size=2, stride=2, bias=True), - ) - for i in range(3) - ]) - - self.head_norm = nn.LayerNorm(dims[-1]) - self.head = nn.Linear(dims[-1], num_classes, bias=True) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - x = self.stem(x) - for i, stage in enumerate(self.stages): - x = stage(x) - if i < len(self.downsamples): - x = self.downsamples[i](x) - x = x.mean([-2, -1]) - x = self.head_norm(x) - return self.head(x) - - -# --------------------------------------------------------------------------- -# ParametrizedModule installation -# --------------------------------------------------------------------------- - -def _conv_fan_in(m: nn.Conv2d) -> int: - # Depthwise: scaling axis is channel count, not spatial fan-in. - if m.groups == m.in_channels == m.out_channels: - return m.in_channels - return (m.in_channels // m.groups) * m.kernel_size[0] * m.kernel_size[1] - - -def install_wrappers(model: ScalableConvNeXt) -> list[str]: - wrapped: list[str] = [] - - # Stem Conv2d — embedding - stem_conv = model.stem[0] - model.stem[0] = ParametrizedModule(stem_conv, width_dim=_conv_fan_in(stem_conv), layer_type="embedding") - wrapped.append("stem.0") - - # Stage blocks: dwconv, pwconv1, pwconv2 - for si, stage in enumerate(model.stages): - for bi, block in enumerate(stage): - base = f"stages.{si}.{bi}" - - dw = block.dwconv - block.dwconv = ParametrizedModule(dw, width_dim=_conv_fan_in(dw), layer_type="hidden") - wrapped.append(f"{base}.dwconv") - - pw1 = block.pwconv1 - block.pwconv1 = ParametrizedModule(pw1, width_dim=pw1.in_features, layer_type="hidden") - wrapped.append(f"{base}.pwconv1") - - pw2 = block.pwconv2 - block.pwconv2 = ParametrizedModule(pw2, width_dim=pw2.in_features, layer_type="hidden") - wrapped.append(f"{base}.pwconv2") - - # Downsample Conv2d — hidden - for di, ds in enumerate(model.downsamples): - conv = ds[1] - model.downsamples[di][1] = ParametrizedModule(conv, width_dim=_conv_fan_in(conv), layer_type="hidden") - wrapped.append(f"downsamples.{di}.1") - - # Head Linear — readout - head = model.head - model.head = ParametrizedModule(head, width_dim=head.in_features, layer_type="readout") - wrapped.append("head") - - return wrapped - - -# --------------------------------------------------------------------------- -# Coord check harness -# --------------------------------------------------------------------------- - -WIDTHS = [32, 48, 64, 96, 128] # base channel widths → dims [w, 2w, 4w, 8w] - - -def _make_model(w: int, parametrized: bool): - model = ScalableConvNeXt(w=w, num_classes=1000) - if not parametrized: - return model, None - install_wrappers(model) - param = Parametrization( - model, - lr_prefactor=1e-3, - optimizer_type="adam", - alignment="full", - sample_input=None, - ) - return model, param.param_groups - - -def _make_input(w: int) -> torch.Tensor: - return torch.randn(2, 3, 224, 224) - - -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(2, 3, 224, 224) - y_fixed = torch.randint(0, 1000, (2,)) - - 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 scalable ConvNeXt (continuous width axis)", - formatter_class=argparse.ArgumentDefaultsHelpFormatter, - ) - parser.add_argument("--parametrized", action="store_true") - parser.add_argument("--steps", type=int, default=3) - parser.add_argument("--seeds", type=int, default=2) - parser.add_argument("--plot", action="store_true") - args = parser.parse_args() - - variant = "parametrized" if args.parametrized else "plain" - print(f"Scalable ConvNeXt 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("base_width", ops, affected, act_stats, WIDTHS) - - if args.plot: - out_name = f"coord_check_convnext_scalable_{variant}.png" - plot_axis("base_width", 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_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/hf_vision_data.py b/experiments/vision/hf_vision_data.py index 580fc49..d92ac77 100644 --- a/experiments/vision/hf_vision_data.py +++ b/experiments/vision/hf_vision_data.py @@ -1,4 +1,4 @@ -"""HF streaming data pipeline for vision experiments.""" +"""HF non-streaming data pipeline for vision experiments.""" from __future__ import annotations @@ -7,7 +7,7 @@ import torch from datasets import load_dataset -from torch.utils.data import DataLoader, IterableDataset +from torch.utils.data import DataLoader, Dataset @dataclass(frozen=True, slots=True) @@ -46,55 +46,47 @@ class DatasetPreset: } -class HFVisionIterableDataset(IterableDataset): - """Bridge HF streaming iterable datasets into torch IterableDataset.""" +class HFVisionDataset(Dataset): + """Map-style wrapper around a downloaded HF Dataset.""" - def __init__( - self, - hf_stream, - *, - preset: DatasetPreset, - transform: Callable, - ) -> None: + def __init__(self, hf_dataset, *, preset: DatasetPreset, transform: Callable) -> None: super().__init__() - self._hf_stream = hf_stream + self._hf = hf_dataset self._preset = preset self._transform = transform - def __iter__(self): - for sample in self._hf_stream: - image = sample.get(self._preset.image_key) - image = image.convert("RGB") - label = sample.get(self._preset.label_key) + def __len__(self) -> int: + return len(self._hf) - x = self._transform(image) - y = torch.tensor(label, dtype=torch.long) - - yield x, y + 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: IterableDataset, + ds: Dataset, *, batch_size: int, num_workers: int, is_train: bool, prefetch_factor: int, ) -> DataLoader: - kwargs = { + kwargs: dict = { "dataset": ds, "batch_size": batch_size, "num_workers": num_workers, "pin_memory": torch.cuda.is_available(), "drop_last": is_train, + "shuffle": is_train, "persistent_workers": (num_workers > 0), } if num_workers > 0: kwargs["prefetch_factor"] = prefetch_factor - - return DataLoader( - **kwargs, - ) + return DataLoader(**kwargs) def build_dataloaders( @@ -103,8 +95,6 @@ def build_dataloaders( batch_size: int, num_workers: int, prefetch_factor: int = 2, - seed: int = 1, - train_shuffle_buffer: int = 10_000, train_transform=None, eval_transform=None, ) -> tuple[DataLoader, DataLoader | None, DatasetPreset]: @@ -115,9 +105,8 @@ def build_dataloaders( preset = DATASET_CONFIGS.get(dataset_name) - train_stream = load_dataset(preset.dataset_id, split=preset.train_split, streaming=False) - train_stream = train_stream.shuffle(buffer_size=train_shuffle_buffer, seed=seed) - train_ds = HFVisionIterableDataset(train_stream, preset=preset, transform=train_transform) + 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, @@ -125,9 +114,9 @@ def build_dataloaders( is_train=True, prefetch_factor=prefetch_factor, ) - - val_stream = load_dataset(preset.dataset_id, split=preset.val_split, streaming=False) - val_ds = HFVisionIterableDataset(val_stream, preset=preset, transform=eval_transform) + + 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, diff --git a/experiments/vision/launch_sweep.py b/experiments/vision/launch_sweep.py index 297212e..a0dcfe8 100755 --- a/experiments/vision/launch_sweep.py +++ b/experiments/vision/launch_sweep.py @@ -50,8 +50,8 @@ "image_size": 224, "drop_path_rate": 0.4, }, - "cnx-t": { - "wall": "24:00:00", + "mlp-s": { + "wall": "6:00:00", "gpus": 1, "methods": ["maxP", "mup-full", "mup-no"], "lrs": ALL_LRS, @@ -59,8 +59,8 @@ "image_size": 224, "drop_path_rate": 0.1, }, - "cnx-s": { - "wall": "36:00:00", + "mlp-m": { + "wall": "12:00:00", "gpus": 1, "methods": ["maxP", "mup-full", "mup-no"], "lrs": ALL_LRS, @@ -68,8 +68,8 @@ "image_size": 224, "drop_path_rate": 0.2, }, - "cnx-b": { - "wall": "48:00:00", + "mlp-b": { + "wall": "12:00:00", "gpus": 1, "methods": ["maxP", "mup-full", "mup-no"], "lrs": TRANSFER_LR, @@ -77,8 +77,8 @@ "image_size": 224, "drop_path_rate": 0.3, }, - "cnx-l": { - "wall": "72:00:00", + "mlp-l": { + "wall": "24:00:00", "gpus": 1, "methods": ["maxP", "mup-full", "mup-no"], "lrs": TRANSFER_LR, diff --git a/experiments/vision/maxp_timm.py b/experiments/vision/maxp_timm.py index beeda4f..c460963 100644 --- a/experiments/vision/maxp_timm.py +++ b/experiments/vision/maxp_timm.py @@ -19,7 +19,10 @@ class ScaleConfig: model_name: str family: str image_size: int - drop_path_rate: float + drop_path_rate: float = 0.0 + hidden: int = 0 + depth: int = 0 + dropout: float = 0.0 SCALE_CONFIGS: dict[str, ScaleConfig] = { @@ -35,21 +38,82 @@ class ScaleConfig: "vit-l": ScaleConfig( model_name="vit_large_patch16_224", family="vit", image_size=224, drop_path_rate=0.4 ), - "cnx-t": ScaleConfig( - model_name="convnextv2_tiny", family="convnext", image_size=224, drop_path_rate=0.1 + "mlp-s": ScaleConfig( + model_name="mlp", family="mlp", image_size=224, hidden=256, depth=4, dropout=0.0 ), - "cnx-s": ScaleConfig( - model_name="convnextv2_small", family="convnext", image_size=224, drop_path_rate=0.2 + "mlp-m": ScaleConfig( + model_name="mlp", family="mlp", image_size=224, hidden=512, depth=6, dropout=0.1 ), - "cnx-b": ScaleConfig( - model_name="convnextv2_base", family="convnext", image_size=224, drop_path_rate=0.3 + "mlp-b": ScaleConfig( + model_name="mlp", family="mlp", image_size=224, hidden=1024, depth=8, dropout=0.2 ), - "cnx-l": ScaleConfig( - model_name="convnextv2_large", family="convnext", image_size=224, drop_path_rate=0.4 + "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) @@ -98,20 +162,6 @@ def forward( x = self.proj(x) x = self.proj_drop(x) return 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 _LinearPatchEmbed(PatchEmbed): @@ -119,7 +169,7 @@ 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) + self.patchify = Patchify(self.patch_size) def forward(self, x): _, _, H, W = x.shape @@ -150,30 +200,25 @@ def create_model( ) -> nn.Module: cfg = SCALE_CONFIGS.get(scale) img_size = image_size if image_size is not None else cfg.image_size - resolved_drop_path = cfg.drop_path_rate + + 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": resolved_drop_path, + "drop_path_rate": cfg.drop_path_rate, + "img_size": img_size, + "attn_layer": _ScaledAttention, + "embed_layer": _LinearPatchEmbed, } - - if cfg.family == "vit": - kwargs["img_size"] = img_size - kwargs["attn_layer"] = _ScaledAttention - kwargs["embed_layer"] = _LinearPatchEmbed - - model = timm.create_model(cfg.model_name, **kwargs) - return model - - -def _conv_width_dim(module: nn.Conv2d) -> int: - # For depthwise convs (groups == in_channels), fan-in stays constant (k*k), - # so use channel width to preserve scaling across model sizes. - if module.groups == module.in_channels and module.in_channels == module.out_channels: - return module.in_channels - k_h, k_w = module.kernel_size - return (module.in_channels // module.groups) * k_h * k_w + return timm.create_model(cfg.model_name, **kwargs) def _get_child(module: nn.Module, key: str) -> nn.Module: @@ -233,46 +278,24 @@ def _install_vit_wrappers(model: nn.Module): _wrap_module(model, "head", embed_dim, "readout", a=0.0) -def _convnext_layer_type(name: str) -> str: - if name.startswith("stem") or name.startswith("downsample_layers.0"): - return "embedding" - if name.startswith("head"): - return "readout" - return "hidden" - - -def _install_convnext_wrappers(model: nn.Module) -> list[str]: - wrapped: list[str] = [] - targets: list[tuple[str, int, str]] = [] - - for name, module in model.named_modules(): - if isinstance(module, ParametrizedModule): - continue - if isinstance(module, nn.Linear): - targets.append((name, _linear_width_dim(module), _convnext_layer_type(name))) - elif isinstance(module, nn.Conv2d): - targets.append((name, _conv_width_dim(module), _convnext_layer_type(name))) - - for module_path, width_dim, layer_type in targets: - if _wrap_module( - model, - module_path, - width_dim=width_dim, - layer_type=layer_type, - ): - wrapped.append(module_path) +def _install_mlp_wrappers(model: ScalableMLP) -> None: + hidden = model.embed.out_features + + for block in model.blocks: + _wrap_module(block, "linear", hidden, "hidden") - return wrapped + _wrap_module(model, "embed", hidden, "embedding") + _wrap_module(model, "head", hidden, "readout", a=0.0) def install_pm_wrappers(model: nn.Module): - if "visiontransformer" in model.__class__.__name__.lower(): + if isinstance(model, ScalableMLP): + _install_mlp_wrappers(model) + elif "visiontransformer" in model.__class__.__name__.lower(): _install_vit_wrappers(model) - elif "convnext" in model.__class__.__name__.lower(): - _install_convnext_wrappers(model) else: raise TypeError( - "Unsupported timm model type. Expected ViT or ConvNeXt-like model." + f"Unsupported model type: {model.__class__.__name__}. Expected ViT or ScalableMLP." ) diff --git a/experiments/vision/run.sh b/experiments/vision/run.sh index e080268..0ed5c1e 100755 --- a/experiments/vision/run.sh +++ b/experiments/vision/run.sh @@ -5,7 +5,7 @@ # bash experiments/vision/run.sh vit-s [extra launch_sweep.py flags ...] set -euo pipefail -SCALE="${1:?Usage: $0 [extra launch_sweep.py flags]}" +SCALE="${1:?Usage: $0 [extra launch_sweep.py flags]}" shift case "$SCALE" in @@ -13,11 +13,11 @@ case "$SCALE" in 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="${TRANSFER_LR:-1e-2}"; SEEDS="1" ;; - cnx-t) 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" ;; - cnx-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" ;; - cnx-b) METHODS="maxP mup-full mup-no"; LRS="${TRANSFER_LR:-1e-2}"; SEEDS="1" ;; - cnx-l) METHODS="maxP mup-full mup-no"; LRS="${TRANSFER_LR:-1e-2}"; SEEDS="1" ;; - *) echo "Unknown scale '$SCALE'. Choose from debug vit-s vit-b vit-l cnx-t cnx-s cnx-b cnx-l."; exit 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" ;; + mlp-b) METHODS="maxP mup-full mup-no"; LRS="${TRANSFER_LR:-1e-2}"; SEEDS="1" ;; + mlp-l) METHODS="maxP mup-full mup-no"; LRS="${TRANSFER_LR:-1e-2}"; 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}" diff --git a/experiments/vision/train.py b/experiments/vision/train.py index ca45c06..dcaad7e 100644 --- a/experiments/vision/train.py +++ b/experiments/vision/train.py @@ -147,7 +147,17 @@ def main() -> None: model = model.to(device) model.train() - model_data_cfg = resolve_model_data_config(model) + 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) @@ -156,7 +166,6 @@ def main() -> None: batch_size=args.batch_size, num_workers=args.num_workers, prefetch_factor=args.prefetch_factor, - seed=args.seed, train_transform=train_transform, eval_transform=eval_transform, ) From ee771a292bd9d0cf7feb58fe4c939950961f4746 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Thu, 30 Apr 2026 16:17:50 +0200 Subject: [PATCH 08/11] Adjust run config --- experiments/vision/run.sh | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/experiments/vision/run.sh b/experiments/vision/run.sh index 0ed5c1e..0d13d7c 100755 --- a/experiments/vision/run.sh +++ b/experiments/vision/run.sh @@ -12,11 +12,11 @@ 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="${TRANSFER_LR:-1e-2}"; 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" ;; - mlp-b) METHODS="maxP mup-full mup-no"; LRS="${TRANSFER_LR:-1e-2}"; SEEDS="1" ;; - mlp-l) METHODS="maxP mup-full mup-no"; LRS="${TRANSFER_LR:-1e-2}"; SEEDS="1" ;; + 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 From 5a8af7e28e8ea9af227d5a5722e47833417d11b2 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Fri, 1 May 2026 00:40:16 +0200 Subject: [PATCH 09/11] Improve training performance --- experiments/vision/hf_vision_data.py | 1 - experiments/vision/train.py | 15 +++++++++------ 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/experiments/vision/hf_vision_data.py b/experiments/vision/hf_vision_data.py index d92ac77..718f679 100644 --- a/experiments/vision/hf_vision_data.py +++ b/experiments/vision/hf_vision_data.py @@ -81,7 +81,6 @@ def make_loader( "num_workers": num_workers, "pin_memory": torch.cuda.is_available(), "drop_last": is_train, - "shuffle": is_train, "persistent_workers": (num_workers > 0), } if num_workers > 0: diff --git a/experiments/vision/train.py b/experiments/vision/train.py index dcaad7e..868c707 100644 --- a/experiments/vision/train.py +++ b/experiments/vision/train.py @@ -99,7 +99,7 @@ def parse_args() -> argparse.Namespace: 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=None) + 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) @@ -282,7 +282,7 @@ def main() -> None: last_log_seen = total_seen if global_step % args.val_interval == 0: - val_interval = evaluate( + eval_metrics = evaluate( model=model, loader=val_loader, device=device, @@ -291,7 +291,7 @@ def main() -> None: num_classes=num_classes, ) else: - val_interval = {} + eval_metrics = {} row: dict[str, Any] = { "step": global_step, @@ -302,7 +302,7 @@ def main() -> None: "perf/samples_per_sec_interval": interval_sps, "lr/lr_prefactor": current_lr_prefactor(optimizer), **collect_layer_lrs(param), - **val_interval, + **eval_metrics, } if dynamic: row.update(collect_alignments(param)) @@ -316,9 +316,12 @@ def main() -> None: 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"[step {global_step}] loss={row['loss/train_loss']:.4f}{val_txt}" - f" sps={interval_sps:.1f}" + 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: From a468c688179c2bc84e0f1eafbf6134525ddd28a2 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Mon, 4 May 2026 21:53:12 +0200 Subject: [PATCH 10/11] Update README and remove paramR submodule Restructure experiment docs: move file trees into dedicated ### Files subsections under LLM and Vision Experiments sections. Add Vision Experiments section with corrected file list (ViT/MLP, split coord_check, add utils.py). Remove stale paramR gitlink. Co-Authored-By: Claude Sonnet 4.6 --- README.md | 78 ++++++++++++++++++++++++++++++++++--------------------- paramR | 1 - 2 files changed, 48 insertions(+), 31 deletions(-) delete mode 160000 paramR diff --git a/README.md b/README.md index 4b81e43..97d67e0 100644 --- a/README.md +++ b/README.md @@ -77,42 +77,12 @@ 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 -├── 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 - -experiments/vision/ # ViT/ConvNeXt-V2 experiments (timm) -├── 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 -├── launch_sweep.py # Generate and submit SLURM jobs for full LR sweep -├── run.sh # Submit a sweep for one vision scale -├── run_debug.sh # Quick tiny-model debug run (CPU/GPU) -└── coord_check.py # Lightweight coord-style diagnostic for wrapped models ``` **Phase 1** (static, at init): discover PMs → reinit weights → solve LP → build param groups. **Phase 2** (dynamic, per step): capture activations → measure alignment → re-solve LP → update per-layer LRs. -### Vision experiment quick start - -```bash -# CPU/GPU debug smoke test (small model + tiny dataset slice) -bash experiments/vision/run_debug.sh --steps 20 - -# Launch full SLURM sweep for one scale -bash experiments/vision/run.sh vit-s -``` - ## ParametrizedModule @@ -316,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 @@ -331,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/paramR b/paramR deleted file mode 160000 index f1934d9..0000000 --- a/paramR +++ /dev/null @@ -1 +0,0 @@ -Subproject commit f1934d99f5df75dba063f36f9a740064ee6482f4 From 5829b75b78855b710d6d91a157b1218affedf8f8 Mon Sep 17 00:00:00 2001 From: Maksymilian Wojnar Date: Mon, 4 May 2026 21:54:26 +0200 Subject: [PATCH 11/11] No conv2d --- maxp/trace.py | 1 - 1 file changed, 1 deletion(-) diff --git a/maxp/trace.py b/maxp/trace.py index 277e975..dc2de58 100644 --- a/maxp/trace.py +++ b/maxp/trace.py @@ -177,7 +177,6 @@ def _collect(obj): torch.addmm: "addmm", F.linear: "linear", F.embedding: "embedding", - F.conv2d: "conv2d", }