diff --git a/docs/design/minimax_h3/COLD_LOADING.md b/docs/design/minimax_h3/COLD_LOADING.md new file mode 100644 index 0000000000..65a4b4cf57 --- /dev/null +++ b/docs/design/minimax_h3/COLD_LOADING.md @@ -0,0 +1,84 @@ +# H3 cold loading and reusable prepared weights + +## Scope and baseline + +Integration base: `fe67339ddf862df0441a4e844fc664d9cafdffd0` (`onecat/main`). +This scope covers original H3, INT8 ConvRot and LightX2V/FastH3 adapters; +checkpoint precision, attention, sampling and output semantics stay unchanged. + +The deployed four-card V100 host has 62 GiB RAM and uses disk-backed masters. +A real original FastH3 request spent 973.19 seconds preparing the service and +104.95 seconds generating. At 604 seconds into loading, worker disk counters +already totalled 99.57 GiB read and 109.08 GiB written. The old path loads row +storage, transforms it to column storage and snapshots it again. Temporary +masters are reconstructed at each service launch. Shared text encoder loading +also repeats across H3 variants. These observations are not a new benchmark. + +## Intended contract + +- Prepare final projection strides before loading, avoiding a second full copy. +- Reuse exact per-component, per-TP-rank prepared CPU storage, including the + shared text encoder, without copying a second complete model for publication. +- Key entries by source/checkpoint identity, precision, topology and relevant + adapter/configuration. Retain original model files and all existing variants. +- Publish atomically only after loading/validation succeeds; validate metadata + and data on restore. Use private mappings to prevent writes changing cache. +- Bound disk use and evict only inactive cache entries under leases. Corrupt or + incomplete entries must rebuild, never silently serve unverified weights. +- Report preparation subphases and distinguish first construction from reuse. + +## Validation status + +Implementation and targeted tests are in progress. No new startup speed or +quality claim is made. GPU validation waits for the user's current generation +to finish; no active user task is interrupted. Real first-load/reload timing +and output comparison are required before promotion. + +AI assistance: OpenAI Codex. + +## Candidate implementation + +`--prepared-weight-cache` enables exact rank-local prepared transformer and +text-encoder storage under `VLLM_CACHE_ROOT/h3-prepared`. The default budget is +128 GiB (`--prepared-weight-cache-gib`). Native callers remain opt-in. Studio +uses the capability probe to enable it independently of generation fast mode. + +The initial fill allocates the final projection strides before loading. A +completed entry includes checksums of all storage files, tensor metadata and +its manifest. Subsequent loads validate it before binding private mappings. +Keys include checkpoint identity, source, Torch/CUDA, processing configuration +and TP rank/topology. Ordinary LoRA sidecars are applied after the base cache; +FastH3 fusion is keyed by its exact adapter. Active entries are protected by +lifetime leases. Only inactive prepared entries can be evicted; original model +files are never removed. Missing, incomplete or corrupt entries are rebuilt. + +Loading reports actual tensor or byte counts separately from component counts. +No overall time percentage or performance prediction is synthesized. + +CPU validation: 80 passed, 3 skipped across prepared weights, host residency, +service, progress, FastH3 and Studio fast-path suites. These cover exact signed +INT8 bytes, FP16 layout, shared storage offsets, cache corruption and active +lease protection. Real GPU startup, output parity and speed remain pending; +this candidate must not be promoted based only on the CPU result. + +## Video VAE loading + +The native video component is a 2.60-billion-element FP32 checkpoint. Its +reference factory initializes parameters and then replaces them with a strict +state-dict load. Prepared loading skips random fills only for Parameters that +are subsequently covered by a successful checkpoint load. Other tensors and +buffers keep their normal initialization; incomplete/custom loaders retry +normally. Temporary initialization hooks are always restored. + +A complete CPU state dict can also be assigned without copying when dtype, +shape, strides and storage offsets match and neither side has tied storage. +Custom parameter metadata, non-strided tensors, explicit assignment options, +precision conversions and incomplete dictionaries retain normal semantics. +The checkpoint's private mappings do not allow writes to change the file. +Audio VAE loading is unchanged. + +An isolated V100-host CPU check of the actual video component produced the +same full parameter-and-buffer SHA256 with and without this optimization: +`3145d6a4c1b6576045991a36cf7b3606ec327c0c36d4152b91cf4498e3a9dfed`. +The final loader selected assignment and took 1.67 seconds in that component +check. This excludes downstream staging and is not a full-model startup time. diff --git a/tests/video/test_h3_host_memory.py b/tests/video/test_h3_host_memory.py index c6065698c4..f084b00711 100644 --- a/tests/video/test_h3_host_memory.py +++ b/tests/video/test_h3_host_memory.py @@ -48,7 +48,7 @@ def test_both_vae_stagers_honor_the_host_policy(monkeypatch, pin_memory, kind): monkeypatch.setattr( vae, "_load_component_config", lambda path: {"sample_rate": 44100} ) - monkeypatch.setattr(vae, "_load_remote_component", lambda *args: remote) + monkeypatch.setattr(vae, "_load_remote_component", lambda *args, **kwargs: remote) monkeypatch.setattr( vae, "PinnedModuleStager", diff --git a/tests/video/test_h3_initialization.py b/tests/video/test_h3_initialization.py new file mode 100644 index 0000000000..bb9ed1bf06 --- /dev/null +++ b/tests/video/test_h3_initialization.py @@ -0,0 +1,161 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import pytest +import torch +from torch import nn + +from vllm.model_executor.models.minimax_h3.initialization import ( + load_without_random_parameter_init, +) + + +def test_only_checkpoint_replaced_parameters_skip_random_fills(): + original = nn.init.kaiming_uniform_ + state = {"weight": torch.ones(4, 3), "bias": torch.full((4,), 2.0)} + + def factory(): + model = nn.Linear(3, 4) + constant = torch.empty(5) + nn.init.uniform_(constant, 3, 3) + model.register_buffer("constant", constant, persistent=False) + model.load_state_dict(state, strict=True) + return model + + model = load_without_random_parameter_init(factory) + assert nn.init.kaiming_uniform_ is original + assert torch.equal(model.weight, state["weight"]) + assert torch.equal(model.constant, torch.full((5,), 3.0)) + assert torch.equal(model(torch.ones(1, 3)), torch.full((1, 4), 5.0)) + + +def test_partial_checkpoint_retries_with_normal_initialization(): + calls = [] + + def factory(): + calls.append(True) + model = nn.Linear(3, 4) + model.load_state_dict({"weight": torch.ones(4, 3)}, strict=False) + return model + + model = load_without_random_parameter_init(factory) + assert len(calls) == 2 + assert torch.isfinite(model.bias).all() + assert torch.all(model.bias.abs() <= 1 / 3**0.5) + + +def test_initializer_and_load_hooks_are_restored_on_error(): + initialize, load = nn.init.uniform_, nn.Module.load_state_dict + + def factory(): + nn.Linear(3, 4).load_state_dict({}, strict=True) + + with pytest.raises(RuntimeError): + load_without_random_parameter_init(factory) + assert nn.init.uniform_ is initialize + assert nn.Module.load_state_dict is load + + +def test_assigning_checkpoint_storage_is_supported(): + def factory(): + model = nn.Linear(3, 4, bias=False) + model.load_state_dict({"weight": torch.ones(4, 3)}, assign=True) + return model + + model = load_without_random_parameter_init(factory) + assert torch.equal(model.weight, torch.ones(4, 3)) + + +def test_matching_checkpoint_is_assigned_without_a_second_copy(tmp_path): + from safetensors.torch import load_file, save_file + + path = tmp_path / "weights.safetensors" + save_file({"weight": torch.arange(12, dtype=torch.float32).reshape(4, 3)}, path) + original = path.read_bytes() + state = load_file(path) + + def factory(): + model = nn.Linear(3, 4, bias=False) + model.load_state_dict(state) + return model + + model = load_without_random_parameter_init(factory) + assert model.weight.data_ptr() == state["weight"].data_ptr() + with torch.no_grad(): + model.weight.add_(1) + assert path.read_bytes() == original # safetensors CPU mappings are private. + + +def test_dtype_conversion_keeps_original_copy_semantics(): + state = {"weight": torch.arange(12, dtype=torch.float16).reshape(4, 3)} + + def factory(): + model = nn.Linear(3, 4, bias=False, dtype=torch.float32) + model.load_state_dict(state) + return model + + model = load_without_random_parameter_init(factory) + assert model.weight.dtype == torch.float32 + assert model.weight.data_ptr() != state["weight"].data_ptr() + assert torch.equal(model.weight, state["weight"].float()) + + +def test_tied_parameters_are_not_detached_by_assignment(): + def factory(): + model = nn.Module() + model.first = nn.Linear(3, 4, bias=False) + model.second = nn.Linear(3, 4, bias=False) + model.second.weight = model.first.weight + model.load_state_dict( + {"first.weight": torch.ones(4, 3), "second.weight": torch.full((4, 3), 2.0)} + ) + return model + + model = load_without_random_parameter_init(factory) + assert model.first.weight is model.second.weight + assert torch.equal(model.first.weight, torch.full((4, 3), 2.0)) + + +def test_assignment_does_not_introduce_new_parameter_aliases(): + def factory(): + model = nn.Sequential(nn.Linear(3, 4, bias=False), nn.Linear(3, 4, bias=False)) + weight = torch.ones(4, 3) + model.load_state_dict({"0.weight": weight, "1.weight": weight}) + return model + + model = load_without_random_parameter_init(factory) + assert model[0].weight.data_ptr() != model[1].weight.data_ptr() + + +def test_storage_marker_expires_after_parameter_conversion(): + from vllm.model_executor.models.minimax_h3.initialization import ( + uses_assigned_checkpoint_storage, + ) + + def factory(): + model = nn.Linear(3, 4, bias=False) + model.load_state_dict({"weight": torch.ones(4, 3)}) + return model + + model = load_without_random_parameter_init(factory) + assert uses_assigned_checkpoint_storage(model) + model.double() + assert not uses_assigned_checkpoint_storage(model) + + +def test_replica_identity_requires_all_workers_and_rejects_replacement(monkeypatch): + from vllm.model_executor.models.minimax_h3 import vae + + monkeypatch.setattr(vae.dist, "is_initialized", lambda: True) + monkeypatch.setattr(vae.dist, "get_world_size", lambda: 2) + values: list[tuple[str, int] | None] = [("same", 1), ("same", 1)] + monkeypatch.setattr( + vae.dist, + "all_gather_object", + lambda result, value: result.__setitem__(slice(None), values), + ) + assert vae._same_checkpoint_replicas(values[0]) + values[1] = None + assert not vae._same_checkpoint_replicas(values[0]) + values[1] = ("replacement", 2) + with pytest.raises(ValueError, match="changed between"): + vae._same_checkpoint_replicas(values[0]) diff --git a/tests/video/test_h3_prepared_weights.py b/tests/video/test_h3_prepared_weights.py new file mode 100644 index 0000000000..69102bf2f6 --- /dev/null +++ b/tests/video/test_h3_prepared_weights.py @@ -0,0 +1,214 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import hashlib +import json + +import pytest +import torch +from torch import nn + +from vllm.model_executor.models.minimax_h3.prepared_weights import ( + PreparedWeights, + checkpoint_identity, +) +from vllm.model_executor.models.minimax_h3.quantization import FP16LinearMethod +from vllm.model_executor.models.minimax_h3.residency import PinnedModuleStager + + +def model(): + module = nn.Module() + module.register_parameter( + "weight", nn.Parameter(torch.empty_strided((3, 4), (1, 3))) + ) + module.register_buffer("scale", torch.ones(3)) + return module + + +def build(root, key="a" * 64, *, limit=4096): + module = model() + cache = PreparedWeights(root, key, module, limit_bytes=limit, reserve_bytes=0) + assert not cache.restore(module) + PinnedModuleStager.map_cpu_weights(module, cache, preserve_parameters=False) + with torch.no_grad(): + module.weight.copy_(torch.arange(12).reshape(3, 4)) + cache.publish(module) + return module, cache + + +def test_reuses_final_strides_and_protects_cache_from_adapter_writes(tmp_path): + module, cache = build(tmp_path) + expected = module.weight.detach().clone() + assert module.weight.stride() == (1, 3) + with torch.no_grad(): + module.weight.add_(100) + cache.close() + restored = model() + reader = PreparedWeights(tmp_path, "a" * 64, restored, reserve_bytes=0) + assert reader.restore(restored) + torch.testing.assert_close(restored.weight, expected, rtol=0, atol=0) + assert restored.weight.stride() == (1, 3) + pointer = restored.weight.data_ptr() + PinnedModuleStager.map_cpu_weights(restored, reader) + assert restored.weight.data_ptr() == pointer + assert not list(reader.fallback.directory.iterdir()) + reader.close() + + +@pytest.mark.parametrize( + "damage", ["data", "metadata", "missing", "truncated", "unfinished"] +) +def test_invalid_entries_rebuild_without_partially_binding(tmp_path, damage): + _, cache = build(tmp_path) + cache.close() + manifest = json.loads((cache.entry / "manifest.json").read_bytes()) + path = cache.entry / manifest["groups"][0]["file"] + if damage == "data": + data = bytearray(path.read_bytes()) + data[0] ^= 0x80 + path.write_bytes(data) + elif damage == "metadata": + (cache.entry / "manifest.json").write_text("{}") + elif damage == "missing": + path.unlink() + elif damage == "truncated": + path.write_bytes(b"") + else: + (cache.entry / "ready.json").unlink() + target = model() + target.weight.data.fill_(7) + before = target.weight.clone() + reader = PreparedWeights(tmp_path, "a" * 64, target, reserve_bytes=0) + assert not reader.restore(target) + torch.testing.assert_close(target.weight, before, rtol=0, atol=0) + assert not (reader.entry / "ready.json").exists() + reader.close() + + +def test_active_cache_is_not_evicted_and_inactive_cache_can_be_reclaimed(tmp_path): + _, first = build(tmp_path, limit=80) + target = model() + second = PreparedWeights( + tmp_path, "b" * 64, target, limit_bytes=80, reserve_bytes=0 + ) + assert not second.restore(target) + with pytest.raises(RuntimeError, match="active caches"): + PinnedModuleStager.map_cpu_weights(target, second, preserve_parameters=False) + assert (first.entry / "ready.json").exists() + first.close() + PinnedModuleStager.map_cpu_weights(target, second, preserve_parameters=False) + assert not first.entry.exists() + second.close() + + +def test_checkpoint_replacement_changes_identity(tmp_path): + path = tmp_path / "weights.safetensors" + path.write_bytes(b"original") + before = checkpoint_identity([path]) + replacement = tmp_path / "replacement" + replacement.write_bytes(b"original") + replacement.replace(path) + assert checkpoint_identity([path]) != before + + +def test_corrupt_binding_cannot_escape_storage_even_with_valid_manifest_hash(tmp_path): + _, cache = build(tmp_path) + cache.close() + path = cache.entry / "manifest.json" + manifest = json.loads(path.read_bytes()) + manifest["groups"][0]["bindings"][0]["offset"] = 2**60 + data = json.dumps(manifest).encode() + path.write_bytes(data) + (cache.entry / "ready.json").write_text( + json.dumps({"sha256": hashlib.sha256(data).hexdigest()}) + ) + reader = PreparedWeights(tmp_path, "a" * 64, model(), reserve_bytes=0) + assert not reader.restore(model()) + reader.close() + + +def test_final_layout_can_be_filled_directly_without_second_allocation(): + module = nn.Linear(7, 5, bias=False, dtype=torch.float16) + module.h3_fp16_weight_layout = "column" + method = FP16LinearMethod() + method.prepare_weights_before_loading(module) + source = torch.arange(35, dtype=torch.bfloat16).reshape(5, 7) + with torch.no_grad(): + module.weight.copy_(source) + pointer = module.weight.data_ptr() + method.process_weights_after_loading(module) + assert module.weight.data_ptr() == pointer + assert module.weight.stride() == (1, 5) + torch.testing.assert_close(module.weight, source.to(torch.float16), rtol=0, atol=0) + + +def test_identity_separates_rank_topology_and_processing_recipe(tmp_path): + from vllm.model_executor.models.minimax_h3.prepared_weights import preparation_key + + path = tmp_path / "weights.safetensors" + path.write_bytes(b"test checkpoint identity") + kwargs = dict( + component="transformer", rank=0, world_size=4, options={"layout": "column"} + ) + baseline = preparation_key([path], **kwargs) + assert baseline == preparation_key([path], **kwargs) + for change in ( + {"rank": 1}, + {"world_size": 2}, + {"component": "text_encoder"}, + {"options": {"layout": "row"}}, + ): + assert preparation_key([path], **(kwargs | change)) != baseline + + +def test_mixed_precision_shared_storage_and_offset_survive_restore(tmp_path): + def mixed(): + module = nn.Module() + raw = torch.arange(24, dtype=torch.float32) + module.register_buffer("first", raw[:12].reshape(3, 4)) + module.register_buffer("second", raw[12:].reshape(3, 4)) + module.register_buffer("quantized", torch.arange(-12, 12, dtype=torch.int8)) + module.register_buffer("half_values", torch.arange(7, dtype=torch.float16)) + return module + + source = mixed() + cache = PreparedWeights(tmp_path, "c" * 64, source, reserve_bytes=0) + assert not cache.restore(source) + cache.publish(source) + cache.close() + target = mixed() + reader = PreparedWeights(tmp_path, "c" * 64, target, reserve_bytes=0) + assert reader.restore(target) + for name, tensor in source.named_buffers(): + torch.testing.assert_close(getattr(target, name), tensor, rtol=0, atol=0) + assert ( + target.first.untyped_storage().data_ptr() + == target.second.untyped_storage().data_ptr() + ) + assert target.second.storage_offset() == 12 + reader.close() + + +def test_int8_final_layout_preserves_signed_bytes_and_scale_without_copy(): + from types import SimpleNamespace + + from vllm.model_executor.models.minimax_h3.quantization import ( + Int8ConvRotLinearMethod, + ) + + module = nn.Module() + module.register_parameter( + "weight", nn.Parameter(torch.empty(5, 7, dtype=torch.int8), requires_grad=False) + ) + module.register_parameter("weight_scale", nn.Parameter(torch.ones(5, 1))) + method = Int8ConvRotLinearMethod( + SimpleNamespace(weight_layout="column"), None, prefix="blocks.0.proj" + ) + method.prepare_weights_before_loading(module) + source = torch.arange(-17, 18, dtype=torch.int8).reshape(5, 7) + module.weight.copy_(source) + pointer = module.weight.data_ptr() + method.process_weights_after_loading(module) + assert module.weight.data_ptr() == pointer + assert module.weight.stride() == (1, 5) + assert torch.equal(module.weight, source) + assert torch.equal(module.weight_scale, torch.ones(5)) diff --git a/vllm/entrypoints/cli/video.py b/vllm/entrypoints/cli/video.py index 8dcd9cba12..3e34876315 100644 --- a/vllm/entrypoints/cli/video.py +++ b/vllm/entrypoints/cli/video.py @@ -98,6 +98,12 @@ def subparser_init(self, subparsers): help="Use disk-backed CPU weights on hosts below 128 GiB RAM", ) mode.add_argument("--host-memory-directory") + mode.add_argument( + "--prepared-weight-cache", + action="store_true", + help="Reuse checksummed, leased H3 CPU weights between service loads", + ) + mode.add_argument("--prepared-weight-cache-gib", type=float, default=128.0) mode.add_argument( "--video-encoder", choices=("libx264", "h264_nvenc"), @@ -169,6 +175,8 @@ def cmd(args): weight_offload=args.weight_offload, host_memory_mode=args.host_memory_mode, host_memory_directory=args.host_memory_directory, + prepared_weight_cache=args.prepared_weight_cache, + prepared_weight_cache_gib=args.prepared_weight_cache_gib, ) if args.video_mode == "serve": from vllm.video.server import serve diff --git a/vllm/media/progress.py b/vllm/media/progress.py index 69f54f1e6e..15af43901d 100644 --- a/vllm/media/progress.py +++ b/vllm/media/progress.py @@ -13,7 +13,7 @@ ) -def report_loading(completed, total, component, *, rank=0, world_size=1): +def report_loading(completed, total, component, *, rank=0, world_size=1, detail=None): """Startup progress travels via the existing owned process log.""" import json @@ -28,12 +28,58 @@ def report_loading(completed, total, component, *, rank=0, world_size=1): "unit": "components", "rank": rank, "world_size": world_size, + **({"detail": detail} if detail is not None else {}), } ), flush=True, ) +class LoadingProgress: + """Bounded-rate preparation detail; does not synchronize a device.""" + + def __init__(self, component, component_index, rank, world_size): + self.component = component + self.index = component_index + self.rank = rank + self.world_size = world_size + self.phase = None + self.started_at = 0.0 + self.last_report = 0.0 + + def __call__(self, phase, completed=0, total=None, unit="items"): + import time + + now = time.time() + changed = phase != self.phase + if changed: + self.phase, self.started_at = phase, now + if not changed and now - self.last_report < 1 and completed != total: + return + self.last_report = now + report_loading( + self.index, + 4, + self.component, + rank=self.rank, + world_size=self.world_size, + detail={ + "phase": phase, + "completed": completed, + "total": total, + "unit": unit, + "started_at": self.started_at, + "updated_at": now, + }, + ) + + def weights(self, values, total): + self("reading_weights", 0, total, "tensors") + for index, value in enumerate(values, 1): + yield value + self("reading_weights", index, total, "tensors") + + @contextmanager def reporting(callback: ProgressCallback | None) -> Iterator[None]: token = _callback.set(callback) diff --git a/vllm/model_executor/models/minimax_h3/config.py b/vllm/model_executor/models/minimax_h3/config.py index 3708331a81..4437e676df 100644 --- a/vllm/model_executor/models/minimax_h3/config.py +++ b/vllm/model_executor/models/minimax_h3/config.py @@ -54,8 +54,18 @@ class H3Config: video_encoder: Literal["libx264", "h264_nvenc"] = "libx264" host_memory_mode: Literal["auto", "pinned", "mmap"] = "auto" host_memory_directory: str | None = None + prepared_weight_cache: bool = False + prepared_weight_cache_gib: float = 128.0 def __post_init__(self) -> None: + if not isinstance(self.prepared_weight_cache, bool): + raise H3InputError("prepared weight cache must be a boolean") + if ( + isinstance(self.prepared_weight_cache_gib, bool) + or not math.isfinite(self.prepared_weight_cache_gib) + or not 0 < self.prepared_weight_cache_gib <= 1024 * 1024 + ): + raise H3InputError("prepared weight cache size must be finite and positive") if self.residual_reduction not in ("native", "peer"): raise H3InputError("residual reduction must be native or peer") if self.residual_reduction == "peer" and not self.residual_sequence_parallel: diff --git a/vllm/model_executor/models/minimax_h3/initialization.py b/vllm/model_executor/models/minimax_h3/initialization.py new file mode 100644 index 0000000000..f18baae567 --- /dev/null +++ b/vllm/model_executor/models/minimax_h3/initialization.py @@ -0,0 +1,150 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Avoid random parameter fills that a strict checkpoint load replaces.""" + +import weakref +from functools import wraps +from threading import RLock + +from torch import Tensor, nn, strided + +from vllm.logger import init_logger + +logger = init_logger(__name__) +_LOCK = RLock() +_RANDOM_INITIALIZERS = ( + "uniform_", + "normal_", + "trunc_normal_", + "kaiming_uniform_", + "kaiming_normal_", + "xavier_uniform_", + "xavier_normal_", + "orthogonal_", +) + + +def _can_assign_checkpoint(module, state_dict): + targets = module.state_dict(keep_vars=True) + if targets.keys() != state_dict.keys(): + return False + storages, source_storages = set(), set() + for name, target in targets.items(): + value = state_dict[name] + if ( + not isinstance(target, Tensor) + or not isinstance(value, Tensor) + or type(target) not in (Tensor, nn.Parameter) + or target.__dict__ + or target.layout != strided + or value.layout != strided + or target.device.type != "cpu" + or value.device.type != "cpu" + or target.dtype != value.dtype + or target.shape != value.shape + or target.stride() != value.stride() + or target.storage_offset() != value.storage_offset() + ): + return False + storage = target.untyped_storage() + if storage.nbytes(): + identity = storage.data_ptr() + if identity in storages: + return False # Preserve tied parameters and storage aliases. + storages.add(identity) + source_identity = value.untyped_storage().data_ptr() + if source_identity in source_storages: + return False + source_storages.add(source_identity) + return True + + +def load_without_random_parameter_init(factory): + """For the isolated H3 loading worker, skip only replaced Parameters. + + Buffers and ordinary tensors still receive their normal initialization. + Complete CPU checkpoints with matching dtypes/layouts and no target aliases + can be assigned directly, avoiding another full per-worker weight copy. + Track actual successful state-dict loads, rather than trusting the factory + to load every parameter. Unsupported/custom partial loaders retry normally. + The process-local patch is restored before returning or propagating errors. + """ + skipped, loaded, assigned = {}, {}, {} + with _LOCK: + originals = {name: getattr(nn.init, name) for name in _RANDOM_INITIALIZERS} + load_state_dict = nn.Module.load_state_dict + + def wrap_initializer(original): + @wraps(original) + def initialize(tensor, *args, **kwargs): + if isinstance(tensor, nn.Parameter): + skipped[id(tensor)] = weakref.ref(tensor) + return tensor + return original(tensor, *args, **kwargs) + + return initialize + + @wraps(load_state_dict) + def load(module, state_dict, *args, **kwargs): + if ( + len(args) < 2 + and "assign" not in kwargs + and _can_assign_checkpoint(module, state_dict) + ): + kwargs["assign"] = True + result = load_state_dict(module, state_dict, *args, **kwargs) + missing = set(result.missing_keys) + for name, parameter in module.named_parameters(): + if name in state_dict and name not in missing: + loaded[id(parameter)] = weakref.ref(parameter) + if kwargs.get("assign"): + assigned[id(parameter)] = weakref.ref(parameter) + return result + + try: + for name, original in originals.items(): + setattr(nn.init, name, wrap_initializer(original)) + nn.Module.load_state_dict = load + model = factory() + finally: + nn.Module.load_state_dict = load_state_dict + for name, original in originals.items(): + setattr(nn.init, name, original) + + uncovered = [ + name + for name, parameter in model.named_parameters() + if id(parameter) in skipped + and skipped[id(parameter)]() is parameter + and ( + id(parameter) not in loaded or loaded[id(parameter)]() is not parameter + ) + ] + if not uncovered: + parameters = list(model.parameters()) + if parameters and all( + id(p) in assigned and assigned[id(p)]() is p for p in parameters + ): + model._h3_assigned_checkpoint_storage = { + name: (p.data_ptr(), p.dtype, p.device, tuple(p.shape), p.stride()) + for name, p in model.named_parameters() + } + return model + logger.warning( + "VAE loader did not replace all skipped parameters; " + "retrying with normal initialization (%s)", + uncovered[:5], + ) + del model + return factory() + + +def uses_assigned_checkpoint_storage(model): + expected = getattr(model, "_h3_assigned_checkpoint_storage", None) + if not expected: + return False + actual = { + name: (p.data_ptr(), p.dtype, p.device, tuple(p.shape), p.stride()) + for name, p in model.named_parameters() + } + return expected == actual diff --git a/vllm/model_executor/models/minimax_h3/pipeline.py b/vllm/model_executor/models/minimax_h3/pipeline.py index 211ba2ce73..799d3ac93b 100644 --- a/vllm/model_executor/models/minimax_h3/pipeline.py +++ b/vllm/model_executor/models/minimax_h3/pipeline.py @@ -26,6 +26,7 @@ from vllm import envs from vllm.distributed import get_tp_group, get_world_group from vllm.logger import init_logger +from vllm.platforms import current_platform from vllm.utils.mem_utils import get_cpu_memory from vllm.video.metrics import DenoiseWorkCounter @@ -95,7 +96,11 @@ from .vae import MiniMaxH3AudioVAE, MiniMaxH3VideoVAE from .vsa import h3_vsa_workspace from .weight_cache import FP16WeightCache -from .weights import iter_checkpoint_weights, resolve_model_root +from .weights import ( + checkpoint_tensor_count, + iter_checkpoint_weights, + resolve_model_root, +) logger = init_logger(__name__) @@ -439,7 +444,7 @@ def __init__(self, config: H3Config, *, shared_weights_dir: str | None = None): ) self.config = config self.partition = config.partition - from vllm.media.progress import report_loading + from vllm.media.progress import LoadingProgress, report_loading def loading(done, component): group = get_tp_group() @@ -552,13 +557,11 @@ def loading(done, component): for module in self.transformer.modules(): if isinstance(module, MiniMaxH3DiTBlock): module.residual_reducer = self._residual_reduction - if self._host_backing is not None: - PinnedModuleStager.map_cpu_weights( - self.transformer, self._host_backing, preserve_parameters=False - ) - weights = iter_checkpoint_weights(transformer_path) - if restore_adaln: - weights = restore_dense_adaln_weights(weights, path / "transformer") + group = get_tp_group() + self._prepared_caches = {} + dit_progress = LoadingProgress( + "transformer", 0, group.rank_in_group, group.world_size + ) fusion = None if isinstance(adapter_spec, FastH3Spec): if adapter_spec.requires_vsa: @@ -569,23 +572,86 @@ def loading(done, component): head_dim=self.transformer.arch.attention_head_dim, device=self.device, ) - weights = fusion.apply(weights) - loaded = self.transformer.load_weights(weights) - if fusion is not None: - fusion.validate_fully_applied(loaded) - required = set(dict(self.transformer.named_parameters())) - required.update(dict(self.transformer.named_buffers())) - missing = required - loaded - if missing: - raise RuntimeError(f"H3 DiT checkpoint missing tensors: {sorted(missing)}") - for layer in self.transformer.modules(): - method = getattr(layer, "quant_method", None) - if method is not None: + if config.prepared_weight_cache: + for layer in self.transformer.modules(): + method = getattr(layer, "quant_method", None) + layer.h3_fp16_weight_layout = config.fp16_weight_layout + prepare = getattr(method, "prepare_weights_before_loading", None) + if prepare is not None: + prepare(layer) + dit_cache = None + if config.prepared_weight_cache and not isinstance(adapter_spec, FlashGenSpec): + files = [ + transformer_path, + path / "model_index.json", + path / "transformer/config.json", + ] + if fusion is not None: + files.append(select_adapter_file(config.lora_path)) + dit_cache = self._prepared_component( + "transformer", + self.transformer, + files, + dit_progress, + options={ + "partition": config.partition, + "attention": config.attention_backend, + "int8_layout": config.int8_weight_layout, + "fp16_layout": config.fp16_weight_layout, + "residual_sp": config.residual_sequence_parallel, + "tf32": torch.backends.cuda.matmul.allow_tf32, + "architecture": current_platform.get_device_capability( + self.device.index + ), + }, + ) + dit_backing = dit_cache or self._host_backing + if dit_cache is None or not dit_cache.restore(self.transformer): + dit_progress("preparing_weight_storage") + if dit_backing is not None: + PinnedModuleStager.map_cpu_weights( + self.transformer, dit_backing, preserve_parameters=False + ) + weights = iter_checkpoint_weights(transformer_path) + total = checkpoint_tensor_count(transformer_path) + if restore_adaln: + weights = restore_dense_adaln_weights(weights, path / "transformer") + total = None + if fusion is not None: + weights = fusion.apply(weights) + if total is not None: + total += sum( + p.assigned is not None for p in fusion.patches.values() + ) + loaded = self.transformer.load_weights(dit_progress.weights(weights, total)) + if fusion is not None: + fusion.validate_fully_applied(loaded) + required = set(dict(self.transformer.named_parameters())) + required.update(dict(self.transformer.named_buffers())) + missing = required - loaded + if missing: + raise RuntimeError( + f"H3 DiT checkpoint missing tensors: {sorted(missing)}" + ) + layers = [ + layer + for layer in self.transformer.modules() + if getattr(layer, "quant_method", None) is not None + ] + dit_progress("preparing_weight_layout", 0, len(layers)) + for index, layer in enumerate(layers, 1): layer.h3_fp16_weight_layout = config.fp16_weight_layout - method.process_weights_after_loading(layer) - if self._host_backing is not None: - PinnedModuleStager.map_cpu_weights(layer, self._host_backing) - self.transformer.post_load_weights() + layer.quant_method.process_weights_after_loading(layer) + if dit_backing is not None: + PinnedModuleStager.map_cpu_weights(layer, dit_backing) + if dit_cache is not None: + dit_cache.release_unused(self.transformer) + dit_progress("preparing_weight_layout", index, len(layers)) + self.transformer.post_load_weights() + if dit_cache is not None: + dit_cache.publish(self.transformer) + else: + self.transformer.validate_restored_host_weights() self.turbo_spec = None if fusion is not None: self.turbo_spec = fusion.spec @@ -599,7 +665,7 @@ def loading(done, component): self.transformer, self.device, pin_memory=config.host_weight_pin_memory, - host_backing=self._host_backing, + host_backing=dit_backing, ) self._weight_cache = FP16WeightCache( self.transformer, @@ -622,16 +688,40 @@ def loading(done, component): load_model=True, encoder_group=self.text_encoder_group, ) - if self._host_backing is not None: - PinnedModuleStager.map_cpu_weights( - self.text_encoder, self._host_backing, preserve_parameters=False + encoder_progress = LoadingProgress( + "text_encoder", 1, group.rank_in_group, group.world_size + ) + encoder_cache = ( + self._prepared_component( + "text_encoder", + self.text_encoder, + [shared / "text_encoder"], + encoder_progress, + options={}, ) - self.text_encoder.load_weights(iter_checkpoint_weights(shared / "text_encoder")) + if config.prepared_weight_cache + else None + ) + encoder_backing = encoder_cache or self._host_backing + if encoder_cache is None or not encoder_cache.restore(self.text_encoder): + encoder_progress("preparing_weight_storage") + if encoder_backing is not None: + PinnedModuleStager.map_cpu_weights( + self.text_encoder, encoder_backing, preserve_parameters=False + ) + self.text_encoder.load_weights( + encoder_progress.weights( + iter_checkpoint_weights(shared / "text_encoder"), + checkpoint_tensor_count(shared / "text_encoder"), + ) + ) + if encoder_cache is not None: + encoder_cache.publish(self.text_encoder) self._encoder_stager = PinnedModuleStager( self.text_encoder, self.device, pin_memory=config.host_weight_pin_memory, - host_backing=self._host_backing, + host_backing=encoder_backing, ) self._dit_layer_stager: LayerwiseModuleStager | None = None self._encoder_layer_stager: LayerwiseModuleStager | None = None @@ -659,6 +749,7 @@ def loading(done, component): load_device=torch.device("cpu"), pin_memory=config.host_weight_pin_memory, shared_weights_dir=shared_weights_dir, + skip_parameter_init=config.prepared_weight_cache, ) self.video_vae.set_parallel_size(config.tensor_parallel_size) loading(3, "audio_vae") @@ -674,6 +765,27 @@ def loading(done, component): self.eval() loading(4, "ready") + def _prepared_component(self, name, module, paths, progress, *, options): + from .prepared_weights import PreparedWeights, preparation_key + + group = get_tp_group() + key = preparation_key( + paths, + component=name, + rank=group.rank_in_group, + world_size=group.world_size, + options=options, + ) + cache = PreparedWeights( + Path(envs.VLLM_CACHE_ROOT) / "h3-prepared", + key, + module, + limit_bytes=int(self.config.prepared_weight_cache_gib * 2**30), + progress=progress, + ) + self._prepared_caches[name] = cache + return cache + def _transformer_for_task(self, task): return self.transformer diff --git a/vllm/model_executor/models/minimax_h3/prepared_weights.py b/vllm/model_executor/models/minimax_h3/prepared_weights.py new file mode 100644 index 0000000000..15ef0f51b1 --- /dev/null +++ b/vllm/model_executor/models/minimax_h3/prepared_weights.py @@ -0,0 +1,419 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Leased, exact CPU weight snapshots built directly in their final storage.""" + +from __future__ import annotations + +import fcntl +import hashlib +import json +import math +import os +import shutil +import tempfile +from contextlib import contextmanager +from pathlib import Path + +import regex as re +import torch + +from .residency import MMapHostWeights, PinnedModuleStager, set_tensor_storage + +FORMAT = 1 + + +def _json(value): + return json.dumps(value, sort_keys=True, separators=(",", ":")).encode() + + +def _targets(module): + return dict(list(module.named_parameters()) + list(module.named_buffers())) + + +def _schema(module): + return { + name: {"shape": list(t.shape), "dtype": str(t.dtype)} + for name, t in _targets(module).items() + } + + +def checkpoint_identity(paths): + """Track installed immutable artifacts; replacing/editing a file invalidates.""" + result = [] + for value in paths: + root = Path(value).resolve() + files = ( + [root] + if root.is_file() + else sorted( + p for p in root.iterdir() if p.suffix in {".json", ".safetensors"} + ) + ) + if not files: + raise ValueError(f"No checkpoint files at {root}") + for path in files: + stat = path.stat() + result.append( + ( + str(path), + stat.st_dev, + stat.st_ino, + stat.st_size, + stat.st_mtime_ns, + stat.st_ctime_ns, + ) + ) + return result + + +def preparation_key(paths, *, component, rank, world_size, options): + source = Path(__file__).parent + files = sorted(source.glob("*.py")) + [ + source.parents[1] / "parameter.py", + source.parents[1] / "layers/linear.py", + source.parents[1] / "layers/sm70_diffusion.py", + ] + source_hash = hashlib.sha256() + for path in files: + source_hash.update(path.name.encode()) + source_hash.update(path.read_bytes()) + value = { + "format": FORMAT, + "component": component, + "rank": rank, + "world_size": world_size, + "files": checkpoint_identity(paths), + "source": source_hash.hexdigest(), + "torch": torch.__version__, + "cuda": torch.version.cuda, + "dtype": str(torch.get_default_dtype()), + "options": options, + } + return hashlib.sha256(_json(value)).hexdigest() + + +class PreparedWeights: + """Only a completed, checksummed entry can be attached by another worker. + + Each worker holds its entry lease for the lifetime of its host masters. + Eviction never touches an active entry or original model artifacts. Builds + use shared writable mappings; publication reattaches private mappings so + adapters or accidental CPU writes cannot change the reusable snapshot. + """ + + def __init__( + self, + root, + key, + module, + *, + limit_bytes=128 * 2**30, + reserve_bytes=2**30, + progress=None, + ): + if not re.fullmatch(r"[0-9a-f]{64}", key): + raise ValueError("Invalid prepared weight cache identity") + self.root = Path(root).resolve() + self.root.mkdir(parents=True, exist_ok=True, mode=0o700) + self.entry = self.root / ("entry-" + key) + self.key = key + self.schema = _schema(module) + self.limit = limit_bytes + self.reserve = reserve_bytes + self.progress = progress or (lambda *args: None) + self.lease = (self.root / (key + ".lock")).open("a+b") + self.files = set() + self.maps = [] + self.owned = set() + self.readonly = False + self.bytes_reserved = 0 + self.fallback = MMapHostWeights(self.root.parent / "h3-host") + + @contextmanager + def capacity(self): + with (self.root / "capacity.lock").open("a+b") as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + yield + + def _entries(self): + return [ + p + for p in self.root.iterdir() + if re.fullmatch(r"entry-[0-9a-f]{64}", p.name) + and p.is_dir() + and not p.is_symlink() + ] + + def _reserved(self, path): + try: + value = json.loads((path / "reservation.json").read_bytes())["bytes"] + if type(value) is not int or value < 0: + raise ValueError("Invalid cache reservation") + return value + except (OSError, ValueError, KeyError, TypeError): + return sum(p.stat().st_size for p in path.iterdir() if p.is_file()) + + def _reservation(self): + (self.entry / "reservation.json").write_bytes( + _json({"bytes": self.bytes_reserved}) + ) + + def _make_room(self, extra): + def enough(): + return ( + sum(self._reserved(p) for p in self._entries()) + extra <= self.limit + and shutil.disk_usage(self.root).free >= extra + self.reserve + ) + + for path in sorted(self._entries(), key=lambda p: p.stat().st_mtime_ns): + if enough(): + return + if path == self.entry: + continue + with (self.root / (path.name[6:] + ".lock")).open("a+b") as lease: + try: + fcntl.flock(lease, fcntl.LOCK_EX | fcntl.LOCK_NB) + except BlockingIOError: + continue + shutil.rmtree(path) + if not enough(): + raise RuntimeError( + "Not enough space for prepared H3 weights; " + "active caches and model files were retained" + ) + + def _digest(self, path, done, total): + digest = hashlib.sha256() + with path.open("rb") as stream: + while chunk := stream.read(8 * 2**20): + digest.update(chunk) + done += len(chunk) + self.progress("checking_prepared_weights", done, total, "bytes") + return digest.hexdigest(), done + + def restore(self, module): + fcntl.flock(self.lease, fcntl.LOCK_SH) + if self._restore(module): + return True + fcntl.flock(self.lease, fcntl.LOCK_UN) + fcntl.flock(self.lease, fcntl.LOCK_EX) + # A previous builder may have published while we waited for the lease. + if self._restore(module): + fcntl.flock(self.lease, fcntl.LOCK_SH) + return True + with self.capacity(): + if self.entry.is_symlink(): + raise ValueError("Prepared weight entries must not be symlinks") + if self.entry.exists(): + shutil.rmtree(self.entry) + self.entry.mkdir(mode=0o700) + self._reservation() + return False + + def _restore(self, module, *, verify=True): + try: + if self.entry.is_symlink() or any( + (self.entry / p).is_symlink() for p in ("ready.json", "manifest.json") + ): + return False + ready = json.loads((self.entry / "ready.json").read_bytes()) + payload = (self.entry / "manifest.json").read_bytes() + if hashlib.sha256(payload).hexdigest() != ready["sha256"]: + return False + manifest = json.loads(payload) + if ( + manifest["format"] != FORMAT + or manifest["key"] != self.key + or manifest["schema"] != self.schema + ): + return False + targets = _targets(module) + names = set() + maps, bindings = [], [] + done = 0 + total = sum(g["bytes"] for g in manifest["groups"]) + self.progress("checking_prepared_weights", 0, total, "bytes") + for group in manifest["groups"]: + size = group["bytes"] + if not isinstance(size, int) or size < 0: + return False + if size: + filename = group["file"] + if not re.fullmatch(r"weights-[a-z0-9_]+\.bin", filename): + return False + path = self.entry / filename + if path.is_symlink() or path.stat().st_size != size: + return False + if verify: + digest, done = self._digest(path, done, total) + if digest != group["sha256"]: + return False + raw = torch.from_file( + str(path), shared=False, size=size, dtype=torch.uint8 + ) + else: + raw = torch.empty(0, dtype=torch.uint8) + maps.append(raw) + for binding in group["bindings"]: + name = binding["name"] + if name not in targets or name in names: + return False + target = targets[name] + dtype = target.dtype + if binding["dtype"] != str(dtype): + return False + shape, stride, offset = ( + binding["shape"], + binding["stride"], + binding["offset"], + ) + if ( + len(shape) != len(stride) + or len(shape) > 16 + or any( + not isinstance(x, int) or x < 0 + for x in [*shape, *stride, offset] + ) + or math.prod(shape) != target.numel() + ): + return False + extent = ( + 0 + if not math.prod(shape) + else offset + + sum((d - 1) * s for d, s in zip(shape, stride)) + + 1 + ) + if extent * dtype.itemsize > size: + return False + value = torch.empty(0, dtype=dtype).set_( + raw.untyped_storage(), offset, shape, stride + ) + bindings.append((target, value)) + names.add(name) + if names != set(targets): + return False + # Do not partially bind a failed/corrupt cache entry. + for target, value in bindings: + set_tensor_storage(target, value) + self.maps = maps + self.owned = { + (m.untyped_storage().data_ptr(), m.untyped_storage().nbytes()) + for m in maps + } + self.readonly = True + os.utime(self.entry, None) + self.progress("reusing_prepared_weights", total, total, "bytes") + return True + except (OSError, ValueError, KeyError, TypeError, RuntimeError, OverflowError): + return False + + def snapshot(self, source, *, preserve=True): + storage = source.untyped_storage() + if (storage.data_ptr(), storage.nbytes()) in self.owned: + return source.detach() + if self.readonly: + # Ordinary LoRA sidecars are installed after the reusable base. + return self.fallback.snapshot(source, preserve=preserve) + if source.device.type != "cpu" or source.dtype != torch.uint8: + raise ValueError("Prepared backing requires CPU storage bytes") + filename = storage.filename + if filename and Path(filename).parent == self.entry: + return source.detach() + if not source.numel(): + return source.detach() + with self.capacity(): + self._make_room(source.numel()) + fd, filename = tempfile.mkstemp( + prefix="weights-", suffix=".bin", dir=self.entry + ) + try: + os.posix_fallocate(fd, 0, source.numel()) + raw = torch.from_file( + filename, shared=True, size=source.numel(), dtype=torch.uint8 + ) + self.bytes_reserved += source.numel() + self.files.add(Path(filename)) + self._reservation() + except BaseException: + Path(filename).unlink(missing_ok=True) + raise + finally: + os.close(fd) + if preserve: + raw.copy_(source) + return raw + + def release_unused(self, module): + live = { + Path(t.untyped_storage().filename) + for t in _targets(module).values() + if t.untyped_storage().filename + } + with self.capacity(): + for path in self.files - live: + self.bytes_reserved -= path.stat().st_size + path.unlink() + self.files.intersection_update(live) + self._reservation() + + def publish(self, module): + if self.readonly: + return + PinnedModuleStager.map_cpu_weights(module, self) + self.release_unused(module) + groups = {} + for name, tensor in _targets(module).items(): + storage = tensor.untyped_storage() + key = (storage.data_ptr(), storage.nbytes()) + if key not in groups: + groups[key] = { + "file": Path(storage.filename).name if storage.nbytes() else None, + "bytes": storage.nbytes(), + "bindings": [], + } + groups[key]["bindings"].append( + { + "name": name, + "shape": list(tensor.shape), + "stride": list(tensor.stride()), + "offset": tensor.storage_offset(), + "dtype": str(tensor.dtype), + } + ) + done, total = 0, sum(g["bytes"] for g in groups.values()) + for group in groups.values(): + if group["bytes"]: + path = self.entry / group["file"] + with path.open("rb") as stream: + os.fsync(stream.fileno()) + group["sha256"], done = self._digest(path, done, total) + payload = _json( + { + "format": FORMAT, + "key": self.key, + "schema": self.schema, + "groups": list(groups.values()), + } + ) + for name, content in [ + ("manifest.json", payload), + ("ready.json", _json({"sha256": hashlib.sha256(payload).hexdigest()})), + ]: + temporary = self.entry / (name + ".tmp") + with temporary.open("wb") as stream: + stream.write(content) + stream.flush() + os.fsync(stream.fileno()) + temporary.replace(self.entry / name) + fd = os.open(self.entry, os.O_RDONLY | os.O_DIRECTORY) + try: + os.fsync(fd) + finally: + os.close(fd) + if not self._restore(module, verify=False): + raise RuntimeError("Published prepared weights could not be restored") + fcntl.flock(self.lease, fcntl.LOCK_SH) + + def close(self): + self.lease.close() diff --git a/vllm/model_executor/models/minimax_h3/quantization.py b/vllm/model_executor/models/minimax_h3/quantization.py index a3e09c78ea..b44fec2c53 100644 --- a/vllm/model_executor/models/minimax_h3/quantization.py +++ b/vllm/model_executor/models/minimax_h3/quantization.py @@ -76,6 +76,16 @@ class FP16LinearMethod(UnquantizedLinearMethod): supports_prepared_fp16 = True supports_rotated_input = False + def prepare_weights_before_loading(self, layer): + if getattr(layer, "h3_fp16_weight_layout", "row") == "column": + rows, columns = layer.weight.shape + layer.weight.data = torch.empty_strided( + (rows, columns), + (1, rows), + dtype=layer.weight.dtype, + device=layer.weight.device, + ) + def process_weights_after_loading(self, layer): if getattr(layer, "h3_fp16_weight_layout", "row") == "column": layer.weight.data = layer.weight.data.t().contiguous().t() @@ -379,6 +389,18 @@ def process_weights_after_loading(self, layer: Module) -> None: layer.weight.data = layer.weight.data.contiguous() layer.weight_scale.data = scale.data.reshape(-1).contiguous() + def prepare_weights_before_loading(self, layer): + if self.quant_config.weight_layout == "column" and self.prefix.startswith( + "blocks." + ): + rows, columns = layer.weight.shape + layer.weight.data = torch.empty_strided( + (rows, columns), + (1, rows), + dtype=layer.weight.dtype, + device=layer.weight.device, + ) + def apply( self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None = None ) -> torch.Tensor: diff --git a/vllm/model_executor/models/minimax_h3/vae.py b/vllm/model_executor/models/minimax_h3/vae.py index ed4024e62b..1f80224f86 100644 --- a/vllm/model_executor/models/minimax_h3/vae.py +++ b/vllm/model_executor/models/minimax_h3/vae.py @@ -64,9 +64,39 @@ def _load_component_config(component_path: str) -> dict[str, Any]: return config +def _checkpoint_file_identity(component_path, config): + if not all(key in config for key in ("source_path", "source_safetensors_path")): + return None + path = ( + Path(component_path) / config["source_path"] / config["source_safetensors_path"] + ).resolve() + stat = path.stat() + return ( + str(path), + stat.st_dev, + stat.st_ino, + stat.st_size, + stat.st_mtime_ns, + stat.st_ctime_ns, + ) + + +def _same_checkpoint_replicas(identity): + if not dist.is_initialized(): + return identity is not None + identities = [None] * dist.get_world_size() + dist.all_gather_object(identities, identity) + known = [value for value in identities if value is not None] + if known and any(value != known[0] for value in known): + raise ValueError("Video VAE checkpoint changed between loading workers") + return len(known) == len(identities) + + def _load_remote_component( component_path: str, config: dict[str, Any], + *, + skip_parameter_init: bool = False, ) -> nn.Module: auto_map = config.get("auto_map") or {} class_reference = auto_map.get("AutoModel") @@ -83,6 +113,17 @@ def _load_remote_component( # anti-aliasing filters call torch.kaiser_window). Callers place the module # explicitly right after this returns, so nothing depends on the context. with torch.device("cpu"): + if skip_parameter_init: + from .initialization import load_without_random_parameter_init + + before = _checkpoint_file_identity(component_path, config) + model = load_without_random_parameter_init( + lambda: component_cls.from_pretrained(component_path) + ) + if before != _checkpoint_file_identity(component_path, config): + raise ValueError("Video VAE checkpoint changed while loading") + model._h3_checkpoint_file_identity = before + return model return component_cls.from_pretrained(component_path) @@ -138,6 +179,7 @@ def __init__( load_device: torch.device | None = None, pin_memory: bool = True, shared_weights_dir: str | None = None, + skip_parameter_init: bool = False, ) -> None: super().__init__() self._device_target = device @@ -145,6 +187,7 @@ def __init__( self.remote = _load_remote_component( component_path, self.config_dict, + skip_parameter_init=skip_parameter_init, ) # Match the reference loader contract before installing inference-only # decoder fast paths. Keyframe encoding remains FP32; decoder Linear @@ -160,7 +203,20 @@ def __init__( pin_memory=pin_memory, ) if shared_weights_dir is not None: - self._stager.share_cpu_storage(Path(shared_weights_dir) / "video") + from .initialization import uses_assigned_checkpoint_storage + + identity = ( + getattr(self.remote, "_h3_checkpoint_file_identity", None) + if uses_assigned_checkpoint_storage(self.remote) + else None + ) + if _same_checkpoint_replicas(identity): + logger.info( + "H3 video VAE reuses private checkpoint mappings; " + "no duplicate host replica snapshot" + ) + else: + self._stager.share_cpu_storage(Path(shared_weights_dir) / "video") self.model = self.remote.model self.use_tiling = True self.use_slicing = False diff --git a/vllm/model_executor/models/minimax_h3/weights.py b/vllm/model_executor/models/minimax_h3/weights.py index 03cb185d66..40750a12f9 100644 --- a/vllm/model_executor/models/minimax_h3/weights.py +++ b/vllm/model_executor/models/minimax_h3/weights.py @@ -56,6 +56,16 @@ def iter_checkpoint_weights(path: str | Path, *, include: set[str] | None = None yield name, checkpoint.get_tensor(name) +def checkpoint_tensor_count(path: str | Path): + path = Path(path) + files = [path] if path.is_file() else sorted(path.glob("*.safetensors")) + total = 0 + for file in files: + with safe_open(file, framework="pt", device="cpu") as checkpoint: + total += len(checkpoint.keys()) + return total + + def write_checkpoint_manifest(root: str | Path, output: str | Path, *, revision: str): root = Path(root) records = [] diff --git a/vllm/video/fastpath.py b/vllm/video/fastpath.py index 211509e299..022b3113a7 100644 --- a/vllm/video/fastpath.py +++ b/vllm/video/fastpath.py @@ -45,4 +45,5 @@ def studio_capabilities(): "quality_status": "not_accepted", "tasks": ["t2va"], } + result["prepared_weight_cache"] = {"format": 1, "available": True} return result