219 lines
8.2 KiB
Python
219 lines
8.2 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import shutil
|
|
import statistics
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parents[1]
|
|
if str(REPO_ROOT) not in sys.path:
|
|
sys.path.insert(0, str(REPO_ROOT))
|
|
|
|
|
|
def _runtime_dependencies():
|
|
try:
|
|
import psutil
|
|
import torch
|
|
from safetensors.torch import load_file, save_file
|
|
except ImportError as exc:
|
|
raise RuntimeError(
|
|
"The real sharded-memory lane requires the accepted repo-local CPU host "
|
|
"runtime with torch, safetensors, and psutil."
|
|
) from exc
|
|
return load_file, psutil, save_file, torch
|
|
|
|
|
|
def _safe_reset(workspace: Path) -> None:
|
|
expected = (REPO_ROOT / ".tmp" / "sharded-loader-memory").resolve()
|
|
resolved = workspace.resolve()
|
|
if resolved != expected:
|
|
raise ValueError("Memory-probe workspace must be the dedicated repository .tmp/sharded-loader-memory path.")
|
|
if resolved.exists():
|
|
shutil.rmtree(resolved)
|
|
resolved.mkdir(parents=True)
|
|
|
|
|
|
def _write_fixture(workspace: Path, shard_mib: int) -> tuple[Path, tuple[Path, ...]]:
|
|
_load_file, _psutil, save_file, torch = _runtime_dependencies()
|
|
elements = shard_mib * 1024 * 1024 // torch.tensor([], dtype=torch.bfloat16).element_size()
|
|
weight_map: dict[str, str] = {}
|
|
shard_paths: list[Path] = []
|
|
total_size = 0
|
|
for index in range(6):
|
|
key = f"weight_{index}"
|
|
shard_name = f"model-{index + 1:05d}-of-00006.safetensors"
|
|
shard_path = workspace / shard_name
|
|
tensor = torch.full((elements,), float(index + 1), dtype=torch.bfloat16)
|
|
save_file({key: tensor}, shard_path)
|
|
total_size += tensor.nbytes
|
|
weight_map[key] = shard_name
|
|
shard_paths.append(shard_path)
|
|
del tensor
|
|
index_path = workspace / "model.safetensors.index.json"
|
|
index_path.write_text(
|
|
json.dumps({"metadata": {"total_size": total_size}, "weight_map": weight_map}),
|
|
encoding="utf-8",
|
|
)
|
|
return index_path, tuple(shard_paths)
|
|
|
|
|
|
def _child(mode: str, index_path: Path, shard_paths: tuple[Path, ...], shard_mib: int) -> None:
|
|
load_file, psutil, _save_file, torch = _runtime_dependencies()
|
|
from LongCat_Video.model_loading_contract import load_indexed_safetensors_state_dict
|
|
|
|
elements = shard_mib * 1024 * 1024 // torch.tensor([], dtype=torch.bfloat16).element_size()
|
|
|
|
class SyntheticModel(torch.nn.Module):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
for index in range(6):
|
|
self.register_parameter(
|
|
f"weight_{index}",
|
|
torch.nn.Parameter(torch.empty(elements, dtype=torch.bfloat16, device="meta")),
|
|
)
|
|
|
|
model = SyntheticModel()
|
|
if mode == "old":
|
|
state_dict = {}
|
|
for shard_path in shard_paths:
|
|
state_dict.update(load_file(shard_path, device="cpu"))
|
|
else:
|
|
state_dict = load_indexed_safetensors_state_dict(
|
|
str(index_path),
|
|
tuple(str(path) for path in shard_paths),
|
|
allowed_dtypes=("torch.bfloat16",),
|
|
expected_state=model.state_dict(),
|
|
)
|
|
result = model.load_state_dict(state_dict, strict=False, assign=True)
|
|
if result.missing_keys or result.unexpected_keys:
|
|
raise RuntimeError("Synthetic state-dict diagnostics changed.")
|
|
del state_dict
|
|
checksum = sum(float(tensor.sum()) for tensor in model.state_dict().values())
|
|
memory = psutil.Process().memory_info()
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"checksum": checksum,
|
|
"final_rss": memory.rss,
|
|
"peak_rss": memory.peak_wset,
|
|
},
|
|
sort_keys=True,
|
|
)
|
|
)
|
|
|
|
|
|
def _verify_real_integrity(
|
|
index_path: Path,
|
|
shard_paths: tuple[Path, ...],
|
|
shard_mib: int,
|
|
) -> None:
|
|
_load_file, _psutil, save_file, torch = _runtime_dependencies()
|
|
from LongCat_Video.model_loading_contract import load_indexed_safetensors_state_dict
|
|
|
|
state_dict = load_indexed_safetensors_state_dict(
|
|
str(index_path),
|
|
tuple(str(path) for path in shard_paths),
|
|
allowed_dtypes=("torch.bfloat16",),
|
|
)
|
|
if tuple(state_dict) != tuple(f"weight_{index}" for index in range(6)):
|
|
raise RuntimeError("Real safetensors keys changed.")
|
|
if any(tensor.dtype != torch.bfloat16 for tensor in state_dict.values()):
|
|
raise RuntimeError("Real safetensors dtype changed.")
|
|
del state_dict
|
|
|
|
elements = shard_mib * 1024 * 1024 // torch.tensor([], dtype=torch.float32).element_size()
|
|
save_file({"weight_0": torch.zeros(elements, dtype=torch.float32)}, shard_paths[0])
|
|
try:
|
|
load_indexed_safetensors_state_dict(
|
|
str(index_path),
|
|
tuple(str(path) for path in shard_paths),
|
|
allowed_dtypes=("torch.bfloat16",),
|
|
)
|
|
except ValueError as exc:
|
|
if "unsupported dtype" not in str(exc):
|
|
raise
|
|
else:
|
|
raise RuntimeError("Real wrong-dtype shard was accepted.")
|
|
|
|
shard_paths[0].write_bytes(b"not-a-safetensors-file")
|
|
try:
|
|
load_indexed_safetensors_state_dict(
|
|
str(index_path),
|
|
tuple(str(path) for path in shard_paths),
|
|
allowed_dtypes=("torch.bfloat16",),
|
|
)
|
|
except ValueError as exc:
|
|
if "Failed to load checkpoint shard" not in str(exc):
|
|
raise
|
|
else:
|
|
raise RuntimeError("Real corrupt shard was accepted.")
|
|
|
|
|
|
def _run_parent(workspace: Path, repeats: int, shard_mib: int) -> int:
|
|
_safe_reset(workspace)
|
|
index_path, shard_paths = _write_fixture(workspace, shard_mib)
|
|
_verify_real_integrity(index_path, shard_paths, shard_mib)
|
|
index_path, shard_paths = _write_fixture(workspace, shard_mib)
|
|
results: dict[str, list[dict[str, float | int]]] = {"old": [], "hardened": []}
|
|
for _repeat in range(repeats):
|
|
for label, child_mode in (("old", "old"), ("hardened", "hardened")):
|
|
command = [
|
|
sys.executable,
|
|
str(Path(__file__).resolve()),
|
|
"--child-mode",
|
|
child_mode,
|
|
"--workspace",
|
|
str(workspace),
|
|
"--shard-mib",
|
|
str(shard_mib),
|
|
]
|
|
completed = subprocess.run(command, check=True, capture_output=True, text=True)
|
|
results[label].append(json.loads(completed.stdout.strip().splitlines()[-1]))
|
|
|
|
old_peak = statistics.median(row["peak_rss"] for row in results["old"])
|
|
hardened_peak = statistics.median(row["peak_rss"] for row in results["hardened"])
|
|
allowed_increase = max(old_peak * 0.05, 32 * 1024 * 1024)
|
|
checksums = {row["checksum"] for rows in results.values() for row in rows}
|
|
report = {
|
|
"allowed_increase_bytes": allowed_increase,
|
|
"hardened_median_peak_rss": hardened_peak,
|
|
"integrity_status": "PASS",
|
|
"old_median_peak_rss": old_peak,
|
|
"repeats": repeats,
|
|
"shard_count": 6,
|
|
"shard_mib": shard_mib,
|
|
"status": "PASS"
|
|
if len(checksums) == 1 and hardened_peak <= old_peak + allowed_increase
|
|
else "FAIL",
|
|
}
|
|
print(json.dumps(report, indent=2, sort_keys=True))
|
|
return 0 if report["status"] == "PASS" else 1
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description="Verify atomic sharded-loader RSS no-regression.")
|
|
parser.add_argument("--workspace", type=Path, default=REPO_ROOT / ".tmp" / "sharded-loader-memory")
|
|
parser.add_argument("--repeats", type=int, default=3)
|
|
parser.add_argument("--shard-mib", type=int, default=64)
|
|
parser.add_argument("--child-mode", choices=("old", "hardened"), help=argparse.SUPPRESS)
|
|
args = parser.parse_args()
|
|
if args.repeats < 3:
|
|
parser.error("--repeats must be at least 3")
|
|
if args.shard_mib < 8:
|
|
parser.error("--shard-mib must be at least 8")
|
|
workspace = args.workspace.resolve()
|
|
if args.child_mode:
|
|
index_path = workspace / "model.safetensors.index.json"
|
|
shard_paths = tuple(workspace / f"model-{index:05d}-of-00006.safetensors" for index in range(1, 7))
|
|
_child(args.child_mode, index_path, shard_paths, args.shard_mib)
|
|
return 0
|
|
return _run_parent(workspace, args.repeats, args.shard_mib)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|