diff --git a/examples/vlm-evaluation/README.md b/examples/vlm-evaluation/README.md index 21c2f66..d8a919e 100644 --- a/examples/vlm-evaluation/README.md +++ b/examples/vlm-evaluation/README.md @@ -81,6 +81,38 @@ automatically) or a specific `step_N` directory. There is **no default task suite** — `--tasks` is required. A representative default benchmark set is still being decided. +## Multi-GPU (data parallel) + +A single benchmark can be **data-parallelized across GPUs**: launch the harness with +`accelerate launch --num_processes N` and lmms-eval shards the benchmark's documents +`[rank::world_size]` across the N processes, gathering results onto rank 0. Each rank loads a +**full replica** of the model on its own GPU, so throughput scales ~N× with identical scores. +`accelerate` ships with lmms-eval, so there is nothing extra to install. + +```bash +accelerate launch --num_processes 4 examples/vlm-evaluation/vlm_eval_harness.py \ + --config configs/train/vlm_jd.toml \ + --checkpoint checkpoints/vlm/step_10000 \ + --tasks mmmu_val \ + --output results/vlm_step_10000.json +``` + +- **No new flags.** DP is entirely launcher-driven — the adapter auto-detects the run from the + `WORLD_SIZE` / `LOCAL_RANK` environment variables the launcher sets and binds each rank to + `cuda:LOCAL_RANK`. Plain `uv run python vlm_eval_harness.py …` (no launcher) is unchanged + single-GPU. +- **Pass `--device cuda` without an index.** Per-rank GPU binding triggers only on the bare + `cuda`; an explicit `--device cuda:0` would pin *every* rank to GPU 0. +- **Replication, not model-parallel.** Every rank holds a full copy, so aggregate GPU memory is + N× — the model must fit on one GPU. Sharded inference for larger models is separate future + work (see [Limitations](#limitations)). +- **`--limit N` stays a global cap** (the per-rank shards union to `N` docs total); `--batch-size` + is per rank. +- **Only rank 0 writes** the `--output` JSON; the other ranks score their shard and exit. +- On the Kempner cluster, request the GPUs in your allocation (`--gres=gpu:N` on one node) and keep + the same `LD_LIBRARY_PATH` / `HF_HOME` environment as a single-GPU run (see + [Cluster environment notes](#cluster-environment-notes)). + ## Flags | Flag | Default | Purpose | @@ -155,10 +187,11 @@ uv run python examples/vlm-evaluation/vlm_eval_harness.py \ Several are tracked follow-ups. -- **Single GPU.** v1 runs on one GPU. Data-parallel - multi-GPU is a localized - future addition; sharded/model-parallel inference for models too large for one - GPU is a larger, separate effort. +- **Data parallel, replicated.** Multi-GPU runs shard the benchmark's documents across GPUs + via `accelerate launch --num_processes N` (see + [Multi-GPU (data parallel)](#multi-gpu-data-parallel)); each rank holds a **full replica** of + the model. Sharded / model-parallel inference for models too large for a single GPU is a + larger, separate effort. - **MoMa is not supported.** The `moma` arch uses non-causal expert-choice routing and cannot autoregressively generate, but eval tasks are generation-only. A MoMa checkpoint fails fast with a clear error. Joint-Decoder diff --git a/examples/vlm-evaluation/adapter.py b/examples/vlm-evaluation/adapter.py index 350da30..9760db2 100644 --- a/examples/vlm-evaluation/adapter.py +++ b/examples/vlm-evaluation/adapter.py @@ -14,7 +14,7 @@ and is arch-agnostic across the generative VLM arches. v1 scope and deliberate choices (see README.md in this directory): -- **Generation: no transformer KV cache, single-GPU, batched.** The decode loop +- **Generation: no transformer KV cache, batched, data-parallel.** The decode loop re-runs the transformer (including the vision encoder + adapter) over the growing sequence each step. There is no transformer KV cache (``Transformer.forward`` forbids combining ``kv_caches`` with any @@ -23,10 +23,12 @@ (``batch_size`` model-arg) by **right-padding** the text to the batch-max length — the same layout training uses (image prefix at ``0..n-1``, text contiguous from ``n``, trailing pads causally masked) — and reading each - row's logits at its own last real position. Single-GPU is the validated - invocation, not a baked-in assumption: rank/world_size come from the lmms - base (defaults 0/1) and model construction sits behind ``_build_model`` so a - data-parallel path is a localized future change. + row's logits at its own last real position. **Data parallelism** is supported: + under ``accelerate launch --num_processes N`` each rank loads a full model + replica (rank/world_size come from an ``accelerate.Accelerator`` built after the + checkpoint load; single-process defaults to 0/1) and lmms-eval shards the + benchmark's documents across ranks, gathering results on rank 0. Sharded / + model-parallel inference for models too large for one GPU is a separate effort. - **Prompt rendering: flatten, no chat template.** KempnerForge pre-training uses no chat template / processor and no ```` placeholder (images are @@ -61,7 +63,9 @@ from __future__ import annotations import json +import os import time +from datetime import timedelta from pathlib import Path from typing import Any @@ -179,12 +183,14 @@ def _check_generative(vlm_config: VLMConfig) -> None: def _load_weights( config: JobConfig, checkpoint: str, device: torch.device, dtype: torch.dtype ) -> VLMWrapper: - """Build a ``VLMWrapper`` and load DCP weights for single-process eval. + """Build a ``VLMWrapper`` and load DCP weights as a full per-rank replica. Accepts either a run directory (resolved to its ``latest``/highest ``step_N`` via ``resolve_resume_path``) or a specific checkpoint directory (used as-is when ``resolve_resume_path`` finds nothing). DCP reshards on - load, so checkpoints saved under FSDP/PP load into the full model. + load, so checkpoints saved under FSDP/PP load into the full model. Under a + data-parallel launch every rank calls this independently and loads the same + full weights (see the ``no_dist=True`` note below). """ ckpt_path = resolve_resume_path(checkpoint) or Path(checkpoint) if not ckpt_path.exists(): @@ -194,10 +200,15 @@ def _load_weights( model = _build_model(config, device, dtype) model.eval() - # Single-process DCP load: build the full (unsharded) model, then load the - # model shards into its state-dict. + # Per-rank full-replica DCP load: build the full (unsharded) model, then load the + # model shards into its state-dict. no_dist=True forces dcp.load's independent + # single-process path (never a collective reshard across the DP process group), so every + # data-parallel rank ends up with a full, identical copy of the weights. In the + # single-GPU path this is a no-op — dcp.load already selects no_dist when + # torch.distributed is not initialized — but stating it makes the load correct regardless + # of whether a process group happens to be live when this runs. state_dict = {"model": model.state_dict()} - dcp.load(state_dict, checkpoint_id=str(ckpt_path)) + dcp.load(state_dict, checkpoint_id=str(ckpt_path), no_dist=True) model.load_state_dict(state_dict["model"]) _log_checkpoint_metadata(ckpt_path) @@ -569,7 +580,20 @@ def __init__( if kwargs: logger.warning(f"Ignoring unsupported model_args: {sorted(kwargs)}") - self._device = torch.device(device) + # Data-parallel (DP) launch detection. Under `accelerate launch --num_processes N` + # the launcher sets WORLD_SIZE / LOCAL_RANK / RANK in the environment for every + # process before any Accelerator or process group exists. Read them here so this + # replica can (a) bind to its own GPU below and (b) create the process group only + # AFTER the checkpoint load (see _load_weights call). With no launcher these default + # to 1 / 0 -> the single-GPU path, byte-identical to before. + env_world_size = int(os.environ.get("WORLD_SIZE", "1")) + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + if env_world_size > 1 and device == "cuda": + # One full model replica per GPU: pin this process to its local device. Pass + # --device cuda (no index) under DP; an explicit cuda:i would collide every rank. + self._device = torch.device(f"cuda:{local_rank}") + else: + self._device = torch.device(device) self._batch_size = int(batch_size) self._default_max_new_tokens = int(max_new_tokens) if self._batch_size < 1: @@ -598,6 +622,26 @@ def __init__( self._frame_size = self._config.data.hf_image_size self._model = _load_weights(self._config, checkpoint, self._device, self._dtype) + if env_world_size > 1: + # Initialize accelerate ONLY under a real multi-process launch, and only AFTER the + # DCP load above so that load runs its independent no_dist path on every rank -> + # one full identical replica per GPU (pure data parallelism). Constructing the + # Accelerator here initializes the torch.distributed process group that + # lmms-eval's evaluator uses to gather per-rank results onto rank 0. + from accelerate import Accelerator, InitProcessGroupKwargs + + # A long NCCL timeout mirrors lmms-eval's own models: with no KV cache, ranks can + # drift far apart in wall-clock (uneven shards, judge/aggregation gaps) and the + # ~30-min default would abort the collective mid-run. + accelerator = Accelerator( + kwargs_handlers=[InitProcessGroupKwargs(timeout=timedelta(weeks=52))] + ) + self.accelerator = accelerator + # Global rank (process_index), not local_process_index: the evaluator shards docs + # by env RANK and indexes an accelerate-gathered, global-rank-ordered tensor as + # gathered_item[lm.rank]. The per-node device already uses LOCAL_RANK above. + self._rank = accelerator.process_index + self._world_size = accelerator.num_processes self._tokenizer = build_tokenizer(self._config.data.tokenizer_path) self._max_seq_len = self._config.model.max_seq_len logger.info( @@ -606,6 +650,13 @@ def __init__( f"dtype={self._dtype}, max_seq_len={self._max_seq_len}" ) + @property + def device(self) -> torch.device: + # lmms-eval's evaluator reads ``lm.device`` when world_size > 1 (to place the + # per-rank instance-count tensor it gathers). The lmms base exposes rank/world_size + # but no device, so the adapter provides it. Single-process: the cuda/cpu passed in. + return self._device + def _decode_subbatch( self, pixel_values: torch.Tensor | None, diff --git a/examples/vlm-evaluation/tests/unit/test_adapter.py b/examples/vlm-evaluation/tests/unit/test_adapter.py index 3f42973..f1f7381 100644 --- a/examples/vlm-evaluation/tests/unit/test_adapter.py +++ b/examples/vlm-evaluation/tests/unit/test_adapter.py @@ -10,6 +10,8 @@ from __future__ import annotations import json +import sys +import types import pytest import torch @@ -877,6 +879,105 @@ def test_explicit_dtype_overrides_config(self, monkeypatch, tiny_vlm_configs, ti assert vlm._dtype == torch.float16 +# --------------------------------------------------------------------------- +# Data-parallel (multi-GPU) wiring +# --------------------------------------------------------------------------- + + +def _install_fake_accelerate(monkeypatch, process_index=0, num_processes=2) -> dict: + """Inject a stub ``accelerate`` so the adapter's lazy + ``from accelerate import Accelerator, InitProcessGroupKwargs`` resolves to a fake that + never initializes a real process group (the real ``Accelerator`` would try to reach NCCL + and hang). Returns a dict capturing the ``Accelerator`` constructor kwargs. + """ + captured: dict = {} + + class _FakeAccelerator: + def __init__(self, kwargs_handlers=None): + captured["kwargs_handlers"] = kwargs_handlers + self.process_index = process_index + self.num_processes = num_processes + self.local_process_index = process_index + + class _FakeInitProcessGroupKwargs: + def __init__(self, timeout=None): + self.timeout = timeout + + fake = types.ModuleType("accelerate") + fake.Accelerator = _FakeAccelerator # type: ignore[attr-defined] + fake.InitProcessGroupKwargs = _FakeInitProcessGroupKwargs # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "accelerate", fake) + return captured + + +class TestDataParallelWiring: + """DP multi-GPU wiring, exercised hermetically on CPU: the launcher env is faked with + ``monkeypatch`` and a stub ``accelerate.Accelerator`` replaces the real one (which would + initialize a live process group and hang). Mirrors ``TestInitGuards``. + """ + + def test_multi_process_env_sets_rank_world_and_accelerator( + self, monkeypatch, tiny_vlm_configs, tiny_vlm_wrapper + ): + _patch_loaders(monkeypatch, _vlm_job_config(tiny_vlm_configs), tiny_vlm_wrapper) + captured = _install_fake_accelerate(monkeypatch, process_index=1, num_processes=2) + monkeypatch.setenv("WORLD_SIZE", "2") + monkeypatch.setenv("LOCAL_RANK", "1") + # device="cpu" keeps the test off CUDA; the Accelerator block is gated on WORLD_SIZE, + # not device, so rank/world_size/accelerator still populate from the fake. + vlm = KempnerForgeVLM(config="x", checkpoint="y", device="cpu", dtype="float32") + # rank must be the GLOBAL process_index: the evaluator indexes gathered_item[lm.rank]. + assert vlm._rank == 1 and vlm.rank == 1 + assert vlm._world_size == 2 and vlm.world_size == 2 + assert isinstance(vlm.accelerator, sys.modules["accelerate"].Accelerator) + # The long-timeout InitProcessGroupKwargs handler is passed through to the Accelerator. + assert captured["kwargs_handlers"] is not None and len(captured["kwargs_handlers"]) == 1 + + def test_multi_process_cuda_binds_local_rank_device( + self, monkeypatch, tiny_vlm_configs, tiny_vlm_wrapper + ): + # device="cuda" + WORLD_SIZE>1 pins this replica to cuda:LOCAL_RANK. No CUDA is + # touched: torch.device is a descriptor and _load_weights is patched to a no-op. + _patch_loaders(monkeypatch, _vlm_job_config(tiny_vlm_configs), tiny_vlm_wrapper) + _install_fake_accelerate(monkeypatch, process_index=1, num_processes=2) + monkeypatch.setenv("WORLD_SIZE", "2") + monkeypatch.setenv("LOCAL_RANK", "1") + vlm = KempnerForgeVLM(config="x", checkpoint="y", device="cuda", dtype="float32") + assert vlm._device == torch.device("cuda:1") + assert vlm.device == torch.device("cuda:1") # the new property + + def test_device_property_returns_underlying_device( + self, monkeypatch, tiny_vlm_configs, tiny_vlm_wrapper + ): + _patch_loaders(monkeypatch, _vlm_job_config(tiny_vlm_configs), tiny_vlm_wrapper) + vlm = KempnerForgeVLM(config="x", checkpoint="y", device="cpu", dtype="float32") + assert vlm.device is vlm._device + assert vlm.device == torch.device("cpu") + + def test_single_process_unchanged_no_accelerate( + self, monkeypatch, tiny_vlm_configs, tiny_vlm_wrapper + ): + # No launcher env -> the adapter must NOT import/construct an Accelerator, and + # rank/world_size/device keep the single-GPU defaults. A booby-trapped stub proves the + # accelerate path is never taken. + monkeypatch.delenv("WORLD_SIZE", raising=False) + monkeypatch.delenv("LOCAL_RANK", raising=False) + boom = types.ModuleType("accelerate") + + def _explode(*args, **kwargs): + raise AssertionError("Accelerator must not be constructed in single-process mode") + + boom.Accelerator = _explode # type: ignore[attr-defined] + boom.InitProcessGroupKwargs = _explode # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "accelerate", boom) + _patch_loaders(monkeypatch, _vlm_job_config(tiny_vlm_configs), tiny_vlm_wrapper) + vlm = KempnerForgeVLM(config="x", checkpoint="y", device="cpu", dtype="float32") + assert vlm._rank == 0 and vlm.rank == 0 + assert vlm._world_size == 1 and vlm.world_size == 1 + assert vlm._device == torch.device("cpu") + assert not hasattr(vlm, "accelerator") + + # --------------------------------------------------------------------------- # _build_model (frames_per_clip wiring for image vs video checkpoints) # --------------------------------------------------------------------------- diff --git a/examples/vlm-evaluation/vlm_eval_harness.py b/examples/vlm-evaluation/vlm_eval_harness.py index d42481b..d137a2a 100644 --- a/examples/vlm-evaluation/vlm_eval_harness.py +++ b/examples/vlm-evaluation/vlm_eval_harness.py @@ -2,18 +2,16 @@ """Run lmms-eval benchmarks on a KempnerForge VLM checkpoint. Evaluates a VLM checkpoint via the ``KempnerForgeVLM`` lmms-eval chat-model -adapter (the sibling ``adapter.py``), on the standard benchmarks lmms-eval -implements as ``generate_until`` tasks (MMMU, MMBench, ScienceQA, SEED, AI2D, -...). The harness constructs the adapter directly and passes the instance to -``simple_evaluate`` — there is no lmms-eval entry-point registration. +adapter, on the standard benchmarks lmms-eval implements as ``generate_until`` +tasks. Requirements (lmms-eval is an OPTIONAL, separately-installed dependency, exactly like lm-eval for text evaluation): uv pip install lmms-eval -v1 is single-GPU; MoMa checkpoints are not supported (see README.md in this -directory). On clusters where importing lmms-eval's evaluator fails with +MoMa checkpoints are not supported (see README.md in this directory). +On clusters where importing lmms-eval's evaluator fails with ``GLIBCXX_... not found``, put a newer libstdc++ on the library path (e.g. ``LD_LIBRARY_PATH=/lib``). @@ -115,8 +113,11 @@ def main() -> None: sys.exit(1) # The adapter imports lmms-eval at module top; the guard above already proved - # it importable. The script's own directory is sys.path[0], so the sibling - # adapter.py resolves as a top-level module. + # it importable. The script's own directory is normally sys.path[0], so the sibling + # adapter.py resolves as a top-level module — but under `accelerate launch` the launcher + # may run this file such that sys.path[0] is not its directory, so insert it explicitly + # so `from adapter import ...` resolves on every rank. + sys.path.insert(0, str(Path(__file__).resolve().parent)) from adapter import KempnerForgeVLM logger.info(f"Running lmms-eval: tasks={args.tasks}, checkpoint={args.checkpoint}") @@ -150,27 +151,31 @@ def main() -> None: cli_args=cli_args, ) - # --- Print results --- - print(f"\n{'=' * 60}") - print("lmms-eval Results") - print(f"{'=' * 60}") - if results is not None and "results" in results: - for task_name, task_results in sorted(results["results"].items()): - print(f"\n {task_name}:") - for metric, value in sorted(task_results.items()): - if isinstance(value, float): - print(f" {metric}: {value:.4f}") - elif metric != "alias": - print(f" {metric}: {value}") - print(f"{'=' * 60}\n") - - # --- Save results --- - if args.output: - output_path = Path(args.output) - output_path.parent.mkdir(parents=True, exist_ok=True) - with open(output_path, "w") as f: - json.dump(results, f, indent=2, default=str) - logger.info(f"Results saved to {output_path}") + # Only rank 0 holds the aggregated results (simple_evaluate returns None on non-zero + # ranks); every other DP rank must skip reporting so it does not print an empty banner, + # dump `None` to --output, or race on the same file. Single-process: rank 0. + if model.rank == 0: + # --- Print results --- + print(f"\n{'=' * 60}") + print("lmms-eval Results") + print(f"{'=' * 60}") + if results is not None and "results" in results: + for task_name, task_results in sorted(results["results"].items()): + print(f"\n {task_name}:") + for metric, value in sorted(task_results.items()): + if isinstance(value, float): + print(f" {metric}: {value:.4f}") + elif metric != "alias": + print(f" {metric}: {value}") + print(f"{'=' * 60}\n") + + # --- Save results --- + if args.output: + output_path = Path(args.output) + output_path.parent.mkdir(parents=True, exist_ok=True) + with open(output_path, "w") as f: + json.dump(results, f, indent=2, default=str) + logger.info(f"Results saved to {output_path}") if __name__ == "__main__":