Add Natural Memory architecture and tooling
This commit is contained in:
WpyQwq
@@ -0,0 +1,423 @@
|
||||
"""Compare the unmodified local Qwen checkpoint with a memory-surgery adapter."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import gc
|
||||
import json
|
||||
import math
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable
|
||||
|
||||
import torch
|
||||
|
||||
from .qwen_integration import QwenMemoryConfig, load_memory_config, load_qwen_base, load_qwen_dynamic, load_tokenizer
|
||||
from .train_qwen_memory import encode_messages, load_records, pad_batch
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--model-path", default=".")
|
||||
parser.add_argument("--data", default="dynamic_memory_lab/data/demo_stream.jsonl")
|
||||
parser.add_argument(
|
||||
"--adapter",
|
||||
default="dynamic_memory_lab/qwen_memory_adapter_surgery_smoke",
|
||||
help="dynamic-memory adapter directory; its memory_config.json selects the surgery mode",
|
||||
)
|
||||
parser.add_argument("--output", default="dynamic_memory_lab/benchmark_qwen.json")
|
||||
parser.add_argument("--max-length", type=int, default=512)
|
||||
parser.add_argument("--repeats", type=int, default=3)
|
||||
parser.add_argument("--warmup", type=int, default=1)
|
||||
parser.add_argument("--max-new-tokens", type=int, default=8)
|
||||
parser.add_argument("--seed", type=int, default=42)
|
||||
parser.add_argument("--no-4bit", action="store_true")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def _adapter_config(adapter_dir: str | Path) -> QwenMemoryConfig:
|
||||
return load_memory_config(adapter_dir)
|
||||
|
||||
|
||||
def _score_output(output: Any, labels: torch.Tensor) -> tuple[float, int, int, int, bool]:
|
||||
shifted_labels = labels[..., 1:]
|
||||
predictions = output.logits[..., :-1, :].argmax(dim=-1)
|
||||
target_positions = shifted_labels != -100
|
||||
token_count = int(target_positions.sum().item())
|
||||
if output.loss is None or token_count == 0:
|
||||
raise RuntimeError("benchmark example has no supervised target tokens")
|
||||
correct_tokens = int((predictions[target_positions] == shifted_labels[target_positions]).sum().item())
|
||||
first_target = target_positions.nonzero(as_tuple=False)[0]
|
||||
first_token_correct = int(
|
||||
predictions[first_target[0], first_target[1]] == shifted_labels[first_target[0], first_target[1]]
|
||||
)
|
||||
sequence_ok = correct_tokens == token_count
|
||||
return float(output.loss.detach().item()), token_count, correct_tokens, first_token_correct, sequence_ok
|
||||
|
||||
|
||||
def _evaluate_base(model: Any, tokenizer: Any, records: list[dict[str, Any]], max_length: int) -> dict[str, float]:
|
||||
device = model.get_input_embeddings().weight.device
|
||||
total_nll = 0.0
|
||||
total_tokens = 0
|
||||
correct_tokens = 0
|
||||
correct_first_tokens = 0
|
||||
correct_sequences = 0
|
||||
for record in records:
|
||||
query = encode_messages(tokenizer, record["query"], max_length)
|
||||
input_ids, attention_mask, labels = pad_batch([query], int(tokenizer.pad_token_id))
|
||||
with torch.inference_mode():
|
||||
output = model(
|
||||
input_ids=input_ids.to(device),
|
||||
attention_mask=attention_mask.to(device),
|
||||
labels=labels.to(device),
|
||||
use_cache=False,
|
||||
)
|
||||
loss, tokens, tokens_correct, first_token_correct, sequence_ok = _score_output(output, labels.to(device))
|
||||
total_nll += loss * tokens
|
||||
total_tokens += tokens
|
||||
correct_tokens += tokens_correct
|
||||
correct_first_tokens += first_token_correct
|
||||
correct_sequences += int(sequence_ok)
|
||||
mean_loss = total_nll / total_tokens
|
||||
return {
|
||||
"loss": mean_loss,
|
||||
"perplexity": math.exp(mean_loss),
|
||||
"token_accuracy": correct_tokens / total_tokens,
|
||||
"first_target_token_accuracy": correct_first_tokens / len(records),
|
||||
"exact_sequence_accuracy": correct_sequences / len(records),
|
||||
"supervised_tokens": total_tokens,
|
||||
}
|
||||
|
||||
|
||||
def _evaluate_dynamic(
|
||||
model: Any,
|
||||
tokenizer: Any,
|
||||
records: list[dict[str, Any]],
|
||||
max_length: int,
|
||||
) -> dict[str, float]:
|
||||
device = model._find_layer_device()
|
||||
pad_id = int(tokenizer.pad_token_id)
|
||||
total_nll = 0.0
|
||||
total_tokens = 0
|
||||
correct_tokens = 0
|
||||
correct_first_tokens = 0
|
||||
correct_sequences = 0
|
||||
for record in records:
|
||||
model.reset_memory()
|
||||
memory = encode_messages(tokenizer, record["memory"], max_length)
|
||||
query = encode_messages(tokenizer, record["query"], max_length)
|
||||
memory_input, memory_mask, _ = pad_batch([memory], pad_id)
|
||||
query_input, query_mask, query_labels = pad_batch([query], pad_id)
|
||||
with torch.inference_mode():
|
||||
memory_output = model(
|
||||
input_ids=memory_input.to(device),
|
||||
attention_mask=memory_mask.to(device),
|
||||
read_memory=False,
|
||||
update_memory=True,
|
||||
return_memory=True,
|
||||
use_cache=False,
|
||||
)
|
||||
output = model(
|
||||
input_ids=query_input.to(device),
|
||||
attention_mask=query_mask.to(device),
|
||||
labels=query_labels.to(device),
|
||||
memory_state=memory_output.memory,
|
||||
read_memory=True,
|
||||
update_memory=False,
|
||||
return_memory=True,
|
||||
use_cache=False,
|
||||
)
|
||||
loss, tokens, tokens_correct, first_token_correct, sequence_ok = _score_output(
|
||||
output, query_labels.to(device)
|
||||
)
|
||||
total_nll += loss * tokens
|
||||
total_tokens += tokens
|
||||
correct_tokens += tokens_correct
|
||||
correct_first_tokens += first_token_correct
|
||||
correct_sequences += int(sequence_ok)
|
||||
model.reset_memory()
|
||||
mean_loss = total_nll / total_tokens
|
||||
return {
|
||||
"loss": mean_loss,
|
||||
"perplexity": math.exp(mean_loss),
|
||||
"token_accuracy": correct_tokens / total_tokens,
|
||||
"first_target_token_accuracy": correct_first_tokens / len(records),
|
||||
"exact_sequence_accuracy": correct_sequences / len(records),
|
||||
"supervised_tokens": total_tokens,
|
||||
}
|
||||
|
||||
|
||||
def _sync(device: torch.device) -> None:
|
||||
if device.type == "cuda":
|
||||
torch.cuda.synchronize(device)
|
||||
|
||||
|
||||
def _measure(name: str, fn: Callable[[], int], repeats: int, warmup: int, device: torch.device) -> dict[str, float]:
|
||||
for _ in range(warmup):
|
||||
fn()
|
||||
_sync(device)
|
||||
started = time.perf_counter()
|
||||
processed = 0
|
||||
for _ in range(repeats):
|
||||
processed += fn()
|
||||
_sync(device)
|
||||
elapsed = time.perf_counter() - started
|
||||
return {
|
||||
"seconds": elapsed / repeats,
|
||||
"tokens_per_second": processed / elapsed,
|
||||
}
|
||||
|
||||
|
||||
def _generation_prompt(tokenizer: Any, messages: list[dict[str, Any]], device: torch.device) -> dict[str, torch.Tensor]:
|
||||
encoded = tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tokenize=True,
|
||||
add_generation_prompt=True,
|
||||
enable_thinking=False,
|
||||
)
|
||||
input_ids = encoded["input_ids"] if hasattr(encoded, "__getitem__") and "input_ids" in encoded else encoded
|
||||
if isinstance(input_ids, torch.Tensor):
|
||||
input_ids = input_ids.tolist()
|
||||
if input_ids and isinstance(input_ids[0], list):
|
||||
input_ids = input_ids[0]
|
||||
ids = torch.tensor([input_ids], dtype=torch.long, device=device)
|
||||
return {"input_ids": ids, "attention_mask": torch.ones_like(ids)}
|
||||
|
||||
|
||||
def _measure_base_generation(
|
||||
model: Any,
|
||||
tokenizer: Any,
|
||||
record: dict[str, Any],
|
||||
repeats: int,
|
||||
warmup: int,
|
||||
max_new_tokens: int,
|
||||
) -> dict[str, float]:
|
||||
device = model.get_input_embeddings().weight.device
|
||||
prompt = _generation_prompt(tokenizer, record["query"][:-1], device)
|
||||
|
||||
def run() -> int:
|
||||
output = model.generate(
|
||||
**prompt,
|
||||
max_new_tokens=max_new_tokens,
|
||||
do_sample=False,
|
||||
use_cache=True,
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
)
|
||||
return int(output.shape[1] - prompt["input_ids"].shape[1])
|
||||
|
||||
return _measure("baseline_generation", run, repeats, warmup, device)
|
||||
|
||||
|
||||
def _measure_dynamic_generation(
|
||||
model: Any,
|
||||
tokenizer: Any,
|
||||
record: dict[str, Any],
|
||||
max_length: int,
|
||||
repeats: int,
|
||||
warmup: int,
|
||||
max_new_tokens: int,
|
||||
) -> dict[str, float]:
|
||||
device = model._find_layer_device()
|
||||
pad_id = int(tokenizer.pad_token_id)
|
||||
memory = encode_messages(tokenizer, record["memory"], max_length)
|
||||
memory_input, memory_mask, _ = pad_batch([memory], pad_id)
|
||||
prompt = _generation_prompt(tokenizer, record["query"][:-1], device)
|
||||
|
||||
def run() -> int:
|
||||
model.reset_memory()
|
||||
with torch.inference_mode():
|
||||
model(
|
||||
input_ids=memory_input.to(device),
|
||||
attention_mask=memory_mask.to(device),
|
||||
read_memory=False,
|
||||
update_memory=True,
|
||||
return_memory=True,
|
||||
use_cache=False,
|
||||
)
|
||||
output = model.generate(
|
||||
**prompt,
|
||||
max_new_tokens=max_new_tokens,
|
||||
do_sample=False,
|
||||
update_memory=False,
|
||||
use_cache=True,
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
)
|
||||
return int(output.shape[1] - prompt["input_ids"].shape[1])
|
||||
|
||||
result = _measure("dynamic_generation", run, repeats, warmup, device)
|
||||
model.reset_memory()
|
||||
return result
|
||||
|
||||
|
||||
def _clean_generated(text: str) -> str:
|
||||
return text.replace(" ", "").replace("\r", "").replace("\n", "").strip()
|
||||
|
||||
|
||||
def _generation_quality_base(model: Any, tokenizer: Any, records: list[dict[str, Any]], max_new_tokens: int) -> dict[str, Any]:
|
||||
device = model.get_input_embeddings().weight.device
|
||||
contains = 0
|
||||
prefixes = 0
|
||||
examples = []
|
||||
for record in records:
|
||||
prompt = _generation_prompt(tokenizer, record["query"][:-1], device)
|
||||
with torch.inference_mode():
|
||||
output = model.generate(
|
||||
**prompt,
|
||||
max_new_tokens=max_new_tokens,
|
||||
do_sample=False,
|
||||
use_cache=True,
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
)
|
||||
generated = _clean_generated(tokenizer.decode(output[0, prompt["input_ids"].shape[1] :], skip_special_tokens=True))
|
||||
expected = str(record["answer"])
|
||||
contains += int(expected in generated)
|
||||
prefixes += int(generated.startswith(expected))
|
||||
if len(examples) < 3:
|
||||
examples.append({"expected": expected, "generated": generated})
|
||||
return {
|
||||
"answer_contains_accuracy": contains / len(records),
|
||||
"answer_prefix_accuracy": prefixes / len(records),
|
||||
"examples": examples,
|
||||
}
|
||||
|
||||
|
||||
def _generation_quality_dynamic(
|
||||
model: Any,
|
||||
tokenizer: Any,
|
||||
records: list[dict[str, Any]],
|
||||
max_length: int,
|
||||
max_new_tokens: int,
|
||||
) -> dict[str, Any]:
|
||||
device = model._find_layer_device()
|
||||
pad_id = int(tokenizer.pad_token_id)
|
||||
contains = 0
|
||||
prefixes = 0
|
||||
examples = []
|
||||
for record in records:
|
||||
model.reset_memory()
|
||||
memory = encode_messages(tokenizer, record["memory"], max_length)
|
||||
memory_input, memory_mask, _ = pad_batch([memory], pad_id)
|
||||
prompt = _generation_prompt(tokenizer, record["query"][:-1], device)
|
||||
with torch.inference_mode():
|
||||
model(
|
||||
input_ids=memory_input.to(device),
|
||||
attention_mask=memory_mask.to(device),
|
||||
read_memory=False,
|
||||
update_memory=True,
|
||||
return_memory=True,
|
||||
use_cache=False,
|
||||
)
|
||||
output = model.generate(
|
||||
**prompt,
|
||||
max_new_tokens=max_new_tokens,
|
||||
do_sample=False,
|
||||
update_memory=False,
|
||||
use_cache=True,
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
)
|
||||
generated = _clean_generated(tokenizer.decode(output[0, prompt["input_ids"].shape[1] :], skip_special_tokens=True))
|
||||
expected = str(record["answer"])
|
||||
contains += int(expected in generated)
|
||||
prefixes += int(generated.startswith(expected))
|
||||
if len(examples) < 3:
|
||||
examples.append({"expected": expected, "generated": generated})
|
||||
model.reset_memory()
|
||||
return {
|
||||
"answer_contains_accuracy": contains / len(records),
|
||||
"answer_prefix_accuracy": prefixes / len(records),
|
||||
"examples": examples,
|
||||
}
|
||||
|
||||
|
||||
def _release(model: Any) -> None:
|
||||
del model
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
if args.repeats < 1 or args.warmup < 0:
|
||||
raise ValueError("repeats must be >= 1 and warmup must be >= 0")
|
||||
torch.manual_seed(args.seed)
|
||||
tokenizer = load_tokenizer(args.model_path)
|
||||
records = load_records(args.data)
|
||||
use_4bit = not args.no_4bit
|
||||
results: dict[str, Any] = {
|
||||
"model_path": str(Path(args.model_path).resolve()),
|
||||
"data": str(Path(args.data).resolve()),
|
||||
"adapter": str(Path(args.adapter).resolve()),
|
||||
"records": len(records),
|
||||
"max_length": args.max_length,
|
||||
"quantization": "4bit_nf4" if use_4bit else "none",
|
||||
}
|
||||
|
||||
print("loading baseline")
|
||||
baseline = load_qwen_base(args.model_path, load_in_4bit=use_4bit)
|
||||
baseline.eval()
|
||||
baseline_device = baseline.get_input_embeddings().weight.device
|
||||
if baseline_device.type == "cuda":
|
||||
torch.cuda.reset_peak_memory_stats(baseline_device)
|
||||
results["baseline"] = {
|
||||
"device": str(baseline_device),
|
||||
"scores": _evaluate_base(baseline, tokenizer, records, args.max_length),
|
||||
"generation_quality": _generation_quality_base(baseline, tokenizer, records, args.max_new_tokens),
|
||||
"generation": _measure_base_generation(
|
||||
baseline, tokenizer, records[0], args.repeats, args.warmup, args.max_new_tokens
|
||||
),
|
||||
}
|
||||
if baseline_device.type == "cuda":
|
||||
results["baseline"]["peak_memory_gb"] = torch.cuda.max_memory_allocated(baseline_device) / 1024**3
|
||||
_release(baseline)
|
||||
baseline = None
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
memory_config = _adapter_config(args.adapter)
|
||||
print(f"loading dynamic mode={memory_config.mode} layers={memory_config.layer_indices}")
|
||||
dynamic = load_qwen_dynamic(
|
||||
args.model_path,
|
||||
memory_config=memory_config,
|
||||
load_in_4bit=use_4bit,
|
||||
)
|
||||
dynamic.load_memory_adapter(args.adapter)
|
||||
dynamic.eval()
|
||||
dynamic_device = dynamic._find_layer_device()
|
||||
if dynamic_device.type == "cuda":
|
||||
torch.cuda.reset_peak_memory_stats(dynamic_device)
|
||||
results["dynamic_memory"] = {
|
||||
"device": str(dynamic_device),
|
||||
"mode": memory_config.mode,
|
||||
"layers": list(dynamic.layer_indices),
|
||||
"scores": _evaluate_dynamic(dynamic, tokenizer, records, args.max_length),
|
||||
"generation_quality": _generation_quality_dynamic(
|
||||
dynamic, tokenizer, records, args.max_length, args.max_new_tokens
|
||||
),
|
||||
"generation": _measure_dynamic_generation(
|
||||
dynamic,
|
||||
tokenizer,
|
||||
records[0],
|
||||
args.max_length,
|
||||
args.repeats,
|
||||
args.warmup,
|
||||
args.max_new_tokens,
|
||||
),
|
||||
"trainable_parameters": sum(parameter.numel() for parameter in dynamic.trainable_parameters),
|
||||
}
|
||||
if dynamic_device.type == "cuda":
|
||||
results["dynamic_memory"]["peak_memory_gb"] = torch.cuda.max_memory_allocated(dynamic_device) / 1024**3
|
||||
_release(dynamic)
|
||||
dynamic = None
|
||||
|
||||
output_path = Path(args.output)
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
output_path.write_text(json.dumps(results, indent=2), encoding="utf-8")
|
||||
|
||||
print(json.dumps(results, indent=2, ensure_ascii=False))
|
||||
print(f"saved={output_path.resolve()}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user