161 lines
6.1 KiB
Python
161 lines
6.1 KiB
Python
import argparse
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[2]
|
|
sys.path.insert(0, str(ROOT))
|
|
|
|
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM, apply_rotary_pos_emb, precompute_freqs_cis
|
|
|
|
|
|
def shape(x):
|
|
return tuple(x.shape)
|
|
|
|
|
|
def max_abs_diff(a, b):
|
|
return (a - b).abs().max().item()
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Trace MiniMind RoPE position slices and KV cache growth.")
|
|
parser.add_argument("--seq_len", type=int, default=5)
|
|
parser.add_argument("--hidden_size", type=int, default=64)
|
|
parser.add_argument("--num_attention_heads", type=int, default=4)
|
|
parser.add_argument("--num_key_value_heads", type=int, default=2)
|
|
parser.add_argument("--num_hidden_layers", type=int, default=2)
|
|
parser.add_argument("--vocab_size", type=int, default=128)
|
|
parser.add_argument("--max_new_tokens", type=int, default=3)
|
|
args = parser.parse_args()
|
|
if args.seq_len < 2:
|
|
raise ValueError("--seq_len must be at least 2 so the script can compare prefix + next-token cache.")
|
|
|
|
torch.manual_seed(42)
|
|
config = MiniMindConfig(
|
|
hidden_size=args.hidden_size,
|
|
num_hidden_layers=args.num_hidden_layers,
|
|
num_attention_heads=args.num_attention_heads,
|
|
num_key_value_heads=args.num_key_value_heads,
|
|
vocab_size=args.vocab_size,
|
|
max_position_embeddings=128,
|
|
dropout=0.0,
|
|
flash_attn=False,
|
|
)
|
|
model = MiniMindForCausalLM(config).eval()
|
|
attn = model.model.layers[0].self_attn
|
|
|
|
print("[Config]")
|
|
print(f"hidden_size={config.hidden_size}")
|
|
print(f"num_attention_heads={config.num_attention_heads}")
|
|
print(f"num_key_value_heads={config.num_key_value_heads}")
|
|
print(f"head_dim={config.head_dim}")
|
|
print(f"max_position_embeddings={config.max_position_embeddings}")
|
|
print(f"rope_theta={config.rope_theta}")
|
|
print()
|
|
|
|
freqs_cos, freqs_sin = precompute_freqs_cis(
|
|
dim=config.head_dim,
|
|
end=config.max_position_embeddings,
|
|
rope_base=config.rope_theta,
|
|
rope_scaling=config.rope_scaling,
|
|
)
|
|
print("[RoPE Table]")
|
|
print(f"freqs_cos.shape={shape(freqs_cos)}")
|
|
print(f"freqs_sin.shape={shape(freqs_sin)}")
|
|
print(f"model.model.freqs_cos.shape={shape(model.model.freqs_cos)}")
|
|
print(f"model.model.freqs_sin.shape={shape(model.model.freqs_sin)}")
|
|
print()
|
|
|
|
q = torch.randn(1, args.seq_len, config.num_attention_heads, config.head_dim)
|
|
k = torch.randn(1, args.seq_len, config.num_key_value_heads, config.head_dim)
|
|
cos = model.model.freqs_cos[: args.seq_len]
|
|
sin = model.model.freqs_sin[: args.seq_len]
|
|
q_rope, k_rope = apply_rotary_pos_emb(q, k, cos, sin)
|
|
|
|
print("[Apply RoPE]")
|
|
print(f"q_before_rope.shape={shape(q)}")
|
|
print(f"k_before_rope.shape={shape(k)}")
|
|
print(f"cos_slice.shape={shape(cos)}")
|
|
print(f"sin_slice.shape={shape(sin)}")
|
|
print(f"q_after_rope.shape={shape(q_rope)}")
|
|
print(f"k_after_rope.shape={shape(k_rope)}")
|
|
print(f"q_changed_max_abs_diff={max_abs_diff(q, q_rope):.8f}")
|
|
print(f"k_changed_max_abs_diff={max_abs_diff(k, k_rope):.8f}")
|
|
print()
|
|
|
|
input_ids = ((torch.arange(args.seq_len, dtype=torch.long) * 4 + 3) % args.vocab_size).unsqueeze(0)
|
|
prefix_ids = input_ids[:, :-1]
|
|
last_id = input_ids[:, -1:]
|
|
|
|
with torch.no_grad():
|
|
full_outputs = model(input_ids, use_cache=False)
|
|
prefix_outputs = model(prefix_ids, use_cache=True)
|
|
prefix_cache = prefix_outputs.past_key_values
|
|
step_outputs = model(last_id, past_key_values=prefix_cache, use_cache=True)
|
|
step_cache = step_outputs.past_key_values
|
|
|
|
prefix_k, prefix_v = prefix_cache[0]
|
|
step_k, step_v = step_cache[0]
|
|
print("[KV Cache Growth]")
|
|
print(f"input_ids={input_ids.tolist()}")
|
|
print(f"prefix_ids={prefix_ids.tolist()}")
|
|
print(f"last_id={last_id.tolist()}")
|
|
print(f"prefix_layer0_k.shape={shape(prefix_k)}")
|
|
print(f"prefix_layer0_v.shape={shape(prefix_v)}")
|
|
print(f"step_layer0_k.shape={shape(step_k)}")
|
|
print(f"step_layer0_v.shape={shape(step_v)}")
|
|
print(f"expected_prefix_cache_len={prefix_ids.shape[1]}")
|
|
print(f"expected_step_cache_len={input_ids.shape[1]}")
|
|
print()
|
|
|
|
full_last_logits = full_outputs.logits[:, -1, :]
|
|
step_last_logits = step_outputs.logits[:, -1, :]
|
|
print("[Full Forward vs Incremental Cache Forward]")
|
|
print(f"full_logits.shape={shape(full_outputs.logits)}")
|
|
print(f"step_logits.shape={shape(step_outputs.logits)}")
|
|
print(f"full_last_logits.shape={shape(full_last_logits)}")
|
|
print(f"step_last_logits.shape={shape(step_last_logits)}")
|
|
print(f"full_vs_incremental_last_logits_max_abs_diff={max_abs_diff(full_last_logits, step_last_logits):.12f}")
|
|
print()
|
|
|
|
with torch.no_grad():
|
|
generated_with_cache = model.generate(
|
|
inputs=prefix_ids,
|
|
max_new_tokens=args.max_new_tokens,
|
|
do_sample=False,
|
|
top_k=0,
|
|
top_p=1.0,
|
|
temperature=1.0,
|
|
eos_token_id=None,
|
|
use_cache=True,
|
|
)
|
|
generated_without_cache = model.generate(
|
|
inputs=prefix_ids,
|
|
max_new_tokens=args.max_new_tokens,
|
|
do_sample=False,
|
|
top_k=0,
|
|
top_p=1.0,
|
|
temperature=1.0,
|
|
eos_token_id=None,
|
|
use_cache=False,
|
|
)
|
|
|
|
print("[Generate With Cache vs Without Cache]")
|
|
print(f"generated_with_cache={generated_with_cache.tolist()}")
|
|
print(f"generated_without_cache={generated_without_cache.tolist()}")
|
|
print(f"cache_and_no_cache_generate_match={torch.equal(generated_with_cache, generated_without_cache)}")
|
|
print()
|
|
|
|
print("[Reading Rule]")
|
|
print("RoPE keeps Q/K shapes unchanged, but changes their values according to absolute position.")
|
|
print("KV cache stores per-layer K/V before repeat_kv: [batch, cached_seq, num_key_value_heads, head_dim].")
|
|
print("During cached decode, seq_len can be 1 while start_pos equals cached_seq.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
|
|
main()
|