From 042e8ee5801d1507f481a0660d225dc3eb96b6a2 Mon Sep 17 00:00:00 2001 From: Junda Chen <32371474+GindaChen@users.noreply.github.com> Date: Thu, 15 Oct 2026 01:35:13 -0700 Subject: [PATCH] [feat] Retain native model weights in a CUDA IPC donor --- docs/contributing/startup_profiling.md | 53 +--- docs/contributing/weight_cache.md | 94 ++++++ examples/benchmark/vae_weight_cache.py | 124 ++++++++ examples/benchmark/wan_pipeline_startup.py | 96 ------ fastvideo/models/loader/component_loader.py | 154 +++++----- fastvideo/models/loader/fsdp_load.py | 181 ++++++----- fastvideo/models/loader/weight_cache.py | 286 ++++++++++++++++++ .../loader/test_component_startup_phases.py | 72 ----- fastvideo/tests/loader/test_weight_cache.py | 153 ++++++++++ mkdocs.yml | 1 + 10 files changed, 823 insertions(+), 391 deletions(-) create mode 100644 docs/contributing/weight_cache.md create mode 100644 examples/benchmark/vae_weight_cache.py delete mode 100644 examples/benchmark/wan_pipeline_startup.py create mode 100644 fastvideo/models/loader/weight_cache.py delete mode 100644 fastvideo/tests/loader/test_component_startup_phases.py create mode 100644 fastvideo/tests/loader/test_weight_cache.py diff --git a/docs/contributing/startup_profiling.md b/docs/contributing/startup_profiling.md index 4131dc604..a26387d17 100644 --- a/docs/contributing/startup_profiling.md +++ b/docs/contributing/startup_profiling.md @@ -79,58 +79,9 @@ page faults and transfers can occur during assignment. These labels describe loader operations, not independent physical disk-I/O and PCIe measurements. Checkpoint-backed special constructors retain their existing paths and only -receive the component-total interval. The transformer and shared text-encoder loaders also expose the boundaries -described below. First-use compilation and kernel warm-up still need separate -measurement. Reusing GPU tensor +receive the component-total interval. Transformer/text-encoder internals, +compilation, and kernel warm-up need further instrumentation. Reusing GPU tensor storage across processes will require a separate compatibility manifest, ownership/lifetime protocol, and loader implementation. The extracted `load_weights()` method still copies from disk; a future IPC loader must bind storage instead of calling `load_state_dict()` and claiming zero-copy attachment. - -## Profile a complete Wan pipeline - -The whole-pipeline example loads the DiT, text encoder, tokenizer, and decoder -VAE, then verifies a fixed tiny request: - -```bash -python examples/benchmark/wan_pipeline_startup.py \ - --model-path /path/to/wan-diffusers-checkpoint \ - --output /tmp/wan-startup.json \ - --tensor-output /tmp/wan-output.pt \ - --synchronize -``` - -It selects one GPU and the in-process executor so the collector sees all loads. -Add `--fsdp` to exercise single-rank FSDP inference; the default is unsharded. -The JSON records resolved offload/FSDP settings and the live transformer -FSDP-module status; assertions require both to match the request. -Both paths disable CPU offload and compilation and use the checkpoint's native -Wan 2.1 1.3B pipeline defaults. The request is 9 frames at 128x128, four steps, -seed 1024, with a fixed prompt. It checks tensor shape and finiteness and can -save the output for an independent original/refactored comparison. It is not a -video-quality or production-throughput benchmark. No video file is written. - -`pipeline.create` includes configuration, worker/device setup and component -loading. Its children include the `component.*` totals. `first_request` measures -the complete tiny request, including any first-use warm-up and output handling; -it does not separate JIT, autotuning, text encoding, denoising, and decoding. -The reported CUDA-initialized flag identifies earlier import side effects. - -Native transformer phases are `transformer.construct` (meta construction), -`transformer.shard_setup`, `transformer.hydrate`, `transformer.finalize`, and -`transformer.compile_setup`. The last two cover post-load conversion/buffers and -compile registration respectively, not the first execution of compiled code. -The sharding and compile-setup intervals may be near zero when disabled. - -Native shared encoder-loader phases are `text_encoder.construct`, -`text_encoder.hydrate`, `text_encoder.device_transfer`, and -`text_encoder.offload_setup`. The shared path also serves image encoders; use the -parent `component.*` name to identify the role. HF passthrough encoders report -`text_encoder.pretrained` as one combined construction/loading operation. - -Hydration intervals contain consumption of the lazy checkpoint iterators, -conversion/fusion, and weight assignment. They deliberately do not call iterator -creation a disk-read measurement. In particular, a CUDA IPC loader would not -necessarily save the entire hydration interval. Keep donor startup, attachment, -compatibility validation, and an independent disk-loaded confirmation in any -future comparison. diff --git a/docs/contributing/weight_cache.md b/docs/contributing/weight_cache.md new file mode 100644 index 000000000..178254062 --- /dev/null +++ b/docs/contributing/weight_cache.md @@ -0,0 +1,94 @@ +# Retaining native weights across process restarts + +A weight-cache donor owns loaded GPU tensors. Consumers create fresh Python +modules on `meta`, validate compatibility, and map those tensors through CUDA +IPC. Exiting a consumer does not unload the donor. This foundation is independent +of an edit runner and does not pickle or transfer a model object. + +The initial contract is Linux, one GPU, native strided tensors, inference only, +and one active consumer at a time. CPU registered tensors, DTensor/FSDP, +quantized representations and training are not supported. Multiple components +and parallel-rank routing need additional adapters and identity fields. + +## Native Wan VAE example + +Use the same GPU, user and IPC namespace for all processes. The model directory +is a local Wan 2.1 1.3B Diffusers VAE component. Supply its immutable revision. +Create a private socket directory: + +```bash +install -d -m 700 /tmp/fv-weights +python examples/benchmark/vae_weight_cache.py donor \ + --model-path /path/to/checkpoint/vae --checkpoint-id IMMUTABLE_REVISION \ + --socket /tmp/fv-weights/vae.sock --output /tmp/donor.json +``` + +After the socket exists, run consumers from another terminal. Each invocation is +an independent process: + +```bash +python examples/benchmark/vae_weight_cache.py consumer \ + --model-path /path/to/checkpoint/vae --checkpoint-id IMMUTABLE_REVISION \ + --socket /tmp/fv-weights/vae.sock --output /tmp/consumer-1.json \ + --tensor-output /tmp/consumer-1.pt +``` + +Repeat with new output filenames. The same donor retains the tensors. The example +reads configuration in each consumer, constructs fresh VAE caches and module +objects, and attaches registered weights without checkpoint hydration. Its FP32 +codec check decodes identical input twice, checks cache reset and weight +fingerprints, and can export output for an independent `disk`-role comparison. +Use Ctrl-C on the donor after consumers finish; graceful shutdown waits for the +active connection to close and writes a final weight fingerprint. + +## Python boundary + +The donor receives an already loaded, idle inference model: + +```python +serve_weight_cache(model.eval(), key, socket_path, stop=stop_event) +``` + +Construction and attachment happen in the consumer: + +```python +client = WeightCacheClient(socket_path, key, device) +with client.load_model(factory) as attached: + output = attached.model(inputs) +``` + +`factory` constructs the component without reading checkpoint weights. It runs +inside a meta-device context. Initialize required non-registered configuration +constants outside that context. Model/request state belongs to the consumer; +only registered parameters and buffers are shared. Non-persistent buffers and +tied-parameter identities are included. Consumer bindings are restored to meta +on context exit, after GPU synchronization. Do not retain mapped tensors outside +the context. + +The donor must outlive every mapped tensor. Shared tensors have a read-only +inference contract; CUDA does not enforce read-only mappings. In-place weight +changes require private copies or a new donor. Donor crashes invalidate consumers +and require recovery. Imports, CUDA contexts, compilation, graphs and request +caches are not transferred or retained by this API. + +## Compatibility and transport + +`WeightCacheKey` requires checkpoint, component, configuration, dtype and +weight-processing identities. The caller must provide the correct immutable +checkpoint identity and increment the processing version when weight preparation +changes. A method-body edit can reuse the same weights. Both peers also check the +Torch version and physical GPU UUID. The complete registered layout checks class, +names, shapes, dtypes, strides, offsets, aliases and buffer persistence before +any bindings change. A mismatch or missing donor raises; no disk fallback exists. + +The Unix socket and its directory must be private to the current user. Each peer +checks Linux peer credentials before decoding trusted same-user descriptors. +Descriptors contain CUDA IPC handles, not GPU weight bytes. New PyTorch reduction +arguments are created for every consumer to preserve per-mapping refcounts. A +per-socket lock prevents a second donor from replacing an active donor's socket. + +The design follows the separation used by vLLM's +[weight-cache donor](https://github.com/vllm-project/vllm/blob/b965ee564d1897fdaeddaeceecc316aa15a45312/vllm/model_executor/model_loader/weight_cache/daemon.py) +and [IPC loader](https://github.com/vllm-project/vllm/blob/b965ee564d1897fdaeddaeceecc316aa15a45312/vllm/model_executor/model_loader/weight_cache/ipc_loader.py). +vLLM additionally handles post-processing, quantization and parallel shards. +Those capabilities are not implied by this initial native-tensor foundation. diff --git a/examples/benchmark/vae_weight_cache.py b/examples/benchmark/vae_weight_cache.py new file mode 100644 index 000000000..e97407529 --- /dev/null +++ b/examples/benchmark/vae_weight_cache.py @@ -0,0 +1,124 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Persistent native Wan VAE weights with independently restartable IPC consumers. + +Run donor in one process; run consumer repeatedly in fresh processes on the same +GPU. The donor must stay alive. The example uses a local, immutable Wan 2.1 1.3B +Diffusers VAE component and FP32. It shares weights, not mutable VAE caches. +""" +import argparse +import hashlib +import json +from pathlib import Path +import signal +import threading +import time + +import torch + +from fastvideo.configs.pipelines import PipelineConfig +from fastvideo.distributed import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel +import fastvideo.envs as envs +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.models.loader.component_loader import VAELoader +from fastvideo.models.loader.utils import set_default_torch_dtype +from fastvideo.models.loader.weight_cache import WeightCacheClient, WeightCacheKey, serve_weight_cache +from fastvideo.models.wan.vae import AutoencoderKLWan +from fastvideo.models.wan.vae_config import WanVAEConfig +from fastvideo.utils import get_open_port + + +def fingerprint(model): + digest = hashlib.sha256() + for name, tensor in [*model.named_parameters(remove_duplicate=False), *model.named_buffers(remove_duplicate=False)]: + digest.update(name.encode()) + digest.update(str((tensor.dtype, tuple(tensor.shape), tuple(tensor.stride()), tensor.storage_offset())).encode()) + digest.update(tensor.detach().contiguous().reshape(-1).view(torch.uint8).cpu().numpy().tobytes()) + return digest.hexdigest() + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("role", choices=["donor", "consumer", "disk"]) + parser.add_argument("--model-path", type=Path, required=True) + parser.add_argument("--checkpoint-id", required=True, help="Immutable checkpoint revision or content digest") + parser.add_argument("--socket", type=Path) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--tensor-output", type=Path) + args = parser.parse_args() + if args.role != "disk" and args.socket is None: + parser.error("--socket is required for donor/consumer") + config_bytes = (args.model_path / "config.json").read_bytes() + architecture = json.loads(config_bytes) + if architecture.get("_class_name") != "AutoencoderKLWan": + raise ValueError("This adapter supports native Wan VAE only") + key = WeightCacheKey(args.checkpoint_id, "vae", hashlib.sha256(config_bytes).hexdigest(), + "torch.float32", "native-wan-fp32-v1") + envs.setdefault_external("MASTER_ADDR", "127.0.0.1") + envs.setdefault_external("MASTER_PORT", str(get_open_port())) + maybe_init_distributed_environment_and_model_parallel(1, 1) + device = torch.device("cuda", torch.cuda.current_device()) + report = {"role": args.role, "checkpoint": args.checkpoint_id, "torch": torch.__version__} + def load_disk(): + runtime = FastVideoArgs(model_path=str(args.model_path), vae_cpu_offload=False, + pipeline_config=PipelineConfig(vae_config=WanVAEConfig(), vae_precision="fp32")) + return VAELoader().load(str(args.model_path), runtime).eval() + native_config = WanVAEConfig() + native_config.update_model_arch(architecture) + def construct(): + return VAELoader().build_model(AutoencoderKLWan, native_config, torch.device("meta")) + def verify(model): + report["weight_sha256_before"] = fingerprint(model) + x = torch.randn((1, 16, 2, 12, 12), generator=torch.Generator().manual_seed(1024)).to(device) + with torch.no_grad(): + first = model.decode(x) + second = model.decode(x) + torch.cuda.synchronize(device) + assert first.shape == (1, 3, 5, 96, 96) and torch.isfinite(first).all() + assert torch.equal(first, second), "Request-local VAE caches were not reset" + report["weight_sha256_after"] = fingerprint(model) + assert report["weight_sha256_before"] == report["weight_sha256_after"] + if args.tensor_output: + args.tensor_output.parent.mkdir(parents=True, exist_ok=True) + torch.save(first.cpu(), args.tensor_output) + try: + if args.role == "donor": + model = load_disk() + report["weight_sha256_before"] = fingerprint(model) + stop = threading.Event() + for sig in (signal.SIGINT, signal.SIGTERM): + signal.signal(sig, lambda *_: stop.set()) + print("Serving donor weights", flush=True) + serve_weight_cache(model, key, args.socket, stop=stop) + report["weight_sha256_after"] = fingerprint(model) + assert report["weight_sha256_before"] == report["weight_sha256_after"] + else: + start = time.perf_counter() + allocated = torch.cuda.memory_allocated(device) + if args.role == "consumer": + with set_default_torch_dtype(torch.float32): + with WeightCacheClient(args.socket, key, device).load_model(construct) as attachment: + report["load_seconds"] = time.perf_counter() - start + report["allocated_weight_delta_bytes"] = torch.cuda.memory_allocated(device) - allocated + report["donor_pid"] = attachment.donor_pid + report["ipc_handles"] = attachment.handles + report["weight_source"] = "cuda_ipc" + verify(attachment.model) + else: + model = load_disk() + torch.cuda.synchronize(device) + report["load_seconds"] = time.perf_counter() - start + report["allocated_weight_delta_bytes"] = torch.cuda.memory_allocated(device) - allocated + report["weight_source"] = "disk" + verify(model) + report["status"] = "ok" + except Exception as error: + report.update(status="error", error=repr(error)) + raise + finally: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, indent=2) + "\n") + cleanup_dist_env_and_memory() + + +if __name__ == "__main__": + main() diff --git a/examples/benchmark/wan_pipeline_startup.py b/examples/benchmark/wan_pipeline_startup.py deleted file mode 100644 index 7e077dce8..000000000 --- a/examples/benchmark/wan_pipeline_startup.py +++ /dev/null @@ -1,96 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Profile local Wan 2.1 1.3B component loading and one tiny verification request. - -One GPU, eager execution, cached local checkpoint. The 9-frame, 128x128, -4-step request verifies execution, not visual quality or production latency. -""" -import argparse -import json -from pathlib import Path -import time - - -def main(): - started = time.perf_counter() - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--model-path", required=True, help="Local Wan 2.1 1.3B Diffusers checkpoint root") - parser.add_argument("--output", required=True, type=Path) - parser.add_argument("--tensor-output", type=Path) - parser.add_argument("--synchronize", action="store_true") - parser.add_argument("--fsdp", action="store_true", help="Exercise the single-rank FSDP inference loader") - args = parser.parse_args() - mark = time.perf_counter() - import torch - torch_seconds = time.perf_counter() - mark - mark = time.perf_counter() - from torch.distributed.fsdp import FSDPModule - from fastvideo import VideoGenerator - from fastvideo.models.loader.startup import collect_startup_timings, startup_phase - fastvideo_seconds = time.perf_counter() - mark - if not torch.cuda.is_available(): - raise RuntimeError("One CUDA GPU is required") - report = {"imports": {"torch_seconds": torch_seconds, "fastvideo_and_helpers_seconds": fastvideo_seconds}, - "model_path": args.model_path, "torch": torch.__version__, "fsdp": args.fsdp, - "cuda_initialized_before_pipeline": torch.cuda.is_initialized(), - "workload": {"frames": 9, "height": 128, "width": 128, "steps": 4, "seed": 1024}} - def synchronize_if_initialized(): - if torch.cuda.is_initialized(): - torch.cuda.synchronize() - generator = None - try: - with collect_startup_timings(synchronize=synchronize_if_initialized if args.synchronize else None) as timings: - with startup_phase("pipeline.create"): - generator = VideoGenerator.from_pretrained( - args.model_path, num_gpus=1, distributed_executor_backend="uni", sp_size=1, tp_size=1, - use_fsdp_inference=args.fsdp, dit_cpu_offload=False, dit_layerwise_offload=False, - vae_cpu_offload=False, - text_encoder_cpu_offload=False, inference_torch_compile=False, - vae_tiling=False, vae_sp=False, - ) - torch.cuda.synchronize() - report["ready_seconds"] = time.perf_counter() - started - effective = generator.fastvideo_args - report["effective_config"] = { - name: getattr(effective, name) for name in ( - "use_fsdp_inference", "dit_layerwise_offload", "dit_cpu_offload", - "text_encoder_cpu_offload", "vae_cpu_offload", "inference_torch_compile", - ) - } - assert effective.use_fsdp_inference == args.fsdp, "Requested FSDP mode was overridden" - assert not any(value for name, value in report["effective_config"].items() - if name != "use_fsdp_inference"), "Offload or compilation unexpectedly enabled" - # This benchmark deliberately uses the in-process executor. - transformer = generator.executor.driver_worker.worker.pipeline.modules["transformer"] - report["transformer_is_fsdp_module"] = isinstance(transformer, FSDPModule) - assert report["transformer_is_fsdp_module"] == args.fsdp, "Live model does not match FSDP request" - with startup_phase("first_request"): - result = generator.generate_video( - "A red cube on a wooden table.", num_frames=9, height=128, width=128, - num_inference_steps=4, seed=1024, save_video=False, return_frames=True, - ) - torch.cuda.synchronize() - samples = result["samples"] - assert isinstance(samples, torch.Tensor) - assert samples.shape == (1, 3, 9, 128, 128) and torch.isfinite(samples).all() - report["verified_seconds"] = time.perf_counter() - started - report["output_shape"] = list(samples.shape) - report["device"] = torch.cuda.get_device_name(0) - report["status"] = "ok" - if args.tensor_output: - args.tensor_output.parent.mkdir(parents=True, exist_ok=True) - torch.save(samples.detach().cpu(), args.tensor_output) - except Exception as error: - report["status"] = "error" - report["error"] = repr(error) - raise - finally: - if "timings" in locals(): - report["startup"] = timings.to_dict() - args.output.parent.mkdir(parents=True, exist_ok=True) - args.output.write_text(json.dumps(report, indent=2) + "\n") - if generator is not None: - generator.executor.shutdown() - - -if __name__ == "__main__": - main() diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index 633f0534a..75b192f08 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -440,46 +440,44 @@ class TextEncoderLoader(ComponentLoader): model_config.quant_config = quant_cls() if getattr(model_cls, "supports_hf_from_pretrained", False): - with startup_phase("text_encoder.pretrained"): - model = model_cls.from_pretrained_local( # type: ignore[attr-defined] - model_path, - model_config, # type: ignore[arg-type] - dtype=PRECISION_TO_TYPE[dtype], - device=target_device, - ) + model = model_cls.from_pretrained_local( # type: ignore[attr-defined] + model_path, + model_config, # type: ignore[arg-type] + dtype=PRECISION_TO_TYPE[dtype], + device=target_device, + ) # HF passthrough encoders return before FastVideo's FSDP # wrapping path, so the text stage needs their placement to # put token tensors on the same device. model._fastvideo_input_device = target_device return model.eval() - with startup_phase("text_encoder.construct"), target_device: + with target_device: model = model_cls(model_config) # type: ignore weights_to_load = {name for name, _ in model.named_parameters()} - with startup_phase("text_encoder.hydrate"): - if (use_text_encoder_override and fastvideo_args.override_text_encoder_safetensors is not None): - if os.path.isdir(checkpoint_path): - override_weights = self._get_all_weights( - model, - checkpoint_path, - to_cpu=bool(cpu_offload), - ) - else: - if self.counter_before_loading_weights == 0.0: - self.counter_before_loading_weights = time.perf_counter() - override_weights = safetensors_weights_iterator( - [checkpoint_path], - to_cpu=use_cpu_offload, - ) - loaded_weights: set[str] = model.load_weights(override_weights) # type: ignore + if (use_text_encoder_override and fastvideo_args.override_text_encoder_safetensors is not None): + if os.path.isdir(checkpoint_path): + override_weights = self._get_all_weights( + model, + checkpoint_path, + to_cpu=bool(cpu_offload), + ) else: - loaded_weights: set[str] = model.load_weights( - self._get_all_weights( - model, - model_path, - to_cpu=cpu_offload, - )) # type: ignore + if self.counter_before_loading_weights == 0.0: + self.counter_before_loading_weights = time.perf_counter() + override_weights = safetensors_weights_iterator( + [checkpoint_path], + to_cpu=use_cpu_offload, + ) + loaded_weights: set[str] = model.load_weights(override_weights) # type: ignore + else: + loaded_weights: set[str] = model.load_weights( + self._get_all_weights( + model, + model_path, + to_cpu=cpu_offload, + )) # type: ignore self.counter_after_loading_weights = time.perf_counter() logger.info( @@ -502,59 +500,57 @@ class TextEncoderLoader(ComponentLoader): checkpoint_quant_config.get_name()) # Explicitly move model to target device after loading weights - with startup_phase("text_encoder.device_transfer"): - model = model.to(target_device) + model = model.to(target_device) - with startup_phase("text_encoder.offload_setup"): - prepare_layerwise = getattr(model, "prepare_layerwise_offload", None) - if envs.FASTVIDEO_H3_ENCODER_LAYERWISE.get() and callable(prepare_layerwise): - if target_device.type != "cpu": - raise ValueError("Layerwise H3 encoder requires text_encoder_cpu_offload=True") - prepare_layerwise(runtime_device) - use_cpu_offload = False - logger.info("Enabled text-only layerwise H3 encoder with CPU token embeddings") + prepare_layerwise = getattr(model, "prepare_layerwise_offload", None) + if envs.FASTVIDEO_H3_ENCODER_LAYERWISE.get() and callable(prepare_layerwise): + if target_device.type != "cpu": + raise ValueError("Layerwise H3 encoder requires text_encoder_cpu_offload=True") + prepare_layerwise(runtime_device) + use_cpu_offload = False + logger.info("Enabled text-only layerwise H3 encoder with CPU token embeddings") - from fastvideo.platforms import current_platform + from fastvideo.platforms import current_platform - if use_cpu_offload and checkpoint_quant_config is not None: - logger.info( - "Skipping FSDP CPU offload for serialized %s text encoder; " - "packed uint8 weights are not FSDP-shardable", - checkpoint_quant_config.get_name(), + if use_cpu_offload and checkpoint_quant_config is not None: + logger.info( + "Skipping FSDP CPU offload for serialized %s text encoder; " + "packed uint8 weights are not FSDP-shardable", + checkpoint_quant_config.get_name(), + ) + elif use_cpu_offload: + pin_cpu_memory = fastvideo_args.pin_cpu_memory and is_pin_memory_available() + # Disable FSDP for MPS as it's not compatible + if current_platform.is_mps(): + logger.info("Disabling FSDP sharding for MPS platform as it's not compatible") + elif current_platform.is_npu(): + mesh = init_device_mesh( + "npu", + mesh_shape=(1, dist.get_world_size()), + mesh_dim_names=("offload", "replicate"), + ) + shard_model( + model, + cpu_offload=True, + reshard_after_forward=True, + mesh=mesh["offload"], + fsdp_shard_conditions=model._fsdp_shard_conditions, + pin_cpu_memory=pin_cpu_memory, + ) + else: + mesh = init_device_mesh( + "cuda", + mesh_shape=(1, dist.get_world_size()), + mesh_dim_names=("offload", "replicate"), + ) + shard_model( + model, + cpu_offload=True, + reshard_after_forward=True, + mesh=mesh["offload"], + fsdp_shard_conditions=model._fsdp_shard_conditions, + pin_cpu_memory=pin_cpu_memory, ) - elif use_cpu_offload: - pin_cpu_memory = fastvideo_args.pin_cpu_memory and is_pin_memory_available() - # Disable FSDP for MPS as it's not compatible - if current_platform.is_mps(): - logger.info("Disabling FSDP sharding for MPS platform as it's not compatible") - elif current_platform.is_npu(): - mesh = init_device_mesh( - "npu", - mesh_shape=(1, dist.get_world_size()), - mesh_dim_names=("offload", "replicate"), - ) - shard_model( - model, - cpu_offload=True, - reshard_after_forward=True, - mesh=mesh["offload"], - fsdp_shard_conditions=model._fsdp_shard_conditions, - pin_cpu_memory=pin_cpu_memory, - ) - else: - mesh = init_device_mesh( - "cuda", - mesh_shape=(1, dist.get_world_size()), - mesh_dim_names=("offload", "replicate"), - ) - shard_model( - model, - cpu_offload=True, - reshard_after_forward=True, - mesh=mesh["offload"], - fsdp_shard_conditions=model._fsdp_shard_conditions, - pin_cpu_memory=pin_cpu_memory, - ) return model.eval() diff --git a/fastvideo/models/loader/fsdp_load.py b/fastvideo/models/loader/fsdp_load.py index 5f995c27e..974bbdb92 100644 --- a/fastvideo/models/loader/fsdp_load.py +++ b/fastvideo/models/loader/fsdp_load.py @@ -21,7 +21,6 @@ from torch.nn.modules.module import _IncompatibleKeys import fastvideo.envs as envs from fastvideo.logger import init_logger from fastvideo.models.loader.lora_patch import DenseLoRAPatch -from fastvideo.models.loader.startup import startup_phase from fastvideo.models.loader.utils import (get_param_names_mapping, hf_to_custom_state_dict) from fastvideo.models.loader.weight_utils import safetensors_weights_iterator from fastvideo.utils import set_mixed_precision_policy, is_pin_memory_available @@ -329,7 +328,7 @@ def maybe_load_fsdp_model( ) logger.info("Loading model with default_dtype: %s", default_dtype) - with startup_phase("transformer.construct"), set_default_dtype(default_dtype), torch.device("meta"): + with set_default_dtype(default_dtype), torch.device("meta"): model = model_cls(**init_params) if pre_fsdp_model_transform is not None: @@ -357,36 +356,35 @@ def maybe_load_fsdp_model( use_fsdp = False logger.info("Disabling FSDP for MPS platform as it's not compatible") - with startup_phase("transformer.shard_setup"): - if use_fsdp: - pin_cpu_memory = pin_cpu_memory and is_pin_memory_available() - world_size = hsdp_replicate_dim * hsdp_shard_dim - if not training_mode and not fsdp_inference: - hsdp_replicate_dim = world_size - hsdp_shard_dim = 1 + if use_fsdp: + pin_cpu_memory = pin_cpu_memory and is_pin_memory_available() + world_size = hsdp_replicate_dim * hsdp_shard_dim + if not training_mode and not fsdp_inference: + hsdp_replicate_dim = world_size + hsdp_shard_dim = 1 - if device_mesh is None and current_platform.is_npu(): - with torch.device("cpu"): - device_mesh = init_device_mesh( - "npu", - # (Replicate(), Shard(dim=0)) - mesh_shape=(hsdp_replicate_dim, hsdp_shard_dim), - mesh_dim_names=("replicate", "shard"), - ) - elif device_mesh is None: + if device_mesh is None and current_platform.is_npu(): + with torch.device("cpu"): device_mesh = init_device_mesh( - "cuda", + "npu", # (Replicate(), Shard(dim=0)) mesh_shape=(hsdp_replicate_dim, hsdp_shard_dim), mesh_dim_names=("replicate", "shard"), ) - shard_model(model, - cpu_offload=cpu_offload, - reshard_after_forward=True, - mp_policy=mp_policy, - mesh=device_mesh, - fsdp_shard_conditions=model._fsdp_shard_conditions, - pin_cpu_memory=pin_cpu_memory) + elif device_mesh is None: + device_mesh = init_device_mesh( + "cuda", + # (Replicate(), Shard(dim=0)) + mesh_shape=(hsdp_replicate_dim, hsdp_shard_dim), + mesh_dim_names=("replicate", "shard"), + ) + shard_model(model, + cpu_offload=cpu_offload, + reshard_after_forward=True, + mp_policy=mp_policy, + mesh=device_mesh, + fsdp_shard_conditions=model._fsdp_shard_conditions, + pin_cpu_memory=pin_cpu_memory) nvfp4_skip_param_names: set[str] = set() if packed_nvfp4_export is not None: @@ -461,88 +459,85 @@ def maybe_load_fsdp_model( } logger.info("Skipping %d AdaLN projection tensors (precomputed modulation tables)", sum("adaln_proj" in n for n in skip_param_names)) - with startup_phase("transformer.hydrate"): - load_model_from_full_model_state_dict( - model, - weight_iterator, - device, - default_dtype, - strict=strict, - cpu_offload=cpu_offload, - param_names_mapping=param_names_mapping_fn, - dense_lora_patch=dense_lora_patch, - skip_param_names=skip_param_names or None, - ) - with startup_phase("transformer.finalize"): - if packed_nvfp4_export is not None: - load_minimax_h3_nvfp4_dit_export(model, packed_nvfp4_export, device=device) - if fp8_targets: - _load_prequantized_fp8(model, weight_dir_list, fp8_targets, device=device) - if hasattr(model, "materialize_non_persistent_buffers"): - model.materialize_non_persistent_buffers(device=device, dtype=default_dtype) - for n, p in chain(model.named_parameters(), model.named_buffers()): - if p.is_meta: - raise RuntimeError(f"Unexpected param or buffer {n} on meta device.") - # Avoid unintended computation graph accumulation during inference - if isinstance(p, torch.nn.Parameter): - p.requires_grad = False + load_model_from_full_model_state_dict( + model, + weight_iterator, + device, + default_dtype, + strict=strict, + cpu_offload=cpu_offload, + param_names_mapping=param_names_mapping_fn, + dense_lora_patch=dense_lora_patch, + skip_param_names=skip_param_names or None, + ) + if packed_nvfp4_export is not None: + load_minimax_h3_nvfp4_dit_export(model, packed_nvfp4_export, device=device) + if fp8_targets: + _load_prequantized_fp8(model, weight_dir_list, fp8_targets, device=device) + if hasattr(model, "materialize_non_persistent_buffers"): + model.materialize_non_persistent_buffers(device=device, dtype=default_dtype) + for n, p in chain(model.named_parameters(), model.named_buffers()): + if p.is_meta: + raise RuntimeError(f"Unexpected param or buffer {n} on meta device.") + # Avoid unintended computation graph accumulation during inference + if isinstance(p, torch.nn.Parameter): + p.requires_grad = False - # Post-load weight quantization. We detect the active scheme by the - # ``quant_method`` attached to each linear layer at construction time - # (via ``QuantizationConfig.get_quant_method``). The loader's - # responsibility is just to materialize the quantized weight buffers - # from the freshly-loaded bf16 weights. No-op when no quantized layers - # are present (lazy imports inside the helper). - _maybe_quantize_model(model, defer_weight_conversion_until_lora_merge=lora_path is not None) + # Post-load weight quantization. We detect the active scheme by the + # ``quant_method`` attached to each linear layer at construction time + # (via ``QuantizationConfig.get_quant_method``). The loader's + # responsibility is just to materialize the quantized weight buffers + # from the freshly-loaded bf16 weights. No-op when no quantized layers + # are present (lazy imports inside the helper). + _maybe_quantize_model(model, defer_weight_conversion_until_lora_merge=lora_path is not None) - with startup_phase("transformer.compile_setup"): - if enable_torch_compile and training_mode: - if not regional_compile: - # Legacy training stack (--enable-torch-compile): preserve the - # whole-model compile behavior and pass kwargs through unchanged. - unsupported = _prepare_model_for_compile(model, regional=False) - if unsupported is not None: - logger.warning("Training torch.compile requested but disabled: %s. Model stays eager.", unsupported) - else: - compile_kwargs = torch_compile_kwargs or {} - logger.info("Enabling whole-model torch.compile with kwargs=%s", compile_kwargs) - model = torch.compile(model, **compile_kwargs) + if enable_torch_compile and training_mode: + if not regional_compile: + # Legacy training stack (--enable-torch-compile): preserve the + # whole-model compile behavior and pass kwargs through unchanged. + unsupported = _prepare_model_for_compile(model, regional=False) + if unsupported is not None: + logger.warning("Training torch.compile requested but disabled: %s. Model stays eager.", unsupported) else: - unsupported = _regional_compile_unsupported_reason( - init_params, - model=model, - training=True, - vsa_tile_size=inference_vsa_tile_size, - ) - if unsupported is None: - unsupported = _prepare_model_for_compile(model, regional=False) - if unsupported is not None: - logger.warning( - "enable_torch_compile requested but disabled: %s. " - "Training continues in eager mode.", unsupported) - else: - attention_count = _enable_regional_attention_compile(model) - logger.info("Enabled attention tracing for %d modules in %s", attention_count, type(model).__name__) - _compile_model_regions(model, torch_compile_kwargs or {}) - elif inference_regional_compile and not training_mode: - # Inference-side counterpart of the #1718 training regional compile: - # per-block fullgraph compile right after the transformer loads, no - # user kwargs needed (fullgraph + emulate_precision_casts injected). + compile_kwargs = torch_compile_kwargs or {} + logger.info("Enabling whole-model torch.compile with kwargs=%s", compile_kwargs) + model = torch.compile(model, **compile_kwargs) + else: unsupported = _regional_compile_unsupported_reason( init_params, model=model, + training=True, vsa_tile_size=inference_vsa_tile_size, ) if unsupported is None: - unsupported = _prepare_model_for_compile(model, regional=True) + unsupported = _prepare_model_for_compile(model, regional=False) if unsupported is not None: logger.warning( - "inference_torch_compile requested but disabled: %s. " - "Inference continues in eager mode.", unsupported) + "enable_torch_compile requested but disabled: %s. " + "Training continues in eager mode.", unsupported) else: attention_count = _enable_regional_attention_compile(model) logger.info("Enabled attention tracing for %d modules in %s", attention_count, type(model).__name__) _compile_model_regions(model, torch_compile_kwargs or {}) + elif inference_regional_compile and not training_mode: + # Inference-side counterpart of the #1718 training regional compile: + # per-block fullgraph compile right after the transformer loads, no + # user kwargs needed (fullgraph + emulate_precision_casts injected). + unsupported = _regional_compile_unsupported_reason( + init_params, + model=model, + vsa_tile_size=inference_vsa_tile_size, + ) + if unsupported is None: + unsupported = _prepare_model_for_compile(model, regional=True) + if unsupported is not None: + logger.warning( + "inference_torch_compile requested but disabled: %s. " + "Inference continues in eager mode.", unsupported) + else: + attention_count = _enable_regional_attention_compile(model) + logger.info("Enabled attention tracing for %d modules in %s", attention_count, type(model).__name__) + _compile_model_regions(model, torch_compile_kwargs or {}) return model diff --git a/fastvideo/models/loader/weight_cache.py b/fastvideo/models/loader/weight_cache.py new file mode 100644 index 000000000..f6e2a20a1 --- /dev/null +++ b/fastvideo/models/loader/weight_cache.py @@ -0,0 +1,286 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Single-GPU, same-user CUDA IPC weights for restartable inference consumers. + +The donor owns registered tensors; each consumer owns new Python/module state. +The donor must outlive every mapped tensor. This prototype serves one active +consumer at a time and never falls back to checkpoint loading. +""" +from collections.abc import Callable, Iterator +from contextlib import contextmanager +from dataclasses import asdict, dataclass, fields +import contextlib +import fcntl +import hashlib +import os +from pathlib import Path +import pickle +import socket +import stat +import struct +import threading +from typing import Any + +import torch +from torch import nn +from torch.multiprocessing.reductions import rebuild_cuda_tensor, reduce_tensor + +_VERSION = 1 +_MAX_MESSAGE = 64 * 1024 * 1024 + + +@dataclass(frozen=True) +class WeightCacheKey: + """Caller supplies immutable checkpoint/config and weight-processing identities. + + Method-body edits may reuse weights; edits affecting names/layout/preparation + need a different processing_version. Tensor compatibility is also checked. + """ + checkpoint: str + component: str + config_hash: str + dtype: str + processing_version: str + + def __post_init__(self): + if any(not isinstance(getattr(self, item.name), str) or not getattr(self, item.name) + for item in fields(self)): + raise ValueError("Every weight-cache identity field must be a nonempty string") + + +@dataclass +class WeightAttachment: + model: nn.Module + donor_pid: int + handles: dict[str, str] + + +def _send(conn: socket.socket, payload: Any) -> None: + data = pickle.dumps(payload, protocol=5) + if len(data) > _MAX_MESSAGE: + raise ValueError("Weight-cache descriptor exceeds message limit") + conn.sendall(struct.pack("!I", len(data)) + data) + + +def _read(conn: socket.socket, length: int) -> bytes: + data = bytearray() + while len(data) < length: + chunk = conn.recv(length - len(data)) + if not chunk: + raise ConnectionError("Weight-cache peer disconnected") + data.extend(chunk) + return bytes(data) + + +def _recv(conn: socket.socket) -> Any: + length, = struct.unpack("!I", _read(conn, 4)) + if length > _MAX_MESSAGE: + raise ValueError("Weight-cache descriptor exceeds message limit") + return pickle.loads(_read(conn, length)) + + +def _same_user(conn: socket.socket) -> None: + # Descriptors contain pickle objects: authenticate the local peer before + # reading any payload. CUDA IPC is currently a Linux-only contract here. + _, uid, _ = struct.unpack("3i", conn.getsockopt(socket.SOL_SOCKET, socket.SO_PEERCRED, 12)) + if uid != os.getuid(): + raise PermissionError("Weight-cache peer is not the socket owner") + + +def _inventory(model: nn.Module) -> tuple[dict[str, torch.Tensor], dict[str, dict], dict[str, str]]: + tensors, specs, aliases, canonical = {}, {}, {}, {} + params = list(model.named_parameters(remove_duplicate=False)) + buffers = list(model.named_buffers(remove_duplicate=False)) + modules = dict(model.named_modules(remove_duplicate=False)) + for kind, values in (("parameter", params), ("buffer", buffers)): + for name, tensor in values: + if type(tensor) not in (torch.Tensor, nn.Parameter) or tensor.layout != torch.strided: + raise ValueError(f"Only native strided tensors are supported: {name}") + if id(tensor) in canonical: + aliases[name] = canonical[id(tensor)] + else: + canonical[id(tensor)] = name + tensors[name] = tensor + parent, _, leaf = name.rpartition(".") + specs[name] = {"kind": kind, "shape": tuple(tensor.shape), "dtype": str(tensor.dtype), + "stride": tuple(tensor.stride()), "offset": tensor.storage_offset(), + "persistent": kind == "parameter" or leaf not in modules[parent]._non_persistent_buffers_set} + return tensors, specs, aliases + + +def _manifest(model: nn.Module) -> dict: + _, specs, aliases = _inventory(model) + return {"model_class": type(model).__module__ + "." + type(model).__qualname__, + "specs": specs, "aliases": aliases} + + +def _validate_manifest(model: nn.Module, manifest: dict) -> None: + if _manifest(model) != manifest: + raise ValueError("Cached model names, shapes, dtypes, strides, aliases or class do not match") + + +def _bind(model: nn.Module, tensors: dict[str, torch.Tensor], manifest: dict) -> None: + """Validate everything before replacing any meta registration.""" + _validate_manifest(model, manifest) + expected = set(manifest["specs"]) - set(manifest["aliases"]) + if set(tensors) != expected: + raise ValueError("Incomplete cached tensor set") + for name, tensor in tensors.items(): + spec = manifest["specs"][name] + if (tuple(tensor.shape), str(tensor.dtype), tuple(tensor.stride()), tensor.storage_offset()) != ( + spec["shape"], spec["dtype"], spec["stride"], spec["offset"]): + raise ValueError(f"Cached tensor layout mismatch: {name}") + modules = dict(model.named_modules(remove_duplicate=False)) + registered = {} + for name in [*tensors, *manifest["aliases"]]: + source = manifest["aliases"].get(name) + spec = manifest["specs"][name] + obj = registered[source] if source is not None else tensors[name] + if source is None and spec["kind"] == "parameter": + obj = nn.Parameter(obj, requires_grad=False) + parent, _, leaf = name.rpartition(".") + if spec["kind"] == "parameter": + modules[parent].register_parameter(leaf, obj) + else: + modules[parent].register_buffer(leaf, obj, persistent=spec["persistent"]) + registered[name] = obj + + +def _gpu_uuid(device: torch.device) -> str: + return str(torch.cuda.get_device_properties(device).uuid) + + +def _private_directory(path: Path) -> None: + path.mkdir(mode=0o700, parents=True, exist_ok=True) + info = path.stat() + if info.st_uid != os.getuid() or stat.S_IMODE(info.st_mode) & 0o077: + raise PermissionError("Weight-cache socket directory must be owned by this user with mode 0700") + + +def serve_weight_cache(model: nn.Module, key: WeightCacheKey, socket_path: str | Path, + *, stop: threading.Event | None = None) -> None: + """Serve an already-loaded, idle inference model until stop is set. + + One active consumer is supported. Graceful stop waits for its connection to + close before releasing the model. SIGKILL/donor crashes invalidate consumers. + Callers own signal handling and must not mutate or offload donor tensors. + """ + if model.training: + raise ValueError("The donor model must be in evaluation mode") + tensors, _, _ = _inventory(model) + devices = {tensor.device for tensor in tensors.values()} + if not tensors or len(devices) != 1 or next(iter(devices)).type != "cuda": + raise ValueError("The donor requires native tensors on one CUDA GPU") + device = next(iter(devices)) + torch.cuda.synchronize(device) + manifest = _manifest(model) + path = Path(socket_path) + _private_directory(path.parent) + stop = stop if stop is not None else threading.Event() + lock_fd = os.open(str(path) + ".lock", os.O_CREAT | os.O_RDWR, 0o600) + try: + fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + except BaseException: + os.close(lock_fd) + raise + server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + try: + path.unlink(missing_ok=True) # Only the lock holder removes a stale socket. + server.bind(str(path)) + os.chmod(path, 0o600) + server.listen() + server.settimeout(0.5) + while not stop.is_set(): + try: + conn, _ = server.accept() + except socket.timeout: + continue + with conn: + try: + _same_user(conn) + conn.settimeout(30) + request = _recv(conn) + if request != {"version": _VERSION, "key": asdict(key), "torch": torch.__version__, + "gpu_uuid": _gpu_uuid(device)}: + _send(conn, {"error": "Weight-cache identity, runtime or GPU mismatch"}) + continue + _send(conn, {"manifest": manifest, "pid": os.getpid()}) + if _recv(conn) != "attach": + continue + # Generate reductions PER CONSUMER: each export acquires its + # own PyTorch CUDA IPC refcounts. Never reuse reduction args. + entries = {name: reduce_tensor(tensor.detach())[1] for name, tensor in tensors.items()} + _send(conn, entries) + conn.settimeout(None) + conn.recv(1) # EOF/release: keep donor storage alive throughout the lease. + except (ConnectionError, OSError, EOFError): + continue + except (ValueError, TypeError, pickle.UnpicklingError) as error: + with contextlib.suppress(OSError): + _send(conn, {"error": str(error)}) + finally: + torch.cuda.ipc_collect() + finally: + server.close() + path.unlink(missing_ok=True) + os.close(lock_fd) + + +class WeightCacheClient: + """Attach fresh module objects to donor-owned GPU storage; no disk fallback.""" + def __init__(self, socket_path: str | Path, key: WeightCacheKey, device: torch.device | None = None): + self.socket_path = Path(socket_path) + self.key = key + self.device = device if device is not None else torch.device("cuda", torch.cuda.current_device()) + + @contextmanager + def load_model(self, factory: Callable[[], nn.Module]) -> Iterator[WeightAttachment]: + """Construct on meta, validate the full layout, and bind IPC tensors. + + Hold this context until all GPU work finishes. Shared tensors are for + read-only inference. Mutable request state must be constructed locally; + training, in-place weight edits, CPU buffers and DTensor are unsupported. + """ + with torch.device("meta"): + model = factory() + meta_tensors, _, _ = _inventory(model) + template = _manifest(model) + if not meta_tensors or any(tensor.device.type != "meta" for tensor in meta_tensors.values()): + raise ValueError("Consumer factory must construct all registered tensors on meta") + _private_directory(self.socket_path.parent) + info = self.socket_path.stat() + if info.st_uid != os.getuid() or not stat.S_ISSOCK(info.st_mode) or stat.S_IMODE(info.st_mode) & 0o077: + raise PermissionError("Weight-cache socket must be private and owned by this user") + conn = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + conn.settimeout(30) + mapped = {} + bound = False + try: + conn.connect(str(self.socket_path)) + _same_user(conn) + _send(conn, {"version": _VERSION, "key": asdict(self.key), "torch": torch.__version__, + "gpu_uuid": _gpu_uuid(self.device)}) + header = _recv(conn) + if "error" in header: + raise ValueError(header["error"]) + _validate_manifest(model, header["manifest"]) + _send(conn, "attach") + entries = _recv(conn) + index = self.device.index if self.device.index is not None else torch.cuda.current_device() + mapped, handles = {}, {} + for name, exported in entries.items(): + args = list(exported) + args[6] = index # Exporter and consumer can have different visible-device ordinals. + mapped[name] = rebuild_cuda_tensor(*args) + handles[name] = hashlib.sha256(args[7] or b"").hexdigest() + _bind(model, mapped, header["manifest"]) + bound = True + yield WeightAttachment(model.eval(), header["pid"], handles) + finally: + # Synchronize even when inference raises before releasing the lease. + try: + torch.cuda.synchronize(self.device) + if bound: + _bind(model, meta_tensors, template) + mapped.clear() + finally: + conn.close() diff --git a/fastvideo/tests/loader/test_component_startup_phases.py b/fastvideo/tests/loader/test_component_startup_phases.py deleted file mode 100644 index 22a942c15..000000000 --- a/fastvideo/tests/loader/test_component_startup_phases.py +++ /dev/null @@ -1,72 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Lazy checkpoint consumption belongs to hydration, not iterator construction.""" -from types import SimpleNamespace - -import torch - -from fastvideo.models.loader import component_loader, fsdp_load -from fastvideo.models.loader.startup import collect_startup_timings, startup_phase - - -class TinyModel(torch.nn.Module): - param_names_mapping = {} - - def __init__(self, config): - super().__init__() - self.config = config - self.weight = torch.nn.Parameter(torch.zeros(2)) - - def load_weights(self, weights): - loaded = set() - for name, tensor in weights: - self.weight.data.copy_(tensor) - loaded.add(name) - return loaded - - -def _weights(*args, **kwargs): - with startup_phase("test.iterator_consumption"): - yield "weight", torch.tensor([3., 5.]) - - -def _assert_hydration(timings, prefix): - phases = timings.to_dict()["phases"] - hydration = next(p for p in phases if p["name"] == f"{prefix}.hydrate") - consumed = next(p for p in phases if p["name"] == "test.iterator_consumption") - assert consumed["parent_id"] == hydration["id"] - assert hydration["status"] == "ok" - assert any(p["name"] == f"{prefix}.construct" for p in phases) - - -def test_transformer_records_lazy_consumption_inside_hydration(monkeypatch): - monkeypatch.setattr(fsdp_load, "safetensors_weights_iterator", _weights) - with collect_startup_timings() as timings: - model = fsdp_load.maybe_load_fsdp_model( - model_cls=TinyModel, init_params={"config": SimpleNamespace(quant_config=None)}, - weight_dir_list=[], device=torch.device("cpu"), hsdp_replicate_dim=1, hsdp_shard_dim=1, - default_dtype=torch.float32, param_dtype=torch.float32, reduce_dtype=torch.float32, - fsdp_inference=False, training_mode=False, - ) - torch.testing.assert_close(model.weight, torch.tensor([3., 5.]), rtol=0, atol=0) - _assert_hydration(timings, "transformer") - - -def test_text_encoder_records_lazy_consumption_inside_hydration(monkeypatch, tmp_path): - monkeypatch.setattr(component_loader, "get_local_torch_device", lambda: torch.device("cpu")) - monkeypatch.setattr(component_loader.ModelRegistry, "resolve_model_cls", lambda _: (TinyModel, None)) - monkeypatch.setattr(component_loader, "_configure_text_encoder_quantization", lambda *args: None) - loader = component_loader.TextEncoderLoader() - monkeypatch.setattr(loader, "_get_all_weights", _weights) - # The existing counter is normally initialized by the real iterator. - loader.counter_before_loading_weights = 0.0 - config = SimpleNamespace(architectures=["TinyModel"], quant_config=None, _fsdp_shard_conditions=[]) - args = SimpleNamespace( - text_encoder_cpu_offload=False, pin_cpu_memory=False, - disable_offload_on_unified_memory=lambda *args, **kwargs: False, - override_text_encoder_quant=None, override_text_encoder_safetensors=None, - ) - with collect_startup_timings() as timings: - model = loader.load_model(str(tmp_path), config, torch.device("cpu"), args, dtype="fp32") - torch.testing.assert_close(model.weight, torch.tensor([3., 5.]), rtol=0, atol=0) - assert not model.training - _assert_hydration(timings, "text_encoder") diff --git a/fastvideo/tests/loader/test_weight_cache.py b/fastvideo/tests/loader/test_weight_cache.py new file mode 100644 index 000000000..706eebf9f --- /dev/null +++ b/fastvideo/tests/loader/test_weight_cache.py @@ -0,0 +1,153 @@ +# SPDX-License-Identifier: Apache-2.0 +"""IPC bindings preserve registration/alias identity and reject stale layouts.""" +from dataclasses import replace +import multiprocessing as mp +from pathlib import Path +import time + +import pytest +import torch + +from fastvideo.models.loader.weight_cache import ( + WeightCacheClient, WeightCacheKey, _bind, _inventory, _manifest, serve_weight_cache, +) + + +class Tiny(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.tensor([2., 3.])) + self.alias = self.weight + self.register_buffer("offset", torch.tensor([5., 7.]), persistent=False) + self.requests = 0 + + def forward(self, x): + self.requests += 1 + return x * self.weight + self.offset + + +def key(): + return WeightCacheKey("immutable-test-revision", "tiny", "config-sha", "torch.float32", "native-v1") + + +def test_binding_preserves_ties_and_nonpersistent_buffers(): + source = Tiny().eval() + tensors, _, _ = _inventory(source) + with torch.device("meta"): + target = Tiny() + _bind(target, tensors, _manifest(source)) + assert target.weight is target.alias + assert target.weight.data_ptr() == source.weight.data_ptr() + assert "offset" not in target.state_dict() + torch.testing.assert_close(target(torch.tensor([11., 13.])), torch.tensor([27., 46.]), rtol=0, atol=0) + assert source.requests == 0 and target.requests == 1 + + +def test_incompatible_layout_is_rejected_before_binding(): + source = Tiny() + with torch.device("meta"): + target = Tiny() + manifest = _manifest(source) + manifest["specs"]["weight"]["shape"] = (3,) + with pytest.raises(ValueError, match="do not match"): + _bind(target, _inventory(source)[0], manifest) + assert all(tensor.is_meta for tensor in target.parameters()) + + +def test_missing_cache_never_falls_back_to_disk(tmp_path): + with pytest.raises(FileNotFoundError): + with WeightCacheClient(tmp_path / "missing.sock", key(), torch.device("cuda:0")).load_model(Tiny): + pytest.fail("A missing donor must not load from disk") + + +def test_identity_requires_every_field(): + with pytest.raises(ValueError, match="nonempty"): + replace(key(), checkpoint="") + + +def _donor(path, stop): + model = Tiny().cuda().eval() + serve_weight_cache(model, key(), path, stop=stop) + torch.testing.assert_close(model.weight.cpu(), torch.tensor([2., 3.]), rtol=0, atol=0) + + +def _consumer(path, queue, mismatched=False): + client_key = replace(key(), checkpoint="wrong-revision") if mismatched else key() + try: + baseline = torch.cuda.memory_allocated() + with WeightCacheClient(path, client_key).load_model(Tiny) as attached: + model = attached.model + assert model.weight is model.alias + result = model(torch.tensor([11., 13.], device="cuda")).cpu().tolist() + assert model.requests == 1 + queue.put({"result": result, "pid": attached.donor_pid, "handles": attached.handles, + "allocated": torch.cuda.memory_allocated() - baseline}) + assert model.weight.is_meta # The model cannot be used beyond its lease. + except ValueError as error: + if not mismatched: + raise + queue.put({"rejected": str(error)}) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA IPC requires one GPU") +def test_donor_survives_two_fresh_consumers_and_rejected_identity(tmp_path): + ctx = mp.get_context("spawn") + stop = ctx.Event() + queue = ctx.Queue() + path = str(tmp_path / "weights.sock") + donor = ctx.Process(target=_donor, args=(path, stop)) + donor.start() + records = [] + try: + deadline = time.monotonic() + 90 + while not Path(path).exists(): + assert donor.is_alive(), "Donor died before readiness" + assert time.monotonic() < deadline, "Donor readiness timeout" + time.sleep(0.1) + for mismatched in (False, True, False): + worker = ctx.Process(target=_consumer, args=(path, queue, mismatched)) + worker.start() + try: + worker.join(90) + assert worker.exitcode == 0 + records.append(queue.get(timeout=10)) + finally: + if worker.is_alive(): + worker.terminate() + worker.join(10) + assert records[0]["result"] == records[2]["result"] == [27., 46.] + assert records[0]["pid"] == records[2]["pid"] == donor.pid + assert records[0]["handles"] == records[2]["handles"] + assert "rejected" in records[1] + assert donor.is_alive() + finally: + stop.set() + donor.join(10) + if donor.is_alive(): + donor.terminate() + donor.join(10) + queue.close() + assert donor.exitcode == 0 + + +def test_incomplete_tensor_set_does_not_partially_bind(): + source = Tiny() + with torch.device("meta"): + target = Tiny() + tensors = _inventory(source)[0] + tensors.pop("offset") + with pytest.raises(ValueError, match="Incomplete"): + _bind(target, tensors, _manifest(source)) + assert target.weight.is_meta and target.offset.is_meta + + +def test_donor_rejects_training_models_before_serving(tmp_path): + with pytest.raises(ValueError, match="evaluation"): + serve_weight_cache(Tiny(), key(), tmp_path / "weights.sock") + assert not (tmp_path / "weights.sock").exists() + + +def test_donor_rejects_cpu_weights_before_serving(tmp_path): + with pytest.raises(ValueError, match="one CUDA GPU"): + serve_weight_cache(Tiny().eval(), key(), tmp_path / "weights.sock") + assert not (tmp_path / "weights.sock").exists() diff --git a/mkdocs.yml b/mkdocs.yml index a06761eb2..101e3c059 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -256,6 +256,7 @@ nav: - Environment Variables: contributing/env_vars.md - Performance Benchmarks: contributing/performance_benchmarks.md - Startup Profiling: contributing/startup_profiling.md + - Retained Weights: contributing/weight_cache.md - Profiling: contributing/profiling.md - Activation Trace: contributing/activation_trace.md - Coding Agents: contributing/coding_agents.md