nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[feat] Retain native model weights in a CUDA IPC donor

Browse Source
This commit is contained in:
Junda Chen
2026-10-15 01:35:13 -07:00
parent 994be72623
commit 042e8ee580
10 changed files with 823 additions and 391 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/contributing/startup_profiling.md
+2 -51
View File
@@ -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.
docs/contributing/weight_cache.md
+94
View File
@@ -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.
examples/benchmark/vae_weight_cache.py
+124
View File
@@ -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()
examples/benchmark/wan_pipeline_startup.py
-96
View File
@@ -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()
fastvideo/models/loader/component_loader.py
+75 -79
View File
@@ -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()
fastvideo/models/loader/fsdp_load.py
+88 -93
View File
@@ -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
fastvideo/models/loader/weight_cache.py
+286
View File
@@ -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()
fastvideo/tests/loader/test_component_startup_phases.py
-72
View File
@@ -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")
fastvideo/tests/loader/test_weight_cache.py
+153
View File
@@ -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()
mkdocs.yml
+1
View File
@@ -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