diff --git a/Makefile b/Makefile index 5d6ec025..6ab6e755 100644 --- a/Makefile +++ b/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),) diff --git a/src/training/fft_trainer_worker.py b/src/training/fft_trainer_worker.py index 6e69e150..c8562026 100644 --- a/src/training/fft_trainer_worker.py +++ b/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. diff --git a/src/training/lora_trainer_worker.py b/src/training/lora_trainer_worker.py index 83561e5e..a51b0905 100644 --- a/src/training/lora_trainer_worker.py +++ b/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.") diff --git a/src/training/prefetch.py b/src/training/prefetch.py new file mode 100644 index 00000000..ec3e49f0 --- /dev/null +++ b/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 diff --git a/tests/test_prefetch.py b/tests/test_prefetch.py new file mode 100644 index 00000000..b23ec00a --- /dev/null +++ b/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()