Merge branch 'main' into feature/default-activation

Author: Ryfernandes

docs/cli/data_generation_offline.md

@@ -51,6 +51,12 @@ python scripts/data_generation_offline.py \
 
 - **`--max-consecutive-errors`** (int, default: value of `--concurrency`) Abort after this many consecutive sample failures (each sample already retried `--max-retries` times). Prevents silently churning through the entire dataset when the server is down. Ignored when `--fail-on-error` is set.
 
+### Multi-Node Arguments
+
+- **`--world-size`** (int, default: `1`) Number of nodes participating in data generation. Each node is assigned a contiguous, non-overlapping chunk of the dataset. This is the number of nodes, not the number of GPUs.
+
+- **`--rank`** (int, default: `0`) Zero-based index of the current node. Must be in the range `[0, world-size)`.
+
 ## Full Example
 
 ```bash

docs/user_guide/tutorials/train_eagle3_offline.md

@@ -154,6 +154,22 @@ output/hidden_states/
 python scripts/launch_vllm.py model -- --data-parallel-size 8
 ```
 
+**Use multiple nodes:**
+
+If you have access to multiple machines, each with its own vLLM server, you can split the dataset across them with `--world-size` and `--rank`. Each node generates a contiguous, non-overlapping chunk of the data into a shared (or later merged) output directory. See the [data_generation_offline.py cli reference](/cli/data_generation_offline.md) for details.
+
+```bash
+# On node 0
+python scripts/data_generation_offline.py \
+  --preprocessed-data ./output --output ./output/hidden_states \
+  --max-samples 5000 --world-size 2 --rank 0
+
+# On node 1
+python scripts/data_generation_offline.py \
+  --preprocessed-data ./output --output ./output/hidden_states \
+  --max-samples 5000 --world-size 2 --rank 1
+```
+
 **Skip validation:**
 
 Validation loads every single generated output file and confirms that the token ids match the sent request and the hidden states length matches expectation. This is generally not required but is a good sanity check. Turn this off to skip loading generated samples during data gen.

pyproject.toml

@@ -53,7 +53,7 @@ dependencies = [
     "torchaudio",
     "torchvision",
     "tqdm>=4.66.3,<=4.67.3",
-    "transformers>=4.56.1,<5.9.0",
+    "transformers>=4.56.1,<5.10.0",
     "typer>=0.12.0",
 ]
 

scripts/data_generation_offline.py

@@ -26,10 +26,15 @@
 from safetensors import safe_open
 from tqdm import tqdm
 
+from speculators.data_generation.offline import (
+    get_existing_hidden_state_indices,
+    get_indices_to_process,
+)
 from speculators.data_generation.vllm_client import (
     DEFAULT_MAX_RETRIES,
     DEFAULT_REQUEST_TIMEOUT,
     generate_hidden_states_async,
+    wait_for_lock_async,
 )
 from speculators.train.data import build_client_item
 from speculators.train.logger import setup_root_logger
@@ -164,72 +169,28 @@ def parse_args():
             "(default: value of --concurrency)"
         ),
     )
+    parser.add_argument(
+        "--world-size",
+        type=int,
+        default=1,
+        help=(
+            "World size for multi-node data generation offline. IMPORTANT: this "
+            "is the number of nodes (not the number of gpus). Defaults to 1"
+        ),
+    )
+    parser.add_argument(
+        "--rank",
+        type=int,
+        default=0,
+        help=(
+            "Rank for multi-node data generation offline. IMPORTANT: this is "
+            "the node index, not an index for a gpu. Must be in range[0, world_size)."
+            " Defaults to 0"
+        ),
+    )
     return parser.parse_args()
 
 
-def get_existing_hidden_state_indices(output_path: Path) -> list[int]:
-    """Find existing `hs_i.safetensors` files (where i is the file index)"""
-
-    existing_file_indices = []
-
-    if not output_path.exists():
-        return existing_file_indices
-
-    for file_path in output_path.iterdir():
-        if file_path.name.startswith("hs_") and file_path.name.endswith(".safetensors"):
-            index_str = file_path.stem[3:]  # Remove "hs_" prefix
-            try:
-                file_index = int(index_str)
-                existing_file_indices.append(file_index)
-            except ValueError:
-                continue
-
-    return sorted(existing_file_indices)
-
-
-def get_indices_to_process(
-    num_samples: int, max_samples: int | None, existing: list[int]
-) -> list[int]:
-    """Determines which indices should be processed. If max_samples is None
-    returns all dataset indices not in existing. Otherwise gets the first
-    `max_samples - len(existing)` samples not already in existing.
-
-    Args:
-        num_samples: Total size of preprocessed dataset
-        max_samples: (Optional) limit for number of samples to process
-        existing: list of ids that have already been processed
-
-    Returns:
-        list of dataset indices to process
-    """
-
-    if len(existing) >= num_samples:
-        logger.info("All samples already processed!")
-        return []
-    if max_samples and len(existing) >= max_samples:
-        logger.info("At least max_samples already processed!")
-        return []
-
-    if len(existing) > 0:
-        logger.info(f"Found {len(existing)} existing samples.")
-
-    existing_s = set(existing)
-    if max_samples is None:
-        return [i for i in range(num_samples) if i not in existing_s]
-
-    num_remaining = min(max_samples, num_samples) - len(existing)
-    to_process = []
-    cur = 0
-    while num_remaining > 0 and cur < num_samples:
-        if cur not in existing_s:
-            to_process.append(cur)
-            num_remaining -= 1
-
-        cur += 1
-
-    return to_process
-
-
 def check_safetensors_file(path: Path, tokens: list[int]):
     with safe_open(path, "pt") as f:
         t_ids = f.get_tensor("token_ids").tolist()
@@ -247,7 +208,7 @@ def check_safetensors_file(path: Path, tokens: list[int]):
             )
 
 
-async def worker(
+async def worker(  # noqa: C901
     client,
     model: str,
     queue: "asyncio.Queue[dict[str, Any]]",
@@ -288,6 +249,10 @@ async def worker(
                     timeout=request_timeout,
                     max_retries=max_retries,
                 )
+            lock_path = hidden_states_path + ".lock"
+            if Path(lock_path).exists():  # noqa: ASYNC240
+                await wait_for_lock_async(lock_path)
+
             async with write_semaphore:  # Limit number of active disk writes
                 await asyncio.to_thread(
                     shutil.move, hidden_states_path, target_hidden_states_path
@@ -374,7 +339,11 @@ async def generate_and_save_hidden_states(args, dataset):
     num_samples = len(dataset)
 
     to_process = get_indices_to_process(
-        num_samples, args.max_samples, existing_file_indices
+        num_samples,
+        args.max_samples,
+        existing_file_indices,
+        args.world_size,
+        args.rank,
     )
     if not to_process:
         return
@@ -441,6 +410,8 @@ async def generate_and_save_hidden_states(args, dataset):
 
 def main():
     args = parse_args()
+    if int(args.rank) < 0 or int(args.rank) >= int(args.world_size):
+        raise ValueError("--rank must be in range [0, world_size)")
     setup_root_logger()
 
     logger.info("EAGLE Offline Data Generation")

src/speculators/data_generation/offline.py

@@ -0,0 +1,74 @@
+import logging
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+
+def get_existing_hidden_state_indices(output_path: Path) -> list[int]:
+    """Find existing `hs_i.safetensors` files (where i is the file index)"""
+
+    existing_file_indices_set: set[int] = set()
+
+    if not output_path.exists():
+        return []
+
+    for file_path in output_path.iterdir():
+        if file_path.name.startswith("hs_") and file_path.name.endswith(".safetensors"):
+            index_str = file_path.stem[3:]  # Remove "hs_" prefix
+            try:
+                file_index = int(index_str)
+                existing_file_indices_set.add(file_index)
+            except ValueError:
+                continue
+
+    return sorted(existing_file_indices_set)
+
+
+def get_indices_to_process(
+    num_samples: int,
+    max_samples: int | None,
+    existing: list[int],
+    world_size: int,
+    rank: int,
+) -> list[int]:
+    """Determines which indices should be processed. If max_samples is None
+    returns all dataset indices not in existing. Otherwise gets the first
+    `max_samples - len(existing)` samples not already in existing.
+
+    Args:
+        num_samples: Total size of preprocessed dataset
+        max_samples: (Optional) limit for number of samples to process
+        existing: list of ids that have already been processed
+        world_size: Number of nodes to generate on
+        rank: The rank of the local node
+
+    Returns:
+        list of dataset indices to process
+    """
+
+    target = min(max_samples, num_samples) if max_samples is not None else num_samples
+
+    if target <= 0:
+        return []
+
+    chunk_size = target // world_size
+    remainder = target % world_size
+    # Distribute remainder across the first `remainder` ranks so chunks differ
+    # by at most 1.
+    start = rank * chunk_size + min(rank, remainder)
+    end = start + chunk_size + (1 if rank < remainder else 0)
+
+    existing_s = set(existing)
+    to_process = [i for i in range(start, end) if i not in existing_s]
+
+    if not to_process:
+        logger.info("All samples for this rank already processed!")
+        return []
+
+    if len(existing_s & set(range(start, end))) > 0:
+        logger.info(
+            f"Found {len(existing_s & set(range(start, end)))} existing samples"
+            f" for rank {rank}."
+        )
+
+    return to_process

src/speculators/data_generation/vllm_client.py

@@ -1,6 +1,8 @@
 import asyncio
+import fcntl
 import functools
 import logging
+import os
 import time
 from typing import TYPE_CHECKING, Any, TypedDict
 
@@ -122,6 +124,47 @@ class ClientItem(TypedDict):
     instead of passing `token_ids` to Completions API."""
 
 
+async def _poll_lock_async(fd, poll_interval):
+    while True:
+        try:
+            fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
+            return
+        except BlockingIOError:
+            await asyncio.sleep(poll_interval)
+
+
+async def wait_for_lock_async(lock_path, timeout=10.0, poll_interval=0.1):
+    fd = os.open(lock_path, os.O_RDONLY)
+    try:
+        await asyncio.wait_for(_poll_lock_async(fd, poll_interval), timeout=timeout)
+    except BaseException:
+        os.close(fd)
+        raise
+    os.close(fd)
+    os.remove(lock_path)
+
+
+def wait_for_lock(lock_path, timeout=10.0, poll_interval=0.1):
+    fd = os.open(lock_path, os.O_RDONLY)
+    try:
+        deadline = time.monotonic() + timeout
+        while True:
+            try:
+                fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
+                break
+            except BlockingIOError:
+                if time.monotonic() >= deadline:
+                    raise TimeoutError(
+                        f"Timed out waiting for lock: {lock_path}"
+                    ) from None
+                time.sleep(poll_interval)
+    except BaseException:
+        os.close(fd)
+        raise
+    os.close(fd)
+    os.remove(lock_path)
+
+
 @with_retries
 async def generate_hidden_states_async(
     client: openai.AsyncClient,

src/speculators/train/data.py

@@ -21,6 +21,7 @@
     DEFAULT_REQUEST_TIMEOUT,
     ClientItem,
     generate_hidden_states,
+    wait_for_lock,
 )
 from speculators.train.noise_transforms import TransformTensors
 
@@ -259,6 +260,11 @@ def _compute_approx_lengths(self) -> list[int]:
     def _maybe_load_hs_file(self, index: int) -> dict[str, torch.Tensor] | None:
         file_idx = self._map_to_file_idx(index)
         candidate_path = self.hidden_states_path / f"hs_{file_idx}.safetensors"
+
+        lock_path = str(candidate_path) + ".lock"
+        if Path(lock_path).exists():
+            wait_for_lock(lock_path)
+
         if candidate_path.exists():
             return load_file(candidate_path)
 
@@ -279,19 +285,22 @@ def _maybe_generate_hs(self, index: int) -> dict[str, torch.Tensor] | None:
                 timeout=self.request_timeout,
                 max_retries=self.max_retries,
             )
-        except Exception as e:  # noqa: BLE001
-            warnings.warn(str(e), stacklevel=1)
-            return None
 
-        loaded_hs = load_file(hs_filepath)
+            loaded_hs = load_file(hs_filepath)
 
-        match self.on_generate:
-            case "cache":
-                file_idx = self._map_to_file_idx(index)
-                target_path = self.hidden_states_path / f"hs_{file_idx}.safetensors"
-                shutil.move(hs_filepath, target_path)
-            case "delete":
-                Path(hs_filepath).unlink()
+            match self.on_generate:
+                case "cache":
+                    file_idx = self._map_to_file_idx(index)
+                    target_path = self.hidden_states_path / f"hs_{file_idx}.safetensors"
+                    shutil.move(hs_filepath, target_path)
+                case "delete":
+                    Path(hs_filepath).unlink()
+        except Exception as e:  # noqa: BLE001
+            warnings.warn(
+                f"Failed to load/cache hidden states for sample {index}: {e}",
+                stacklevel=1,
+            )
+            return None
 
         return loaded_hs