134 lines
6.1 KiB
Python
134 lines
6.1 KiB
Python
#!/usr/bin/env python3
|
|
"""Capture a real LLM2Vec prompt fixture from the upstream Kimodo encoder."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
from pathlib import Path
|
|
import sys
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
|
|
def sha256(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as stream:
|
|
for chunk in iter(lambda: stream.read(1024 * 1024), b""):
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--upstream", type=Path, required=True)
|
|
parser.add_argument("--base", type=Path, required=True)
|
|
parser.add_argument("--mntp-adapter", type=Path, required=True)
|
|
parser.add_argument("--supervised-adapter", type=Path, required=True)
|
|
parser.add_argument("--prompt", required=True)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
parser.add_argument("--device", default="cuda")
|
|
parser.add_argument("--debug-layer", type=int, default=0, choices=range(32), help="capture module checkpoints for this transformer layer")
|
|
return parser.parse_args()
|
|
|
|
|
|
def as_f32(value: torch.Tensor) -> np.ndarray:
|
|
return value.detach().to(device="cpu", dtype=torch.float32).contiguous().numpy()
|
|
|
|
|
|
def main() -> None:
|
|
opt = args()
|
|
if not (opt.upstream / "kimodo").is_dir():
|
|
raise SystemExit("--upstream is not a Kimodo checkout")
|
|
if any(not path.is_dir() for path in (opt.base, opt.mntp_adapter, opt.supervised_adapter)):
|
|
raise SystemExit("all model paths must be existing directories")
|
|
sys.path.insert(0, str(opt.upstream.resolve()))
|
|
from kimodo.model.llm2vec import LLM2Vec # pylint: disable=import-outside-toplevel
|
|
from kimodo.model.llm2vec.llm2vec import batch_to_device # pylint: disable=import-outside-toplevel
|
|
from peft import PeftModel # pylint: disable=import-outside-toplevel
|
|
|
|
torch.manual_seed(42)
|
|
encoder = LLM2Vec.from_pretrained(
|
|
str(opt.base), peft_model_name_or_path=str(opt.mntp_adapter), merge_peft=True, torch_dtype=torch.bfloat16,
|
|
)
|
|
encoder.model = PeftModel.from_pretrained(encoder.model, str(opt.supervised_adapter))
|
|
encoder.to(opt.device).eval()
|
|
for parameter in encoder.parameters():
|
|
parameter.requires_grad = False
|
|
|
|
prepared = encoder.prepare_for_tokenization(encoder._convert_to_str("", opt.prompt))
|
|
features = encoder.tokenize([prepared])
|
|
recorded_features = {name: value.detach().cpu().numpy() for name, value in features.items()}
|
|
features = batch_to_device(features, opt.device)
|
|
layer_values: list[torch.Tensor | None] = [None] * 32
|
|
embedding_values: list[torch.Tensor] = []
|
|
layer0_values: dict[str, torch.Tensor] = {}
|
|
layer_hooks = []
|
|
debug_modules = {
|
|
"debug_input_norm": f"layers.{opt.debug_layer}.input_layernorm",
|
|
"debug_q": f"layers.{opt.debug_layer}.self_attn.q_proj",
|
|
"debug_k": f"layers.{opt.debug_layer}.self_attn.k_proj",
|
|
"debug_v": f"layers.{opt.debug_layer}.self_attn.v_proj",
|
|
"debug_o": f"layers.{opt.debug_layer}.self_attn.o_proj",
|
|
"debug_post_norm": f"layers.{opt.debug_layer}.post_attention_layernorm",
|
|
"debug_gate": f"layers.{opt.debug_layer}.mlp.gate_proj",
|
|
"debug_up": f"layers.{opt.debug_layer}.mlp.up_proj",
|
|
"debug_down": f"layers.{opt.debug_layer}.mlp.down_proj",
|
|
}
|
|
for name, module in encoder.model.named_modules():
|
|
if name.endswith("embed_tokens"):
|
|
layer_hooks.append(module.register_forward_hook(lambda _m, _i, o: embedding_values.append(o)))
|
|
for output_name, suffix in debug_modules.items():
|
|
if name.endswith(suffix):
|
|
layer_hooks.append(module.register_forward_hook(
|
|
lambda _m, _i, o, output_name=output_name: layer0_values.__setitem__(output_name, o[0] if isinstance(o, tuple) else o)
|
|
))
|
|
for index in range(32):
|
|
if name.endswith(f"layers.{index}"):
|
|
layer_hooks.append(module.register_forward_hook(
|
|
lambda _m, _i, o, index=index: layer_values.__setitem__(index, o[0] if isinstance(o, tuple) else o)
|
|
))
|
|
try:
|
|
with torch.inference_mode():
|
|
reps = encoder.model(**features)
|
|
pooled = encoder.get_pooling(features, reps.last_hidden_state)
|
|
finally:
|
|
for hook in layer_hooks:
|
|
hook.remove()
|
|
if len(embedding_values) != 1 or any(value is None for value in layer_values) or set(layer0_values) != set(debug_modules):
|
|
raise RuntimeError("did not observe exactly one embedding and every Llama layer")
|
|
arrays: dict[str, np.ndarray] = {
|
|
"input_ids": recorded_features["input_ids"].astype(np.int64),
|
|
"attention_mask": recorded_features["attention_mask"].astype(np.int64),
|
|
"embed_mask": recorded_features["embed_mask"].astype(np.int64),
|
|
"token_embeddings": as_f32(embedding_values[0]),
|
|
"final_hidden_state": as_f32(reps.last_hidden_state),
|
|
"pooled_embedding": as_f32(pooled),
|
|
}
|
|
for index, value in enumerate(layer_values):
|
|
arrays[f"layer_{index:02d}_output"] = as_f32(value) # type: ignore[arg-type]
|
|
for name, value in layer0_values.items():
|
|
arrays[name] = as_f32(value)
|
|
opt.output.parent.mkdir(parents=True, exist_ok=True)
|
|
np.savez_compressed(opt.output.with_suffix(".npz"), **arrays)
|
|
metadata = {
|
|
"fixture_format": 1,
|
|
"prompt": opt.prompt,
|
|
"prepared_text": prepared,
|
|
"device": opt.device,
|
|
"torch": torch.__version__,
|
|
"cuda": torch.version.cuda,
|
|
"base_sha256": {path.name: sha256(path) for path in opt.base.glob("*.safetensors")},
|
|
"mntp_adapter_sha256": sha256(opt.mntp_adapter / "adapter_model.safetensors"),
|
|
"supervised_adapter_sha256": sha256(opt.supervised_adapter / "adapter_model.safetensors"),
|
|
"npz_sha256": sha256(opt.output.with_suffix(".npz")),
|
|
"debug_layer": opt.debug_layer,
|
|
}
|
|
opt.output.with_suffix(".json").write_text(json.dumps(metadata, indent=2) + "\n", encoding="utf-8")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|