Feat/multi gpu eval

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

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 ``<image>`` 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,

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)
 # ---------------------------------------------------------------------------