[feat] Reuse Wan inference weights through normal loaders
This commit is contained in:
Junda Chen
@@ -1,87 +0,0 @@
|
||||
# Profiling model startup
|
||||
|
||||
Measure startup separately from inference latency before choosing what to retain
|
||||
between development runs. The initial instrumentation covers component totals
|
||||
and the native VAE safetensors loader, including Wan. It does not implement a
|
||||
weight donor or CUDA IPC.
|
||||
|
||||
## Run the Wan VAE profile
|
||||
|
||||
Use a prepared CUDA environment and a local Wan 2.1 1.3B Diffusers VAE component directory:
|
||||
|
||||
```bash
|
||||
python examples/benchmark/vae_startup.py \
|
||||
--model-path /path/to/checkpoint/vae \
|
||||
--output /tmp/vae-startup.json \
|
||||
--tensor-output /tmp/vae-output.pt \
|
||||
--synchronize
|
||||
```
|
||||
|
||||
The script measures torch and FastVideo/helper imports separately, device and
|
||||
process-group setup, VAE loading, and the first decode of a fixed seeded latent
|
||||
input. It preserves launcher-provided rendezvous settings; otherwise it chooses a
|
||||
local rendezvous port. It uses one GPU, FP32, and the default cached Wan VAE path.
|
||||
Imports and configuration can themselves initialize CUDA. The
|
||||
`cuda_initialized_before_device_setup` field records whether the explicit
|
||||
device-setup interval begins with an existing context; import time includes
|
||||
any initialization side effects.
|
||||
It verifies output shape and finiteness; those checks alone do not establish
|
||||
parity. Compare the saved tensors against an independent unchanged-source run.
|
||||
|
||||
Run the script in independent processes to measure repeated startup. Record the
|
||||
source revision, checkpoint revision, hardware, runtime versions, cache state,
|
||||
and process wall time alongside the JSON. The script excludes checkpoint
|
||||
downloads. Cached local files do not establish a cold filesystem-cache run.
|
||||
|
||||
## Collect component timings in a worker
|
||||
|
||||
```python
|
||||
from fastvideo.models.loader.startup import collect_startup_timings
|
||||
|
||||
with collect_startup_timings() as timings:
|
||||
pipeline = build_pipeline(...) # Your existing construction call.
|
||||
report = timings.to_dict()
|
||||
```
|
||||
|
||||
Collection is opt-in and local to the current context/process. Install a
|
||||
collector inside each spawned worker to measure its component loads; wrapping
|
||||
only a multiprocess driver does not collect worker timings. Lazy components are
|
||||
recorded when they actually load, if collection is still active.
|
||||
|
||||
Each phase has an ID, parent ID, start offset, duration, and success/error status.
|
||||
Parent durations include child durations: do not sum both. Recording duration
|
||||
also includes gaps between phases. A failed phase propagates its exception and
|
||||
records an error; the collector restores the previous context on exit.
|
||||
|
||||
Default durations measure host wall time, including asynchronous submission.
|
||||
For GPU completion timing, supply a synchronization callback. The example uses a
|
||||
callback that synchronizes only after CUDA is initialized, so the first callback
|
||||
does not move context initialization out of the device-setup interval.
|
||||
Synchronization changes overlap and adds overhead. Compare end-to-end latency
|
||||
with an uninstrumented run; do not interpret instrumented timings as free.
|
||||
|
||||
## Native VAE boundaries
|
||||
|
||||
`VAELoader.load()` retains the existing disk-loading behavior. Internally it now
|
||||
separates two operations:
|
||||
|
||||
- `build_model()` constructs the native module and moves its initial tensors to
|
||||
the target device. The caller supplies the resolved class, architecture config,
|
||||
and dtype context.
|
||||
- `load_weights()` reads safetensors, preserves checkpoint remapping and loading
|
||||
strictness, assigns weights, applies existing overlays/memory-format changes,
|
||||
and switches the module to evaluation mode.
|
||||
|
||||
Phases are `vae.construct`, `vae.device_transfer`, `vae.checkpoint_read`,
|
||||
`vae.weight_assignment`, and `vae.finalize`, nested inside `component.vae` when
|
||||
called through `PipelineComponentLoader`. Checkpoint reads can map files lazily;
|
||||
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. 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.
|
||||
@@ -6,9 +6,83 @@ IPC. Exiting a consumer does not unload the donor. This foundation is independen
|
||||
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.
|
||||
and one active consumer pipeline at a time. The normal inference loader supports
|
||||
Wan 2.1 T2V 1.3B DiT, UMT5 text encoder and VAE together. CPU registered tensors,
|
||||
DTensor/FSDP, offload, compilation, quantization, LoRA and training are unsupported.
|
||||
These combinations raise rather than falling back to disk.
|
||||
|
||||
## Complete Wan inference through the normal loader
|
||||
|
||||
The donor example loads DiT, T5 and VAE once through `VideoGenerator`. One process
|
||||
owns all three components and serves one socket per component. Tokenizers and
|
||||
schedulers remain local; each consumer builds new model objects and request state.
|
||||
|
||||
Start the donor on the same GPU, user and IPC namespace as the consumers:
|
||||
|
||||
```bash
|
||||
python examples/benchmark/wan_weight_cache.py donor \
|
||||
--model-path /path/to/Wan2.1-T2V-1.3B-Diffusers \
|
||||
--checkpoint-id IMMUTABLE_REVISION --cache-dir /tmp/fv-wan-weights \
|
||||
--output /tmp/wan-donor.json
|
||||
```
|
||||
|
||||
Wait for `DiT, T5 and VAE donor sockets ready`, then run independent consumers:
|
||||
|
||||
```bash
|
||||
python examples/benchmark/wan_weight_cache.py consumer \
|
||||
--model-path /path/to/Wan2.1-T2V-1.3B-Diffusers \
|
||||
--checkpoint-id IMMUTABLE_REVISION --cache-dir /tmp/fv-wan-weights \
|
||||
--output /tmp/wan-consumer-1.json --tensor-output /tmp/wan-consumer-1.pt
|
||||
```
|
||||
|
||||
Repeat with new output filenames. Use role `disk` for fresh-loading parity checks.
|
||||
Each worker runs two identical seeded 9-frame, 128x128, four-step requests. The
|
||||
in-process executor records full component fingerprints and IPC-handle witnesses.
|
||||
`--executor mp` also exercises the default executor selection; FastVideo resolves
|
||||
single-GPU `mp` to `UniprocExecutor`, which the report records explicitly.
|
||||
Stop the donor with Ctrl-C after consumers finish; it waits for active leases.
|
||||
|
||||
For application code, use the ordinary public inference entry point:
|
||||
|
||||
```python
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
local_checkpoint, num_gpus=1, tp_size=1, sp_size=1,
|
||||
weight_cache_dir="/tmp/fv-wan-weights",
|
||||
weight_cache_checkpoint="IMMUTABLE_REVISION",
|
||||
dit_cpu_offload=False, dit_layerwise_offload=False,
|
||||
text_encoder_cpu_offload=False, vae_cpu_offload=False,
|
||||
lazy_module_load=False, inference_torch_compile=False,
|
||||
)
|
||||
try:
|
||||
result = generator.generate_video(prompt, save_video=False)
|
||||
finally:
|
||||
generator.shutdown()
|
||||
```
|
||||
|
||||
The same fields have CLI options `--weight-cache-dir` and
|
||||
`--weight-cache-checkpoint`. A checkpoint identity without a directory keeps disk
|
||||
loading and records the donor keys. Supplying the directory makes attachment
|
||||
mandatory for all three components. There is no silent hydration fallback.
|
||||
|
||||
`PipelineComponentLoader` retains normal config resolution and class dispatch.
|
||||
Its native loader adapters attach after configuration but before checkpoint
|
||||
hydration. Identity includes the resolved architecture, config constants, dtype,
|
||||
attention backend request and loader processing version. The caller must provide
|
||||
an immutable checkpoint revision or content identity; the cache does not hash
|
||||
all checkpoint files during attachment. Recreate the donor when the checkpoint,
|
||||
layout, or weight-processing implementation changes.
|
||||
|
||||
The pipeline owns an `ExitStack` of attachment leases. Pipeline close and worker
|
||||
shutdown synchronize and detach all components; partial load failure releases
|
||||
already attached components. Direct `PipelineComponentLoader` users must supply
|
||||
`weight_cache_leases=stack` and keep that stack open throughout inference. Runtime
|
||||
LoRA merge/change APIs reject shared weights. Arbitrary user tensor mutations
|
||||
remain outside the read-only contract and cannot be prevented by CUDA IPC.
|
||||
|
||||
The donor reserves the combined weight memory for its lifetime. Attached workers
|
||||
still need their own CUDA contexts, communication state and activation memory.
|
||||
Keep the donor and sequential workers inside the same GPU allocation; this does
|
||||
not retain GPU memory across allocation teardown or machine restart.
|
||||
|
||||
## Native Wan VAE example
|
||||
|
||||
|
||||
@@ -1,86 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Profile a local Wan 2.1 1.3B VAE load and first decode on one CUDA GPU.
|
||||
|
||||
Run by filename so import timing begins before importing torch/FastVideo.
|
||||
The checkpoint must already be downloaded; this does not measure Hub downloads.
|
||||
"""
|
||||
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 VAE component directory")
|
||||
parser.add_argument("--output", type=Path, required=True)
|
||||
parser.add_argument("--tensor-output", type=Path, help="Optional decoded tensor for independent parity checks")
|
||||
parser.add_argument("--synchronize", action="store_true", help="Synchronize CUDA at phase boundaries")
|
||||
args = parser.parse_args()
|
||||
mark = time.perf_counter()
|
||||
import torch
|
||||
torch_import = time.perf_counter() - mark
|
||||
mark = time.perf_counter()
|
||||
import fastvideo.envs as envs
|
||||
from fastvideo.utils import get_open_port
|
||||
from fastvideo.configs.pipelines import PipelineConfig
|
||||
from fastvideo.distributed import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel
|
||||
from fastvideo.fastvideo_args import FastVideoArgs
|
||||
from fastvideo.models.loader.component_loader import PipelineComponentLoader
|
||||
from fastvideo.models.loader.startup import collect_startup_timings, startup_phase
|
||||
from fastvideo.models.wan.vae_config import WanVAEConfig
|
||||
fastvideo_import = time.perf_counter() - mark
|
||||
if not torch.cuda.is_available():
|
||||
raise RuntimeError("This profile requires one CUDA GPU")
|
||||
mark = time.perf_counter()
|
||||
runtime = FastVideoArgs(model_path=args.model_path,
|
||||
pipeline_config=PipelineConfig(vae_config=WanVAEConfig(), vae_precision="fp32"))
|
||||
runtime.vae_cpu_offload = False
|
||||
configuration_seconds = time.perf_counter() - mark
|
||||
result = {"imports": {"torch_seconds": torch_import, "fastvideo_and_helpers_seconds": fastvideo_import},
|
||||
"model_path": args.model_path, "torch": torch.__version__,
|
||||
"configuration_seconds": configuration_seconds,
|
||||
"cuda_initialized_before_device_setup": torch.cuda.is_initialized()}
|
||||
distributed_ready = False
|
||||
def synchronize_if_initialized():
|
||||
# Do not move CUDA context initialization out of the device_setup phase.
|
||||
if torch.cuda.is_initialized():
|
||||
torch.cuda.synchronize()
|
||||
try:
|
||||
with collect_startup_timings(synchronize=synchronize_if_initialized if args.synchronize else None) as timings:
|
||||
with startup_phase("device_setup"):
|
||||
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)
|
||||
distributed_ready = True
|
||||
vae = PipelineComponentLoader.load_module("vae", args.model_path, "diffusers", runtime)
|
||||
torch.cuda.synchronize()
|
||||
result["ready_seconds"] = time.perf_counter() - started
|
||||
latents = torch.randn((1, 16, 2, 12, 12), generator=torch.Generator().manual_seed(1024)).cuda()
|
||||
with startup_phase("first_forward"), torch.no_grad():
|
||||
output = vae.decode(latents)
|
||||
torch.cuda.synchronize()
|
||||
assert output.shape == (1, 3, 5, 96, 96) and torch.isfinite(output).all()
|
||||
result["verified_seconds"] = time.perf_counter() - started
|
||||
result["output_shape"] = list(output.shape)
|
||||
result["device"] = torch.cuda.get_device_name(0)
|
||||
if args.tensor_output:
|
||||
args.tensor_output.parent.mkdir(parents=True, exist_ok=True)
|
||||
torch.save(output.cpu(), args.tensor_output)
|
||||
result["status"] = "ok"
|
||||
except Exception as error:
|
||||
result["status"] = "error"
|
||||
result["error"] = repr(error)
|
||||
raise
|
||||
finally:
|
||||
if "timings" in locals():
|
||||
result["startup"] = timings.to_dict()
|
||||
args.output.parent.mkdir(parents=True, exist_ok=True)
|
||||
args.output.write_text(json.dumps(result, indent=2) + "\n")
|
||||
if distributed_ready:
|
||||
cleanup_dist_env_and_memory()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,149 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Retain Wan DiT/T5/VAE weights; verify fresh normal-loader inference workers.
|
||||
|
||||
Use a local Wan 2.1 T2V 1.3B Diffusers checkpoint on one GPU. Start a donor,
|
||||
then run consumers in separate processes with the same private cache directory.
|
||||
The tiny request checks execution/parity, not video quality or throughput.
|
||||
"""
|
||||
import argparse
|
||||
from dataclasses import asdict
|
||||
import hashlib
|
||||
import json
|
||||
from pathlib import Path
|
||||
import signal
|
||||
import threading
|
||||
import time
|
||||
|
||||
import torch
|
||||
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.models.loader.inference_weight_cache import COMPONENTS
|
||||
from fastvideo.models.loader.weight_cache import serve_weight_cache
|
||||
|
||||
|
||||
def fingerprints(modules):
|
||||
result = {}
|
||||
for component in COMPONENTS:
|
||||
digest = hashlib.sha256()
|
||||
model = modules[component]
|
||||
for name, tensor in [*model.named_parameters(remove_duplicate=False),
|
||||
*model.named_buffers(remove_duplicate=False)]:
|
||||
digest.update(name.encode())
|
||||
digest.update(str((tuple(tensor.shape), tensor.dtype, tuple(tensor.stride()))).encode())
|
||||
digest.update(tensor.detach().contiguous().reshape(-1).view(torch.uint8).cpu().numpy().tobytes())
|
||||
result[component] = digest.hexdigest()
|
||||
return result
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("role", choices=("donor", "consumer", "disk"))
|
||||
parser.add_argument("--model-path", required=True)
|
||||
parser.add_argument("--checkpoint-id", required=True)
|
||||
parser.add_argument("--cache-dir", type=Path, required=True)
|
||||
parser.add_argument("--output", type=Path, required=True)
|
||||
parser.add_argument("--tensor-output", type=Path)
|
||||
parser.add_argument("--executor", choices=("uni", "mp"), default="uni")
|
||||
args = parser.parse_args()
|
||||
report = {"role": args.role, "torch": torch.__version__, "executor": args.executor}
|
||||
generator = None
|
||||
stop = threading.Event()
|
||||
threads = []
|
||||
failures = []
|
||||
started = time.perf_counter()
|
||||
try:
|
||||
# Import/runtime allocations are separate from retained weight storage.
|
||||
torch.cuda.init()
|
||||
allocated = torch.cuda.memory_allocated()
|
||||
generator = VideoGenerator.from_pretrained(
|
||||
args.model_path, num_gpus=1, distributed_executor_backend=args.executor,
|
||||
tp_size=1, sp_size=1, hsdp_replicate_dim=1, hsdp_shard_dim=1,
|
||||
dit_cpu_offload=False, dit_layerwise_offload=False, vae_cpu_offload=False,
|
||||
text_encoder_cpu_offload=False, image_encoder_cpu_offload=False,
|
||||
lazy_module_load=False, inference_torch_compile=False,
|
||||
vae_tiling=False, vae_sp=False,
|
||||
weight_cache_dir=str(args.cache_dir) if args.role == "consumer" else None,
|
||||
weight_cache_checkpoint=args.checkpoint_id,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
report["ready_seconds"] = time.perf_counter() - started
|
||||
report["pipeline_allocation_delta_bytes"] = torch.cuda.memory_allocated() - allocated
|
||||
# FastVideo resolves both uni and single-GPU mp to UniprocExecutor.
|
||||
# Record the effective executor instead of assuming the requested name
|
||||
# implies a spawned worker.
|
||||
report["effective_executor"] = type(generator.executor).__name__
|
||||
modules = None
|
||||
if report["effective_executor"] == "UniprocExecutor":
|
||||
pipeline = generator.executor.driver_worker.worker.pipeline
|
||||
modules = pipeline.modules
|
||||
report["keys"] = {name: asdict(modules[name]._fastvideo_weight_cache_key) for name in COMPONENTS}
|
||||
report["attachments"] = {name: modules[name]._fastvideo_weight_cache_witness for name in COMPONENTS}
|
||||
report["fingerprints_before"] = fingerprints(modules)
|
||||
if args.role == "donor":
|
||||
if modules is None:
|
||||
raise ValueError("The donor must use the uni executor")
|
||||
args.cache_dir.mkdir(mode=0o700, parents=True, exist_ok=True)
|
||||
for sig in (signal.SIGINT, signal.SIGTERM):
|
||||
signal.signal(sig, lambda *_: stop.set())
|
||||
def serve(component):
|
||||
try:
|
||||
torch.cuda.set_device(0)
|
||||
serve_weight_cache(modules[component], modules[component]._fastvideo_weight_cache_key,
|
||||
args.cache_dir / f"{component}.sock", stop=stop)
|
||||
except BaseException as error:
|
||||
failures.append(error)
|
||||
stop.set()
|
||||
for component in COMPONENTS:
|
||||
thread = threading.Thread(target=serve, args=(component,))
|
||||
thread.start()
|
||||
threads.append(thread)
|
||||
deadline = time.monotonic() + 30
|
||||
while not all((args.cache_dir / f"{name}.sock").exists() for name in COMPONENTS):
|
||||
if stop.wait(.1) or time.monotonic() > deadline:
|
||||
raise RuntimeError(f"Donor failed to become ready: {failures}")
|
||||
report["donor_ready_seconds"] = time.perf_counter() - started
|
||||
print("DiT, T5 and VAE donor sockets ready", flush=True)
|
||||
while not stop.wait(.5):
|
||||
pass
|
||||
if failures:
|
||||
raise failures[0]
|
||||
else:
|
||||
outputs = []
|
||||
for _ in range(2):
|
||||
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,
|
||||
)
|
||||
samples = result["samples"].detach().cpu()
|
||||
assert samples.shape == (1, 3, 9, 128, 128) and torch.isfinite(samples).all()
|
||||
outputs.append(samples)
|
||||
assert torch.equal(outputs[0], outputs[1]), "Repeated seeded requests differ"
|
||||
report["output_shape"] = list(outputs[0].shape)
|
||||
if args.tensor_output:
|
||||
torch.save(outputs[0], args.tensor_output)
|
||||
# Close servers before the final fingerprint so no consumer can still
|
||||
# execute on these tensors during the donor's final verification.
|
||||
stop.set()
|
||||
for thread in threads:
|
||||
thread.join()
|
||||
if modules is not None:
|
||||
report["fingerprints_after"] = fingerprints(modules)
|
||||
assert report["fingerprints_before"] == report["fingerprints_after"]
|
||||
report["status"] = "ok"
|
||||
except BaseException as error:
|
||||
report["status"] = "error"
|
||||
report["error"] = repr(error)
|
||||
raise
|
||||
finally:
|
||||
stop.set()
|
||||
for thread in threads:
|
||||
thread.join()
|
||||
if generator is not None:
|
||||
generator.executor.shutdown()
|
||||
report["total_seconds"] = time.perf_counter() - started
|
||||
args.output.parent.mkdir(parents=True, exist_ok=True)
|
||||
args.output.write_text(json.dumps(report, indent=2) + "\n")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -162,6 +162,11 @@ class FastVideoArgs:
|
||||
trust_remote_code: bool = False
|
||||
revision: str | None = None
|
||||
|
||||
# Persistent native Wan inference weights. Directory selects IPC attachment;
|
||||
# identity alone records donor keys during ordinary disk loading.
|
||||
weight_cache_dir: str | None = None
|
||||
weight_cache_checkpoint: str | None = None
|
||||
|
||||
# Parallelism
|
||||
num_gpus: int = 1
|
||||
tp_size: int = -1
|
||||
@@ -569,6 +574,15 @@ class FastVideoArgs:
|
||||
help="The path of the model weights. This can be a local folder or a Hugging Face repo ID.",
|
||||
)
|
||||
|
||||
parser.add_argument("--weight-cache-dir",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Private directory containing native Wan inference weight-donor sockets")
|
||||
parser.add_argument("--weight-cache-checkpoint",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Immutable checkpoint revision/content identity for retained weights")
|
||||
|
||||
# Running mode
|
||||
parser.add_argument(
|
||||
"--mode",
|
||||
|
||||
@@ -7,7 +7,7 @@ import os
|
||||
import time
|
||||
from abc import ABC, abstractmethod
|
||||
from collections.abc import Generator, Iterable
|
||||
from contextlib import nullcontext
|
||||
from contextlib import ExitStack, nullcontext
|
||||
from copy import deepcopy
|
||||
from typing import Any, cast
|
||||
|
||||
@@ -39,7 +39,8 @@ from fastvideo.models.loader.text_encoder_quantization import (
|
||||
_process_quantized_text_encoder_weights,
|
||||
_resolve_text_encoder_checkpoint_path,
|
||||
)
|
||||
from fastvideo.models.loader.startup import startup_phase
|
||||
from fastvideo.models.loader.inference_weight_cache import (attach_cached_model, component_weight_cache,
|
||||
validate_cached_class)
|
||||
from fastvideo.models.loader.utils import set_default_torch_dtype
|
||||
from fastvideo.models.loader.weight_utils import (
|
||||
filter_duplicate_safetensors_files,
|
||||
@@ -439,6 +440,11 @@ class TextEncoderLoader(ComponentLoader):
|
||||
quant_cls = get_quantization_config(fastvideo_args.override_text_encoder_quant)
|
||||
model_config.quant_config = quant_cls()
|
||||
|
||||
cached = attach_cached_model(model_cls, model_config, PRECISION_TO_TYPE[dtype],
|
||||
lambda: model_cls(model_config))
|
||||
if cached is not None:
|
||||
return cached
|
||||
|
||||
if getattr(model_cls, "supports_hf_from_pretrained", False):
|
||||
model = model_cls.from_pretrained_local( # type: ignore[attr-defined]
|
||||
model_path,
|
||||
@@ -837,6 +843,7 @@ class VAELoader(ComponentLoader):
|
||||
class_name = getattr(fastvideo_args, "_model_index_class_names", {}).get("vae")
|
||||
assert class_name is not None, (
|
||||
"Model config does not contain a _class_name attribute. Only diffusers format is supported.")
|
||||
validate_cached_class(class_name)
|
||||
fastvideo_args.model_paths["vae"] = model_path
|
||||
|
||||
from fastvideo.platforms import current_platform
|
||||
@@ -949,6 +956,11 @@ class VAELoader(ComponentLoader):
|
||||
vae_config = fastvideo_args.pipeline_config.vae_config
|
||||
vae_config.update_model_arch(config)
|
||||
vae_cls, _ = ModelRegistry.resolve_model_cls(class_name)
|
||||
cached = attach_cached_model(
|
||||
vae_cls, vae_config, torch.get_default_dtype(),
|
||||
lambda: self.build_model(vae_cls, vae_config, torch.device("meta")))
|
||||
if cached is not None:
|
||||
return cached
|
||||
vae = self.build_model(vae_cls, vae_config, target_device)
|
||||
|
||||
return self.load_weights(vae, model_path, class_name=class_name, config=config, target_device=target_device)
|
||||
@@ -959,10 +971,8 @@ class VAELoader(ComponentLoader):
|
||||
Checkpoint-backed special constructors in load() retain their existing
|
||||
behavior. This boundary covers the native safetensors path, including Wan.
|
||||
"""
|
||||
with startup_phase("vae.construct"):
|
||||
vae = vae_cls(config)
|
||||
with startup_phase("vae.device_transfer"):
|
||||
vae = vae.to(target_device)
|
||||
vae = vae_cls(config)
|
||||
vae = vae.to(target_device)
|
||||
return vae
|
||||
|
||||
def load_weights(self, vae: nn.Module, model_path: str, *, class_name: str,
|
||||
@@ -988,9 +998,8 @@ class VAELoader(ComponentLoader):
|
||||
# Common case: a single `.safetensors` checkpoint file.
|
||||
# Some models may be sharded into multiple files; in that case we merge.
|
||||
loaded = {}
|
||||
with startup_phase("vae.checkpoint_read"):
|
||||
for sf_file in safetensors_list:
|
||||
loaded.update(safetensors_load_file(sf_file))
|
||||
for sf_file in safetensors_list:
|
||||
loaded.update(safetensors_load_file(sf_file))
|
||||
|
||||
# LTX-2 CausalVideoAutoencoder needs per_channel_statistics remapping
|
||||
if class_name == "CausalVideoAutoencoder" and "vae" in config:
|
||||
@@ -1019,17 +1028,15 @@ class VAELoader(ComponentLoader):
|
||||
# strictly so missing/unexpected keys are surfaced early.
|
||||
# Kandinsky6SRVAE: a partially loaded KVAE decodes plausible-looking garbage, so it is strict as well.
|
||||
strict_load = class_name in {"AutoencoderKL", "AutoencoderKLMiniMaxH3", "Kandinsky6SRVAE"}
|
||||
with startup_phase("vae.weight_assignment"):
|
||||
vae.load_state_dict(loaded, strict=strict_load)
|
||||
with startup_phase("vae.finalize"):
|
||||
if class_name == "AutoencoderKLMiniMaxH3" and int8_convrot_path is not None:
|
||||
from fastvideo.models.vaes.minimax_h3_int8_convrot import overlay_minimax_h3_int8_convrot_decoder
|
||||
overlay_minimax_h3_int8_convrot_decoder(vae, int8_convrot_path)
|
||||
if (class_name == "AutoencoderKLWan" and getattr(vae.config, "use_light_vae", False)
|
||||
and target_device.type == "cuda" and hasattr(vae, "optimize_memory_format")):
|
||||
vae.optimize_memory_format()
|
||||
vae.load_state_dict(loaded, strict=strict_load)
|
||||
if class_name == "AutoencoderKLMiniMaxH3" and int8_convrot_path is not None:
|
||||
from fastvideo.models.vaes.minimax_h3_int8_convrot import overlay_minimax_h3_int8_convrot_decoder
|
||||
overlay_minimax_h3_int8_convrot_decoder(vae, int8_convrot_path)
|
||||
if (class_name == "AutoencoderKLWan" and getattr(vae.config, "use_light_vae", False)
|
||||
and target_device.type == "cuda" and hasattr(vae, "optimize_memory_format")):
|
||||
vae.optimize_memory_format()
|
||||
|
||||
return vae.eval()
|
||||
return vae.eval()
|
||||
|
||||
|
||||
class AudioDecoderLoader(ComponentLoader):
|
||||
@@ -1235,6 +1242,17 @@ class TransformerLoader(ComponentLoader):
|
||||
|
||||
model_cls, _ = ModelRegistry.resolve_model_cls(cls_name)
|
||||
|
||||
default_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.dit_precision]
|
||||
def construct_cached_transformer():
|
||||
with set_default_torch_dtype(default_dtype):
|
||||
return model_cls(config=dit_config, hf_config=hf_config)
|
||||
record_resolved_attention_backend(dit_config)
|
||||
cached = attach_cached_model(model_cls, dit_config, default_dtype, construct_cached_transformer,
|
||||
extra_config=hf_config)
|
||||
if cached is not None:
|
||||
fastvideo_args.pipeline_config.dit_config = dit_config
|
||||
return cached
|
||||
|
||||
# Find all safetensors files
|
||||
safetensors_list = glob.glob(os.path.join(str(model_path), "*.safetensors"))
|
||||
if not safetensors_list:
|
||||
@@ -1614,6 +1632,7 @@ class PipelineComponentLoader:
|
||||
component_model_path: str,
|
||||
transformers_or_diffusers: str,
|
||||
fastvideo_args: FastVideoArgs,
|
||||
weight_cache_leases: ExitStack | None = None,
|
||||
):
|
||||
"""
|
||||
Load a pipeline module.
|
||||
@@ -1654,5 +1673,10 @@ class PipelineComponentLoader:
|
||||
attention_context = (_component_attention_backend_scope(resolved, component=module_name)
|
||||
if resolved is not None else nullcontext())
|
||||
|
||||
with startup_phase(f"component.{module_name}"), attention_context:
|
||||
return loader.load(component_model_path, fastvideo_args)
|
||||
with attention_context:
|
||||
with component_weight_cache(module_name, fastvideo_args, weight_cache_leases) as cache:
|
||||
model = loader.load(component_model_path, fastvideo_args)
|
||||
if cache is not None:
|
||||
model._fastvideo_weight_cache_key = cache.key
|
||||
model._fastvideo_weight_cache_witness = cache.witness
|
||||
return model
|
||||
|
||||
@@ -0,0 +1,132 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Opt-in native Wan loader adapters for persistent CUDA IPC weights."""
|
||||
from collections.abc import Callable
|
||||
from contextlib import ExitStack, contextmanager
|
||||
from contextvars import ContextVar
|
||||
from dataclasses import asdict, dataclass, is_dataclass
|
||||
from enum import Enum
|
||||
import hashlib
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from fastvideo.models.loader.weight_cache import WeightCacheClient, WeightCacheKey
|
||||
|
||||
COMPONENTS = ("transformer", "text_encoder", "vae")
|
||||
_CLASSES = {"transformer": "WanTransformer3DModel", "text_encoder": "UMT5EncoderModel",
|
||||
"vae": "AutoencoderKLWan"}
|
||||
|
||||
|
||||
def validate_weight_cache_args(args) -> None:
|
||||
"""Reject modes that can mutate, relocate or reinterpret shared storage."""
|
||||
if not args.inference_mode or str(args.mode) not in ("inference", "ExecutionMode.INFERENCE"):
|
||||
raise ValueError("Weight reuse supports inference only")
|
||||
if type(args.pipeline_config).__name__ != "WanT2V480PConfig":
|
||||
raise ValueError("Weight reuse currently supports the native Wan 2.1 T2V 1.3B pipeline only")
|
||||
if any(getattr(args, name) != 1 for name in ("num_gpus", "tp_size", "sp_size",
|
||||
"hsdp_replicate_dim", "hsdp_shard_dim")):
|
||||
raise ValueError("Weight reuse requires one GPU and single-rank parallelism")
|
||||
unsupported = ("use_fsdp_inference", "dit_cpu_offload", "dit_layerwise_offload", "vae_cpu_offload",
|
||||
"text_encoder_cpu_offload", "lazy_module_load", "enable_torch_compile",
|
||||
"enable_torch_compile_text_encoder", "enable_torch_compile_vae", "inference_torch_compile",
|
||||
"lora_path", "init_weights_from_safetensors", "init_weights_from_safetensors_2",
|
||||
"override_text_encoder_safetensors", "override_text_encoder_quant",
|
||||
"override_transformer_cls_name", "override_transformer_safetensors")
|
||||
enabled = [name for name in unsupported if getattr(args, name, None)]
|
||||
if enabled:
|
||||
raise ValueError(f"Weight reuse does not support: {', '.join(enabled)}")
|
||||
configs = [args.pipeline_config.dit_config, *args.pipeline_config.text_encoder_configs]
|
||||
if any(getattr(config, "quant_config", None) is not None for config in configs):
|
||||
raise ValueError("Weight reuse does not support quantized components")
|
||||
if args.pipeline_config.vae_config.use_light_vae:
|
||||
raise ValueError("Weight reuse does not support light-VAE memory-format conversion")
|
||||
if not args.weight_cache_checkpoint:
|
||||
raise ValueError("Weight reuse requires an immutable weight_cache_checkpoint identity")
|
||||
|
||||
|
||||
def _json_default(value):
|
||||
if is_dataclass(value):
|
||||
return asdict(value)
|
||||
if isinstance(value, torch.Tensor):
|
||||
if value.is_meta:
|
||||
raise ValueError("Configuration constants must be initialized outside meta construction")
|
||||
return {"dtype": str(value.dtype), "shape": list(value.shape), "values": value.cpu().tolist()}
|
||||
if isinstance(value, (set, frozenset)):
|
||||
return sorted(value)
|
||||
if isinstance(value, Enum):
|
||||
return value.value
|
||||
if isinstance(value, (torch.dtype, torch.device, Path)):
|
||||
return str(value)
|
||||
if callable(value):
|
||||
return value.__module__ + "." + value.__qualname__
|
||||
raise TypeError(f"Unsupported weight configuration identity: {type(value)}")
|
||||
|
||||
|
||||
@dataclass
|
||||
class _LoadContext:
|
||||
component: str
|
||||
args: Any
|
||||
leases: ExitStack | None
|
||||
key: WeightCacheKey | None = None
|
||||
witness: dict | None = None
|
||||
|
||||
|
||||
_CURRENT: ContextVar[_LoadContext | None] = ContextVar("inference_weight_cache_load", default=None)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def component_weight_cache(component, args, leases):
|
||||
"""The normal loader owns dispatch; the pipeline owns attachment lifetimes.
|
||||
|
||||
A checkpoint identity without a cache directory records donor keys while
|
||||
retaining normal disk hydration. With a directory, attachment is mandatory.
|
||||
"""
|
||||
if not (getattr(args, "weight_cache_dir", None) or getattr(args, "weight_cache_checkpoint", None)):
|
||||
yield None
|
||||
return
|
||||
validate_weight_cache_args(args)
|
||||
if component not in (*COMPONENTS, "scheduler", "tokenizer"):
|
||||
raise ValueError(f"Unsupported weight-cache pipeline component: {component}")
|
||||
if component not in COMPONENTS:
|
||||
yield None
|
||||
return
|
||||
if args.weight_cache_dir and leases is None:
|
||||
raise ValueError("IPC loading requires a caller-owned weight_cache_leases ExitStack")
|
||||
context = _LoadContext(component, args, leases)
|
||||
token = _CURRENT.set(context)
|
||||
try:
|
||||
yield context
|
||||
if context.key is None:
|
||||
raise ValueError(f"No native weight-cache adapter for {component}")
|
||||
finally:
|
||||
_CURRENT.reset(token)
|
||||
|
||||
|
||||
def validate_cached_class(class_name):
|
||||
context = _CURRENT.get()
|
||||
if context is not None and class_name != _CLASSES[context.component]:
|
||||
raise ValueError(f"Unsupported cached {context.component} class: {class_name}")
|
||||
|
||||
|
||||
def attach_cached_model(model_cls, config, dtype, factory: Callable[[], nn.Module], *, extra_config=None):
|
||||
"""Called after config resolution and before any checkpoint hydration."""
|
||||
context = _CURRENT.get()
|
||||
if context is None:
|
||||
return None
|
||||
validate_cached_class(model_cls.__name__)
|
||||
if getattr(config, "quant_config", None) is not None:
|
||||
raise ValueError("Weight reuse does not support checkpoint-selected quantization")
|
||||
identity = {"class": model_cls.__module__ + "." + model_cls.__qualname__,
|
||||
"config": config, "extra": extra_config}
|
||||
encoded = json.dumps(identity, default=_json_default, sort_keys=True).encode()
|
||||
context.key = WeightCacheKey(context.args.weight_cache_checkpoint, context.component,
|
||||
hashlib.sha256(encoded).hexdigest(), str(dtype), "native-wan-loader-v1")
|
||||
if not context.args.weight_cache_dir:
|
||||
return None
|
||||
client = WeightCacheClient(Path(context.args.weight_cache_dir) / f"{context.component}.sock", context.key)
|
||||
attached = context.leases.enter_context(client.load_model(factory))
|
||||
context.witness = {"donor_pid": attached.donor_pid, "handles": attached.handles}
|
||||
return attached.model
|
||||
@@ -1,84 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Opt-in, process-local startup timings without importing or initializing CUDA."""
|
||||
from collections.abc import Callable, Iterator
|
||||
from contextlib import contextmanager
|
||||
from contextvars import ContextVar
|
||||
from dataclasses import dataclass, field
|
||||
import os
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
|
||||
@dataclass
|
||||
class StartupTimings:
|
||||
"""Nested wall-clock intervals; parent durations include their children.
|
||||
|
||||
Without a synchronization callback these measure host submission time only.
|
||||
Passing a callback (for example torch.cuda.synchronize) waits at both phase
|
||||
boundaries and changes execution/overlap; use a separate uninstrumented run
|
||||
for an end-to-end latency comparison.
|
||||
"""
|
||||
synchronize: Callable[[], None] | None = field(default=None, repr=False)
|
||||
phases: list[dict[str, Any]] = field(default_factory=list)
|
||||
_start: float = field(default_factory=time.perf_counter, repr=False)
|
||||
_end: float | None = field(default=None, repr=False)
|
||||
_next_id: int = field(default=0, repr=False)
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {
|
||||
"schema_version": 1,
|
||||
"pid": os.getpid(),
|
||||
"clock": "perf_counter",
|
||||
"timing_mode": "synchronized_boundaries" if self.synchronize else "host",
|
||||
"recording_seconds": (self._end if self._end is not None else time.perf_counter()) - self._start,
|
||||
"phases": sorted(self.phases, key=lambda item: item["id"]),
|
||||
}
|
||||
|
||||
|
||||
_active: ContextVar[StartupTimings | None] = ContextVar("fastvideo_startup_timings", default=None)
|
||||
_stack: ContextVar[tuple[tuple[StartupTimings, int], ...]] = ContextVar("fastvideo_startup_stack", default=())
|
||||
|
||||
|
||||
@contextmanager
|
||||
def collect_startup_timings(*, synchronize: Callable[[], None] | None = None) -> Iterator[StartupTimings]:
|
||||
"""Collect only work in this context/process. Spawned workers need their own collector."""
|
||||
timings = StartupTimings(synchronize=synchronize)
|
||||
token = _active.set(timings)
|
||||
try:
|
||||
yield timings
|
||||
finally:
|
||||
timings._end = time.perf_counter()
|
||||
_active.reset(token)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def startup_phase(name: str) -> Iterator[None]:
|
||||
"""Record a named phase when collection is enabled; otherwise just execute it."""
|
||||
timings = _active.get()
|
||||
if timings is None:
|
||||
yield
|
||||
return
|
||||
if timings.synchronize is not None:
|
||||
timings.synchronize()
|
||||
phase_id = timings._next_id
|
||||
timings._next_id += 1
|
||||
stack = _stack.get()
|
||||
parent = next((index for owner, index in reversed(stack) if owner is timings), None)
|
||||
token = _stack.set((*stack, (timings, phase_id)))
|
||||
start = time.perf_counter()
|
||||
status = "error"
|
||||
try:
|
||||
yield
|
||||
if timings.synchronize is not None:
|
||||
timings.synchronize()
|
||||
status = "ok"
|
||||
finally:
|
||||
timings.phases.append({
|
||||
"id": phase_id,
|
||||
"parent_id": parent,
|
||||
"name": name,
|
||||
"start_seconds": start - timings._start,
|
||||
"duration_seconds": time.perf_counter() - start,
|
||||
"status": status,
|
||||
})
|
||||
_stack.reset(token)
|
||||
@@ -9,7 +9,6 @@ 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
|
||||
@@ -164,6 +163,10 @@ def serve_weight_cache(model: nn.Module, key: WeightCacheKey, socket_path: str |
|
||||
close before releasing the model. SIGKILL/donor crashes invalidate consumers.
|
||||
Callers own signal handling and must not mutate or offload donor tensors.
|
||||
"""
|
||||
# Keep the ordinary loader importable on platforms without Linux IPC.
|
||||
# Only the explicitly requested donor needs flock.
|
||||
import fcntl
|
||||
|
||||
if model.training:
|
||||
raise ValueError("The donor model must be in evaluation mode")
|
||||
tensors, _, _ = _inventory(model)
|
||||
|
||||
@@ -9,6 +9,7 @@ import argparse
|
||||
import os
|
||||
from abc import ABC, abstractmethod
|
||||
from collections.abc import Iterator
|
||||
from contextlib import ExitStack
|
||||
from functools import partial
|
||||
from typing import Any, cast
|
||||
|
||||
@@ -143,6 +144,7 @@ class ComposedPipelineBase(ABC):
|
||||
self._stages: list[PipelineStage] = []
|
||||
self._stage_name_mapping: dict[str, PipelineStage] = {}
|
||||
self._trace_mgr = None
|
||||
self._weight_cache_leases = ExitStack()
|
||||
|
||||
if required_config_modules is not None:
|
||||
self._required_config_modules = required_config_modules
|
||||
@@ -170,7 +172,11 @@ class ComposedPipelineBase(ABC):
|
||||
# Load modules directly in initialization
|
||||
logger.info("Loading pipeline modules...")
|
||||
with self.profiler_controller.region("profiler_region_model_loading"):
|
||||
self.modules = self.load_modules(fastvideo_args, loaded_modules)
|
||||
try:
|
||||
self.modules = self.load_modules(fastvideo_args, loaded_modules)
|
||||
except BaseException:
|
||||
self._weight_cache_leases.close()
|
||||
raise
|
||||
|
||||
def set_trainable(self) -> None:
|
||||
# Only train DiT
|
||||
@@ -627,6 +633,7 @@ class ComposedPipelineBase(ABC):
|
||||
component_model_path=component_model_path,
|
||||
transformers_or_diffusers=transformers_or_diffusers,
|
||||
fastvideo_args=fastvideo_args,
|
||||
weight_cache_leases=self._weight_cache_leases,
|
||||
)
|
||||
|
||||
if self._lazy_module_load_enabled(fastvideo_args) and module_name in self._lazy_module_names:
|
||||
@@ -787,6 +794,9 @@ class ComposedPipelineBase(ABC):
|
||||
raise NotImplementedError("if training_mode is True, the pipeline must implement this method")
|
||||
|
||||
def close(self) -> None:
|
||||
leases = getattr(self, "_weight_cache_leases", None)
|
||||
if leases is not None:
|
||||
leases.close()
|
||||
detach_activation_trace(getattr(self, "_trace_mgr", None))
|
||||
self._trace_mgr = None
|
||||
|
||||
|
||||
@@ -424,6 +424,8 @@ class LoRAPipeline(ComposedPipelineBase):
|
||||
accumulate: Add this adapter to an already merged pure low-rank adapter.
|
||||
"""
|
||||
|
||||
if getattr(self.fastvideo_args, "weight_cache_dir", None):
|
||||
raise ValueError("LoRA edits cannot mutate shared IPC weights")
|
||||
if not math.isfinite(strength):
|
||||
raise ValueError(f"LoRA strength must be finite, got {strength}")
|
||||
|
||||
@@ -592,6 +594,8 @@ class LoRAPipeline(ComposedPipelineBase):
|
||||
NVFP4 layers are only supported while their BF16 weight is retained: a purged
|
||||
BF16 weight cannot be unmerged and remerged.
|
||||
"""
|
||||
if getattr(self.fastvideo_args, "weight_cache_dir", None):
|
||||
raise ValueError("LoRA edits cannot mutate shared IPC weights")
|
||||
if _has_nvfp4_weights_without_bf16(self.trainable_transformer_modules):
|
||||
# TODO(David): Merge from the packed FP4 weight instead of rejecting.
|
||||
raise RuntimeError(
|
||||
@@ -616,6 +620,8 @@ class LoRAPipeline(ComposedPipelineBase):
|
||||
|
||||
NVFP4 is rejected: its FP4 buffers cannot be rebuilt from a purged BF16 weight.
|
||||
"""
|
||||
if getattr(self.fastvideo_args, "weight_cache_dir", None):
|
||||
raise ValueError("LoRA edits cannot mutate shared IPC weights")
|
||||
if _has_quantized_nvfp4_weights(self.trainable_transformer_modules):
|
||||
# TODO(David): Preserve BF16 weights and requantize NVFP4 weights after LoRA unmerge.
|
||||
raise RuntimeError(
|
||||
|
||||
@@ -0,0 +1,162 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Normal-loader IPC dispatch fails closed and scopes leases to the pipeline."""
|
||||
from contextlib import ExitStack, contextmanager
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from fastvideo.fastvideo_args import ExecutionMode, FastVideoArgs
|
||||
from fastvideo.models.loader.component_loader import ComponentLoader, PipelineComponentLoader
|
||||
from fastvideo.models.loader import inference_weight_cache as cache
|
||||
from fastvideo.models.wan.pipeline_config import WanT2V480PConfig
|
||||
|
||||
|
||||
def args(**overrides):
|
||||
config = dict(model_path="unused", pipeline_config=WanT2V480PConfig(), num_gpus=1,
|
||||
tp_size=1, sp_size=1, hsdp_replicate_dim=1, hsdp_shard_dim=1,
|
||||
dit_cpu_offload=False, dit_layerwise_offload=False, vae_cpu_offload=False,
|
||||
text_encoder_cpu_offload=False, image_encoder_cpu_offload=False, lazy_module_load=False,
|
||||
inference_torch_compile=False, weight_cache_checkpoint="immutable-revision")
|
||||
config.update(overrides)
|
||||
return FastVideoArgs(**config)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("option,value", [
|
||||
("inference_mode", False), ("use_fsdp_inference", True), ("num_gpus", 2),
|
||||
("dit_cpu_offload", True), ("dit_layerwise_offload", True), ("text_encoder_cpu_offload", True),
|
||||
("vae_cpu_offload", True), ("lora_path", "adapter"), ("enable_torch_compile", True),
|
||||
("init_weights_from_safetensors", "override"), ("weight_cache_checkpoint", ""),
|
||||
])
|
||||
def test_unsupported_modes_rejected_before_disk_loading(option, value):
|
||||
runtime = args(weight_cache_dir="unused")
|
||||
setattr(runtime, option, value)
|
||||
with pytest.raises(ValueError):
|
||||
with cache.component_weight_cache("transformer", runtime, ExitStack()):
|
||||
pytest.fail("Unsupported mode reached a loader")
|
||||
|
||||
|
||||
def test_training_mode_rejected_even_if_inference_boolean_is_stale():
|
||||
runtime = args()
|
||||
runtime.mode = ExecutionMode.FINETUNING
|
||||
with pytest.raises(ValueError, match="inference only"):
|
||||
cache.validate_weight_cache_args(runtime)
|
||||
|
||||
|
||||
def test_ipc_requires_explicit_lifetime_owner():
|
||||
with pytest.raises(ValueError, match="ExitStack"):
|
||||
with cache.component_weight_cache("vae", args(weight_cache_dir="unused"), None):
|
||||
pytest.fail("IPC cannot escape an unmanaged lifetime")
|
||||
|
||||
|
||||
def test_unknown_adapter_cannot_silently_load_from_disk():
|
||||
with pytest.raises(ValueError, match="No native"):
|
||||
with cache.component_weight_cache("vae", args(), None):
|
||||
pass
|
||||
|
||||
|
||||
def test_context_restores_after_failed_loader():
|
||||
with pytest.raises(RuntimeError):
|
||||
with cache.component_weight_cache("vae", args(), None):
|
||||
raise RuntimeError("Failed constructor")
|
||||
assert cache.attach_cached_model(None, None, None, lambda: None) is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("component", cache.COMPONENTS)
|
||||
def test_normal_dispatch_attaches_all_components_without_disk_hydration(monkeypatch, component):
|
||||
runtime = args(weight_cache_dir="/private/cache")
|
||||
cls = type(cache._CLASSES[component], (torch.nn.Linear,), {})
|
||||
calls = []
|
||||
|
||||
class Loader:
|
||||
def load(self, path, runtime):
|
||||
model = cache.attach_cached_model(cls, {"width": 2}, torch.float32, lambda: cls(2, 2))
|
||||
assert model is not None, "Would have reached disk hydration"
|
||||
return model
|
||||
|
||||
class Client:
|
||||
def __init__(self, path, key):
|
||||
assert str(path) == f"/private/cache/{component}.sock"
|
||||
self.key = key
|
||||
|
||||
@contextmanager
|
||||
def load_model(self, factory):
|
||||
with torch.device("meta"):
|
||||
model = factory()
|
||||
calls.append("attach")
|
||||
try:
|
||||
yield SimpleNamespace(model=model, donor_pid=123, handles={"weight": "handle"})
|
||||
finally:
|
||||
calls.append("release")
|
||||
|
||||
monkeypatch.setattr(ComponentLoader, "for_module_type", lambda *a: Loader())
|
||||
monkeypatch.setattr(cache, "WeightCacheClient", Client)
|
||||
leases = ExitStack()
|
||||
model = PipelineComponentLoader.load_module(component, "unused", "diffusers", runtime,
|
||||
weight_cache_leases=leases)
|
||||
assert calls == ["attach"]
|
||||
assert model._fastvideo_weight_cache_key.component == component
|
||||
assert model._fastvideo_weight_cache_witness["donor_pid"] == 123
|
||||
leases.close()
|
||||
assert calls == ["attach", "release"]
|
||||
|
||||
|
||||
def test_key_changes_for_dtype_configuration_and_checkpoint():
|
||||
cls = type("AutoencoderKLWan", (torch.nn.Linear,), {})
|
||||
def key(config, dtype=torch.float32, checkpoint="revision"):
|
||||
with cache.component_weight_cache("vae", args(weight_cache_checkpoint=checkpoint), None) as context:
|
||||
assert cache.attach_cached_model(cls, config, dtype, lambda: None) is None
|
||||
return context.key
|
||||
baseline = key({"width": 2})
|
||||
assert baseline == key({"width": 2})
|
||||
assert baseline != key({"width": 3})
|
||||
assert baseline != key({"width": 2}, torch.bfloat16)
|
||||
assert baseline != key({"width": 2}, checkpoint="other")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("method", ["set_lora_adapter", "merge_lora_weights", "unmerge_lora_weights"])
|
||||
def test_runtime_lora_apis_cannot_modify_shared_weights(method):
|
||||
from fastvideo.pipelines.lora_pipeline import LoRAPipeline
|
||||
pipeline = SimpleNamespace(fastvideo_args=args(weight_cache_dir="unused"))
|
||||
kwargs = {"lora_nickname": "test"} if method == "set_lora_adapter" else {}
|
||||
with pytest.raises(ValueError, match="cannot mutate"):
|
||||
getattr(LoRAPipeline, method)(pipeline, **kwargs)
|
||||
|
||||
|
||||
def test_pipeline_close_releases_leases_and_is_idempotent():
|
||||
from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
|
||||
closed = []
|
||||
leases = ExitStack()
|
||||
leases.callback(lambda: closed.append(True))
|
||||
pipeline = SimpleNamespace(_weight_cache_leases=leases, _trace_mgr=None)
|
||||
ComposedPipelineBase.close(pipeline)
|
||||
ComposedPipelineBase.close(pipeline)
|
||||
assert closed == [True]
|
||||
|
||||
|
||||
def test_partial_pipeline_load_failure_releases_existing_attachments(monkeypatch):
|
||||
from contextlib import nullcontext
|
||||
from fastvideo.pipelines import composed_pipeline_base as base
|
||||
released = []
|
||||
pipeline = SimpleNamespace()
|
||||
def failing_load(*args):
|
||||
pipeline._weight_cache_leases.callback(lambda: released.append("first-component"))
|
||||
raise ValueError("Later component incompatible")
|
||||
pipeline.load_modules = failing_load
|
||||
monkeypatch.setattr(base, "maybe_init_distributed_environment_and_model_parallel", lambda *a: None)
|
||||
monkeypatch.setattr(base, "get_world_group", lambda: SimpleNamespace(local_rank=0))
|
||||
monkeypatch.setattr(base, "get_or_create_profiler", lambda *a: SimpleNamespace(region=lambda *a: nullcontext()))
|
||||
runtime = SimpleNamespace(inference_mode=False, tp_size=1, sp_size=1)
|
||||
with pytest.raises(ValueError, match="Later component"):
|
||||
base.ComposedPipelineBase.__init__(pipeline, "unused", runtime, required_config_modules=[])
|
||||
assert released == ["first-component"]
|
||||
|
||||
|
||||
def test_unsupported_vae_special_constructor_rejected_before_checkpoint_access(monkeypatch):
|
||||
from fastvideo.models.loader import component_loader
|
||||
runtime = args(weight_cache_dir="unused")
|
||||
monkeypatch.setattr(component_loader, "get_diffusers_config",
|
||||
lambda **kwargs: {"_class_name": "LingBotWorld2WanVAE"})
|
||||
with ExitStack() as leases, pytest.raises(ValueError, match="Unsupported cached vae class"):
|
||||
PipelineComponentLoader.load_module("vae", "unused", "diffusers", runtime,
|
||||
weight_cache_leases=leases)
|
||||
@@ -1,64 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
import pytest
|
||||
|
||||
from fastvideo.models.loader.startup import collect_startup_timings, startup_phase
|
||||
|
||||
|
||||
def test_nested_intervals_are_inclusive_and_serializable():
|
||||
import json
|
||||
syncs = []
|
||||
with collect_startup_timings(synchronize=lambda: syncs.append(True)) as timings:
|
||||
with startup_phase("load"):
|
||||
with startup_phase("read"):
|
||||
pass
|
||||
result = timings.to_dict()
|
||||
outer, inner = result["phases"]
|
||||
assert outer["parent_id"] is None
|
||||
assert inner["parent_id"] == outer["id"]
|
||||
assert outer["duration_seconds"] >= inner["duration_seconds"] >= 0
|
||||
assert result["recording_seconds"] >= outer["duration_seconds"]
|
||||
assert result["timing_mode"] == "synchronized_boundaries"
|
||||
assert len(syncs) == 4
|
||||
json.dumps(result)
|
||||
|
||||
|
||||
def test_error_is_recorded_and_collector_is_restored():
|
||||
with collect_startup_timings() as outer:
|
||||
with pytest.raises(ValueError, match="checkpoint"):
|
||||
with collect_startup_timings() as failed:
|
||||
with startup_phase("read"):
|
||||
raise ValueError("checkpoint")
|
||||
with startup_phase("recovery"):
|
||||
pass
|
||||
assert [p["name"] for p in outer.phases] == ["recovery"]
|
||||
assert failed.phases[0]["status"] == "error"
|
||||
assert outer.phases[0]["parent_id"] is None
|
||||
with startup_phase("disabled"):
|
||||
pass
|
||||
assert len(outer.phases) == 1
|
||||
|
||||
|
||||
def test_disabled_collection_does_not_read_the_clock(monkeypatch):
|
||||
from fastvideo.models.loader import startup
|
||||
def unexpected():
|
||||
raise AssertionError("disabled collection measured work")
|
||||
monkeypatch.setattr(startup.time, "perf_counter", unexpected)
|
||||
with startup_phase("disabled"):
|
||||
pass
|
||||
|
||||
|
||||
def test_sync_failure_does_not_leave_a_stale_parent():
|
||||
calls = 0
|
||||
def synchronize():
|
||||
nonlocal calls
|
||||
calls += 1
|
||||
if calls == 2:
|
||||
raise RuntimeError("CUDA failure")
|
||||
with collect_startup_timings(synchronize=synchronize) as timings:
|
||||
with pytest.raises(RuntimeError, match="CUDA failure"):
|
||||
with startup_phase("failed"):
|
||||
pass
|
||||
with startup_phase("next"):
|
||||
pass
|
||||
assert timings.phases[0]["status"] == "error"
|
||||
assert timings.phases[1]["parent_id"] is None
|
||||
@@ -8,7 +8,6 @@ from safetensors.torch import save_file
|
||||
|
||||
from fastvideo.models.loader import component_loader
|
||||
from fastvideo.models.loader.component_loader import PipelineComponentLoader, VAELoader
|
||||
from fastvideo.models.loader.startup import collect_startup_timings
|
||||
from fastvideo.models.loader.utils import set_default_torch_dtype
|
||||
|
||||
|
||||
@@ -40,8 +39,7 @@ def test_native_load_matches_separate_construction_and_hydration(tmp_path, monke
|
||||
else:
|
||||
save_file(tensors, str(tmp_path / "model.safetensors"))
|
||||
loader = VAELoader()
|
||||
with collect_startup_timings() as timings:
|
||||
combined = PipelineComponentLoader.load_module("vae", str(tmp_path), "diffusers", args)
|
||||
combined = PipelineComponentLoader.load_module("vae", str(tmp_path), "diffusers", args)
|
||||
with set_default_torch_dtype(torch.float32):
|
||||
separate = loader.build_model(TinyVAE, config, torch.device("cpu"))
|
||||
assert torch.count_nonzero(separate.weight) == 0
|
||||
@@ -53,13 +51,6 @@ def test_native_load_matches_separate_construction_and_hydration(tmp_path, monke
|
||||
assert not combined.training and not separate.training
|
||||
for name, tensor in separate.state_dict().items():
|
||||
torch.testing.assert_close(tensor, tensors[name], rtol=0, atol=0)
|
||||
assert {p["name"] for p in timings.phases} == {
|
||||
"component.vae", "vae.construct", "vae.device_transfer", "vae.checkpoint_read", "vae.weight_assignment",
|
||||
"vae.finalize",
|
||||
}
|
||||
|
||||
root_phase = next(p for p in timings.phases if p["name"] == "component.vae")
|
||||
assert all(p["parent_id"] == root_phase["id"] for p in timings.phases if p is not root_phase)
|
||||
|
||||
|
||||
def test_construction_does_not_read_checkpoint(monkeypatch):
|
||||
|
||||
@@ -196,8 +196,7 @@ class Worker:
|
||||
logger.info("Worker %d shutting down...", self.rank, local_main_process_only=False)
|
||||
# Clean up resources
|
||||
if hasattr(self, 'pipeline') and self.pipeline is not None:
|
||||
# Clean up pipeline resources if needed
|
||||
pass
|
||||
self.pipeline.close()
|
||||
|
||||
# Destroy the distributed environment
|
||||
cleanup_dist_env_and_memory(shutdown_ray=False)
|
||||
|
||||
@@ -255,7 +255,6 @@ nav:
|
||||
- Testing: contributing/testing.md
|
||||
- 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
|
||||
|
||||
Reference in New Issue
Block a user