264 lines
11 KiB
Python
264 lines
11 KiB
Python
"""
|
|
STARFlow-V checkpoint loader.
|
|
|
|
Supports exactly three checkpoint layouts:
|
|
* a single ``.pth``/``.pt`` file (loaded with ``torch.load(..., weights_only=True)``,
|
|
which refuses to unpickle anything but tensors/primitives)
|
|
* a single ``.safetensors`` file
|
|
* a sharded safetensors checkpoint (``model.safetensors.index.json`` + shards),
|
|
such as produced by scripts/convert_checkpoint.py for checkpoints >5GB
|
|
|
|
Memory safety
|
|
-------------
|
|
Loading the real ~7B checkpoint is done without ever materializing a
|
|
float32 skeleton or a duplicated whole-model state dict:
|
|
|
|
* ``load_starflow_v_checkpoint`` builds the skeleton *directly* in the
|
|
target dtype (e.g. bf16), avoiding a transient fp32 copy of the weights.
|
|
Full ``meta`` init is
|
|
deliberately not used because the RoPE module computes ``self.freqs`` as
|
|
a plain fp32 attribute at construction time, which a meta context would
|
|
silently turn into a meta tensor.
|
|
* ``load_state_dict_into_model`` streams the checkpoint one shard at a
|
|
time via ``safetensors.safe_open`` (lazy per-tensor reads), casts each
|
|
tensor to the target dtype, and installs it with
|
|
``load_state_dict(..., assign=True)`` so the loaded tensor *replaces*
|
|
the corresponding skeleton parameter instead of being copied into it.
|
|
Peak extra RAM is therefore ~one shard, not the whole checkpoint.
|
|
|
|
Key/shape checking stays strict (validated up-front from the safetensors
|
|
header without reading tensor data) except for a minimal, explicit
|
|
allowlist of buffers upstream itself marks as deprecated/unused
|
|
(``Model.var``, ``MetaBlock.attn_mask``, and the RoPE modules' deprecated
|
|
``freqs_cos``/``freqs_sin`` buffers), which may be present or absent.
|
|
|
|
Prefix normalization is limited to a small explicit allowlist of known
|
|
wrapper prefixes (``module.``, ``_orig_mod.``, ``model.``) and is only
|
|
applied when *every* key in the checkpoint shares the same prefix -- this
|
|
avoids silently mangling a checkpoint that happens to have a real top-level
|
|
key coincidentally starting with one of these strings.
|
|
|
|
``load_state_dict_into_model`` operates on an already-constructed model, so
|
|
it can be exercised in tests against small synthetic models without ever
|
|
instantiating the full ~7B-parameter STARFlow-V skeleton.
|
|
``load_starflow_v_checkpoint`` is the convenience entry point used by the
|
|
ComfyUI loader node, which builds the real model from STARFlowVConfig.
|
|
"""
|
|
import json
|
|
import os
|
|
from typing import Optional
|
|
|
|
import torch
|
|
|
|
from ..config import STARFlowVConfig
|
|
from ..model import build_starflow_v_model
|
|
|
|
_DEPRECATED_KEY_SUFFIXES = (
|
|
"var",
|
|
"attn_mask",
|
|
"feat_rope.freqs_cos", "feat_rope.freqs_sin",
|
|
"feat_rope_gen.freqs_cos", "feat_rope_gen.freqs_sin",
|
|
)
|
|
|
|
_KNOWN_WRAPPER_PREFIXES = ("module.", "_orig_mod.", "model.")
|
|
|
|
|
|
class CheckpointFormatError(RuntimeError):
|
|
pass
|
|
|
|
|
|
def _is_deprecated_key(key: str) -> bool:
|
|
return any(key == suf or key.endswith('.' + suf) for suf in _DEPRECATED_KEY_SUFFIXES)
|
|
|
|
|
|
def _prefix_map(raw_keys) -> dict:
|
|
"""Return a mapping raw_key -> normalized_key. If *every* raw key shares
|
|
one known wrapper prefix, that prefix is stripped; otherwise identity."""
|
|
if not raw_keys:
|
|
return {}
|
|
for prefix in _KNOWN_WRAPPER_PREFIXES:
|
|
if all(k.startswith(prefix) for k in raw_keys):
|
|
return {k: k[len(prefix):] for k in raw_keys}
|
|
return {k: k for k in raw_keys}
|
|
|
|
|
|
# Backwards-compatible helper retained for callers/tests that import it.
|
|
def _normalize_prefix(state_dict: dict) -> dict:
|
|
norm = _prefix_map(list(state_dict.keys()))
|
|
return {norm[k]: v for k, v in state_dict.items()}
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Streaming checkpoint sources.
|
|
#
|
|
# Each opener returns (shapes, groups, load_group, close):
|
|
# shapes : dict raw_key -> tuple(shape) (cheap; no tensor data read)
|
|
# groups : list of (group_id, [raw_keys]) (one group == one shard)
|
|
# load_group : (group_id, [raw_keys]) -> dict raw_key -> tensor
|
|
# close : () -> None (release file handles)
|
|
# --------------------------------------------------------------------------- #
|
|
def _open_pth(path: str):
|
|
state_dict = torch.load(path, map_location="cpu", weights_only=True)
|
|
if not isinstance(state_dict, dict):
|
|
raise CheckpointFormatError(f"{path}: expected a state_dict (dict), got {type(state_dict)}")
|
|
for k, v in state_dict.items():
|
|
if not isinstance(v, torch.Tensor):
|
|
raise CheckpointFormatError(
|
|
f"{path}: non-tensor value for key {k!r} ({type(v)}); refusing to load")
|
|
shapes = {k: tuple(v.shape) for k, v in state_dict.items()}
|
|
groups = [("__all__", list(state_dict.keys()))]
|
|
|
|
def load_group(_gid, keys):
|
|
return {k: state_dict[k] for k in keys}
|
|
|
|
def close():
|
|
state_dict.clear()
|
|
|
|
return shapes, groups, load_group, close
|
|
|
|
|
|
def _open_single_safetensors(path: str):
|
|
from safetensors import safe_open
|
|
handle = safe_open(path, framework="pt", device="cpu")
|
|
keys = list(handle.keys())
|
|
shapes = {k: tuple(handle.get_slice(k).get_shape()) for k in keys}
|
|
groups = [("__all__", keys)]
|
|
|
|
def load_group(_gid, gkeys):
|
|
return {k: handle.get_tensor(k) for k in gkeys}
|
|
|
|
def close():
|
|
pass # safe_open handle is released when it goes out of scope
|
|
|
|
return shapes, groups, load_group, close
|
|
|
|
|
|
def _open_sharded_safetensors(index_path: str):
|
|
from safetensors import safe_open
|
|
with open(index_path, "r") as f:
|
|
index = json.load(f)
|
|
weight_map = index["weight_map"]
|
|
shard_dir = os.path.dirname(index_path)
|
|
|
|
shard_keys = {}
|
|
for key, shard_file in weight_map.items():
|
|
shard_keys.setdefault(shard_file, []).append(key)
|
|
|
|
handles = {}
|
|
shapes = {}
|
|
for shard_file, keys in shard_keys.items():
|
|
h = safe_open(os.path.join(shard_dir, shard_file), framework="pt", device="cpu")
|
|
handles[shard_file] = h
|
|
present = set(h.keys())
|
|
for key in keys:
|
|
if key not in present:
|
|
raise CheckpointFormatError(f"key {key!r} missing from shard {shard_file}")
|
|
shapes[key] = tuple(h.get_slice(key).get_shape())
|
|
|
|
groups = [(shard_file, keys) for shard_file, keys in shard_keys.items()]
|
|
|
|
def load_group(shard_file, gkeys):
|
|
h = handles.get(shard_file)
|
|
out = {k: h.get_tensor(k) for k in gkeys}
|
|
# get_tensor returns owning copies, so release this shard's file
|
|
# handle (and its mmap) immediately to keep peak RSS ~one shard.
|
|
handles.pop(shard_file, None)
|
|
return out
|
|
|
|
def close():
|
|
handles.clear()
|
|
|
|
return shapes, groups, load_group, close
|
|
|
|
|
|
def _open_checkpoint(checkpoint_path: str):
|
|
if os.path.isdir(checkpoint_path):
|
|
index_path = os.path.join(checkpoint_path, "model.safetensors.index.json")
|
|
if os.path.exists(index_path):
|
|
return _open_sharded_safetensors(index_path)
|
|
single = os.path.join(checkpoint_path, "model.safetensors")
|
|
if os.path.exists(single):
|
|
return _open_single_safetensors(single)
|
|
raise CheckpointFormatError(f"no model.safetensors[.index.json] found in {checkpoint_path}")
|
|
|
|
if checkpoint_path.endswith(".index.json"):
|
|
return _open_sharded_safetensors(checkpoint_path)
|
|
if checkpoint_path.endswith(".safetensors"):
|
|
return _open_single_safetensors(checkpoint_path)
|
|
if checkpoint_path.endswith(".pth") or checkpoint_path.endswith(".pt"):
|
|
return _open_pth(checkpoint_path)
|
|
raise CheckpointFormatError(f"unrecognized checkpoint format for {checkpoint_path}")
|
|
|
|
|
|
def load_state_dict_into_model(model: torch.nn.Module, checkpoint_path: str,
|
|
dtype: torch.dtype = torch.bfloat16) -> torch.nn.Module:
|
|
shapes_raw, groups, load_group, close = _open_checkpoint(checkpoint_path)
|
|
try:
|
|
raw_keys = list(shapes_raw.keys())
|
|
norm = _prefix_map(raw_keys) # raw_key -> normalized_key
|
|
shapes = {norm[k]: shapes_raw[k] for k in raw_keys}
|
|
|
|
model_state = model.state_dict()
|
|
expected_keys = set(model_state.keys())
|
|
ckpt_keys = set(shapes.keys())
|
|
|
|
missing = {k for k in (expected_keys - ckpt_keys) if not _is_deprecated_key(k)}
|
|
unexpected = {k for k in (ckpt_keys - expected_keys) if not _is_deprecated_key(k)}
|
|
if missing:
|
|
raise CheckpointFormatError(
|
|
f"checkpoint is missing {len(missing)} expected key(s), e.g. {sorted(missing)[:5]}")
|
|
if unexpected:
|
|
raise CheckpointFormatError(
|
|
f"checkpoint has {len(unexpected)} unexpected key(s), e.g. {sorted(unexpected)[:5]}")
|
|
|
|
for key in ckpt_keys & expected_keys:
|
|
expected_shape = tuple(model_state[key].shape)
|
|
got_shape = shapes[key]
|
|
if expected_shape != got_shape:
|
|
raise CheckpointFormatError(
|
|
f"shape mismatch for key {key!r}: expected {expected_shape}, got {got_shape}")
|
|
|
|
# Stream shard-by-shard: cast to target dtype and *assign* (replace)
|
|
# each tensor so the skeleton parameter is dropped rather than copied.
|
|
loaded = set()
|
|
for gid, gkeys in groups:
|
|
tensors = load_group(gid, gkeys)
|
|
batch = {}
|
|
for rk in gkeys:
|
|
nk = norm[rk]
|
|
if nk not in expected_keys: # deprecated key not present in model -> skip
|
|
continue
|
|
t = tensors[rk]
|
|
if t.is_floating_point():
|
|
t = t.to(dtype)
|
|
batch[nk] = t
|
|
loaded.add(nk)
|
|
if batch:
|
|
model.load_state_dict(batch, strict=False, assign=True)
|
|
del tensors, batch
|
|
|
|
still_missing = {k for k in (expected_keys - loaded) if not _is_deprecated_key(k)}
|
|
if still_missing:
|
|
raise CheckpointFormatError(
|
|
f"after streaming load, {len(still_missing)} expected key(s) never loaded, "
|
|
f"e.g. {sorted(still_missing)[:5]}")
|
|
finally:
|
|
close()
|
|
|
|
model = model.to(dtype)
|
|
model.eval()
|
|
for p in model.parameters():
|
|
p.requires_grad = False
|
|
return model
|
|
|
|
|
|
def load_starflow_v_checkpoint(checkpoint_path: str, cfg: Optional[STARFlowVConfig] = None,
|
|
dtype: torch.dtype = torch.bfloat16,
|
|
local_attn_window: Optional[int] = None) -> torch.nn.Module:
|
|
cfg = cfg or STARFlowVConfig()
|
|
# Build the skeleton directly in the target floating dtype to avoid a
|
|
# transient fp32 skeleton for the ~7B model.
|
|
build_dtype = dtype if (isinstance(dtype, torch.dtype) and dtype.is_floating_point) else None
|
|
model = build_starflow_v_model(cfg, local_attn_window=local_attn_window, dtype=build_dtype)
|
|
return load_state_dict_into_model(model, checkpoint_path, dtype=dtype)
|