8b fixes (#12)

* transformers bump + fa2

* patched data_loader

* add patched data loaders to runners

* load models w/ fa2 properly

* add train / eval dataset handling

* fmt

* gitignore update for dataset cache

* Add load_tokenizer helper and SFT_EXCLUDED_KEYS

load_tokenizer loads only the tokenizer without model weights, for
use during pre-tokenization on the driver. SFT_EXCLUDED_KEYS lists
config keys that must be filtered when building TrainingArguments
for plain Trainer (replacing SFTConfig).

* Add pre-tokenization functions for SFT and DPO

tokenize_sft / tokenize_and_pack_sft: distributed tokenization via
Ray .map(), with optional BFD packing via TRL's pack_dataset.

tokenize_dpo / tokenize_dpo_dataset: replicates DPOTrainer's
tokenization pipeline (extract prompt, apply chat template, tokenize,
append eos, truncate) as a Ray .map() step.

* Rewrite SFT to use plain Trainer with pre-tokenized data

Replace SFTTrainer/SFTConfig with Trainer/TrainingArguments since
tokenization now happens centrally on the driver. Use TRL's
DataCollatorForLanguageModeling for labels, padding, and padding-free
mode. Get eval shard from Ray instead of materializing full eval on
every worker. Fix import paths (configs -> training_configs) and
extract run_name_template before filtering config keys.

* Pre-tokenized DPO with PreTokenizedDPOTrainer subclass

* Centralize tokenization on driver via Ray .map()

* Add tokenization caching with parquet persistence

* Simplify checkpoint callback to skip Ray checkpoint duplication

* Fix checkpoint callback to call _rename_checkpoint on save

* Use available memory instead of total for Ray object store

* Make tokenization caching opt-in via cache_dataset flag

* check if cuda available before committing time to ray workflow

* consistent naming for cached dataset

* _resolve_model_id already has this logic

* Make tokenization caching opt-in via cache_dataset flag

* Clear metrics after checkpoint report to avoid re-reporting stale values

* Avoid full dataset materialization for non-packing SFT runs

---------

Co-authored-by: EdoardoMosca <edoardo@liquid.ai>

Author: alay2shah

.gitignore

@@ -23,3 +23,6 @@ scripts/
 
 # Quick testing files
 *.ipynb
+
+# Tokenization cache
+.cache/

pyproject.toml

@@ -12,7 +12,9 @@ dependencies = [
     "peft>=0.15.2",
     "deepspeed>=0.17.1",
     "torch>=2.7.1",
-    "transformers @ git+https://github.com/huggingface/transformers.git@0c9a72e4576fe4c84077f066e585129c97bfd4e6",
+    "transformers>=5.0.0",
+    "flash-attn>=2.8.0",
+    "numpy>=2.4.1",
     "trl>=0.18.2",
     "rich>=14.1.0",
     "pillow>=11.3.0",
@@ -31,8 +33,8 @@ leap-finetune = "leap_finetune:main"
 requires = ["hatchling"]
 build-backend = "hatchling.build"
 
-[tool.hatch.metadata]
-allow-direct-references = true
+[tool.uv]
+no-build-isolation-package = ["flash-attn"]
 
 [dependency-groups]
 dev = [

src/leap_finetune/data_loaders/dataset_loader.py

@@ -21,6 +21,7 @@ class DatasetLoader:
     split: str = "train"
     test_size: float = 0.2
     subset: str | None = None
+    cache_dataset: bool = False
     # Optional preprocessing function: takes Ray Dataset, returns Ray Dataset
     # Applied before validation - use for custom filtering, transforms, joins, etc.
     preprocess_fn: Callable | None = field(default=None, repr=False)

src/leap_finetune/data_loaders/ray_data_utils.py

@@ -1,43 +1,151 @@
+import hashlib
+import json
+import logging
+
 import ray.data
 from datasets import Dataset
 from rich.console import Console
 
+from leap_finetune.utils.constants import TOKENIZATION_CACHE_DIR
+
 from .dataset_loader import DatasetLoader
+from .tokenize_data import tokenize_and_pack_sft, tokenize_dpo_dataset
 from .validate_loader import get_row_filter, normalize_columns
 
+logger = logging.getLogger(__name__)
+
+
+# === Tokenization Cache ===
+
+
+def _build_cache_key(
+    loader: DatasetLoader,
+    shuffle_seed: int,
+    tokenizer_id: str,
+    training_config: dict,
+) -> tuple[str, dict]:
+    """Build a deterministic cache key from all parameters affecting tokenized output."""
+    dataset_type = loader.dataset_type
+
+    key = {
+        "dataset_path": loader.dataset_path,
+        "subset": loader.subset,
+        "split": loader.split,
+        "limit": loader.limit,
+        "test_size": loader.test_size,
+        "dataset_type": dataset_type,
+        "tokenizer": tokenizer_id,
+        "shuffle_seed": shuffle_seed,
+    }
+
+    if dataset_type == "sft":
+        key["max_length"] = training_config.get("max_length", 2048)
+        key["packing"] = training_config.get("packing", False)
+    elif dataset_type == "dpo":
+        key["max_prompt_length"] = training_config.get("max_prompt_length")
+        key["max_completion_length"] = training_config.get("max_completion_length")
+
+    canonical = json.dumps(key, sort_keys=True)
+    fingerprint = hashlib.sha256(canonical.encode()).hexdigest()[:16]
+    return fingerprint, key
+
+
+def _try_load_cache(
+    fingerprint: str,
+) -> tuple[ray.data.Dataset, ray.data.Dataset] | None:
+    """Load cached train/eval parquet if the cache directory exists."""
+    cache_dir = TOKENIZATION_CACHE_DIR / fingerprint
+    train_dir = cache_dir / "train"
+    eval_dir = cache_dir / "eval"
+
+    if not train_dir.exists() or not eval_dir.exists():
+        return None
+
+    try:
+        train_ds = ray.data.read_parquet(str(train_dir))
+        eval_ds = ray.data.read_parquet(str(eval_dir))
+        return train_ds, eval_ds
+    except Exception:
+        logger.warning("Failed to read tokenization cache, will re-tokenize")
+        return None
+
+
+def _save_cache(
+    fingerprint: str,
+    train_ds: ray.data.Dataset,
+    eval_ds: ray.data.Dataset,
+    key_dict: dict,
+) -> None:
+    """Write train/eval datasets as parquet into the cache directory."""
+    cache_dir = TOKENIZATION_CACHE_DIR / fingerprint
+    cache_dir.mkdir(parents=True, exist_ok=True)
+
+    train_ds.write_parquet(str(cache_dir / "train"))
+    eval_ds.write_parquet(str(cache_dir / "eval"))
+
+    (cache_dir / "fingerprint.json").write_text(json.dumps(key_dict, indent=2))
+
 
 def create_ray_datasets(
     loader: DatasetLoader,
     shuffle_seed: int = 42,
+    tokenizer=None,
+    training_config: dict | None = None,
 ) -> tuple[ray.data.Dataset, ray.data.Dataset]:
     """
     Create validated, shuffled, split Ray Datasets from a DatasetLoader.
 
-    Pipeline: quick_validate → load → [preprocess] → filter → normalize → shuffle → split
+    Pipeline: quick_validate → load → [preprocess] → filter → normalize → shuffle → split → [tokenize/pack]
 
-    Uses Ray Data native operations (filter/map) - no pandas/arrow imports needed.
+    When tokenizer is provided, tokenization and optional packing happen
+    centrally before sharding, producing equal-length shards (±1 row).
+    Tokenized results are cached as Parquet for subsequent runs.
     """
     console = Console()
 
+    # === Check tokenization cache ===
+    use_pretokenize = tokenizer is not None and training_config is not None
+    can_cache = (
+        use_pretokenize and loader.cache_dataset and loader.preprocess_fn is None
+    )
+    fingerprint = None
+    key_dict = None
+
+    if can_cache:
+        fingerprint, key_dict = _build_cache_key(
+            loader, shuffle_seed, tokenizer.name_or_path, training_config
+        )
+        cached = _try_load_cache(fingerprint)
+        if cached is not None:
+            train_ds, eval_ds = cached
+            train_count = train_ds.count()
+            eval_count = eval_ds.count()
+            console.print(
+                f"[green]✓ Cache hit[/green] [dim]({fingerprint})[/dim]: "
+                f"{train_count + eval_count:,} samples "
+                f"(train: {train_count:,}, eval: {eval_count:,})"
+            )
+            return train_ds, eval_ds
+        logger.info(
+            "Tokenization cache miss (%s), will tokenize and cache", fingerprint
+        )
+
+    # === Full pipeline: load → filter → normalize → shuffle → split → tokenize ===
+
+    loader.quick_validate()
     ds = loader.to_ray_dataset()
 
-    # Apply user preprocessing if provided (before validation)
     if loader.preprocess_fn is not None:
         console.print("[dim]Applying preprocessing...[/dim]")
         ds = loader.preprocess_fn(ds)
 
-    # Filter invalid rows using Ray's native filter (pure Python, Ray handles Arrow)
     row_filter = get_row_filter(loader.dataset_type)
     ds = ds.filter(row_filter)
 
-    # Normalize column names
     normalizer = normalize_columns(loader.dataset_type)
     ds = ds.map(normalizer)
 
-    # Shuffle before split
     ds = ds.random_shuffle(seed=shuffle_seed)
-
-    # Materialize to get count
     total_count = ds.count()
 
     if total_count == 0:
@@ -46,7 +154,6 @@ def create_ray_datasets(
             f"the expected format for dataset_type='{loader.dataset_type}'"
         )
 
-    # Calculate split sizes
     eval_count = max(1, int(total_count * loader.test_size))
     train_count = total_count - eval_count
 
@@ -56,14 +163,51 @@ def create_ray_datasets(
             f"with test_size={loader.test_size}"
         )
 
-    # Split into train/eval
     train_ds, eval_ds = ds.split_at_indices([train_count])
 
     console.print(
         f"[green]✓ Dataset ready:[/green] {total_count:,} samples "
         f"(train: {train_count:,}, eval: {eval_count:,})"
     )
 
+    # === Pre-tokenize if tokenizer provided ===
+    if use_pretokenize:
+        dataset_type = loader.dataset_type
+
+        if dataset_type == "sft":
+            max_length = training_config.get("max_length", 2048)
+            packing = training_config.get("packing", False)
+
+            console.print(
+                f"[dim]Tokenizing SFT (max_length={max_length}, packing={packing})...[/dim]"
+            )
+            train_ds = tokenize_and_pack_sft(train_ds, tokenizer, max_length, packing)
+            eval_ds = tokenize_and_pack_sft(
+                eval_ds, tokenizer, max_length, packing=False
+            )
+
+        elif dataset_type == "dpo":
+            max_prompt_length = training_config.get("max_prompt_length")
+            max_completion_length = training_config.get("max_completion_length")
+
+            console.print("[dim]Tokenizing DPO...[/dim]")
+            train_ds = tokenize_dpo_dataset(
+                train_ds, tokenizer, max_prompt_length, max_completion_length
+            )
+            eval_ds = tokenize_dpo_dataset(
+                eval_ds, tokenizer, max_prompt_length, max_completion_length
+            )
+
+        # === Save to cache ===
+        if can_cache and fingerprint is not None:
+            try:
+                _save_cache(fingerprint, train_ds, eval_ds, key_dict)
+                console.print(f"[dim]Cached tokenized data ({fingerprint})[/dim]")
+            except Exception:
+                logger.warning(
+                    "Failed to write tokenization cache, continuing without cache"
+                )
+
     return train_ds, eval_ds