Optimize training data paths and enable local one-GPU training (#41)

* Use shared trainer loops for single-GPU SFT

* Enable local single-GPU VLM DPO

* Enable local single-GPU MoE LoRA training

* Add local single-GPU E2E coverage

* Clarify MoE dataset setup comment

* Run E2E tests through default dispatcher

* Optimize Ray training data paths

* Simplify local trainer eligibility

---------

Co-authored-by: Alay Shah <alay.shah@liquid.ai>

Author: alay2shah

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Liquid AI
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -203,8 +203,10 @@ Launch training:
 uv run leap-finetune job_configs/sft_example.yaml
 ```
 
-Training uses Ray Train and Accelerate for distributed execution. SFT, DPO, VLM SFT, and VLM DPO automatically use the native Hugging Face Trainer when exactly one GPU
-is visible; set `LEAP_LAUNCHER=ray` to force Ray. Unless
+Training uses Ray Train and Accelerate for distributed execution. SFT, DPO, VLM
+SFT, and VLM DPO automatically use the native Hugging Face Trainer when exactly
+one GPU is visible. Text and VLM GRPO use the native Trainer only for one-GPU
+`vllm_mode: colocate` runs. Set `LEAP_LAUNCHER=ray` to force Ray. Unless
 `output_dir` is set, results are written to
 `outputs/{project_name}/{run_name}/`. Each run gets a unique name based on the
 model, dataset, learning rate, and timestamp.
@@ -262,7 +264,9 @@ uvx --from . leap-finetune /absolute/path/to/config.yaml
 
 You can also start a run from Python. This uses the same backend dispatch as
 the CLI: configs with `slurm`, `modal`, or `kuberay` submit remotely; other
-configs run local training and require visible CUDA devices. SFT, DPO, VLM SFT, and VLM DPO use the native Trainer on one GPU; other modes use Ray.
+configs run local training and require visible CUDA devices. SFT, DPO, VLM SFT,
+and VLM DPO use the native Trainer on one GPU; GRPO uses it only for one-GPU
+colocated vLLM; other modes use Ray.
 
 ```python
 from leap_finetune import run_config

src/leap_finetune/data_loading/image_loader.py

@@ -1,5 +1,8 @@
 import io
 import logging
+import os
+from collections import OrderedDict
+from threading import RLock
 
 import requests
 from PIL import Image, ImageFile
@@ -10,6 +13,44 @@
 
 logger = logging.getLogger(__name__)
 
+_IMAGE_CACHE_MAX_ITEMS = 32
+_IMAGE_CACHE: OrderedDict[tuple[str, int, int], Image.Image] = OrderedDict()
+_IMAGE_CACHE_LOCK = RLock()
+
+
+def _load_cached_path_image(src) -> Image.Image:
+    """Load a local image with a small per-process cache.
+
+    Return a copy so callers may close their image without invalidating the
+    cached object. File size and mtime invalidate entries when a path changes.
+    """
+    path = os.fspath(src)
+    try:
+        stat = os.stat(path)
+    except OSError:
+        with Image.open(path) as image:
+            return image.convert("RGB")
+
+    key = (path, stat.st_mtime_ns, stat.st_size)
+    with _IMAGE_CACHE_LOCK:
+        cached = _IMAGE_CACHE.pop(key, None)
+        if cached is not None:
+            _IMAGE_CACHE[key] = cached
+            return cached.copy()
+
+    with Image.open(path) as image:
+        decoded = image.convert("RGB")
+
+    with _IMAGE_CACHE_LOCK:
+        for old_key in list(_IMAGE_CACHE):
+            if old_key[0] == path:
+                _IMAGE_CACHE.pop(old_key).close()
+        _IMAGE_CACHE[key] = decoded
+        while len(_IMAGE_CACHE) > _IMAGE_CACHE_MAX_ITEMS:
+            _, evicted = _IMAGE_CACHE.popitem(last=False)
+            evicted.close()
+        return decoded.copy()
+
 
 def load_image(src) -> Image.Image:
     """Load image from various sources and return PIL Image in RGB."""
@@ -24,7 +65,26 @@ def load_image(src) -> Image.Image:
         resp.raise_for_status()
         return Image.open(resp.raw).convert("RGB")
     # file path -> PIL
-    return Image.open(src).convert("RGB")
+    return _load_cached_path_image(src)
+
+
+def get_image_size(src) -> tuple[int, int]:
+    """Read image dimensions without decoding the full local image."""
+    if isinstance(src, (bytes, bytearray)):
+        with Image.open(io.BytesIO(src)) as image:
+            return image.size
+    if isinstance(src, str) and src.startswith(("http://", "https://")):
+        resp = requests.get(
+            src, stream=True, headers={"User-Agent": "leap-finetune"}, timeout=15
+        )
+        try:
+            resp.raise_for_status()
+            with Image.open(resp.raw) as image:
+                return image.size
+        finally:
+            resp.close()
+    with Image.open(src) as image:
+        return image.size
 
 
 def is_image_loadable(src: str) -> bool:

src/leap_finetune/data_loading/length_grouping.py

@@ -41,3 +41,27 @@ def get_length_grouped_sampler(
         model_input_name="input_ids",
         generator=generator,
     )
+
+
+def get_tile_count_grouped_sampler(
+    dataset: Dataset | None,
+    batch_size: int,
+    *,
+    generator: torch.Generator | None = None,
+):
+    """Group VLM examples with similar processed image tile counts."""
+    if dataset is None:
+        return None
+    column_names = getattr(dataset, "column_names", None)
+    if column_names is None or "_vlm_tile_count" not in column_names:
+        return None
+    counts = list(dataset["_vlm_tile_count"])
+    if not counts:
+        return None
+    if len(set(counts)) == 1:
+        return None
+    return LengthGroupedSampler(
+        batch_size=batch_size,
+        lengths=[max(1, int(count)) for count in counts],
+        generator=generator,
+    )

src/leap_finetune/data_loading/ray_data_utils.py

@@ -319,7 +319,7 @@ def _tokenize_datasets(
                 eval_ds,
                 tokenizer,
                 max_length,
-                packing=False,
+                packing=packing,
                 assistant_only_loss=assistant_only_loss,
                 completion_only_loss=completion_only_loss,
                 drop_overlength=drop_overlength,

src/leap_finetune/data_loading/tokenize_data.py

@@ -4,8 +4,7 @@
 import ray
 import ray.data
 import torch
-from datasets import Dataset
-import pyarrow as pa
+from datasets import Dataset, Features, Sequence, Value
 from rich.console import Console
 from trl.data_utils import maybe_apply_chat_template, maybe_extract_prompt
 from trl.data_utils import pack_dataset
@@ -268,26 +267,20 @@ def tokenize_and_pack_sft(
 
     # === 2. Pack or truncate ===
     if packing:
-        # Keep packing Arrow-native; the previous row list doubled peak memory.
-        if hasattr(ds, "to_arrow_refs"):
-            tables = ray.get(ds.to_arrow_refs())
-            if not tables:
-                return ds
-            table = pa.concat_tables(tables)
-            columns = [
-                name
-                for name in ("input_ids", "assistant_masks", "completion_mask")
-                if name in table.column_names
-            ]
-            hf_ds = Dataset(table.select(columns))
-        else:
-            hf_ds = Dataset.from_generator(ds.iter_rows)
-            columns = [
-                name
-                for name in ("input_ids", "assistant_masks", "completion_mask")
-                if name in hf_ds.column_names
-            ]
-            hf_ds = hf_ds.select_columns(columns)
+        # Packing requires full materialization into an HF Dataset
+        rows = []
+        features_dict = {"input_ids": Sequence(Value("int64"))}
+        for row in ds.iter_rows():
+            packed_row = {"input_ids": row["input_ids"]}
+            if "assistant_masks" in row:
+                packed_row["assistant_masks"] = row["assistant_masks"]
+                features_dict["assistant_masks"] = Sequence(Value("int64"))
+            if "completion_mask" in row:
+                packed_row["completion_mask"] = row["completion_mask"]
+                features_dict["completion_mask"] = Sequence(Value("int64"))
+            rows.append(packed_row)
+        features = Features(features_dict)
+        hf_ds = Dataset.from_list(rows, features=features)
         console.print(f"[dim]Tokenized {len(hf_ds):,} rows[/dim]")
         console.print(f"[dim]Packing sequences (BFD, max_length={max_length})...[/dim]")
         hf_ds = pack_dataset(hf_ds, seq_length=max_length, strategy="bfd")
@@ -347,10 +340,15 @@ def tokenize_dpo(
 
     # Column names must match TRL v1's DPO data collator:
     # prompt_ids, chosen_ids, rejected_ids (changed from *_input_ids in TRL 0.x)
+    # Include the actual collated sequence length so length grouping works for DPO.
     return {
         "prompt_ids": list(prompt_input_ids),
         "chosen_ids": list(chosen_input_ids),
         "rejected_ids": list(rejected_input_ids),
+        "length": max(
+            len(prompt_input_ids) + len(chosen_input_ids),
+            len(prompt_input_ids) + len(rejected_input_ids),
+        ),
     }
 
 
@@ -374,12 +372,7 @@ def tokenize_dpo_dataset(
             "max_completion_length": max_completion_length,
         },
     )
-
-    arrow_refs = ds.to_arrow_refs()
-    if not arrow_refs:
-        return ds
-
-    tables = ray.get(arrow_refs)
-    row_count = sum(table.num_rows for table in tables)
-    console.print(f"[dim]Tokenized {row_count:,} DPO rows[/dim]")
-    return ray.data.from_arrow(pa.concat_tables(tables))
+    # Keep tokenization lazy in Ray. The driver must not gather all tokenized
+    # DPO rows with ray.get(); Ray Train shards the dataset first, and each
+    # worker materializes only its local shard for the HF trainer.
+    return ds

src/leap_finetune/data_loading/vlm_batching.py

@@ -0,0 +1,91 @@
+import logging
+from collections.abc import Iterator
+
+from datasets import Dataset
+
+from leap_finetune.data_loading.image_loader import get_image_size
+
+logger = logging.getLogger(__name__)
+VLM_TILE_COUNT_COLUMN = "_vlm_tile_count"
+
+
+def _image_sources(value) -> Iterator[str]:
+    if isinstance(value, dict):
+        if value.get("type") == "image" and isinstance(value.get("image"), str):
+            yield value["image"]
+            return
+        for child in value.values():
+            yield from _image_sources(child)
+    elif isinstance(value, (list, tuple)):
+        for child in value:
+            yield from _image_sources(child)
+
+
+def _row_image_sources(row: dict) -> Iterator[str]:
+    for key in ("messages", "prompt", "chosen", "rejected"):
+        if key in row:
+            yield from _image_sources(row[key])
+    if isinstance(row.get("image"), str):
+        yield row["image"]
+    if isinstance(row.get("images"), (list, tuple)):
+        yield from (image for image in row["images"] if isinstance(image, str))
+
+
+def _tile_count(image_processor, height: int, width: int) -> int:
+    """Match LFM2-VL's grid decision without running pixel preprocessing."""
+    if not getattr(image_processor, "do_image_splitting", False):
+        return 1
+
+    is_too_large = getattr(image_processor, "_is_image_too_large", None)
+    get_grid_layout = getattr(image_processor, "_get_grid_layout", None)
+    if is_too_large is None or get_grid_layout is None:
+        return 1
+
+    kwargs = {
+        "max_image_tokens": int(getattr(image_processor, "max_image_tokens", 256)),
+        "encoder_patch_size": int(getattr(image_processor, "encoder_patch_size", 16)),
+        "downsample_factor": int(getattr(image_processor, "downsample_factor", 2)),
+        "max_pixels_tolerance": float(
+            getattr(image_processor, "max_pixels_tolerance", 2.0)
+        ),
+    }
+    if not is_too_large(height=height, width=width, **kwargs):
+        return 1
+
+    _, _, _, _, tiles = get_grid_layout(
+        height=height,
+        width=width,
+        min_tiles=int(getattr(image_processor, "min_tiles", 2)),
+        max_tiles=int(getattr(image_processor, "max_tiles", 10)),
+        tile_size=int(getattr(image_processor, "tile_size", 512)),
+    )
+    if getattr(image_processor, "use_thumbnail", True) and tiles > 1:
+        tiles += 1
+    return int(tiles)
+
+
+def estimate_vlm_tile_count(row: dict, processor) -> int:
+    """Estimate processed visual tiles for one normalized VLM row."""
+    image_processor = getattr(processor, "image_processor", None)
+    if image_processor is None:
+        return 0
+
+    count = 0
+    for source in _row_image_sources(row):
+        try:
+            width, height = get_image_size(source)
+            count += _tile_count(image_processor, height, width)
+        except Exception:
+            logger.debug(
+                "Could not estimate VLM tile count for %s", source, exc_info=True
+            )
+            count += 1
+    return count
+
+
+def add_vlm_tile_counts(dataset: Dataset | None, processor) -> Dataset | None:
+    """Add local, non-training metadata used by the VLM batch sampler."""
+    if dataset is None or VLM_TILE_COUNT_COLUMN in dataset.column_names:
+        return dataset
+    counts = [estimate_vlm_tile_count(row, processor) for row in dataset]
+    return dataset.add_column(VLM_TILE_COUNT_COLUMN, counts)

src/leap_finetune/distribution/local_trainer.py

@@ -12,7 +12,9 @@
 )
 from leap_finetune.training import TRAINING_LOOPS
 
-_LOCAL_TYPES = frozenset({"sft", "dpo", "vlm_sft", "vlm_dpo"})
+_LOCAL_TYPES = frozenset(
+    {"sft", "dpo", "vlm_sft", "vlm_dpo", "grpo", "vlm_grpo", "moe_sft", "moe_dpo"}
+)
 
 
 def should_use_local(job_config: dict) -> bool:
@@ -21,11 +23,15 @@ def should_use_local(job_config: dict) -> bool:
         return False
     training_type = job_config["training_type"]
     is_moe = is_moe_model_from_name(job_config["model_name"])
-    if training_type not in _LOCAL_TYPES and training_type not in {
-        "moe_sft",
-        "moe_dpo",
-    }:
+    if training_type not in _LOCAL_TYPES:
         return False
+    if training_type in {"grpo", "vlm_grpo"}:
+        train_config = job_config.get("training_config") or {}
+        if train_config.get("vllm_mode", "colocate") != "colocate":
+            return False
+        rollout_config = job_config.get("grpo_rollout") or {}
+        if int(rollout_config.get("tensor_parallel_size", 1) or 1) != 1:
+            return False
     if training_type.startswith("moe_") and not is_moe:
         return False
     if is_moe:
@@ -90,6 +96,7 @@ def local_trainer(job_config: dict):
         "model_config": job_config.get("model_config"),
         "benchmark_configs": job_config.get("benchmark_configs"),
         "rewards": job_config.get("rewards"),
+        "rl_env": job_config.get("rl_env"),
         "async_eval": job_config.get("async_eval"),
         "config_dir": job_config.get("config_dir"),
     }