diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index b123b20..f5abaa0 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -44,7 +44,7 @@ jobs: python-version: ${{ matrix.python-version }} - name: Run tests - run: uv run --frozen --extra pytorch python -m pytest tests --cov --cov-config=pyproject.toml --cov-report=xml + run: uv run --frozen --extra pytorch --extra vit python -m pytest tests --cov --cov-config=pyproject.toml --cov-report=xml env: # Ray 2.58+ otherwise creates worker environments from the bare # `uv run` command and drops the optional PyTorch dependencies. diff --git a/README.md b/README.md index 0f66866..48f69a4 100644 --- a/README.md +++ b/README.md @@ -76,6 +76,13 @@ python -m tinyexp.examples.mnist_exp The commands below assume that the environment containing TinyExp is active. For a source checkout, run `source .venv/bin/activate` first. +A larger bundled example trains DeiT-S with tensor parallelism, ported line-by-line +against the official DeiT implementation (see `docs/vit_tp.md`): + +```bash +python -m tinyexp.examples.vit_tp_exp mode=bench # throughput/memory check on 2 GPUs +``` + Run MNIST with config override: ```bash diff --git a/docs/vit_tp.md b/docs/vit_tp.md new file mode 100644 index 0000000..ef8323d --- /dev/null +++ b/docs/vit_tp.md @@ -0,0 +1,154 @@ +# DeiT-S with Tensor Parallelism (`vit_tp_exp`) + +`tinyexp/examples/vit_tp_exp.py` trains and evaluates **DeiT-S** with **tensor +parallelism** (TP), implemented against the official DeiT sources line by line +(对拍): only formatting-level changes are allowed, and every deviation from upstream is +listed below and proven numerically equivalent by tests. This page is the complete +plan-and-record home for the example: ported sources, sanctioned deviations, usage, +and the four-layer cross-check record (numerics, eval accuracy, full-training +accuracy, throughput). + +## Why DeiT-S (and not DeiT-Ti) + +Tensor parallelism shards attention on **head boundaries**, so `num_heads` must be +divisible by the TP size. DeiT-Ti has 3 heads — structurally unusable at TP=2. DeiT-S has +6 heads (3 per rank at TP=2), 22,050,664 parameters, and the official 300-epoch recipe +reaches **79.8% top-1** on ImageNet val, which is the accuracy cross-check target. + +## Ported sources (pinned) + +| Component | Source | +| --- | --- | +| Model (`VisionTransformer`/`Block`/`Attention`/`Mlp`/`PatchEmbed`) | timm `vision_transformer` implementation configured for timm 0.3.2 DeiT behavior; `Attention` is adapted locally to split QKV for TP | +| Factory hyper-parameters, checkpoint URL | facebookresearch/deit `models.py` @ `7e160fe4` | +| Train/eval loops, recipe defaults, transforms, `RASampler`, `MetricLogger` | deit `engine.py`/`main.py`/`datasets.py`/`samplers.py`/`utils.py` (same commit) | +| Data/loss/optimizer helpers (`Mixup`, `create_transform`, `NativeScaler`, `create_scheduler`, …) | timm >= 1.0 (same algorithms) | + +TP itself is `torch.distributed.tensor.parallel` (`ColwiseParallel`/`RowwiseParallel`) +wrapped by `TPAccelerator`, following the classic Megatron block layout: colwise on +`attn.q/k/v` + `mlp.fc1`, rowwise on `attn.proj` + `mlp.fc2`, one all-reduce per block +pass. Attention runs per-head locally with no communication inside, using PyTorch's +`scaled_dot_product_attention` (FlashAttention backend on supported CUDA devices). + +## Sanctioned deviations (each covered by a test) + +- **D1** — the official fused `attn.qkv` Linear is split into `q`/`k`/`v` so colwise + sharding lands on head boundaries. `convert_fused_qkv_to_split` / + `convert_split_qkv_to_fused` translate weights; checkpoints always store the official + fused layout and load official DeiT checkpoints directly. +- **D2** — `Attention.forward` reshapes use `-1` (local-shape aware), as torch's + `use_local_output=True` TP pattern requires; at tp=1 the shapes resolve to exactly the + official ones. +- **D3** — `mode=bench` is an added capability (synthetic throughput/memory check); it + does not touch the training semantics. +- **D4** — EMA (official default on) is disabled at tp>1: an EMA copy holds full weights + and per-step gathering from DTensors is expensive. +- **D5/D6** — under TP all ranks consume identical batches and share one RNG seed (TP + collectives require identical inputs; the official `seed + rank` is the DDP + data-sharding context). +- **D7** — optional Redis byte cache for the training set + (`redis_cfg.redis_cache_enabled`, default **off**; ray-managed Redis, no external + wrapper needed): the cache stores raw JPEG bytes verbatim, so samples are + bit-identical to `datasets.ImageFolder`; the first epoch fills the cache from disk, + later epochs read from memory. `safe_set` failures (e.g. maxmemory full) fall back to + disk reads — correctness never depends on the cache. + +## Usage + +Full recipe (2 GPUs, TP=2, official defaults: 300 epochs, AdamW wd 0.05, cosine with +5-epoch warmup, mixup 0.8/cutmix 1.0, repeated augmentation): + +```bash +python -m tinyexp.examples.vit_tp_exp +``` + +Eval-only against the official checkpoint (accuracy cross-check, expects 79.8 ± 0.1): + +```bash +python -m tinyexp.examples.vit_tp_exp mode=eval \ + module_cfg.pretrained_from=https://dl.fbaipublicfiles.com/deit/deit_small_patch16_224-cd65a155.pth +``` + +Throughput / memory benchmark (compare against DDP by overriding the accelerator): + +```bash +python -m tinyexp.examples.vit_tp_exp mode=bench +python -m tinyexp.examples.vit_tp_exp mode=bench accelerator_cfg.accelerator=ddp +``` + +CPU smoke test (no ImageNet, ray 2 workers, gloo TP=2): + +```bash +python -m tinyexp.examples.vit_tp_exp dataloader_cfg.fake_data=true dataloader_cfg.input_size=32 \ + module_cfg.img_size=32 module_cfg.num_classes=10 epochs=1 max_train_steps=4 \ + ray_cfg.ray_num_gpus_per_worker=0.0 ray_cfg.ray_num_cpus_per_worker=4 +``` + +Full 300-epoch training on many GPUs (e.g. 8×H200/H800 single node, DDP — DeiT-S has +6 heads and cannot shard over TP=8), with the Redis byte cache for slow shared +filesystems (D7; the ray mixin starts the Redis shards itself): + +```bash +python -m tinyexp.examples.vit_tp_exp accelerator_cfg.accelerator=ddp \ + ray_cfg.ray_num_worker=8 redis_cfg.redis_cache_enabled=true redis_cfg.redis_cache_max_memory=300 +``` + +Recipe note (global batch / lr): the official linear rule `lr × global_batch / 512` +is applied verbatim — the default 2×256=512 gives lr 5e-4 (the same scaling point as +the official 1024 @ 1e-3), and the 8×256=2048 H200 run below used lr 2e-3. + +## Data preparation + +ImageNet-1k in the standard `ImageFolder` layout, located via `IMAGENET_HOME` or +`dataloader_cfg.data_root` (default `./data/imagenet/`): + +``` +/train//*.JPEG # 1.28M images, 1000 class dirs (full training) +/val//*.JPEG # 50k images (eval-only cross-check) +``` + +## Cross-check record + +Four layers, each with a hard pass criterion; all hit. + +- **L1 — numerical equivalence** (CI, CPU+gloo): TP=2 forward logits and one-step + gradients match the single-process run at fp32 `allclose(atol=1e-5)` + (`tests/examples/test_vit_tp_exp_tp.py`), and the split-qkv forward matches the + fused official one at `atol=1e-6`. +- **L2 — eval accuracy** (official checkpoint, full 50k val): **79.82% top-1 / + 94.95% top-5 at tp=1**, **79.81% / 94.94% at tp=2** (official reference 79.8; + tp=1 vs tp=2 differ by 0.01%). +- **L3 — full-training accuracy** (300-epoch official recipe from scratch, 8×H200 + DDP via rjob, global batch 2048 / lr 2e-3, Redis byte cache on): **best top-1 + 79.84% / top-5 94.99%** in 17h56m (~3.5 min/epoch incl. full-val each epoch) — + paper reports 79.8 / ~95.0. Trajectory: ep45 58.6 → ep113 69.7 → ep181 74.1 → + ep249 78.4 → ep299 79.84. Artifacts: `output/vit_tp_l3_h200/{last,best}.ckpt` + (official fused-qkv layout), per-epoch `log.txt`. +- **L4 — throughput/memory** (bench mode, table below). + +Two engine hardenings landed with this example, each held by a regression test: +`DDPAccelerator.reduce_sum` stages CPU tensors through the device so nccl +metric-sync works at epoch boundaries, and `store_and_run_exp` resolves the +canonical importable twin of a `python -m` `__main__` class so ray ships it by +reference (robust against cluster agents that replace `builtins.print` after +import). + +Redis cache behavior measured on the L3 run: `RASampler(num_repeats=3)` visits +~1/3 unique indices per epoch, so the cache fills over ~3 epochs to the full +1,281,167-key train set (6 shards × 213,528); once warm, data time drops to +~0.6 ms/step and steps run at ~0.25 s. + +## Measured on 2× RTX 4080 (DeiT-S, 224, AMP) + +| Config | batch/rank | step latency | effective img/s | peak mem/rank | +| --- | --- | --- | --- | --- | +| DDP=2 | 256 (its ceiling; 288 OOM) | 294.7 ms | 1737 | 13.44 GB | +| TP=2 | 256 | 841.0 ms | 304 | **8.61 GB** | +| TP=2 | 384 (its ceiling; 512 OOM) | 1262.5 ms | 304 | 12.87 GB | + +Honest reading: at DeiT-S scale TP is **not** a throughput win (communication dominates; +~17% of DDP). What TP buys is per-device memory — a 1.5× larger per-rank batch ceiling +(384 vs 256) on the same 16 GB cards. The correctness gate is met: evaluating the official checkpoint on the full 50k +ImageNet val yields **79.82% top-1 / 94.95% top-5 at tp=1** and **79.81% / 94.94% at +tp=2** (official reference: 79.8%); numerical equivalence of TP=2 vs single-rank +forward/backward is asserted in CI (fp32 `allclose`, `tests/examples/test_vit_tp_exp_tp.py`). diff --git a/mkdocs.yml b/mkdocs.yml index 04626da..3d08b25 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -10,6 +10,7 @@ copyright: Maintained by zengarden. nav: - Home: index.md - Running Modes: running-modes.md + - DeiT-S Tensor Parallel Example: vit_tp.md - Philosophy: philosophy.md - Modules: modules.md plugins: diff --git a/pyproject.toml b/pyproject.toml index 56d2355..4c482c3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,6 +54,16 @@ pytorch = [ # TorchVision pulls Pillow; use a release with CPython 3.14 wheels. "pillow>=11.3; python_version >= '3.14'", ] +vit = [ + # vit_tp_exp example: data pipeline (Mixup/RandAugment), losses, scheduler, and + # DropPath/trunc_normal_ used by the ported DeiT model. The model itself is + # ported verbatim from timm 0.3.2 into tinyexp/examples/vit_tp_exp.py. + # timm does not pull torch, and the example imports torch/torchvision directly, + # so pull the pytorch extra via a self-referential extra instead of duplicating + # its per-interpreter version pins. + "TinyExp[pytorch]", + "timm>=1.0", +] [project.urls] Homepage = "https://zengarden.github.io/TinyExp/" @@ -161,6 +171,12 @@ ignore = [ [tool.ruff.lint.per-file-ignores] "tests/*" = ["S101"] "tinyexp/utils/redis_utils.py" = ["S603"] +# Ported verbatim from facebookresearch/deit and timm 0.3.2 (VIT_TOD.md §3.1): keep the +# upstream asserts/exception messages/int() wrappers instead of restyling official code. +"tinyexp/dataset/ra_sampler.py" = ["S101", "TRY003", "RUF046"] +"tinyexp/examples/vit_tp_exp.py" = ["S101", "TRY003"] +# The RASampler test oracle recomputes the official index math verbatim. +"tests/dataset/test_ra_sampler.py" = ["RUF046"] [tool.ruff.format] preview = true diff --git a/tests/dataset/test_ra_sampler.py b/tests/dataset/test_ra_sampler.py new file mode 100644 index 0000000..a481455 --- /dev/null +++ b/tests/dataset/test_ra_sampler.py @@ -0,0 +1,81 @@ +"""Tests for the ported RASampler (docs/vit_tp.md). + +The oracle recomputes the official index math from facebookresearch/deit ``samplers.py`` +inline, so any accidental change to the ported sampling logic shows up as a mismatch. +""" + +from __future__ import annotations + +import math + +import pytest +import torch + +from tinyexp.dataset.ra_sampler import RASampler + + +class _LenOnlyDataset: + def __init__(self, length: int) -> None: + self.length = length + + def __len__(self) -> int: + return self.length + + +def _official_indices(dataset_len: int, num_replicas: int, rank: int, epoch: int, num_repeats: int) -> list[int]: + """Verbatim re-implementation of the official RASampler iteration math.""" + g = torch.Generator() + g.manual_seed(epoch) + indices = torch.randperm(dataset_len, generator=g) + indices = torch.repeat_interleave(indices, repeats=num_repeats, dim=0).tolist() + num_samples = int(math.ceil(dataset_len * num_repeats / num_replicas)) + total_size = num_samples * num_replicas + padding_size = total_size - len(indices) + if padding_size > 0: + indices += indices[:padding_size] + indices = indices[rank:total_size:num_replicas] + num_selected = int(math.floor(dataset_len // 256 * 256 / num_replicas)) + return indices[:num_selected] + + +@pytest.mark.parametrize("num_repeats", [1, 3]) +@pytest.mark.parametrize("rank", [0, 1]) +def test_matches_official_index_math(num_repeats: int, rank: int) -> None: + dataset = _LenOnlyDataset(1000) + sampler = RASampler(dataset, num_replicas=2, rank=rank, shuffle=True, num_repeats=num_repeats) + sampler.set_epoch(7) + + assert list(sampler) == _official_indices(1000, 2, rank, epoch=7, num_repeats=num_repeats) + assert len(sampler) == 384 # floor(1000 // 256 * 256 / 2) + + +def test_epoch_changes_permutation_deterministically() -> None: + dataset = _LenOnlyDataset(1000) + sampler = RASampler(dataset, num_replicas=2, rank=0) + first = list(sampler) + + sampler.set_epoch(0) + assert list(sampler) == first # default epoch is 0 + sampler.set_epoch(1) + assert list(sampler) != first + sampler.set_epoch(0) + assert list(sampler) == first + + +def test_no_index_exceeds_num_repeats_copies() -> None: + dataset = _LenOnlyDataset(1000) + rank0 = RASampler(dataset, num_replicas=2, rank=0, num_repeats=3) + rank1 = RASampler(dataset, num_replicas=2, rank=1, num_repeats=3) + rank0.set_epoch(2) + rank1.set_epoch(2) + + counts: dict[int, int] = {} + for index in [*rank0, *rank1]: + counts[index] = counts.get(index, 0) + 1 + assert counts # sanity: the merged stream is non-empty + assert max(counts.values()) <= 3 + + +def test_num_repeats_must_be_positive() -> None: + with pytest.raises(ValueError): + RASampler(_LenOnlyDataset(10), num_replicas=2, rank=0, num_repeats=0) diff --git a/tests/examples/test_vit_tp_exp_run.py b/tests/examples/test_vit_tp_exp_run.py new file mode 100644 index 0000000..e14c0eb --- /dev/null +++ b/tests/examples/test_vit_tp_exp_run.py @@ -0,0 +1,148 @@ +"""Run-level tests for the vit_tp_exp example (docs/vit_tp.md). + +The train/eval/bench paths are exercised end to end on CPU with fake data (single +process, launcher=mp); the TP-specific numerical equivalence lives in +test_vit_tp_exp_tp.py. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +pytest.importorskip("timm", reason="tinyexp vit extra (timm) is required for the vit_tp_exp example") + +import torch + +from tinyexp.examples.vit_tp_exp import VitTpExp + + +def _tiny_exp(tmp_path: Path, mode: str = "train", **overrides) -> VitTpExp: + defaults = { + "mode": mode, + "launcher": "mp", + "output_root": str(tmp_path), + "exp_name": "vit_tp_test", + "epochs": 1, + "accelerator_cfg": VitTpExp.AcceleratorCfg(accelerator="cpu"), + "module_cfg": VitTpExp.ModuleCfg(img_size=32, num_classes=10), + "dataloader_cfg": VitTpExp.DataloaderCfg( + fake_data=True, + fake_data_len=16, + input_size=32, + train_batch_size_per_device=4, + val_batch_size_per_device=4, + num_workers=0, + pin_mem=False, + ), + "bench_cfg": VitTpExp.BenchCfg(warmup_steps=1, measure_steps=2), + } + defaults.update(overrides) + return VitTpExp(**defaults) + + +def test_train_smoke_produces_official_layout_checkpoint(tmp_path: Path) -> None: + exp = _tiny_exp(tmp_path) + exp.run() + + run_dir = Path(exp.get_run_dir()) + stats_lines = (run_dir / "log.txt").read_text().strip().splitlines() + assert len(stats_lines) == 1 + stats = json.loads(stats_lines[0]) + assert stats["epoch"] == 0 + assert "train_loss" in stats and "test_acc1" in stats + + checkpoint = torch.load(run_dir / "last.ckpt", map_location="cpu", weights_only=False) + model_state = checkpoint["model_state_dict"] + # official fused-qkv layout: byte-compatible with facebookresearch/deit checkpoints + assert "blocks.0.attn.qkv.weight" in model_state + assert not any(key.endswith(".attn.q.weight") for key in model_state) + assert checkpoint["optimizer_state_dict"] + + +def test_eval_mode_reproduces_training_accuracy(tmp_path: Path) -> None: + train_exp = _tiny_exp(tmp_path) + train_exp.run() + best_metric = torch.load(Path(train_exp.get_run_dir()) / "best.ckpt", map_location="cpu", weights_only=False)[ + "best_metric" + ] + + eval_exp = _tiny_exp(tmp_path, mode="eval", resume_from=str(Path(train_exp.get_run_dir()) / "best.ckpt")) + eval_exp.run() + + # deterministic protocol: same checkpoint, sequential full val -> same number + assert eval_exp._run_result.startswith(f"eval acc@1={best_metric:.2f}%") + + +def test_eval_mode_requires_checkpoint_source(tmp_path: Path) -> None: + exp = _tiny_exp(tmp_path, mode="eval") + with pytest.raises(ValueError, match="resume_from"): + exp.run() + + +def test_bench_mode_reports_result(tmp_path: Path) -> None: + exp = _tiny_exp(tmp_path, mode="bench") + exp.run() + assert exp._run_result.startswith("bench[cpu world=1]") + + +def test_ray_export_survives_builtins_print_patch(tmp_path: Path) -> None: + """Platform agents may replace ``builtins.print`` between the ``python -m`` + module execution and the ray actor export; shipping the ``__main__``-defined + exp class by value then dies with "Can't pickle : + it's not the same object as builtins.print". ``store_and_run_exp`` must + resolve the canonical importable twin so both the actor class and the + structured-config metadata ship by reference.""" + import subprocess + import sys + + driver = ( + "import builtins, importlib.util, sys\n" + "import ray\n" + "_real_init = ray.init\n" + "def _init_then_patch_print(*args, **kwargs):\n" + " context = _real_init(*args, **kwargs)\n" + " # platform-agent style patch AFTER the exp module ran as __main__\n" + " # (classes captured the real print), BEFORE actor export.\n" + " builtins.print = lambda *a, _r=builtins.print, **k: _r(*a, **k)\n" + " return context\n" + "ray.init = _init_then_patch_print\n" + "# mirror ``python -m``: the running __main__ carries the module spec\n" + 'sys.modules["__main__"].__spec__ = importlib.util.find_spec("tinyexp.examples.vit_tp_exp")\n' + "import runpy\n" + 'sys.argv = ["vit_tp_exp.py"] + sys.argv[1:]\n' + 'runpy.run_module("tinyexp.examples.vit_tp_exp", run_name="__main__")\n' + 'print("RESULT:ok")\n' + ) + run_dir = tmp_path / "run" + result = subprocess.run( # noqa: S603 + [ + sys.executable, + "-c", + driver, + "mode=bench", + "launcher=ray", + "accelerator_cfg.accelerator=cpu", + "module_cfg.img_size=32", + "module_cfg.num_classes=10", + "dataloader_cfg.input_size=32", + "dataloader_cfg.train_batch_size_per_device=4", + "bench_cfg.warmup_steps=1", + "bench_cfg.measure_steps=1", + "ray_cfg.ray_num_worker=1", + "ray_cfg.ray_num_gpus_per_worker=0.0", + "ray_cfg.ray_num_cpus_per_worker=2", + f"hydra.run.dir={run_dir}", + f"output_root={tmp_path}", + ], + cwd=str(Path(__file__).resolve().parents[2]), + text=True, + capture_output=True, + check=False, + timeout=180, + ) + assert result.returncode == 0, f"stdout:\n{result.stdout}\n\nstderr:\n{result.stderr}" + assert "RESULT:ok" in result.stdout + assert "PicklingError" not in result.stderr diff --git a/tests/examples/test_vit_tp_exp_tp.py b/tests/examples/test_vit_tp_exp_tp.py new file mode 100644 index 0000000..c702704 --- /dev/null +++ b/tests/examples/test_vit_tp_exp_tp.py @@ -0,0 +1,128 @@ +"""L1 numerical-equivalence tests for the TP code path (docs/vit_tp.md L1 / §9 M1). + +Two gloo workers on CPU run the tiny ViT with ``TPAccelerator`` (tp=2) and the result +must match the single-process unparallelized forward/backward within fp32 tolerance. +""" + +from __future__ import annotations + +import os +import socket + +import pytest + +pytest.importorskip("timm", reason="tinyexp vit extra (timm) is required for the vit_tp_exp example") + +import torch +import torch.multiprocessing as mp +from torch.distributed.tensor import DTensor + +from tinyexp.examples.vit_tp_exp import VisionTransformer, build_tp_parallelize_plan +from tinyexp.tiny_engine.accelerator import TPAccelerator + +ATOL = 1e-5 + + +def _tiny_model() -> VisionTransformer: + torch.manual_seed(42) + model = VisionTransformer( + img_size=32, patch_size=16, embed_dim=64, depth=2, num_heads=4, num_classes=10, qkv_bias=True + ) + return model.eval() + + +def _free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return sock.getsockname()[1] + + +def _reference(state_dict: dict, x: torch.Tensor) -> dict: + model = _tiny_model() + model.load_state_dict(state_dict) + with torch.no_grad(): + logits = model(x).clone() + loss = model(x).square().mean() + loss.backward() + grads = {name: param.grad.detach().clone() for name, param in model.named_parameters()} + return {"logits": logits, "grads": grads} + + +def _tp_worker(rank: int, world_size: int, port: int, state_dict: dict, x: torch.Tensor, out_path: str) -> None: + # Force the CPU/gloo path of TPAccelerator even on machines with GPUs. + os.environ["CUDA_VISIBLE_DEVICES"] = "" + os.environ["RANK"] = str(rank) + os.environ["WORLD_SIZE"] = str(world_size) + os.environ["LOCAL_RANK"] = str(rank) + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = str(port) + + accelerator = TPAccelerator() + try: + model = _tiny_model() + model.load_state_dict(state_dict) + plan = build_tp_parallelize_plan(model, tp_size=world_size) + model = accelerator.prepare_model(model, parallelize_plan=plan) + model.eval() + + with torch.no_grad(): + logits = model(x).clone() + + # drop_rate/drop_path_rate are 0 in the tiny model, so train() stays deterministic. + model.train() + loss = model(x).square().mean() + accelerator.backward(loss) + grads = {} + for name, param in model.named_parameters(): + grad = param.grad + if isinstance(grad, DTensor): + grad = grad.full_tensor() + grads[name] = grad.detach().cpu() + + if accelerator.is_main_process: + torch.save({"logits": logits.cpu(), "grads": grads}, out_path) + accelerator.wait_for_everyone() + finally: + accelerator.destroy() + + +def test_tp2_forward_and_grads_match_single_process(tmp_path) -> None: + state_dict = _tiny_model().state_dict() + x = torch.randn(4, 3, 32, 32) + reference = _reference(state_dict, x) + + out_path = tmp_path / "tp2_result.pt" + mp.spawn( + _tp_worker, + args=(2, _free_port(), state_dict, x, str(out_path)), + nprocs=2, + join=True, + ) + + tp_result = torch.load(out_path) + torch.testing.assert_close(tp_result["logits"], reference["logits"], atol=ATOL, rtol=0) + assert set(tp_result["grads"]) == set(reference["grads"]) + for name, ref_grad in reference["grads"].items(): + torch.testing.assert_close(tp_result["grads"][name], ref_grad, atol=ATOL, rtol=0) + + +def test_build_tp_plan_rejects_indivisible_heads() -> None: + model = _tiny_model() # num_heads=4 + with pytest.raises(ValueError, match="divisible"): + build_tp_parallelize_plan(model, tp_size=3) + + +def test_tp_accelator_world_size_1_is_degenerate(monkeypatch) -> None: + for var in ("RANK", "WORLD_SIZE", "LOCAL_RANK"): + monkeypatch.delenv(var, raising=False) + accelerator = TPAccelerator() + try: + model = _tiny_model() + x = torch.randn(2, 3, 32, 32) + expected = model(x).detach().clone() + prepared = accelerator.prepare_model(model) # no plan required at world_size 1 + got = prepared(x.to(accelerator.device)).cpu() + # fp32 rounding differs between the cpu reference and the accelerator device. + torch.testing.assert_close(got, expected, atol=1e-4, rtol=0) + finally: + accelerator.destroy() diff --git a/tests/examples/test_vit_tp_exp_unit.py b/tests/examples/test_vit_tp_exp_unit.py new file mode 100644 index 0000000..354a00c --- /dev/null +++ b/tests/examples/test_vit_tp_exp_unit.py @@ -0,0 +1,211 @@ +"""Unit tests for the ported DeiT model (docs/vit_tp.md). + +The ported split-qkv ``Attention`` is checked against a verbatim copy of the official +timm 0.3.2 fused ``Attention`` (the oracle below), and the full model is cross-checked +against ``timm.create_model("deit_small_patch16_224")`` plus the official checkpoint. +""" + +from __future__ import annotations + +import copy +from collections import OrderedDict + +import pytest + +pytest.importorskip("timm", reason="tinyexp vit extra (timm) is required for the vit_tp_exp example") + +import torch +import torch.nn as nn + +from tinyexp.examples.vit_tp_exp import ( + DEIT_SMALL_PATCH16_224_CKPT_URL, + Attention, + VisionTransformer, + convert_fused_qkv_to_split, + convert_split_qkv_to_fused, + deit_small_patch16_224, +) + +# Official DeiT-S size (paper reports 22.1M). +DEIT_S_NUM_PARAMS = 22050664 + + +class FusedAttention(nn.Module): + """Oracle: timm 0.3.2 ``timm/models/vision_transformer.py::Attention``, verbatim.""" + + def __init__( + self, + dim: int, + num_heads: int = 8, + qkv_bias: bool = False, + qk_scale: float | None = None, + attn_drop: float = 0.0, + proj_drop: float = 0.0, + ) -> None: + super().__init__() + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = qk_scale or head_dim**-0.5 + + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim) + self.proj_drop = nn.Dropout(proj_drop) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + B, N, C = x.shape + qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) + q, k, v = qkv[0], qkv[1], qkv[2] # make torchscript happy (cannot use tensor as tuple) + + attn = (q @ k.transpose(-2, -1)) * self.scale + attn = attn.softmax(dim=-1) + attn = self.attn_drop(attn) + + x = (attn @ v).transpose(1, 2).reshape(B, N, C) + x = self.proj(x) + x = self.proj_drop(x) + return x + + +def _to_fused_reference(model: VisionTransformer) -> VisionTransformer: + """Return a deep copy of ``model`` whose split-qkv attentions are fused oracles.""" + fused = copy.deepcopy(model) + for name, module in list(fused.named_modules()): + if not isinstance(module, Attention): + continue + new = FusedAttention(module.q.in_features, num_heads=module.num_heads, qkv_bias=module.q.bias is not None) + with torch.no_grad(): + new.qkv.weight.copy_(torch.cat([module.q.weight, module.k.weight, module.v.weight], dim=0)) + if module.q.bias is not None: + new.qkv.bias.copy_(torch.cat([module.q.bias, module.k.bias, module.v.bias], dim=0)) + new.proj.weight.copy_(module.proj.weight) + new.proj.bias.copy_(module.proj.bias) + parent = fused.get_submodule(".".join(name.split(".")[:-1])) + setattr(parent, name.split(".")[-1], new) + return fused + + +def _tiny_vit() -> VisionTransformer: + torch.manual_seed(0) + model = VisionTransformer( + img_size=32, patch_size=16, embed_dim=64, depth=2, num_heads=4, num_classes=10, qkv_bias=True + ) + return model.eval() + + +def test_deit_s_param_count_matches_official() -> None: + model = deit_small_patch16_224() + assert sum(p.numel() for p in model.parameters()) == DEIT_S_NUM_PARAMS + + +def test_split_qkv_forward_and_grad_match_fused_reference() -> None: + model = _tiny_vit() + reference = _to_fused_reference(model) + x = torch.randn(4, 3, 32, 32) + + loss_ours = model(x).square().mean() + loss_ref = reference(x).square().mean() + torch.testing.assert_close(loss_ours, loss_ref, atol=1e-6, rtol=0) + + loss_ours.backward() + loss_ref.backward() + torch.testing.assert_close( + model.patch_embed.proj.weight.grad, reference.patch_embed.proj.weight.grad, atol=1e-6, rtol=0 + ) + torch.testing.assert_close( + model.blocks[1].mlp.fc2.weight.grad, reference.blocks[1].mlp.fc2.weight.grad, atol=1e-6, rtol=0 + ) + + +def test_converter_round_trip_is_lossless() -> None: + state_dict = _tiny_vit().state_dict() + round_trip = convert_fused_qkv_to_split(convert_split_qkv_to_fused(state_dict)) + assert set(round_trip) == set(state_dict) + for key, value in state_dict.items(): + assert torch.equal(value, round_trip[key]) + + +def test_timm_reference_cross_check() -> None: + timm = pytest.importorskip("timm") + reference = timm.create_model("deit_small_patch16_224", pretrained=False).eval() + ours = deit_small_patch16_224().eval() + + assert sum(p.numel() for p in ours.parameters()) == sum(p.numel() for p in reference.parameters()) + + # Cross-load in both directions: key/shape layouts must agree exactly. + ours.load_state_dict(convert_fused_qkv_to_split(reference.state_dict()), strict=True) + reference.load_state_dict(convert_split_qkv_to_fused(ours.state_dict()), strict=True) + + x = torch.randn(2, 3, 224, 224) + with torch.no_grad(): + torch.testing.assert_close(ours(x), reference(x), atol=1e-5, rtol=0) + + +def test_official_checkpoint_cross_check() -> None: + """Load the official DeiT-S checkpoint through the converter and compare forwards.""" + try: + checkpoint = torch.hub.load_state_dict_from_url( + url=DEIT_SMALL_PATCH16_224_CKPT_URL, map_location="cpu", check_hash=True + ) + except Exception as error: # offline environments skip this test + pytest.skip(f"could not download official checkpoint: {error}") + + official_state: OrderedDict = checkpoint["model"] + + reference = deit_small_patch16_224().eval() + fused_reference = _to_fused_reference(reference) + # The fused oracle must strict-load the official checkpoint: proves the ported + # structure matches the official one key-for-key. + fused_reference.load_state_dict(official_state, strict=True) + reference.load_state_dict(convert_fused_qkv_to_split(official_state), strict=True) + + x = torch.randn(2, 3, 224, 224) + with torch.no_grad(): + torch.testing.assert_close(reference(x), fused_reference(x), atol=1e-5, rtol=0) + + +def test_patch_embed_rejects_wrong_input_size() -> None: + model = _tiny_vit() + with pytest.raises(AssertionError): + model(torch.randn(1, 3, 64, 64)) + + +def test_train_dataset_redis_gating(tmp_path, monkeypatch) -> None: + """docs/vit_tp.md D7: the Redis byte cache is opt-in, train-only, and never applies to fake data.""" + from PIL import Image + from torchvision import datasets as tv_datasets + + from tinyexp.examples.vit_tp_exp import RedisCachedImageFolder, VitTpExp + + for split in ("train", "val"): + for cls in ("a", "b"): + cls_dir = tmp_path / split / cls + cls_dir.mkdir(parents=True) + Image.new("RGB", (8, 8)).save(cls_dir / "x.jpg") + + exp = VitTpExp() + exp.dataloader_cfg.data_root = str(tmp_path) + + # Default: redis disabled -> the verbatim official ImageFolder pipeline. + assert not exp.redis_cfg.redis_cache_enabled + ds = exp.dataloader_cfg._build_dataset(is_train=True, redis_cfg=exp.redis_cfg) + assert type(ds) is tv_datasets.ImageFolder + + class _FakeRedisManager: + def __init__(self, *args, **kwargs) -> None: + pass + + monkeypatch.setattr("tinyexp.utils.redis_utils.RedisClientManager", _FakeRedisManager) + exp.redis_cfg.redis_cache_enabled = True + + ds = exp.dataloader_cfg._build_dataset(is_train=True, redis_cfg=exp.redis_cfg) + assert type(ds) is RedisCachedImageFolder + + # The val set is never redis-cached (byte-identical semantics matter most there). + ds = exp.dataloader_cfg._build_dataset(is_train=False, redis_cfg=exp.redis_cfg) + assert type(ds) is tv_datasets.ImageFolder + + # Synthetic smoke data wins over the cache. + exp.dataloader_cfg.fake_data = True + ds = exp.dataloader_cfg._build_dataset(is_train=True, redis_cfg=exp.redis_cfg) + assert isinstance(ds, torch.utils.data.TensorDataset) diff --git a/tests/test_tinyexp.py b/tests/test_tinyexp.py index bc2e240..4b01d13 100644 --- a/tests/test_tinyexp.py +++ b/tests/test_tinyexp.py @@ -187,3 +187,44 @@ def _instance(cls): assert recorded["name"] == "cfg" assert isinstance(recorded["node"], _BadExpClass) assert recorded["node"].exp_class == expected_path + + +def test_canonical_exp_class_prefers_importable_spec_name(monkeypatch: pytest.MonkeyPatch) -> None: + """``python -m pkg.mod`` runs the module as __main__; the recorded exp must use + the canonical importable class so ray ships it by reference (regression: + by-value export crashed with "Can't pickle " when a + platform agent patched builtins.print).""" + import types + + from tinyexp import _resolve_canonical_exp_class + + # A class from a normally imported module has no canonical twin to swap in. + assert _resolve_canonical_exp_class(_CfgExp) is None + + # A __main__ class whose running module has a spec name resolves canonically. + from tinyexp.examples import vit_tp_exp + + original_module = vit_tp_exp.VitTpExp.__module__ + dummy_main = types.ModuleType("__main__") + dummy_main.__spec__ = types.SimpleNamespace(name="tinyexp.examples.vit_tp_exp") # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "__main__", dummy_main) + try: + vit_tp_exp.VitTpExp.__module__ = "__main__" + resolved = _resolve_canonical_exp_class(vit_tp_exp.VitTpExp) + # In this single-import simulation the spec-name lookup returns the very + # class object we aliased to __main__; under ``python -m`` it is a + # distinct twin of the same file. Either way it must be a TinyExp subclass + # with the right qualname. + assert resolved is not None + assert issubclass(resolved, TinyExp) + assert resolved.__qualname__ == "VitTpExp" + finally: + vit_tp_exp.VitTpExp.__module__ = original_module + + # A plain script (no spec) has no canonical twin. + dummy_main.__spec__ = None # type: ignore[attr-defined] + try: + vit_tp_exp.VitTpExp.__module__ = "__main__" + assert _resolve_canonical_exp_class(vit_tp_exp.VitTpExp) is None + finally: + vit_tp_exp.VitTpExp.__module__ = original_module diff --git a/tests/tine_engine/test_ddp_accelerator.py b/tests/tine_engine/test_ddp_accelerator.py index b756698..8ae6168 100644 --- a/tests/tine_engine/test_ddp_accelerator.py +++ b/tests/tine_engine/test_ddp_accelerator.py @@ -36,6 +36,18 @@ def test_reduce_sum(self): expected_result = torch.tensor([expected_val], device=device, dtype=torch.float32) assert torch.equal(res, expected_result) + + # nccl cannot reduce cpu tensors directly; reduce_sum must transparently + # stage them on the device and return the result on the input device + # (regression: metric sync after an epoch crashed with "No backend type + # associated with device type cpu" under DDP+nccl). + cpu_tensor = torch.tensor([float(self.accelerator.rank)], dtype=torch.float64) + cpu_res = self.accelerator.reduce_sum(cpu_tensor) + expected_cpu = torch.tensor([expected_val], dtype=torch.float64) + assert cpu_res.device.type == "cpu" + assert torch.equal(cpu_res, expected_cpu) + # The caller's tensor is never mutated in place. + assert torch.equal(cpu_tensor, torch.tensor([float(self.accelerator.rank)], dtype=torch.float64)) return True diff --git a/tinyexp/__init__.py b/tinyexp/__init__.py index 07185db..7ac622e 100644 --- a/tinyexp/__init__.py +++ b/tinyexp/__init__.py @@ -217,6 +217,35 @@ def info(self, message): # type: ignore[no-untyped-def] raise UnknownLauncherError(cfg.launcher) +def _resolve_canonical_exp_class(exp_class: type[TinyExp]) -> Optional[type[TinyExp]]: + """Return the importable twin of a ``__main__``-defined experiment class, if any. + + ``python -m pkg.mod`` executes the module as ``__main__``; the executed class + object would make ray/cloudpickle serialize the whole class graph *by value* + (via the actor export and the structured-config metadata shipped with + ``set_cfg``), which is fragile: any builtin the class graph touches (e.g. a + platform-agent-patched ``builtins.print``) can break the by-reference + identity check ("it's not the same object as builtins.print"). When the same + class is importable under the running module's spec name (``pkg.mod.Exp``), + prefer that twin so everything ships by reference. + """ + if exp_class.__module__ != "__main__": + return None + main_spec = getattr(sys.modules.get("__main__"), "__spec__", None) + spec_name = getattr(main_spec, "name", None) + if not spec_name or spec_name == "__main__": + return None + import importlib + + try: + canonical = getattr(importlib.import_module(spec_name), exp_class.__qualname__) + except (ImportError, AttributeError): + return None + if isinstance(canonical, type) and issubclass(canonical, TinyExp): + return canonical + return None + + def store_and_run_exp(exp_class: type[TinyExp]) -> None: """ Extract the config from the exp_class and store it in the ConfigStore(hydra config store). @@ -229,9 +258,12 @@ def store_and_run_exp(exp_class: type[TinyExp]) -> None: None: This function does not return anything. """ - # this is the hack for hydra to find the experiment class - exp_class_path = f"{exp_class.__module__}.{exp_class.__qualname__}" - exp_cfg = exp_class() + # this is the hack for hydra to find the experiment class; prefer the + # canonical importable twin of a __main__ class so ray workers ship it by + # reference (see _resolve_canonical_exp_class) + canonical_class = _resolve_canonical_exp_class(exp_class) or exp_class + exp_class_path = f"{canonical_class.__module__}.{canonical_class.__qualname__}" + exp_cfg = canonical_class() exp_cfg.exp_class = exp_class_path # store the experiment configuration in the ConfigStore and launch the experiment diff --git a/tinyexp/dataset/ra_sampler.py b/tinyexp/dataset/ra_sampler.py new file mode 100644 index 0000000..e883725 --- /dev/null +++ b/tinyexp/dataset/ra_sampler.py @@ -0,0 +1,80 @@ +"""Repeated-augmentation sampler for distributed training. + +Ported verbatim from facebookresearch/deit ``samplers.py`` at commit +``7e160fe43f0252d17191b71cbb5826254114ea5b`` (pinned in ``docs/vit_tp.md``). Only +formatting and type annotations differ from the original; the sampling logic is +unchanged. +""" + +from __future__ import annotations + +import math +from collections.abc import Iterator + +import torch +import torch.distributed as dist + + +class RASampler(torch.utils.data.Sampler[int]): + """Sampler that restricts data loading to a subset of the dataset for distributed, + with repeated augmentation. + It ensures that different each augmented version of a sample will be visible to a + different process (GPU) + Heavily based on torch.utils.data.DistributedSampler + """ + + def __init__( + self, + dataset: torch.utils.data.Dataset, + num_replicas: int | None = None, + rank: int | None = None, + shuffle: bool = True, + num_repeats: int = 3, + ) -> None: + if num_replicas is None: + if not dist.is_available(): + raise RuntimeError("Requires distributed package to be available") + num_replicas = dist.get_world_size() + if rank is None: + if not dist.is_available(): + raise RuntimeError("Requires distributed package to be available") + rank = dist.get_rank() + if num_repeats < 1: + raise ValueError("num_repeats should be greater than 0") + self.dataset = dataset + self.num_replicas = num_replicas + self.rank = rank + self.num_repeats = num_repeats + self.epoch = 0 + self.num_samples = int(math.ceil(len(self.dataset) * self.num_repeats / self.num_replicas)) + self.total_size = self.num_samples * self.num_replicas + self.num_selected_samples = int(math.floor(len(self.dataset) // 256 * 256 / self.num_replicas)) + self.shuffle = shuffle + + def __iter__(self) -> Iterator[int]: + if self.shuffle: + # deterministically shuffle based on epoch + g = torch.Generator() + g.manual_seed(self.epoch) + indices = torch.randperm(len(self.dataset), generator=g) + else: + indices = torch.arange(start=0, end=len(self.dataset)) + + # add extra samples to make it evenly divisible + indices = torch.repeat_interleave(indices, repeats=self.num_repeats, dim=0).tolist() + padding_size: int = self.total_size - len(indices) + if padding_size > 0: + indices += indices[:padding_size] + assert len(indices) == self.total_size + + # subsample + indices = indices[self.rank : self.total_size : self.num_replicas] + assert len(indices) == self.num_samples + + return iter(indices[: self.num_selected_samples]) + + def __len__(self) -> int: + return self.num_selected_samples + + def set_epoch(self, epoch: int) -> None: + self.epoch = epoch diff --git a/tinyexp/examples/vit_tp_exp.py b/tinyexp/examples/vit_tp_exp.py new file mode 100644 index 0000000..ab93394 --- /dev/null +++ b/tinyexp/examples/vit_tp_exp.py @@ -0,0 +1,1348 @@ +"""DeiT-S experiment with tensor parallelism, ported against the official DeiT. + +This module is written by strict cross-checking (对拍) against upstream sources; only +formatting-level changes are allowed (see ``docs/vit_tp.md``). The ported sources are: + +- Model classes (``Mlp``/``Block``/``PatchEmbed``/``VisionTransformer``): timm's + ``vision_transformer`` implementation, configured to retain the timm 0.3.2 DeiT + behavior. ``Attention`` is the only model layer adapted locally for TP: its fused + projection is split into ``q``/``k``/``v``. +- ``deit_small_patch16_224`` factory hyper-parameters and checkpoint URL: + facebookresearch/deit ``models.py`` at commit ``7e160fe43f0252d17191b71cbb5826254114ea5b``. +- Train/eval loops, recipe defaults, transforms and RASampler wiring: deit ``engine.py``, + ``main.py``, ``datasets.py`` (same commit). +- Logging statistics (``SmoothedValue``/``MetricLogger``): deit ``utils.py`` (detr lineage). + +Sanctioned deviation (docs/vit_tp.md D1): the official fused ``attn.qkv`` Linear is split +into separate ``attn.q``/``attn.k``/``attn.v`` Linears so that ``ColwiseParallel`` can +shard attention on head boundaries; ``convert_fused_qkv_to_split`` / +``convert_split_qkv_to_fused`` translate weights between the two layouts, and the unit +tests prove numerical equivalence of the split forward against the fused original. + +Sanctioned deviation (docs/vit_tp.md D2): ``Attention.forward`` reshapes use ``-1`` +(local-shape aware) to stay correct under colwise sharding; at tp=1 the resolved shapes +are exactly the official ones (guarded by the fused-oracle unit tests). + +Sanctioned deviation (docs/vit_tp.md D7): the training set can optionally be served +through a Redis raw-byte cache (``redis_cfg.redis_cache_enabled``, default off) for +slow shared filesystems; samples stay bit-identical to ``datasets.ImageFolder``. + +Cross-check results (full record in ``docs/vit_tp.md``): TP=2 is numerically +equivalent to single-rank (CI); the official checkpoint evaluates to 79.82/79.81 +top-1 at tp=1/tp=2 (official 79.8); the full 300-epoch recipe trained from scratch +(8xH200 DDP via rjob, Redis cache on) reached **79.84% top-1 / 94.99% top-5** +(paper: 79.8 / ~95.0). + +Usage (full recipe, 2 GPUs, TP=2):: + + # ImageNet must use the standard ImageFolder layout with train/ and val/ + # class directories. The same variable is used by the ResNet example. + export IMAGENET_HOME=/path/to/imagenet + + python -m tinyexp.examples.vit_tp_exp + +Eval-only against the official checkpoint (accuracy cross-check):: + + export IMAGENET_HOME=/path/to/imagenet + python -m tinyexp.examples.vit_tp_exp mode=eval module_cfg.pretrained_from= + +Full 300-epoch training on 8 GPUs with DDP and the optional Redis byte cache:: + + export IMAGENET_HOME=/path/to/imagenet + python -m tinyexp.examples.vit_tp_exp accelerator_cfg.accelerator=ddp \\ + ray_cfg.ray_num_worker=8 redis_cfg.redis_cache_enabled=true redis_cfg.redis_cache_max_memory=300 + +Throughput / memory benchmark, e.g. TP vs DDP:: + + python -m tinyexp.examples.vit_tp_exp mode=bench + python -m tinyexp.examples.vit_tp_exp mode=bench accelerator_cfg.accelerator=ddp + +CPU smoke (no ImageNet on disk, ray 2 workers, gloo TP=2):: + + python -m tinyexp.examples.vit_tp_exp dataloader_cfg.fake_data=true dataloader_cfg.input_size=32 \\ + module_cfg.img_size=32 module_cfg.num_classes=10 epochs=1 max_train_steps=4 \\ + ray_cfg.ray_num_gpus_per_worker=0.0 ray_cfg.ray_num_cpus_per_worker=4 + +Not ported: ``HybridEmbed``, distilled models, distillation, ``--cosub``, +``--ThreeAugment``, ``--bce-loss``, ``--attn-only``, INAT/CIFAR datasets, submitit, +position-embedding interpolation for non-224 finetuning. +""" + +from __future__ import annotations + +import datetime +import io +import json +import math +import os +import sys +import time +import weakref +from collections import defaultdict, deque +from dataclasses import dataclass, field +from functools import partial +from types import SimpleNamespace +from typing import Any, Callable + +import numpy as np +import torch +import torch.nn as nn +from PIL import Image +from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD, Mixup, create_transform +from timm.layers import trunc_normal_ +from timm.loss import LabelSmoothingCrossEntropy, SoftTargetCrossEntropy +from timm.models.vision_transformer import Attention as TimmAttention +from timm.models.vision_transformer import Block as TimmBlock +from timm.models.vision_transformer import VisionTransformer as TimmVisionTransformer +from timm.scheduler import create_scheduler +from timm.utils import ModelEma, NativeScaler, accuracy, get_state_dict +from torch.distributed.tensor.parallel import ColwiseParallel, ParallelStyle, RowwiseParallel +from torchvision import datasets, transforms + +from tinyexp import TinyExp, store_and_run_exp +from tinyexp.dataset.ra_sampler import RASampler +from tinyexp.exceptions import UnknownAcceleratorTypeError +from tinyexp.exp_mixins import CheckpointCfgMixin, LoggerCfgMixin, RayCfgMixin, RedisCfgMixin, WandbCfgMixin +from tinyexp.tiny_engine.accelerator import AcceleratorProtocol, TPAccelerator + +# Official DeiT-S checkpoint (facebookresearch/deit models.py). +DEIT_SMALL_PATCH16_224_CKPT_URL = "https://dl.fbaipublicfiles.com/deit/deit_small_patch16_224-cd65a155.pth" + + +class _SplitQkv(nn.Module): + """Compose split projections into the fused tensor expected by timm Attention.""" + + def __init__(self, attention: Attention) -> None: + super().__init__() + self._attention = weakref.ref(attention) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + attention = self._attention() + q, k, v = attention.q(x), attention.k(x), attention.v(x) + # ColwiseParallel returns local plain tensors. timm's inherited forward + # then reshapes this local QKV tensor using the local head count. + attention.attn_dim = q.shape[-1] + attention.num_heads = attention.attn_dim // attention.head_dim + return torch.cat((q, k, v), dim=-1) + + +class Attention(TimmAttention): + """timm Attention with split projections for tensor parallelism.""" + + def __init__( + self, + dim: int, + num_heads: int = 8, + qkv_bias: bool = False, + qk_scale: float | None = None, + attn_drop: float = 0.0, + proj_drop: float = 0.0, + proj_bias: bool = True, + **kwargs: Any, + ) -> None: + kwargs.pop("qk_norm", None) + kwargs.pop("scale_norm", None) + super().__init__( + dim=dim, + num_heads=num_heads, + qkv_bias=qkv_bias, + qk_norm=False, + scale_norm=False, + proj_bias=proj_bias, + attn_drop=attn_drop, + proj_drop=proj_drop, + **kwargs, + ) + del self.qkv + self.q = nn.Linear(dim, self.attn_dim, bias=qkv_bias) + self.k = nn.Linear(dim, self.attn_dim, bias=qkv_bias) + self.v = nn.Linear(dim, self.attn_dim, bias=qkv_bias) + self.qkv = _SplitQkv(self) + if qk_scale is not None: + self.scale = qk_scale + + +class Block(TimmBlock): + """timm block that keeps the old Attention call signature when unmasked.""" + + def forward( + self, + x: torch.Tensor, + attn_mask: torch.Tensor | None = None, + is_causal: bool = False, + ) -> torch.Tensor: + if attn_mask is None and not is_causal: + x = x + self.drop_path1(self.ls1(self.attn(self.norm1(x)))) + else: + x = x + self.drop_path1(self.ls1(self.attn(self.norm1(x), attn_mask=attn_mask, is_causal=is_causal))) + x = x + self.drop_path2(self.ls2(self.mlp(self.norm2(x)))) + return x + + +class VisionTransformer(TimmVisionTransformer): + """DeiT-compatible timm ViT with the TP attention adapter above.""" + + def __init__( + self, + img_size: int = 224, + patch_size: int = 16, + in_chans: int = 3, + num_classes: int = 1000, + embed_dim: int = 768, + depth: int = 12, + num_heads: int = 12, + mlp_ratio: float = 4.0, + qkv_bias: bool = False, + qk_scale: float | None = None, + drop_rate: float = 0.0, + attn_drop_rate: float = 0.0, + drop_path_rate: float = 0.0, + norm_layer: type[nn.Module] = nn.LayerNorm, + **kwargs: Any, + ) -> None: + super().__init__( + img_size=img_size, + patch_size=patch_size, + in_chans=in_chans, + num_classes=num_classes, + global_pool="token", + embed_dim=embed_dim, + depth=depth, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + drop_path_rate=drop_path_rate, + pos_drop_rate=drop_rate, + proj_drop_rate=drop_rate, + attn_drop_rate=attn_drop_rate, + norm_layer=norm_layer, + block_fn=Block, + attn_layer=Attention, + weight_init="skip", + fc_norm=False, + **kwargs, + ) + if qk_scale is not None: + for block in self.blocks: + block.attn.scale = qk_scale + self._init_old_weights() + + def _init_old_weights(self) -> None: + if self.pos_embed is not None: + trunc_normal_(self.pos_embed, std=0.02) + if self.cls_token is not None: + trunc_normal_(self.cls_token, std=0.02) + self.apply(self._init_weights) + + @staticmethod + def _init_weights(module: nn.Module) -> None: + if isinstance(module, nn.Linear): + trunc_normal_(module.weight, std=0.02) + if module.bias is not None: + nn.init.constant_(module.bias, 0) + elif isinstance(module, nn.LayerNorm): + nn.init.constant_(module.bias, 0) + nn.init.constant_(module.weight, 1.0) + + @torch.jit.ignore + def no_weight_decay(self) -> set[str]: + return {"pos_embed", "cls_token"} + + +def deit_small_patch16_224(pretrained: bool = False, **kwargs: Any) -> VisionTransformer: + """Build DeiT-S with the official hyper-parameters (facebookresearch/deit models.py). + + With ``pretrained=True`` the official checkpoint is loaded through the qkv-split + converter of docs/vit_tp.md D1. + """ + model = VisionTransformer( + patch_size=16, + embed_dim=384, + depth=12, + num_heads=6, + mlp_ratio=4, + qkv_bias=True, + norm_layer=partial(nn.LayerNorm, eps=1e-6), + **kwargs, + ) + if pretrained: + checkpoint = torch.hub.load_state_dict_from_url( + url=DEIT_SMALL_PATCH16_224_CKPT_URL, map_location="cpu", check_hash=True + ) + model.load_state_dict(convert_fused_qkv_to_split(checkpoint["model"])) + return model + + +def build_tp_parallelize_plan(module: VisionTransformer, tp_size: int) -> dict[str, ParallelStyle]: + """Build the Megatron-style colwise/rowwise plan for the ported VisionTransformer (docs/vit_tp.md). + + Head divisibility is the hard structural requirement of TP attention: ColwiseParallel + shards the ``q``/``k``/``v`` output dim, and a shard boundary must not fall inside a + head (that is exactly why the fused ``qkv`` Linear had to be split, docs/vit_tp.md D1). + """ + num_heads = module.blocks[0].attn.num_heads + if tp_size < 1: + raise ValueError(f"tp_size must be >= 1, got {tp_size}") + if num_heads % tp_size != 0: + raise ValueError(f"num_heads={num_heads} must be divisible by tp_size={tp_size} to shard on head boundaries") + plan: dict[str, ParallelStyle] = {} + for i in range(len(module.blocks)): + plan.update( + { + f"blocks.{i}.attn.q": ColwiseParallel(), + f"blocks.{i}.attn.k": ColwiseParallel(), + f"blocks.{i}.attn.v": ColwiseParallel(), + # Rowwise proj consumes the Shard(-1) head chunk and all-reduces back to + # a replicated output: one forward all-reduce per attention. + f"blocks.{i}.attn.proj": RowwiseParallel(), + f"blocks.{i}.mlp.fc1": ColwiseParallel(), + f"blocks.{i}.mlp.fc2": RowwiseParallel(), + } + ) + return plan + + +def convert_fused_qkv_to_split(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """Convert an official (fused-qkv) DeiT/ViT state dict into the split layout of ``Attention``. + + The fused ``attn.qkv`` weight has shape ``(3 * dim, dim)`` with rows ordered + ``[q, k, v]`` (timm 0.3.2 reshapes the ``(B, N, 3 * dim)`` output with the factor of + 3 leading), so contiguous thirds map to the three split Linears. + """ + converted: dict[str, torch.Tensor] = {} + for key, value in state_dict.items(): + if ".attn.qkv." in key: + prefix, suffix = key.split(".attn.qkv.") + third = value.shape[0] // 3 + converted[f"{prefix}.attn.q.{suffix}"] = value[:third].clone() + converted[f"{prefix}.attn.k.{suffix}"] = value[third : 2 * third].clone() + converted[f"{prefix}.attn.v.{suffix}"] = value[2 * third :].clone() + else: + converted[key] = value + return converted + + +def convert_split_qkv_to_fused(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """Inverse of :func:`convert_fused_qkv_to_split`: emit official checkpoint layout.""" + converted = dict(state_dict) + for key in state_dict: + if not key.endswith(".attn.q.weight"): + continue + prefix = key[: -len(".attn.q.weight")] + for stem in ("weight", "bias"): + q = state_dict[f"{prefix}.attn.q.{stem}"] + k = state_dict[f"{prefix}.attn.k.{stem}"] + v = state_dict[f"{prefix}.attn.v.{stem}"] + converted[f"{prefix}.attn.qkv.{stem}"] = torch.cat([q, k, v], dim=0) + del converted[f"{prefix}.attn.q.{stem}"] + del converted[f"{prefix}.attn.k.{stem}"] + del converted[f"{prefix}.attn.v.{stem}"] + return converted + + +class _FusedQkvModelProxy: + """Adapter that reads/writes the model in the official fused-qkv checkpoint layout. + + tinyexp checkpoints therefore stay byte-compatible with facebookresearch/deit + checkpoints (``model_state_dict`` holds fused ``attn.qkv`` keys) regardless of the + accelerator in use. + """ + + def __init__(self, module: nn.Module, accelerator: AcceleratorProtocol | None = None) -> None: + self.module = module + self.accelerator = accelerator + + def state_dict(self) -> dict[str, torch.Tensor]: + dump = getattr(self.accelerator, "dump_model_to_state_dict", None) if self.accelerator is not None else None + if dump is not None: + state = dump(self.module) + else: # CPUAccelerator / plain module: state_dict is already whole. + state = {key: value.cpu() for key, value in self.module.state_dict().items()} + return convert_split_qkv_to_fused(state) + + def load_state_dict(self, state_dict: dict[str, torch.Tensor], strict: bool = True) -> Any: + # Match the official finetune behavior: ignore classifier weights when + # the checkpoint and configured class count differ. + state_dict = dict(state_dict) + model_state = self.module.state_dict() + for key in ("head.weight", "head.bias"): + if key in state_dict and key in model_state and state_dict[key].shape != model_state[key].shape: + print(f"Removing key {key} from pretrained checkpoint") + del state_dict[key] + return self.module.load_state_dict(convert_fused_qkv_to_split(state_dict), strict=strict) + + +def _tensors_to_plain(state: Any) -> Any: + """Recursively convert DTensors to plain whole cpu tensors. + + ``full_tensor()`` is a collective: under TP every rank must call this on the same + object graph at the same time. + """ + from torch.distributed.tensor import DTensor + + if isinstance(state, DTensor): + return state.full_tensor().cpu() + if isinstance(state, dict): + return {key: _tensors_to_plain(value) for key, value in state.items()} + if isinstance(state, list): + return [_tensors_to_plain(item) for item in state] + if isinstance(state, tuple): + return tuple(_tensors_to_plain(item) for item in state) + if isinstance(state, torch.Tensor): + return state.cpu() + return state + + +class _PrecomputedStateDict: + """Duck-typed stand-in so CheckpointCfg.save_checkpoint stores an already-dumped dict.""" + + def __init__(self, state_dict: dict) -> None: + self._state_dict = state_dict + + def state_dict(self) -> dict: + return self._state_dict + + +# ---------------------- metric logging (deit utils.py, verbatim) ---------------------- # + + +class SmoothedValue: + """Track a series of values and provide access to smoothed values over a + window or the global series average. + """ + + def __init__(self, window_size: int = 20, fmt: str | None = None) -> None: + if fmt is None: + fmt = "{median:.4f} ({global_avg:.4f})" + self.deque: deque = deque(maxlen=window_size) + self.total = 0.0 + self.count = 0 + self.fmt = fmt + + def update(self, value: float, n: int = 1) -> None: + self.deque.append(value) + self.count += n + self.total += value * n + + @property + def median(self) -> float: + d = torch.tensor(list(self.deque)) + return d.median().item() + + @property + def avg(self) -> float: + d = torch.tensor(list(self.deque)) + return d.mean().item() + + @property + def global_avg(self) -> float: + return self.total / self.count + + @property + def max(self) -> float: + return max(self.deque) + + @property + def value(self) -> float: + return self.deque[-1] + + def __str__(self) -> str: + return self.fmt.format( + median=self.median, avg=self.avg, global_avg=self.global_avg, max=self.max, value=self.value + ) + + +class MetricLogger: + def __init__(self, delimiter: str = "\t", log_fn: Callable[..., None] = print) -> None: + self.meters: defaultdict = defaultdict(SmoothedValue) + self.delimiter = delimiter + self.log_fn = log_fn + + def update(self, **kwargs: Any) -> None: + for key, value in kwargs.items(): + self.meters[key].update(value) + + def __getattr__(self, attr: str) -> SmoothedValue: + if attr in self.meters: + return self.meters[attr] + raise AttributeError(f"'{type(self).__name__}' object has no attribute '{attr}'") + + def __str__(self) -> str: + loss_str = [] + for name, meter in self.meters.items(): + loss_str.append(f"{name}: {meter!s}") + return self.delimiter.join(loss_str) + + def synchronize_between_processes(self, accelerator: AcceleratorProtocol) -> None: + """detr-style sync: all-reduce (total, count) per meter, then average.""" + if getattr(accelerator, "world_size", 1) < 2: + return + for meter in self.meters.values(): + pair = accelerator.reduce_sum(torch.tensor([meter.total, meter.count], dtype=torch.float64)) + meter.total = float(pair[0].item()) / accelerator.world_size + meter.count = float(pair[1].item()) / accelerator.world_size + + def add_meter(self, name: str, meter: SmoothedValue) -> None: + self.meters[name] = meter + + def log_every(self, iterable: Any, print_freq: int, header: str | None = None) -> Any: + if not header: + header = "" + start_time = time.time() + end = time.time() + iter_time = SmoothedValue(fmt="{avg:.4f}") + data_time = SmoothedValue(fmt="{avg:.4f}") + space_fmt = ":" + str(len(str(len(iterable)))) + "d" + if torch.cuda.is_available(): + log_msg = self.delimiter.join( + [ + header, + "[{0" + space_fmt + "}/{1}]", + "eta: {eta}", + "{meters}", + "time: {time}", + "data: {data}", + "max mem: {memory:.0f}", + ] + ) + else: + log_msg = self.delimiter.join( + [ + header, + "[{0" + space_fmt + "}/{1}]", + "eta: {eta}", + "{meters}", + "time: {time}", + "data: {data}", + ] + ) + i = 0 + for obj in iterable: + data_time.update(time.time() - end) + yield obj + iter_time.update(time.time() - end) + if i % print_freq == 0 or i == len(iterable) - 1: + eta_seconds = iter_time.global_avg * (len(iterable) - i) + eta_string = str(datetime.timedelta(seconds=int(eta_seconds))) + if torch.cuda.is_available(): + self.log_fn( + log_msg.format( + i, + len(iterable), + eta=eta_string, + meters=str(self), + time=str(iter_time), + data=str(data_time), + memory=torch.cuda.max_memory_allocated() / 1024.0 / 1024.0, + ) + ) + else: + self.log_fn( + log_msg.format( + i, len(iterable), eta=eta_string, meters=str(self), time=str(iter_time), data=str(data_time) + ) + ) + i += 1 # noqa: SIM113 (official utils.py counter) + end = time.time() + total_time = time.time() - start_time + total_time_str = str(datetime.timedelta(seconds=int(total_time))) + self.log_fn(f"{header} Total time: {total_time_str} ({total_time / max(len(iterable), 1):.4f} s / it)") + + +class RedisCachedImageFolder: + """ImageFolder wrapper caching raw JPEG bytes in Redis (resnet_exp pattern). + + Authorized deviation (docs/vit_tp.md D7): IO layer only — cache stores the exact + file bytes and the decode/transform pipeline is unchanged, so samples are + bit-identical to ``datasets.ImageFolder``. First epoch fills the cache from + disk (misses), later epochs read from Redis (hits). + """ + + def __init__( + self, + redis_host: str, + redis_ports: list[int], + root: str, + transform=None, + target_transform=None, + redis_world_size: int = 1, + ): + self.root = root + self.transform = transform + self.target_transform = target_transform + self.dataset = datasets.ImageFolder(root) + + self.cache_misses = 0 + self.cache_hits = 0 + from tinyexp.utils.redis_utils import RedisClientManager + + self.redis_client_manager = RedisClientManager(redis_host, redis_ports, redis_world_size) + + def __getitem__(self, index): + path, target = self.dataset.samples[index] + cache_key = index + + file_data = self.redis_client_manager.safe_get(cache_key) + if file_data is None: + self.cache_misses += 1 + with open(path, "rb") as f: + file_data = f.read() + self.redis_client_manager.safe_set(cache_key, file_data) + else: + self.cache_hits += 1 + + image = Image.open(io.BytesIO(file_data)).convert("RGB") + + if self.transform is not None: + image = self.transform(image) + if self.target_transform is not None: + target = self.target_transform(target) + return image, target + + def __len__(self): + return len(self.dataset) + + +# ------------------------------ experiment (docs/vit_tp.md) ------------------------------ # +# The Exp wiring below maps facebookresearch/deit main.py/engine.py/datasets.py onto the +# tinyexp Cfg-component layout; recipe defaults are the official argparse defaults. + + +def _set_cudnn_benchmark(value: bool = True) -> None: + """official main.py: ``import torch.backends.cudnn as cudnn; cudnn.benchmark = True``. + + The local import is also what makes this picklable for Ray: a direct + ``torch.backends.cudnn`` attribute chain is captured by cloudpickle and + ``CudnnModule`` cannot be pickled. + """ + import torch.backends.cudnn as cudnn + + cudnn.benchmark = value + + +@dataclass(repr=False) +class VitTpExp(TinyExp, RayCfgMixin, RedisCfgMixin, CheckpointCfgMixin, WandbCfgMixin, LoggerCfgMixin): + mode: str = "train" # train / eval / bench + launcher: str = "ray" + epochs: int = 300 # official --epochs + max_train_steps: int = -1 # smoke runs only; official has no step cap + seed: int = 0 # official --seed + # DDP offsets the seed per rank (official: seed + rank) for data shuffling. TP ranks + # consume the *same* batches and must keep identical RNG state, so all ranks share + # the seed under accelerator="tp" (docs/vit_tp.md D6). + + @dataclass + class RayCfg(RayCfgMixin.RayCfg): + ray_num_worker: int = 2 + ray_num_gpus_per_worker: float = 1.0 + ray_num_cpus_per_worker: int = 12 # main process plus 10 dataloader workers + + ray_cfg: RayCfg = field(default_factory=RayCfg) + + @dataclass + class RedisCfg(RedisCfgMixin.RedisCfg): + # Opt-in (resnet_exp defaults to True): the official deit pipeline is the + # verbatim default; the Redis byte cache is enabled per run for slow + # filesystems (docs/vit_tp.md D7). No effect on fake_data runs. + redis_cache_enabled: bool = False + + redis_cfg: RedisCfg = field(default_factory=RedisCfg) + + @dataclass + class AcceleratorCfg: + accelerator: str = "tp" + + def build_accelerator(self) -> AcceleratorProtocol: + from tinyexp.tiny_engine.accelerator import CPUAccelerator, DDPAccelerator + + if self.accelerator == "cpu": + return CPUAccelerator() + if self.accelerator == "ddp": + return DDPAccelerator() + if self.accelerator == "tp": + return TPAccelerator() + raise UnknownAcceleratorTypeError(self.accelerator) + + accelerator_cfg: AcceleratorCfg = field(default_factory=AcceleratorCfg) + + @dataclass + class ModuleCfg: + # factory kwargs from facebookresearch/deit models.py (depth 12 shared by all) + model_name: str = "deit_small_patch16_224" # also: deit_tiny_/deit_base_patch16_224 + img_size: int = 224 + num_classes: int = 1000 + drop_rate: float = 0.0 # official --drop + drop_path_rate: float = 0.1 # official --drop-path + # official checkpoint (URL or path) loaded through the D1 converter; mismatched + # classifier heads are dropped (official --finetune branch) + pretrained_from: str = "" + + def build_module(self) -> VisionTransformer: + factory = { + "deit_tiny_patch16_224": 192, + "deit_small_patch16_224": 384, + "deit_base_patch16_224": 768, + } + if self.model_name not in factory: + raise ValueError(f"unknown model_name {self.model_name}") + heads = {"deit_tiny_patch16_224": 3, "deit_small_patch16_224": 6, "deit_base_patch16_224": 12}[ + self.model_name + ] + return VisionTransformer( + img_size=self.img_size, + patch_size=16, + embed_dim=factory[self.model_name], + depth=12, + num_heads=heads, + mlp_ratio=4, + qkv_bias=True, + norm_layer=partial(nn.LayerNorm, eps=1e-6), + drop_rate=self.drop_rate, + drop_path_rate=self.drop_path_rate, + num_classes=self.num_classes, + ) + + module_cfg: ModuleCfg = field(default_factory=ModuleCfg) + + @dataclass + class DataloaderCfg: + data_root: str = os.environ.get("IMAGENET_HOME", "./data/imagenet/") + input_size: int = 224 + # 256/rank * 2 ranks = 512 global batch -> lr 5e-4 via the official linear rule + # (docs/vit_tp.md), the same scaling point as the official 1024 @ lr 1e-3. + train_batch_size_per_device: int = 256 + val_batch_size_per_device: int = 384 # int(1.5 * batch), official main.py + num_workers: int = 10 + color_jitter: float = 0.3 + auto_augment: str = "rand-m9-mstd0.5-inc1" + train_interpolation: str = "bicubic" + reprob: float = 0.25 + remode: str = "pixel" + recount: int = 1 + repeated_aug: bool = True + num_repeats: int = 3 # RASampler default + dist_eval: bool = False # official default; forced off under TP (D5) + eval_crop_ratio: float = 0.875 + # synthetic tensors instead of ImageNet: smoke tests and mode=bench, no disk data + fake_data: bool = False + fake_data_len: int = 64 + pin_mem: bool = True + + def _build_transform(self, is_train: bool) -> transforms.Compose: + # official datasets.py build_transform (IMNET branch) + resize_im = self.input_size > 32 + if is_train: + # this should always dispatch to transforms_imagenet_train + transform = create_transform( + input_size=self.input_size, + is_training=True, + color_jitter=self.color_jitter, + auto_augment=self.auto_augment, + interpolation=self.train_interpolation, + re_prob=self.reprob, + re_mode=self.remode, + re_count=self.recount, + ) + if not resize_im: + # replace RandomResizedCropAndInterpolation with RandomCrop + transform.transforms[0] = transforms.RandomCrop(self.input_size, padding=4) + return transform + t = [] + if resize_im: + size = int(self.input_size / self.eval_crop_ratio) + t.append( + # official passes interpolation=3 (BICUBIC) + transforms.Resize(size, interpolation=transforms.InterpolationMode.BICUBIC), + ) + t.append(transforms.CenterCrop(self.input_size)) + t.append(transforms.ToTensor()) + t.append(transforms.Normalize(IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD)) + return transforms.Compose(t) + + def _build_dataset(self, is_train: bool, redis_cfg=None) -> torch.utils.data.Dataset: + if self.fake_data: + return torch.utils.data.TensorDataset( + torch.randn(self.fake_data_len, 3, self.input_size, self.input_size), + torch.randint(0, 2, (self.fake_data_len,)), + ) + # official datasets.py build_dataset IMNET branch + root = os.path.join(self.data_root, "train" if is_train else "val") + transform = self._build_transform(is_train) + if is_train and redis_cfg is not None and redis_cfg.redis_cache_enabled: + # docs/vit_tp.md D7: Redis byte cache for slow filesystems; samples are + # bit-identical to datasets.ImageFolder (bytes are cached verbatim). + return RedisCachedImageFolder( + redis_host=redis_cfg.redis_cluster_host, + redis_ports=list(redis_cfg.redis_cluster_ports), + root=root, + transform=transform, + redis_world_size=int(redis_cfg.redis_rendezvous_world_size), + ) + return datasets.ImageFolder(root, transform=transform) + + def build_train_dataloader( + self, accelerator: AcceleratorProtocol, replicate_data: bool, redis_cfg=None + ) -> torch.utils.data.DataLoader: + dataset = self._build_dataset(is_train=True, redis_cfg=redis_cfg) + num_replicas = 1 if replicate_data else accelerator.world_size + rank = 0 if replicate_data else accelerator.rank + if num_replicas > 1 and self.repeated_aug: + sampler: torch.utils.data.Sampler = RASampler( + dataset, num_replicas=num_replicas, rank=rank, shuffle=True, num_repeats=self.num_repeats + ) + elif num_replicas > 1: + sampler = torch.utils.data.DistributedSampler( + dataset, num_replicas=num_replicas, rank=rank, shuffle=True + ) + else: + sampler = torch.utils.data.RandomSampler(dataset) + return torch.utils.data.DataLoader( + dataset, + sampler=sampler, + batch_size=self.train_batch_size_per_device, + num_workers=self.num_workers, + pin_memory=self.pin_mem, + drop_last=True, + ) + + def build_val_dataloader( + self, accelerator: AcceleratorProtocol, replicate_data: bool + ) -> torch.utils.data.DataLoader: + dataset = self._build_dataset(is_train=False) + if self.dist_eval and not replicate_data: + sampler: torch.utils.data.Sampler = torch.utils.data.DistributedSampler( + dataset, num_replicas=accelerator.world_size, rank=accelerator.rank, shuffle=False + ) + else: + # official non-dist-eval: every rank iterates the full validation set + sampler = torch.utils.data.SequentialSampler(dataset) + return torch.utils.data.DataLoader( + dataset, + sampler=sampler, + batch_size=self.val_batch_size_per_device, + num_workers=self.num_workers, + pin_memory=self.pin_mem, + drop_last=False, + ) + + dataloader_cfg: DataloaderCfg = field(default_factory=DataloaderCfg) + + @dataclass + class LossCfg: + # official mixup/smoothing argparse defaults + mixup: float = 0.8 + cutmix: float = 1.0 + mixup_prob: float = 1.0 + mixup_switch_prob: float = 0.5 + mixup_mode: str = "batch" + smoothing: float = 0.1 + + def build_mixup_fn(self, num_classes: int) -> Mixup | None: + mixup_active = self.mixup > 0 or self.cutmix > 0.0 + if not mixup_active: + return None + return Mixup( + mixup_alpha=self.mixup, + cutmix_alpha=self.cutmix, + cutmix_minmax=None, # official default; not exposed (cloudpickle/OmegaConf annotation limits) + prob=self.mixup_prob, + switch_prob=self.mixup_switch_prob, + mode=self.mixup_mode, + label_smoothing=self.smoothing, + num_classes=num_classes, + ) + + def build_criterion(self, num_classes: int) -> nn.Module: + mixup_active = self.mixup > 0 or self.cutmix > 0.0 + criterion = LabelSmoothingCrossEntropy() + if mixup_active: + # smoothing is handled with mixup label transform + criterion = SoftTargetCrossEntropy() + elif self.smoothing: + criterion = LabelSmoothingCrossEntropy(smoothing=self.smoothing) + else: + criterion = nn.CrossEntropyLoss() + return criterion + + loss_cfg: LossCfg = field(default_factory=LossCfg) + + @dataclass + class OptimizerCfg: + lr: float = 5e-4 # official --lr, linearly scaled below (official --unscale-lr to disable) + weight_decay: float = 0.05 + opt_eps: float = 1e-8 + clip_grad: float = 0.0 # official default is None (no clipping); <=0 disables here + + def build_optimizer( + self, module: nn.Module, dataloader, accelerator: AcceleratorProtocol + ) -> torch.optim.Optimizer: + linear_scaled_lr = self.lr * dataloader.batch_size * accelerator.world_size / 512.0 + # timm create_optimizer adamw path: params listed in module.no_weight_decay() + # get zero weight decay + skip = module.no_weight_decay() if hasattr(module, "no_weight_decay") else set() + parameters = [ + { + "params": [p for n, p in module.named_parameters() if n not in skip and p.requires_grad], + "weight_decay": self.weight_decay, + }, + { + "params": [p for n, p in module.named_parameters() if n in skip and p.requires_grad], + "weight_decay": 0.0, + }, + ] + return torch.optim.AdamW(parameters, lr=linear_scaled_lr, eps=self.opt_eps) + + optimizer_cfg: OptimizerCfg = field(default_factory=OptimizerCfg) + + @dataclass + class LrSchedulerCfg: + # official scheduler argparse defaults; fed to timm create_scheduler verbatim + sched: str = "cosine" + warmup_epochs: int = 5 + warmup_lr: float = 1e-6 + min_lr: float = 1e-5 + cooldown_epochs: int = 10 + decay_epochs: float = 30 + decay_rate: float = 0.1 + patience_epochs: int = 10 + lr_noise_pct: float = 0.67 + lr_noise_std: float = 1.0 + + def build_lr_scheduler(self, optimizer: torch.optim.Optimizer, epochs: int): + sched_args = SimpleNamespace( + sched=self.sched, + epochs=epochs, + min_lr=self.min_lr, + warmup_lr=self.warmup_lr, + warmup_epochs=self.warmup_epochs, + cooldown_epochs=self.cooldown_epochs, + decay_epochs=self.decay_epochs, + decay_rate=self.decay_rate, + patience_epochs=self.patience_epochs, + lr_noise=None, # official default; not exposed (cloudpickle/OmegaConf annotation limits) + lr_noise_pct=self.lr_noise_pct, + lr_noise_std=self.lr_noise_std, + seed=0, + cycle_mul=1.0, + cycle_limit=1, + t_in_epochs=True, + scale_mode="cycle", + ) + lr_scheduler, _ = create_scheduler(sched_args, optimizer) + return lr_scheduler + + lr_scheduler_cfg: LrSchedulerCfg = field(default_factory=LrSchedulerCfg) + + @dataclass + class EmaCfg: + model_ema: bool = True # official default on; forced off at tp>1 (docs/vit_tp.md D4) + model_ema_decay: float = 0.99996 + + ema_cfg: EmaCfg = field(default_factory=EmaCfg) + + @dataclass + class BenchCfg: + # docs/vit_tp.md D3: throughput/memory check on synthetic batches (no dataset IO) + warmup_steps: int = 10 + measure_steps: int = 50 + + bench_cfg: BenchCfg = field(default_factory=BenchCfg) + + # ------------------------------ execution part ------------------------------ # + + def run(self) -> None: + accelerator = self.accelerator_cfg.build_accelerator() + try: + self._run(accelerator) + finally: + accelerator.destroy() + + def _run(self, accelerator: AcceleratorProtocol) -> None: + run_dir = self.get_run_dir() + logger = self.logger_cfg.build_logger( + save_dir=run_dir, distributed_rank=accelerator.rank, filename="deit_train.log" + ) + cfg_dict = self.print_cfg(logger) + + if accelerator.device.type == "cuda": + _set_cudnn_benchmark() # official main.py + # official seeds seed + rank; TP keeps all ranks on one seed (docs/vit_tp.md D6) + seed = self.seed + (accelerator.rank if self.accelerator_cfg.accelerator == "ddp" else 0) + torch.manual_seed(seed) + np.random.seed(seed) + + if self.mode == "train": + self._train(accelerator=accelerator, logger=logger, cfg_dict=cfg_dict, run_dir=run_dir) + elif self.mode == "eval": + self._run_eval(accelerator=accelerator, logger=logger) + elif self.mode == "bench": + self._bench(accelerator=accelerator, logger=logger) + else: + raise NotImplementedError(f"Mode {self.mode} is not implemented") + + def _prepare_for_train_or_eval(self, accelerator): + """Components shared by train and eval, in official main.py order.""" + replicate_data = self.accelerator_cfg.accelerator == "tp" # TP ranks eat identical batches (D5/D6) + dataloader_val = self.dataloader_cfg.build_val_dataloader(accelerator, replicate_data) + + # official args.nb_classes (IMNET branch hard-codes 1000) + nb_classes = self.module_cfg.num_classes + mixup_fn = self.loss_cfg.build_mixup_fn(nb_classes) + criterion = self.loss_cfg.build_criterion(nb_classes) + + model = self.module_cfg.build_module() + if self.module_cfg.pretrained_from: + path_or_url = self.module_cfg.pretrained_from + if path_or_url.startswith("https"): + checkpoint = torch.hub.load_state_dict_from_url(path_or_url, map_location="cpu", check_hash=True) + else: + checkpoint = torch.load(path_or_url, map_location="cpu", weights_only=False) + checkpoint_model = checkpoint["model"] if "model" in checkpoint else checkpoint.get("model_state_dict") + _FusedQkvModelProxy(model).load_state_dict(checkpoint_model, strict=False) + model.to(accelerator.device) + + return replicate_data, dataloader_val, mixup_fn, criterion, model + + def _parallelize(self, accelerator, model, optimizer=None): + if isinstance(accelerator, TPAccelerator): + plan = build_tp_parallelize_plan(accelerator.unwrap_model(model), tp_size=accelerator.world_size) + return accelerator.prepare(model, optimizer, parallelize_plan=plan) + return accelerator.prepare(model, optimizer) + + def _train(self, accelerator, logger, cfg_dict, run_dir: str) -> None: # noqa: C901 + replicate_data, data_loader_val, mixup_fn, criterion, model = self._prepare_for_train_or_eval(accelerator) + data_loader_train = self.dataloader_cfg.build_train_dataloader(accelerator, replicate_data, self.redis_cfg) + + model_ema = None + if self.ema_cfg.model_ema: + if accelerator.world_size > 1 and self.accelerator_cfg.accelerator == "tp": + logger.info("model_ema disabled under tp>1 (docs/vit_tp.md D4)") + else: + # Important to create EMA model after cuda() but before TP/DDP wrapper + model_ema = ModelEma(model, decay=self.ema_cfg.model_ema_decay, device="", resume="") + + n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad) + logger.info(f"number of params: {n_parameters}") + + optimizer = self.optimizer_cfg.build_optimizer(model, data_loader_train, accelerator) + loss_scaler = NativeScaler(device=accelerator.device.type) # GradScaler auto-disables on cpu + lr_scheduler = self.lr_scheduler_cfg.build_lr_scheduler(optimizer, epochs=self.epochs) + + model, optimizer = self._parallelize(accelerator, model, optimizer) + model_proxy = _FusedQkvModelProxy(accelerator.unwrap_model(model), accelerator) + + start_epoch = 0 + max_accuracy = 0.0 + if self.resume_from: + checkpoint = self.checkpoint_cfg.load_checkpoint( + self.resume_from, + model=model_proxy, + optimizer=optimizer, + scheduler=lr_scheduler, + scaler=loss_scaler, + map_location=accelerator.device, + ) + start_epoch = int(checkpoint.get("epoch", -1)) + 1 + max_accuracy = float(checkpoint.get("best_metric") or 0.0) + extra_state = checkpoint.get("extra_state") or {} + if model_ema is not None and extra_state.get("model_ema_state_dict") is not None: + model_ema.ema.load_state_dict(extra_state["model_ema_state_dict"]) + lr_scheduler.step(start_epoch) + + if self.wandb_cfg.enable_wandb and accelerator.is_main_process: + self.wandb_cfg.build_wandb(accelerator=accelerator, project="TinyExp", config=cfg_dict, name="vit_tp_exp") + + logger.info(f"Start training for {self.epochs} epochs") + start_training_time = time.time() + global_step = 0 + for epoch in range(start_epoch, self.epochs): + train_sampler = getattr(data_loader_train, "sampler", None) + if train_sampler is not None and hasattr(train_sampler, "set_epoch"): + train_sampler.set_epoch(epoch) + + train_stats = self._train_one_epoch( + accelerator, + logger, + model, + criterion, + data_loader_train, + optimizer, + accelerator.device, + epoch, + loss_scaler, + self.optimizer_cfg.clip_grad if self.optimizer_cfg.clip_grad > 0 else None, + model_ema, + mixup_fn, + ) + global_step += len(data_loader_train) + lr_scheduler.step(epoch) + + # official order: last checkpoint every epoch, then evaluate and keep best + self._save_checkpoint( + accelerator, + run_dir, + self.checkpoint_cfg.last_ckpt_name, + model_proxy, + optimizer, + lr_scheduler, + loss_scaler, + model_ema, + epoch, + global_step, + max_accuracy, + ) + + test_stats = self._evaluate(accelerator, logger, model, accelerator.device, data_loader_val) + logger.info( + f"Accuracy of the network on the {len(data_loader_val.dataset)} test images: {test_stats['acc1']:.1f}%" + ) + + if max_accuracy < test_stats["acc1"]: + max_accuracy = test_stats["acc1"] + self._save_checkpoint( + accelerator, + run_dir, + self.checkpoint_cfg.best_ckpt_name, + model_proxy, + optimizer, + lr_scheduler, + loss_scaler, + model_ema, + epoch, + global_step, + max_accuracy, + ) + logger.info(f"Max accuracy: {max_accuracy:.2f}%") + + log_stats = { + **{f"train_{k}": v for k, v in train_stats.items()}, + **{f"test_{k}": v for k, v in test_stats.items()}, + "epoch": epoch, + "n_parameters": n_parameters, + } + if accelerator.is_main_process: + with open(os.path.join(run_dir, "log.txt"), "a") as f: + f.write(json.dumps(log_stats) + "\n") + if self.wandb_cfg.enable_wandb: + import wandb + + wandb.log(log_stats) + + if 0 < self.max_train_steps <= global_step: + break + + total_time = time.time() - start_training_time + logger.info(f"Training time {datetime.timedelta(seconds=int(total_time))}") + + def _save_checkpoint( + self, + accelerator, + run_dir, + name, + model_proxy, + optimizer, + lr_scheduler, + loss_scaler, + model_ema, + epoch, + global_step, + max_accuracy, + ) -> None: + # DTensor -> full_tensor gathers are collectives: every rank must run the dumps, + # only the main process writes the file (otherwise TP ranks deadlock). + model_state = model_proxy.state_dict() + optimizer_state = _tensors_to_plain(optimizer.state_dict()) + if not accelerator.is_main_process: + return + self.checkpoint_cfg.save_checkpoint( + run_dir=run_dir, + name=name, + model=_PrecomputedStateDict(model_state), + optimizer=_PrecomputedStateDict(optimizer_state), + scheduler=lr_scheduler, + epoch=epoch, + global_step=global_step, + best_metric=max_accuracy, + exp_name=self.exp_name, + exp_class=self.exp_class, + extra_state={ + "scaler_state_dict": loss_scaler.state_dict(), + # official main.py: get_state_dict(model_ema); ema only exists at world 1 + "model_ema_state_dict": get_state_dict(model_ema) if model_ema is not None else None, + }, + ) + + def _train_one_epoch( + self, + accelerator, + logger, + model, + criterion, + data_loader, + optimizer, + device, + epoch, + loss_scaler, + clip_grad=None, + model_ema=None, + mixup_fn=None, + ): + """Verbatim port of deit engine.train_one_epoch (cosub/bce branches cut, V1).""" + model.train() + metric_logger = MetricLogger(log_fn=logger.info) + metric_logger.add_meter("lr", SmoothedValue(window_size=1, fmt="{value:.6f}")) + header = f"Epoch: [{epoch}]" + print_freq = 10 + + for samples, targets in metric_logger.log_every(data_loader, print_freq, header): + samples = samples.to(device, non_blocking=True) + targets = targets.to(device, non_blocking=True) + + if mixup_fn is not None: + samples, targets = mixup_fn(samples, targets) + + with torch.amp.autocast(device.type, enabled=device.type == "cuda"): + outputs = model(samples) + loss = criterion(outputs, targets) + + loss_value = loss.item() + + if not math.isfinite(loss_value): + logger.error(f"Loss is {loss_value}, stopping training") + sys.exit(1) + + optimizer.zero_grad() + + # this attribute is added by timm on one optimizer (adahessian) + is_second_order = hasattr(optimizer, "is_second_order") and optimizer.is_second_order + loss_scaler( + loss, + optimizer, + clip_grad=clip_grad, + parameters=model.parameters(), + create_graph=is_second_order, + ) + + if model_ema is not None: + model_ema.update(model) + + metric_logger.update(loss=loss_value) + metric_logger.update(lr=optimizer.param_groups[0]["lr"]) + # gather the stats from all processes + metric_logger.synchronize_between_processes(accelerator) + logger.info(f"Averaged stats: {metric_logger}") + return {k: meter.global_avg for k, meter in metric_logger.meters.items()} + + @torch.no_grad() + def _evaluate(self, accelerator, logger, model, device, val_dataloader) -> dict: + """Verbatim port of deit engine.evaluate.""" + criterion = nn.CrossEntropyLoss() + + metric_logger = MetricLogger(log_fn=logger.info) + header = "Test:" + + # switch to evaluation mode + model.eval() + + for images, target in metric_logger.log_every(val_dataloader, 10, header): + images = images.to(device, non_blocking=True) + target = target.to(device, non_blocking=True) + + # compute output + with torch.amp.autocast(device.type, enabled=device.type == "cuda"): + output = model(images) + loss = criterion(output, target) + + acc1, acc5 = accuracy(output, target, topk=(1, 5)) + + batch_size = images.shape[0] + metric_logger.update(loss=loss.item()) + metric_logger.meters["acc1"].update(acc1.item(), n=batch_size) + metric_logger.meters["acc5"].update(acc5.item(), n=batch_size) + # gather the stats from all processes + metric_logger.synchronize_between_processes(accelerator) + logger.info( + f"* Acc@1 {metric_logger.acc1.global_avg:.3f} Acc@5 {metric_logger.acc5.global_avg:.3f} loss {metric_logger.loss.global_avg:.3f}" + ) + + return {k: meter.global_avg for k, meter in metric_logger.meters.items()} + + def _run_eval(self, accelerator, logger) -> float: + source = self.resume_from or self.module_cfg.pretrained_from + if not source: + raise ValueError("mode=eval requires resume_from (tinyexp ckpt) or module_cfg.pretrained_from") + _, data_loader_val, _, _, model = self._prepare_for_train_or_eval(accelerator) + if self.resume_from and not self.module_cfg.pretrained_from: + proxy = _FusedQkvModelProxy(model, accelerator) + self.checkpoint_cfg.load_checkpoint(self.resume_from, model=proxy, map_location=accelerator.device) + model = self._parallelize(accelerator, model) + test_stats = self._evaluate(accelerator, logger, model, accelerator.device, data_loader_val) + logger.info( + f"Accuracy of the network on the {len(data_loader_val.dataset)} test images: {test_stats['acc1']:.1f}%" + ) + self._run_result = f"eval acc@1={test_stats['acc1']:.2f}% acc@5={test_stats['acc5']:.2f}%" + return test_stats["acc1"] + + def _bench(self, accelerator, logger) -> None: + """Synthetic throughput/memory check (docs/vit_tp.md L4). No dataset involved.""" + model = self.module_cfg.build_module() + model.to(accelerator.device) + if accelerator.device.type == "cuda": + _set_cudnn_benchmark() + model = self._parallelize(accelerator, model) + criterion = nn.CrossEntropyLoss() + + batch_size = self.dataloader_cfg.train_batch_size_per_device + images = torch.randn(batch_size, 3, self.dataloader_cfg.input_size, self.dataloader_cfg.input_size) + target = torch.randint(0, self.module_cfg.num_classes, (batch_size,)) + + model.train() + for _ in range(self.bench_cfg.warmup_steps): + images_d = images.to(accelerator.device, non_blocking=True) + target_d = target.to(accelerator.device, non_blocking=True) + with torch.amp.autocast(accelerator.device.type, enabled=accelerator.device.type == "cuda"): + loss = criterion(model(images_d), target_d) + model.zero_grad(set_to_none=True) + accelerator.backward(loss) + + if accelerator.device.type == "cuda": + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + start = time.perf_counter() + for _ in range(self.bench_cfg.measure_steps): + images_d = images.to(accelerator.device, non_blocking=True) + target_d = target.to(accelerator.device, non_blocking=True) + with torch.amp.autocast(accelerator.device.type, enabled=accelerator.device.type == "cuda"): + loss = criterion(model(images_d), target_d) + model.zero_grad(set_to_none=True) + accelerator.backward(loss) + if accelerator.device.type == "cuda": + torch.cuda.synchronize() + elapsed = time.perf_counter() - start + + steps = self.bench_cfg.measure_steps + images_per_s_per_rank = batch_size * steps / elapsed + # DDP ranks process disjoint batches (global throughput = rank * world), TP + # ranks cooperate on the same batch (global throughput = one rank's rate). + effective_global = images_per_s_per_rank * ( + accelerator.world_size if self.accelerator_cfg.accelerator != "tp" else 1 + ) + peak_mem_gb = 0.0 + if accelerator.device.type == "cuda": + peak_mem_gb = torch.cuda.max_memory_allocated() / 1024**3 + result = ( + f"bench[{self.accelerator_cfg.accelerator} world={accelerator.world_size}] " + f"batch/rank={batch_size} step={elapsed / steps * 1000:.1f}ms " + f"img/s/rank={images_per_s_per_rank:.0f} effective_global_img/s={effective_global:.0f} " + f"peak_mem/rank={peak_mem_gb:.2f}GB" + ) + logger.info(result) + self._run_result = result + accelerator.wait_for_everyone() + + def get_ray_run_result(self) -> str | None: + return getattr(self, "_run_result", None) + + +if __name__ == "__main__": + store_and_run_exp(VitTpExp) diff --git a/tinyexp/tiny_engine/accelerator/__init__.py b/tinyexp/tiny_engine/accelerator/__init__.py index 8ecc5f2..6b6eaba 100644 --- a/tinyexp/tiny_engine/accelerator/__init__.py +++ b/tinyexp/tiny_engine/accelerator/__init__.py @@ -4,6 +4,7 @@ from .base_accelerator import AcceleratorProtocol, BaseAccelerator from .cpu_accelerator import CPUAccelerator from .ddp_accelerator import DDPAccelerator +from .tp_accelerator import TPAccelerator with _contextlib.suppress(ImportError): from .hf_accelerator import HFAccelerator diff --git a/tinyexp/tiny_engine/accelerator/ddp_accelerator.py b/tinyexp/tiny_engine/accelerator/ddp_accelerator.py index fb44e9d..3a6f911 100644 --- a/tinyexp/tiny_engine/accelerator/ddp_accelerator.py +++ b/tinyexp/tiny_engine/accelerator/ddp_accelerator.py @@ -135,9 +135,13 @@ def reduce_sum(self, tensor: torch.Tensor) -> torch.Tensor: world_size = self.world_size if world_size < 2: return tensor - tensor = tensor.clone() - dist.all_reduce(tensor, op=dist.ReduceOp.SUM) - return tensor + # NCCL only reduces device-resident tensors; move over and back so + # cpu-side metric tensors also reduce correctly. + device_tensor = tensor.to(self.device) + if device_tensor is tensor: + device_tensor = tensor.clone() + dist.all_reduce(device_tensor, op=dist.ReduceOp.SUM) + return device_tensor.to(tensor.device) def reduce_mean(self, tensor: torch.Tensor) -> torch.Tensor: return self.reduce_sum(tensor) / self.world_size diff --git a/tinyexp/tiny_engine/accelerator/tp_accelerator.py b/tinyexp/tiny_engine/accelerator/tp_accelerator.py new file mode 100644 index 0000000..c35dd95 --- /dev/null +++ b/tinyexp/tiny_engine/accelerator/tp_accelerator.py @@ -0,0 +1,158 @@ +"""Tensor-parallel accelerator built on ``torch.distributed.tensor.parallel`` (docs/vit_tp.md). + +V1 scope is pure tensor parallelism: ``tp_size == world_size``. Every rank receives the +same input batch; the sharded Linear layers compute partial results that are combined by +an all-reduce inside each transformer block, so the block outputs (and thus the model +outputs) are replicated on all ranks. This is the classic Megatron-style block layout: + +- ``attn.q``/``attn.k``/``attn.v`` (ColwiseParallel): shard the output dim on head + boundaries. Attention is per-head independent, so each rank runs its own heads' + ``softmax(q k^T) v`` with no communication inside attention. +- ``attn.proj`` (RowwiseParallel): input is Shard(-1) (its local heads' channel chunk), + the rowwise matmul produces partial sums and one forward all-reduce restores a + replicated output. +- ``mlp.fc1`` (ColwiseParallel) + ``mlp.fc2`` (RowwiseParallel): the same pattern for + the MLP, one all-reduce per block pass. + +The caller supplies the ``parallelize_plan`` (the model layer layout is not the +accelerator's business); see ``tinyexp/examples/vit_tp_exp.py::build_tp_parallelize_plan``. +""" + +from __future__ import annotations + +from typing import Any + +import torch +import torch.distributed as dist +from torch.distributed.device_mesh import DeviceMesh, init_device_mesh +from torch.distributed.tensor import DTensor +from torch.distributed.tensor.parallel import parallelize_module + +from .base_accelerator import BaseAccelerator + + +class TPAccelerator(BaseAccelerator): + """Accelerator that shards transformer Linear layers across the whole world.""" + + def __init__(self) -> None: + super().__init__() + if torch.cuda.is_available(): + if self.device.index is None: + self.device = torch.device("cuda", torch.cuda.current_device()) + torch.cuda.set_device(self.device) + self.backend = "nccl" + device_type = "cuda" + else: + # CPU + gloo keeps the TP code path testable in CI. + self.device = torch.device("cpu") + self.backend = "gloo" + device_type = "cpu" + + if self.world_size > 1: + if not dist.is_initialized(): + self._init_process_group() + self._process_group_initialized = True + # A 1D whole-world mesh reuses the default process group (already + # initialized with the backend selected above). torch 2.7's + # init_device_mesh has no backend kwarg. + self.mesh: DeviceMesh | None = init_device_mesh(device_type, (self.world_size,)) + else: + self.mesh = None + self.sync_gradients = True + + def _init_process_group(self) -> None: + dist.init_process_group(backend=self.backend, init_method="env://") + + def destroy(self) -> None: + """Destroy the process group once, if this accelerator owns it.""" + if self._destroyed: + return + if self._process_group_initialized: + if dist.is_initialized(): + dist.destroy_process_group() + self._process_group_initialized = False + self._destroyed = True + + def unwrap_model(self, model: Any) -> Any: + # parallelize_module patches modules in place; there is no wrapper to undo. + return model.module if hasattr(model, "module") else model + + def prepare(self, model: Any, optimizer: Any = None, parallelize_plan: dict | None = None) -> Any: + ret_model = self.prepare_model(model, parallelize_plan=parallelize_plan) + if optimizer is not None: + ret_optimizer = self.prepare_optimizer(optimizer) + return ret_model, ret_optimizer + return ret_model + + def prepare_model(self, module: Any, parallelize_plan: dict | None = None) -> Any: + module = module.to(self.device) + if self.world_size < 2: + # Degenerate path: a single rank is mathematically the unparallelized model. + return module + if parallelize_plan is None: + raise ValueError("TPAccelerator requires a parallelize_plan at world_size > 1") # noqa: TRY003 + module = parallelize_module(module, self.mesh, parallelize_plan) + # ColwiseParallel/RowwiseParallel (use_local_output=True, the torch default) + # annotate inputs/outputs at module boundaries; the tensors between a colwise + # output and its paired rowwise input are plain local tensors sharded on the + # last dim, which the model's forward must account for with local-shape math + # (see the reshape(-1) adjustments in vit_tp_exp.Attention, docs/vit_tp.md D2). + return module + + def prepare_optimizer(self, optimizer: Any) -> Any: + # Optimizers step DTensor parameters directly; DTensor handles the cross-shard + # gradient bookkeeping during backward. + return optimizer + + def backward(self, loss: torch.Tensor) -> None: + loss.backward() + + def wait_for_everyone(self) -> None: + if self.world_size < 2: + return + dist.barrier() + + def reduce_sum(self, tensor: torch.Tensor) -> torch.Tensor: + if self.world_size < 2: + return tensor + # NCCL only reduces device-resident tensors; move over and back. ``to`` + # is a no-op when the input is already on this device, while all_reduce + # is in-place, so copy only in that aliasing case to preserve the input. + device_tensor = tensor.to(self.device) + if device_tensor is tensor: + device_tensor = tensor.clone() + dist.all_reduce(device_tensor, op=dist.ReduceOp.SUM) + return device_tensor.to(tensor.device) + + def reduce_mean(self, tensor: torch.Tensor) -> torch.Tensor: + return self.reduce_sum(tensor) / self.world_size + + def dump_model_to_state_dict(self, module: Any) -> dict: + """ + dump model to a plain cpu state_dict, gathering DTensor shards to full tensors + """ + model_state = module.state_dict() + model_state_cpu = type(model_state)() + for key, val in model_state.items(): + if isinstance(val, DTensor): + val = val.full_tensor() + model_state_cpu[key] = val.cpu() + return model_state_cpu + + @property + def is_main_process(self) -> bool: + """True for one process per server.""" + return self.rank == 0 + + @property + def is_local_main_process(self) -> bool: + """True for one process per server.""" + return self.local_rank == 0 + + @property + def is_last_process(self) -> bool: + return self.rank == self.world_size - 1 + + def print(self, *args: Any, **kwargs: Any) -> None: + if self.is_local_main_process: + print(*args, **kwargs) diff --git a/uv.lock b/uv.lock index 15f58c5..cead19e 100644 --- a/uv.lock +++ b/uv.lock @@ -209,7 +209,7 @@ resolution-markers = [ "python_full_version < '3.10' and sys_platform != 'linux'", ] dependencies = [ - { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "colorama", marker = "python_full_version < '3.10' and sys_platform == 'win32'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/b9/2e/0090cbf739cee7d23781ad4b89a9894a41538e4fcf4c31dcdd705b78eb8b/click-8.1.8.tar.gz", hash = "sha256:ed53c9d8990d83c2a27deae68e4ee337473f6330c040a31d4225c9574d16096a", size = 226593, upload-time = "2024-12-21T18:38:44.339Z" } wheels = [ @@ -231,7 +231,7 @@ resolution-markers = [ "python_full_version == '3.10.*' and sys_platform != 'linux'", ] dependencies = [ - { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "colorama", marker = "python_full_version >= '3.10' and sys_platform == 'win32'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/60/6c/8ca2efa64cf75a977a0d7fac081354553ebe483345c734fb6b6515d96bbc/click-8.2.1.tar.gz", hash = "sha256:27c491cc05d968d271d5a1db13e3b5a184636d9d930f148c50b038f0d0646202", size = 286342, upload-time = "2025-05-20T23:19:49.832Z" } wheels = [ @@ -331,7 +331,7 @@ name = "cuda-bindings" version = "13.3.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "cuda-pathfinder" }, + { name = "cuda-pathfinder", marker = "python_full_version >= '3.10' and sys_platform == 'linux'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/a9/21/8464d133752951c154feafb3b65c297e7d80f301183d220bec4c830f1441/cuda_bindings-13.3.1-cp310-cp310-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:120fcc53d57903df529c3486962c56528cba5b7d6c57c99537320ed9922c8b86", size = 6073403, upload-time = "2026-05-29T23:11:36.22Z" }, @@ -366,43 +366,43 @@ wheels = [ [package.optional-dependencies] cublas = [ - { name = "nvidia-cublas", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, - { name = "nvidia-cuda-nvrtc", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-cublas", marker = "(python_full_version >= '3.10' and platform_machine == 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cuda-nvrtc", marker = "(python_full_version >= '3.10' and platform_machine == 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] cudart = [ - { name = "nvidia-cuda-runtime", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-cuda-runtime", marker = "(python_full_version >= '3.10' and platform_machine == 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] cufft = [ - { name = "nvidia-cufft", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, - { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-cufft", marker = "(python_full_version >= '3.10' and platform_machine == 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-nvjitlink", marker = "(python_full_version >= '3.10' and platform_machine == 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] cufile = [ - { name = "nvidia-cufile", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-cufile", marker = "(python_full_version >= '3.10' and platform_machine == 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] cupti = [ - { name = "nvidia-cuda-cupti", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-cuda-cupti", marker = "(python_full_version >= '3.10' and platform_machine == 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] curand = [ - { name = "nvidia-curand", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-curand", marker = "(python_full_version >= '3.10' and platform_machine == 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] cusolver = [ - { name = "nvidia-cublas", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, - { name = "nvidia-cusolver", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, - { name = "nvidia-cusparse", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, - { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-cublas", marker = "(python_full_version >= '3.10' and platform_machine == 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cusolver", marker = "(python_full_version >= '3.10' and platform_machine == 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-cusparse", marker = "(python_full_version >= '3.10' and platform_machine == 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-nvjitlink", marker = "(python_full_version >= '3.10' and platform_machine == 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] cusparse = [ - { name = "nvidia-cusparse", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, - { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-cusparse", marker = "(python_full_version >= '3.10' and platform_machine == 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nvidia-nvjitlink", marker = "(python_full_version >= '3.10' and platform_machine == 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] nvjitlink = [ - { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-nvjitlink", marker = "(python_full_version >= '3.10' and platform_machine == 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] nvrtc = [ - { name = "nvidia-cuda-nvrtc", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-cuda-nvrtc", marker = "(python_full_version >= '3.10' and platform_machine == 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] nvtx = [ - { name = "nvidia-nvtx", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-nvtx", marker = "(python_full_version >= '3.10' and platform_machine == 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [[package]] @@ -428,7 +428,7 @@ name = "exceptiongroup" version = "1.3.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions" }, + { name = "typing-extensions", marker = "python_full_version < '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/0b/9f/a65090624ecf468cdca03533906e7c69ed7588582240cfe7cc9e770b50eb/exceptiongroup-1.3.0.tar.gz", hash = "sha256:b241f5885f560bc56a59ee63ca4c6a8bfa46ae4ad651af316d4e81817bb9fd88", size = 29749, upload-time = "2025-05-10T17:42:51.123Z" } wheels = [ @@ -572,7 +572,7 @@ name = "importlib-metadata" version = "8.7.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "zipp" }, + { name = "zipp", marker = "python_full_version < '3.10'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/76/66/650a33bd90f786193e4de4b3ad86ea60b53c89b669a5c7be931fac31cdb0/importlib_metadata-8.7.0.tar.gz", hash = "sha256:d13b81ad223b890aa16c5471f2ac3056cf76c5f10f82d6f9292f0b415f389000", size = 56641, upload-time = "2025-04-27T15:29:01.736Z" } wheels = [ @@ -1239,7 +1239,7 @@ name = "nvidia-cublas" version = "13.1.1.3" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cuda-nvrtc" }, + { name = "nvidia-cuda-nvrtc", marker = "python_full_version >= '3.10' and sys_platform == 'linux'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/a7/a1/0bd24ee8c8d03adac032fd2909426a00c88f8c57961b1277ded97f91119f/nvidia_cublas-13.1.1.3-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:b7a210458267ac818974c53038fbec2e969d5c99f305ab15c72522fa9f001dd5", size = 542848918, upload-time = "2026-04-08T18:46:22.985Z" }, @@ -1310,7 +1310,7 @@ name = "nvidia-cudnn-cu12" version = "9.10.2.21" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas-cu12" }, + { name = "nvidia-cublas-cu12", marker = "python_full_version < '3.10' and sys_platform == 'linux'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/ba/51/e123d997aa098c61d029f76663dedbfb9bc8dcf8c60cbd6adbe42f76d049/nvidia_cudnn_cu12-9.10.2.21-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:949452be657fa16687d0930933f032835951ef0892b37d2d53824d1a84dc97a8", size = 706758467, upload-time = "2025-06-06T21:54:08.597Z" }, @@ -1321,7 +1321,7 @@ name = "nvidia-cudnn-cu13" version = "9.20.0.48" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas" }, + { name = "nvidia-cublas", marker = "python_full_version >= '3.10' and sys_platform == 'linux'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/56/c5/83384d846b2fd17c44bd499b36c75a45ed4f095fbbb2252294e89cea5c5c/nvidia_cudnn_cu13-9.20.0.48-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:e31454ae00094b0c55319d9d15b6fa2fc50a9e1c0f5c8c80fb75258234e731e1", size = 444574296, upload-time = "2026-03-09T19:28:27.751Z" }, @@ -1333,7 +1333,7 @@ name = "nvidia-cufft" version = "12.0.0.61" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink" }, + { name = "nvidia-nvjitlink", marker = "python_full_version >= '3.10' and sys_platform == 'linux'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/8b/ae/f417a75c0259e85c1d2f83ca4e960289a5f814ed0cea74d18c353d3e989d/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:2708c852ef8cd89d1d2068bdbece0aa188813a0c934db3779b9b1faa8442e5f5", size = 214053554, upload-time = "2025-09-04T08:31:38.196Z" }, @@ -1345,7 +1345,7 @@ name = "nvidia-cufft-cu12" version = "11.3.3.83" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink-cu12" }, + { name = "nvidia-nvjitlink-cu12", marker = "python_full_version < '3.10' and sys_platform == 'linux'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/1f/13/ee4e00f30e676b66ae65b4f08cb5bcbb8392c03f54f2d5413ea99a5d1c80/nvidia_cufft_cu12-11.3.3.83-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:4d2dd21ec0b88cf61b62e6b43564355e5222e4a3fb394cac0db101f2dd0d4f74", size = 193118695, upload-time = "2025-03-07T01:45:27.821Z" }, @@ -1390,9 +1390,9 @@ name = "nvidia-cusolver" version = "12.0.4.66" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas" }, - { name = "nvidia-cusparse" }, - { name = "nvidia-nvjitlink" }, + { name = "nvidia-cublas", marker = "python_full_version >= '3.10' and sys_platform == 'linux'" }, + { name = "nvidia-cusparse", marker = "python_full_version >= '3.10' and sys_platform == 'linux'" }, + { name = "nvidia-nvjitlink", marker = "python_full_version >= '3.10' and sys_platform == 'linux'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/c8/c3/b30c9e935fc01e3da443ec0116ed1b2a009bb867f5324d3f2d7e533e776b/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:02c2457eaa9e39de20f880f4bd8820e6a1cfb9f9a34f820eb12a155aa5bc92d2", size = 223467760, upload-time = "2025-09-04T08:33:04.222Z" }, @@ -1404,9 +1404,9 @@ name = "nvidia-cusolver-cu12" version = "11.7.3.90" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas-cu12" }, - { name = "nvidia-cusparse-cu12" }, - { name = "nvidia-nvjitlink-cu12" }, + { name = "nvidia-cublas-cu12", marker = "python_full_version < '3.10' and sys_platform == 'linux'" }, + { name = "nvidia-cusparse-cu12", marker = "python_full_version < '3.10' and sys_platform == 'linux'" }, + { name = "nvidia-nvjitlink-cu12", marker = "python_full_version < '3.10' and sys_platform == 'linux'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/85/48/9a13d2975803e8cf2777d5ed57b87a0b6ca2cc795f9a4f59796a910bfb80/nvidia_cusolver_cu12-11.7.3.90-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:4376c11ad263152bd50ea295c05370360776f8c3427b30991df774f9fb26c450", size = 267506905, upload-time = "2025-03-07T01:47:16.273Z" }, @@ -1417,7 +1417,7 @@ name = "nvidia-cusparse" version = "12.6.3.3" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink" }, + { name = "nvidia-nvjitlink", marker = "python_full_version >= '3.10' and sys_platform == 'linux'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/f8/94/5c26f33738ae35276672f12615a64bd008ed5be6d1ebcb23579285d960a9/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:80bcc4662f23f1054ee334a15c72b8940402975e0eab63178fc7e670aa59472c", size = 162155568, upload-time = "2025-09-04T08:33:42.864Z" }, @@ -1429,7 +1429,7 @@ name = "nvidia-cusparse-cu12" version = "12.5.8.93" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink-cu12" }, + { name = "nvidia-nvjitlink-cu12", marker = "python_full_version < '3.10' and sys_platform == 'linux'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/c2/f5/e1854cb2f2bcd4280c44736c93550cc300ff4b8c95ebe370d0aa7d2b473d/nvidia_cusparse_cu12-12.5.8.93-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1ec05d76bbbd8b61b06a80e1eaf8cf4959c3d4ce8e711b65ebd0443bb0ebb13b", size = 288216466, upload-time = "2025-03-07T01:48:13.779Z" }, @@ -2110,14 +2110,14 @@ resolution-markers = [ "python_full_version < '3.10' and sys_platform != 'linux'", ] dependencies = [ - { name = "click", version = "8.1.8", source = { registry = "https://pypi.org/simple" } }, - { name = "filelock" }, - { name = "jsonschema" }, - { name = "msgpack" }, - { name = "packaging" }, - { name = "protobuf" }, - { name = "pyyaml" }, - { name = "requests" }, + { name = "click", version = "8.1.8", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" }, + { name = "filelock", marker = "python_full_version < '3.10'" }, + { name = "jsonschema", marker = "python_full_version < '3.10'" }, + { name = "msgpack", marker = "python_full_version < '3.10'" }, + { name = "packaging", marker = "python_full_version < '3.10'" }, + { name = "protobuf", marker = "python_full_version < '3.10'" }, + { name = "pyyaml", marker = "python_full_version < '3.10'" }, + { name = "requests", marker = "python_full_version < '3.10'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/6a/ad/59270b7d1003152ef231b65c38c3721066fc970b2a2475314e7c8ee81990/ray-2.51.2-cp310-cp310-macosx_12_0_arm64.whl", hash = "sha256:eb9b995de9ba3110373f00e77dda86f6a55a80a58114b1eae5e6daf1f5697338", size = 68040029, upload-time = "2025-11-29T00:28:25.435Z" }, @@ -2156,14 +2156,14 @@ resolution-markers = [ "python_full_version == '3.10.*' and sys_platform != 'linux'", ] dependencies = [ - { name = "click", version = "8.2.1", source = { registry = "https://pypi.org/simple" } }, - { name = "filelock" }, - { name = "jsonschema" }, - { name = "msgpack" }, - { name = "packaging" }, - { name = "protobuf" }, - { name = "pyyaml" }, - { name = "requests" }, + { name = "click", version = "8.2.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.10'" }, + { name = "filelock", marker = "python_full_version >= '3.10'" }, + { name = "jsonschema", marker = "python_full_version >= '3.10'" }, + { name = "msgpack", marker = "python_full_version >= '3.10'" }, + { name = "packaging", marker = "python_full_version >= '3.10'" }, + { name = "protobuf", marker = "python_full_version >= '3.10'" }, + { name = "pyyaml", marker = "python_full_version >= '3.10'" }, + { name = "requests", marker = "python_full_version >= '3.10'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/2f/41/760f9a076085d043d9fbb8f840436e1d21e47682c976d3620efbf02a3b9e/ray-2.58.0-cp310-cp310-macosx_12_0_arm64.whl", hash = "sha256:a2235a5a17e865be053e8aff6d9e9e36d1edb45ad8162f57f7cdfd0202a5a397", size = 67217855, upload-time = "2026-08-23T04:45:46.075Z" }, @@ -2687,6 +2687,24 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/40/44/4a5f08c96eb108af5cb50b41f76142f0afa346dfa99d5296fe7202a11854/tabulate-0.9.0-py3-none-any.whl", hash = "sha256:024ca478df22e9340661486f85298cff5f6dcdba14f3813e8830015b9ed1948f", size = 35252, upload-time = "2022-10-06T17:21:44.262Z" }, ] +[[package]] +name = "timm" +version = "1.0.29" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "huggingface-hub" }, + { name = "pyyaml" }, + { name = "safetensors" }, + { name = "torch", version = "2.8.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" }, + { name = "torch", version = "2.13.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.10'" }, + { name = "torchvision", version = "0.23.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" }, + { name = "torchvision", version = "0.28.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.10'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/cf/72/58b8363e22508418d5fac897ac68e2c33f01f84d9d471edc38cbb6523dc6/timm-1.0.29.tar.gz", hash = "sha256:1af8d12bb7e15c5f96e98d60b7b8319cd7f31730778bf1af58e912a7ce171576", size = 2497559, upload-time = "2026-08-28T15:07:31.007Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5e/b7/19267ec60740ef1899d265272bff3623f19fe7f7bf1022fb098766db2134/timm-1.0.29-py3-none-any.whl", hash = "sha256:4266486e355c9beb9432e689de580d09e2efe7fa92a9a79e76c03e7763eac724", size = 2633555, upload-time = "2026-08-28T15:07:29.348Z" }, +] + [[package]] name = "tinyexp" version = "0.1.4" @@ -2717,6 +2735,15 @@ pytorch = [ { name = "torchvision", version = "0.23.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" }, { name = "torchvision", version = "0.28.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.10'" }, ] +vit = [ + { name = "accelerate" }, + { name = "pillow", version = "12.3.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.14'" }, + { name = "timm" }, + { name = "torch", version = "2.8.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" }, + { name = "torch", version = "2.13.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.10'" }, + { name = "torchvision", version = "0.23.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" }, + { name = "torchvision", version = "0.28.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.10'" }, +] [package.dev-dependencies] dev = [ @@ -2747,6 +2774,8 @@ requires-dist = [ { name = "redis" }, { name = "rpds-py", marker = "python_full_version >= '3.14'", specifier = ">=0.26" }, { name = "tabulate" }, + { name = "timm", marker = "extra == 'vit'", specifier = ">=1.0" }, + { name = "tinyexp", extras = ["pytorch"], marker = "extra == 'vit'" }, { name = "torch", marker = "python_full_version >= '3.14' and extra == 'pytorch'", specifier = ">=2.13" }, { name = "torch", marker = "python_full_version < '3.14' and extra == 'pytorch'" }, { name = "torchvision", marker = "python_full_version >= '3.14' and extra == 'pytorch'", specifier = ">=0.28" }, @@ -2754,7 +2783,7 @@ requires-dist = [ { name = "tqdm" }, { name = "wandb" }, ] -provides-extras = ["pytorch"] +provides-extras = ["pytorch", "vit"] [package.metadata.requires-dev] dev = [ @@ -2817,27 +2846,27 @@ resolution-markers = [ "python_full_version < '3.10' and sys_platform != 'linux'", ] dependencies = [ - { name = "filelock" }, - { name = "fsspec" }, - { name = "jinja2" }, - { name = "networkx", version = "3.2.1", source = { registry = "https://pypi.org/simple" } }, - { name = "nvidia-cublas-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-cuda-cupti-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-cuda-nvrtc-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-cuda-runtime-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-cudnn-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-cufft-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-cufile-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-curand-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-cusolver-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-cusparse-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-cusparselt-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-nccl-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-nvjitlink-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-nvtx-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "sympy" }, - { name = "triton", version = "3.4.0", source = { registry = "https://pypi.org/simple" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "typing-extensions" }, + { name = "filelock", marker = "python_full_version < '3.10'" }, + { name = "fsspec", marker = "python_full_version < '3.10'" }, + { name = "jinja2", marker = "python_full_version < '3.10'" }, + { name = "networkx", version = "3.2.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" }, + { name = "nvidia-cublas-cu12", marker = "python_full_version < '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-cupti-cu12", marker = "python_full_version < '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-nvrtc-cu12", marker = "python_full_version < '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-runtime-cu12", marker = "python_full_version < '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cudnn-cu12", marker = "python_full_version < '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cufft-cu12", marker = "python_full_version < '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cufile-cu12", marker = "python_full_version < '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-curand-cu12", marker = "python_full_version < '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cusolver-cu12", marker = "python_full_version < '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cusparse-cu12", marker = "python_full_version < '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cusparselt-cu12", marker = "python_full_version < '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-nccl-cu12", marker = "python_full_version < '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-nvjitlink-cu12", marker = "python_full_version < '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-nvtx-cu12", marker = "python_full_version < '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "sympy", marker = "python_full_version < '3.10'" }, + { name = "triton", version = "3.4.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "typing-extensions", marker = "python_full_version < '3.10'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/63/28/110f7274254f1b8476c561dada127173f994afa2b1ffc044efb773c15650/torch-2.8.0-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:0be92c08b44009d4131d1ff7a8060d10bafdb7ddcb7359ef8d8c5169007ea905", size = 102052793, upload-time = "2025-08-06T14:53:15.852Z" }, @@ -2881,21 +2910,21 @@ resolution-markers = [ "python_full_version == '3.10.*' and sys_platform != 'linux'", ] dependencies = [ - { name = "cuda-bindings", marker = "python_full_version < '3.15' and sys_platform == 'linux'" }, - { name = "cuda-toolkit", extra = ["cublas", "cudart", "cufft", "cufile", "cupti", "curand", "cusolver", "cusparse", "nvjitlink", "nvrtc", "nvtx"], marker = "sys_platform == 'linux'" }, - { name = "filelock" }, - { name = "fsspec" }, - { name = "jinja2" }, - { name = "networkx", version = "3.4.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "cuda-bindings", marker = "python_full_version >= '3.10' and python_full_version < '3.15' and sys_platform == 'linux'" }, + { name = "cuda-toolkit", extra = ["cublas", "cudart", "cufft", "cufile", "cupti", "curand", "cusolver", "cusparse", "nvjitlink", "nvrtc", "nvtx"], marker = "python_full_version >= '3.10' and sys_platform == 'linux'" }, + { name = "filelock", marker = "python_full_version >= '3.10'" }, + { name = "fsspec", marker = "python_full_version >= '3.10'" }, + { name = "jinja2", marker = "python_full_version >= '3.10'" }, + { name = "networkx", version = "3.4.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.10.*'" }, { name = "networkx", version = "3.5", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "nvidia-cudnn-cu13", marker = "sys_platform == 'linux'" }, - { name = "nvidia-cusparselt-cu13", marker = "sys_platform == 'linux'" }, - { name = "nvidia-nccl-cu13", marker = "sys_platform == 'linux'" }, - { name = "nvidia-nvshmem-cu13", marker = "sys_platform == 'linux'" }, - { name = "setuptools" }, - { name = "sympy" }, - { name = "triton", version = "3.7.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.15' and sys_platform == 'linux'" }, - { name = "typing-extensions" }, + { name = "nvidia-cudnn-cu13", marker = "python_full_version >= '3.10' and sys_platform == 'linux'" }, + { name = "nvidia-cusparselt-cu13", marker = "python_full_version >= '3.10' and sys_platform == 'linux'" }, + { name = "nvidia-nccl-cu13", marker = "python_full_version >= '3.10' and sys_platform == 'linux'" }, + { name = "nvidia-nvshmem-cu13", marker = "python_full_version >= '3.10' and sys_platform == 'linux'" }, + { name = "setuptools", marker = "python_full_version >= '3.10'" }, + { name = "sympy", marker = "python_full_version >= '3.10'" }, + { name = "triton", version = "3.7.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.10' and python_full_version < '3.15' and sys_platform == 'linux'" }, + { name = "typing-extensions", marker = "python_full_version >= '3.10'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/7f/e7/19894fdb51c7dbaf94f5a79bb0871da0992e8e4241e579cb006da46d2e58/torch-2.13.0-cp310-cp310-macosx_14_0_arm64.whl", hash = "sha256:94f0de129916f77b8dc2c7a8eff644cfeddfe59e39c9f55e9f6e17543410281d", size = 111178962, upload-time = "2026-07-08T16:05:49.855Z" }, @@ -2933,9 +2962,9 @@ resolution-markers = [ "python_full_version < '3.10' and sys_platform != 'linux'", ] dependencies = [ - { name = "numpy", version = "2.0.2", source = { registry = "https://pypi.org/simple" } }, - { name = "pillow", version = "11.2.1", source = { registry = "https://pypi.org/simple" } }, - { name = "torch", version = "2.8.0", source = { registry = "https://pypi.org/simple" } }, + { name = "numpy", version = "2.0.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" }, + { name = "pillow", version = "11.2.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" }, + { name = "torch", version = "2.8.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/4d/49/5ad5c3ff4920be0adee9eb4339b4fb3b023a0fc55b9ed8dbc73df92946b8/torchvision-0.23.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:7266871daca00ad46d1c073e55d972179d12a58fa5c9adec9a3db9bbed71284a", size = 1856885, upload-time = "2025-08-06T14:57:55.024Z" }, @@ -2979,11 +3008,11 @@ resolution-markers = [ "python_full_version == '3.10.*' and sys_platform != 'linux'", ] dependencies = [ - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.14'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.10' and python_full_version < '3.14'" }, { name = "numpy", version = "2.5.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.14'" }, - { name = "pillow", version = "11.2.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.14'" }, + { name = "pillow", version = "11.2.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.10' and python_full_version < '3.14'" }, { name = "pillow", version = "12.3.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.14'" }, - { name = "torch", version = "2.13.0", source = { registry = "https://pypi.org/simple" } }, + { name = "torch", version = "2.13.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.10'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/b4/df/1ba039ad6cfe6e69209c36766b9b6e8c6fe92481c6d4e4ca52296f5f699d/torchvision-0.28.0-cp310-cp310-macosx_14_0_arm64.whl", hash = "sha256:2a1ef4b6f4bf5828b48cfad97372c8982db906830884b2868ba5c3df937a7d81", size = 1856019, upload-time = "2026-07-08T16:07:59.283Z" }, @@ -3069,8 +3098,8 @@ resolution-markers = [ "python_full_version < '3.10' and sys_platform == 'linux'", ] dependencies = [ - { name = "importlib-metadata" }, - { name = "setuptools" }, + { name = "importlib-metadata", marker = "python_full_version < '3.10' and sys_platform == 'linux'" }, + { name = "setuptools", marker = "python_full_version < '3.10' and sys_platform == 'linux'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/62/ee/0ee5f64a87eeda19bbad9bc54ae5ca5b98186ed00055281fd40fb4beb10e/triton-3.4.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7ff2785de9bc02f500e085420273bb5cc9c9bb767584a4aa28d6e360cec70128", size = 155430069, upload-time = "2025-07-30T19:58:21.715Z" }, @@ -3124,7 +3153,7 @@ resolution-markers = [ "python_full_version < '3.10' and sys_platform != 'linux'", ] dependencies = [ - { name = "typing-extensions" }, + { name = "typing-extensions", marker = "python_full_version < '3.10'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/55/e3/70399cb7dd41c10ac53367ae42139cf4b1ca5f36bb3dc6c9d33acdb43655/typing_inspection-0.4.2.tar.gz", hash = "sha256:ba561c48a67c5958007083d386c3295464928b01faa735ab8547c5692e87f464", size = 75949, upload-time = "2025-10-01T02:14:41.687Z" } wheels = [ @@ -3146,7 +3175,7 @@ resolution-markers = [ "python_full_version == '3.10.*' and sys_platform != 'linux'", ] dependencies = [ - { name = "typing-extensions" }, + { name = "typing-extensions", marker = "python_full_version >= '3.10'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/a3/26/b09b8010994eccc3c09092e6b34058f36a460eea2d4c3e8b910c695975a0/typing_inspection-0.4.4.tar.gz", hash = "sha256:547274fa6b0a561ccf549cc9524b999a578e737d015d8709d021f9d0d13bea47", size = 76928, upload-time = "2026-08-12T12:37:25.997Z" } wheels = [