141 lines
4.6 KiB
Python
141 lines
4.6 KiB
Python
"""Verify merged LoRA model matches finetuned model numerically.
|
|
|
|
Usage:
|
|
python verify_lora.py \
|
|
--merged merged_model \
|
|
--ft FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import logging
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
from safetensors.torch import load_file
|
|
|
|
LOG = logging.getLogger(__name__)
|
|
|
|
|
|
def configure_logging(level: str = "INFO"):
|
|
handler = logging.StreamHandler()
|
|
handler.setFormatter(logging.Formatter("%(asctime)s %(levelname)s %(message)s", datefmt="%Y-%m-%d %H:%M:%S"))
|
|
LOG.addHandler(handler)
|
|
LOG.setLevel(level)
|
|
|
|
|
|
def load_transformer_weights(model_path: str | Path) -> dict:
|
|
"""Load transformer weights from model directory."""
|
|
model_path = Path(model_path)
|
|
|
|
if not model_path.exists() or not model_path.is_dir():
|
|
from huggingface_hub import snapshot_download
|
|
LOG.info(f"Downloading {model_path} from HuggingFace Hub...")
|
|
model_path = Path(snapshot_download(repo_id=str(model_path), ignore_patterns=["*.onnx", "*.msgpack"]))
|
|
|
|
transformer_dir = model_path / "transformer"
|
|
|
|
if not transformer_dir.exists():
|
|
raise FileNotFoundError(f"Transformer directory not found: {transformer_dir}")
|
|
|
|
weight_files = sorted(transformer_dir.glob("*.safetensors"))
|
|
if not weight_files:
|
|
raise FileNotFoundError(f"No safetensors files in {transformer_dir}")
|
|
|
|
LOG.info(f"Loading {len(weight_files)} file(s) from {transformer_dir}")
|
|
|
|
state_dict = {}
|
|
for f in weight_files:
|
|
if "custom" in f.name:
|
|
continue
|
|
state_dict.update(load_file(str(f)))
|
|
|
|
return state_dict
|
|
|
|
|
|
def compare_models(merged_sd: dict, ft_sd: dict) -> dict:
|
|
"""Compare merged and finetuned model weights."""
|
|
|
|
common_keys = set(merged_sd.keys()) & set(ft_sd.keys())
|
|
merged_only = set(merged_sd.keys()) - set(ft_sd.keys())
|
|
ft_only = set(ft_sd.keys()) - set(merged_sd.keys())
|
|
|
|
LOG.info(f"Common keys: {len(common_keys)}")
|
|
if merged_only:
|
|
LOG.warning(f"Keys only in merged: {len(merged_only)}")
|
|
if ft_only:
|
|
LOG.warning(f"Keys only in finetuned: {len(ft_only)}")
|
|
|
|
results = []
|
|
for key in sorted(common_keys):
|
|
merged_param = merged_sd[key]
|
|
ft_param = ft_sd[key]
|
|
|
|
if merged_param.shape != ft_param.shape:
|
|
LOG.error(f"{key}: shape mismatch {merged_param.shape} vs {ft_param.shape}")
|
|
continue
|
|
|
|
diff = (merged_param.float() - ft_param.float()).abs()
|
|
max_abs = diff.max().item()
|
|
mean_abs = diff.mean().item()
|
|
|
|
merged_norm = merged_param.float().norm().item()
|
|
rel_mean = (mean_abs / merged_norm * 100) if merged_norm > 0 else 0
|
|
|
|
results.append({
|
|
"key": key,
|
|
"shape": tuple(merged_param.shape),
|
|
"max_abs": max_abs,
|
|
"mean_abs": mean_abs,
|
|
"rel_mean": rel_mean
|
|
})
|
|
|
|
return results
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Verify merged LoRA model matches finetuned model")
|
|
parser.add_argument("--merged", required=True, help="Merged model directory")
|
|
parser.add_argument("--ft", required=True, help="Finetuned model ID or path")
|
|
parser.add_argument("--log-level", default="INFO", help="Logging level")
|
|
args = parser.parse_args()
|
|
|
|
configure_logging(args.log_level)
|
|
|
|
LOG.info(f"Loading merged model: {args.merged}")
|
|
merged_sd = load_transformer_weights(args.merged)
|
|
LOG.info(f"Loaded {len(merged_sd)} parameters")
|
|
|
|
LOG.info(f"Loading finetuned model: {args.ft}")
|
|
ft_sd = load_transformer_weights(args.ft)
|
|
LOG.info(f"Loaded {len(ft_sd)} parameters")
|
|
|
|
LOG.info("Comparing models...")
|
|
results = compare_models(merged_sd, ft_sd)
|
|
|
|
results.sort(key=lambda x: x["max_abs"], reverse=True)
|
|
|
|
LOG.info(f"\nTop 10 mismatches by max_abs_error:")
|
|
for i, r in enumerate(results[:10], 1):
|
|
LOG.info(f"{i:2d}. {r['key']}")
|
|
LOG.info(
|
|
f" shape={r['shape']}, max_abs={r['max_abs']:.3e}, mean_abs={r['mean_abs']:.3e}, rel_mean={r['rel_mean']:.4f}%"
|
|
)
|
|
|
|
overall_mean = sum(r["mean_abs"] for r in results) / len(results)
|
|
overall_max = max(r["max_abs"] for r in results)
|
|
|
|
LOG.info(f"\nOverall metrics:")
|
|
LOG.info(f" Layers compared: {len(results)}")
|
|
LOG.info(f" Mean(mean_abs): {overall_mean:.3e}")
|
|
LOG.info(f" Max(max_abs): {overall_max:.3e}")
|
|
|
|
if overall_mean < 1e-4:
|
|
LOG.info("\nVerification PASSED: Merge is numerically accurate")
|
|
else:
|
|
LOG.warning(f"\nVerification WARNING: Mean error {overall_mean:.3e} > 1e-4")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|