perf(trainer): read the checkpoint into the page cache before loading it

Makefile

@@ -19,7 +19,7 @@ HOST           ?= 127.0.0.1
 PORT           ?= 9003
 # The fully qualified base URL used by local CLI tools and clients
 BASE_URL       ?= http://$(HOST):$(PORT)
-UNIT_TESTS ?= tests.test_session_lifecycle tests.test_proto_codec tests.test_forward_only tests.test_fft_batch_failure tests.test_api_server_paths tests.test_accel_timeslicer tests.test_trainer_optimizer_correctness tests.test_trainer_logprob_chunks tests.test_worker_manager tests.test_scheduler_worker_manager tests.test_estimator tests.test_redis_store tests.test_cluster_eval_script tests.test_perf_compare_script tests.test_delta_weight_sync tests.test_diffing_backends tests.test_sampler_weight_rotation tests.test_sampler_batch_limits tests.test_commands tests.test_weight_sync_config tests.test_exclusive_trainer tests.test_delete_model tests.test_automodel_worker tests.test_worker_thread_device tests.test_trainer_data_parallel tests.test_openrl_status tests.test_skill_docs tests.test_trainer_ready tests.test_sampling_store tests.test_lora_snapshots tests.test_pool_accept tests.test_sampling_dispatcher tests.test_vllm_sampler tests.test_vllm_options tests.test_dashboard
+UNIT_TESTS ?= tests.test_session_lifecycle tests.test_proto_codec tests.test_forward_only tests.test_fft_batch_failure tests.test_api_server_paths tests.test_accel_timeslicer tests.test_trainer_optimizer_correctness tests.test_trainer_logprob_chunks tests.test_worker_manager tests.test_scheduler_worker_manager tests.test_estimator tests.test_redis_store tests.test_cluster_eval_script tests.test_perf_compare_script tests.test_delta_weight_sync tests.test_diffing_backends tests.test_sampler_weight_rotation tests.test_sampler_batch_limits tests.test_commands tests.test_weight_sync_config tests.test_exclusive_trainer tests.test_delete_model tests.test_automodel_worker tests.test_worker_thread_device tests.test_trainer_data_parallel tests.test_openrl_status tests.test_skill_docs tests.test_trainer_ready tests.test_sampling_store tests.test_lora_snapshots tests.test_prefetch tests.test_pool_accept tests.test_sampling_dispatcher tests.test_vllm_sampler tests.test_vllm_options tests.test_dashboard
 # Only forward BASE_URL to e2e when the user supplied it. The Makefile default
 # is for local CLI usage; e2e should start its own backend by default.
 TRAINING_TEST_BASE_URL ?= $(if $(filter environment command line,$(origin BASE_URL)),$(BASE_URL),)

src/training/fft_trainer_worker.py

@@ -15,6 +15,7 @@
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedModel
 
+from training.prefetch import prefetch
 from training.trainer_worker import BaseTrainerWorker, Datum
 from training.types import FFTConfig
 
@@ -86,6 +87,7 @@ def load_base_model(self, base_model_name: str) -> None:
     self.tokenizer = AutoTokenizer.from_pretrained(base_model_name)
     dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float32
 
+    prefetch(base_model_name)
     self.model = AutoModelForCausalLM.from_pretrained(base_model_name, dtype=dtype, device_map=target_device)
     print("Successfully loaded full fine-tuning model.")
 
@@ -315,6 +317,7 @@ def load_from_state(self, model_id: str, state_path: str, restore_optimizer: boo
     dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float32
     num_gpus = torch.cuda.device_count() if torch.cuda.is_available() else 0
     target_device = "auto" if num_gpus > 1 else self.device
+    prefetch(state_path)
     self.model = AutoModelForCausalLM.from_pretrained(state_path, dtype=dtype, device_map=target_device)
     self.prepare_model_for_training()
     # No sampler save holds the loaded weights.

src/training/lora_trainer_worker.py

@@ -13,6 +13,7 @@
 from peft import PeftModelForCausalLM, get_peft_model
 from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedModel
 
+from training.prefetch import prefetch
 from training.trainer_worker import BaseTrainerWorker, Datum
 from training.types import LoraConfig
 
@@ -50,6 +51,7 @@ def load_base_model(self, base_model_name: str) -> None:
     self.tokenizer = AutoTokenizer.from_pretrained(base_model_name)
     dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float32
 
+    prefetch(base_model_name)
     self.base_model = AutoModelForCausalLM.from_pretrained(base_model_name, dtype=dtype, device_map=self.device)
     print("Successfully loaded.")
 

src/training/prefetch.py

@@ -0,0 +1,121 @@
+"""Read a checkpoint's safetensors into the page cache before from_pretrained maps them.
+
+from_pretrained memory-maps the files and faults them in a few pages at a time,
+each read paying a round trip on a network filesystem. This uses vLLM's loader
+prefetch policy (network filesystems only, 8 threads, 16 MiB reads, skipped when
+the files exceed 90% of available memory) and splits files into chunks, so a
+single-file checkpoint is read by every thread.
+"""
+
+import os
+import time
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+
+from huggingface_hub import snapshot_download
+from huggingface_hub.errors import HFValidationError, LocalEntryNotFoundError
+
+NETWORK_FS = {"nfs", "nfs4", "lustre"}
+THREADS = 8
+BLOCK_BYTES = 16 * 1024 * 1024
+# One task per chunk, so a single large file is read by every thread.
+CHUNK_BYTES = 256 * 1024 * 1024
+MEMINFO = Path("/proc/meminfo")
+CGROUP = Path("/sys/fs/cgroup")
+
+
+def checkpoint_files(name_or_path: str) -> list[Path]:
+  """The safetensors of a checkpoint directory, or of a model id's Hugging Face
+  cache snapshot. A model that is not cached yet is left to the loader."""
+  directory = Path(name_or_path)
+  if not directory.is_dir():
+    try:
+      directory = Path(snapshot_download(name_or_path, local_files_only=True, allow_patterns=["*.safetensors"]))
+    except (LocalEntryNotFoundError, HFValidationError):
+      return []
+  # is_file follows the cache's symlinks into blobs/.
+  return [p for p in directory.glob("*.safetensors") if p.is_file()]
+
+
+def fs_type(path: Path) -> str:
+  """The filesystem type of the mount that holds path, from /proc/mounts."""
+  best, best_type = "", ""
+  try:
+    resolved = str(path.resolve())
+    with open("/proc/mounts") as f:
+      for line in f:
+        parts = line.split()
+        if len(parts) < 3:
+          continue
+        mount = parts[1]
+        if (resolved == mount or resolved.startswith(mount.rstrip("/") + "/")) and len(mount) >= len(best):
+          best, best_type = mount, parts[2]
+  except OSError:
+    return ""
+  return best_type
+
+
+def available_bytes() -> int | None:
+  """MemAvailable, capped by the room left under the cgroup memory limit, or None
+  when unknown. Cgroup usage is the working set (current minus inactive_file), as
+  kubelet counts it, because the pod's own page cache can be reclaimed."""
+  try:
+    avail = next(int(line.split()[1]) * 1024 for line in MEMINFO.read_text().splitlines() if line.startswith("MemAvailable:"))
+  except (OSError, StopIteration):
+    return None
+  try:
+    limit = (CGROUP / "memory.max").read_text().strip()
+    if limit != "max":
+      stat = dict(line.split() for line in (CGROUP / "memory.stat").read_text().splitlines())
+      working_set = int((CGROUP / "memory.current").read_text()) - int(stat["inactive_file"])
+      avail = min(avail, max(0, int(limit) - working_set))
+  except OSError:
+    pass
+  return avail
+
+
+def read_range(path: Path, offset: int, length: int) -> int:
+  """Read [offset, offset+length) with plain reads; the kernel keeps the pages."""
+  done = 0
+  fd = os.open(path, os.O_RDONLY)
+  try:
+    while done < length:
+      chunk = os.pread(fd, min(BLOCK_BYTES, length - done), offset + done)
+      if not chunk:
+        break
+      done += len(chunk)
+  finally:
+    os.close(fd)
+  return done
+
+
+def prefetch(name_or_path: str) -> int:
+  """Warm the page cache for name_or_path's safetensors and return the bytes read.
+  Errors are logged, not raised: without the prefetch the load is only slower."""
+  try:
+    return _prefetch(name_or_path)
+  except Exception as e:
+    print(f"[prefetch] failed, loading without it: {e!r}")
+    return 0
+
+
+def _prefetch(name_or_path: str) -> int:
+  sizes = {path: path.stat().st_size for path in checkpoint_files(name_or_path)}
+  if not sizes:
+    return 0
+  kind = fs_type(next(iter(sizes)))
+  if kind not in NETWORK_FS:
+    return 0
+  total = sum(sizes.values())
+  avail = available_bytes()
+  if avail is not None and total > 0.9 * avail:
+    print(f"[prefetch] {total / 2**30:.1f} GiB of checkpoint files would not fit in {avail / 2**30:.1f} GiB of available memory; not prefetching")
+    return 0
+  tasks = [(path, offset, min(CHUNK_BYTES, size - offset)) for path, size in sizes.items() for offset in range(0, size, CHUNK_BYTES)]
+  print(f"[prefetch] reading {len(sizes)} file(s), {total / 2**30:.2f} GiB, from {kind} into the page cache with {THREADS} threads")
+  started = time.monotonic()
+  with ThreadPoolExecutor(max_workers=THREADS) as pool:
+    read = sum(pool.map(read_range, *zip(*tasks)))
+  seconds = time.monotonic() - started
+  print(f"[prefetch] done: {read / 2**30:.2f} GiB in {seconds:.1f} s ({read / 2**20 / max(seconds, 1e-6):.0f} MiB/s)")
+  return read

tests/test_prefetch.py

@@ -0,0 +1,64 @@
+"""training.prefetch warms the page cache for a checkpoint before the loader
+maps it. These tests cover its decisions."""
+
+import io
+import os
+import tempfile
+import unittest
+from contextlib import redirect_stdout
+from pathlib import Path
+from unittest.mock import patch
+
+from training import prefetch
+
+TOTAL = 4 * 1024 * 1024 + 17
+
+
+class PrefetchTest(unittest.TestCase):
+  def setUp(self):
+    self.enterContext(redirect_stdout(io.StringIO()))
+    self.dir = Path(self.enterContext(tempfile.TemporaryDirectory()))
+    (self.dir / "model-00001-of-00002.safetensors").write_bytes(os.urandom(3 * 1024 * 1024))
+    (self.dir / "model-00002-of-00002.safetensors").write_bytes(os.urandom(1024 * 1024 + 17))
+    (self.dir / "pytorch_model.bin").write_bytes(b"unused")
+    (self.dir / "config.json").write_text("{}")
+
+  def test_lists_only_safetensors(self):
+    files = prefetch.checkpoint_files(str(self.dir))
+    self.assertEqual(sorted(f.name for f in files), ["model-00001-of-00002.safetensors", "model-00002-of-00002.safetensors"])
+
+  def test_an_uncached_model_id_is_left_to_the_loader(self):
+    self.assertEqual(prefetch.checkpoint_files("org/not-downloaded"), [])
+
+  def test_skips_local_filesystems(self):
+    with patch.object(prefetch, "fs_type", return_value="ext4"):
+      self.assertEqual(prefetch.prefetch(str(self.dir)), 0)
+
+  def test_reads_every_chunk_on_a_network_filesystem(self):
+    with patch.object(prefetch, "fs_type", return_value="nfs"), patch.object(prefetch, "CHUNK_BYTES", 1024 * 1024):
+      self.assertEqual(prefetch.prefetch(str(self.dir)), TOTAL)
+
+  def test_skips_when_the_files_would_not_fit_in_memory(self):
+    with patch.object(prefetch, "fs_type", return_value="nfs"), patch.object(prefetch, "available_bytes", return_value=1024 * 1024):
+      self.assertEqual(prefetch.prefetch(str(self.dir)), 0)
+
+  def test_a_read_error_does_not_fail_the_load(self):
+    with patch.object(prefetch, "fs_type", return_value="nfs"), patch.object(prefetch, "read_range", side_effect=OSError("Stale file handle")):
+      self.assertEqual(prefetch.prefetch(str(self.dir)), 0)
+
+  def test_available_memory_counts_the_cgroup_working_set(self):
+    (self.dir / "meminfo").write_text(f"MemTotal: 100 kB\nMemAvailable: {64 * 1024 * 1024} kB\n")
+    cgroup = self.dir / "cgroup"
+    cgroup.mkdir()
+    gib = 2**30
+    (cgroup / "memory.max").write_text(f"{16 * gib}\n")
+    (cgroup / "memory.current").write_text(f"{15 * gib}\n")
+    (cgroup / "memory.stat").write_text(f"anon {5 * gib}\ninactive_file {10 * gib}\n")
+    with patch.object(prefetch, "MEMINFO", self.dir / "meminfo"), patch.object(prefetch, "CGROUP", cgroup):
+      self.assertEqual(prefetch.available_bytes(), 11 * gib)
+      (cgroup / "memory.max").write_text("max\n")
+      self.assertEqual(prefetch.available_bytes(), 64 * gib)
+
+
+if __name__ == "__main__":
+  unittest.main()